From 6a22b601e2fd46837e975c4c8fe5500b44bf88ce Mon Sep 17 00:00:00 2001 From: fms-zth <2579013254@qq.com> Date: Fri, 11 Sep 2026 18:13:49 +0800 Subject: [PATCH] feat(dcu-kernel-auto-opt): harness planner, gate policy and DKAO tooling Content imported from the DCU dev container tree (/root/zth_agent/MetaInfer), 3-way merged onto the merged PR #26 head so the next-app alignment and CI fixes stay in place. Scope is limited to this task plugin. Added: - orchestrator: planner.py (+ planner_policy/manifest/gates), gate_policy.py, gpu_preflight.py, harness_io.py, predictions.py, validation_budget.py, variant_promote.py - harness_default/: planner_catalog.yaml, planner_policy.yaml, manifest.yaml, gates.yaml - reports/: best-variant tables (csv/json/xlsx) and the summary scripts - skills/: int8-w8a8-gemm foundations/decode/prefill skill packs - tools/baseline/: Triton INT8 benchmark harness, baseline graph JSONs - 11 new test modules (planner, gate policy, gpu preflight, harness io, predictions, validation budget, variant promote, plan render, prompt schema, planner wiring, planner policy) - retuned variant kernels for glm52/TP8 and hy3/TP4, refreshed WORKFLOW.md Changed: - gen_and_opt_pipeline.py, w8a8_pipeline.py, real_pipeline.py, pipeline.py, skill_store.py, variant_store.py, server/routes.py, static (detail view, shape input, css), form.yaml, bridge/dsh/dsh_agent.py, api/assets helpers - PR-side fixes kept: init_or_resume(task_id) alignment (73c213c) and the CPU-only-CI torch guards (3d70f58) / hermetic dsh tests (6a8f235) Excluded: *.bak-* kernel backups. --- .../tasks/dcu_kernel_auto_opt/WORKFLOW.md | 52 +- .../dcu_kernel_auto_opt/assets/w8a8_bench.py | 25 +- .../bridge/dsh/dsh_agent.py | 244 ++ .../bridge/dsh/tests/test_dsh_agent.py | 14 +- metainfer/tasks/dcu_kernel_auto_opt/form.yaml | 24 + .../harness_default/README.md | 38 + .../harness_default/gates.yaml | 38 + .../harness_default/manifest.yaml | 36 + .../harness_default/planner_catalog.yaml | 57 + .../harness_default/planner_policy.yaml | 68 + .../orchestrator/gate_policy.py | 256 ++ .../orchestrator/gen_and_opt_pipeline.py | 95 +- .../orchestrator/gpu_preflight.py | 238 ++ .../orchestrator/harness_io.py | 95 + .../orchestrator/planner.py | 461 ++++ .../orchestrator/predictions.py | 141 + .../orchestrator/skill_store.py | 50 +- .../orchestrator/validation_budget.py | 95 + .../orchestrator/variant_promote.py | 211 ++ .../orchestrator/variant_store.py | 117 + .../orchestrator/w8a8_pipeline.py | 175 +- .../reports/_build_summary.py | 251 ++ .../reports/_diag_o_proj_ab.sh | 24 + .../dcu_kernel_auto_opt/reports/_make_xlsx.py | 257 ++ .../dkao_optimized_operators_bestvariant.csv | 34 + .../dkao_optimized_operators_bestvariant.json | 728 ++++++ .../dkao_optimized_operators_bestvariant.xlsx | Bin 0 -> 16612 bytes .../dcu_kernel_auto_opt/server/routes.py | 22 +- .../skills/int8-w8a8-gemm-decode/SKILL.md | 376 +++ .../references/m2-decode-search.md | 166 ++ .../references/wqkv-a-m16-split10-570gbps.md | 116 + .../int8-w8a8-gemm-foundations/SKILL.md | 339 +++ .../skills/int8-w8a8-gemm-prefill/SKILL.md | 760 ++++++ .../references/wqkv-a-m3072-prefetch-isa.md | 144 ++ .../SKILL.md | 38 + .../dcu_kernel_auto_opt/static/dkao-detail.js | 13 + .../static/dkao-shape-input.js | 53 +- .../tasks/dcu_kernel_auto_opt/static/dkao.css | 35 + .../tests/test_gate_policy.py | 158 ++ .../tests/test_gpu_preflight.py | 121 + .../tests/test_harness_io.py | 105 + .../tests/test_plan_render.py | 102 + .../dcu_kernel_auto_opt/tests/test_planner.py | 204 ++ .../tests/test_planner_policy.py | 62 + .../tests/test_planner_wiring.py | 84 + .../tests/test_predictions.py | 94 + .../tests/test_prompt_schema.py | 47 + .../tests/test_skill_store.py | 34 + .../tests/test_validation_budget.py | 54 + .../tests/test_variant_promote.py | 149 ++ .../tests/test_variant_store.py | 28 + .../tools/baseline/bench_triton_tp4_m4096.py | 277 ++ .../tools/baseline/bench_triton_tp8_m4096.py | 233 ++ .../tools/baseline/int8_utils.py | 602 +++++ .../tools/baseline/tp8_m4096_graph.json | 662 +++++ .../baseline/triton_tp4_m4096_eager.json | 281 ++ .../baseline/triton_tp4_m4096_graph.json | 281 ++ .../tools/planner_parity.py | 257 ++ .../deepseek-v4/TP8/M16/indexer_wq_b.hip | 1 + .../deepseek-v4/TP8/M16/shared_down_proj.hip | 1 + .../TP8/M16/shared_gate_up_proj.hip | 1 + .../deepseek-v4/TP8/M16/wo_b.hip | 1 + .../deepseek-v4/TP8/M16/wq_b.hip | 1 + .../deepseek-v4/TP8/M16/wqkv_a.hip | 1 + .../glm52/TP8/M16/fused_qkv_a_proj.hip | 1 + .../int8w8a8-gemm/glm52/TP8/M16/kv_b_proj.hip | 1 + .../int8w8a8-gemm/glm52/TP8/M16/o_proj.hip | 1 + .../int8w8a8-gemm/glm52/TP8/M16/q_b_proj.hip | 1 + .../glm52/TP8/M16/shared_down_proj.hip | 1 + .../glm52/TP8/M16/shared_gate_up_proj.hip | 1 + .../glm52/TP8/M4096/fused_qkv_a_proj.hip | 1556 ++++++----- .../glm52/TP8/M4096/kv_b_proj.hip | 2298 ++++++----------- .../glm52/TP8/M4096/q_b_proj.hip | 1229 +++++---- .../int8w8a8-gemm/hy3/TP4/M16/o_proj.hip | 545 ++-- .../int8w8a8-gemm/hy3/TP4/M16/qkv_proj.hip | 1 + .../hy3/TP4/M16/shared_down_proj.hip | 494 ++-- .../hy3/TP4/M16/shared_gate_up_proj.hip | 1 + .../int8w8a8-gemm/hy3/TP8/M16/o_proj.hip | 1414 +++++----- .../int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip | 1936 ++++++++------ .../hy3/TP8/M16/shared_down_proj.hip | 1 + .../hy3/TP8/M16/shared_gate_up_proj.hip | 1 + .../int8w8a8-gemm/hy3/TP8/M4096/o_proj.hip | 920 ++++--- .../hy3/TP8/M4096/shared_down_proj.hip | 1469 ++++++++--- .../minimax-m3/TP4/M4096/qkv_proj.hip | 1 + .../TP4/M4096/qkv_proj_and_indexer_qk.hip | 1 + .../TP4/M4096/shared_gate_up_proj.hip | 1 + .../minimax-m3/TP8/M16/o_proj.hip | 1 + .../minimax-m3/TP8/M16/qkv_proj.hip | 1 + .../TP8/M16/qkv_proj_and_indexer_qk.hip | 1 + .../minimax-m3/TP8/M16/shared_down_proj.hip | 1 + .../TP8/M16/shared_gate_up_proj.hip | 1 + 91 files changed, 16201 insertions(+), 5405 deletions(-) create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/harness_default/README.md create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/harness_default/gates.yaml create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/harness_default/manifest.yaml create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_catalog.yaml create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_policy.yaml create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gate_policy.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_preflight.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/orchestrator/harness_io.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/orchestrator/planner.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/orchestrator/predictions.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/orchestrator/validation_budget.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_promote.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/_build_summary.py create mode 100755 metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_o_proj_ab.sh create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/_make_xlsx.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.csv create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.json create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.xlsx create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/SKILL.md create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/references/m2-decode-search.md create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/references/wqkv-a-m16-split10-570gbps.md create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-foundations/SKILL.md create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-prefill/SKILL.md create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-prefill/references/wqkv-a-m3072-prefetch-isa.md create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-quantized-gemm-optimization/SKILL.md create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_gate_policy.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_harness_io.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_plan_render.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_policy.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_wiring.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_predictions.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_prompt_schema.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_validation_budget.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_promote.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp4_m4096.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m4096.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/int8_utils.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m4096_graph.json create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_eager.json create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_graph.json create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/planner_parity.py diff --git a/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md b/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md index f32dc992..53f8912c 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md +++ b/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md @@ -88,13 +88,13 @@ M/N/K 是否合法（M 范围、K%32==0、N%16==0、(K,N) 是否在 TP4/TP8 表 `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、 +`tools/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`。 +`tools/baseline/bench_triton_tp4_m4096.py`（脚本在 task 目录内，`int8_utils.py` 与之同目录）。 ### 2.4 Parallel explore（worker 生命周期，w8a8_pipeline.py） @@ -180,8 +180,8 @@ TP8 另有 wq_b/wo_b (1024,4096)、gate_up (4096,512)、down (256,4096)、indexe `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。 + （`tools/baseline/bench_triton_tp8_m4096.py`，int8_utils.matmul_kernel + CUDA-graph replay， + 结果存 `tools/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。 @@ -254,8 +254,11 @@ TP8 另有 wq_b/wo_b (1024,4096)、gate_up (4096,512)、down (256,4096)、indexe `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 测试里有覆盖）。 + **规范 skill 以 `metainfer/tasks/dcu_kernel_auto_opt/skills/` 内的副本为种子**（int8-w8a8 家族， + 2026-08-27 起随插件一起维护）：新机器上 `sync_skill_libraries()`（或 WebUI 同步按钮）会自动 + 把缺失的 skill 补种进 `~/.dsh/skills/`，再镜像到 `~/.claude/skills/`；已有的库 skill 不被覆盖。 + baseline 测量工具在 `tools/baseline/`（`bench_triton_tp4/tp8_m4096.py` + `int8_utils.py`）。 + kernel-repos 默认在 MetaInfer 同级（`METAINFER_KERNEL_REPOS` 可改），不随插件目录走。 5. 改动任何行为后：更新本文件相关段落 + 跑 tests + 用真实任务验证（优先在 zth_meta 里）。 ## 7. 不确定性标注 @@ -264,3 +267,40 @@ TP8 另有 wq_b/wo_b (1024,4096)、gate_up (4096,512)、down (256,4096)、indexe - “恢复流程”是手工驱动（复用 `_synthesize_final_candidate`），不是 UI 一键重试；UI 是否提供重试以 `server/routes.py` 实际实现为准。 - 本文档不替代 skill 里的性能调优细节（tile 选择、LDS、DUMMA API、hipprof 用法），那些看对应 skill。 + +## 8. AHE 接入准备（M1，2026-09-09，纯新增/默认零行为变化） + +设计文档：`MetaInfer/docs/ahe_dkao_integration_plan.md`、`ahe_dkao_design.md`、 +`dkao_harness_eval_protocol.md`；AHE 官方参考在 `/root/zth_agent/ahe-ref`。 + +新增组件（都在本插件内，运行管线默认不消费、不改变行为）： + +| 文件 | 作用 | wired | +|---|---|---| +| `harness_default/manifest.yaml` | 可演化 harness 组件清单 | false | +| `harness_default/gates.yaml` | gate 规范值（漂移守卫 tests/test_harness_io.py） | false | +| `harness_default/planner_catalog.yaml` | 方案目录（14 个 plan id） | false | +| `orchestrator/harness_io.py` | 定位/读取/播种 harness（`METAINFER_HARNESS_ROOT` 可覆盖） | partial | +| `orchestrator/planner.py` | 状态条件化方案选择器 v0（P0 修复→P1 预算/plateau→P2 瓶颈→P3 覆盖→P4 兜底）+ `render_plan()` 渲染成轮次指令文字 | false | +| `tools/planner_parity.py` | 离线 parity 只读分析（历史轮次状态回放 planner vs 菜单，报告 JSON） | tool | +| `orchestrator/predictions.py` | 内层决策钩子：结构化 prediction 解析/核对（hit/miss/na） | partial | +| w8a8_pipeline.py（改动） | 轮记录加 `prediction_checked` / `plan_id`（proposal 带结构化字段才写） | partial | +| gen_and_opt_pipeline.py（改动） | Generate staging 写 `harness_snapshot/` + scaffold_manifest.harness（revision+digests） | partial | + +- 运行时菜单（prompts.py::w8a8_round_strategy）**默认不变**；设 `METAINFER_PLANNER=1` 时 + 轮次指令改由 planner 渲染（`_round_strategy_text`，w8a8_pipeline.py），用于 A/B 对比。 +- worker prompt 模板现提示可选 `plan_id` / `prediction`（expected_us_range/direction）， + agent 自主决定是否带；带时轮记录会写入 `prediction_checked` 与 `plan_id`（决策钩子活化）。 +- 受控 parity（tests/test_plan_render.py）：fresh lane/plateau/ISA/faster_wrong 状态下 + planner 渲染文字与菜单关键方向词一致。离线 corpus parity（tools/planner_parity.py） + 曾抓出 planner "空 history 误判 fix_build" bug（已修），其精确分歧率受关键字启发式 + 与历史缺 PMC 影响，仅作 sanity，不作准绳。 +- 测试：`tests/test_harness_io.py`、`test_planner.py`、`test_predictions.py`、 + `test_plan_render.py`、`test_prompt_schema.py`、`test_planner_wiring.py`； + 全量 269 passed。 + +## 9. harness_evolve 外循环插件（同仓、平级 task，2026-09-09） + +新增 `metainfer/tasks/harness_evolve/`（自动发现即可见，headless CLI 用法与边界见其 +`README.md`）。它把本插件的 DKAO 任务实例当评测单元跑 AHE 外层闭环（dry-run 已通； +`dkao-cli` evaluator 与真实 Evolve Agent 需在 worker29 后续迭代验证）。 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py index 51a3b879..6e4b54a5 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py @@ -191,9 +191,28 @@ def main() -> int: 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) + # Defaults follow the harness budget env when set (see + # orchestrator/validation_budget.py): METAINFER_BENCH_WARMUPS / _SAMPLES / + # _REPLAYS. Explicit CLI flags still win, so callers keep control. + def _env_default(name: str, fallback: int) -> int: + try: + value = int(os.environ.get(name, "")) + except (TypeError, ValueError): + return fallback + return value if value > 0 else fallback + + parser.add_argument( + "--warmups", type=int, + default=_env_default("METAINFER_BENCH_WARMUPS", 100), + ) + parser.add_argument( + "--samples", type=int, + default=_env_default("METAINFER_BENCH_SAMPLES", 30), + ) + parser.add_argument( + "--replays-per-sample", type=int, + default=_env_default("METAINFER_BENCH_REPLAYS", 100), + ) parser.add_argument("--reference-cache-dir", type=Path) parser.add_argument("--probe", action="store_true") parser.add_argument("--self-test", action="store_true") 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 index ce292b9d..0bfbd0fc 100755 --- a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/dsh_agent.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/dsh_agent.py @@ -189,6 +189,228 @@ def emit_system(session_id: str, model: str) -> None: }) +def _resolve_dsh_bin() -> Optional[str]: + """Locate the dsh CLI (npm global lives in /usr/local/bin).""" + import shutil + + found = shutil.which("dsh") + if found: + return found + for cand in ("/usr/local/bin/dsh", "/usr/bin/dsh"): + if os.path.isfile(cand) and os.access(cand, os.X_OK): + return cand + return None + + +def _with_dsh_bin_on_path() -> Dict[str, str]: + """Env for the headless child with the dsh bin dir prepended to PATH.""" + env = dict(os.environ) + dsh_bin = _resolve_dsh_bin() + if dsh_bin: + bindir = os.path.dirname(dsh_bin) + path = env.get("PATH") or "" + entries = [p for p in path.split(":") if p and p != bindir] + env["PATH"] = ":".join([bindir] + entries + ["/usr/bin", "/bin"]) + return env + + +#: dsh-base tool rows -> the agent-facing tool names each one provides. +#: Used to build a profile patch that disables every row outside a whitelist. +_TOOL_ROWS: Dict[str, set] = { + "tool-fs": {"read", "write", "edit", "multiedit", "notebookedit"}, + "tool-fs-search": {"glob", "grep", "search"}, + "tool-str-replace-editor": {"strreplace", "str_replace", "edit"}, + "tool-bash": {"bash", "bashoutput", "killshell", "run"}, + "tool-pwsh": {"powershell"}, + "tool-skill": {"skill"}, + "tool-web": {"webfetch", "websearch", "web"}, + "tool-todo": {"todowrite", "todoread", "todo"}, + "tool-jobs": {"joboutput", "jobkill", "joblist", "jobs"}, + "tool-workflow": {"workflow"}, + "tool-ralph": {"ralph"}, + "tool-goal": {"goalcreate", "goalupdate", "getgoal", "goal"}, + "tool-subagent": {"task", "subagent"}, + "tool-subagent-control": {"sendmessage", "interruptagent"}, + "tool-subagent-list-agents": {"listagents"}, + "tool-subagent-fork": {"subagentfork"}, + "tool-subagent-report": set(), +} + + +def _parse_tool_list(raw: Optional[str]) -> set: + if not raw: + return set() + return { + part.strip().lower().replace("-", "").replace("_", "") + for part in raw.replace(" ", ",").split(",") + if part.strip() + } + + +def _tool_patch_file(args: Any) -> Optional[str]: + """Write a dsh profile patch disabling tool rows outside the whitelist. + + Returns the patch path (caller deletes it) or None when no gating was + requested. ``--tools`` is a whitelist; ``--disallowedTools`` additionally + removes rows. Rows not listed in the dsh-base bundle are never targeted, + so the patch stays valid across profiles. + """ + allowed = _parse_tool_list(getattr(args, "tools", None)) + denied = _parse_tool_list(getattr(args, "disallowed_tools", None)) + if not allowed and not denied: + return None + lines = [ + "# generated by MetaInfer dsh_agent headless backend " + "(tool gating from --tools/--disallowedTools)", + ] + disabled_any = False + for row, tools in _TOOL_ROWS.items(): + if allowed: + keep = bool(tools & allowed) + else: + keep = True + if tools & denied: + keep = False + if not keep: + lines.append(f"- id: {row}") + lines.append(" disabled: true") + disabled_any = True + if not disabled_any: + return None + import tempfile + fd, path = tempfile.mkstemp(suffix=".patch.yml", prefix="dsh-tools-") + with os.fdopen(fd, "w", encoding="utf-8") as fh: + fh.write("\n".join(lines) + "\n") + return path + + +def _headless_backend_main(args: Any, prompt: str) -> int: + """Run one task through the local `dsh --profile headless` CLI. + + The full prompt is written to a temp brief file because prompts here can + exceed the kernel argv limit (one DKAO brief is up to ~300 KB) and the + headless app only accepts its task on argv. We ask the agent to Read the + brief and follow it. A keepalive line is written to stdout while the + sub-agent is still working so SubAgentManager's no-output watchdog never + mistakes a long task for a stuck process; the final assistant text is + emitted as the last message, matching the ccb stream-json contract. + """ + import subprocess + import tempfile + import threading + import time + + dsh_bin = _resolve_dsh_bin() + if not dsh_bin: + sys.stderr.write( + "dsh_agent: headless backend requires the `dsh` CLI on PATH\n" + ) + return 1 + + model = map_model(args.model) + session_id = args.resume or args.session_id or f"session-{os.urandom(8).hex()}" + + fd, brief = tempfile.mkstemp(suffix=".md", prefix="dsh-task-brief-") + tool_patch: Optional[str] = None + try: + with os.fdopen(fd, "w", encoding="utf-8") as fh: + fh.write(prompt) + task = ( + f"You are a coding sub-agent inside a kernel-optimization pipeline. " + f"Read the complete task brief at {brief} with the Read tool and " + "follow it to completion, autonomously, without asking the user. " + "When the task is finished, output exactly the final result text " + "and stop." + ) + tool_patch = _tool_patch_file(args) + cmd = [dsh_bin, "--profile", "headless"] + if tool_patch: + cmd += ["--patch", tool_patch] + cmd.append(task) + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + env=_with_dsh_bin_on_path(), + ) + except Exception as exc: # noqa: BLE001 + if tool_patch: + try: + os.unlink(tool_patch) + except OSError: + pass + sys.stderr.write(f"dsh_agent: failed to start headless run: {exc!r}\n") + return 1 + + out_lines: List[str] = [] + err_lines: List[str] = [] + + def pump(pipe, sink: List[str]) -> None: + assert pipe is not None + for line in pipe: + sink.append(line) + + t_out = threading.Thread(target=pump, args=(proc.stdout, out_lines), daemon=True) + t_err = threading.Thread(target=pump, args=(proc.stderr, err_lines), daemon=True) + t_out.start() + t_err.start() + + emitted_system = False + last_beat = time.time() + beat_secs = float(os.environ.get("DSH_AGENT_BEAT_SECONDS", "45")) + while proc.poll() is None: + if not emitted_system: + emit_system(session_id, model) + emitted_system = True + if time.time() - last_beat >= beat_secs: + sys.stdout.write("# dsh-agent keepalive\n") + sys.stdout.flush() + last_beat = time.time() + time.sleep(1) + + t_out.join(timeout=10) + t_err.join(timeout=10) + + if not emitted_system: + emit_system(session_id, model) + + def _cleanup_temps() -> None: + for path in (brief, tool_patch): + if path: + try: + os.unlink(path) + except OSError: + pass + + if proc.returncode != 0: + tail = "".join(err_lines)[-2000:] + sys.stderr.write( + f"dsh_agent: headless run failed (rc={proc.returncode}): {tail}\n" + ) + _cleanup_temps() + return 1 + final_text = "".join(out_lines).strip() + if not final_text: + sys.stderr.write("dsh_agent: headless run produced no output\n") + _cleanup_temps() + return 1 + _cleanup_temps() + emit({ + "type": "assistant", + "session_id": session_id, + "message": {"content": [{"type": "text", "text": final_text}]}, + }) + emit({ + "type": "result", + "session_id": session_id, + "result": final_text, + "finish_reason": "completed", + "usage": {}, + }) + return 0 + + def extract_text_blocks(content: Any) -> List[str]: """Pull text blocks from an assistant message content array.""" if not isinstance(content, list): @@ -219,6 +441,11 @@ def build_parser() -> argparse.ArgumentParser: p.add_argument("--resume") p.add_argument("--session-id") p.add_argument("--max-turns") + # Tool gating (used by the HE evolve agent). The headless backend maps + # these ccb-style tool lists onto a dsh profile patch that disables the + # tool rows outside the whitelist. + p.add_argument("--tools") + p.add_argument("--disallowedTools", dest="disallowed_tools") # Anything else (claude-specific) is ignored. return p @@ -231,6 +458,23 @@ def main(argv: Optional[List[str]] = None) -> int: sys.stderr.write("dsh_agent: empty prompt on stdin\n") return 1 + # Backend selection. ``headless`` runs the local `dsh --profile headless` + # CLI (one autonomous agent task, prints the final result); ``sdk`` uses + # the deepseek-harness python SDK + node carrier below. ``auto`` prefers + # headless when the dsh CLI is on PATH (the SDK node carrier needs a + # repo-built runtime closure that may be absent on worker hosts). + backend = os.environ.get("DSH_AGENT_BACKEND", "auto").strip().lower() + if backend == "auto": + backend = "headless" if _resolve_dsh_bin() else "sdk" + if backend == "headless": + return _headless_backend_main(args, prompt) + if backend != "sdk": + sys.stderr.write( + f"dsh_agent: unknown DSH_AGENT_BACKEND {backend!r} " + "(use sdk | headless | auto)\n" + ) + return 1 + try: from deepseek_harness import DeepSeekHarness, DeepSeekHarnessConfig except ImportError as exc: # pragma: no cover - environment check 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 index 0580f2b3..fb0e2302 100644 --- 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 @@ -96,16 +96,24 @@ def __init__(self, text): self.buffer = io.BytesIO(text.encode("utf-8")) -def run_wrapper(module, argv, prompt): +def run_wrapper(module, argv, prompt, backend="sdk"): """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. + + ``backend`` selects the DSH agent backend through ``DSH_AGENT_BACKEND`` and + is restored afterwards. """ out, err = io.StringIO(), io.StringIO() old_stdin = sys.stdin + old_backend = os.environ.get("DSH_AGENT_BACKEND") + if backend: + os.environ["DSH_AGENT_BACKEND"] = backend + elif old_backend is not None: + os.environ.pop("DSH_AGENT_BACKEND", None) saved_key = os.environ.get("TENCENT_API_KEY") os.environ["TENCENT_API_KEY"] = saved_key or "ci-test-key" sys.stdin = FakeStdin(prompt) @@ -118,6 +126,10 @@ def run_wrapper(module, argv, prompt): os.environ.pop("TENCENT_API_KEY", None) else: os.environ["TENCENT_API_KEY"] = saved_key + if old_backend is not None: + os.environ["DSH_AGENT_BACKEND"] = old_backend + else: + os.environ.pop("DSH_AGENT_BACKEND", None) return code, out.getvalue(), err.getvalue() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/form.yaml b/metainfer/tasks/dcu_kernel_auto_opt/form.yaml index 96215dc2..3f1066a4 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/form.yaml +++ b/metainfer/tasks/dcu_kernel_auto_opt/form.yaml @@ -117,6 +117,30 @@ - label: "GLM5.2" description: "GLM-5.2 INT4/INT8 mix via SGLang v0.5.15" +- key: validation_scope + question: "最终串行验证范围：全量 API shapes（默认，含未优化 shape 的回归）或仅本次任务 shape（更快）" + header: "Final validation" + required: false + form: select + default: "All API shapes" + options: + - label: "All API shapes" + description: "验证冻结 API 的每个 shape（含 fallback 回归）——最稳，但每题约 50-90 分钟" + - label: "Task shapes only" + description: "只验证本任务优化的 shape——快得多，适合只关心本次 shape 性能的场景" + +- key: bench_profile + question: "benchmark 采样密度：full 用资产默认(100/30/100)，quick 用较快且稳定的采样(30/30/50)" + header: "Bench profile" + required: false + form: select + default: "full" + options: + - label: "full" + description: "默认高密度采样" + - label: "quick" + description: "轻量采样，串行验证更快（median/p90 略有噪声）" + - key: shape_assignment_mode question: "Choose whether the control plane or you assign shapes to physical GPUs." header: "Shape assignment" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/harness_default/README.md b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/README.md new file mode 100644 index 00000000..a5a0a41b --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/README.md @@ -0,0 +1,38 @@ +# harness_default —— DKAO 可演化 harness 的种子目录 + +这是 AHE「组件可观测性」的第一步落地（M1 切片 1，纯新增、不改运行时行为）。 + +## 现状与用途 + +DKAO 的 harness（gates 阈值、prompts/菜单、planner、skills、工具集、中间件策略、记忆） +目前散在 Python 代码与静态文件里。AHE 只能演化"文件化、可版本化、可校验"的组件， +因此我们先把它们**播种成文件**，放在本目录： + +``` +harness_default/ +├── manifest.yaml # 组件清单（AHE Evolve 唯一可写空间的入口描述） +├── gates.yaml # 验收/plateau/ISA 门的规范值（seed 阶段） +├── planner_catalog.yaml # 优化方案选择机制：方案目录（★ 见 orchestrator/planner.py v0） +├── README.md # 本文件 +└── (后续切片加入) + ├── systemprompt/ # prompts 模板（coordinator/bootstrap/worker/synthesis） + ├── tools/ # cordis 组合与 per-role 工具白名单 + ├── middleware/ # resume/compaction/重试策略参数 + ├── skills/ # SKILL.md 种子 + └── memory/ # LongTermMEMORY（measured 事实） +``` + +## seed 阶段语义（重要） + +- **wired: false**：运行时管线仍读 Python 常量（config.py:105、w8a8_pipeline.py:78-80 等）， + 本目录文件**尚未被加载器接线**——因此默认行为零变化。 +- 一致性由测试守卫：`tests/test_harness_io.py` 断言 gates.yaml 的值与 Python 常量相等， + 防止两份来源漂移；接线（renderer/loader）在后续切片落地，落地时删掉该守卫并让 + 运行时以本目录为唯一来源。 +- AHE 外循环（harness_evolve）将以本目录为种子，拷贝出可写 workspace（git 仓库）， + 每轮评测 pin 一个 workspace 快照；Evolve Agent 只改 workspace。 + +## 读取方式 + +见 `orchestrator/harness_io.py`（harness_root 解析、load_manifest/load_gates、seed_workspace）。 +可用环境变量 `METAINFER_HARNESS_ROOT` 覆盖根目录（默认本目录）。 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/harness_default/gates.yaml b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/gates.yaml new file mode 100644 index 00000000..79bbb718 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/gates.yaml @@ -0,0 +1,38 @@ +schema_version: 1 +component: gates +status: seed +wired: true +source_refs: + config_py: metainfer/tasks/dcu_kernel_auto_opt/orchestrator/config.py + w8a8_pipeline_py: metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py +gates: + round_acceptance_improvement_percent: 1.0 + p90_guard: no_p90_regression + shadow: + enabled: true + min_improvement_percent: 0.3 + max_exclusive_percent: 1.0 + plateau: + recent_valid_rounds: 3 + max_regression_percent: 2.0 + window_upper_exclusive_percent: 2.0 + isa_gate: + required_valid_isa_guided_rounds: 2 + required_hip_rounds_rule: "max(1, max_iterations - 2)" + inline_asm_requires_compiler_limitation: true + task_budget: + default_max_iterations: 10 + minimum_improvement_percent_semantics: final validated result versus fixed baseline +source_annotations: |- + round_acceptance_improvement_percent: config.py:105 ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + p90_guard: w8a8_pipeline.py:422 evaluate_candidate_acceptance (candidate_p90 <= best_p90) + shadow.min_improvement_percent: w8a8_pipeline.py:80 _SHADOW_MIN_IMPROVEMENT_PERCENT + plateau.max_regression_percent: w8a8_pipeline.py:79 _PLATEAU_MAX_REGRESSION_PERCENT + plateau.window_upper_exclusive_percent: w8a8_pipeline.py:154 value < 2.0 + isa_gate.required_valid_isa_guided_rounds: w8a8_pipeline.py:78 _REQUIRED_VALID_ISA_GUIDED_ROUNDS + isa_gate.required_hip_rounds_rule: w8a8_pipeline.py:111 required_hip_rounds = max(1, max_iterations - 2) +notes: >- + These values ARE read at runtime (orchestrator/gate_policy.py) whenever the + manifest marks this component wired; anything missing falls back to the + built-in defaults, which equal the historical Python constants. Editing this + file now changes behaviour, and the Evolve loop can version it. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/harness_default/manifest.yaml b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/manifest.yaml new file mode 100644 index 00000000..3512f59a --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/manifest.yaml @@ -0,0 +1,36 @@ +schema_version: 1 +harness_name: dcu_kernel_auto_opt_default +revision: seed-v0.1 +seed_source: metainfer/tasks/dcu_kernel_auto_opt/harness_default +components: + manifest: + path: manifest.yaml + kind: meta + role: component inventory (Evolve readable system map); wired flags + decide which components are applied at runtime + wired: true + gates: + path: gates.yaml + kind: data + role: acceptance / plateau / ISA gating canonical values (read at runtime + through orchestrator/gate_policy.py) + wired: true + drift_guard: tests/test_harness_io.py::test_gates_match_python_constants + planner_catalog: + path: planner_catalog.yaml + kind: data + role: stable plan ids + worker-facing plan descriptions + wired: true + drift_guard: tests/test_planner.py::test_catalog_loads_from_yaml_seed + planner_policy: + path: planner_policy.yaml + kind: data + role: evolvable state-to-plan policy (repair / phase / bottleneck / coverage / fallback) + wired: true + drift_guard: tests/test_planner_policy.py +layout_note: >- + Future components land here as files: systemprompt/, planner/ (optimization + plan selection mechanism), tools/, middleware/, skills/, memory/. Together + they form the evolvable harness workspace that the harness_evolve outer loop + versions (git) and evolves per the evaluation protocol v2 + (docs/dkao_harness_eval_protocol.md). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_catalog.yaml b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_catalog.yaml new file mode 100644 index 00000000..7eb8302c --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_catalog.yaml @@ -0,0 +1,57 @@ +# planner 方案目录（v0, seed） +# +# DKAO 优化方案选择机制的「方案目录」：把"每一轮可以试什么方案"枚举成稳定 id， +# 供 orchestrator/planner.py（状态条件化选择器 v0）与后续 harness 演化使用。 +# +# wired: false —— 运行时仍由 prompts.py::w8a8_round_strategy 按轮数菜单出文字， +# 本目录与 planner.py 当前是纯新增的"选择机制 v0"，用于沉淀/测试/演进，尚未接线。 +# 规则分层见 orchestrator/planner.py：P0 修复 -> P1 预算/阶段 -> P2 瓶颈签名 +# -> P3 覆盖/防循环 -> P4 兜底（legacy 菜单近似）。 + +catalog: + repair_faster_wrong: + focus: 修复"更快但错"候选（保留高速架构，只修最小正确性缺陷） + family: repair + fix_build: + focus: 修复上一轮编译/API/语法错误（最小改动，不换架构） + family: repair + retry_same: + focus: 基础设施失败后的原样重试（不算一轮） + family: repair + bootstrap_correctness: + focus: 迭代 0：先拿到正确 kernel（标量/DUMMA 最小可用） + family: bootstrap + establish_arch: + focus: 建立首个正确 DUMMA 几何（不追求性能） + family: explore + architecture_explore: + focus: 探索 launch 几何族（tile/waves/blocks，覆盖 CU） + family: explore + grid_splitk: + focus: grid 并行度不足时评估 split-K（含非 2 幂候选与 combine 成本） + family: explore + pipeline_tune: + focus: 流水线：staging 深度/双缓冲/prefetch/发射分组 + family: optimize + memory_layout: + focus: 访存布局：LDS padding/swizzle、packed 布局、bank 冲突 + family: optimize + occupancy_resource: + focus: 占用率/资源：LDS 尺寸、VGPR live range、waves per block + family: optimize + epilogue_fusion: + focus: epilogue：scale/bf16/存储融合、去 workspace 中间 pass + family: optimize + isa_guided_hip: + focus: ISA 引导的 HIP/DUMMA 代码塑形（需控制面 plateau 门） + family: isa + conditional_inline_asm: + focus: 最小 inline asm（需前一轮确认编译器限制 + 目标指令） + family: isa + consolidate: + focus: 收尾：固定胜出架构做最终 HIP-only 打磨 + family: optimize + +guards: >- + 运行时未接线，改动本目录不影响现有行为；tests/test_planner.py 覆盖 + P0-P4 分层判定与 catalog 加载。 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_policy.yaml b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_policy.yaml new file mode 100644 index 00000000..77278127 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_policy.yaml @@ -0,0 +1,68 @@ +schema_version: 1 +component: planner_policy +wired: true + +repair_priority: + faster_wrong: repair_faster_wrong + infrastructure_failure: retry_same + build_failure: fix_build + no_valid_kernel: establish_arch + +phase_gates: + consolidate_when_rounds_left_lte: 1 + plateau_recent_valid_rounds: 3 + plateau_min_improvement_percent: -2.0 + plateau_max_improvement_percent_exclusive: 2.0 + isa_required_valid_rounds_rule: "max(1, max_iterations - 2)" + isa_plan: isa_guided_hip + inline_asm_plan: conditional_inline_asm + consolidate_plan: consolidate + +bottleneck_to_plan: + occupancy_limited: occupancy_resource + bank_conflicts: memory_layout + lds_wait: pipeline_tune + l2_low: memory_layout + grid_limited: grid_splitk + +coverage: + max_consecutive_same_plan: 2 + max_trials_per_plan: 2 + +fallback: + fresh: establish_arch + prefill: + - memory_layout + - pipeline_tune + - architecture_explore + - memory_layout + - epilogue_fusion + - pipeline_tune + - occupancy_resource + - consolidate + m16: + - architecture_explore + - grid_splitk + - pipeline_tune + - memory_layout + - pipeline_tune + - pipeline_tune + - occupancy_resource + - consolidate + small: + - establish_arch + - memory_layout + - occupancy_resource + - architecture_explore + - memory_layout + - pipeline_tune + - pipeline_tune + - consolidate + +# No-evidence fallback: when PMC profiling yields no bottleneck signature the +# selector has nothing to condition on, so it defers to the hand-tuned legacy +# round menu (a known-good baseline). The Evolve agent may evolve this row. +uncertainty: + enabled: true + plan: legacy_menu + min_valid_rounds: 1 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gate_policy.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gate_policy.py new file mode 100644 index 00000000..a7dd251b --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gate_policy.py @@ -0,0 +1,256 @@ +"""Runtime gate values, read from the harness workspace (gates.yaml). + +Until now ``gates.yaml`` was a *mirror* of Python constants: the file existed +and a drift test kept it equal, but nothing read it, so an Evolve change to +gates had no effect at all (``wired: false`` in the manifest). This module +makes the component real: + +* values come from ``/gates.yaml`` (the AHE candidate snapshot, + or the seed for production tasks); +* a component marked ``wired: false`` in ``manifest.yaml`` is *ignored* — the + manifest is therefore meaningful too (it decides which components apply); +* anything missing falls back to the built-in defaults, which are exactly the + values the pipelines hard-coded before, so behaviour is unchanged until a + harness actually edits the file. + +Cache keys include both files' mtimes so a mid-run harness swap is picked up. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Dict, Optional, Tuple + +_BUILTIN_GATES: Dict[str, Any] = { + "round_acceptance_improvement_percent": 1.0, + "p90_guard": "no_p90_regression", + "shadow": { + "enabled": True, + "min_improvement_percent": 0.3, + "max_exclusive_percent": 1.0, + }, + "plateau": { + "recent_valid_rounds": 3, + "max_regression_percent": 2.0, + "window_upper_exclusive_percent": 2.0, + }, + "isa_gate": { + "required_valid_isa_guided_rounds": 2, + "required_hip_rounds_rule": "max(1, max_iterations - 2)", + "inline_asm_requires_compiler_limitation": True, + }, + "task_budget": { + "default_max_iterations": 10, + }, +} + +_CACHE: Dict[Tuple[str, float, float], Dict[str, Any]] = {} +_ENV_ROOT = "METAINFER_HARNESS_ROOT" + + +def builtin_gates() -> Dict[str, Any]: + """The defaults that were hard-coded before this file was wired.""" + import copy + + return copy.deepcopy(_BUILTIN_GATES) + + +def _root(root: Optional[Path] = None) -> Path: + if root is not None: + return Path(root) + env = os.environ.get(_ENV_ROOT) + if env: + return Path(env).expanduser() + from .harness_io import default_harness_dir + + return default_harness_dir() + + +def _yaml(path: Path) -> Dict[str, Any]: + if not path.is_file(): + return {} + try: + import yaml + data = yaml.safe_load(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return data if isinstance(data, dict) else {} + + +def _merge(base: Dict[str, Any], overlay: Dict[str, Any]) -> Dict[str, Any]: + out = dict(base) + for key, value in (overlay or {}).items(): + if isinstance(value, dict) and isinstance(out.get(key), dict): + out[key] = _merge(out[key], value) + else: + out[key] = value + return out + + +def component_wired(name: str, root: Optional[Path] = None) -> bool: + """Whether ``manifest.yaml`` marks this component as actually applied.""" + manifest = _yaml(_root(root) / "manifest.yaml") + components = manifest.get("components") + if not isinstance(components, dict): + return True # no manifest -> assume wired + entry = components.get(name) + if not isinstance(entry, dict): + return True + return bool(entry.get("wired", True)) + + +def gates(root: Optional[Path] = None) -> Dict[str, Any]: + """Effective gate values for this harness workspace.""" + workspace = _root(root) + gates_path = workspace / "gates.yaml" + manifest_path = workspace / "manifest.yaml" + try: + key = (str(workspace), + gates_path.stat().st_mtime if gates_path.is_file() else 0.0, + manifest_path.stat().st_mtime if manifest_path.is_file() else 0.0) + except OSError: + key = (str(workspace), 0.0, 0.0) + cached = _CACHE.get(key) + if cached is not None: + return cached + if not component_wired("gates", workspace): + value = builtin_gates() + else: + data = _yaml(gates_path).get("gates") + value = _merge(builtin_gates(), data if isinstance(data, dict) else {}) + _CACHE[key] = value + return value + + +def reset_cache() -> None: + _CACHE.clear() + + +def _number(path: str, default: float, root: Optional[Path] = None) -> float: + node: Any = gates(root) + for part in path.split("."): + if not isinstance(node, dict): + return default + node = node.get(part) + try: + return float(node) + except (TypeError, ValueError): + return float(default) + + +def _flag(path: str, default: bool, root: Optional[Path] = None) -> bool: + node: Any = gates(root) + for part in path.split("."): + if not isinstance(node, dict): + return default + node = node.get(part) + if node is None: + return default + if isinstance(node, bool): + return node + return str(node).strip().lower() in {"1", "true", "yes", "on"} + + +# ------------------------------------------------------------- accessors ---- + +def round_acceptance_improvement_percent(root: Optional[Path] = None) -> float: + return _number("round_acceptance_improvement_percent", 1.0, root) + + +def p90_guard_mode(root: Optional[Path] = None) -> str: + """``no_p90_regression`` (default) or a numeric tolerance like ``1.02``.""" + node: Any = gates(root).get("p90_guard") + return str(node if node is not None else "no_p90_regression") + + +def p90_tolerance(root: Optional[Path] = None) -> float: + """Multiplier on the best p90 that still counts as a pass (1.0 = strict).""" + mode = p90_guard_mode(root).strip().lower() + if mode in {"no_p90_regression", "strict", ""}: + return 1.0 + try: + value = float(mode) + except ValueError: + return 1.0 + return value if value >= 1.0 else 1.0 + + +def shadow_enabled(root: Optional[Path] = None) -> bool: + return _flag("shadow.enabled", True, root) + + +def shadow_min_improvement_percent(root: Optional[Path] = None) -> float: + return _number("shadow.min_improvement_percent", 0.3, root) + + +def shadow_max_exclusive_percent(root: Optional[Path] = None) -> float: + return _number("shadow.max_exclusive_percent", 1.0, root) + + +def plateau_recent_valid_rounds(root: Optional[Path] = None) -> int: + return int(_number("plateau.recent_valid_rounds", 3, root)) + + +def plateau_max_regression_percent(root: Optional[Path] = None) -> float: + return _number("plateau.max_regression_percent", 2.0, root) + + +def plateau_window_upper_exclusive_percent(root: Optional[Path] = None) -> float: + return _number("plateau.window_upper_exclusive_percent", 2.0, root) + + +def isa_required_valid_rounds(root: Optional[Path] = None) -> int: + return int(_number("isa_gate.required_valid_isa_guided_rounds", 2, root)) + + +def inline_asm_requires_compiler_limitation(root: Optional[Path] = None) -> bool: + return _flag("isa_gate.inline_asm_requires_compiler_limitation", True, root) + + +def required_hip_rounds(max_iterations: int, root: Optional[Path] = None) -> int: + """Honour the documented rule but allow a literal override. + + ``required_hip_rounds_rule`` may stay the documented expression + ``max(1, max_iterations - 2)`` or be replaced by an integer literal. + """ + node: Any = gates(root).get("isa_gate") + rule = (node or {}).get("required_hip_rounds_rule") \ + if isinstance(node, dict) else None + if isinstance(rule, (int, float)): + return max(1, int(rule)) + try: + literal = int(str(rule).strip()) + return max(1, literal) + except (TypeError, ValueError): + return max(1, int(max_iterations) - 2) + + +def default_max_iterations(root: Optional[Path] = None) -> int: + return int(_number("task_budget.default_max_iterations", 10, root)) + + +def snapshot(state_dir: Path, root: Optional[Path] = None) -> Optional[Path]: + """Persist the effective gate values a run used (AHE mechanism evidence). + + The harness_evolve mechanism gate compares this record with the candidate + harness it staged, which is how a gates.yaml edit becomes verifiable + instead of merely "unobserved". + """ + import json + + workspace = _root(root) + payload = { + "harness_root": str(workspace), + "harness_revision": _yaml(workspace / "manifest.yaml").get("revision"), + "gates_wired": component_wired("gates", workspace), + "gates": gates(workspace), + } + try: + target = Path(state_dir) / "gates_effective.json" + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(json.dumps(payload, ensure_ascii=False, indent=2, + sort_keys=True), encoding="utf-8") + except OSError: + return None + return target 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 index daa10c5a..de293088 100644 --- 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 @@ -47,6 +47,7 @@ generate_kernel_prompt, shape_balanced_assignment, ) +from .harness_io import harness_root, load_manifest, seed_workspace 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 @@ -58,6 +59,10 @@ evaluate_final_target, snapshot_accepted_kernel_artifact, ) +from . import gate_policy as _gates +from .validation_budget import ( + resolve_bench_kwargs, resolve_validation_scope, +) from .w8a8_baselines import fixed_triton_graph_baseline @@ -554,6 +559,21 @@ def _render_prebuilt_dispatch( ]) +def _validation_shape_list( + optimized: list, fallback: list, scope: str +) -> list: + """Shapes the final serial validation must cover for this scope. + + ``api`` keeps the full-regression sweep (production default); ``task`` + validates only the optimized shapes, which is what a single-shape AHE + child actually needs. + """ + shapes = list(optimized) + if scope == "api": + shapes += list(fallback) + return shapes + + def _final_performance_gate( *, shape_id: str, @@ -815,8 +835,33 @@ def _prepare_worktrees( "initial_kernel": "pending_parallel_explore_child_generation", "created_at": time.time(), } + # Harness snapshot (M1 slice 5): pin the evolvable-harness workspace + # revision + file digests into the fresh task repo. Default = built-in + # harness_default/ seed; METAINFER_HARNESS_ROOT overrides to an evolved + # workspace. Pure bookkeeping — must never fail staging. + try: + _harness_root = harness_root() + _harness_dst = seed / "harness_snapshot" + seed_workspace(_harness_dst, root=_harness_root) + scaffold_manifest["harness"] = { + "source": str(_harness_root), + "revision": ( + load_manifest(_harness_root).get("revision") or "seed" + ), + "files": { + str(p.relative_to(_harness_dst)): file_digest(p) + for p in _harness_dst.rglob("*") + if p.is_file() + }, + } + except Exception as _snap_exc: # noqa: BLE001 + self.store.append_timeline( + "harness_snapshot_skipped", {"error": str(_snap_exc)} + ) write_json(seed / "scaffold_manifest.json", scaffold_manifest) _run(["git", "add", "scaffold_manifest.json"], cwd=seed) + if (seed / "harness_snapshot").is_dir(): + _run(["git", "add", "harness_snapshot"], cwd=seed) if _run( ["git", "diff", "--cached", "--name-only"], cwd=seed ).stdout.strip(): @@ -2114,11 +2159,27 @@ def _synthesize_final_candidate( ) validation: Dict[str, Dict[str, Any]] = {} - validation_shapes = [ + _answers = self.req.get("answers") + if not isinstance(_answers, dict): + _answers = self.req + validation_scope = resolve_validation_scope(_answers) + bench_kwargs = resolve_bench_kwargs(_answers) + optimized_shapes = [ {"id": shape.id, **shape.params} for shape in config.shapes.values() if shape.id in optimized_ids - ] + fallback_shapes + ] + validation_shapes = _validation_shape_list( + optimized_shapes, fallback_shapes, validation_scope + ) + self.store.append_timeline( + "final_validation_plan", + { + "scope": validation_scope, + "shapes": [str(shape["id"]) for shape in validation_shapes], + "bench": dict(bench_kwargs), + }, + ) try: for shape in validation_shapes: shape_id = str(shape["id"]) @@ -2127,11 +2188,16 @@ def _synthesize_final_candidate( if key != "id" } if shape_id in optimized_ids: - metrics = runner.benchmark(params) + metrics = runner.benchmark(params, **bench_kwargs) else: - metrics = runner.benchmark( - params, warmups=2, samples=3 - ) + # Regression shapes stay on the light sampling; only the + # replay count follows the configured budget. + regression_kwargs = {"warmups": 2, "samples": 3} + if "replays_per_sample" in bench_kwargs: + regression_kwargs["replays_per_sample"] = ( + bench_kwargs["replays_per_sample"] + ) + metrics = runner.benchmark(params, **regression_kwargs) validation[shape_id] = metrics if not metrics.get("passed"): raise RuntimeError( @@ -2148,7 +2214,9 @@ def _synthesize_final_candidate( shape_id=shape_id, best_median=best_median, metrics=metrics, - benchmark=lambda: runner.benchmark(params), + benchmark=lambda: runner.benchmark( + params, **bench_kwargs + ), max_retries=_PERF_GATE_MAX_RETRIES, retry_interval_s=_PERF_GATE_RETRY_INTERVAL_S, store=self.store, @@ -2202,7 +2270,9 @@ def _synthesize_final_candidate( "fallback_regression_shapes": [ str(shape["id"]) for shape in fallback_shapes ], - "all_api_shapes_validated": True, + "all_api_shapes_validated": validation_scope == "api", + "validation_scope": validation_scope, + "bench_budget": dict(bench_kwargs), "hip_recompiled_by_main": False, "serial_validation_gpu": serial_gpu, } @@ -2218,6 +2288,13 @@ def _synthesize_final_candidate( # ------------------------------------------------------------------ # def run(self, *, dry_run: bool = False) -> Dict[str, Any]: + # Record the gate values this run actually used, so the harness_evolve + # mechanism gate can verify a gates.yaml change (not just observe it). + try: + _gates.snapshot(getattr(self.store, "task_dir", None) + or self.workspace_dir.parent) + except Exception: # noqa: BLE001 - evidence is best effort + pass """Full generate-then-optimize pipeline. 1. PREPARE: parse config (may lack GPU assignments), create seed repo. @@ -2386,7 +2463,7 @@ def _plan(config: OptimizerConfig) -> Dict[str, Any]: "final validated result versus fixed baseline" ), "round_acceptance_improvement_percent": ( - ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + _gates.round_acceptance_improvement_percent() ), "shape_scope": config.shape_scope, "assignment_mode": config.assignment_mode, diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_preflight.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_preflight.py new file mode 100644 index 00000000..723a8a5c --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_preflight.py @@ -0,0 +1,238 @@ +"""GPU admission pre-check shared by production DKAO tasks and AHE children. + +Two different questions are answered separately, because they need different +rules: + +* **Can we run here at all?** Only VRAM pressure decides that. Our W8A8 + kernels need a few GB at most, so the operating rule is simple: a device is + usable while its used-VRAM percentage is at or below + ``vram_limit_percent`` (default 85). Above that the caller prefers another + device, and only then waits. +* **Will the measurement be trustworthy?** Busy compute (high HCU%, high + board power) or foreign processes hold the device (KFD lists pids invisible + inside this container) do NOT block the run — they only mark the question as + ``measurement_suspect`` so a polluted reading cannot silently decide an + A/B comparison. + +The module is deliberately dependency-free: it parses ``hy-smi`` text and the +KFD sysfs listing, so it works inside the task containers without extra tools. +""" + +from __future__ import annotations + +import os +import re +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +#: A device is usable while at least this much VRAM stays free. Absolute free +#: memory is the right unit for us (one W8A8 worker needs ~1-3GB), whereas a +#: percentage punishes big cards that hold idle reservations. +DEFAULT_MIN_FREE_GB = 4.0 +#: Assumed total VRAM when the driver sysfs does not expose it (K500SM_AI). +DEFAULT_TOTAL_GB = 65.5 +#: Above this average HCU% the device is considered shared with other work. +DEFAULT_UTIL_SUSPECT_PERCENT = 20.0 +#: Above this board power (W) another workload is almost certainly running. +DEFAULT_POWER_SUSPECT_W = 230.0 + +KFD_PROC_ROOT = Path("/sys/class/kfd/kfd/proc") +_SMI_ROW = re.compile( + r"^(?P\d+)\s+" # device index + r"(?P[\d.]+)C\s+" + r"(?P[\d.]+)W\s+" + r"\S+\s+" # perf mode + r"[\d.]+W\s+" # power cap + r"(?P[\d.]+)%\s+" + r"(?P[\d.]+)%" +) + + +def parse_hy_smi(text: str) -> Dict[int, Dict[str, float]]: + """Parse ``hy-smi`` output into ``{index: {vram_percent, util_percent, power_w, temp_c}}``.""" + out: Dict[int, Dict[str, float]] = {} + for line in (text or "").splitlines(): + match = _SMI_ROW.match(line.strip()) + if not match: + continue + index = int(match.group("index")) + out[index] = { + "vram_percent": float(match.group("vram")), + "util_percent": float(match.group("util")), + "power_w": float(match.group("power")), + "temp_c": float(match.group("temp")), + } + return out + + +def read_hy_smi(timeout_s: float = 10.0) -> str: + import subprocess + + for binary in ("hy-smi", "rocm-smi"): + try: + proc = subprocess.run( + [binary], capture_output=True, text=True, timeout=timeout_s, + ) + except (OSError, subprocess.SubprocessError): + continue + if proc.returncode == 0 and proc.stdout.strip(): + return proc.stdout + return "" + + +def foreign_kfd_pids(kfd_root: Path = KFD_PROC_ROOT) -> List[int]: + """GPU-holding pids that are NOT visible inside this container. + + ``/sys/class/kfd`` is host-wide, so an entry whose ``/proc/`` does not + exist here belongs to another container (or the host) — i.e. foreign load. + """ + out: List[int] = [] + try: + entries = list(Path(kfd_root).iterdir()) + except OSError: + return out + for entry in entries: + try: + pid = int(entry.name) + except ValueError: + continue + if not Path(f"/proc/{pid}").exists(): + out.append(pid) + return sorted(out) + + +def sample_gpu_state(gpu_ids: Optional[List[int]] = None, *, + samples: int = 3, interval_s: float = 3.0, + smi_reader=read_hy_smi) -> Dict[int, Dict[str, float]]: + """Average a few hy-smi samples per device (drops idle flicker).""" + acc: Dict[int, Dict[str, List[float]]] = {} + for attempt in range(max(1, samples)): + parsed = parse_hy_smi(smi_reader()) + for index, row in parsed.items(): + if gpu_ids is not None and index not in gpu_ids: + continue + bucket = acc.setdefault(index, {"vram_percent": [], "util_percent": [], + "power_w": [], "temp_c": []}) + for key, value in row.items(): + bucket.setdefault(key, []).append(value) + if attempt + 1 < samples: + time.sleep(max(0.0, interval_s)) + out: Dict[int, Dict[str, float]] = {} + for index, bucket in acc.items(): + out[index] = { + key: (sum(values) / len(values) if values else 0.0) + for key, values in bucket.items() + } + return out + + +def total_vram_gb(card_index: int = 0) -> float: + """Total VRAM of one card in GB (driver sysfs, else the known default).""" + candidates = sorted(Path("/sys/class/drm").glob("card[0-9]/device/mem_info_vram_total")) + if candidates: + pick = candidates[card_index] if card_index < len(candidates) else candidates[0] + try: + return float(pick.read_text(encoding="utf-8").strip()) / (1024 ** 3) + except (OSError, ValueError): + pass + return DEFAULT_TOTAL_GB + + +def check_gpu(gpu_id: int, state: Dict[str, float], *, + min_free_gb: float = DEFAULT_MIN_FREE_GB, + total_gb: Optional[float] = None, + util_suspect_percent: float = DEFAULT_UTIL_SUSPECT_PERCENT, + power_suspect_w: float = DEFAULT_POWER_SUSPECT_W, + foreign_pids: Optional[List[int]] = None, + ) -> Dict[str, Any]: + """Decide usability (free VRAM >= ``min_free_gb``) and suspicion (busy).""" + vram = float(state.get("vram_percent") or 0.0) + util = float(state.get("util_percent") or 0.0) + power = float(state.get("power_w") or 0.0) + foreign = list(foreign_pids or []) + total = float(total_gb if total_gb is not None else total_vram_gb()) + free_gb = total * max(0.0, 100.0 - vram) / 100.0 + usable = free_gb >= float(min_free_gb) + reasons: List[str] = [] + if util > float(util_suspect_percent): + reasons.append(f"busy: HCU {util:.0f}% > {util_suspect_percent:.0f}%") + if power > float(power_suspect_w): + reasons.append(f"high power draw {power:.0f}W > {power_suspect_w:.0f}W") + # Foreign KFD entries alone are NOT evidence of contention: idle holders + # (long-lived runtimes in other containers) keep their handle while the + # device sits at 0% util / idle power, which is harmless for us. They are + # reported as notes, and the busy/power signals above decide suspicion. + notes: List[str] = [] + if foreign: + notes.append(f"{len(foreign)} foreign KFD pid(s) hold this GPU") + if not usable: + reasons.append( + f"only {free_gb:.1f}GB VRAM free (< {float(min_free_gb):.1f}GB)") + return { + "gpu": gpu_id, + "usable": usable, + "free_gb": round(free_gb, 2), + "total_gb": round(total, 1), + "min_free_gb": float(min_free_gb), + "vram_percent": round(vram, 1), + "util_percent": round(util, 1), + "power_w": round(power, 1), + "temp_c": round(float(state.get("temp_c") or 0.0), 1), + "foreign_pids": foreign[:8], + "measurement_suspect": bool(any( + r for r in reasons if "VRAM free" not in r)), + "reasons": reasons, + "notes": notes, + } + + +def preflight_gpus(gpu_ids: List[int], *, samples: int = 3, + interval_s: float = 3.0, enabled: bool = True, + min_free_gb: float = DEFAULT_MIN_FREE_GB, + **kwargs: Any) -> Dict[str, Any]: + """Return ``{gpu_id: check}`` plus a preferred order (clean devices first).""" + if not enabled: + return {"enabled": False, "gpus": {}, "preferred": list(gpu_ids), + "suspect_ids": []} + states = sample_gpu_state(gpu_ids, samples=samples, interval_s=interval_s) + foreign = foreign_kfd_pids() + checks: Dict[int, Dict[str, Any]] = {} + for gpu_id in gpu_ids: + checks[gpu_id] = check_gpu( + gpu_id, states.get(gpu_id) or {}, min_free_gb=min_free_gb, + foreign_pids=foreign, **kwargs, + ) + usable = [g for g in gpu_ids if checks[g]["usable"]] + clean = [g for g in usable if not checks[g]["measurement_suspect"]] + suspect = [g for g in usable if checks[g]["measurement_suspect"]] + busy_but_usable = [g for g in gpu_ids if not checks[g]["usable"]] + # Devices over the VRAM limit stay usable in principle (configuration + # choice), but when any device is within the limit they are pushed to the + # back so concurrent children do not collide with a full device. + preferred = (clean + suspect + busy_but_usable) if not usable \ + else (clean + suspect) + return { + "enabled": True, + "gpus": checks, + "preferred": preferred, + "clean_ids": clean, + "suspect_ids": suspect, + "over_limit_ids": busy_but_usable, + "foreign_kfd_pids": foreign[:8], + "sampled_at": time.time(), + "samples": samples, + } + + +def preflight_enabled(answers: Optional[Dict[str, Any]] = None) -> bool: + """``METAINFER_GPU_PREFLIGHT=0`` or a form switch can disable the probe.""" + env = str(os.environ.get("METAINFER_GPU_PREFLIGHT") or "").strip().lower() + if env in {"0", "false", "no", "off"}: + return False + if env in {"1", "true", "yes", "on"}: + return True + value = (answers or {}).get("gpu_preflight") + if value is None: + return True + return str(value).strip().lower() not in {"0", "false", "no", "off", ""} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/harness_io.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/harness_io.py new file mode 100644 index 00000000..82b09fc6 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/harness_io.py @@ -0,0 +1,95 @@ +"""Harness workspace IO: locate the evolvable-harness seed, read components, +and copy a seed workspace to a destination. + +This is the first piece of AHE "component observability" for dcu_kernel_auto_opt: +harness components (gates, and later prompts/planner/skills/tools/middleware/ +memory) are seeded as files under ``harness_default/`` so they can become a +versioned, evolvable workspace driven by the ``harness_evolve`` outer loop. + +Slice-1 semantics (important): this module is **read-only/additive**. +Runtime pipelines still read the current Python constants (``config.py``, +``w8a8_pipeline.py``); ``harness_io`` only exposes the seed for inspection and +copying. Consistency tests guard the YAML seed against the Python constants so +the two sources cannot drift before the loader/renderer wiring lands in a later +slice (wiring will make the harness workspace the single source of truth). +""" + +from __future__ import annotations + +import os +import shutil +from pathlib import Path +from typing import Any, Dict + +import yaml + +#: Name of the seed directory inside the dcu_kernel_auto_opt plugin tree. +HARNESS_DEFAULT_DIRNAME = "harness_default" +#: Optional env override for the harness root (used by AHE experiments to point +#: at an evolved workspace instead of the built-in seed). +ENV_HARNESS_ROOT = "METAINFER_HARNESS_ROOT" + + +def plugin_dir() -> Path: + """Return the dcu_kernel_auto_opt plugin root (parent of orchestrator/).""" + return Path(__file__).resolve().parents[1] + + +def default_harness_dir() -> Path: + """Built-in seed directory (this plugin's harness_default/).""" + return plugin_dir() / HARNESS_DEFAULT_DIRNAME + + +def harness_root() -> Path: + """Resolve the active harness root. + + Prefers ``METAINFER_HARNESS_ROOT`` (absolute or relative-to-plugin path), + otherwise falls back to the built-in ``harness_default/`` seed. + """ + override = os.environ.get(ENV_HARNESS_ROOT, "").strip() + if override: + path = Path(override).expanduser() + if not path.is_absolute(): + path = plugin_dir() / path + return path.resolve() + return default_harness_dir().resolve() + + +def _load_yaml(path: Path) -> Dict[str, Any]: + if not path.is_file(): + return {} + with path.open("r", encoding="utf-8") as handle: + data = yaml.safe_load(handle) + return data if isinstance(data, dict) else {} + + +def load_manifest(root: Path | None = None) -> Dict[str, Any]: + """Load the component inventory (manifest.yaml) of a harness workspace.""" + root = (root or harness_root()) + return _load_yaml(root / "manifest.yaml") + + +def load_gates(root: Path | None = None) -> Dict[str, Any]: + """Load the gates component (acceptance/plateau/ISA values) of a workspace.""" + root = (root or harness_root()) + return _load_yaml(root / "gates.yaml") + + +def seed_workspace(dst: Path, root: Path | None = None) -> Path: + """Copy the harness seed tree into ``dst`` (creating dirs as needed). + + Used by task staging / AHE experiments to materialize a fresh, evolvable + harness workspace snapshot from the seed. + """ + src = (root or harness_root()) + if not src.is_dir(): + raise FileNotFoundError(f"harness seed not found: {src}") + dst = Path(dst) + dst.mkdir(parents=True, exist_ok=True) + for item in src.iterdir(): + target = dst / item.name + if item.is_dir(): + shutil.copytree(item, target, dirs_exist_ok=True) + else: + shutil.copy2(item, target) + return dst diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/planner.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/planner.py new file mode 100644 index 00000000..e4ef12c1 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/planner.py @@ -0,0 +1,461 @@ +"""State-conditioned optimization plan selector (v0, not yet wired). + +DKAO currently chooses the next optimization direction with the round menus in +``prompts.py::w8a8_round_strategy`` (keyed by iteration number). This module is +the first step toward a *state-conditioned* plan selector that decides "what to +try this round" from the current measured state (history + PMC + budget), per +the design in ``docs/dkao_harness_eval_protocol.md`` / M1 slice-2. + +Slice-2 semantics: +- **Additive and inert**: nothing in the runtime pipelines calls this module yet. + ``prompts.py``/``w8a8_pipeline.py`` keep their current behaviour unchanged. +- The plan catalog lives as data in ``harness_default/planner_catalog.yaml`` + (read through ``harness_io`` with ``METAINFER_HARNESS_ROOT`` override) so it is + already an evolvable harness component; this module only consumes it. +- Selection layering: P0 repair -> P1 budget/phase -> P2 bottleneck signature -> + P3 coverage / anti-loop -> P4 fallback (legacy round-menu approximation). + Rules are deterministic pure functions, tested without GPU. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, List, Optional + +import yaml + +from .harness_io import harness_root + +# --------------------------------------------------------------------------- +# Catalog (data in harness_default/planner_catalog.yaml; fallback built-in) +# --------------------------------------------------------------------------- +BUILTIN_PLAN_IDS = [ + "repair_faster_wrong", "fix_build", "retry_same", + "bootstrap_correctness", "establish_arch", "architecture_explore", + "grid_splitk", "pipeline_tune", "memory_layout", "occupancy_resource", + "epilogue_fusion", "isa_guided_hip", "conditional_inline_asm", "consolidate", +] + + +def _builtin_catalog() -> Dict[str, Dict[str, Any]]: + return {pid: {"focus": pid.replace("_", " ")} for pid in BUILTIN_PLAN_IDS} + + +def _wired(name: str, root: Optional[Path]) -> bool: + """A component marked wired:false in the manifest is not applied at all.""" + try: + from . import gate_policy as _gates + return _gates.component_wired(name, root) + except Exception: # noqa: BLE001 + return True + + +def catalog(root: Optional[Path] = None) -> Dict[str, Dict[str, Any]]: + """Load the plan catalog from the harness workspace (YAML) or built-in.""" + base = (root or harness_root()) / "planner_catalog.yaml" + if not _wired("planner_catalog", root): + return _builtin_catalog() + if base.is_file(): + with base.open("r", encoding="utf-8") as fh: + data = yaml.safe_load(fh) or {} + entries = data.get("catalog") + if isinstance(entries, dict) and entries: + return {pid: dict(meta) for pid, meta in entries.items()} + return _builtin_catalog() + + +def _builtin_policy() -> Dict[str, Any]: + return { + "repair_priority": { + "faster_wrong": "repair_faster_wrong", + "infrastructure_failure": "retry_same", + "build_failure": "fix_build", + "no_valid_kernel": "establish_arch", + }, + "phase_gates": { + "consolidate_when_rounds_left_lte": 1, + "plateau_recent_valid_rounds": 3, + "plateau_min_improvement_percent": -2.0, + "plateau_max_improvement_percent_exclusive": 2.0, + "isa_plan": "isa_guided_hip", + "inline_asm_plan": "conditional_inline_asm", + "consolidate_plan": "consolidate", + }, + "bottleneck_to_plan": { + "occupancy_limited": "occupancy_resource", + "bank_conflicts": "memory_layout", + "lds_wait": "pipeline_tune", + "l2_low": "memory_layout", + "grid_limited": "grid_splitk", + }, + "coverage": {"max_consecutive_same_plan": 2, + "max_trials_per_plan": 2}, + "fallback": { + "fresh": "establish_arch", + "prefill": ["memory_layout", "pipeline_tune", "architecture_explore", + "memory_layout", "epilogue_fusion", "pipeline_tune", + "occupancy_resource", "consolidate"], + "m16": ["architecture_explore", "grid_splitk", "pipeline_tune", + "memory_layout", "pipeline_tune", "pipeline_tune", + "occupancy_resource", "consolidate"], + "small": ["establish_arch", "memory_layout", "occupancy_resource", + "architecture_explore", "memory_layout", "pipeline_tune", + "pipeline_tune", "consolidate"], + }, + } + + +def policy(root: Optional[Path] = None) -> Dict[str, Any]: + """Load the evolvable state->plan policy from the harness workspace.""" + base = (root or harness_root()) / "planner_policy.yaml" + if not _wired("planner_policy", root): + return _builtin_policy() + if base.is_file(): + with base.open("r", encoding="utf-8") as fh: + data = yaml.safe_load(fh) or {} + if isinstance(data, dict): + # only policy sections; metadata keys are ignored by selector + return data + return _builtin_policy() + + +# --------------------------------------------------------------------------- +# Context helpers +# --------------------------------------------------------------------------- +def _valid(record: Dict[str, Any]) -> bool: + metrics = record.get("metrics") or {} + return bool( + record.get("build_success") is True + and record.get("correctness_passed") is True + and metrics.get("graph_capture_passed") is True + ) + + +_INFRA_TOKENS = ("timeout", "timed out", "killed", "no result", "exit 143") + + +def _is_infra_failure(record: Dict[str, Any]) -> bool: + reason = str(record.get("failure_reason") or "").lower() + return any(token in reason for token in _INFRA_TOKENS) + + +def derive_ctx( + history: List[Dict[str, Any]], + pmc: Optional[Dict[str, Any]] = None, + iteration: int = 1, + max_iterations: int = 10, + *, + compiler_limitation_confirmed: bool = False, + plan_tags: Optional[List[str]] = None, +) -> Dict[str, Any]: + """Fold history + PMC into the compact state vector used by the selector. + + Alignment notes (v0, documented approximations): + - ``valid`` mirrors ``isa_round_policy``'s definition of a valid experiment + (build+correctness+graph-capture ok). + - plateau uses the last three *valid* improvements within [-2%, 2%). + - infra/build/faster-wrong signals mirror the repair branches already in + ``w8a8_round_strategy``. + - plan_tags: per-round plan ids from proposal (schema lands in a later + slice). Absent -> coverage guards are inert. + """ + pol = policy() + phase = pol.get("phase_gates") or {} + recent_n = int(phase.get("plateau_recent_valid_rounds", 3)) + plateau_min = float(phase.get("plateau_min_improvement_percent", -2.0)) + plateau_max = float( + phase.get("plateau_max_improvement_percent_exclusive", 2.0) + ) + valid = [r for r in history if _valid(r)] + recent = valid[-recent_n:] + improvements = [ + float((r.get("acceptance") or {}).get("improvement_percent", float("inf"))) + for r in recent + ] + plateau = ( + len(recent) == recent_n + and all(plateau_min <= v < plateau_max for v in improvements) + ) + faster_wrong = any( + r.get("build_success") is True + and r.get("correctness_passed") is False + and isinstance(r.get("speedup"), (int, float)) + and float(r["speedup"]) > 1.0 + for r in history + ) + last = history[-1] if history else None + last_present = last is not None + last_build_ok = bool(last and last.get("build_success") is True) + last_infra = bool(last and _is_infra_failure(last)) + + rounds_used = iteration + tried: Dict[str, int] = {} + consecutive = 0 + tags = plan_tags or [] + for t in tags: + tried[t] = tried.get(t, 0) + 1 + if tags: + for t in reversed(tags): + if t == tags[-1]: + consecutive += 1 + else: + break + + return { + "iteration": iteration, + "max_iterations": max_iterations, + "rounds_left": max(0, max_iterations - rounds_used + 1), + "valid_hip_rounds": len(valid), + "recent_improvements": improvements, + "plateau": plateau, + "faster_wrong": faster_wrong, + "last_present": last_present, + "last_build_ok": last_build_ok, + "last_infra": last_infra, + "compiler_limitation_confirmed": compiler_limitation_confirmed, + "pmc": pmc or {}, + "tried_counts": tried, + "consecutive_same": consecutive, + "plan_tags": tags, + } + + +# --------------------------------------------------------------------------- +# Bottleneck signature (P2) +# --------------------------------------------------------------------------- +def bottleneck_tags(pmc: Dict[str, Any]) -> List[str]: + """Map available PMC fields to bottleneck tags (heuristic, v0). + + Field names follow the counters referenced across the worker pipeline + (occupancy resources, lds bank conflicts / waits, grid vs CU, L2 hit). + Missing/None evidence yields []. Calibration against the real hipprof + schema happens when this selector is wired. + """ + tags: List[str] = [] + cu = pmc.get("device_cu_count") + grid = pmc.get("grid_blocks") + if isinstance(cu, (int, float)) and isinstance(grid, (int, float)): + if cu > 0 and grid < 2.0 * cu: + tags.append("grid_limited") + + lds_conflicts = pmc.get("lds_bank_conflicts") + lds_inst = pmc.get("lds_instructions") + if ( + isinstance(lds_conflicts, (int, float)) + and isinstance(lds_inst, (int, float)) + and lds_inst > 0 + and lds_conflicts / lds_inst > 1.5 + ): + tags.append("bank_conflicts") + + lds_wait = pmc.get("lds_wait_instructions") + if ( + isinstance(lds_wait, (int, float)) + and isinstance(lds_inst, (int, float)) + and lds_inst > 0 + and lds_wait >= lds_inst + ): + tags.append("lds_wait") + + occupancy = pmc.get("waves_per_cu") or pmc.get("active_waves_per_cu") + target = pmc.get("target_waves_per_cu") + if isinstance(occupancy, (int, float)) and isinstance(target, (int, float)): + if target > 0 and occupancy < target: + tags.append("occupancy_limited") + + l2 = pmc.get("l2_hit_rate") + if isinstance(l2, (int, float)) and l2 < 70.0: + tags.append("l2_low") + + # no tags -> caller falls back to P3/P4 + return tags + + +# --------------------------------------------------------------------------- +# Selector +# --------------------------------------------------------------------------- +_BOTTLENECK_TO_PLAN = { + "occupancy_limited": "occupancy_resource", + "bank_conflicts": "memory_layout", + "lds_wait": "pipeline_tune", + "l2_low": "memory_layout", + "grid_limited": "grid_splitk", +} + + +def _legacy_fallback(ctx: Dict[str, Any]) -> str: + """P4 fallback approximating today's round menus by M regime + iteration. + + Mirrors the *intent* of ``w8a8_round_strategy``'s small_m / m16 / large_m + portfolios (explicitly approximate; runtime is unchanged until wiring). + """ + pol = policy() + fallback = pol.get("fallback") or _builtin_policy()["fallback"] + m = int((ctx.get("shape") or {}).get("M", 0) or 0) + it = int(ctx.get("iteration", 1)) + if it <= 1: + return str(fallback.get("fresh") or "establish_arch") + regime = "prefill" if m >= 128 else ("m16" if m >= 16 else "small") + table = list(fallback.get(regime) or _builtin_policy()["fallback"][regime]) + if not table: + return "establish_arch" + # Cycle the portfolio instead of clamping to its last row: clamping used to + # pin long runs to whatever the table ended with (consolidate), which + # turned every remaining round into a no-op (observed on 9-8-8 iteration 3: + # 9 of 11 rounds were consolidate). + idx = (max(0, it - 2)) % len(table) + return str(table[idx]) + + +def choose_plan(ctx: Dict[str, Any]) -> str: + """Deterministically pick this round's plan id from the state vector. + + Layering: P0 repair -> P1 budget/phase gates -> P2 bottleneck signature + -> P3 coverage guards -> P4 legacy fallback. + """ + pol = policy() + repair = pol.get("repair_priority") or _builtin_policy()["repair_priority"] + phase = pol.get("phase_gates") or _builtin_policy()["phase_gates"] + bottleneck_map = pol.get("bottleneck_to_plan") or _BOTTLENECK_TO_PLAN + coverage = pol.get("coverage") or _builtin_policy()["coverage"] + + # P0: evolvable repair priorities. + if ctx.get("faster_wrong"): + return str(repair["faster_wrong"]) + if ctx.get("last_present"): + if ctx.get("last_infra"): + return str(repair["infrastructure_failure"]) + if not ctx.get("last_build_ok"): + return str(repair["build_failure"]) + if ctx.get("valid_hip_rounds", 0) == 0: + return str(repair["no_valid_kernel"]) + + # P1: evolvable budget / ISA phase gates. + if ctx.get("rounds_left", 1) <= int( + phase.get("consolidate_when_rounds_left_lte", 1) + ): + return str(phase.get("consolidate_plan") or "consolidate") + if ctx.get("plateau") and ctx.get("valid_hip_rounds", 0) >= max( + 1, int(ctx.get("max_iterations", 10)) - 2 + ): + if ctx.get("compiler_limitation_confirmed"): + return str(phase.get("inline_asm_plan") or "conditional_inline_asm") + return str(phase.get("isa_plan") or "isa_guided_hip") + + # P1.5: uncertainty fallback. With no bottleneck evidence the selector has + # nothing to condition on; the legacy menu is a known-good default there, + # and treating it as an available plan keeps policy evolution from ever + # being worse than the hand-tuned menu (it becomes an option HE can pick). + uncertainty = pol.get("uncertainty") or {} + if uncertainty.get("enabled"): + has_bottleneck = bool(bottleneck_tags(ctx.get("pmc") or {})) + if (not has_bottleneck + and ctx.get("valid_hip_rounds", 0) >= int( + uncertainty.get("min_valid_rounds", 1))): + return str(uncertainty.get("plan") or "legacy_menu") + + # P2: evolvable bottleneck -> plan mapping. + candidates: List[str] = [] + for tag in bottleneck_tags(ctx.get("pmc") or {}): + plan = bottleneck_map.get(tag) + if plan and plan not in candidates: + candidates.append(str(plan)) + + # P3: evolvable coverage / anti-loop limits. + tried = ctx.get("tried_counts") or {} + consecutive = int(ctx.get("consecutive_same", 0)) + max_consecutive = int(coverage.get("max_consecutive_same_plan", 2)) + max_trials = int(coverage.get("max_trials_per_plan", 2)) + for plan in candidates: + if consecutive >= max_consecutive and plan == candidates[0]: + continue + if tried.get(plan, 0) < max_trials: + return plan + + # P4: fallback + return _legacy_fallback(ctx) + + +def choose_plan_from_history( + history: List[Dict[str, Any]], + pmc: Optional[Dict[str, Any]] = None, + iteration: int = 1, + max_iterations: int = 10, + **overrides: Any, +) -> str: + """Convenience: derive ctx from history + pmc, then select the plan.""" + # Forward ctx-shaping overrides into derive_ctx (plan_tags/consecutive and + # compiler gate must be part of the derived state, not late-added). + derived_kwargs = {} + for key in ("plan_tags", "compiler_limitation_confirmed"): + if key in overrides: + derived_kwargs[key] = overrides.pop(key) + ctx = derive_ctx( + history, + pmc, + iteration=iteration, + max_iterations=max_iterations, + **derived_kwargs, + ) + ctx.update(overrides) + return choose_plan(ctx) + + +# --------------------------------------------------------------------------- +# Rendering (v0): plan -> worker-facing mandate text +# --------------------------------------------------------------------------- +_FAMILY_EXTRA: Dict[str, str] = { + "repair": ( + "Preserve the fast/previous architecture and make the smallest " + "correction for the reported issue; do not start a redesign." + ), + "bootstrap": ( + "Correctness first: a compile-clean kernel with exact int32 math " + "matters more than performance this round." + ), + "explore": ( + "Make ONE bounded architecture experiment; report grid_blocks, " + "waves_per_block and estimated active CUs in proposal.json." + ), + "optimize": ( + "One bounded mechanism per round with a falsifiable prediction; " + "keep all other paths byte-identical." + ), + "isa": ( + "ISA evidence is advisory; only shape compiler output through " + "HIP/DUMMA/intrinsics. Raw inline asm follows the control-plane gate." + ), +} + + +def render_plan( + plan_id: str, + *, + ctx: Optional[Dict[str, Any]] = None, + cat: Optional[Dict[str, Dict[str, Any]]] = None, +) -> str: + """Render a selected plan id into a worker-facing mandate (deterministic). + + v0 wording is intentionally concise and derives from the catalog's focus + + family discipline. When this selector replaces ``w8a8_round_strategy``, the + round prompt will embed this text; until then it is only used for parity + tests and offline exploration. + """ + entries = cat if cat is not None else catalog() + meta = entries.get(plan_id) or {} + focus = str(meta.get("focus") or plan_id.replace("_", " ")) + family = str(meta.get("family") or "") + lines = [ + f"Mandatory decision for this round: {plan_id}.", + f"Focus: {focus}.", + ] + extra = _FAMILY_EXTRA.get(family) + if extra: + lines.append(extra) + if ctx: + lines.append( + f"Round context: iteration {ctx.get('iteration', '?')} / " + f"{ctx.get('max_iterations', '?')}, " + f"{ctx.get('rounds_left', '?')} rounds left." + ) + return "\n".join(lines) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/predictions.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/predictions.py new file mode 100644 index 00000000..3e2807f6 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/predictions.py @@ -0,0 +1,141 @@ +"""Inner decision-observability helpers (M1 slice 3, additive & inert). + +The worker round record already stores the agent's ``hypothesis`` prose and +``profile_evidence``. This module adds *optional structured prediction* support: + +- proposal may carry a top-level ``prediction`` object: + { + "expected_us_range": [lo, hi], # optional numeric range for median_us + "direction": "improve|regress|flat", # optional relative to current best + "at_risk": true/false, # optional: change may regress P90/shapes + } +- after a round is measured, ``check_prediction`` returns a compact verdict + (``hit`` / ``miss`` / ``na``) plus the evidence used, or ``None`` when the + proposal declared no structured prediction (fully backward compatible). + +Runtime behaviour is unchanged: prompts do not yet instruct agents to emit +``prediction``; the hook in ``w8a8_pipeline`` only adds fields to the round +record when a structured prediction is present. ``plan_id`` capture (feeding +the planner's coverage guards) is handled separately in the pipeline hook. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + + +_INFRA_TOKENS = ("timeout", "timed out", "killed", "no result", "exit 143") + + +def parse_prediction(proposal: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Extract the structured prediction (if any) from a proposal dict.""" + prediction = proposal.get("prediction") + if not isinstance(prediction, dict): + return None + return dict(prediction) + + +def _is_infra_failure(failure_reason: Any) -> bool: + reason = str(failure_reason or "").lower() + return any(token in reason for token in _INFRA_TOKENS) + + +def check_prediction( + proposal: Dict[str, Any], + *, + candidate_us: float, + best_us: float, + passed: bool, + p90_guard_passed: bool, + failure_reason: Any = None, +) -> Optional[Dict[str, Any]]: + """Compare a declared structured prediction against the measured round. + + Returns None when the proposal declared nothing (or it is not dict-shaped). + Otherwise returns: + { + "declared": {...}, + "checked": "hit" | "miss" | "na", + "candidate_us": float, + "best_us": float, + "reason": str, + } + """ + declared = parse_prediction(proposal) + if declared is None: + return None + + na = not passed and _is_infra_failure(failure_reason) + if na: + return { + "declared": declared, + "checked": "na", + "candidate_us": candidate_us, + "best_us": best_us, + "reason": "infrastructure failure; prediction not judged", + } + + reasons: List[str] = [] + expected_range = declared.get("expected_us_range") + direction = declared.get("direction") + + range_ok: Optional[bool] = None + if isinstance(expected_range, (list, tuple)) and len(expected_range) == 2: + try: + lo, hi = float(expected_range[0]), float(expected_range[1]) + except (TypeError, ValueError): + lo = hi = None + if lo is not None: + range_ok = (lo <= candidate_us <= hi) + if not range_ok: + reasons.append( + f"median {candidate_us:.1f}us outside declared range " + f"[{lo:.1f}, {hi:.1f}]" + ) + + dir_ok: Optional[bool] = None + if isinstance(direction, str) and direction in ("improve", "regress", "flat"): + # relative comparison against the current best median (v0: flat = within + # +-2% of best, matching the plateau band used elsewhere) + if best_us > 0 and candidate_us != float("inf"): + delta = (best_us - candidate_us) / best_us * 100.0 + if direction == "improve": + dir_ok = delta >= 2.0 and passed + if not dir_ok: + reasons.append(f"expected improvement, got delta {delta:.1f}%") + elif direction == "regress": + dir_ok = delta <= -2.0 + if not dir_ok: + reasons.append(f"expected regression, got delta {delta:.1f}%") + else: # flat + dir_ok = -2.0 <= delta < 2.0 + if not dir_ok: + reasons.append(f"expected flat, got delta {delta:.1f}%") + + if not passed: + reasons.append("correctness failed") + + checks = [c for c in (range_ok, dir_ok) if c is not None] + if checks: + hit = passed and all(checks) + else: + # structured object present but nothing checkable declared + hit = passed + + if not hit: + reasons.append("no p90 guard evidence" if not p90_guard_passed else "") + reasons = [r for r in reasons if r] + + return { + "declared": declared, + "checked": "hit" if hit else "miss", + "candidate_us": candidate_us, + "best_us": best_us, + "reason": "; ".join(reasons) or "ok", + } + + +def plan_tag(proposal: Dict[str, Any]) -> Optional[str]: + """Optional plan_id from proposal (feeds planner coverage guards later).""" + pid = proposal.get("plan_id") + return pid if isinstance(pid, str) and pid.strip() else None diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/skill_store.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/skill_store.py index 99a38082..1c28ab7c 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/skill_store.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/skill_store.py @@ -268,6 +268,44 @@ def ccb_skills_root() -> Path: ) +def bundled_skills_root() -> Path: + """Skills bundled with this task plugin (seed content). + + The canonical skills this plugin curates (the INT8 W8A8 GEMM family) are + vendored under ``metainfer/tasks/dcu_kernel_auto_opt/skills/`` so a fresh + checkout carries them; ``seed_bundled_skills()`` installs any that are + missing from the live dsh library. + """ + return Path(__file__).resolve().parents[1] / "skills" + + +def seed_bundled_skills() -> Dict[str, Any]: + """Copy bundled plugin skills into the dsh library when missing. + + Idempotent and never overwrites: the live library is the authoritative + copy once an operator publishes/fuses over it, so only absent skills are + seeded. Mirrors to ccb are handled by ``sync_skill_libraries()``. + """ + source = bundled_skills_root() + root = dsh_skills_root() + if not source.exists(): + return {"added": [], "root": str(root)} + root.mkdir(parents=True, exist_ok=True) + added: list[str] = [] + for src_dir in sorted(source.glob("*/")): + if not (src_dir / "SKILL.md").is_file(): + continue + name = src_dir.name + dst = root / name + if dst.exists(): + continue + shutil.copytree( + src_dir, dst, ignore=shutil.ignore_patterns("*.bak-*") + ) + added.append(name) + return {"added": added, "root": str(root)} + + def existing_skills_root() -> Path: """The library publish/fuse write into (dsh, the authoritative library).""" return dsh_skills_root() @@ -358,12 +396,15 @@ def _write_sync_summary(workspace_dir: Path, summary: Dict[str, Any]) -> None: def sync_skill_libraries(*, workspace_dir: Path | None = None) -> Dict[str, Any]: - """Mirror the authoritative dsh library into the ccb library (one-way). + """Seed bundled plugin skills, then mirror dsh into ccb (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. + First installs any bundled plugin skills (metainfer/.../dcu_kernel_auto_opt/ + skills) that are missing from the dsh library, then 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. """ + seeded = seed_bundled_skills() source = dsh_skills_root() target = ccb_skills_root() added: list[str] = [] @@ -393,6 +434,7 @@ def sync_skill_libraries(*, workspace_dir: Path | None = None) -> Dict[str, Any] "added": added, "updated": updated, "skipped": skipped, + "seeded": sorted(seeded.get("added") or []), "ccb_only": sorted(_skill_name_set(target) - _skill_name_set(source)), "ts": time.time(), } diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/validation_budget.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/validation_budget.py new file mode 100644 index 00000000..189d996d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/validation_budget.py @@ -0,0 +1,95 @@ +"""Validation budget knobs shared by production DKAO tasks and AHE children. + +Two dimensions of the final serial validation are configurable without +touching the algorithmic path: + +* **scope** — which shapes the final validation must cover: ``api`` (every + shape of the frozen operator API, incl. un-optimized fallback shapes) or + ``task`` (only the shapes this task optimized). AHE evaluates a single + question per child, so the api-wide regression sweep is pure overhead + there; production tasks keep the full sweep by default. +* **bench profile** — sampling density of each benchmark: ``full`` (the + bench asset defaults) or ``quick`` (small warmup/sample/replay counts), + or explicit ``bench_warmups``/``bench_samples``/``bench_replays`` numbers. + +Resolution order (highest first): explicit environment variable, then the +task's ``answers`` (the WebUI form), then the conservative default. +""" + +from __future__ import annotations + +import os +from typing import Any, Dict, Mapping, Optional + +ENV_VALIDATE_SCOPE = "METAINFER_VALIDATE_SCOPE" +ENV_BENCH_WARMUPS = "METAINFER_BENCH_WARMUPS" +ENV_BENCH_SAMPLES = "METAINFER_BENCH_SAMPLES" +ENV_BENCH_REPLAYS = "METAINFER_BENCH_REPLAYS" + +VALID_SCOPES = ("api", "task") + +#: quick profile: still ~2x cheaper than the asset defaults (100/30/100) but +#: with enough warmup/sample depth that medians stay comparable across rounds +#: (a 10-warmup profile proved too noisy for decode-sized kernels). +QUICK_BENCH = {"warmups": 30, "samples": 30, "replays_per_sample": 50} + + +def _answers(answers: Optional[Mapping[str, Any]]) -> Mapping[str, Any]: + return answers if isinstance(answers, Mapping) else {} + + +def resolve_validation_scope(answers: Optional[Mapping[str, Any]] = None) -> str: + """``api`` (default) or ``task`` (optimal for single-shape evaluation).""" + env = str(os.environ.get(ENV_VALIDATE_SCOPE) or "").strip().lower() + if env in VALID_SCOPES: + return env + raw = str(_answers(answers).get("validation_scope") or "").strip().lower() + if raw in VALID_SCOPES: + return raw + if raw in {"task only", "task-only", "selected shapes only", + "task shapes only"}: + return "task" + if raw in {"all api shapes", "all", "api shapes"}: + return "api" + return "api" + + +def _int_or_none(value: Any) -> Optional[int]: + try: + parsed = int(value) + except (TypeError, ValueError): + return None + return parsed if parsed > 0 else None + + +def _positive_env(name: str) -> Optional[int]: + return _int_or_none(os.environ.get(name)) + + +def resolve_bench_kwargs( + answers: Optional[Mapping[str, Any]] = None, +) -> Dict[str, int]: + """Benchmark overrides for this task (``{}`` = asset defaults).""" + ans = _answers(answers) + profile = str(ans.get("bench_profile") or "").strip().lower() + kwargs: Dict[str, int] = {} + + warmups = _positive_env(ENV_BENCH_WARMUPS) + samples = _positive_env(ENV_BENCH_SAMPLES) + replays = _positive_env(ENV_BENCH_REPLAYS) + if warmups is None: + warmups = _int_or_none(ans.get("bench_warmups")) + if samples is None: + samples = _int_or_none(ans.get("bench_samples")) + if replays is None: + replays = _int_or_none(ans.get("bench_replays")) + + if profile in {"quick", "fast"}: + kwargs.update(QUICK_BENCH) + if warmups is not None: + kwargs["warmups"] = warmups + if samples is not None: + kwargs["samples"] = samples + if replays is not None: + kwargs["replays_per_sample"] = replays + return kwargs diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_promote.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_promote.py new file mode 100644 index 00000000..daa6366d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_promote.py @@ -0,0 +1,211 @@ +"""Kernel promotion: turn an evaluated task's accepted kernel into a variant. + +Both DKAO tasks and harness_evolve (HE) children write the same workspace +layout (``workers/*/accepted//kernel.hip`` + ``manifest.json``), so +one implementation can serve both. This module adds the conservative gates the +automatic path needs on top of :func:`variant_store.add_variant`: + +* correctness must not be known-failed; +* the candidate must beat the existing variant for the same shape by at least + ``min_improvement_percent`` (a strictly slower candidate is always rejected, + by ``add_variant`` itself); +* every write is backed up (``add_variant``'s ``backup=True``) and reported + with old/new medians, improvement and the backup path so a promotion can be + audited and rolled back. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict, Optional + +from .variant_store import ( + _parse_variant_header, + add_variant, + derive_variant_meta, + variant_path, +) + +#: shape-id prefix -> the model label used by the variant taxonomy. +MODEL_LABEL_BY_PREFIX = { + "hy3": "Hy3 (Hunyuan 3)", + "minimax": "MiniMax M3", + "glm": "GLM5.2", + "dsv4": "DeepSeek V4 Flash", + "deepseek": "DeepSeek V4 Flash", +} + + +def model_label_for(shape_id: str, fallback: str = "") -> str: + """Map a shape id (or model name) to a known variant-taxonomy label.""" + key = (shape_id or "").strip().lower() + for prefix, label in MODEL_LABEL_BY_PREFIX.items(): + if key.startswith(prefix): + return label + return fallback + + +def accepted_kernel_for(workspace_dir: Path, + shape_id: str) -> Optional[Dict[str, Any]]: + """Locate the accepted kernel + manifest for one shape in any worker lane.""" + root = Path(workspace_dir) / "workers" + if not root.is_dir(): + return None + for worker_root in sorted(root.glob("worker_*")): + candidate = worker_root / "accepted" / shape_id / "kernel.hip" + if not candidate.is_file(): + continue + manifest: Dict[str, Any] = {} + try: + manifest = json.loads( + (candidate.parent / "manifest.json").read_text(encoding="utf-8") + ) + except (OSError, ValueError): + manifest = {} + return { + "path": candidate, + "manifest": manifest, + "metrics": dict(manifest.get("metrics") or {}), + "commit": str(manifest.get("commit") or ""), + "shape": dict(manifest.get("shape") or {}), + } + return None + + +def _baseline_us(workspace_dir: Path, shape_id: str, + shape_params: Dict[str, Any]) -> Optional[float]: + """Baseline for the recorded speedup: task table first, fixed table next.""" + report = workspace_dir / "final_report.json" + if report.is_file(): + try: + initial = (json.loads(report.read_text(encoding="utf-8")) + .get("initial_metrics") or {}) + except (OSError, ValueError): + initial = {} + entry = initial.get(shape_id) + if isinstance(entry, dict) and entry.get("median_us"): + return float(entry["median_us"]) + if isinstance(entry, (int, float)): + return float(entry) + try: + from .w8a8_baselines import fixed_triton_graph_baseline + params = dict(shape_params or {}) + return float(fixed_triton_graph_baseline(shape_id, params)["median_us"]) + except Exception: # noqa: BLE001 - baseline is advisory + return None + + +def promote_variant( + *, + workspace_dir: Path, + answers: Dict[str, Any], + shape_id: str, + source_task: str = "", + correctness_ok: Optional[bool] = None, + min_improvement_percent: float = 0.0, + tp: Optional[int] = None, + m: Optional[int] = None, + model_label: Optional[str] = None, + dry_run: bool = False, +) -> Dict[str, Any]: + """Promote one accepted kernel into the shared variant tree. + + Returns a structured outcome with ``action`` one of ``added`` | ``updated`` + | ``skipped`` | ``rejected`` | ``no_kernel`` (or ``would-add`` / + ``would-update`` when ``dry_run``), plus old/new medians and improvement. + """ + workspace_dir = Path(workspace_dir) + found = accepted_kernel_for(workspace_dir, shape_id) + if found is None: + return {"ok": False, "action": "no_kernel", "shape": shape_id, + "reason": f"no accepted kernel under {workspace_dir}"} + metrics = dict(found["metrics"]) + new_median = metrics.get("median_us") + if correctness_ok is False: + return {"ok": False, "action": "skipped", "shape": shape_id, + "reason": "correctness failed for this candidate", + "new_median_us": new_median} + if new_median is None: + return {"ok": False, "action": "skipped", "shape": shape_id, + "reason": "accepted manifest carries no median_us"} + + answers_eff = dict(answers or {}) + label = model_label or model_label_for( + shape_id, str(answers_eff.get("model") or "")) + if label: + answers_eff["model"] = label + meta = derive_variant_meta(answers_eff, shape_id) + if tp is not None: + meta["tp"] = int(tp) + if m is not None: + meta["m"] = int(m) + + target = variant_path(meta) + old_median: Optional[float] = None + if target.is_file(): + try: + header = _parse_variant_header( + target.read_text(encoding="utf-8", errors="replace")) + value = header.get("median_us") + old_median = float(value) if value is not None else None + except OSError: + old_median = None + + improvement: Optional[float] = None + if old_median and new_median: + improvement = (old_median - float(new_median)) / old_median * 100.0 + if improvement < float(min_improvement_percent): + return { + "ok": True, "action": "skipped", "shape": shape_id, + "reason": (f"improvement {improvement:.2f}% < required " + f"{float(min_improvement_percent):.2f}%"), + "old_median_us": old_median, "new_median_us": float(new_median), + "improvement_percent": improvement, "path": str(target), + "meta": meta, + } + + baseline = _baseline_us(workspace_dir, shape_id, found["shape"]) + if baseline and metrics.get("median_us"): + metrics["baseline_us"] = baseline + metrics["speedup"] = float(baseline) / float(metrics["median_us"]) + + if dry_run: + return { + "ok": True, + "action": "would-update" if target.is_file() else "would-add", + "shape": shape_id, "reason": "dry-run: no write performed", + "old_median_us": old_median, "new_median_us": float(new_median), + "improvement_percent": improvement, + "baseline_us": metrics.get("baseline_us"), + "speedup": metrics.get("speedup"), + "path": str(target), "meta": meta, + } + + try: + result = add_variant( + meta=meta, + kernel_source=found["path"].read_text(encoding="utf-8", + errors="replace"), + commit=found["commit"], + metrics=metrics, + source_task=source_task, + backup=True, + reject_slower_than_existing=True, + ) + except ValueError as exc: + return {"ok": False, "action": "rejected", "shape": shape_id, + "reason": str(exc), "old_median_us": old_median, + "new_median_us": float(new_median), + "improvement_percent": improvement, "path": str(target), + "meta": meta} + return { + "ok": True, "action": result["action"], "shape": shape_id, + "reason": "", + "old_median_us": old_median, "new_median_us": float(new_median), + "improvement_percent": improvement, + "baseline_us": metrics.get("baseline_us"), + "speedup": metrics.get("speedup"), + "path": result["path"], "backup": result.get("backup"), + "meta": meta, + } diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_store.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_store.py index d1774e03..879a1c38 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_store.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/variant_store.py @@ -277,3 +277,120 @@ def _parse_variant_header(text: str) -> Dict[str, Any]: elif key in _HEADER_TEXT_FIELDS: fields[key] = value return fields + + +def backfill_variant_baselines() -> Dict[str, Any]: + """Enrich variant headers that lack ``baseline_us``/``speedup``. + + Older variants (added before the add-route baseline fallback, or from + stopped tasks without a final report) only carry ``median_us``/``p90_us``. + This resolves each variant's ``(K, N)`` from its model/TP operator + catalog, reads the fixed Triton baseline, and appends ``baseline_us`` and + ``speedup`` to the header comment block. Idempotent: files that already + carry a speedup, or for which no median / catalog entry / baseline can be + resolved, are skipped untouched. + """ + from metainfer.tasks.dcu_kernel_auto_opt.api.int8w8a8gemm.int8_w8a8_gemm_api import ( # noqa: PLC0415 + GLM52_TP4_OPERATOR_KN, + GLM52_TP8_OPERATOR_KN, + HY3_TP4_OPERATOR_KN, + HY3_TP8_OPERATOR_KN, + MINIMAX_TP4_OPERATOR_KN, + MINIMAX_TP8_OPERATOR_KN, + TP4_OPERATOR_KN, + TP8_OPERATOR_KN, + ) + from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.w8a8_baselines import ( # noqa: PLC0415 + fixed_triton_graph_baseline, + ) + + tables: Dict[tuple, Dict[str, tuple]] = { + ("deepseek-v4", "TP4"): TP4_OPERATOR_KN, + ("deepseek-v4", "TP8"): TP8_OPERATOR_KN, + ("hy3", "TP4"): HY3_TP4_OPERATOR_KN, + ("hy3", "TP8"): HY3_TP8_OPERATOR_KN, + ("minimax-m3", "TP4"): MINIMAX_TP4_OPERATOR_KN, + ("minimax-m3", "TP8"): MINIMAX_TP8_OPERATOR_KN, + ("glm52", "TP4"): GLM52_TP4_OPERATOR_KN, + ("glm52", "TP8"): GLM52_TP8_OPERATOR_KN, + } + root = variant_root() + updated: list[str] = [] + skipped: Dict[str, int] = { + "has_speedup": 0, "no_median": 0, "no_catalog": 0, "no_baseline": 0, + } + if not root.exists(): + return {"updated": updated, "skipped": skipped} + for file in sorted(root.rglob("*.hip")): + if file.name.startswith("w8a8_gemm_variants"): + continue + rel = file.relative_to(root).parts + if len(rel) < 4: + continue + model, tp, m_dir = rel[1], rel[2], rel[3] + header = _parse_variant_header( + file.read_text(encoding="utf-8", errors="replace") + ) + if header.get("baseline_us") is not None or header.get("speedup") is not None: + skipped["has_speedup"] += 1 + continue + median = header.get("median_us") + if median is None: + skipped["no_median"] += 1 + continue + table = tables.get((model, tp)) + if table is None: + skipped["no_catalog"] += 1 + continue + kn = next( + ( + entry + for key, entry in table.items() + if key.replace(".", "_") == file.stem + ), + None, + ) + if kn is None or not m_dir.startswith("M"): + skipped["no_catalog"] += 1 + continue + k, n = kn + m = int(m_dir[1:]) + try: + baseline = fixed_triton_graph_baseline( + header.get("shape") or file.stem, + { + "tp_size": int(tp[2:]), + "M": m, + "N": n, + "K": k, + }, + ).get("median_us") + except ValueError: + baseline = None + if baseline is None: + skipped["no_baseline"] += 1 + continue + speedup = float(baseline) / float(median) + new_line = f"// baseline_us={float(baseline):.6g} speedup={speedup:.6g}" + text = file.read_text(encoding="utf-8", errors="replace") + lines = text.splitlines(keepends=True) + # Insert immediately after the ``// @@variant`` header line. Kernel + # source comments also start with ``//`` and there is no blank line + # between the header and the kernel, so only the @@variant anchor is + # reliable. + insert_at: int | None = None + for index, line in enumerate(lines): + if line.startswith("// @@variant"): + insert_at = index + break + if insert_at is None: + skipped["no_catalog"] += 1 + continue + file.write_text( + "".join(lines[: insert_at + 1]) + + new_line + "\n" + + "".join(lines[insert_at + 1:]), + encoding="utf-8", + ) + updated.append(str(file.relative_to(root))) + return {"updated": updated, "skipped": skipped} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py index 41d3d61a..4ec2e402 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py @@ -26,6 +26,100 @@ from .gpu_binding import bind_worker_gpu from .experience_store import load_verified_experience from .guidance import claim_next_guidance +from .predictions import check_prediction, plan_tag +from . import gate_policy as _gates +from .planner import choose_plan_from_history, render_plan + + +_ENV_PLANNER = "METAINFER_PLANNER" + + +def _planner_enabled() -> bool: + return os.environ.get(_ENV_PLANNER, "").strip().lower() in { + "1", "true", "yes", + } + + +def _round_strategy_text( + shape: Dict[str, Any], + iteration: int, + history: list, + pmc_evidence: Dict[str, Any], + isa_policy: Dict[str, Any], + *, + plan_sink: "Path | None" = None, + shape_id: str = "", +) -> str: + """Round mandate text: legacy menu by default; planner (v0) when + METAINFER_PLANNER=1. Behaviour is unchanged unless the env var is set.""" + max_iterations = int((isa_policy or {}).get("max_iterations") or 10) + if not _planner_enabled(): + return w8a8_round_strategy( + shape, + iteration, + history, + pmc_evidence, + max_iterations=max_iterations, + isa_policy=isa_policy, + ) + tags = [ + r.get("plan_id") for r in (history or []) + if isinstance(r.get("plan_id"), str) and r.get("plan_id") + ] + phase = (isa_policy or {}).get("phase") + plan_id = choose_plan_from_history( + (history or []), + pmc_evidence, + iteration=iteration, + max_iterations=max_iterations, + shape=shape, + plan_tags=tags, + compiler_limitation_confirmed=(phase == "conditional_inline_asm"), + ) + if plan_id == "legacy_menu": + # Deferred to the hand-tuned menu: render exactly what the non-planner + # path would have produced, so this choice can never be worse than the + # legacy baseline. + if plan_sink is not None: + try: + plan_sink.parent.mkdir(parents=True, exist_ok=True) + with plan_sink.open("a", encoding="utf-8") as fh: + fh.write(json.dumps({ + "ts": time.time(), "iteration": iteration, + "shape_id": shape_id or str(shape.get("id") or ""), + "plan_id": "legacy_menu", "source": "planner", + }) + "\n") + except OSError: + pass + return w8a8_round_strategy( + shape, iteration, history, pmc_evidence, + max_iterations=max_iterations, isa_policy=isa_policy, + ) + + # Record the planner's OWN decision (distinct from the plan id an agent + # may report for itself) so the AHE mechanism gate can verify that a + # planner_policy change really drove the round mandate. + if plan_sink is not None: + try: + plan_sink.parent.mkdir(parents=True, exist_ok=True) + with plan_sink.open("a", encoding="utf-8") as fh: + fh.write(json.dumps({ + "ts": time.time(), + "iteration": iteration, + "shape_id": shape_id or str(shape.get("id") or ""), + "plan_id": plan_id, + "source": "planner", + }) + "\n") + except OSError: + pass + return render_plan( + plan_id, + ctx={ + "iteration": iteration, + "max_iterations": max_iterations, + "rounds_left": max(0, max_iterations - iteration + 1), + }, + ) from .isa_analysis import ( analyze_inline_asm_source, evaluate_inline_asm_gate, @@ -108,7 +202,7 @@ def isa_round_policy( 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) + required_hip_rounds = _gates.required_hip_rounds(max_iterations) base = { "phase": "hip_only", "skill_allowed": False, @@ -116,7 +210,7 @@ def isa_round_policy( "plateau": False, "max_iterations": max_iterations, "required_valid_hip_rounds": required_hip_rounds, - "required_valid_isa_guided_rounds": _REQUIRED_VALID_ISA_GUIDED_ROUNDS, + "required_valid_isa_guided_rounds": _gates.isa_required_valid_rounds(), "reason": "At least eight HIP-only rounds are required.", } valid = [ @@ -149,9 +243,10 @@ def isa_round_policy( for record in recent ] plateau = ( - len(recent) == 3 + len(recent) == _gates.plateau_recent_valid_rounds() and all( - -_PLATEAU_MAX_REGRESSION_PERCENT <= value < 2.0 + -_gates.plateau_max_regression_percent() <= value + < _gates.plateau_window_upper_exclusive_percent() for value in improvements ) ) @@ -185,7 +280,7 @@ def isa_round_policy( == "isa_guided_hip" ] policy["valid_isa_guided_rounds"] = len(valid_isa) - if len(valid_isa) < _REQUIRED_VALID_ISA_GUIDED_ROUNDS: + if len(valid_isa) < _gates.isa_required_valid_rounds(): return policy previous = valid_isa[-1] @@ -232,7 +327,7 @@ def phase_extension_reason( 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) + required_hip = _gates.required_hip_rounds(max_iterations) valid_hip = [ record for record in valid if (record.get("isa_policy") or {}).get("phase", "hip_only") @@ -257,10 +352,10 @@ def phase_extension_reason( if (record.get("isa_policy") or {}).get("phase") == "isa_guided_hip" ] - if len(valid_isa) < _REQUIRED_VALID_ISA_GUIDED_ROUNDS: + if len(valid_isa) < _gates.isa_required_valid_rounds(): return ( "need " - f"{_REQUIRED_VALID_ISA_GUIDED_ROUNDS - len(valid_isa)} more valid " + f"{_gates.isa_required_valid_rounds() - len(valid_isa)} more valid " "ISA-guided HIP experiment(s)" ) @@ -433,7 +528,7 @@ def evaluate_candidate_acceptance( 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 + p90_guard_passed = candidate_p90 <= best_p90 * _gates.p90_tolerance() accepted = ( passed and candidate_us < best_us @@ -452,8 +547,10 @@ def evaluate_candidate_acceptance( shadow_eligible = ( passed and not accepted - and _SHADOW_MIN_IMPROVEMENT_PERCENT <= improvement - < minimum_improvement_percent + and _gates.shadow_enabled() + and _gates.shadow_min_improvement_percent() <= improvement + < min(minimum_improvement_percent, + _gates.shadow_max_exclusive_percent()) and p90_guard_passed and improves_shadow ) @@ -1226,7 +1323,7 @@ def _plan(config: OptimizerConfig) -> Dict[str, Any]: "final validated result versus fixed baseline" ), "round_acceptance_improvement_percent": ( - ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + _gates.round_acceptance_improvement_percent() ), "shape_scope": config.shape_scope, "assignment_mode": config.assignment_mode, @@ -1898,6 +1995,7 @@ def _run_worker( verified_experience, comparison_target, isa_policy=isa_policy, continuation=shape_session_id is not None, + attempt_limit=attempt_limit, ) prompt_file = root / "logs" / ( f"{shape_id}-iteration-{iteration}.prompt.txt" @@ -2345,7 +2443,7 @@ def preflight() -> list[str]: metrics=metrics, best_metrics=best_metrics, minimum_improvement_percent=( - ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT + _gates.round_acceptance_improvement_percent() ), shadow_metrics=shadow_metrics, ) @@ -2424,6 +2522,21 @@ def preflight() -> list[str]: ), "timestamp": time.time(), } + # Inner decision-observability hook (additive, inert unless the + # proposal carries a structured `prediction` / `plan_id`): + _prediction = check_prediction( + proposal, + candidate_us=candidate_us, + best_us=float(best_metrics["median_us"]), + passed=passed, + p90_guard_passed=p90_guard_passed, + failure_reason=failure_reason, + ) + if _prediction is not None: + experiment["prediction_checked"] = _prediction + _plan_id = plan_tag(proposal) + if _plan_id is not None: + experiment["plan_id"] = _plan_id candidate_files = archive_iteration_candidate( source, iteration_dir, changed ) @@ -2711,6 +2824,7 @@ def _worker_prompt( comparison_baseline: Dict[str, Any] | None = None, isa_policy: Dict[str, Any] | None = None, continuation: bool = False, + attempt_limit: int | None = None, ) -> str: history = history or [] verified_experience = verified_experience or [] @@ -2724,6 +2838,14 @@ def _worker_prompt( max_iterations=10, history=history, ) + # The planner's budget gate must see the budget the worker will + # actually run with (``attempt_limit`` grows with repair/replacement + # rounds), otherwise ``rounds_left`` collapses to <=1 early and the + # selector degrades into consolidate for the rest of the run. + if attempt_limit is not None and int(attempt_limit) > 0: + if int(isa_policy.get("max_iterations") or 0) != int(attempt_limit): + isa_policy = dict(isa_policy) + isa_policy["max_iterations"] = int(attempt_limit) if not isa_policy["skill_allowed"]: pmc_evidence = dict(pmc_evidence) pmc_evidence.pop("isa", None) @@ -2825,10 +2947,10 @@ def _worker_prompt( }, ''' 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, + round_strategy = _round_strategy_text( + shape, iteration, history, pmc_evidence, isa_policy, + plan_sink=root / "planner_plans.jsonl", + shape_id=shape_id, ) prompt_best = _compact_metrics_for_prompt(best) prompt_baseline = _compact_metrics_for_prompt(comparison_baseline) @@ -2885,7 +3007,8 @@ def _worker_prompt( 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 +profile_evidence, architecture, optional isa_optimization when required, +optional plan_id and prediction, and files_changed=["csrc/w8a8_gemm_hip.hip"]. """ return f"""You are {assignment.worker_id}, a shape-specialized native @@ -3052,12 +3175,26 @@ def _worker_prompt( 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: +diff, and then write `{root / 'source' / 'proposal.json'}` as strict JSON. + +Optional decision-observability fields (recommended when you can commit to +them): top-level `plan_id` names the plan family this round belongs to; top-level +`prediction` declares a falsifiable expectation — `expected_us_range` is the +median-us interval you expect after this change, and/or `direction` is your +expected movement relative to the current best median (`improve|flat|regress`). +The control plane logs both and checks `prediction` against the measured round; +omit fields you are not confident about. ```json {{ "iteration": {iteration}, "hypothesis": "one falsifiable bottleneck hypothesis and the focused change", + "plan_id": "occupancy_resource or the plan this round belongs to", + "prediction": {{ + "expected_us_range": [340.0, 355.0], + "direction": "improve", + "at_risk": false + }}, "profile_evidence": {{ "observed_best": {json.dumps(prompt_best)}, "path": "scalar_lds or dumma_m16n16k32", diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/_build_summary.py b/metainfer/tasks/dcu_kernel_auto_opt/reports/_build_summary.py new file mode 100644 index 00000000..1de3823e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_build_summary.py @@ -0,0 +1,251 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Aggregate "best variant performance" per optimized operator shape across the +7 successfully finished (status=success + final_report.json) real DKAO runs. + +Input : nodes/worker29/workspaces//final_report.json +Output: reports/dkao_optimized_operators_bestvariant.{csv,xlsx} +""" +import csv +import datetime as _dt +import json +import os + +ROOT = "/root/zth_agent/MetaInfer/nodes/worker29/workspaces" +OUT_DIR = os.path.dirname(os.path.abspath(__file__)) + +# (task_id, display model) +RUNS = [ + ("hy3-dsh-tp4-m16-2-368c654c", "Hy3 (gfx928 INT8 W8A8)"), + ("hy3-tp4-m4096-dsh-test1-bfdf3002", "Hy3 (gfx928 INT8 W8A8)"), + ("hy3-dsh-tp8-m4096-1-0ebb994d", "Hy3 (gfx928 INT8 W8A8)"), + ("hy3-dsh-tp8-m16-9-8-0161e718", "Hy3 (gfx928 INT8 W8A8)"), + ("minimax-dsh-tp8-m16-1-78402260", "MiniMax-M3 (gfx928 INT8 W8A8)"), + ("minimaxm3-dsh-tp4-m4096-1-0c2f84a9", "MiniMax-M3 (gfx928 INT8 W8A8)"), + ("minimaxm3-dsh-tp8-m4096-1-b0482833", "MiniMax-M3 (gfx928 INT8 W8A8)"), + ("glm5-2-dsh-tp8-m4096-1-e6a280a2", "GLM5.2 (gfx928 INT8 W8A8)"), +] + +STATE_VALIDATED = "最终验收通过" +STATE_WORKER_ONLY = "仅worker验收(未过最终验收)" + + +def short_model(repo: str) -> str: + r = os.path.basename(repo or "") + for key in ("hy3", "glm5.2", "glm", "minimax", "dsv4", "deepseek"): + if r.lower().startswith(key): + return r + return r or "?" + + +def dt(ts): + try: + return _dt.datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M") + except Exception: + return "?" + + +def collect(task_id, model_label): + ws = os.path.join(ROOT, task_id) + rep = json.load(open(os.path.join(ws, "final_report.json"))) + cfg = rep.get("config") or {} + plan_shapes = cfg.get("shapes") or [] + if not plan_shapes: # fallback: report.config may nest plan differently + plan_shapes = (rep.get("config") or {}).get("shapes") or [] + final_target = (rep.get("final_target") or {}).get("shapes") or {} + fv = rep.get("final_validation") or {} + workers = rep.get("workers") or {} + wv = rep.get("worker_validation") or {} + rows = [] + for s in plan_shapes: + sid = s.get("id") + op = s.get("operator") + tp = s.get("tp_size") + ft = final_target.get(sid) + fr = fv.get(sid) + if ft is None: + # some reports store target by sid only; try direct median from fv + print(f" [warn] {task_id}: no final_target entry for {sid}") + continue + baseline_us = ft.get("baseline_us") + final_us = ft.get("final_us") + improv = ft.get("improvement_percent") + target_met = ft.get("target_met") + passed = fr.get("passed") if fr else None + p90 = fr.get("p90_us") if fr else None + shape = (fr or {}).get("shape") or {} + M = shape.get("M", s.get("M")) + N = shape.get("N", s.get("N")) + K = shape.get("K", s.get("K")) + # worker best (accepted best during parallel exploration) across workers + wbest = None + wbest_meta = None + for wid, wd in workers.items(): + sh = (wd.get("shapes") or {}).get(sid) + if not sh: + continue + m = sh.get("metrics") or {} + mu = m.get("median_us") + art = sh.get("artifact") or {} + if mu is not None and (wbest is None or mu < wbest): + wbest = mu + wbest_meta = (wid, art.get("source"), m.get("p90_us")) + wv_rec = wv.get(sid) + if wv_rec: + wm = wv_rec.get("metrics") or {} + wmu = wm.get("median_us") + if wmu is not None and (wbest is None or wmu < wbest): + wbest = wmu + wbest_meta = (wv_rec.get("worker_id"), None, wm.get("p90_us")) + # sanity + if fr and final_us is not None: + delta = abs((fr.get("median_us") or 0) - final_us) + if delta > 1e-6 * max(1.0, final_us): + print(f" [note] {task_id}/{sid}: final_us={final_us} vs fv.median={fr.get('median_us')}") + tops = None + if final_us and M and N and K: + tops = 2.0 * M * N * K / (final_us * 1e-6) / 1e12 + speedup = (baseline_us / final_us) if (baseline_us and final_us) else None + rows.append({ + "task_id": task_id, + "model": model_label, + "kernel_repo": os.path.basename((cfg.get("kernel_repo") or "")) or short_model(cfg.get("kernel_repo")), + "tp": tp, + "operator": op, + "M": M, "N": N, "K": K, + "dtype": "INT8 W8A8", + "baseline_us": baseline_us, + "worker_best_us": wbest, + "best_variant_source": (f"{wbest_meta[0]} {wbest_meta[1]}" if wbest_meta else ""), + "final_us": final_us, + "final_p90_us": p90, + "speedup_x": speedup, + "improvement_pct": improv, + "target_met": target_met, + "passed": passed, + "logical_tops": tops, + "finished_at": dt(rep.get("finished_at") or rep.get("started_at")), + "validation_state": STATE_VALIDATED, + }) + return rows + + +def collect_worker_only(task_id, model_label): + """Rows for a run that never produced final_report.json: use the best + worker-accepted variant measured during parallel exploration, and label the + row as not-final-validated.""" + ws = os.path.join(ROOT, task_id) + plan = json.load(open(os.path.join(ws, "plan.json"))) + tp_default = (plan.get("shapes") or [{}])[0].get("tp_size") + baseline = {} + try: + baseline = (json.load(open(os.path.join(ws, "shared_baseline", "results.json"))) + or {}).get("shapes") or {} + except OSError: + pass + best = {} + for wid in sorted(os.listdir(os.path.join(ws, "workers"))): + rp = os.path.join(ws, "workers", wid, "result.json") + if not os.path.isfile(rp): + continue + try: + data = json.load(open(rp)) + except ValueError: + continue + for sid, rec in (data.get("shapes") or {}).items(): + met = rec.get("metrics") or {} + med = met.get("median_us") + if med is None: + continue + cur = best.get(sid) + if cur is None or med < cur["median_us"]: + best[sid] = {"median_us": med, "p90_us": met.get("p90_us"), + "samples": len(met.get("latency_samples_us") or []), + "worker": wid, "candidate": rec.get("candidate")} + rows = [] + for s in plan.get("shapes") or []: + sid = s.get("id") + rec = best.get(sid) + if not rec: + print(f" [warn] {task_id}: no worker measurement for {sid}") + continue + M, N, K = s.get("M"), s.get("N"), s.get("K") + base = baseline.get(sid) or {} + baseline_us = base.get("baseline_us") or base.get("median_us") + final_us = rec["median_us"] + tops = (2.0 * M * N * K / (final_us * 1e-6) / 1e12) if (M and N and K) else None + rows.append({ + "task_id": task_id, + "model": model_label, + "kernel_repo": os.path.basename((plan.get("kernel_repo") or "")) or short_model(plan.get("kernel_repo")), + "tp": s.get("tp_size", tp_default), + "operator": s.get("operator"), + "M": M, "N": N, "K": K, + "dtype": "INT8 W8A8", + "baseline_us": baseline_us, + "worker_best_us": final_us, + "best_variant_source": f"{rec['worker']} accepted/ (n={rec['samples']} samples)", + "final_us": final_us, + "final_p90_us": rec.get("p90_us"), + "speedup_x": (baseline_us / final_us) if (baseline_us and final_us) else None, + "improvement_pct": ((baseline_us / final_us - 1) * 100) if (baseline_us and final_us) else None, + "target_met": None, + "passed": None, + "logical_tops": tops, + "finished_at": "", + "validation_state": STATE_WORKER_ONLY, + }) + return rows + + +def main(): + all_rows = [] + for tid, label in RUNS: + print(f"== {tid} ==") + rows = [] + rp = os.path.join(ROOT, tid, "final_report.json") + if os.path.isfile(rp) and os.path.getsize(rp) > 0: + try: + if json.load(open(rp)).get("status") == "success": + rows = collect(tid, label) + except ValueError as exc: + print(f" [warn] unreadable final_report ({exc})") + if rows: + print(f" final-validated shapes={len(rows)}") + else: + rows = collect_worker_only(tid, label) + print(f" worker-only shapes={len(rows)} (no success final_report)") + all_rows.extend(rows) + print(f"TOTAL rows = {len(all_rows)}") + + headers = [ + "task_id", "model", "kernel_repo", "tp", "operator", "M", "N", "K", "dtype", + "baseline_us", "worker_best_us", "best_variant_source", + "final_us", "final_p90_us", "speedup_x", "improvement_pct", + "target_met", "passed", "logical_tops", "finished_at", "validation_state", + ] + csv_path = os.path.join(OUT_DIR, "dkao_optimized_operators_bestvariant.csv") + with open(csv_path, "w", newline="", encoding="utf-8-sig") as f: + w = csv.DictWriter(f, fieldnames=headers, extrasaction="ignore") + w.writeheader() + w.writerows(all_rows) + print("wrote", csv_path) + + # quick console summary + print("\n%-20s %-16s %2s %-24s %6s %9s %9s %9s %s" % ( + "task", "repo", "tp", "op", "M", "baseline", "final_us", "speedup", "state")) + for r in all_rows: + print("%-20s %-16s %2s %-24s %6s %9.2f %9.2f %8.2fx %s" % ( + r["task_id"][:20], r["kernel_repo"][:16], r["tp"], r["operator"], + r["M"], r["baseline_us"], r["final_us"], r["speedup_x"], + r["validation_state"])) + + import json as _j + with open(os.path.join(OUT_DIR, "dkao_optimized_operators_bestvariant.json"), "w") as f: + _j.dump(all_rows, f, indent=2, ensure_ascii=False, default=str) + return all_rows + + +if __name__ == "__main__": + main() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_o_proj_ab.sh b/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_o_proj_ab.sh new file mode 100755 index 00000000..995cb57e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_o_proj_ab.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# A/B diagnostic for the 9.8 run o_proj gate failure (worker build vs final build). +set -u +WS=/root/zth_agent/MetaInfer/nodes/worker29/workspaces/hy3-dsh-tp8-m16-9-8-0161e718 +GPU=${GPU:-1} + +run_bench () { + local root=$1 label=$2 + echo "=== $label ===" + echo "root=$root gpu=$GPU" + cd "$root/source" || return 1 + env HIP_VISIBLE_DEVICES="$GPU" MAX_JOBS=2 PYTHONDONTWRITEBYTECODE=1 \ + PYTORCH_ROCM_ARCH=gfx928 \ + TORCH_EXTENSIONS_DIR="$root/cache/torch" \ + TRITON_CACHE_DIR="$root/cache/triton" \ + XDG_CACHE_HOME="$root/cache/xdg" \ + TMPDIR="$root/cache/tmp" \ + python3 "$root/source/w8a8_bench.py" --source "$root/source" \ + --m 16 --n 4096 --k 1024 \ + --reference-cache-dir "$root/cache/references" 2>&1 | tail -20 +} + +run_bench "$WS/workers/worker_1" "WORKER_1 explore-time build (bench target of 11.044us)" +run_bench "$WS/final" "FINAL synthesized prebuilt-object build (gate measured 11.745us)" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/_make_xlsx.py b/metainfer/tasks/dcu_kernel_auto_opt/reports/_make_xlsx.py new file mode 100644 index 00000000..72e4d538 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_make_xlsx.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Build the multi-sheet .xlsx summary of best-variant performances.""" +import csv +import json +import os + +from openpyxl import Workbook +from openpyxl.styles import Alignment, Border, Font, PatternFill, Side +from openpyxl.utils import get_column_letter + +OUT_DIR = os.path.dirname(os.path.abspath(__file__)) +CSV = os.path.join(OUT_DIR, "dkao_optimized_operators_bestvariant.csv") + +HDR = [ + ("task_id", "任务ID"), + ("model", "模型(任务)"), + ("kernel_repo", "kernel_repo"), + ("tp", "TP"), + ("operator", "算子"), + ("M", "M"), + ("N", "N"), + ("K", "K"), + ("dtype", "数据类型"), + ("baseline_us", "Triton基线(µs)"), + ("worker_best_us", "探索期Worker最优(µs)"), + ("best_variant_source", "最优variant来源"), + ("final_us", "最终验收最优variant(µs)"), + ("final_p90_us", "最终P90(µs)"), + ("speedup_x", "加速比(vs基线×)"), + ("improvement_pct", "提升(%)"), + ("target_met", "达标(≥3%)"), + ("passed", "正确性"), + ("logical_tops", "计算性能(TOPS)"), + ("finished_at", "任务完成时间"), + ("validation_state", "验收状态"), +] +COL_KEYS = [h[0] for h in HDR] +COL_LABELS = [h[1] for h in HDR] +FMT = { # key -> openpyxl number format + "baseline_us": "0.00", + "worker_best_us": "0.00", + "final_us": "0.00", + "final_p90_us": "0.00", + "speedup_x": "0.00", + "improvement_pct": "0.0", + "logical_tops": "0.0", +} + +thin = Side(style="thin", color="D0D0D0") +border = Border(left=thin, right=thin, top=thin, bottom=thin) +head_fill = PatternFill("solid", fgColor="1F4E78") +head_font = Font(bold=True, color="FFFFFF", size=11) +alt_fill = PatternFill("solid", fgColor="EAF1F8") + + +def read_rows(): + rows = list(csv.DictReader(open(CSV, encoding="utf-8-sig"))) + for r in rows: + for k in ("tp", "M", "N", "K", "baseline_us", "worker_best_us", + "final_us", "final_p90_us", "speedup_x", "improvement_pct", + "logical_tops"): + v = r.get(k) + if v in (None, ""): + r[k] = None + else: + try: + r[k] = float(v) if "." in v or "e" in v.lower() else int(float(v)) + except ValueError: + pass + return rows + + +def sheet_from_rows(ws, rows, title_row=True, num_keys=None): + num_keys = num_keys or FMT + ws.append(COL_LABELS) + for c in range(1, len(COL_LABELS) + 1): + cell = ws.cell(row=1, column=c) + cell.fill, cell.font = head_fill, head_font + cell.alignment = Alignment(horizontal="center", vertical="center") + cell.border = border + ws.freeze_panes = "A2" + ws.auto_filter.ref = f"A1:{get_column_letter(len(COL_LABELS))}{1 + len(rows)}" + for i, r in enumerate(rows, start=2): + for j, k in enumerate(COL_KEYS, start=1): + v = r.get(k) + cell = ws.cell(row=i, column=j, value=(v if not isinstance(v, float) else round(v, 4))) + cell.border = border + if k in num_keys and isinstance(v, (int, float)): + cell.number_format = num_keys[k] + cell.alignment = Alignment(horizontal="right") + elif k in ("target_met", "passed"): + cell.alignment = Alignment(horizontal="center") + cell.value = "✔" if v is True or str(v).lower() == "true" else str(v) + elif isinstance(v, (int, float)): + cell.alignment = Alignment(horizontal="right") + else: + cell.alignment = Alignment(horizontal="left", vertical="center") + if i % 2 == 0: + cell.fill = alt_fill + widths = { + "task_id": 34, "model": 30, "kernel_repo": 24, "tp": 6, "operator": 26, + "M": 8, "N": 8, "K": 8, "dtype": 12, "baseline_us": 13, + "worker_best_us": 15, "best_variant_source": 40, "final_us": 16, + "final_p90_us": 12, "speedup_x": 12, "improvement_pct": 10, + "target_met": 10, "passed": 10, "logical_tops": 14, "finished_at": 18, + "validation_state": 26, + } + for j, k in enumerate(COL_KEYS, start=1): + ws.column_dimensions[get_column_letter(j)].width = widths.get(k, 14) + ws.row_dimensions[1].height = 24 + + +def main(): + rows = read_rows() + wb = Workbook() + + ws = wb.active + ws.title = "汇总-最优variant" + sheet_from_rows(ws, rows) + + # --- by model & TP summary --- + ws2 = wb.create_sheet("按模型×TP汇总") + groups = {} + for r in rows: + key = (r["model"], r["kernel_repo"], r["tp"]) + groups.setdefault(key, []).append(r) + head2 = ["模型(任务)", "kernel_repo", "TP", "算子数", "平均加速比×", + "最小加速比×", "最大加速比×", "平均最终µs", "平均TOPS", "平均提升%"] + ws2.append(head2) + for c in range(1, len(head2) + 1): + cell = ws2.cell(row=1, column=c) + cell.fill, cell.font = head_fill, head_font + cell.alignment = Alignment(horizontal="center") + cell.border = border + for i, ((model, repo, tp), rr) in enumerate(sorted(groups.items()), start=2): + sp = [r["speedup_x"] for r in rr if r["speedup_x"] is not None] + fin = [r["final_us"] for r in rr if r["final_us"] is not None] + tops = [r["logical_tops"] for r in rr if r["logical_tops"] is not None] + imp = [r["improvement_pct"] for r in rr if r["improvement_pct"] is not None] + vals = [model, repo, tp, len(rr), + sum(sp) / len(sp), min(sp), max(sp), + sum(fin) / len(fin), + sum(tops) / len(tops), + sum(imp) / len(imp)] + ws2.append(vals) + for j, v in enumerate(vals, start=1): + cell = ws2.cell(row=i, column=j) + cell.border = border + if j >= 5 and j != 4: + cell.number_format = "0.00" + if j == 4: + cell.alignment = Alignment(horizontal="right") + if i % 2 == 0: + for j in range(1, len(head2) + 1): + ws2.cell(row=i, column=j).fill = alt_fill + for j, w in enumerate([26, 24, 6, 9, 13, 13, 13, 14, 12, 12], start=1): + ws2.column_dimensions[get_column_letter(j)].width = w + ws2.freeze_panes = "A2" + ws2.auto_filter.ref = f"A1:{get_column_letter(len(head2))}{ws2.max_row}" + + # --- by operator (across runs) --- + ws3 = wb.create_sheet("按算子汇总") + ops = {} + for r in rows: + ops.setdefault(r["operator"], []).append(r) + head3 = ["算子", "出现次数", "平均加速比×", "最大加速比×", "平均TOPS", + "覆盖模型×TP×M"] + ws3.append(head3) + for c in range(1, len(head3) + 1): + cell = ws3.cell(row=1, column=c) + cell.fill, cell.font = head_fill, head_font + cell.alignment = Alignment(horizontal="center") + cell.border = border + for i, (op, rr) in enumerate(sorted(ops.items()), start=2): + sp = [r["speedup_x"] for r in rr if r["speedup_x"] is not None] + tops = [r["logical_tops"] for r in rr if r["logical_tops"] is not None] + cov = ", ".join(sorted({f"{r['kernel_repo']}/TP{r['tp']}/M{r['M']}" for r in rr})) + vals = [op, len(rr), sum(sp) / len(sp), max(sp), + sum(tops) / len(tops), cov] + ws3.append(vals) + for j, v in enumerate(vals, start=1): + cell = ws3.cell(row=i, column=j) + cell.border = border + if j in (3, 4, 5): + cell.number_format = "0.00" + if j in (2, 3, 4, 5): + cell.alignment = Alignment(horizontal="right") + if i % 2 == 0: + for j in range(1, len(head3) + 1): + ws3.cell(row=i, column=j).fill = alt_fill + for j, w in enumerate([24, 10, 13, 13, 12, 70], start=1): + ws3.column_dimensions[get_column_letter(j)].width = w + ws3.freeze_panes = "A2" + + # --- source runs --- + ws4 = wb.create_sheet("数据来源任务") + src = {} + for r in rows: + src.setdefault(r["task_id"], r) + head4 = ["任务ID", "模型", "kernel_repo", "TP", "优化算子数", "完成时间", "状态"] + ws4.append(head4) + for c in range(1, len(head4) + 1): + cell = ws4.cell(row=1, column=c) + cell.fill, cell.font = head_fill, head_font + cell.alignment = Alignment(horizontal="center") + cell.border = border + for i, (tid, r) in enumerate(sorted(src.items()), start=2): + ws4.append([tid, r["model"], r["kernel_repo"], r["tp"], + sum(1 for x in rows if x["task_id"] == tid), + r["finished_at"], "success(final_report)"]) + for j in range(1, len(head4) + 1): + ws4.cell(row=i, column=j).border = border + if i % 2 == 0: + for j in range(1, len(head4) + 1): + ws4.cell(row=i, column=j).fill = alt_fill + for j, w in enumerate([34, 30, 24, 6, 12, 18, 22], start=1): + ws4.column_dimensions[get_column_letter(j)].width = w + ws4.freeze_panes = "A2" + + # --- notes --- + ws5 = wb.create_sheet("字段说明") + notes = [ + ["说明", "本表汇总 dcu_kernel_auto_opt 任务下所有已完成并通过最终串行验收(status=success, final_report.json)的真实优化任务，", + "取每个 (模型×TP×算子×M) 最优已验收 variant 的最终性能；每个 shape 一行。"], + ["硬件/算子", "worker29, 4×K500SM_AI / gfx928 (Hygon DCU, CDNA 系); 算子 = INT8 W8A8 GEMM (HIP C++, DUMMA Tensor Core), dtype int8 × fp32 scale → bf16 out"], + ["Triton基线(µs)", "固定 Triton Graph 基线中位延迟（fixed table / 实测 graph replay），与 custom 同协议（GPU event、warmup、graph replay median）"], + ["探索期Worker最优(µs)", "并行探索阶段各 worker 在验收轮记录到的最快 median（accepted 轮）。部分行略快于最终验收值，属重测协议差异(最终门限 ≤1.05×worker best)"], + ["最终验收最优variant(µs)", "串行最终验收重新测量并全部 shape 通过(正确性+性能门)的 median —— 本表采用的“最优 variant 性能”主列"], + ["加速比(×)", "= Triton基线 / 最终验收median"], + ["提升(%)", "= (Triton基线/最终验收median − 1) × 100%；任务门槛 minimum_improvement_percent=3.0%（对 plan shape 的最终验收 vs 固定基线）"], + ["计算性能(TOPS)", "= 2×M×N×K / final_us，纯算法逻辑 INT8 运算速率（非实测带宽）"], + ["达标/正确性", "target_met: 最终验收 ≥3% 提升门槛；passed: 最终串行验证正确性通过(CPU int64 exact reference)"], + ["验收状态", "'最终验收通过' = 该任务串行最终验收全部 shape 通过并写出 final_report.json (status=success)；'仅worker验收(未过最终验收)' = 该任务未产出成功 final_report，本行取的是 worker 并行探索阶段已验收最优 variant 的 median（未经最终门限确认）"], + ["覆盖", "Hy3: TP4 M16/M4096, TP8 M4096; MiniMax-M3: TP4 M4096, TP8 M16/M4096; GLM5.2: TP8 M4096 —— 共 33 个优化 shape，全部为各任务自身 plan 的优化算子（不含默认42-shape回归项）"], + ["生成时间", "2026-09（数据截至各任务 final_report.json 快照）"], + ] + for i, row in enumerate(notes, start=1): + ws5.append(row) + for j in range(1, 4): + c = ws5.cell(row=i, column=j) + c.alignment = Alignment(wrap_text=True, vertical="top") + if i == 1: + for j in range(1, 3): + ws5.cell(row=i, column=j).font = Font(bold=True) + ws5.column_dimensions["A"].width = 24 + ws5.column_dimensions["B"].width = 60 + ws5.column_dimensions["C"].width = 70 + + xlsx = os.path.join(OUT_DIR, "dkao_optimized_operators_bestvariant.xlsx") + wb.save(xlsx) + print("wrote", xlsx) + print("sheets:", wb.sheetnames) + + +if __name__ == "__main__": + main() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.csv b/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.csv new file mode 100644 index 00000000..10929acb --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.csv @@ -0,0 +1,34 @@ +task_id,model,kernel_repo,tp,operator,M,N,K,dtype,baseline_us,worker_best_us,best_variant_source,final_us,final_p90_us,speedup_x,improvement_pct,target_met,passed,logical_tops,finished_at +hy3-dsh-tp4-m16-2-368c654c,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP4-m16-2,4,qkv_proj,16,2560,4096,INT8 W8A8,80.19,27.359989881515503,worker_0 accepted/hy3_tp4_qkv_proj_m16/kernel.hip,27.722550630569458,27.860960960388184,2.89259098373059,189.259098373059,True,True,12.103659741538994,2026-08-26 14:59 +hy3-dsh-tp4-m16-2-368c654c,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP4-m16-2,4,o_proj,16,4096,2048,INT8 W8A8,54.617,15.513530373573303,worker_1 accepted/hy3_tp4_o_proj_m16/kernel.hip,15.558480024337769,15.603289604187012,3.510432890267166,251.0432890267166,True,True,17.25332137715848,2026-08-26 14:59 +hy3-dsh-tp4-m16-2-368c654c,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP4-m16-2,4,shared_gate_up_proj,16,768,4096,INT8 W8A8,73.262,11.581599712371826,worker_2 accepted/hy3_tp4_shared_gate_up_proj_m16/kernel.hip,11.451255083084106,11.476860046386719,6.397726665631898,539.7726665631898,True,True,8.790590661865588,2026-08-26 14:59 +hy3-dsh-tp4-m16-2-368c654c,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP4-m16-2,4,shared_down_proj,16,4096,384,INT8 W8A8,23.216,9.17439490556717,worker_3 accepted/hy3_tp4_shared_down_proj_m16/kernel.hip,9.128850102424622,9.246439933776855,2.5431461508863897,154.31461508863896,True,True,5.513470747715742,2026-08-26 14:59 +hy3-tp4-m4096-dsh-test1-bfdf3002,Hy3 (gfx928 INT8 W8A8),hy3_TP4_M4096_dsh-test1,4,qkv_proj,4096,2560,4096,INT8 W8A8,24045.12,756.7258453369141,worker_0 accepted/hy3_tp4_qkv_proj_m4096/kernel.hip,756.7258453369141,760.15869140625,31.77520650070368,3077.520650070368,True,True,113.51448671844342,2026-08-22 05:44 +hy3-tp4-m4096-dsh-test1-bfdf3002,Hy3 (gfx928 INT8 W8A8),hy3_TP4_M4096_dsh-test1,4,o_proj,4096,4096,2048,INT8 W8A8,19881.949,803.8980102539062,worker_1 accepted/hy3_tp4_o_proj_m4096/kernel.hip,803.8980102539062,805.1364135742188,24.731929605996175,2373.1929605996174,True,True,85.48282973644304,2026-08-22 05:44 +hy3-tp4-m4096-dsh-test1-bfdf3002,Hy3 (gfx928 INT8 W8A8),hy3_TP4_M4096_dsh-test1,4,shared_gate_up_proj,4096,768,4096,INT8 W8A8,6678.282,220.71203231811523,worker_2 accepted/hy3_tp4_shared_gate_up_proj_m4096/kernel.hip,220.71203231811523,221.53764724731445,30.25789726939083,2925.7897269390833,True,True,116.75758455641257,2026-08-22 05:44 +hy3-tp4-m4096-dsh-test1-bfdf3002,Hy3 (gfx928 INT8 W8A8),hy3_TP4_M4096_dsh-test1,4,shared_down_proj,4096,4096,384,INT8 W8A8,4381.289,214.83811378479004,worker_3 accepted/hy3_tp4_shared_down_proj_m4096/kernel.hip,214.83811378479004,215.11651992797852,20.393443802010253,1939.3443802010254,True,True,59.97493489868936,2026-08-22 05:44 +hy3-dsh-tp8-m4096-1-0ebb994d,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M4096-1,8,qkv_proj,4096,1280,4096,INT8 W8A8,19381.383,418.2191276550293,worker_0 accepted/hy3_tp8_qkv_proj_m4096/kernel.hip,318.4573459625244,326.0519790649414,60.86021643313191,5986.021643313191,True,True,134.86789833717404,2026-08-27 07:23 +hy3-dsh-tp8-m4096-1-0ebb994d,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M4096-1,8,o_proj,4096,4096,1024,INT8 W8A8,14151.792,322.98728942871094,worker_1 accepted/hy3_tp8_o_proj_m4096/kernel.hip,282.8514766693115,283.3587074279785,50.0325901304917,4903.25901304917,True,True,121.47625592271805,2026-08-27 07:23 +hy3-dsh-tp8-m4096-1-0ebb994d,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M4096-1,8,shared_gate_up_proj,4096,384,4096,INT8 W8A8,16712.231,115.95122337341309,worker_2 accepted/hy3_tp8_shared_gate_up_proj_m4096/kernel.hip,115.34959316253662,115.53359985351562,144.8833111743286,14388.331117432861,True,True,111.70305446889755,2026-08-27 07:23 +hy3-dsh-tp8-m4096-1-0ebb994d,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M4096-1,8,shared_down_proj,4096,4096,192,INT8 W8A8,10762.616,171.18735313415527,worker_3 accepted/hy3_tp8_shared_down_proj_m4096/kernel.hip,171.81203842163086,172.1600341796875,62.64180379251588,6164.180379251588,True,True,37.497086951439755,2026-08-27 07:23 +minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,qkv_proj,16,1280,6144,INT8 W8A8,110.478,33.20169925689697,worker_0 accepted/minimax_tp8_qkv_proj_m16/kernel.hip,33.19578528404236,33.328590393066406,3.328073098879459,232.8073098879459,True,True,7.581029876132359,2026-08-28 17:17 +minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,qkv_proj_and_indexer_qk,16,1536,6144,INT8 W8A8,111.486,39.72720980644226,worker_1 accepted/minimax_tp8_qkv_proj_and_indexer_qk_m16/kernel.hip,39.73748564720154,39.8534893989563,2.8055625106680915,180.55625106680915,True,True,7.599622449219237,2026-08-28 17:17 +minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,o_proj,16,6144,1024,INT8 W8A8,38.655,11.368839740753174,worker_2 accepted/minimax_tp8_o_proj_m16/kernel.hip,11.77299976348877,12.421009540557861,3.2833602969975013,228.33602969975013,True,True,17.10070466699301,2026-08-28 17:17 +minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,shared_gate_up_proj,16,768,6144,INT8 W8A8,104.527,17.114474773406982,worker_3 accepted/minimax_tp8_shared_gate_up_proj_m16/kernel.hip,17.31948971748352,17.501089572906494,6.0352239993816355,503.5223999381635,True,True,8.718209743071991,2026-08-28 17:17 +minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,shared_down_proj,16,6144,384,INT8 W8A8,23.152,8.95284503698349,worker_3 accepted/minimax_tp8_shared_down_proj_m16/kernel.hip,9.341430068016052,10.175809860229492,2.478421379962978,147.8421379962978,True,True,8.082003660070677,2026-08-28 17:17 +minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,qkv_proj,4096,2304,6144,INT8 W8A8,31556.192,1559.944839477539,worker_0 accepted/minimax_tp4_qkv_proj_m4096/kernel.hip,1561.9083404541016,1577.305908203125,20.20361322279994,1920.361322279994,True,True,74.24514869949742,2026-08-23 18:45 +minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,qkv_proj_and_indexer_qk,4096,2560,6144,INT8 W8A8,35918.768,1081.603660583496,worker_1 accepted/minimax_tp4_qkv_proj_and_indexer_qk_m4096/kernel.hip,1089.0154266357422,1102.6435089111328,32.9827908048673,3198.27908048673,True,True,118.31698222866211,2026-08-23 18:45 +minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,o_proj,4096,6144,2048,INT8 W8A8,29702.473,765.9097671508789,worker_2 accepted/minimax_tp4_o_proj_m4096/kernel.hip,752.4849700927734,753.6545562744141,39.47251331324,3847.251331324,True,True,136.98508169710212,2026-08-23 18:45 +minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,shared_gate_up_proj,4096,1536,6144,INT8 W8A8,19776.488,986.5248489379883,worker_3 accepted/minimax_tp4_shared_gate_up_proj_m4096/kernel.hip,962.3474502563477,965.7874298095703,20.55025759639306,1955.025759639306,True,True,80.33419874225937,2026-08-23 18:45 +minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,shared_down_proj,4096,6144,768,INT8 W8A8,11759.473,636.2568092346191,worker_3 accepted/minimax_tp4_shared_down_proj_m4096/kernel.hip,634.1561698913574,634.6337509155273,18.543496946524407,1754.3496946524406,True,True,60.95455267843923,2026-08-23 18:45 +minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,qkv_proj,4096,1280,6144,INT8 W8A8,41000.618,475.5491065979004,worker_0 accepted/minimax_tp8_qkv_proj_m4096/kernel.hip,474.8045539855957,475.41736602783203,86.35262163311907,8535.262163311907,True,True,135.68637642417067,2026-08-29 03:28 +minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,qkv_proj_and_indexer_qk,4096,1536,6144,INT8 W8A8,35604.698,531.2225532531738,worker_1 accepted/minimax_tp8_qkv_proj_and_indexer_qk_m4096/kernel.hip,519.466609954834,520.4890060424805,68.54087888939718,6754.087888939718,True,True,148.82460170966874,2026-08-29 03:28 +minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,o_proj,4096,6144,1024,INT8 W8A8,19897.517,397.44619369506836,worker_2 accepted/minimax_tp8_o_proj_m4096/kernel.hip,394.50483322143555,394.9032211303711,50.43668752426038,4943.668752426038,True,True,130.64379245024566,2026-08-29 03:28 +minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,shared_gate_up_proj,4096,768,6144,INT8 W8A8,31882.257,319.16600227355957,worker_3 accepted/minimax_tp8_shared_gate_up_proj_m4096/kernel.hip,311.9345474243164,312.3681640625,102.20816278048036,10120.816278048036,True,True,123.91928365478229,2026-08-29 03:28 +minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,shared_down_proj,4096,6144,384,INT8 W8A8,22389.162,312.4720001220703,worker_3 accepted/minimax_tp8_shared_down_proj_m4096/kernel.hip,311.00934982299805,311.1349678039551,71.98871034823276,7098.871034823276,True,True,62.14396076195009,2026-08-29 03:28 +glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,fused_qkv_a_proj,4096,2624,6144,INT8 W8A8,68661.417,937.5793075561523,worker_0 accepted/glm_tp8_fused_qkv_a_proj_m4096/kernel.hip,938.0712509155273,939.3240356445312,73.19424503521313,7219.4245035213125,True,True,140.78913965554716,2026-08-31 11:23 +glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,q_b_proj,4096,2048,2048,INT8 W8A8,14385.571,220.13651847839355,worker_1 accepted/glm_tp8_q_b_proj_m4096/kernel.hip,215.435791015625,215.74857711791992,66.77428542482365,6577.428542482365,True,True,159.48946182998893,2026-08-31 11:23 +glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,kv_b_proj,4096,3584,512,INT8 W8A8,6939.939,137.40396976470947,worker_1 accepted/glm_tp8_kv_b_proj_m4096/kernel.hip,134.75635528564453,134.93956565856934,51.499901324055074,5049.990132405507,True,True,111.55233090221004,2026-08-31 11:23 +glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,o_proj,4096,6144,2048,INT8 W8A8,54775.797,759.9214935302734,worker_2 accepted/glm_tp8_o_proj_m4096/kernel.hip,755.8016204833984,756.3575744628906,72.47377554571302,7147.377554571302,True,True,136.3839562001365,2026-08-31 11:23 +glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,shared_gate_up_proj,4096,512,6144,INT8 W8A8,21254.589,273.8345146179199,worker_3 accepted/glm_tp8_shared_gate_up_proj_m4096/kernel.hip,272.7110004425049,272.99739837646484,77.93814318275388,7693.814318275387,True,True,94.49491855548743,2026-08-31 11:23 +glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,shared_down_proj,4096,6144,256,INT8 W8A8,17700.262,278.44139099121094,worker_3 accepted/glm_tp8_shared_down_proj_m4096/kernel.hip,273.6263179779053,274.4647216796875,64.6877176537867,6368.77176537867,True,True,47.089410050974806,2026-08-31 11:23 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.json b/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.json new file mode 100644 index 00000000..7d4d11bb --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.json @@ -0,0 +1,728 @@ +[ + { + "task_id": "hy3-dsh-tp4-m16-2-368c654c", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP4-m16-2", + "tp": 4, + "operator": "qkv_proj", + "M": 16, + "N": 2560, + "K": 4096, + "dtype": "INT8 W8A8", + "baseline_us": 80.19, + "worker_best_us": 27.359989881515503, + "best_variant_source": "worker_0 accepted/hy3_tp4_qkv_proj_m16/kernel.hip", + "final_us": 27.722550630569458, + "final_p90_us": 27.860960960388184, + "speedup_x": 2.89259098373059, + "improvement_pct": 189.259098373059, + "target_met": true, + "passed": true, + "logical_tops": 12.103659741538994, + "finished_at": "2026-08-26 14:59" + }, + { + "task_id": "hy3-dsh-tp4-m16-2-368c654c", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP4-m16-2", + "tp": 4, + "operator": "o_proj", + "M": 16, + "N": 4096, + "K": 2048, + "dtype": "INT8 W8A8", + "baseline_us": 54.617, + "worker_best_us": 15.513530373573303, + "best_variant_source": "worker_1 accepted/hy3_tp4_o_proj_m16/kernel.hip", + "final_us": 15.558480024337769, + "final_p90_us": 15.603289604187012, + "speedup_x": 3.510432890267166, + "improvement_pct": 251.0432890267166, + "target_met": true, + "passed": true, + "logical_tops": 17.25332137715848, + "finished_at": "2026-08-26 14:59" + }, + { + "task_id": "hy3-dsh-tp4-m16-2-368c654c", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP4-m16-2", + "tp": 4, + "operator": "shared_gate_up_proj", + "M": 16, + "N": 768, + "K": 4096, + "dtype": "INT8 W8A8", + "baseline_us": 73.262, + "worker_best_us": 11.581599712371826, + "best_variant_source": "worker_2 accepted/hy3_tp4_shared_gate_up_proj_m16/kernel.hip", + "final_us": 11.451255083084106, + "final_p90_us": 11.476860046386719, + "speedup_x": 6.397726665631898, + "improvement_pct": 539.7726665631898, + "target_met": true, + "passed": true, + "logical_tops": 8.790590661865588, + "finished_at": "2026-08-26 14:59" + }, + { + "task_id": "hy3-dsh-tp4-m16-2-368c654c", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP4-m16-2", + "tp": 4, + "operator": "shared_down_proj", + "M": 16, + "N": 4096, + "K": 384, + "dtype": "INT8 W8A8", + "baseline_us": 23.216, + "worker_best_us": 9.17439490556717, + "best_variant_source": "worker_3 accepted/hy3_tp4_shared_down_proj_m16/kernel.hip", + "final_us": 9.128850102424622, + "final_p90_us": 9.246439933776855, + "speedup_x": 2.5431461508863897, + "improvement_pct": 154.31461508863896, + "target_met": true, + "passed": true, + "logical_tops": 5.513470747715742, + "finished_at": "2026-08-26 14:59" + }, + { + "task_id": "hy3-tp4-m4096-dsh-test1-bfdf3002", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3_TP4_M4096_dsh-test1", + "tp": 4, + "operator": "qkv_proj", + "M": 4096, + "N": 2560, + "K": 4096, + "dtype": "INT8 W8A8", + "baseline_us": 24045.12, + "worker_best_us": 756.7258453369141, + "best_variant_source": "worker_0 accepted/hy3_tp4_qkv_proj_m4096/kernel.hip", + "final_us": 756.7258453369141, + "final_p90_us": 760.15869140625, + "speedup_x": 31.77520650070368, + "improvement_pct": 3077.520650070368, + "target_met": true, + "passed": true, + "logical_tops": 113.51448671844342, + "finished_at": "2026-08-22 05:44" + }, + { + "task_id": "hy3-tp4-m4096-dsh-test1-bfdf3002", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3_TP4_M4096_dsh-test1", + "tp": 4, + "operator": "o_proj", + "M": 4096, + "N": 4096, + "K": 2048, + "dtype": "INT8 W8A8", + "baseline_us": 19881.949, + "worker_best_us": 803.8980102539062, + "best_variant_source": "worker_1 accepted/hy3_tp4_o_proj_m4096/kernel.hip", + "final_us": 803.8980102539062, + "final_p90_us": 805.1364135742188, + "speedup_x": 24.731929605996175, + "improvement_pct": 2373.1929605996174, + "target_met": true, + "passed": true, + "logical_tops": 85.48282973644304, + "finished_at": "2026-08-22 05:44" + }, + { + "task_id": "hy3-tp4-m4096-dsh-test1-bfdf3002", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3_TP4_M4096_dsh-test1", + "tp": 4, + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 768, + "K": 4096, + "dtype": "INT8 W8A8", + "baseline_us": 6678.282, + "worker_best_us": 220.71203231811523, + "best_variant_source": "worker_2 accepted/hy3_tp4_shared_gate_up_proj_m4096/kernel.hip", + "final_us": 220.71203231811523, + "final_p90_us": 221.53764724731445, + "speedup_x": 30.25789726939083, + "improvement_pct": 2925.7897269390833, + "target_met": true, + "passed": true, + "logical_tops": 116.75758455641257, + "finished_at": "2026-08-22 05:44" + }, + { + "task_id": "hy3-tp4-m4096-dsh-test1-bfdf3002", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3_TP4_M4096_dsh-test1", + "tp": 4, + "operator": "shared_down_proj", + "M": 4096, + "N": 4096, + "K": 384, + "dtype": "INT8 W8A8", + "baseline_us": 4381.289, + "worker_best_us": 214.83811378479004, + "best_variant_source": "worker_3 accepted/hy3_tp4_shared_down_proj_m4096/kernel.hip", + "final_us": 214.83811378479004, + "final_p90_us": 215.11651992797852, + "speedup_x": 20.393443802010253, + "improvement_pct": 1939.3443802010254, + "target_met": true, + "passed": true, + "logical_tops": 59.97493489868936, + "finished_at": "2026-08-22 05:44" + }, + { + "task_id": "hy3-dsh-tp8-m4096-1-0ebb994d", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP8-M4096-1", + "tp": 8, + "operator": "qkv_proj", + "M": 4096, + "N": 1280, + "K": 4096, + "dtype": "INT8 W8A8", + "baseline_us": 19381.383, + "worker_best_us": 418.2191276550293, + "best_variant_source": "worker_0 accepted/hy3_tp8_qkv_proj_m4096/kernel.hip", + "final_us": 318.4573459625244, + "final_p90_us": 326.0519790649414, + "speedup_x": 60.86021643313191, + "improvement_pct": 5986.021643313191, + "target_met": true, + "passed": true, + "logical_tops": 134.86789833717404, + "finished_at": "2026-08-27 07:23" + }, + { + "task_id": "hy3-dsh-tp8-m4096-1-0ebb994d", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP8-M4096-1", + "tp": 8, + "operator": "o_proj", + "M": 4096, + "N": 4096, + "K": 1024, + "dtype": "INT8 W8A8", + "baseline_us": 14151.792, + "worker_best_us": 322.98728942871094, + "best_variant_source": "worker_1 accepted/hy3_tp8_o_proj_m4096/kernel.hip", + "final_us": 282.8514766693115, + "final_p90_us": 283.3587074279785, + "speedup_x": 50.0325901304917, + "improvement_pct": 4903.25901304917, + "target_met": true, + "passed": true, + "logical_tops": 121.47625592271805, + "finished_at": "2026-08-27 07:23" + }, + { + "task_id": "hy3-dsh-tp8-m4096-1-0ebb994d", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP8-M4096-1", + "tp": 8, + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 384, + "K": 4096, + "dtype": "INT8 W8A8", + "baseline_us": 16712.231, + "worker_best_us": 115.95122337341309, + "best_variant_source": "worker_2 accepted/hy3_tp8_shared_gate_up_proj_m4096/kernel.hip", + "final_us": 115.34959316253662, + "final_p90_us": 115.53359985351562, + "speedup_x": 144.8833111743286, + "improvement_pct": 14388.331117432861, + "target_met": true, + "passed": true, + "logical_tops": 111.70305446889755, + "finished_at": "2026-08-27 07:23" + }, + { + "task_id": "hy3-dsh-tp8-m4096-1-0ebb994d", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP8-M4096-1", + "tp": 8, + "operator": "shared_down_proj", + "M": 4096, + "N": 4096, + "K": 192, + "dtype": "INT8 W8A8", + "baseline_us": 10762.616, + "worker_best_us": 171.18735313415527, + "best_variant_source": "worker_3 accepted/hy3_tp8_shared_down_proj_m4096/kernel.hip", + "final_us": 171.81203842163086, + "final_p90_us": 172.1600341796875, + "speedup_x": 62.64180379251588, + "improvement_pct": 6164.180379251588, + "target_met": true, + "passed": true, + "logical_tops": 37.497086951439755, + "finished_at": "2026-08-27 07:23" + }, + { + "task_id": "minimax-dsh-tp8-m16-1-78402260", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimax-dsh-TP8-m16-1", + "tp": 8, + "operator": "qkv_proj", + "M": 16, + "N": 1280, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 110.478, + "worker_best_us": 33.20169925689697, + "best_variant_source": "worker_0 accepted/minimax_tp8_qkv_proj_m16/kernel.hip", + "final_us": 33.19578528404236, + "final_p90_us": 33.328590393066406, + "speedup_x": 3.328073098879459, + "improvement_pct": 232.8073098879459, + "target_met": true, + "passed": true, + "logical_tops": 7.581029876132359, + "finished_at": "2026-08-28 17:17" + }, + { + "task_id": "minimax-dsh-tp8-m16-1-78402260", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimax-dsh-TP8-m16-1", + "tp": 8, + "operator": "qkv_proj_and_indexer_qk", + "M": 16, + "N": 1536, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 111.486, + "worker_best_us": 39.72720980644226, + "best_variant_source": "worker_1 accepted/minimax_tp8_qkv_proj_and_indexer_qk_m16/kernel.hip", + "final_us": 39.73748564720154, + "final_p90_us": 39.8534893989563, + "speedup_x": 2.8055625106680915, + "improvement_pct": 180.55625106680915, + "target_met": true, + "passed": true, + "logical_tops": 7.599622449219237, + "finished_at": "2026-08-28 17:17" + }, + { + "task_id": "minimax-dsh-tp8-m16-1-78402260", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimax-dsh-TP8-m16-1", + "tp": 8, + "operator": "o_proj", + "M": 16, + "N": 6144, + "K": 1024, + "dtype": "INT8 W8A8", + "baseline_us": 38.655, + "worker_best_us": 11.368839740753174, + "best_variant_source": "worker_2 accepted/minimax_tp8_o_proj_m16/kernel.hip", + "final_us": 11.77299976348877, + "final_p90_us": 12.421009540557861, + "speedup_x": 3.2833602969975013, + "improvement_pct": 228.33602969975013, + "target_met": true, + "passed": true, + "logical_tops": 17.10070466699301, + "finished_at": "2026-08-28 17:17" + }, + { + "task_id": "minimax-dsh-tp8-m16-1-78402260", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimax-dsh-TP8-m16-1", + "tp": 8, + "operator": "shared_gate_up_proj", + "M": 16, + "N": 768, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 104.527, + "worker_best_us": 17.114474773406982, + "best_variant_source": "worker_3 accepted/minimax_tp8_shared_gate_up_proj_m16/kernel.hip", + "final_us": 17.31948971748352, + "final_p90_us": 17.501089572906494, + "speedup_x": 6.0352239993816355, + "improvement_pct": 503.5223999381635, + "target_met": true, + "passed": true, + "logical_tops": 8.718209743071991, + "finished_at": "2026-08-28 17:17" + }, + { + "task_id": "minimax-dsh-tp8-m16-1-78402260", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimax-dsh-TP8-m16-1", + "tp": 8, + "operator": "shared_down_proj", + "M": 16, + "N": 6144, + "K": 384, + "dtype": "INT8 W8A8", + "baseline_us": 23.152, + "worker_best_us": 8.95284503698349, + "best_variant_source": "worker_3 accepted/minimax_tp8_shared_down_proj_m16/kernel.hip", + "final_us": 9.341430068016052, + "final_p90_us": 10.175809860229492, + "speedup_x": 2.478421379962978, + "improvement_pct": 147.8421379962978, + "target_met": true, + "passed": true, + "logical_tops": 8.082003660070677, + "finished_at": "2026-08-28 17:17" + }, + { + "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh_TP4_M4096-1", + "tp": 4, + "operator": "qkv_proj", + "M": 4096, + "N": 2304, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 31556.192, + "worker_best_us": 1559.944839477539, + "best_variant_source": "worker_0 accepted/minimax_tp4_qkv_proj_m4096/kernel.hip", + "final_us": 1561.9083404541016, + "final_p90_us": 1577.305908203125, + "speedup_x": 20.20361322279994, + "improvement_pct": 1920.361322279994, + "target_met": true, + "passed": true, + "logical_tops": 74.24514869949742, + "finished_at": "2026-08-23 18:45" + }, + { + "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh_TP4_M4096-1", + "tp": 4, + "operator": "qkv_proj_and_indexer_qk", + "M": 4096, + "N": 2560, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 35918.768, + "worker_best_us": 1081.603660583496, + "best_variant_source": "worker_1 accepted/minimax_tp4_qkv_proj_and_indexer_qk_m4096/kernel.hip", + "final_us": 1089.0154266357422, + "final_p90_us": 1102.6435089111328, + "speedup_x": 32.9827908048673, + "improvement_pct": 3198.27908048673, + "target_met": true, + "passed": true, + "logical_tops": 118.31698222866211, + "finished_at": "2026-08-23 18:45" + }, + { + "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh_TP4_M4096-1", + "tp": 4, + "operator": "o_proj", + "M": 4096, + "N": 6144, + "K": 2048, + "dtype": "INT8 W8A8", + "baseline_us": 29702.473, + "worker_best_us": 765.9097671508789, + "best_variant_source": "worker_2 accepted/minimax_tp4_o_proj_m4096/kernel.hip", + "final_us": 752.4849700927734, + "final_p90_us": 753.6545562744141, + "speedup_x": 39.47251331324, + "improvement_pct": 3847.251331324, + "target_met": true, + "passed": true, + "logical_tops": 136.98508169710212, + "finished_at": "2026-08-23 18:45" + }, + { + "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh_TP4_M4096-1", + "tp": 4, + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 1536, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 19776.488, + "worker_best_us": 986.5248489379883, + "best_variant_source": "worker_3 accepted/minimax_tp4_shared_gate_up_proj_m4096/kernel.hip", + "final_us": 962.3474502563477, + "final_p90_us": 965.7874298095703, + "speedup_x": 20.55025759639306, + "improvement_pct": 1955.025759639306, + "target_met": true, + "passed": true, + "logical_tops": 80.33419874225937, + "finished_at": "2026-08-23 18:45" + }, + { + "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh_TP4_M4096-1", + "tp": 4, + "operator": "shared_down_proj", + "M": 4096, + "N": 6144, + "K": 768, + "dtype": "INT8 W8A8", + "baseline_us": 11759.473, + "worker_best_us": 636.2568092346191, + "best_variant_source": "worker_3 accepted/minimax_tp4_shared_down_proj_m4096/kernel.hip", + "final_us": 634.1561698913574, + "final_p90_us": 634.6337509155273, + "speedup_x": 18.543496946524407, + "improvement_pct": 1754.3496946524406, + "target_met": true, + "passed": true, + "logical_tops": 60.95455267843923, + "finished_at": "2026-08-23 18:45" + }, + { + "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh-TP8-M4096-1", + "tp": 8, + "operator": "qkv_proj", + "M": 4096, + "N": 1280, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 41000.618, + "worker_best_us": 475.5491065979004, + "best_variant_source": "worker_0 accepted/minimax_tp8_qkv_proj_m4096/kernel.hip", + "final_us": 474.8045539855957, + "final_p90_us": 475.41736602783203, + "speedup_x": 86.35262163311907, + "improvement_pct": 8535.262163311907, + "target_met": true, + "passed": true, + "logical_tops": 135.68637642417067, + "finished_at": "2026-08-29 03:28" + }, + { + "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh-TP8-M4096-1", + "tp": 8, + "operator": "qkv_proj_and_indexer_qk", + "M": 4096, + "N": 1536, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 35604.698, + "worker_best_us": 531.2225532531738, + "best_variant_source": "worker_1 accepted/minimax_tp8_qkv_proj_and_indexer_qk_m4096/kernel.hip", + "final_us": 519.466609954834, + "final_p90_us": 520.4890060424805, + "speedup_x": 68.54087888939718, + "improvement_pct": 6754.087888939718, + "target_met": true, + "passed": true, + "logical_tops": 148.82460170966874, + "finished_at": "2026-08-29 03:28" + }, + { + "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh-TP8-M4096-1", + "tp": 8, + "operator": "o_proj", + "M": 4096, + "N": 6144, + "K": 1024, + "dtype": "INT8 W8A8", + "baseline_us": 19897.517, + "worker_best_us": 397.44619369506836, + "best_variant_source": "worker_2 accepted/minimax_tp8_o_proj_m4096/kernel.hip", + "final_us": 394.50483322143555, + "final_p90_us": 394.9032211303711, + "speedup_x": 50.43668752426038, + "improvement_pct": 4943.668752426038, + "target_met": true, + "passed": true, + "logical_tops": 130.64379245024566, + "finished_at": "2026-08-29 03:28" + }, + { + "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh-TP8-M4096-1", + "tp": 8, + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 768, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 31882.257, + "worker_best_us": 319.16600227355957, + "best_variant_source": "worker_3 accepted/minimax_tp8_shared_gate_up_proj_m4096/kernel.hip", + "final_us": 311.9345474243164, + "final_p90_us": 312.3681640625, + "speedup_x": 102.20816278048036, + "improvement_pct": 10120.816278048036, + "target_met": true, + "passed": true, + "logical_tops": 123.91928365478229, + "finished_at": "2026-08-29 03:28" + }, + { + "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", + "model": "MiniMax-M3 (gfx928 INT8 W8A8)", + "kernel_repo": "minimaxM3-dsh-TP8-M4096-1", + "tp": 8, + "operator": "shared_down_proj", + "M": 4096, + "N": 6144, + "K": 384, + "dtype": "INT8 W8A8", + "baseline_us": 22389.162, + "worker_best_us": 312.4720001220703, + "best_variant_source": "worker_3 accepted/minimax_tp8_shared_down_proj_m4096/kernel.hip", + "final_us": 311.00934982299805, + "final_p90_us": 311.1349678039551, + "speedup_x": 71.98871034823276, + "improvement_pct": 7098.871034823276, + "target_met": true, + "passed": true, + "logical_tops": 62.14396076195009, + "finished_at": "2026-08-29 03:28" + }, + { + "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", + "model": "GLM5.2 (gfx928 INT8 W8A8)", + "kernel_repo": "glm5.2-dsh-TP8-M4096-1", + "tp": 8, + "operator": "fused_qkv_a_proj", + "M": 4096, + "N": 2624, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 68661.417, + "worker_best_us": 937.5793075561523, + "best_variant_source": "worker_0 accepted/glm_tp8_fused_qkv_a_proj_m4096/kernel.hip", + "final_us": 938.0712509155273, + "final_p90_us": 939.3240356445312, + "speedup_x": 73.19424503521313, + "improvement_pct": 7219.4245035213125, + "target_met": true, + "passed": true, + "logical_tops": 140.78913965554716, + "finished_at": "2026-08-31 11:23" + }, + { + "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", + "model": "GLM5.2 (gfx928 INT8 W8A8)", + "kernel_repo": "glm5.2-dsh-TP8-M4096-1", + "tp": 8, + "operator": "q_b_proj", + "M": 4096, + "N": 2048, + "K": 2048, + "dtype": "INT8 W8A8", + "baseline_us": 14385.571, + "worker_best_us": 220.13651847839355, + "best_variant_source": "worker_1 accepted/glm_tp8_q_b_proj_m4096/kernel.hip", + "final_us": 215.435791015625, + "final_p90_us": 215.74857711791992, + "speedup_x": 66.77428542482365, + "improvement_pct": 6577.428542482365, + "target_met": true, + "passed": true, + "logical_tops": 159.48946182998893, + "finished_at": "2026-08-31 11:23" + }, + { + "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", + "model": "GLM5.2 (gfx928 INT8 W8A8)", + "kernel_repo": "glm5.2-dsh-TP8-M4096-1", + "tp": 8, + "operator": "kv_b_proj", + "M": 4096, + "N": 3584, + "K": 512, + "dtype": "INT8 W8A8", + "baseline_us": 6939.939, + "worker_best_us": 137.40396976470947, + "best_variant_source": "worker_1 accepted/glm_tp8_kv_b_proj_m4096/kernel.hip", + "final_us": 134.75635528564453, + "final_p90_us": 134.93956565856934, + "speedup_x": 51.499901324055074, + "improvement_pct": 5049.990132405507, + "target_met": true, + "passed": true, + "logical_tops": 111.55233090221004, + "finished_at": "2026-08-31 11:23" + }, + { + "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", + "model": "GLM5.2 (gfx928 INT8 W8A8)", + "kernel_repo": "glm5.2-dsh-TP8-M4096-1", + "tp": 8, + "operator": "o_proj", + "M": 4096, + "N": 6144, + "K": 2048, + "dtype": "INT8 W8A8", + "baseline_us": 54775.797, + "worker_best_us": 759.9214935302734, + "best_variant_source": "worker_2 accepted/glm_tp8_o_proj_m4096/kernel.hip", + "final_us": 755.8016204833984, + "final_p90_us": 756.3575744628906, + "speedup_x": 72.47377554571302, + "improvement_pct": 7147.377554571302, + "target_met": true, + "passed": true, + "logical_tops": 136.3839562001365, + "finished_at": "2026-08-31 11:23" + }, + { + "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", + "model": "GLM5.2 (gfx928 INT8 W8A8)", + "kernel_repo": "glm5.2-dsh-TP8-M4096-1", + "tp": 8, + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 512, + "K": 6144, + "dtype": "INT8 W8A8", + "baseline_us": 21254.589, + "worker_best_us": 273.8345146179199, + "best_variant_source": "worker_3 accepted/glm_tp8_shared_gate_up_proj_m4096/kernel.hip", + "final_us": 272.7110004425049, + "final_p90_us": 272.99739837646484, + "speedup_x": 77.93814318275388, + "improvement_pct": 7693.814318275387, + "target_met": true, + "passed": true, + "logical_tops": 94.49491855548743, + "finished_at": "2026-08-31 11:23" + }, + { + "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", + "model": "GLM5.2 (gfx928 INT8 W8A8)", + "kernel_repo": "glm5.2-dsh-TP8-M4096-1", + "tp": 8, + "operator": "shared_down_proj", + "M": 4096, + "N": 6144, + "K": 256, + "dtype": "INT8 W8A8", + "baseline_us": 17700.262, + "worker_best_us": 278.44139099121094, + "best_variant_source": "worker_3 accepted/glm_tp8_shared_down_proj_m4096/kernel.hip", + "final_us": 273.6263179779053, + "final_p90_us": 274.4647216796875, + "speedup_x": 64.6877176537867, + "improvement_pct": 6368.77176537867, + "target_met": true, + "passed": true, + "logical_tops": 47.089410050974806, + "finished_at": "2026-08-31 11:23" + } +] \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.xlsx b/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.xlsx new file mode 100644 index 0000000000000000000000000000000000000000..f97cf0db2c976a0ac6e277f77a96a236065ee91a GIT binary patch literal 16612 zcmZ{L19T+{&fB=xIqR0!XOyA)lUenjFHzzOq2X17n#x`Z@R*aW{^o`dk&dUyjpFI;h@vo{|T z76P5ja6G70T@~xpn z?gV`Oh`FU@s2Ie#AdqPRA$s8zmq~k!Ir!+b?j-2#pD;q96p-vW0#b-4bzqBA$nPGQ zoXvoy7zXs$ZT^-72PZ$5*8yE6WEJR29k9o0Ct4bWVD04EQrm5FZHNs>iqsUGIK8y% zPxqEwZ=1L3I!n`46nF~itF?Uft{Z)S>i9iv1ji`7FLfPt0sQCe{Ls!M zLI48*$Rq;*K>nH?cWXLla}yhre;pbA%+I-|meVRbitlChw7c1Z4fG&uKw|^4lZESA z!;j!|S-cQL^%!dp9@*GUnv!y|T&+Y4I5qtNiA8+ZK` zy0yi&Cc4dSHAZ4Y0s3#n)mPn6h0?meG%W!7*nLKjk^0sHqadW*QN|Xvx3w_oh!#&8 z);3gwEv$1f#EkWQbg9kxr)_VJjWP{)*GO!5Jd)dWae!W(X$tq0y#H+fB_ zjyZ7{18nKB)baf$?TYtgHN$@^_x5*Nn(?Vx4b9Ao`v;ErB!RX{Jmk-zLNLN|ab~Gc zx_UqKpZ|;=U3F%*!Q9@xc4u!Kn~k!^2H93<%kIs(v+gPofNaROm^NO1SgndLDvME) zn|t)xI?XW8go@FUo3nD;nEM_vSUK$u+j%V))>joVSU&B1YC%f&US(BqjAkw#J= z|BSS}ceOxcqEFaEEt0X^*a=_yA zwv4I7H>ab$pT&twGJ=!#?bMLhuS8k~VPv*k$c~)YW`f>R4-fH$B#cv+&rqkb+-7*H zDK$D%Q+wIbjzn!659&S-a$>ebOQqfQgiA@H#3Db3XO8`!sbz3~$ARKbG62^Qh&jORNSzA&6q`WI7!<%~-#J7mCh2sI;%= zr>dg@rebU{m$6U}EV?ImUsq;w;~*nVBte4-gr*XcB9(IZ%7`O>P!+CzWX|zlCD9ST zNXExn)DsmulTI8lVfKhtW4VNx&vJWyfIerkKw-yY;Q*Oo)r%4YTeMx_Rug|uuk1yn z667`f#)|3q3oW40p)^=u6KQb9U3Rts@lB(GC6QNu_-rEHWpXhd4lzNO?d!jjQE`Y} zxS{Nl>F0S4>LBmgx_uHoM}cFe@6EoU3ryD3-yg;1(IxlbNsIZ88JWTJ(={MFf>kYl zmW46%mf$i%fnK-Z^6q=OL$Rd{w#!!qdI5gUtg`i%KLvJk_|f`I6<)U7@#M}nu@U*0 z^~^|bNuW+dpGCJzCM3hAr7ui*PtOni<1hd z7a(qaiyFP6?1wxu)*^u%Q=T(ni$uy6#IikBx|d8|+OKXQZ|2t8JFd=k=B1u`=GHrG zn1E&yYI2l1X-w+;ManLDS37Pmi#rCxWyYO=)999i!}9u|#WKlcr{KxI8oSdf1~-lYhybf-@9;lSgoo8=5xN zClu!4rI%)wJSfkMVq;xFOfA}^-@J59jk0a5Pm=>4RILDXqo|_Kpfw0g6_@8N9W8QX zf0(MRw>y)tyH;fK%eI=r?VAmuB1f-98nRjW=%z_+j|$F|iq=g{rzKAVs_fKEo~g7o+b=t zMK@bCFp$(*;TlLgGCJoAhZ&hwo`kM1sct>`H}Cc%V(|N1hkAE9LRG#G8@n&IiCp5+ znH~d(fX*U7M9J7-Xy}Z4@)rqlV90c50EflDemAo8&BGf^l zvbwV*NGOptFrWnJj=PtaiI#5BOE1&`Ho?k_#tZ19LJ>VYOpZd!$|Sc{tEA8&Ydgvl zZJ3>;`ci$Kwplb}pKH4FJPqhPqPO`!wsGsjoH9I{dqpZf)faJ}8Mip>*ElzUbq|oK zg$jIod?F+L5z7rbTdU=sNiw+_ysiWAGu{^03LZALVZ7(dH4BZY5c#CQz3~#jjH0#f znDt^&3@eLjD+gB@U_AWQZ((vx*|)+((lMD9qAECIpzH(|`1kqN=3&hrV41PwL3J5^ zG_vOJ$wG0je_X=W^!9)MeIC^X_@53WCgF5b0~G+^Djfg-<6jQM&Cb!v$=t-m*@^C- z(?4FMQpd^(t101~irVH}^&tqqJ)qWFZG-qxQ&Yh^czr?))h348mRKRzoNGm|3t%8@ zegjxPR8j?LjMF?5kR3cGn`!$MFKQo8uF0k2wp*_LmZihfH|dblFnEt`U8kYX7*t+nMTcr}&8;#kT{0xPs~c# z$k(U-rLeON8gO61?!koRQ#TgqtZ(LU!17mfY!^JUTD{6sdB*X!_cdl3s1t*RKHEe&m~X&K7OBUo`o*z3*Vo5)wX=t&-t4!K1{E!O)oar zyA_)-DG}G5)EC){pXcxHgP{z^KF<#*%hs6SV2`(9ho4g)t1vfXVC}g9>&$UE0oEr1 zHkjno-K^X9VATfTGO*)C0f4aHz?-&LjR=>l?rPzWSnc5M1X!5L8DOZ^Iq9S^c4jes zq(!A8wt2s$Pm6w_t(BtLSBcCiho2imv#-UV(>31GYzMQktxPuiE74+Bn{XWiO?o;d zp0`7@Ur)$1D$z}^vu*mVji8%9Wbo=aa=kk8ypm0^y=qeUW>uPgJsosD&oLHW%HEH` zbqisH4m!lz#ay$AS9@l^a7S(pnm)Tjg6)%CzqgD%Mdjk(pMn@Jo*$>Li7H3*9}jzD zq=Qo{7_!SQFr^{v?~q{ca)bwT$7z1^V_`2BmKuDt;e%a8SaM>*8GZ+}f{^8lnaPe1 zN#R%3&cRZ7BgM{%fHQm48+(D~j@emwN*Z`{6XyAHyM>;ouD2q>Z`I=fbcf{wLe_Q# z(bolww3BUDJ9iQ37g5|IednAgYVmn6i3(XCxQjQQEg#3vK+W^KD++vTGv)GIW@FE1 z2B@E!e15?lJL#ASW3$}`^y4p*ltxQxhLk;WSe&vtqd7XU`H{tWu=O*&9j)2Yr&P!u zoLHhe#g}sS*+lbkZ1*MlDN8=$xq-ZIcbCd>4iBQgsx;vXepyzsZaO)Y4eNL97fPuS z>SDqp;{{n_(S8qS9+_p7&6q>ytox4Qt118^=a&vAh8+33A&XzkXLhuTPc4vdiz2X^ zY>4U-@HVkV_6!ZWMGJYaRcd8=|9tC&eeZs#T+GLt?yQy@4ljak@rZ7IHSW91L(|Fo z-Tgte{)t!DtDn0g5BV-Ljt#p)d9M@1FF3?Uf-Sz2N3t!wD!^mekGI{|g5f>T~m(FkK9U5+8)Lxl_UsfCi##0_j^$3_Jx8(KSPOP z@$`+4johI|3BF^p`+}_wdMVkW>I6R0E&s0@jHLFePV)1(DurKvPh6naktg% z2)ARd2l6hUrI*Gq4QAxWn3ZPj#jl~KmNW?U1!2g{oHHk-x0 zF`06Izp%7Y*7@?VOpelwMT~;6LmlVO6i73PI~$zHnLG1xjFPAdl+{g7kam5Q@M1xV zl@+H>7A`=80Gp5ucLORMG{^y+P(U!KQA7zY>{Kx0X`?)lnFzqXn3WhYNkf*FX;!tP zpI3AoVVR3=T8w@G>lE=dijF^{fKHqsna#51U0sU8gO4!`c;z<~nC{j91*3pp&Msad z9q7#wQ30#>qPfwqZx-1zD zQ|lT=EDZ?@W^LVii1yDcC||J2OesN^)^L^9teAMyAYUVYP-EWDNn-nTwsb0FD!Cm` z9f$7+#+$gJ_V1i9GPbQU)s9O5Bjx#~OG`u$>KcgBNv^;^DV%_5Gm5RFwRGikFfh`TLl zH{uf-4Uhw786YS$i_b|@hW25-je7M^!b+uM9Jd&D21dq{CRXd99ElJl1BO5wmdHm2 z1Sm2U>ub2}g};PAaj`4(yo8))8#&u?5VVu=LU1SWmt?|@Hs4`6g;gnO7jSL&YazEux#4Wk*i73+PPCOR3jsuCfD)W( z%)+x*vm|_*>s!PI8jtKX77gN1`^Q^L*&U0;v{V=sdc2_STs&W=_kqajExYDQ&+6xt zv{u`r)@a3BZ+xh3d;7e1g%Sg110_OT-&gKp)6=nB1b)FgvwJ z{HKy`F)VPrQh9NUZ(*P@1b_{h9%BakM!??p>V-E1#FX^0Y}2$H;*UkhfjRO5nl#cy zML$MXA;V7{9%ajt{M&2pBYO|2LVyzBmIV1n-S7&NfXEjzM->)iK)zU$}vFLYEyk^zs2!e}^fLB&}>I-%|{ zgsVDrjZsKZ7Y$_MU!ZGT@r?y-k|{D!+$AWX%y6%rGPc(@`g0y&X|=>(m5ZOl9J2i=5)(ocdmn6l7^B2xrg2NC_8~j~YLDwDl(N zs!qFleXbAj3i41eR4AeYjvv1TkA@0}ppN#Qa4>R7HyJ;bj35HX#FF z)&~9Uw4eU%32p@x6y~Rbam$6-^7hM!6v^?!2hM~{NC*5pDv~*S#ZNfM920GVjCTC! z!q9ShU-^q6E%V~>J%Cdxkr^hk9CU~DAI+>gb!=Uc4aikN%>}@X=&6o@77jhxio~=c zq|448Xpe!(phhLwp!qdEG=1!q)WK}Z41jjI&_IV8COJ%PKYuG0P}yuQ5V?!M|Cp9n zqL;NN#ja?BDFuuOAxWx0pjA3g6FZ(#cYf@4{fC4HJN91v7>;$ zfHxa~x&~1iM4KEz&Vd5L4o9kBcGAsH(Mv8^YIUxNCDa9~V5i@L5zmLbg$hZH#Q953 z5{JEJ#^cM>&B$5(E(P7CX;j*A^p{E5Ibu#+y|R3SM@bhPt2lxDq%Eq75hQv{(fH9g zmDI!k3Z(T?C9hm(l!8Q2AgUyB)|g>6qTD-Q`AMrGA<7X;lPEMU+ z6Mw^n?%EegfCwOxYV{O*&yB-0B4&3is1}{O;e?RkG;~~{#8ac-_P#=e6qmoIM9Lh& zh$3c?0g1HxWag$Xe(O!w%qy#x&BIsoFc zd_n;$)%~NzDxcL$t4dnyma|OeYbpGqAk-*feKyi#+DJ|cEY*c5?IEZ9*+bs&lVi$W zmJ0T_F6yMRGN%xRdfzIBtCzUNDyIu%p~c}MNs9$340nm+w10zSPaQoAGjH~8C?nPl z!lHXZhG-a(b!`@71~5z;irAP&2vVzs>CLM6#AXdu)LYHt6OaECcQ$A~3I&Uu^(qaK zhsK?v|2UQ#(*{XwHXX}PqOYP0#?>MOyac~E4$lk3!9>ibG!F!I(}TNCIw2trM?U2lx!%)*Ru8T zRI&w6EU4q6d8oU$;OB@!lID*D_!_TH(Vv>%(9947mP16O{}K;_FkU9Ga2K4216WPl zvERDGHzS7RtBe_ed%Q{us^#v!_2OTH98fTyYVTi>%{^`6kAA4q(=aZxp%S$_d!r%J zAQ$y`v@2>8_Gmt$*8&o;DUN*W`c`IfP5@JvT#e0iRL%lX0>xkXJG<~Ff${+$tpY}5 zyeXKV_8_VLP(w2H!W;(=enJy(4WI~ElXN(tJ_1K#$lJKGV)oRqkuwt_-#r~{ zVj*Rs*^qI3zdkr^LYlK4+g5h(A)hLe=>(F}-ZxVD+1~T%FZbt|4$R$$)MixKq}o49 zF^>ULsl2Si;ZF3(6l@>BKK@%%fBPu~5@ukG>0;U~CVn`Qg8f1WsF}AJj~5$PVlM2RhY?Y$R0^<-(M93e z$d$?G-Q&LHxlJYty9Wfy>uNGm{%gy4Y{yRGU29Nn8>rr$e-2^Jt*S@vzUaay1fk~F zR91(xM97l+1ku+ZXg0`f&ZCg}b6G+-nFnlkS)5`&7&Nn;q5ck&WlFFTWBDUN-kaG8j5h^`8NZK&kc41z9kvB? z^OC$*)~Xq9SG8xz(#XYs|Ixn@Insz5Y#1Y=Q?O$SB+g42+TWu%@; zIq~hd?MEk)UW@MDIf=Y5PYAz)#R%Z0ZpcQtC>Q#owK-ZFe$Gs~nuNQ)>(>wYjR*XRkUh({MjFX>}CL$V6aU{ zt6-t7XnCltHY|yV%p$Ipif8NOM()L~PTZ+HiH9GBC!jvluksu5J{usc?|!?|4LhNLc5yc+pi$zCN~^%Y3O{`KA7RdV3JtoZ3h| zejd7{xg>Y18Hi*r5!Mt_=A58VX3r`1;Y`d|>QAWfRyyBE^W~eMMy|ype)wKtTJlxy zrN3}!+qmTu4l&5ArKt~gI_ZxdTaEGxvq=G|J9U?)+kaL-sCCpUR3Vdz?PgMcCilkT z{mIgb<-_&+*_!2b_47uV?nBDG47c=w&e_i1?(5LX(u-r`8>gT9$;JX(E*oxRtystY z+v)j5*uoA2#IC~St>G`t_3BS|t*NS06OiDoEo!v#b9jOSi8tp;ug|B?{|v`!T1U)W zfdT+B2>AHNo2@y7&^>6tvK)mMz?1d`fGealR# zS-mZdUHcjshEm!1K%0XfD=R(+KYE7XMV@Uu4tP%)?p~z3>)ANcs$f(tJU)^Voj_MjY zy}80q`8~b5jFXE3Bd(E{pGA*fB;tTdqnhPB^*{vq__k*#PscOkCa)i4uw7Y!} zKp*P~A$pn<6KDhHCuP|+qW%So*A#J6RV~5-bJxJW;_1awdpqXjd#4akb}RUt9+eK& z!iiJhOgvnFtq(p%jCS9sbpJppRM(#xt_{N6@OsD3H&Y{^JEnIppp+_?0A@hT@28DZ zxJeVOu1{ZI%J606y|+1a#O(H#EM89A#NT+Stj~s98V$yp-Wvx@R|_GND;=537wU?| zZ7}88G-A26zL#sen~xKk8b7Qxk8Gmm1|!y)8T!ph3+82^#BkO^GYJp|!ks4-z=~y) z^GvZEzw<{E$f}cq)~}NrPKViPscLc7d+>f>Pcje!=U+okuBT@wUWiKv5!I4=*`@Q=3r=I_^6f1Dti|i zfv-w#mXp+cU_Z$~h5$hZ1t)D;RPrll7?A`an6&pB1i)co02*x(2#^8PcQps9%%W&0 zbdhg{2$I<5IYpZGWjewOh$#`MAacG@NRBd)&yEpzvJq%veB{@8^24sVaa@9kiGqYA zXAsQr7Ns@XOhO{UP#}5axj>kV@bD2PIRLGu=?jSq znMpy>LKq1vibUC;O7o5J8n-DJSg!ngPHV3sT;VuTA+H#IE)E8$NCkPw`%TV*T4N%g2qmHF(py+UqlVaNqDU~9CzOxg~{B?+^oRw<}Ja9N*Tv7$P9h)5o8N$kwRH_u}f+q z79GCxGn>Z**20q_r!hB#qp^19+4sgk5C?}6{))yI0@)G2njmJOr^YBXJnqgTm%6%Z zdNNpiqo-4QD}SEZss26mlyhO|%0he)zA_V+Mr7I)4%6UTMjQmJTvJQCy2Jzmv>^gj zNfRCz3`S@89Unm+vzHckypxWX4GK<)m++rM>pBDD>jGg`WWgVP|*Jzss#3&S`=;%mF_LwUt{tJ7M0D7zg z#bf=}upWg`^8_9-{N8{E`vH9jL~c?DLwHZ2%>CK!J7}YW@8S9g5^1c{^!)DId@)6t zh$(GIkvT2Z5VLFU-;y@f1~gLCI-_+HLQ~6j8`>uh#)u!k;@J@ainMBU$ZBjVwIzDo zm7AuEscyU-*CzD-m&GCqJX^tT3a3p2@1K8FeZ)Y1+aCe}0EqZ1`S@Rwc1(Y~#J>8P z?V2dUJ8B@{jdHwK7jc6BAbJCT3{>qJFI2;?1a+kX!77@zFDH>;x>rcNI3W~pAku!8 z`6H{|%v25gGV-QV6U!*5SSCd+(cb3#K?qyT&8l9atYkSITQo(aMS_|;rM!7Hg;Zlc zM+OZ!Sh!tcmOJzNsaw4xCoQgZ%7LjW*g~IaV$#%4Y1K1Q0ysTgZ&@{~Q{Dc#YgUaP zRz`knlalU{Ty#;+MmLl+H#y8>5>G@Q_N7^04#LaQbz#$Io3Vu^((Cc+)(|Og+GA4Y z%DZD?mQz(f>YUEhVoEwCfDZ|5BufTw0%n}_w~c{Kt-~(8REE0yFY2f+Y26{_AAnkE zV~u9(>=s$AC|wP=ux-xk1rL{croI}tP~01&$sAO5e6t}sa)PK}(Dl-|jIS-r@5Y){ z&9`o9#wzbvXP8>Oz1tVLfh!IAX+5M98CEy*j_u&xM-cuE06V}(hyXT#cRzvDq4jzF zu>r5px_SJ$Q-aI^^A0)*%1vLmYe$^5(b{sH7q=%i=EoNL!9(ZPOCMHK=re()Ww*NW z)#cx>-?;=!QSgRpL7Qx;JFRrQF`WhCRC`;dftJ)pAl%fdLOAzstd~jp>atODqg$y4 zGrN?tEe&*nM*F z3&4VIG$XJGW+93u5{Fch)+3tSif%q4q8Ex+##EwVKa@^b3LguYS6c}cGX|;jb7oS? zrwVPr_>FO~`|<(O(U70Y%J3AA(S3TZpo>+(*<)pt`$*0HJ#7)wKoVmITOxLejfnOI zBHE#!!x^Fgkzg~E7M$?|OjToulB>gUKtWPNu1vSptb&YalyD$vPn%y(UYsFzmRB$# zT!Ip?x49q1qqpsdh&iGL$b~2`B4wJNM-dzem+(GJ|6O{`fOS=$fslT9hw9Dp2*VMs z@%F{x0;aN*+V(02(prSVS||%0J?B243r;T!=63YRs(VH^2YGfoi}}go*@r{H&bwyQ zy&5ezocPQxQphZ@QpjEo+fG#b4l(UCJb)8KRoJSqf$V!gQdjPduvj@e4S$mWcD~A; zrL22Hq!5c3d9kE~sCiqVFZkx%e^JVKxXQ~@I_R+rc*|LEs! zWWvVmSU{v@s&DpQlt((34 z6ItSyhHS0)3M8uo__ch_nGVYbm%P_78C(VeNMv!RB7~b)PrQI{*j-7d$LfrdN!5wG zf}`RFx!ZndLiTs)9~zeUua|0M*PE%pkMSpz6lxZ_pg%N#l&73(?EZ6$Oz~EUPx?iEU`YQf`7!??zvMJ4yL@(( z(9g&L$^0MiDBeBn7DyJWwp_`i<>sp^&+760#3k!!nMGN8U5&ZSCYkBp5iOawH>*<| z2x(_gC$*c7wB1fA2lw+UPzv99s&#EV&=$-@Mw~oIY0IdQ9Gh0%T{N!ED--MHg*xRr zs3jfESZ9XN7OWM|gwyxD6#_g{O%(1Od*%%JTykzbx5esKpj^5$WJb;VE}Ne;XenC` zT_{h9cRjLDI7-~PT{?TOVU?~qx3_{XPlIf7OzWs4#eO-pWpj4sCpp3CUW_u&ECJ<_ zcipOM-!Lb84g=0e63clL^c79kA><9udmJkY3*ttW@32kXvv`b}RNN_kAg+fusr_m%qZ!>&1SXpkC+HT;ZevO*>AP>ElBK|jk- zm&UGHU6_^0pvEzl*o>ZWY8woL{D=pcC(Y3X;&cUFC*9Ts;z5#Ib)St2I}#V1CIxv1 zaU~9DglymoF;atA3zo09L{m#_i4_kKmP;oFL$O%j&0b&$BbxNX)X64w5QGhBn#{6bTI>V~hADq9M#nL@G=*3I zln!#=e8Q09$W z5N7?uX{wvCUFua>6J(@V-rJj z#Mg&Fa3JiX6W^Iu##*o98GiyE*M~T(s6;i_a3Jx)p5gvULI+B&*;*nz^Z~54Cb+w2 zM-I&RUlG|Nw(K1Nozxv-XB}n`$G7$F$4?>M@UeHSgUS4qJ@4tD#Fe&rgjq35hj{`V zK)hH)CW2q$Frp_o?V)9K#J44@#2&DUxKak6)JY;rEP7z+I0{KG9gWY-xol0`nzxdm z!?z9!;Y&F8b!$E^Ki=X0lUq)AtggYp0RVWuDx&_EjTOruZi!L1amrx-di}=SYDkox z%W)I7=T78wvpAgWkj~KNGYzL9)3&(^$V~O~qk@DOB)*Gcz^`nW;(a}fyf)H(qrUvb zdSuZtIyyJg!H|)NPk#95nUys3ENXUT`UCy6zuU$EJt;$w44TcN1B!f*QJu_!H9dNx ze$0?@vdJ{r^ren}-?c+F@=(q*nJ@eMaEb2~O-CEs!&gDBT>{s;PR8kk380&Rkx6`& zi)xWaC{zbGPF8#(g#+?>xS!&?dgBMt>8MTiX9JraS}~EeTY;LVZAAnKExc!$8kH5~ z!?m7?^L&K`S>xx@;-~#&5L1&xY4P)kSS$TwF%gR(bcb{!)h&(RTLo2E@GrJ0wZd1C zSYPUWBId0WQa!aq1MpMk+oioypgb0V9W~NvOQ3UU2)&H%GRSd`6SsqPPC04DZ|KiN zlxOBOrl0KA+q~EW_GA;**j*Vwy2bP}(wm=y!FiOvpIFJkqwd%k^j2G-x1;^V0?5ja zf;}L+hy#nt7zx9=EyjeRQO^$t#k7natrCr-T_bE@JyO!|_<|roR1ta88_{8&;17`T7mxt_>xP7#t#@k6gz@7YCe%p_Ybykv0fj40pr^FhM+6kJ z#P!#uS(|-{Qe8;Ec3FwzyPa)wH@qw6d%cuvSICkC+HGkd-K(Lk4Y{;9KoP$qTWgSU z!F4WQo4NJpBEQp`I|nxpUC0GoKdwxyprhop?JEhRaz4FzgF`v?5Z<92pHP74A23ZX-qtAbmu%8sekc@p`HQ(|ce=Qxmbg8Ap+UR&hVPo8BvO8q3Zj7k{M6A@S_-4g;KXj&ND#T5gGxs zBsnV=6~*~5T1c=u!Hkc@QgLnuS-BuVGnV7pqpxHG2su}NLJ>W6Nu2L!oEdbq%Ipi{lJNBx>q)f+*BpvHMwVDqIR#QE*_)B3~*n-P1vB2BxBsGFK; z*`QwVxDxla@=(cAaDC} zxzI&2cy<{3O$eaj7A*`0*TRGAxtHJzNKKnvWXPJD48Ou3BRe4dW2&41%$36}3-XZuJ&=kr6gVhr}>$Ck%RNVV8Qq4O)GCjB4HH=xpl~faLMA_ zd7-|g%cU?9NDJ-NRDFW%x>b^v&lxNt0lt%!bnp-qoSwgs>exlT3G{i4ty?=Ygq>1w zJ%mLP6?f3IF6OexVCew1?4xqmT4fg2QOVNeX>$$@y5-EB+60$20X z^Sy-4ScSx@xRx2Oa)?;P441Vr!lZfKBx?b^% zxdrN2q}eB#}U)}sssZ0o?Lnd?geTn1#`f3uGIHCB}TRft58~J z)};;+c1#$k9=b3WZfUoZ44QF^%bAksiFpINV&rUU5LI>^{?8M%3 z*4x|z0aYzDqy{*{ZZAffC%L_VN+xGyTi8=rsPuV@@yc}qcNn%B9nMyQzC~@FT)Gj> z;Jq$}rT?}As!>|XS=8e%SOOaczA8wC24)2$dPWqhtiyk^$cmH<1#u_^kRQ8RkU#Td z7N}&!v?sLN71kZ1F!x!j!$UK}$F)$+OV;gG_r?~KQhKs;`54^8daONTWwXz7^1F1Y zn-0MNxJ>o1J;8bft$a4*e8&4v-$^Ml=VAp60Py?Ecf$UwlGw@F!`j5@kHg$h|82WU zkKmg;3_lHU%*Y6+n2KsK%c^4OcUDaRr|!5j&VX}wd{{4pBv~b?rP>-GdB)}Ewv{zi zyCG5P7PwLMZElkl*@J_`sxnHhQ;&9uI)6WeTHThpRQcONQMsz`SE?`(iT}rySdhSx zS<;~~>j>Kl_grC9QiitNNDZ3D#Wf27Sl@Z$ zHvJW$vi-+Y*5NGJyV^h48%;H*)i66c)8P-B6<+C&)y0$PRmhBmAM#d&Ym~8?cgk6~ ze^GnOThcJiiL>jBWUgADBia<0MR`izt)lJsVEfkUjlqVh@_G`UsY?};bw{t?UPRaj zX$7djCRURm@1RA$NyDoyv)WWv0ajGD+onsa$(RkgbStc-*Ke?5bqrUjU3r!A*-1mR zoPU00qis>7o);gQEMx%cd95pAR~KV4(L$DWAUj8y4gxRF?OAiZGMiE&`7rEiq*Fz^ zQLZn3%zCdR@JBXUR1<^v{Ul8~X(O{x6Jsl$;o4I^7rKmISTVhPCSiF(Kly+a=WV-~ z8iOo*Ka)Es+-JRK34>#_fXi1^x< zN^`xoa&wDHx6XSrDN_%JT&s97h{=lPep{7|gm0Ud9Mf&nurK7mV?pbd*x-3)W+83K z+2EoY{9)oLxX_vCPNy<6q3At}#5+;sz648LfThgBFUW0yzJe}b5Pb@{(L_SqW#aey z_UbM#fO~}n?e`%}m|f=$%xME(dN7Q^TzbAwC2k{T+~U6G7%mng6edvyxP`)qbM)J} zn27OWu32)ZK>B`0%%Ln8Luxw%2)~U*dRzbQH~-?peaxRW?Sk1=z}B)q%Mhp)+mUL; zueF1Uxa|DH^Db5f-vEFzKr39k{hIJfN^PYPx+@IVYfz)D@*w8;XK{%7iH5$k>plTlcG7@N{NLE+%YN10M z^fwP4ZlrZw3;uKVg2t#)jCTY@3(bXQK=DSBSc}Pc!K!&`T$Zp|e&IDEp&CtUUVhqd zrcE8jXvIziYT?>TS%o}nbaDcA{?(k)4uTfqQcQ4H6SIDdkGC)^BKwj=$y%81bNe5z z+Y&Rmgz9MNA|rEoIE8*VLl1Ft$JIuJ4W3#@&kp9wP~X>$NTMAK8$MV9N8rzV|B8v) z{HP5C{Zi8JD|-h2FZ%h@$Y5w^XZ44GekIA+&G93AMFTW6P9iveq6#|@3xlE%gs+w? zS=z*&DlE0lT#!=m+^b$I%tK4U`_BU-rGmw63pQV?VfIbw+Zfqnp%4^gE;1Q)Tr%aa z*u(~}-v%$&KNt-rOA%JS(hMp1f-tC*=fSTl3lbQW$&kC@ETr-SqXHhKar}}rnTVip zoV7sL)HzN9&h1$uZ{pD(*g!^t z2U$w8c~+!9NU+ph1}mNUDP>m}$8%CjonP7{0UIv0NGVD$TpW zTJJf>%?Lbbo3T#x!9Shj&KKEKeT)2`4l9I1JZtD{rEq@*lQ93aV*X_v|8w2^V<`hk zGH`+P2%^=3o&chN3=G|)uv)vI>74Stav75t*AIErIT6rop(?jBbB`Sp%x(H)EeycN zgNpPI8~iQP;}0PGl16kC%xGc9MT}Lc+#+g0$z{nbf(>hh+lw#?##W~GKALYK=4E9H z?yM)9yWM)-SM*gRF~fw-l4?YUK!dj5kiEbAu*ZwYZ#w=SUDNGe6pSZBv7tlp61;v8 zj37(;>iWD2hVw%q^mP@Fn#k;Uxr@E$?ieV1M?TE$5*=y+cK<1I4fLPt%Wr&+Q2_z~ zTz(nMpVm<=Av;@V6I*9JWeP`>INJvsWyXX7%j;gY2Fs$XcpsQkTC*NgTA# ztgnL-upubaTshP!Vwq&a2Q&n+o^k9;yw(ID%KXh~fIB4nH1cMEEUE$L3?i6tbu`Mb zp;GN&QwxrOhk>X77&*^H%=(OX;Dow1qYvB|h4?46!4Kof(6$}d9f)mQb~GS>N`I4O zhg0*s_)N>T`R~M9+g=t!yRgBdY_N6?ZJBw`c1;f%67t%T?d;kT@?Lz-423z}|LV$d zm5Wh+7)IaV>w5W`TRk<1Kg#WmzoW+gMwuu7uPFb-kpD*cJ3jPp6ndKfitz2cgp_{ z1<;!LAC!OZoBtO5+nxVg6o~ylocrH`f4j$j3!Z)nel3arNqT<^|IL*D7RKTFNBDo? m&A)N}rksD{U~~Uh0+N>k{c3Li^F9F!p!)T0;m-5t>i+=NuXZy4 literal 0 HcmV?d00001 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py b/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py index 6f3b845e..10404739 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py @@ -44,6 +44,7 @@ derive_variant_meta, list_variant_index, ) +from ..orchestrator.w8a8_baselines import fixed_triton_graph_baseline PLUGIN_TYPE = "dcu-kernel-auto-opt" @@ -810,19 +811,34 @@ async def add_variant_route(task_id: str, request: Request) -> Dict[str, Any]: 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 {} + meta = derive_variant_meta(answers, shape_id) + # Baseline (for the speedup multiple) comes from the task's fixed + # table first (final_report); when the task stopped before REPORT, + # fall back to the shared fixed Triton baseline table so the variant + # still records a meaningful baseline_us/speedup. baseline = None + initial = _load( + workspace_dir / "final_report.json", {} + ).get("initial_metrics") or {} b = initial.get(shape_id) if isinstance(b, dict): baseline = b.get("median_us") elif b is not None: baseline = b + if baseline is None: + shape_params = dict(manifest.get("shape") or {}) + if meta.get("tp") is not None: + shape_params["tp_size"] = meta["tp"] + try: + baseline = fixed_triton_graph_baseline( + shape_id, shape_params + ).get("median_us") + except ValueError: + baseline = None 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: diff --git a/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/SKILL.md b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/SKILL.md new file mode 100644 index 00000000..70559901 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/SKILL.md @@ -0,0 +1,376 @@ +--- +name: int8-w8a8-gemm-decode +description: > + INT8 W8A8 GEMM decode kernels for M<=32 on Hygon K500SM_AI/gfx928: split-K + partial kernels with fused atomic combine, small-M zero-pad path, and a + µs-scale acceptance protocol with noise tolerance. Load together with + int8-w8a8-gemm-foundations (contract, layout, DUMMA rules, Graph safety, + benchmarking). For M>32 use int8-w8a8-gemm-prefill. +--- +# INT8 W8A8 GEMM — Decode (M <= 32) — gfx928 / SGLang TP4 + +Worker-29 lineage, DTK 26.04; K500SM_AI/gfx928, 120 CUs, wavefront 64, +64 KiB LDS/CU, 65,536 VGPRs/CU. Validated evidence covers M<=16; M in +(16, 32] is a boundary gap — reuse the M=16 recipes below and measure (see +§3). This skill is the decode half of the phase split; load +`int8-w8a8-gemm-foundations` for the operator/layout contract, DUMMA and +epilogue rules, Graph-safe interface, SGLang integration, and shared +guardrails. Baselines are fixed user-supplied Triton CUDA-graph numbers that +were never re-measured in this lineage — do not treat any baseline as optimal; +"speedup" is always vs that fixed table. + +For the four Hy3 TP4 M=16 shapes (2026-08-23, hy3-dsh-tp4-m16-1-7f1fb1d1) the +source of truth is the authoritative control-plane fact ledger plus each +worker's `runs//experiments.jsonl` — pending worker SKILL.md tables may +be stale and must not override the ledger. **All four shapes have accepted +candidates; no plateau was proven anywhere (`plateau=false`): every accepted +number is best-found, not a proven optimum.** A shape whose worker timed out / +produced no accepted candidate stays **unoptimized** — never infer success, a +speedup, or optimality from a failed or timed-out worker. + +Operator: `x_q[M,K] int8 @ W[K,N] int8 -> int32 -> * x_scale[M,1] * +weight_scale[N,1].T -> bf16[M,N]`. Timed region = GEMM + scales + bf16 +epilogue + split-K combine inside the CUDA Graph; excludes quantization, +weight packing/preprocessing, allocation, JIT, capture. Protocol: unprofiled +CUDA-graph replay, 100 warmup / 30 samples x 100 replays, hot cache; report +median AND p90; Graph vs Graph only. + +## 1. M = 16 decode (validated) + +Use gfx928 INT8 DUMMA: + +```cpp +DUFragment +DUFragment +DUFragment +``` + +Best validated TP4 dispatch: + +| K,N | Path | +|---|---| +| 4096,1536 | non-uniform split-K=10, 4 waves, stage-K=64, double-buffer prefetch | +| 4096,1024 | split-K=16 | +| 1024,8192 | two-wave staged kernel, K-stage=128 | +| 2048,4096 | two-wave staged kernel, K-stage=128 | +| 512,4096 | two-wave staged kernel, K-stage=128 | + +The non-split kernel computes a 16x32 output tile with two wavefronts. It +loads A/B cooperatively, performs m16n16k32 DUMMA, and writes each accumulator +fragment directly from registers through a fused scale/bf16 epilogue. + +Split-K requirements: + +- Allocate int32 workspace before Graph capture. +- Every launch overwrites every partial tile; no workspace clear is needed. +- Measure partial GEMM plus combine together. +- Do not report only the main kernel as total operator latency. +- Do not restrict split-K to powers of two. Sweep block count near an integer + multiple of the physical CU count, while keeping every K boundary aligned + to the DUMMA/staging unit. + +For the fully validated `M=16,K=4096,N=1536` optimization sequence, byte +accounting, ISA evidence, rejected variants, and the 570 GB/s result, read +[references/wqkv-a-m16-split10-570gbps.md](references/wqkv-a-m16-split10-570gbps.md). + +### 1.1 Hy3 TP4 M=16 shapes (authoritative, 2026-08-23) + +Each accepted kernel is selected by an exact (m,n,k) shape guard placed +**before** generic tiled/scalar paths so decode, M=3072/4096 prefill, and +generic arms stay untouched. `HIP_KERNEL_NAME(templated<...>)` + +`hipLaunchKernelGGL` are valid on this toolchain and used by every accepted +kernel here. Scalar bootstrap kernels (qkv 476.8 us, o_proj 511.4 us, +gate_up 331.0 us, down_proj 39.4 us) are correctness-first placeholders, not +baselines. + +| shape_id | M,N,K | Accepted kernel (final) | Best median / p90 us | baseline us | speedup | +|---|---|---|---|---|---| +| hy3_tp4_qkv_proj_m16 | 16,2560,4096 | `w8a8_dumma_m16_n64_splitk_partial_kernel<6, FUSED=true>` (iter 23) | 26.046 / 26.090 | 80.19 | 3.08x | +| hy3_tp4_o_proj_m16 | 16,4096,2048 | `w8a8_gemm_m16_dumma_packedb_kernel` (iter 16) | 18.221 / 18.242 | 54.617 | 3.00x | +| hy3_tp4_shared_gate_up_proj_m16 | 16,768,4096 | `m16_dumma_runtime_splitk_lds_staged_colmajor_B` (iter 8) | 11.310 / 11.325 | 73.262 | 6.48x | +| hy3_tp4_shared_down_proj_m16 | 16,4096,384 | `w8a8_dumma_m16n16k32_sk2_kernel` (iter 15) | 10.848 / 11.088 | 23.216 | 2.14x | + +Accepted architectures (reusable recipes): + +- **qkv (N=2560,K=4096)**: block-N=64, 256 thr = 4 waves (one 16x16 + quadrant/wave), grid = 40 tiles x split-K=6 = **240 blocks = 2 blocks/CU**; + stage-K=64 double-buffered A+B LDS (10,240 B/block); one 16-B cooperative + load/lane/stage; load-bearing tid-linear B lane roles (`b_krow = tid>>2`, + `b_nchunk = tid&3` = 16 unique 128-B lines per wave-load); one + `__syncthreads`/stage; non-uniform 64-aligned K slicing ([0,K) exactly once, + preserves bit-exact int32 order); **fused last-arrival combine tail** — + monotonic global counters (`arrived % SPLIT_K == SPLIT_K-1`), + `__threadfence` + `atomicAdd`, last arriver sums the L2-hot planes + ascending, scales, stores one 8-B bf16; counters in the last 160 B of the + 16-plane contract workspace, one-time async memset before first launch → + **one kernel launch per replay**. Accepted chain: 476.8 (scalar) → 140.63 → + 131.71 → 124.92 → 54.21 → 36.88 → 31.52 → 28.48 → 26.33 → **26.05** us. +- **o_proj (N=4096,K=2048)**: 256 blocks x 128 thr = 2 waves, one 16x16 N + tile/block; in-block split-K=2, 8 uniform stages of K=128 (4 m16n16k32 + steps, 8 dwordx2 loads)/wave; **register-only K-loop transport with depth-1 + prefetch** (next stage's 8 dwordx2 loads issued before the current 4-MMAC + burst, rotated in after; compiler wait lands at next fill); packed-B + fragment-slot layout `packed[n_tile][k_step][lane][8]` produced once + out-of-timed-region for (k,n)==(2048,4096); A stays logical row-major + (32 KiB, L2-hot); LDS = two 1 KiB int32 partial planes only; one END-of-K + barrier; 44 VGPR/24 SGPR/2048 B LDS, 0 spills. +- **gate_up (N=768,K=4096)**: 16x64 block tile, 2 waves x 2 N-tiles/wave + sharing one A fragment, 12 N-groups; runtime split-K 1..16 default **16** → + 192 blocks (1.6/CU, all 120 CUs); whole 256-row split slice staged once + (A 16x272, B 64x272; `kStagedRowStride=272=256+16` bank skew), batched + 16-B/thread loads → `ds_write_b128`, one barrier, then **zero-barrier + LDS-only K loop** (unroll 2: 2 ds_read2_b32 + 2 ds_read2_b64 + 4 v_mmac); + packed `[N,K]` n-major weight for (k,n)==(4096,768) with `col_major` B + + one 8-B LDS read per fragment; S=16 combine specialization = 96 blocks x + 128 thr with all 16 plane loads issued into a register array before the + ascending sum; GEMM 123 VGPR/68 SGPR/25,856 B LDS; combine in the timed + Graph. +- **down_proj (N=4096,K=384)**: grid = N/16 = 256 blocks x 128 thr = 2 waves, + split-K=2 along K (K-half per wave, 6 unrolled steps); A staged 16x192 + row-major, 208-B padded stride (192+16), dwordx4 loads → ds_write_b128; B + staged n-major from packed `[N][K]` transpose `P[n*384+k]=W[k*4096+n]`, + 208-B padded column stride, one ds_read_b64 per step (was 8 ds_read_u8 + + ~11 VALU); **fused per-block LDS combine**: wave 0 publishes int32 partial + to `s_part`, one barrier, wave 1 adds s=0 then s=1 (exact int32 order) and + emits scaled bf16 from registers; scales prefetched into LDS `s_scale[32]` + at kernel top with stores **deferred until after both staging loops** + (removes prologue vmcnt window); 43 VGPR/18 SGPR/14,464 B LDS. + +Cross-shape rules that held (measured): + +1. **M=16 grid parallelism is scarce — split-K is the occupancy lever.** + Sweep non-power-of-two split counts near integer multiples of 120 CUs + ({2,3,4,5,6,8,9,10,12,16}); winners were 6 (qkv), 2 (o_proj, down_proj), + 16 (gate_up). Keep every K boundary aligned to the staging/DUMMA unit. +2. **Direct global fragment loads are poison first**: library byte-load + + reassembly + per-load `vmcnt` before each mmac cost 1.6-3x. Fix: one + aligned 8/16-B vector load per lane per fragment/step, via LDS staging + (gate_up, down_proj, qkv) or register-only transport (o_proj). +3. **Pack the cold once-read weight (B) outside the timed region/Graph** + into the fragment's layout ([N,K] n-major transpose, per-tile/ + fragment-slot packs) so each fragment is one contiguous dwordx2/b64 read; + leave small L2-hot A logical. Packing alone was small; it is the + prerequisite for vectorized B loads. +4. **Zero in-loop barriers for decode**; one END-of-K barrier for the LDS + int32 partial-plane combine is cheap and order-independent (bit-identical). +5. **Split-K combine is part of the operator wall** — never time the partial + kernel alone. Fuse it where it wins (qkv last-arrival tail; down_proj + per-block LDS) but **not universally**: gate_up's atomic-accumulator + + finalize fusion regressed 43.8%. +6. **LDS bank skew**: pad row/column strides to 16-B-aligned non-power-of-two + values (272=256+16, 208=192+16, 80/72/68/144 elsewhere); strides ≡ 0 + (mod 128 B) alias every row onto one bank phase. +7. **Exact int32 accumulation order (k-ascending per split, ascending split + sum) is preserved across all accepted variants → outputs are bit-identical + (0 mismatches).** +8. **Launch geometry before micro-opt** (down_proj: 256x1-wave 68.9 us → + 128x2-wave 15.4 us); verify the exact code object of the timed run + (digest-matched), not a stale sibling; source-level prefetch is not + overlap without ISA proof (`global_load_dwordx4 -> v_mmac -> vmcnt(0) -> + ds_write_b128`). + +### 1.2 Fallback routing (must stay byte-identical and building) + +- qkv: FUSED=false partial + quadrant combine kernel for workspaces < 16 + planes; generic scalar fallback for every other (m,n,k) including paired + M=2. +- o_proj: scalar fallback decodes the packed-B layout for (k,n)==(2048,4096); + M=2 fallback exact. +- gate_up: scalar fallback decodes the [N,K] pack for (k,n)==(4096,768); + identity copy otherwise; generic single-buffer tiled arms kept untuned. +- down_proj: unsplit two-wave DUMMA (grid 128) for other M=16 shapes; generic + scalar with a `b_transposed` flag for M=2. +- Routing rule: exact-shape guard → tuned kernel; any other shape → generic + tiled path or scalar fallback, never the tuned kernel. + +## 2. 1 <= M < 16 + +The implemented compatibility path zero-pads A to 16 rows in LDS and masks +output rows. It is correct and Graph-safe, but performance must be measured +for each M. Do not assume the M=16 dispatch is optimal for M=1/2/4/8. + +For the validated `M=2,K=4096,N=1536` split-K=8 route, ISA/PMC evidence, +bandwidth accounting, failed B-staging/prefetch/finalize variants, the other +M=2 baselines, and acceptance gates, read +[references/m2-decode-search.md](references/m2-decode-search.md). Do not route +other M=2 shapes to a custom kernel until that reference's exactness and +repeated Graph speedup gates pass; retain Triton otherwise. + +## 3. Boundary: 16 < M <= 32 + +No accepted candidate has been measured in (16, 32]. Reuse the M=16 recipe +(split-K + fused combine family) and measure per exact M before committing a +route. Keep a decode-compatible arm and a staged prefill arm for the same +(K,N) until measured; the decode direct-B-from-global pattern does not +transfer to larger M, and prefill's no-split-K rule does not apply here. + +## 4. Validated evidence + +### M=2 decode, 2026-08-03 + +Fair preallocated-output baseline, hot cache, 50 warmups, 50 samples, +20 launches/sample; exact output check passed: + +| Case | Eager ms | Graph ms | +|---|---|---:|---:| +| wqkv_a | 0.099411 | 0.066924 | +| wq_b / indexer.wq_b | 0.095939 | 0.047181 | +| wo_b | 0.096399 | 0.054453 | +| shared gate_up_proj | 0.096759 | 0.060389 | +| shared down_proj | 0.094487 | 0.018975 | + +For `wqkv_a (K=4096,N=1536)`, the accepted HIP dispatch is split-K=8, +StageK=512, two waves/block sharing one padded A stage, direct row-major B +loads, direct partial stores to `[8,2,N]`, and a 64-thread vectorized combine. +Three final Graph median/P90 runs were `20.004/21.592`, `20.024/20.448`, and +`19.944/20.336 us`; changed-input Graph replay was exact bf16. This is about +3.35x faster than the 66.924 us Triton Graph baseline. The other four rows +remain baseline-only and require their own search. + +### M=16 optimized HIP Graph + +Formal hot-cache run: 50 warmups, 100 samples, 20 launches/sample. +Split-K results include the combine kernel: + +| Case | HIP path | HIP Graph ms | Triton Graph ms | Speedup | +|---|---|---|---|---:|---:| +| wqkv_a | split-K=10 w4 stage64 prefetch | **0.017344** | 0.066597 | **3.84x** | +| wq_b | staged2 K128 | 0.030972 | 0.048049 | 1.55x | +| wo_b | staged2 K128 | 0.050917 | 0.054617 | 1.07x | +| shared gate_up_proj | split-K=16 | 0.018440 | 0.060965 | 3.31x | +| shared down_proj | staged2 K128 | 0.014568 | 0.021372 | 1.47x | + +The updated `wqkv_a` row uses 100 warmups, 300 samples, and 20 Graph replays +per event sample. Its P90 is 0.017616 ms and exact bf16 replay with changed +input contents passed. The other rows retain their earlier protocol and must +not be interpreted as having been rerun in the same session. + +### Hy3 TP4 M=16 (2026-08-23, hy3-dsh-tp4-m16-1-7f1fb1d1) + +The authoritative table (accepted kernels, median/p90, baselines, recipes, +accepted chains) lives in §1.1. Earlier drafts of this table listed down_proj +as (16, 6144, 768) — that was stale; the ledger and worker_3 evidence confirm +`hy3_tp4_shared_down_proj_m16 = (16, 4096, 384)` (`w8a8_dumma_m16n16k32_sk2_ +kernel`, 10.848 us median / 11.088 us p90, 2.14x vs the 23.216 us baseline). + +## 5. Decode acceptance: µs-scale noise protocol + +Decode kernels are µs-scale, so acceptance gates are dominated by measurement +noise, not kernel quality: + +- Report median AND p90; never min. Compare Graph vs Graph only. +- The `final <= 1.05 x worker best` gate can trip on thermal/GPU-state drift at + µs scale. Observed failures: `hy3_tp4_shared_gate_up_proj_m16` final 13.009 us + vs worker best 11.310 us (15% gap) and `minimax_tp8_qkv_proj_m16` final + 73.611 us vs worker best 33.202 us (2.2x gap) despite unchanged kernels — + both tasks stopped on the gate. The control plane now re-measures on a gate + failure: up to 3 re-measures, 300 s apart (`_PERF_GATE_MAX_RETRIES=3`, + `_PERF_GATE_RETRY_INTERVAL_S=300`), accepting the best (min) median; only an + every-attempt failure is reported as a regression. A tolerance band (e.g. + ±10% at M<=16) is also justified. +- Split-K candidates: measure partial GEMM **plus combine together**; never + accept on partial-kernel time alone. +- Expect small speedups: the decode launch/latency floor (~5-10 us) and the + already-decent Triton decode baselines cap headroom at roughly 3-6x. + A flat or small gain can still be the right answer at M<=32. +- Protocol: unprofiled CUDA-graph replay, 100 warmup / 30 samples x 100 + replays, hot cache; report median AND p90 (p90 guard vs current best); + Graph vs Graph only. Timed region = GEMM + scales + bf16 epilogue + + split-K combine inside the Graph (see intro). +- One smallest focused change per round. Only after acceptance explain with + hipprof PMC (`--pmc --pmc-type 3`) and require predicted counters to move, + else the mechanism is ambiguous. `lds_wait ≈ lds_instructions` is NOT a + bottleneck signal (true in accepted kernels too). Profiled duration is + never the score. Derived TOPS/GB/s are logical rates (`2MNK/median`), not + measured HBM traffic. +- After a killed/aborted round, restore the accepted source to its recorded + digest before the next experiment. + +### 5.1 Integration and correctness gates (Graph-safe) + +- Build/load the extension before capture; preallocate output and workspace + before capture; launch on `at::cuda::getCurrentCUDAStream(device)`; no + allocation/sync/JIT/preprocessing/CPU reads inside capture; replay may + update contents, never addresses/shapes. Route M<=16 through + `gemm_out_optimized`, M>16 through the prefill path; unsupported + conditions return None to keep the original SGLang path (foundations holds + the shared Graph-safe contract). +- Correctness: exact bf16 vs reference (0 mismatches expected — int32 order + preserved), extreme/random int8, M tails, Graph capture + changed contents + + replay. Paired M=2 fallback validation must decode the packed layout: + qkv 531 us, o_proj 789.4 us, gate_up 0 mismatch, down_proj 85.9 us — all + mismatch 0. + +## 6. Rejected / guardrail evidence (decode) + +- Do not assume power-of-two split-K. The validated wqkv_a winner uses + non-uniform split10; split4 and split16 both regressed for different reasons. +- Do not fuse split-K finalization merely to remove a Graph node. Atomic + last-arriver finalization regressed, and DTK 26.04 failed to capture the + cooperative-launch candidate into the Graph. (Exception: the hy3 qkv_proj_m16 + winner's fused monotonic-arrival combine is accepted on this lineage — it + passed exact Graph replay.) +- Do not assume fewer tiles per block is better: for decode, insufficient + block count can leave roughly 120 CUs underfilled. +- Historical M=2 trials rejected uint16/uint32 reinterpret vectorization, + manual ILP/unrolling, row fusion, and collapsing grid parallelism for two + shapes. Re-profile before revisiting. +- For M=2 wqkv_a, final ISA contains `v_mmac_i32_16x16x32_i8`, but the + row-major B fragment loader expands each K=512 split to 128 + `global_load_ubyte` instructions. Coalesced B-to-LDS staging regressed to + 32.712 us, paired B-fragment prefetch to 20.488 us, uneven split10 to + 21.381 us, and atomic last-arriver finalize to 21.892 us. Do not replace + the accepted route with these merely because their source looks more + parallel or more vectorized. +- Do not describe the M=2 event-derived 316 GB/s logical tensor rate or + 326 GB/s workspace-inclusive program rate as peak HBM. PMC reported about + 98,560 external 64B read requests (about 6.31 MB) for the main kernel, but + the real compute path also executes a fixed M16 DUMMA tile: 14 of 16 rows + are padding. A 500+ GB/s pure-memory proxy is not a valid compute baseline. +- Hy3 TP4 per-shape rejected (measured; do not re-litigate without new + evidence): + - **qkv** (flat/regressed): packed-B + A-only LDS (55.7), LDS bank skew ldm + 64→68 (36.8), split-K 12 (40.0), stage-K 128 (42.1), deeper prefetch + (32.4), block-N=128 split-K=12 (32.1), per-64-column packed B (26.5), + split-K 9 (28.9), cold-plane L2 warm (28.4). Killed rounds 12/15/20/21 + (infra — prove nothing). + - **o_proj**: in-block split-K=4 on the byte-load kernel (94.2), grid + split-K + combine kernel (94.6), in-loop-barrier A-only LDS (209.3), + depth-2 prefetch (20.16), 8-wave uniform (19.60), stage-K=128 at 6 waves + imbalanced (19.33), epilogue plane bank-skew (19.39), 1-wave occupancy + cliff (19.60), stage-K=256 (19.17). No compile failures in this lineage. + - **gate_up**: finer one-wave 16x32 grid (13.62), packed A-fragment loader — + lighter ISA regressed 44% (20.19; ISA counts are not a performance + proxy), CU-aligned split-K=10 120 blocks (12.99), split-K=8 24 N-groups + (13.23), nontemporal glc/slc B loads (13.11), the whole combine CU-spread + family (13.0-20.7), atomic last-arriver + finalize 3-launch tail (20.11), + exact 2-blocks/CU 240-block grid (15.22). Killed round 7 (agent timeout — + proves nothing). + - **down_proj**: register prefetch depth-2 (flat — compiler kept vmcnt + before mmac), split-K 4/3 occupancy probes (15.0/13.6), stride-4 bank + repack (capacity-invalid, caught by correctness and reverted), four + per-8-k-group B planes (flat 11.72). More co-residency was monotonically + worse at this shape (latency/issue-bound, not occupancy-starved). All 15 + rounds built and passed correctness. +- A compile error proves only that candidate failed to compile (e.g. gate_up + iter 5's undeclared `kPackedK`, fixed in-round), never that DUMMA, INT8, a + HIP API, templates, or inline asm are unsupported. Killed rounds (missing + proposal, ~900 s timeout, restore failures) are infra/agent failures that + prove nothing. +- Do not transfer any winner (split-K default, stage size, layout, occupancy) + to another M / (K,N) / TP size / architecture without re-measuring; decode + vs prefill (M=3072/4096) conclusions do not cross M regimes. + +## 7. Next optimization questions (decode) + +- M=2: keep the validated wqkv_a dispatch; execute the phased search for the + other four `(K,N)` families and add static routes only after they clear the + Graph acceptance gate. +- M=1/4/8: build per-M dispatch only if Graph latency improves. +- M in (16, 32]: measure the M=16 recipe per exact M before committing a route + (see §3); the decode acceptance protocol in §5 applies. +- Hy3 TP4 M=16: all four shapes are **optimized** (accepted, correct, + Graph-passed) with `plateau=false` — accepted numbers are best-found, not + proven floors; continue per-shape HIP search. Baselines (80.19 / 54.617 / + 73.262 / 23.216 us) are fixed user-supplied Triton Graph numbers, never + re-measured, **never optimal**. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/references/m2-decode-search.md b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/references/m2-decode-search.md new file mode 100644 index 00000000..3187997c --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/references/m2-decode-search.md @@ -0,0 +1,166 @@ +# M=2 decode search and validated wqkv_a route on gfx928 + +Use this reference only for `M=2` W8A8 decode. Optimize Graph replay first; +use eager timing only as a secondary launch/binding check. + +## Fair Triton Graph baselines + +| K,N | Graph baseline | +|---|---:| +| 4096,1536 | 66.924 us | +| 1024,8192 | 47.181 us | +| 2048,4096 | 54.453 us | +| 4096,1024 | 60.389 us | +| 512,4096 | 18.975 us | + +Do not choose scalar or DUMMA from M alone. Padded DUMMA wastes 14 of 16 +rows; scalar code gives up matrix throughput. Benchmark both under the same +Graph scope. + +## Candidate A: specialized padded DUMMA + +1. Start with one wave per N16 tile. +2. Keep `A[16,K_STAGE]` in LDS, but initialize padded rows 2..15 once outside + the K loop and copy only the two real rows per stage. +3. Sweep `K_STAGE={64,128,256}` and reject resource/occupancy regressions. +4. Load packed row-major B directly first; add B LDS only after measured reuse + or load-stall evidence. +5. Store only accumulator rows 0 and 1 through the verified direct epilogue. + Lanes with `(lane&15)<2` write; all lanes still execute DUMMA. +6. Fuse scales, bf16 conversion, and output writeback. + +## Candidate B: two-row skinny dot + +Try only when padded DUMMA lacks a robust win. Let one thread own N columns, +compute two int32 accumulators, and reuse each B value across the two rows. +Keep adjacent lanes on adjacent N columns. Start with one wave and 64 columns +per block. Do not use packed casts or `sdot4` until DTK ISA and an appropriate +load-time packed weight layout are verified. Reject row fusion if reduced grid +parallelism outweighs B reuse. + +## Starting grids + +| K,N | N16 blocks | First candidates | +|---|---:|---| +| 4096,1536 | 96 | 1-wave DUMMA; split-K=2/4 | +| 1024,8192 | 512 | 1-wave; 2-wave A reuse; no split initially | +| 2048,4096 | 256 | 1-wave; no split initially | +| 4096,1024 | 64 | 1-wave; split-K=2/4/8 | +| 512,4096 | 256 | 1-wave, stage64/128; avoid split initially | + +For split-K, store only `[split_k,2,N] int32`, not 16 padded rows. Include the +combine kernel in total timing. + +## Validated wqkv_a route: M=2, K=4096, N=1536 + +Worker29, K500SM_AI/gfx928, DTK 26.04, preallocated output/workspace, 50 +warmups, 100 samples, and 20 operations/event produced these retained Graph +results: + +| Candidate | Median us | P90 us | Status | +|---|---:|---:|---| +| generic M<16 padded DUMMA | 197.585 | 197.793 | baseline | +| direct M2 Stage64 | 92.066 | 92.146 | retained | +| direct M2 Stage128 | 69.781 | 70.225 | retained | +| direct M2 Stage256 | 41.625 | 42.161 | retained | +| direct M2 Stage512 | 40.213 | 40.624 | best no-workspace route | +| direct M2 Stage1024 | 40.884 | 41.368 | rejected | +| direct M2 Stage2048 | 41.480 | 41.905 | rejected | +| split2 Stage512 | 26.564 | 26.736 | retained | +| split4 Stage256 | 22.684 | 23.312 | retained | +| split8 Stage512, one wave, combine256 | 20.544 | 21.360 | superseded | +| split8 Stage512, two waves, combine64 | 20.280 | 20.696 | selected architecture | +| B coalesced global-to-LDS staging | 32.712 | 32.912 | rejected | +| paired B-fragment prefetch | 20.488 | 20.856 | rejected | +| uneven split10 | 21.381 | 21.681 | rejected | +| fused atomic last-arriver finalize | 21.892 | 22.008 | rejected | + +After static dispatch integration, three independent selected-path Graph +median/P90 runs were: + +```text +20.004 / 21.592 us +20.024 / 20.448 us +19.944 / 20.336 us +``` + +This clears the 66.924 us Triton Graph baseline by about 3.35x. Exact bf16 +comparison passed for every retained stage and for split8 after changing all +captured input tensor contents before Graph replay. + +### Selected algorithm + +1. Split the 128 K32 MMAC tiles uniformly eight ways; launch + `96 N16 tiles x 8 splits`. +2. Put two adjacent N16 waves in each block. Let both waves share one + `A[16,512]` LDS stage. +3. Zero padded A rows 2..15 once per block and copy only rows 0..1 with + aligned 16-byte HIP vector loads. +4. Load row-major B directly through `du_load_matrix_sync`; execute 16 + `m16n16k32` MMAC instructions per split. +5. Store only two accumulator rows into `[8,2,N] int32` workspace using the + verified direct fragment lane mapping. +6. Use a separate 64-thread vectorized combine kernel for scaling and bf16 + output. Allocate output and workspace before Graph capture. + +The public no-workspace fallback uses direct Stage512. The optimized Graph +API uses split8/two-wave/combine64. + +### ISA and PMC evidence + +The exact accepted code object confirms: + +```text +16 v_mmac_i32_16x16x32_i8 +128 global_load_ubyte +16 ds_read2_b32 +2 ds_write_b128 +1 global_load_dwordx4 +4 global_store_dword +77 s_waitcnt +2 s_barrier +``` + +PMC reports 64 architecture VGPR, 32 SGPR, 8192 bytes LDS, zero scratch, and +about 98,560 external 64-byte read requests per main-kernel invocation +(approximately 6.31 MB). The byte loads are not proof that DUMMA is absent: +the MMAC is present. They arise from the row-major B fragment mapping, where +each lane's eight K values are separated by leading dimension N. + +Do not write raw MMAC merely to satisfy an inline-assembly goal. DUMMA already +emits the target instruction. Raw VMEM is also not justified here: previous +gfx928 raw global-load constraints require independent correctness proof, and +the B elements needed by one lane are not a contiguous dword. Prefer a +prepacked weight experiment only when one-time preprocessing and an explicit +packed-layout contract are allowed. + +### Bandwidth labels + +At about 19.99 us, unique logical tensor bytes imply roughly 316 GB/s; +including split-workspace writes and reads implies roughly 326 GB/s. Label +both as event-derived program rates, not physical peak HBM bandwidth. The +kernel performs real computation and the fixed DUMMA M16 tile pads 14 of 16 +rows, so comparison with a 500-800 GB/s pure-memory proxy is invalid. + +To pursue 500+ GB/s without removing computation, change one of the contracts: + +- prepack B into a fragment/lane-friendly layout outside timed execution; +- batch/fuse multiple M2 requests so the physical M16 tile has more live rows; +- fuse adjacent operators so split/combine and launch costs are amortized. + +Remeasure fallback layout and end-to-end TTFT after any weight preprocessing. + +## Experiment and acceptance order + +```text +Triton baseline -> generic padded DUMMA -> M2-specialized DUMMA +-> stage/wave sweep -> selected split-K -> skinny dot if needed +-> PMC explanation -> SGLang Graph replay +``` + +Accept a dispatch only after exact bf16 random/extreme tests, changed-content +Graph replay, at least 5% median Graph improvement in three independent runs, +no material P90 regression, no scratch/spill, combine-inclusive timing, and no +eager/end-to-end regression that erases the Graph gain. Otherwise retain +Triton. Historical vectorization/ILP failures apply only to those measured +implementations, not to all tiny-M kernels. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/references/wqkv-a-m16-split10-570gbps.md b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/references/wqkv-a-m16-split10-570gbps.md new file mode 100644 index 00000000..04ed0206 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-decode/references/wqkv-a-m16-split10-570gbps.md @@ -0,0 +1,116 @@ +# wqkv_a M16 split10: 570 GB/s logical request bandwidth + +## Contents + +- Evidence and metric contract +- Final kernel structure +- Why non-uniform split-K=10 wins +- Logical-byte derivation +- Optimization sequence and rejected evidence + +## Evidence and metric contract + +This case was validated on worker29 K500SM_AI/gfx928, DTK 26.04: + +```text +M=16, K=4096, N=1536 +int8 A x int8 B -> int32 -> scales -> bf16 +``` + +Use unprofiled GPU events with CUDA/HIP Graph, 100 warmups, 300 samples, and +20 operator calls per event sample. Time partial plus combine. Exact eager, +Graph, and changed-input-at-captured-address bf16 checks passed. + +The reported 570.0 GB/s is `known program-level global requests / event time`. +It is not a hardware-counter measurement of physical HBM traffic. Hot Graph +replay can hit caches. To claim HBM bandwidth, use a cache-cold ring whose +footprint exceeds cache and verify HBM bytes with appropriate PMC counters. + +## Final kernel structure + +- Use block-N=64: four wave64 compute four 16x16 tiles and share A. +- Use stage-K=64 and two LDS buffers. +- Issue stage `i+1` aligned 16-byte A/B global loads before the two current + `v_mmac_i32_16x16x32_i8` operations. +- Place `s_waitcnt vmcnt(0)` immediately before writing the alternate LDS + buffer. Keep the cross-wave barrier; the validated synchronization variant + uses a raw first `s_barrier` and compiler-managed second barrier. +- Write split-major int32 partials and use a separate 256-thread combine + kernel for ten partials, scales, bf16 conversion, and output. + +Do not infer a working software pipeline from HIP source alone. Require final +ISA to show next-stage `global_load_dwordx4` before current MMAC and its wait +near `ds_write_b128` into alternate LDS. + +Final partial-kernel metadata: + +```text +40 VGPR, 27 SGPR, 10240 B LDS +private/scratch 0, VGPR spill 0, SGPR spill 0 +workgroup 256, wavefront 64 +``` + +## Why non-uniform split-K=10 wins + +K contains 64 stage-K=64 units. Divide them as `4x7 + 6x6`, so every split +boundary remains stage/DUMMA aligned. With `N/64=24` N tiles: + +```text +partial blocks = 24 * 10 = 240 +device CUs = 120 +grid batches = exactly 2 blocks/CU +``` + +This avoids the 192-block split8 grid's second batch containing only 72 +blocks. It also avoids split16's extra workspace/combine traffic. Generalize +the method, not the literal split: scan +`ceil(N/blockN) * splitK` around integer multiples of the target CU count, +while measuring block lifetime, residency, partial bytes, and combine cost. + +## Logical-byte derivation + +```text +B read once 4096*1536 = 6,291,456 B +A reread for each of 24 N blocks 16*4096*24 = 1,572,864 B +10 partial planes written 10*16*1536*4 = 983,040 B +10 partial planes read 10*16*1536*4 = 983,040 B +bf16 output 16*1536*2 = 49,152 B +fp32 scales (16+1536)*4 = 6,208 B +total 9,885,760 B +``` + +At the default optimized API Graph median of 17.344 us: + +```text +9,885,760 B / 17.344 us = 570.0 GB/s +2*M*N*K / 17.344 us = 11.608 INT8 TOPS +P90 = 17.616 us +``` + +The 300-sample direct tuner result was 17.432 us median / 17.888 us P90. +The same-session split8/w4 baseline was 18.312 us / 18.512 us. + +## Optimization sequence and rejected evidence + +1. Reduce stage-K to 64 to shorten wait/barrier granularity. +2. Double-buffer LDS and overlap next-stage VMEM with current MMAC. +3. Move from three to four waves: A logical reads fall from 32 copies to 24, + and the partial kernel uses a 16x64 block tile. +4. Use the validated raw-first-barrier synchronization mode. +5. Replace uniform split8 with stage-aligned non-uniform split10 to match the + 120-CU grid. + +Retain these negative results as search boundaries: + +| Candidate | Graph median | Reason | +|---|---:|---| +| waves8, split8 | 18.848 us | only 96 blocks; insufficient CU coverage | +| waves4, split4 | 22.338 us | only 96 blocks despite lower partial traffic | +| waves3/4, split16 | 20.604--22.144 us | excess partial/combine traffic | +| atomic last-arriver fused finalize | 18.864 us | atomic and tail imbalance erase launch saving | +| cooperative grid finalize | unusable in Graph | DTK 26.04 captured an empty Graph | + +Do not optimize only the partial kernel. Separately measure partial and +combine for diagnosis, then accept only full-operator Graph timing. Preserve +working and rejected symbols so later DTK versions can revisit Graph capture +or synchronization without losing evidence. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-foundations/SKILL.md b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-foundations/SKILL.md new file mode 100644 index 00000000..97a52282 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-foundations/SKILL.md @@ -0,0 +1,339 @@ +--- +name: int8-w8a8-gemm-foundations +description: > + Shared foundations for INT8 W8A8 GEMM on Hygon K500SM_AI/gfx928: operator + and layout contracts, DUMMA/epilogue rules, Graph-safe PyTorch bindings, + SGLang integration, fair benchmarking, and correctness gates. Always load + this together with the phase skill: int8-w8a8-gemm-decode (M<=32) or + int8-w8a8-gemm-prefill (M>32). +--- + +# INT8 W8A8 GEMM — Shared Foundations (gfx928 / SGLang TP4) + +This skill records the implementation and evidence validated on worker29 with +DTK 26.04. It supersedes generic rules when the exact TP4 shapes below match. +Do not transfer measured latency or routing to another architecture, TP size, +DTK version, or weight layout without rerunning the benchmark. + +## 0. Phase routing (read first) + +The INT8 W8A8 GEMM skill family is split by execution phase because the two +regimes have non-transferable recipes, different bottlenecks, and different +acceptance protocols: + +| M | Skill to load | Kernel family | Bottleneck | +|---|---|---|---| +| **M <= 32 (decode)** | `int8-w8a8-gemm-decode` | split-K partial + fused combine, small tiles | launch/latency-bound, µs-scale | +| **M > 32 (prefill)** | `int8-w8a8-gemm-prefill` | 2D M-tile, staged LDS, no split-K | LDS/compute-bound, ms-scale | + +Load **this foundations skill together with the phase skill**. Recipes do NOT +transfer across the boundary: the M=16 decode kernel's direct-B-from-global +win does not transfer to prefill, and prefill's "no split-K" rule does not +apply to decode. M in (32, 128] is a measured-boundary gap (validated decode +covers M<=16, validated prefill starts at M=3072): route by measurement and +keep both a staged prefill arm and a decode-compatible arm until measured. + +## 1. Operator Contract + +```text +x_q[M,K] int8 @ weight[K,N] int8 + -> accumulator[M,N] int32 + -> accumulator * x_scale[M,1] * weight_scale[N,1].T + -> out[M,N] bf16 +``` + +The timed GEMM includes integer accumulation, both scales, bf16 conversion, +and output writeback. It excludes activation quantization, weight +preprocessing, tensor allocation, output clearing, JIT compilation, and +Graph capture. + +Target TP4 shapes: + +| Layer | 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 | + +All target K values are divisible by 128 and all N values by 64. + +The M=4096 chunked-prefill family adds four fixed-M shapes (see +`int8-w8a8-gemm-prefill`): `hy3_tp4_qkv_proj_m4096` (M,N,K = 4096,2560,4096), +`hy3_tp4_o_proj_m4096` (4096,4096,2048), +`hy3_tp4_shared_gate_up_proj_m4096` (4096,768,4096), and +`hy3_tp4_shared_down_proj_m4096` (4096,4096,384). + +## 2. Shape and Layout Contract + +SGLang data flow: + +```text +hidden states x[M,K] bf16, stride=(K,1) + -> per_token_quant_int8 +x_q[M,K] int8, stride=(K,1) +x_scale[M,1] fp32, contiguous +``` + +Checkpoint and runtime weight flow: + +```text +checkpoint weight[N,K] int8, contiguous, stride=(K,1) + -> layer.weight = weight.t() +SGLang Triton weight[K,N], non-contiguous, stride=(1,K) + -> one-time contiguous() after loading +HIP weight[K,N], contiguous, stride=(N,1) +``` + +Keep both layouts: + +- Preserve `layer.weight` as the original transposed view for safe fallback. +- Register a non-persistent contiguous HIP buffer once after checkpoint load. +- Never replace `layer.weight.data` with the contiguous copy. Doing so changes + the Triton fallback layout and was observed to regress TTFT. +- Weight preprocessing is allowed by the task and must remain outside timed + execution and Graph capture. + +For the M=4096 lineage the packed weight layout is shape-specific and opaque +(o_proj: `[N/64, K, 64]` int8 panels; gate_up: `[K,N]→[N,K]` transpose; +down_proj: n-major `packed[n][k]`; qkv: n-major `packed[n][k]` for +(k,n)==(4096,2560)); packing is one-time, out-of-timed-region, out-of-Graph, +keeps the same byte count and buffer (graph-stable addresses unchanged), and +correctness is always checked against the raw logical `[K,N]` weight. + +## 3. DUMMA and Epilogue Rules + +gfx928 uses wavefront=64. Block size must be a multiple of 64. + +For INT8, the supported primitive is: + +```text +int8 x int8 -> int32, m16n16k32 +``` + +API rules: + +1. `du_fill_fragment(acc, 0)` before the K loop. +2. Load matrix A/B with `du_load_matrix_sync`. +3. Accumulate with `du_mma_sync`. +4. `du_store_matrix_sync` accepts an accumulator fragment, not a raw array. +5. If materializing C in LDS, pass a scalar pointer such as + `&smem_c[0][0]` and synchronize before other lanes read it. +6. A direct epilogue may use the verified gfx928 accumulator ownership: + `row=lane&15`, `col_mod4=lane>>4`, `frag.x[i]` maps to columns + `col_mod4+4*i`. Re-derive this mapping when changing architecture or + fragment type. + +Prefer direct fragment epilogues when they are verified: they remove the +accumulator LDS round trip and one barrier. For the final non-split `wo_b` +kernel, the resource report was 40 VGPR, 32 SGPR, 6144 bytes LDS, and zero +scratch. + +### 3.1 ISA use for memory-compute overlap + +Use ISA as an acceptance test after HIP/DUMMA architecture choices have been +measured, not as a substitute for tile and stage exploration: + +1. Save the exact gfx928 code object used by the timing run. +2. Locate the exact template symbol; do not disassemble a stale sibling. +3. Check load/MMAC/wait/LDS-store order in the steady-state stage. +4. Read VGPR, SGPR, LDS, scratch, and spill fields from code-object metadata. +5. Reject a candidate if the compiler removes a required load, moves the wait + before MMAC, spills prefetched values, or changes cache semantics silently. +6. Use raw inline assembly only for a minimal operation whose constraint and + hazard behavior is already verified on the current DTK/gfx928 toolchain. +7. Re-run exact correctness, Graph replay with changed contents, median, and + P90 after every ISA-level change. + +For large-M prefill, prefer HIP/DUMMA for loads and MMAC. The accepted raw asm +is limited to `s_waitcnt vmcnt(0)` before alternate-LDS writes and the +validated LDS-ready sequence `s_waitcnt lgkmcnt(0); s_barrier`. Do not copy +raw global-load/store or raw MMAC templates without a lane-layout microtest. + +## 4. Graph-Safe PyTorch Interface + +Preferred API: + +```python +gemm_out(x_q, packed_weight, x_scale, weight_scale, out) +gemm_out_optimized( + x_q, packed_weight, x_scale, weight_scale, out, workspace +) +gemm_out_prefill(x_q, packed_weight, x_scale, weight_scale, out) +``` + +Requirements: + +- Build/load the extension before capture. +- Allocate output and workspace before capture for the low-level interface. +- Launch on `at::cuda::getCurrentCUDAStream(device)`. +- Do not create streams, synchronize, allocate device memory, autotune, + preprocess weights, inspect GPU values on the CPU, or change tensor + addresses inside capture. +- The launcher should contain only shape validation, static dispatch, and + asynchronous kernel launches. +- A Graph replay may update tensor contents but not captured addresses or + shapes. + +Python exposure: + +```text +w8a8_gemm.py + -> torch.utils.cpp_extension.load (before capture) + -> TORCH_LIBRARY / TORCH_LIBRARY_IMPL registration + -> torch.ops.zth_w8a8.gemm_out* + -> HIP launcher on current PyTorch stream +``` + +## 5. SGLang Integration + +Integration flow: + +```text +CompressedTensorsW8A8Int8.process_weights_after_loading() + -> prepare_dcu_w8a8_layer() + -> preserve transposed fallback weight + -> create contiguous HIP weight/workspace once + +CompressedTensorsW8A8Int8.apply_weights() + -> per_token_quant_int8(x) + -> try_dcu_w8a8(...) + M<=32: gemm_out_optimized + M>32 : gemm_out_prefill + unsupported: return None + -> existing Triton fallback +``` + +Files on worker29: + +```text +/workspace/int8w8a8gemm/w8a8_gemm.py +/workspace/int8w8a8gemm/csrc/bindings.cpp +/workspace/int8w8a8gemm/csrc/w8a8_gemm_hip.hip +/workspace/sglang-v0.5.10_dpsk_v4/python/sglang/srt/layers/quantization/dcu_w8a8_gemm.py +/workspace/sglang-v0.5.10_dpsk_v4/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py +``` + +Enable before server start: + +```bash +export PYTHONPATH="/workspace/int8w8a8gemm:${PYTHONPATH:-}" +export SGLANG_USE_ZTH_W8A8_GEMM=1 +``` + +Unsupported conditions must return `None` and retain the original SGLang +path. Current custom path excludes non-TP4 target K/N shapes, bias, non-bf16 +output, wrong dtype/rank/scale shape, non-contiguous HIP buffers, and weights +that were not prepared. + +## 6. Fair Benchmarking + +Keep three scopes separate: + +1. Existing Python operator: may include `torch.zeros`, allocation, and extra + fill kernels. +2. Fair `gemm_out`: preallocated output; times all kernels needed to produce + final bf16 output. +3. PMC kernel sample: diagnostic only; profiler timing is not the final score. + +Formal fair baseline: + +```text +input quantization excluded +weight preprocessing excluded +output allocation/clear excluded +GEMM included +scale + bf16 epilogue included +split-K combine included +``` + +Use unprofiled GPU Events for acceptance. Report eager and Graph separately: + +```text +Triton eager vs HIP eager +Triton Graph vs HIP Graph +``` + +Never compare HIP Graph against Triton eager. Use identical warmups, samples, +launches/sample, inputs, output ownership, and cache policy. + +Recommended stable hot-cache protocol: + +```text +warmups=50 +samples=50 or 100 +launches_per_sample=20 +median + p90 + min +``` + +Use PMC after timing to explain VGPR/LDS/scratch, bank conflicts, MFMA use, +and cache behavior. Do not treat lower VGPR or bank conflicts alone as a +performance win. `lds_wait ≈ lds_instructions` alone is not a bottleneck +signal (true in accepted kernels too). Profiled durations are perturbed — +never use PMC time as the score. Derived TOPS / bandwidth numbers are +`2*M*N*K / median` style logical rates, not measured HBM traffic; use PMC +vmem_read for real global reads. Phase-specific acceptance protocols live in +the decode / prefill skills. + +## 7. Correctness and Acceptance (shared) + +Before accepting a variant: + +1. Compare against the reference for all five TP4 shapes. +2. Test extreme int8 values and random values. +3. Test M=1/2/4/8/16, M tails such as 65, and actual prefill M=3072. +4. Require exact bf16 equality when comparing implementations with identical + int32 accumulation order; otherwise define and justify tolerance. +5. Capture Graph, update static input contents, replay, and compare again. +6. Verify current stream behavior and absence of hidden synchronization. +7. Run fair unprofiled median/P90. +8. Then use hipprof/PMC to explain the result. +9. Run SGLang fallback tests and a representative end-to-end workload. + +After a killed/aborted round, restore the accepted source to its recorded +digest before the next experiment. `hipLaunchKernelGGL` + +`HIP_KERNEL_NAME(templated<...>)` is valid and used by every accepted kernel; +compile failures of a proposal are candidate/infra failures, never evidence +that templates/DUMMA/INT8 are unsupported. + +Relevant tests: + +```bash +HIP_VISIBLE_DEVICES=0 python /workspace/int8w8a8gemm/tests/test_correctness_graph.py --full --m 16 +HIP_VISIBLE_DEVICES=0 python /workspace/int8w8a8gemm/tests/test_sglang_integration.py +HIP_VISIBLE_DEVICES=0 python /workspace/int8w8a8gemm/tests/benchmark_prefill.py --case tp4_decode --m 3072 +``` + +## 8. Shared Guardrails (rejected evidence) + +- Do not mutate the original SGLang transposed weight into contiguous storage; + it silently changes fallback performance. +- Do not count deletion of `torch.zeros` as GEMM kernel acceleration. +- Do not use PMC time as final latency. +- Do not perform JIT, autotuning, allocation, preprocessing, CPU reads, or + synchronization during Graph capture. +- Do not call event-derived logical bytes/time "physical HBM bandwidth". + Report it as program-level logical request bandwidth unless PMC and a + cache-cold protocol prove actual HBM bytes. +- A source change with unchanged/stale profile data is not evidence. Confirm + the code object/source digest and rerun unprofiled timing. +- Reference cache: for M>=3072 the CPU int64 exact reference takes ~10+ + minutes (torch.mm int64 at ~0.2 GFLOPS). The control plane pre-seeds it + once per shape (`w8a8_bench.py --prepare-reference`, 3600s budget) before + the timed benchmark; never re-run a correctness check without the reference + cache present, or the 900s subprocess budget is blown. + +## 9. End-to-end evidence + +TP4 test service with `num_hidden_layers=5`, chunked prefill size 3072, +16 prompts, 20k input, 64 output, max concurrency 16, warmup requests 0: + +| Metric | Original | HIP decode+prefill | Change | +|---|---:|---:|---:| +| duration | 79.03 s | 52.20 s | -33.9% | +| mean TTFT | 41.81 s | 27.09 s | -35.2% | +| mean TPOT | 590.3 ms | 398.1 ms | -32.6% | +| input throughput | 4049 tok/s | 6130 tok/s | 1.51x | + +This is a five-layer validation service, not full 43-layer model performance. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-prefill/SKILL.md b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-prefill/SKILL.md new file mode 100644 index 00000000..0378fe7d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-prefill/SKILL.md @@ -0,0 +1,760 @@ +--- +name: int8-w8a8-gemm-prefill +description: > + INT8 W8A8 GEMM prefill kernels for M>32 on Hygon K500SM_AI/gfx928: 2D M-tile + staged-LDS DUMMA kernels, no split-K, occupancy/LDS tuning, and the M=3072 / + M=4096 validated evidence. Load together with int8-w8a8-gemm-foundations + (contract, layout, DUMMA rules, Graph safety, benchmarking). For M<=32 use + int8-w8a8-gemm-decode. +--- +# INT8 W8A8 GEMM — Prefill (M > 32) — gfx928 / SGLang TP4 + TP8 + +Worker-29 lineage, DTK 26.04; K500SM_AI/gfx928, 120 CUs, wavefront 64, +64 KiB LDS/CU, 65,536 VGPRs/CU. Validated evidence covers M=3072 (TP4) and +M=4096 in three lineages: TP4 (`hy3_tp4_*_m4096`, §2), TP8 +(`hy3_tp8_*_m4096`, §3) and GLM-5.2 TP8 (`glm5-2-dsh-tp8-m4096-1`, §4). +M in (32, 3072) [TP4] / (32, 4096) [TP8 / GLM-5.2] are +measured-boundary gaps — reuse the staged recipes below and measure. This +skill is the prefill half of the phase split; load `int8-w8a8-gemm-foundations` +for the operator/layout contract, DUMMA and epilogue rules, Graph-safe +interface, SGLang integration, and shared guardrails. For each M=4096 lineage +the source of truth is its authoritative control-plane fact ledger (TP4: +2026-08-22 re-verified values; TP8: 2026-08-27; GLM-5.2: +`glm5-2-dsh-tp8-m4096-1-e6a280a2`): pending worker SKILL.md +tables may be stale and must not override the ledger. Baselines are fixed +user-supplied Triton CUDA-graph numbers that were never re-measured — do not +treat any baseline as measured or optimal; "speedup" is always vs that fixed +table. + +## 1. M = 3072 prefill (validated) + +For `M=3072,K=4096,N=1536`, use the validated large-M path: + +```text +block tile = 64x64 +K stage = 64 +threads = 256 = 4 wavefronts +wave ownership = one 32x32 quadrant per wave +DUMMA unit = m16n16k32 +LDS = two buffers of A[64,64] + B[64,64] = 16 KB total +prefetch = two 16-byte vectors/thread into VGPR +epilogue = direct fragment -> scale -> bf16 store +tail M = zero-filled A loads + masked stores +``` + +Launch grid: + +```text +grid.x = N / 64 +grid.y = ceil(M / 64) +``` + +Issue stage `i+1` global loads before the stage `i` MMAC sequence. Keep the +prefetched vectors in VGPR, execute the current eight MMAC instructions, wait +for VMEM immediately before writing the alternate LDS buffer, then use an +LDS-ready barrier. Do not infer this overlap from source order: require the +gfx928 ISA to show `global_load_dwordx4 -> v_mmac_i32_* -> vmcnt(0) -> +ds_write_b128` in that order. + +Current optimized prefill constraints are `M>32`, `K%64==0`, `N%64==0`, +contiguous inputs/scales/output, bf16 output, and no bias. + +Do not use split-K for large M unless profiling proves grid parallelism is +insufficient. MxN already supplies many blocks, and reduction overhead is +normally unnecessary. + +For the full M=3072 experiment table, ISA procedure, resource evidence, +bandwidth accounting, and rejected stage sizes, read +[references/wqkv-a-m3072-prefetch-isa.md](references/wqkv-a-m3072-prefetch-isa.md). + +## 2. M = 4096 chunked prefill — TP4 (hy3_tp4_*_m4096) + +Validated on the worker-29 lineage, one worker/GPU per shape ID, DTK 26.04. +Operator contract is identical to the foundations skill; weight +packing/preprocessing runs outside the timed region and outside Graph +capture. **Source of truth = the authoritative control-plane fact ledger** +(2026-08-22 re-verified values); pending worker SKILL.md tables may be stale +and must not override the ledger. + +| shape_id | M,N,K | Worker | Status | Accepted best (median / p90 us) | TOPS | vs baseline | +|---|---|---|---|---|---|---| +| hy3_tp4_qkv_proj_m4096 | 4096, 2560, 4096 | worker_0 | optimized (iter 18) | 756.73 / 760.16 | 113.51 | 31.78x (24045.12) | +| hy3_tp4_o_proj_m4096 | 4096, 4096, 2048 | worker_1 | optimized (iter 10) | 803.90 / 805.14 | 85.48 | 24.73x (19881.949) | +| hy3_tp4_shared_gate_up_proj_m4096 | 4096, 768, 4096 | worker_2 | optimized (iter 13) | 220.71 / 221.54 | 116.76 | 30.26x (6678.282) | +| hy3_tp4_shared_down_proj_m4096 | 4096, 4096, 384 | worker_3 | optimized (iter 14) | 214.84 / 215.12 | 59.97 | 20.39x (4381.289) | + +> The **2026-08-22 re-verified values**: each accepted kernel object was +> re-benchmarked with correctness ON (mismatch=0, max_abs_error=0.0) using the +> same protocol (100 warmup / 30 samples × 100 replays, CUDA-graph replay, hot +> cache). TOPS = 2MNK/median; bandwidth = algorithmic bytes / median (not +> measured HBM traffic). Speedup = fixed user-supplied Triton baseline / +> median. Iteration-chain values in §6 are the historical in-run measurements +> of the same accepted kernels. + +- All four shapes have accepted candidates: **no GPU worker is unavailable and + no shape in this lineage is unoptimized**. Rule for any future shape: a + shape whose worker timed out or produced no accepted candidate stays + **unoptimized** — never infer success, a speedup, or optimality from a + failed/timed-out worker, and never call its baseline optimal when no + candidate was accepted. +- M=4096 prefill: **no split-K** in any accepted candidate — the M×N + output-tile grid already dwarfs 120 CUs, and the 16 MiB workspace has no + useful int32 partial-plane capacity. +- Routing: exact `(m,n,k)` guard per shape → tuned DUMMA kernel; every other + shape falls back. Place the exact-shape guard before generic tiled paths and + keep it shape-exact so decode (M≤32), M=3072, and generic arms are + untouched. Each source carries a generic scalar int8/int32 fallback that + decodes the packed layout when the `(k,n)` matches: o_proj `(2048,4096)`, + qkv `(4096,2560)`, down_proj n-major, gate_up `(4096,768)`; gate_up also + keeps generic single-buffer tiled arms (`<64,128,128>`, `<128,64,128>`, + `<64,64,128>`) that were never re-tuned. Untuned paths must remain + byte-identical to the accepted source and must still build/pass. +- All speedups are vs the fixed user-supplied Triton CUDA-graph baseline table + above; those baselines were never re-measured and must not be called + optimal. +- Decode (M≤32), M=3072, and any other shape are outside this lineage and + must keep their own validated routes/fallbacks. + +## 3. M = 4096 chunked prefill — TP8 (hy3_tp8_*_m4096) + +Validated on the worker-29 lineage (hy3-dsh-tp8-m4096-1), one worker/GPU per +shape ID, DTK 26.04. **Source of truth = the authoritative control-plane fact +ledger (2026-08-27)**; pending worker SKILL.md tables may be stale and must +not override the ledger (e.g. down_proj's draft names round 8 as final while +the ledger also accepted rounds 10/12/13; gate_up's draft calls round 11 "not +accepted" though the ledger accepts it). All four TP8 M=4096 prefill shapes +have accepted candidates (4/4); no GPU worker is unavailable. ~30 rounds in +this lineage ended as `compile_or_agent_failure` (killed ~900 s, missing +proposal, restore/resume failures): these are candidate/infra failures and +prove nothing — never infer success, a speedup, or optimality from a +timed-out/failed worker, and a compile error is never evidence that DUMMA, +INT8, a HIP API, templates, or inline asm are unsupported. Recipes do NOT +transfer across TP size: the TP4 numbers in §2 are a different lineage with +different shapes. + +### 3.1 Routing and accepted evidence (authoritative ledger) + +| shape_id | M,N,K | accepted best (iter) | median / p90 us | logical TOPS | speedup vs baseline | +|---|---|---|---|---|---| +| hy3_tp8_qkv_proj_m4096 | 4096,1280,4096 | 6 | 418.22 / 431.55 | 102.7 | 46.34x (19381.383) | +| hy3_tp8_o_proj_m4096 | 4096,4096,1024 | 12 | 322.99 / 325.38 | 106.4 | 43.82x (14151.792) | +| hy3_tp8_shared_gate_up_proj_m4096 | 4096,384,4096 | 18 | 115.95 / 116.23 | 111.1 | 144.13x (16712.231) | +| hy3_tp8_shared_down_proj_m4096 | 4096,4096,192 | 10 | 170.80 / 171.03 | 37.7 | 63.01x (10762.616) | + +All four: correctness passed (mismatch 0 / max_abs_error 0.0), Graph capture ++ changed-contents replay passed. No proven optimum on any shape (no plateau +established): accepted numbers are the best found, not ceilings. No split-K +in any accepted candidate — the MxN tile grid dwarfs 120 CUs; the only +measured split (gate_up S=2) cost ~40-70 us of combine. + +### 3.2 TP8 pack layouts and fallback guards (lineage-specific) + +Foundations records the TP4 layouts; the TP8 layouts below are distinct and +shape-specific. Checkpoint weight is [N,K] (stride K); SGLang Triton view is +W.t() — preserve it for fallback and register one one-time contiguous HIP +buffer; never replace `layer.weight.data` with the contiguous copy. Packing +is one-time, out-of-timed-region, out-of-Graph, same byte count, same +graph-stable buffer; correctness is always checked against the raw logical +[K,N] weight: + +- qkv (k=4096,n=1280): n-major `packed[n][k]` +- o_proj (k=1024,n=4096): n-major `packed[n*K+kk] == raw[kk*N+n]` +- gate_up (k=4096,n=384): [N,K] transpose (exact (K,N)==(4096,384) only) +- down_proj (k=192,n=4096): n-major `packed[n][k]` + +Validate every pack kernel with a byte-exact simulation of the index +relation. o_proj iter 10's mis-vectorized transpose silently left most of +the buffer unwritten (mismatch 16,775,923 / max_abs_error 86.5): a fast but +wrong measurement is not a win; the iter-11 repair (per-thread 16-byte +gather from 16 raw rows + aligned int4 store) restored mismatch 0. + +Fallback routing (must stay byte-identical and building): exact-shape guard +first -> tuned DUMMA kernel; every other shape -> generic tiled arms / scalar +fallback (`w8a8_gemm_scalar_fallback_kernel`), which decodes the packed +layout when the (k,n) matches and is otherwise identity. Guards (unchanged +across the lineage): qkv `(m >= 128 && n == 1280 && k == 4096)`; o_proj +`(m,n,k) == (4096,4096,1024)`; gate_up `(m == 4096 && k == 4096 && n == 384)`; +down_proj `(m >= 128 && k == 192 && n % 128 == 0)`. M tails and paired shapes +must still pass. Occupancy: verify from the exact code object (vgpr_count, +sgpr, LDS bytes, scratch/spills); `__launch_bounds__` minBlocks is a hint +only. Math: blockDim multiple of 64; blocks x LDS <= 64 KiB; VGPR x threads x +blocks <= 65,536. + +### 3.3 Accepted architectures (per shape) + +- **qkv — `w8a8_dumma_prefill_packedb_tiled_kernel<64,128>` (iter 6)**: + 64x128 tile, 256 thr = 4 waves, 32x64 quadrant/wave = 8 m16n16k32 int32 + acc; grid (N/128)x(M/64) = 10x64 = 640 blocks; single-buffered 64-K LDS + stage (A[64,64] stride 68 B odd-word, B[128,64] n-major stride 80 B, + 14,592 B/block), 2 __syncthreads/stage, peeled final stage (127 + barriers/block); hoisted 12 fragment loads (4 ds_read2_b32 A + 4 + ds_read2_b64 B via load_b_frag8) then 16-MMAC burst; register-scale + coalesced epilogue (4x4 shuffle transpose, 2+4 scale loads prefetched, one + 8-byte store/lane/fragment); 2 blocks/CU, VGPR-bound (arch 96 / code-object + 89 / sgpr 26 / no spills). Chain: 448.86 (iter 1) -> 418.22 (iter 6). +- **o_proj — `w8a8_dumma_128x64x64_kernel` (iter 12)**: 128x64 tile, 256 thr + = 4 waves, 64x32 quadrant/wave = 8 acc; grid dim3(64,32) = 2048 blocks; + single-buffered 64-K stage A[128,80]+B[64,80] = 15,360 B, 2 + __syncthreads/stage (32/block); 3 int4 global loads/thread/stage -> + ds_write_b128 -> barrier -> 6 ds_read2_b64 -> 16 v_mmac; n-major [N,K] pack + (repaired iter-11 gather) + load_fragment8 on BOTH operands (one 8-byte LDS + read per fragment, zero reassembly VALU); coalesced 8-byte epilogue (iter 6: + 4x4 __shfl_xor transpose, ds_bpermute lowering accepted at the block tail); + 3 blocks/CU per the ledger (worker code-object re-verification: 61 VGPR/22 + SGPR/0 scratch — confirm residency from the exact object). Chain: 419.05 + (iter 6) -> 348.27 (iter 11) -> 322.99 (iter 12). +- **gate_up — `w8a8_gemm_prefill_tiled_kernel_w8<64,128,64>` (iter 18)**: + 64x128 tile, 512 thr = 8 waves, 32x32 quadrant/wave = 4 acc, 512 + v_mmac/wavefront; grid 3x64 = 192 blocks; K stage 64 double-buffered, one + __syncthreads/stage (65/block); B packed transposed [N,K] (exact + (K,N)==(4096,384) only), col_major fragments via load_fragment8 (one + contiguous 8-byte LDS read); LDS 30,720 B/block, 2 blocks/CU, code object + 53 VGPR/26 SGPR/0 scratch. Chain: 169.46 (iter 2) -> 118.47 (iter 9: w4, + 256 thr/4 waves/32x64 quadrant/8 acc — biggest single win) -> 117.50 + (iter 13, load-all-then-MMAC-all burst) -> 115.95 (iter 18, w8); iters + 11/14/15 also accepted (117.6-117.7). +- **down_proj — `w8a8_dumma_prefill_kernel<64,128,64>` (iter 10)**: 64x128 + tile, 512 thr = 8 waves, 32x32 quadrant/wave = 4 acc; grid (N/128, M/64) = + 2048 blocks; K=192 in 3 x kStageK=64 double-buffered stages, one + barrier/stage + prologue (4 dynamic; iter-12 variant drops the dead last + barrier -> 3); A stride 88 B (kLdsPad 24), B n-major packed stride 72 B + (kLdsPadB 8), ds_write2_b64 staging; 29,696 B/block, 2 blocks/CU = 16 + resident waves, 55-57 VGPR; load_frag8 (ds_read_b64, merged to 4 + ds_read2_b64/wave-stage) replacing du_load_matrix_sync; fused + direct-fragment epilogue with iter-10 register-batched weight_scale loads + (16 -> 8 scale loads/thread); left-associative ((float)acc * xs) * ws fp32, + __float2bfloat16. Chain: 237.78 (iter 1) -> 234.23 (iter 4) -> 230.91 + (iter 6) -> 171.19 (iter 8, load_frag8, -25.9%) -> 170.80 (iter 10); + iters 12/13 also accepted (170.93/170.94). + +### 3.4 Rejected ideas (measured, correct where noted; do not re-litigate) + +- qkv: 64x64 flip 509.14; 128x64 466.05; k8 swizzle 451.86 (neutral — + conflicts not binding); explicit load_a_frag8 561.88 (CONFOUNDED: VGPR + 96->78 silently moved 2->3 blocks/CU) and 896.63 (pinned — the library + loader's "dead" identity VALU is load-bearing latency fill; do NOT remove + A-side VALU here); 512-thread blocks 481.55; one-stage-ahead staging + prefetch 1097.63 (DTK scattered 34 vmcnt waits); double-buffer/1-barrier + 457.14; A stride 68->80 768.34 (odd 17-word stride load-bearing); B + st-pair merge 659.43; tail-guard consolidation 448.17. +- o_proj: iter-10 pack bug 327.36 (CORRECTNESS FAILED — never count a fast + wrong result); 128x128 854.05; 512 thr 491.51; B stride 80->72 719.73 / + A-stride 80->72 437.15 (conflict floor not on the critical path); epilogue + scale hoisting 741.15; double-buffer prefetch 891.29 (confounded: real vgpr + 85 + 30,720 B = 2 blocks/CU, 50% occupancy loss); prefetch + + __launch_bounds__(256,4) 659.5 (15 spills + 48-B scratch — + register-blocked in pure HIP); last-barrier drop 491.30 (flat/unstable + window). +- gate_up: BM=32 grid 768 586.23 (doubled B staging); 128x64 flip 181.44; + 64x64 3 blocks/CU 223.90; stage-top publish 181.21 and 121.40 at w4 + (+2.5%); split-K S=2 240.19; deeper-MMAC-pipe prediction REFUTED (halving + per-wavefront MMAC count gave only +2.17% — ~2-deep per-SIMD pipe). +- down_proj: 2-stage prefetch + 1 block/CU 617.07 (co-residency load-bearing); + staging thread remap 182.74 (-6.3%: 4x L2 requests); ws sharing 16->8 + 172.17 (p90 guard fail); scales staged in LDS 179.52 (-4.9%); split-3 + infeasible by arithmetic (LDS 89,088 B > 64 KiB, VGPR 86,016 > 65,536). +- Killed rounds (qkv 2,10,13,15,16,20; o_proj 2,9,13,15,22; gate_up 1,5,12, + 16,17; down_proj 2,5,7,11,15,17) prove nothing about the toolchain. + +### 3.5 Cross-shape rules that held (measured) + +1. Pack B once into the exact (k,n) layout (n-major / [N,K]) outside + timing+Graph; together with 8-byte fragment loads this is the dominant + lever on every shape. +2. `load_fragment8` (one ds_read2_b64 per fragment, identical x[0..7] bytes = + identical v_mmac operand) is the winning fragment-load form — EXCEPT qkv's + A operand, where the library row-major loader's identity VALU is + load-bearing latency fill (removal regressed even with residency pinned; + qkv's B-side load_b_frag8 is fine). +3. LDS bank skew: 16-byte-aligned non-power-of-two row strides (68/72/80/88); + strides ≡ 0 (mod 128 B) alias every row onto one bank phase. When + stride % 16 == 8, stage with two int64 halves (ds_write2_b64), not + ds_write_b128. +4. Residency is per-shape and load-bearing: qkv 2 blocks/CU (VGPR-bound), + o_proj 3-4, gate_up 2, down_proj 2 = 16 waves/CU. Verify the exact code + object before trusting any delta; occupancy confounds invalidate verdicts. +5. These kernels are LDS-latency/issue-bound, not HBM- or MMAC-bound (MMAC + pipe ~2-3%): bank-conflict elimination, barrier removal, deeper register + prefetch, and occupancy changes stop paying once strides and pack layout + are fixed. +6. Epilogue: coalesced 8-byte stores (4x fewer vmem_write) and scale-load + hoisting/register batching are the winning post-MMAC moves; staging scales + in LDS lost (down_proj -4.9%). +7. Bit-identical discipline: preserve k0-outer/kk-inner int32 order and the + element-to-slot fragment mapping; variants then pass mismatch 0 / + max_abs_error 0.0 with no tolerance debate. +8. Recipes do not transfer across M regimes or per-shape constants: tile + aspect (down_proj 64x128 wins at K=192 vs 128x64 at TP4 K=384), stage + size (64 saturates), and occupancy must be re-measured per exact shape. + +Acceptance for this lineage follows the shared protocol in §7 plus the +TP8-specific notes there. + +## 4. M = 4096 chunked prefill — GLM-5.2 TP8 (glm5-2-dsh-tp8-m4096-1) + +Synthesis of the `glm5-2-dsh-tp8-m4096-1-e6a280a2` lineage: 4 workers, 6 +exact-shape IDs, all INT8 W8A8 GEMM prefill (M=4096) on Hygon K500SM_AI / +gfx928 (120 CUs, wavefront 64, 64 KiB LDS/CU, 65,536 VGPR/CU, DTK 26.04). +**Source of truth = the authoritative control-plane fact ledger**; pending +worker SKILL.md tables may be stale and must not override it. All 6/6 exact +M=4096 shapes are optimized (accepted, correct, Graph-passed); no GPU worker +is unavailable (`{}`). Baselines are fixed user-supplied Triton CUDA-graph +numbers that were never re-measured — never call them optimal, and no accepted +candidate is a proven optimum (`plateau=false` everywhere: accepted numbers +are best-found, not ceilings). ~30 rounds across the lineage ended +`compile_or_agent_failure` (agent killed ~900 s, missing proposal, restore +failures): these prove nothing — a compile error is evidence only that that +candidate failed to compile, never that DUMMA, INT8, a HIP API, templates, or +inline asm are unsupported. `hipLaunchKernelGGL(HIP_KERNEL_NAME(templated<...>))` +is valid and used by every accepted kernel. + +### 4.1 Routing and fallback guards + +Exact-shape guard first -> tuned DUMMA kernel; every other `(m,n,k)` falls +through to generic tiled arms / scalar fallback. Guards (unchanged across the +lineage): fused_qkv_a `m >= 128 && n == 2624 && k == 6144`; kv_b +`k == 512 && n == 3584 && m >= 64 && m % 64 == 0`; q_b exact `(k,n) == +(2048,2048)`; o_proj exact `(m,n,k) == (4096,6144,2048)`; shared_down exact +`(m,n,k) == (4096,6144,256)`; shared_gate_up exact `(m,n,k) == +(4096,512,6144)`. Untuned arms (generic single-buffered `<64,128,128>` / +`<128,64,128>` / `<64,64,128>` tiled instantiations, pack kernels, scalar +fallback) must stay byte-identical and still build/pass; M tails and paired +API shapes keep working. New mechanisms are added as defaulted template +parameters behind `if constexpr` so retention instantiations are +token-identical. Operator contract: no bias, no workspace (`(void)workspace` — +no combine pass exists), **no split-K** (MxN tile grid dwarfs 120 CUs on every +shape; the only measured split in the TP8 family cost ~40-70 us of combine); +scale order left-associative `(float(dot) * x_scale[row]) * weight_scale[col]`, +then `__float2bfloat16` (plain RNE without exec-masked inf/NaN fixup is +accepted where bit-identical for finite inputs). + +### 4.2 Pack layouts (shape-specific, validated byte-exact) + +Checkpoint weight is `[N,K]` (stride K); the SGLang Triton view is `W.t()` — +preserve it for fallback, register one one-time contiguous HIP buffer, never +replace `layer.weight.data`. Every pack is one-time, out-of-timed, out-of-Graph, +same byte count, same graph-stable buffer, and must be validated by a +byte-exact simulation of the index relation before timing (a fast kernel +reading an incompletely written pack is a correctness failure, never a win; +the sibling lineage's mis-vectorized transpose left 16.7M elements unwritten). +The scalar fallback must decode the pack for the exact (k,n). + +| shape (K, N) | pack layout | +|---|---| +| fused_qkv_a (6144, 2624) | B-panel `packed[(n>>6)*(k*64) + (kk>>4)*1024 + (n&63)*16 + (kk&15)]` — 41 panels of [K,64], 16-B aligned contiguous k-runs/column | +| kv_b (512, 3584) | tile-contiguous stage-major `packed[(((k0*(n/128)+nt)*8+kc)*128+row)*8+b]` — each (stage, n-tile) B tile one contiguous 8192-B region | +| q_b (2048, 2048) | swizzled 64-k-stage-major plane `[kc][n][8]` — lane-linear 8-byte fragment reads, zero LDS bank conflicts | +| o_proj (2048, 6144) | n-major `packed[n*K+kk] == raw[kk*N+n]` | +| shared_down (256, 6144) | panel `[N/128][K/64][128][64]` (n-major 64-k-byte rows/panel) for fully coalesced 128-B staging lines | +| shared_gate_up (6144, 512) | `[K,N] -> [N,K]` n-major transpose | + +### 4.3 Accepted evidence (authoritative ledger; median / p90 us, CUDA-graph replay) + +| shape_id | M, N, K | accepted best (iter) | kernel | median / p90 us | logical TOPS | speedup vs fixed baseline | +|---|---|---|---|---|---|---| +| glm_tp8_fused_qkv_a_proj_m4096 | 4096, 2624, 6144 | 5 | `w8a8_dumma_prefill_packedb_db_kernel<128,64,64>` | 937.58 / 939.48 | 140.86 | 73.23x (68661.417) | +| glm_tp8_kv_b_proj_m4096 | 4096, 3584, 512 | 17 | `w8a8_dumma_prefill_tile_kernel<64,128,32,32,true,24>` | 137.40 / 137.85 | 109.40 | 50.51x (6939.939) | +| glm_tp8_q_b_proj_m4096 | 4096, 2048, 2048 | 7 | `w8a8_dumma_prefill_128x128_kernel` | 220.14 / 220.50 | 156.08 | 65.35x (14385.571) | +| glm_tp8_o_proj_m4096 | 4096, 6144, 2048 | 9 | `w8a8_dumma_256x64x64_packedb_kernel` | 759.92 / 760.81 | 135.64 | 72.08x (54775.797) | +| glm_tp8_shared_down_proj_m4096 | 4096, 6144, 256 | 5 (9 also accepted) | `w8a8_dumma_prefill_tiled_kernel<64,128,64,...>` | 278.44 / 279.19 | 46.28 | 63.57x (17700.262) | +| glm_tp8_shared_gate_up_proj_m4096 | 4096, 512, 6144 | 16 | `w8a8_dumma_prefill_tiled_kernel<64,64,64,true,true,true,true,true>` | 273.83 / 275.35 | 94.11 | 77.62x (21254.589) | + +All six: correctness passed (mismatch 0 / max_abs_error 0.0), Graph capture + +changed-contents replay passed. No split-K anywhere; exactly one kernel per +GEMM; workspace untouched. + +### 4.4 Accepted architectures (condensed recipes) + +- **fused_qkv_a (iter 5)**: 128x64 tile, 256 thr = 4 waves, 64x32 quadrant/ + wave = 8 int32 acc; grid 41x32 = 1312 blocks (M=4096=32x128, N=2624=41x64 + exact fit — no tail); 96 ascending 64-K stages, double-buffered LDS + `a_tile[2][128*72] + b_tile[2][64*72]` = 27,648 B/block -> **2 blocks/CU = + 8 waves/CU**; 93 VGPR / 31 SGPR / 0 spills; stage-(s+1) loads prefetched + into 3 int4 staging VGPR at stage top, flush after MMAC burst, ONE + barrier/stage; B-panel pack. Chain: 5697.3 (iter 1) -> 2276.9 (iter 2, + n-major pack) -> 1070.2 (iter 3, double buffer) -> **937.6** (iter 5, + B-panel pack). +- **kv_b (iter 17)**: 64x128 tile, 512 thr = 8 waves, 32x32 quadrant/wave = + 4 int32 acc; grid 28x64 = 1792 blocks; K=512 = 8 x 64-K stages, + single-buffered (two barriers/stage); A stride 88 (kALdsPad=24, zero- + conflict reads) + B tile-contiguous stage-major, ds_write2_b64 staging; + LDS 13,824 B/block -> **2 blocks/CU = 16 waves/CU**, 61 VGPR; + `kPrefetchNext` issues stage-s+1 grouped A+B loads AFTER the s MMAC burst, + before the trailing barrier (8-VGPR payload live only across barrier + + back-edge — never carry staging payloads across the burst). Chain: + 209.41 -> 169.26 (128x64, 3 blocks/CU) -> 149.54 (64x128 flip, 16 waves/CU) + -> 146.54 (tile-contiguous pack) -> 142.70 (plain RNE bf16) -> 139.57 + (grouped A+B staging) -> **137.40** (consolidation of prefetch + A skew). +- **q_b (iter 7)**: 128x128 tile, 256 thr = 4 waves, 64x64 quadrant/wave = + 16 int32 acc; grid 16x32 = 512 blocks; single-buffered 64-K stage + A[128,80] + B plane `[kc][n][8]` = 18,432 B/block -> **1 block/CU = 4 + waves/CU** (best occupancy on this shape; more resident waves regressed); + 141 VGPR; `load_frag8` on BOTH operands (iter 7, -11%); register-batched + coalesced epilogue (iter 6, -9.1%). Chain: 310.22 (n-major pack) -> 275.79 + (128x128) -> 272.06 (swizzled plane) -> 247.28 (epilogue) -> **220.14** + (A load_frag8). +- **o_proj (iter 9)**: 256x64 tile, 256 thr = 4 waves, 128x32 quadrant/wave = + 16 int32 acc; grid 96x16 = 1536 blocks; single-buffered 64-K stage + A[256,88] + B[64,72] = 27,136 B/block -> **2 blocks/CU = 8 waves/CU**; + 151 VGPR / 34 SGPR; 5 `global_load_dwordx4` -> 5 `ds_write2_b64` (strides + 8 mod 16) -> 10 `ds_read2_b64` -> 32 v_mmac per wave/stage; n-major pack + + `load_frag8` B; register-batched scales epilogue (iter 6, +9.0%). Chain: + 1149.26 (128x64 DB) -> 848.53 (256x64, +35.4% biggest win) -> 839.96 (pack + + load_frag8) -> 770.49 (epilogue batching) -> **759.92** (A stride 88 + conflict-free, +1.4%). +- **shared_down (iter 5; iter 9 also accepted)**: 64x128 tile, 512 thr = 8 + waves, 32x32 quadrant/wave = 4 int32 acc; grid 48x64 = 3072 blocks; K=256 = + 4 x 64-K stages single-buffered (2 barriers/stage) — single buffer beat + double buffer on this K=256 shape (~1.7%); LDS A[64,88] + B[128,72] = + 14,848 B/block -> 2 blocks/CU; panel pack (fully coalesced 128-B staging + lines) + load_frag8 + kSkew8 strides. Iter 9 adds kFragReuse + + `__launch_bounds__(512,3)` (3 blocks/CU = 24 waves/CU at 42 VGPR) — + essentially tied (278.52/279.11), accepted per ledger. Chain: 286.53 -> + 281.65 -> **278.44**. +- **shared_gate_up (iter 16)**: 64x64 tile, 256 thr = 4 waves, 32x32 + quadrant/wave = 4 int32 acc; grid 8x64 = 512 blocks; K=6144 = 96 x 64-K + stages double-buffered, ONE barrier/stage (97/block); LDS A[64,88] + + B[64,72] x2 = 20,480 B/block -> **3 blocks/CU = 12 waves/CU**, 54 VGPR; + `kPayload` lead-2 loop-carried payload: stage s+2 global loads issued at + the top of iteration s (no wait), payload published into the idle LDS + buffer before the stage-s MMAC burst, vmcnt wait lands AFTER the 8 x v_mmac + burst immediately before the barrier (removes the exposed prefetch wait); + `load_frag8` (iter 10, -9.1%, body 175->118 instr) + kSkew8 strides + (iter 11, conflicts 9.44M->3.15M). Chain: 432.52 -> 368.76 -> 331.64 + (double buffer) -> 301.57 (load_frag8) -> 293.30 (skew) -> **273.83** + (payload). + +### 4.5 Rejected ideas (measured; do not re-litigate without new evidence) + +- **fused_qkv_a**: `<128,128,64>` square tile, 1 block/CU 3491.7 (occupancy + halving dominates); branch-free bf16 RNE epilogue 2554.8 (epilogue NOT the + bottleneck — "ISA-neutral by construction" must be measured); prefetch + distance 2 via two staging sets 4992.1 (register blowup); BM=64, 3 blocks/CU + 1148.0; StageK 128 / 512 thr / 1 block/CU 1055.5; rasterization flip + `dim3(32,41)` 3354.3; bottom-of-stage prefetch issue 3233.0; load-all-then- + MMAC-all 938.8 (neutral, p90 guard fail). +- **kv_b**: double-buffered one-barrier stage 197.24 (28,672 B -> 2 blocks/CU + occupancy loss dominates VMEM-stall removal); double-buffer pipelines with + the 8-VGPR payload carried across the burst (72/72/68 VGPR -> 1 block/CU) + 153-156; 64x64 tile + minBlocks=5 probe (vgpr 48 honored, 20 waves/CU) + 491.10; ISA-guided dead last-stage barrier removal 137.87 (p90 guard fail). +- **q_b**: 512-thread w8 variant 256.09 (+50% per-stage LDS read bytes); + `#pragma unroll 1` -> 2 blocks/CU 266.57 (occupancy REFUTED on this shape — + 1 block/CU = 4 waves/CU stays best); loop-carried payload + stage-top + publish 221.57 (+0.65%, p90 fail — VMEM exposure refuted at 1 wave/SIMD); + A stride 88 skew 223.73 (conflict floor not on the critical path here); + 128-K stage depth 357.77 (stage 64 saturates). +- **o_proj**: 64x128 tile 1256.05 (per-byte B ds_read_u8 class); direct-A from + global 1715.91 (cooperative A+B staging load-bearing); one-stage-ahead + prefetch 786.08 (-2.0%, falsified — VMEM latency already hidden by + co-residency); load-all-then-MMAC-all 760.88 (flat, DTK already hoists); + wave-private staging + barrier removal 973.16 (-21.9%, 2x staging traffic); + disjoint 64x64 quadrants 796.73 (-4.6% despite -20% LDS fragment-read + traffic — LDS-read-pipe hypothesis falsified). +- **shared_down**: 64x64 tile 297.15 (+33% A staging traffic); kRNEEpilogue + branch-free RNE 320.88 (+15.2% — epilogue VALU not binding); kSPayload + register payload 281.19 (p90 fail — load-wait relocation alone does not pay + at K=256); kStageK=128 279.61; kBatchEpilogue scale hoisting 284.61; + kLoadPair back-to-back staging loads 281.24. +- **shared_gate_up**: 32x64 tile / 2-wave blocks 749.87 (confounded, wave + count per CU load-bearing); q64 128-thr 32x64 quadrant 629.02 (+90%); + fragment-prefetch pipeline (persistent fragments across the K loop) 337.97; + 64x128 w8 / 512 thr 395.39 (-16.1%, loses grid blocks and tail); kStageK=32 + / 6 blocks/CU 447.37 (-34.4% — more occupancy with phase-locked co-resident + blocks does not pay). +- Killed rounds (fused_qkv_a 9,10,12,15,17,18; kv_b 11,14,15,18,19,21; q_b + 1,2,12-15,17; o_proj 7,11,13,14; shared_down 1,7; shared_gate_up 5,9,13,14, + 15,17) are infra/agent failures — no evidence about hardware, APIs, or + candidate ideas. + +### 4.6 Cross-shape rules that held (measured) + +1. **Per-shape B packing once, outside timing+Graph, into the shape-exact + layout** is the dominant lever on every shape; together with 8-byte + fragment loads it removes the per-byte LDS read + reassembly VALU class. + Validate packs byte-exact; the scalar fallback must decode the pack for + the exact (k,n). +2. **`load_frag8` wins** (identical operand bytes -> bit-identical int32) + except where the library A-side identity VALU is load-bearing latency fill + (verify per shape). +3. **LDS bank skew**: 16-byte-aligned non-power-of-two row strides + (68/72/80/88/136); `stride % 16 == 8` -> `ds_write2_b64` halves. Fix + conflicts only when they sit on the critical path (payoff is per-shape: + +0.6% kv_b, -1.6% q_b, +1.4% o_proj). +4. **64-K stage depth saturates** (128-K lost on q_b 357.77 and fused_qkv_a + 1055.5; 32-K lost on gate_up 447.37). Double buffering wins only when + residency is preserved (fused_qkv_a, gate_up); it lost whenever it cost + blocks/CU (kv_b, down_proj). +5. **Residency is per-shape and load-bearing**: fused_qkv_a/o_proj 2 + blocks/CU = 8 waves/CU; kv_b/down_proj 2 blocks/CU = 16 waves/CU; gate_up + 3 blocks/CU = 12 waves/CU; q_b 1 block/CU = 4 waves/CU. Verify from the + exact code object; occupancy confounds invalidate verdicts. +6. These kernels are **LDS-latency/issue-bound, not HBM- or MMAC-bound** + (MMAC pipe ~2-3%). Once pack layout + strides + pipeline are fixed, + barrier removal, occupancy probes, epilogue-scale hoisting and load-pair + micro-scheduling stop paying. +7. **No split-K** for M=4096; launcher is a pure dispatch on the caller + stream; packing/scales stay out of timing and Graph. +8. Recipes/constants do NOT transfer across shapes or M regimes: tile aspect + (down_proj 64x128 at K=256 vs q_b 128x128 at K=2048 vs o_proj 256x64 at + N=6144), occupancy, stage depth and pipeline form must be re-measured per + exact (M,N,K). + +### 4.7 Boundary gaps + +- 32 < M < 4096: **no accepted candidate measured** in this lineage. Reuse + the staged recipes above and measure per exact M before committing a route; + keep both a staged prefill arm and a decode-compatible arm until measured. +- M <= 32 (decode): **UNOPTIMIZED in this lineage** — do not claim any decode + speedup from this lineage, never infer success from the prefill evidence, + and route M <= 32 work to `int8-w8a8-gemm-decode` + + `int8-w8a8-gemm-foundations` (split-K partial kernels + fused combine, + µs-scale noise-tolerant acceptance). Decode conclusions do not cross M + regimes (the M=16 direct-global-operand win regressed 7.9x on M=4096 + gate_up in the sibling lineage). +- No plateau proven on any of the six shapes; continue per-shape HIP search + with the §7 protocol (one bounded mechanism per round, falsifiable + prediction, p90 guard vs current best, exact code-object verification). + Next candidates: fused_qkv_a's p90-failed 938.84 us load-all-then-MMAC-all + variant and kv_b's shadow mechanisms (stride skew / prefetch, ~138.2-138.8 + us) are within ~1% of the accepted best. Revalidate on the full model / + 8-card deployment before claiming production gains; never present the fixed + Triton baselines as measured or optimal. + +## 5. Boundary: 32 < M < 3072 (TP4) / 32 < M < 4096 (TP8 / GLM-5.2) + +No accepted candidate has been measured in (32, 3072) for TP4 or (32, 4096) +for TP8 or the GLM-5.2 lineage. Reuse the staged recipes above (2D M-tile, +64-K double-buffered stage, no split-K) and measure per exact M before +committing a route. Keep the M×N output-tile grid well above 120 CUs; if it +is not, split-K may become necessary (measure it — do not assume either way). +TP8's optimized evidence is exactly M=4096 (hy3 4/4 shapes, §3; GLM-5.2 6/6 +shapes, §4); do not extend those accepted numbers to other M values. + +## 6. Validated evidence (TP4) + +### M=3072 prefill microbenchmark + +Exact bf16 comparison passed for all five TP4 shapes, using the real +transposed Triton weight view and a preprocessed contiguous HIP weight: + +| Case | Triton ms | HIP ms | Speedup | +|---|---|---:|---:|---:| +| wqkv_a | 9.487898 | 0.857572 | 11.06x | +| wq_b | 15.437104 | 1.149457 | 13.43x | +| wo_b | 14.680174 | 1.102976 | 13.31x | +| shared gate_up_proj | 5.985047 | 0.623089 | 9.61x | +| shared down_proj | 4.198581 | 0.303484 | 13.83x | + +The very large gain is layout-sensitive. It must not be claimed for a Triton +baseline already using an equally optimized packed layout. + +The later wqkv_a ISA retune supersedes the `0.857572 ms` HIP row above for the +exact shape `(3072,1536,4096)`: the selected stage64 double-buffer prefetch +kernel measured `0.598983 ms` median and `0.605783 ms` P90 in an alternating +thermal-state trial, or `64.534 INT8 TOPS`, with exact bf16 output. A sustained +100-warmup/300-sample run measured `0.615996 ms` median and `0.621276 ms` P90, +showing device frequency/thermal drift that must be reported rather than +hidden by the best short run. + +### M=4096 chunked prefill (hy3_tp4_*_m4096) + +Accepted candidates per shape; baselines are the fixed user-supplied Triton +table in §2. Numbers below are the **2026-08-22 re-verified values**; the +iteration chains are the historical in-run measurements of the same accepted +kernels. + +**hy3_tp4_qkv_proj_m4096 (worker_0)** — final accepted iter 18 (kernel +`w8a8_dumma_prefill_64x128_packedb_kernel`): 64x128 tile, 256 threads = 4 +waves, 32x64 quadrant/wave = 8 m16n16k32 int32 accumulators; grid +(N/128)x(M/64) = 20x64 = 1280 blocks; single-buffered 64-K stage with 2 +`__syncthreads`/stage; A staged row-major stride 68 B, B staged n-major stride +80 B; LDS 14,592 B/block, **3 blocks/CU** (arch_vgpr 80); one-time n-major +pack `packed[n][k]` for (k,n)==(4096,2560) outside the timed region; coalesced +4-bf16-per-lane epilogue stores. Accepted chain (us): 1296.50 (iter 1) → +1108.58 (iter 2, double-buffered 64-K) → 880.99 (iter 3, single-buffer +control) → 770.14 (iter 5, n-major packed B + 3 blocks/CU) → 721.87 (iter 6, +coalesced epilogue stores) → **756.73** (iter 18, explicit A-fragment loader +replacing the library `du_load_matrix_sync` row_major byte-reassembly VALU). +Plateau NOT proven (iter 19's reason: require 8 valid HIP rounds within ±2% of +best — continue HIP-only work). + +**hy3_tp4_o_proj_m4096 (worker_1)** — final accepted iter 10 (kernel +`w8a8_dumma_128x64x64_packed_kernel`): weight packed once into `[N/64, K, 64]` +int8 panels (each 64-K stage = one coalesced 4 KiB stream); 128x64 tile, 256 +threads = 4 waves, 64x32 quadrant/wave = eight m16n16k32 int32 accumulators; +grid dim3(64,32) = 2048 blocks; 64-K double-buffered LDS, A and B both +80-byte row strides (16-byte-aligned, five bank phases), 30,720 B/block, 2 +blocks/CU; loop-carried int4 register payload published into the idle buffer +at the **top** of the stage (before the MMAC burst) with global prefetch one +full extra stage ahead; one `__syncthreads`/stage; fused fragment→scale→bf16 +epilogue. Accepted chain (us): 981.53 (iter 1) → 834.14 (iter 3, A LDS stride +64→80; PMC: conflicts 77.6M→31.5M, 8-way→2-way) → **803.90** (iter 10, +stage-top publish). + +**hy3_tp4_shared_gate_up_proj_m4096 (worker_2)** — final accepted iter 13 +(`w8a8_gemm_prefill_tiled_kernel<64,128,64,2>`): 64x128 tile, 512 threads = 8 +waves, 32x32 quadrant/wave = 4 accumulators; grid dim3(6,64) = 384 blocks; +K stage 64 double-buffered, one barrier/stage (65/block), 30,720 B LDS, 2 +blocks/CU; `pack_weight` transposes [K,N]→[N,K] for this shape only (outside +timed region); B staged [N,K], `col_major` fragments; iter 13's +`load_fragment8` writes the 8 consecutive bytes straight into fragment +storage (one ds_read2_b64 per fragment), removing the per-byte mask/OR +reassembly (VALU 16.08M→6.25M). Accepted chain (us): 436.21 (iter 1) → +392.58 (iter 3, double-buffered K) → 280.05 (iter 10, [N,K] pack + col_major) +→ **220.71** (iter 13, load_fragment8). Bottleneck chain was LDS-wait-bound +throughout: bank-skew → software pipeline → B-fragment aliasing → +byte-reassembly VALU. + +**hy3_tp4_shared_down_proj_m4096 (worker_3)** — final accepted iter 14 +(`w8a8_dumma_prefill_kernel<128,64,64>`): 128x64 tile, 512 threads = 8 waves, +32x32 quadrant/wave; grid (N/64, M/128) = 2048 blocks; K stage 64 +double-buffered, 29,696 B/block (A stride 80, B stride 72), **2 blocks/CU at +57 VGPR — first clean 2-block residency at constant tile** (VGPR-bound; +`__launch_bounds__` minBlocks was ignored at 74 VGPR); iter 14 packs weight +**n-major `packed[n][k]`** so B staging is 2× ds_write_b64/thread (was 16× +ds_write_b8 + ~12 VALU extracts); kk body 2×-unrolled with two independent +fragment sets, all 8 ds_read2 before the 8 v_mmac. Accepted chain (us): +266.74 → 259.23 (double-buffer + VGPR prefetch) → 255.46 (n-major B in LDS + +col_major) → 245.29 (stage 32→64) → 236.39 (unrolled kk) → **214.84** (n-major +pack). This lineage ended `plateau=false` — the accepted 214.84 us is the best +found, not a proven optimum. + +## 7. Prefill acceptance protocol (M=4096 lineages) + +Unprofiled CUDA-graph replay, 100 warmup / 30 samples × 100 replays, hot +cache; report median AND p90 (p90 guard vs current best), never min; compare +Graph vs Graph only; re-verified runs must have correctness ON (mismatch=0, +max_abs_error=0.0). PMC (`hipprof --pmc --pmc-type 3`) explains results after +acceptance and must move the predicted counters, else the mechanism is +ambiguous. Report thermal/frequency drift explicitly (see the M=3072 +wqkv_a retune) rather than hiding it with the best short run. + +Additionally for M=4096: + +- Exact bf16 vs reference (mismatch 0); preserve the exact int32 accumulation + order (k0-outer, kk-inner) and the element-to-slot fragment mapping so + variants are bit-identical. +- Graph capture + changed input contents + replay must pass; after any change, + re-capture the Graph with changed input contents and replay again. A Graph + replay may update contents, never addresses/shapes. M tails and the scalar + fallback must still pass. +- Verify resources from the exact code object + launch record (arch_vgpr vs + code-object vgpr_count, LDS bytes, scratch/spills, grid/workgroup); + `__launch_bounds__` minBlocks is a hint and can be ignored (down_proj + iter 7). Occupancy math: blockDim multiple of 64; 2-3 blocks × LDS ≤ 64 KiB; + VGPR × threads × blocks ≤ 512 KiB. +- Use the pre-seeded CPU int64 exact reference (it is slow); never run + correctness without the reference cache present. +- Report thermal/frequency drift (max samples ~1.5x median are common; + implied clock ~1.45 GHz) instead of hiding it with the best short run. +- `lds_wait ~= lds_instructions` alone is not a bottleneck signal (true in + accepted kernels too); profiled time is never the score — PMC must move the + predicted counters. + +TP8-specific notes (hy3_tp8_*_m4096 and GLM-5.2 glm_tp8_*_m4096): + +- Restore the accepted source to its recorded digest before the next + experiment; make ONE bounded mechanism per round with a falsifiable + prediction (expected direction + expected PMC/ISA deltas; a >2% regression + falsifies the mechanism). +- Verify the exact code object of the timed run (digest-matched, never a + stale sibling symbol); audit the exact template symbol for load -> MMAC -> + wait -> LDS-store order, VGPR/LDS/scratch/spills, no removed required + loads, no spilled prefetches. `__launch_bounds__` minBlocks is a hint, not + a contract. +- Report thermal/frequency drift: max samples ~1.5-2.1x median are routine + (GLM-5.2 lineage), not just the ~1.5x seen on TP4. +- Continue per-shape HIP search — no TP8 shape has a proven optimum (qkv's + plateau rule of 8 valid HIP rounds within ±2% was never met; down_proj's + best is best-found; all six GLM-5.2 shapes are `plateau=false`). + +## 8. Rejected / guardrail evidence (TP4 prefill) + +TP8 M=4096 rejected ideas live in §3.4. TP4 rejected evidence: + +**M=4096 (hy3_tp4_*_m4096) rejected ideas** (measured, correct where noted — +do not repeat without new evidence): + +- qkv: 128x64 tile flip (1152.10); one-stage-earlier global prefetch (725.77, + −0.54% — global-load latency already hidden by co-residency); minBlocks=4 / + 72→64 VGPR squeeze (749.80 — 4 spill slots, confounded); 512-thread 8-wave + blocks (733.01, −1.52%); B stride 80→72 conflict fix (752.79 — B conflicts + sit at the ~4-cycle 32-bank pigeonhole floor); A stride 68→80 (732.59); + sync-group retile (920.53); iter 19 accepted-build variant (713.84, p90 + guard failed). +- o_proj: 128x128 tile/1 block per CU (840.44, loses 2nd resident block); + 64x128 tile flip (959.38); register-rotated kk fragment prefetch (845.00); + kStageK 64→32 / 3 blocks per CU (866.60); four-A-fragment WAR removal + (830.13, flat); per-wave staging with zero barriers (967.17 — duplicated + staging traffic swamps barrier savings); dual loop-carried payloads / + two-stage lead (827.76, flat); eight A fragments at stage top (872.61). +- gate_up: direct-A arm with B-only staging (3444.26 — 7.9x slower; the + M=16 decode kernel's direct-B-from-global win does **not** transfer across M + regimes); 128x128 tile / 1024 threads (519.86, 5 VGPR spill slots + + 192-block grid); kStage 64→32 / 3 blocks per CU twice (407.85, 413.60 — + more occupancy with phase-locked co-resident blocks does not pay). +- down_proj: K stage 64→128 (363.76 — stage-depth axis saturates at 64); + 64x64 active tile occupancy probe (270.34 — confounded: shrunk tile doubles + per-output B staging); B-staging "register diet" (304.22 — flat batched + staging issue order is load-bearing; the restructure caused scratch spills). +- Common pattern: these kernels are LDS-latency/issue-bound, not HBM-bound; + barrier removal, occupancy changes, and deeper register prefetch all failed + once LDS strides and pack layout were fixed. Do not re-litigate without + profiling the exact shape. +- Killed rounds produced no candidate (`compile_or_agent_failure`: agent + killed ~900 s, missing proposal.json, restore/resume failures): qkv + 8,9,10,15,17,20; o_proj 2,5,9,11,12,13,18,19; gate_up 5-7,9,11,14,15; + down_proj 3,4,6,8,13,15. These are **infra/agent failures and prove + nothing** — not that DUMMA, INT8, a HIP API, templates, or inline asm are + unsupported. After a killed round, restore the accepted source to its + recorded digest before the next experiment. +- **LDS bank-skew rule**: pad row strides to 16-byte-aligned non-power-of-two + values (e.g., 80 = 64+16, 72 = 64+8, 68 = 64+4, 144 = 128+16); strides ≡ 0 + (mod 128 B) alias every row onto one bank phase (up to 8-16-way fragment + conflicts). +- For compute prefill, report unique API bytes and tiled global requests + separately. The M=3072 winner has only about 47.30 GB/s of unique API bytes + but about 1.024 TB/s of tiled A/B+output requests; cache can serve repeated + tile traffic, so neither number alone is physical HBM bandwidth. +- Do not assume a deeper Graph helps large-M compute. For the M=3072 winner, + one node beat 4/8/16-node cold-ring Graphs because launch overhead was + already negligible relative to the roughly 0.6 ms kernel. +- M=4096 baselines are the fixed user-supplied Triton table values; they were + not re-measured and must not be called optimal. All four shapes now have + accepted candidates; a future shape whose worker timed out or produced no + accepted candidate stays unoptimized — never infer a speedup or optimality + for it. + +## 9. Next optimization questions (prefill) + +- Prefill shapes other than wqkv_a: repeat the stage32/64/128 and single/double + LDS search; do not copy the wqkv_a stage64 dispatch without measurement. +- Prefill tile search: compare 64x64 against larger/smaller tiles while + preserving enough grid blocks and auditing VGPR/LDS/scratch in the exact + code object. +- Measure hot and production-like cache behavior separately. +- Revalidate on the full model/8-card deployment before claiming production + DeepSeek-V4 gains. +- M=4096 per-shape search should continue: **no proven optimum** in either + lineage — the control plane ended `plateau=false` for TP4 qkv and + down_proj, o_proj/gate_up explicitly disclaim optimality, and no TP8 shape + established a plateau. Accepted numbers are the best found, not proven + ceilings. For TP4 qkv, iter 19's plateau rule (8 valid HIP rounds within + ±2% of best) was not met — continue HIP-only work on that shape. TP4 + down_proj's 214.84 us is the best found, not a proven optimum; TP8 qkv + (iter 6), o_proj (iter 12), gate_up (iter 18), down_proj (iter 10) are + likewise best-found. +- TP8 (hy3_tp8_*_m4096): continue per-shape HIP search with the §7 protocol + (one bounded mechanism per round, falsifiable prediction); boundary + 32 < M < 4096: reuse the staged recipes and measure per exact M before + committing a route. +- GLM-5.2 TP8 M=4096 (glm5-2-dsh-tp8-m4096-1, §4): continue per-shape HIP + search on all six shapes with the §7 protocol (no plateau proven) — e.g. + fused_qkv_a's p90-failed 938.84 us load-all-then-MMAC-all variant and + kv_b's shadow mechanisms (stride skew / prefetch, ~138.2-138.8 us) are + within ~1% of the accepted best. GLM-5.2 decode (M <= 32) is unoptimized in + that lineage: run the decode skill's split-K + fused-combine family on the + GLM-5.2 TP8 shapes. Revalidate on the full model / 8-card deployment before + claiming production gains. +- M in (32, 3072): reuse the staged recipes and measure (see §5). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-prefill/references/wqkv-a-m3072-prefetch-isa.md b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-prefill/references/wqkv-a-m3072-prefetch-isa.md new file mode 100644 index 00000000..d4c75624 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-gemm-prefill/references/wqkv-a-m3072-prefetch-isa.md @@ -0,0 +1,144 @@ +# wqkv_a M=3072: DUMMA prefetch and ISA evidence + +## Contents + +- [Scope and result](#scope-and-result) +- [Selected pipeline](#selected-pipeline) +- [Why stage64 won](#why-stage64-won) +- [ISA audit](#isa-audit) +- [Bandwidth accounting](#bandwidth-accounting) +- [Reproduction and artifacts](#reproduction-and-artifacts) + +## Scope and result + +Validated on worker29, K500SM_AI/gfx928, DTK 26.04: + +```text +M=3072, K=4096, N=1536 +int8 A x int8 B -> int32 -> x_scale * weight_scale -> bf16 +``` + +Every accepted candidate matched the full bf16 torch float32-GEMM reference +exactly and passed HIP Graph replay. The selected result was: + +```text +tile 64x64 +workgroup 256 threads / four wave64 +stage-K 64 +LDS 16 KiB, two alternating 8 KiB buffers +VGPR / SGPR 63 / 43 +scratch / spills 0 / 0 +median / P90 598.983 / 605.783 us +INT8 throughput 64.534 TOPS +``` + +A longer 100-warmup/300-sample run measured 615.996 us median and 621.276 us +P90. Preserve both the alternating comparison and sustained result; the gap +is thermal/frequency drift, not a correctness change. + +## Selected pipeline + +Each block owns a 64x64 output tile. Four waves each own one 32x32 quadrant, +represented by four m16n16 accumulators. For every stage: + +```text +issue next A/B global_load_dwordx4 into VGPR +load current A/B fragments from current LDS buffer +execute eight v_mmac_i32_16x16x32_i8 per wave +s_waitcnt vmcnt(0) +write prefetched int4 vectors to alternate LDS with ds_write_b128 +s_waitcnt lgkmcnt(0) +s_barrier +swap current/alternate LDS buffers +``` + +Each thread carries two 16-byte next-stage vectors for stage64. Keep those +registers live only across the current MMAC group; extending their lifetime +raises VGPR pressure without improving overlap. + +## Why stage64 won + +| Candidate | Median us | P90 us | INT8 TOPS | Diagnosis | +|---|---:|---:|---:|---| +| original stage128, single LDS | 841.748 | 844.148 | 45.922 | no VMEM/MMAC overlap | +| double LDS stage32 | 819.212 | 832.112 | 47.185 | twice as many stage barriers | +| double LDS stage64 | 602.054 | 610.554 | 64.205 | best latency/occupancy balance | +| double LDS stage128 | 883.311 | 885.871 | 43.761 | 85 VGPR + 32 KiB LDS pressure | +| stage64 + raw LDS-ready barrier | 598.983 | 605.783 | 64.534 | selected | + +Do not infer that shallower stages are always better. stage32 reduced LDS and +prefetch registers but doubled synchronization frequency. stage128 doubled +work per barrier but expanded LDS and live prefetch state enough to regress. + +One-node Graph was best. Ordinary stage64 measured approximately +603.035/610.055/613.035/614.690 us per call at Graph depths 1/4/8/16. + +## ISA audit + +Use the code object extracted from the same JIT extension that produced the +timing. Match the exact symbol containing `ILi64ELb1E`; `Li64` is stage64 and +`Lb1` is the accepted raw-ready-barrier specialization. + +Required steady-state evidence: + +1. A next-stage `global_load_dwordx4` appears before the current MMAC group. +2. Eight `v_mmac_i32_16x16x32_i8` sites appear before the load's + `s_waitcnt vmcnt(0)`. +3. `ds_write_b128` consumes the prefetched VGPR tuple after the VMEM wait. +4. `s_waitcnt lgkmcnt(0); s_barrier` protects the LDS-ready boundary. +5. Metadata reports 63 VGPR, 43 SGPR, 16,384 bytes LDS, zero private segment, + and zero spills. + +Reject source-only claims. In the pure-memory precursor, LLVM commoned two +nominal A reads; timing looked faster, but ISA showed only five loads instead +of the required six. An opaque VGPR index was needed to retain the second +ordinary GLOBAL load. The general rule is to count required memory sites and +inspect their cache modifiers before accepting bytes/time. + +Use raw asm only after proving a compiler limitation or synchronization +redundancy. Do not hand-write raw global memory or raw MMAC here: tuple +constraints, EXEC/VCC/SCC effects, cache flags, and fragment lane ownership +must be independently verified. + +## Bandwidth accounting + +Keep these metrics separate: + +```text +unique API bytes/call 28,329,984 +tiled A/B requests 603,979,776 +output stores 9,437,184 +tiled operand+output request bytes 613,416,960 +``` + +At 598.983 us, unique API bytes give about 47.30 GB/s, while tiled +operand+output requests give about 1.024 TB/s. The first reflects useful tensor +footprint; the second shows that the global request pipeline remains busy +while MMAC executes. Repeated tile requests can hit cache, so neither is a +physical-HBM claim without PMC plus a cache-cold protocol. + +The pure-memory M=3072 proxy reached 805.651 GB/s using 56,659,968 explicit +read+write bytes per call and a 16-node cold-ring Graph. Do not compare that +copy-program metric directly with real GEMM's unique-byte metric. + +## Reproduction and artifacts + +```bash +cd /workspace/ISA_test_codex/wqkv_a_compute_tuning +HIP_VISIBLE_DEVICES=0 python benchmark_prefill_m3072.py \ + --stage-k=65 --warmups=50 --samples=150 --nodes=1 --check +``` + +`stage-k=65` is only the retained tuning identifier for stage64 plus the raw +LDS-ready barrier. The public `gemm_out_prefill(...)` dispatches the same +kernel. + +Artifacts: + +```text +/workspace/ISA_test_codex/wqkv_a_compute_tuning/w8a8_gemm_compute.hip +/workspace/ISA_test_codex/wqkv_a_compute_tuning/M3072_RESULTS.md +/workspace/ISA_test_codex/wqkv_a_compute_tuning/wqkv_a_m3072_compute_gfx928.co +/workspace/ISA_test_codex/wqkv_a_compute_tuning/wqkv_a_m3072_compute_stage64_raw.isa +/workspace/ISA_test_codex/wqkv_a_compute_tuning/wqkv_a_m3072_compute.notes +``` diff --git a/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-quantized-gemm-optimization/SKILL.md b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-quantized-gemm-optimization/SKILL.md new file mode 100644 index 00000000..2df82b28 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/skills/int8-w8a8-quantized-gemm-optimization/SKILL.md @@ -0,0 +1,38 @@ +--- +name: int8-w8a8-quantized-gemm-optimization +description: > + Router for the INT8 W8A8 GEMM skill family on Hygon K500SM_AI/gfx928. + This skill is split by execution phase into three skills — load + int8-w8a8-gemm-foundations plus the phase skill: int8-w8a8-gemm-decode + (M<=32) or int8-w8a8-gemm-prefill (M>32). This router exists for backward + compatibility; it carries no tuning content itself. +--- + +# INT8 W8A8 GEMM — skill family router + +The former single "int8-w8a8-quantized-gemm-optimization" skill was split by +execution phase (2026-08-23) because decode and prefill have non-transferable +recipes, different bottlenecks, and different acceptance protocols. Load the +following instead: + +| M | Load | +|---|---| +| M <= 32 (decode) | `int8-w8a8-gemm-decode` + `int8-w8a8-gemm-foundations` | +| M > 32 (prefill) | `int8-w8a8-gemm-prefill` + `int8-w8a8-gemm-foundations` | + +- **int8-w8a8-gemm-foundations** — operator/layout contract, DUMMA and + epilogue rules, Graph-safe PyTorch interface, SGLang integration, fair + benchmarking, shared correctness gates and guardrails, E2E evidence. +- **int8-w8a8-gemm-decode** — split-K partial kernels + fused combine, + M=1..16 validated recipes, M=16 decode, M > 32 -> prefill**. M in (32, 128] +(decode side) and (32, 3072) (prefill side) are measured-boundary gaps — route +by measurement and keep both an arm per regime until measured. + +No tuning content lives in this router; always load the foundations skill plus +the matching phase skill. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js index d2f29238..43b942bb 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js @@ -877,6 +877,19 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) { ? "更新 variant" : "加入 variant"} + ${(() => { + const v = variantIndex ? variantIndex.get(shape) : null; + if (!v) { + return html`
当前: 无
`; + } + if (v.speedup != null) { + return html`
当前: 超基线 ${metric(v.speedup, 2)}x
`; + } + if (v.median_us != null) { + return html`
当前: ${metric(v.median_us)} µs
`; + } + return html`
当前: —
`; + })()} `)} 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 index c40b4da7..57fffc1a 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-shape-input.js +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-shape-input.js @@ -439,6 +439,24 @@ function ShapeCatalog(_a) { `; } +// Quick select/deselect helper shown above subset (Selected shapes only) +// catalogs. A subset task starts with NOTHING selected on purpose: the user +// ticks the shapes that should enter this task (or hits "Select all"). +function SubsetActions(_a) { + var count = _a.count, total = _a.total, onAll = _a.onAll, onNone = _a.onNone; + return html` +
+ ${count} / ${total} selected + + + Only ticked shapes enter this task. +
+ ${count === 0 ? html`

No shape selected — the task starts empty. Tick at least one shape above (or Select all).

` : null} + `; +} + function ApiDefaultsPreview(_a) { var catalog = _a.catalog; var tpSizes = catalog @@ -502,7 +520,9 @@ function AiShapeSubset(_a) { var catalog = modelCatalog(model); var _b = useState(function () { var parsed = parseShapeRecords(value || ""); - return selectedShapes(parsed, true, catalog, model); + // Start EMPTY (no fallback to the full catalog): only shapes the user + // ticks (or Select-all) enter this task. + return selectedShapes(parsed, false, catalog, model); }), selected = _b[0], setSelected = _b[1]; var emit = useCallback(function (next) { @@ -515,6 +535,10 @@ function AiShapeSubset(_a) { useEffect(function () { emit(selected); }, []); + var count = catalog.filter(function (shape) { + return selected[shape.id]; + }).length; + return html`
@@ -523,6 +547,15 @@ function AiShapeSubset(_a) { shapes stay on the trusted fallback and are checked again before the candidate is published.
+ <${SubsetActions} + count=${count} + total=${catalog.length} + onAll=${function () { + var next = {}; + catalog.forEach(function (shape) { next[shape.id] = true; }); + emit(next); + }} + onNone=${function () { emit({}); }} /> <${ShapeCatalog} shapes=${catalog} selected=${selected} @@ -573,7 +606,8 @@ function ManualGpuAssignment(_a) { var _b = useState(function () { var parsedShapes = parseShapeRecords(value || ""); return subset - ? selectedShapes(parsedShapes, true, catalog, model) + // Start EMPTY for a manual subset: only ticked shapes enter this task. + ? selectedShapes(parsedShapes, false, catalog, model) : selectedMap(catalog.map(function (shape) { return shape.id; })); @@ -620,6 +654,21 @@ function ManualGpuAssignment(_a) { appears exactly once; empty GPU cards are allowed.
${subset && html` + <${SubsetActions} + count=${catalog.filter(function (s) { return selected[s.id]; }).length} + total=${catalog.length} + onAll=${function () { + var nextSelected = {}; + var nextOwners = Object.assign({}, owners); + catalog.forEach(function (shape) { + nextSelected[shape.id] = true; + if (nextOwners[shape.id] == null) { + nextOwners[shape.id] = operatorGroupedOwners(catalog)[shape.id]; + } + }); + emitState(nextOwners, nextSelected); + }} + onNone=${function () { emitState(owners, {}); }} /> <${ShapeCatalog} shapes=${catalog} selected=${selected} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css index 55f7e00f..94b628aa 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css @@ -907,3 +907,38 @@ color: var(--muted); font-size: var(--fs-xs); } + +/* Subset (Selected shapes only): quick select-all / clear bar */ +.dkao-subset-actions { + display: flex; + align-items: center; + gap: 10px; + flex-wrap: wrap; + margin: 10px 0 4px; + padding: 8px 10px; + border: 1px dashed var(--border); + border-radius: var(--radius-sm); + background: var(--bg-elev); +} +.dkao-subset-count { + font-weight: 700; + color: var(--fg); + font-size: var(--fs-xs); +} +.dkao-subset-actions .dkao-shape-tab { + padding: 4px 12px; + border: 1px solid var(--border); + border-radius: var(--radius-sm); + background: var(--panel-2); + color: var(--fg); + cursor: pointer; + font-size: var(--fs-xs); +} +.dkao-subset-actions .dkao-shape-tab:disabled { + opacity: 0.5; + cursor: not-allowed; +} +.dkao-subset-hint { + color: var(--muted); + font-size: var(--fs-xxs); +} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gate_policy.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gate_policy.py new file mode 100644 index 00000000..8cb57b72 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gate_policy.py @@ -0,0 +1,158 @@ +"""gates.yaml is wired: the runtime reads it, the manifest decides if it applies.""" + +from __future__ import annotations + +import shutil + +import pytest +import yaml + +from ..orchestrator import gate_policy as gp +from ..orchestrator import planner +from ..orchestrator.harness_io import default_harness_dir + + +@pytest.fixture(autouse=True) +def _clear_cache(): + gp.reset_cache() + yield + gp.reset_cache() + + +def _workspace(tmp_path, *, gates_wired=True, policy_wired=True): + root = tmp_path / "harness" + shutil.copytree(default_harness_dir(), root) + manifest = yaml.safe_load((root / "manifest.yaml").read_text(encoding="utf-8")) + manifest["components"]["gates"]["wired"] = gates_wired + manifest["components"]["planner_policy"]["wired"] = policy_wired + (root / "manifest.yaml").write_text(yaml.safe_dump(manifest, sort_keys=False), + encoding="utf-8") + return root + + +def _edit_gates(root, **values): + data = yaml.safe_load((root / "gates.yaml").read_text(encoding="utf-8")) + for key, value in values.items(): + node = data["gates"] + parts = key.split(".") + for part in parts[:-1]: + node = node.setdefault(part, {}) + node[parts[-1]] = value + (root / "gates.yaml").write_text(yaml.safe_dump(data, sort_keys=False), + encoding="utf-8") + + +def test_seed_matches_builtin_defaults(): + """The seed file mirrors the historical constants (no surprise changes).""" + seed = yaml.safe_load( + (default_harness_dir() / "gates.yaml").read_text(encoding="utf-8") + )["gates"] + defaults = gp.builtin_gates() + assert seed["round_acceptance_improvement_percent"] == ( + defaults["round_acceptance_improvement_percent"]) + assert seed["shadow"]["min_improvement_percent"] == ( + defaults["shadow"]["min_improvement_percent"]) + assert seed["plateau"]["max_regression_percent"] == ( + defaults["plateau"]["max_regression_percent"]) + assert seed["isa_gate"]["required_valid_isa_guided_rounds"] == ( + defaults["isa_gate"]["required_valid_isa_guided_rounds"]) + assert seed["task_budget"]["default_max_iterations"] == ( + defaults["task_budget"]["default_max_iterations"]) + + +def test_editing_gates_changes_runtime_values(tmp_path, monkeypatch): + root = _workspace(tmp_path) + _edit_gates(root, **{ + "round_acceptance_improvement_percent": 2.5, + "p90_guard": "1.05", + "shadow.enabled": False, + "shadow.min_improvement_percent": 0.7, + "plateau.max_regression_percent": 3.5, + "isa_gate.required_valid_isa_guided_rounds": 4, + "task_budget.default_max_iterations": 6, + }) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + + assert gp.round_acceptance_improvement_percent() == 2.5 + assert gp.p90_tolerance() == 1.05 + assert gp.shadow_enabled() is False + assert gp.shadow_min_improvement_percent() == 0.7 + assert gp.plateau_max_regression_percent() == 3.5 + assert gp.isa_required_valid_rounds() == 4 + assert gp.default_max_iterations() == 6 + # the pipeline helpers see the same values + from ..orchestrator import w8a8_pipeline as W + assert W._gates.round_acceptance_improvement_percent() == 2.5 + + +def test_unwired_gates_component_is_ignored(tmp_path, monkeypatch): + """wired:false means the file is inert — the manifest is meaningful.""" + root = _workspace(tmp_path, gates_wired=False) + _edit_gates(root, **{"round_acceptance_improvement_percent": 9.0}) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + assert gp.round_acceptance_improvement_percent() == 1.0 + assert gp.component_wired("gates") is False + + +def test_unwired_planner_policy_falls_back_to_builtin(tmp_path, monkeypatch): + root = _workspace(tmp_path, policy_wired=False) + data = yaml.safe_load((root / "planner_policy.yaml").read_text(encoding="utf-8")) + data["fallback"]["m16"] = ["consolidate"] * 8 + (root / "planner_policy.yaml").write_text(yaml.safe_dump(data, sort_keys=False), + encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + effective = planner.policy() + assert effective["fallback"]["m16"] != ["consolidate"] * 8 + assert effective["fallback"] == planner._builtin_policy()["fallback"] + + +def test_required_hip_rounds_rule_and_literal(tmp_path, monkeypatch): + root = _workspace(tmp_path) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + assert gp.required_hip_rounds(10) == 8 # documented expression + _edit_gates(root, **{"isa_gate.required_hip_rounds_rule": 3}) + gp.reset_cache() + assert gp.required_hip_rounds(10) == 3 # literal override + + +def test_p90_guard_mode_parsing(tmp_path, monkeypatch): + root = _workspace(tmp_path) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + assert gp.p90_guard_mode() == "no_p90_regression" + assert gp.p90_tolerance() == 1.0 + _edit_gates(root, **{"p90_guard": "not-a-number"}) + gp.reset_cache() + assert gp.p90_tolerance() == 1.0 # never loosens by accident + + +def test_snapshot_records_what_a_run_used(tmp_path, monkeypatch): + """Children record their effective gates so the mechanism gate can verify.""" + import json + + root = _workspace(tmp_path) + _edit_gates(root, **{"round_acceptance_improvement_percent": 3.25}) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + state = tmp_path / "state" + out = gp.snapshot(state) + assert out is not None and out.is_file() + payload = json.loads(out.read_text(encoding="utf-8")) + assert payload["gates"]["round_acceptance_improvement_percent"] == 3.25 + assert payload["gates_wired"] is True + assert payload["harness_root"] == str(root) + + +def test_snapshot_marks_unwired_gates(tmp_path, monkeypatch): + import json + + root = _workspace(tmp_path, gates_wired=False) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + out = gp.snapshot(tmp_path / "state2") + payload = json.loads(out.read_text(encoding="utf-8")) + assert payload["gates_wired"] is False + assert payload["gates"]["round_acceptance_improvement_percent"] == 1.0 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py new file mode 100644 index 00000000..f3f0b80c --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py @@ -0,0 +1,121 @@ +"""GPU admission pre-check: usage rule + pollution marking.""" + +from __future__ import annotations + +import pytest + +from ..orchestrator import gpu_preflight as gp + +SMI_SAMPLE = """ +================================= System Management Interface ================================== +================================================================================================ +HCU Temp AvgPwr Perf PwrCap VRAM% HCU% Mode +0 51.0C 145.0W manual 300.0W 0% 0.0% Normal +1 49.0C 147.0W manual 300.0W 12% 4.5% Normal +2 57.0C 251.0W manual 300.0W 40% 97.0% Normal +3 58.0C 254.0W manual 300.0W 92% 81.7% Normal +======================================== End of SMI Log ========================================= +""" + + +def test_parse_hy_smi_rows(): + parsed = gp.parse_hy_smi(SMI_SAMPLE) + assert set(parsed) == {0, 1, 2, 3} + assert parsed[0]["vram_percent"] == 0.0 + assert parsed[2]["util_percent"] == 97.0 + assert parsed[3]["power_w"] == 254.0 + assert parsed[1]["temp_c"] == 49.0 + + +def test_usability_is_free_vram_not_percentage(): + """The rule is `free >= 4GB`, so a 91%-used card is still usable.""" + ok = gp.check_gpu(0, {"vram_percent": 91.0, "util_percent": 0.0}, + total_gb=64.0) + too_full = gp.check_gpu(1, {"vram_percent": 94.0, "util_percent": 0.0}, + total_gb=64.0) + assert ok["usable"] is True and ok["free_gb"] == 5.76 + assert ok["measurement_suspect"] is False + assert too_full["usable"] is False and too_full["free_gb"] == 3.84 + assert any("VRAM free" in r for r in too_full["reasons"]) + # a full device is a capacity problem, not a pollution problem + assert too_full["measurement_suspect"] is False + assert gp.DEFAULT_MIN_FREE_GB == 4.0 + + +def test_busy_marks_suspect_without_blocking(): + busy = gp.check_gpu(2, {"vram_percent": 40.0, "util_percent": 97.0, + "power_w": 251.0}, foreign_pids=[12345]) + assert busy["usable"] is True # usage is allowed (VRAM ok) + assert busy["measurement_suspect"] is True + assert any("busy" in r for r in busy["reasons"]) + assert any("power" in r for r in busy["reasons"]) + # foreign pids are informational (idle holders are harmless) + assert any("foreign" in n for n in busy["notes"]) + + +def test_idle_foreign_holders_are_not_suspicion(): + """A device at 0% util / idle power is clean even if other containers hold it.""" + idle = gp.check_gpu(0, {"vram_percent": 0.0, "util_percent": 0.0, + "power_w": 144.0}, foreign_pids=[1, 2, 3]) + assert idle["usable"] is True + assert idle["measurement_suspect"] is False + assert idle["reasons"] == [] + assert any("foreign" in n for n in idle["notes"]) + + +def test_preflight_prefers_clean_then_suspect_then_over_limit(monkeypatch): + monkeypatch.setattr(gp, "sample_gpu_state", + lambda ids=None, **kw: { + 0: {"vram_percent": 1.0, "util_percent": 0.0}, + 1: {"vram_percent": 5.0, "util_percent": 55.0, + "power_w": 250.0}, + 2: {"vram_percent": 30.0, "util_percent": 3.0}, + 3: {"vram_percent": 97.0, "util_percent": 10.0}, + }) + monkeypatch.setattr(gp, "foreign_kfd_pids", lambda *a, **k: []) + monkeypatch.setattr(gp, "total_vram_gb", lambda *a, **k: 64.0) + plan = gp.preflight_gpus([0, 1, 2, 3]) + assert plan["clean_ids"] == [0, 2] + assert plan["suspect_ids"] == [1] + assert plan["over_limit_ids"] == [3] + # usable devices first (clean before suspect); the over-limit device drops + # out while any device is within the limit + assert plan["preferred"] == [0, 2, 1] + assert 3 not in plan["preferred"] + + +def test_over_limit_devices_return_when_nothing_else_is_usable(monkeypatch): + monkeypatch.setattr(gp, "sample_gpu_state", + lambda ids=None, **kw: { + 0: {"vram_percent": 97.0, "util_percent": 5.0}, + 1: {"vram_percent": 99.0, "util_percent": 5.0}, + }) + monkeypatch.setattr(gp, "foreign_kfd_pids", lambda *a, **k: []) + monkeypatch.setattr(gp, "total_vram_gb", lambda *a, **k: 64.0) + plan = gp.preflight_gpus([0, 1]) + assert plan["preferred"] == [0, 1] + + +def test_foreign_kfd_pids_detects_invisible_pids(tmp_path): + (tmp_path / "1").write_text("", encoding="utf-8") # our own pid + (tmp_path / "99999999").write_text("", encoding="utf-8") # foreign + pids = gp.foreign_kfd_pids(tmp_path) + assert 99999999 in pids + assert 1 not in pids + + +def test_preflight_enabled_env_and_answers(monkeypatch): + monkeypatch.delenv("METAINFER_GPU_PREFLIGHT", raising=False) + assert gp.preflight_enabled({}) is True # default on + assert gp.preflight_enabled({"gpu_preflight": "false"}) is False + monkeypatch.setenv("METAINFER_GPU_PREFLIGHT", "1") + assert gp.preflight_enabled({"gpu_preflight": "false"}) is True + monkeypatch.setenv("METAINFER_GPU_PREFLIGHT", "off") + assert gp.preflight_enabled({}) is False + + +def test_disabled_preflight_returns_default_order(): + plan = gp.preflight_gpus([0, 1, 2, 3], enabled=False) + assert plan["enabled"] is False + assert plan["preferred"] == [0, 1, 2, 3] + assert plan["gpus"] == {} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_harness_io.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_harness_io.py new file mode 100644 index 00000000..013375af --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_harness_io.py @@ -0,0 +1,105 @@ +"""Tests for orchestrator.harness_io (M1 slice 1: harness seed + reader).""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import yaml + +from ..orchestrator import w8a8_pipeline +from ..orchestrator.config import ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT +from ..orchestrator.harness_io import ( + default_harness_dir, + harness_root, + load_gates, + load_manifest, + plugin_dir, + seed_workspace, +) + + +def test_plugin_dir_and_default_seed_exist(): + root = plugin_dir() + assert root.name == "dcu_kernel_auto_opt" + seed = default_harness_dir() + assert seed.is_dir() + assert (seed / "manifest.yaml").is_file() + assert (seed / "gates.yaml").is_file() + + +def test_harness_root_defaults_to_seed(monkeypatch): + monkeypatch.delenv("METAINFER_HARNESS_ROOT", raising=False) + assert harness_root() == default_harness_dir().resolve() + + +def test_harness_root_env_override(monkeypatch, tmp_path): + marker = tmp_path / "marker.txt" + marker.write_text("x", encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(tmp_path)) + assert harness_root() == tmp_path.resolve() + + +def test_manifest_lists_gates_component(): + manifest = load_manifest() + assert manifest.get("schema_version") == 1 + components = manifest.get("components", {}) + assert "gates" in components + # gates is read at runtime now (orchestrator/gate_policy.py), so it is wired + assert components["gates"]["wired"] is True + assert "gate_policy" in (components["gates"].get("role") or "") or ( + "runtime" in (components["gates"].get("role") or "")) + + +def test_gates_seed_matches_builtin_defaults(): + """Drift guard: the seed gates.yaml must equal the built-in defaults. + + The runtime now reads gates.yaml through orchestrator/gate_policy.py (see + tests/test_gate_policy.py); this test keeps the *seed file* honest so a + fresh harness starts exactly where the hard-coded behaviour was. + """ + from ..orchestrator import gate_policy as gp + + gates = load_gates().get("gates", {}) + defaults = gp.builtin_gates() + assert gates["round_acceptance_improvement_percent"] == ( + defaults["round_acceptance_improvement_percent"]) + assert gates["shadow"]["min_improvement_percent"] == ( + defaults["shadow"]["min_improvement_percent"]) + assert gates["shadow"]["max_exclusive_percent"] == ( + defaults["shadow"]["max_exclusive_percent"]) + assert gates["plateau"]["max_regression_percent"] == ( + defaults["plateau"]["max_regression_percent"]) + assert gates["plateau"]["recent_valid_rounds"] == ( + defaults["plateau"]["recent_valid_rounds"]) + assert gates["isa_gate"]["required_valid_isa_guided_rounds"] == ( + defaults["isa_gate"]["required_valid_isa_guided_rounds"]) + assert gates["task_budget"]["default_max_iterations"] == ( + defaults["task_budget"]["default_max_iterations"]) + + +def test_seed_manifest_marks_data_components_wired(): + manifest = load_manifest() + components = manifest.get("components") or {} + assert components["gates"]["wired"] is True + assert components["planner_policy"]["wired"] is True + assert components["planner_catalog"]["wired"] is True + assert components["manifest"]["wired"] is True + + +def test_seed_workspace_copies_tree(tmp_path): + dst = seed_workspace(tmp_path / "ws") + assert dst.is_dir() + for name in ("manifest.yaml", "gates.yaml", "README.md"): + assert (dst / name).is_file() + # content equality with the seed + with (dst / "gates.yaml").open(encoding="utf-8") as fh: + copied = yaml.safe_load(fh) + original = load_gates() + assert copied == original + + +def test_seed_workspace_missing_root_raises(tmp_path): + with pytest.raises(FileNotFoundError): + seed_workspace(tmp_path / "out", root=tmp_path / "does-not-exist") diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plan_render.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plan_render.py new file mode 100644 index 00000000..20475fb3 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plan_render.py @@ -0,0 +1,102 @@ +"""Tests for planner rendering + controlled-state parity with the legacy menu. + +Controlled-state parity (not corpus heuristic): on states where the planner's +legacy-aligned branches deliberately mirror the current round menus, assert the +mandate texts agree on key direction tokens. +""" + +from __future__ import annotations + +from ..orchestrator.planner import catalog, choose_plan_from_history, render_plan +from ..orchestrator.prompts import w8a8_round_strategy +from ..orchestrator.w8a8_pipeline import isa_round_policy + + +def _valid(improvement: float = 0.5) -> dict: + return { + "build_success": True, + "correctness_passed": True, + "speedup": 1.0, + "metrics": {"graph_capture_passed": True}, + "acceptance": {"improvement_percent": improvement}, + } + + +def test_render_all_catalog_plans_nonempty_and_deterministic(): + cat = catalog() + ctx = {"iteration": 3, "max_iterations": 10, "rounds_left": 8} + seen = {} + for pid in cat: + text = render_plan(pid, ctx=ctx, cat=cat) + assert pid in text and len(text) > 10 + seen[pid] = text + # deterministic + for pid, text in seen.items(): + assert render_plan(pid, ctx=ctx, cat=cat) == text + + +def test_fresh_lane_m16_matches_menu_on_dumma(): + # Fresh m16 lane: planner (establish_arch) and legacy menu (round 1 DUMMA + # bootstrap instruction) both direct at DUMMA. + plan = choose_plan_from_history([], iteration=1, shape={"M": 16}) + menu = w8a8_round_strategy( + {"M": 16, "N": 1536, "K": 4096}, + 1, + [], + {}, + max_iterations=10, + isa_policy=isa_round_policy( + iteration=1, max_iterations=10, history=[] + ), + ) + rendered = render_plan(plan, cat=catalog()) + assert "DUMMA" in rendered and "DUMMA" in menu + + +def test_plateau_open_matches_menu_on_isa(): + history = [_valid(0.5) for _ in range(5)] + [ + _valid(v) for v in (0.5, 0.9, 1.2) + ] + shape = {"M": 16, "N": 1536, "K": 4096} + iteration, max_iterations = 9, 10 + policy = isa_round_policy( + iteration=iteration, max_iterations=max_iterations, history=history + ) + assert policy.get("phase") == "isa_guided_hip" + + plan = choose_plan_from_history( + history, iteration=iteration, max_iterations=max_iterations, + shape=shape, + ) + menu = w8a8_round_strategy( + shape, iteration, history, {}, + max_iterations=max_iterations, isa_policy=policy, + ) + rendered = render_plan(plan, cat=catalog()) + assert plan == "isa_guided_hip" + assert "ISA" in rendered and "ISA" in menu + + +def test_faster_wrong_matches_menu_on_repair(): + history = [_valid(0.5), { + "build_success": True, + "correctness_passed": False, + "speedup": 1.9, + "metrics": {"graph_capture_passed": True}, + }] + shape = {"M": 16, "N": 1536, "K": 4096} + iteration, max_iterations = 4, 10 + plan = choose_plan_from_history( + history, iteration=iteration, max_iterations=max_iterations, + shape=shape, + ) + policy = isa_round_policy( + iteration=iteration, max_iterations=max_iterations, history=history + ) + menu = w8a8_round_strategy( + shape, iteration, history, {}, + max_iterations=max_iterations, isa_policy=policy, + ) + rendered = render_plan(plan, cat=catalog()) + assert plan == "repair_faster_wrong" + assert "repair" in rendered.lower() and "repair" in menu.lower() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner.py new file mode 100644 index 00000000..ffe79f98 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner.py @@ -0,0 +1,204 @@ +"""Tests for orchestrator.planner (M1 slice 2: state-conditioned selector v0).""" + +from __future__ import annotations + +from ..orchestrator import planner + + +def _valid_record(improvement: float = 0.5, *, speedup: float = 1.0) -> dict: + return { + "build_success": True, + "correctness_passed": True, + "speedup": speedup, + "metrics": {"graph_capture_passed": True}, + "acceptance": {"improvement_percent": improvement}, + } + + +def _wrong_fast_record() -> dict: + return { + "build_success": True, + "correctness_passed": False, + "speedup": 1.8, + "metrics": {"graph_capture_passed": True}, + } + + +def _build_fail_record() -> dict: + return {"build_success": False, "failure_reason": "compile error line 42"} + + +def _infra_record() -> dict: + return { + "build_success": False, + "failure_reason": "timed out after 900s", + } + + +def test_catalog_loads_from_yaml_seed(): + cat = planner.catalog() + for pid in ("repair_faster_wrong", "isa_guided_hip", "consolidate"): + assert pid in cat + + +def test_p0_faster_wrong_beats_everything(): + history = [_valid_record(), _wrong_fast_record()] + assert planner.choose_plan_from_history(history, iteration=3) == ( + "repair_faster_wrong" + ) + + +def test_empty_history_is_not_build_failure(): + # A fresh lane (no previous round) must not be routed to fix_build. + assert planner.choose_plan_from_history([]) == "establish_arch" + + +def test_p0_infra_retry_and_build_fix(): + assert planner.choose_plan_from_history([_valid_record(), _infra_record()]) == ( + "retry_same" + ) + assert planner.choose_plan_from_history( + [_valid_record(), _build_fail_record()] + ) == "fix_build" + + +def test_p1_last_round_consolidates(): + history = [_valid_record() for _ in range(3)] + assert planner.choose_plan_from_history( + history, iteration=10, max_iterations=10 + ) == "consolidate" + + +def test_p1_plateau_opens_isa_then_asm(): + # 8 valid HIP rounds + last three improvements inside [-2, 2) prove + # plateau with the ISA gate open (required_hip_rounds = max_iterations - 2). + history = [_valid_record(0.5) for _ in range(5)] + [ + _valid_record(v) for v in (0.5, 0.9, 1.2) + ] + plan = planner.choose_plan_from_history( + history, iteration=9, max_iterations=10 + ) + assert plan == "isa_guided_hip" + plan2 = planner.choose_plan_from_history( + history, + iteration=9, + max_iterations=10, + compiler_limitation_confirmed=True, + ) + assert plan2 == "conditional_inline_asm" + + +def test_p2_occupancy_signature_selects_resource_round(): + history = [_valid_record() for _ in range(2)] + pmc = {"waves_per_cu": 12, "target_waves_per_cu": 16} + assert planner.choose_plan_from_history( + history, pmc=pmc, iteration=5 + ) == "occupancy_resource" + + +def test_p2_bank_conflict_signature(): + history = [_valid_record() for _ in range(2)] + pmc = {"lds_instructions": 1000, "lds_bank_conflicts": 2000} + assert planner.choose_plan_from_history( + history, pmc=pmc, iteration=5 + ) == "memory_layout" + + +def test_p4_fallback_matches_legacy_menu_intent(monkeypatch): + # The no-evidence fallback (uncertainty -> legacy_menu) is on by default; + # disable it here so the P4 portfolio table itself is exercised. + pol = dict(planner._builtin_policy()) + pol["uncertainty"] = {"enabled": False} + monkeypatch.setattr(planner, "policy", lambda root=None: pol) + history = [_valid_record() for _ in range(2)] + plan = planner.choose_plan_from_history( + history, iteration=2, max_iterations=10, shape={"M": 4096} + ) + assert plan == "memory_layout" # large_m legacy menu, step after bootstrap + + +def test_coverage_guard_breaks_same_plan_repetition(): + history = [_valid_record() for _ in range(2)] + pmc = { + "waves_per_cu": 12, + "target_waves_per_cu": 16, + "l2_hit_rate": 40.0, + } + # two prior identical occupancy rounds -> P3 should route to the next + # candidate (memory_layout from l2_low) instead of repeating occupancy. + plan = planner.choose_plan_from_history( + history, + pmc=pmc, + iteration=6, + plan_tags=["occupancy_resource", "occupancy_resource"], + ) + assert plan == "memory_layout" + + +def test_fallback_portfolio_cycles_instead_of_pinning_the_tail(): + """Long runs must not be pinned to the table's last row (consolidate). + + 9-8-8 iteration 3 showed 9 of 11 planner picks were ``consolidate`` + because the fallback indexed past the table end and clamped. The + portfolio now cycles, so a long run keeps revisiting exploration plans. + """ + from ..orchestrator.planner import _legacy_fallback + + picks = [ + _legacy_fallback({"shape": {"M": 16}, "iteration": it}) + for it in range(2, 14) + ] + assert len(set(picks)) > 2, picks + tail = picks[-4:] + assert len(set(tail)) > 1, tail + assert tail.count("consolidate") < len(tail), tail + + +def test_uncertainty_falls_back_to_legacy_menu(monkeypatch): + """No bottleneck evidence -> defer to the hand-tuned menu.""" + from ..orchestrator import planner as P + + pol = dict(P._builtin_policy()) + pol["uncertainty"] = {"enabled": True, "plan": "legacy_menu", + "min_valid_rounds": 1} + monkeypatch.setattr(P, "policy", lambda root=None: pol) + monkeypatch.setattr(P, "bottleneck_tags", lambda pmc: []) + ctx = {"valid_hip_rounds": 2, "pmc": {}, "shape": {"M": 16}, + "iteration": 3, "max_iterations": 11, "rounds_left": 9, + "tried_counts": {}, "consecutive_same": 0, + "last_present": False, "faster_wrong": False} + assert P.choose_plan(ctx) == "legacy_menu" + + +def test_uncertainty_is_skipped_when_a_bottleneck_is_known(monkeypatch): + from ..orchestrator import planner as P + + pol = dict(P._builtin_policy()) + pol["uncertainty"] = {"enabled": True, "plan": "legacy_menu"} + pol["bottleneck_to_plan"] = {"lds_bank_conflicts": "memory_layout"} + monkeypatch.setattr(P, "policy", lambda root=None: pol) + monkeypatch.setattr(P, "bottleneck_tags", + lambda pmc: ["lds_bank_conflicts"]) + ctx = {"valid_hip_rounds": 2, "pmc": {"available": True}, + "shape": {"M": 16}, "iteration": 3, "max_iterations": 11, + "rounds_left": 9, "tried_counts": {}, "consecutive_same": 0, + "last_present": False, "faster_wrong": False} + assert P.choose_plan(ctx) == "memory_layout" + + +def test_legacy_menu_plan_renders_the_legacy_menu(tmp_path, monkeypatch): + import json as _json + from ..orchestrator import w8a8_pipeline as W + + monkeypatch.setenv("METAINFER_PLANNER", "1") + monkeypatch.setattr(W, "choose_plan_from_history", + lambda *a, **k: "legacy_menu") + sink = tmp_path / "planner_plans.jsonl" + text = W._round_strategy_text( + {"M": 16, "N": 4096, "K": 2048}, 2, [], {"available": False}, + {"max_iterations": 5, "skill_allowed": True}, + plan_sink=sink, shape_id="hy3_tp4_o_proj_m16", + ) + assert text and len(text) > 50 + row = _json.loads(sink.read_text(encoding="utf-8").splitlines()[-1]) + assert row["plan_id"] == "legacy_menu" and row["source"] == "planner" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_policy.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_policy.py new file mode 100644 index 00000000..56e4a4cd --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_policy.py @@ -0,0 +1,62 @@ +"""Prove planner_policy.yaml is a live, evolvable state->plan component.""" + +from __future__ import annotations + +import shutil +from pathlib import Path + +import yaml + +from ..orchestrator import planner +from ..orchestrator.harness_io import default_harness_dir + + +def _valid(): + return {"build_success": True, "correctness_passed": True, + "metrics": {"graph_capture_passed": True}, + "acceptance": {"improvement_percent": 5.0}} + + +def test_policy_seed_loads(): + pol = planner.policy() + assert pol["repair_priority"]["faster_wrong"] == "repair_faster_wrong" + assert pol["bottleneck_to_plan"]["occupancy_limited"] == "occupancy_resource" + + +def test_workspace_policy_change_changes_selected_plan(monkeypatch, tmp_path): + root = tmp_path / "harness" + shutil.copytree(default_harness_dir(), root) + policy_path = root / "planner_policy.yaml" + data = yaml.safe_load(policy_path.read_text(encoding="utf-8")) + # AHE-like change: route occupancy pressure to pipeline_tune instead of + # occupancy_resource. The selector must obey the evolved workspace file. + data["bottleneck_to_plan"]["occupancy_limited"] = "pipeline_tune" + policy_path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + + plan = planner.choose_plan_from_history( + [_valid(), _valid()], + pmc={"waves_per_cu": 8, "target_waves_per_cu": 16}, + iteration=4, + max_iterations=10, + shape={"M": 4096}, + ) + assert plan == "pipeline_tune" + + +def test_policy_fallback_order_is_evolvable(monkeypatch, tmp_path): + root = tmp_path / "harness" + shutil.copytree(default_harness_dir(), root) + path = root / "planner_policy.yaml" + data = yaml.safe_load(path.read_text(encoding="utf-8")) + data["fallback"]["prefill"][0] = "occupancy_resource" + # exercise the evolvable portfolio itself, not the no-evidence fallback + data.setdefault("uncertainty", {})["enabled"] = False + path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + + plan = planner.choose_plan_from_history( + [_valid(), _valid()], pmc={}, iteration=2, max_iterations=10, + shape={"M": 4096}, + ) + assert plan == "occupancy_resource" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_wiring.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_wiring.py new file mode 100644 index 00000000..9261414f --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_wiring.py @@ -0,0 +1,84 @@ +"""Tests for planner wiring env toggle (METAINFER_PLANNER).""" + +from __future__ import annotations + +from ..orchestrator.w8a8_pipeline import ( + _planner_enabled, + _round_strategy_text, + isa_round_policy, +) +from ..orchestrator.prompts import w8a8_round_strategy + + +def _isa_policy(iteration: int, max_iterations: int, history) -> dict: + return isa_round_policy( + iteration=iteration, max_iterations=max_iterations, history=history + ) + + +def test_planner_disabled_by_default(monkeypatch): + monkeypatch.delenv("METAINFER_PLANNER", raising=False) + assert _planner_enabled() is False + + +def test_env_toggle_enables(monkeypatch): + monkeypatch.setenv("METAINFER_PLANNER", "1") + assert _planner_enabled() is True + + +def test_disabled_returns_exact_legacy_text(monkeypatch): + monkeypatch.delenv("METAINFER_PLANNER", raising=False) + shape = {"M": 16, "N": 1536, "K": 4096} + history = [] + policy = _isa_policy(1, 10, history) + got = _round_strategy_text(shape, 1, history, {}, policy) + expected = w8a8_round_strategy( + shape, 1, history, {}, + max_iterations=10, isa_policy=policy, + ) + assert got == expected + + +def test_enabled_renders_planner_mandate(monkeypatch): + monkeypatch.setenv("METAINFER_PLANNER", "1") + shape = {"M": 16, "N": 1536, "K": 4096} + history = [] + policy = _isa_policy(1, 10, history) + got = _round_strategy_text(shape, 1, history, {}, policy) + assert got.startswith("Mandatory decision for this round:") + assert "establish_arch" in got + + +def test_round_strategy_text_records_planner_decision(tmp_path, monkeypatch): + """The planner's own choice is written to planner_plans.jsonl (hard evidence).""" + import json as _json + from ..orchestrator.w8a8_pipeline import _round_strategy_text + + monkeypatch.setenv("METAINFER_PLANNER", "1") + sink = tmp_path / "planner_plans.jsonl" + text = _round_strategy_text( + {"M": 16, "N": 4096, "K": 2048}, 1, [], {"available": False}, + {"max_iterations": 3, "skill_allowed": True}, + plan_sink=sink, shape_id="hy3_tp4_o_proj_m16", + ) + assert text + assert sink.is_file() + row = _json.loads(sink.read_text(encoding="utf-8").splitlines()[0]) + assert row["source"] == "planner" + assert row["plan_id"] + assert row["iteration"] == 1 + assert row["shape_id"] == "hy3_tp4_o_proj_m16" + + +def test_round_strategy_text_writes_nothing_without_planner(tmp_path, monkeypatch): + from ..orchestrator.w8a8_pipeline import _round_strategy_text + + monkeypatch.delenv("METAINFER_PLANNER", raising=False) + sink = tmp_path / "planner_plans.jsonl" + text = _round_strategy_text( + {"M": 16, "N": 4096, "K": 2048}, 1, [], {"available": False}, + {"max_iterations": 3, "skill_allowed": True}, + plan_sink=sink, shape_id="hy3_tp4_o_proj_m16", + ) + assert text + assert not sink.exists() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_predictions.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_predictions.py new file mode 100644 index 00000000..6b047a5b --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_predictions.py @@ -0,0 +1,94 @@ +"""Tests for orchestrator.predictions (M1 slice 3: inner decision hook).""" + +from __future__ import annotations + +from ..orchestrator.predictions import check_prediction, parse_prediction, plan_tag + + +def test_no_structured_prediction_returns_none(): + assert check_prediction( + {"hypothesis": "plain prose"}, + candidate_us=100.0, + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) is None + assert parse_prediction({"hypothesis": "x"}) is None + + +def test_infra_failure_is_na(): + out = check_prediction( + {"prediction": {"direction": "improve"}}, + candidate_us=float("inf"), + best_us=100.0, + passed=False, + p90_guard_passed=False, + failure_reason="timed out after 900s", + ) + assert out is not None and out["checked"] == "na" + + +def test_range_hit_and_miss(): + proposal = {"prediction": {"expected_us_range": [90.0, 105.0]}} + hit = check_prediction( + proposal, + candidate_us=98.0, + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) + assert hit["checked"] == "hit" + miss = check_prediction( + proposal, + candidate_us=120.0, + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) + assert miss["checked"] == "miss" + assert "outside declared range" in miss["reason"] + + +def test_direction_improve_requires_band(): + ok = check_prediction( + {"prediction": {"direction": "improve"}}, + candidate_us=100.0, + best_us=110.0, # ~9% better + passed=True, + p90_guard_passed=True, + ) + assert ok["checked"] == "hit" + flat = check_prediction( + {"prediction": {"direction": "flat"}}, + candidate_us=109.0, # ~0.9% delta -> inside flat band + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) + assert flat["checked"] == "hit" + wrong = check_prediction( + {"prediction": {"direction": "improve"}}, + candidate_us=109.0, # only ~0.9% better -> below the 2% band + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) + assert wrong["checked"] == "miss" + + +def test_correctness_failure_means_miss(): + out = check_prediction( + {"prediction": {"expected_us_range": [90.0, 105.0]}}, + candidate_us=98.0, + best_us=110.0, + passed=False, + p90_guard_passed=False, + ) + assert out["checked"] == "miss" + assert "correctness failed" in out["reason"] + + +def test_plan_tag_helpers(): + assert plan_tag({"hypothesis": "x"}) is None + assert plan_tag({"plan_id": "occupancy_resource"}) == "occupancy_resource" + assert plan_tag({"plan_id": " "}) is None diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_prompt_schema.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_prompt_schema.py new file mode 100644 index 00000000..a3dd5acd --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_prompt_schema.py @@ -0,0 +1,47 @@ +"""Prompt-schema tests: worker prompt now carries optional decision fields.""" + +from __future__ import annotations + +from pathlib import Path + +from ..orchestrator.config import WorkerAssignment +from ..orchestrator.w8a8_pipeline import RealW8A8OptimizationPipeline + + +def _round_prompt(**overrides): + args = dict( + assignment=WorkerAssignment("worker_0", 0, ["m2"]), + shape_id="m2", + shape={"M": 2, "N": 16, "K": 32}, + best={"median_us": 10.0}, + root=Path("/tmp/worker"), + iteration=1, + guidance=None, + history=[], + pmc_evidence={}, + ) + args.update(overrides) + return RealW8A8OptimizationPipeline._worker_prompt(**args) + + +def test_first_turn_prompt_mentions_decision_fields(): + prompt = _round_prompt() + assert '"plan_id"' in prompt + assert '"prediction"' in prompt + assert '"expected_us_range"' in prompt + assert '"direction": "improve"' in prompt + assert "checks `prediction` against the measured round" in prompt + + +def test_continuation_prompt_mentions_optional_fields(): + prompt = _round_prompt( + continuation=True, + history=[{ + "iteration": 1, + "accepted": True, + "build_success": True, + "correctness_passed": True, + "metrics": {"median_us": 9.0}, + }], + ) + assert "optional plan_id and prediction" in prompt 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 index 54333e68..9eda815c 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_skill_store.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_skill_store.py @@ -3,19 +3,31 @@ import pytest from ..orchestrator.config import load_config +from ..orchestrator import skill_store as skill_store_mod from ..orchestrator.skill_store import ( _apply_fuse_decision, _parse_fuse_decision, + bundled_skills_root, dsh_skills_root, generate_merged_skill, generate_worker_skill, list_skill_library, publish_skill, rollback_skill, + seed_bundled_skills, sync_skill_libraries, ) +@pytest.fixture(autouse=True) +def _isolated_bundled_skills(tmp_path, monkeypatch): + """Point the plugin's bundled-skills root at an empty tmp dir by default, + so seeding never pollutes library assertions in unrelated tests.""" + bundled = tmp_path / "bundled-skills" + bundled.mkdir(exist_ok=True) + monkeypatch.setattr(skill_store_mod, "bundled_skills_root", lambda: bundled) + + def test_skill_library_lists_and_publishes_without_overwrite( tmp_path, monkeypatch ): @@ -84,6 +96,28 @@ def test_sync_skill_libraries_mirrors_dsh_to_ccb(tmp_path, monkeypatch): assert summary2["skipped"] == ["alpha", "beta"] +def test_seed_bundled_skills_copies_missing_only(tmp_path, monkeypatch): + dsh_root = tmp_path / "dsh" + bundled = tmp_path / "bundled" + (bundled / "zeta").mkdir(parents=True) + (bundled / "zeta" / "SKILL.md").write_text("# zeta seed\n", encoding="utf-8") + (bundled / "eta").mkdir(parents=True) + (bundled / "eta" / "SKILL.md").write_text("# eta seed\n", encoding="utf-8") + # dsh already has eta (different content) -> must not be overwritten + (dsh_root / "eta").mkdir(parents=True) + (dsh_root / "eta" / "SKILL.md").write_text("# eta live\n", encoding="utf-8") + monkeypatch.setenv("DSH_SKILLS_DIR", str(dsh_root)) + monkeypatch.setattr(skill_store_mod, "bundled_skills_root", lambda: bundled) + + seeded = seed_bundled_skills() + assert seeded["added"] == ["zeta"] + assert (dsh_root / "zeta" / "SKILL.md").read_text() == "# zeta seed\n" + # an existing live skill is never clobbered by the bundled seed + assert (dsh_root / "eta" / "SKILL.md").read_text() == "# eta live\n" + # idempotent + assert seed_bundled_skills()["added"] == [] + + def test_rollback_skill_restores_backup(tmp_path, monkeypatch): dsh_root = tmp_path / "dsh" ccb_root = tmp_path / "ccb" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_validation_budget.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_validation_budget.py new file mode 100644 index 00000000..fbd26361 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_validation_budget.py @@ -0,0 +1,54 @@ +"""Validation-budget knobs (scope + bench sampling) shared by DKAO and AHE.""" + +from __future__ import annotations + +import pytest + +from ..orchestrator.gen_and_opt_pipeline import _validation_shape_list +from ..orchestrator.validation_budget import ( + ENV_BENCH_SAMPLES, ENV_BENCH_WARMUPS, ENV_VALIDATE_SCOPE, QUICK_BENCH, + resolve_bench_kwargs, resolve_validation_scope, +) + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + for name in (ENV_VALIDATE_SCOPE, ENV_BENCH_WARMUPS, ENV_BENCH_SAMPLES, + "METAINFER_BENCH_REPLAYS"): + monkeypatch.delenv(name, raising=False) + + +def test_scope_defaults_to_full_api(monkeypatch): + assert resolve_validation_scope({}) == "api" + assert resolve_validation_scope(None) == "api" + + +def test_scope_reads_form_labels(monkeypatch): + assert resolve_validation_scope({"validation_scope": "Task shapes only"}) == "task" + assert resolve_validation_scope({"validation_scope": "All API shapes"}) == "api" + assert resolve_validation_scope({"validation_scope": "task"}) == "task" + + +def test_scope_env_wins(monkeypatch): + monkeypatch.setenv(ENV_VALIDATE_SCOPE, "task") + assert resolve_validation_scope({"validation_scope": "All API shapes"}) == "task" + + +def test_bench_kwargs_default_and_quick(monkeypatch): + assert resolve_bench_kwargs({}) == {} + assert resolve_bench_kwargs({"bench_profile": "quick"}) == QUICK_BENCH + assert resolve_bench_kwargs({"bench_profile": "full"}) == {} + + +def test_bench_kwargs_explicit_and_env_priority(monkeypatch): + out = resolve_bench_kwargs({"bench_profile": "quick", "bench_samples": 7}) + assert out["samples"] == 7 and out["warmups"] == QUICK_BENCH["warmups"] + monkeypatch.setenv(ENV_BENCH_SAMPLES, "5") + assert resolve_bench_kwargs({"bench_samples": 7})["samples"] == 5 + + +def test_validation_shape_list_scopes(): + optimized = [{"id": "a"}] + fallback = [{"id": "b"}, {"id": "c"}] + assert [s["id"] for s in _validation_shape_list(optimized, fallback, "task")] == ["a"] + assert [s["id"] for s in _validation_shape_list(optimized, fallback, "api")] == ["a", "b", "c"] diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_promote.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_promote.py new file mode 100644 index 00000000..cd23a824 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_promote.py @@ -0,0 +1,149 @@ +"""Kernel promotion gates (variant tree write) — no GPU calls.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from ..orchestrator import variant_store as vs +from ..orchestrator.variant_promote import ( + accepted_kernel_for, model_label_for, promote_variant, +) + +ANSWERS = {"operator": "Quantized GEMM", "dtype": "INT8 W8A8"} +SHAPE = "hy3_tp4_o_proj_m16" + + +@pytest.fixture() +def variant_root(tmp_path, monkeypatch): + root = tmp_path / "variant" + root.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(vs, "variant_root", lambda: root) + return root + + +def _workspace(tmp_path: Path, median_us: float, *, shape: str = SHAPE, + commit: str = "deadbeef") -> Path: + ws = tmp_path / "ws" + leaf = ws / "workers" / "worker_0" / "accepted" / shape + leaf.mkdir(parents=True, exist_ok=True) + (leaf / "kernel.hip").write_text("// kernel source\n", encoding="utf-8") + (leaf / "manifest.json").write_text(json.dumps({ + "commit": commit, + "shape": {"M": 16, "N": 4096, "K": 2048}, + "metrics": {"median_us": median_us, "p90_us": median_us * 1.02}, + }), encoding="utf-8") + return ws + + +def _seed_variant(variant_root: Path, median_us: float) -> Path: + meta = vs.derive_variant_meta(ANSWERS | {"model": "Hy3 (Hunyuan 3)"}, SHAPE) + target = vs.variant_path(meta) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(vs.section_header(meta, commit="old", metrics={ + "median_us": median_us, "p90_us": median_us}) + "// old kernel\n// @@end\n", + encoding="utf-8") + return target + + +def test_adds_when_no_existing_variant(tmp_path, variant_root): + ws = _workspace(tmp_path, 20.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=3.0, + source_task="he-test") + assert out["action"] == "added" + assert Path(out["path"]).is_file() + assert "he-test" in Path(out["path"]).read_text(encoding="utf-8") + + +def test_updates_when_faster_than_threshold(tmp_path, variant_root): + _seed_variant(variant_root, 100.0) + ws = _workspace(tmp_path, 90.0) # 10% faster + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=3.0) + assert out["action"] == "updated" + assert out["improvement_percent"] == pytest.approx(10.0, abs=1e-6) + assert out["backup"] and Path(out["backup"]).is_file() + assert "median_us=90" in Path(out["path"]).read_text(encoding="utf-8") + + +def test_skips_improvement_below_threshold(tmp_path, variant_root): + _seed_variant(variant_root, 100.0) + ws = _workspace(tmp_path, 98.5) # only 1.5% faster + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=3.0) + assert out["action"] == "skipped" + assert "required 3.00%" in out["reason"] + # existing variant untouched (still the old median) + assert "median_us=100" in Path(out["path"]).read_text(encoding="utf-8") + + +def test_slower_candidate_is_skipped_by_default_gate(tmp_path, variant_root): + """With the default gate (min 0%) any regression is skipped up front.""" + _seed_variant(variant_root, 100.0) + ws = _workspace(tmp_path, 120.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=0.0) + assert out["action"] == "skipped" + assert "required 0.00%" in out["reason"] + + +def test_add_variant_guard_rejects_slower_when_gate_is_loosened(tmp_path, variant_root): + """A loosened threshold still cannot overwrite with a slower kernel.""" + _seed_variant(variant_root, 100.0) + ws = _workspace(tmp_path, 101.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=-10.0) + assert out["action"] == "rejected" + assert "slower" in out["reason"] + + +def test_skips_when_correctness_failed(tmp_path, variant_root): + ws = _workspace(tmp_path, 10.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=False, min_improvement_percent=3.0) + assert out["action"] == "skipped" + assert "correctness" in out["reason"] + + +def test_no_kernel_found(tmp_path, variant_root): + out = promote_variant(workspace_dir=tmp_path / "empty", answers=ANSWERS, + shape_id=SHAPE) + assert out["action"] == "no_kernel" + + +def test_dry_run_does_not_write(tmp_path, variant_root): + ws = _workspace(tmp_path, 10.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=3.0, + dry_run=True) + assert out["action"] == "would-add" + assert not Path(out["path"]).exists() + + +def test_meta_taxonomy_and_overrides(tmp_path, variant_root): + ws = _workspace(tmp_path, 10.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True) + meta = out["meta"] + assert (meta["family"], meta["model"], meta["tp"], meta["m"], + meta["operator_name"]) == ("int8w8a8-gemm", "hy3", 4, 16, "o_proj") + + # shapes without tp/M in the id take them from the caller (pool contract) + ws2 = _workspace(tmp_path / "two", 10.0, shape="m16_wqkv_a") + out2 = promote_variant(workspace_dir=ws2, answers=ANSWERS, + shape_id="m16_wqkv_a", correctness_ok=True, + tp=8, m=16, model_label="DeepSeek V4 Flash") + assert out2["meta"]["tp"] == 8 and out2["meta"]["m"] == 16 + assert "/TP8/M16/" in out2["path"] + + +def test_model_label_and_accepted_lookup(tmp_path): + assert model_label_for("minimax_tp8_o_proj_m16") == "MiniMax M3" + assert model_label_for("glm_tp8_o_proj_m4096") == "GLM5.2" + ws = _workspace(tmp_path, 5.0) + found = accepted_kernel_for(ws, SHAPE) + assert found and found["metrics"]["median_us"] == 5.0 + assert accepted_kernel_for(ws, "nope") is None 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 index 6310e5e2..0f5a29c4 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_store.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_store.py @@ -196,3 +196,31 @@ def test_reject_slower_guard_skips_when_median_unknown(tmp_path, monkeypatch): ) assert result["action"] == "updated" assert "NEW" in vs.variant_path(meta).read_text() + + +def test_backfill_variant_baselines(tmp_path, monkeypatch): + monkeypatch.setattr(vs, "variant_root", lambda: tmp_path) + target = ( + tmp_path / "int8w8a8-gemm" / "deepseek-v4" / "TP4" / "M16" + / "wqkv_a.hip" + ) + target.parent.mkdir(parents=True) + target.write_text( + "// @@variant shape=tp4_wqkv_a_m16 commit=abc added=2026-08-26\n" + "// median_us=20.0 p90_us=21.0\n" + "// source=task-1\n" + "// @@end\nkernel\n", + encoding="utf-8", + ) + r = vs.backfill_variant_baselines() + # tp4 wqkv_a m16 fixed Triton baseline is 66.597 us -> speedup ~3.33 + assert r["updated"] == ["int8w8a8-gemm/deepseek-v4/TP4/M16/wqkv_a.hip"] + text = target.read_text() + assert "baseline_us=66.597" in text + assert "speedup=3.329" in text + assert text.index("baseline_us") < text.index("median_us") + # idempotent and surfaced by the index + assert vs.backfill_variant_baselines()["updated"] == [] + v = vs.list_variant_index()[0] + assert v["speedup"] is not None + assert abs(v["baseline_us"] - 66.597) < 1e-3 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp4_m4096.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp4_m4096.py new file mode 100644 index 00000000..6da463c1 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp4_m4096.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +"""Fair Triton W8A8 GEMM baseline for the six TP4 logical shapes at M=4096. + +Timed region launches only lmslim-style int8_utils.matmul_kernel into a +preallocated bf16 output (output allocation/clear excluded), matching the +MetaInfer fair-baseline scope. Uses GPU events, hot cache, and reports +median/P90/min in microseconds plus effective INT8 TOPS. + +The Triton config replicates matmul_int8's built-in default for M > 1024: + BLOCK_SIZE_M=256, BLOCK_SIZE_N=256, BLOCK_SIZE_K=64, + GROUP_SIZE_M=8, SPLIT_K=1, num_stages=0, num_warps=8 +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import math +import statistics +import sys +from pathlib import Path + +import torch +import triton + + +UTILS_PATH = str(Path(__file__).resolve().parent / "int8_utils.py") +M = 4096 + +# (shape_id, operator, K, N) -- the six TP4 logical operators. +TP4_CASES = [ + ("tp4_wqkv_a_m4096", "wqkv_a", 4096, 1536), + ("tp4_wq_b_m4096", "wq_b", 1024, 8192), + ("tp4_indexer_wq_b_m4096", "indexer.wq_b", 1024, 8192), + ("tp4_wo_b_m4096", "wo_b", 2048, 4096), + ("tp4_shared_gate_up_proj_m4096", "shared_gate_up_proj", 4096, 1024), + ("tp4_shared_down_proj_m4096", "shared_down_proj", 512, 4096), +] + +# Exact default config used by matmul_int8 when M > 1024. +CONFIG = { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8, +} + + +def load_int8_utils(path: str): + spec = importlib.util.spec_from_file_location( + "baseline_int8_utils", path + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def launch(out, a, a_scale, b, b_scale, m, n, k, utils) -> None: + grid = ( + triton.cdiv(m, CONFIG["BLOCK_SIZE_M"]) + * triton.cdiv(n, CONFIG["BLOCK_SIZE_N"]), + CONFIG["SPLIT_K"], + ) + utils.matmul_kernel[grid]( + a, + a_scale, + b, + b_scale, + out, + m, + n, + k, + a.stride(0), + a.stride(1), + a_scale.stride(0), + b.stride(0), + b.stride(1), + b_scale.stride(0), + out.stride(0), + out.stride(1), + **CONFIG, + ) + + +def percentile(values, fraction): + ordered = sorted(values) + return ordered[max(0, math.ceil(fraction * len(ordered)) - 1)] + + +def measure(fn, warmups, samples, launches_per_sample): + for _ in range(warmups): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + times_us = [] + for _ in range(samples): + start.record() + for _ in range(launches_per_sample): + fn() + end.record() + end.synchronize() + times_us.append( + start.elapsed_time(end) * 1000.0 / launches_per_sample + ) + return { + "median_us": statistics.median(times_us), + "p90_us": percentile(times_us, 0.90), + "min_us": min(times_us), + "max_us": max(times_us), + "samples_us": [round(v, 3) for v in times_us], + } + + +@torch.no_grad() +def reference_us(a, b, a_scale, b_scale, device): + dot = torch.mm( + a.to("cpu", dtype=torch.int64), + b.to("cpu", dtype=torch.int64), + ) + scaled = ( + dot.to(torch.float32) + * a_scale.to("cpu", dtype=torch.float32) + * b_scale.to("cpu", dtype=torch.float32).T + ) + return scaled.to(torch.bfloat16).to(device) + + +def run_case(utils, case_id, operator, k, n, warmups, samples, + launches_per_sample, skip_check, mode): + device = torch.device("cuda:0") + torch.manual_seed(20260724 + M + n + k) + a = torch.randint(-127, 128, (M, k), device=device, dtype=torch.int8) + b = torch.randint(-127, 128, (k, n), device=device, dtype=torch.int8) + a_scale = torch.rand((M, 1), device=device, dtype=torch.float32) + 0.01 + b_scale = torch.rand((n, 1), device=device, dtype=torch.float32) + 0.01 + out = torch.empty((M, n), device=device, dtype=torch.bfloat16) + + fn = lambda: launch(out, a, a_scale, b, b_scale, M, n, k, utils) + + # One untimed JIT compile + correctness sanity pass. + fn() + torch.cuda.synchronize() + errors = None + if not skip_check: + ref = reference_us(a, b, a_scale, b_scale, device) + diff = (out.float() - ref.float()).abs() + denom = ref.float().abs().clamp_min(1e-6) + errors = { + "max_abs": float(diff.max().item()), + "max_rel": float((diff / denom).max().item()), + "mismatch_count": int((out != ref).sum().item()), + } + del ref + + result = { + "shape_id": case_id, + "operator": operator, + "M": M, + "N": n, + "K": k, + "config": CONFIG, + "errors": errors, + } + + if mode in ("eager", "both"): + result["eager"] = measure( + fn, warmups, samples, launches_per_sample + ) + if mode in ("graph", "both"): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + torch.cuda.synchronize() + result["graph"] = measure( + graph.replay, warmups, samples, launches_per_sample + ) + del graph + + del a, b, a_scale, b_scale, out + return result + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--utils", default=UTILS_PATH) + parser.add_argument("--warmups", type=int, default=10) + parser.add_argument("--samples", type=int, default=20) + parser.add_argument("--launches-per-sample", type=int, default=5) + parser.add_argument("--mode", choices=["eager", "graph", "both"], + default="both") + parser.add_argument("--skip-check", action="store_true") + parser.add_argument("--json", type=str, default="") + args = parser.parse_args() + + utils = load_int8_utils(args.utils) + print( + "BENCHMARK triton=int8_utils.matmul_kernel output=preallocated " + f"quantization=excluded allocation=excluded epilogue=included " + f"M={M} config={json.dumps(CONFIG)} mode={args.mode}", + flush=True, + ) + print( + f"{'shape_id':28s} {'N':5s} {'K':5s} {'eager_us':>10s} " + f"{'graph_us':>10s} {'tops':>9s}", + flush=True, + ) + results = [] + for case_id, operator, k, n in TP4_CASES: + res = run_case( + utils, + case_id, + operator, + k, + n, + args.warmups, + args.samples, + args.launches_per_sample, + args.skip_check, + args.mode, + ) + results.append(res) + eager_us = res.get("eager", {}).get("median_us") + graph_us = res.get("graph", {}).get("median_us") + tops = ( + 2.0 * M * n * k / ((graph_us or eager_us) * 1.0e-6) / 1.0e12 + if (graph_us or eager_us) else float("nan") + ) + print( + f"{case_id:28s} {n:5d} {k:5d} " + f"{(eager_us if eager_us is not None else float('nan')):10.3f} " + f"{(graph_us if graph_us is not None else float('nan')):10.3f} " + f"{tops:9.3f}", + flush=True, + ) + for key in ("eager", "graph"): + if key in res: + m_ = res[key] + print( + f" {key}: median={m_['median_us']:.3f}us " + f"p90={m_['p90_us']:.3f}us min={m_['min_us']:.3f}us " + f"max={m_['max_us']:.3f}us errors={res['errors']}", + flush=True, + ) + if args.json: + with open(args.json, "w", encoding="utf-8") as handle: + json.dump( + { + "scope": "gemm_out", + "triton_source": args.utils, + "M": M, + "warmups": args.warmups, + "samples": args.samples, + "launches_per_sample": args.launches_per_sample, + "mode": args.mode, + "results": results, + }, + handle, + indent=2, + ensure_ascii=False, + ) + print(f"wrote {args.json}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m4096.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m4096.py new file mode 100644 index 00000000..3dc94b39 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m4096.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +"""Fair Triton W8A8 GEMM baselines for the model-catalog TP8 shapes at M=4096. + +Same protocol as bench_triton_tp4_m4096.py / the 2026-08-12 catalog entries: +lmslim-style int8_utils.matmul_kernel into a preallocated bf16 output (output +allocation/clear excluded), GPU events, hot cache, CUDA-Graph replay, +warmups=10 / samples=20 / launches_per_sample=5. Reports median/P90/min in +microseconds. + +Triton config replicates matmul_int8's built-in default for M > 1024: + BLOCK_SIZE_M=256, BLOCK_SIZE_N=256, BLOCK_SIZE_K=64, + GROUP_SIZE_M=8, SPLIT_K=1, num_stages=0, num_warps=8 +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import math +import statistics +import sys +from pathlib import Path + +import torch +import triton + + +UTILS_PATH = str(Path(__file__).resolve().parent / "int8_utils.py") +M = 4096 + +# (shape_id, operator, K, N) -- Hy3 / MiniMax M3 / GLM5.2 TP8 at M=4096. +TP8_CASES = [ + ("hy3_tp8_qkv_proj_m4096", "qkv_proj", 4096, 1280), + ("hy3_tp8_o_proj_m4096", "o_proj", 1024, 4096), + ("hy3_tp8_shared_gate_up_proj_m4096", "shared_gate_up_proj", 4096, 384), + ("hy3_tp8_shared_down_proj_m4096", "shared_down_proj", 192, 4096), + ("minimax_tp8_qkv_proj_m4096", "qkv_proj", 6144, 1280), + ("minimax_tp8_qkv_proj_and_indexer_qk_m4096", "qkv_proj_and_indexer_qk", 6144, 1536), + ("minimax_tp8_o_proj_m4096", "o_proj", 1024, 6144), + ("minimax_tp8_shared_gate_up_proj_m4096", "shared_gate_up_proj", 6144, 768), + ("minimax_tp8_shared_down_proj_m4096", "shared_down_proj", 384, 6144), + ("glm52_tp8_fused_qkv_a_proj_m4096", "fused_qkv_a_proj", 6144, 2624), + ("glm52_tp8_q_b_proj_m4096", "q_b_proj", 2048, 2048), + ("glm52_tp8_kv_b_proj_m4096", "kv_b_proj", 512, 3584), + ("glm52_tp8_o_proj_m4096", "o_proj", 2048, 6144), + ("glm52_tp8_shared_gate_up_proj_m4096", "shared_gate_up_proj", 6144, 512), + ("glm52_tp8_shared_down_proj_m4096", "shared_down_proj", 256, 6144), +] + +# Exact default config used by matmul_int8 when M > 1024. +CONFIG = { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8, +} + + +def load_int8_utils(path: str): + spec = importlib.util.spec_from_file_location( + "baseline_int8_utils", path + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def launch(out, a, a_scale, b, b_scale, m, n, k, utils) -> None: + grid = ( + triton.cdiv(m, CONFIG["BLOCK_SIZE_M"]) + * triton.cdiv(n, CONFIG["BLOCK_SIZE_N"]), + CONFIG["SPLIT_K"], + ) + utils.matmul_kernel[grid]( + a, + a_scale, + b, + b_scale, + out, + m, + n, + k, + a.stride(0), + a.stride(1), + a_scale.stride(0), + b.stride(0), + b.stride(1), + b_scale.stride(0), + out.stride(0), + out.stride(1), + **CONFIG, + ) + + +def percentile(values, fraction): + ordered = sorted(values) + return ordered[max(0, math.ceil(fraction * len(ordered)) - 1)] + + +def measure(fn, warmups, samples, launches_per_sample): + for _ in range(warmups): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + times_us = [] + for _ in range(samples): + start.record() + for _ in range(launches_per_sample): + fn() + end.record() + end.synchronize() + times_us.append( + start.elapsed_time(end) * 1000.0 / launches_per_sample + ) + return { + "median_us": statistics.median(times_us), + "p90_us": percentile(times_us, 0.90), + "min_us": min(times_us), + "max_us": max(times_us), + "samples_us": [round(v, 3) for v in times_us], + } + + +def run_case(utils, case_id, operator, k, n, warmups, samples, + launches_per_sample, mode): + device = torch.device("cuda:0") + torch.manual_seed(20260724 + M + n + k) + a = torch.randint(-127, 128, (M, k), device=device, dtype=torch.int8) + b = torch.randint(-127, 128, (k, n), device=device, dtype=torch.int8) + a_scale = torch.rand((M, 1), device=device, dtype=torch.float32) + 0.01 + b_scale = torch.rand((n, 1), device=device, dtype=torch.float32) + 0.01 + out = torch.empty((M, n), device=device, dtype=torch.bfloat16) + + fn = lambda: launch(out, a, a_scale, b, b_scale, M, n, k, utils) + + # One untimed JIT compile pass (correctness is the production Triton path). + fn() + torch.cuda.synchronize() + + result = { + "shape_id": case_id, + "operator": operator, + "M": M, + "N": n, + "K": k, + "config": CONFIG, + } + if mode in ("eager", "both"): + result["eager"] = measure( + fn, warmups, samples, launches_per_sample + ) + if mode in ("graph", "both"): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + torch.cuda.synchronize() + result["graph"] = measure( + graph.replay, warmups, samples, launches_per_sample + ) + del graph + + del a, b, a_scale, b_scale, out + return result + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--utils", default=UTILS_PATH) + parser.add_argument("--warmups", type=int, default=10) + parser.add_argument("--samples", type=int, default=20) + parser.add_argument("--launches-per-sample", type=int, default=5) + parser.add_argument("--mode", choices=["eager", "graph", "both"], + default="graph") + parser.add_argument("--json", type=str, default="") + args = parser.parse_args() + + utils = load_int8_utils(args.utils) + print( + "BENCHMARK triton=int8_utils.matmul_kernel output=preallocated " + f"quantization=excluded allocation=excluded epilogue=included " + f"M={M} config={json.dumps(CONFIG)} mode={args.mode}", + flush=True, + ) + print( + f"{'shape_id':42s} {'N':5s} {'K':5s} {'graph_us':>10s} " + f"{'p90_us':>10s} {'tops':>9s}", + flush=True, + ) + results = [] + for case_id, operator, k, n in TP8_CASES: + res = run_case( + utils, + case_id, + operator, + k, + n, + args.warmups, + args.samples, + args.launches_per_sample, + args.mode, + ) + results.append(res) + graph_us = (res.get("graph") or res.get("eager") or {}).get("median_us") + p90_us = (res.get("graph") or res.get("eager") or {}).get("p90_us") + tops = ( + 2.0 * M * n * k / (graph_us * 1e-6) / 1e12 + if graph_us else None + ) + print( + f"{case_id:42s} {n:5d} {k:5d} " + f"{graph_us:>10.3f} {p90_us:>10.3f} " + f"{(f'{tops:.3f}' if tops else '—'):>9s}", + flush=True, + ) + if args.json: + with open(args.json, "w", encoding="utf-8") as fh: + json.dump(results, fh, indent=2) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/int8_utils.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/int8_utils.py new file mode 100644 index 00000000..11886719 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/int8_utils.py @@ -0,0 +1,602 @@ +import functools +import json +import logging +import os +from typing import Any, Dict, List, Optional, Tuple +import json + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice +import lmslim.envs as lsenvs + +use_lightop = lsenvs.LMSLIM_USE_LIGHTOP +device_name = lsenvs.LMSLIM_GPU_NAME + +@triton.jit +def _per_token_quant_int8( + x_ptr, + xq_ptr, + scale_ptr, + stride_x, + stride_xq, + N, + BLOCK: tl.constexpr, +): + # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282 + row_id = tl.program_id(0) + + cols = tl.arange(0, BLOCK) + mask = cols < N + + x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask, + other=0.0).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10) + scale_x = absmax / 127 + x_q = x * (127 / absmax) + x_q = libdevice.nearbyint(x_q).to(tl.int8) + + tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask) + tl.store(scale_ptr + row_id, scale_x) + + +def per_token_quant_int8(x): + M = x.numel() // x.shape[-1] + N = x.shape[-1] + x_q = torch.empty_like(x, device=x.device, dtype=torch.int8) + scales = torch.empty(x.shape[:-1] + (1, ), + device=x.device, + dtype=torch.float32) + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + + #assert x.is_contiguous() + _per_token_quant_int8[(M, )]( + x, + x_q, + scales, + stride_x=x.stride(-2), + stride_xq=x_q.stride(-2), + N=N, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=1, + ) + + return x_q, scales + + +@triton.jit +def _per_token_group_quant_int8( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + # Stride of input + group_size, + # M, + # K, + # # Collums of input + # N, + SIZE, + # Avoid to divide zero + eps, + # Information for int8 + int8_min, + int8_max, + # Meta-parameters + BLOCK: tl.constexpr, + s_num : tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group + quantization on a tensor. + + This function converts the tensor values into int8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * BLOCK + y_q_ptr += g_id * BLOCK + y_s_ptr += g_id * s_num + + cols = tl.arange(0, BLOCK) # N <= BLOCK + s_cols = tl.arange(0, s_num) + mask = g_id * BLOCK + cols < SIZE + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + y = tl.reshape(y, (s_num, 128)) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps) + y_s = (_absmax / int8_max).reshape(s_num, 1) + y_q = tl.clamp(y / y_s, int8_min, int8_max).to(y_q_ptr.dtype.element_ty) + + y_q = tl.reshape(y_q, (s_num*128)) + y_s = tl.reshape(y_s, (s_num)) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr + s_cols, y_s.to(y_s_ptr.dtype.element_ty)) + + +def per_token_group_quant_int8( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: torch.dtype = torch.int8, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Function to perform per-token-group quantization on an input tensor `x`. + + It converts the tensor values into signed int8 values and returns the + quantized tensor along with the scaling factor used for quantization. + + Args: + x: The input tenosr with ndim >= 2. + group_size: The group size used for quantization. + eps: The minimum to avoid dividing zero. + dtype: The dype of output tensor. Note that only `torch.int8` + is supported for now. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the + scaling factor for quantization. + """ + assert (x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + #assert x.is_contiguous(), "`x` is not contiguous" + + iinfo = torch.iinfo(dtype) + int8_max = iinfo.max + int8_min = iinfo.min + + x_q = torch.empty_like(x, device=x.device, dtype=dtype) + N = group_size + + m = x.shape[0] + if m<=16: + config={"BLOCK":128,"s_num":1,"num_warps":1,"num_stages":1} + elif m<=256: + config={"BLOCK":1024,"s_num":8,"num_warps":4,"num_stages":1} + else: + config={"BLOCK":2048,"s_num":16,"num_warps":4,"num_stages":2} + + grid = lambda META: ( + triton.cdiv(x.numel(), META['BLOCK']), + ) + + x_s = torch.empty( + x.shape[:-1] + (x.shape[-1] // group_size, ), + device=x.device, + dtype=torch.float32, + ) + + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + _per_token_group_quant_int8[grid]( + x, + x_q, + x_s, + group_size, + # M, + # K, + # N, + x.numel(), + eps, + int8_min=int8_min, + int8_max=int8_max, + **config + ) + + return x_q, x_s + + +@triton.jit +def _w8a8_block_int8_matmul( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, + and store the result in output tensor `C`. + """ + + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + + # offs_bsn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_bsn = pid_n * BLOCK_SIZE_N // group_n + + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # a_ptrs = A + (offs_am[:, None] * stride_am) + # b_ptrs = B + (offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + # offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + a = tl.load(a_ptrs, + mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, + other=0.0) + b = tl.load(b_ptrs, + mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, + other=0.0) + + + accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :] + + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +def matmul_block_int8( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype = torch.float16, + bias:Optional[torch.Tensor]=None, + config:Optional[dict] = None +) -> torch.Tensor: + """This function performs matrix multiplication with block-wise + quantization. + + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should be + 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + + Returns: + torch.Tensor: The result of matmul. + """ + + block_n, block_k = block_size[0], block_size[1] + M = A.numel() // A.shape[-1] + N, K = B.shape + + C_shape = A.shape[:-1] + (N, ) + C = A.new_empty(C_shape, dtype=output_dtype) + + if config is None: + if M<=64: + config = { + "BLOCK_SIZE_M": 16, #64 + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_stages": 0, + } + elif M<128: + config = { + "BLOCK_SIZE_M": 32, #64 + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_stages": 0, + } + elif M<=256: + config = { + "BLOCK_SIZE_M": 64, #64 + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_stages": 0, + } + else : + config = { + "BLOCK_SIZE_M": 64, #64 + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 0, + } + + def grid(META): + return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * + triton.cdiv(N, META["BLOCK_SIZE_N"]), ) + + + _w8a8_block_int8_matmul[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + **config, + ) + + if bias is not None: + C = C + bias + + return C + +def apply_w8a8_block_int8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + config: Optional[dict] = None, +) -> torch.Tensor: + assert input_scale is None + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = per_token_group_quant_int8(input_2d, block_size[1]) + + output = matmul_block_int8( + q_input, weight, x_scale, weight_scale, block_size, + output_dtype=input.dtype,config=config, + ) + + if bias is not None: + output = output + bias + return output.to(dtype=input.dtype).view(*output_shape) + + +def input_to_int8( + x: torch.Tensor, dtype: torch.dtype = torch.int8 +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function quantizes input values to + int8 values with tensor-wise quantization. + """ + iinfo = torch.iinfo(dtype) + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) + int8_min, int8_max = iinfo.min, iinfo.max + scale = int8_max / amax + x_scl_sat = (x * scale).clamp(min=int8_min, max=int8_max) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + + +def block_dequant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: List[int], +) -> torch.Tensor: + """This function conducts block-wise dequantization. + The inputs are block-wise quantization tensor `x_q_block`, + block-wise quantization scale and the block size. + The outputs are dequantized tensor. + """ + block_n, block_k = block_size[0], block_size[1] + n, k = x_q_block.shape + n_tiles = (n + block_n - 1) // block_n + k_tiles = (k + block_k - 1) // block_k + assert n_tiles == x_s.shape[0] + assert k_tiles == x_s.shape[1] + + x_dq_block = x_q_block.to(torch.float32) + + for i in range(k_tiles): + for j in range(n_tiles): + x_dq_block[ + j * block_n : min((j + 1) * block_n, n), + i * block_k : min((i + 1) * block_k, k), + ] *= x_s[j][i] + + return x_dq_block + +@triton.jit +def matmul_kernel( + # Pointers to matrices + a_ptr, as_ptr, b_ptr, bs_ptr, c_ptr, + # Matrix dimensions + M, N, K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` + # by to get the element one row down (A has M rows). + stride_am, stride_ak, + stride_asm, + stride_bk, stride_bn, + stride_bsn, + stride_cm, stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, SPLIT_K: tl.constexpr, +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + # See above `L2 Cache Optimizations` section for details. + pid = tl.program_id(axis=0) + pid_sp_k = tl.program_id(axis=1) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + # See above `Pointer Arithmetics` section for details + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = pid_sp_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + as_ptrs = as_ptr + offs_am * stride_asm + bs_ptrs = bs_ptr + offs_bn * stride_bsn + a_scale = tl.load(as_ptrs, mask=offs_am < M, other=0.0) + b_scale = tl.load(bs_ptrs, mask=offs_bn < N, other=0.0) + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.int32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K * SPLIT_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K * SPLIT_K, other=0.0) + # We accumulate along the K dimension. + accumulator += tl.dot(a, b) + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak + b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk + # You can fuse arbitrary activation functions here + # while the accumulator is still in FP32! + if SPLIT_K == 1: + c = (accumulator.to(tl.float32) * a_scale[:, None] * b_scale[None, :]).to(c_ptr.dtype.element_ty) + else: + c = (accumulator.to(tl.float32) * a_scale[:, None] * b_scale[None, :])#.to(c_ptr.dtype.element_ty) + # ----------------------------------------------------------- + # Write back the block of the output matrix C with masks. + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + if SPLIT_K == 1: + tl.store(c_ptrs, c, mask=c_mask) + else: + tl.atomic_add(c_ptrs, c, mask=c_mask) + +def matmul_int8(a, a_scale, b, b_scale,out_dtype,config): + # Check constraints. + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + M, K = a.shape + K, N = b.shape + + if config is None: + #print("no config") + if M<=32: + config= {'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4,'SPLIT_K': 1,'num_stages':0, 'num_warps':4} + elif M<=64: + config= {'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4,'SPLIT_K': 1,'num_stages':0, 'num_warps':4} + elif M<=128: + config= {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4,'SPLIT_K': 1,'num_stages':0, 'num_warps':4} + elif M<=1024: + config= {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4,'SPLIT_K': 1,'num_stages':0, 'num_warps':8} + else: + config= {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8,'SPLIT_K': 1,'num_stages':0, 'num_warps':8} + if config['SPLIT_K']==1: + c = torch.zeros((M, N), device=a.device, dtype=out_dtype) + else: + c = torch.zeros((M, N), device=a.device, dtype=torch.float32) + + grid = lambda META: ( + triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), + META['SPLIT_K'], + ) + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + + + matmul_kernel[grid]( + a, a_scale, b, b_scale, c, + M, N, K, + a.stride(0), a.stride(1), + a_scale.stride(0), + b.stride(0), b.stride(1), + b_scale.stride(0), + c.stride(0), c.stride(1), + **config + ) + + if config['SPLIT_K']!=1: + c=c.to(out_dtype) + return c + + diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m4096_graph.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m4096_graph.json new file mode 100644 index 00000000..ccbf2f8d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m4096_graph.json @@ -0,0 +1,662 @@ +[ + { + "shape_id": "hy3_tp8_qkv_proj_m4096", + "operator": "qkv_proj", + "M": 4096, + "N": 1280, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 19381.382751464844, + "p90_us": 20706.951904296875, + "min_us": 15527.166748046875, + "max_us": 33025.71716308594, + "samples_us": [ + 33025.717, + 19930.211, + 15664.891, + 17588.31, + 18619.633, + 18823.849, + 21526.091, + 15859.666, + 19238.075, + 20260.184, + 15527.167, + 19733.803, + 19524.69, + 19786.154, + 20706.952, + 17930.505, + 20039.809, + 19157.535, + 16103.629, + 19555.281 + ] + } + }, + { + "shape_id": "hy3_tp8_o_proj_m4096", + "operator": "o_proj", + "M": 4096, + "N": 4096, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 14151.79214477539, + "p90_us": 15978.352355957031, + "min_us": 10805.253601074219, + "max_us": 17007.850646972656, + "samples_us": [ + 14512.738, + 17007.851, + 11326.995, + 10805.254, + 13926.711, + 14670.845, + 14319.626, + 14161.777, + 13924.182, + 12676.867, + 13472.328, + 11711.237, + 14111.505, + 14375.624, + 14141.808, + 16150.25, + 14195.662, + 14111.92, + 15978.352, + 14732.059 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 384, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 16712.230682373047, + "p90_us": 17514.361572265625, + "min_us": 15547.166442871094, + "max_us": 18294.686889648438, + "samples_us": [ + 16706.294, + 18294.687, + 16718.167, + 17544.826, + 17514.362, + 16727.061, + 15547.166, + 16738.069, + 16576.123, + 16353.539, + 16856.978, + 16953.551, + 16529.98, + 16421.12, + 16027.951, + 16203.271, + 16812.18, + 16739.862, + 16485.535, + 16286.533 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 4096, + "K": 192, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 10762.615585327148, + "p90_us": 12987.09716796875, + "min_us": 9852.87094116211, + "max_us": 13697.439575195312, + "samples_us": [ + 13056.47, + 10556.271, + 11678.119, + 9955.428, + 10320.246, + 10068.512, + 13697.44, + 11907.071, + 10484.849, + 10555.374, + 10817.062, + 12246.387, + 10072.864, + 10708.17, + 11191.416, + 12355.088, + 12987.097, + 10983.04, + 9852.871, + 9935.653 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_m4096", + "operator": "qkv_proj", + "M": 4096, + "N": 1280, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 41000.61798095703, + "p90_us": 41881.70166015625, + "min_us": 40260.833740234375, + "max_us": 42669.35729980469, + "samples_us": [ + 40708.112, + 41134.851, + 41381.177, + 41881.702, + 41801.389, + 42043.939, + 42669.357, + 40753.168, + 40260.834, + 41577.524, + 40839.377, + 40447.772, + 40937.708, + 41451.099, + 40329.794, + 40462.268, + 41187.875, + 40845.648, + 40568.027, + 41063.528 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_and_indexer_qk_m4096", + "operator": "qkv_proj_and_indexer_qk", + "M": 4096, + "N": 1536, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 35604.698181152344, + "p90_us": 45884.57946777344, + "min_us": 27749.08447265625, + "max_us": 46322.772216796875, + "samples_us": [ + 46322.772, + 45884.579, + 45077.631, + 45431.604, + 45714.41, + 44870.535, + 45646.893, + 46203.641, + 36398.19, + 27957.782, + 32282.748, + 31533.688, + 35115.964, + 29992.844, + 32868.619, + 27749.084, + 32278.049, + 31587.637, + 36093.433, + 28834.07 + ] + } + }, + { + "shape_id": "minimax_tp8_o_proj_m4096", + "operator": "o_proj", + "M": 4096, + "N": 6144, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 19897.51739501953, + "p90_us": 20350.270080566406, + "min_us": 16113.107299804688, + "max_us": 22143.16864013672, + "samples_us": [ + 18625.305, + 16587.906, + 19008.813, + 19965.579, + 19805.585, + 20406.717, + 20196.259, + 20116.774, + 22143.169, + 19389.856, + 17160.941, + 20140.356, + 19829.456, + 20350.27, + 20321.152, + 19153.064, + 16113.107, + 18237.416, + 20135.303, + 20016.586 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 768, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 31882.257080078125, + "p90_us": 32810.479736328125, + "min_us": 28815.545654296875, + "max_us": 33757.51953125, + "samples_us": [ + 28815.546, + 32810.48, + 33757.52, + 33088.87, + 31372.415, + 32189.896, + 31768.692, + 31452.863, + 31608.185, + 32180.679, + 31995.822, + 31029.648, + 32477.213, + 32050.54, + 31519.229, + 31712.759, + 32021.133, + 31480.319, + 31442.175, + 32122.824 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 6144, + "K": 384, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 22389.16244506836, + "p90_us": 23217.950439453125, + "min_us": 19888.59405517578, + "max_us": 26559.368896484375, + "samples_us": [ + 26559.369, + 23217.95, + 23532.243, + 22366.364, + 22783.244, + 22659.154, + 22893.546, + 20970.54, + 22374.107, + 21844.141, + 21321.919, + 21402.269, + 22404.218, + 19888.594, + 23014.03, + 22565.503, + 20996.107, + 22041.0, + 22622.835, + 20145.161 + ] + } + }, + { + "shape_id": "glm52_tp8_fused_qkv_a_proj_m4096", + "operator": "fused_qkv_a_proj", + "M": 4096, + "N": 2624, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 68661.41662597656, + "p90_us": 70210.54077148438, + "min_us": 47223.052978515625, + "max_us": 70451.49536132812, + "samples_us": [ + 69346.32, + 68621.863, + 68886.328, + 70219.086, + 68263.092, + 70210.541, + 70451.495, + 68671.973, + 69810.4, + 69454.895, + 69113.562, + 68650.861, + 69821.857, + 58982.715, + 47223.053, + 49626.874, + 51567.096, + 52752.594, + 53318.494, + 49815.668 + ] + } + }, + { + "shape_id": "glm52_tp8_q_b_proj_m4096", + "operator": "q_b_proj", + "M": 4096, + "N": 2048, + "K": 2048, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 14385.5712890625, + "p90_us": 15009.501647949219, + "min_us": 11206.52847290039, + "max_us": 16653.09295654297, + "samples_us": [ + 14075.645, + 12226.59, + 11618.449, + 14454.031, + 14636.586, + 14437.842, + 14322.932, + 14592.075, + 16653.093, + 15187.801, + 11206.528, + 14112.669, + 14517.839, + 15009.502, + 14548.204, + 12725.101, + 11365.851, + 13823.718, + 14589.165, + 14333.301 + ] + } + }, + { + "shape_id": "glm52_tp8_kv_b_proj_m4096", + "operator": "kv_b_proj", + "M": 4096, + "N": 3584, + "K": 512, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 6939.938735961914, + "p90_us": 7771.638488769531, + "min_us": 5538.338470458984, + "max_us": 8975.7568359375, + "samples_us": [ + 6990.929, + 8975.757, + 8890.351, + 5578.849, + 5567.842, + 5553.794, + 7100.013, + 7378.083, + 7270.152, + 7218.505, + 6882.964, + 7771.638, + 6567.775, + 6888.949, + 5586.017, + 5567.042, + 5538.338, + 6640.636, + 7402.626, + 7460.033 + ] + } + }, + { + "shape_id": "glm52_tp8_o_proj_m4096", + "operator": "o_proj", + "M": 4096, + "N": 6144, + "K": 2048, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 54775.79650878906, + "p90_us": 56661.138916015625, + "min_us": 32334.09423828125, + "max_us": 57705.450439453125, + "samples_us": [ + 36803.061, + 38816.177, + 32334.094, + 36348.23, + 36142.572, + 57260.956, + 56661.139, + 57705.45, + 54939.215, + 56308.258, + 55439.417, + 54385.504, + 55639.667, + 55714.325, + 55518.042, + 53696.216, + 54612.378, + 54527.069, + 54463.104, + 55575.513 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 512, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 21254.589080810547, + "p90_us": 21705.372619628906, + "min_us": 20394.827270507812, + "max_us": 23713.912963867188, + "samples_us": [ + 20852.121, + 20782.909, + 20957.654, + 21016.469, + 21695.837, + 20957.654, + 20914.935, + 21597.601, + 22041.522, + 23713.913, + 21585.057, + 21181.616, + 21223.471, + 21285.707, + 21705.373, + 21578.625, + 21093.587, + 20394.827, + 21346.121, + 21464.134 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 6144, + "K": 256, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 17700.262451171875, + "p90_us": 19782.623291015625, + "min_us": 15626.028442382812, + "max_us": 22184.974670410156, + "samples_us": [ + 17591.882, + 19158.612, + 19004.921, + 16231.48, + 17128.282, + 19256.914, + 19557.191, + 19782.623, + 22184.975, + 19914.011, + 17808.643, + 18663.495, + 15845.126, + 17278.868, + 17352.115, + 17899.68, + 16849.059, + 15626.028, + 17149.53, + 17040.669 + ] + } + } +] \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_eager.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_eager.json new file mode 100644 index 00000000..48a063d7 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_eager.json @@ -0,0 +1,281 @@ +{ + "scope": "gemm_out", + "triton_source": "/tmp/baseline_int8_utils.py", + "M": 4096, + "warmups": 10, + "samples": 20, + "launches_per_sample": 5, + "mode": "eager", + "results": [ + { + "shape_id": "tp4_wqkv_a_m4096", + "operator": "wqkv_a", + "M": 4096, + "N": 1536, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 13287.435150146484, + "p90_us": 13315.707397460938, + "min_us": 13170.518493652344, + "max_us": 13707.110595703125, + "samples_us": [ + 13226.106, + 13308.22, + 13230.266, + 13292.188, + 13315.707, + 13190.199, + 13218.904, + 13296.187, + 13313.756, + 13294.81, + 13282.683, + 13296.956, + 13353.532, + 13170.518, + 13203.448, + 13707.111, + 13215.833, + 13229.721, + 13219.096, + 13292.923 + ] + } + }, + { + "shape_id": "tp4_wq_b_m4096", + "operator": "wq_b", + "M": 4096, + "N": 8192, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 20623.723602294922, + "p90_us": 20699.789428710938, + "min_us": 20575.30517578125, + "max_us": 21157.14569091797, + "samples_us": [ + 20627.628, + 20606.154, + 20615.851, + 20721.645, + 20621.803, + 20619.339, + 20625.644, + 21157.146, + 20654.027, + 20667.628, + 20578.441, + 20637.163, + 20699.789, + 20575.305, + 20658.06, + 20611.241, + 20650.603, + 20599.562, + 20597.418, + 20594.986 + ] + } + }, + { + "shape_id": "tp4_indexer_wq_b_m4096", + "operator": "indexer.wq_b", + "M": 4096, + "N": 8192, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 20654.634857177734, + "p90_us": 20685.16387939453, + "min_us": 20580.072021484375, + "max_us": 21134.840393066406, + "samples_us": [ + 20639.98, + 20622.218, + 20637.993, + 20662.732, + 20598.09, + 20663.66, + 20658.06, + 20632.042, + 20633.29, + 20651.21, + 20665.291, + 21134.84, + 20685.164, + 20580.072, + 20664.265, + 20590.665, + 20696.204, + 20669.418, + 20670.186, + 20642.058 + ] + } + }, + { + "shape_id": "tp4_wo_b_m4096", + "operator": "wo_b", + "M": 4096, + "N": 4096, + "K": 2048, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 19915.618896484375, + "p90_us": 19963.380432128906, + "min_us": 19849.97100830078, + "max_us": 20379.873657226562, + "samples_us": [ + 19945.428, + 19849.971, + 20379.874, + 19900.435, + 19941.397, + 19923.093, + 19873.17, + 19899.539, + 19913.683, + 19905.396, + 19963.38, + 19887.508, + 19876.372, + 19917.555, + 19924.883, + 20352.672, + 19895.027, + 19882.227, + 19918.611, + 19960.02 + ] + } + }, + { + "shape_id": "tp4_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 1024, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 8774.919891357422, + "p90_us": 8838.265228271484, + "min_us": 8717.046356201172, + "max_us": 9269.158172607422, + "samples_us": [ + 8721.815, + 8792.313, + 8838.265, + 8745.751, + 8775.352, + 8793.529, + 8774.488, + 8747.095, + 8717.046, + 8845.21, + 8749.047, + 8780.152, + 8805.529, + 8770.008, + 8790.744, + 8806.841, + 8736.246, + 8773.912, + 9269.158, + 8771.128 + ] + } + }, + { + "shape_id": "tp4_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 4096, + "K": 512, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 5582.3015213012695, + "p90_us": 5592.509841918945, + "min_us": 5555.356979370117, + "max_us": 5619.230270385742, + "samples_us": [ + 5585.566, + 5592.51, + 5555.357, + 5586.398, + 5571.326, + 5581.853, + 5582.75, + 5586.045, + 5586.43, + 5582.973, + 5576.606, + 5571.358, + 5567.805, + 5563.549, + 5576.189, + 5619.23, + 5587.838, + 5566.749, + 5568.093, + 5595.134 + ] + } + } + ] +} \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_graph.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_graph.json new file mode 100644 index 00000000..a6464631 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_graph.json @@ -0,0 +1,281 @@ +{ + "scope": "gemm_out", + "triton_source": "/tmp/baseline_int8_utils.py", + "M": 4096, + "warmups": 10, + "samples": 20, + "launches_per_sample": 5, + "mode": "graph", + "results": [ + { + "shape_id": "tp4_wqkv_a_m4096", + "operator": "wqkv_a", + "M": 4096, + "N": 1536, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 13247.252655029297, + "p90_us": 13320.054626464844, + "min_us": 13144.914245605469, + "max_us": 13643.904113769531, + "samples_us": [ + 13336.472, + 13262.518, + 13314.262, + 13197.523, + 13320.055, + 13643.904, + 13300.246, + 13269.365, + 13284.662, + 13253.653, + 13234.421, + 13224.02, + 13218.292, + 13228.885, + 13147.346, + 13240.852, + 13214.676, + 13229.396, + 13144.914, + 13285.077 + ] + } + }, + { + "shape_id": "tp4_wq_b_m4096", + "operator": "wq_b", + "M": 4096, + "N": 8192, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 20590.225219726562, + "p90_us": 20648.818969726562, + "min_us": 20520.016479492188, + "max_us": 21081.182861328125, + "samples_us": [ + 20596.434, + 21081.183, + 20520.016, + 20586.673, + 20577.713, + 20597.105, + 20556.367, + 20648.819, + 20589.33, + 20599.345, + 20590.61, + 20578.609, + 20576.752, + 20544.048, + 21078.076, + 20597.136, + 20589.841, + 20524.046, + 20621.681, + 20624.913 + ] + } + }, + { + "shape_id": "tp4_indexer_wq_b_m4096", + "operator": "indexer.wq_b", + "M": 4096, + "N": 8192, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 20602.671813964844, + "p90_us": 20697.555541992188, + "min_us": 20566.477966308594, + "max_us": 21020.53985595703, + "samples_us": [ + 20609.616, + 20639.314, + 20574.065, + 20697.556, + 20587.505, + 21020.54, + 20601.36, + 20644.786, + 20570.096, + 20603.824, + 20603.375, + 20586.96, + 20610.159, + 20601.968, + 20566.478, + 20597.328, + 20613.904, + 20600.912, + 20981.018, + 20566.574 + ] + } + }, + { + "shape_id": "tp4_wo_b_m4096", + "operator": "wo_b", + "M": 4096, + "N": 4096, + "K": 2048, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 19881.948852539062, + "p90_us": 19916.25213623047, + "min_us": 19814.553833007812, + "max_us": 20356.455993652344, + "samples_us": [ + 19870.557, + 19886.205, + 19901.534, + 19921.021, + 19859.677, + 19874.237, + 19911.804, + 19833.531, + 19829.051, + 20356.456, + 19863.837, + 19898.973, + 19877.693, + 19857.532, + 19838.747, + 19896.861, + 19889.693, + 19896.861, + 19814.554, + 19916.252 + ] + } + }, + { + "shape_id": "tp4_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 1024, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 8790.192031860352, + "p90_us": 8815.552520751953, + "min_us": 8717.342376708984, + "max_us": 9243.083953857422, + "samples_us": [ + 8815.553, + 8757.567, + 8770.112, + 9243.084, + 8726.078, + 8750.815, + 8791.904, + 8795.233, + 8717.342, + 8759.872, + 8832.194, + 8806.433, + 8799.168, + 8779.68, + 8723.199, + 8802.465, + 8798.4, + 8788.48, + 8806.753, + 8740.222 + ] + } + }, + { + "shape_id": "tp4_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 4096, + "K": 512, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 5545.821189880371, + "p90_us": 5561.069488525391, + "min_us": 5528.492736816406, + "max_us": 6035.929489135742, + "samples_us": [ + 5570.702, + 5552.142, + 5551.213, + 5546.317, + 5534.701, + 5537.453, + 5546.029, + 5541.453, + 5545.197, + 5557.133, + 5537.933, + 5528.493, + 5561.069, + 5544.845, + 5552.333, + 5551.566, + 5545.613, + 5533.453, + 5544.685, + 6035.929 + ] + } + } + ] +} \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/planner_parity.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/planner_parity.py new file mode 100644 index 00000000..046255fd --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/planner_parity.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python3 +"""Offline parity analysis: planner (v0) vs the legacy round menu. + +Read-only tool over historical task repos under ``kernel-repos`` (and optionally +workspace ``iterations``). For each recorded round it reconstructs the state +(prior same-worker/same-shape records + pmc evidence), reproduces what the +legacy menu (``w8a8_round_strategy``) would mandate, asks what the +state-conditioned ``planner`` would pick, maps both to coarse plan families, and +reports agreement/divergence by M regime and iteration bucket, plus samples. + +Caveats (read the output with these in mind): +- State reconstruction is approximate: pmc evidence is frequently absent in + records, so the planner often falls back to its P4 legacy approximation. +- The menu text -> family mapping is a keyword heuristic. +- This tool never writes into the repos; it only prints / dumps a report. + +Usage: + python3 tools/planner_parity.py [--kernel-repos ROOT] [--limit-repos N] + [--out /tmp/planner_parity_report.json] +""" + +from __future__ import annotations + +import argparse +import collections +import json +import re +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +# Make the plugin importable when run as a script from the plugin tree. +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) # .../dcu_kernel_auto_opt +sys.path.insert(0, str(Path(__file__).resolve().parents[4])) # MetaInfer/ + +from metainfer.tasks.dcu_kernel_auto_opt.orchestrator import planner as planner_mod # noqa: E402 +from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.planner import choose_plan_from_history # noqa: E402 +from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.prompts import w8a8_round_strategy # noqa: E402 +from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.w8a8_pipeline import isa_round_policy # noqa: E402 + +_MENU_FAMILY_RULES: List[Tuple[str, str]] = [ + # Explicit markers first so generic words inside round prose do not win. + ("faster but incorrect", "repair"), + ("faster_wrong", "repair"), + ("repair", "repair"), + ("conditional inline-asm", "inline_asm"), + ("inline asm", "inline_asm"), + ("ISA-guided", "isa_guided"), + ("ISA", "isa"), + ("split-K", "grid_splitk"), + ("split_k", "grid_splitk"), + ("consolidat", "consolidate"), + ("epilogue", "epilogue"), + ("Architecture round", "architecture"), + ("Architecture/pipeline round", "architecture"), + ("macro-tile", "architecture"), + ("tile-shape", "architecture"), + ("geometry", "architecture"), + ("tile", "architecture"), + ("architecture", "architecture"), + ("packing", "memory_layout"), + ("packed", "memory_layout"), + ("swizzle", "memory_layout"), + ("bank", "memory_layout"), + ("layout", "memory_layout"), + ("double buff", "pipeline"), + ("prefetch", "pipeline"), + ("staging", "pipeline"), + ("barrier", "pipeline"), + ("pipeline", "pipeline"), + ("occupancy", "occupancy"), + ("waves per block", "occupancy"), + ("VGPR", "occupancy"), + ("register", "occupancy"), + ("bootstrap", "bootstrap"), + ("scalar", "bootstrap"), +] + +_PLAN_TO_FAMILY = { + "repair_faster_wrong": "repair", + "fix_build": "repair", + "retry_same": "repair", + "bootstrap_correctness": "bootstrap", + "establish_arch": "architecture", + "architecture_explore": "architecture", + "grid_splitk": "grid_splitk", + "pipeline_tune": "pipeline", + "memory_layout": "memory_layout", + "occupancy_resource": "occupancy", + "epilogue_fusion": "epilogue", + "isa_guided_hip": "isa_guided", + "conditional_inline_asm": "inline_asm", + "consolidate": "consolidate", +} + + +def menu_family(text: str) -> str: + lowered = text.lower() + for needle, family in _MENU_FAMILY_RULES: + if needle in lowered: + return family + return "unknown" + + +def iter_number_from_path(path: Path) -> int: + m = re.search(r"iteration(\d+)", path.name) + if m: + return int(m.group(1)) + # parent dirs like .../iteration13/iteration.json + m = re.search(r"iteration(\d+)", path.parent.name) + return int(m.group(1)) if m else -1 + + +def load_records(root: Path, limit_repos: Optional[int]) -> List[Tuple[str, Dict[str, Any]]]: + """Return (repo_dir_name, record) for every candidate iteration.json.""" + records: List[Tuple[str, Dict[str, Any]]] = [] + repos = sorted(p for p in root.iterdir() if p.is_dir()) + if limit_repos: + repos = repos[:limit_repos] + for repo in repos: + for path in sorted( + p for p in repo.rglob("iteration.json") + if "candidates" in p.parts + ): + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + continue + if isinstance(data, dict): + records.append((repo.name, data)) + return records + + +def analyze(root: Path, limit_repos: Optional[int]) -> Dict[str, Any]: + records = load_records(root, limit_repos) + chains: Dict[Tuple[str, str, str], List[Dict[str, Any]]] = collections.defaultdict(list) + for repo, rec in records: + worker = str(rec.get("worker_id", "?")) + shape = rec.get("shape_id", "?") + iteration = iter_number_from_path(Path(str(rec.get("record_path", "iteration0")))) if False else int(rec.get("iteration", 0)) + chains[(repo, worker, shape)].append( + {"record": rec, "iteration": iteration} + ) + + stats = collections.Counter() + buckets: Dict[str, collections.Counter] = collections.defaultdict(collections.Counter) + samples: List[Dict[str, Any]] = [] + + for (repo, worker, shape_id), entries in chains.items(): + entries.sort(key=lambda e: e["iteration"]) + history: List[Dict[str, Any]] = [] + for entry in entries: + rec = entry["record"] + iteration = entry["iteration"] + shape = rec.get("shape", {}) or {} + m = int(shape.get("M", 0) or 0) + regime = "prefill(M>=128)" if m >= 128 else ("m16" if m >= 16 else "small(M<16)") + pmc = rec.get("pmc_evidence") or {} + + try: + policy = isa_round_policy( + iteration=iteration, + max_iterations=10, + history=history, + ) + menu_text = w8a8_round_strategy( + shape, iteration, history, pmc, + max_iterations=10, isa_policy=policy, + ) + except Exception as exc: # noqa: BLE001 + stats["menu_error"] += 1 + history.append(rec) + continue + + menu_fam = menu_family(menu_text) + try: + plan_id = choose_plan_from_history( + history, pmc, iteration=iteration, max_iterations=10, + shape=shape, + compiler_limitation_confirmed=( + policy.get("phase") == "conditional_inline_asm" + ), + ) + except Exception as exc: # noqa: BLE001 + stats["planner_error"] += 1 + history.append(rec) + continue + + plan_fam = _PLAN_TO_FAMILY.get(plan_id, "unknown") + agree = (menu_fam == plan_fam) + stats["total"] += 1 + stats["agree" if agree else "disagree"] += 1 + bucket_key = f"{regime}|iter{min(iteration, 12)}" + buckets[bucket_key]["agree" if agree else "disagree"] += 1 + + if not agree and len(samples) < 40: + samples.append({ + "repo": repo, + "shape_id": shape_id, + "iteration": iteration, + "regime": regime, + "menu_family": menu_fam, + "plan_family": plan_fam, + "plan_id": plan_id, + "hypothesis": str(rec.get("hypothesis") or "")[:200], + }) + history.append(rec) + + total = stats.get("total", 0) + return { + "corpus": {"repos_scanned": len(set(r for r, _ in records)) if records else 0, + "records": len(records)}, + "stats": dict(stats), + "agree_rate": (stats.get("agree", 0) / total) if total else None, + "buckets": { + k: dict(v) for k, v in sorted(buckets.items()) + }, + "samples": samples, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--kernel-repos", type=Path, + default=Path(__file__).resolve().parents[2].parent.parent.parent / "kernel-repos") + parser.add_argument("--limit-repos", type=int, default=None) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + if not args.kernel_repos.is_dir(): + print(f"kernel-repos not found at {args.kernel_repos}", file=sys.stderr) + return 2 + + report = analyze(args.kernel_repos, args.limit_repos) + total = report["stats"].get("total", 0) + print(f"corpus: {report['corpus']['records']} records / {report['corpus']['repos_scanned']} repos") + print(f"agree={report['stats'].get('agree',0)} disagree={report['stats'].get('disagree',0)} " + f"agree_rate={report['agree_rate']:.1%}" if report["agree_rate"] is not None + else f"agree={report['stats'].get('agree',0)} disagree={report['stats'].get('disagree',0)}") + print("--- buckets (regime|iteration): agree/disagree ---") + for k, v in list(report["buckets"].items())[:40]: + print(f" {k}: {v}") + print("--- divergence samples ---") + for s in report["samples"][:10]: + print(f" [{s['repo']} {s['shape_id']} it{s['iteration']} {s['regime']}] " + f"menu={s['menu_family']} plan={s['plan_family']} ({s['plan_id']})") + print(f" hypo: {s['hypothesis'][:150]}") + + if args.out: + args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2), + encoding="utf-8") + print(f"report -> {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) 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 index 143caa0b..c6a575d2 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=tp8_indexer_wq_b_m16 commit=9b5b710d0a2c4cd7ac22e2037dfe8735059eec0b added=2026-08-26 +// baseline_us=48.274 speedup=2.70443 // 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) 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 index 487c6b3a..2cce19a8 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=tp8_shared_down_proj_m16 commit=119c281feefbb681f88dba89897484424bb07937 added=2026-08-26 +// baseline_us=14.849 speedup=1.75665 // 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. 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 index 15e74f81..4c63671f 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=tp8_shared_gate_up_proj_m16 commit=ea0f87b2e586e93141907abbdfc541b1585f4c50 added=2026-08-26 +// baseline_us=59.21 speedup=5.75973 // 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. 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 index f420f456..c67c8445 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=tp8_wo_b_m16 commit=eb0fbbad60c58b1cf09e6ceff6cb9aaa68a2c1cb added=2026-08-26 +// baseline_us=32.137 speedup=3.27594 // 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). 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 index ff0598e7..0b19eb88 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=tp8_wq_b_m16 commit=bfaef2fc76d0e4a095b2ffae2b466f2863511d43 added=2026-08-26 +// baseline_us=32.137 speedup=2.19066 // 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) 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 index bfa665ac..a383ad61 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=tp8_wqkv_a_m16 commit=b67e227c20b3349d46020f9f52cd2484b2142ff3 added=2026-08-26 +// baseline_us=66.591 speedup=3.8052 // median_us=17.5 p90_us=17.52 // source=hy3-dsh-tp8-m16-1-adf021ab // MetaInfer W8A8 GEMM - worker_0 HIP implementation (gfx928). 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 index a7e83daf..7748f78d 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_fused_qkv_a_proj_m16 commit=bf306c40c48713a81efbd7ace430c9cf21cbb429 added=2026-08-31 +// baseline_us=97.99 speedup=2.33476 // 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). 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 index 33985f73..93be3411 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_kv_b_proj_m16 commit=a150485d1a63d39ca9251b13d3da304013431efd added=2026-08-31 +// baseline_us=18.833 speedup=2.1881 // median_us=8.607 p90_us=9.766 // source=glm5-2-dsh-tp8-m16-test1-cb4d262c // csrc/w8a8_gemm_hip.hip 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 index 3ec28b39..caab0383 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_o_proj_m16 commit=c1cfe1975c2fcb8d4bc23c86e8026afae225a4c4 added=2026-08-31 +// baseline_us=66.98 speedup=2.79083 // median_us=24 p90_us=24.02 // source=glm5-2-dsh-tp8-m16-test1-cb4d262c // MetaInfer INT8 W8A8 GEMM for Hygon K500SM_AI / gfx928. 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 index 8d888111..ad036b39 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_q_b_proj_m16 commit=8a9a28e409277a6f266a59e0c053a3d2bb516f12 added=2026-08-31 +// baseline_us=42.035 speedup=2.97698 // median_us=14.12 p90_us=14.14 // source=glm5-2-dsh-tp8-m16-test1-cb4d262c // csrc/w8a8_gemm_hip.hip 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 index fdf64fd2..b841e219 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_shared_down_proj_m16 commit=42d49afdb24af21a5480786c88ca11a84d57f44a added=2026-08-31 +// baseline_us=18.177 speedup=2.20621 // 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). 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 index 3455553e..b33439e2 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_shared_gate_up_proj_m16 commit=fa6ea50630e35036b22141edc48f72975fca041d added=2026-08-31 +// baseline_us=99.318 speedup=9.17911 // 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). 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 index aaa4a2a8..f5387ae1 100644 --- 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 @@ -1,8 +1,8 @@ -// @@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). +// @@variant shape=glm_tp8_fused_qkv_a_proj_m4096 commit=99b9a733eddfc7196f3d186b80affe4324351ac4 added=2026-09-11 +// median_us=883.7 p90_us=885.5 speedup=77.7 baseline_us=6.866e+04 +// source=test-9-11-1-b01955b3/iteration_001/glm_tp8_fused_qkv_a_proj_m4096 +// MetaInfer W8A8 INT8 GEMM -- round-1 architecture establishment for +// Hygon K500SM_AI / gfx928 (worker_0, physical GPU 0). // // Assigned shape: // glm_tp8_fused_qkv_a_proj_m4096 : M=4096, N=2624, K=6144 @@ -11,255 +11,446 @@ // 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). +// Round 1 (mandatory establish_arch) -- packed_weight architecture family: // -// 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. +// Bottleneck hypothesis, from the trusted PMC of the current best source +// (the correctness-first bootstrap: w8a8_dumma_prefill_tiled_kernel +// <128, 64, 128>, identity [K, N] weight, library du_load_matrix_sync +// loaders, LDS row strides 128 (A) / 64 (B), grid 41 x 32 = 1312 blocks, +// 256 threads, 2 blocks/CU): +// median 2920.88 us / p90 2962.82 us, lds_instructions 21,999,616, +// lds_bank_conflicts 274,071,552 (~12.5 per LDS instruction), +// lds_wait_instructions 50,181,819, valu_instructions 86,560,512, +// vmem_read_instructions 1,574,400, and a counter-derived operator HBM +// bandwidth of only 50.6 GB/s (147.8 MB per replay at a 90.7% L2 hit +// rate). The exact-shape stage loop is LDS-issue/conflict bound, not +// MMA- or HBM-bound: both LDS row strides alias every row onto one bank +// phase (128 mod 32 == 0, 64 mod 32 == 0), and the row_major matrix_b +// loader of du_load_matrix_sync lowers each m16n16k32 int8 B fragment to +// 8 ds_read_u8 at 16-byte spacing plus a byte-reassembly chain. // -// 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). +// Focused change (ONE architecture family -- packed_weight -- entirely +// inside csrc/w8a8_gemm_hip.hip): +// (1) launch_pack_w8a8_weight 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, same buffer allocation, so every captured +// address is unchanged). Every other (K, N) keeps the identity +// device-to-device copy. +// (2) the exact-shape guard dispatches the new +// w8a8_dumma_prefill_packedb_kernel<128, 64, 128>. It keeps the +// EXACT bootstrap geometry -- grid (N/64) x (M/128) = 41 x 32 = 1312 +// blocks, 256 threads = 4 gfx928 wavefronts, one 64 x 32 quadrant of +// eight m16n16k32 int32 accumulator fragments per wave, the same +// single-buffered 128-K LDS stage with two __syncthreads per stage, +// the same k0-outer / kk-inner ascending int32 accumulation order and +// the same coalesced bf16 epilogue -- and changes only the operand +// path to the packed-weight architecture: +// * B is staged n-major (b_tile[n_row][k]) from the packed [N, K] +// weight, so a col_major matrix_b fragment is 8 contiguous K +// bytes per lane and every staging vector is one coalesced +// 16-byte global read; +// * both LDS tiles use the bank-safe row stride StageK + 8 = 136 +// bytes (34 dwords): 8-byte aligned with an odd 8-byte count +// (136/8 = 17), so the 16 rows of each 16-lane ds_read_b64 phase +// land on 16 distinct bank pairs (the conflict-free 4-phase +// floor), and 136 % 16 == 8 keeps the two-half int64 staging +// form; +// * both fragment loaders are the explicit load_frag8: ONE aligned +// 8-byte load per lane per fragment (row = lane & 15, k-quarter = +// lane >> 4), replacing the byte-reassembly library path with a +// single ds_read_b64 while producing identical fragment values. +// (3) LDS/block grows 24,576 -> 26,112 B (A[128, 136] + B[64, 136]); +// 2 x 26,112 = 52,224 <= 65,536 B keeps TWO resident blocks/CU = +// 8 waves/CU, so grid_blocks (1312), waves_per_block (4) and the +// active CU count (120) are unchanged. +// (4) the scalar fallback decodes the n-major layout when +// (k, n) == (6144, 2624) with ascending kk, so the paired M=2 / +// M=16 API shapes with the same (N, K) stay byte-exact; the generic +// <64, 64, 128> identity arm is untouched. // -// Iteration 3 (mandated pipeline round: single vs double buffering across K -// tiles; retain double buffering only on ISA/PMC evidence of reduced VMEM +// Correctness: the same weight elements reach the same fragment slots in +// the same ascending k-chunk sequence (k0 outer over 48 stages of 128 K, +// kk inner ascending over four m16n16k32 steps), so every int32 +// accumulator -- and therefore every bf16 output byte -- is identical to +// the bootstrap kernel and to the scalar reference. +// +// Round 2 (mandated operand-reuse round: cooperative A+B LDS staging vs direct +// loads, quantified A/B reuse per macro-tile, vectorized coalesced global loads +// with a bank-safe LDS layout): +// +// Operand-reuse accounting for the accepted <128, 64, 128> macro-tile +// (M=4096, N=2624, K=6144): +// * per block the unique operand bytes are A = 128 x 6144 = 786,432 and +// B = 6144 x 64 = 393,216; the 41 N-blocks re-read every A row-slab +// (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): 1.55 GB/replay of staged +// global traffic; +// * inside one block each staged 16-B vector is re-read from LDS by the two +// wave columns (A) / the two wave rows (B), i.e. exactly 2x wave-level +// reuse per staged byte, so the LDS read traffic is 2x the staged bytes +// (4.72 GB/replay, ~23 B/cycle/CU at the measured 1.271 GHz -- far below +// the LDS bandwidth), and the counter-derived HBM rate is only +// 108.4 GB/s (147.8 MB/replay at a 90.7% L2 hit rate). The kernel is +// therefore NOT HBM- or LDS-bandwidth bound, so a direct per-wave global +// fragment path is rejected up front: it would issue 4x more vmem +// instructions (4 waves x 48 stages of 8-B fragment loads) and every 8-B +// fragment load would be a row-strided multi-line access instead of the +// coalesced 16-B staging vectors of the cooperative path. Cooperative +// A+B LDS staging is RETAINED -- the round improves it instead. +// +// Remaining exposed cost of the retained path: the VMEM round trip in front +// of every stage consume. The accepted source (PMC digest daf403b3..., +// profiles/.../iteration2/pmc.json: profiled 1368.793 us, grid 1312 blocks, +// 2 blocks/CU, 26,112 B LDS, lds_instructions 7,892,992 with +// lds_wait 4,535,391, lds_bank_conflicts 0, vmem_read 1,574,400, +// grbm_gui_active == grbm_count) issues per 128-K stage 24 wave-level +// global_load_dwordx4 + 24 ds_write2_b64 and then, per the exact-source ISA, +// strictly serializes global_load_dwordx4 -> s_waitcnt vmcnt(0) -> +// ds_write2_b64 -> s_barrier -> 24x ds_read_b64 -> 32x v_mmac -> s_barrier: +// with 8 resident waves/CU the GPU is 100% GRBM GUI-active yet issues only +// ~4.6% of its slots, i.e. the stage boundary -- not HBM, LDS bandwidth or +// MMA throughput -- sets the pace. +// +// Focused change (ONE mechanism -- the K stage pipeline -- entirely inside +// csrc/w8a8_gemm_hip.hip): the exact-shape guard now dispatches +// w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64>. It keeps the SAME +// macro-tile (grid 41 x 32 = 1312 blocks, 256 threads = 4 wavefronts, 64 x 32 +// quadrant of eight m16n16k32 int32 accumulators per wave), the SAME packed +// n-major B (16-B coalesced staging vectors), the SAME explicit load_frag8 +// fragment reads, the SAME bank-safe LDS stride form and the SAME coalesced +// epilogue -- and changes only how the K stages are staged: +// * the single-buffered 128-K stage (TWO __syncthreads per stage, full +// VMEM latency exposed in front of every consume) becomes a +// DOUBLE-buffered 64-K stage with ONE __syncthreads per stage: the +// stage-(s+1) global loads (2 int4 A + 1 int4 B per thread) are issued +// into staging VGPRs BEFORE the stage-s consume and flushed into the idle +// LDS buffer after it, so the stage-(s+1) VMEM latency overlaps the +// stage-s LDS/v_mmac work. +// +// 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. kLdsStride = 72 = +// StageK + 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 the accepted kernel -- +// and 72 % 16 == 8 keeps the two-half ds_write2_b64 staging form. 12 extra +// VGPR hold the prefetch payload (~72 -> ~84, still <= the 128/thread budget +// of 2 blocks/CU; 0 spills expected). +// +// 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 the accepted 48 x 128-K stages, with the identical +// element-to-slot fragment mapping, so every int32 accumulator -- and every +// bf16 output byte -- is bit-identical to the accepted kernel and to the +// scalar reference. The n-major pack, the scalar-fallback n-major decode, the +// generic <64, 64, 128> identity arm and every shape guard are untouched. +// +// Round 3 (mandated pipeline round: compare single and 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, +// +// * Mandatory comparison (exact trusted PMC/ISA of the two accepted arms): +// arm median us LDS B/blk blocks/CU VGPR +// single 128-K <128,64,128> 1363.078 26,112 2/8 waves 72 +// double 64-K <128,64,64> 1153.970 27,648 2/8 waves 96 +// A single 64-K arm is not built: with NO prefetch it can only add +// barriers (2 per 64-K stage = 192 per block vs the 128-K arm's 96) while +// exposing the same VMEM round trip, so it is strictly dominated and +// rejected by construction. +// RETAINED -- double buffering. Evidence: median 1363.078 -> 1153.970 us +// (-15.3%) and p90 1364.825 -> 1164.522 us (-14.7%); the exact-source PMC +// of the double-buffered arm (digest ee609c3f..., iteration3/pmc.json) +// shows lds_wait_instructions 4,562,210 -> 3,590,796 (-21.3%) with +// grbm_gui_active == grbm_count (100% GUI-active), LDS 27,648 B/block +// (2 x 27,648 = 55,296 <= 65,536 so the trusted occupancy-probe split 2 +// is unchanged = 8 waves/CU), arch_vgpr 96 with 0 spills (<= the +// 128/thread budget of 2 blocks/CU) and sgpr 32. The ISA confirms the +// mechanism: the stage-(s+1) global_load_dwordx4s are issued BEFORE the +// stage-s ds_read_b64 / v_mmac burst and the vmcnt(0) wait lands on the +// ds_write2_b64 flush AFTER the MMACs, while the 128-K arm serializes +// global_load_dwordx4 -> s_waitcnt vmcnt(0) -> ds_write2_b64 -> s_barrier +// with zero compute overlap. So the prefetch converted exposed VMEM +// latency into overlapped work without trading away occupancy. +// LDS-growth verdict: the +1,536 B/block stays inside the 2-block budget +// and lds_wait FELL (the LDS path is not harmed in the net); the new +// lds_bank_conflicts 12,091,392 is exactly 6 x 4 x 4 x 96 x 1312 = one +// conflict per 16-lane phase of every merged ds_read2_b64, i.e. the +// bank-preserving alias of the +16-row second fragment offset +// (16 x 72 = 1152 B = 9 x 128 B; for ANY 8-aligned uniform row stride S, +// 16 rows = 4S dwords with 4S % 32 == 0). The unmerged ds_read_b64 +// control at the same 72/136-byte strides measured 0 conflicts, so this +// is an alias of the merged read form, not a layout defect. +// Every other K-tile pipeline is rejected inside the 65,536 B budget: +// StageK=96 double buffer = 39,936 B/block, StageK=128 double buffer = +// 52,224 B/block and 3 x 64-K = 41,472 B/block all exceed the 32,768 +// B/block that 2 resident blocks/CU permit (each would drop to 1 +// block/CU = 4 waves/CU, a 2x occupancy loss, explicitly excluded), and a +// 32-K stage leaves only 128 B staging vectors per block (half the 256 +// threads idle) while carrying a single m16n16k32 step per stage (too +// little consume work to amortize the barrier). StageK=64 x 2 buffers is +// therefore the unique 2-resident-block K pipeline and is retained. +// +// * Focused change (ONE mechanism -- the B operand's global-memory layout -- +// entirely inside csrc/w8a8_gemm_hip.hip). The pipeline round's own PMC +// identifies the next exposed cost: the accepted source reads +// 170,221,184 B/replay at only an 86.938% L2 hit rate (3,250,848 misses) +// and 166.147 GB/s counter-derived HBM, while the unique operands are only +// 25.2 MB (A) + 16.1 MB (B). The B side causes it: the n-major pack +// packed[n*K + kk] lays a [64, 64] B tile out as 64 separate 64-byte +// column chunks at a 6,144-byte stride, so every one of the 8 warp-level +// 16-byte staging loads touches 32 half-used 128-byte lines. +// launch_pack_w8a8_weight now packs the exact (k, n) == (6144, 2624) +// weight once, outside the timed region and out of Graph capture, into 41 +// B 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] +// (byte-exact bijection 41 x 384 x 64 x 16 = k*n = 16,121,856; same byte +// count, same buffer and same captured addresses; every other (K, N) +// keeps the identity copy). Each [64, 64] B tile of a 64-K stage is then +// FOUR contiguous 1-KiB slabs (one per 16-k-run) and every staging vector +// is a 16-byte-aligned contiguous k-run of one column, so a warp's 32 +// x 16-byte B staging loads touch 4 FULLY-USED 128-byte lines instead of +// 32 half-used ones (8x fewer line requests per warp-load; the tile's line +// footprint is 32 fully-used lines instead of 64 half-used ones). +// The ONLY kernel change is the two B global-load sites (prologue + +// steady-state prefetch) of the dispatched +// w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64>: byte +// (n0>>6)*(k*64) + ((k0>>4)+k16)*1024 + n_row*16 replaces +// (n0+n_row)*k + k0 + k16*16. The LDS staging layout b_tile[2][64][72] +// (two-half ds_write2_b64), the col_major load_frag8 consumes, the +// k0-outer/kk-inner ascending int32 accumulation order, the coalesced +// epilogue, the grid (41 x 32 = 1312 blocks), the block size (256 = 4 +// wavefronts), the retained double-buffered 64-K pipeline (one +// __syncthreads per stage), LDS/block 27,648 B (2 blocks/CU = 8 waves/CU) +// and the VGPR budget are ALL unchanged, so this is a pure global-layout +// A/B. +// * Correctness: the same weight elements reach the same fragment slots in +// the same ascending k-chunk sequence, so every int32 accumulator -- and +// every bf16 output byte -- is bit-identical to the accepted kernel and to +// the scalar reference. The scalar fallback for (k, n) == (6144, 2624) +// decodes the panel layout in ascending 16-byte k-runs (the paired M=2 / +// M=16 API shapes with the same (N, K) stay byte-exact); the generic +// identity arm, the round-2 n-major pack kernel (kept compiled as the +// layout control) and every shape guard are untouched. Packing stays +// outside the timed region and out of Graph capture. +// * Predicted deltas: vmem_read_instructions unchanged (1,574,400), LDS +// counters unchanged (lds_instructions 4,901,632, lds_wait 3,590,796, +// conflicts at the merged-read alias floor 12,091,392), L2 misses strictly +// below 3,250,848 with the hit rate up, operator read bytes strictly below +// 170,221,184 (B line fetches halve), median < 1153.970 us. +// * Superseded by Round 5: the B-panel pack and the b_tile[2][64][72] LDS +// form below are replaced by the fragment-interleaved swizzle. +// +// Round 5 (mandated packing round: one weight packing/swizzle that makes each +// DUMMA B tile vector-loadable and LDS-bank-safe): +// +// * Baseline: the accepted double-buffered <128, 64, 64> kernel (official +// median 937.431 us / p90 939.452 us; exact-source PMC +// profiles/.../iteration4/pmc.json: profiled 889.453 us, grid 41x32 = +// 1312 blocks, 2 blocks/CU = 8 waves/CU, 27,648 B LDS, arch_vgpr 96 with +// 0 spills, lds_instructions 4,901,632 = 9.7 per wave-stage (6 merged +// ds_read2_b64 + 3 ds_write2_b64), lds_wait 3,828,676, +// lds_bank_conflicts 12,091,392 = 6 x 4 x 4 x 96 x 1312, vmem_read +// 1,574,400, L2 hit 90.39%). The surviving exposed LDS cost is the +// b_tile row-major layout: each col_major B fragment is one ds_read_b64 +// at a 72-byte row stride, and the compiler merges the fragment pairs +// into ds_read2_b64 whose +16-row second offset (16 x 72 = 1152 B = +// 9 x 128 B) aliases the same bank pairs -- one conflict per 16-lane +// phase, exactly the measured 12,091,392 conflicts. +// * Mechanism (ONE packing/swizzle, entirely in this file): replace the +// Round-3 B-panel pack with a fragment-interleaved swizzle. 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 as [k8][n][8] (512 B per fragment block), +// 32-k major inside the 64-K stage: +// packed[((kk>>5)*n16g + (col>>4))*512 + +// (((kk>>3)&3)*16 + (col&15))*8 + (kk&7)] = raw[kk*n + col], +// with n16g = n>>4 (164 for the exact shape). Byte-exact bijection +// 192 x 164 x 512 = k*n = 16,121,856: same byte count, same buffer and +// same captured addresses (graph-stable), and packing still runs once, +// outside the timed region and out of Graph capture. +// * Kernel-side effect: every 16-byte B staging vector is now 2 n-rows x +// 8 B of one k8 of one fragment block, so the staging store is ONE +// 16-byte-aligned int4 (ds_write_b128) flat copy per thread (256 chunks +// = one per thread per 64-K stage; a warp reads 8 fully-used 128-byte +// lines) and every col_major B fragment is load_frag8_swz: lane +// (n = lane&15, k8 = lane>>4) reads p + k8*128 + n*8, i.e. one +// ds_read_b64 whose 16-lane phase spans exactly one 128-byte k8 block -- +// all 32 banks exactly once per phase, conflict-free by construction. +// The A tile keeps its bank-safe 72-byte stride, and the double-buffered +// 64-K pipeline, the k0-outer/kk-inner ascending int32 order, the +// coalesced epilogue, the grid (41x32 = 1312 blocks), the block +// (256 threads = 4 wavefronts) and the occupancy (2 blocks/CU = 8 +// waves/CU) are all unchanged, so this is a pure packing/swizzle A/B. +// * LDS: b_tile shrinks from 2 x [64, 72] = 9,216 B to 2 x 4,096 B; the +// block total is 18,432 (A) + 8,192 (B) = 26,624 B, so 2 x 26,624 = +// 53,248 <= 65,536 B keeps 2 resident blocks/CU = 8 waves/CU. +// * Correctness: the same weight elements reach the same fragment slots in +// the same k0-outer/kk-inner ascending k-chunk sequence, so every int32 +// accumulator -- and every bf16 output byte -- is bit-identical to the +// accepted kernel and to the scalar reference. The scalar fallback for +// (k, n) == (6144, 2624) decodes the swizzle element-wise in ascending +// kk (the paired M=2 / M=16 API shapes stay byte-exact); the generic +// identity arm, the Round-2 n-major pack control and every shape guard +// are untouched. +// * Predicted deltas: vmem_read_instructions unchanged (~1,574,400), +// LDS/block 27,648 -> 26,624 B, lds_bank_conflicts strictly below +// 12,091,392, lds_wait below 3,828,676, median < 937.431 us. +// Falsifiable: retained only if median_us < 937.4308013916016, +// 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). +// LDS == 26,624 B / VGPR <= 128 (2 blocks/CU), and the PMC shows +// lds_bank_conflicts and lds_wait strictly below the baseline with +// vmem_read not growing; any regression restores the iteration-3 source. +// +// Round 6 (mandated epilogue round: fuse per-row/per-column scales, bf16 +// conversion and the final coalesced store into the compute kernel; remove any +// unnecessary workspace/combine pass): +// +// * Fusion audit (nothing to remove, so the fused epilogue is made cheaper). +// The accepted source ALREADY fuses everything: store_prefill_fragment_ +// coalesced applies x_scale[row](per-row) * weight_scale[col](per-column), +// converts to bf16 and writes one coalesced 8-byte store per lane, and the +// exact-shape operator replay is a SINGLE dispatch -- the trusted PMC +// (profiles/glm_tp8_fused_qkv_a_proj_m4096/iteration6/pmc.json, digest +// 7e5731dc..., compile key 120c807e...) reports operator_aggregate +// kernel_count == 1, scratch_bytes == 0, vmem_write_instructions 41,984 = +// 8 x 4 x 1312 and write_bytes_per_operator_replay 21,500,384 = m*n*2, i.e. +// there is no workspace, no split-K combine pass and no second Graph node +// to remove; workspace_bytes is unused and workspace_split_k_capacity is +// zero for this shape. The remaining epilogue cost is therefore the +// conversion itself. +// * Measured cost of the fused epilogue (exact gfx928 code object of the +// current best, iteration6/current-best-isa/isa.txt). The dispatched +// kernel is ~1,920 instructions; the steady-state K-stage loop body is 107 +// instructions (one `s_branch 65428` back-edge at 0x6BB8 -> 0x6A0C) and +// the prologue is ~259, so the fused epilogue is ~1,490 static +// instructions -- roughly 13% of the dynamic instruction stream over +// 96 stages. The dominant term is that this DTK lowers EACH of the 32 +// per-lane __float2bfloat16 calls to a ~13-instruction EXEC-masked +// Inf/NaN sequence: +// v_and_b32 (b & 0x7f800000) / v_cmp_ne_u32 / s_and_saveexec_b64 / +// s_xor_b64 / v_bfe_u32 (bit 16) / v_add3_u32 (b + lsb + 0x7fff) / +// s_andn2_saveexec_b64 / v_mov_b32 0 / v_or_b32 (b | 0x10000) / +// v_cmp_eq_u32_sdwa / v_cndmask_b64 / s_or_b64 (restore EXEC) +// so ~420 instructions per lane plus 4 x (3 shifts + 3 ors) of packing, +// with 32 dependent SALU EXEC save/restore pairs that cannot be software- +// pipelined and serialise the wavefront (the epilogue is the only place +// VALU issues drop while the exact-shape kernel is 100% GRBM GUI-active). +// * Focused change (ONE mechanism -- the bf16 conversion of the fused +// epilogue -- entirely inside csrc/w8a8_gemm_hip.hip). Replace the library +// conversion + shift/or pack in store_prefill_fragment_coalesced with the +// branch-free round-to-nearest-even integer identity +// bf16_rne_high16(v) = (b + 0x7fff + ((b >> 16) & 1)) >> 16, +// b = __float_as_uint(v), +// which the disassembly proves is exactly what the normal (finite) path of +// the library conversion computes (v_bfe bit 16, v_add3 b + lsb + 0x7fff). +// Adding 0x7fff + the retained bit-16 LSB cannot carry out of the sign, the +// +((b>>16)&1) term rounds exact ties to an even LSB (RNE), max-float still +// rounds to +/-Inf and 0.0/-0.0 and subnormal encodings are preserved, so +// every 16-bit payload is identical. Every value converted here is finite +// by construction (|int32 dot| < 2^31 times finite dequant scales), so the +// removed Inf/NaN special case is unreachable. Four values now pack into +// one 8-byte word with three VALU ops (v_bfe/v_and + v_add3 + v_lshrrev) +// instead of ~13 EXEC-bracketed instructions each, and the 32 EXEC +// save/restore chains disappear. The epilogue structure, the per-row +// x_scale gather, the per-column float4 weight_scale gather, the register +// 4x4 shuffle transpose, the multiply order float(dot) * x_scale * +// weight_scale, the 8-byte-aligned coalesced store address, the grid +// (41 x 32 = 1312), the block (256 = 4 wavefronts), the double-buffered +// 64-K pipeline, the swizzled packed B, LDS/block 26,624 B (2 blocks/CU = +// 8 waves/CU) and the 65,536 B workspace budget are ALL unchanged, so this +// is a pure conversion A/B. The scalar fallback keeps __float2bfloat16 (it +// is not on the timed path); the generic identity arm and every shape +// guard are untouched. +// * Predicted deltas: grid_blocks 1312, LDS/block 26,624 B, 2 blocks/CU = +// 8 waves/CU, vmem_read 1,574,400 and vmem_write 41,984 unchanged; +// valu_instructions strictly below 23,920,384 and the epilogue's static +// EXEC/SALU count down by >1,000; median < 918.191 us with no p90 +// regression. Falsifiable: retained only if median_us <= 909.0 us +// (>= 1% round acceptance) and p90_us <= 919.530 us, mismatch_count == 0, +// max_abs_error == 0.0, graph_capture_passed == true, the exact code +// object shows 0 spills / LDS == 26,624 B / VGPR <= 128 (2 blocks/CU) and +// no __float2bfloat16 EXEC sequence remains in the epilogue; otherwise the +// iteration-5 source is restored. +// +// Round 9 (late ISA-diagnosis round; the control-plane plateau gate is CLOSED +// for this round -- phase hip_only, plateau false, raw inline asm and ISA +// Skills forbidden -- so this stays a HIP-only round): // -// 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, +// * Bottleneck evidence (exact trusted gfx928 ISA of the current best, +// profiles/glm_tp8_fused_qkv_a_proj_m4096/iteration6/current-best-isa/ +// isa.txt, digest 7e5731dc...; the round-6 source differs from that +// profiled official best only in the bf16 conversion, so the steady-state +// loop analysed here is the dispatched loop). The steady-state stage body +// is 107 instructions (0x6A0C..0x6BB8, one back-edge `s_branch 65428` at +// 0x6BB8 -> 0x6A0C) and roughly 14 of them are pure A-row tail +// predication: +// 0x6A0C..0x6A20 six unconditional `v_mov_b32` zero-fills of the two +// int4 staging vectors (v33/v34 and v37..v40), +// 0x6A24..0x6A34 `s_and_saveexec_b64` / `s_cbranch_execz` / +// `global_load_dwordx4` / `s_or_b64` for A vector 0, +// 0x6A38..0x6A50 two more `v_mov_b32` zero-fills (v35/v36) plus the +// same EXEC save/restore + uniform branch for A +// vector 1. +// The guard source `a_pre[v] = (g_row < m) ? load : int4{0,0,0,0}` is +// never false for the dispatched M=4096 (M / BM = 4096 / 128 = 32 exact), +// yet its zero-fills execute unconditionally in every one of the 96 +// stages and its two EXEC read-modify-writes serialize the wavefront in +// front of the very global loads that prefetch the next tile. That is +// ~14/107 = 13% of the hot loop's instruction stream, on top of the +// epilogue's per-fragment `row >= m` EXEC bracket (s_and_saveexec / +// s_xor_b64 / s_andn2_saveexec_b64 / s_cbranch_execz / s_or_b64 around +// each of the eight unrolled fragment stores), for a kernel the PMC shows +// 100% GRBM GUI-active yet issuing only ~0.42 instructions per cycle per +// CU (lds_wait_instructions 5,302,705 against lds_instructions +// 4,901,632). +// * Focused change (ONE mechanism -- exact-shape full-tile specialization, +// entirely inside csrc/w8a8_gemm_hip.hip): the dispatched kernel gains a +// compile-time `bool kFullTile` template parameter. The exact-shape guard +// keeps `(m >= 128, n == 2624, k == 6144)` and the same grid (N/64 x +// ceil(M/128)) and block (256 = 4 wavefronts), and selects +// `kFullTile == true` exactly when the caller's m is a multiple of +// BM = 128 -- i.e. when every block tile lies fully inside the output (the +// assigned M=4096 case). In that instantiation the A staging loads +// (prologue AND steady state) drop the `(g_row < m)` select / zero-fill +// and the epilogue drops the per-fragment `row >= m` early return, so the +// loop body keeps only the address math and the three coalesced +// `global_load_dwordx4`. Every other m keeps the guarded instantiation +// byte-for-byte (zero-fill, epilogue row mask, tail-block coverage), and +// the paired M=2/M=16 API shapes with the same (N, K) still take the +// scalar fallback because the `m >= 128` guard is untouched. +// * Why this is the remaining lever and not a new pipeline: rounds 7 +// (deeper register prefetch / stage pairs) and 8 (smaller macro-tile, more +// resident waves) both regressed, so the double-buffered 64-K stage (ONE +// __syncthreads per stage, LDS 26,624 B/block, 2 blocks/CU = 8 waves/CU), +// the 128 x 64 macro-tile, the fragment-interleaved packed B swizzle, the +// load_frag8 / load_frag8_swz consumes, the k0-outer / kk-inner ascending +// m16n16k32 int32 accumulation order and the fused branch-free bf16 +// epilogue are ALL retained unchanged. The round removes instructions that +// are provably dead for this shape without touching any resource, any +// memory access order or any MMA input. +// * Correctness: for m % 128 == 0 no block tile crosses m, so the removed +// guard predicate is never false and the removed zero-fill never reaches an +// accumulator; the surviving assignments, the ascending k-chunk sequence +// and the exact `float(dot) * x_scale * weight_scale` multiply order are +// bit-identical to the guarded path (mismatch 0 / max_abs_error 0.0 +// expected). The guarded instantiation preserves the exact current +// semantics for every m % 128 != 0. +// * Predicted deltas: steady-state loop 107 -> ~93 instructions (the eight +// zero-fills and the two EXEC save/restore + uniform-branch pairs removed), +// valu_instructions strictly below 23,920,384, grid_blocks 1312, +// block_threads 256, LDS/block 26,624 B, arch_vgpr <= 128 with 0 spills, +// vmem_read_instructions 1,574,400, vmem_write_instructions 41,984 and the +// epilogue store count all unchanged. Falsifiable: retained only if +// median_us < 912.7147674560547 (the current shadow base) with no p90 +// regression, 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 +// LDS == 26,624 B / VGPR <= 128 (2 blocks/CU) and the kFullTile symbol's +// loop body no longer contains the A-guard +// `s_and_saveexec`/`s_cbranch_execz`/zero-fill sequence; any median/p90 +// regression, mismatch, spill or occupancy loss restores the round-6 // source. // // Include order is fixed by the control plane: hip_runtime, hip_bfloat16, @@ -273,7 +464,7 @@ namespace { -constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kWaveSize = 64; // gfx928 native wavefront size constexpr int kTileM = 16; constexpr int kTileN = 16; constexpr int kTileK = 32; // INT8 DUMMA unit: m16n16k32, int32 accumulation @@ -286,8 +477,7 @@ 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 +// Coalesced accumulator epilogue for gfx928 INT8 m16n16k32: 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 @@ -298,9 +488,43 @@ using namespace du::dumma; // 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 +// (col0 % 4 == 0) and the 8-byte store (n even -> offset a multiple of 8) are +// aligned. +// +// Round 6: the fused bf16 conversion itself is made branch-free. The exact +// gfx928 code object of the previous best lowers each of the 32 per-lane +// __float2bfloat16 calls to a ~13-instruction sequence that brackets an +// Inf/NaN special case with EXEC save/restore: +// v_and_b32 (b & 0x7f800000) / v_cmp_ne_u32 / s_and_saveexec_b64 / s_xor_b64 +// / v_bfe_u32 (bit 16) / v_add3_u32 (b + lsb + 0x7fff) / s_andn2_saveexec_b64 +// / v_mov_b32 0 / v_or_b32 (b | 0x10000) / v_cmp_eq_u32_sdwa / v_cndmask_b64 +// / s_or_b64 (restore EXEC) +// The EXEC manipulation cannot be software-pipelined and serialises the +// wavefront. For every FINITE fp32 the library conversion is exactly +// round-to-nearest-even at bit 16, which is the branch-free integer identity +// bf16 = (b + 0x7fff + ((b >> 16) & 1)) >> 16, b = __float_as_uint(v) +// (adding 0x7fff + the retained bit-16 LSB never carries out of the sign, and +// the +((b>>16)&1) term breaks exact ties toward an even LSB, so max-float +// still rounds to +/-Inf and -0.0/0.0 are preserved). Every value converted +// here is finite by construction: |int32 dot| < 2^31 and the dequant scales +// are finite, so no Inf/NaN can reach the conversion. This yields the exact +// same 16-bit payloads as __float2bfloat16 while dropping the EXEC sequence +// and the separate shift/or pack (four bf16 pairs become three VALU ops per +// 8-byte word instead of 4 x 13 + 3). +__device__ __forceinline__ unsigned int bf16_rne_high16(float v) { + const unsigned int b = __float_as_uint(v); + return (b + 0x7fffu + ((b >> 16) & 1u)) >> 16; +} + +// Round 9 adds the compile-time `kGuarded` flag: kGuarded == false is the +// exact-shape full-tile specialization (m % BM == 0), for which every row of +// every block tile is strictly inside the output, so the `row >= m` early +// return -- which the exact code object lowers to an EXEC save/restore pair +// plus a uniform branch around each of the eight unrolled fragment stores -- +// is removed at compile time. The default kGuarded == true keeps the exact +// current semantics for every other m (including the generic arm and the +// single-buffered siblings, which never pass the flag). +template __device__ __forceinline__ void store_prefill_fragment_coalesced( const AccFragment& frag, const float* __restrict__ x_scale, @@ -312,8 +536,10 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced( int base_col, int lane) { const int row = base_row + (lane & 15); - if (row >= m) { - return; + if (kGuarded) { + if (row >= m) { + return; + } } const int c4 = lane >> 4; // 0..3 const int x0 = frag.x[0]; @@ -352,53 +578,105 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced( 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) | + bf16_rne_high16(v0) | + (bf16_rne_high16(v1) << 16)) | (static_cast( - static_cast(__float2bfloat16(v3))) - << 48); + bf16_rne_high16(v2) | + (bf16_rne_high16(v3) << 16)) + << 32); *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 +// 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 +// -> the 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. `lane` is +// threadIdx.x % kWaveSize, so no wave-id builtin is required. +template +__device__ __forceinline__ void load_frag8( + Frag& f, + const int8_t* __restrict__ p, + unsigned ldm, + int lane) { + const unsigned row = static_cast(lane) & 0xfu; + const unsigned kq = static_cast(lane) >> 4; + const int64_t v = *reinterpret_cast(p + row * ldm + (kq << 3)); + reinterpret_cast(f.x)[0] = v; +} + +// Round-5 fragment-interleaved B swizzle helpers. +// +// Each 32-k x 16-n DUMMA B fragment block is 512 B = [k8 0..3][n 0..15][8 B] +// (n-interleaved k8 groups). The 16-byte chunk (k32, n16, k8, n_pair) holds +// the two 8-byte k-runs of n-rows 2*n_pair and 2*n_pair+1 of one k8 group, +// and sits at byte (k32 * n16g + n16) * 512 + k8 * 128 + n_pair * 16 of the +// packed [n>>4][k>>5][512 B] array (n16g = n>>4). Global elements satisfy +// packed element (col, kk) at ((kk>>5)*n16g + (col>>4))*512 + +// (((kk>>3)&3)*16 + (col&15))*8 + (kk&7) -- the exact inverse of this chunk +// map, so staging is a flat 16-byte copy and the fragment load reads the +// same bytes the Round-3 layout delivered. +__device__ __forceinline__ int64_t packed_b_swz_offset( + int k32, + int n16, + int n16g, + int k8, + int n_pair) { + return (static_cast(k32) * n16g + n16) * 512 + + static_cast(k8) * 128 + (n_pair << 4); +} + +// Round-5 col_major B fragment loader on the swizzled tile: lane (n = lane&15, +// k8 = lane>>4) reads p + k8*128 + n*8. One ds_read_b64 per fragment and each +// 16-lane phase spans exactly one 128-byte k8 block (each bank exactly once), +// so the read is conflict-free by construction. Byte placement is identical to +// load_frag8, so the v_mmac operands and the int32 accumulation stay +// bit-exact. +template +__device__ __forceinline__ void load_frag8_swz( + Frag& f, + const int8_t* __restrict__ p, + int lane) { + const unsigned n = static_cast(lane) & 0xfu; + const unsigned k8 = static_cast(lane) >> 4; + const int64_t v = + *reinterpret_cast(p + (k8 << 7) + (n << 3)); + reinterpret_cast(f.x)[0] = v; +} + +// Generic native INT8 DUMMA m16n16k32 tiled prefill kernel with the identity +// [K, N] weight layout (packed_weight == raw logical weight for every (K, N) +// other than the exact assigned pair). 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)). +// 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 +// 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 +// flat 16-byte 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. +// 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 above. 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 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] + hip_bfloat16* __restrict__ out, // [M, N] int m, int n, int k) { @@ -509,56 +787,42 @@ w8a8_dumma_prefill_tiled_kernel( } } -// 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). +// Round-1 packed-weight architecture kernel, kept as the single-buffered +// 128-K sibling of the dispatched <128, 64, 64> double-buffered kernel below +// (Round 2 moved the exact-shape dispatch to that sibling; Round 3 keeps this +// kernel compiled but undispatched as the n-major layout / single-buffer +// control, and only the double-buffered kernel reads the packed B; Round 5 +// replaces that packed B layout with the fragment-interleaved swizzle, so this +// n-major control no longer matches the pack either). It is defined for +// the n-major packed_weight[n * K + kk] = raw[kk * N + n] layout produced by +// launch_pack_w8a8_weight. SAME geometry as the bootstrap +// w8a8_dumma_prefill_tiled_kernel<128, 64, 128>: +// grid (N/64) x (M/128) = 41 x 32 = 1312 blocks, 256 threads = 4 wavefronts, +// one 64 x 32 quadrant of eight m16n16k32 int32 accumulators per wave, +// single-buffered 128-K LDS stage with TWO __syncthreads per stage, same +// k0-outer/kk-inner ascending int32 accumulation order, same coalesced +// fragment epilogue. The two differences, both inside the operand path: +// * B is staged n-major (b_tile[n_row][kk] from the packed [N, K] weight) +// so col_major B fragments are 8 contiguous K bytes per lane and staging +// is one coalesced 16-byte vector per (n_row, k16) pair; +// * A and B LDS row strides are 136 B (StageK + 8): 8-byte aligned with an +// odd 8-byte count (136/8 = 17) so the 16 rows of each 16-lane ds_read_b64 +// phase land on 16 distinct bank pairs (the conflict-free 4-phase floor), +// and 136 % 16 == 8 keeps the two-half int64 staging form. // 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. +// exact int32 accumulation (k0-outer over 48 stages, kk-inner ascending over +// four m16n16k32 steps) are bit-identical to the bootstrap 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 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] + hip_bfloat16* __restrict__ out, // [M, N] int m, int n, int k) { @@ -568,9 +832,10 @@ w8a8_dumma_prefill_packedb_kernel( "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). + // Bank-safe LDS row stride: StageK + 8 bytes = 34 dwords for StageK = 128. + // 8-aligned with an odd 8-byte count -> every lane's 8-byte fragment is a + // ds_read_b64 at the 4-phase conflict-free floor (see the kernel comment + // above); 8 mod 16 keeps the two-half staging store form. 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"); @@ -602,14 +867,14 @@ w8a8_dumma_prefill_packedb_kernel( } } - // Flat 16-byte staging vectors per stage (dispatch guarantees - // k % StageK == 0 and n % BN == 0). 8 vectors per 128-K row. + // Flat 16-byte staging vectors per stage (the dispatch guarantees + // k % StageK == 0 and n % BN == 0). 8 vectors per StageK row for 128. + 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 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 + // Stage A[BM, StageK] row-major with the bank-safe row stride: 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) { @@ -621,18 +886,17 @@ w8a8_dumma_prefill_packedb_kernel( 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). + // Two 8-byte halves: the row stride is 8 mod 16, so odd rows are not + // 16-byte aligned and a single int4 LDS 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. + // 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; @@ -658,14 +922,16 @@ w8a8_dumma_prefill_packedb_kernel( load_frag8( a_frag[j], a_tile + (local_row + j * kTileM) * kLdsStride + kk, - kLdsStride); + static_cast(kLdsStride), + lane); } #pragma unroll for (int i = 0; i < kWaveN16; ++i) { load_frag8( b_frag[i], b_tile + (local_col + i * kTileN) * kLdsStride + kk, - kLdsStride); + static_cast(kLdsStride), + lane); } #pragma unroll for (int j = 0; j < kWaveM16; ++j) { @@ -675,8 +941,8 @@ w8a8_dumma_prefill_packedb_kernel( } } } - // Single buffer: every wavefront must finish reading LDS before the - // next stage's stores overwrite it. + // Single buffer: every wavefront must finish reading LDS before the next + // stage's stores overwrite it. __syncthreads(); } @@ -693,51 +959,74 @@ w8a8_dumma_prefill_packedb_kernel( } } -// 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 +// Round 9 exact-shape A-staging load. For the dispatched M=4096 every 128-row +// block tile lies fully inside the output, so the `(g_row < m)` row guard is +// dead: the exact code object lowers it to eight unconditional zero-fill +// `v_mov_b32` plus two EXEC save/restore (`s_and_saveexec_b64` / +// `s_cbranch_execz` / `s_or_b64`) uniform-branch pairs per K stage, i.e. ~14 +// of the 107 steady-state loop instructions and an EXEC serialization in +// front of the very `global_load_dwordx4`s that prefetch the next tile. The +// kFullTile instantiation removes the select (one coalesced load per staging +// vector, no zero-fill), while kFullTile == false keeps the zero-fill + row +// mask semantics identical to the current kernel for every m % 128 != 0. +template +__device__ __forceinline__ int4 load_a_stage_vec( + const int8_t* __restrict__ x_q, + int64_t offset, + int m, + int g_row) { + if (kFullTile) { + return *reinterpret_cast(x_q + offset); + } + return (g_row < m) ? *reinterpret_cast(x_q + offset) + : int4{0, 0, 0, 0}; +} + +// Round-2 double-buffered packed-B prefill kernel, dispatched ONLY for +// (m >= 128, n == 2624, k == 6144). SAME macro-tile and SAME operand path as +// the accepted w8a8_dumma_prefill_packedb_kernel<128, 64, 128>: grid +// (N/64) x (M/128) = 41 x 32 = 1312 blocks, 256 threads = 4 wavefronts, one +// 64 x 32 quadrant of eight m16n16k32 int32 accumulators per wave, packed +// n-major B, explicit load_frag8 reads, coalesced epilogue. The ONLY change is +// the K pipeline (see the Round-2 section of the file header): +// * 64-K DOUBLE-buffered stages (2 x 13,824 = 27,648 B LDS/block, still +// 2 blocks/CU = 8 waves/CU): the stage-(s+1) global loads are issued into +// 3 int4 staging registers per thread BEFORE the stage-s consume and +// flushed into the idle LDS buffer after it, so the VMEM latency of stage +// s+1 overlaps the LDS/v_mmac work of stage s (one __syncthreads per +// stage instead of two); +// * k0-outer over 96 ASCENDING 64-K stages with kk-inner ascending over two +// m16n16k32 steps, i.e. the same (k0 + kk) k-chunk sequence +// 0, 32, 64, 96, ... as the accepted 48 x 128-K stages, so every int32 +// accumulator is bit-identical. +// kLdsStride = StageK + 8 = 72 B: 8-aligned with an odd 8-byte count (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 conflict-free 4-phase floor -- and +// 72 % 16 == 8 keeps the two-half ds_write2_b64 staging form. +// +// Round 5 additionally reads the packed B operand from the fragment- +// interleaved swizzle (see w8a8_pack_bswz_i8_kernel and the Round-5 header +// section): the LDS b_tile is now the SAME [k32][n16][k8][n][8] layout as the +// pack (2 x 256 16-byte chunks per 64-K stage), so staging is one flat +// 16-byte copy per thread and each col_major B fragment is one conflict-free +// ds_read_b64 (load_frag8_swz). The int32 accumulation order is untouched. +// +// Round 9 adds the compile-time `kFullTile` flag: it is true exactly when the +// caller's m is a multiple of BM, so every block tile is fully inside the +// output. In that instantiation the A staging loads drop the row guard (see +// load_a_stage_vec) and the epilogue drops the per-fragment row mask (see +// store_prefill_fragment_coalesced); false keeps the guarded form +// byte-for-byte. Geometry, resources, operand path and accumulation order are +// identical in both instantiations. +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 int8_t* __restrict__ x_q, // [M, K] row-major + const int8_t* __restrict__ packed_w, // [n>>4][k>>5][512 B] swizzled B + // fragment blocks (round-5 layout) + const float* __restrict__ x_scale, // [M] const float* __restrict__ weight_scale, // [N] - hip_bfloat16* __restrict__ out, // [M, N] + hip_bfloat16* __restrict__ out, // [M, N] int m, int n, int k) { @@ -747,19 +1036,19 @@ w8a8_dumma_prefill_packedb_db_kernel( "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"); + // Round 5: the packed B weight is the fragment-interleaved swizzle (see + // w8a8_pack_bswz_i8_kernel), defined for BN == 64 and StageK == 64 (one + // 32-k x 16-n fragment block per 512 B; 2 k32 blocks x 4 n16 groups x 4 k8 + // groups x 8 n-pairs = 256 16-byte chunks per stage, one per thread). Only + // <128, 64, 64> is ever dispatched by the exact-shape guard, so the chunk + // index math below is exact. + static_assert(BN == 64 && StageK == 64, + "swizzled B pack requires BN == 64 and StageK == 64"); const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; @@ -769,9 +1058,12 @@ w8a8_dumma_prefill_packedb_db_kernel( 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 + const int n16g = n >> 4; // 16-n row groups per 32-k block of the pack __shared__ __align__(16) int8_t a_tile[2][BM * kLdsStride]; - __shared__ __align__(16) int8_t b_tile[2][BN * kLdsStride]; + // Round-5 swizzled B tile: flat [k32][n16][k8][n_pair][16 B] = 4,096 B per + // 64-K stage, byte-identical to the packed global layout of that stage. + __shared__ __align__(16) int8_t b_tile[2][BN * StageK]; DUFragment a_frag[kWaveM16]; @@ -787,10 +1079,10 @@ w8a8_dumma_prefill_packedb_db_kernel( } } - // Flat 16-byte staging vectors per stage. For the dispatched <128,64,64> + // 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. + // 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)); @@ -799,12 +1091,16 @@ w8a8_dumma_prefill_packedb_db_kernel( static_assert(kAVectors % kThreadsPerBlock == 0 && kBVectors % kThreadsPerBlock == 0, "staging vectors must divide evenly across the block"); + // Round 5: the swizzled B stage is exactly one 16-byte chunk per thread, so + // thread `vec` = tid is also the flat [k32][n16][k8][n_pair] chunk index. + static_assert(kBVecPerThread == 1 && kBVectors == kThreadsPerBlock, + "swizzled B staging must be one 16-B chunk per thread"); 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. + // ---- 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 @@ -813,25 +1109,22 @@ w8a8_dumma_prefill_packedb_db_kernel( 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}; + a_pre[v] = load_a_stage_vec( + x_q, static_cast(g_row) * k + k0 + k16 * 16, m, g_row); } #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); + // Round 5: the 16-byte chunk (k32 = vec>>7, n16 = (vec>>5)&3, + // k8 = (vec>>3)&3, n_pair = vec&7) of the swizzled pack sits at byte + // ((k0>>5)+k32)*n16g*512 + ((n0>>4)+n16)*512 + k8*128 + n_pair*16; + // consecutive lanes read consecutive 16-byte chunks, so a warp's + // staging loads touch 8 fully-used 128-byte lines. + b_pre[v] = *reinterpret_cast( + packed_w + + packed_b_swz_offset((k0 >> 5) + (vec >> 7), + (n0 >> 4) + ((vec >> 5) & 3), n16g, + (vec >> 3) & 3, vec & 7)); } #pragma unroll for (int v = 0; v < kAVecPerThread; ++v) { @@ -847,26 +1140,22 @@ w8a8_dumma_prefill_packedb_db_kernel( #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]; + // Round 5: the swizzled tile is a flat copy of the packed stage, so the + // commit is ONE 16-byte-aligned int4 (ds_write_b128) per thread. + *reinterpret_cast(b_tile[0] + vec * 16) = b_pre[v]; } } __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. + // ---- Steady state: for stage s, 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 (so the vmcnt wait for stage s+1 lands after the MMAC + // burst of stage s, overlapping the VMEM latency of s+1 with the stage-s + // LDS/v_mmac work). 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 + // Prefetch stage s+1: issue the global loads (held in VGPRs while the // stage-s consume runs). const int k0 = (s + 1) * StageK; #pragma unroll @@ -875,29 +1164,28 @@ w8a8_dumma_prefill_packedb_db_kernel( 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}; + a_pre[v] = load_a_stage_vec( + x_q, static_cast(g_row) * k + k0 + k16 * 16, m, g_row); } #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); + // Round 5: the 16-byte chunk (k32 = vec>>7, n16 = (vec>>5)&3, + // k8 = (vec>>3)&3, n_pair = vec&7) of the swizzled pack sits at byte + // ((k0>>5)+k32)*n16g*512 + ((n0>>4)+n16)*512 + k8*128 + n_pair*16; + // consecutive lanes read consecutive 16-byte chunks, so a warp's + // staging loads touch 8 fully-used 128-byte lines. + b_pre[v] = *reinterpret_cast( + packed_w + + packed_b_swz_offset((k0 >> 5) + (vec >> 7), + (n0 >> 4) + ((vec >> 5) & 3), n16g, + (vec >> 3) & 3, vec & 7)); } // 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 @@ -905,14 +1193,19 @@ w8a8_dumma_prefill_packedb_db_kernel( load_frag8( a_frag[j], a_tile[s & 1] + (local_row + j * kTileM) * kLdsStride + kk, - kLdsStride); + static_cast(kLdsStride), + lane); } #pragma unroll for (int i = 0; i < kWaveN16; ++i) { - load_frag8( + // Round 5: fragment block (k32 = kk>>5, n16 = wave_col*kWaveN16 + i) + // of the swizzled tile -- 512 B per 16-n x 32-k DUMMA B fragment, one + // conflict-free ds_read_b64 per lane. + load_frag8_swz( b_frag[i], - b_tile[s & 1] + (local_col + i * kTileN) * kLdsStride + kk, - kLdsStride); + b_tile[s & 1] + + ((kk >> 5) * (BN / 16) + wave_col * kWaveN16 + i) * 512, + lane); } #pragma unroll for (int j = 0; j < kWaveM16; ++j) { @@ -940,13 +1233,8 @@ w8a8_dumma_prefill_packedb_db_kernel( #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]; + // Round 5: flat 16-byte swizzled-tile commit (ds_write_b128). + *reinterpret_cast(b_tile[buf] + vec * 16) = b_pre[v]; } } __syncthreads(); @@ -957,7 +1245,6 @@ w8a8_dumma_prefill_packedb_db_kernel( { 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 @@ -965,14 +1252,19 @@ w8a8_dumma_prefill_packedb_db_kernel( load_frag8( a_frag[j], a_tile[s & 1] + (local_row + j * kTileM) * kLdsStride + kk, - kLdsStride); + static_cast(kLdsStride), + lane); } #pragma unroll for (int i = 0; i < kWaveN16; ++i) { - load_frag8( + // Round 5: fragment block (k32 = kk>>5, n16 = wave_col*kWaveN16 + i) + // of the swizzled tile -- 512 B per 16-n x 32-k DUMMA B fragment, one + // conflict-free ds_read_b64 per lane. + load_frag8_swz( b_frag[i], - b_tile[s & 1] + (local_col + i * kTileN) * kLdsStride + kk, - kLdsStride); + b_tile[s & 1] + + ((kk >> 5) * (BN / 16) + wave_col * kWaveN16 + i) * 512, + lane); } #pragma unroll for (int j = 0; j < kWaveM16; ++j) { @@ -990,7 +1282,9 @@ w8a8_dumma_prefill_packedb_db_kernel( for (int j = 0; j < kWaveM16; ++j) { #pragma unroll for (int i = 0; i < kWaveN16; ++i) { - store_prefill_fragment_coalesced( + // Round 9: kGuarded == !kFullTile -- the full-tile instantiation skips + // the dead per-fragment `row >= m` EXEC bracket (see the helper above). + store_prefill_fragment_coalesced( acc[j][i], x_scale, weight_scale, out, m, n, base_row + j * kTileM, base_col + i * kTileN, lane); } @@ -999,23 +1293,37 @@ w8a8_dumma_prefill_packedb_db_kernel( // 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. +// they never reach the m >= 128 tiled/packed paths). 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 a manual uint16 +// bf16 store. +// +// Round 1: for the exact (k, n) == (6144, 2624) the packed_weight buffer +// holds a packed layout (produced once, outside the timed region), so the +// fallback decodes that layout there with ascending kk (every other (K, N) +// keeps the identity [K, N] row-major copy). The int32 accumulation stays +// kk-ascending 0..k-1 in both arms; integer addition is exact and commutative, +// so the int32 dot -- and the paired M=2 / M=16 API shapes -- remain +// byte-exact against the identity-layout reference. Round 5 replaces the +// exact-shape pack with the fragment-interleaved swizzle (see +// w8a8_pack_bswz_i8_kernel), so the exact-shape decode below walks +// packed[((kk>>5)*n16g + (col>>4))*512 + +// (((kk>>3)&3)*16 + (col&15))*8 + (kk&7)], n16g = n>>4, +// element-by-element in ascending kk -- identical elements, identical int32 +// accumulators. __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 int8_t* __restrict__ a, // [M, K] row-major + const int8_t* __restrict__ b, // packed weight (swizzled/identity) + const float* __restrict__ x_scale, // [M] const float* __restrict__ weight_scale, // [N] - hip_bfloat16* __restrict__ out, // [M, 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; + const bool packed_swz = (k == kExactK && n == kExactN); for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < total; @@ -1023,38 +1331,29 @@ __global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( 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]); - } + if (packed_swz) { + // Round-5 swizzle: one 8-byte k-run per (n, k8) inside a 512-byte + // 32-k x 16-n fragment block; walk it in ascending kk. + const int64_t n16g = static_cast(n) >> 4; + for (int kk = 0; kk < k; ++kk) { + const int64_t off = + ((static_cast(kk >> 5) * n16g + (col >> 4)) << 9) + + ((((kk >> 3) & 3) * 16 + (col & 15)) << 3) + (kk & 7); + acc += static_cast(a_row[kk]) * + static_cast(b[off]); } } 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]); + 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); + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + reinterpret_cast(out)[idx] = + static_cast(__float2bfloat16(scaled)); } } @@ -1073,14 +1372,13 @@ __global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( } } -// 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.) +// Round-1 exact-shape n-major pack: 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] so every col_major B fragment is 8 +// contiguous K bytes in the packed buffer and every staging vector is one +// coalesced 16-byte read. Same byte count and buffer as the identity pack, so +// the captured addresses are unchanged. One thread per output byte; runs once +// during weight preparation. __global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, @@ -1098,36 +1396,52 @@ __global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( } } -// 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( +// Round-5 exact-shape fragment-interleaved B swizzle: for (k, n) == +// (6144, 2624) the weight is packed once, outside the timed region and out of +// Graph capture, into 512-byte 32-k x 16-n DUMMA B fragment blocks laid out +// [k32][n16][k8][n][8]: +// packed[((kk >> 5) * n16g + (n >> 4)) * 512 + +// (((kk >> 3) & 3) * 16 + (n & 15)) * 8 + (kk & 7)] = raw[kk * n + n], +// where n16g = n >> 4 (164 for the exact shape). This is a byte-exact +// bijection 192 x 164 x 512 = k*n = 16,121,856: same byte count and buffer as +// the identity/n-major/panel packs, so captured addresses are unchanged. Each +// thread packs one 16-byte chunk = the 8-byte k-runs of two n-rows of one k8 +// group (a strided 2x8 gather on the read side; packing is outside timing, so +// the gather is off the critical path). Consecutive chunks are consecutive +// packed bytes, so the kernel's B staging is a flat 16-byte copy and every +// col_major B fragment read is one conflict-free ds_read_b64 (see +// load_frag8_swz). The round-2 n-major pack kernel above stays compiled as the +// layout control. +__global__ __launch_bounds__(256) void w8a8_pack_bswz_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 chunks = (static_cast(k) * n) >> 4; 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]; + 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(k32) * 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 = packed + (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] = raw[static_cast(k0 + i) * n + n0]; + dst16[8 + i] = raw[static_cast(k0 + i) * n + n0 + 1]; + } } } @@ -1149,7 +1463,9 @@ __global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( // 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. +// touches the default stream. The workspace is unused for this shape +// (single-buffered, no split-K): m*n*4 = 42,990,080 B exceeds the 16 MiB +// workspace budget, so workspace_split_k_capacity is zero. extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -1170,96 +1486,77 @@ extern "C" void launch_w8a8_gemm( 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. + // (M=4096, N=2624, K=6144). The m >= 128 guard keeps the paired M=2/M=16 + // API shapes with the same (N, K) on the scalar fallback below. Round 2 + // keeps the accepted packed-weight architecture (the exact-fit 2-D macro-tile + // <128, 64, ...> with the packed weight produced once, out of the + // timed region, by launch_pack_w8a8_weight, the bank-safe LDS row strides and + // the explicit 8-byte load_frag8 fragment loads) and changes ONLY the K-stage + // pipeline: the dispatch goes to the 64-K double-buffered prefetch kernel + // <128, 64, 64> instead of the single-buffered <128, 64, 128> one, so the + // stage-(s+1) VMEM latency overlaps the stage-s LDS/v_mmac work. Round 3 + // retains that double-buffered 64-K pipeline unchanged (it is the unique + // 2-resident-block K pipeline inside the 65,536 B LDS budget: 2 x 27,648 = + // 55,296 <= 65,536) and changes only the packed B layout consumed by the two + // B global-load sites -- the weight is now the B-panel layout, so a warp's B + // staging loads touch 4 fully-used 128-byte lines instead of 32 half-used + // ones. Round 5 keeps the kernel, geometry and occupancy identical and + // changes only the packed B layout again: the B-panel pack is replaced by the + // fragment-interleaved swizzle, so the B stage is a flat 16-byte copy and + // every col_major B fragment is one conflict-free ds_read_b64 + // (load_frag8_swz). M = 32*128 and N = 41*64 are both exact, so no tail + // predication is introduced. Round 9 keeps that kernel, geometry and + // occupancy and only selects a compile-time tail-predication specialization: + // when m is a multiple of BM = 128 every block tile is fully inside the + // output, so the instantiation with kFullTile == true drops the dead A-row + // zero-fill/select in both staging sites and the dead epilogue row mask + // (see the Round-9 header section); every other m keeps the guarded + // instantiation byte-for-byte. 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); + if ((m & 127) == 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64, true>), + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64, false>), + grid, + block, + 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. + // Generic native INT8 DUMMA prefill path for every other large-M shape with + // N % 64 == 0 and K % 128 == 0 (identity [K, N] layout). The A row guard + // zero-fills rows past M and the epilogue masks them, so only the exact + // shape above avoids predication. if (m >= 128 && (n % 64) == 0 && (k % 128) == 0) { const dim3 grid( static_cast(n / 64), @@ -1282,8 +1579,8 @@ extern "C" void launch_w8a8_gemm( return; } - // Scalar int8/int32 fallback for every unmatched (m, n, k) and all - // small-M API cases. + // 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; @@ -1311,15 +1608,15 @@ extern "C" void launch_w8a8_gemm( 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. +// Stable host launch symbol consumed by csrc/bindings.cpp. Round 5: the exact +// (k, n) == (6144, 2624) int8 weight is packed once (outside the timed region +// and out of Graph capture) into the fragment-interleaved swizzle +// packed[((kk>>5)*n16g + (n>>4))*512 + (((kk>>3)&3)*16 + (n&15))*8 + (kk&7)] +// = raw[kk*n + n] with n16g = n>>4, consumed by the packed-weight prefill +// kernel's flat B stage copy and by the scalar fallback's element decode; every +// other (K, N) keeps the identity copy of the logical [K, N] row-major weight. +// The [N] fp32 scales are always copied unchanged. Same byte count (k*n and n), +// same allocations, so the captured addresses stay graph-stable. extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -1330,14 +1627,10 @@ extern "C" void launch_pack_w8a8_weight( 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; + // Round 5: one thread packs one 16-byte chunk (16 bytes = 2 n-rows x 8 B). + const int64_t chunks = weight_count >> 4; + int64_t blocks = (chunks + kPackThreads - 1) / kPackThreads; if (blocks > 4096) { blocks = 4096; } @@ -1345,7 +1638,7 @@ extern "C" void launch_pack_w8a8_weight( blocks = 1; } hipLaunchKernelGGL( - w8a8_pack_panel_i8_kernel, + w8a8_pack_bswz_i8_kernel, dim3(static_cast(blocks)), dim3(kPackThreads), 0, @@ -1393,4 +1686,3 @@ extern "C" void launch_pack_w8a8_weight( } } // @@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 index 7de05137..af1b458c 100644 --- 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 @@ -1,242 +1,320 @@ -// @@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 +// @@variant shape=glm_tp8_kv_b_proj_m4096 commit=547eec13f53e3cee7da02bbdd17763f46e71761e added=2026-09-11 +// median_us=120.1 p90_us=120.9 speedup=57.78 baseline_us=6940 +// source=test-9-11-1-b01955b3/iteration_001/glm_tp8_kv_b_proj_m4096 // 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 +// Worker: worker_2 (physical GPU 2). Assigned shape: // 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 9 (consolidate round): the exact-shape arm keeps its winning +// 128x128 macro-tile, its 512-thread / eight-wavefront geometry, its +// 16,384-B single-buffered swizzled-A + plane-B staging, its direct 8-byte +// fragment loads, its k0-outer/kk-inner int32 order and its coalesced +// 8-byte epilogue. The one consolidated defect is the bf16 CONVERSION +// lowering: the exact gfx928 code object of the current best +// (compile_cache_key 0ece2e0cdbddfe2f512228e2bde5e23189f88fc22d01ae5ae49daf4a97a80f77, +// source_hip_digest 95dab478d6509cacbb6ca0888d90635197ed82d9202aead00ed49f50c4c7c9ac) +// shows hip's __float2bfloat16 emitting, for EVERY element each lane +// converts, an exec-masked inf/NaN fixup - v_and 0x7f800000, v_cmp_ne, +// s_and_saveexec, v_bfe, v_add3, v_or 0x10000, v_cndmask, s_andn2_saveexec +// (code-object addresses 0x5EEC..0x77E0 of +// w8a8_dumma_prefill_128x128_nmajor_kernel) - i.e. ~9 vector instructions +// plus two scalar EXEC manipulations per element against the 3-instruction +// rounding the value actually needs. With M*N = 14.68M converted elements +// and 6,931,456 measured valu_instructions, the block-tail epilogue is the +// single largest VALU consumer in the kernel, and the +// s_and_saveexec/s_andn2_saveexec pairs serialize the wavefront there where +// the only other work left is the second co-resident block. This round +// replaces ONLY that conversion, ONLY in the exact-shape epilogue, with a +// plain RNE form (bf16_rne_u16: u += 0x7fff + ((u >> 16) & 1); u >> 16) - +// the same lineage-validated fix the same-shape sibling reference +// references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip applies on this +// exact DTK 26.04 / gfx928. The formula is bit-identical to hip's own +// non-masked branch for every finite float, and int32 accumulators times +// finite scales can only produce finite floats, so mismatch_count 0 / +// max_abs_error 0 is expected on the unchanged pre-seeded reference data. +// Everything else - the dense swizzle, the tile-contiguous pack and its +// scalar decode, the fragment placement, the k0/kk accumulation order, the +// exact (k, n) == (512, 3584) guard, the identity-layout 128x64 arm, the +// pack kernels and the generic scalar fallback - is byte-identical. // -// 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 10 (consolidate round): final HIP-only polish on the same winning +// architecture. The one consolidated defect is the dead per-fragment tail-M +// masking in the exact-shape epilogue. The exact gfx928 code object of the +// current best (compile_cache_key 0f7ddb1a6c31e6f3cdd7772a8760e5b20dcea6da9ffca5d556e35790d7e25113, +// source_hip_digest 8c2468b4dfc603de4dd4a957ebf58d891830cbac940a07c58aed7c52030a8def, +// arch_vgpr 64, 0 scratch, 16,384 B LDS, 896 blocks) shows the block tail +// split into eight exec-masked basic blocks +// (s_and_saveexec / s_cbranch_execz / s_or_b64 around each of the eight +// fragment epilogues) even though the dispatch guard for this arm is +// m >= 128 with m % 128 == 0, so every row of every 128-row tile is in bounds +// and neither the (row >= m) early return in the store helper nor the two +// (rs < m) x_scale ternaries in the kernel can ever fire. Because each +// fragment epilogue sits in its own basic block, the compiler also repeats the +// fragment-independent shuffle scaffolding per fragment (v_mbcnt_lo / +// v_mbcnt_hi lane-id, the two __shfl_xor partner indices, the two comparison +// and cndmask pairs, ~13 VALU per fragment x 8). This round removes ONLY that +// dead masking from the exact-shape arm: the store helper drops its +// (row >= m) early return and its now-unused m argument, and the caller loads +// x_scale[rs0] / x_scale[rs1] directly. Nothing else changes: the 128x128 +// macro-tile, the 512-thread / eight-wavefront block, the 16,384-B +// single-buffered swizzled-A + plane-B staging, the direct 8-byte fragment +// loads, the k0-outer/kk-inner int32 order, the 4x4 shfl_xor transpose and +// the single 8-byte bf16 store per lane are all untouched, so the written +// bytes are bit-identical (removing a mask that is never asserted cannot +// change any output). The 128x64 identity-layout arm keeps its own +// store_fragment_scaled bounds check and the generic scalar fallback keeps its +// linear bounds check, both byte-identical. Falsifiable signature: the eight +// exec-masked guard blocks disappear from the 128x128 code object, the +// per-fragment duplicate lane-id / shuffle-index VALU is hoisted, and the +// dynamic valu_instructions fall from 5,512,192 without changing +// lds_instructions 1,089,536, lds_bank_conflicts 458,752, +// vmem_write_instructions 57,344 or the 16 v_mmac per stage. // -// 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 13 (consolidate round, 13/13): final HIP-only polish on the +// unchanged winning exact-shape architecture, one bounded mechanism. The +// exact current-best code object (source_hip_digest +// 8010314fdb6ba5ec0da98d662fc3911763b088245fa7c316bf0986572499207d, +// compile_cache_key 3fbcc466cfb5087e6f7057b3f26684e598046f28d0f71e0d794962688901c1c7, +// arch_vgpr 64, scratch 0, 16,384 B static LDS, 896 blocks x 512 threads) puts +// the whole 64-K stage wall on the EXPOSED cooperative-staging round trip: the +// two global_load_dwordx4 A/B chunks are issued only after the previous stage's +// trailing s_barrier and are then drained within ~8 instructions by +// s_waitcnt vmcnt(1) / vmcnt(0) in front of the three ds_write_b64/b128 +// commits (0x5D1C..0x5D5C), so no part of the A+B L2 latency overlaps any +// arithmetic or the barrier wait. Iteration 7 raised the staging depth to 2 to +// overlap exactly that round trip, but it doubled LDS to 32,768 B/block and +// gave up the 2-blocks/CU residency point, so the prefetch mechanism itself was +// never measured. This round re-introduces ONLY the single-buffered register +// prefetch: the A row base and the B per-stage stride are hoisted out of the K +// loop, the next stage's A and B int4 chunks are loaded BEFORE the trailing +// barrier and consumed by the next stage's ds_write commits, while LDS stays +// 16,384 B, __launch_bounds__(kNMBlockThreads, 2), the 64-VGPR residency point +// and the two barriers per stage are all untouched, and the last stage re-reads +// its own chunk (k1 == k0) instead of reading past the last K block. The +// swizzled A tile, the tile-contiguous plane-B layout, the direct 8-byte +// fragment loads, the k0-outer/kk-inner ascending int32 order, the coalesced +// RNE epilogue, the exact (k, n) == (512, 3584) dispatch guard, the 128x64 +// identity-layout arm, the pack kernels and the generic scalar fallback are +// byte-identical, so every stored bit is unchanged. Falsifiable signature: the +// global_load_dwordx4 pair for stage s+1 is scheduled before the trailing +// s_barrier and its vmcnt drain moves to the next stage's top, while arch_vgpr +// stays 64 with vgpr_spill_count 0 and scratch_bytes 0; if the machine +// scheduler keeps the loads where they already are the object is unchanged and +// the round is flat. // -// 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 5 (packing round): the exact (k, n) == (512, 3584) weight is +// re-packed once, outside the timed/Graph region, from the n-major [N, K] +// transpose to a TILE-CONTIGUOUS 64-K-stage 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, +// i.e. every (64-K stage, 128-N tile) DUMMA B tile is ONE contiguous 8192-B +// global region laid out as [kc = 8-byte k sub-chunk][128 N rows][8 bytes], +// which is exactly the plane order the B fragments are read from in LDS. The +// GEMM's cooperative B staging therefore becomes a plain 2 x 16-B contiguous +// vector copy per thread (global_load_dwordx4 -> ds_write_b128) with +// consecutive lanes on consecutive 16-B slots: the global reads are fully +// coalesced (the n-major [N, K] layout made a wave touch 16 rows 512 B apart) +// and the 16-B LDS stores have no bank conflicts. The B fragment load becomes +// lane-linear (kc*(128*8) + row*8, no swizzle arithmetic) and each 16-lane +// phase covers all 32 banks exactly once, so B is now both vector-loadable and +// LDS-bank-safe; the swizzled 8-byte-group map is retained for A only (x_q is +// a runtime activation and cannot be pre-packed). The int32 accumulation is +// untouched: identical per-lane fragment placement (row = row0 + (lane & 15), +// k group kc = (lane >> 4) + (kk >> 3), 8 consecutive natural-order k bytes) +// and identical k0-outer/kk-inner ascending order, so every output byte is +// bit-identical. Falsifiable signature: the B half of the 917,504 measured +// lds_bank_conflicts disappears (32 -> 16 per stage per wave), the B staging's +// 4 x ds_write_b64 become 2 x ds_write_b128, lds_instructions fall +// 688,128 -> 630,784, and the median improves only if the global/LDS staging +// latency chain, not the 16-MMAC burst, is the limiter. // -// 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 6 (mandated epilogue round): the fused scale/bf16 epilogue has +// been inside the compute kernel since iteration 1 and the workspace is +// already ignored ((void)workspace in launch_w8a8_gemm), so the remaining +// epilogue inefficiency of the exact (512, 3584) arm is the STORE pattern: +// store_fragment_scaled issues one 2-byte bf16 store per lane per element - +// 64 scattered stores per wave, each wavefront store touching 16 rows x 4 +// columns so every 32-B store sector is only 25% utilized (the bootstrapped +// PMC counts vmem_write_instructions == 229,376 == 896 x 4 x 16 x 4 exactly). +// The new epilogue is register-batched: the per-row x_scale is loaded once per +// 16-row row-block and the four per-column scales of each fragment's owned +// columns are one aligned float4 load, so the epilogue issues 4 scalar + 16 +// vector vmem loads per lane instead of 80 scalar loads. It then transposes +// the four-element groups within each 4-lane column group in registers (two +// 2x2 shfl_xor steps, 16 then 32 - the lineage-validated pattern ported from +// references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip, which lowers 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 .. +3, converts them to +// bf16 and writes ONE 8-byte store per lane per fragment: 16 coalesced stores +// per wave instead of 64 scattered ones at 100% sector efficiency, i.e. +// vmem_write 229,376 -> 57,344. The scale multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the __float2bfloat16 +// rounding are unchanged, so the stored bits are identical; no LDS, no +// barrier, and no workspace/combine pass is added. Everything else - the A +// swizzle, the tile-contiguous plane-B staging, the k0-outer/kk-inner +// ascending int32 order, the exact (k, n) == (512, 3584) guard, the +// identity-layout 128x64 arm, the pack kernels and the generic scalar +// fallback - is byte-identical. // -// 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 8 (resource round): the exact-shape 128x128 n-major arm keeps its +// winning tile family, its tile-contiguous plane-B pack, its A swizzle, its +// direct 8-byte fragment loads, its k0-outer/kk-inner int32 order and its +// coalesced epilogue, and changes only the occupancy resources: waves per +// block, the VGPR live range and the per-thread LDS staging payload. +// * Measured occupancy of the current best (iteration-6 source): arch_vgpr +// 120, static LDS 16,384 B/block, __launch_bounds__(256, 2) -> 2 blocks/CU +// x 4 waves = 8 waves/CU = 2 waves/SIMD with 16 int32 accumulator VGPR +// fragments + 8 operand fragments + 4 live int4 staging registers per +// thread. The PMC says the kernel is issue-stalled, not LDS-bandwidth- or +// MMAC-bound: 195,964 cycles for 7,168 block-stages, 974,848 +// lds_instructions (4.97 per GPU cycle over 120 CUs) and 1,404,478 +// lds_wait_instructions against a theoretical LDS issue rate of 1/cycle/CU +// - i.e. the SIMDs are starved of co-resident waves, not of LDS slots. +// * Focused change: the same 128x128 macro-tile is covered by EIGHT +// wavefronts of 64 lanes (512 threads, kNMBlockThreads = 8*64) instead of +// four; each wave now owns a 32x64 quadrant = kAccM 2 A fragments x +// kAccN 4 B fragments = 8 m16n16k32 int32 accumulators, i.e. 32 +// accumulator VGPR per thread instead of 64. The staging payload per +// thread halves to exactly one int4 of A (row = tid >> 2, 16-byte K group +// (tid & 3) * 16) and one int4 of B (offset tid * 16), and the arm is +// pinned by __launch_bounds__(kNMBlockThreads, 2) to the 64-VGPR / +// 2-blocks/CU point (65536 VGPR per CU / (2 x 512 lanes) = 64), i.e. +// 2 blocks x 8 waves = 16 waves/CU = 4 waves/SIMD, twice the currently +// measured wave count. LDS stays 16,384 B/block, so 2 blocks use 32,768 B +// of the 65,536-B per-CU LDS (no exact-fit cliff like iteration 7's +// 32,768 B/block double buffer). +// * Why exactly this point: the trusted occupancy-probe candidate list for +// the measured 120-CU device is [2, 3, 4, 5, 6, 7, 8] waves per block; 8 +// is the largest value that still tiles the 128x128 macro-tile with +// rectangular m16n16k32 quadrants (4 M row groups x 2 N column groups) and +// it gives the top of the trusted residency range (8 waves/CU -> 16 +// waves/CU = 4/SIMD). The direct sibling evidence is the lineage reference +// references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip (same exact +// M=4096/N=3584/K=512 shape, same DTK 26.04/gfx928): its 8-wave/512-thread +// 16-waves/CU recipe (config 3, arch_vgpr <= 64, 2 blocks/CU) is the +// recorded 137.4 us winner of that family versus 169.3 us for its +// 12-waves/CU 128x64 form, which is the same occupancy axis this round +// probes while keeping our better 1:1 macro-tile and 0.75-fragment/MMAC +// quadrant instead of that reference's 1.0-fragment/MMAC 32x32 quadrants. +// * Falsifiers: the exact code object must show vgpr_spill_count == 0, +// scratch_bytes == 0 and arch_vgpr <= 64 (else residency falls to 1 +// block/CU = 8 waves/CU, cancelling the whole point); the LDS-read volume +// per MMAC rises 0.5 -> 0.75 fragments, so the median improves only if the +// doubled wave count hides more LDS/barrier latency than the +50% fragment +// reads cost. The int32 accumulation is untouched (exact integer +// additions of the same products), every fragment placement is the same +// 8-byte per-lane operand, and the exact (k, n) == (512, 3584) guard, the +// identity-layout 128x64 arm, the pack kernels and the generic scalar +// fallback 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 4 (tile-shape round): the exact-shape n-major arm moves from the +// 128x64 macro-tile (M:N = 2:1) to a 128x128 macro-tile (M:N = 1:1) for the +// exact (M, N, K) = (4096, 3584, 512) shape; the four-wavefront/256-thread +// block and the direct 8-byte swizzled fragment path are unchanged. Each +// wave now owns a 64x64 quadrant built from sixteen m16n16k32 int32 +// accumulators (kAccM = 4 A fragments x kAccN = 4 B fragments = 16 MMACs per +// kk) instead of a 64x32 quadrant with eight. Why the aspect/area change is +// the lever for this exact M/N/K: +// * Per-block-stage cooperative-load bytes are fixed by the tile shape: +// 128*64 (A) + 128*64 (B) = 16 KiB for 128x128 versus 8 KiB + 4 KiB = +// 12 KiB for 128x64, but the 128x128 tile holds 2x the outputs, so the +// global/L2 read traffic for the whole operator drops from +// M*N*K*(1/128 + 1/64) = 176.2 MB to M*N*K*(1/128 + 1/128) = 117.4 MB +// (-33%). Concretely: A is re-read once per N-tile column (56 -> 28) so +// A traffic halves 117.4 -> 58.7 MB, while B keeps its 32 M-tile re-reads +// at 58.7 MB. The 128x64 arm's 2:1 A:B read imbalance becomes 1:1, which +// is the aspect balance this round is for; the smaller A (2 MiB) and B +// (1.75 MiB) both stay L2-resident either way. +// * LDS fragment reads per MMAC drop 6/8 -> 8/16 (-33%): the 64x32 quadrant +// needed 4 A + 2 B = 6 direct 8-byte reads for 8 MMACs, the 64x64 +// quadrant needs 4 A + 4 B = 8 reads for 16 MMACs. +// * The measured per-stage wall is dominated by the exposed A+B L2 round +// trip in front of the MMAC burst (exact code object: 3 x +// global_load_dwordx4 -> s_waitcnt vmcnt(2)/(1)/(0) -> 6 x ds_write_b64, +// per 64-K stage). Halving the macro-tile count (grid 1792 -> 896) halves +// the number of exposed stage walls for the same total work. +// * 896 blocks is still ~3.7 waves of the 240 co-resident blocks +// (2 blocks/CU x 120 CU), i.e. enough blocks to occupy every CU; no +// split-K and the workspace stays unused. +// Residency: LDS 16,384 B/block (2 x 16,384 = 32,768 <= 64 KiB) and the +// sixteen accumulators cost 64 VGPR, so the arm is pinned to 2 blocks/CU = +// 8 waves/CU by an explicit __launch_bounds__(256, 2) register cap at +// 128 VGPR (floor(16384 VGPR/SIMD / (128 x 64 lanes)) = 2 waves/SIMD); +// falsified if the exact code object shows > 128 VGPR or any spill -> 1 +// block/CU = 4 waves/CU. +// Everything else is byte-identical to iteration 2/3: the dense 8-byte-group +// swizzle, load_frag8_swizzled placement, the one-time n-major packed B, the +// k0-outer/kk-inner ascending int32 order, the exact (k, n) == (512, 3584) +// guard, the identity-layout 128x64 arm, the pack kernels and the generic +// scalar fallback. // -// 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). +// Iteration 2 (operand-reuse round): keep the cooperative A+B LDS staging +// (the reuse argument below favours it over per-fragment global loads) but +// replace the two structural costs the bootstrap PMC exposes on the assigned +// (M, N, K) = (4096, 3584, 512) shape: +// * lds_instructions 2,465,792 and lds_wait_instructions 7,165,860: +// the installed du_load_matrix_sync int8 loaders lower to a per-byte +// ds_read_u8 gather plus mask/OR reassembly VALU for EVERY fragment. +// * lds_bank_conflicts 6,881,280: the 80-byte padded row stride has bank +// period 8 over a fragment's 16 rows, so lanes r and r+8 collide. +// Reuse accounting for one 128x64 macro-tile / 64-K stage (the reason direct +// global operand loads are rejected): +// A[128,64] = 8 KiB staged; each A fragment row-block feeds two B fragments +// (kAccN = 2) and two kk sub-steps, i.e. every staged A byte is consumed +// by 4 or more MMAC operand reads; a direct global fragment load would +// re-fetch it per fragment (4x redundant global traffic). +// B[64,64] = 4 KiB staged; each B fragment feeds all four A row-blocks +// (kAccM = 4), i.e. 4x reuse. So cooperative staging is retained and the +// real lever is the *read* path and the LDS conflict structure. +// Change (one mechanism, both tiles): +// * Both LDS tiles become dense 64-byte-row swizzled 8-byte-group layouts +// (no pad). group(row, c) = (((row & 15) + 4*(c & 3)) & 15) + 16*c + +// 128*(row >> 4); the writers emit two 8-byte LDS stores per 16-byte +// global chunk at group(row, 2j) / group(row, 2j+1) and the fragment +// loader reads the swizzle back. Every 16-lane fragment load then touches +// 16 distinct bank pairs (zero conflicts) instead of the period-8 +// collision. +// * Every A and B fragment read becomes ONE 8-byte LDS load per lane +// (load_frag8_swizzled), reproducing du_load_matrix_sync's per-lane +// (row = row0 + (lane & 15), c = (lane >> 4) + (kk >> 3), 8 consecutive +// k bytes in natural order) placement, so the int32 accumulation is +// bit-identical. +// * B is read col_major from a one-time n-major transpose +// packed[n*k + kk] = raw[kk*n + n] (produced outside the timed/Graph +// region by launch_pack_w8a8_weight for exactly (k, n) == (512, 3584)); +// this makes each lane's 8 k bytes contiguous. The identity [K, N] copy +// is kept for every other (k, n). +// Iteration-2 geometry (superseded for the n-major arm by the iteration-4 +// 128x128 tile above; the identity-layout arm still uses it): one 128x64 +// tile/block, four gfx928 wavefronts of 64 lanes (256 threads), each wave +// owning a 64x32 quadrant (kAccM = 4, kAccN = 2, eight m16n16k32 int32 +// accumulators), 2 barriers per 64-K stage, no split-K, grid +// dim3(N/64, M/128) = (56, 32) = 1792 blocks, LDS = 8 KiB + 4 KiB = 12,288 +// B/block. +// * The fused int32 * x_scale * weight_scale -> bf16 epilogue keeps the +// verified gfx928 int8 m16n16k32 accumulator ownership (row = lane & 15, +// col = (lane >> 4) + 4*i), multiply order (float(dot) * x_scale[row] * +// weight_scale[col]) and bf16 rounding of the harness reference. +// * The identity-layout 128x64 DUMMA arm is retained for every non-packed +// large shape, and the generic scalar int8/int32 fallback covers every +// unmatched (m, n, k), including all small-M API cases (M=2, M=16) and +// M < 128 with the same (K, N): one thread per output element, exact +// int32 dot over K, then float(acc) * x_scale[row] * weight_scale[col] -> +// bf16. For exactly (512, 3584) the fallback decodes the n-major pack. +// * launch_pack_w8a8_weight runs outside the timed/Graph region, allocates +// nothing, keeps the same byte count and graph-stable addresses: the +// n-major transpose for exactly (k, n) == (512, 3584), identity for every +// other (k, n), fp32 scale always copied identity. +// +// The timed operator performs no allocation, compilation, autotuning, weight +// packing, host/device synchronization, or default-stream launch; it uses only +// the caller-provided out and workspace and launches on PyTorch's current HIP +// stream via hipLaunchKernelGGL. #include #include @@ -251,71 +329,55 @@ 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 kWaveSize = 64; // gfx928 native wavefront, not 32 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; +constexpr int kPad = 16; // 16 B LDS row padding (64 -> 80 B rows) +constexpr int kAStride = kStageK + kPad; +constexpr int kBStride = kBlockN + kPad; +constexpr int kBlockThreads = 4 * kWaveSize; // 256 = four gfx928 wavefronts -// 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; +// Iteration-4 tile-shape round: the exact-shape n-major arm uses a 128x128 +// macro-tile (M:N = 1:1) instead of the 128x64 identity-arm geometry. +// Iteration-8 resource round: the same 128x128 macro-tile is covered by EIGHT +// wavefronts of 64 lanes (512 threads), each owning a 32x64 quadrant built +// from kAccM = 2 A fragments x kAccN = 4 B fragments = eight m16n16k32 int32 +// accumulators (32 VGPR of accumulator per thread instead of 64). The block +// is pinned to 2 blocks/CU by __launch_bounds__(512, 2), i.e. a 64-VGPR cap: +// 2 blocks x 8 waves = 16 waves/CU = 4 waves/SIMD instead of the +// measured 2 blocks x 4 waves = 8 waves/CU = 2 waves/SIMD. +constexpr int kNMBlockM = 128; +constexpr int kNMBlockN = 128; +constexpr int kNMQuadM = 32; // per-wave M extent (4 row groups) +constexpr int kNMQuadN = 64; // per-wave N extent (2 col groups) +constexpr int kNMBlockThreads = 8 * kWaveSize; // 512 = eight gfx928 wavefronts -inline bool is_swizzled_pack(int k, int n) { - return (k == kPackSwizzleK1 && n == kPackSwizzleN1) || - (k == kPackSwizzleK2 && n == kPackSwizzleN2); -} +// Exact (k, n) pair whose weight buffer uses the iteration-5 tile-contiguous +// 64-K-stage pack packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] = +// raw[kk*n + col] (identity [K, N] for every other pair). +constexpr int kPackK = 512; +constexpr int kPackN = 3584; +constexpr int kPackTileN = 128; // N rows per tile-contiguous pack region +constexpr int kPackStageK = 64; // k bytes per pack stage + +using bf16_t = hip_bfloat16; -// 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; + return k == kPackK && n == kPackN; } -// 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. +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding match +// the harness reference ((dot.float() * x_scale) * weight_scale.T -> bf16), so +// the stored bf16 bits are identical. // --------------------------------------------------------------------------- template -__device__ __forceinline__ void store_prefill_fragment( +__device__ __forceinline__ void store_fragment_scaled( const AccFragment& frag, const float* __restrict__ x_scale, const float* __restrict__ weight_scale, @@ -340,45 +402,75 @@ __device__ __forceinline__ void store_prefill_fragment( } // --------------------------------------------------------------------------- -// 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. +// Bank-safe dense 8-byte-group LDS map for the exact-shape swizzled arm. +// group(row, c) is the 8-byte group index of tile position (row, k 8-byte +// group c). For a fragment load the 16 lanes of a group share row0 and kk, so +// row & 15 ranges over 0..15 and c is fixed; (row & 15) -> ((row & 15) + +// 4*(c & 3)) & 15 is a permutation, giving 16 distinct bank pairs and zero +// bank conflicts over the whole fragment read. +// --------------------------------------------------------------------------- +__device__ __forceinline__ int swizzle_group(int row, int c) { + return (((row & 15) + 4 * (c & 3)) & 15) + 16 * c + 128 * (row >> 4); +} + +// One 8-byte LDS store of each half of a 16-byte global chunk. +__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); +} + +// Direct 8-byte fragment load that reproduces du_load_matrix_sync's per-lane +// placement exactly: row = row0 + (lane & 15), k group c = (lane >> 4) + +// (kk >> 3), eight consecutive k bytes in natural order. Writing them +// straight into the fragment storage lets the compiler emit one 8-byte LDS +// read per fragment with none of the library's per-byte gather/reassembly +// VALU, so the int32 accumulation is bit-identical. +__device__ __forceinline__ void load_frag8_swizzled( + DUFragmentBase& frag, + const signed char* __restrict__ p, + int lane, + int row0, + int kk) { + const int row = row0 + (lane & 15); + const int c = (lane >> 4) + (kk >> 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(p + swizzle_group(row, c) * 8); +} + +// Iteration-5 (packing round) B fragment loader for the tile-contiguous plane +// LDS layout b_tile[kc][n][8] (kc = 8-byte k sub-chunk; plane stride +// kNMBlockN*8 = 1024 B). Lane l reads the same 8 consecutive k bytes as +// load_frag8_swizzled's placement (n row = row0 + (lane & 15), k group +// kc = (lane >> 4) + (kk >> 3)) but from the lane-linear address +// kc * (kNMBlockN * 8) + row * 8, +// i.e. ((lane >> 4) * 1024) + ((lane & 15) * 8) relative to row0's plane +// origin, so each 16-lane group of a 128-B LDS phase spans all 32 banks +// exactly once (zero bank conflicts) and the compiler emits one ds_read_b64. +__device__ __forceinline__ void load_frag8_plane( + DUFragmentBase& frag, + const signed char* __restrict__ p, + int lane, + int row0, + int kk) { + const int row = row0 + (lane & 15); + const int kc = (lane >> 4) + (kk >> 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast( + p + kc * (kNMBlockN * 8) + row * 8); +} + // --------------------------------------------------------------------------- -// 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. +// Iteration 9 (consolidate round): plain round-to-nearest-even fp32 -> bf16 +// conversion, bit-identical to hip's __float2bfloat16 for every finite input +// (the only values this GEMM can produce: int32 accumulators times finite +// scales), but without the exec-masked inf/NaN fixup sequence the DTK 26.04 +// gfx928 lowering emits around it (see the header note). The canonical RNE +// bit trick: add 0x7fff plus the retained LSB so ties round to even, then +// truncate. Same form as the sibling reference +// references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip on this DTK. // --------------------------------------------------------------------------- __device__ __forceinline__ unsigned short bf16_rne_u16(float x) { union { @@ -390,21 +482,52 @@ __device__ __forceinline__ unsigned short bf16_rne_u16(float x) { return static_cast(u >> 16); } -template -__device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( +// --------------------------------------------------------------------------- +// Iteration-6 (epilogue round): register-batched COALESCED direct-fragment +// epilogue for the exact-shape arm, ported from the lineage-validated sibling +// reference references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip. The +// per-element variant (store_fragment_scaled) issues one 2-byte bf16 store per +// lane per element: four scattered stores per fragment, 64 per wave, each +// wavefront store touching 16 rows x 4 columns so every 32-B store sector is +// only 25% utilized (the bootstrap PMC's 229,376 vmem_write instructions = +// 896 blocks x 4 waves x 16 fragments x 4). This epilogue transposes the +// four-element groups within each 4-lane column group (lanes r, r+16, r+32, +// r+48 - a 4x4 transpose done as two 2x2 steps with shfl_xor 16 then 32, the +// lineage-validated pattern that lowers to ds_bpermute at the block tail +// where the LDS pipe is idle), so lane (r, c4) ends up holding the four +// CONTIGUOUS columns base_col + 4*c4 .. +3, converts them to bf16 and writes +// ONE 8-byte store (100% store sector efficiency; 16 stores per wave instead +// of 64). 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. +// xs is batched by the caller (one load per 16-row row-block) and ws is the +// float4 of the four per-column scales at col0, 16-B aligned because base_col +// is a multiple of 16, so the epilogue issues 2 scalar + 8 vector vmem loads +// per lane instead of 80 scalar loads. Only the exact-shape n-major arm calls +// this helper, and its dispatch guard is m >= 128 with m % 128 == 0, so every +// row of every block is in bounds and no tail-M mask is needed: the shuffles +// never mix active and inactive lanes, and n*2 is a multiple of 8, so the +// 8-byte store is aligned. +// +// Iteration 9 (consolidate round): the bf16 conversion below is a plain RNE +// (bf16_rne_u16) instead of __float2bfloat16. The exact code object of the +// current best showed the DTK 26.04 lowering wrapping every element in an +// exec-masked inf/NaN fixup (~9 vector instructions + two s_and_saveexec / +// s_andn2_saveexec scalar manipulations) whose only effect vs RNE is on +// inf/NaN inputs; int32 accumulators times finite scales are always finite, +// and for finite floats the two forms produce the identical 16-bit pattern. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void store_fragment_coalesced_scaled( const AccFragment& frag, float xs, - const float4& ws, + 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]; @@ -431,165 +554,62 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( 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. + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte aligned: + // base_col is a multiple of 16 and n is a multiple of 4); ws is the float4 + // of per-column scales at col0. const int col0 = base_col + 4 * c4; + 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; - 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); + const uint64_t 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); + *reinterpret_cast( + out + static_cast(row) * n + col0) = packed; } -// --------------------------------------------------------------------------- -// 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. +// B[64,64] from the raw [K, N] row-major weight into a single-buffered LDS +// stage. Two barriers per 64-K stage: one after the cooperative load, one +// before the next stage overwrites the buffer. 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) { +__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 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; 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). + // Single-buffered stage: A[128, 80] + B[64, 80] = 15,360 B/block (four + // blocks fit the 64 KiB per-CU 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> + DUFragment b_frag0, b_frag1; DUFragment acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; @@ -602,21 +622,18 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne 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 + // 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 the 16-byte K-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). + // B: 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). 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) + const int b_k = tid >> 2; // K row 0..63 + const int b_nc = (tid & 3) * 16; // N-column group 0/16/32/48 for (int k0 = 0; k0 < k; k0 += kStageK) { *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = @@ -625,43 +642,23 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne *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); - } + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + k0 + + stage_col); + *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. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk = 0 then + // 32), matching the reference int32 accumulation exactly. #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(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, @@ -689,101 +686,84 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne 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); + store_fragment_scaled(acc00, x_scale, weight_scale, out, base_row, base_col, + m, n, lane); + store_fragment_scaled(acc01, x_scale, weight_scale, out, base_row, + base_col + kTileN, m, n, lane); + store_fragment_scaled(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_fragment_scaled(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, lane); + store_fragment_scaled(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, lane); + store_fragment_scaled(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, m, n, lane); + store_fragment_scaled(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, lane); + store_fragment_scaled(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. +// Exact-shape (k, n) == (512, 3584) large-M prefill: 128x128 macro-tile per +// block, EIGHT gfx928 wavefronts of 64 lanes (512 threads). Each wave owns a +// 32x64 quadrant = kAccM 2 A fragments x kAccN 4 B fragments = eight +// m16n16k32 int32 accumulators (8 MMACs per kk). Cooperative A+B LDS staging +// is retained (each A row-block feeds 4 B fragments x 2 kk, each B fragment +// feeds 2 A row-blocks x 2 kk) and every fragment read is one direct 8-byte +// LDS load. A comes from the identity [M, K] x_q (row-major, stride k) and +// keeps the dense swizzled 8-byte-group LDS layout; B comes from the +// iteration-5 one-time tile-contiguous 64-K-stage pack +// packed[(((k0*(n/128)+nt)*8+kc)*128+row)*8+b] = raw[kk*n+col] and uses the +// plane LDS layout [kc][128][8], so its staging is a plain contiguous 16-B +// vector copy and its fragment loads are lane-linear and bank-conflict-free. +// Two barriers per 64-K stage, single-buffered, int32 accumulation resident +// through the whole K loop, no split-K (grid dim3(28, 32) = 896 blocks over +// 120 CUs, ~3.7 waves of the 240 co-resident blocks). Dispatch guarantees +// k == 512, n == 3584, m >= 128 and m % 128 == 0, so every global load is +// in-bounds and 16-byte aligned. The two-argument __launch_bounds__ caps the +// register file at the 64-VGPR 2-blocks/CU residency point for 512-thread +// blocks (65536 VGPR per CU / (2 x 512 lanes) = 64): 2 blocks x 8 waves = +// 16 waves/CU = 4 waves/SIMD instead of the measured 120-VGPR 128x128 +// four-wave form at 2 waves/SIMD. The iteration-6 epilogue is kept +// register-only: 2 batched per-row scale loads, one aligned float4 per-column +// scale load per fragment and one 8-byte coalesced store per lane per +// fragment, so it adds no LDS, no barrier and no occupancy pressure. // --------------------------------------------------------------------------- -__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 - +__global__ __launch_bounds__(kNMBlockThreads, 2) void +w8a8_dumma_prefill_128x128_nmajor_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) { + // Iteration 10 (consolidate round): the dispatch guard for this arm is + // is_tile_major_pack(k, n) && m >= kNMBlockM && m % kNMBlockM == 0, so the + // 128-row tile is always fully in bounds and no tail-M masking is needed. + (void)m; const int tid = static_cast(threadIdx.x); - const int wave = tid / kWaveSize; + const int wave = tid / kWaveSize; // 0..7 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; + const int wave_row = wave >> 1; // 0..3 + const int wave_col = wave & 1; // 0..1 + const int m0 = static_cast(blockIdx.y) * kNMBlockM; + const int n0 = static_cast(blockIdx.x) * kNMBlockN; + const int local_row = wave_row * kNMQuadM; // 0, 32, 64, 96 + const int local_col = wave_col * kNMQuadN; // 0 or 64 - // 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]; + // A 8 KiB (dense swizzled 8-byte groups) + B 8 KiB (iteration-5 + // tile-contiguous plane [kc][128][8]) = 16,384 B/block. + __shared__ __align__(16) int8_t a_tile[kNMBlockM * kStageK]; + __shared__ __align__(16) int8_t b_tile[kNMBlockN * kStageK]; DUFragment - a_frag0, a_frag1, a_frag2, a_frag3; + a_frag0, a_frag1; 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; + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13; du_fill_fragment(acc00, 0); du_fill_fragment(acc01, 0); du_fill_fragment(acc02, 0); @@ -792,102 +772,59 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern 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). + // Cooperative staging (iteration-8 resource round): with 512 threads each + // thread owns exactly ONE 16-byte chunk of A and ONE 16-byte chunk of B, so + // the staging payload per thread halves (1 int4 + 1 int4 instead of + // 2 + 2) - this is what keeps the live-range budget compatible with the + // 64-VGPR / 2-blocks/CU occupancy point. + // A: thread tid owns row tid >> 2 and the 16-byte K column group + // (tid & 3) * 16; the swizzled 8-byte-group LDS store is unchanged. + // B (iteration-5 packing round): the tile-contiguous pack lays each + // (64-K stage, 128-N tile) B tile out as one contiguous 8192-B region + // already in the LDS plane order [kc][128 N rows][8 bytes], so staging + // is a plain 1 x 16-B contiguous vector copy per thread. + const int vector_byte_offset = tid * static_cast(sizeof(int4)); + const int stage_row = vector_byte_offset / kStageK; // 0..127 + const int stage_col = vector_byte_offset - stage_row * kStageK; // 0,16,32,48 + const int b_tile_off = vector_byte_offset; // 0..8176 + + // Iteration 13 (consolidate round): one-stage register prefetch of the + // cooperative staging payload. The A row base and the B per-stage stride are + // loop-invariant, so the next 64-K A/B chunks are loaded while the current + // stage's fragment reads and MMACs are still draining; at the consumer side + // the load is already in flight across the trailing __syncthreads() instead + // of being issued and drained within a handful of instructions at the stage + // top. The loaded values, their order, the LDS bytes and the int32 + // accumulation are unchanged, so every stored bit is unchanged. + const int8_t* const a_stage_base = + x_q + static_cast(m0 + stage_row) * k + stage_col; + const int64_t b_stage_step = + static_cast(n / kNMBlockN) * (kNMBlockN * kStageK); + const int64_t b_stage_base = + static_cast(n0 / kNMBlockN) * (kNMBlockN * kStageK); + int4 a_chunk = *reinterpret_cast(a_stage_base); + int4 b_chunk = + *reinterpret_cast(weight + b_stage_base + b_tile_off); + 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); - } + store_chunk16_swizzled(a_tile, stage_row, stage_col >> 4, a_chunk); + *reinterpret_cast(b_tile + b_tile_off) = b_chunk; __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. + // Each wave consumes its 32x64 quadrant: eight m16n16k32 MMACs per kk. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk = 0 then + // 32), matching the reference int32 accumulation exactly (the DUMMA + // accumulator is exact int32 addition, so the 2 x 4 fragment partition + // over the same products is bit-identical). #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); + load_frag8_plane(b_frag0, b_tile, lane, local_col, kk); + load_frag8_plane(b_frag1, b_tile, lane, local_col + kTileN, kk); + load_frag8_plane(b_frag2, b_tile, lane, local_col + 2 * kTileN, kk); + load_frag8_plane(b_frag3, b_tile, lane, local_col + 3 * 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(acc02, a_frag0, b_frag2, acc02); @@ -896,626 +833,83 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern 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]; + // Iteration 13 (consolidate round): issue the next 64-K stage's A and B + // chunks before the trailing barrier so their L2 round trip overlaps the + // barrier wait and the next stage's commit instead of being drained + // immediately at the stage top. The last stage re-reads its own chunk + // (k1 == k0) rather than running past the final K block, so every + // prefetch address stays in bounds. { - // 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(); + const int k1 = (k0 + kStageK < k) ? (k0 + kStageK) : k0; + a_chunk = *reinterpret_cast(a_stage_base + k1); + b_chunk = *reinterpret_cast( + weight + static_cast(k1 / kStageK) * b_stage_step + + b_stage_base + b_tile_off); } - } 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(); - } + // 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); - } - } + // --------------------------------------------------------------------- + // Iteration-6 (epilogue round): register-batched coalesced epilogue, kept + // unchanged; the iteration-8 eight-wave quadrant only changes how many + // fragments each wave owns (2 row-blocks x 4 col-blocks = 8 fragments). + // Iteration 10 (consolidate round): the dead per-fragment tail-M guard is + // removed from this arm. The dispatch guard is m >= 128 with m % 128 == 0, + // so base_row + (lane & 15) is always < m and the two (rs < m) x_scale + // ternaries and the (row >= m) early return in the store helper never fire; + // as written they forced each of the eight fragment epilogues into its own + // exec-masked basic block (eight s_and_saveexec / s_or_b64 / s_cbranch + // groups and eight redundant lane-id / shuffle-index computations in the + // iteration-10 code object). Removing them keeps the exact 128x128 tile, + // the eight-wave geometry, the 16,384-B single-buffered staging, the LDS + // fragment loads, the int32 accumulation and the coalesced 8-byte store, + // and lets the eight independent fragment epilogues issue as one + // straight-line block. The 128x64 generic arm and the scalar fallback keep + // their own bounds checks unchanged. + // --------------------------------------------------------------------- + const int gcol_base = n0 + local_col; + const int rs0 = m0 + local_row + (lane & 15); + const float xs0 = x_scale[rs0]; + const int rs1 = rs0 + kTileM; + const float xs1 = x_scale[rs1]; + store_fragment_coalesced_scaled(acc00, xs0, weight_scale, out, + m0 + local_row, gcol_base, n, lane); + store_fragment_coalesced_scaled(acc01, xs0, weight_scale, out, + m0 + local_row, gcol_base + kTileN, n, lane); + store_fragment_coalesced_scaled(acc02, xs0, weight_scale, out, + m0 + local_row, gcol_base + 2 * kTileN, n, + lane); + store_fragment_coalesced_scaled(acc03, xs0, weight_scale, out, + m0 + local_row, gcol_base + 3 * kTileN, n, + lane); + store_fragment_coalesced_scaled(acc10, xs1, weight_scale, out, + m0 + local_row + kTileM, gcol_base, n, lane); + store_fragment_coalesced_scaled(acc11, xs1, weight_scale, out, + m0 + local_row + kTileM, gcol_base + kTileN, + n, lane); + store_fragment_coalesced_scaled(acc12, xs1, weight_scale, out, + m0 + local_row + kTileM, + gcol_base + 2 * kTileN, n, lane); + store_fragment_coalesced_scaled(acc13, xs1, weight_scale, out, + m0 + local_row + kTileM, + gcol_base + 3 * kTileN, 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. +// over K, then fused x_scale * weight_scale, then bf16 store. kPackedNMajor +// == true decodes the iteration-5 tile-contiguous 64-K-stage pack +// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] = raw[kk*n + col] for +// exactly (k, n) == (512, 3584); false reads the identity [K, N] row-major +// layout. Handles every unmatched (m, n, k), including all small-M API cases +// (M=2, M=16), the paired same-(K, N) shapes, and M < 128 with the exact +// packed pair. // --------------------------------------------------------------------------- -template +template __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, @@ -1535,39 +929,21 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( 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); - } + if constexpr (kPackedNMajor) { + // Tile-contiguous pack: 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. + const int n_row = col & (kPackTileN - 1); + const int nt = col / kPackTileN; + for (int kk = 0; kk < k; ++kk) { + const int k0 = kk / kPackStageK; + const int kc = (kk >> 3) & 7; + const int b = kk & 7; + const int64_t tile = + static_cast(k0) * (n / kPackTileN) + nt; + const int8_t* p = + weight + ((tile * 8 + kc) * kPackTileN + n_row) * 8 + b; + acc += static_cast(a_row[kk]) * static_cast(*p); } } else { // identity [K, N] row-major: column col is strided by n. @@ -1583,37 +959,20 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( } // --------------------------------------------------------------------------- -// 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. +// Iteration-5 packing round: one-time tile-contiguous 64-K-stage permutation +// for exactly (k, n) == (512, 3584): +// 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) DUMMA B tile is ONE contiguous 8192-B global +// region laid out as [kc = 8-byte k sub-chunk][128 N rows][8 bytes], which is +// exactly the plane order the GEMM stages into LDS, so the staging is a plain +// contiguous 16-B vector copy and the B fragment reads are lane-linear and +// bank-conflict-free. 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 and the same graph-stable buffer +// address. // --------------------------------------------------------------------------- -__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( +__global__ __launch_bounds__(256) void w8a8_pack_tile_major_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, int k, @@ -1624,46 +983,31 @@ __global__ __launch_bounds__(256) void w8a8_pack_swizzle_kernel( 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 k0 = kk / kPackStageK; const int kc = (kk >> 3) & 7; const int b = kk & 7; - packed[(((static_cast(k0) * 8 + kc) * n + col) * 8) + b] = - raw[linear]; + const int nt = col / kPackTileN; + const int row = col & (kPackTileN - 1); + const int64_t tile = static_cast(k0) * (n / kPackTileN) + nt; + packed[((tile * 8 + kc) * kPackTileN + row) * 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( +// --------------------------------------------------------------------------- +// Bootstrap identity device-to-device packing for every (k, n) except the +// exact tile-major-packed pair. Runs once outside the timed +// region and outside Graph capture, keeps the same byte count and +// graph-stable buffer addresses, and is the generic fallback for every +// (k, n). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, - int k, - int n) { + int64_t count) { 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]; + if (linear < count) { + packed[linear] = raw[linear]; } } @@ -1685,187 +1029,89 @@ __global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( // 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) { +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); - } + // Explicit dispatch. For exactly (k, n) == (512, 3584) the weight is the + // one-time tile-contiguous 64-K-stage pack: large-M uses the 128x128 + // plane-B direct-load DUMMA arm (iteration-8 resource round: eight + // wavefronts / 512 threads per block, the wavefront=64 geometry of the + // exact-shape arm), every other m (small-M API cases and M < 128) decodes + // the packed layout in the scalar fallback. Every other (m, n, k) keeps + // the bootstrap behaviour: the identity-layout 128x64 DUMMA arm (four + // wavefronts / 256 threads) for large shapes, the scalar fallback + // otherwise. + if (is_tile_major_pack(k, n) && m >= kNMBlockM && m % kNMBlockM == 0) { + const dim3 grid(static_cast(n / kNMBlockN), + static_cast(m / kNMBlockM)); + const dim3 block(kNMBlockThreads); + hipLaunchKernelGGL(w8a8_dumma_prefill_128x128_nmajor_kernel, grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else if (is_tile_major_pack(k, n)) { + 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 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); - } + hipLaunchKernelGGL(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); - } + 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) { +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, + // Iteration-5 packing round: one-time tile-contiguous 64-K-stage pack so + // the DUMMA arm stages each B tile with contiguous 16-B vector loads and + // LDS stores and loads B fragments lane-linearly with zero bank conflicts. + const dim3 tile_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_tile_major_kernel, tile_grid, block, 0, stream, raw_weight, packed_weight, k, n); } else { - hipLaunchKernelGGL(w8a8_pack_identity_kernel, - weight_grid, block, 0, stream, + 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(); + const dim3 scale_grid(static_cast((n + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, scale_grid, block, 0, + stream, weight_scale, packed_weight_scale, n); } // @@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 index 72932e37..b327ab35 100644 --- 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 @@ -1,143 +1,267 @@ -// @@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 +// @@variant shape=glm_tp8_q_b_proj_m4096 commit=9510ab4d00ea72a47db8821d2b56b35395c23ba5 added=2026-09-11 +// median_us=209.5 p90_us=210.2 speedup=68.68 baseline_us=1.439e+04 +// source=test-9-11-1-b01955b3/iteration_001/glm_tp8_q_b_proj_m4096 // 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 +// Worker worker_3 (physical GPU 3), assigned shapes: +// glm_tp8_q_b_proj_m4096 : M=4096, N=2048, K=2048 // // 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 +// Iteration 1 bootstrap (correctness-first): +// * Large-M DUMMA path (m >= 128 && m % 128 == 0 && n % 64 == 0 && +// k % 64 == 0): native INT8 DUMMA m16n16k32 tiled kernel, one 128x64 +// output tile per block, four gfx928 wavefronts of 64 lanes (256 +// threads). The block cooperatively stages A[128,64] (row-major, +// 16-byte int4 loads) and B[64,64] (raw [K, N] row-major, 16-byte int4 +// loads) into a single-buffered LDS stage (128*80 + 64*80 = 15,360 B +// with 16-byte per-row padding for bank skew) and performs two +// __syncthreads per K stage. Each wave owns a 64x32 quadrant built from +// eight m16n16k32 int32 accumulator fragments (four M x two N). Int32 +// accumulation order is k0-outer over 64-K stages, kk-inner (kk=0 then +// kk=32). The fused epilogue computes +// float(dot) * x_scale[row] * weight_scale[col] and stores bf16 directly // 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). +// Grid = (N/64, M/128); for the assigned shape that is 32 x 32 = 1024 +// blocks, far more than the 120 CUs, so no split-K is needed. +// * launch_pack_w8a8_weight is the bootstrap identity device-to-device copy +// for both the int8 weight and the fp32 scale (for every (k, n) except the +// exact pair that iteration 2 packs n-major below), so the GEMM consumes +// the raw [K, N] row-major weight directly. +// * Generic scalar int8/int32 fallback (one thread per output element) for +// every unmatched (m, n, k), including all small-M API cases (M=2, M=16) +// and any M in (0, 128) sharing the assigned (N, K). // -// 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 2 (operand-reuse round): the exact (k, n) == (2048, 2048) +// q_b_proj pair keeps the 128x64 macro-tile and the cooperative A+B LDS +// staging (each staged byte is consumed by two of the four waves and each +// staged fragment by two MMACs, so cooperative staging halves global operand +// traffic versus direct-to-register global fragment loads: ~402 MiB instead of +// ~805 MiB for the assigned shape). The staged B path is made direct, +// vectorized and bank-safe instead: +// * The weight is packed once (outside the timed region, outside Graph +// capture, same byte count and same graph-stable buffer) to n-major +// packed[n * K + k] == raw[k * N + n], so 16 consecutive k bytes of one n +// row are one aligned, perfectly coalesced int4 global read. +// * The B LDS tile is stored n-major with a 72-byte (18-word) bank-skewed +// row stride; each lane's 8-byte col_major fragment slot then starts on +// bank (18 * (lane & 15) + 2 * (lane >> 4)) mod 32, and 18 * r mod 32 is a +// permutation of the 16 even bank residues, so a 16-lane LDS phase covers +// all 32 banks exactly once (zero conflicts) and every staged access is +// 8-byte aligned. +// * Each B fragment is loaded with one 8-byte LDS read (load_fragment8) +// instead of the row_major loader's eight strided ds_read_u8 byte gathers +// plus mask/OR reassembly VALU. The 8 operand bytes are identical in the +// same fragment slots, so the int32 accumulation is still bit-identical. +// The exact (k, n) == (2048, 2048) pair takes this arm; every other (k, n) +// keeps the identity pack and the byte-identical raw [K, N] k-major arm. +// +// Iteration 4 (tile-aspect round): for the exact (k, n) == (2048, 2048) +// q_b_proj arm the 128x64 macro-tile is widened along N to a 128x128 tile +// (w8a8_dumma_prefill_128x128_nmajor_kernel), giving an M:N tile aspect of +// 1:1. The 128x64 best kernel re-reads A once per N/64 = 32 column tiles +// (~268 MiB A traffic) against B's ~134 MiB (A:B = 2:1) and issues 6 LDS +// fragment reads per 8 MMACs in each 64x32 wave quadrant; the 128x128 tile +// halves the A-side re-reads (A reuse 32 -> 16, ~134 MiB) while B reuse stays +// 32 (~134 MiB), balancing per-byte A:B tile traffic at 1:1 and cutting total +// staged tile traffic ~402 -> ~268 MiB (-33%). Each wave now owns a 64x64 +// quadrant (four A x four B fragments, sixteen int32 accumulators) and issues +// 8 LDS fragment reads per 16 MMACs (-33% per MMAC). The n-major stride-72 B +// layout and the direct load_fragment8 reads from iteration 2 are kept +// unchanged; the block stages A[128,64] (two int4/thread, 80-byte-strided +// rows) and the packed B[128,64] (two 16-byte n-row chunks/thread into the +// 72-byte-strided n-major tile) into a 19,456 B single buffer, so two +// blocks/CU still co-reside. Accumulation order (k0-outer, kk-inner) and the +// fragment element-to-slot mapping are unchanged, so the int32 accumulation is +// bit-identical; the scalar fallback and the 128x64 raw [K, N] arm for every +// other (k, n) are 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. +// weight is packed FRAGMENT-MAJOR instead of n-major. One DUMMA matrix_b +// tile (m16n16k32 -> 16 n x 32 k int8 = 512 contiguous bytes) is one +// graph-stable, 16-byte-aligned block, so: +// * the packer writes a pure permutation of the same K * N bytes into the +// same buffer: byte (`n`, `k`) of the tile (`n/16`, `k/32`) sits at +// block * 512 + (k % 32) / 8 * 128 + (n % 16) * 8 + (k % 8); +// * the cooperative stage copies each packed 16-byte chunk with ONE aligned +// int4 global read and ONE aligned int4 LDS store (no 8-byte halves and +// no shift/or recombine VALU), and every staged 16-byte chunk is +// 16-byte aligned by construction; +// * each wave-level B fragment is one 8-byte LDS read per lane at +// (lane & 15) * 8 + (lane >> 4) * 128, i.e. word 2 * (lane & 15) plus a +// 32-word k-group stride inside the 512-byte tile, so a 16-lane LDS +// phase (lane & 15 = 0..15) covers all 32 banks exactly once with zero +// padding (bank-conflict-free, fully vector-loadable). +// The fragment bytes, their (k, n) slots and the k0-outer/kk-inner int32 +// accumulation order are unchanged, so correctness is bit-identical; the +// scalar fallback decodes the same permutation for the unmatched-M cases of +// the exact (k, n) pair. // // 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. +// never used ((void)workspace; no split-K, no combine pass anywhere in the +// call chain), so the remaining epilogue inefficiency of the exact +// (k, n) == (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 touching 16 rows x 4 +// columns, so every 32-byte store 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; a 4x4 register transpose in two 2x2 steps +// with __shfl_xor 16 then 32), so lane (r, c4) owns the four CONTIGUOUS +// columns 4*c4 .. +3 and writes ONE aligned 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. The +// 128x64 arm keeps the byte-identical per-element epilogue and the generic +// scalar fallback is untouched. +// +// Iteration 7 (compute-pipeline round): the 128x128 packed arm's per-K-stage +// critical path still begins with four cooperative 16-byte global staging loads +// whose latency is drained immediately by the s_waitcnt vmcnt(3)/(2)/(1)/(0) +// sequence that precedes the four ds_write_b128 (the code object shows the load +// and its store in the same step, so the L2/HBM round trip sits inside the +// barrier-to-barrier wall of every one of the 32 K stages; the PMC LDS wait +// count and the ~0.05 MMAC/cycle/SIMD issue rate show the warp stalling far +// more than issuing). The staging loop now carries a ONE-K-STAGE register +// lookahead: this thread's four staging chunks (A rows stage_row and +// stage_row + 64, packed B tiles b_blk and b_blk + 8) are loaded for K stage +// s + 1 right AFTER the publish barrier of stage s and BEFORE stage s's twelve +// fragment ds_reads and thirty-two v_mmac, so their latency is covered by the +// whole stage-s compute burst and the trailing WAR barrier, and the stage-s + 1 +// ds_write_b128 no longer waits on a just-issued load. The payload is four +// int4 (16 VGPRs) live across the compute; LDS size, barrier count (still two +// per K stage), staged bytes, the k0-outer/kk-inner int32 accumulation order +// and every fragment address are unchanged, so the int32 accumulation and the +// stored bits are bit-identical. The 128x64 packed/k-major arms and the +// generic scalar fallback are untouched. +// +// Iteration 9 (A-fragment direct-load round): the 128x128 packed arm reads its +// four A fragments per kk with the library du_load_matrix_sync row_major +// loader. du_mma.h defines that loader with the same +// x[i] = p[(lane & 15) * ldm + ((lane >> 4) << 3) + i] +// eight-contiguous-bytes mapping as the matrix_b col_major loader (the comment +// on load_fragment8_block below records exactly that mapping for both), but on +// this gfx928 code object the A fragments come out as pairs of ds_read2_b32 - +// two 4-byte LDS reads plus a byte reassembly chain per 8-byte fragment - and +// they sit between the LDS read and the dependent v_mmac, while the B side, +// direct since iterations 2 and 5, reads each fragment with one wide LDS +// access. Iteration 9 replaces the four library A loads with the direct +// 8-byte load_fragment8 on the A tile (row stride kAStride = 80 B), so each +// lane's x[0..7] is filled by one aligned 64-bit LDS read at +// (lane & 15) * kAStride + ((lane >> 4) << 3). The operand bytes, the +// k0-outer/kk-inner int32 accumulation order and every fragment address are +// unchanged, so the int32 accumulation stays bit-identical; only the +// redundant reassembly VALU and its lgkmcnt wait states leave the LDS->MMAC +// critical path. The A-tile layout, both barriers per K step, the staged +// bytes, the one-K-stage staging lookahead, the coalesced epilogue, the +// 128x64 packed/k-major arms and the generic scalar fallback are untouched. +// +// Iteration 10 (consolidation round, final HIP-only polish on the fixed +// 128x128 packed architecture): the arm's last measured LDS inefficiency is +// the A tile's 80-byte (20-word) row stride. The exact-source PMC of the +// accepted iteration-9 source (digest 4f7eee6f..., arch_vgpr 152, workgroup +// 256, grid 512, lds_bytes 18,432) reports lds_instructions 1,048,576 = the +// eight 8-byte fragment reads per kk (4 A + 4 B) x 32 stages x 4 waves x 512 +// blocks, and lds_bank_conflicts 3,145,728 = 3 x lds_instructions exactly. +// The B side is conflict-free by construction (fragment-major tiles: a lane's +// 8 bytes start at word 2*(lane&15), so a 16-lane phase touches all 32 banks +// once), so all 3,145,728 conflicts are A-side, i.e. 6 per A wave-instruction: +// with a 20-word stride the start word of row r is 20*r mod 32, which has +// period 8 (20*8 == 160 == 0 mod 32 == 0 mod 128 B), so lanes r and r+8 always +// land on the same bank pair and each 16-lane LDS phase runs at 2x the minimum +// (the exact aliasing the iteration-7 and iteration-9 hypotheses recorded and +// the sibling reference also left in place). A 16-byte-aligned row start can +// only ever produce 8 distinct bank residues mod 32 (starts are multiples of +// 4 words), so no padding of the existing 80-byte layout can remove it; the +// fix has to be an 8-byte-granular row stride. Focused change, 128x128 packed +// arm only: the A tile row stride becomes kAStride128 = 72 (18 words). 18 has +// gcd(18, 32) = 2, so 18*r mod 32 has period 16 and r = 0..15 is a permutation +// of the 16 even bank residues: every lane's 8-byte read occupies one distinct +// even bank and its following odd bank, i.e. all 32 banks exactly once per +// 16-lane phase (zero conflicts, the 4-cycle minimum for a 512-B fragment). +// Because 72 is only 8-byte aligned, the cooperative A staging keeps its two +// aligned int4 global loads per thread but publishes each 16-byte chunk with +// two aligned 8-byte (`int2`) LDS stores at +0/+8, so all LDS addresses stay +// naturally aligned; the whole-kernel LDS stage shrinks 18,432 -> 17,408 B +// (A 10,240 -> 9,216 B) and every staged byte keeps its exact (row, k) slot. +// The A fragment operand bytes, the 128-byte global A reads, the packed-B +// path, both barriers per K step, the one-K-stage staging lookahead, the +// coalesced epilogue, the k0-outer/kk-inner int32 accumulation order and the +// 128x64 packed/k-major arms plus the generic scalar fallback are untouched, +// so the int32 accumulation and the stored bits are bit-identical. // -// 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 12 (consolidation round, final HIP-only polish on the fixed +// 128x128 packed architecture): the last measured LDS inefficiency is the +// A-staging publish mapping. The exact-source PMC of the accepted +// iteration-10 source (source_hip_digest 9386d564..., arch_vgpr 152, workgroup +// 256, grid 512, lds_bytes 17,408) reports lds_bank_conflicts 1,048,576 with +// lds_instructions 1,048,576, and that conflict total is exactly the A publish: +// four 8-byte staging halves per thread per K stage (two int2 halves for each +// of the two payloads) x 4 LDS phases x 1 extra cycle x 32 stages x 4 waves x +// 512 blocks = 1,048,576, while the B publish (linear tid * 16), both fragment +// ds_reads and the coalesced epilogue are conflict-free. 18 has gcd(18, 32) +// = 2, so the 72-byte A row stride makes the fragment reads conflict-free, but +// the publish mapping (thread tid owns A row tid / 4 with 16-byte group +// (tid % 4) * 16) puts the four rows of one 16-lane phase at 18 r mod 32 for +// r = 0..3 - only four of the sixteen even bank residues - so every 8-byte +// staging half needs two LDS cycles instead of one; no 16-byte-aligned stride +// can remove it (such strides are multiples of four words, so a phase reaches +// at most eight residues). Focused change, 128x128 packed arm only: the A +// publish mapping becomes +// row = 16 * (tid / 64) + 2 * ((tid % 64) / 16) +// + 8 * (((tid / 4) % 4) & 1) + (((tid / 4) % 4) >> 1) +// col = 16 * (tid % 4) +// so the four rows of each 16-lane phase are (2p, 2p + 8, 2p + 1, 2p + 9) and +// their publish word starts 18 r mod 32 are exactly the sixteen even residues +// for both the +0 and the +8 half: each phase touches all 32 LDS banks once +// (one cycle, zero conflicts). Every 4-lane group still publishes one A row's +// 64 contiguous bytes, so the global A read set, its L1 line/sector pattern +// and the 128-byte row reads are identical to the previous tid-linear mapping; +// the A tile layout (72-byte rows), both barrier-separated fragment reads, the +// B staging, the one-K-stage staging lookahead, the coalesced epilogue, the +// k0-outer/kk-inner int32 accumulation order and the 128x64 arms plus the +// generic scalar fallback are untouched, so the int32 accumulation and the +// stored bits are bit-identical. +// +// Iteration 18 (consolidation round, final HIP-only polish on the fixed +// 128x128 packed architecture): the last measured fixed-region waste of the +// arm is the bf16 conversion in the fused epilogue. The exact gfx928 code +// object of the accepted iteration-12 source shows every __float2bfloat16(v) +// inside store_prefill_fragment_coalesced_scaled lowered to the Inf/NaN-guarded +// software sequence of amd_hip_bfloat16.h float_to_bfloat16: v_and +// 0x7f800000, v_cmp_ne, s_and_saveexec_b64, s_xor_b64, v_bfe_u32 (round bit), +// v_add3_u32 (0x7fff), s_andn2_saveexec_b64, v_mov_b32 0, v_or_b32 0x10000, +// v_cmp_eq_u32_sdwa, v_cndmask_b32, s_or_b64 exec - roughly a dozen +// VALU/scalar instructions and two exec-mask branches for each of the 64 +// values a wave converts (16 fragments x 4 elements), and those branches split +// the epilogue into many small basic blocks. That Inf/NaN arm is unreachable +// on this arm: its operand is an int32 dot product (|dot| <= 2048 * 127 * 127 +// < 2^25) times the finite x_scale and weight_scale values of this operator +// family, so every converted value is finite and the library's finite path is +// just the round-to-nearest-even increment u += 0x7fff + ((u >> 16) & 1) +// followed by u >> 16. Focused change, 128x128 packed arm only: the coalesced +// epilogue converts its four accumulators per fragment with a straight-line +// bf16_rne_finite() helper (four integer instructions, no s_mov_b32 +// 0x7f800000 constant, no exec mask, no branch), which stores bit-identical +// bf16 for every finite input (the 8-byte coalesced stores are untouched). The +// K loop, the LDS layouts, both barriers per K step, the one-K-stage staging +// lookahead, the in-loop scale batch, the fragment mapping, the store pattern +// and the 128x64 packed/k-major arms plus the generic scalar fallback (which +// keep the library __float2bfloat16 for their per-element stores) are +// byte-identical. +// +// The timed operator performs no allocation, compilation, autotuning, weight +// packing, host synchronization, device synchronization, or default-stream +// launch; both launchers run on the caller-provided PyTorch HIP stream and are +// CUDA/HIP-Graph safe. #include #include @@ -155,27 +279,50 @@ constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 constexpr int kWaveSize = 64; // gfx928 native wavefront constexpr int kBlockM = 128; constexpr int kBlockN = 64; +// Iteration 4 tile-aspect arm: the exact (k, n) == (2048, 2048) packed pair +// takes a 128-wide N tile (M:N tile aspect 1:1) instead of kBlockN == 64. +constexpr int kBlockN128 = 128; 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 +// Fragment-major packed B (iteration 5). One DUMMA matrix_b tile +// (m16n16k32 -> 16 n x 32 k int8) is kTileN * kTileK = 512 contiguous bytes: +// byte (n, k) of tile (n/16, k/32) sits at +// (k % 32) / 8 * 128 + (n % 16) * 8 + (k % 8). +// Every 16-byte chunk of a tile is therefore one aligned int4 LDS store / +// global read, and a lane's 8-byte fragment slot inside the tile is at +// (lane & 15) * 8 + (lane >> 4) * 128 +// i.e. word 2 * (lane & 15) plus a 32-word k-group stride, so the 16-lane LDS +// phase (lane & 15 = 0..15) covers all 32 banks exactly once: bank-conflict- +// free with zero padding. +constexpr int kBFragBytes = kTileN * kTileK; // 512 +// Legacy raw [K, N] k-major arm: 64 -> 80-byte row stride (16-byte aligned so +// the cooperative staging stays one 128-bit LDS store). +constexpr int kBStride = kBlockN + 16; +constexpr int kAStride = kStageK + 16; // 64 -> 80-byte A row stride +// Iteration 10 (consolidation round): the 128x128 packed arm uses a 72-byte +// (18-word) A row stride. 72 is 8-byte aligned (so every 8-byte fragment read +// and every split int2 staging store is naturally aligned) but NOT 16-byte +// aligned, and 18*r mod 32 for r = 0..15 is exactly the 16 even bank residues, +// so a 16-lane 8-byte fragment phase touches all 32 LDS banks once instead of +// the 2-way aliasing of the 20-word (80-byte) stride. The 128x64 arm keeps +// the accepted 80-byte stride via kAStride above. +constexpr int kAStride128 = kStageK + 8; // 64 -> 72-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; +// Exact (k, n) pair whose packed weight buffer is fragment-major for this +// worker's assigned shape (glm_tp8_q_b_proj_m4096). Every other (k, n) keeps +// the identity device-to-device pack and the raw [K, N] k-major arm. +constexpr int kPackFragMajorK = 2048; +constexpr int kPackFragMajorN = 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. +// Verified gfx928 int8 m16n16k32 accumulator ownership: lane & 15 selects the +// row, lane >> 4 selects col % 4, and frag.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 match the harness reference +// ((dot.float() * x_scale) * ws.T).to(bfloat16), so the stored bits match. // --------------------------------------------------------------------------- template __device__ __forceinline__ void store_prefill_fragment( @@ -202,28 +349,48 @@ __device__ __forceinline__ void store_prefill_fragment( } } +// --------------------------------------------------------------------------- +// Iteration 18 (consolidation round): straight-line round-to-nearest-even +// float -> bfloat16 for the coalesced epilogue. This is exactly the finite +// path of amd_hip_bfloat16.h float_to_bfloat16 +// u += 0x7fff + ((u >> 16) & 1); return u >> 16; +// with the Inf/NaN special case omitted; the 128x128 packed q_b_proj arm only +// ever converts finite values (int32 dot x finite scales), so the bf16 bits +// are identical to __float2bfloat16 while the exec-masked compare/branch pair +// and the s_mov_b32 0x7f800000 / v_or_b32 0x10000 fixups leave the epilogue. +// The per-element 128x64 arm and the generic scalar fallback keep the library +// conversion byte-identically. +// --------------------------------------------------------------------------- +__device__ __forceinline__ unsigned short bf16_rne_finite(float v) { + unsigned int u = __float_as_uint(v); + u += 0x7fffu + ((u >> 16) & 1u); // round to nearest, ties to even + return static_cast(u >> 16); +} + // --------------------------------------------------------------------------- // 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. +// epilogue for the exact (k, n) == (2048, 2048) q_b_proj 128x128 arm. The +// per-element variant above issues one 2-byte bf16 store per lane per +// element: 64 scattered stores per wave per fragment set, each store touching +// 16 rows x 4 columns, so every 32-byte store sector is only 25% utilized +// (PMC: 131,072 vmem_write_instructions = 512 blocks x 4 waves x 64). +// +// This epilogue transposes the four int32 values inside each 4-lane column +// group (lanes r, r+16, r+32, r+48 - a 4x4 register transpose, two 2x2 steps +// with __shfl_xor 16 then 32, which lowers to ds_bpermute at the block tail +// where the LDS pipe is idle), so lane (r, c4) ends up owning the four +// CONTIGUOUS columns base_col + 4*c4 .. +3. It converts those to bf16, packs +// four bf16 into one 8-byte value and writes ONE 8-byte store per lane (100% +// store sector efficiency; the previous 64 stores become 16 per lane). +// +// Only the int32 accumulator values are re-routed between lanes; the +// per-element multiply order (float(dot) * x_scale[row] * weight_scale[col]) +// and the __float2bfloat16 rounding are unchanged, so the stored bits are +// identical. xs and ws come from caller registers batched on the last K +// stage, so the epilogue issues no interleaved vmem loads. base_col is a +// multiple of 16 and n*2 is a multiple of 8, so the 8-byte store is aligned; +// the 128x128 arm is only dispatched for m % 128 == 0, so the row >= m guard +// is always false and the shuffles never mix active and inactive lanes. // --------------------------------------------------------------------------- template __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( @@ -238,9 +405,9 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( int lane) { const int row = base_row + (lane & 15); if (row >= m) { - return; // tail-M masking: padded rows never write (wavefront-uniform) + return; // tail-M masking: never fires for m % 128 == 0 } - const int c4 = lane >> 4; // 0..3 + const int c4 = lane >> 4; // 0..3 (column group inside the m16n16 tile) const int x0 = frag.x[0]; const int x1 = frag.x[1]; const int x2 = frag.x[2]; @@ -266,77 +433,98 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( 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 + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 of row `row`; 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; + // Iteration 18: straight-line RNE conversion (finite path), bit-identical + // to the library __float2bfloat16 for this arm's finite values. 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); + static_cast(bf16_rne_finite(v0)) | + (static_cast(bf16_rne_finite(v1)) << 16) | + (static_cast(bf16_rne_finite(v2)) << 32) | + (static_cast(bf16_rne_finite(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. +// Direct 8-byte staged-fragment load. du_mma.h defines both the matrix_a +// row_major and the matrix_b col_major int8 loaders as +// x[i] = p[(lane & 15) * ldm + ((lane >> 4) << 3) + i], +// i.e. eight CONTIGUOUS bytes per lane, but the generic slot-by-slot loader +// lowers to per-byte LDS byte gathers plus a mask/shift/OR reassembly chain. +// Writing those same 8 bytes with one 64-bit LDS read keeps the fragment +// storage bit-identical (the int32 accumulation is unchanged) and shortens the +// LDS -> MMAC critical path. +// +// Iteration 5 form: `block` is the base of one 512-byte fragment-major DUMMA +// matrix_b tile, whose layout puts the lane's 8 bytes at +// (lane & 15) * 8 + (lane >> 4) * 128 +// (word 2 * (lane & 15) plus a 32-word k-group stride). A 16-lane LDS phase +// therefore touches each of the 32 banks exactly once, and the access is a +// single aligned 64-bit read. Operand bytes are byte-for-byte identical to +// the n-major layout, so correctness is unaffected. // --------------------------------------------------------------------------- -__device__ __forceinline__ void load_fragment8( +__device__ __forceinline__ void load_fragment8_block( DUFragmentBase& frag, - const signed char* __restrict__ base, - int stride, + const signed char* __restrict__ block, int lane) { - const int off = (lane & 15) * stride + ((lane >> 4) << 3); + const int off = (lane & 15) * 8 + ((lane >> 4) << 7); *reinterpret_cast(&frag.x[0]) = - *reinterpret_cast(base + off); + *reinterpret_cast(block + 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. +// Iteration 9 (A-fragment direct-load round): direct 8-byte staged A-fragment +// load, the A-side twin of load_fragment8_block above. du_mma.h defines the +// matrix_a row_major int8 loader with the same +// x[i] = p[(lane & 15) * ldm + ((lane >> 4) << 3) + i] +// mapping as the matrix_b col_major loader, i.e. eight contiguous memory-order +// bytes per lane: the identical mapping the load_fragment8_block comment above +// records for both loaders. The library wrapper passes base = the fragment +// origin and ldm = the LDS row stride, so the direct read below (base + off, +// off = (lane & 15) * stride + ((lane >> 4) << 3)) addresses exactly the same +// bytes. It fills the same x[0..7] storage with one 64-bit LDS read, letting +// the compiler feed v_mmac directly instead of materializing the value through +// a pair of ds_read2_b32 plus the byte reassembly chain; the operand bytes and +// therefore the int32 accumulation are bit-identical. `base + off` is 8-byte +// aligned for the 128x128 arm: a_tile is 16-byte aligned, the row stride is +// 72 (kAStride128, iteration 10) or 80 (kAStride, the 128x64 arm), local_row +// is a multiple of 64, kk is 0 or 32 and every term is a multiple of 8. // --------------------------------------------------------------------------- -__device__ __forceinline__ void load_fragment8_plane( +__device__ __forceinline__ void load_fragment8( DUFragmentBase& frag, const signed char* __restrict__ base, + int stride, int lane) { - const int off = ((lane >> 4) << 10) + ((lane & 15) << 3); + const int off = (lane & 15) * stride + ((lane >> 4) << 3); *reinterpret_cast(&frag.x[0]) = *reinterpret_cast(base + off); } +// Byte offset of the 512-byte fragment-major DUMMA B tile for output-column +// fragment n_frag (= n / 16) and k fragment k_frag (= k / 32). Used by both +// the cooperative staging (global read) and by the scalar fallback decode. +__device__ __forceinline__ int64_t b_frag_block_offset(int n_frag, + int k_frag, + int k) { + return (static_cast(n_frag) * (k / kTileK) + k_frag) * + kBFragBytes; +} + // --------------------------------------------------------------------------- -// 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). +// Compile-time B staged-layout selector. row_major / col_major are tag types +// in du::dumma, so the layout template argument has to be selected as a type. +// kNMajorB == true (fragment-major packed B, exact (k, n) == (2048, 2048)): +// the LDS tile is a contiguous array of 512-byte DUMMA B tiles, each +// staged with one aligned int4 read/store and each read with the +// bank-conflict-free load_fragment8_block. +// kNMajorB == false (raw [K, N] row-major B): k-major tile with the accepted +// library loader (unchanged). // --------------------------------------------------------------------------- template struct b_frag_layout { @@ -348,22 +536,27 @@ struct b_frag_layout { }; // --------------------------------------------------------------------------- -// 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. +// Large-M prefill: one 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. Dispatch guarantees m % 128 == 0, n % 64 == 0, k % 64 == 0, so every +// global load / store is in-bounds and 16-byte aligned. +// +// B is templated on its staged layout: +// * kNMajorB == true (exact (k, n) == (2048, 2048)): B is the +// fragment-major packed weight; each thread stages one aligned 16-byte +// chunk of the 8 contiguous 512-byte DUMMA B tiles (4 n fragments x 2 k +// fragments) with one coalesced int4 global read and one aligned int4 LDS +// store, and each col_major fragment is one bank-conflict-free 8-byte LDS +// read per lane (load_fragment8_block). +// * kNMajorB == false: B is the raw [K, N] identity layout staged k-major +// with the library row_major loader (unchanged arm). +// +// LDS budget: A[128, 80] = 10,240 B + B (8 * 512 = 4,096 B on the packed arm, +// B[64,80] = 5,120 B on the k-major arm) <= 15,360 B/block, well inside the +// 64 KiB per-CU LDS capacity (4 blocks/CU by LDS alone). // --------------------------------------------------------------------------- template __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kernel( @@ -375,6 +568,9 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne int m, int n, int k) { + constexpr int kBBlocksB = + kNMajorB ? (kBlockN / kTileN) * (kStageK / kTileK) : kStageK; + constexpr int kBStrideB = kNMajorB ? kBFragBytes : kBStride; const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; const int lane = tid % kWaveSize; @@ -385,11 +581,8 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne 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]; + __shared__ __align__(16) int8_t b_tile[kBBlocksB * kBStrideB]; DUFragment a_frag0, a_frag1; @@ -407,21 +600,14 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne 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). + // Cooperative staging: A[128,64] is 512 int4s (two per thread, rows + // [0,64) and [64,128)); B[64,64] is 256 int4s (one per thread). Thread tid + // owns A row tid>>2 and 16-byte column group (tid & 3) * 16, and B row + // tid>>2 with 16-byte group (tid & 3) * 16 (k-major K row for the identity + // arm, packed fragment-major B tile for the packed arm). 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) = @@ -433,39 +619,51 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne 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) = + // Fragment-major packed B: thread tid stages 16-byte chunk tid of the + // 8 contiguous 512-byte DUMMA B tiles (4 n fragments x 2 k fragments): + // one aligned int4 global read and one aligned int4 LDS store at LDS + // byte offset tid * 16 (contiguous across the wave). + const int b_blk = tid >> 5; // 0..7 packed B tile index + const int b_slot = tid & 31; // 0..31 16-byte chunk inside the tile + *reinterpret_cast(b_tile + b_blk * kBFragBytes + b_slot * 16) = *reinterpret_cast( - weight + static_cast(n0 + b_k) * k + k0 + b_nc); + weight + b_frag_block_offset(n0 / kTileN + (b_blk >> 1), + k0 / kTileK + (b_blk & 1), k) + + b_slot * 16); } 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) = + // Raw [K, N] row-major B: 16 consecutive n bytes of the (k0 + b_n) row. + const int b_n = tid >> 2; // B row (K) + const int b_nc = (tid & 3) * 16; // 16-byte group inside that row + *reinterpret_cast(b_tile + b_n * kBStrideB + b_nc) = *reinterpret_cast( - weight + static_cast(k0 + b_k) * n + n0 + b_nc); + weight + static_cast(k0 + b_n) * 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. + // 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); + // Fragment-major LDS tile: the wave's two B fragments are the two + // contiguous 512-byte DUMMA B tiles (local_col/16 + i, kk/32), each + // read with one aligned bank-conflict-free 8-byte LDS read per lane. + load_fragment8_block( + b_frag0, + b_tile + ((local_col / kTileN) * 2 + (kk / kTileK)) * kBFragBytes, + lane); + load_fragment8_block( + b_frag1, + b_tile + (((local_col / kTileN) + 1) * 2 + (kk / kTileK)) * + kBFragBytes, + lane); } else { - du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, - kBStride); + du_load_matrix_sync(b_frag0, b_tile + kk * kBStrideB + local_col, + kBStrideB); du_load_matrix_sync(b_frag1, - b_tile + kk * kBStride + local_col + kTileN, - kBStride); + b_tile + kk * kBStrideB + local_col + kTileN, + kBStrideB); } du_load_matrix_sync(a_frag0, a_tile + local_row * kAStride + kk, kAStride); @@ -505,37 +703,42 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne 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); + 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); + 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. +// Iteration 4 (tile-aspect round): 128x128 macro-tile for the exact +// (k, n) == (2048, 2048) q_b_proj arm (fragment-major packed weight), four +// wavefronts of 64 lanes, each wave owning a 64x64 quadrant built from four A +// x four B m16n16k32 fragments (sixteen int32 accumulators). Widening the +// N-tile 64 -> 128 makes the M:N tile aspect 1:1: the A-side global re-read +// halves (A reuse 32 -> 16, ~268 -> ~134 MiB) while B reuse stays 32 +// (~134 MiB), so total staged tile traffic drops ~402 -> ~268 MiB (-33%), and +// each wave issues 8 LDS fragment reads per 16 MMACs instead of 6 per 8. +// +// The block stages A[128,64] (two int4/thread, rows tid>>2 and tid>>2 + 64; +// iteration 10 makes the A row stride 72 B = kAStride128, so each 16-byte +// chunk is published with two aligned int2 LDS stores) and the fragment-major +// packed B[128,64] (16 contiguous 512-byte DUMMA B tiles: two 16-byte +// chunks/thread, one aligned int4 global read and one aligned int4 LDS store +// each) into one single-buffered LDS stage: A[128,72] (9,216 B) + B +// (16 * 512 = 8,192 B) = 17,408 B/block, well inside the 64 KiB LDS budget. +// A fragments use the iteration-9 direct 8-byte load_fragment8 (the 72-byte +// stride makes each 16-lane 8-byte phase touch all 32 banks exactly once, i.e. +// zero A-side bank conflicts); B fragments use the bank-conflict-free +// fragment-major load_fragment8_block. Two __syncthreads per K stage. +// Dispatch guarantees m % 128 == 0, n % 128 == 0 and k % 64 == 0 for this +// arm, so every global load/store is in-bounds and 16-byte aligned. +// Accumulation order (k0-outer over 64-K stages, kk-inner 0 then 32) and the +// fragment element-to-slot mapping are identical to the 128x64 arm, so the +// int32 accumulation is bit-identical. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kernel( +__global__ __launch_bounds__(kBlockThreads) void +w8a8_dumma_prefill_128x128_nmajor_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, const float* __restrict__ x_scale, @@ -544,43 +747,32 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern 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 - + constexpr int kBBlocksB = + (kBlockN128 / kTileN) * (kStageK / kTileK); // 16 packed B tiles 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 wave_row = wave >> 1; // 0..1 + const int wave_col = wave & 1; // 0..1 + const int m0 = static_cast(blockIdx.y) * kBlockM; const int n0 = static_cast(blockIdx.x) * kBlockN128; - const int local_row = wave_row * (kBlockM128 / 2); // 0 or 64 + const int local_row = wave_row * (kBlockM / 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 + // 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. + // its 16 accumulator fragments; they are loaded exactly once on the last K + // stage before the final protective __syncthreads, 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]; + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride128]; + __shared__ __align__(16) int8_t b_tile[kBBlocksB * kBFragBytes]; DUFragment a_frag0, a_frag1, a_frag2, a_frag3; @@ -606,93 +798,157 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern 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). + // Cooperative staging: A[128,64] is 512 int4s (two per thread: M rows + // stage_row and stage_row + 64, 16-byte k group stage_col) and the packed + // B[128,64] is 16 contiguous 512-byte DUMMA B tiles = 512 16-byte chunks + // (two per thread: tile b_blk = tid>>5 and b_blk + 8, 16-byte slot + // tid & 31). Each chunk is one aligned int4 global read; the B chunks are + // also one aligned int4 LDS store, while the A chunks (72-byte strided rows, + // iteration 10) are two aligned int2 LDS stores at +0/+8. Iteration 12 + // remaps which lane publishes which A chunk (stage_row / stage_col below) so + // that every 16-lane publish phase covers all 32 LDS banks exactly once while + // each 4-lane group still keeps one A row's 64 contiguous bytes. + // + // Iteration 7 (compute-pipeline round): the four chunks are loaded ONE K + // STAGE AHEAD. The prologue loads K stage 0 into the payload registers; in + // the loop the payload of stage s is stored to LDS first, then (for + // k0 + kStageK < k) stage s + 1 is loaded into the same registers BEFORE the + // fragment ds_reads and the 32 v_mmac of stage s, so the global-load latency + // is covered by the stage-s compute burst plus the trailing WAR barrier + // rather than being drained by a just-issued `s_waitcnt vmcnt(N)` at the + // stage top. The payload stays four int4 (16 VGPRs); the LDS stage, both + // barriers per K step, every staged byte and every fragment address are + // unchanged. + // Iteration 12 (consolidation round): A-staging lane mapping. The four rows + // of a 16-lane LDS phase are (2p, 2p + 8, 2p + 1, 2p + 9), so the publish + // word starts 18 * row mod 32 of the phase are exactly the sixteen even bank + // residues (4p + {0, 2, 4, ..., 30}): each 8-byte staging half touches all + // 32 LDS banks once instead of twice. + const int a_phase = lane >> 4; // 0..3 16-lane LDS phase + const int a_slot = (lane >> 2) & 3; // 0..3 4-lane row group + const int stage_row = wave * 16 + 2 * a_phase + 8 * (a_slot & 1) + + (a_slot >> 1); // 0..63 + const int stage_col = (lane & 3) << 4; // 0,16,32,48 + const int b_blk = tid >> 5; // 0..7 packed B tile index + const int b_slot = tid & 31; // 0..31 16-byte chunk inside the tile + + const int64_t a_row0 = static_cast(m0 + stage_row) * k; + const int64_t a_row1 = + static_cast(m0 + stage_row + kBlockM / 2) * k; + const int64_t b_step = + 4 * (static_cast(k / kTileK) * kBFragBytes); // +4 n fragments + int4 a_payload0 = make_int4(0, 0, 0, 0); + int4 a_payload1 = make_int4(0, 0, 0, 0); + int4 b_payload0 = make_int4(0, 0, 0, 0); + int4 b_payload1 = make_int4(0, 0, 0, 0); + if (k > 0) { + a_payload0 = *reinterpret_cast(x_q + a_row0 + stage_col); + a_payload1 = *reinterpret_cast(x_q + a_row1 + stage_col); + const int64_t b_off0 = + b_frag_block_offset(n0 / kTileN + (b_blk >> 1), (b_blk & 1), k); + b_payload0 = + *reinterpret_cast(weight + b_off0 + b_slot * 16); + b_payload1 = + *reinterpret_cast(weight + b_off0 + b_step + b_slot * 16); + } + 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); - } + // Publish the payload prefetched one K stage ago (prologue for k0 == 0). + // Iteration 10: the 128x128 arm's A rows are 72-byte strided, so each + // 16-byte chunk is published with two 8-byte (`int2`) LDS stores at +0 and + // +8 (row*72 + stage_col is 8-byte aligned for every row); the staged + // bytes keep their exact (row, k) slots. + int8_t* const a_dst0 = a_tile + stage_row * kAStride128 + stage_col; + int8_t* const a_dst1 = + a_tile + (stage_row + kBlockM / 2) * kAStride128 + stage_col; + *reinterpret_cast(a_dst0) = make_int2(a_payload0.x, a_payload0.y); + *reinterpret_cast(a_dst0 + 8) = + make_int2(a_payload0.z, a_payload0.w); + *reinterpret_cast(a_dst1) = make_int2(a_payload1.x, a_payload1.y); + *reinterpret_cast(a_dst1 + 8) = + make_int2(a_payload1.z, a_payload1.w); + *reinterpret_cast(b_tile + b_blk * kBFragBytes + b_slot * 16) = + b_payload0; + *reinterpret_cast(b_tile + (b_blk + 8) * kBFragBytes + + b_slot * 16) = b_payload1; __syncthreads(); + if (k0 + kStageK < k) { + // Prefetch K stage s + 1 before consuming stage s: four aligned int4 + // global reads whose latency overlaps the fragment reads and MMACs. + const int k1 = k0 + kStageK; + a_payload0 = + *reinterpret_cast(x_q + a_row0 + k1 + stage_col); + a_payload1 = + *reinterpret_cast(x_q + a_row1 + k1 + stage_col); + // Tiles 0..7 cover n fragments 0..3 (k fragment = b_blk & 1); tiles + // 8..15 cover n fragments 4..7 with the same two k fragments. + const int64_t b_off0 = + b_frag_block_offset(n0 / kTileN + (b_blk >> 1), + k1 / kTileK + (b_blk & 1), k); + b_payload0 = + *reinterpret_cast(weight + b_off0 + b_slot * 16); + b_payload1 = + *reinterpret_cast(weight + b_off0 + b_step + + b_slot * 16); + } + // 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. + // 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) { - load_fragment8(a_frag0, a_tile + local_row * kTileStride + kk, - kTileStride, lane); + // Iteration 9 (A-fragment direct-load round): the four A fragments are + // read with the direct 8-byte load_fragment8 instead of the library + // du_load_matrix_sync row_major loader. The loader has the identical + // x[i] = p[(lane & 15) * stride + ((lane >> 4) << 3) + i] mapping + // (eight contiguous memory-order bytes per lane), but on this gfx928 + // code object it lowers to a pair of ds_read2_b32 plus a byte + // reassembly chain between the LDS read and the dependent v_mmac, + // whereas the B side (already direct since iterations 2/5) reads its + // fragment with one wide LDS access. Filling the same x[0..7] with one + // aligned 64-bit LDS read keeps the operand bytes bit-identical (the + // int32 accumulation is unchanged) and removes the reassembly VALU and + // its lgkmcnt wait states from the LDS->MMAC critical path. + // + // Iteration 10 (consolidation round): the A tile row stride is now + // kAStride128 = 72 (18 words) instead of 80, so the start word of row r + // is 18*r mod 32, a permutation of the 16 even bank residues: a 16-lane + // 8-byte fragment phase occupies each of the 32 banks exactly once and + // the A-side 2-way bank aliasing (80*8 == 640 == 0 mod 128) is gone. + load_fragment8(a_frag0, a_tile + local_row * kAStride128 + kk, + kAStride128, lane); load_fragment8(a_frag1, - a_tile + (local_row + kTileM) * kTileStride + kk, - kTileStride, lane); + a_tile + (local_row + kTileM) * kAStride128 + kk, + kAStride128, lane); load_fragment8(a_frag2, - a_tile + (local_row + 2 * kTileM) * kTileStride + kk, - kTileStride, lane); + a_tile + (local_row + 2 * kTileM) * kAStride128 + kk, + kAStride128, 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); + a_tile + (local_row + 3 * kTileM) * kAStride128 + kk, + kAStride128, lane); + // Fragment-major LDS tile: the wave's four B fragments are four + // contiguous 512-byte DUMMA B tiles (local_col/16 + i, kk/32), each read + // with one aligned bank-conflict-free 8-byte LDS read per lane. + load_fragment8_block( + b_frag0, + b_tile + ((local_col / kTileN) * 2 + (kk / kTileK)) * kBFragBytes, + lane); + load_fragment8_block( + b_frag1, + b_tile + (((local_col / kTileN) + 1) * 2 + (kk / kTileK)) * + kBFragBytes, + lane); + load_fragment8_block( + b_frag2, + b_tile + (((local_col / kTileN) + 2) * 2 + (kk / kTileK)) * + kBFragBytes, + lane); + load_fragment8_block( + b_frag3, + b_tile + (((local_col / kTileN) + 3) * 2 + (kk / kTileK)) * + kBFragBytes, + 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); @@ -711,13 +967,12 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern 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. + // Last stage only: prefetch the epilogue's per-row x_scale and per-column + // weight_scale values into registers. k0 is block-uniform, so the + // barrier below is still reached by every thread; the row < m guard keeps + // the x_scale reads in bounds (never fires for m % 128 == 0). The vmem + // latency overlaps the barrier below 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; @@ -732,16 +987,17 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern weight_scale + base_col + j * kTileN + 4 * c4); } } - // Protect the LDS buffers from the next stage's cooperative overwrite. + + // Protect the LDS buffer 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 + // 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); @@ -793,14 +1049,13 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern } // --------------------------------------------------------------------------- -// 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. +// Generic scalar fallback: one thread per output element, exact int32 dot over +// K, then the fused x_scale[row] * weight_scale[col] epilogue and one bf16 +// store. Handles every unmatched (m, n, k), including all small-M API cases +// (M=2, M=16) and any M < 128 with the same (N, K) as the tiled path. +// kNMajorPack == true decodes the fragment-major packed weight of the exact +// (k, n) == (2048, 2048) pair; otherwise the weight is the raw [K, N] +// identity layout. // --------------------------------------------------------------------------- template __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( @@ -820,23 +1075,26 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( } 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; + int32_t acc = 0; 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]. + // Fragment-major pack: decode the 512-byte DUMMA B tile offset of + // (n = col, k = kk) exactly as w8a8_pack_fragment_major_i8_kernel wrote it. + const int n_frag = col / kTileN; + const int n_rem = col % kTileN; + const int k_frags = k / kTileK; 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); + const int k_frag = kk / kTileK; + const int kk8 = (kk % kTileK) / 8; + const int kk8r = kk % 8; + const int64_t idx = + ((static_cast(n_frag) * k_frags + k_frag) * kBFragBytes) + + (kk8 * 128) + (n_rem * 8) + kk8r; + acc += static_cast(a_row[kk]) * + static_cast(weight[idx]); } } else { - // identity [K, N] row-major: column col is strided by n. + // Identity [K, N]: 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]) * @@ -849,54 +1107,63 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( } // --------------------------------------------------------------------------- -// 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. +// Iteration 5 fragment-major permutation for the exact (k, n) == (2048, 2048) +// q_b_proj weight: one DUMMA matrix_b tile (16 n x 32 k int8 = kBFragBytes = +// 512 B) is one contiguous, 16-byte-aligned block, +// packed[((n / 16) * (K / 32) + k / 32) * 512 +// + (k % 32) / 8 * 128 + (n % 16) * 8 + (k % 8)] +// == raw[k * N + n], +// i.e. a pure permutation of the same K * N bytes into the same graph-stable +// buffer. Runs once outside the timed region and outside Graph capture. The +// layout lets every 16-byte staged chunk become one aligned int4 global read +// plus one aligned int4 LDS store, and every per-lane 8-byte fragment read +// covers all 32 LDS banks exactly once per 16-lane phase. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( +__global__ __launch_bounds__(256) void w8a8_pack_fragment_major_i8_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, - int64_t count) { + int k, + int n) { const int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (linear < count) { - packed[linear] = raw[linear]; + const int64_t total = static_cast(n) * k; + if (linear >= total) { + return; } + const int k_frags = k / kTileK; // K / 32 + const int blk = static_cast(linear / kBFragBytes); + const int off = + static_cast(linear - static_cast(blk) * kBFragBytes); + const int kk8 = off >> 7; // (k % 32) / 8, 0..3 + const int rem = off & 127; // 0..127 + const int n_rem = rem >> 3; // n % 16 + const int kk8r = rem & 7; // k % 8 + const int n_frag = blk / k_frags; // n / 16 + const int k_frag = blk - n_frag * k_frags; // k / 32 + const int n_idx = (n_frag * kTileN) + n_rem; + const int k_idx = (k_frag * kTileK) + (kk8 * 8) + kk8r; + packed[linear] = raw[static_cast(k_idx) * n + n_idx]; } -// 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( +// --------------------------------------------------------------------------- +// Bootstrap identity device-to-device packing (outside the timed region and +// outside Graph capture) for every (k, n) except the exact fragment-major pair +// above. +// The packed buffers keep the same byte count and the same graph-stable +// addresses, so the GEMM consumes the raw [K, N] row-major layout directly. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, - int k, - int n) { + int64_t count) { 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]; + if (linear < count) { + packed[linear] = raw[linear]; } } -__global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( +__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( const float* __restrict__ raw, float* __restrict__ packed, int count) { @@ -910,9 +1177,9 @@ __global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( // --------------------------------------------------------------------------- // Stable host launch symbols consumed by csrc/bindings.cpp. -// Both launchers are pure dispatch: no allocation, no packing, no +// Both launchers are pure dispatch: no allocation, no packing decision, no // synchronization, no default-stream launch; they run on the caller-provided -// PyTorch stream and are CUDA-Graph safe. +// PyTorch stream and are CUDA/HIP-Graph safe. // --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, @@ -930,32 +1197,39 @@ extern "C" void launch_w8a8_gemm( (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). + // Explicit dispatch. The assigned large-prefill shape (M=4096, N=2048, + // K=2048) takes the native INT8 DUMMA tiled path; every other (m, n, k) - + // including all small-M API cases (M=2, M=16) and M < 128 - takes the + // scalar fallback. The geometry guard is exact so the paired small-M API + // shape with the same (N, K) can never enter the tiled kernel. Within the + // tiled path, only the exact (k, n) == (2048, 2048) pair - whose weight the + // packer permuted fragment-major - takes the packed arm: iteration 4 gives + // it the 128x128 tile-aspect arm when n is a multiple of 128, and the + // 128x64 arm otherwise. Every other pair keeps the byte-identical raw + // [K, N] identity 128x64 arm. + const bool packed_pair = (k == kPackFragMajorK && n == kPackFragMajorN); 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) { + if (packed_pair && n % kBlockN128 == 0) { // 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)); + // tile (M:N aspect 1:1) with the iteration-5 fragment-major packed B + // layout; grid (N/128, M/128) = 512 blocks for M=4096. n == 2048 is + // divisible by 128 by the exact guard, and the outer tiled guard already + // implies m % 128 == 0 and k % 64 == 0. + const dim3 grid128(static_cast(n / kBlockN128), + static_cast(m / kBlockM)); hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_128x128_kernel), + HIP_KERNEL_NAME(w8a8_dumma_prefill_128x128_nmajor_kernel), grid128, block, 0, stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else if (packed_pair) { + 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 { hipLaunchKernelGGL( HIP_KERNEL_NAME(w8a8_dumma_prefill_128x64_kernel), @@ -967,7 +1241,7 @@ extern "C" void launch_w8a8_gemm( 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) { + if (packed_pair) { hipLaunchKernelGGL( HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), grid, block, 0, stream, @@ -992,18 +1266,21 @@ extern "C" void launch_pack_w8a8_weight( 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). + + // Iteration 5 pack: fragment-major permutation of the exact + // (k, n) == (2048, 2048) q_b_proj weight (one 512-byte DUMMA B tile per + // (n/16, k/32) pair); identity device-to-device copy for every other + // (k, n). Same byte count as the raw weight either way, so the buffer size + // / address is stable across Graph capture and replay. 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, + if (k == kPackFragMajorK && n == kPackFragMajorN) { + hipLaunchKernelGGL(w8a8_pack_fragment_major_i8_kernel, weight_grid, block, 0, stream, raw_weight, packed_weight, k, n); } else { - hipLaunchKernelGGL(w8a8_pack_identity_kernel, + hipLaunchKernelGGL(w8a8_pack_identity_i8_kernel, weight_grid, block, 0, stream, raw_weight, packed_weight, weight_count); } @@ -1011,9 +1288,9 @@ extern "C" void launch_pack_w8a8_weight( const int64_t scale_count = n; const dim3 scale_grid( static_cast((scale_count + kBlock - 1) / kBlock)); - hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + hipLaunchKernelGGL(w8a8_pack_identity_f32_kernel, scale_grid, block, 0, stream, - weight_scale, packed_weight_scale, n); + weight_scale, packed_weight_scale, static_cast(n)); (void)hipGetLastError(); } // @@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 index fc602161..30ff5116 100644 --- 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 @@ -1,96 +1,58 @@ -// @@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 +// @@variant shape=hy3_tp4_o_proj_m16 commit=71d1281fd4cc5068bc425a5f9a49d5b8d40e90ec added=2026-09-11 +// median_us=14.73 p90_us=14.79 speedup=3.707 baseline_us=54.62 +// source=test-9-11-1-b01955b3/iteration_001/hy3_tp4_o_proj_m16 // 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) +// 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. +// Round 1 (mandate: establish_arch - first correct DUMMA geometry). The +// scalar bootstrap is replaced by a native gfx928 DUMMA INT8 m16n16k32 +// architecture for the exact assigned shape: +// * grid = N / (16 * 2) = 128 blocks, one block per pair of adjacent 16x16 +// N tiles (128 >= 120 device CUs), 128 threads = 2 x 64-lane wavefronts; +// * in-block split-K=2: wave w owns 8 interleaved K=128 stages (K=1024, +// stages w, w + 2, ...) of BOTH tiles as two independent int32 accumulator +// fragments, so the A fragments loaded per K stage are shared across the +// two tiles from registers; combined with the group-major packed-B layout +// the block's aggregate walk is one 8 KiB window per iteration and the whole +// 128-block grid tiles those windows into a single dense advancing stream; +// * B is pre-packed once (untimed, outside Graph capture) into a +// tile-pair-interleaved [k_step_group][n_tile_pair][k_step_within][lane][16] +// layout: one lane's 16 bytes hold tile0's 8 B fragment bytes followed by +// tile1's 8 B fragment bytes for the same (k_step, lane), so a single +// aligned 16-byte dwordx4 load covers both tiles and every B byte is read +// exactly once. Grouping the kStepsPerGroup = 8 k_steps one K-loop +// iteration consumes outermost makes the 128 blocks' concurrent 8 KiB +// packed-B windows tile ONE contiguous monotonic group region (1 MiB at +// n = 4096) instead of 128 windows 64 KiB apart, so the whole grid presents +// the DRAM/L2 queues a single dense advancing stream; +// * register-only K=128-stage transport with depth-1 prefetch (4 A dwordx2 +// shared by both tiles + 4 B dwordx4 = 8 loads, 6 KiB in flight per wait); +// * the only barrier is one END-of-K __syncthreads() before the fused +// partial combine + scale -> bf16 epilogue; barriers per K step = 0. // -// 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). +// The timed operator (launch_w8a8_gemm) only launches kernels on the +// caller-provided stream and touches only caller-provided buffers: no +// allocation, no compilation, no packing, no host or device synchronization, +// no default-stream fallback. In-block split-K needs no cross-block partials, +// so the caller workspace is deliberately untouched. // -// 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. +// Scalar fallback: one thread per output element for every shape other than +// the exact M=16 assignment, including the paired M=2 shape that shares the +// same (N, K); for (n, k) == (4096, 2048) it decodes the packed tile-pair +// layout elementwise, so the paired validation stays exact. +// +// pack_weight: for the exact (k, n) == (2048, 4096) o_proj weight it is the +// one-time permutation into the packed fragment-slot layout (byte count +// unchanged, 8 MiB); 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). +// hip_bfloat16.h, then du_mma.h. #include #include @@ -109,11 +71,10 @@ 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. +// Each block owns TWO adjacent 16x16 N tiles (T=2) and 2 x 64-thread +// wavefronts; wave w accumulates the int32 dot over its 8-of-16 interleaved +// K=128 stages (K=1024) 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 @@ -122,30 +83,49 @@ 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). +// Tile-pair-interleaved packed-B slot: 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 one aligned 16-byte (dwordx4) +// load per lane covers one k_step of both tiles: 64 lanes * 16 B = 1 KiB +// contiguous per load. k_steps are ordered GROUP-major: the kStepsPerGroup = 8 +// k_steps one K-loop iteration of a block consumes are adjacent across pairs, +// and the per-pair slot stride inside a group is kStepsPerGroup * +// kPackedPairBytes (see kStepsPerGroup below). constexpr int kPackedPairBytes = kDummaBlockThreads * 16; // 1024 +// Padded row stride (in int32 lanes) of one 16x16 int32 combine plane. +// du_store_matrix_sync stores the int8 m16n16k32 accumulator fragment with +// row = lane & 15 and col = (lane >> 4) + 4 * i (x[i] at +4i dwords), so with +// the natural stride 16 every lane lands on bank 16*(row & 1) + (col & 7): +// only 8 of the 32 LDS banks are touched and the store serializes (~28 extra +// bank-conflict cycles per store, 28,672 per operator replay - the only +// non-compulsory counter left in the current best). +// Stride 24 = 16 + 8 spreads the same 128 dwords exactly 4-per-bank (the +// hardware minimum) for the store, and keeps the epilogue's (row, row + 8) +// read pair 8 * 24 = 192 = 3 * 64 dwords apart, so the compiler can still emit +// one ds_read2st64 per pair and those reads also stay at the 4-phase minimum. +// Same int32 values, same row/col meaning, order-independent combine. +constexpr int kCombineStride = kDummaTileN + 8; // 24 + // 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 kStageK = 128; // K elements per stage constexpr int kStageSteps = kStageK / kDummaTileK; // 4 DUMMA steps per stage +// Packed-B group: the kStepsPerGroup = 8 k_steps (2 waves x kStageSteps) that +// one K-loop iteration of a block consumes. This is the outermost packed +// dimension, so at loop iteration s the whole grid reads one contiguous +// group_stride region: block b's 8 KiB window sits at pair-slot b inside it. +// Dense monotonic footprint instead of 128 windows 64 KiB apart. +constexpr int kStepsPerGroup = + kDummaWavesPerBlock * (kStageK / kDummaTileK); // 8 + // --------------------------------------------------------------------------- // 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). +// weight[kk, n] from the group-major +// [k_step_group][n_tile_pair][k_step_within][lane][16] tile-pair-interleaved +// layout (see the pack kernel for the exact permutation). // --------------------------------------------------------------------------- __device__ __forceinline__ int8_t packed_b_element(const int8_t* __restrict__ packed, int n, int kk) { @@ -157,18 +137,23 @@ packed_b_element(const int8_t* __restrict__ packed, int n, int kk) { 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) * + // Group-major packed layout [k_step_group][n_tile_pair][k_step_within][lane]. + const int n_pairs = kTargetN / (2 * kDummaTileN); + const int group = k_step / kStepsPerGroup; + const int within = k_step % kStepsPerGroup; + return packed[static_cast(group) * n_pairs * kStepsPerGroup * kPackedPairBytes + - lane * 16 + sub * 8 + i]; + static_cast(pair) * kStepsPerGroup * kPackedPairBytes + + static_cast(within) * 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 +// 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 group-major +// [k_step_group][n_tile_pair][k_step_within][lane][16] layout, decoded // elementwise. __global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( const int8_t* __restrict__ a, // [m, k] row-major @@ -203,55 +188,59 @@ __global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( } // --------------------------------------------------------------------------- -// 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. +// 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 +// 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). +// bytes). A stays the logical row-major 32 KiB activation (L2-hot) and +// feeds BOTH accumulator tiles. +// * B (tile-pair-interleaved, group-major layout): per k_step one aligned +// 16-byte load from the [k_step_group][n_tile_pair][k_step_within][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. +// * The caller's k0_base K-slice offset plus the k0_stride stage stride make +// each split wave walk its own packed k_step sequence; group_stride is the +// byte size of one k_step_group across all n_tile_pairs. The four k_steps +// of one stage are consecutive inside their group, so all four 16-byte +// loads are 16 B aligned (kPackedPairBytes apart) 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, + int64_t k0_base, int64_t k0_stride, int64_t group_stride, 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 k0 = k0_base + static_cast(s) * k0_stride; 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 + + // Group-major packed-B slot: k_step decides the group (outermost) and the + // position inside it; b_base is this block's n_tile_pair offset in a group. + const int64_t k_step = k0 / kDummaTileK; + const int64_t pair_off = (k_step / kStepsPerGroup) * group_stride + + (k_step % kStepsPerGroup) * 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); + 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 @@ -262,22 +251,14 @@ __device__ __forceinline__ void load_stage_fragments_t2( 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 +// Float32 -> bf16 round-to-nearest-even WITHOUT the __float2bfloat16 NaN/Inf +// branch. 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 +// 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)); @@ -293,48 +274,44 @@ __device__ __forceinline__ hip_bfloat16 bf16_rne_finite(float f) { // 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. +// 16x16 N tiles; 128 >= 120 device CUs) x 128 threads = 2 x 64-lane +// wavefronts. +// * In-block split-K=2: wave w owns 8 x K=128 stages (K=1024) and +// accumulates BOTH tiles over that slice - two independent accumulator +// fragments per wave. The two waves' stage ownership is INTERLEAVED +// (wave w owns stages w, w + 2, ...) and the packed-B layout is +// group-major, so at each K-loop iteration the two waves cover one +// contiguous 8 KiB window and the 128 blocks' windows tile one dense +// group region that advances monotonically (a single dense DRAM stream). +// 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) covering an +// 8-v_mmac burst (4 steps x 2 tiles). Every B byte is read exactly once; A +// is read once per block. +// * REGISTER-ONLY K-loop transport with depth-1 prefetch: the current +// stage's 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). +// * LDS holds only the four int32 partial planes (2 waves x 2 tiles, each +// 16 x kCombineStride = 16 x 24 int32, 6 KiB/block, padded so the fragment +// store and the epilogue read are bank-conflict-free); 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 over each +// wave's own stage sequence (the two waves own interleaved K=128 stages), +// then the two waves' partials are summed in LDS +// (order-independent int32, no overflow: max |dot| = 2048*127*127 << +// 2^31). The packed layout is coherent across the pack kernel, the DUMMA +// loads and the scalar fallback decode (same (tile, k_step, lane, byte) +// -> logical (k, n) mapping), 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. +// 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 int8_t* __restrict__ weight, // packed group-major B fragment slots const float* __restrict__ x_scale, const float* __restrict__ weight_scale, hip_bfloat16* __restrict__ out, @@ -363,28 +340,29 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( 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; + // Group-major packed-B layout: one k_step_group (the kStepsPerGroup = 8 + // k_steps a K-loop iteration consumes) holds ALL n_tile_pairs, one + // kStepsPerGroup * kPackedPairBytes slot each, so the 128 blocks' concurrent + // 8 KiB windows tile one contiguous group_stride region (1 MiB at n = 4096) + // and consecutive loop iterations advance that dense region monotonically. + // b_base is this block's pair slot inside every group. + const int64_t n_pairs = static_cast(n) / (2 * kDummaTileN); + const int64_t group_stride = + n_pairs * kStepsPerGroup * kPackedPairBytes; const int64_t b_base = - static_cast(blockIdx.x) * k_steps * kPackedPairBytes; + static_cast(blockIdx.x) * kStepsPerGroup * kPackedPairBytes; - // Per-lane A fragment ownership is loop-invariant (du_mma.hpp lines - // 447-460): row = lane & 0xf, col = (lane >> 4) << 3. + // Per-lane A fragment ownership is loop-invariant: 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. + // Prefetch the four per-thread scale values BEFORE the stage-0 fragment + // loads so the fused epilogue has zero exposed global latency on the block's + // finishing tail: the four vmem reads (x_scale[row], x_scale[row2], + // weight_scale[n0+col], weight_scale[n0+16+col] - each used twice) move 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. 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); @@ -393,31 +371,45 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( 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]; + // Four bank-conflict-free int32 partial planes: [wave][tile][16 rows x + // kCombineStride] (16 x 24 int32 = 1.5 KiB per plane, 6 KiB static LDS per + // block). The padded row stride is the only change versus the natural + // 16-column layout: the same fragment values are stored and the same 2-wave + // order-independent int32 sums are read back, so the result is bit-identical. + __shared__ __align__(16) int32_t + acc_planes[kDummaWavesPerBlock * kDummaTilesPerBlock * kDummaTileM * + kCombineStride]; // 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. + // loads are unchanged, and the end-of-K int32 combine is order-independent. + // + // Stage ownership is INTERLEAVED: wave w owns stages w, w + 2, ..., so the + // two waves of a block walk ADJACENT K=128 stages (4 KiB apart) in lockstep + // instead of two 32 KiB-apart streams. Combined with the group-major packed + // layout, at loop iteration s the block reads its 8 KiB window + // [s * group_stride, (s+1) * group_stride) at pair-slot offset, and because + // blockIdx.x tiles those slots the WHOLE GRID reads one contiguous dense + // group_stride region per iteration that advances monotonically - a single + // dense DRAM/L2 stream instead of 128 windows 64 KiB apart. Every wave still + // covers exactly 8 stages = K=1024 and the per-wave k order is still + // ascending; the bytes, the addresses (as a permutation) and the LDS combine + // are unchanged, and the int32 sums are order-independent. 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). + const int n_stages = total_stages / kDummaWavesPerBlock; // 8 stages of K=128 + const int64_t k0_base = static_cast(wave) * kStageK; + const int64_t k0_stage_stride = + static_cast(kDummaWavesPerBlock) * kStageK; + + // Prologue: the wave's first stage 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); + k0_stage_stride, group_stride, 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) @@ -429,8 +421,9 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( 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); + k0_base, k0_stage_stride, group_stride, 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 @@ -449,9 +442,9 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( 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. + // Rotate: the prefetched stage s+1 becomes the current stage. No LDS WOR + // 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; @@ -469,36 +462,41 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( } // 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. + // 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 with the + // padded kCombineStride leading dimension, which spreads the fragment lanes + // over all 32 banks instead of 8. du::dumma::du_store_matrix_sync( - acc_planes + (wave * kDummaTilesPerBlock) * kDummaTileM * kDummaTileN, - acc0, kDummaTileN, du::dumma::mem_row_major); + acc_planes + (wave * kDummaTilesPerBlock) * kDummaTileM * kCombineStride, + acc0, kCombineStride, 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); + kCombineStride, + acc1, kCombineStride, 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. + // 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, and + // bf16_rne_finite replaces the branchy __float2bfloat16 conversion. 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]; + w * kDummaTilesPerBlock * kDummaTileM * kCombineStride; + // Element (row, col) lives at row * kCombineStride + col; the second + // element of each thread is (row + 8, col), exactly 8 * kCombineStride = + // 192 dwords away (an integral ds_read2st64 offset), and tile 1's plane is + // one full 16 x kCombineStride plane further on. + const int lo = plane_base + row * kCombineStride + col; + const int hi = lo + (kDummaTileM / 2) * kCombineStride; + const int tile1 = kDummaTileM * kCombineStride; + t0_lo += acc_planes[lo]; + t0_hi += acc_planes[hi]; + t1_lo += acc_planes[lo + tile1]; + t1_hi += acc_planes[hi + tile1]; } 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); @@ -541,16 +539,19 @@ 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. +// the logical [K, N] int8 layout into the group-major +// packed[k_step_group][n_tile_pair][k_step_within][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) (du_mma.hpp matrix_b row_major lane mapping: row = lane & 0xf, +// col = (lane >> 4) << 3, x[i] = p[(col + i) * ldm + row]). k_step_group is +// k_step / kStepsPerGroup (the 8 k_steps one K-loop iteration consumes), so all +// n_tile_pairs of a group are adjacent: the grid's concurrent 8 KiB windows +// tile one dense monotonic group region. 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, @@ -590,18 +591,23 @@ __global__ __launch_bounds__(kScalarThreads) void w8a8_pack_o_proj_m16_kernel( slot.y = static_cast(v0 >> 32); slot.z = static_cast(v1 & 0xffffffffu); slot.w = static_cast(v1 >> 32); + // Group-major destination: [k_step_group][n_tile_pair][k_step_within][lane]. + const int64_t n_pairs = static_cast(n) / (2 * kDummaTileN); + const int group = k_step / kStepsPerGroup; + const int within = k_step % kStepsPerGroup; *reinterpret_cast( - packed + (static_cast(pair) * (k / kDummaTileK) + k_step) * + packed + static_cast(group) * n_pairs * kStepsPerGroup * kPackedPairBytes + - lane * 16) = slot; + static_cast(pair) * kStepsPerGroup * kPackedPairBytes + + static_cast(within) * 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: +// 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. +// writes. Returns nothing; the binding returns the caller-provided out. extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -621,12 +627,11 @@ extern "C" void launch_w8a8_gemm( 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, + // 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 + // 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)). + // M=2 shape with the same (N, K). const dim3 grid(static_cast( kTargetN / (kDummaTileN * kDummaTilesPerBlock))); const dim3 block(kDummaBlockThreads * kDummaWavesPerBlock); @@ -636,7 +641,7 @@ extern "C" void launch_w8a8_gemm( return; } - // Generic scalar fallback for every unmatched (m, n, k). This also covers + // 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( @@ -648,10 +653,11 @@ extern "C" void launch_w8a8_gemm( } // 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. +// 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, @@ -661,8 +667,9 @@ extern "C" void launch_pack_w8a8_weight( 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). + // Exact o_proj M=16 weight: permute into the group-major + // [k_step_group][n_tile_pair][k_step_within][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( 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 index 6d9acfa9..c42c939c 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=hy3_tp4_qkv_proj_m16 commit=7d6959345400cb268cfd9289338d0891db2297f9 added=2026-08-26 +// baseline_us=80.19 speedup=3.07831 // 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). 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 index 4a5f7a29..456eadd0 100644 --- 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 @@ -1,53 +1,100 @@ -// @@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). +// @@variant shape=hy3_tp4_shared_down_proj_m16 commit=49c8cf883f92c0d9a791c78c94e84604b8eafc17 added=2026-09-10 +// median_us=8.489 p90_us=8.888 speedup=2.735 baseline_us=23.22 +// source=test-9-8-8-803e5cee/iteration_001/hy3_tp4_shared_down_proj_m16 +// INT8 W8A8 GEMM implementation for Hygon K500SM_AI / gfx928. // -// Assigned shape (worker_3, physical GPU 3): -// hy3_tp4_shared_down_proj_m16 : M=16, N=4096, K=384 +// Parallel-explore child worker_3, physical GPU 3. +// Assigned shape: 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). +// Iteration 11 = HIP-only occupancy-consolidation round. The accepted +// iteration-3 kernel (median 9.092 us, one 64-lane wave per 16x16 tile, +// grid 256 = 256 wavefronts over the measured 120 CUs -> only ~2.13 resident +// waves/CU, i.e. ~0.53 waves per SIMD) is per-wave memory-latency bound: +// iteration 2 doubled resident wavefronts with in-block split-K=2 and moved +// nothing (0.148%, 512 waves = only ~1.07 waves/SIMD - still no inter-wave +// latency hiding), iteration 4 halved the grid to 128 (two N-tiles per wave +// sharing A) and regressed -25.1% (12.140 us), iteration 5 replaced the +// per-step B global waits with a block-head B-only LDS stage and regressed +// -18.97% (11.220 us), iteration 7 split the accumulator into two independent +// 6-deep v_mmac chains (regressed -2.67%, 9.341 us - the serial mma chain is +// NOT the binding constraint), and iteration 10 reordered all 12 cold-B loads +// into a one-burst register preload and regressed -38% (14.719 us - a +// front-loaded load burst starves the interleaved schedule). The fresh +// current-best ISA (arch_vgpr 24, sgpr 32, 0 scratch, 0 LDS, grid 256) shows +// each wave issues ~220 VALU + 24 global_load_dwordx2 + 12 v_mmac around 27 +// s_waitcnt and that the packed [N][K] B stream is cold per replay; with only +// ~0.5 resident wavefronts per SIMD each wave's per-step B-load vmcnt wait is +// exposed with no co-resident wave to fill it. // -// * 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. +// This round consolidates the two best-measured mechanisms into one kernel: +// the iteration-3 packed-B fast path (direct col_major 8-byte fragment loads, +// exact int32 accumulation) is launched with the iteration-2 in-block split-K +// combine (validated, bit-exact LDS int32 plane merge) raised to split-K=4 +// from the trusted occupancy-probe set {2,3,4,6} that divides the 12-step +// K=384 loop: blockDim 256 = four 64-lane wavefronts, grid 256, each wave +// accumulates one K quarter (96 = 3 m16n16k32 steps, K ascending) and one +// __syncthreads() orders the four int32 partial planes before all 256 lanes +// sum them in ascending-K order and store coalesced bf16. 256 blocks x 4 +// waves = 1024 resident wavefronts (~8.5 waves/CU = ~2.1 waves/SIMD) so each +// SIMD can overlap one wave's B-load/mma stalls with a co-resident wave's +// work, while per-wave serial K depth drops 12 -> 3 steps. HBM traffic is +// unchanged (each block's [16 n x 384 k] B pack slice is still read exactly +// once, A is read once per tile from L2, the 4 KiB partial planes stay +// on-chip), the epilogue store becomes fully coalesced row-major, and the +// int32 result is bit-identical to iteration 3 (no overflow: each partial is +// bounded by 96*127*127 << 2^31, and the quarter sums are combined in +// ascending-K wave order), so the scaled bf16 output stays bit-identical. // -// Exact int32 accumulation order (k-ascending, no split-K) is preserved, so -// the bf16 output is expected bit-identical to the scalar bootstrap. +// Iteration 3 = data-layout + epilogue round. The accepted iteration-1/2 +// one-wave and two-wave split-K=2 DUMMA kernels (median 21.656 / 21.624 us) +// both load the B fragment directly from the *logical row-major [K][N]* +// weight with du_load_matrix_sync(matrix_b, row_major). In that layout each +// lane's 8 fragment bytes are strided by N, so the compiler expands every B +// fragment load into ~8 scalar `global_load_ubyte` instructions plus +// reassembly, with a vmcnt wait burst before each v_mmac (PMC: 50,432 vmem +// reads / 27.36 us profiled; static ISA shows 16+ ubyte loads + 13-16 +// waitcnts per K step). Iteration-2 measured that doubling resident +// wavefronts via in-block split-K=2 changes nothing (0.148%) -> the binding +// constraint is per-wave load vectorization / issue cost, NOT grid +// parallelism. This round removes both remaining structural costs: // -// 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. +// * B 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 M=16 DUMMA arm then loads +// B fragments with du_load_matrix_sync(matrix_b, col_major) over P: each +// lane's 8 fragment bytes are contiguous in P -> one 8-byte load per lane +// per m16n16k32 step instead of the strided byte-load expansion. +// * The epilogue drops the LDS int32 materialization + __syncthreads + +// LDS read-back entirely and instead uses the verified gfx928 int8 +// m16n16k32 accumulator ownership (lane%16 = row, lane/16 = col%4, +// frag.x[i] = column col%4 + 4*i): every lane fuses the fp32 scale and +// bf16 conversion directly on its four register accumulators. Kernel +// static LDS becomes 0 B and there are no barriers anywhere. // -// 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. +// Geometry stays the proven iteration-1 shape: one 64-lane wavefront owns one +// 16x16 output tile, grid = N/16 = 256 blocks, explicit int32 accumulation +// over K ascending in 12 fully-unrolled m16n16k32 steps (bit-exact int8 dot +// -> bit-identical bf16 output, 0 mismatches expected), no cross-wave +// barrier. The scalar kernel below is retained as the generic fallback for +// every unmatched (m, n, k), especially the paired M=2 API shape with the +// same (N, K); it decodes the n-major pack for (k, n) == (384, 4096) via a +// b_transposed flag and reads the identity row-major [K][N] layout for every +// other pair. Shape-specialized launches stay behind exact (m, n, k) +// guards. +// +// Two stable host symbols are provided (declared in csrc/bindings.cpp): +// launch_w8a8_gemm(...) -- timed / CUDA-Graph-captured GEMM +// launch_pack_w8a8_weight(...) -- out-of-timed-region weight packing +// The pack is the [N][K] n-major transpose for the exact (k, n) == (384,4096) +// pair and an identity device copy otherwise; packing never occurs inside +// gemm_out. +// +// Graph-safety: gemm_out launches only on the caller-provided stream, touches +// only caller-provided tensors/workspace, and performs no allocation, +// compilation, autotuning, packing, or host/device synchronization. +// Known-good include order for this DTK: hip_runtime.h must precede +// hip_bfloat16.h and du_mma.h (du_mma.h is not self-contained otherwise). #include #include #include @@ -56,39 +103,164 @@ namespace { -// One block per 256 threads; 256 is a multiple of the gfx928 wavefront (64). +// gfx928 wavefront is 64 lanes; block size must be a multiple of 64. +constexpr int kWaveSize = 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 +// DUMMA INT8 hardware tile consumed by one wavefront per mma step. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; -// Exact pack guard for the assigned (K, N) pair: (384, 4096). +// Exact guard constants 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, +// gfx928 DUMMA INT8 m16n16k32 in-block split-K=4 kernel (assigned-shape fast +// path, iteration 11). +// +// acc[m, n] = sum_{kk=0}^{k-1} int32(a[m* k + kk]) * int32(b[kk*n + n0+n]) +// out[m, n0+n] = bf16( float(acc[m, n]) * x_scale[m] * weight_scale[n0+n] ) +// +// Block geometry: blockDim = 256 (four 64-lane wavefronts), grid = n/16 +// output tiles, split-K = 4 along K. Wave w accumulates the K quarter +// [w*96, (w+1)*96) in ascending kTileK steps (K=384 -> 96 per wave = 3 +// du_mma_sync steps), reading A row_major from the logical x_q and B +// col_major from the exact-shape [N][K] n-major pack exactly as the accepted +// iteration-3 one-wave kernel does at the same k offsets (one contiguous +// 8-byte global load per lane per step). All four waves then publish their +// int32 partial to LDS with du_store_matrix_sync(mem_row_major), a single +// __syncthreads() orders the planes, and every one of the 256 lanes sums the +// four partials in ascending-K wave order (partial 0 = K [0,96), ..., partial +// 3 = K [288,384) -> bit-exact full-K int8 dot, identical to the iteration-3 +// 12-step chain because no intermediate sum overflows int32) and writes its +// one scaled bf16 element with a coalesced row-major store. No barrier +// exists inside the K loop; the sole __syncthreads() is reached by every +// block thread between the partial store and the combine. +// +// Occupancy consolidation: 256 blocks x 4 waves = 1024 resident wavefronts +// (~8.5 waves/CU ~= 2.1 waves/SIMD on the measured 120 CUs) instead of the +// iteration-3 ~0.53 waves/SIMD, so a wave's per-step cold-B vmcnt waits can +// be overlapped by the co-resident wave on the same SIMD, and each wave's +// serial mma depth is 3 steps instead of 12. The 4 KiB LDS partial planes +// are on-chip (no HBM traffic change) and the kernel stays one launch, one +// barrier, graph-safe, behind the exact (m, n, k) guard. +// --------------------------------------------------------------------------- +constexpr int kSplitK = 4; // trusted occupancy-probe split that divides the + // 12-step K loop and reaches ~2 waves/SIMD +constexpr int kStepsPerWave = (kPackK / kTileK) / kSplitK; // 12 / 4 = 3 + +__global__ __launch_bounds__(kSplitK * kWaveSize) void +w8a8_dumma_m16n16k32_sk4_kernel( + const int8_t* __restrict__ a, // [16, 384] row-major int8 activations + const int8_t* __restrict__ b, // packed [4096, 384] n-major int8 weight + const float* __restrict__ x_scale, // [16] fp32 per-row scale + const float* __restrict__ weight_scale, // [4096] fp32 per-col scale + __hip_bfloat16* __restrict__ out, // [16, 4096] bf16 result + int n) { + const int tid = static_cast(threadIdx.x); // 0..255 (four 64-lane waves) + const int wave = tid / kWaveSize; // hardware wave id: 0..3 + const int n0 = static_cast(blockIdx.x) * kTileN; // this tile's columns + + __shared__ __align__(16) int32_t part[kSplitK][kTileM * kTileN]; // 4 KiB + + using namespace du::dumma; + DUFragment a_frag; + DUFragment b_frag; + DUFragment acc_frag; + 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 + k + (lane%16)*384 + (lane>>4)*8 + i], + // i.e. each lane's 8 fragment bytes are contiguous in P (8-byte aligned) -> + // one 8-byte global load per lane per step. This launch is guarded to the + // exact (m, n, k) == (16, 4096, 384) shape, so each wave's K quarter + // [wave*96, wave*96 + 96) is a whole number of kTileK steps. + const int64_t b_tile_base = static_cast(n0) * kPackK; + const int kStart = wave * kStepsPerWave * kTileK; + + // K quarter, ascending, exactly as the iteration-3 one-wave kernel at the + // same k offsets (per-wave partial is bit-identical to that chain's prefix). +#pragma unroll + for (int s = 0; s < kStepsPerWave; ++s) { + const int koff = kStart + s * kTileK; + // a + koff: 16 A rows (stride 384), columns = K slice [koff, koff + 32). + du_load_matrix_sync(a_frag, a + koff, kPackK); + du_load_matrix_sync(b_frag, b + b_tile_base + koff, kPackK); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Publish this wave's int32 partial for K quarter [kStart, kStart + 96) + // (ascending) into its own LDS plane. Row-major [16][16] plane layout with + // ldm = 16 matches du_store_matrix_sync(mem_row_major) exactly as in the + // accepted split-K=2 lineage. + du_store_matrix_sync(&part[wave][0], acc_frag, kTileN, mem_row_major); + __syncthreads(); + + // All 256 lanes combine the four partials in ascending split order (wave 0 + // = K [0, 96) ... wave 3 = K [288, 384) -> bit-exact full-K int32 dot) and + // emit one scaled bf16 element each with a coalesced row-major store: + // adjacent 16-lane groups write consecutive columns of one output row. + const int row = tid / kTileN; // 0..15 + const int col = tid - row * kTileN; // 0..15 + const int32_t acc = + part[0][tid] + part[1][tid] + part[2][tid] + part[3][tid]; + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[n0 + col]; + out[static_cast(row) * n + n0 + col] = __float2bfloat16(scaled); +} + +void launch_dumma_m16n16_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + int n, + hipStream_t stream) { + const unsigned int blocks = static_cast(n / kTileN); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16n16k32_sk4_kernel), + dim3(blocks), + dim3(kSplitK * kWaveSize), + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast<__hip_bfloat16*>(out), + n); +} + +// --------------------------------------------------------------------------- +// Scalar int8 dot-product kernel (generic fallback for every unmatched shape). +// +// acc(row,col) = sum_{kk=0}^{k-1} int32(a[row*k + kk]) * int32(b[kk*n + col]) +// out[row*n + col] = bf16( (float)acc * x_scale[row] * weight_scale[col] ) +// +// x_scale is the [M,1] fp32 per-token scale (contiguous -> M floats); +// weight_scale is the [N,1] fp32 per-output-channel scale (-> N floats). +// b is the packed buffer produced by launch_pack_w8a8_weight: for the exact +// (k, n) == (384, 4096) pair it is the [N][K] n-major transpose +// P[n*384 + k] = W[k*4096 + n] (logical column col starts at P[col*384] with +// unit stride); every other pair keeps the identity row-major [K][N] layout +// (logical column col has stride n). The b_transposed branch is hoisted out +// of the K loop by the compiler. +// +// The maximum assigned K keeps the exact int8 dot well inside int32 range; +// int32 -> float is exact here, so converting only at the very end after the +// two fp32 scale multiplies matches the reference +// (A.float() @ B.float()) * x_scale * weight_scale.T -> bfloat16. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 activations + const int8_t* __restrict__ b, // packed weight buffer (see above) + const float* __restrict__ x_scale, // [M] fp32 per-row scale + const float* __restrict__ weight_scale, // [N] fp32 per-col scale + __hip_bfloat16* __restrict__ out, // [M, N] bf16 result int m, int n, int k, @@ -97,92 +269,67 @@ __global__ __launch_bounds__(kScalarBlockThreads) void w8a8_scalar_gemm_kernel( static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const int64_t total = static_cast(m) * n; if (linear >= total) { - return; + return; // tail guard for shapes where M*N is not a block multiple } - 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 int8_t* __restrict__ a_row = a + static_cast(row) * k; + const int8_t* __restrict__ b_col = + b + (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]; + // Scale in fp32 exactly as the reference does (dot first, then x_scale, + // then weight_scale), 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); +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, + int b_transposed, + hipStream_t stream) { + const int64_t total = static_cast(m) * static_cast(n); + if (total == 0) { + return; } + const unsigned int blocks = static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_kernel), + dim3(blocks), + dim3(kScalarBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast<__hip_bfloat16*>(out), + m, + n, + k, + b_transposed); } +// --------------------------------------------------------------------------- +// Packing helpers (run out of the timed region / Graph). +// --------------------------------------------------------------------------- + // 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 @@ -199,9 +346,9 @@ w8a8_identity_copy_kernel(const T* __restrict__ src, // 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; +// 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. +// 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, @@ -209,7 +356,8 @@ w8a8_pack_b_transpose_kernel(const int8_t* __restrict__ src, 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; + const int64_t stride = + static_cast(gridDim.x) * static_cast(blockDim.x); for (int64_t t = static_cast(blockIdx.x) * blockDim.x + static_cast(threadIdx.x); t < total; t += stride) { @@ -227,12 +375,16 @@ w8a8_pack_b_transpose_kernel(const int8_t* __restrict__ src, } // 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 +// --------------------------------------------------------------------------- +// Optional out-of-timed-region weight packing entry point (registered as +// torch.ops.zth_w8a8.pack_weight through the trusted csrc/bindings.cpp). +// 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. +// copy. Asynchronous on the caller's stream: no host/device +// synchronization. +// --------------------------------------------------------------------------- extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -267,9 +419,12 @@ extern "C" void launch_pack_w8a8_weight( 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. +// --------------------------------------------------------------------------- +// Timed GEMM entry point (registered as torch.ops.zth_w8a8.gemm_out through +// the trusted csrc/bindings.cpp). Must stay Graph-safe: launches only on the +// current PyTorch stream, no allocation/sync/packing/autotuning, and uses +// only caller-provided `out` and `workspace`. +// --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -282,40 +437,33 @@ extern "C" void launch_w8a8_gemm( int n, int k, hipStream_t stream) { - // The one-wave DUMMA arm performs no split-K, so the workspace is unused. + // No inter-block split-K in this architecture (the in-block split-K=4 + // partials are combined on-chip in LDS), so the workspace is unused; it is + // still accepted (and required to be preallocated) for API compatibility. (void)workspace; (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + // 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]). + // any other (k, n) keeps the identity pack (row-major [K][N]). Decoding + // must follow the same rule in the scalar fallback below. 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. + // Explicit assigned-shape dispatch point. The M=16 / N=4096 / K=384 + // assigned shape runs on the DUMMA m16n16k32 split-K=4 path. The generic + // scalar fallback below serves every unmatched (m, n, k), including the + // paired M=2 API shape with the same (N, K). 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); + launch_dumma_m16n16_gemm(a, b, x_scale, weight_scale, out, n, stream); 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); + // Generic scalar fallback for every unmatched (m, n, k), including M=2: + // decodes the n-major pack when b_transposed, row-major otherwise. + launch_scalar_gemm(a, b, x_scale, weight_scale, out, m, n, k, b_transposed, + stream); } // @@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 index 7d26ecac..b98d2cc7 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=hy3_tp4_shared_gate_up_proj_m16 commit=f52225ed98ffd619fa1318d3eca7e56addd32f82 added=2026-08-26 +// baseline_us=73.262 speedup=6.47763 // median_us=11.31 p90_us=11.32 // source=hy3-dsh-tp4-m16-1-7f1fb1d1 // MetaInfer W8A8 INT8 GEMM for gfx928 (K500SM_AI). 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 index 40e1cdef..b5950e08 100644 --- 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 @@ -1,208 +1,170 @@ -// @@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). +// @@variant shape=hy3_tp8_o_proj_m16 commit=0b2614f457656239125b5d510d60e65275d45e24 added=2026-09-09 +// median_us=11.04 p90_us=11.16 speedup=2.91 baseline_us=32.14 +// source=hy3-dsh-tp8-m16-9-8-0161e718 +// MetaInfer W8A8 INT8 GEMM HIP implementation. // -// 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. +// Worker: worker_1 (physical GPU 1) +// Task: hy3-dsh-tp8-m16-9-8-0161e718, parallel-explore child +// Shape: hy3_tp8_o_proj_m16 M=16, N=4096, K=1024 // -// 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. +// Iteration 0 = correctness bootstrap: one scalar int8 dot-product kernel, +// one thread -> one output element (no DUMMA). Official median 78.99 us. +// +// Iteration 1 = minimal native gfx928 DUMMA m16n16k32 INT8 tile with the +// installed API: one 64-thread wavefront per block, one 16x16 N +// tile per block (grid = N/16 = 256 blocks; M = 16 fits exactly one M-tile), +// 32 strictly ascending du_mma_sync steps over K = 1024 in a single int32 +// accumulator fragment, zero barriers anywhere (single wavefront per block; +// the epilogue LDS store->read round-trip is intra-wavefront program order), +// exact fp32 scale chain (float(acc) * x_scale[m] * weight_scale[n]) and +// round-to-nearest-even bf16 store, one launch per replay on the caller +// stream, no workspace use. Official median 18.2465 us (1.76x vs the fixed +// Triton 32.137 us Graph baseline), but the grid is only 256 resident +// one-wave blocks = 0.53 wavefronts/SIMD (256 waves over 120 CUs x 4 +// SIMDs): about half of every SIMD's issue slots are empty, so the kernel +// is co-residency / load-latency bound, not compute bound. +// +// Iteration 2 (accepted official best 15.118930/15.126930 us; architecture +// mandate: one complete launch geometry among 1/2/4 waves per block x 1/2/4 adjacent N tiles): choose +// 2 wavefronts per block x 1 adjacent N tile per block = in-block +// split-K=2. Grid stays kNTiles = 256 blocks (>= 120 device CUs, ~2.13 +// blocks/CU); each 128-thread block holds two 64-thread wavefronts. Wave 0 +// accumulates the contiguous ascending K slice [0, 512), wave 1 accumulates +// [512, 1024) (16 m16n16k32 steps each, direct 8-B vectorized global +// fragment loads as in iteration 1). Each wave publishes its int32 partial +// tile to a per-block LDS plane pair (2 x 1 KiB), one END-of-K +// __syncthreads(), then wave 1 adds the planes (s=0 then s=1; int32 +// addition is exact, so the combined total is bit-identical to iteration +// 1's single full-K accumulator) and runs the unchanged scale/RNE-bf16 +// epilogue. This doubles the wavefront count to 512 (~1.07/SIMD, 4-5 +// resident waves/CU) without any extra global traffic, needs no workspace, +// and mirrors the validated N=4096 M=16 o_proj winner of the TP4 decode +// lineage (256 blocks x 2 waves, one tile/block, one END-of-K barrier). +// The exact-shape weight pack (untimed, launch_pack_w8a8_weight) switches +// from the identity copy to an n-major transpose packed[n*K + k] = +// raw[k*N + n] for exactly (k, n) == (1024, 4096) so each lane's 8-byte +// matrix_b col_major fragment is one contiguous 8-B vectorized global load +// (identity [K, N] would scatter it one byte per k row); the generic scalar +// fallback decodes that n-major layout for the exact (k, n) pair so every +// other (m, n, k) -- including the paired M=2 API shape with the same +// (N, K) -- stays byte-exact. Identity packing is retained for every other +// (K, N). +// +// Iteration 9 (this round): the accepted iteration-2 kernel (official +// 15.118930/15.126930 us) is replaced by the exact-shape reference lineage's +// validated final architecture for this geometry (hy3-dsh-tp8-m16-2-dc53295a, +// official 12.43/12.59 us on the same fixed 32.137 us Triton baseline): +// grid-level split-K=4 (1024 one-wave blocks, ~2.13 resident wavefronts/ +// SIMD), each block accumulating one contiguous ascending K quarter with A +// AND B both staged through the depth-1 double-buffered LDS window pipeline +// (one 16-B uint4 per lane per window issued one window ahead; per-step +// fragment loads read LDS at the 80-B padded stride), publishing its int32 +// tile to a workspace plane, with the split-K combine FUSED into the same +// kernel as a per-tile last-arrival tail (monotonic counters in the +// workspace tail) so the operator stays one launch per replay. In-lineage +// experiments that failed to beat iteration 2 are reverted and NOT replayed: +// iteration 5 (wave-1 tail scale prefetch, 15.2968/15.3360 us), iteration 7 +// (rolled depth-1 register stage transport, 19.2353/19.2721 us), iteration 8 +// (in-block split-K=4 occupancy probe, 15.4697/15.5025 us). +// +// Iteration 17 (this round): B-stream-locality round on the accepted +// iteration-9 kernel (official 11.320700/11.428699 us, digest 71ac0c9a, +// verified == live source before this edit; git HEAD 2dc3fd7 clean). Rounds +// 11 (publish tail surgery, 12.7361/12.7681), 14 (A register-resident, +// 12.3341/12.4781) and 15 (grid split-K 8, 14.8350/14.8526) were rejected +// and rolled back byte-for-byte, so this is a NEW bounded experiment from +// the pristine accepted source. Round-14/15 evidence leaves ONE untested +// axis of the accepted architecture: the physical layout of the once-read +// 4-MiB packed-B stream. The accepted kernel's exact-shape pack is the +// n-major transpose packed[n*K+k] = raw[k*N+n] (column stride K = 1024), so +// every one of a block's four 1-KiB window prefetches issues 16 FULL 64-B +// L2 lines strided 1 KiB apart (each block's B quarter is 16 scattered +// column chunks of 256 B over 16 KiB of address space), and the replay-wide +// 4-MiB B stream is served as ~65k single lines with no run structure at +// the DRAM row level (~377 GB/s effective read for a per-replay 4.27-MB +// footprint in 11.32 us). This round REPACKS the exact-shape buffer (same +// 4-MiB buffer, same byte count, one-time untimed pack, byte-permutation +// only) so each (split, tile) block owns one contiguous 4-KiB slot holding +// its whole B slice: slot = split*256 + tile, then window (1 KiB), then +// column (64 B), then k. Each window prefetch then reads 16 SEQUENTIAL +// 64-B lines (one contiguous 1-KiB run per block per matrix per window), +// block slots are issued in launch order (blockIdx = split*256 + tile), and +// the replay B stream decomposes into 1024 advancing 4-KiB sequential runs +// instead of 4096 scattered 1-KiB column runs -- strictly better DRAM +// row-buffer / prefetch structure for byte-identical traffic. LDS staging, +// fragment order, plane/workspace protocol, epilogue chain, launch geometry +// and every byte VALUE the kernel consumes are unchanged (the slot is +// exactly the old 16 columns x 256 k of the block in the same ascending-k +// order), so int32 accumulation is bit-identical to iteration 9; the +// generic scalar fallback and the paired M=2 path decode the same new pack +// (w8a8_pack_b_off), keeping every output byte exact. Falsifiable +// prediction: official unprofiled Graph-replay median below 11.207493 us +// (>= 1% under 11.320700 us) with p90 <= 1.05 x 11.428699 = 11.999134 us, +// 0 mismatches, Graph capture passed and the paired M=2 fallback validation +// unchanged; flat 11.2-11.6 us falsifies DRAM run-locality as the limiter +// (kernel is latency-chain bound instead) and directs the next round to +// per-wave latency reduction at unchanged geometry/traffic. +// +// Iteration 19 (this round): two-tile-per-wave A-sharing round on the +// shadow-active iteration-17 source (digest 19ec89c8, official Graph-replay +// median 11.272130/11.289730 us on the same fixed 32.137 us Triton baseline; +// git HEAD 2dc3fd7 = accepted iteration 9). Rounds 11 (publish tile_lds +// removal, 12.7361/12.7681), 14 (A register-resident with 8-B global +// fragment preloads, 12.3341/12.4781), 15 (grid split-K 8, 2048 blocks, +// 14.8350/14.8526), 17 (B block-contiguous repack, 11.2721/11.2897 -- only +// +0.43%, not accepted) and 18 (split-K 8 + A register-resident at 3588 B +// LDS, 15.0071/16.1807) bound this round. The surviving limiter reading: +// the replay wall is dominated by per-replay L2 line-REQUEST service, not by +// DRAM bytes, LDS path instructions or occupancy. Every falsified round kept +// the per-replay line-request count at 131,072 read requests (A 65,536 = 4 +// splits x 256 tile-blocks x one 4-KiB/64-line A slice each, B 65,536 +// compulsory once-read lines), and round 14 is the one negative control that +// MOVED the request count: its 8 x 8-B dwordx2 A preloads re-requested each +// A 64-B line twice (request count up ~50%) and the wall went up ~9%. +// This round halves the A-side requests by pairing the two adjacent 16-col +// tiles of each (split, pair) in ONE 64-lane wave that stages and loads its +// 4-KiB A slice ONCE and feeds two 16x16 accumulator tiles from the same +// a_lds window stream (grid 512 = 4 splits x 128 tile pairs; b_lds gains a +// second per-stage 1280-B plane for the paired tile; per-window MMA work +// doubles 2 -> 4 so the depth-1 window-prefetch latency cover per wave +// grows; per-block plane publish and arrival machinery now serve two tiles). +// Per-replay read requests drop to ~98,304 (-25%: A 32,768, B 65,536) at +// byte-identical per-(tile, split) traffic, DRAM bytes, workspace protocol, +// int32 accumulation order and every output byte. Falsifiable prediction: +// official unprofiled Graph-replay median below 11.207493 us (>= 1% under +// 11.320700 us) with p90 <= 11.999134 us, 0 mismatches, Graph capture passed +// and the paired M=2 fallback validation unchanged; point expectation +// ~9.6-10.9 us if request service is the wall (the tp4 o_proj M16 lineage's +// 512-wavefront final kernel streamed its once-read B at ~460-520 GB/s vs +// ~372-450 GB/s here at 1024 waves -- fewer waves with more per-wave work is +// the proven-faster regime on this exact N=4096 geometry); flat ~11.2-11.6 +// us or worse falsifies request service as the limiter and shows that >= ~2 +// resident waves/SIMD are required to hide the per-wave latency chain, +// directing the next round back to grid-1024 per-wave latency/tail work. +// +// Logical operator (fixed 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] ) +// with x_q [M, K] int8, packed_weight the exact-shape n-major [N, K] int8 +// transpose (identity [K, N] for every other shape), x_scale [M, 1] fp32, +// weight_scale [N, 1] fp32 and out [M, N] bf16. +// +// Host symbols (called from csrc/bindings.cpp, which owns TORCH_LIBRARY +// registration and the current-stream lookup): +// launch_w8a8_gemm(...) +// launch_pack_w8a8_weight(...) +// +// Graph-safety: the gemm launch uses only caller-provided tensors and the +// passed hipStream_t (PyTorch's current stream); it never allocates, packs, +// synchronizes, or touches the default stream. The exact-shape arm uses the +// caller workspace as int32 partial planes plus a 1-KiB monotonic +// arrival-counter tail that the launcher zeros once per workspace pointer on +// its first eager use (before Graph capture; never during capture/replay). +// pack_weight runs outside the timed region and outside Graph capture. +// +// Include order is the known-good DTK order: 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; du_mma_common.h is not a public header). #include #include @@ -212,619 +174,641 @@ namespace { -// gfx928 wavefront is 64; block dimensions must be multiples of 64. -constexpr int kScalarBlockThreads = 256; +// --------------------------------------------------------------------------- +// Geometry / tuning constants +// --------------------------------------------------------------------------- + +// Generic scalar fallback block size (multiple of the 64-lane wavefront). +constexpr int kScalarThreads = 256; + +// gfx928 wavefront / DUMMA tile geometry. constexpr int kWaveSize = 64; +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; + +// Exact assigned decode shape: hy3_tp8_o_proj_m16 (worker_1 / GPU 1). +// (K, N) = (1024, 4096) is also a valid paired M=2 API shape; the exact-shape +// guard below deliberately includes m == kExactM so those paired shapes keep +// reaching the generic scalar fallback (which decodes the exact-shape n-major +// pack there). +constexpr int kExactM = 16; +constexpr int kExactN = 4096; +constexpr int kExactK = 1024; + +// Exact-shape DUMMA launch plan (iteration 9): one 16x16 N tile per block, +// grid-level split-K = 4 along K (uniform contiguous ascending K quarters), +// one 64-thread wavefront per block, fused per-tile last-arrival combine +// (the exact-shape reference lineage's validated final architecture for this +// geometry: 12.43/12.59 us in hy3-dsh-tp8-m16-2-dc53295a). Each block stages +// its A+B slice through a depth-1 double-buffered LDS window pipeline (see +// the kernel below); the int32 partial planes and the per-tile arrival +// counters live in the caller workspace. +constexpr int kNTiles = kExactN / kTileN; // 256 tiles (>= 120 device CUs) +constexpr int kSkSplitK = 4; // grid split factor along K +constexpr int kSkKPerBlock = kExactK / kSkSplitK; // 256 k per block +constexpr int kSkTilesPerBlock = 2; // iteration 19: adjacent N tiles per wave +constexpr int kSkTPairs = kNTiles / kSkTilesPerBlock; // 128 tile pairs per split +constexpr int kSkGrid = kSkSplitK * kSkTPairs; // 512 one-wave blocks (1.07/SIMD) +constexpr int kWindowKRows = 64; // k-rows per LDS window (2 m16n16k32 steps) +constexpr int kWindowKStride = 80; // 64 k-rows + 16 pad; 16-B aligned, + // non-power-of-two bank skew (<=2-way) +constexpr int kWindowBytes = kTileN * kWindowKStride; // 1280 B per stage +constexpr int kNumLdsStages = 2; // depth-1 double buffering +constexpr int kSkWindows = kSkKPerBlock / kWindowKRows; // 4 windows per block + +// Iteration-17 B-pack slot geometry (see w8a8_pack_b_off / the pack kernel): +// the exact-shape packed buffer stores every (split, tile) block's whole B +// slice as one contiguous 4-KiB slot (4 windows x 16 columns x 64 k), so +// each window prefetch is one contiguous 1-KiB run (16 sequential 64-B +// lines) per matrix instead of 16 lines strided by the old column stride +// K = 1024. +constexpr int kBPackSliceBytes = 1 << 12; // 4096 B per slot +constexpr int kBPackWindowBytes = kTileN * kWindowKRows; // 1024 B per window +constexpr int kBPackColBytes = kWindowKRows; // 64 B per column per window +constexpr int64_t kPlaneInts = + static_cast(kExactM) * kExactN; // 65536 int32 per split plane +constexpr int64_t kPlaneBytes = kPlaneInts * sizeof(int32_t); // 256 KiB +constexpr int64_t kCounterInts = kNTiles; // one arrival counter per tile -// 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 helpers +// --------------------------------------------------------------------------- + +// float32 -> bfloat16 with round-to-nearest-even, stored as its raw uint16 +// bit pattern. This is bit-identical to __float2bfloat16() / torch's bf16 +// cast for finite inputs and keeps the epilogue independent of any one +// bfloat16 class/family: out tensors are plain 16-bit bf16 storage. +__device__ __forceinline__ uint16_t float_to_bf16_rn(const float v) { + const uint32_t bits = __float_as_uint(v); + // Add 0x7FFF plus the LSB of the kept half (ties round to even), then keep + // the high 16 bits. + const uint32_t rounded = (bits + 0x7FFFu + ((bits >> 16) & 1u)) >> 16; + return static_cast(rounded); } -__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 + +// --------------------------------------------------------------------------- +// Exact-shape block-contiguous B pack decode (iteration 17) +// --------------------------------------------------------------------------- +// Byte offset of column `col` / k-row `kk` inside the exact-shape packed +// buffer. The iteration-17 pack layout is a pure byte permutation of the +// iteration-1/9 n-major transpose packed[n*K+k] = raw[k*N+n] (same buffer, +// same byte count, same logical (col, k) byte values): +// packed[ split*256 + tile ][ window ][ col&15 ][ kk&63 ] +// with split = kk >> 8 (K quarter 0..3), window = (kk >> 6) & 3 (1-KiB +// window inside the block's 4-KiB slice), tile = col >> 4 (0..255) and +// column c = col & 15. Slot = split*256 + tile holds one block's whole B +// slice: 4 windows x 1024 B; a window holds 16 columns x 64 contiguous k. +__device__ __forceinline__ int64_t w8a8_pack_b_off(const int col, + const int kk) { + const int64_t split = static_cast(kk) >> 8; // 0..3 + const int64_t w = (static_cast(kk) >> 6) & 3; // 0..3 + const int64_t tile = static_cast(col) >> 4; // 0..255 + const int64_t c = static_cast(col) & 15; // 0..15 + return (split * kNTiles + tile) * kBPackSliceBytes + + w * kBPackWindowBytes + c * kBPackColBytes + + (static_cast(kk) & (kBPackColBytes - 1)); } // --------------------------------------------------------------------------- -// 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]). +// Generic scalar correctness kernel (fallback for every unmatched shape) // --------------------------------------------------------------------------- -__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) { + +// Scalar int8 dot-product GEMM, one thread per output element. Grid: +// blockIdx.x tiles the N dimension with kScalarThreads columns per block, +// blockIdx.y selects the M row, so adjacent lanes (threadIdx.x) read +// adjacent addresses along the fastest-varying N dimension (coalesced B +// reads) and every lane of a block broadcasts the same A row. +// For the exact (k, n) == (1024, 4096) the packed weight buffer holds the +// iteration-17 block-contiguous pack (one contiguous 4-KiB slot per +// (split, tile), see w8a8_pack_b_off), so the fallback decodes that layout +// there (this keeps the paired M=2 API shape with the same (N, K) +// byte-exact; the decode is a pure permutation of the iteration-1 n-major +// transpose it supersedes); every other (k, n) keeps the identity pack +// (raw row-major [K, N]). +__global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 + const int8_t* __restrict__ b, // [K, N] (identity) or [N, K] (exact pair) + const float* __restrict__ x_scale, // [M] fp32 + const float* __restrict__ weight_scale, // [N] fp32 + uint16_t* __restrict__ out, // [M, N] bf16 raw storage + const int m, + const int n, + const int k) { + const int col = static_cast(blockIdx.x) * kScalarThreads + + static_cast(threadIdx.x); + const int row = static_cast(blockIdx.y); + if (row >= m || col >= n) { return; } - const int row = static_cast(linear / n); - const int col = static_cast(linear - static_cast(row) * n); + // Exact k-ascending int32 accumulation. Assigned K keeps the exact int8 + // dot well inside int32 range. 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; + const bool packed_nmajor = (k == kExactK && n == kExactN); 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; + // Iteration-17 block-contiguous pack decode (w8a8_pack_b_off): the + // exact-shape buffer was repacked (same bytes, new permutation) so each + // block's B slice is one contiguous 4-KiB slot; the fallback reads the + // same logical (col, kk) byte values as the iteration-1/9 n-major + // transpose it supersedes. for (int kk = 0; kk < k; ++kk) { acc += static_cast(a_row[kk]) * - static_cast(b_pack_col[kk]); + static_cast(b[w8a8_pack_b_off(col, kk)]); } } else { + const int8_t* b_col = b + col; // advances by n every k step for (int kk = 0; kk < k; ++kk) { acc += static_cast(a_row[kk]) * - static_cast(b_col[static_cast(kk) * n]); + static_cast(*b_col); + b_col += n; } } + + // Convert to float only for the scaled product, then store bf16: + // (int32 dot) * x_scale[m] * weight_scale[n] == ((dot * xs) * ws). const float scaled = static_cast(acc) * x_scale[row] * weight_scale[col]; - out[linear] = __float2bfloat16(scaled); + out[static_cast(row) * n + col] = float_to_bf16_rn(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, +// Exact-shape DUMMA decode kernel (iteration 19: iteration-17 source plus +// two-tile-per-wave A-sharing; grid-level split-K = 4, A+B LDS-window +// staged, fused per-tile last-arrival combine) +// --------------------------------------------------------------------------- +// Architecture lineage: iteration 9 ported this geometry's validated +// reference kernel (hy3-dsh-tp8-m16-2-dc53295a, 12.43/12.59 us) -- one 16x16 +// N tile per ONE-wave block, SPLIT_K = 4 grid-level K quarters (1024 blocks, +// ~2.13 resident waves/SIMD, zero intra-pipeline barriers), A and B both +// staged through a depth-1 double-buffered LDS window pipeline, split-K +// combine FUSED into the kernel (per-tile last-arrival tail) -- and reached +// the official best 11.320700/11.428699 us; iteration 17 repacked B into +// block-contiguous 4-KiB slots (shadow 11.272130/11.289730 us, +0.43% only). +// +// Iteration 19 geometry: each of the 128 tile PAIRS per split (two adjacent +// 16-col tiles) is computed by ONE 64-lane wave, so the grid is 512 one-wave +// blocks = 4 splits x 128 pairs (1.07 waves/SIMD). The wave stages and +// loads its 4-KiB A slice ONCE per window and feeds two 16x16 accumulator +// tiles from the same a_lds window stream (a_frag shared, b_lds holds a +// second per-stage 1280-B plane for the paired tile); per-window MMA work +// doubles (2 -> 4 m16n16k32 steps) so each window's compute covers its own +// depth-1 prefetch latency longer; per-block prologue, plane-publish and +// arrival machinery now serve two tiles. This is the first lineage round +// that reduces the per-replay L2 line-request count: A-side requests halve +// (65,536 -> 32,768: one 64-line A-slice request set per 2 tiles instead of +// per tile) and total read requests drop ~25% to ~98,304 (B stays 65,536 +// compulsory once-read lines; DRAM bytes, workspace protocol, plane layout, +// counter semantics and every output byte are unchanged). +// +// Window pipeline per tile (byte-for-byte the accepted data path): tile t of +// a pair owns the contiguous ascending K slice [256*split, 256*split + 256) +// and k-rows are covered window by window ([64w, 64w + 64), w = 0..3 = two +// m16n16k32 steps per tile). Lane l prefetches ONE 16-B uint4 per matrix per +// window -- 16 contiguous k-bytes at k-offset 16*(lane>>4) of the window: +// A from the logical row-major ldm-K=1024 activation at row (lane&15), B +// from the iteration-17 block-contiguous pack slots (split*256 + 2*pair)<<12 +// and +4096 (each (split, tile) slot = one contiguous 4-KiB run; the pair's +// two slots are adjacent) at column (lane&15) -- issued one full window +// ahead of consumption, stored into +// a_lds[w&1][(lane&15)*80 + 16*(lane>>4)] and +// b_lds[w&1][t&1][(lane&15)*80 + 16*(lane>>4)] (16 x 64-k window planes at +// the 80-B padded k-stride = 64+16 bank skew). The per-step fragment loads +// (du_mma.hpp: lane l reads 8 contiguous bytes at row (l&15)*ldm + +// (l>>4)*8 + i) read LDS at ldm = 80 -- no global vmcnt in the mmac chain. +// +// Exact int32 accumulation: per-tile dots stay strictly k-ascending over the +// windowed m16n16k32 steps (max |partial| = 256*127*127 << 2^31), so each +// published plane tile is bit-identical to the iteration-9/17 per-(split, +// tile) blocks; the last-arrival tail sums the SPLIT_K = 4 planes in +// ascending split order (s = 0..3); int32 addition is exact, so every output +// byte is bit-identical to the accepted kernels (0 mismatches expected). +// +// Combine + epilogue (fused tail, counters != nullptr): the wave publishes +// BOTH tiles' int32 results to workspace planes `split` (per-tile contiguous +// layout workspace[split][tile][256], one 16-B vector store per lane per +// tile, through two API LDS planes tile_lds0/tile_lds1 with +// du_store_matrix_sync mem_row_major ldm 16), then lane 0 __threadfence()s +// (release), increments BOTH tiles' monotonic arrival counters (atomicAdd), +// __threadfence()s (acquire) and records for each tile whether the +// pre-increment value satisfies (arrived % SPLIT_K) == SPLIT_K - 1 in +// shared s_last0/s_last1. Each tile's LAST arrival of the replay then reads +// its 4 row-major int32 chunks of each plane in ascending split order and +// runs the UNCHANGED accepted scale/RNE-bf16 epilogue +// float(total) * x_scale[row] * weight_scale[n0 + col] for its 4 columns +// (row = lane>>2, cbase = 4*(lane&3)); the other three splits of each tile +// exit. Counters are monotonic (each replay adds exactly SPLIT_K per tile), +// so no per-replay reset is needed; the launcher zeroes them once per +// workspace pointer on first eager use, before Graph capture. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void w8a8_tail_emit_tile( + const int32_t* __restrict__ partials, const int tile, 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 float* __restrict__ weight_scale, uint16_t* __restrict__ out) { + // Fused per-tile last-arrival epilogue (byte-for-byte the iteration-9 tail): + // 16-B aligned reads of the per-tile-contiguous plane layout + // workspace[split][tile][256]; plane stride kPlaneInts/4 uint4s. Only the + // tile's last arrival calls this (guarded by the caller), so the planes it + // reads are all released (acquire fence in the caller) and L2-fresh. 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 + const int row = lane >> 2; // 0..15 + const int cbase = (lane & 3) * 4; // 0, 4, 8, 12 within the tile + const uint4* plane = reinterpret_cast(partials) + + static_cast(tile) * (kTileM * kTileN / 4) + + row * 4 + (lane & 3); + int32_t acc[4] = {0, 0, 0, 0}; +#pragma unroll + for (int s = 0; s < kSkSplitK; ++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]; + const int n0 = tile * kTileN; + uint16_t* out_p = + out + static_cast(row) * kExactN + n0 + cbase; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float scaled = + static_cast(acc[i]) * xs * weight_scale[n0 + cbase + i]; + out_p[i] = float_to_bf16_rn(scaled); + } +} +__global__ __launch_bounds__(kWaveSize) void w8a8_dumma_m16_n16_sk4_lds_fused_kernel( + const int8_t* __restrict__ a, // [M=16, K=1024] row-major int8 + const int8_t* __restrict__ b, // packed [N, K] n-major int8 + int32_t* __restrict__ partials, // workspace [kSkSplitK][kPlaneInts] + int32_t* __restrict__ counters, // workspace tail, 256 int32 + const float* __restrict__ x_scale, // [16] fp32 + const float* __restrict__ weight_scale, // [N] fp32 + uint16_t* __restrict__ out) { // [16, N] bf16 raw storage + const int lane = static_cast(threadIdx.x); // 0..63 + const int pair = static_cast(blockIdx.x) % kSkTPairs; // 0..127 + const int split = static_cast(blockIdx.x) / kSkTPairs; // 0..3 (K quarter) + const int tile0 = pair * kSkTilesPerBlock; // even N tile of the pair + const int tile1 = tile0 + 1; // adjacent odd N tile + const int kStart = split * kSkKPerBlock; // contiguous ascending K quarter + + // Double-buffered LDS windows (private to this single-wave block; no + // barrier inside the pipeline: store->read ordering is per-wave program + // order, enforced by the compiler's lgkmcnt waits). b_lds holds one + // 1280-B window plane per stage per tile of the pair; a_lds is SHARED by + // the two tiles (one A window load/stage feeds two 16x16 accumulators). + __shared__ __align__(16) int8_t a_lds[kNumLdsStages][kWindowBytes]; + __shared__ __align__(16) int8_t b_lds[kNumLdsStages][kSkTilesPerBlock] + [kWindowBytes]; + + 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; + du::dumma::DUFragment + b0_frag[2]; + du::dumma::DUFragment + b1_frag[2]; + 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); + + // Per-lane base pointers: A slice = activation row (lane&15). B = the + // pair's two contiguous 4-KiB slots of the iteration-17 block-contiguous + // pack at (split*256 + tile0) << 12 and +4096; inside a window (1 KiB), + // lane (lane&15) owns column c = lane&15 whose 64 k-bytes sit at c << 6, + // so one window prefetch is one contiguous 1-KiB run per tile (a pair + // covers 8 KiB of sequential pack slots per block). Byte values are + // identical to the iteration-9/17 layouts. + const int8_t* b_slot = + b + (static_cast(split) * kNTiles + tile0) * kBPackSliceBytes + + (lane & 15) * kBPackColBytes; + const int8_t* a_row = a + static_cast(lane & 15) * kExactK; + + // Prologue: prefetch window 0 for A and for each tile's B plane + // (16 contiguous k-bytes per lane per matrix per tile). uint4 a_prefetch = *reinterpret_cast( a_row + kStart + 16 * (lane >> 4)); + uint4 b0_prefetch = *reinterpret_cast( + b_slot + 16 * (lane >> 4)); + uint4 b1_prefetch = *reinterpret_cast( + b_slot + kBPackSliceBytes + 16 * (lane >> 4)); + #pragma unroll - for (int w = 0; w < numWindows; ++w) { + for (int w = 0; w < kSkWindows; ++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)); + *reinterpret_cast( + &b_lds[w & 1][0][(lane & 15) * kWindowKStride + 16 * (lane >> 4)]) = + b0_prefetch; + *reinterpret_cast( + &b_lds[w & 1][1][(lane & 15) * kWindowKStride + 16 * (lane >> 4)]) = + b1_prefetch; + if (w + 1 < kSkWindows) { a_prefetch = *reinterpret_cast( a_row + kStart + (w + 1) * kWindowKRows + 16 * (lane >> 4)); + b0_prefetch = *reinterpret_cast( + b_slot + (w + 1) * kBPackWindowBytes + 16 * (lane >> 4)); + b1_prefetch = *reinterpret_cast( + b_slot + kBPackSliceBytes + (w + 1) * kBPackWindowBytes + + 16 * (lane >> 4)); } - const int8_t* bwin = &b_lds[w & 1][0]; + // Two ascending m16n16k32 steps of this window PER TILE, fragments from + // LDS. The A fragments are loaded once and shared by both tiles' MMAs, + // so one A window global load + LDS stage serves two 16x16 tiles: the + // per-replay A-side line requests halve vs one tile per wave. Each + // tile's step order stays k-ascending (step 0 = lower k32 of the window, + // step 1 = upper k32), bit-identical to the iteration-9 per-(split, + // tile) blocks. const int8_t* awin = &a_lds[w & 1][0]; + const int8_t* bwin0 = &b_lds[w & 1][0][0]; + const int8_t* bwin1 = &b_lds[w & 1][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, + du::dumma::du_load_matrix_sync(a_frag[1], awin + kTileK, + kWindowKStride); + du::dumma::du_load_matrix_sync(b0_frag[0], bwin0, kWindowKStride); + du::dumma::du_load_matrix_sync(b0_frag[1], bwin0 + kTileK, kWindowKStride); - du::dumma::du_load_matrix_sync(b_frag[1], bwin + kDummaTileK, + du::dumma::du_load_matrix_sync(b1_frag[0], bwin1, kWindowKStride); + du::dumma::du_load_matrix_sync(b1_frag[1], bwin1 + kTileK, 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); + du::dumma::du_mma_sync(acc0_frag, a_frag[0], b0_frag[0], acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag[0], b1_frag[0], acc1_frag); + du::dumma::du_mma_sync(acc0_frag, a_frag[1], b0_frag[1], acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag[1], b1_frag[1], acc1_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, + // Publish this wave's two int32 partial tiles to workspace plane `split` + // (per-tile contiguous plane layout workspace[split][tile][256], row-major + // via the API LDS planes), one 16-B vector store per lane per tile. + __shared__ __align__(16) int32_t tile_lds0[kTileM * kTileN]; + __shared__ __align__(16) int32_t tile_lds1[kTileM * kTileN]; + const int e = lane * 4; // this lane's 4 consecutive row-major int32 + du::dumma::du_store_matrix_sync(&tile_lds0[0], acc0_frag, kTileN, 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. + *reinterpret_cast( + partials + split * kPlaneInts + + static_cast(tile0) * (kTileM * kTileN) + e) = + make_uint4(tile_lds0[e], tile_lds0[e + 1], tile_lds0[e + 2], + tile_lds0[e + 3]); + du::dumma::du_store_matrix_sync(&tile_lds1[0], acc1_frag, kTileN, + du::dumma::mem_row_major); + *reinterpret_cast( + partials + split * kPlaneInts + + static_cast(tile1) * (kTileM * kTileN) + e) = + make_uint4(tile_lds1[e], tile_lds1[e + 1], tile_lds1[e + 2], + tile_lds1[e + 3]); + + // Fused last-arrival combine tail (unchanged per-tile protocol, now for + // two tiles). Only a tile's LAST arrival of this replay (pre-increment + // counter value % SPLIT_K == SPLIT_K - 1) sums its 4 planes in ascending + // split order and emits the scaled bf16 output, so each tile's combine + // runs as soon as its splits land while sibling blocks are still + // computing and the plane reads hit L2-fresh data. if (counters != nullptr) { __syncthreads(); - __shared__ int s_is_last; + __shared__ int s_last0; + __shared__ int s_last1; 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); + __threadfence(); // release: this wave's plane stores (both tiles) are + // visible to the observers of the arrival atomics + const int arrived0 = atomicAdd(&counters[tile0], 1); + const int arrived1 = atomicAdd(&counters[tile1], 1); + __threadfence(); // acquire: the plane reads below see every sibling + // block's store released before its arrival atomic + s_last0 = ((arrived0 % kSkSplitK) == kSkSplitK - 1); + s_last1 = ((arrived1 % kSkSplitK) == kSkSplitK - 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); - } + if (s_last0) { + w8a8_tail_emit_tile(partials, tile0, x_scale, weight_scale, out); + } + if (s_last1) { + w8a8_tail_emit_tile(partials, tile1, x_scale, weight_scale, out); } } } // --------------------------------------------------------------------------- -// 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. +// Pack kernels (both run outside the timed region, in launch_pack_w8a8_weight) // --------------------------------------------------------------------------- -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; +// Identity device-to-device copy (generic fallback for every (K, N) and the +// scale copy for all shapes). +template +__global__ __launch_bounds__(kScalarThreads) void w8a8_identity_copy_kernel( + const T* __restrict__ src, T* __restrict__ dst, const int64_t numel) { + const int64_t stride = + static_cast(gridDim.x) * static_cast(blockDim.x); + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < numel; i += stride) { + dst[i] = src[i]; } +} - 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); +void launch_identity_copy_int8(const int8_t* src, + int8_t* dst, + const int64_t numel, + hipStream_t stream) { + constexpr int64_t kMaxBlocks = 65535; + int64_t blocks = (numel + kScalarThreads - 1) / kScalarThreads; + if (blocks > kMaxBlocks) { + blocks = kMaxBlocks; + } + if (blocks < 1) { + blocks = 1; } + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(static_cast(blocks)), dim3(kScalarThreads), 0, stream, + src, dst, numel); } -// --------------------------------------------------------------------------- -// 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]; +void launch_identity_copy_f32(const float* src, + float* dst, + const int64_t numel, + hipStream_t stream) { + constexpr int64_t kMaxBlocks = 65535; + int64_t blocks = (numel + kScalarThreads - 1) / kScalarThreads; + if (blocks > kMaxBlocks) { + blocks = kMaxBlocks; } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(static_cast(blocks)), dim3(kScalarThreads), 0, stream, + src, dst, numel); } -// --------------------------------------------------------------------------- -// 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( +// Exact-shape block-contiguous pack (iteration 17): for (K, N) == +// (1024, 4096) the weight is transposed one-time (outside the timed region +// / Graph) into the iteration-17 slot layout consumed by the DUMMA kernel +// and decoded by the scalar fallback (w8a8_pack_b_off): every (split, tile) +// block owns one contiguous 4-KiB slot = 4 windows x (16 columns x 64 k), +// with the k-byte run of each column inside a window contiguous so the +// kernel's 16-B per-lane window prefetches are one contiguous 1-KiB run per +// matrix per window (16 sequential 64-B lines). Same byte count and buffer +// as the identity/n-major pack, so captured addresses are unchanged; the +// mapping is a pure byte permutation of the iteration-1 n-major transpose. +// One thread per output byte; runs once during weight prep. Only correct / +// invoked for (k, n) == (1024, 4096) (the launcher guard below). +__global__ __launch_bounds__(kScalarThreads) void w8a8_pack_nmajor_bytes_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, - int k, - int n) { + const int k, + const int n) { const int64_t total = static_cast(k) * n; const int64_t stride = - static_cast(gridDim.x) * blockDim.x; + static_cast(gridDim.x) * static_cast(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]; + // idx = slot * 4096 + w * 1024 + c * 64 + kk + const int64_t slot = idx / kBPackSliceBytes; // split*256+tile (0..1023) + const int64_t rem = idx % kBPackSliceBytes; + const int64_t w = rem / kBPackWindowBytes; // 0..3 window in the slot + const int64_t c = (rem / kBPackColBytes) & 15; // 0..15 col in the tile + const int64_t kk = rem % kBPackColBytes; // k inside the 64-k window + const int k_src = + static_cast(((slot / kNTiles) << 8) | (w << 6) | kk); // 0..1023 + const int col = static_cast(((slot % kNTiles) << 4) | c); // 0..4095 + packed[idx] = raw[static_cast(k_src) * 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*); +// Generic scalar launch shared by the exact-shape-adjacent fallback path. +void launch_w8a8_gemm_scalar(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + const int m, + const int n, + const int k, + hipStream_t stream) { + const dim3 block(kScalarThreads); + const dim3 grid((static_cast(n) + kScalarThreads - 1) / + kScalarThreads, + static_cast(m)); + hipLaunchKernelGGL(w8a8_gemm_scalar_kernel, grid, block, 0, stream, a, b, + x_scale, weight_scale, + reinterpret_cast(out), m, n, k); +} } // namespace // --------------------------------------------------------------------------- -// Host launchers (stable symbols consumed by csrc/bindings.cpp). +// Host launch symbols (stable contract with 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 +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) { + // Timed operator constraints: no allocation, compilation, autotuning, + // packing, host/device synchronization, or default-stream launch. Only the + // caller-provided out and the caller-provided workspace (int32 partial + // planes + arrival-counter tail, iteration 9) may be written; the one-time + // counter hipMemsetAsync below is a stream op issued before Graph capture + // (first eager use), never during a capture/replay call. + // + // Exact assigned decode shape arm (hy3_tp8_o_proj_m16). Iteration 19 runs + // the grid-level split-K=4 kernel w8a8_dumma_m16_n16_sk4_lds_fused_kernel + // (512 one-wave blocks = 4 splits x 128 two-tile pairs, A+B LDS-window + // staged, two adjacent 16-col tiles per wave sharing one A window stream, + // fused last-arrival combine; iteration-9 reference architecture plus the + // iteration-19 A-sharing request-halving round). The guard keeps + // `m == kExactM`, so paired shapes with the same (N, K) (e.g. the M=2 API + // shape) fall through to the generic scalar fallback below (which decodes + // the exact-shape block-contiguous pack). + if (m == kExactM && n == kExactN && k == kExactK) { + // Split-K=4 fused path: 4 int32 planes (4 x 256 KiB) + 256 per-tile + // arrival counters (1 KiB) in the caller workspace. The fixed API + // validates workspaces for this (M=16, N=4096) shape at the 16-plane / + // 4-MiB capacity, so the fused single-launch path is always taken; an + // undersized workspace falls through to the generic scalar fallback and + // never runs the tuned pipeline. 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; + constexpr int64_t kFusedBytes = + static_cast(kSkSplitK) * kPlaneBytes + + kCounterInts * static_cast(sizeof(int32_t)); + if (workspace_bytes >= kFusedBytes) { + int32_t* counters = partials + kSkSplitK * kPlaneInts; + // Counters are monotonic (each replay adds exactly kSkSplitK per tile + // counter), so zeroing once per workspace pointer -- on its first + // eager use, before Graph capture -- is sufficient; the zero is never + // issued inside capture or on replay. static const void* s_fused_counters_ws = nullptr; if (s_fused_counters_ws != workspace) { - hipMemsetAsync(counters, 0, static_cast(counters_bytes), - stream); + (void)hipMemsetAsync( + counters, 0, static_cast(kCounterInts) * sizeof(int32_t), + 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); + w8a8_dumma_m16_n16_sk4_lds_fused_kernel, + dim3(static_cast(kSkGrid)), dim3(kWaveSize), 0, stream, + a, b, partials, counters, x_scale, weight_scale, + reinterpret_cast(out)); 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); + // Generic scalar fallback for every unmatched (m, n, k). + launch_w8a8_gemm_scalar(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) { +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-shape n-major pack (iteration 1): only (k, n) == (1024, 4096) gets + // the transposed fragment layout consumed by the exact-shape DUMMA kernel + // (and decoded by the scalar fallback for the same (k, n)); every other + // (K, N) keeps the byte-identical identity copy so the generic scalar + // fallback stays correct for all other shapes. 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); + constexpr int64_t kMaxPackBlocks = 65535; + int64_t weight_blocks = (weight_bytes + kScalarThreads - 1) / kScalarThreads; + if (weight_blocks > kMaxPackBlocks) { + weight_blocks = kMaxPackBlocks; + } + if (weight_blocks < 1) { + weight_blocks = 1; + } + if (k == kExactK && n == kExactN) { hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_pack_nmajor_bytes_kernel), - dim3(pack_blocks), - dim3(kPackThreads), - 0, - stream, - raw_weight, - packed_weight, - k, - n); + w8a8_pack_nmajor_bytes_kernel, + dim3(static_cast(weight_blocks)), dim3(kScalarThreads), 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); + launch_identity_copy_int8(raw_weight, packed_weight, weight_bytes, + stream); } - - 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); + launch_identity_copy_f32(weight_scale, packed_weight_scale, + static_cast(n), stream); } // @@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 index 7ff2b06e..e0c06dd2 100644 --- 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 @@ -1,757 +1,1123 @@ -// @@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. +// @@variant shape=hy3_tp8_qkv_proj_m16 commit=f2b771e350c9fea3b13d59b5fec4219786976bcc added=2026-09-09 +// median_us=19.3 p90_us=19.32 speedup=3.875 baseline_us=74.78 +// source=hy3-dsh-tp8-m16-9-8-0161e718 +// W8A8 INT8 GEMM HIP implementation 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. +// Worker: worker_0 (physical GPU 0), shape hy3_tp8_qkv_proj_m16 +// M = 16, N = 1280, K = 4096 (logical weight layout [K, N]). // -// 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 0 (bootstrap): correctness-first scalar kernel (one thread per +// output element) served every shape and measured 336 us on this shape. // -// 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): +// Iteration 1 (DUMMA bootstrap, this round): replace the scalar kernel for +// the exact guard (m, n, k) == (16, 1280, 4096) with the minimal native +// gfx928 DUMMA INT8 m16n16k32 tile per the round mandate: +// - one 64-lane wavefront per block, one 16x16 output tile per block +// (grid = N/16 = 80 blocks), no cross-wave barrier; +// - A and B fragments are loaded DIRECTLY from their row-major global +// tensors every 32-K step (no LDS staging, no split-K, no workspace, +// no combine kernel in the timed region); +// - explicit int32 accumulation: du_fill_fragment(acc, 0) then one +// du_mma_sync per ascending 32-K step (exact int32 dot, same sum as +// the scalar kernel since int32 addition is exact and order-free); +// - single-wave LDS epilogue: du_store_matrix_sync materializes the +// 16x16 int32 tile in LDS (library-defined ownership), one +// __syncthreads() reached by all 64 threads, then the 64 lanes scale +// with x_scale[row] * weight_scale[n0+col] (scalar-path float order) +// and store bf16 through the manual round-to-nearest-even conversion. +// This establishes the exact DUMMA API family below so later rounds can +// tune geometry (wavefronts/block, LDS staging, split-K planes). // -// 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). +// Iteration 2 (architecture round, this round): measure grid parallelism +// before polishing. The iteration-1 geometry (80 one-wave blocks over 120 +// CUs, 0.67 blocks/CU) is far below the two-blocks-per-CU latency-hiding +// target, and the measured 130.43 us median (vs 74.782 us Triton Graph +// baseline) is dominated by underfilled CUs and per-block serial load->mma +// chains (128 direct-load 32-K steps per block). This round replaces the +// exact-guard dispatch with a grid-level split-K=3 geometry, still one +// 64-lane wavefront per block and one 16x16 N tile per wave (zero in-block +// barriers): +// - grid = (N/16) * SPLIT_K = 80 * 3 = 240 blocks = exactly 2 blocks per +// device CU (120 CUs), split id on gridDim.y, tile id on gridDim.x; +// - the 4,096-long K axis is partitioned into three legal 32-aligned, +// non-uniform, ascending slices of {1376, 1376, 1344} rows (43/43/42 +// m16n16k32 steps), each block du_fill_fragment(acc,0) then runs its +// own ascending 32-K DUMMA chain over its slice only (bit-exact int32 +// partial dot: |partial| <= 6.6e7 << 2^31, order-free); +// - each block materializes its 16x16 int32 partial (LDS epilogue, same +// single-wave store/barrier pattern as iteration 1) into its own plane +// of the caller's int32 workspace: partials[split][16][N]; every launch +// overwrites every plane, so no workspace clear is needed; +// - a small combine kernel (80 blocks x 64 lanes, second launch on the +// same stream inside the timed API) sums the SPLIT_K planes ascending +// per output element (exact int32, identical to the scalar k-ascending +// sum) and applies x_scale[row] * weight_scale[col] then bf16 RNE, +// keeping every output bit-identical to the CPU reference. +// Split-K combine is part of the operator wall and stays inside the timed +// region/Graph. Workspace need = SPLIT_K*M*N*4 = 245,760 B, well under the +// contract's 16-plane 1,310,720 B allocation for this shape. // -// Both kernels launch on the caller's stream inside the timed Graph, so the -// combine cost is included in the measured operator latency. +// Iteration 3 (architecture round, this round): continue the mandated +// split-K occupancy sweep. Iteration 2's falsifiable claim held (240 +// one-wave blocks = 2 blocks/CU cut the wall from 130.43 to 68.65 us +// median, 1.089x over the 74.782 us Triton Graph baseline), but the +// measured logical rate (86 GB/s, 2.44 TOPS) is still far below the +// headroom of this decode family, so the sweep goes finer inside the +// trusted CU-aligned probe set {2,3,4,5,6,8,9,12}: SPLIT_K=6 -> grid = +// 80 N tiles x 6 K-slices = 480 one-wave zero-in-block-barrier blocks = +// exactly 4 blocks per device CU, and each block's serial 32-K chain +// halves to {22,22,21,21,21,21} steps = rows {704,704,672,672,672,672} +// (the existing generic slice formula; disjoint, ascending, 32-aligned, +// exact cover of [0,4096)). The partial and combine kernels are +// UNCHANGED code -- both were already generic over gridDim.y / +// split_count -- so this round only moves the launcher geometry constant +// kSplitK from 3 to 6. Workspace need = 6*M*N*4 = 491,520 B, still +// inside the caller's 16-plane 1,310,720 B allocation; the combine kernel +// (80 blocks x 64 lanes) sums the six planes ascending per output element +// (exact int32 -> bit-identical to the scalar k-ascending sum) and stays +// in the timed Graph. // -// 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 4 (architecture/pipeline round, this round): multi-N-tile +// reuse. The iteration-3 code object (PMC digest 26d6df45...) shows the +// poison direct-load pattern per 32-K step: 16 global_load_ubyte (8 per +// A/B fragment) with ~10 scattered s_waitcnt vmcnt before byte reassembly +// and one v_mmac_i32_16x16x32_i8, and no cross-iteration load overlap, so +// at 4 resident one-wave blocks/CU every block's serial chain is one +// ~load-latency step per m16n16k32 mma (52.19 us median; the wall is the +// per-step latency of a one-mma step, not raw issue or HBM traffic). +// This round keeps the same split-K partial + combine architecture and +// the same 4 blocks/CU co-residency but widens each one-wave block to TWO +// adjacent 16-wide N tiles (32 columns): per 32-K step the wave loads the +// A fragment ONCE and issues two du_mma_sync (independent accumulators), +// so the x_q A bytes [16 x 32] of the shared step are reused across the +// two output tiles (per-output A global loads/reassembly halve) while +// each B byte is still read exactly once (no B reuse). To keep 480 +// one-wave blocks = exactly 4/CU over 120 CUs the grid becomes 40 N +// groups x SPLIT_K=12 K-slices ({11x8, 10x4} 32-K steps = rows +// {352x8, 320x4}, disjoint ascending exact cover of [0,4096), every +// boundary a DUMMA-k multiple; existing generic slice formula). Per-block +// serial chain drops to <= 11 steps of TWO mmas, so if the iteration-3 +// wall is per-step load->mma latency the per-output step latency halves +// and the unprofiled Graph median/P90 must drop materially below +// 52.194/52.269 us. Workspace need = 12*M*N*4 = 983,040 B, still inside +// the caller's 16-plane 1,310,720 B allocation; the combine kernel (80 +// blocks x 64 lanes) is UNCHANGED code and sums the twelve planes +// ascending per output element (exact int32 -> bit-identical to the +// scalar k-ascending sum), staying in the timed Graph. The only edits +// are the widened per-wave N tile in the partial kernel plus the launcher +// geometry; the workspace-less fallback keeps the iteration-1 one-wave +// direct kernel byte-for-byte and the scalar fallback still serves every +// unmatched shape including the paired M=2 API shape. // -// 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): +// Iteration 5 (packed-weight/staging round, this round): B-only LDS slice +// staging. The iteration-4 code object (PMC digest 1a4e30d0..., vmem_read +// 122,880 = 24 byte-granular global_load_ubyte per 32-K step x 5120 +// block-steps) shows the remaining wall is the poison direct-global B half +// of the per-step load->wait->reassembly->mma chain: per step the wave +// still pays 16 global byte loads (B0/B1) plus ~10 scattered s_waitcnt +// vmcnt on the MMA critical path even though every byte is L2-hot in the +// timed Graph replay. Per the lineage rule, direct-global fragment loads +// are poison: this round removes the B half by staging each block's whole +// 32-column B slice (<= 352 rows x 32 B = 11,264 B) into LDS ONCE at block +// start with 16-B int4 cooperative loads (64 lanes x 16 B per wavefront +// instruction, 4-deep source batching so one vmcnt drain covers 4 +// in-flight round trips), one same-wave __syncthreads(), then every +// 32-K step reads the B0/B1 fragments from LDS with du_load_matrix_sync +// (~30-cycle LDS latency, no global round trip) at row stride 32 -- the +// staged tile is an exact byte mirror of the block's global B strip, so +// fragment contents, the k-ascending int32 order and the partial planes +// are bit-identical to iteration 4. A stays direct-global on purpose: +// this round isolates the B-side hypothesis (A-only staging and the +// packed [N,K] layout remain untested axes). Geometry is UNCHANGED (40 +// N groups x SPLIT_K=12 = 480 one-wave blocks = exactly 4 blocks/CU over +// 120 CUs; per-block LDS 11,264 B + 2 x 1,024 B acc tiles = 13,312 B, +// 65,536/13,312 = 4.9, so the 4/CU co-residency is preserved and LDS does +// not become the occupancy limiter). If the per-step wall is the B-load +// half of the global wait ladder, the official unprofiled Graph median/P90 +// must drop materially below 41.154/41.233 us (partial-kernel profiled +// 36.32 us); if instead the median stays flat or regresses, the wall at 4 +// waves/CU is A-side latency or per-CU issue (16 vs 24 loads immaterial), +// and the next round must stage A too (LDS-only K loop) or raise +// co-residency -- never a scalar-fallback change. Exactness is preserved +// for every shape: the (m,n,k)==(16,1280,4096) guard is unchanged; the +// staged bytes are a copy of the same global bytes the direct loader +// read, the workspace-less fallback keeps the iteration-1 direct kernel +// byte-for-byte (still correct for this shape), the scalar fallback still +// serves every unmatched shape including the paired M=2 API shape, and no +// pack/layout/env/stream/config change is made; all launches stay on the +// caller stream inside the timed API/Graph. // -// 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. +// Iteration 6 (infrastructure failure, no valid artifact): the round-6 +// attempt was killed by the agent infrastructure before it produced a +// candidate result; no round-6 source change was accepted and the session +// was reset to the iteration-5 accepted best source below. The failure does +// not count as a completed optimization round. // -// 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 7 (staging round, this round): A+B whole-slice LDS staging +// (LDS-only K loop). The iteration-5 result (39.404/39.485 us official +// Graph median/P90 vs iteration-4 41.154/41.233) falsified its own claim: +// the predicted 25-33 us for the B-half hypothesis did not appear, so per +// the iteration-5 contingency text this round stages the shared A operand +// too. The iteration-5 PMC (source digest 5fb42aff..., grid 480 = 40 N +// groups x SPLIT_K=12, vmem_read 46,080, l2_misses 98,879, counter-derived +// HBM ~187.9 GB/s << peak) still shows 8 byte-granular global_load_ubyte of +// the shared A fragment plus ~8 scattered s_waitcnt vmcnt before shift/or +// reassembly on the MMA path every 32-K step (~40,960 of the 46,080 +// vmem_read over 5,120 block-steps): pure per-step A-side global latency +// serialization, not traffic. This round stages BOTH operands once per +// block with 16-B int4 cooperative loads (A slice: 16 m-rows x 256 B + 16 B +// pad/row = 4,352 B at row stride 272 for bank skew; B strip: 256 x 32 = +// 8,192 B), one same-wave __syncthreads(), then an LDS-only K loop (A +// fragment at row stride 272, B0/B1 at row stride 32 via +// du_load_matrix_sync) with zero per-step global loads/waits. Staged tiles +// are exact byte mirrors of the global operands, so fragment bytes, the +// k-ascending int32 order and the partial planes are bit-identical to +// iteration 5. Geometry forced by LDS capacity: a whole-slice A stage does +// not fit at SPLIT_K=12 with 4 resident blocks/CU (A 16 x 352 = 5,632 B + +// B 11,264 B + acc 2,048 B = 18,944 B -> floor(65,536/18,944) = 3), so the +// guard launch moves to SPLIT_K=16: sixteen UNIFORM 8-step slices (128 == +// 16 x 8, rows {256 x 16}; every boundary a DUMMA-k multiple), grid = 40 N +// groups x 16 = 640 one-wave blocks ~ 5.33/CU over 120 CUs, per-block LDS +// 4,352 + 8,192 + 2,048 = 14,592 B -> floor(65,536/14,592) = 4 resident +// blocks/CU, i.e. the iteration-5 co-residency is preserved. Workspace +// need = 16 int32 planes = 1,310,720 B = the caller's full fixed 16-plane +// allocation for this shape, still gated by the launcher's workspace_bytes +// >= kWorkspaceNeeded check (undersized workspaces keep the workspace-less +// fallback). The combine kernel (80 blocks x 64 lanes, second same-stream +// launch in the timed API/Graph) is UNCHANGED code, generic over +// split_count = 16, and sums the sixteen planes ascending per output +// element (exact int32 -> bit-identical to the scalar k-ascending sum), +// staying in the timed Graph. All fallbacks, guards, pack/layout/env/ +// stream/config behavior are unchanged. // -// 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 8 (occupancy/resource round, this round): raise the waves-per- +// block occupancy limiter from 1 to 2 wavefronts per block (blockDim 64 -> +// 128). PMC of the accepted iteration-7 code object (source digest +// b64e3209..., grid 640, workgroup 64, arch_vgpr 40, sgpr 48, scratch 0, +// LDS/block 14,592 B -> floor(65,536/14,592) = 4 resident blocks/CU = +// 5.33 waves/CU by grid) shows the residual wall is per-SIMD single-wave +// latency exposure, not issue or HBM traffic: each resident wave is alone +// on its SIMD (4 blocks x 1 wave -> 4 SIMDs), so every per-step +// s_waitcnt lgkmcnt after the staged-fragment du_load_matrix_sync and every +// du_mma accumulator dependency stalls that SIMD with no second independent +// wave to issue (VALU issue utilization ~3-6% of the 26.56 us profiled +// kernel; only 12 one-time staging global loads per block, vmem_read +// 7,680). The fix keeps the iteration-7 geometry byte-for-byte in every +// respect that affects partials (SPLIT_K=16, 40 x 16 = 640 blocks, 8-step +// slices rows {256 x 16}, staged exact byte mirrors, k-ascending int32 +// order, 16-plane workspace) and splits the one-time staging ACROSS the two +// waves (wave 0 copies the A slice, wave 1 copies the B strip; exactly the +// same global bytes read once -- no repeated HBM reads traded for +// occupancy), then each wave runs its OWN 8-step accumulator chain for one +// n16 tile of the block's 32 columns (one du_mma per step per wave; total +// MMA count unchanged). LDS per block stays 14,592 B -> still 4 resident +// blocks/CU, but now 8 resident waves/CU = 2 per SIMD, so a stalling wave +// is covered by its co-resident sibling. The only deliberately added cost +// is LDS re-reads: the shared A fragment is now loaded by both waves +// (predicted lds_instructions 99,840 -> ~125-135k with conflicts scaling at +// the measured ~2.56/instr); LDS is not the wall at ~3-6% issue +// utilization. The combine kernel, guards, workspace-less fallback, scalar +// fallback and pack identity are all unchanged. // -// 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. +// Iteration 9 (combine restructure round, this round): attack the OTHER half +// of the operator wall. PMC of the accepted iteration-8 code object (source +// digest beb25b76..., official Graph median/P90 28.884/28.933 us) profiles +// BOTH timed dispatches of the operator replay: the 640-block partial kernel +// at 19.361 us AND the 80-block combine kernel at 16.960 us -- the combine is +// 46.7% of the 36.321 us profiled operator aggregate, not a small tail. Its +// structure explains the cost: 80 one-wave 64-lane blocks (<= 80 of 480 +// SIMDs busy, ~40 CUs fully idle) whose lanes each run FOUR output elements +// through a DEPENDENT per-plane chain -- one 4-B global load + s_waitcnt +// vmcnt + int32 add per plane, 16 serial DRAM round trips per element (the +// exact code object's static body is one load/wait/add per plane iteration) +// -- against plane data that is L2-cold at combine start (combine +// l2_hit_rate 6.9%, l2_misses 21,596 ~= the whole 1.31 MB of partial planes, +// evicted by the partial kernel's 5.37 MB streaming read): DRAM-latency +// serialization on only 1.34 MB of reads, plus underfill. This round +// restructures ONLY the combine kernel (the partial kernel, SPLIT_K=16 +// 640-block geometry, staged byte mirrors, k-ascending int32 partials, +// 16-plane 1,310,720 B workspace contract, workspace-less fallback, exact +// guard and scalar fallback are byte-for-byte unchanged): the guard launcher +// switches the combine grid from 80 x 1 one-wave blocks to a 2-D grid of 80 +// n-tiles x 4 four-row m-groups = 320 one-wave blocks, and the kernel's +// split_count == 16 specialization gives every lane exactly ONE output +// element that issues ALL sixteen plane loads as INDEPENDENT reads into a +// 16-int register array (one vmcnt drain, 16-way memory-level parallelism +// per lane), then sums the registers in ascending split order s = 0..15 -- +// the identical int32 sequence as the old ascending loop -- and applies the +// unchanged float scale order and f32_to_bf16_rne, so every output element +// stays bit-identical (this mirrors the validated gate_up S=16 combine +// recipe: issue all 16 plane loads into a register array before the +// ascending sum). Same stream, same two launches, no allocation or sync +// inside the timed API/Graph. // -// 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). +// Iterations 10-15 (five-wave occupancy arc on this guard, all archived at +// iterations/hy3_tp8_qkv_proj_m16/iterationN): iteration 10 (direct-register +// epilogue on the two-wave mapping: acc tiles + epilogue barrier removed, +// LDS 14,592 -> 12,544 B = 5 resident blocks/CU = 10 waves/CU) measured +// 20.827/21.585 us and was NOT accepted; iteration 11 raised waves-per-block +// 2 -> 5 (blockDim 320, grid 16 x 16 = 256 blocks, LDS 29,952 B = 2 resident +// blocks/CU = 10 waves/CU) but FAILED exact correctness (its source declared +// no __launch_bounds__ while the launcher requests 320 threads, so the +// compiled descriptor capped the workgroup at 256 and the fifth wavefront's +// tile was absent: 4,091/20,480 mismatches, first at (m,n)=(0,64)); the +// iteration-11 repair was completed in iteration 13 by declaring +// __launch_bounds__(320) (exact correctness passed, 0 mismatches; measured +// 23.764/24.086 us); iteration 14 (ping/pong fragment double-buffer) passed +// correctness but the toolchain spilled the second fragment set to the +// private segment on every K step (.private_segment_fixed_size 48) -> 39.472 +// us, closing the register-pipelining axis; iteration 15 replaced the +// five-wave epilogue __syncthreads() with __threadfence_block() (the acc +// tiles are wave-private and the waves' plane stores are disjoint, so the +// cross-wave join was pure skew cost) and measured 20.613/20.641 us, correct +// but still not below the accepted two-wave 20.179/20.216. // -// 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). +// Iteration 16 (this round): close the five-wave occupancy axis with the +// measured occupancy evidence and apply the repair arc's validated epilogue +// discovery to the accepted mapping. Occupancy ledger of the partial +// kernel, all correct and bit-identical: 8 resident waves/CU (4 x 2-wave +// blocks, LDS 14,592 B) = 20.179/20.216 us (iteration 9, accepted); the two +// independent 10 resident-waves/CU configurations -- iteration 10 (5 x +// 2-wave blocks) 20.827/21.585 and iteration 15 (2 x 5-wave blocks) +// 20.613/20.641 -- BOTH trail the 8-wave configuration, i.e. per-SIMD wave +// density 2.5 (some SIMDs carry 3 waves while others idle) is worse than a +// clean 2.0 on this kernel, and the five-wave mapping (2 or 1 resident +// 5-wave blocks/CU = 10 or 5 waves/CU only) can never reach the 8-wave +// sweet spot. The accepted iteration-9 source below still carries the SAME +// provably-unnecessary epilogue __syncthreads() cross-wave join whose +// five-wave form cost ~3.15 us (iteration 13 23.764 -> iteration 15 20.613 +// with the join removed). Fix (the ONLY code change of this round; the +// two-wave geometry, blockDim 128, grid 40 x 16 = 640 blocks, staging +// partition, stage byte mirrors, 8-step K loop, combine kernel and its +// 320-block launch, workspace, guards and fallbacks are byte-for-byte the +// accepted iteration-9 candidate, verified by diff): replace the partial +// kernel's epilogue __syncthreads() with __threadfence_block(), exactly as +// iteration 15 validated on the five-wave form of this same kernel. +// acc_tile0/acc_tile1 are each written by du_store_matrix_sync and read +// back only by the SAME wave (w == 0 -> acc_tile0, w == 1 -> acc_tile1), and +// the two waves' plane stores touch disjoint 16-column ranges, so no +// epilogue stage has any cross-wave dependency; __threadfence_block() still +// orders each wave's own outstanding LDS writes before its read-back +// ds_reads (wavefront-wide lgkmcnt drain) while costing zero cross-wave +// wait. The staging __syncthreads() above stays (wave 0 stages the A slice +// that wave 1 also reads). Every partial byte and combine output byte is +// unchanged, so the output is bit-identical to iteration 9. // -// 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. +// Iteration 18 (this round): compile-time 8-step K-loop specialization of +// the partial kernel, applied to the current checked-in source (the +// iteration-16 candidate; working-tree digest cbc40932... = the iteration-16 +// archive, verified by diff). Round 17 drafted exactly this edit and +// verified its codegen locally, but the agent session was killed by the +// round infrastructure before a valid proposal landed -- no source change +// was ever applied, no build/benchmark ran, and per policy that +// infrastructure failure does not count toward the valid-round ledger, so +// this round re-submits the same HIP-only change as a fresh valid round. +// The iteration-11 repair item is closed since iteration 13 (archived +// iteration-11 five-wave source + __launch_bounds__(320) ONLY, exact-correct +// with 0 mismatches; iterations 13-16 completed that arc and iteration 16 +// re-based on the accepted two-wave mapping). Status of the current source: +// the iteration-16 epilogue barrier-decoupled two-wave partial kernel +// (blockDim 128, grid 40 x 16 = 640 blocks, LDS 14,592 B = 4 resident +// blocks/CU = 8 resident waves/CU) measured official Graph median/P90 +// 19.9885/20.0133 us (min 19.9557 / max 21.3734, correctness 0 mismatches, +// Graph capture and paired M=2 fallback passed), improving the accepted +// 20.179/20.216 by 0.95%; it is kept as the experimental base ("official +// best is unchanged until cumulative improvement reaches the normal +// acceptance threshold"), and this round keeps building on that source. +// The partial kernel dominates the operator wall, and the exact-source code +// object shows its K loop is ROLLED: the generic loop is runtime-bounded +// (my_steps from the runtime k and gridDim.y args), so the backend emits one +// static body of [ds_read2_b32 (A) + 8 x ds_read_u8 (B) + ~25-VALU byte +// reassembly + 4 lgkmcnt drains + one v_mmac_i32_16x16x32_i8] with a +// loop-carried WAR on the single fragment register set: step s+1's LDS +// reads cannot issue until step s's v_mmac has consumed the same fragment +// registers, so every step's full LDS latency sits serially on the +// accumulator chain with only 2 resident waves/SIMD to cover it (reused PMC +// of this exact geometry: lds_instructions 107,520, lds_wait_instructions +// 19,597, lds_bank_conflicts 276,480). The guard launch is the kernel's +// ONLY call site and always runs the uniform 8-step slice (k == 4096, +// gridDim.y == 16 -> my_steps == 8 == kStageMaxSteps, rows == 256 == +// kStageMaxRows, staged capacity exactly filled), so this round specializes +// that case: when my_steps == kStageMaxSteps the K loop runs with a +// COMPILE-TIME trip count of 8 under #pragma unroll; the backend then fully +// unrolls the chain, renames the fragment registers per step and issues +// each step's LDS reads after the previous step's v_mmac has consumed its +// registers, overlapping the next step's LDS round trip with the current +// v_mmac's execution (no loop-carried WAR; verified in the local exact-flag +// code object of this edit: 8 static v_mmac + 72 static ds_reads in the +// specialized path vs 1 static v_mmac / 10 static ds_reads in the rolled +// predecessor, vgpr 34 / sgpr 39 / private_segment 0 / group_segment 14,592 +// B and exactly ONE s_barrier unchanged -> the 4-blocks/CU = 8-waves/CU +// occupancy is preserved). Both branches issue the identical ascending-s +// du_mma sequence into the same accumulator (v_mmac acc operand v[1:4] +// throughout) over the same staged byte mirrors, so every partial byte and +// every combine output byte is bit-identical to iteration 16 and exact +// correctness must pass with 0 mismatches; the generic rolled loop stays +// verbatim as the else path for any non-guard (my_steps != 8) +// configuration. Geometry, staging partition, stage byte mirrors, the +// blockDim 128 / grid 40 x 16 = 640 blocks launch, the epilogue +// (acc-tile store + __threadfence_block + read-back plane store), the +// combine kernel and its 320-block launch, the 16-plane 1,310,720 B +// workspace need, the workspace-less iteration-1 fallback, the exact +// (m,n,k)==(16,1280,4096) guard and the scalar fallback for every unmatched +// shape (including paired M=2) are all unchanged. // -// 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. +// Iteration 20 (DS-side staging-latency round, accepted): the exact-guard +// launch always stages the full 512-chunk B strip (chunks == 512 == 8 x 64 +// lanes), so wave 1's B copy now runs as ONE straight-line 8-deep batched +// group (all eight 16-B source loads issue before the first LDS store, one +// vmcnt drain covers the whole strip) instead of two serialized 4-deep +// iterations -- the second global round trip is removed from every block's +// staging join critical path (official Graph median/P90 19.967 -> 19.541/ +// 19.558 us, accepted). The generic rolled loop stays verbatim as the else +// path for any non-guard (chunks != 512) configuration. // -// 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. +// Iteration 21 (this round): DS-side bank-skew round for the B-fragment +// ds_read_u8 structure (pre-registered in iterations 18 and 20: the next +// increment must attack the PMC-documented DS ledger of this exact object -- +// lds_bank_conflicts 276,480 = ~2.57 replays per LDS instruction, dominated +// by the 8 byte-granular B-fragment reads per step per wave, 128 per block). +// The row-major B fragment loader reads each lane's 8 bytes at addresses +// base + (8*(lane>>4) + i) * ldm for i = 0..7, so the four row groups +// g = lane>>4 of the fragment's 32 k rows sit 8*ldm bytes apart; with the +// stage_b rows packed at the old ldm == 32, every group sat 256 B (== 0 +// mod 128 B) apart and each ds_read_u8 replayed the SAME LDS bank 4x per +// wavefront. This round pads each staged B row to kStageBRowStride = 36 B +// (32 data + 4 pad): the groups then sit 288 B (== 32 mod 128 B) apart and +// hit banks {b, b+8, b+16, b+24} -- the B-fragment reads become +// conflict-free while every staged data byte, the fragment contents, the +// k-ascending int32 order and the partial planes stay bit-identical (the +// pads are never written or read; only the loader's ldm argument and the +// wave-1 store offsets change). Wave-1 staging stores become four 4-B +// words per 16-B chunk (36*r is only 4-B aligned) at int32 offsets +// 9*r + 4*h + {0,1,2,3}; LDS store bandwidth is unchanged. LDS per block +// grows 14,592 -> 15,616 B (stage_b 8,192 -> 9,216): 4 x 15,616 = 62,464 +// <= 65,536, so the 4-resident-blocks/CU = 8-resident-waves/CU occupancy +// sweet spot is preserved. The workspace-less fallback, the exact +// (m,n,k)==(16,1280,4096) guard, the scalar fallback for every unmatched +// shape (including paired M=2), the combine kernel and its 320-block +// launch, the pack identity and the 16-plane 1,310,720 B workspace need are +// all unchanged. The iteration-11 repair item stays closed since +// iteration 13 (five-wave archive + __launch_bounds__(320) ONLY, exact- +// correct; the five-wave axis closed on measurement: best 20.613/20.641 us +// vs the accepted two-wave mapping 19.541/19.558 us). // -// 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). +// Iterations 22-25 (rounds 22 and 23 measured and rejected, tree reverted to +// the iteration-21 source e7ae27d0...; round 24 was killed mid-round before +// any measurement -- its archive draft was never compiled or run): round 22 +// removed the acc-tile LDS round trip from the partial epilogue at pinned +// occupancy (direct-register plane stores) and REGRESSED (+0.82 us -> +// 20.276/20.298 us) because the register-distributed acc values force 16-B +// store granules per wavefront instruction instead of the read-back path's +// four full 64-B lines -- rejected and reverted; round 23 widened each +// combine block to a 32-column n-strip with dwordx2 plane loads (160 blocks) +// and measured FLAT (median 19.4504 us vs the 19.4528 shadow, within its own +// +/-0.05 us falsification band, plus a 41.1 us max outlier) -- rejected and +// reverted. Round 23's pre-registered falsification branch (a) therefore +// committed THIS round to a combine block mapping whose per-plane reads run +// in longer contiguous runs per row (or a partial-plane store-order change), +// never occupancy / fragment duplication / raw asm / a scalar-fallback +// change. +// +// Iteration 25 (this round): combine mapping v3 -- one output row per +// block. The combine byte stream is fixed (every plane word is read +// exactly once, 64-B sectors fully consumed), so v1 (iteration 9: blocks of +// 4 rows x 16 cols) and v2 (iteration 23: 4 rows x 32 cols) both issue +// every per-plane wavefront load as FOUR 64-B (v1) or 128-B (v2) runs at +// the 5,120-B m-row stride -- four different DRAM rows, i.e. up to four row +// activations per load instruction -- which is exactly why widening the +// per-row run 64 -> 128 B in v2 could not move the combine (it never +// reduced the rows-per-instruction scatter). The guard combine launch now +// uses one-wave blocks that each own ONE m-row x 64 consecutive columns +// (kDummaCombineColsPerBlock = 64; grid 20 x 16 = 320 blocks, the same +// 320-block count and the same one-output-element-per-lane, 16-independent- +// plane-load structure as iteration 9): every per-plane wavefront load then +// reads 64 lanes x 4 B = 256 B CONTIGUOUS = four consecutive 64-B lines of +// ONE DRAM row (single run, ~4x fewer DRAM row activations for the same +// byte stream). Each element's ascending-s int32 sum, the float scale +// order and every output bf16 word are unchanged element-for-element, so +// the output is bit-identical to the iteration-21 shadow; the partial +// kernel, the 16-plane workspace contract, the exact guard, the +// workspace-less fallback and the scalar fallback (paired M=2 included) are +// byte-for-byte untouched. +// +// launch_pack_w8a8_weight (optional, out of the timed region) is an +// identity device-to-device copy: the logical [K, N] int8 weight and the +// [N, 1] fp32 scales are copied unchanged, so the DUMMA direct loader +// reads the exact logical [K, N] row-major layout. Later rounds may +// change the packed layout here together with the matching GEMM +// interpretation. +// +// This file is compiled with hipcc by torch.utils.cpp_extension; it must +// stay free of torch/ATen dependencies. Only the two extern "C" launchers +// below are the host contract consumed by csrc/bindings.cpp. +// +// Include order is the known-good DTK order: HIP runtime first, then the +// bfloat16 header, then du_mma.h (consumed by the DUMMA kernel below). #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; +// IEEE round-to-nearest-even conversion of an fp32 value to a bf16 bit +// pattern, implemented with pure integer ops. This is exactly the +// conversion applied by the CPU reference (`scaled.to(torch.bfloat16)`), +// independent of which bf16 type family the DTK headers provide. +__device__ __forceinline__ uint16_t f32_to_bf16_rne(float f) { + constexpr uint32_t kRoundBias = 0x7FFFu; + uint32_t u; + __builtin_memcpy(&u, &f, sizeof(u)); + const uint32_t lsb = (u >> 16) & 1u; // tie goes to even + u += kRoundBias + lsb; + return static_cast(u >> 16); +} -// 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}; +} // namespace // --------------------------------------------------------------------------- -// Scalar GEMM kernel: one thread computes one output element (generic -// fallback for every unmatched shape, including the paired M=2 shape). +// Scalar generic GEMM: out[m, n] = bf16(int32(sum_k a[m,k]*b[k,n]) +// * x_scale[m] * weight_scale[n]). +// One thread per output element. Adjacent lanes walk adjacent columns n +// (the fastest-changing output/weight dimension). gfx928 wavefront = 64; +// the launcher always uses blockDim 128 (a multiple of 64). // --------------------------------------------------------------------------- -// 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, +__global__ void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, // x_q [M, K] row-major + const int8_t* __restrict__ b, // packed_weight [K, N] row-major + const float* __restrict__ x_scale, // [M] (logical [M, 1]) + const float* __restrict__ w_scale, // [N] (logical [N, 1]) + uint16_t* __restrict__ out, // [M, N] bf16 payload 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) { + const int tid = + static_cast(blockIdx.x) * static_cast(blockDim.x) + + static_cast(threadIdx.x); + const int total = m * n; + if (tid >= total) { return; } - const int row = static_cast(linear / n); - const int col = static_cast(linear - static_cast(row) * n); + const int row = tid / n; // M is small for decode; N is the fast axis + const int col = tid - 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); + const int8_t* b_col = b + col; // stride n along k + int32_t acc = 0; +#pragma unroll 4 for (int kk = 0; kk < k; ++kk) { acc += static_cast(a_row[kk]) * - static_cast(b_col[static_cast(kk) * b_stride]); + static_cast(b_col[static_cast(kk) * n]); } - // 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); + const float scaled = static_cast(acc) * x_scale[row] * w_scale[col]; + out[tid] = f32_to_bf16_rne(scaled); } // --------------------------------------------------------------------------- -// 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]) +// Exact-guard DUMMA path for hy3_tp8_qkv_proj_m16 (M=16, N=1280, K=4096). +// +// Minimal native gfx928 INT8 tile: DUFragment fragments with a DUFragment +// accumulator, one 64-thread wavefront per block, one 16x16 output tile per +// block (grid = N/16 = 80 blocks). Every K step is an explicit int32 +// du_mma_sync over a 32-aligned ascending k chunk, so the int32 dot is +// identical to the scalar kernel's sum (exact integer addition, no +// overflow: |out| <= K*127*127 ~= 6.6e7 << 2^31). +// +// The gfx928 int8 m16n16k32 accelerator fragment ownership is opaque, so +// the epilogue materializes the tile with du_store_matrix_sync into a +// single-wave LDS staging array (library-defined mem_row_major placement, +// 256 x int32 = 1,024 B) and only afterwards applies the per-element float +// scaling and the bf16 RNE store. BlockDim is one wavefront, so the single +// __syncthreads() below is a same-wave barrier (all 64 threads reach it) +// and there is no cross-wave synchronization anywhere in the kernel. // --------------------------------------------------------------------------- -// 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, +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; +constexpr int kDummaWaveSize = 64; + +// Combine constants. Iteration 9: the exact-guard launcher (SPLIT_K == 16) +// launched the combine kernel as 80 n-tiles x four-row m-groups of one-wave +// blocks and the kernel specialized split_count == kDummaCombineSplits with +// one output element per lane. Iteration 25 (mapping v3): the guard launch +// becomes 20 column-groups x 16 m-rows of one-wave blocks -- each block owns +// ONE output row x kDummaCombineColsPerBlock = 64 consecutive columns (4 x +// kDummaN), so every per-plane wavefront load reads a single 256-B +// contiguous run (four consecutive 64-B lines inside one DRAM row) instead +// of four 64/128-B runs at the 5,120-B row stride (see the kernel comment +// below for the DRAM row-activation mechanism). +constexpr int kDummaCombineSplits = 16; // == the guard launcher kSplitK +constexpr int kDummaCombineColsPerBlock = 64; // n-columns per combine block + +__global__ void w8a8_dumma_m16_tile_kernel( + const int8_t* __restrict__ a, // x_q [M, K] row-major + const int8_t* __restrict__ b, // packed_weight [K, N] row-major + const float* __restrict__ x_scale, // [M] + const float* __restrict__ w_scale, // [N] + uint16_t* __restrict__ out, // [M, N] bf16 payload 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); // blockDim == 64 + const int n0 = static_cast(blockIdx.x) * 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; + __shared__ __align__(16) int32_t acc_tile[kDummaM * kDummaN]; 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 + signed char, du::dumma::row_major> 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. + // A is [16, k] (row stride k) and B is [k, 16] (row stride n); both stay + // row-major in their global tensors, so each 32-K step loads straight + // from global with the library row_major loaders. + for (int k0 = 0; k0 < k; k0 += kDummaK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); du::dumma::du_load_matrix_sync( - b_frag, b + static_cast(n0) * k + kk, k); + b_frag, b + static_cast(k0) * n + n0, n); 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); -} + du::dumma::du_store_matrix_sync(acc_tile, acc_frag, kDummaN, + du::dumma::mem_row_major); + __syncthreads(); -// 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; + // 256 elements / 64 lanes: each lane scales and stores four outputs. +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; linear += kDummaWaveSize) { + const int row = linear / kDummaN; + const int col = linear - row * kDummaN; + const int out_col = n0 + col; + const float scaled = static_cast(acc_tile[linear]) * + x_scale[row] * w_scale[out_col]; + out[static_cast(row) * n + out_col] = f32_to_bf16_rne(scaled); + } } // --------------------------------------------------------------------------- -// 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]) +// Split-K partial kernel (iteration 2 probe, widened in iteration 4 to +// multi-N-tile reuse, B-staged in iteration 5): same m16n16k32 int8 DUMMA +// tile as iteration 1, but gridDim = (N/(16*kTilesPerWave), SPLIT_K) so +// every one-wave block owns kTilesPerWave adjacent 16-wide N tiles AND one +// ascending 32-aligned K slice, and writes its exact int32 partials into +// workspace planes partials[split][16][n]. Iteration 4: kTilesPerWave = 2, +// so per 32-K step the wave loads the A fragment once and shares it between +// the two du_mma_sync calls (A bytes of the step reused across both output +// tiles; every B byte is still read exactly once). Iteration 5 (staging +// round, B-only): each block first stages its whole 32-column B slice +// (my_steps*32 rows x 32 B, <= 352 x 32 = 11,264 B under the SPLIT_K=12 +// guard launch) into an LDS tile stage_b that is an EXACT byte mirror of +// the block's global B strip (16-B int4 cooperative loads, 4-deep source +// batching so one vmcnt drain covers 4 in-flight round trips), then after +// one same-wave __syncthreads() the K loop reads the B0/B1 fragments from +// LDS with du_load_matrix_sync (row stride 32) -- same fragment bytes as +// the iteration-4 direct global loads, so the int32 partials stay +// bit-identical -- while A stays on the direct-global path (this round +// isolates the B side only). BlockDim is one 64-lane wavefront: the two +// barriers are the staging barrier above and the single same-wave +// __syncthreads() after the LDS materialization of both partial tiles. +// LDS per block: 11,264 B (stage_b) + 2 x 1,024 B (acc tiles) = 13,312 B +// -> 65,536 / 13,312 = 4.9, so the 4-blocks-per-CU launch co-residency is +// preserved (LDS is not the occupancy limiter). +// +// Iteration 7 (this round): A+B whole-slice LDS staging. Each one-wave +// block now stages BOTH operands of its whole K slice once: stage_a holds +// the 16 m-rows x 256 slice bytes at row stride 272 (256 data + 16 pad for +// bank skew, 4,352 B) and stage_b holds the 256 x 32 B strip (8,192 B); +// both use 16-B int4 cooperative loads with the same 4-deep batching, one +// same-wave __syncthreads(), and then the K loop reads every fragment from +// LDS (A at row stride 272, B0/B1 at row stride 32) -- zero per-step global +// loads and zero vmcnt waits on the MMA path. Geometry moved to SPLIT_K=16 +// (uniform 8-step slices, 40 x 16 = 640 blocks): per-block LDS 4,352 + +// 8,192 + 2,048 = 14,592 B -> floor(65,536/14,592) = 4, so the 4-blocks/CU +// co-residency of iteration 5 is preserved. Barriers per block stay two +// (staging, epilogue). Staged tiles are exact byte mirrors, so partials +// are bit-identical to iteration 5 and the combine kernel is unchanged. +// +// Iteration 8 (occupancy/resource round, this round): tune the waves-per- +// block occupancy limiter. PMC of the accepted iteration-7 code object +// (source digest b64e3209..., grid 640, workgroup 64, arch_vgpr 40, sgpr 48, +// scratch 0, LDS/block 14,592 B -> floor(65,536/14,592) = 4 resident blocks/ +// CU, 5.33 waves/CU by grid) shows a per-SIMD latency-exposure wall: every +// block is ONE wavefront, so at 4 resident blocks/CU the four waves land one +// per SIMD and each SIMD has no second independent wave to issue while the +// resident wave stalls on its per-step s_waitcnt lgkmcnt (lds_wait 8,359) +// after the staged-fragment du_load_matrix_sync and on its du_mma accum +// chain (VALU issue utilization ~3-6% of the 26.56 us profiled kernel; +// vmem_read 7,680 = only the 12 one-time staging loads per block). Fix: +// blockDim 64 -> 128, i.e. TWO 64-lane wavefronts per block, and the block's +// 32-column strip is split by ownership: wave w computes ONLY n-tile w (cols +// [n0+16w, n0+16w+16)) over the shared 8-step slice with its own independent +// accumulator chain, one du_mma per step per wave (same total MMA count). +// LDS per block is UNCHANGED (stage_a 4,352 + stage_b 8,192 + 2 x 1,024 acc +// tiles = 14,592 B), so 4 blocks/CU still fit but now carry 8 resident waves +// = 2 per SIMD: one wave's LDS->MMA stall is filled by the other. The +// one-time staging is partitioned, NOT duplicated: wave 0 copies the shared +// A slice and wave 1 copies the B strip (their iteration-7 loops unchanged), +// so exactly the same global bytes are read once (no repeated HBM reads for +// occupancy) and the staging barrier becomes a cross-wave __syncthreads(). +// Deliberate added cost: the A fragment is re-read from LDS by each wave +// (lds_instructions 99,840 -> ~125-135k predicted; LDS pipe/conflicts are +// not the wall at ~3-6% issue utilization -- the trade is LDS re-reads, not +// HBM traffic). Geometry, planes, combine kernel, guards and fallbacks are +// otherwise unchanged. // --------------------------------------------------------------------------- -// 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, +__global__ void w8a8_dumma_m16_splitk_partial_kernel( + const int8_t* __restrict__ a, // x_q [16, k] row-major + const int8_t* __restrict__ b, // packed_weight [k, n] row-major + int32_t* __restrict__ partials, // [gridDim.y][16][n] int32 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)); - } - } + constexpr int kTilesPerWave = 2; // two adjacent 16-wide N tiles per block + // Iteration 8 (occupancy round): each block is TWO 64-lane wavefronts + // (blockDim 128). Wave w owns N tile w of the block's 32 columns (cols + // [n0 + 16*w, n0 + 16*w + 16)) and runs an independent 8-step accumulator + // chain over the shared K slice for that tile alone. + const int tid = static_cast(threadIdx.x); // blockDim == 128 + const int w = tid >> 6; // wave id inside block 0/1 + const int lane = tid & (kDummaWaveSize - 1); // lane id within the wave + const int group = static_cast(blockIdx.x); // 32-wide n-group id + const int split = static_cast(blockIdx.y); // k-slice id + const int n0 = group * (kDummaN * kTilesPerWave); - // 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)); - } - } + // Stage-buffer sizing is tied to the ONLY launch site of this kernel: the + // exact (m,n,k)==(16,1280,4096) guard with gridDim.y == 16 (iteration 7), + // so total_steps == 128 and every split owns exactly 8 steps = 256 rows + // (128 == 16 * 8; the generic slice formula below degenerates to uniform + // 8-step slices). Compile-time capacities: 256 rows of B per block and + // 16 m-rows x (256 + 16 pad) bytes of A per block. + constexpr int kStageSplitCount = 16; // gridDim.y of the guard launch + constexpr int kStageMaxSteps = + (128 + kStageSplitCount - 1) / kStageSplitCount; // 8 (uniform) + constexpr int kStageMaxRows = kStageMaxSteps * kDummaK; // 256 rows x 32 + constexpr int kStageARowStride = + kStageMaxRows + 16; // 272 B per A m-row: 256 data + 16 pad for banks + // Iteration 21 (DS-side bank-skew round): each staged B k-row now occupies + // kStageBRowStride = 36 B (32 data bytes + 4 pad bytes). The row-major B + // fragment loader reads a lane's 8 fragment bytes at addresses + // base + (8*(lane>>4) + i) * ldm, i = 0..7: the four lane groups + // g = lane>>4 sit 8*ldm bytes apart, so with the old ldm == 32 every group + // landed 256 B (== 0 mod 128 B) apart and each of the 8 ds_read_u8 per + // fragment replayed the SAME LDS bank 4x (the PMC-documented + // lds_bank_conflicts 276,480 = ~2.57 replays/instr on this object, ~89% of + // it from the B-fragment byte reads: 128 per block). ldm == 36 puts the + // groups 288 B (== 32 mod 128 B, 8 banks) apart, so the four groups hit + // banks {b, b+8, b+16, b+24} -- fully conflict-free -- while the 4-B pad + // per row keeps every row start 4-B aligned for the staging stores. + // stage_b grows 8,192 -> 9,216 B and the block total becomes 4,352 + // (stage_a) + 9,216 (stage_b) + 2,048 (acc tiles) = 15,616 B/block: + // 4 x 15,616 = 62,464 <= 65,536, so the 4-resident-blocks/CU = + // 8-resident-waves/CU occupancy sweet spot is preserved unchanged. + constexpr int kStageBRowStride = kDummaN * 2 + 4; // 36 B per staged B row - __syncthreads(); + __shared__ __align__(16) int32_t acc_tile0[kDummaM * kDummaN]; + __shared__ __align__(16) int32_t acc_tile1[kDummaM * kDummaN]; + __shared__ __align__(16) int8_t stage_a[kDummaM * kStageARowStride]; + __shared__ __align__(16) int8_t stage_b[kStageMaxRows * kStageBRowStride]; 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; + signed char, du::dumma::row_major> + b_frag; // this wave's B half of the strip (b_col == 16*w bytes) du::dumma::DUFragment - acc_frag; + acc_frag; // this wave's accumulator for its n16 tile 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); + // 32-aligned non-uniform K slicing over the gridDim.y splits: split s owns + // my_steps ascending 32-K steps beginning at step_begin, slices are + // disjoint and cover [0, k) exactly once in ascending order, so the int32 + // partial sums and the ascending combine below reproduce the scalar + // k-ascending sum bit-for-bit. + const int total_steps = k / kDummaK; + const int split_count = static_cast(gridDim.y); + const int base_steps = total_steps / split_count; + const int rem_steps = total_steps - base_steps * split_count; + const int my_steps = base_steps + (split < rem_steps ? 1 : 0); + const int step_begin = + split * base_steps + (split < rem_steps ? split : rem_steps); + const int k_base = step_begin * kDummaK; // first global k row of slice + const int rows = my_steps * kDummaK; // slice rows: my_steps x 32 + + // Stage the block's operands (iteration 8: PARTITIONED across the two + // wavefronts, not duplicated -- the union reads exactly the same global + // bytes as the iteration-7 single wave, so no byte is fetched twice and + // HBM traffic is unchanged). Wave 0 copies the shared A slice: the 16 + // m-rows x rows bytes of A (global columns [k_base, k_base + rows)), + // stored m-major in stage_a at row stride kStageARowStride (an exact byte + // mirror of global a[m][k_base .. k_base + rows)). Wave 1 copies the + // block's 32-column B strip: my_steps*32 rows x 32 B, stored row-major in + // stage_b at the 36-B row pitch kStageBRowStride (32 data + 4 pad bytes + // per row; iteration 21 bank skew -- an exact byte mirror of the global + // data rows [k_base, k_base + rows) x columns [n0, n0 + 32), pads never + // written or read). Each copy uses the + // unchanged iteration-7 cooperative code: 64 lanes x 16 B per wavefront + // instruction, chunks in LDS int4 units, loads batched 4-deep in source so + // one vmcnt drain covers 4 in-flight global round trips (the two waves' + // round trips overlap further across the SIMDs). Global m-rows are + // k-major at stride k and global B rows are n-major at stride n (k % 16 == + // 0 and n % 16 == 0 under the guard; k_base and n0 are 32-aligned), so + // every 16-B global load is 16-B aligned (the wave-1 LDS stores are + // 4-B-granular since iteration 21: 36*row is only 4-B aligned). + const int a_chunks = (rows * kDummaM) / 16; // rows/16 int4 per m-row + const int chunks = rows * (kDummaN * 2) / 16; // rows*2 int4 chunks + if (w == 0) { + const int per_row = rows / 16; + const int4* __restrict__ a4 = reinterpret_cast(a); + int4* __restrict__ ald4 = reinterpret_cast(stage_a); + const int kb4 = k_base >> 4; // global int4 column of the slice start + const auto ld_chunk = [&](int e) { + const int m = e / per_row; + const int t = e - m * per_row; + return a4[static_cast(m) * (k >> 4) + kb4 + t]; + }; + const auto st_chunk = [&](int e, const int4& v) { + const int m = e / per_row; + const int t = e - m * per_row; + ald4[m * (kStageARowStride / 16) + t] = v; + }; + int e = lane; + for (; e + 3 * kDummaWaveSize < a_chunks; e += 4 * kDummaWaveSize) { + const int4 v0 = ld_chunk(e); + const int4 v1 = ld_chunk(e + kDummaWaveSize); + const int4 v2 = ld_chunk(e + 2 * kDummaWaveSize); + const int4 v3 = ld_chunk(e + 3 * kDummaWaveSize); + st_chunk(e, v0); + st_chunk(e + kDummaWaveSize, v1); + st_chunk(e + 2 * kDummaWaveSize, v2); + st_chunk(e + 3 * kDummaWaveSize, v3); + } + for (; e < a_chunks; e += kDummaWaveSize) { + st_chunk(e, ld_chunk(e)); + } + } else { + const int4* __restrict__ b4 = reinterpret_cast(b); + const auto ld_chunk = [&](int e) { + const int r = e >> 1; + // Global int4 index of (global row k_base + r, col half e&1). + return b4[static_cast(k_base + r) * (n >> 4) + (n0 >> 4) + + (e & 1)]; + }; + // Iteration 21: stage_b rows now live at the 36-B pitch + // kStageBRowStride (32 data bytes + 4 pad per row), so a 16-B chunk + // (row r = e >> 1, col half h = e & 1 of that row) can no longer be + // stored with one 16-B-aligned ds_write_b128 (36*r is 16-B aligned only + // for r % 4 == 0). Each chunk is stored as four consecutive int32 words + // at word offset 9*r + 4*h + {0,1,2,3} == (36*r + 16*h)/4 + {0,1,2,3}: + // every data byte lands at exactly the same byte offset (row*36 + ncol) + // the loader reads back, so the staged data byte mirror of the global B + // strip is unchanged -- only the pad bytes (row*36 + 32..35) are never + // written and never read (the fragment loader only touches ncol < 32). + // LDS store bandwidth is unchanged (4 x 4-B words per chunk); the four + // words of one chunk are consecutive, so a lane's store stream stays + // dense per row half. + int32_t* __restrict__ lds32 = reinterpret_cast(stage_b); + const auto st_chunk = [&](int e, const int4& v) { + const int word0 = (e >> 1) * (kStageBRowStride >> 2) + (e & 1) * 4; + lds32[word0 + 0] = v.x; + lds32[word0 + 1] = v.y; + lds32[word0 + 2] = v.z; + lds32[word0 + 3] = v.w; + }; + // Iteration 20 (guard-slice staging latency specialization): the guard + // launch is the ONLY call site of this kernel and always stages the + // full 512-chunk B strip (chunks == rows*2 == 512 == 8 x 64 lanes, rows + // == 256 == kStageMaxRows, staged capacity exactly filled). The + // runtime-bounded 4-deep loop below then compiles to TWO serialized + // 4-deep iterations (each iteration issues four 16-B source loads and + // drains them with vmcnt waits before the next iteration's loads can + // issue -- two full global round trips on the block's staging critical + // path, one per iteration), and wave 0's A copy finishes in a single + // 4-deep group and then idles at the __syncthreads() join while wave 1 + // completes its second round trip. When chunks == 8*kDummaWaveSize the + // loop now runs as ONE straight-line 8-deep batched group: all eight + // 16-B source loads per lane are issued before the first LDS store, so + // one drain window covers the whole strip and the second round trip is + // removed from the join's critical path. Every chunk id [0, 512) is + // still loaded and stored exactly once by exactly one lane to the same + // stage_b byte (the chunk set, the store order and the byte mirror are + // identical to the rolled loop), so every staged byte, every partial + // and every combine output byte is bit-identical; the generic rolled + // loop stays verbatim as the else path for any non-guard (chunks != + // 512) configuration, which never executes at the only launch site. + if (chunks == 8 * kDummaWaveSize) { + const int4 v0 = ld_chunk(lane); + const int4 v1 = ld_chunk(lane + kDummaWaveSize); + const int4 v2 = ld_chunk(lane + 2 * kDummaWaveSize); + const int4 v3 = ld_chunk(lane + 3 * kDummaWaveSize); + const int4 v4 = ld_chunk(lane + 4 * kDummaWaveSize); + const int4 v5 = ld_chunk(lane + 5 * kDummaWaveSize); + const int4 v6 = ld_chunk(lane + 6 * kDummaWaveSize); + const int4 v7 = ld_chunk(lane + 7 * kDummaWaveSize); + st_chunk(lane, v0); + st_chunk(lane + kDummaWaveSize, v1); + st_chunk(lane + 2 * kDummaWaveSize, v2); + st_chunk(lane + 3 * kDummaWaveSize, v3); + st_chunk(lane + 4 * kDummaWaveSize, v4); + st_chunk(lane + 5 * kDummaWaveSize, v5); + st_chunk(lane + 6 * kDummaWaveSize, v6); + st_chunk(lane + 7 * kDummaWaveSize, v7); + } else { + int e = lane; + for (; e + 3 * kDummaWaveSize < chunks; e += 4 * kDummaWaveSize) { + const int4 v0 = ld_chunk(e); + const int4 v1 = ld_chunk(e + kDummaWaveSize); + const int4 v2 = ld_chunk(e + 2 * kDummaWaveSize); + const int4 v3 = ld_chunk(e + 3 * kDummaWaveSize); + st_chunk(e, v0); + st_chunk(e + kDummaWaveSize, v1); + st_chunk(e + 2 * kDummaWaveSize, v2); + st_chunk(e + 3 * kDummaWaveSize, v3); + } + for (; e < chunks; e += kDummaWaveSize) { + st_chunk(e, ld_chunk(e)); + } + } } + __syncthreads(); // stage_a and stage_b ready for every lane of both waves - // 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); -} + // LDS-only K loop (iteration 8): every 32-K step wave w reads the A + // fragment from the staged A mirror at row stride kStageARowStride and its + // own B half from the staged B strip at the 36-B row pitch + // kStageBRowStride (byte column b_col; iteration 21 bank skew) + // -- the same fragment bytes the iteration-7 loads produced, with no + // global round trip and no vmcnt wait on the MMA path -- then issues ONE + // du_mma on its private accumulator. The two resident wavefronts carry + // independent 8-step chains, so a SIMD holding both can issue one wave's + // instructions while the other stalls on its s_waitcnt lgkmcnt / du_mma + // dependency (per-SIMD latency exposure was the iteration-7 wall). + const int b_col = w * kDummaN; // byte offset of this wave's n16 tile -// 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); + // Iteration 18 (compile-time 8-step K-loop specialization): the guard + // launch is the ONLY call site of this kernel and always runs the uniform + // 8-step slice (k == 4096, gridDim.y == 16 -> my_steps == 8 == + // kStageMaxSteps, rows == 256 == kStageMaxRows, staged capacity exactly + // filled). The runtime-bounded generic loop below compiles to a ROLLED + // body (verified in the exact-source code object: one static v_mmac per + // step with per-step ds_read2_b32 + 8 x ds_read_u8 + ~25-VALU byte + // reassembly + 4 lgkmcnt drains on the accumulator chain and a + // loop-carried WAR on the single fragment set), so every step's full LDS + // latency sits serially on the MMA path with only 2 resident waves/SIMD + // to cover it (PMC of the exact object: lds_instructions 107,520, + // lds_wait_instructions 19,597). Making the guard trip count + // compile-time 8 lets the backend fully unroll the chain, rename the + // fragment registers per step and hoist later steps' LDS reads across the + // current step's v_mmac (local exact-flag codegen check of this edit: 8 + // static v_mmac in the specialized path with the next steps' LDS reads + // issued before the current v_mmac; vgpr 34 / private_segment 0 / + // group_segment 14,592 B unchanged -> the 4-blocks/CU = 8-waves/CU + // occupancy is preserved). Both branches issue the identical + // ascending-s du_mma sequence into the same accumulator over the same + // staged byte mirrors, so every partial byte and every combine output + // byte is bit-identical to iteration 16; the generic rolled loop stays + // verbatim as the else path for any non-guard (my_steps != 8) + // configuration. + if (my_steps == kStageMaxSteps) { +#pragma unroll + for (int s = 0; s < kStageMaxSteps; ++s) { + du::dumma::du_load_matrix_sync( + a_frag, stage_a + static_cast(s) * kDummaK, + kStageARowStride); + du::dumma::du_load_matrix_sync( + b_frag, + stage_b + + static_cast(s) * kDummaK * kStageBRowStride + b_col, + kStageBRowStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + } else { + for (int s = 0; s < my_steps; ++s) { + du::dumma::du_load_matrix_sync( + a_frag, stage_a + static_cast(s) * kDummaK, + kStageARowStride); + du::dumma::du_load_matrix_sync( + b_frag, + stage_b + + static_cast(s) * kDummaK * kStageBRowStride + b_col, + kStageBRowStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + } + + // Epilogue (iteration 16, barrier-decoupled): each wave materializes its + // own tile into its own LDS acc buffer (wave 0 -> acc_tile0, wave 1 -> + // acc_tile1, 1,024 B each), then its 64 lanes read back that same buffer + // and store the 256 int32 partials (4 stores per lane) to + // plane[split][m][n0 + b_col + col]. acc_tile{w} is written and read ONLY + // by wave w, and the two waves' plane stores touch disjoint 16-column + // ranges, so nothing in the epilogue needs cross-wave convergence. The + // former epilogue __syncthreads() was a 128-thread cross-wave join whose + // release was gated by the slower of the block's two waves on a kernel + // that runs only 8 resident waves/CU (4 x 2-wave blocks) to hide the skew + // -- the same join whose five-wave form measured ~3.15 us of pure cost + // (iteration 13 23.764/24.086 -> iteration 15 20.613/20.641 us with the + // join removed, exact correctness passing in both). It is replaced by + // __threadfence_block(), which only orders THIS wave's own outstanding LDS + // writes (the du_store_matrix_sync ds_writes of its private acc tile) + // before its read-back ds_reads and costs no cross-wave wait. The staging + // __syncthreads() above stays (wave 0 stages the A slice and wave 1 stages + // the B strip; every wave reads bytes the other wave staged). Every + // partial byte, the k-ascending int32 order and every plane store address + // are unchanged, so the output is bit-identical to iteration 9 and exact + // correctness must pass with 0 mismatches. + if (w == 0) { + du::dumma::du_store_matrix_sync(acc_tile0, acc_frag, kDummaN, + du::dumma::mem_row_major); + } else { + du::dumma::du_store_matrix_sync(acc_tile1, acc_frag, kDummaN, + du::dumma::mem_row_major); + } + __threadfence_block(); // iteration 16: wave-local store->read ordering + // only; no cross-wave join (was __syncthreads()) + + int32_t* plane = partials + static_cast(split) * kDummaM * n; + const int32_t* tile = (w == 0) ? acc_tile0 : acc_tile1; + const int out_col0 = n0 + b_col; +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; + linear += kDummaWaveSize) { + const int row = linear / kDummaN; + const int col = linear - row * kDummaN; + plane[static_cast(row) * n + out_col0 + col] = tile[linear]; + } } // --------------------------------------------------------------------------- -// 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). +// Split-K combine kernel (iteration 2, restructured in iteration 9, mapping +// v3 in iteration 25): sums the gridDim.y int32 partial planes ascending per +// output element (exact int32, no overflow, bit-identical to the scalar +// k-ascending dot), applies x_scale[row] * weight_scale[col] in the scalar +// float order and stores bf16 through f32_to_bf16_rne. Iteration 9 added +// the guard specialization (split_count == 16): one output element per lane, +// all sixteen plane loads issued as INDEPENDENT reads into a 16-int register +// array (single vmcnt drain instead of a dependent load->wait->add chain of +// 16 DRAM round trips per element), then the ascending-s register sum and +// the scaled bf16 store. Iteration 25 (this round) keeps that per-lane +// structure and changes ONLY the block-to-output mapping: every one-wave +// block owns ONE m-row x kDummaCombineColsPerBlock = 64 consecutive columns +// (guard grid = 20 x 16 = 320 one-wave blocks, same block count as +// iteration 9). In the iteration-9/23 mappings each block covered 4 rows x +// 16/32 columns, so every per-plane wavefront load issued four 64/128-B runs +// at the 5,120-B m-row stride -- four different DRAM rows, i.e. up to four +// row activations per load instruction -- which is why widening the per-row +// run 64 -> 128 B in iteration 23 could not move the combine (it never +// reduced the rows-per-instruction scatter). With ONE row per block every +// per-plane load instruction is 64 lanes x 4 B = 256 B CONTIGUOUS = four +// consecutive 64-B lines of one DRAM row (a single run: ~4x fewer DRAM row +// activations for the identical 1.31 MB byte stream, identical 64-B sectors, +// identical wavefront-load count). The element set is unchanged, so every +// element's ascending-s int32 sum, the float scale order and every output +// bf16 word are bit-identical to the pre-iteration-25 combine; the generic +// loop below stays verbatim for every other call site / split count / grid +// shape. // --------------------------------------------------------------------------- -__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]; +__global__ void w8a8_dumma_m16_combine_kernel( + const int32_t* __restrict__ partials, // [split_count][16][n] int32 + const float* __restrict__ x_scale, // [16] + const float* __restrict__ w_scale, // [n] + uint16_t* __restrict__ out, // [16, n] bf16 payload + int n, + int split_count) { + const int lane = static_cast(threadIdx.x); // blockDim == 64 + const int n0 = static_cast(blockIdx.x) * kDummaN; + + // --------------------------------------------------------------------- + // Iteration 25 specialization (mapping v3; guard launch only: + // split_count == 16 and gridDim.y == kDummaM == 16 -> blockIdx.x = + // 64-column group in [0, 20), blockIdx.y = one m-row in [0, 16)). Each + // lane owns output element (row, col) with row = blockIdx.y and col = + // blockIdx.x * 64 + lane: the sixteen independent 4-B plane reads (plane + // stride kDummaM * n int32) batch into one vmcnt drain, and because the + // 64 lanes of the block are 64 CONSECUTIVE columns of ONE row, every + // per-plane wavefront load covers 256 contiguous bytes (4 consecutive + // 64-B lines of one DRAM row) instead of four 5,120-B-strided runs; the + // ascending-s register sum then reproduces the exact int32 sequence for + // this element. + // --------------------------------------------------------------------- + if (split_count == kDummaCombineSplits && gridDim.y == kDummaM) { + const int row = static_cast(blockIdx.y); + const int col = + static_cast(blockIdx.x) * kDummaCombineColsPerBlock + lane; + const int64_t elem = static_cast(row) * n + col; + int32_t v[kDummaCombineSplits]; +#pragma unroll + for (int s = 0; s < kDummaCombineSplits; ++s) { + v[s] = partials[static_cast(s) * (kDummaM * n) + elem]; + } + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kDummaCombineSplits; ++s) { + acc += v[s]; // ascending split order: identical int32 sum + } + const float scaled = + static_cast(acc) * x_scale[row] * w_scale[col]; + out[elem] = f32_to_bf16_rne(scaled); + return; } - const float scaled = - static_cast(sum) * x_scale[row] * weight_scale[out_col]; - out[row * n + out_col] = __float2bfloat16(scaled); +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; + linear += kDummaWaveSize) { + const int row = linear / kDummaN; + const int col = linear - row * kDummaN; + const int out_col = n0 + col; + int32_t acc = 0; + for (int s = 0; s < split_count; ++s) { + acc += partials[(static_cast(s) * kDummaM + row) * n + + out_col]; + } + const float scaled = + static_cast(acc) * x_scale[row] * w_scale[out_col]; + out[static_cast(row) * n + out_col] = f32_to_bf16_rne(scaled); + } } // --------------------------------------------------------------------------- -// Weight pack op (outside the timed region). +// Identity device-to-device packing (bootstrap): grid-stride byte copies on +// the caller stream. Valid for every (K, N), including unmatched shapes. // --------------------------------------------------------------------------- -// 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( +__global__ 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]; + 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]; } } -__global__ __launch_bounds__(kCopyBlockThreads) void -w8a8_identity_copy_f32_kernel( +__global__ 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 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]; } - 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). +// Host launcher consumed by csrc/bindings.cpp (gemm_out, timed region). +// The exact (16, 1280, 4096) guard dispatches the split-K DUMMA partial + +// combine kernels (two wavefronts per partial block since iteration 8); +// the scalar fallback below launches 128 threads per 128 outputs. // --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, @@ -765,128 +1131,118 @@ extern "C" void launch_w8a8_gemm( int n, int k, hipStream_t stream) { - auto* out_bf16 = static_cast(out); + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0 || a == nullptr || b == nullptr || + x_scale == nullptr || weight_scale == nullptr || out == nullptr) { + return; + } - // 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)). + // --------------------------------------------------------------------- + // Exact-shape DUMMA guard: hy3_tp8_qkv_proj_m16 (M=16, N=1280, K=4096). + // Lineage: iteration 4 widened each one-wave block to TWO adjacent + // 16-wide N tiles (32 columns) sharing one A-fragment load per 32-K step + // at SPLIT_K=12 (grid 40 x 12 = 480 blocks = 4/CU over 120 CUs); + // iteration 5 staged each block's 32-column B slice into LDS once and + // read the B0/B1 fragments from LDS per step, A direct-global + // (39.404/39.485 us official Graph median/P90, from 41.154/41.233). + // Iteration 7: BOTH operands are staged per K slice -- the A fragment + // also comes from LDS, so the K loop is global-load-free -- which forces + // SPLIT_K=16: sixteen UNIFORM 8-step slices (128 == 16 x 8, rows + // {256 x 16}, exact ascending cover of [0, 4096), every boundary a + // DUMMA-k multiple), grid = 40 N groups x 16 = 640 one-wave blocks + // (~5.33 waves/CU, 4 resident blocks/CU by LDS: 14,592 B/block -> + // 65,536/14,592 = 4.4), each block writing exact int32 partial planes + // (cols [n0, n0+16) and [n0+16, n0+32) of plane[split]), then one + // UNCHANGED 80-block combine kernel (generic over split_count = 16) that + // sums the planes ascending and emits the scaled bf16 output. Staged + // tiles are exact byte mirrors of the global operands, so fragment + // bytes, K-slice bounds, the k-ascending int32 order and the partial + // plane layout are unchanged: the combine+scale output is bit-identical + // to iteration 5. + // Iteration 8 (this round): occupancy/resource change only -- the same + // 640-block SPLIT_K=16 geometry now runs TWO 64-lane wavefronts per block + // (blockDim 128): wave w computes tile w (cols [n0+16w, n0+16w+16)) of + // the block's 32-column strip with its own independent accumulator chain + // over the shared slice. LDS per block is UNCHANGED (14,592 B -> 4 + // resident blocks/CU), so co-residency rises from 4 to 8 resident waves + // per CU = 2 per SIMD, hiding each wave's per-step LDS->MMA stall behind + // the other wave. Staging is partitioned across the two waves (wave 0 = + // A slice, wave 1 = B strip; identical global bytes read exactly once). + // The combine kernel, the 16-plane workspace need (1,310,720 B), the + // k-ascending partial order and the bit-identical output are unchanged. + // Both launches stay on the caller stream inside the timed API/Graph; + // the combine cost is part of the operator wall. Requires 16 int32 + // partial planes (1,310,720 B) of the caller's workspace -- the full + // fixed 16-plane 1,310,720 B allocation for this shape; when the + // workspace is missing/undersized the exact guard falls back to the + // iteration-1 one-wave direct geometry (still no workspace). Every + // unmatched (m, n, k) -- including the paired M=2 API shapes with the + // same (N, K) -- keeps reaching the scalar fallback below. + // Iteration 9 (this round): combine-side restructure only. PMC of the + // accepted iteration-8 code object profiles the combine kernel at 16.960 + // us vs 19.361 us for the partial kernel (46.7% of the 36.321 us profiled + // operator aggregate) with a 6.9% L2 hit rate and only 80 one-wave blocks, + // so the guard now launches the SAME combine kernel on a 2-D grid (80 + // n-tiles x 4 four-row m-groups = 320 one-wave blocks) and the kernel's + // split_count==16 specialization gives each lane one output element that + // batches all sixteen plane reads into a register array before the + // ascending sum. Partials, output bytes, bit-exactness and the partial + // kernel launch are unchanged. + // --------------------------------------------------------------------- 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; - } - } - } + constexpr int kDummaThreads = 64; // one wavefront (tile/combine) + constexpr int kPartialThreads = 2 * kDummaThreads; // 2 waves per block + constexpr int kTilesPerWave = 2; // two 16-wide N tiles per block (n32) + constexpr int kSplitK = 16; // 40 groups x 16 = 640 blocks, 4/CU + const int ntiles = n / kDummaN; // 1280 / 16 = 80 combine tiles + const int ngroups = ntiles / kTilesPerWave; // 80 / 2 = 40 n32 groups + const int64_t kWorkspaceNeeded = + static_cast(kSplitK) * kDummaM * n * sizeof(int32_t); + auto* out_bf16 = reinterpret_cast(out); + if (workspace != nullptr && workspace_bytes >= kWorkspaceNeeded) { + auto* partials = reinterpret_cast(workspace); + hipLaunchKernelGGL(w8a8_dumma_m16_splitk_partial_kernel, + dim3(ngroups, kSplitK), dim3(kPartialThreads), 0, + stream, a, b, partials, n, k); + // Iteration 25 (mapping v3): 2-D combine grid = 20 column-groups x 16 + // one-row blocks = 320 one-wave blocks (one output row x 64 consecutive + // columns per block, one output element per lane, 16-deep independent + // plane-load register batch in the kernel's split_count==16 path; every + // per-plane wavefront load is one 256-B contiguous run of one DRAM row + // instead of four 5,120-B-strided runs). + hipLaunchKernelGGL(w8a8_dumma_m16_combine_kernel, + dim3(n / kDummaCombineColsPerBlock, kDummaM), + dim3(kDummaThreads), 0, stream, + partials, x_scale, weight_scale, out_bf16, n, + kSplitK); + } else { + // Workspace-less fallback: iteration-1 one-wave direct geometry. + hipLaunchKernelGGL(w8a8_dumma_m16_tile_kernel, + dim3(ntiles), dim3(kDummaThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, n, k); } - - // 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); + // --------------------------------------------------------------------- + // Scalar generic fallback: every shape outside the exact guard above. + // --------------------------------------------------------------------- + + constexpr int kBlockThreads = 128; // multiple of the gfx928 wavefront 64 + const int total = m * n; + const int grid = (total + kBlockThreads - 1) / kBlockThreads; + auto* out_bf16 = reinterpret_cast(out); + hipLaunchKernelGGL(w8a8_scalar_gemm_kernel, + dim3(grid), dim3(kBlockThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); } +// --------------------------------------------------------------------------- +// Host launcher for the optional out-of-timed-region weight packing. +// Bootstrap: identity device-to-device copy of the int8 [K, N] weight and +// the fp32 [N, 1] scales; valid for every (K, N), unmatched shapes included. +// --------------------------------------------------------------------------- extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -895,47 +1251,27 @@ extern "C" void launch_pack_w8a8_weight( 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); + constexpr int kPackBlockThreads = 256; + + const int64_t i8_count = static_cast(k) * n; + if (i8_count > 0 && raw_weight != nullptr && packed_weight != nullptr) { + const int i8_grid = + static_cast((i8_count + kPackBlockThreads - 1) / + kPackBlockThreads); + hipLaunchKernelGGL(w8a8_identity_copy_i8_kernel, + dim3(i8_grid), dim3(kPackBlockThreads), 0, stream, + raw_weight, packed_weight, i8_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)); + const int64_t f32_count = static_cast(n); + if (f32_count > 0 && weight_scale != nullptr && + packed_weight_scale != nullptr) { + const int f32_grid = + static_cast((f32_count + kPackBlockThreads - 1) / + kPackBlockThreads); + hipLaunchKernelGGL(w8a8_identity_copy_f32_kernel, + dim3(f32_grid), dim3(kPackBlockThreads), 0, stream, + weight_scale, packed_weight_scale, f32_count); + } } // @@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 index 01194376..5d133848 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=hy3_tp8_shared_down_proj_m16 commit=6f8e396e4eb7b316acecbb928ea5db9f2f3ffe4d added=2026-08-27 +// baseline_us=17.312 speedup=2.24948 // median_us=7.696 p90_us=8.174 // source=hy3-dsh-tp8-m16-2-dc53295a // INT8 W8A8 GEMM implementation for gfx928 (K500SM_AI). 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 index 3b0c1c88..1f40f800 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=hy3_tp8_shared_gate_up_proj_m16 commit=2321c42b8618aae23700bc2cf3532b59c8f6c8d8 added=2026-08-27 +// baseline_us=66.527 speedup=6.40915 // median_us=10.38 p90_us=10.42 // source=hy3-dsh-tp8-m16-2-dc53295a // MetaInfer W8A8 INT8 GEMM - HIP implementation (gfx928 / K500SM_AI). 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 index 8cf04a9f..5be78d23 100644 --- 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 @@ -1,132 +1,85 @@ -// @@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) +// @@variant shape=hy3_tp8_o_proj_m4096 commit=5a89acbee264bee5d61f74c281509b86c3fcd4e3 added=2026-09-06 +// median_us=266.8 p90_us=267.2 speedup=53.05 baseline_us=1.415e+04 +// source=hy3-tp8-dsh-m16-and-m4096-2-f7fa9503 // INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. // -// Worker: worker_1 (physical GPU 1), assigned shape: +// Worker: worker_1 (physical GPU 1). Assigned shapes: +// hy3_tp8_o_proj_m16 : M=16, N=4096, K=1024 // 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)): +// +// * Large-prefill path (exact (m, n, k) == (4096, 4096, 1024), m4096 +// iteration-1 architecture round - DUMMA throughput baseline with the +// accepted lineage steady state): // 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. +// output tile per block; four wavefronts (256 threads, gfx928 wavefront +// = 64); each wave owns a 64x32 quadrant built from eight m16n16k32 +// int32 accumulator fragments. The block cooperatively vector-loads +// A[128,64] (x_q rows m0..m0+127, row-major stride k) and B[64,64] +// (packed [N, K] weight rows n0..n0+63, n-major stride k) into one +// single-buffered LDS stage (16,384 B total: a_tile 128*88 + b_tile +// 64*80; A rows padded 64 -> 88 B and staged as two 8-B halves per +// 16-B vector, B rows padded 64 -> 80 B); two __syncthreads per stage +// (one after the cooperative load, one before the next stage overwrites +// LDS). Each DUMMA operand fragment is filled by the direct 8-byte +// load_fragment8 LDS read (one contiguous 8-B read per lane per +// fragment on the m-major A tile / n-major B tile - no library loader +// byte-reassembly VALU and no per-byte ds_read_u8), two A fragments +// live per MMAC burst, and the epilogue is fused: int32 accumulator * +// x_scale[m] * weight_scale[n] -> bf16 through the coalesced fragment +// store (4x4 lane transpose + one 8-byte store per lane, 100% store +// sector efficiency). +// Grid dim3(64, 32) = 2048 output-tile blocks dwarfs the 120 CUs, so no +// split-K is used (nothing is ever written to the workspace). // -// 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). +// * Decode path (exact (m, n, k) == (16, 4096, 1024), iteration 5 +// architecture round): fused single-kernel split-K=2 - one 16x16 N tile +// per block, grid 256 x 128 threads (2.13 blocks/CU on the measured 120 +// CUs, two co-resident 64-lane wavefronts per block), wave w +// accumulating the ascending-k half w (K/2 = 512, 16 m16n16k32 steps) +// of the tile with the iteration-3 per-wave body (direct global +// fragment loads: matrix_a row_major off x_q, matrix_b col_major off +// the packed [N, K] weight, depth-1 double-buffered prefetch, no LDS in +// the K loop). The split combine is fused into this one kernel: wave 0 +// publishes its int32 accumulator fragment to a 1-KiB shared plane, one +// END-of-K __syncthreads makes it visible, and wave 1 adds plane 0 + +// its own upper-K-half fragment (exact int32, ascending split order, +// bit-identical to the unsplit ascending-k dot) and applies the fused +// x_scale * weight_scale -> bf16 epilogue. One launch on the caller +// stream inside the timed Graph, no caller workspace and no second +// combine kernel: the exact-shape Graph has a single node and no +// ~1-MiB workspace write+read round trip. Iteration 3's two-launch +// workspace partial+combine variant and iteration 2's unsplit two-wave +// kernel are superseded (the fused kernel needs no workspace, so no +// workspace-fit fallback branch remains). // -// 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. +// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), +// including the paired M=2 API shape with the same (K, N) = (1024, +// 4096) and any other M tail: one thread per output element, exact +// int32 dot over K, packed [N, K] weight decode (column n = one +// contiguous k-run), fused x_scale * weight_scale, bf16 store. // -// 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. +// * launch_pack_w8a8_weight: out-of-Graph, out-of-timed device transpose +// of the raw [K, N] weight into the packed [N, K] buffer +// (packed[n][k] = weight[k][n]) and an identity copy of weight_scale +// into packed_weight_scale for every (k, n). The [N, K] packed layout +// is what every GEMM consumer below reads: for the gfx928 INT8 +// m16n16k32 matrix_b fragment each lane owns 8 consecutive k bytes of +// one n row, so the [K, N] identity layout would force 8 scattered +// 1-byte global loads per lane per k32 step while the [N, K] layout +// makes them one contiguous 8-B vector load per lane per step. +// Packing never happens inside the timed GEMM and stays outside Graph +// capture; the packed buffer keeps the same byte count (K*N) and the +// same allocated address as raw_weight. // -// 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. +// Timed operator rules honored: no allocation, compilation, autotuning, +// weight packing, host/device synchronization, or default-stream launch; the +// GEMM runs on the caller-provided PyTorch current HIP stream; only `out` is +// written by the timed kernels (the fused decode split-K=2 kernel combines +// in block-local LDS, so the caller workspace is left untouched by every +// path). #include #include @@ -145,81 +98,109 @@ constexpr int kWaveSize = 64; // gfx928 native wavefront constexpr int kTargetM = 4096; constexpr int kTargetN = 4096; constexpr int kTargetK = 1024; +constexpr int kDecodeM = 16; // assigned decode shape M 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; +constexpr int kPad = 16; // B: 64 -> 80-byte LDS row stride (bank skew) +constexpr int kAPad = 24; // A: 64 -> 88-byte LDS row stride (22 dwords/row); + // 88 % 16 == 8, so odd rows are only 8-B aligned + // and the A stage is written as two 8-B halves + // (ds_write_b64 class). 22 dwords/row maps the 16 + // rows of a fragment window onto 16 distinct 4-B + // word phases within each fragment q phase (the + // 20-dword 80-B rows of the accepted port alias + // rows r and r+8, doubling the LDS bank pressure + // of every A-fragment read). +constexpr int kAStride = kStageK + kAPad; // A tile row stride (m-major) +constexpr int kBStride = kBlockN + kPad; // B tile row stride (n-major: row + // = n, k bytes along the padded row) +constexpr int kBlockThreads = 4 * kWaveSize; // 256 threads = 4 wavefronts +// Decode launch geometry (iteration 3 grid, iteration 5 fused combine): grid +// 256 = one 16x16 N tile per block reaches 2.13 blocks/CU on the measured 120 +// CUs - the finest legal grid at the two-blocks-per-CU latency-hiding target, +// since the 256 whole 16-col N tiles (4096/16) admit no finer whole-tile +// split - and keeps two co-resident 64-lane wavefronts per 128-thread block +// (the iteration-2 latency-hiding mechanism: one wave's global-load stalls +// overlap the other wave's MMAC/address work). Wave w of block b accumulates +// the ascending-k segment [w*512, (w+1)*512) of tile b. Iteration 3 combined +// the two per-tile halves through two caller-workspace int32 planes plus a +// second combine+scale kernel launch; iteration 5 fuses that combine into the +// same kernel through a 1-KiB block-local shared plane and one END-of-K +// __syncthreads (wave 0 publishes, wave 1 sums plane 0 + its own upper half +// and stores scaled bf16), removing the second Graph node and the workspace +// round trip entirely. +constexpr int kDecodeThreads = kWaveSize * 2; // 128 = 2 wavefronts +constexpr int kDecodeSplitK = 2; // K segments per tile +constexpr int kDecodeSegK = kTargetK / kDecodeSplitK; // 512 +constexpr int kDecodeSegSteps = kDecodeSegK / kTileK; // 16 x k32 steps +constexpr int kNumDecodeTiles = kTargetN / kTileN; // 256 16-col tiles +constexpr int kDecodeLaneInts = 4; // int32 accumulator values per lane +constexpr int kDecodeTileInts = + kWaveSize * kDecodeLaneInts; // 256 int32 = 1 KiB shared plane per block 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. +// block cooperatively stages A[128,64] from x_q (m-major, row stride k) and +// B[64,64] from the packed n-major [N, K] weight into a single-buffered LDS +// stage. Two barriers per K stage. DUMMA operand fragments are filled by the +// direct 8-byte LDS fill (load_fragment8) and the fused scale -> bf16 +// epilogue stores through the coalesced fragment store. // --------------------------------------------------------------------------- -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 +// 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 +// per-element store is four 2-byte scalar stores per lane whose wavefront +// addresses touch each 32-B sector at 25% utilization (vmem_write 262,144 +// for this shape). 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: 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 +// bf16 (8 B), and writes ONE 8-byte store per lane (vmem_write 65,536: 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 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 +// 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 (col0 % 4 == 0) and the 8-byte store // ((row*n + col0) % 4 == 0 elements -> byte offset % 8 == 0) are aligned. +// This is the exact mechanism validated on the accepted hy3 TP8 o_proj M4096 +// lineage (iteration 6, vmem_write_instructions 262,144 -> 65,536). + +// Plain round-to-nearest-even fp32 -> bf16 conversion (bit 16 of the fp32 +// pattern), without the NaN/Inf signaling-bit fixup of hip's +// __float2bfloat16. The arithmetic (u + 0x7fff + ((u >> 16) & 1)) >> 16 is +// exactly the finite-domain path of DTK's float_2_bfloatraw (exponent bits +// != all-ones), so every finite fp32 input - including +/-0, denormals, and +// values whose RNE carries into/out of the bf16 exponent field - produces +// the identical 16-bit pattern as __float2bfloat16; the outputs can differ +// only for NaN/Inf inputs, which this kernel's epilogue can never produce +// (|int32 dot| <= K*127^2 and both scale inputs are finite, so the scaled +// accumulator is always finite). Replacing the library conversion drops the +// exec-masked special-case chain the DTK otherwise emits per element in the +// epilogue tail (v_and exponent mask + v_cmp_ne + s_and_saveexec + s_xor + +// v_bfe + v_add3 + s_andn2_saveexec + v_mov/v_or + v_cmp_eq_sdwa + +// v_cndmask + s_or in the accepted iteration-5 ISA: ~12 extra SALU/VALU per +// element x 32 elements per wave, all executed once per block tail) for 3 +// VALU per element, with stored bits unchanged (verified against the +// float_2_bfloatraw source: the else-if branch it removes is reachable only +// from Inf/NaN inputs). +__device__ __forceinline__ unsigned short bf16_rne_finite_bits(float f) { + union { + float fp32; + unsigned int u32; + } u = {f}; + u.u32 += 0x7fffu + ((u.u32 >> 16) & 1u); // round half to even at bit 16 + return static_cast(u.u32 >> 16); +} + template __device__ __forceinline__ void store_prefill_fragment_coalesced( const AccFragment& frag, @@ -266,37 +247,37 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced( const int col0 = base_col + 4 * c4; const float xs = x_scale[row]; const float4 ws = *reinterpret_cast(weight_scale + col0); + // Left-associative scale order, matching the logical reference + // (A.float() @ B.float()) * x_scale[m] * weight_scale[n].T. 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); + static_cast(bf16_rne_finite_bits(v0)) | + (static_cast(bf16_rne_finite_bits(v1)) << 16) | + (static_cast(bf16_rne_finite_bits(v2)) << 32) | + (static_cast(bf16_rne_finite_bits(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). +// Direct 8-byte LDS fragment fill for the DUMMA m16n16k32 signed char +// operands on a padded tile whose row direction carries the fragment's +// 16-row index and whose columns hold the K bytes (A m-major tile row = m, +// B n-major tile row = n): the fragment slot of lane (r = lane & 15, +// kq = lane >> 4) is the eight CONSECUTIVE k-bytes at (lane & 15) * stride + +// (lane >> 4) * 8. The library du_load_matrix_sync int8 loaders make the +// compiler emit per-byte ds_read_u8 loads plus a mask/OR byte-reassembly +// chain per loaded dword that is an arithmetic identity on the v_mmac +// operand (dead issue and dead latency on the load -> v_mmac dependency +// path). Writing the same 8 bytes straight into the fragment storage (x[0..7] +// of the DUFragmentBase) 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 M4096 +// lineage (accepted iterations 10/12: 348.27 -> 322.99 us median, VALU +// ~22.56M -> ~15.7M, SQ_WAIT_INST_LDS drop, arch_vgpr 82 -> <= 64). template __device__ __forceinline__ void load_fragment8( Frag& frag, @@ -308,7 +289,7 @@ __device__ __forceinline__ void load_fragment8( *reinterpret_cast(base + off); } -__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, const float* __restrict__ x_scale, @@ -320,23 +301,29 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( 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 wave_row = wave >> 1; // quadrant rows [0,1] x 64 + const int wave_col = wave & 1; // quadrant cols [0,1] x 32 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. + // Single-buffered K-stage: A[128, 88] + B[64, 80] = 16,384 B/block. + // Round-8 residency: the one-stage-ahead staging payload (+12 VGPR) moves + // the code object from the accepted exact-64-VGPR 4-blocks/CU point (64 x + // 1024 threads = 65,536 VGPR exactly) to 80 VGPR => 3 blocks/CU (80 x 256 + // x 3 = 61,440 <= 65,536 VGPR; LDS 16,384 x 3 = 49,152 B; <= 85 VGPR keeps + // 3 blocks, no spills at the offline-verified 80). __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; - __shared__ __align__(16) int8_t b_tile[kBlockN * kBStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + // DUMMA m16n16k32 signed char operand fragments are filled by + // load_fragment8 with the eight consecutive k-bytes each lane owns on the + // padded tile (A m-major rows, B n-major rows of the [N, K] packed + // weight): no library loader, no per-byte ds_read_u8, no byte-reassembly + // VALU, one 8-byte LDS read straight into the v_mmac operand. Only two A + // fragments are live per MMAC burst, keeping the register footprint at the + // 4-blocks/CU occupancy point. DUFragment a_frag0, a_frag1; DUFragment @@ -352,55 +339,112 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( 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); + // Cooperative staging (identity input layouts, [N, K] packed output): + // * A[128,64] is 512 16-B chunks -> two per thread: rows [0,64) and + // [64,128). Each 16-B chunk is loaded from global as two 8-B halves + // and written to LDS as the adjacent 8-B pair (ds_write2_b64 class) + // because the 88-B A row stride is 8 (mod 16): odd rows start at an + // 8-B (not 16-B) aligned LDS offset, so a 16-B aligned LDS store would + // be misaligned. The 88-B stride (22 dwords/row) also maps the 16 rows + // of a fragment window onto 16 distinct word phases within each q + // phase, removing the r/r+8 alias of the old 80-B rows from the + // A-fragment reads. + // * B[64,64] of the packed [N, K] weight (packed[n][k]) is 256 int4s -> + // one int4 per thread; thread (tid>>2) selects the stage-local n row + // and (tid & 3)*16 the k offset, so global reads are 64-B contiguous + // per n row of the packed buffer and LDS stores are 16-B aligned along + // the padded n-major b_tile row (stride 80 unchanged). + const int row_a = tid >> 2; // 0..63 (first A row half) + const int col_a = (tid & 3) << 4; // 0,16,32,48 (8-B half at +8) + const int row_b = tid >> 2; // 0..63 (stage-local n row) + const int col_b = (tid & 3) << 4; // 0,16,32,48 (k offset) + + // Round-8 change: the cooperative staging below is software-pipelined one + // k64 stage ahead. Each thread keeps the 48 B/lane it stages for stage + // s+1 (two 16-B A row chunks + one 16-B B chunk = 12 payload VGPR) in + // registers and issues those global loads right after stage s's LDS + // stores/visibility barrier, so the full global round trip (~500+ cycles + // on gfx928) overlaps stage s's LDS fragment reads and MMACs instead of + // being exposed between the stage barriers and the LDS stores of the next + // stage (in the accepted single-buffered code the stage-s loads are issued + // at the top of the same stage, immediately followed by the vmcnt waits, + // so ~one full global latency is on the per-stage critical path of every + // wave, 16 times per block). The single LDS buffer layout, the exact + // staged bytes, the store order, the two __syncthreads per stage and the + // fragment-read/MMAC stream are unchanged - only the issue TIME of the + // global loads moves one stage earlier (their results wait in registers), + // so every int32 partial and every stored bf16 bit stays identical. The + // 12 payload VGPR take the code object past the exact 64-VGPR + // four-blocks/CU point to the 3-blocks/CU point (predicted arch_vgpr <= + // 80: 80 * 256 threads * 3 blocks = 61,440 <= 65,536 VGPR, LDS 16,384 B * + // 3 = 49,152 B): 12 resident waves per CU instead of 16, each no longer + // stalling per stage on the global round trip - the predicted occupancy + // arm of the experiment (no spill/scratch allowed; > 80 arch_vgpr would + // drop residency to 2 blocks/CU and falsify the mechanism). + const int64_t a_g0_base = + static_cast(m0 + row_a) * k + col_a; // A row 0..63, bytes + const int64_t a_g1_base = + static_cast(m0 + row_a + kBlockM / 2) * k + col_a; // row 64.. + const int64_t b_g_base = + static_cast(n0 + row_b) * k + col_b; // packed B n row, bytes + const int a_lds0 = row_a * kAStride + col_a; + const int a_lds1 = (row_a + kBlockM / 2) * kAStride + col_a; + const int b_lds = row_b * kBStride + col_b; + + // Loop-carried staging payload (48 B/lane): holds the stage-k0 chunks + // until they are stored to LDS at the top of the iteration, then is + // refilled with stage (k0 + kStageK) so the round trip overlaps the MMACs + // below. Same addresses and byte patterns as the accepted staging loads. + int k0 = 0; + unsigned long long a0_lo = 0, a0_hi = 0, a1_lo = 0, a1_hi = 0; + int4 b_chunk = make_int4(0, 0, 0, 0); + if (k0 < k) { // prologue: issue the stage-0 loads (only these are exposed) + a0_lo = *reinterpret_cast(x_q + a_g0_base); + a0_hi = *reinterpret_cast(x_q + a_g0_base + 8); + a1_lo = *reinterpret_cast(x_q + a_g1_base); + a1_hi = *reinterpret_cast(x_q + a_g1_base + 8); + b_chunk = *reinterpret_cast(weight + b_g_base); + } + for (; k0 < k; k0 += kStageK) { + // Barrier A (the accepted code's post-compute barrier, moved to the top + // of the next stage): every wave has finished reading the previous + // stage from the single LDS buffer, so the stores below may overwrite + // it. The stage-k0 global loads were issued one full stage earlier and + // landed during the previous stage's MMACs. __syncthreads(); + *reinterpret_cast(a_tile + a_lds0) = a0_lo; + *reinterpret_cast(a_tile + a_lds0 + 8) = a0_hi; + *reinterpret_cast(a_tile + a_lds1) = a1_lo; + *reinterpret_cast(a_tile + a_lds1 + 8) = a1_hi; + *reinterpret_cast(b_tile + b_lds) = b_chunk; + __syncthreads(); // Barrier B: stage k0 is visible to every wave + + if (k0 + kStageK < k) { // issue stage k0 + kStageK (overlaps the MMACs + const int64_t kn = k0 + kStageK; + a0_lo = + *reinterpret_cast(x_q + a_g0_base + kn); + a0_hi = *reinterpret_cast(x_q + a_g0_base + + kn + 8); + a1_lo = + *reinterpret_cast(x_q + a_g1_base + kn); + a1_hi = *reinterpret_cast(x_q + a_g1_base + + kn + 8); + b_chunk = *reinterpret_cast(weight + b_g_base + kn); + } - // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per 32-k + // chunk, operands filled by the direct 8-byte loads (A m-major stride + // kAStride, B n-major stride kBStride). The kk order (0 then 32), the + // per-accumulator MMAC order and the ascending-k int32 accumulation are + // unchanged, so every accumulator partial is bit-identical to the + // library-loader path. #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(b_frag1, + b_tile + (local_col + kTileN) * kBStride + kk, kBStride, + lane); load_fragment8(a_frag0, a_tile + local_row * kAStride + kk, kAStride, lane); load_fragment8(a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, @@ -420,45 +464,196 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( 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(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); + 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); + 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); + 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); + 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); + 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); + base_row + 3 * kTileM, base_col + kTileN, m, n, + lane); +} + +// Device helper for the fused split-K combine: load a lane's 4 int32 +// accumulator values as one aligned 16-B vector (LDS plane offsets are +// 16-B aligned: 1-KiB plane, 16-B lane stride). +__device__ __forceinline__ void load4_ints(const int32_t* src, + int32_t out[kDecodeLaneInts]) { + const int4 v = *reinterpret_cast(src); + out[0] = v.x; + out[1] = v.y; + out[2] = v.z; + out[3] = v.w; +} + +// --------------------------------------------------------------------------- +// Decode fused split-K=2 kernel (iteration 5, architecture round): one 16x16 +// N tile per block, grid 256 = 2.13 blocks/CU on the measured 120 CUs (the +// finest legal grid at the two-blocks-per-CU target; the 256 whole 16-col N +// tiles admit no finer whole-tile split), block 128 threads = two co-resident +// 64-lane wavefronts. Wave w of block b accumulates the k segment +// [w*512, (w+1)*512) of tile b in ascending k order - 16 m16n16k32 steps with +// the exact iteration-3 per-wave body: same lane mapping, direct global +// fragment loads (matrix_a row_major off x_q, matrix_b col_major off the +// packed [N, K] weight, so each lane's 8 B bytes are one contiguous 8-B +// vector load per k32 step instead of 8 scattered byte loads), depth-1 +// double-buffered fragment prefetch, no LDS and no barrier inside the K loop. +// The per-block split-K combine is fused into this single kernel: wave 0 +// publishes its int32 accumulator fragment (4 int32/lane = one 16-B vector, +// lane-contiguous, 1 KiB total) to a block-local shared plane, one END-of-K +// __syncthreads (executed unconditionally by both wavefronts after the +// symmetric K loop) makes it visible, and wave 1 reads the plane back, adds +// plane 0 + its own upper-K-half fragment with exact int32 arithmetic - the +// ascending-split sum is bit-identical to the unsplit ascending-k int32 dot, +// since K=1024 keeps every partial and the total far below int32 overflow - +// and applies the fused reference epilogue chain +// ((float)acc * x_scale[row]) * weight_scale[col] -> bf16 with the verified +// fragment->out mapping. One launch on the caller stream inside the timed +// Graph: the exact-shape decode is now a single Graph node with no caller +// workspace write, no workspace read and no second combine launch (iteration +// 3's workspace partial+combine pair and iteration 2's unsplit two-wave +// kernel are superseded). 1-KiB LDS/block does not dent occupancy, and the +// one END-of-K barrier is the accepted down_proj lineage pattern at the same +// 256-block / two-wave geometry (w8a8_dumma_m16n16k32_sk2_kernel). M=16 fills +// the m16 tile exactly so every row guard is trivially satisfied; every other +// (m, n, k) including the paired M=2 API shape stays on the scalar fallback. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDecodeThreads) void +w8a8_dumma_m16n16_sk2_fused_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 wave = + static_cast(threadIdx.x) >> 6; // k segment 0..kDecodeSplitK-1 + const int lane = + static_cast(threadIdx.x) & (kWaveSize - 1); // lane 0..63 + const int tile = static_cast(blockIdx.x); // one tile per block + const int n0 = tile * kTileN; + const int k0 = wave * kDecodeSegK; // segment base (0 or 512) + + // 1-KiB in-block int32 partial plane for the fused split-K combine: wave 0 + // publishes, wave 1 reads after the single END-of-K barrier. + __shared__ __align__(16) int32_t s_part[kDecodeTileInts]; + + DUFragment + a_frag[2]; + // matrix_b col_major: the packed weight is [N, K] (n-major rows of length + // k), so the library's col_major matrix_b loader addresses fragment slot + // (k, n) at packed[n * ldm + k]. Each lane's 8 fragment bytes (8 + // consecutive k of one n row) are then one contiguous, 8-B-aligned vector + // load per k32 step instead of 8 scattered byte loads with ldm = n over + // the identity [K, N] layout. + DUFragment + b_frag[2]; + DUFragment acc; + du_fill_fragment(acc, 0); + + // Epilogue scale prefetch (iteration 17): the fused combine/scale epilogue + // lives on the post-barrier critical tail of wave 1, where each lane + // performs 4 unique weight_scale loads + 1 x_scale load before its FMA + // chain. weight_scale is touched exactly once per kernel replay (each + // 16-col tile block reads its own 16 floats of the 4096-entry array), so + // those loads miss L1/L2 on the tail. Issue them before the K loop of the + // segment instead: identical values and the identical left-associative + // multiply chain keep the scaled-bf16 output bit-exact, while the global + // latency now hides under the 16 MMAC steps instead of stalling the + // END-of-K combine. Wave 0 never executes the branch and never reads the + // prefetched registers. + float xs; + float ws[kDecodeLaneInts]; + if (wave == 1) { + const int row = lane & 15; + const int col_mod4 = lane >> 4; + xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < kDecodeLaneInts; ++i) { + ws[i] = weight_scale[n0 + col_mod4 + 4 * i]; + } + } + + // Depth-1 software pipeline over the segment (unchanged iteration-3 body: + // the global loads for step s+1 are issued just before the mmac of step s, + // so one full global round trip overlaps each mmac; kDecodeSegSteps = 16 + // constant-trip iterations let the scheduler hoist loads while registers + // allow). + du_load_matrix_sync(a_frag[0], x_q + k0, kTargetK); + du_load_matrix_sync(b_frag[0], + weight + static_cast(n0) * kTargetK + k0, + kTargetK); + int buf = 0; +#pragma unroll + for (int s = 0; s < kDecodeSegSteps; ++s) { + const int nxt = buf ^ 1; + if (s + 1 < kDecodeSegSteps) { + du_load_matrix_sync(a_frag[nxt], x_q + k0 + (s + 1) * kTileK, + kTargetK); + du_load_matrix_sync(b_frag[nxt], + weight + static_cast(n0) * kTargetK + k0 + + (s + 1) * kTileK, + kTargetK); + } + du_mma_sync(acc, a_frag[buf], b_frag[buf], acc); + buf = nxt; + } + + if (wave == 0) { + // Publish the lower-K-half partial: lane-contiguous fragment order + // acc.x[0..3] = one 16-B vector at s_part[lane * 4] (1 KiB plane). + *reinterpret_cast(&s_part[lane * kDecodeLaneInts]) = + make_int4(acc.x[0], acc.x[1], acc.x[2], acc.x[3]); + } + // One END-of-K barrier: both wavefronts execute it unconditionally after + // their symmetric K loops; makes wave 0's plane visible to wave 1. + __syncthreads(); + + if (wave == 1) { + // Wave 1 owns the upper-K-half partial in `acc`; add the lower-K-half + // plane in ascending split order (plane 0 + plane 1 = full ascending-k + // int32 dot, bit-identical to the unsplit dot by construction). The + // epilogue scales were prefetched into xs/ws before the K loop above. + int32_t p0[kDecodeLaneInts]; + load4_ints(s_part + lane * kDecodeLaneInts, p0); + const int row = lane & 15; + if (row >= m) { + return; // tail-M guard (the exact shape always has m == 16) + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < kDecodeLaneInts; ++i) { + const int col = n0 + col_mod4 + 4 * i; + const float scaled = static_cast(p0[i] + acc.x[i]) * xs * ws[i]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); + } + } } // --------------------------------------------------------------------------- // 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. +// over K (k ascending), then fused x_scale * weight_scale, then bf16 store. +// Handles every unmatched (m, n, k), including the paired M=2 API shape +// (with the same (K, N) as the assigned decode/prefill shapes) and any M +// tail. Reads the packed [N, K] weight layout for every (k, n) - column n +// is one contiguous k-run at weight + n*k - exactly as launch_pack_w8a8_weight +// fills it. // --------------------------------------------------------------------------- __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( const int8_t* __restrict__ x_q, @@ -479,20 +674,11 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( 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]); - } + // Packed weight is [N, K] (packed[n][k]), so column n is one contiguous + // k-run starting at packed + n * k (element (k, n) at [n * k + k]). + 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]); } const float scaled = static_cast(acc) * x_scale[row] * weight_scale[col]; @@ -500,74 +686,26 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( } // --------------------------------------------------------------------------- -// 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. +// Bootstrap weight packing (outside the timed region and outside Graph +// capture): out-of-Graph device transpose raw [K, N] -> packed [N, K] +// (packed[n][k] = raw[k][n]) for every (k, n), plus an identity scale copy. +// The packed buffer keeps the same byte count (K*N) and the same allocated +// address as raw_weight, so the layout is graph-stable. Every GEMM consumer +// in this file (decode DUMMA, prefill DUMMA, scalar fallback) decodes the +// [N, K] layout; the scale buffer stays n-indexed. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(256) void w8a8_pack_o_proj_panels_kernel( +__global__ __launch_bounds__(256) void w8a8_pack_transpose_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, + int n, int64_t count) { const int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (linear < count) { - packed[linear] = raw[linear]; + const int64_t k_row = linear / n; + const int64_t n_col = linear - k_row * n; + packed[n_col * k + k_row] = raw[linear]; } } @@ -585,9 +723,9 @@ __global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( // --------------------------------------------------------------------------- // 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. +// Both launchers are pure dispatch on the caller-provided HIP stream: no +// allocation, no packing, no synchronization, no default-stream launch; they +// are CUDA/HIP-Graph safe. // --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, @@ -601,29 +739,41 @@ extern "C" void launch_w8a8_gemm( 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) { + // Explicit dispatch: + // * assigned decode shape (M=16, N=4096, K=1024, iteration 5): native + // INT8 DUMMA fused split-K=2 - one kernel, 256 blocks (one 16-col N + // tile per block) x 128 threads (wave w = ascending-k half w of the + // tile). Wave 0 publishes its int32 partial to a 1-KiB block-local + // shared plane, one END-of-K __syncthreads, wave 1 sums plane 0 + its + // own upper-K-half fragment (exact int32) and stores scaled bf16. A + // single launch inside the timed Graph at 2.13 blocks/CU; no caller + // workspace and no second combine kernel (workspace left untouched, so + // the iteration-3 workspace-fit fallback no longer exists). + // * assigned large-prefill shape (M=4096, N=4096, K=1024): the DUMMA + // 128x64 single-buffered path (workspace untouched). + // * every other (m, n, k) - including the paired M=2 API shape with the + // same (K, N) - takes the scalar fallback on the packed [N, K] weight + // layout (column n contiguous over k). + if (m == kDecodeM && n == kTargetN && k == kTargetK) { + const dim3 grid(kNumDecodeTiles); // 256 blocks, one tile per block + const dim3 block(kDecodeThreads); // 128 threads = 2 wavefronts + hipLaunchKernelGGL(w8a8_dumma_m16n16_sk2_fused_kernel, grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else 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, + hipLaunchKernelGGL(w8a8_dumma_prefill_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); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, m, n, + k); } (void)hipGetLastError(); } @@ -638,29 +788,21 @@ extern "C" void launch_pack_w8a8_weight( 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); - } + + // Out-of-timed device transpose of the raw [K, N] weight into the packed + // [N, K] buffer (same byte count, same address: graph-stable). All GEMM + // consumers in this file decode [N, K] (see launch_w8a8_gemm). + const int64_t weight_count = static_cast(k) * n; + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_transpose_kernel, weight_grid, block, 0, + stream, raw_weight, packed_weight, k, n, 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); + 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/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_down_proj.hip index 6c5d9375..f02ed9ff 100644 --- 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 @@ -1,65 +1,72 @@ -// @@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) +// @@variant shape=hy3_tp8_shared_down_proj_m4096 commit=303c7ee331a658a2b081df0b08f962d51259464c added=2026-09-06 +// median_us=89.46 p90_us=89.63 speedup=120.3 baseline_us=1.076e+04 +// source=hy3-tp8-dsh-m16-and-m4096-2-f7fa9503 +// MetaInfer W8A8 INT8 GEMM backend for Hygon K500SM_AI / gfx928 (worker_3). // -// 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. +// Assigned shapes (physical GPU 3): +// hy3_tp8_shared_down_proj_m16 M=16, N=4096, K=192 +// hy3_tp8_shared_down_proj_m4096 M=4096, N=4096, K=192 // -// 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. +// Logical operator (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] ) // -// 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. +// Pack layout: identity for every weight except the assigned TP8 down-proj +// weight ((k, n) == (192, 4096)), which is packed n-major +// packed[n*k + kk] == raw[kk*n + n_local] (each logical [K, N] column +// becomes a contiguous 192-byte run; same byte count, same buffer). +// launch_pack_w8a8_weight performs the transpose once, out of the timed +// region and Graph capture. Every consumer of the (k, n) == (192, 4096) +// packed buffer (the M=16 kernel below, the M >= 128 tiled DUMMA path, and +// the scalar fallback including the paired M=2 shape) decodes the matching +// n-major layout; all other (k, n) keep the identity pack. +// +// Strategy: +// * exact assigned shape (M=16, N=4096, K=192) -> native INT8 DUMMA +// m16n16k32 kernel with a bounded whole-window LDS prefetch: one +// 64-lane wavefront per block, one 16x16 output tile per block, the wave +// prefetches the full K=192 A/B operand window into LDS once with six +// 16-byte cooperative vector loads issued back to back before any LDS +// store (one __syncthreads; the prologue then pays one in-flight memory +// round trip per block instead of six serialized load->wait->write +// chains), then the 6 m16n16k32 steps consume +// A fragments row_major from a_tile and B fragments col_major from the +// packed n-major stage b_tile via du_load_matrix_sync (no vmem inside +// the K loop; every per-lane fragment byte run is one contiguous 8-byte +// LDS read), explicit int32 accumulation, direct fragment -> +// ((acc * x_scale) * weight_scale) -> bf16 epilogue, no barriers inside +// the K loop. +// * down-proj family (M >= 128, K == 192, N % kActiveBlockN == 0) -> the +// 2-D macro-tile DUMMA throughput kernel: a templated staged pipeline +// over a (kBlockM x kBlockN) output tile with the three macro-tiles of +// the family (64x64, 64x128, 128x64) compiled and a compile-time active +// selection (kActiveTileIndex, see below); the tile-1 64x128 path also +// carries a compile-time pipeline selector (kActivePipeline) between the +// double-buffered form (w8a8_dumma_prefill_kernel) and the round-3 +// single-buffered control (w8a8_dumma_prefill_sb_kernel), both staying +// compiled. K = 192 +// is pipelined in kPrefillStageK = 64 chunks (3 stages): the global +// int4 loads of stage s+1 issue before the MMAC sequence of stage s and +// land in the alternate LDS buffer right after it (one barrier per +// stage), so the vmcnt(0) wait and the ds_write overlap compute. The +// packed n-major weight ((k,n)==(192,4096), P[n*K+kk]) is staged +// n-major into LDS (b_tile[n][k], padded 72-byte rows) and consumed +// through col_major matrix_b fragments; A is staged row-major with +// padded 88-byte rows and consumed row_major. Both operand fragment +// paths use explicit 8-byte LDS reads (load_frag8 -> ds_read_b64, the +// 4-phase conflict-free floor at these even dword strides) instead of +// the library ds_read2_b32 lowering (2-way conflict in every phase at +// ANY even dword row stride). int32 accumulation stays resident over +// the whole K loop in four 16x16 accumulator fragments per wave and the +// epilogue is a direct fragment -> ((acc * x_scale) * weight_scale) -> +// bf16 store. No split-K, no raw asm, no speculative pipeline. +// * Every other shape (including the paired M=2 API shapes with the same +// (N, K)) -> generic scalar int8/int32 fallback. The M=16 kernel keeps +// an exact (m, n, k) == (16, 4096, 192) guard so paired M=2 API shapes +// and unmatched shapes still reach the scalar fallback. +// +// Headers are included in the mandated order for this DTK: HIP runtime first +// (du_mma.h is not self-contained), then hip_bfloat16.h, then du_mma.h. #include #include @@ -69,131 +76,165 @@ 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: +constexpr int kInt8TileM = 16; // DUMMA m16n16k32 (int8 -> int32) tiles +constexpr int kInt8TileN = 16; +constexpr int kInt8TileK = 32; + +constexpr int kWaveSize = 64; // gfx928 native wavefront size + +// Large-M tiled kernel geometry (K=192 = 3 x kPrefillStageK stages). +// +// Compile-time selection of the active 2-D macro-tile of the prefill +// family: // 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; +// All three instantiations of the templated kernel are compiled so the +// family stays measurable; only the active one is launched inside the timed +// region. The active tile is 64x128: with the (kBlockM x kBlockN) tile, +// 8 wavefronts (512 threads) and the double-buffered stage-64 pipeline the +// per-block LDS footprints are +// 64x64 : 2*(64*88 + 64*72) = 20,480 B (256 threads / 4 waves) +// 64x128 : 2*(64*88 + 128*72) = 29,696 B (512 threads / 8 waves) +// 128x64 : 2*(128*88 + 64*72) = 31,744 B (512 threads / 8 waves) +// and 2 blocks/CU stay LDS-resident for all three (40,960 / 59,392 / 63,488 +// B per CU <= 64 KiB), i.e. 16 resident wavefronts per CU for the 512-thread +// tiles. For the assigned M=4096 x N=4096 grid the 64x128 tile yields +// grid (32, 64) = 2048 blocks and re-reads each A 64-row tile across 32 +// N-blocks and each B 128-column tile across 64 M-blocks (both operand +// streams stay L2-resident). +constexpr int kPrefillStageK = 64; // K depth staged per barrier pair +constexpr int kAssignedK = 192; // the only K the optimized guard allows +constexpr int kActiveTileIndex = 1; // 0 -> 64x64, 1 -> 64x128, 2 -> 128x64 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). +// Pipeline form of the tile-1 (64x128) prefill path, compared in round 3: +// 0 -> single-buffered control (w8a8_dumma_prefill_sb_kernel, ACTIVE this +// round): LDS 14,848 B/block, prologue + 2 barriers/stage boundary = +// 5 barriers/block for K=192, one-stage-ahead VGPR prefetch kept, +// occupancy unchanged at 2 blocks/CU (VGPR-bound); +// 1 -> accepted double-buffered pipeline (w8a8_dumma_prefill_kernel, +// retained and still compiled): LDS 29,696 B/block, one barrier per +// stage (4/block incl. its dead last-stage barrier). +// Tiles 0 and 2 always use the double-buffered form (compile-time retention +// instantiations; only tile 1 is launched for the assigned M=4096 grid). +constexpr int kActivePipeline = 0; + +// LDS row padding of the staged A tile (row-major, one row per output 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: row stride kAStride = 88 B = 22 dwords, and {22*r mod 32} +// for r = 0..15 is a permutation of the 16 even bank-pairs, so an 8-byte +// ds_read_b64 fragment read runs at the conflict-free 4-phase floor. +// Stride 20 dwords (pad 16) aliases rows r and r+8 onto identical bank +// pairs at any even dword stride, which is exactly why the ds_read2_b32 +// lowering doubles every phase. +constexpr int kLdsPadA = 24; +constexpr int kAStride = kPrefillStageK + kLdsPadA; // 88, % 8 == 0 + +// LDS row padding of the transposed (n-major) B tile: one row per output +// column, kPrefillStageK contiguous K bytes + 8 B pad -> stride 72 +// (% 8 == 0 for the col_major 8-byte fragment reads and the two-half +// ds_write_b64 B staging; % 16 == 8, so the staging store is the explicit +// two-int64 form, never a 16-byte-misaligned int4 store). {18*r mod 32} +// for r = 0..15 is likewise a permutation of the 16 even bank-pairs. constexpr int kLdsPadB = 8; -constexpr int kBTStride = kPrefillStageK + kLdsPadB; +constexpr int kBTStride = kPrefillStageK + kLdsPadB; // 72, % 8 == 0 + +constexpr int kScalarThreads = 256; // scalar fallback threads per block +constexpr int kPackThreads = 256; // pack copy threads per block 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. +// float32 -> bfloat16 with round-to-nearest-even, stored as raw uint16 bits. +// The output tensor is at::kBFloat16 (2 bytes/element); this conversion is +// bit-identical to the hardware RNE cvt used by the reference path for all +// finite inputs and keeps the store independent of bf16 header type details. +__device__ __forceinline__ uint16_t bf16_bits_from_f32(float f) { + uint32_t bits; + __builtin_memcpy(&bits, &f, sizeof(bits)); + const uint32_t rounding_bias = 0x7FFFu + ((bits >> 16) & 1u); + return static_cast((bits + rounding_bias) >> 16); +} + +// Store one 16x16 int32 accumulator fragment scaled to bf16 with the three +// tail-gated weight_scale operands (wv1..wv3) arriving as already-loaded +// registers. +// +// Exact-source ISA of the accepted kernel (round-6 code object, +// compile_cache_key 1df09b96..., bundle +// profiles/.../iteration7/current-best-isa/isa.txt): the compiler already +// hoists x_scale[row] into the staging prologue (global_load_dword at 0x4710, +// before the first s_waitcnt vmcnt) and issues wv0 immediately after the last +// ds_write_b128 (0x4780), but wv1..wv3 are only issued between v_mmac 1..3 +// (0x4844 / 0x484C / 0x4888) and are still in flight when the epilogue +// consumes them: real s_waitcnt vmcnt waits (0x4ACC vmcnt(4), 0x4AD4 vmcnt(3), +// 0x4B20 / 0x4B48 / 0x4B6C vmcnt(3)) sit between the four +// global_store_short_d16_hi, i.e. the per-block store tail can serialize on +// ~L2-latency arrivals of per-lane weight_scale dwords issued only ~3 v_mmacs +// earlier. The exact-shape kernel below therefore loads wv1..wv3 in its +// staging batch (same three vmem read instructions as the generic epilogue, +// moved to the front of the wave's memory FIFO so they land during the +// staging round trip plus the six-step K chain) and passes them into this +// helper, which converts, scales and stores each element purely from +// registers. Math stays byte-identical to the generic fragment stores +// (store_prefill_fragment / the reference formula): row = lane & 15 is +// always < m == 16, and each element keeps the same left-associative +// ((float)acc_i * xv) * wv_i fp32 evaluation order with the same RNE bf16 +// rounding; xv and wv0 keep their proven round-6 issue points (loaded here, +// where the round-6 compiler scheduled them), so the delta is exactly the +// three loads whose issue the round-6 code object placed inside the mmac +// chain. template -__device__ __forceinline__ void store_prefill_fragment( +__device__ __forceinline__ void store_acc_fragment_tail_preloaded( const AccFragment& frag, const float* __restrict__ x_scale, const float* __restrict__ weight_scale, - __hip_bfloat16* __restrict__ out, - int m, + float wv1, + float wv2, + float wv3, + uint16_t* __restrict__ out, int n, - int base_row, int base_col, int lane) { - const int row = base_row + (lane & 15); - if (row >= m) { - return; - } + const int row = lane & 15; // m == 16 guard: always in range 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); - } + const float xv = x_scale[row]; + const float wv0 = weight_scale[base_col + col_mod4]; + const int64_t out_base = + static_cast(row) * n + base_col + col_mod4; + // Elements i = 0..3 sit at columns base_col + col_mod4 + 4*i; each keeps the + // same left-associative ((float)acc_i * xv) * wv_i evaluation order and RNE + // bf16 rounding as the generic fragment stores. + const float s0 = static_cast(frag.x[0]) * xv; + out[out_base] = bf16_bits_from_f32(s0 * wv0); + const float s1 = static_cast(frag.x[1]) * xv; + out[out_base + 4] = bf16_bits_from_f32(s1 * wv1); + const float s2 = static_cast(frag.x[2]) * xv; + out[out_base + 8] = bf16_bits_from_f32(s2 * wv2); + const float s3 = static_cast(frag.x[3]) * xv; + out[out_base + 12] = bf16_bits_from_f32(s3 * wv3); } // --------------------------------------------------------------------------- -// 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. +// Explicit 8-byte fragment loader for the int8 m16k32 operand fragments. +// +// Both library loaders (row_major matrix_a, col_major matrix_b) place the +// same 8 contiguous bytes per lane into the same x[0..7] slots - lane l -> +// row (l & 15), k-quarter (l >> 4) - 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 one k-quarter with rows 0-15 of the +// other) into one 32-bank cycle, so the two k-quarter 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 - a 2-way conflict in +// every phase (8 phases per read) at ANY even dword row stride. 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, so fragment register contents (and therefore +// the exact int32 k-ascending MMAC accumulation) are identical to the +// library loaders. // --------------------------------------------------------------------------- template __device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, @@ -205,54 +246,124 @@ __device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, reinterpret_cast(f.x)[0] = v; } +// Store one 16x16 int32 accumulator fragment scaled to bf16 (prefill +// epilogue). gfx928 int8 m16n16k32 accumulator ownership (validated +// against du_store_matrix_sync in this DTK): lane&15 selects the row, +// lane>>4 selects col%4, and frag.x[i] holds the columns col%4 + 4*i. +// Per element the fp32 scaling keeps the same left-associative +// ((float)acc_i * x_scale) * weight_scale order and the same RNE bf16 +// rounding as the reference path and as the M=16 kernel. +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + uint16_t* __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; // tail-M rows: never touch x_scale/out out of bounds + } + const int col_mod4 = lane >> 4; + const float xs = x_scale[row]; + const int64_t out_base = + static_cast(row) * n + base_col + col_mod4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float scaled = static_cast(frag.x[i]) * xs * + weight_scale[base_col + col_mod4 + 4 * i]; + out[out_base + 4 * i] = bf16_bits_from_f32(scaled); + } +} + +// Round-6 fused-epilogue transpose write, iteration-14 register-batched +// form: scale one 16x16 accumulator fragment with its row scale (xv) and +// its four column scales (wv0..wv3, arriving as already-loaded registers) +// and store the RNE bf16 result into the block's row-major out_tile at the +// element's final (row, col) position. The caller reads the four column +// scales of each 16-column fragment base (tile_col0 and +// tile_col0 + kInt8TileN) ONCE per wave and shares them between the two +// fragments of that base (acc00/acc10 and acc01/acc11), so the per-wave +// epilogue issues 8 ds_read_b32 instead of the round-8 form's 16 in-loop +// reads and no LDS load sits between the last v_mmac and the element store +// tail. Per element the fp32 math keeps the same left-associative +// ((float)acc_i * xv) * wv_i evaluation order and the same RNE bf16 +// rounding as store_prefill_fragment, and every wv_i is the bit-identical +// float the in-loop form read (same buffer s_wscale, same index +// base + c4 + 4*i, no write to s_wscale between the prologue landing and +// the epilogue), so the bytes this lane publishes are identical; only the +// destination differs from a scattered 2-byte global store (an LDS +// transpose slot that the coalesced sweep below reads back). +template +__device__ __forceinline__ void transpose_bf16_fragment_lane( + const AccFragment& frag, float xv, int tile_row0, int tile_col0, + int r_lane, int c4, float wv0, float wv1, float wv2, float wv3, + uint16_t* __restrict__ out_tile, int out_stride) { + const float s0 = static_cast(frag.x[0]) * xv; + out_tile[(tile_row0 + r_lane) * out_stride + tile_col0 + c4] = + bf16_bits_from_f32(s0 * wv0); + const float s1 = static_cast(frag.x[1]) * xv; + out_tile[(tile_row0 + r_lane) * out_stride + tile_col0 + c4 + 4] = + bf16_bits_from_f32(s1 * wv1); + const float s2 = static_cast(frag.x[2]) * xv; + out_tile[(tile_row0 + r_lane) * out_stride + tile_col0 + c4 + 8] = + bf16_bits_from_f32(s2 * wv2); + const float s3 = static_cast(frag.x[3]) * xv; + out_tile[(tile_row0 + r_lane) * out_stride + tile_col0 + c4 + 12] = + bf16_bits_from_f32(s3 * wv3); +} + // --------------------------------------------------------------------------- -// 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. +// Large-M prefill path (down-proj family: M >= 128, K == kAssignedK, N a +// multiple of the active macro-tile N). One block computes a kBlockM x +// kBlockN output tile with (kBlockM/32) x (kBlockN/32) wavefronts (one +// 32x32 quadrant per wave, four 16x16 int32 accumulators kept resident +// over the whole K loop). A and the packed n-major B are cooperatively +// staged into a double-buffered, bank-padded LDS in kPrefillStageK = 64 +// chunks with VGPR 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 the ds_write overlap +// compute (one barrier per stage; prologue + 3 stages = 4 barriers per +// block for K = 192). A stays row-major in LDS (stride kAStride = 88); B +// is n-major at pack time (packed[n][k], one contiguous K run per output +// column) and is staged n-major (stride kBTStride = 72, K-contiguous), +// consumed through col_major matrix_b fragments; both operand paths use +// the explicit 8-byte load_frag8 ds_read_b64 reads above. Each stage is +// consumed as kStageK = 2 x kInt8TileK fully unrolled steps into two +// independent fragment register sets, so the second step's LDS latency +// overlaps the first step's MMACs. The epilogue is the fused direct +// fragment -> ((acc * x_scale) * weight_scale) -> bf16 store. // --------------------------------------------------------------------------- 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 int8_t* __restrict__ a, + const int8_t* __restrict__ b, const float* __restrict__ x_scale, const float* __restrict__ weight_scale, - __hip_bfloat16* __restrict__ out, + uint16_t* __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, + static_assert(kStageK % kInt8TileK == 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"); + static_assert(kStageK == 2 * kInt8TileK, + "unrolled kk body assumes kStageK == 2 * kInt8TileK"); const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; - const int lane = tid % kWaveSize; + const int lane = tid - wave * kWaveSize; const int wave_row = wave / kWavesN; const int wave_col = wave - wave_row * kWavesN; const int m0 = static_cast(blockIdx.y) * kBlockM; @@ -268,28 +379,30 @@ w8a8_dumma_prefill_kernel( 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. + // Transposed (n-major) B tile: logical B[k0+kk][n0+n_local] lives at + // b_tile[n_local][kk], so col_major fragment loads read 8 contiguous K + // bytes per lane (one ds_read_b64). The staged runs are written from the + // n-major packed weight as two 8-byte halves per thread. __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBTStride]; - DUFragment + DUFragment a_frag0, a_frag1, a_frag0_1, a_frag1_1; - DUFragment + DUFragment b_frag0, b_frag1, b_frag0_1, b_frag1_1; - DUFragment + 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 + // 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. + // staging stores are the two-half int64 form. int4 a_pref{0, 0, 0, 0}; int4 b_pref{0, 0, 0, 0}; { @@ -299,21 +412,21 @@ w8a8_dumma_prefill_kernel( const int global_row = m0 + row; if (global_row < m) { a_pref = *reinterpret_cast( - x_q + global_row * k + seg * sizeof(int4)); + a + static_cast(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); + b + (static_cast(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. + // lowers to one ds_write2_b64 per thread. const int64_t* src64 = reinterpret_cast(&a_pref); int64_t* dst64 = reinterpret_cast( a_tile[0] + row * kAStride + seg * sizeof(int4)); @@ -337,7 +450,7 @@ w8a8_dumma_prefill_kernel( 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 + // 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; @@ -347,7 +460,8 @@ w8a8_dumma_prefill_kernel( const int global_row = m0 + row; if (global_row < m) { a_pref = *reinterpret_cast( - x_q + global_row * k + k1 + seg * sizeof(int4)); + a + static_cast(global_row) * k + k1 + + seg * sizeof(int4)); } else { a_pref = int4{0, 0, 0, 0}; } @@ -356,41 +470,35 @@ w8a8_dumma_prefill_kernel( 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); + b + (static_cast(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 + // per wave per kInt8TileK step. kStageK = 64 = 2 kInt8TileK 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, + // Second 32-k step (kk = kInt8TileK): independent registers, issued + // while the first step's data is still in flight. + load_frag8(a_frag0_1, a_tile[buf] + local_row * kAStride + kInt8TileK, 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, + a_tile[buf] + (local_row + 16) * kAStride + kInt8TileK, + kAStride); + // col_major: element B[k0+kk+k][n0+n] lives at b_tile[n][kk+k], so the + // loader's p[row*ldm + k + i] reads 8 contiguous bytes per lane. + load_frag8(b_frag0_1, b_tile[buf] + local_col * kBTStride + kInt8TileK, kBTStride); load_frag8(b_frag1_1, - b_tile[buf] + (local_col + 16) * kBTStride + kTileK, + b_tile[buf] + (local_col + 16) * kBTStride + kInt8TileK, kBTStride); du_mma_sync(acc00, a_frag0, b_frag0, acc00); du_mma_sync(acc01, a_frag0, b_frag1, acc01); @@ -401,7 +509,7 @@ w8a8_dumma_prefill_kernel( 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 + // 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. @@ -427,121 +535,707 @@ w8a8_dumma_prefill_kernel( } // 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. + // parity-flipped buffer is overwritten two stages later. Every thread + // reaches 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 + 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); + 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 + kInt8TileN, lane); + store_prefill_fragment(acc10, x_scale, weight_scale, out, m, n, + base_row + kInt8TileM, base_col, lane); + store_prefill_fragment(acc11, x_scale, weight_scale, out, m, n, + base_row + kInt8TileM, base_col + kInt8TileN, 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. +// Round-3 pipeline control of the same down-proj family tile: the SINGLE- +// buffered staged pipeline (one LDS buffer per operand instead of the +// parity pair of w8a8_dumma_prefill_kernel above), kept bit-identical and +// structurally identical otherwise: same (kBlockM x kBlockN) tile and wave +// -> quadrant mapping, same one-stage-ahead int4 VGPR prefetch (stage s+1 +// global loads issued before the stage-s fragment loads/MMACs), same +// 8-byte load_frag8 fragment consumption with the identical k-ascending +// int32 MMAC order. The single buffer is overwritten in place, so every +// stage boundary needs a WAR __syncthreads (retire the stage-s fragment +// reads) before the ds_write of stage s+1 and a publish __syncthreads after +// it: prologue + 2 x (kStages - 1) = 5 barriers per block for K = 192 vs +// the double-buffered prologue + kStages = 4 (of which the last-stage loop +// barrier is dead); the final stage's MMACs run straight into the epilogue. +// The SB epilogue is the round-6 FUSED form (scale windows staged in LDS in +// the prologue, per-fragment RNE bf16 transpose through a padded out_tile, +// one publish barrier, fully-coalesced 8-byte sweep stores) - same per- +// element fp32 math and RNE bf16 bytes as the direct round-3/4 fragment +// stores of the DB arm, only the store form differs. Static LDS grows to +// 32,512 B/block for 64x128 (5,632 A + 9,216 B + 256 s_xscale + 512 +// s_wscale + 16,896 out_tile; see the in-body comment), so residency stays +// 2 blocks/CU = 16 waves/CU (2 x 32,512 = 65,024 <= 65,536; the kernel +// stays VGPR-bound at ~56-64 vgpr x 512 threads x 2 blocks <= 65,536 +// VGPRs), which keeps occupancy from confounding the pipeline comparison. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(256) void w8a8_scalar_fallback_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, +template +__global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize, 2) void +w8a8_dumma_prefill_sb_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, + uint16_t* __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) { + int k) { + constexpr int kWavesN = kBlockN / 32; + 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 % kInt8TileK == 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 * kInt8TileK, + "unrolled kk body assumes kStageK == 2 * kInt8TileK"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid - wave * 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"); + // No runtime stage count: the launch guard fixes k == kAssignedK == 192, + // so the kernel body below is the compile-time 3-stage form (k is still + // read at runtime inside run_stage for the global A/B row stride of the + // stage-prefetch addresses; its value is always 192 under the guard). + + // Single-buffered stage: A rows on the padded 88-B stride and the packed + // n-major B rows on the padded 72-B stride, exactly as the double-buffered + // kernel, minus the parity [2] dimension. + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kBTStride]; + + // Round-6 fused-epilogue staging: + // * s_xscale/s_wscale - the block's per-row x_scale[kBlockM] and + // per-column weight_scale[kBlockN] windows, copied cooperatively into + // LDS inside the stage-0 staging batch (192 dword loads issued with the + // operand vectors, latency hidden under the prologue barrier and the + // whole K chain). The epilogue then scales and converts every + // accumulator element from LDS-resident scales: no global load and no + // vmcnt wait can sit between the last v_mmac and the final store + // (the round-5 exact-source ISA shows the per-fragment scale dword + // loads issued only ~3 v_mmacs before the store tail, with real + // s_waitcnt vmcnt(n) waits interleaved between the four + // global_store_short_d16_hi groups of each fragment). + // * out_tile - the block's (kBlockM x kBlockN) bf16 output tile in + // row-major layout with row stride kOutStride = 132 elements (264 B, + // % 16 == 8 -> every row start stays 8-byte aligned for the coalesced + // ds_read_b64 sweep, while the +4-element pad spreads the 16 fragment + // rows over 16 distinct LDS bank words so the fragment-layout 2-byte + // transpose writes run at the ~1-phase floor). Each wave converts its + // four fragments to RNE bf16 and transposes them through out_tile; one + // extra barrier later every thread stores one contiguous 32-B run of + // one output row as four 8-byte vector stores (fully coalesced lines, + // ~8x fewer L2 write line-touches than the old per-fragment + // global_store_short_d16_hi tail, which scattered each 64-lane store + // instruction over 16 rows' 8-B segments at the 8,192-B row stride). + // Static LDS stays 32,512 B/block (5,632 A + 9,216 B + 256 + 512 scales + + // 16,896 out), so 2 blocks/CU = 16 waves/CU remain resident + // (2 x 32,512 = 65,024 <= 65,536); occupancy remains VGPR-bound. + __shared__ __align__(16) float s_xscale[kBlockM]; + __shared__ __align__(16) float s_wscale[kBlockN]; + constexpr int kOutStride = kBlockN + 4; // 132 bf16 per row (264 B), % 8 == 0 + static_assert((kOutStride * 2) % 8 == 0, + "out_tile rows must stay 8-byte aligned for the b64 sweep"); + __shared__ __align__(16) uint16_t out_tile[kBlockM * kOutStride]; + + 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); + + int4 a_pref{0, 0, 0, 0}; + int4 b_pref{0, 0, 0, 0}; + + // Prologue: fetch stage 0 into VGPRs, land it in the single buffer with + // the two-half int64 form (kAStride/kBTStride are 8 mod 16, so a single + // 16-byte store would be misaligned on odd rows), publish with one + // barrier. Identical source bytes and store addresses as stage 0 of the + // double-buffered kernel. + { + // Round-6 fused-epilogue staging: the block's x_scale[kBlockM] and + // weight_scale[kBlockN] windows are loaded here, at the FRONT of the + // prologue batch, so their vmcnt results land while the A/B staging + // stores and the publish barrier below drain (the fused epilogue then + // scales/converts every accumulator element purely from LDS: no global + // load and no vmcnt wait can sit between the last v_mmac and the final + // store, replacing the round-3/4 tail that serialized the four + // store_short_d16_hi groups on per-lane weight_scale dword arrivals). + float xscale_v = 0.0f; + float wscale_v = 0.0f; + if (tid < kBlockM) { + const int g_row = m0 + tid; + if (g_row < m) { + xscale_v = x_scale[g_row]; // tail-M rows keep 0, never used + } + } + if (tid < kBlockN) { + wscale_v = weight_scale[n0 + tid]; + } + 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( + a + static_cast(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( + b + (static_cast(n0) + n_local) * k + kchunk); + } + 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 + 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 + n_local * kBTStride + kchunk); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + // Land the staged scale windows (loaded at the front of this batch, + // above) so the prologue publish barrier makes them visible to every + // wave's fused epilogue. + if (tid < kBlockM) { + s_xscale[tid] = xscale_v; + } + if (tid < kBlockN) { + s_wscale[tid] = wscale_v; + } + __syncthreads(); + } + + // Iteration-15 exact-K stage unroll: the launch guard fixes + // k == kAssignedK == 192 == 3 x kPrefillStageK for every launch of this + // kernel, so the three stage passes below are written out at compile time + // for that exact K. Each stage keeps the byte-identical round-3/14 + // structure - stage s+1 global int4 prefetch issued before the stage-s + // fragment loads/MMACs, WAR __syncthreads, in-place landing of the + // prefetched stage, publish __syncthreads at the first two boundaries, + // straight into the epilogue after the last stage - with the exact int32 + // k-ascending MMAC order, operand staging bytes, fragment slots, the 6 + // barriers per block and the a_pref/b_pref VGPR prefetch registers reused + // per stage unchanged, so the published bytes are bit-identical. The + // delta vs the rolled loop is only the per-stage control the round-8 code + // object carries (trip-count increment s19, k offset s22, per-pass + // compare/branch and the loop back-edge): with the count compile-time + // known, all of it disappears and the stage-2 MMACs flow straight into the + // fused epilogue. + // + // Iteration-16 exact-K specialization (terminal step of the round-15 + // unroll): the runtime dispatch below ran run_stage(0..2) for the guarded + // K and kept a rolled fallback loop for any other k. That fallback can + // never execute - the host guard use_tiled_dumma requires + // k == kAssignedK == 192 for every launch of this kernel - yet the timed + // code object still carries its full extra stage copy (8 du_mma_sync, 8 + // ds_read_b64 fragment loads, 2 ds_write2_b64 landings, 2 s_barrier and + // the loop control: v_mmac 32 total vs the 24 the executed path runs, + // 8 s_barrier in .text vs the 6 executed) plus the uniform + // s_cmp_eq_u32/s_cbranch on k == kAssignedK, and the register allocator + // still has to allocate one merged CFG containing the loop's live ranges + // and phi copies. This round removes the dispatch and the dead rolled + // copy: the three compile-time stages below run unconditionally (k is + // still read at runtime for the global A/B row stride k of the stage + // prefetch addresses, always 192 under the guard). The executed path - + // stage bodies, prefetch/land pairs, the 6 barriers, the fused epilogue + // and the sweep - is source-identical, so the published bytes stay + // bit-identical; the code object loses only dead code, dead barriers and + // the k compare, and the allocator's problem loses the fallback merge. + // Round-15's linearization raised the object from 57 to 65 VGPRs (sgpr + // 28 -> 32) at the 64-VGPR 2-blocks/CU register boundary every prior + // guard priced; the falsifiable question of this round is how much of + // that +8 was the never-taken fallback path's register/phi cost (then the + // specialized object returns to <= 64 VGPRs) versus inherent ILP of the + // straight-line three-stage schedule (then it stays at 65 and this round + // is the honest null at parity). + const auto run_stage = [&](int s, bool has_next) { + // Issue the global loads for stage s+1 now, at the same point as the + // double-buffered kernel, so their latency hides behind the fragment + // loads and MMACs below. Out-of-range A rows are zero-filled (tail-M + // guard), byte-identical to the rolled zero fill. + if (has_next) { + const int k1 = (s + 1) * 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( + a + static_cast(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( + b + (static_cast(n0) + n_local) * k + k1 + kchunk); + } + } + + // Consume stage s from the single buffer: identical fragment loads and + // MMACs to the rolled loop (same padded strides, same element-to-slot + // mapping, same k-ascending int32 order). The second kInt8TileK step + // uses an independent register set so its LDS latency overlaps the first + // step's MMACs. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + load_frag8(a_frag0, a_tile + local_row * kAStride, kAStride); + load_frag8(a_frag1, a_tile + (local_row + 16) * kAStride, kAStride); + load_frag8(b_frag0, b_tile + local_col * kBTStride, kBTStride); + load_frag8(b_frag1, b_tile + (local_col + 16) * kBTStride, kBTStride); + load_frag8(a_frag0_1, a_tile + local_row * kAStride + kInt8TileK, + kAStride); + load_frag8(a_frag1_1, + a_tile + (local_row + 16) * kAStride + kInt8TileK, kAStride); + load_frag8(b_frag0_1, b_tile + local_col * kBTStride + kInt8TileK, + kBTStride); + load_frag8(b_frag1_1, + b_tile + (local_col + 16) * kBTStride + kInt8TileK, + 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); + + if (has_next) { + // WAR barrier: every stage-s fragment read in the whole block must + // retire before the single buffer is overwritten in place with stage + // s+1. + __syncthreads(); + // Land the prefetched stage s+1. The compiler inserts the vmcnt(0) + // wait here via the a_pref/b_pref data dependency, i.e. after the + // MMACs and the WAR barrier, never before the stage-s burst. + 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 + 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 + n_local * kBTStride + kchunk); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + // Publish: make the stage s+1 writes visible before the next stage's + // fragment reads. + __syncthreads(); + } + }; + // Iteration-16 exact-K specialization: the launch guard fixes + // k == kAssignedK == 192 == 3 x kPrefillStageK for every launch of this + // kernel, so the three compile-time stages run unconditionally below + // (the round-15 rolled fallback for any other k, which could never + // execute and whose dead stage copy, dead barriers and loop control + // stayed in the timed code object, is removed). + run_stage(0, true); + run_stage(1, true); + run_stage(2, false); + + // Round-6 fused epilogue (SB arm only): every accumulator element is + // scaled with the prologue-staged LDS windows (s_xscale/s_wscale) and + // converted to RNE bf16 with the identical left-associative fp32 math as + // the direct fragment stores of rounds 3/4, but each element is transposed + // into the block's padded row-major out_tile instead of a scattered + // 2-byte global store. One extra barrier later a fully-coalesced sweep + // stores the whole tile: each thread copies one contiguous 32-B run of one + // output row to global as four 8-byte vector stores, so the per-block + // epilogue global-store tail collapses from 256 scattered + // store_short_d16_hi instructions (each 64-lane instruction touching 16 + // rows' 8-B segments at the 8,192-B row stride, with vmcnt waits for + // in-flight per-lane scale dwords between the four groups) into 32 fully + // coalesced 8-byte store instructions with no global-load dependency at + // all. Every published byte is bit-identical to the round-3/4 store form. + const int tile_row0 = wave_row * 32; // n0/m0-relative quadrant origin + const int tile_col0 = wave_col * 32; + const int r_lane = lane & 15; + const int c4 = lane >> 4; + // Iteration-14 register-batched column scales: the four s_wscale dwords + // of each 16-column fragment base are read ONCE per wave, back to back, + // and shared by the two fragments of that base (acc00/acc10 use base + // tile_col0, acc01/acc11 use base tile_col0 + kInt8TileN). The round-8 + // form re-read them inside every fragment call's store loop (16 + // ds_read_b32/wave with one load->use dependency per element); this form + // issues 8 ds_read_b32/wave ahead of the element math. Every read is + // unguarded: the index stays inside s_wscale[kBlockN] for every lane + // (tile_col0 + kInt8TileN + c4 + 12 <= 96 + 16 + 15 <= 127) and tail-M + // blocks simply discard the values under the row guards below, so the + // published bytes stay bit-identical to the round-8 form. + const int g_row0 = m0 + tile_row0 + r_lane; // rows of acc00/acc01 + const int g_row1 = m0 + tile_row0 + kInt8TileM + r_lane; // acc10/acc11 + const float wv00 = s_wscale[tile_col0 + c4]; + const float wv01 = s_wscale[tile_col0 + c4 + 4]; + const float wv02 = s_wscale[tile_col0 + c4 + 8]; + const float wv03 = s_wscale[tile_col0 + c4 + 12]; + const float xv0 = s_xscale[tile_row0 + r_lane]; + if (g_row0 < m) { // row tail guard, as in the round-3/4 fragment stores + transpose_bf16_fragment_lane(acc00, xv0, tile_row0, tile_col0, r_lane, + c4, wv00, wv01, wv02, wv03, out_tile, + kOutStride); + } + const float xv1 = s_xscale[tile_row0 + kInt8TileM + r_lane]; + if (g_row1 < m) { + transpose_bf16_fragment_lane(acc10, xv1, tile_row0 + kInt8TileM, + tile_col0, r_lane, c4, wv00, wv01, wv02, + wv03, out_tile, kOutStride); + } + // The tile_col0 batch is consumed; only now read the second 4-dword batch + // (column base tile_col0 + kInt8TileN, shared by acc11/acc01) so at most + // four scale registers are live at once in the tail. + const float wv10 = s_wscale[tile_col0 + kInt8TileN + c4]; + const float wv11 = s_wscale[tile_col0 + kInt8TileN + c4 + 4]; + const float wv12 = s_wscale[tile_col0 + kInt8TileN + c4 + 8]; + const float wv13 = s_wscale[tile_col0 + kInt8TileN + c4 + 12]; + if (g_row1 < m) { + transpose_bf16_fragment_lane(acc11, xv1, tile_row0 + kInt8TileM, + tile_col0 + kInt8TileN, r_lane, c4, wv10, + wv11, wv12, wv13, out_tile, kOutStride); + } + if (g_row0 < m) { + transpose_bf16_fragment_lane(acc01, xv0, tile_row0, + tile_col0 + kInt8TileN, r_lane, c4, wv10, + wv11, wv12, wv13, out_tile, kOutStride); + } + // Publish: every wave's out_tile transpose writes must be visible to the + // coalesced sweep below (the 6th barrier of the block). + __syncthreads(); + // Coalesced sweep: thread tid stores one contiguous 32-B run (16 bf16 = + // four 8-byte vectors) of one output row. kThreadsPerRow consecutive + // threads cover one row, so every 8-byte store instruction in the sweep + // is fully coalesced. Global row stride n is a multiple of kBlockN by the + // launch guard and c is a multiple of 16 elements, so both the LDS source + // (row stride kOutStride * 2 B, 8-byte aligned by the static_assert above) + // and the global destination stay 8-byte aligned. Tail-M rows are skipped + // (their out_tile slots were never transposed). + { + constexpr int kElemsPerThread = (kBlockM * kBlockN) / kThreads; // 16 + constexpr int kThreadsPerRow = kBlockN / kElemsPerThread; // 8 + static_assert(kBlockN % kElemsPerThread == 0, + "sweep requires an integral number of threads per row"); + static_assert((kElemsPerThread * 2) % 8 == 0, + "sweep copy must stay 8-byte aligned"); + const int row_local = tid / kThreadsPerRow; + const int seg = tid - row_local * kThreadsPerRow; + const int g_row = m0 + row_local; + if (g_row < m) { + const int c = seg * kElemsPerThread; + const int64_t* __restrict__ src64 = + reinterpret_cast(out_tile + + row_local * kOutStride + c); + int64_t* __restrict__ dst64 = reinterpret_cast( + out + static_cast(g_row) * n + n0 + c); +#pragma unroll + for (int j = 0; j < (kElemsPerThread * 2) / 8; ++j) { + dst64[j] = src64[j]; + } + } + } +} + +// Exact assigned shape hy3_tp8_shared_down_proj_m16 (M=16, N=4096, K=192): +// one-wave INT8 DUMMA m16n16k32 kernel with a bounded whole-window LDS +// prefetch of the block's operand tiles and the packed n-major weight. +// +// The down-proj weight ((k,n)==(192,4096)) is packed once, out of the timed +// region, as P[n*K+kk] = W[kk*N+n]: each logical n column becomes a +// contiguous 192-byte run, so this block's private B window (n0..n0+15) is +// one contiguous 3,072-byte global region. One 64-lane wavefront per block +// computes one 16x16 output tile (all 16 rows of A, columns n0..n0+15). +// Before the K loop the wave cooperatively prefetches the whole K=192 window +// into LDS with wide 16-byte vector loads -- A[16,192] (3,072 B; the same +// L2-hot bytes every block reads) into a bank-skewed 16x208 padded +// row-major stage, and the packed B strip (3,072 B, fetched from global +// exactly once, fully coalesced) into a 16x208 padded n-major stage. The +// six 16-byte cooperative loads (A vectors lane/lane+64/lane+128 and the B +// strip vectors lane/lane+64/lane+128) are issued back to back into distinct +// registers before any ds_write, so all six global reads run in flight and +// the wave pays one memory round trip for the whole window instead of the +// six serialized global_load -> s_waitcnt vmcnt(0) -> ds_write_b128 chains +// the round-5 exact-source ISA shows (register reuse forced a full wait per +// chunk) -- then +// the six m16n16k32 steps load their fragments with du_load_matrix_sync / +// du_mma_sync straight from LDS (int32 accumulation). A is consumed +// row_major and B is consumed col_major on its n-major stage, so each +// lane's 8 fragment bytes are one contiguous k-run: every per-lane fragment +// read is a single 8-byte LDS read (ds_read_b64) with no byte reassembly, +// the pattern the validated TP4 down-proj lineage (same N=4096 grid=256 +// K-half=192 geometry, 208-byte strides) measured. No vmem instruction or +// global-latency wait remains inside the K loop; the only __syncthreads is +// the single one after the staging stores (one wavefront, none in the +// loop), and the epilogue goes straight from the verified gfx928 +// accumulator fragment ownership to the scaled bf16 store. P is a pure +// transpose of the identity pack (P[n*K+kk] == W[kk*N+n]), so every +// fragment byte equals the byte the round-4 k-major stage delivered for the +// same (kk, n) and the k-ascending int32 accumulation order is unchanged. +__global__ __launch_bounds__(kWaveSize) void +w8a8_gemm_dumma_m16n16k32_packedb_lds_kernel( + const int8_t* __restrict__ a, // [16, K] row-major + const int8_t* __restrict__ b, // packed [N, K] n-major down-proj + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + uint16_t* __restrict__ out, // [16, N] bf16 bits + int n, + int k) { + const int lane = static_cast(threadIdx.x); // wavefront lane 0..63 + const int n0 = static_cast(blockIdx.x) * kInt8TileN; + (void)k; // exact-shape guard fixes k == 192 == kLdsK (compile-time bound) + + // Whole-window LDS stage (guard guarantees k == 192). Both strides are + // 16-byte multiples so every staging store and fragment load stays + // 16-byte aligned; 208 = 192 + 16 is the validated TP4 down_proj skew, + // used here for the A rows and for the B n-columns alike. + constexpr int kLdsK = 192; // K of the guarded shape + constexpr int kLdsAStride = 208; // A row stride in bytes (192 + 16 pad) + constexpr int kLdsBStride = 208; // B n-column stride in bytes (192 + 16 pad) + __shared__ __align__(16) int8_t a_tile[kInt8TileM * kLdsAStride]; + __shared__ __align__(16) int8_t b_tile[kInt8TileN * kLdsBStride]; + + // A: all 16 rows x 192 bytes are contiguous in global memory (3,072 B = + // 192 int4 vectors) and B: this block's private strip is the 16 packed n + // rows n0..n0+15, one contiguous 3,072 B region (P row n is 192 contiguous + // k bytes). In both windows lane l owns vectors l, l+64 and l+128, i.e. + // six fully coalesced 16-byte loads per lane cover the whole prologue. + // All six loads are issued into distinct registers before any LDS store: + // the round-5 exact-source ISA shows the previous per-vector loops + // degraded into six serialized global_load -> s_waitcnt vmcnt(0) -> + // ds_write_b128 round trips per block (three for A, three for B), each + // paying the full L2/HBM latency on the block's critical path. Batching + // keeps every source byte and every padded tile slot identical while + // collapsing the prologue to one in-flight memory round trip. + const int4 a_v0 = *reinterpret_cast(a + lane * 16); + const int4 a_v1 = + *reinterpret_cast(a + (lane + kWaveSize) * 16); + const int4 a_v2 = + *reinterpret_cast(a + (lane + 2 * kWaveSize) * 16); + const int8_t* b_strip = b + static_cast(n0) * kLdsK; + const int4 b_v0 = *reinterpret_cast(b_strip + lane * 16); + const int4 b_v1 = + *reinterpret_cast(b_strip + (lane + kWaveSize) * 16); + const int4 b_v2 = + *reinterpret_cast(b_strip + (lane + 2 * kWaveSize) * 16); + // wv1..wv3 scale preload, issued inside the staging batch right behind the + // six operand vectors (same three vmem read instructions per block as the + // generic epilogue performs; only the issue point moves to the front of the + // wave's vmem FIFO). col_mod4 = lane >> 4; M == 16 fills the whole tile, so + // these are the exact columns the epilogue scales, and every index stays + // in range (n0 + col_mod4 + 12 <= 4095 < N). See + // store_acc_fragment_tail_preloaded for the round-6 ISA evidence that + // wv1..wv3 were issued between v_mmac 1..3 and waited on between the four + // bf16 stores. + const int col_mod4 = lane >> 4; + const float wv1 = weight_scale[n0 + col_mod4 + 4]; + const float wv2 = weight_scale[n0 + col_mod4 + 8]; + const float wv3 = weight_scale[n0 + col_mod4 + 12]; + constexpr int kChunksPerRow = kLdsK / 16; // 12 x 16-B chunks per 192-B row + // Chunk vec -> padded-tile int4 slot (row = vec / 12, k16 = (vec % 12)*16); + // both tiles use the same 208-byte skew so the mapping is shared. + const auto stage_slot = [](int vec) -> int { + const int row = vec / kChunksPerRow; + return (row * kLdsAStride + (vec - row * kChunksPerRow) * 16) / 16; + }; + reinterpret_cast(a_tile)[stage_slot(lane)] = a_v0; + reinterpret_cast(a_tile)[stage_slot(lane + kWaveSize)] = a_v1; + reinterpret_cast(a_tile)[stage_slot(lane + 2 * kWaveSize)] = a_v2; + reinterpret_cast(b_tile)[stage_slot(lane)] = b_v0; + reinterpret_cast(b_tile)[stage_slot(lane + kWaveSize)] = b_v1; + reinterpret_cast(b_tile)[stage_slot(lane + 2 * kWaveSize)] = b_v2; + __syncthreads(); + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + // K = 192 = 6 x kInt8TileK DUMMA steps; k ascending, int32 accumulation. + // Fragment loads hit LDS (du_load_matrix_sync on a __shared__ pointer, + // the same pattern the staged prefill kernel above validates), so the K + // loop carries no vmem and no global-latency wait. The bound is the + // compile-time kLdsK (the exact-shape guard fixes k == 192), so the six + // steps fully unroll and the scheduler can issue step i+1's ds_reads while + // step i's v_mmac is still in the pipe, instead of a rolled loop in which + // each iteration's reads can only issue after the previous step retires. + for (int k0 = 0; k0 < kLdsK; k0 += kInt8TileK) { + du_load_matrix_sync(a_frag, a_tile + k0, kLdsAStride); + // col_major on the n-major stage: byte (k, n) sits at n*208 + k, so the + // k0 step's fragment starts at b_tile + k0 and each lane's 8 bytes are + // one contiguous run in its column (one ds_read_b64 per lane). + du_load_matrix_sync(b_frag, b_tile + k0, kLdsBStride); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Direct fragment -> ((acc * x_scale) * weight_scale) -> bf16 store with + // wv1..wv3 supplied from the prologue preload: no vmem load issued inside + // the mmac chain remains to gate the store tail (xv and wv0 load here at + // their proven round-6 issue points). M == 16 fills the whole tile; row is + // always in range. + store_acc_fragment_tail_preloaded(acc_frag, x_scale, weight_scale, wv1, wv2, + wv3, out, n, n0, lane); +} + +// Generic scalar int8/int32 fallback: one thread per output element, exact +// int32 accumulation over K, fp32 scaling, bf16 store. Serves every +// unmatched (m, n, k) and all small-M API cases (including the assigned +// M=16 shape at bootstrap). Requires only m >= 1 and the API-validated +// alignment properties; it never reads out of bounds. +__global__ __launch_bounds__(kScalarThreads) +void w8a8_gemm_scalar_fallback_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + uint16_t* __restrict__ out, + 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 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; + const int8_t* a_row = a + static_cast(row) * k; 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]); + if (k == 192 && n == 4096) { + // The (k, n) == (192, 4096) down-proj packed buffer is n-major + // P[n*K+kk] = W[kk*N+n]: column col is one contiguous 192-byte run. + const int8_t* b_col = b + static_cast(col) * k; +#pragma unroll 4 + 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; +#pragma unroll 4 + 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); + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[static_cast(row) * n + col] = bf16_bits_from_f32(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; +// Generic element-wise device-to-device copy used by the identity pack. +template +__global__ void w8a8_element_copy_kernel(const T* __restrict__ src, + 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]; } - 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]; +// One-time n-major transpose pack for the TP8 down-proj weight +// ((k, n) == (192, 4096)): packed[n*k + kk] = raw[kk*n + n_local]. Every +// logical [K, N] column of the weight becomes a contiguous k-run of the +// packed buffer (same byte count, same graph-stable buffer). Weight +// preprocessing; runs outside the timed region and Graph capture. +__global__ __launch_bounds__(kPackThreads) +void w8a8_pack_nmajor_downproj_kernel(const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, int k, + int n) { + const int64_t total = static_cast(n) * k; + 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 n_local = static_cast(idx / k); + const int kk = static_cast(idx - static_cast(n_local) * k); + packed[idx] = raw[static_cast(kk) * n + n_local]; } } -__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]; +void launch_pack_identity_impl(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 int64_t scale_count = n; + const dim3 block(kPackThreads); + const unsigned weight_blocks = static_cast( + (weight_count + kPackThreads - 1) / kPackThreads); + const unsigned scale_blocks = static_cast( + (scale_count + kPackThreads - 1) / kPackThreads); + if (weight_count > 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_element_copy_kernel), + dim3(weight_blocks), block, 0, stream, raw_weight, packed_weight, + weight_count); + } + if (scale_count > 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_element_copy_kernel), dim3(scale_blocks), + block, 0, stream, weight_scale, packed_weight_scale, scale_count); } } } // 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, @@ -556,15 +1250,35 @@ extern "C" void launch_w8a8_gemm( 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) { + auto* out_ptr = static_cast(out); + + // Exact assigned decode shape hy3_tp8_shared_down_proj_m16. The exact + // (m, n, k) == (16, 4096, 192) guard keeps the paired M=2 API shape with + // the same (N, K) on the generic scalar fallback below. + if (m == 16 && n == 4096 && k == 192) { + // One 64-lane wavefront per block, one 16x16 output tile per block: + // grid = 4096 / 16 = 256 blocks, whole-window LDS prefetch of A and of + // the packed n-major B strip, one staging barrier per block, col_major + // B fragment consumption (one 8-byte LDS read per lane per step). + const dim3 block(kWaveSize); + const dim3 grid(static_cast(n / kInt8TileN)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_dumma_m16n16k32_packedb_lds_kernel), grid, + block, 0, stream, a, b, x_scale, weight_scale, out_ptr, n, k); + return; + } + + // Down-proj family guard (assigned M=4096 shape). The M < 128 shapes - + // including the paired M=2 API shape with the same (N, K) - deliberately + // fall through to the scalar kernel below. The guard is shape-exact on K + // (k == kAssignedK == 192, the only K whose weight is packed n-major) and + // requires N to be a multiple of the active macro-tile N. The active + // macro-tile is fixed at compile time (kActiveTileIndex), so the timed + // region never tunes; the other two family instantiations stay compiled + // for compile-time A/B. + const bool use_tiled_dumma = + m >= 128 && k == kAssignedK && (n % kActiveBlockN) == 0; + if (use_tiled_dumma) { const dim3 grid( static_cast(n / kActiveBlockN), static_cast((m + kActiveBlockM - 1) / kActiveBlockM)); @@ -573,39 +1287,44 @@ extern "C" void launch_w8a8_gemm( 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); + grid, block, 0, stream, a, b, x_scale, weight_scale, out_ptr, 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); + if (kActivePipeline == 0) { + // Round-3 single-buffered pipeline control (64x128 tile): same + // bytes, MMAC order and occupancy as the double-buffered arm, LDS + // halved to 14,848 B/block, one extra barrier per stage boundary. + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_sb_kernel<64, 128, kPrefillStageK>), + grid, block, 0, stream, a, b, x_scale, weight_scale, out_ptr, m, + n, k); + } else { + // Accepted double-buffered pipeline (retention arm). + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_kernel<64, 128, kPrefillStageK>), + grid, block, 0, stream, a, b, x_scale, weight_scale, out_ptr, 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); + HIP_KERNEL_NAME( + w8a8_dumma_prefill_kernel<128, 64, kPrefillStageK>), + grid, block, 0, stream, a, b, x_scale, weight_scale, out_ptr, 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); + // Generic scalar fallback (one grid row per M, columns chunked). + const dim3 scalar_grid( + static_cast((n + kScalarThreads - 1) / kScalarThreads), + static_cast(m)); + const dim3 scalar_block(kScalarThreads); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, scalar_grid, + scalar_block, 0, stream, a, b, x_scale, weight_scale, + out_ptr, m, n, k); } extern "C" void launch_pack_w8a8_weight( @@ -616,39 +1335,33 @@ extern "C" void launch_pack_w8a8_weight( 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 assigned TP8 down-proj weight ((k, n) == (192, 4096)) is packed + // n-major (P[n*k + kk] = raw[kk*n + n_local], one contiguous k-run per n + // column) so the M=16 kernel and the M>=128 tiled DUMMA path can consume + // B fragments col_major with one vector LDS read per lane per step. Every + // other (k, n) keeps the logical [K, N] row-major identity pack (weight + // and scale copies), which the generic fallback decodes. + if (k == 192 && n == 4096) { + const int64_t weight_count = static_cast(k) * n; + const int64_t scale_count = n; + const dim3 block(kPackThreads); + if (weight_count > 0) { + const unsigned weight_blocks = static_cast( + (weight_count + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL(HIP_KERNEL_NAME(w8a8_pack_nmajor_downproj_kernel), + dim3(weight_blocks), block, 0, stream, raw_weight, + packed_weight, k, n); + } + if (scale_count > 0) { + const unsigned scale_blocks = static_cast( + (scale_count + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL(HIP_KERNEL_NAME(w8a8_element_copy_kernel), + dim3(scale_blocks), block, 0, stream, weight_scale, + packed_weight_scale, scale_count); + } + return; } - - // 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)); + launch_pack_identity_impl(raw_weight, weight_scale, packed_weight, + packed_weight_scale, k, n, stream); } // @@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 index 81b98bb3..4d72a3d5 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp4_qkv_proj_m4096 commit=8e400dcec14217e9b03ccc7512c4d99149820b06 added=2026-08-24 +// baseline_us=31556.2 speedup=20.2283 // median_us=1560 p90_us=1575 // source=minimaxm3-dsh-tp4-m4096-1-0c2f84a9 // MetaInfer W8A8 INT8 GEMM bootstrap for gfx928 (K500SM_AI). 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 index eca651cf..0df41638 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp4_qkv_proj_and_indexer_qk_m4096 commit=d6cbef3829c738228f5dab5255b89c384c66c6ef added=2026-08-24 +// baseline_us=35918.8 speedup=33.1966 // median_us=1082 p90_us=1089 // source=minimaxm3-dsh-tp4-m4096-1-0c2f84a9 // MetaInfer W8A8 INT8 GEMM bootstrap for gfx928 (K500SM_AI). 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 index f9b6e0ad..4db5a3fa 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp4_shared_gate_up_proj_m4096 commit=5c8bfa17fe45a29749afdfc9124782e9068e415c added=2026-08-24 +// baseline_us=19776.5 speedup=20.0471 // 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). 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 index f771caa3..e41630bf 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_o_proj_m16 commit=35384b7737eacb9fb8946dc1c5f984e8bf364aae added=2026-08-28 +// baseline_us=38.655 speedup=3.39974 // 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). 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 index bf0cbf03..3434210f 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_qkv_proj_m16 commit=b86aaad82e2a39632cd692ec7fe88dcaeeadaaed added=2026-08-28 +// baseline_us=110.478 speedup=3.32765 // median_us=33.2 p90_us=33.31 // source=minimax-dsh-tp8-m16-1-78402260 // MetaInfer W8A8 INT8 GEMM — gfx928 (K500SM_AI) bootstrap. 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 index 3f4a48e3..ed05e4cb 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_qkv_proj_and_indexer_qk_m16 commit=e72025dfcfa50ff3bf3fcb00b87663e65f4e0fc4 added=2026-08-28 +// baseline_us=111.486 speedup=2.80609 // median_us=39.73 p90_us=39.81 // source=minimax-dsh-tp8-m16-1-78402260 // csrc/w8a8_gemm_hip.hip 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 index f5930cb7..fe0eade0 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_shared_down_proj_m16 commit=2d8b403746ea59e59c41eabce07da2e6b2447072 added=2026-08-28 +// baseline_us=23.152 speedup=2.58595 // 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. 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 index 4400c8bd..8f28c5a6 100644 --- 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 @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_shared_gate_up_proj_m16 commit=4358dea23da7468d3868ce258b971667a50c139e added=2026-08-28 +// baseline_us=104.527 speedup=6.10912 // 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.