diff --git a/README.md b/README.md index 68366e8d..3b94c399 100644 --- a/README.md +++ b/README.md @@ -79,13 +79,77 @@ ## Quick start -### Step 1: Install ccb (MetaInfer currently uses the open-source Claude Code CLI; other coding agents are not yet supported — contributions welcome) +### Step 1: Install a coding agent CLI + +MetaInfer launches phase-specific coding agents through a shared +`SubAgentManager`. The backend is selected by `METAINFER_AGENT_BACKEND`. + +#### Claude/ccb backend + +Install ccb if you want to use the Claude-compatible backend: -Open-source ccb repository: https://github.com/claude-code-best/claude-code ``` npm i -g claude-code-best ``` +Run MetaInfer with: + +```bash +METAINFER_AGENT_BACKEND=claude ./serve.py +``` + +#### Codex backend + +Install the Codex CLI and configure it before starting MetaInfer. MetaInfer +does not store or pass API keys; it invokes `codex exec` and reuses your local +Codex login, environment, and config such as `$CODEX_HOME/config.toml`. + +```bash +codex login +codex doctor +``` + +Then start MetaInfer with the Codex backend: + +```bash +METAINFER_AGENT_BACKEND=codex ./serve.py +``` + +If the Codex binary is not named `codex` or is not on `PATH`, set: + +```bash +METAINFER_AGENT_BACKEND=codex METAINFER_CODEX_BIN=/path/to/codex ./serve.py +``` + +#### pi backend + +Install the [pi](https://pi.dev) coding agent CLI and configure a provider +before starting MetaInfer. MetaInfer does not store or pass API keys; it +invokes `pi --mode json -p` and reuses your local `~/.pi` settings and +environment variables (`PI_PROVIDER`, `PI_MODEL`, provider API keys, etc.). + +```bash +npm i -g @earendil-works/pi-coding-agent +pi auth status # verify a provider is configured +``` + +Then start MetaInfer with the pi backend: + +```bash +METAINFER_AGENT_BACKEND=pi ./serve.py +``` + +If the pi binary is not named `pi` or is not on `PATH`, set: + +```bash +METAINFER_AGENT_BACKEND=pi METAINFER_PI_BIN=/path/to/pi ./serve.py +``` + +The pi backend maps the orchestrator's `effort` knob onto pi's `--thinking` +level (`low` / `medium` / `high` / `max`). Session continuation uses +`--session --continue`; the session id is captured from pi's leading +`session` event so later turns can resume it. + ### Step 2: Install MetaInfer ```bash git clone https://github.com/MetaInfer/MetaInfer.git diff --git a/README_CN.md b/README_CN.md index 7654c80a..bd05ef5d 100644 --- a/README_CN.md +++ b/README_CN.md @@ -79,13 +79,76 @@ ## 快速开始 -### 第一步，安装ccb。（本项目目前使开源的claude code版本，暂不支持其他coding agent，欢迎贡献代码以支持更多coding agent） +### 第一步，安装一个 coding agent CLI + +MetaInfer 通过统一的 `SubAgentManager` 拉起各阶段 coding agent。 +后端通过 `METAINFER_AGENT_BACKEND` 选择。 + +#### Claude/ccb 后端 + +如果使用 Claude-compatible 后端，先安装 ccb： -开源ccb项目地址：https://github.com/claude-code-best/claude-code ``` npm i -g claude-code-best ``` +启动时指定： + +```bash +METAINFER_AGENT_BACKEND=claude ./serve.py +``` + +#### Codex 后端 + +如果使用 Codex 后端，需要先在本机安装并配置 Codex CLI。MetaInfer +不会保存或传递 API key；它只调用 `codex exec`，复用你本机已有的 +Codex 登录状态、环境变量和 `$CODEX_HOME/config.toml` 等配置。 + +```bash +codex login +codex doctor +``` + +然后用 Codex 后端启动 MetaInfer： + +```bash +METAINFER_AGENT_BACKEND=codex ./serve.py +``` + +如果 Codex 二进制不叫 `codex`，或者不在 `PATH` 里，可以显式指定： + +```bash +METAINFER_AGENT_BACKEND=codex METAINFER_CODEX_BIN=/path/to/codex ./serve.py +``` + +#### pi 后端 + +如果使用 [pi](https://pi.dev) 后端，需要先在本机安装 pi 编码代理 CLI +并配置好 provider。MetaInfer 不会保存或传递 API key；它只调用 +`pi --mode json -p`，复用你本机已有的 `~/.pi` 设置和环境变量 +（`PI_PROVIDER`、`PI_MODEL`、各 provider 的 API key 等）。 + +```bash +npm i -g @earendil-works/pi-coding-agent +pi auth status # 确认已配置某个 provider +``` + +然后用 pi 后端启动 MetaInfer： + +```bash +METAINFER_AGENT_BACKEND=pi ./serve.py +``` + +如果 pi 二进制不叫 `pi`，或者不在 `PATH` 里，可以显式指定： + +```bash +METAINFER_AGENT_BACKEND=pi METAINFER_PI_BIN=/path/to/pi ./serve.py +``` + +pi 后端会把 orchestrator 的 `effort` 旋钮映射到 pi 的 `--thinking` 等级 +（`low` / `medium` / `high` / `max`）。会话续接使用 `--session --continue`； +session id 从 pi 流首的 `session` 事件中捕获，供后续轮次续接。 + ### 第二步，安装MetaInfer ```bash git clone https://github.com/MetaInfer/MetaInfer.git diff --git a/metainfer/orchestrator/_bootstrap.py b/metainfer/orchestrator/_bootstrap.py index 22b1d68e..c8679805 100644 --- a/metainfer/orchestrator/_bootstrap.py +++ b/metainfer/orchestrator/_bootstrap.py @@ -174,6 +174,9 @@ def restore() -> None: def make_subagent_manager( *, claude_bin: str, + codex_bin: Optional[str] = None, + pi_bin: Optional[str] = None, + agent_backend: Optional[str] = None, model: Optional[str], permission_mode: str, effort: str, @@ -181,6 +184,7 @@ def make_subagent_manager( snapshot_file: Path, max_concurrent: int = 4, budget: Any = None, + state_dir: Optional[Path] = None, ) -> SubAgentManager: """Build a SubAgentManager with the standard settings shared by every orchestrator. Per-orchestrator customization happens via @@ -189,9 +193,32 @@ def make_subagent_manager( ``budget`` (optional) wires the per-task :class:`TokenBudget` so every agent launch is gated + every result's cost is recorded. + + ``state_dir`` (optional) scopes the pi backend's session storage to a + per-task directory (``/pi-sessions``) so orchestrator + sessions are invisible to a user running ``pi --continue`` from their + shell. Other backends ignore it. See ``SubAgentManager.pi_session_dir``. """ + resolved_agent_backend = ( + agent_backend + or os.environ.get("METAINFER_AGENT_BACKEND") + or "claude" + ) + resolved_codex_bin = ( + codex_bin + or os.environ.get("METAINFER_CODEX_BIN") + or "codex" + ) + resolved_pi_bin = ( + pi_bin + or os.environ.get("METAINFER_PI_BIN") + or "pi" + ) return SubAgentManager( claude_bin=claude_bin, + codex_bin=resolved_codex_bin, + pi_bin=resolved_pi_bin, + agent_backend=resolved_agent_backend, default_model=model, permission_mode=permission_mode, effort=effort, @@ -199,6 +226,7 @@ def make_subagent_manager( snapshot_file=snapshot_file, max_concurrent=max_concurrent, budget=budget, + pi_session_dir=(Path(state_dir) / "pi-sessions") if state_dir else None, ) diff --git a/metainfer/orchestrator/subagent_manager.py b/metainfer/orchestrator/subagent_manager.py index 3c79c722..5d8c92a1 100644 --- a/metainfer/orchestrator/subagent_manager.py +++ b/metainfer/orchestrator/subagent_manager.py @@ -33,6 +33,32 @@ # Exit codes / signals that indicate infrastructure rather than logic failure. # 124 = timeout (coreutils convention), 137 = SIGKILL (128+9), 143 = SIGTERM (128+15). _INFRA_EXIT_CODES = {124, 137, 143} +AgentBackend = Literal["claude", "codex", "pi"] + + +# Claude Code "effort" levels map onto pi's ``--thinking`` levels 1:1 for +# the overlapping set (low / medium / high / max). pi additionally offers +# ``off`` / ``minimal`` / ``xhigh`` but those have no Claude equivalent, so +# we pass the value through unchanged and let pi validate it. +_PI_THINKING_ALIASES = { + "low": "low", + "medium": "medium", + "high": "high", + "max": "max", +} + + +def _normalize_agent_backend(value: str) -> AgentBackend: + v = (value or "claude").strip().lower() + if v in {"claude", "claude-code", "ccb"}: + return "claude" + if v in {"codex", "openai-codex"}: + return "codex" + if v in {"pi", "pi-coding-agent", "earendil"}: + return "pi" + raise ValueError( + f"invalid agent backend {value!r}; expected 'claude', 'codex', or 'pi'" + ) # --------------------------------------------------------------------------- # @@ -149,6 +175,9 @@ class SubAgentManager: def __init__( self, claude_bin: str = "ccb", + codex_bin: str = "codex", + pi_bin: str = "pi", + agent_backend: AgentBackend = "claude", default_model: Optional[str] = None, max_concurrent: int = 4, permission_mode: str = "bypassPermissions", @@ -157,8 +186,22 @@ def __init__( snapshot_file: Optional[Path] = None, budget: Any = None, budget_source: str = "orchestrator", + # pi backend only: isolate orchestrator session files into a + # per-task directory so they are invisible to a user running + # ``pi --continue`` / ``pi --resume`` from their shell (which + # defaults to ~/.pi). Without this, pi indexes sessions by cwd, + # so a user cd'ing into an iteration workdir — or passing a + # partial session id lifted from events.jsonl — could attach to + # a production orchestrator session mid-flight and corrupt the + # context that later ``--session --continue`` turns rely on. + # When set, ``_build_pi_command`` prepends ``--session-dir`` so + # both the initial launch and every resume read/write here. + pi_session_dir: Optional[Path] = None, ) -> None: + self.agent_backend = _normalize_agent_backend(agent_backend) self.claude_bin = claude_bin + self.codex_bin = codex_bin + self.pi_bin = pi_bin self.default_model = default_model self.max_concurrent = max_concurrent # Per-task token / cost budget (TokenBudget instance or None to @@ -208,6 +251,9 @@ def __init__( self.snapshot_file: Optional[Path] = ( Path(snapshot_file) if snapshot_file else None ) + self.pi_session_dir: Optional[Path] = ( + Path(pi_session_dir).resolve() if pi_session_dir else None + ) self._handles: Dict[str, AgentHandle] = {} self._results: Dict[str, AgentResult] = {} self._ctrl_lock = threading.Lock() @@ -616,6 +662,13 @@ def dump_snapshot(self) -> None: # ------------------------------------------------------------------ # def _build_command(self, spec: AgentSpec) -> List[str]: + if self.agent_backend == "codex": + return self._build_codex_command(spec) + if self.agent_backend == "pi": + return self._build_pi_command(spec) + return self._build_claude_command(spec) + + def _build_claude_command(self, spec: AgentSpec) -> List[str]: cmd = [ self.claude_bin, "-p", # print (non-interactive) mode; prompt read from stdin @@ -658,6 +711,86 @@ def _build_command(self, spec: AgentSpec) -> List[str]: cmd += list(spec.extra_args) return cmd + def _build_codex_command(self, spec: AgentSpec) -> List[str]: + # Codex CLI owns auth + user configuration. Do not pass API keys + # or --ignore-user-config; child processes inherit the current + # CODEX_HOME / OPENAI_* environment exactly like a normal user run. + model = spec.model or self.default_model + if spec.resume_session_id: + cmd = [ + self.codex_bin, + "exec", + "resume", + "--json", + "--skip-git-repo-check", + ] + if model: + cmd += ["--model", model] + cmd += list(spec.extra_args) + cmd += [spec.resume_session_id, "-"] + else: + cmd = [ + self.codex_bin, + "exec", + "--json", + "--color", "never", + "--skip-git-repo-check", + "--sandbox", "workspace-write", + "-C", str(spec.workdir), + "--add-dir", "/tmp", + ] + for d in self.extra_add_dirs: + cmd += ["--add-dir", str(d)] + if model: + cmd += ["--model", model] + cmd += list(spec.extra_args) + cmd += ["-"] + return cmd + + def _build_pi_command(self, spec: AgentSpec) -> List[str]: + # pi (https://pi.dev) owns auth + user configuration via ~/.pi + # settings and environment variables (PI_PROVIDER, PI_MODEL, + # provider API keys, etc.). We invoke `pi --mode json -p` so the + # orchestrator gets a stream of JSONL events on stdout (captured + # into events.jsonl just like claude/codex), with the prompt piped + # via stdin — pi merges piped stdin into the initial user message. + # + # pi has no ``--add-dir`` equivalent: its built-in tools (read, + # bash, edit, write, grep, find, ls) can already reach any path, + # so ``extra_add_dirs`` is a no-op here (cwd scoping via Popen is + # enough). We keep the param for interface parity. + model = spec.model or self.default_model + cmd = [ + self.pi_bin, + "--mode", "json", + "-p", # print (non-interactive) mode; prompt read from stdin + ] + # Isolate orchestrator sessions from the user's ~/.pi index so a + # production task can't be hijacked by ``pi --continue`` / ``pi + # --resume`` / ``pi --session `` from a shell sharing + # the same cwd. See __init__ docstring on ``pi_session_dir``. + if self.pi_session_dir is not None: + self.pi_session_dir.mkdir(parents=True, exist_ok=True) + cmd += ["--session-dir", str(self.pi_session_dir)] + if model: + cmd += ["--model", model] + # Thinking level maps from the Claude "effort" knob. The overlapping + # set (low / medium / high / max) is identical; unknown values are + # passed through and pi validates them. + if self.effort: + cmd += ["--thinking", _PI_THINKING_ALIASES.get(self.effort, self.effort)] + # Session continuation. ``--session --continue`` resumes an + # existing session; ``--session-id `` pins the UUID for the + # first turn so the caller knows what to resume later. If neither + # is set, pi mints a fresh session and the manager captures its + # id from the ``session`` event at the top of the stream. + if spec.resume_session_id: + cmd += ["--session", spec.resume_session_id, "--continue"] + elif spec.session_id: + cmd += ["--session-id", spec.session_id] + cmd += list(spec.extra_args) + return cmd + def _build_env(self, spec: AgentSpec) -> Dict[str, str]: env = dict(os.environ) env.update(spec.env_overrides) @@ -673,6 +806,8 @@ def _build_env(self, spec: AgentSpec) -> Dict[str, str]: env["METAINFER_ROOT"] = str(_paths.root_dir()) # Keep the agent from going interactive env.setdefault("DISABLE_INTERACTIVITY", "1") + if self.agent_backend in ("codex", "pi"): + return env # bypassPermissions under EUID=0 normally trips a hard exit # ("--dangerously-skip-permissions cannot be used with root/sudo # privileges"). ccb skips that check when IS_SANDBOX=1, which is @@ -703,40 +838,60 @@ def _materialize_result( except json.JSONDecodeError: continue final_text = "" - for ev in reversed(events): - if ev.get("type") == "assistant" and isinstance(ev.get("message"), dict): - content = ev["message"].get("content") - if isinstance(content, list): - for blk in content: - if isinstance(blk, dict) and blk.get("type") == "text": - final_text = blk.get("text", "") - break - if final_text: - break - if ev.get("type") == "result": - final_text = ev.get("result", "") or final_text - break - # Session id: emitted on the very first ``system`` event of the - # stream and again on every ``result`` event. Prefer the result's - # value (it's the final, post-turn session id; for ``--resume`` - # invocations this matches the resumed-from id and confirms the - # continuation actually happened). session_id = None - for ev in events: - sid = ev.get("session_id") - if sid: - session_id = sid - if ev.get("type") == "result": - break - # Pull the cost / usage block from the final ``result`` event. - # This is what the token budget circuit breaker keys off. None - # when the agent was killed / crashed before emitting result. usage: Optional[Dict[str, Any]] = None - for ev in reversed(events): - if ev.get("type") == "result" and isinstance(ev, dict): - if "usage" in ev or "total_cost_usd" in ev: - usage = ev + if self.agent_backend == "pi": + final_text, session_id, usage = self._extract_pi_result(events) + else: + for ev in reversed(events): + if ev.get("type") == "item.completed": + item = ev.get("item") + if isinstance(item, dict) and item.get("type") == "agent_message": + text = item.get("text") + if isinstance(text, str) and text: + final_text = text + break + if ev.get("type") == "assistant" and isinstance(ev.get("message"), dict): + content = ev["message"].get("content") + if isinstance(content, list): + for blk in content: + if isinstance(blk, dict) and blk.get("type") == "text": + final_text = blk.get("text", "") + break + if final_text: + break + if ev.get("type") == "result": + final_text = ev.get("result", "") or final_text break + # Session id: emitted on the very first ``system`` event of the + # stream and again on every ``result`` event. Prefer the result's + # value (it's the final, post-turn session id; for ``--resume`` + # invocations this matches the resumed-from id and confirms the + # continuation actually happened). + for ev in events: + sid = ev.get("session_id") + if not sid: + sid = ev.get("thread_id") + if sid: + session_id = sid + if ev.get("type") == "result": + break + # Pull the cost / usage block from the final ``result`` event. + # This is what the token budget circuit breaker keys off. None + # when the agent was killed / crashed before emitting result. + for ev in reversed(events): + if ev.get("type") == "result" and isinstance(ev, dict): + if "usage" in ev or "total_cost_usd" in ev: + usage = ev + break + if ev.get("type") == "turn.completed" and isinstance(ev, dict): + if "usage" in ev: + usage = ev + break + if usage is not None and session_id: + if not usage.get("session_id") and not usage.get("thread_id"): + usage = dict(usage) + usage["thread_id"] = session_id error = None failure_mode: Optional[Literal["infra", "logic", "budget"]] = None success = (rc == 0) and not handle.killed @@ -765,6 +920,91 @@ def _materialize_result( usage=usage, ) + def _extract_pi_result( + self, events: List[Dict[str, Any]] + ) -> tuple[str, Optional[str], Optional[Dict[str, Any]]]: + """Parse pi's JSONL event stream into (final_text, session_id, usage). + + pi emits a different schema than claude/codex stream-json: + + * Session id lives on the leading ``session`` event's ``id`` field + (not a top-level ``session_id`` / ``thread_id``). + * Assistant text lives on ``message_end`` / ``turn_end`` / + ``agent_end`` events under ``message.content`` (a list of blocks; + concatenate the ``text`` blocks of the last assistant message). + * Usage lives on the same events under ``message.usage`` with pi's + own field names (``input`` / ``output`` / ``cacheRead`` / + ``cacheWrite`` / ``cost.total``). We normalize it into the + canonical shape that :func:`usage_from_result_event` already + understands so the token-budget circuit breaker works unchanged. + """ + session_id: Optional[str] = None + for ev in events: + if ev.get("type") == "session" and isinstance(ev.get("id"), str): + session_id = ev["id"] + break + + final_text = "" + usage: Optional[Dict[str, Any]] = None + # Walk newest -> oldest. The last assistant message in stream order + # is the first one we hit in reverse, which is what we want for both + # final_text and the final usage tally. + for ev in reversed(events): + etype = ev.get("type") + if etype not in ("message_end", "turn_end", "agent_end"): + continue + # agent_end carries the full message list; pick the last assistant. + if etype == "agent_end": + msgs = ev.get("messages") + if not isinstance(msgs, list) or not msgs: + continue + msg = None + for m in reversed(msgs): + if isinstance(m, dict) and m.get("role") == "assistant": + msg = m + break + if msg is None: + continue + else: + msg = ev.get("message") + if not isinstance(msg, dict) or msg.get("role") != "assistant": + continue + + # Final assistant text: concatenate text blocks. + if not final_text: + content = msg.get("content") + if isinstance(content, list): + parts = [ + blk.get("text", "") + for blk in content + if isinstance(blk, dict) and blk.get("type") == "text" + ] + text = "".join(p for p in parts if isinstance(p, str)) + if text: + final_text = text + + # Usage: pi reports cumulative usage on every assistant message. + # The last assistant message carries the final tally. + if usage is None: + u = msg.get("usage") + if isinstance(u, dict): + cost = u.get("cost") if isinstance(u.get("cost"), dict) else {} + usage = { + "type": "result", + "usage": { + "input_tokens": int(u.get("input", 0) or 0), + "output_tokens": int(u.get("output", 0) or 0), + "cache_read_input_tokens": int(u.get("cacheRead", 0) or 0), + "cache_creation_input_tokens": int(u.get("cacheWrite", 0) or 0), + }, + "total_cost_usd": float(cost.get("total", 0.0) or 0.0), + "session_id": session_id, + } + + if final_text and usage is not None: + break + return final_text, session_id, usage + def _write_status( self, spec: AgentSpec, diff --git a/metainfer/orchestrator/tests/test_subagent_manager.py b/metainfer/orchestrator/tests/test_subagent_manager.py new file mode 100644 index 00000000..f87ea94e --- /dev/null +++ b/metainfer/orchestrator/tests/test_subagent_manager.py @@ -0,0 +1,352 @@ +"""Focused tests for SubAgentManager backend command/parse behavior.""" + +from __future__ import annotations + +import json +import tempfile +import time +from pathlib import Path + +from metainfer.orchestrator.subagent_manager import ( + AgentHandle, + AgentSpec, + SubAgentManager, +) +from metainfer.orchestrator._bootstrap import make_subagent_manager + + +class _FakeProcess: + def __init__(self, returncode: int = 0) -> None: + self.returncode = returncode + + +def _spec(tmp: Path, *, resume_session_id: str | None = None) -> AgentSpec: + workdir = tmp / "work" + log_dir = tmp / "logs" + workdir.mkdir() + log_dir.mkdir() + prompt = tmp / "prompt.txt" + prompt.write_text("do work", encoding="utf-8") + return AgentSpec( + name="agent1", + role="tester", + prompt_file=prompt, + workdir=workdir, + log_dir=log_dir, + resume_session_id=resume_session_id, + ) + + +def test_codex_command_uses_exec_json_and_current_config(): + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + kb = tmp / "kb" + kb.mkdir() + mgr = SubAgentManager( + agent_backend="codex", + codex_bin="codex-dev", + extra_add_dirs=[kb], + default_model=None, + ) + cmd = mgr._build_command(_spec(tmp)) + assert cmd[:5] == ["codex-dev", "exec", "--json", "--color", "never"] + assert "--skip-git-repo-check" in cmd + assert "--sandbox" in cmd + assert "workspace-write" in cmd + assert "-C" in cmd + assert str(tmp / "work") in cmd + assert "--add-dir" in cmd + assert "/tmp" in cmd + assert str(kb.resolve()) in cmd + assert "--model" not in cmd + assert "--ignore-user-config" not in cmd + assert cmd[-1] == "-" + + +def test_codex_resume_command_threads_session(): + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + mgr = SubAgentManager(agent_backend="codex", codex_bin="codex-dev", + default_model="gpt-5-codex") + cmd = mgr._build_command(_spec(tmp, resume_session_id="thread-123")) + assert cmd == [ + "codex-dev", "exec", "resume", "--json", + "--skip-git-repo-check", "--model", "gpt-5-codex", + "thread-123", "-", + ] + + +def test_materialize_codex_jsonl_result(): + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + spec = _spec(tmp) + events = [ + {"type": "thread.started", "thread_id": "thread-123"}, + {"type": "turn.started"}, + { + "type": "item.completed", + "item": {"type": "agent_message", "text": "done"}, + }, + { + "type": "turn.completed", + "usage": { + "input_tokens": 10, + "cached_input_tokens": 8, + "output_tokens": 2, + "reasoning_output_tokens": 1, + }, + }, + ] + ef = spec.events_file(1) + ef.write_text( + "Reading additional input from stdin...\n" + + "\n".join(json.dumps(e) for e in events) + + "\n", + encoding="utf-8", + ) + handle = AgentHandle( + spec=spec, + attempt=1, + process=_FakeProcess(0), # type: ignore[arg-type] + started_at=time.time() - 1, + last_output_at=time.time(), + ) + mgr = SubAgentManager(agent_backend="codex") + result = mgr._materialize_result(handle, spec, 1) + assert result.success + assert result.final_text == "done" + assert result.session_id == "thread-123" + assert result.usage == {**events[-1], "thread_id": "thread-123"} + + +def test_bootstrap_factory_reads_codex_env(monkeypatch): + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + monkeypatch.setenv("METAINFER_AGENT_BACKEND", "codex") + monkeypatch.setenv("METAINFER_CODEX_BIN", "codex-dev") + mgr = make_subagent_manager( + claude_bin="ccb", + model=None, + permission_mode="bypassPermissions", + effort="max", + extra_add_dirs=[], + snapshot_file=tmp / "agents.json", + ) + try: + assert mgr.agent_backend == "codex" + assert mgr.codex_bin == "codex-dev" + finally: + mgr.shutdown() + + +def test_pi_command_uses_json_print_mode(): + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + kb = tmp / "kb" + kb.mkdir() + mgr = SubAgentManager( + agent_backend="pi", + pi_bin="pi-dev", + extra_add_dirs=[kb], + default_model="anthropic/claude-sonnet", + effort="high", + ) + cmd = mgr._build_command(_spec(tmp)) + assert cmd[:4] == ["pi-dev", "--mode", "json", "-p"] + assert "--model" in cmd + assert "anthropic/claude-sonnet" in cmd + assert "--thinking" in cmd + assert "high" in cmd + # session-id pinning only when spec.session_id is set; fresh run + # leaves it to pi to mint (captured from the stream later). + assert "--session-id" not in cmd + assert "--continue" not in cmd + + +def test_pi_command_resume_threads_session(): + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + mgr = SubAgentManager(agent_backend="pi", pi_bin="pi-dev", + default_model="anthropic/claude-sonnet") + cmd = mgr._build_command(_spec(tmp, resume_session_id="01abc-def")) + assert "--session" in cmd + assert "01abc-def" in cmd + assert "--continue" in cmd + assert "--model" in cmd + + +def test_pi_session_dir_isolates_from_user_pi_index(): + """Orchestrator sessions must land in a per-task --session-dir so a + user running ``pi --continue`` from their shell (which indexes + ~/.pi) cannot attach to a production orchestrator session mid-flight. + """ + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + session_dir = tmp / "pi-sessions" + mgr = SubAgentManager( + agent_backend="pi", pi_bin="pi-dev", + pi_session_dir=session_dir, + ) + cmd = mgr._build_command(_spec(tmp)) + assert "--session-dir" in cmd + idx = cmd.index("--session-dir") + assert cmd[idx + 1] == str(session_dir.resolve()) + # the dir is created lazily so pi can write into it + assert session_dir.is_dir() + + +def test_pi_session_dir_applied_on_resume_too(): + """Resume turns must keep using the isolated session dir so the + continued session stays out of the user's ~/.pi index. + """ + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + session_dir = tmp / "pi-sessions" + mgr = SubAgentManager( + agent_backend="pi", pi_bin="pi-dev", + pi_session_dir=session_dir, + ) + cmd = mgr._build_command(_spec(tmp, resume_session_id="01abc-def")) + assert "--session-dir" in cmd + assert "--session" in cmd + assert "01abc-def" in cmd + assert "--continue" in cmd + + +def test_pi_session_dir_absent_when_not_configured(): + """Legacy callers (no state_dir) keep the old behavior: pi picks its + own default session dir. No regression for backends that never set + pi_session_dir. + """ + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + mgr = SubAgentManager(agent_backend="pi", pi_bin="pi-dev") + cmd = mgr._build_command(_spec(tmp)) + assert "--session-dir" not in cmd + + +def test_bootstrap_factory_wires_pi_session_dir_from_state_dir(monkeypatch): + """make_subagent_manager derives pi_session_dir from state_dir so every + task's orchestrator sessions are isolated under /pi-sessions. + """ + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + state_dir = tmp / "state" + state_dir.mkdir() + monkeypatch.setenv("METAINFER_AGENT_BACKEND", "pi") + mgr = make_subagent_manager( + claude_bin="ccb", + model=None, + permission_mode="bypassPermissions", + effort="max", + extra_add_dirs=[], + snapshot_file=tmp / "agents.json", + state_dir=state_dir, + ) + try: + assert mgr.agent_backend == "pi" + assert mgr.pi_session_dir == (state_dir / "pi-sessions").resolve() + finally: + mgr.shutdown() + + +def test_bootstrap_factory_no_state_dir_leaves_pi_session_dir_none(monkeypatch): + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + monkeypatch.setenv("METAINFER_AGENT_BACKEND", "pi") + mgr = make_subagent_manager( + claude_bin="ccb", + model=None, + permission_mode="bypassPermissions", + effort="max", + extra_add_dirs=[], + snapshot_file=tmp / "agents.json", + ) + try: + assert mgr.pi_session_dir is None + finally: + mgr.shutdown() + + +def test_pi_normalize_accepts_aliases(): + from metainfer.orchestrator.subagent_manager import _normalize_agent_backend + assert _normalize_agent_backend("pi") == "pi" + assert _normalize_agent_backend("PI") == "pi" + assert _normalize_agent_backend("pi-coding-agent") == "pi" + assert _normalize_agent_backend("earendil") == "pi" + + +def test_materialize_pi_jsonl_result(): + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + spec = _spec(tmp) + events = [ + {"type": "session", "version": 3, "id": "01a00ecf-aaaa-bbbb-cccc-dddddddddddd", + "timestamp": "2026-08-17T08:21:39.437Z", "cwd": "/tmp"}, + {"type": "agent_start"}, + {"type": "turn_start"}, + {"type": "message_end", "message": {"role": "user", "content": []}}, + {"type": "message_end", "message": { + "role": "assistant", + "content": [{"type": "text", "text": "hello world"}], + "usage": {"input": 581, "output": 3, "cacheRead": 0, + "cacheWrite": 0, "reasoning": 0, "totalTokens": 584, + "cost": {"input": 0, "output": 0, "cacheRead": 0, + "cacheWrite": 0, "total": 0.0123}}, + "stopReason": "stop"}}, + {"type": "turn_end", "message": {"role": "assistant", + "content": [{"type": "text", "text": "hello world"}], + "usage": {"input": 581, "output": 3, "cacheRead": 0, + "cacheWrite": 0, "totalTokens": 584, + "cost": {"total": 0.0123}}}, "toolResults": []}, + {"type": "agent_end", "messages": [ + {"role": "user", "content": []}, + {"role": "assistant", "content": [{"type": "text", "text": "hello world"}], + "usage": {"input": 581, "output": 3, "cacheRead": 0, + "cacheWrite": 0, "cost": {"total": 0.0123}}}]}, + {"type": "agent_settled"}, + ] + ef = spec.events_file(1) + ef.write_text( + "\n".join(json.dumps(e) for e in events) + "\n", + encoding="utf-8", + ) + handle = AgentHandle( + spec=spec, + attempt=1, + process=_FakeProcess(0), # type: ignore[arg-type] + started_at=time.time() - 1, + last_output_at=time.time(), + ) + mgr = SubAgentManager(agent_backend="pi") + result = mgr._materialize_result(handle, spec, 1) + assert result.success + assert result.final_text == "hello world" + assert result.session_id == "01a00ecf-aaaa-bbbb-cccc-dddddddddddd" + assert result.usage is not None + assert result.usage["usage"]["input_tokens"] == 581 + assert result.usage["usage"]["output_tokens"] == 3 + assert result.usage["usage"]["cache_read_input_tokens"] == 0 + assert result.usage["usage"]["cache_creation_input_tokens"] == 0 + assert result.usage["total_cost_usd"] == 0.0123 + assert result.usage["session_id"] == "01a00ecf-aaaa-bbbb-cccc-dddddddddddd" + + +def test_bootstrap_factory_reads_pi_env(monkeypatch): + with tempfile.TemporaryDirectory() as td: + tmp = Path(td) + monkeypatch.setenv("METAINFER_AGENT_BACKEND", "pi") + monkeypatch.setenv("METAINFER_PI_BIN", "/usr/local/bin/pi") + mgr = make_subagent_manager( + claude_bin="ccb", + model=None, + permission_mode="bypassPermissions", + effort="max", + extra_add_dirs=[], + snapshot_file=tmp / "agents.json", + ) + try: + assert mgr.agent_backend == "pi" + assert mgr.pi_bin == "/usr/local/bin/pi" + finally: + mgr.shutdown() diff --git a/metainfer/orchestrator/tests/test_token_budget.py b/metainfer/orchestrator/tests/test_token_budget.py index bf5d66b8..6af06b93 100644 --- a/metainfer/orchestrator/tests/test_token_budget.py +++ b/metainfer/orchestrator/tests/test_token_budget.py @@ -176,6 +176,24 @@ def test_usage_from_result_event(): assert rec2.total_cost_usd == 0.0 +def test_usage_from_codex_turn_completed_event(): + ev = { + "type": "turn.completed", + "thread_id": "thread-123", + "usage": { + "input_tokens": 100, + "cached_input_tokens": 80, + "output_tokens": 20, + }, + } + rec = usage_from_result_event(ev, agent="codex-a", source="orchestrator") + assert rec.input_tokens == 100 + assert rec.output_tokens == 20 + assert rec.cache_read_input_tokens == 80 + assert rec.total_cost_usd == 0.0 + assert rec.session_id == "thread-123" + + def test_reset_clears_everything(): with tempfile.TemporaryDirectory() as td: b = TokenBudget(td, max_cost_usd=1.0) diff --git a/metainfer/orchestrator/token_budget.py b/metainfer/orchestrator/token_budget.py index 7b0c8f86..f4c24ca3 100644 --- a/metainfer/orchestrator/token_budget.py +++ b/metainfer/orchestrator/token_budget.py @@ -567,17 +567,32 @@ def usage_from_result_event( usage = event.get("usage") if isinstance(event, dict) else None if not isinstance(usage, dict): usage = {} + cache_read = usage.get("cache_read_input_tokens") + if cache_read is None: + cache_read = usage.get("cached_input_tokens", 0) + # pi reports usage under its own camelCase names (input / output / + # cacheRead / cacheWrite). Tolerate them so a raw pi event also works. + input_tokens = usage.get("input_tokens", usage.get("input", 0)) + output_tokens = usage.get("output_tokens", usage.get("output", 0)) + cache_creation = usage.get( + "cache_creation_input_tokens", usage.get("cacheWrite", 0) + ) + total_cost = event.get("total_cost_usd") + if total_cost is None: + cost = usage.get("cost") + if isinstance(cost, dict): + total_cost = cost.get("total", 0.0) return UsageRecord( agent=str(agent), source=str(source), phase=phase, ended_at=time.time(), - input_tokens=int(usage.get("input_tokens", 0) or 0), - output_tokens=int(usage.get("output_tokens", 0) or 0), - cache_read_input_tokens=int(usage.get("cache_read_input_tokens", 0) or 0), - cache_creation_input_tokens=int(usage.get("cache_creation_input_tokens", 0) or 0), - total_cost_usd=float(event.get("total_cost_usd", 0.0) or 0.0), - session_id=event.get("session_id"), + input_tokens=int(input_tokens or 0), + output_tokens=int(output_tokens or 0), + cache_read_input_tokens=int(cache_read or 0), + cache_creation_input_tokens=int(cache_creation or 0), + total_cost_usd=float(total_cost or 0.0), + session_id=event.get("session_id") or event.get("thread_id"), ) diff --git a/metainfer/tasks/calc_value/orchestrator/orchestrator.py b/metainfer/tasks/calc_value/orchestrator/orchestrator.py index d50bd96b..518e2fbe 100644 --- a/metainfer/tasks/calc_value/orchestrator/orchestrator.py +++ b/metainfer/tasks/calc_value/orchestrator/orchestrator.py @@ -251,6 +251,7 @@ def run_with_requirements( snapshot_file=paths["agents_file"], max_concurrent=5, budget=budget, + state_dir=state_dir, ) # Wire the hard-exhausted callback NOW that the manager exists. # When the hard threshold is crossed, every in-flight agent gets diff --git a/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md b/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md index f32dc992..ea1cfe68 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md +++ b/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md @@ -50,7 +50,7 @@ metainfer/tasks/dcu_kernel_auto_opt/ ### 2.1 新建任务（表单字段，见 form.yaml） 关键字段：`operator`（Quantized GEMM）、`kernel_language`（HIP C++）、`target_hardware`（K500SM_AI/gfx928）、 -`dtype`（INT8 W8A8）、`agent_framework`（ccb / dsh）、`agent_model`（ccb: Opus/Sonnet；dsh: deepseek-v4-flash）、`execution_mode`（Mock / Real INT8 W8A8 GEMM / +`dtype`（INT8 W8A8）、`agent_framework`（ccb / dsh）、`agent_model`（ccb: Opus/Sonnet；dsh: deepseek-flash-4.1 / deepseek-v4-flash）、`execution_mode`（Mock / Real INT8 W8A8 GEMM / Generate & optimize / Infra smoke）、`target_repo_path`、`shape_assignment_mode`（AI automatic / Manual by GPU）、 `shape_scope`（All API shapes / Selected shapes only）、`shape_config`、`max_iterations`、`minimum_improvement_percent`、`extra_notes`。 @@ -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） @@ -177,11 +177,25 @@ TP8 另有 wq_b/wo_b (1024,4096)、gate_up (4096,512)、down (256,4096)、indexe 注意 TP1 目录前端可选、契约仍只收 tp_size 4/8，选 TP1 会同样在 prepare 被拒。 - **模型目录 TP8 大 prefill 边界（2026-08-27）**：Hy3/MiniMax M3/GLM5.2 的 TP8 形状支持 M=4096 优化（"Selected shapes only" 手动选择，如 minimax 任务）。契约加 - `MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES=(4096,)`（**不进** `DEFAULT_OPTIMIZATION_SHAPES`， + `MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES`（**不进** `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。 + `mValues`（用 `W8A8_MODEL_TP8_M_VALUES` 常量）；基线表 15 条 `(8,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。 +- **模型目录 TP8 短 decode 边界（2026-09-20）**：Hy3/MiniMax M3/GLM5.2 的 TP8 形状再加 + M=8，即 `MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES=(8, 4096)`；前端 + `W8A8_MODEL_TP8_M_VALUES = [2, 8, 16, 3072, 4096]`（DeepSeek TP1/TP8 仍是 + `W8A8_M_VALUES=[2,16,3072]`，TP4 仍是 `W8A8_TP4_M_VALUES=[2,16,3072,4096]`）。 + 基线表 15 条 `(8,8,…)` 在 worker29/gfx928 实测三次 + （`tools/baseline/bench_triton_tp8_m8.py --m 8 --mode graph`，协议同 2026-08-12 目录： + CUDA-graph replay、hot cache、warmups=10/samples=20/launches=5、BM16/BN32/BK256）， + 公布值取「每次 pass median 再取 median」；原始数据 + `tools/baseline/tp8_m8_graph_run{1,2,3}.json`，汇总 `tp8_m8_graph_merged.json`。 + 12/15 形状三次 pass 极差 ≤1.6%（共享卡上 pass 2 的 minimax shared_down_proj 被同租户 + 干扰到 33.76 µs，取 median 后为 19.98 µs）。DeepSeek TP8 不加 M=8。 +- **面板/契约/基线三层一致性测试**：`tests/test_shape_catalog_panel.py` 用 node 真实求值 + `static/dkao-shape-input.js` 的 `MODEL_WORKLOADS`，断言面板给出的每个 TP4/TP8 shape + 都能过契约校验且能查到基线（否则任务要到 baseline 阶段才报 ValueError）。 - `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 +268,12 @@ 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_m4096.py`、 + `bench_triton_tp8_m4096.py`、`bench_triton_tp8_m8.py` + `int8_utils.py`）。 + kernel-repos 默认在 MetaInfer 同级（`METAINFER_KERNEL_REPOS` 可改），不随插件目录走。 5. 改动任何行为后：更新本文件相关段落 + 跑 tests + 用真实任务验证（优先在 zth_meta 里）。 ## 7. 不确定性标注 @@ -264,3 +282,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/api/int8w8a8gemm/int8_w8a8_gemm_api.py b/metainfer/tasks/dcu_kernel_auto_opt/api/int8w8a8gemm/int8_w8a8_gemm_api.py index 2951f48c..5f52266f 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/api/int8w8a8gemm/int8_w8a8_gemm_api.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/api/int8w8a8gemm/int8_w8a8_gemm_api.py @@ -193,14 +193,19 @@ def _merge_kn_tables( # the original three M values, so each TP4 operator gets exactly one M=4096 # shape while TP8 defaults are unchanged. TP4_EXTRA_OPTIMIZATION_M_VALUES: Final[tuple[int, ...]] = (4096,) -# Model-catalog TP8 large-prefill boundary added on 2026-08-27. The DeepSeek -# TP8 default workload keeps the original three M values; Hy3 / MiniMax M3 / -# GLM5.2 TP8 operators are additionally optimizable at M=4096 via -# "Selected shapes only" (their (K,N) pairs are already in the allowed sets -# and their Triton CUDA-graph baselines are measured). This constant is not -# applied by _default_optimization_shapes(), so DEFAULT_OPTIMIZATION_SHAPES +# Model-catalog TP8 extra M values. The DeepSeek TP8 default workload keeps +# the original three M values; Hy3 / MiniMax M3 / GLM5.2 TP8 operators are +# additionally optimizable at these M values via "Selected shapes only" +# (their (K,N) pairs are already in the allowed sets and their Triton +# CUDA-graph baselines are measured). This constant is not applied by +# _default_optimization_shapes(), so DEFAULT_OPTIMIZATION_SHAPES # (DeepSeek-only) and its serial-validation fallback scope are unchanged. -MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES: Final[tuple[int, ...]] = (4096,) +# M=4096: model-catalog large-prefill boundary added on 2026-08-27. +# M=8: model-catalog short decode boundary added on 2026-09-20; the +# Triton Graph baselines for the 15 catalog (K, N) pairs were +# measured on worker29/gfx928 with tools/baseline/ +# bench_triton_tp8_m8.py (raw passes in tp8_m8_graph_run{1,2,3}.json). +MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES: Final[tuple[int, ...]] = (8, 4096) def _default_optimization_shapes() -> tuple[dict[str, int | str], ...]: 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..af7af3ab 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py @@ -8,8 +8,11 @@ import json import math import os +import re import statistics +import subprocess import sys +import time from pathlib import Path try: @@ -18,6 +21,159 @@ torch = None # type: ignore[assignment] +#: Device admission rule for anything that measures: VRAM <= 90% and HCU == 0. +#: This harness is the only thing that runs the candidate kernel, and an agent +#: can invoke it directly from its own shell, so the rule is enforced here too — +#: not only in the orchestrator that happens to launch us. A number taken while +#: another workload shares the card is noise (the same unchanged shape has been +#: seen at 884us -> 3047us), so we refuse to produce one. +GATE_EXIT_CODE = 75 +_GATE_DRM_ROOT = Path("/sys/class/drm") +_GATE_SMI_ROW = re.compile( + r"^(?P\d+)\s+[\d.]+C\s+[\d.]+W\s+\S+\s+[\d.]+W\s+" + r"(?P[\d.]+|N/A|n/a)%\s+(?P[\d.]+|N/A|n/a)%" +) + + +def _gate_number(value): + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _gate_hy_smi() -> dict: + """``{gpu: {vram_percent, util_percent}}`` from hy-smi; empty if unreadable.""" + for binary in ("hy-smi", "rocm-smi"): + try: + proc = subprocess.run([binary], capture_output=True, text=True, + timeout=10) + except (OSError, subprocess.SubprocessError): + continue + if proc.returncode != 0 or not proc.stdout.strip(): + continue + out = {} + for line in proc.stdout.splitlines(): + match = _GATE_SMI_ROW.match(line.strip()) + if not match: + continue + out[int(match.group("index"))] = { + "vram_percent": _gate_number(match.group("vram").strip("%")), + "util_percent": _gate_number(match.group("util").strip("%")), + } + if out: + return out + return {} + + +def _gate_sysfs() -> dict: + """Driver sysfs fallback: works when the management interface says N/A.""" + out = {} + cards = [] + try: + entries = sorted(_GATE_DRM_ROOT.iterdir(), key=lambda p: p.name) + except OSError: + return out + for entry in entries: + if not re.fullmatch(r"card\d+", entry.name): + continue + device = entry / "device" + try: + uevent = (device / "uevent").read_text(encoding="utf-8") + except OSError: + continue + if not any(driver in uevent for driver in ("hycu", "amdgpu", "hydcu")): + continue + if not (device / "mem_info_vram_total").is_file(): + continue + cards.append(device) + for index, device in enumerate(cards): + row = {} + try: + row["util_percent"] = float( + (device / "gpu_busy_percent").read_text().strip()) + except (OSError, ValueError): + pass + try: + used = float((device / "mem_info_vram_used").read_text().strip()) + total = float((device / "mem_info_vram_total").read_text().strip()) + if total > 0: + row["vram_percent"] = round(100.0 * used / total, 2) + except (OSError, ValueError, ZeroDivisionError): + pass + if row: + out[index] = row + return out + + +def _gate_violation(device: int) -> str: + """``""`` when the device passes, else why it does not.""" + from_smi = _gate_hy_smi().get(device) or {} + complete = (from_smi.get("vram_percent") is not None + and from_smi.get("util_percent") is not None) + readings = from_smi if complete else dict(_gate_sysfs().get(device) or {}) + if not complete: + for key in ("vram_percent", "util_percent"): + if readings.get(key) is None and from_smi.get(key) is not None: + readings[key] = from_smi[key] + vram = readings.get("vram_percent") + util = readings.get("util_percent") + if vram is None or util is None: + return "device state unavailable (no HCU/VRAM reading)" + if vram > 90.0: + return f"VRAM {vram:.0f}% > 90%" + if util > 0.0: + return f"HCU {util:.0f}% > 0%" + return "" + + +def enforce_measurement_gate(device: int | None = None) -> None: + """Wait for ``VRAM <= 90% and HCU == 0``; exit if it never happens. + + ``METAINFER_GPU_PREFLIGHT=0`` is the deliberate opt-out. Waiting is + deliberately bounded (``METAINFER_GATE_WAIT_SECONDS`` x + ``METAINFER_GATE_MAX_WAITS``, default 30 min x 48 = 24 h): a shared card + must never silently turn into a measurement, and waiting forever is not an + answer either. + """ + if str(os.environ.get("METAINFER_GPU_PREFLIGHT", "")).strip().lower() in { + "0", "false", "no", "off"}: + return + if device is None: + raw = str(os.environ.get("HIP_VISIBLE_DEVICES") + or os.environ.get("ROCR_VISIBLE_DEVICES") or "0") + device = int(raw.split(",")[0].strip() or 0) + try: + wait_s = int(os.environ.get("METAINFER_GATE_WAIT_SECONDS", "") or 1800) + except ValueError: + wait_s = 1800 + try: + max_waits = int(os.environ.get("METAINFER_GATE_MAX_WAITS", "") or 48) + except ValueError: + max_waits = 48 + for attempt in range(1, max_waits + 1): + reason = _gate_violation(int(device)) + if not reason: + return + if attempt > max_waits: + break + print(json.dumps({ + "measurement_gate": "blocked", "device": int(device), + "attempt": attempt, "max_waits": max_waits, + "wait_seconds": wait_s, "reason": reason, + }), file=sys.stderr) + if attempt == max_waits: + break + time.sleep(max(0.0, float(wait_s))) + reason = _gate_violation(int(device)) or "device never became idle" + print(json.dumps({ + "measurement_gate": "give_up", "device": int(device), + "checks": max_waits, "reason": reason, + }), file=sys.stderr) + raise SystemExit(GATE_EXIT_CODE) + + + def load_module(module_path: Path, name: str): spec = importlib.util.spec_from_file_location( name, module_path @@ -191,9 +347,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") @@ -312,6 +487,11 @@ def main() -> int: }, sort_keys=True)) return 0 + # Everything below initialises the device and measures on it. CPU-only runs + # (--self-test, --prepare-reference) returned above, so they are never + # blocked; every device-touching run has to pass the gate first. + enforce_measurement_gate() + source = args.source.resolve() fixed_contract = source / "int8_w8a8_gemm_api.py" if fixed_contract.is_file(): diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/README.md b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/README.md index 22aace2d..2444e929 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/README.md +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/README.md @@ -32,8 +32,9 @@ prompt 从 stdin 传入，stdout 输出逐行 stream-json 事件 WebUI new-task 表单中 dcu-kernel-auto-opt 增加 **Agent framework** 选择： - `ccb` → Claude Code（模型 `Sonnet` / `Opus`） -- `dsh` → DeepSeek Harness（模型 `deepseek-v4-flash`，即 - `deepseek/deepseek-v4-flash-0731`） +- `dsh` → DeepSeek Harness： + - `deepseek-flash-4.1` → `deepseek/deepseek-flash`（**默认**，4.1 Flash） + - `deepseek-v4-flash` → `deepseek/deepseek-v4-flash-0731`（旧 pinned 构建，可复现历史运行） orchestrator CLI 在 `agent_framework=dsh` 时自动把 `claude_bin` 指向本目录的 `dsh_agent.py`（`--claude-bin` 显式传入时优先），无需改动 MetaInfer 共享代码。 @@ -55,7 +56,8 @@ runtime 只内置 `deepseek-official` adapter（`llm-deepseek`），它通过 - provider：`deepseek-official`（runtime adapter；TokenHub 经 base_url 覆盖） - baseURL：`https://tokenhub.tencentmaas.com/plan/v3` -- 模型：`deepseek/deepseek-v4-flash-0731`（默认） +- 模型：`deepseek/deepseek-flash`（默认，= 表单里的 `deepseek-flash-4.1`）； + 旧 pinned 构建 `deepseek/deepseek-v4-flash-0731`（= `deepseek-v4-flash`） - API key：`~/.dsh/.credentials.yaml` 的 `TENCENT_API_KEY` `dsh_agent.py` 的环境变量： @@ -63,7 +65,7 @@ runtime 只内置 `deepseek-official` adapter（`llm-deepseek`），它通过 | 变量 | 默认 | 说明 | |---|---|---| | `DSH_AGENT_PROVIDER` | `deepseek-official` | DSH runtime provider（dev runtime 仅内置此 adapter） | -| `DSH_AGENT_MODEL` | `deepseek/deepseek-v4-flash-0731` | 模型覆盖 | +| `DSH_AGENT_MODEL` | `deepseek/deepseek-flash` | 默认模型覆盖（只改 `deepseek-flash-4.1` 这个默认标签的 id） | | `DSH_AGENT_BASE_URL` | `https://tokenhub.tencentmaas.com/plan/v3` | 模型端点（`DEEPSEEK_BASE_URL` 优先） | | `TENCENT_API_KEY` / `DEEPSEEK_API_KEY` | 凭据文件兜底 | API key | | `DSH_AGENT_CORDIS` | 本目录 `cordis.yml` | 自定义组合路径 | 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..e4aaa77b 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 @@ -27,7 +27,8 @@ DSH_AGENT_PROVIDER provider name for the DSH runtime (default: deepseek-official — the only adapter shipped by the dev-checkout runtime) - DSH_AGENT_MODEL model override (default deepseek/deepseek-v4-flash-0731) + DSH_AGENT_MODEL model override (default deepseek/deepseek-flash, + the 4.1 Flash label deepseek-flash-4.1) DSH_AGENT_BASE_URL model endpoint; falls back to DEEPSEEK_BASE_URL (default: https://tokenhub.tencentmaas.com/plan/v3) TENCENT_API_KEY preferred API key (matches ~/.dsh/settings.yaml) @@ -71,9 +72,11 @@ def default_provider() -> str: def default_model() -> str: + """Default DSH model id (the 4.1 Flash build) unless DSH_AGENT_MODEL pins + another one. Matches orchestrator/config.py DSH_DEFAULT_MODEL_ID.""" return ( - os.environ.get("DSH_AGENT_MODEL", "deepseek/deepseek-v4-flash-0731") - .strip() or "deepseek/deepseek-v4-flash-0731" + os.environ.get("DSH_AGENT_MODEL", "deepseek/deepseek-flash") + .strip() or "deepseek/deepseek-flash" ) @@ -144,13 +147,20 @@ def default_api_key() -> str: # --------------------------------------------------------------------------- # -# Model mapping (legacy Claude labels -> this host's DSH model ids) +# Model mapping (legacy Claude labels / form labels -> this host's DSH ids) # --------------------------------------------------------------------------- # +#: Form labels produced by both MetaInfer task frontends (dcu-kernel-auto-opt's +#: agent_model field and harness-evolve's evolve_model field). Kept in sync with +#: orchestrator/config.py (DSH_MODEL_IDS) — the orchestrator resolves the label +#: itself, so this map only serves callers that hand the wrapper a bare label. _MODEL_MAP = { - "opus": "deepseek-v4-pro", - "sonnet": "deepseek-v4-flash", - "haiku": "deepseek-v4-flash", + "deepseek-flash-4.1": "deepseek/deepseek-flash", + "deepseek-flash": "deepseek/deepseek-flash", + "deepseek-v4-flash": "deepseek/deepseek-v4-flash-0731", + "opus": "deepseek/deepseek-v4-flash-0731", + "sonnet": "deepseek/deepseek-v4-flash-0731", + "haiku": "deepseek/deepseek-v4-flash-0731", } @@ -160,11 +170,8 @@ def map_model(requested: Optional[str]) -> str: key = requested.strip().lower() if key in _MODEL_MAP: return _MODEL_MAP[key] - if key == "deepseek-v4-flash": - # Bare user-facing label -> the pinned host model id. - return default_model() if key.startswith("deepseek"): - # Full model id (e.g. deepseek/deepseek-v4-flash-0731): pass through. + # Full model id (e.g. deepseek/deepseek-flash): pass through. return key # Unknown label: keep the caller's intent but stay on a known model id. return default_model() @@ -189,6 +196,235 @@ 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. + + Model selection caveat: the ``dsh --profile headless`` app takes no + ``--model`` flag, so it runs whatever the user's DSH settings document + (``~/.dsh/settings.yaml`` -> ``agent-default-model``) selects — the host + default, not ``--model``. Emit the requested model in the system event so + the mismatch is visible instead of silent. The SDK backend below does honour + ``--model``. + """ + 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 +455,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 +472,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/run_dsh_task.sh b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/run_dsh_task.sh index 49196165..ff344660 100755 --- a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/run_dsh_task.sh +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/run_dsh_task.sh @@ -12,7 +12,8 @@ # # Optional env: # DSH_AGENT_MAX_TOKENS per-request output cap (default 65536) -# DSH_AGENT_MODEL model override (default deepseek/deepseek-v4-flash-0731) +# DSH_AGENT_MODEL model override (default deepseek/deepseek-flash, +# the 4.1 Flash build) # METAINFER_GPU_IDS restrict GPUs, e.g. 4,5,6,7 set -euo pipefail diff --git a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/smoke_sdk.py b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/smoke_sdk.py index 94f29330..39868d26 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/smoke_sdk.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/bridge/dsh/tests/smoke_sdk.py @@ -30,7 +30,7 @@ def main() -> int: WORK.mkdir(parents=True, exist_ok=True) cfg = DeepSeekHarnessConfig( provider="deepseek-official", - model="deepseek-v4-flash", + model="deepseek/deepseek-flash", max_tokens=2048, cwd=str(WORK), session_root=str(SESSION_ROOT), 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..bd2d9b53 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 @@ -5,7 +5,8 @@ replaced by a fake whose behavior is scripted per test. Verify: 1. arg parsing tolerates every SubAgentManager flag (incl. unknown ones) - 2. model mapping: sonnet/haiku -> deepseek-v4-flash, opus -> deepseek-v4-pro + 2. model mapping: form labels deepseek-flash-4.1 / deepseek-v4-flash and the + legacy Claude labels resolve to real DSH model ids 3. the emitted stream contains system -> assistant* -> result, exit 0 4. a resume failure (finish_reason == "error") falls back to a fresh session and the stream's first session_id is the fallback one @@ -96,16 +97,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 +127,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() @@ -154,10 +167,40 @@ def test_model_mapping(self): mod = install_fake_sdk([ ("s1", "ok", "completed"), ]) - # opus -> deepseek-v4-pro + # Legacy Claude labels stay routable (mapped onto the pinned V4 build). run_wrapper(mod, [a if a != "sonnet" else "opus" for a in BASE_ARGS], "hi") cfg = FakeDeepSeekHarness.instances[0].config - self.assertEqual(cfg.model, "deepseek-v4-pro") + self.assertEqual(cfg.model, "deepseek/deepseek-v4-flash-0731") + + def test_form_labels_map_to_dsh_model_ids(self): + """Form labels from both frontends resolve to real platform ids.""" + cases = { + "deepseek-flash-4.1": "deepseek/deepseek-flash", # new default + "deepseek-v4-flash": "deepseek/deepseek-v4-flash-0731", + # Full ids pass through untouched. + "deepseek/deepseek-flash": "deepseek/deepseek-flash", + } + for label, expected in cases.items(): + mod = install_fake_sdk([("s1", "ok", "completed")]) + args = [a if a != "sonnet" else label for a in BASE_ARGS] + code, _out, err = run_wrapper(mod, args, "hi") + self.assertEqual(code, 0, err) + cfg = FakeDeepSeekHarness.instances[0].config + self.assertEqual(cfg.model, expected, label) + + def test_default_model_is_flash_41(self): + """With no --model the wrapper defaults to the 4.1 Flash build.""" + mod = install_fake_sdk([("s1", "ok", "completed")]) + args = [a for a in BASE_ARGS if a not in ("--model", "sonnet")] + old = os.environ.pop("DSH_AGENT_MODEL", None) + try: + code, _out, err = run_wrapper(mod, args, "hi") + finally: + if old is not None: + os.environ["DSH_AGENT_MODEL"] = old + self.assertEqual(code, 0, err) + cfg = FakeDeepSeekHarness.instances[0].config + self.assertEqual(cfg.model, "deepseek/deepseek-flash") def test_resume_fallback(self): mod = install_fake_sdk([ diff --git a/metainfer/tasks/dcu_kernel_auto_opt/form.yaml b/metainfer/tasks/dcu_kernel_auto_opt/form.yaml index 96215dc2..ae870e84 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/form.yaml +++ b/metainfer/tasks/dcu_kernel_auto_opt/form.yaml @@ -66,10 +66,10 @@ - label: "ccb" description: "Claude Code (claude-sonnet-5 / claude-opus-5)" - label: "dsh" - description: "DeepSeek Harness (deepseek-v4-flash)" + description: "DeepSeek Harness (deepseek-flash-4.1 / deepseek-v4-flash)" - key: agent_model - question: "Model used by the selected agent framework (ccb: Sonnet/Opus; dsh: deepseek-v4-flash)." + question: "Model used by the selected agent framework (ccb: Sonnet/Opus; dsh: deepseek-flash-4.1 / deepseek-v4-flash)." header: "Agent model" required: true multi: false @@ -81,8 +81,10 @@ description: "Pinned to Claude Opus 5 (claude-opus-5) for difficult kernel architecture and correctness work" - label: "Sonnet" description: "Pinned to Claude Sonnet 5 (claude-sonnet-5) for faster, lower-cost optimization runs" + - label: "deepseek-flash-4.1" + description: "DeepSeek Harness 4.1 Flash (deepseek/deepseek-flash) — current default DSH model" - label: "deepseek-v4-flash" - description: "DeepSeek Harness model (deepseek-v4-flash)" + description: "DeepSeek Harness V4 Flash, pinned build (deepseek/deepseek-v4-flash-0731)" - key: execution_mode question: "Task execution mode for this run." @@ -117,6 +119,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" @@ -200,3 +226,16 @@ header: "Notes" required: false form: textarea + +- key: execution_profile + question: "执行画像：production = 生产任务（走 broker 的最高优先级、排在 HE 之前；仍然必须过 VRAM<=90%/HCU==0 的准入闸门，因为共享卡上的测量是脏数据）；default = 现状" + header: "Profile" + required: false + form: select + default: "default" + options: + - label: "default" + description: "与现状一致（调控/占据模式按 gpu_mode 决定）" + - label: "production" + description: "生产模式：优先级 100 插队、优先派发，仍过闸门、不用等 HE 的轮次" + 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..67372060 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/README.md @@ -0,0 +1,61 @@ +# harness_default —— DKAO 可演化 harness 的种子目录 + +这是 AHE「组件可观测性」的落地目录：DKAO 的 harness（验收门阈值、轮次指令文案、 +planner 策略、以及后续的 skills/工具/中间件/记忆）**文件化**在这里，成为可版本化、 +可校验、可演化的组件。AHE 外循环（`harness_evolve`）以本目录为种子，拷贝出可写 +workspace（git 仓库），每轮评测 pin 一个快照；Evolve Agent 只改 workspace。 + +## 目录内容与接线状态 + +``` +harness_default/ +├── manifest.yaml # 组件清单 + 每个组件的 wired 开关（AHE 可写空间的入口描述） +├── gates.yaml # 验收/plateau/ISA 门阈值 ← 已接线 gate_policy.gates() +├── planner_policy.yaml # 状态→方案选择规则（P0-P4 分层） ← 已接线 planner.policy() +├── planner_catalog.yaml # 方案目录：稳定 plan id + 面向 worker 的描述 ← 已接线 planner.catalog() +├── systemprompt/ +│ └── round_strategy.yaml # 每轮给 worker 的指令文案（P0 修复/P1 阶段/各 M regime 组合） +│ # ← 已接线 round_strategy.load_round_strategy() +├── README.md # 本文件 +└── (后续切片加入) + ├── systemprompt/ # 其余 prompt 模板（coordinator/bootstrap/worker/synthesis） + ├── tools/ # cordis 组合与 per-role 工具白名单 + ├── middleware/ # resume/compaction/重试策略参数 + ├── skills/ # SKILL.md 种子 + └── memory/ # LongTermMEMORY（measured 事实） +``` + +## 语义（重要） + +**`wired` 是总开关。** `manifest.yaml` 里每个组件标 `wired: true` 时，运行时读该文件； +标 `wired: false`（或删掉该组件项）时，运行时回落到 Python 内置默认值 —— 也就是说 +**编辑一个未接线的组件不会改变任何行为**，而 `harness_evolve` 的 +`_wired_scope_violations()` 会把这种"改了但不生效"的改动报出来。 + +**回落是逐键的。** 文件缺失、YAML 损坏、某个键写错类型或写成空串，都只让**那一个键** +回落到内置文本，其余照常生效；`RoundStrategy.notes` / 加载器会记录原因，但**绝不抛异常** +（一个手改坏的 harness 是性能问题，不该让一轮评测崩掉）。 + +**决策逻辑不在文件里。** 以 `round_strategy.yaml` 为例：走哪个分支由实测状态（history / +ISA 阶段 / PMC 证据 / M regime / 轮次预算）在 `orchestrator/round_strategy.py` + +`prompts.py` 里判定，文件只承载每个分支的**文案**。这样"这轮让 worker 试什么"可演化， +而选择状态机本身不可被改写。 + +**一致性由测试守卫**（drift guard，见 manifest 每个组件的 `drift_guard`）： +每个 YAML 种子的取值都被断言等于 Python 内置默认值，因此两条来源不会无声漂移。 +`tests/test_round_strategy.py` 另外用 5400 个渲染用例的 sha256 指纹把文案钉死 —— 任何 +意外改词都会让该测试失败。**有意改文案时**：同步改 `BUILTIN_STRATEGY` 与种子文件， +并有意重算指纹。 + +## 读取方式 + +- 根目录解析：`orchestrator/harness_io.harness_root()` —— 优先环境变量 + `METAINFER_HARNESS_ROOT`（AHE 用它指向本轮候选快照），否则回落本目录。 +- 组件读取：`gate_policy.gates()` / `planner.policy()` / `planner.catalog()` / + `round_strategy.load_round_strategy()`；通用 YAML 读取用 + `harness_io.load_component_yaml()`（返回 `(data, error)`，不抛异常）。 +- 播种：`harness_io.seed_workspace(dst)` 或 CLI 路径上的任务建仓逻辑会把整棵树拷进 + kernel repo 的 `harness_snapshot/`，并把 `source` / `revision` / 逐文件 digest 写进 + `scaffold_manifest.json`，用于"哪一轮用了哪版 harness"的追溯。 +- 缓存：各加载器按 `(路径, 文件 mtime, manifest mtime)` 缓存，所以运行中途换 harness + 会被下一轮读到（测试亦用 `reset_cache()` 隔离）。 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..955405d9 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/manifest.yaml @@ -0,0 +1,46 @@ +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 + systemprompt: + path: systemprompt/round_strategy.yaml + kind: data + role: >- + per-round mandate wording handed to the worker (round_strategy.yaml is + the only file today); the branch decision stays in + orchestrator/round_strategy.py, only the text is data + wired: true + drift_guard: tests/test_round_strategy.py +layout_note: >- + Future components land here as files: systemprompt/ (worker-facing + instruction wording; systemprompt/round_strategy.yaml is seeded), + 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..2a8868ee --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_catalog.yaml @@ -0,0 +1,59 @@ +# planner 方案目录（v0, seed） +# +# DKAO 优化方案选择机制的「方案目录」：把"每一轮可以试什么方案"枚举成稳定 id， +# 供 orchestrator/planner.py（状态条件化选择器 v0）与后续 harness 演化使用。 +# +# wired: true（以 manifest.yaml 为准）—— orchestrator/planner.py::catalog() 会读本文件， +# 但其选择结果只在 METAINFER_PLANNER=1 时决定轮次指令；planner 关闭时轮次菜单仍由 +# prompts.py::w8a8_round_strategy 出文字（该文案已外部化为 +# systemprompt/round_strategy.yaml，见 manifest 的 systemprompt 组件）。 +# 规则分层见 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/harness_default/systemprompt/round_strategy.yaml b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/systemprompt/round_strategy.yaml new file mode 100644 index 00000000..4fa491e9 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/systemprompt/round_strategy.yaml @@ -0,0 +1,204 @@ +schema_version: 1 +component: systemprompt +status: seed +wired: true +role: >- + Per-round mandate wording handed to the optimization worker ("what to try + this round"). The decision logic stays in + orchestrator/round_strategy.py + prompts.py: measured history, ISA phase, PMC + evidence and the M regime select the branch. This file only carries the text + of each branch, so the Evolve loop can version and evolve the instructions + without making the round-selection state machine writable. +loader: orchestrator/round_strategy.load_round_strategy +read_at_runtime: >- + orchestrator/prompts.py::w8a8_round_strategy, reached from + orchestrator/w8a8_pipeline.py (planner disabled, and the planner's + legacy_menu deferral). DKAO reads it through METAINFER_HARNESS_ROOT. +format_fields: >- + Templates may use these str.format fields only: {iteration}, {speedup}, + {artifact_dir}, {completed}, {split_candidates}, {grid_blocks}, {cu_count}. + An unknown field makes that one suffix fall back to the built-in text + instead of raising, so a malformed harness degrades to a weaker instruction + rather than breaking a round. +drift_guard: tests/test_round_strategy.py::test_seed_matches_builtin_defaults + +# P0 repairs: a broken previous attempt outranks any exploration plan. +repair: + faster_wrong: >- + Highest priority: repair the faster but incorrect candidate from iteration + {iteration} (measured speedup {speedup}x). Read its archived source at + `{artifact_dir}` and preserve the fast mapping. Fix only the smallest + correctness defect: signed int8 unpacking, tail bounds, scale indexing, + bf16 conversion, or a race. Do not replace it with an unrelated + architecture. + infrastructure_failure: >- + The preceding attempt failed in the agent infrastructure. Return to the + accepted best source and start a new bounded experiment; do not repair or + replay a partially written candidate. This failure does not count as a + completed optimization round. + build_failure: >- + Repair the immediately preceding candidate from iteration {iteration} at + `{artifact_dir}`. Keep its strategy and make only the minimum + compile/API/syntax correction; do not start another redesign this round. + +# P1 phase gates: the ISA-gated rounds, chosen by the control plane. +phases: + isa_guided_hip: >- + ISA-guided HIP round. This is successful ISA experiment {completed} of at + least 2. Select one measured memory or compute bottleneck, compare the + exact primary-kernel ISA with the preceding code object, and make one + HIP/DUMMA/intrinsic code-shaping change. Raw inline asm remains forbidden. + Record a compiler limitation only when the before/after binary proves it + and name the exact target instructions. + conditional_inline_asm: >- + Conditional inline-asm experiment. Target only the compiler limitation and + exact instructions verified by the immediately preceding ISA-guided HIP + round. Keep the asm block minimal, preserve complete constraints/clobbers, + and reject it unless the candidate ISA, exact correctness, median, P90, and + resources all validate. Do not write raw global/buffer/flat loads or MMAC. + +# Evidence-qualified suffixes appended to a portfolio mandate. +append: + split_candidates: >- + Trusted occupancy-probe split candidates for the measured CU count are + {split_candidates}. They include non-power-of-two values where useful, are + not a whitelist, and must fit the workspace and stage-alignment + constraints. Explore outside this set when evidence supports it. + grid_warning: >- + Trusted control-plane warning: current grid has {grid_blocks} blocks for + {cu_count} CUs, below the two-blocks-per-CU latency-hiding target. Before + micro-optimization, benchmark a finer one-wave zero-barrier grid or + multiple legal split-K candidates including combine cost. + +# Portfolios: the mandate for each round index, per M regime. +# small = M < 16, m16 = 16 <= M < 128, large = M >= 128. +# late_start = max(9, max_iterations - 1): rounds before it use the numbered +# slot below; from it on, slot 9 is used while iteration < max_iterations and +# slot 10 on the final round. The small regime stops at slot 8 and keeps that +# consolidation mandate to the end. +portfolios: + small: + "1": Vectorize contiguous K loads with exact signed-int8 semantics. + "2": >- + Increase instruction-level parallelism with independent int32 + accumulators or adjacent N outputs; avoid per-K-tile LDS barriers. + "3": >- + Stage all of tiny A once with at most one barrier, or tune unroll one + step if whole-A staging is not cheaper. + "4": >- + Change one launch variable only: waves per block, N columns per wave, or + unroll factor. + "5": >- + HIP-only memory round: change one vector-load width, contiguous N + mapping, or whole-A reuse decision. Raw inline asm is forbidden. + "6": >- + HIP-only pipeline round: reduce one dependency chain or barrier using + ordinary HIP/intrinsics. Raw inline asm is forbidden. + "7": >- + HIP-only resource round: tune one block size, unroll factor, or live + range while preserving coalescing. Raw inline asm is forbidden. + "8": >- + HIP-only consolidation round: revisit the fastest correct archived + mapping and make one final architecture/codegen improvement. Raw inline + asm is forbidden. + m16: + "1": >- + Establish a minimal 16x16x32 DUMMA tile with the exact API below, one + wave per output tile and explicit int32 accumulation. If the seed already + has a correct DUMMA kernel, preserve it and instead test the smallest + one-wave-per-block, one-N-tile geometry with no cross-wave barrier; do + not spend the round reimplementing it. + "2": >- + Architecture round: measure grid parallelism before polishing. Explore + one complete launch geometry among 1/2/4 waves per block and 1/2/4 + adjacent N tiles. Prefer enough independent blocks to cover at least all + device CUs; report grid_blocks, waves_per_block and estimated_active_cus + in proposal.json. + "3": >- + Architecture round: if the unsplit grid has fewer than two blocks per + device CU and K >= 1024, implement and measure split-K=2 plus at least + one CU-aligned candidate (which may be non-power-of-two), or test a + one-wave zero-barrier geometry that reaches the same parallelism. Write + int32 partials into the caller workspace and include the combine+scale + kernel in the timed Graph. + "4": >- + Architecture/pipeline round: explore one of multi-N-tile reuse, A-only + staging, or bounded register/LDS prefetch. Retain enough blocks to cover + all CUs, state which A/B bytes are reused, and measure whether the change + improves normal median/P90. + "5": >- + HIP-only packed-weight/staging round: compare one packed layout, A-only + staging, or B-only staging design. Raw inline asm is forbidden. + "6": >- + Pipeline round: choose exactly one staging family from direct, A-only + LDS, B-only LDS, or A+B LDS using L2/VMEM evidence. Use coalesced 8- or + 16-byte cooperative loads and report HBM/LDS byte changes; do not claim + asynchronous overlap without evidence. + "7": >- + Pipeline round: compare single buffering with double buffering only when + K>=1024, L2 hit rate is below 70%, and the doubled LDS budget stays below + 48 KiB. Count barriers per K step. + "8": >- + HIP-only resource round: tune one occupancy limiter using actual PMC + evidence: waves per block, VGPR live range, LDS footprint, or spill + removal. Do not trade repeated HBM reads for occupancy. + "9": >- + Late ISA-diagnosis round. Only if the control-plane plateau gate is open, + use one selected ISA Skill and trusted disassembly to shape compiler + output through HIP/DUMMA/intrinsics. Raw inline asm remains forbidden. + Otherwise continue HIP-only exploration. + "10": >- + Final conditional inline-asm round. Raw asm is allowed only when the + control plane confirms a HIP plateau and the prior ISA-guided round + recorded one concrete compiler limitation plus target instructions. + Otherwise make one HIP-only consolidation change. + large: + "1": >- + Establish a correct DUMMA throughput baseline using a 2-D macro-tile. + Benchmark 64x64, 64x128, and 128x64 block tiles; record waves per block, + VGPRs, LDS bytes, occupancy and TOPS. + "2": >- + Operand-reuse round: compare direct loads with cooperative A+B LDS + staging. Quantify A/B reuse per macro-tile and use vectorized coalesced + global loads with a bank-safe LDS layout. + "3": >- + Pipeline round: compare single and double buffering across K tiles. + Retain double buffering only when ISA/PMC evidence shows reduced VMEM + stalls without harmful LDS or occupancy growth. + "4": >- + Tile-shape round: tune M-tile versus N-tile aspect ratio for this exact + M/N/K, balancing B reuse, A reuse and enough blocks to occupy every CU. + Do not inherit decode launch geometry. + "5": >- + Packing round: test one weight packing/swizzle that makes each DUMMA B + tile vector-loadable and LDS-bank-safe. Include packing outside timing + and validate the graph-stable packed layout. + "6": >- + Epilogue round: fuse per-row and per-column scales, bf16 conversion and + the final coalesced store into the compute kernel; remove any + unnecessary workspace/combine pass. + "7": >- + Compute-pipeline round: tune DUMMA issue grouping, prefetch distance and + accumulator independence using ISA stall evidence. Raw inline asm + remains forbidden. + "8": >- + Resource round: tune waves per block, VGPR live ranges and LDS footprint + from measured occupancy. Recheck the best tile family with normal + median/P90 measurements. + "9": >- + Late ISA-diagnosis round. Only if the control-plane plateau gate is open, + use one selected ISA Skill and trusted disassembly to shape compiler + output through HIP/DUMMA/intrinsics. Raw inline asm remains forbidden. + Otherwise continue HIP-only exploration. + "10": >- + Final conditional inline-asm round. Raw asm is allowed only when the + control plane confirms a HIP plateau and the prior ISA-guided round + recorded one concrete compiler limitation plus target instructions. + Otherwise make one HIP-only consolidation change. +notes: >- + These texts are byte-identical to the Python literals they replaced + (orchestrator/round_strategy.py BUILTIN_STRATEGY), so an un-evolved harness + behaves exactly as before. Editing this file now changes the round + instructions the worker receives, and the Evolve loop can version it. + A key omitted here keeps its built-in text; an unknown regime or slot is + reported in RoundStrategy.notes and ignored. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/config.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/config.py index d908ba13..b32662d6 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/config.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/config.py @@ -3,7 +3,7 @@ from __future__ import annotations import os -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Mapping @@ -24,25 +24,57 @@ # dsh routes claude_bin to bridge/dsh/dsh_agent.py (see resolve_claude_bin). CCB_FRAMEWORK = "ccb" DSH_FRAMEWORK = "dsh" -DSH_DEFAULT_MODEL_ID = "deepseek/deepseek-v4-flash-0731" + +#: DSH model label -> platform model id, in form order. The first entry is the +#: framework default. The label is what the New Task form shows (and what the +#: frontend model widget in ``static/dkao-agent-fields.js`` mirrors); the id is +#: what the agent binary receives through ``--model``. +#: +#: ``deepseek-flash-4.1`` is the 4.1 revision of the DeepSeek Flash line, served +#: by the TokenHub gateway as ``deepseek/deepseek-flash`` (the id this host's +#: DSH profile lists as ``dsv4-flash4.1``). ``deepseek-v4-flash`` stays +#: selectable for runs that must reproduce the older pinned build. +DSH_MODEL_IDS: Dict[str, str] = { + "deepseek-flash-4.1": "deepseek/deepseek-flash", + "deepseek-v4-flash": "deepseek/deepseek-v4-flash-0731", +} +DSH_DEFAULT_MODEL_LABEL = "deepseek-flash-4.1" +DSH_DEFAULT_MODEL_ID = DSH_MODEL_IDS[DSH_DEFAULT_MODEL_LABEL] + + +def _dsh_model_override() -> str: + """The DSH_AGENT_MODEL pin, stripped, or "" when unset.""" + return (os.environ.get("DSH_AGENT_MODEL") or "").strip() def dsh_model_id() -> str: - """Model id used for the dsh framework (env DSH_AGENT_MODEL overrides).""" - override = os.environ.get("DSH_AGENT_MODEL") - return (override.strip() if override and override.strip() - else DSH_DEFAULT_MODEL_ID) + """Default model id for the dsh framework, honoring the DSH_AGENT_MODEL pin. + + ``DSH_AGENT_MODEL`` is the host-level override the launcher exports: when it + is set, every dsh model label resolves to it. That keeps a host pinned to one + verified build whatever the form submits, which is what a mixed-version + deployment wants — the catalog still tells the operator which labels exist, + and an unset variable leaves each label on its own id. + """ + return _dsh_model_override() or DSH_DEFAULT_MODEL_ID def agent_framework_models(framework: str) -> Dict[str, str]: """Label -> model id for one agent framework.""" if framework == DSH_FRAMEWORK: - return {"deepseek-v4-flash": dsh_model_id()} + if _dsh_model_override(): + # A host pin routes every dsh label to that one build; the catalog + # keeps its labels so an old requirements file still resolves. + return {label: _dsh_model_override() + for label in DSH_MODEL_IDS} + return dict(DSH_MODEL_IDS) return dict(CLAUDE_MODELS) def agent_framework_default_model(framework: str) -> str: - return "deepseek-v4-flash" if framework == DSH_FRAMEWORK else "Opus" + if framework == DSH_FRAMEWORK: + return DSH_DEFAULT_MODEL_LABEL + return "Opus" def resolve_agent_framework(req: Mapping[str, Any]) -> str: @@ -180,6 +212,18 @@ class OptimizerConfig: # Agent framework (ccb | dsh) that produced claude_model; the pipeline # uses it to pick the agent binary via resolve_claude_bin. agent_framework: str = CCB_FRAMEWORK + #: 占据模式 (occupy, default): workers keep the devices the assignment mode + #: picked (manual pins worker_N -> GPU N). 调控模式 (scheduled): devices are + #: leased from the shared GPU broker so several tasks can share the machine + #: without measuring on a contended card. + gpu_mode: str = "occupy" + gpu_leases: Dict[str, Any] = field(default_factory=dict) + #: "default" | "production". Production means: an operator-launched real + #: task that outranks the evolving harness in the device queue + #: (``PRIORITY_PRODUCTION``) and keeps waiting for a device that passes the + #: admission gate. It never skips the gate itself — a measurement taken on a + #: shared card is stale data regardless of who asked for it. + execution_profile: str = "default" def _answers(req: Mapping[str, Any]) -> Mapping[str, Any]: @@ -397,6 +441,37 @@ def load_config(req: Mapping[str, Any]) -> OptimizerConfig: if threshold < 0: raise ValueError("minimum_improvement_percent must be non-negative") + #: 生产模式 (``execution_profile: production``): an operator-launched real + #: task. It is *not* a way around the admission gate — a measurement taken on + #: a shared card is stale data whatever the task is — it is a **priority** + #: statement: production work goes through the broker at + #: ``PRIORITY_PRODUCTION`` so the evolving harness (priority 50) yields to + #: it, and it keeps waiting for a device that passes the gate instead of + #: giving up. + profile = str(answers.get("execution_profile") or "").strip().lower() + if profile in {"", "default", "harness", "evolve", "ahe"}: + profile = "default" + elif profile in {"production", "prod"}: + profile = "production" + else: + raise ValueError( + "execution_profile must be 'default' or 'production'") + production = profile == "production" + + gpu_mode = str(answers.get("gpu_mode") or "occupy").strip().lower() + if production: + # production always goes through the broker (that is what carries the + # priority); pinning devices would step around the queue entirely. + gpu_mode = "scheduled" + if gpu_mode not in {"occupy", "scheduled"}: + raise ValueError("gpu_mode must be occupy or scheduled") + lease_info: Dict[str, Any] = {} + if gpu_mode == "scheduled": + # 调控模式 / 生产模式: devices come from the shared broker, one per worker + assignments, lease_info = lease_gpus_for_task( + assignments, answers, production=production) + assignment_mode = "scheduled" + return OptimizerConfig( operator=str(answers.get("operator") or "Custom operator"), dtype=str(answers.get("dtype") or "Other"), @@ -412,9 +487,95 @@ def load_config(req: Mapping[str, Any]) -> OptimizerConfig: mock_iterations=iterations, minimum_improvement_percent=threshold, agent_framework=agent_framework, + execution_profile=profile, + gpu_mode=gpu_mode, + gpu_leases=lease_info, ) +def rebalance_for_gpus(assignments: List[WorkerAssignment], + gpus: List[int]) -> List[WorkerAssignment]: + """One worker per leased device, shapes spread round-robin. + + With fewer leases than planned workers the shapes are merged into fewer + lanes instead of putting two workers on one device — the whole point of + scheduled mode is that a device is never shared while measuring. + """ + ordered = sorted({int(g) for g in gpus}) + if not ordered: + return [] + shape_ids = [sid for assignment in assignments for sid in assignment.shape_ids] + buckets: Dict[int, List[str]] = {g: [] for g in ordered} + for index, shape_id in enumerate(shape_ids): + buckets[ordered[index % len(ordered)]].append(shape_id) + rebuilt = [ + WorkerAssignment(worker_id=f"worker_{gpu}", gpu=gpu, shape_ids=sids) + for gpu, sids in sorted(buckets.items()) if sids + ] + return rebuilt or [WorkerAssignment(worker_id=f"worker_{ordered[0]}", + gpu=ordered[0], shape_ids=shape_ids)] + + +def lease_gpus_for_task(assignments: List[WorkerAssignment], + answers: Mapping[str, Any], + *, production: bool = False, + ) -> tuple[List[WorkerAssignment], Dict[str, Any]]: + """调控模式 / 生产模式: take GPU leases and lay the workers out on them. + + Both ask the broker at ``PRIORITY_PRODUCTION`` so a real task outranks the + evolving harness; ``production`` additionally keeps waiting for a device that + passes the admission gate instead of failing after a fixed number of checks, + and records that intent in the lease metadata. + """ + from metainfer.orchestrator.gpu_broker import ( + PRIORITY_PRODUCTION, GpuBroker, + ) + + broker = GpuBroker() + task_id = str(answers.get("task_id") or "dkao-task") + holder = f"dkao:{task_id}" + want = max(1, len(assignments)) + idle_wait = float(answers.get("gpu_lease_wait_seconds") or 1800) + max_waits = int(answers.get("gpu_lease_max_waits") or 48) + if production: + # waiting longer is the point: a production task yields the queue to + # nobody, but it still refuses to measure on a device that fails the gate + max_waits = max(max_waits, int(answers.get("production_max_waits") or 48)) + leases: List[Dict[str, Any]] = [] + for round_no in range(1, max_waits + 1): + leases = broker.acquire( + holder, want=want, priority=PRIORITY_PRODUCTION, + timeout_s=(0.0 if round_no == 1 else idle_wait), poll_s=30, + pid=os.getpid(), + meta={"mode": "scheduled", "task": task_id, "round": round_no, + "profile": "production" if production else "default"}, + ) + if leases: + break + if not leases: + raise RuntimeError( + ("production: no device passed the admission gate (VRAM<=90% and " + f"HCU==0) after {max_waits} checks; the gate is not skipped for " + "production tasks because a shared card yields stale numbers") + if production else + ("gpu_mode=scheduled: no GPU lease became available after " + f"{max_waits} checks; switch to gpu_mode=occupy to pin devices") + ) + gpus = [int(lease["gpu"]) for lease in leases] + rebuilt = rebalance_for_gpus(assignments, gpus) + info = { + "mode": "scheduled", + "profile": "production" if production else "default", + "priority": PRIORITY_PRODUCTION, + "holder": holder, + "gpus": gpus, + "lease_seconds": int(answers.get("gpu_lease_ttl_s") or 0), + "mapping": {a.worker_id: a.gpu for a in rebuilt}, + "requested_workers": len(assignments), + } + return rebuilt, info + + def replace_assignments( config: OptimizerConfig, assignments: List[WorkerAssignment], @@ -435,6 +596,9 @@ def replace_assignments( mock_iterations=config.mock_iterations, minimum_improvement_percent=config.minimum_improvement_percent, agent_framework=config.agent_framework, + execution_profile=config.execution_profile, + gpu_mode=config.gpu_mode, + gpu_leases=dict(getattr(config, "gpu_leases", {}) or {}), ) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/contention.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/contention.py new file mode 100644 index 00000000..023d2414 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/contention.py @@ -0,0 +1,187 @@ +"""Discard a worker round that was measured while its device was shared. + +The parent harness freezes a question's child when the device it holds is taken +over, and records that window in ``/contention.json``. The +child, on the way back up, must then do two things for that round: + +1. **not keep the number** — a measurement taken while another workload shared + the card is noise (the same unchanged shape has read 884us and 3047us), and + if it is kept it also pollutes the iteration history the planner reasons + about; +2. **go back to the previous good state** — the accepted kernel commit, the + official best metrics and any shadow candidate from that round are undone, so + the round can simply be re-measured instead of starting the whole question + over. + +The record is *quarantined*, never deleted: the failed number stays available as +evidence of the contamination. + +Re-planning comes for free: ``plan.json`` is written once from static config, +while the per-round planner inputs (``isa_policy``, ``pmc_profile_decision``, +round mandates) are derived from ``experiments.jsonl`` at the top of every +iteration. Removing the contaminated row therefore means the next round is +planned against clean history. +""" + +from __future__ import annotations + +import json +import os +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +CONTENTION_FILE = "contention.json" +QUARANTINE_DIR = "contaminated" + + +def contention_file(child_state_dir: Optional[Path]) -> Optional[Path]: + """Locate the parent's contention record for this child.""" + candidates: List[Path] = [] + if child_state_dir: + candidates.append(Path(child_state_dir) / CONTENTION_FILE) + env = str(os.environ.get("METAINFER_CHILD_STATE_DIR") or "").strip() + if env: + candidates.append(Path(env) / CONTENTION_FILE) + for path in candidates: + try: + if path.is_file(): + return path + except OSError: + continue + return None + + +def read_contention(child_state_dir: Optional[Path] = None) -> Dict[str, Any]: + """The contention window recorded for this child (empty when none).""" + path = contention_file(child_state_dir) + if path is None: + return {} + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return data if isinstance(data, dict) else {} + + +def iteration_is_contaminated(iteration: int, window: Dict[str, Any], *, + iteration_dir: Optional[Path] = None, + now: Optional[float] = None) -> tuple[bool, str]: + """Whether ``iteration`` was measured inside the contention window. + + The window is ``[frozen_from, recovered_at]``: the child was frozen the + moment contention was seen, so any round whose work lands in that interval + (or which was still unfinished when the window closed) is suspect. A window + that is still open contaminates the round that was running when it opened. + """ + if not window: + return False, "no contention recorded" + frozen_from = window.get("frozen_from") + if frozen_from is None: + return False, "window has no start" + recovered_at = window.get("recovered_at") + stamp = float(now if now is not None else time.time()) + # where the round's work sits on disk (mtime is what we can trust: the + # records themselves carry no timestamp) + mtime = None + if iteration_dir is not None: + try: + mtime = Path(iteration_dir).stat().st_mtime + except OSError: + mtime = None + reference = mtime if mtime is not None else stamp + if reference < float(frozen_from): + return False, "round finished before the device was taken over" + if recovered_at is None: + # the window is still open: this round ran while the device was shared + return True, "device was taken over while this round was running" + if reference <= float(recovered_at): + return True, "round lands inside the window while the device was shared" + return False, "round was measured after the device recovered" + + +def _quarantine(experiments_path: Path, iteration: int, record: Dict[str, Any], + meta: Dict[str, Any], *, state_dir: Optional[Path] = None, + ) -> Optional[Path]: + """Move one round's record out of the live history into quarantine.""" + target_dir = Path(state_dir) if state_dir else experiments_path.parent + target_dir = target_dir / QUARANTINE_DIR + target_dir.mkdir(parents=True, exist_ok=True) + target = target_dir / f"iteration-{int(iteration):03d}.json" + payload = { + "reason": "device_shared_during_measurement", + "discarded_at": time.time(), + "iteration": int(iteration), + "metrics": { + k: (record.get("metrics") or {}).get(k) + for k in ("median_us", "p90_us", "passed", "mismatch_count") + }, + **meta, + "record": record, + } + try: + target.write_text(json.dumps(payload, indent=2, ensure_ascii=False), + encoding="utf-8") + except OSError: + return None + return target + + +def discard_contaminated_round(*, experiments_path: Path, iteration: int, + state_dir: Optional[Path] = None, + iteration_dir: Optional[Path] = None, + window: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Remove a contaminated round from the live history (quarantined). + + Returns ``{"discarded": bool, "reason": str, ...}``. The caller is expected + to restore the kernel source and metrics to the previous good state (the + discard itself only touches the history and the quarantined copy). + """ + experiments_path = Path(experiments_path) + window = window if window is not None else read_contention(state_dir) + contaminated, why = iteration_is_contaminated(iteration, window, + iteration_dir=iteration_dir) + if not contaminated: + return {"discarded": False, "reason": why} + if not experiments_path.is_file(): + return {"discarded": False, "reason": "no experiment history to prune"} + + rows: List[Dict[str, Any]] = [] + dropped: List[Dict[str, Any]] = [] + for line in experiments_path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + try: + row = json.loads(line) + except ValueError: + rows.append({"_raw": line}) + continue + if int(row.get("iteration") or 0) == int(iteration): + dropped.append(row) + else: + rows.append(row) + + quarantined = [] + for record in dropped: + path = _quarantine(experiments_path, iteration, record, { + "window": window, "reason_detail": why, + }, state_dir=state_dir) + if path is not None: + quarantined.append(str(path)) + + try: + experiments_path.write_text( + "".join(json.dumps(r) + "\n" for r in rows if "_raw" not in r + ) + "".join(r["_raw"] + "\n" for r in rows if "_raw" in r), + encoding="utf-8") + except OSError as exc: # noqa: BLE001 + return {"discarded": False, "reason": f"cannot rewrite history: {exc}"} + return { + "discarded": True, + "reason": why, + "rounds_removed": len(dropped), + "quarantined": quarantined, + "history_remaining": len([r for r in rows if "_raw" not in r]), + "window": window, + } 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..3e0bc0cf 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 @@ -14,7 +14,7 @@ import time from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -from typing import Any, Dict, List +from typing import Any, Callable, Dict, List, Optional, Sequence from metainfer.orchestrator.state import StateStore from metainfer.orchestrator.subagent_manager import AgentSpec, SubAgentManager @@ -47,6 +47,8 @@ generate_kernel_prompt, shape_balanced_assignment, ) +from .gpu_binding import bind_worker_gpu +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 @@ -54,10 +56,16 @@ RealW8A8OptimizationPipeline, W8A8Runner, _SOURCE_ONLY_AGENT_ARGS, + _gate_preflight_enabled, _sha256_file, + ensure_any_measurement_gate, 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 @@ -204,11 +212,48 @@ def _last_json_object(text: str) -> Dict[str, Any]: raise ValueError(f"no JSON object in output: {text[-1000:]}") +def _gate_state_dir_for(pipeline: Any) -> "Path | None": + """Where the Generate probe's gate rows go: the task's own state dir. + + Without it a 30-minute wait on a busy device leaves no trace in the task's + audit trail, which is exactly the "why did nothing happen" question this + gate has to answer. + """ + raw = getattr(getattr(pipeline, "store", None), "task_dir", None) + return Path(raw) if raw else None + + +def _probe_gate_devices(config: Any) -> List[int]: + """Devices the Generate probe may use: this task's own GPUs first. + + The probe only needs *a* device that passes the gate (a device is visible and + HIP Graph is available). Its own assignment is preferred so a task never + competes for a card it does not own, but an assigned card that is currently + shared must not stall the whole task while a sibling device sits idle. + """ + assigned: List[int] = [] + for item in list(getattr(config, "assignments", None) or []): + try: + gpu = int(getattr(item, "gpu", -1)) + except (TypeError, ValueError): + continue + if gpu >= 0 and gpu not in assigned: + assigned.append(gpu) + try: + from metainfer.orchestrator.gpu_broker import GpuBroker + known = [int(d) for d in GpuBroker().devices] + except Exception: # noqa: BLE001 - fall back to the task's own devices + known = list(assigned) or [0] + return assigned + [d for d in known if d not in assigned] + + def _validate_generate_scaffold( source: Path, *, contract_sha256: str, shapes: Dict[str, Dict[str, Any]], + probe_devices: Optional[Sequence[int]] = None, + state_dir: "Path | None" = None, ) -> Dict[str, Any]: """Run trusted, implementation-free Generate preflight checks.""" kernel_path = source / _GENERATED_KERNEL_FILE @@ -248,11 +293,22 @@ def _validate_generate_scaffold( ) probe_env = dict(os.environ) - probe_env.update({ - "HIP_VISIBLE_DEVICES": "0", - "ROCR_VISIBLE_DEVICES": "0", - "PYTHONDONTWRITEBYTECODE": "1", - }) + probe_gpu = 0 + # The probe initialises a device, so it follows the same admission rule as a + # timed measurement (VRAM <= 90% and HCU == 0). It is not pinned to GPU 0: + # any device that passes the gate can answer "is a device visible and is HIP + # Graph available", and sleeping 24 h on a busy GPU 0 while GPU 1 is idle is + # how a task stalls without doing anything wrong. + if _gate_preflight_enabled(): + probe_gpu = ensure_any_measurement_gate( + list(probe_devices or [0]), state_dir=state_dir, env=probe_env, + site="generate_probe") + # Binding goes through the task's one policy: HIP_VISIBLE_DEVICES only. + # Setting ROCR_VISIBLE_DEVICES to the same *non-zero* index filters twice + # and hides the device ("No HIP GPUs are available" at device 1, while + # device 0 happened to survive). + bind_worker_gpu(probe_env, probe_gpu) + probe_env["PYTHONDONTWRITEBYTECODE"] = "1" probe_result = _run( [ "python3", str(harness), @@ -554,6 +610,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 +886,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(): @@ -882,7 +978,8 @@ def _generate_kernel_repo( sid: shape.params for sid, shape in config.shapes.items() } fixed_assignment: Dict[str, Dict[str, Any]] - if config.assignment_mode == "manual": + if (config.assignment_mode == "manual" + and getattr(config, "gpu_mode", "occupy") == "occupy"): actual = { item.worker_id: item.gpu for item in config.assignments } @@ -936,6 +1033,8 @@ def _generate_kernel_repo( agent_source_dir, contract_sha256=trusted_contract_digest, shapes=shapes_for_prompt, + probe_devices=_probe_gate_devices(config), + state_dir=_gate_state_dir_for(self), ) preflight_path = seed / "generation_preflight.json" write_json(preflight_path, preflight) @@ -2114,11 +2213,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 +2242,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 +2268,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 +2324,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, } @@ -2217,7 +2341,54 @@ def _synthesize_final_candidate( # Main pipeline # ------------------------------------------------------------------ # + def _register_gpu_leases(self, config: Any) -> Optional[Callable[[], Any]]: + """调控模式: announce leased devices and return them when we exit. + + The lease also has a TTL, so a hard crash still frees the devices; this + hook keeps the common paths prompt (a device should not sit idle for two + hours after a task finishes). + """ + leases = dict(getattr(config, "gpu_leases", {}) or {}) + if not leases: + return None + holder = str(leases.get("holder") or "") + self.store.append_timeline("gpu_leases_acquired", { + "holder": holder, + "gpus": leases.get("gpus"), + "mapping": leases.get("mapping"), + "requested_workers": leases.get("requested_workers"), + "mode": leases.get("mode"), + # production leases are visible as such, so the GPU view can tell a + # real operator task from an evolving-harness round + "profile": leases.get("profile"), + "priority": leases.get("priority"), + }) + + def _release() -> list: + try: + from metainfer.orchestrator.gpu_broker import GpuBroker + freed = GpuBroker().release_prefix(holder) + if freed: + self.store.append_timeline("gpu_leases_released", { + "holder": holder, "gpus": freed, + }) + return freed + except Exception: # noqa: BLE001 - release must never raise + return [] + + import atexit + atexit.register(_release) + self.release_gpu_leases = _release + return _release + 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. @@ -2240,6 +2411,7 @@ def run(self, *, dry_run: bool = False) -> Dict[str, Any]: # ---- PREPARE -------------------------------------------------- # self._phase(phases.PREPARE) config = load_config(self.req) + self._register_gpu_leases(config) self._validate_contract(config) self._prepare_worktrees(config, task_id) plan = self._plan(config) @@ -2386,7 +2558,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..d5a92e4b --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_preflight.py @@ -0,0 +1,449 @@ +"""GPU admission pre-check shared by production DKAO tasks and AHE children. + +The operating rule is one gate: **a device may only be used for a measurement +when VRAM <= 90% and HCU == 0** (``usable``). Anything else — a full card, a +card shared with another workload, a card whose state we cannot read — is not +usable, and the caller waits and re-checks instead of measuring on it. + +A second, weaker signal is reporting only: board power and foreign KFD pids are +recorded for diagnosis, and a device that passes the gate but still looks +suspicious is marked ``measurement_suspect`` so a polluted reading cannot +silently decide an A/B comparison. + +Telemetry has two sources, because ``hy-smi`` alone is not enough: some driver +configurations print ``N/A`` for HCU% (the container cannot open the management +interface), and a parser that silently yields nothing would make *every* card +look idle. ``sample_gpu_state`` therefore falls back to the driver sysfs +(``gpu_busy_percent`` / ``mem_info_vram_used``) and reports ``unavailable`` when +neither source answers. + +The module is deliberately dependency-free: it parses ``hy-smi`` text and the +KFD/sysfs listings, 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 + +#: Operating rule requested by the operator: a device may only be used when +#: VRAM <= ``DEFAULT_VRAM_LIMIT_PERCENT`` **and** the device is truly idle +#: (HCU == 0). The percentage form is kept for reference. +DEFAULT_VRAM_LIMIT_PERCENT = 90.0 +#: HCU tolerance: 0.0 means "strictly idle". +DEFAULT_UTIL_TOLERANCE = 0.0 +#: Kept for reference/reporting (VRAM <= 90% implies >= ~6.4GB free on 64GB). +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") +#: Driver sysfs used when ``hy-smi`` cannot report HCU%/VRAM. +DRM_ROOT = Path("/sys/class/drm") +#: DRM drivers that mean "this card is a Hygon HCU". +_HCU_DRM_DRIVERS = ("hycu", "amdgpu", "hydcu") +#: ``N/A`` (or ``-``) where a number is expected: the device is present but the +#: management interface cannot be read. Such a field is recorded as unknown +#: instead of dropping the whole row, which used to make every card look idle. +_UNKNOWN = r"(?:N/A|n/a|-{1,2})" +_SMI_ROW = re.compile( + r"^(?P\d+)\s+" # device index + r"(?P[\d.]+)C\s+" + r"(?:{u}|(?P[\d.]+)W)\s+" # average power + r"\S+\s+" # perf mode + r"(?:{u}|[\d.]+W)\s+" # power cap + r"(?:(?P[\d.]+)%|{u})\s+" # VRAM% + r"(?:(?P[\d.]+)%|{u})" # HCU% + .format(u=_UNKNOWN) +) + + +def parse_hy_smi( + text: str, *, device_ids: Optional[List[int]] = None, +) -> Dict[int, Dict[str, float]]: + """Parse ``hy-smi`` into ``{index: {vram_percent, util_percent, power_w, temp_c}}``. + + A field printed as ``N/A`` is *omitted* from that device's row: the caller + then sees "unknown" rather than a fabricated ``0.0``, which is what the + gate needs to refuse the device instead of silently admitting it. + """ + 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")) + if device_ids is not None and index not in device_ids: + continue + row: Dict[str, float] = {} + for key in ("vram", "util", "power", "temp"): + raw = match.group(key) + if raw is None: + continue + row[{"vram": "vram_percent", "util": "util_percent", + "power": "power_w", "temp": "temp_c"}[key]] = float(raw) + out[index] = row + return out + + +def smi_is_incomplete( + text: str, *, device_ids: Optional[List[int]] = None, +) -> bool: + """True when ``hy-smi`` names a device but prints ``N/A`` for HCU%. + + That is the signature of "the management interface cannot be read from + here" (``Open mkfd failed`` inside a container without the HCU device). On + those hosts the *whole* utilisation block is a placeholder — including a + ``0%`` VRAM column — so the driver sysfs is the more trustworthy source + and takes precedence over the rest of the same row. + """ + rows = 0 + for line in (text or "").splitlines(): + match = _SMI_ROW.match(line.strip()) + if not match: + continue + index = int(match.group("index")) + if device_ids is not None and index not in device_ids: + continue + rows += 1 + if match.group("util") is None: + return True + return rows == 0 + + +def read_sysfs_gpu_state( + drm_root: Path = DRM_ROOT, +) -> Dict[int, Dict[str, float]]: + """Driver sysfs fallback: ``{index: {vram_percent, util_percent}}``. + + ``/sys/class/drm`` is host-wide and does not depend on the management + interface, so it still answers when ``hy-smi`` prints ``N/A``. Cards are + filtered to the HCU driver and ordered by card index (``card1`` is the + first HCU on the K500SM_AI machines, since ``card0`` is the BMC VGA), which + matches the ``hy-smi``/KFD device numbering the tasks use. + """ + cards: List[Path] = [] + try: + entries = list(drm_root.iterdir()) + except OSError: + return {} + for entry in sorted(entries, key=lambda p: p.name): + match = re.fullmatch(r"card(\d+)", entry.name) + if not match: + continue + device = entry / "device" + try: + uevent = (device / "uevent").read_text(encoding="utf-8") + except OSError: + continue + if not any(driver in uevent for driver in _HCU_DRM_DRIVERS): + continue + if not (device / "mem_info_vram_total").is_file(): + continue + cards.append(device) + out: Dict[int, Dict[str, float]] = {} + for index, device in enumerate(cards): + row: Dict[str, float] = {} + busy = _read_int(device / "gpu_busy_percent") + if busy is not None: + row["util_percent"] = float(busy) + used = _read_int(device / "mem_info_vram_used") + total = _read_int(device / "mem_info_vram_total") + if used is not None and total: + row["vram_percent"] = round(100.0 * float(used) / float(total), 2) + if row: + out[index] = row + return out + + +def _read_int(path: Path) -> Optional[int]: + try: + return int((path.read_text(encoding="utf-8") or "").strip()) + except (OSError, ValueError): + return None + + +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, + sysfs_reader=read_sysfs_gpu_state, + ) -> Dict[int, Dict[str, float]]: + """Average a few hy-smi samples per device (drops idle flicker). + + ``hy-smi`` is the primary source. When it cannot report HCU% (it prints + ``N/A`` on driver/container combinations where the management interface is + unreadable), the missing fields are taken from the driver sysfs instead of + being treated as zero: a fabricated ``util=0`` would admit a card that is + actively busy, which is exactly the failure this gate exists to prevent. + A device neither source can describe is reported with ``unavailable=True`` + so the gate refuses it. + """ + wanted = None if gpu_ids is None else {int(g) for g in gpu_ids} + acc: Dict[int, Dict[str, List[float]]] = {} + smi_texts: List[str] = [] + for attempt in range(max(1, samples)): + text = smi_reader() + smi_texts.append(text or "") + for index, row in parse_hy_smi(text).items(): + if wanted is not None and index not in wanted: + 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)) + # ``hy-smi`` printouts that say N/A for HCU% come from a host where the + # management interface is unreadable: the rest of that row (including a + # "0%" VRAM column) is a placeholder, so sysfs wins. + smi_unreliable = any( + smi_is_incomplete(text, device_ids=gpu_ids) for text in smi_texts) + out: Dict[int, Dict[str, float]] = {} + for index, bucket in acc.items(): + row: Dict[str, float] = {} + for key, values in bucket.items(): + if values and not (smi_unreliable and key == "vram_percent"): + row[key] = sum(values) / len(values) + if bucket.get("vram_percent") and not smi_unreliable: + row["vram_max_percent"] = max(bucket["vram_percent"]) + if bucket.get("util_percent"): + row["util_max_percent"] = max(bucket["util_percent"]) + if row: + row["state_source"] = "hy-smi" + out[index] = row + need_fallback = ( + not out + or smi_unreliable + or any( + row.get("util_percent") is None or row.get("vram_percent") is None + for row in out.values() + ) + or any( + index not in out + or out[index].get("util_percent") is None + or out[index].get("vram_percent") is None + for index in (wanted or ()) + ) + ) + fallback: Dict[int, Dict[str, float]] = {} + if sysfs_reader is not None and need_fallback: + try: + fallback = dict(sysfs_reader() or {}) + except Exception: # noqa: BLE001 - the fallback is best effort + fallback = {} + for index, source in fallback.items(): + if wanted is not None and index not in wanted: + continue + row = out.setdefault(index, {}) + filled: List[str] = [] + for key in ("util_percent", "vram_percent"): + if source.get(key) is None: + continue + if row.get(key) is not None and not (smi_unreliable + and key == "vram_percent"): + continue + row[key] = float(source[key]) + peak = ("util_max_percent" if key == "util_percent" + else "vram_max_percent") + row.setdefault(peak, float(source[key])) + filled.append(key) + if filled: + had_smi = row.get("power_w") is not None or row.get("temp_c") is not None + row["state_source"] = "hy-smi+sysfs" if had_smi else "sysfs" + if wanted is not None: + for index in sorted(wanted): + out.setdefault(index, {}) + for index, row in out.items(): + missing = [k for k in ("vram_percent", "util_percent") if k not in row] + if missing: + row["unavailable"] = 1.0 + row["missing_fields"] = missing + 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, + vram_limit_percent: float = DEFAULT_VRAM_LIMIT_PERCENT, + util_tolerance: float = DEFAULT_UTIL_TOLERANCE, + 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]: + """Usable only when VRAM <= limit *and* the device is idle (HCU == 0). + + A device whose state could not be read is **not** usable: "we cannot tell" + must never be treated as "idle", otherwise a busy card silently admits + measurements. Callers then wait and re-check (30 min, up to 48 times). + """ + unavailable = bool(state.get("unavailable")) + vram = float(state.get("vram_max_percent") + if state.get("vram_max_percent") is not None + else state.get("vram_percent") or 0.0) + util_peak = float(state.get("util_max_percent") + if state.get("util_max_percent") is not None + else state.get("util_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 + vram_ok = vram <= float(vram_limit_percent) + idle_ok = util_peak <= float(util_tolerance) + usable = bool(vram_ok and idle_ok and not unavailable) + reasons: List[str] = [] + if unavailable: + missing = ",".join(state.get("missing_fields") or []) or "all" + reasons.append(f"state unavailable: no HCU/VRAM reading ({missing})") + # The gate is exactly "VRAM <= 90% and HCU == 0" (operator rule). Power draw + # is recorded for diagnosis only: a card idling at HCU 0 while another + # process holds a handle is still safe to measure on, and treating power as + # a blocker would refuse cards the rule admits. + notes: List[str] = [] + if foreign: + notes.append(f"{len(foreign)} foreign KFD pid(s) hold this GPU") + if not vram_ok: + reasons.append( + f"VRAM {vram:.0f}% > {float(vram_limit_percent):.0f}% limit") + if not idle_ok: + reasons.append( + f"device busy: HCU {util_peak:.1f}% > {float(util_tolerance):.1f}%") + return { + "gpu": gpu_id, + "usable": usable, + "free_gb": round(free_gb, 2), + "total_gb": round(total, 1), + "vram_limit_percent": float(vram_limit_percent), + "util_tolerance": float(util_tolerance), + "vram_peak_percent": round(vram, 1), + "util_peak_percent": round(util_peak, 1), + "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], + "unavailable": unavailable, + "state_source": str(state.get("state_source") or "unknown"), + "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, + vram_limit_percent: float = DEFAULT_VRAM_LIMIT_PERCENT, + util_tolerance: float = DEFAULT_UTIL_TOLERANCE, + **kwargs: Any) -> Dict[str, Any]: + """Return ``{gpu_id: check}`` plus a preferred order (clean devices first). + + ``clean_ids`` is the gate result: the only devices a measurement may run + on. ``blocked_ids`` are the devices that failed it (busy, full, or + unreadable) — callers wait and re-check rather than falling back to them. + """ + if not enabled: + return {"enabled": False, "gpus": {}, "preferred": list(gpu_ids), + "suspect_ids": [], "clean_ids": [], "blocked_ids": [], + "over_limit_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 {}, + vram_limit_percent=vram_limit_percent, + util_tolerance=util_tolerance, + 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"]] + blocked = [g for g in gpu_ids if not checks[g]["usable"]] + # Only gate-passing devices are ever offered for a measurement: a busy or + # unreadable card is left out so the caller waits (30 min, x48) instead of + # producing a measurement that another workload has polluted. + preferred = clean + suspect + return { + "enabled": True, + "gpus": checks, + "preferred": preferred, + "clean_ids": clean, + "suspect_ids": suspect, + "blocked_ids": blocked, + "over_limit_ids": [g for g in blocked + if checks[g]["vram_peak_percent"] + > float(vram_limit_percent)], + "gate_reasons": {str(g): checks[g]["reasons"] for g in blocked}, + "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..04833baa --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/harness_io.py @@ -0,0 +1,137 @@ +"""Harness workspace IO: locate the evolvable-harness seed, read components, +and copy a seed workspace to a destination. + +This is the AHE "component observability" layer for dcu_kernel_auto_opt: +harness components (gates, planner policy/catalog, systemprompt, and later +skills/tools/middleware/memory) are seeded as files under ``harness_default/`` +so they form a versioned, evolvable workspace driven by the ``harness_evolve`` +outer loop. + +Wiring status: ``gates.yaml`` is read through :mod:`gate_policy`, +``planner_policy.yaml``/``planner_catalog.yaml`` through :mod:`planner`, and +``systemprompt/round_strategy.yaml`` through :mod:`round_strategy`. Each of +those falls back to built-in defaults when its file is absent, so an untouched +harness behaves exactly as it did before the component was externalised. +Consistency tests guard every YAML seed against its Python defaults so the two +cannot drift silently. +""" + +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 env_root() -> Path: + """The harness root as the *environment* defines it (no explicit argument). + + Same resolution as :func:`harness_root`, named separately so components + that already accept an explicit ``root`` argument (``gate_policy``, + ``planner``, ``round_strategy``) can share one implementation and one + cache key. + """ + return harness_root() + + +def load_component_file(relative: str, root: Path | None = None) -> Path: + """Path of one component file inside a harness workspace.""" + return (root or harness_root()) / str(relative) + + +def load_component_yaml(relative: str, root: Path | None = None, + ) -> tuple[Dict[str, Any], str]: + """Read one YAML component. Returns ``(data, error)``; never raises. + + A component file that is absent yields ``({}, "")`` — "not present" is a + normal state that makes the caller fall back to its built-in defaults. A + file that exists but cannot be parsed yields ``({}, "")`` so the + caller can report it instead of silently ignoring a hand-edited mistake. + """ + path = load_component_file(relative, root) + if not path.is_file(): + return {}, "" + try: + with path.open("r", encoding="utf-8") as handle: + data = yaml.safe_load(handle) + except (OSError, ValueError, yaml.YAMLError) as exc: + return {}, f"{relative}: {exc}" + if data is None: + return {}, "" + if not isinstance(data, dict): + return {}, f"{relative}: expected a mapping at the top level" + return data, "" + + +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/orchestrator.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/orchestrator.py index 8d74875c..ccf9056d 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/orchestrator.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/orchestrator.py @@ -77,6 +77,7 @@ def run_with_requirements( extra_add_dirs=[workspace_dir], snapshot_file=state_dir / "agents.json", max_concurrent=4, + state_dir=state_dir, ) restore = install_subagent_shutdown_handlers(manager, pid_file=pid_file) try: 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..372853d8 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/planner.py @@ -0,0 +1,465 @@ +"""State-conditioned optimization plan selector. + +DKAO chooses the next optimization direction in one of two ways: + +* ``METAINFER_PLANNER`` unset — the round menus in + ``prompts.py::w8a8_round_strategy`` (keyed by iteration number; the wording of + those menus is now data in ``systemprompt/round_strategy.yaml``); +* ``METAINFER_PLANNER=1`` — this module, choosing from the current measured + state (history + PMC + budget), per ``docs/dkao_harness_eval_protocol.md`` / + M1 slice-2. ``w8a8_pipeline`` reaches it through ``choose_plan_from_history`` + and records every decision in the worker's ``planner_plans.jsonl``. + +Semantics: +- **Opt-in**: the pipeline consults this module only when ``METAINFER_PLANNER`` + is set; with it unset, behaviour is byte-identical to the legacy menus. +- The plan catalog lives as data in ``harness_default/planner_catalog.yaml`` + (read through ``harness_io`` with ``METAINFER_HARNESS_ROOT`` override) so it is + 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/prompts.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/prompts.py index fe80a0c7..cc56ee10 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/prompts.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/prompts.py @@ -7,6 +7,12 @@ from pathlib import Path from typing import Any, Dict, Optional +from .round_strategy import ( + APPEND_GRID_WARNING as APPEND_GRID_WARNING_TEXT, + APPEND_SPLIT_CANDIDATES as APPEND_SPLIT_CANDIDATES_TEXT, + LATE_START_FLOOR, + load_round_strategy, +) HARNESS_PATH = ( Path(__file__).resolve().parent.parent / "assets" / "w8a8_bench.py" @@ -156,6 +162,23 @@ def w8a8_strategy_guidance( one launch geometry onto all regimes.""" +def _strategy_append(suffix: str, text: str, **fields: Any) -> str: + """Render one evidence suffix, formatted and separated by a leading space. + + A harness that renames a format field (``{cu_count}`` -> ``{cus}``) must + yield a weaker suffix, never a crash: a malformed harness is a performance + problem, not a reason to abort a round. The separator lives here so the + built-in text and the harness template render byte-identically. + """ + if not suffix: + return "" + try: + rendered = str(suffix).format(**fields) + except (KeyError, IndexError, ValueError): + rendered = text + return " " + rendered if rendered else "" + + def w8a8_round_strategy( shape: Dict[str, Any], iteration: int, @@ -164,11 +187,20 @@ def w8a8_round_strategy( max_iterations: int = 10, isa_policy: Dict[str, Any] | None = None, ) -> str: - """Choose an architecture-first round plan from measured evidence.""" + """Choose an architecture-first round plan from measured evidence. + + The *decision* is code (measured history, phase, PMC evidence and the M + region pick the branch); the *wording* of each branch is data, loaded from + the harness workspace's ``systemprompt/round_strategy.yaml`` and falling + back to the historical built-in texts. See + :mod:`orchestrator.round_strategy`. + """ history = history or [] pmc_evidence = pmc_evidence or {} isa_policy = isa_policy or {} + strategy = load_round_strategy() + faster_wrong = [ record for record in history if record.get("correctness_passed") is False @@ -179,14 +211,11 @@ def w8a8_round_strategy( candidate = max( faster_wrong, key=lambda record: float(record["speedup"]) ) - return ( - "Highest priority: repair the faster but incorrect candidate from " - f"iteration {candidate.get('iteration')} (measured speedup " - f"{candidate.get('speedup')}x). Read its archived source at " - f"`{candidate.get('artifact_dir')}` and preserve the fast mapping. " - "Fix only the smallest correctness defect: signed int8 unpacking, " - "tail bounds, scale indexing, bf16 conversion, or a race. Do not " - "replace it with an unrelated architecture." + return strategy.text( + "repair.faster_wrong", + iteration=candidate.get("iteration"), + speedup=candidate.get("speedup"), + artifact_dir=candidate.get("artifact_dir"), ) if history: @@ -195,198 +224,41 @@ def w8a8_round_strategy( if any(token in failure for token in ( "timeout", "timed out", "killed", "no result", "exit 143", )): - return ( - "The preceding attempt failed in the agent infrastructure. " - "Return to the accepted best source and start a new bounded " - "experiment; do not repair or replay a partially written " - "candidate. This failure does not count as a completed " - "optimization round." - ) + return strategy.text("repair.infrastructure_failure") if latest.get("build_success") is False: - return ( - "Repair the immediately preceding candidate from iteration " - f"{latest.get('iteration')} at `{latest.get('artifact_dir')}`. " - "Keep its strategy and make only the minimum compile/API/syntax " - "correction; do not start another redesign this round." + return strategy.text( + "repair.build_failure", + iteration=latest.get("iteration"), + artifact_dir=latest.get("artifact_dir"), ) phase = str(isa_policy.get("phase") or "hip_only") if phase == "isa_guided_hip": completed = int(isa_policy.get("valid_isa_guided_rounds") or 0) - return ( - "ISA-guided HIP round. This is successful ISA experiment " - f"{completed + 1} of at least 2. Select one measured memory or " - "compute bottleneck, compare the exact primary-kernel ISA with " - "the preceding code object, and make one HIP/DUMMA/intrinsic " - "code-shaping change. Raw inline asm remains forbidden. Record " - "a compiler limitation only when the before/after binary proves " - "it and name the exact target instructions." - ) + return strategy.text("phases.isa_guided_hip", completed=completed + 1) if phase == "conditional_inline_asm": - return ( - "Conditional inline-asm experiment. Target only the compiler " - "limitation and exact instructions verified by the immediately " - "preceding ISA-guided HIP round. Keep the asm block minimal, " - "preserve complete constraints/clobbers, and reject it unless " - "the candidate ISA, exact correctness, median, P90, and resources " - "all validate. Do not write raw global/buffer/flat loads or MMAC." - ) + return strategy.text("phases.conditional_inline_asm") - small_m = { - 1: "Vectorize contiguous K loads with exact signed-int8 semantics.", - 2: ( - "Increase instruction-level parallelism with independent int32 " - "accumulators or adjacent N outputs; avoid per-K-tile LDS barriers." - ), - 3: ( - "Stage all of tiny A once with at most one barrier, or tune unroll " - "one step if whole-A staging is not cheaper." - ), - 4: ( - "Change one launch variable only: waves per block, N columns per " - "wave, or unroll factor." - ), - 5: ( - "HIP-only memory round: change one vector-load width, contiguous " - "N mapping, or whole-A reuse decision. Raw inline asm is forbidden." - ), - 6: ( - "HIP-only pipeline round: reduce one dependency chain or barrier " - "using ordinary HIP/intrinsics. Raw inline asm is forbidden." - ), - 7: ( - "HIP-only resource round: tune one block size, unroll factor, or " - "live range while preserving coalescing. Raw inline asm is forbidden." - ), - 8: ( - "HIP-only consolidation round: revisit the fastest correct archived " - "mapping and make one final architecture/codegen improvement. Raw " - "inline asm is forbidden." - ), - } - m16 = { - 1: ( - "Establish a minimal 16x16x32 DUMMA tile with the exact API below, " - "one wave per output tile and explicit int32 accumulation. If the " - "seed already has a correct DUMMA kernel, preserve it and instead " - "test the smallest one-wave-per-block, one-N-tile geometry with no " - "cross-wave barrier; do not spend the round reimplementing it." - ), - 2: ( - "Architecture round: measure grid parallelism before polishing. " - "Explore one complete launch geometry among 1/2/4 waves per block " - "and 1/2/4 adjacent N tiles. Prefer enough independent blocks to " - "cover at least all device CUs; report grid_blocks, waves_per_block " - "and estimated_active_cus in proposal.json." - ), - 3: ( - "Architecture round: if the unsplit grid has fewer than two " - "blocks per device CU and K >= 1024, implement and measure " - "split-K=2 plus at least one CU-aligned candidate (which may be " - "non-power-of-two), or test a one-wave zero-barrier geometry that " - "reaches the same parallelism. Write int32 partials into the " - "caller workspace and include the combine+scale kernel in the " - "timed Graph." - ), - 4: ( - "Architecture/pipeline round: explore one of multi-N-tile reuse, " - "A-only staging, or bounded register/LDS prefetch. Retain enough " - "blocks to cover all CUs, state which A/B bytes are reused, and " - "measure whether the change improves normal median/P90." - ), - 5: ( - "HIP-only packed-weight/staging round: compare one packed layout, " - "A-only staging, or B-only staging design. Raw inline asm is forbidden." - ), - 6: ( - "Pipeline round: choose exactly one staging family from direct, " - "A-only LDS, B-only LDS, or A+B LDS using L2/VMEM evidence. Use " - "coalesced 8- or 16-byte cooperative loads and report HBM/LDS byte " - "changes; do not claim asynchronous overlap without evidence." - ), - 7: ( - "Pipeline round: compare single buffering with double buffering " - "only when K>=1024, L2 hit rate is below 70%, and the doubled LDS " - "budget stays below 48 KiB. Count barriers per K step." - ), - 8: ( - "HIP-only resource round: tune one occupancy limiter using actual " - "PMC evidence: waves per block, VGPR live range, LDS footprint, " - "or spill removal. Do not trade repeated HBM reads for occupancy." - ), - 9: ( - "Late ISA-diagnosis round. Only if the control-plane plateau gate " - "is open, use one selected ISA Skill and trusted disassembly to " - "shape compiler output through HIP/DUMMA/intrinsics. Raw inline " - "asm remains forbidden. Otherwise continue HIP-only exploration." - ), - 10: ( - "Final conditional inline-asm round. Raw asm is allowed only when " - "the control plane confirms a HIP plateau and the prior ISA-guided " - "round recorded one concrete compiler limitation plus target " - "instructions. Otherwise make one HIP-only consolidation change." - ), - } - large_m = { - 1: ( - "Establish a correct DUMMA throughput baseline using a 2-D " - "macro-tile. Benchmark 64x64, 64x128, and 128x64 block tiles; " - "record waves per block, VGPRs, LDS bytes, occupancy and TOPS." - ), - 2: ( - "Operand-reuse round: compare direct loads with cooperative " - "A+B LDS staging. Quantify A/B reuse per macro-tile and use " - "vectorized coalesced global loads with a bank-safe LDS layout." - ), - 3: ( - "Pipeline round: compare single and double buffering across K " - "tiles. Retain double buffering only when ISA/PMC evidence shows " - "reduced VMEM stalls without harmful LDS or occupancy growth." - ), - 4: ( - "Tile-shape round: tune M-tile versus N-tile aspect ratio for " - "this exact M/N/K, balancing B reuse, A reuse and enough blocks " - "to occupy every CU. Do not inherit decode launch geometry." - ), - 5: ( - "Packing round: test one weight packing/swizzle that makes each " - "DUMMA B tile vector-loadable and LDS-bank-safe. Include packing " - "outside timing and validate the graph-stable packed layout." - ), - 6: ( - "Epilogue round: fuse per-row and per-column scales, bf16 " - "conversion and the final coalesced store into the compute " - "kernel; remove any unnecessary workspace/combine pass." - ), - 7: ( - "Compute-pipeline round: tune DUMMA issue grouping, prefetch " - "distance and accumulator independence using ISA stall evidence. " - "Raw inline asm remains forbidden." - ), - 8: ( - "Resource round: tune waves per block, VGPR live ranges and LDS " - "footprint from measured occupancy. Recheck the best tile family " - "with normal median/P90 measurements." - ), - 9: m16[9], - 10: m16[10], - } m = int(shape.get("M", 0)) if m >= 128: - portfolio = large_m + regime = "large" elif m < 16: - portfolio = small_m + regime = "small" else: - portfolio = m16 - late_start = max(9, max_iterations - 1) + regime = "m16" + late_start = max(LATE_START_FLOOR, max_iterations - 1) if m < 16: if iteration < late_start: - return portfolio.get(iteration, portfolio[8]) - return portfolio[9] if iteration < max_iterations else portfolio[10] + return strategy.mandate(regime, iteration) or strategy.slot(regime, 8) + # The <16 portfolio has no dedicated late slot: it keeps its final + # consolidation round to the end (it used to raise KeyError here for + # every max_iterations in 9..10 with iteration >= late_start). + return strategy.slot(regime, 8) if iteration < late_start: - decision = portfolio.get(iteration, portfolio[8]) + decision = strategy.mandate(regime, iteration) or strategy.slot(regime, 8) else: - decision = portfolio[9] if iteration < max_iterations else portfolio[10] + decision = strategy.slot( + regime, LATE_START_FLOOR if iteration < max_iterations else 10) grid_blocks = pmc_evidence.get("grid_blocks") cu_count = pmc_evidence.get("device_cu_count") split_candidates = split_k_candidate_set( @@ -394,26 +266,20 @@ def w8a8_round_strategy( cu_count=int(cu_count) if isinstance(cu_count, (int, float)) else 0, ) if iteration >= 2 and split_candidates: - decision += ( - " Trusted occupancy-probe split candidates for the measured CU " - f"count are {split_candidates}. They include non-power-of-two " - "values where useful, are not a whitelist, and must fit the " - "workspace and stage-alignment constraints. Explore outside this " - "set when evidence supports it." - ) + decision += _strategy_append( + strategy.append.get("split_candidates"), + APPEND_SPLIT_CANDIDATES_TEXT, + split_candidates=split_candidates) if ( iteration >= 2 and isinstance(grid_blocks, (int, float)) and isinstance(cu_count, (int, float)) and grid_blocks < 2 * cu_count ): - decision += ( - f" Trusted control-plane warning: current grid has {grid_blocks} " - f"blocks for {cu_count} CUs, below the two-blocks-per-CU latency-" - "hiding target. Before micro-optimization, benchmark a finer " - "one-wave zero-barrier grid or multiple legal split-K candidates " - "including combine cost." - ) + decision += _strategy_append( + strategy.append.get("grid_warning"), + APPEND_GRID_WARNING_TEXT, + grid_blocks=grid_blocks, cu_count=cu_count) return decision diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/round_strategy.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/round_strategy.py new file mode 100644 index 00000000..25e9feb2 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/round_strategy.py @@ -0,0 +1,454 @@ +"""Round-strategy texts: the per-round mandate menu, as evolvable data. + +``w8a8_round_strategy`` in :mod:`prompts` used to embed every round mandate as +a Python string literal, so nothing about it could be versioned, reviewed or +evolved. This module is the *data* half of that function: + +* :data:`BUILTIN_STRATEGY` holds the historical texts, byte-identical to the + literals they replace, so an un-evolved harness behaves exactly as before; +* :func:`load_round_strategy` reads ``/systemprompt/ + round_strategy.yaml`` when the ``systemprompt`` component is marked wired in + ``manifest.yaml``, and falls back to the built-in defaults otherwise + (missing file, unwired component, malformed YAML — the same layering + ``gates.yaml`` and ``planner_policy.yaml`` already use). + +The decision *logic* stays in code: which branch applies is determined by +measured state (history, phase, PMC evidence, M, iteration budget), and only +the wording of each branch is data. That split is deliberate — it keeps +"what we tell the worker to do this round" evolvable without making the +round-selection state machine itself writable. + +Templates use ``str.format`` fields. The only fields any template may use are +``iteration``, ``speedup``, ``artifact_dir``, ``completed``, ``split_candidates``, +``grid_blocks`` and ``cu_count``. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Tuple + +from .harness_io import env_root, load_component_yaml + +#: Component name in ``manifest.yaml``; ``wired: false`` makes this file inert. +COMPONENT = "systemprompt" +#: Path of the strategy menu inside a harness workspace. +RELATIVE_PATH = "systemprompt/round_strategy.yaml" + +#: ``max_iterations`` below this makes the late-round branch unreachable. +LATE_START_FLOOR = 9 + +# --------------------------------------------------------------------------- # +# Built-in defaults (the historical literals, verbatim) +# --------------------------------------------------------------------------- # + +REPAIR_UNKNOWN = ( + "Highest priority: repair the faster but incorrect candidate from " + "iteration {iteration} (measured speedup " + "{speedup}x). Read its archived source at " + "`{artifact_dir}` and preserve the fast mapping. " + "Fix only the smallest correctness defect: signed int8 unpacking, " + "tail bounds, scale indexing, bf16 conversion, or a race. Do not " + "replace it with an unrelated architecture." +) + +REPAIR_INFRASTRUCTURE = ( + "The preceding attempt failed in the agent infrastructure. " + "Return to the accepted best source and start a new bounded " + "experiment; do not repair or replay a partially written " + "candidate. This failure does not count as a completed " + "optimization round." +) + +REPAIR_BUILD = ( + "Repair the immediately preceding candidate from iteration " + "{iteration} at `{artifact_dir}`. " + "Keep its strategy and make only the minimum compile/API/syntax " + "correction; do not start another redesign this round." +) + +ISA_GUIDED = ( + "ISA-guided HIP round. This is successful ISA experiment " + "{completed} of at least 2. Select one measured memory or " + "compute bottleneck, compare the exact primary-kernel ISA with " + "the preceding code object, and make one HIP/DUMMA/intrinsic " + "code-shaping change. Raw inline asm remains forbidden. Record " + "a compiler limitation only when the before/after binary proves " + "it and name the exact target instructions." +) + +CONDITIONAL_INLINE_ASM = ( + "Conditional inline-asm experiment. Target only the compiler " + "limitation and exact instructions verified by the immediately " + "preceding ISA-guided HIP round. Keep the asm block minimal, " + "preserve complete constraints/clobbers, and reject it unless " + "the candidate ISA, exact correctness, median, P90, and resources " + "all validate. Do not write raw global/buffer/flat loads or MMAC." +) + +#: Evidence-qualified suffixes appended after a portfolio mandate. The leading +#: space is added by the caller, never stored here: the same text is used both +#: as the built-in fallback and as the template loaded from the harness file, +#: and the two paths have to render identically. +APPEND_SPLIT_CANDIDATES = ( + "Trusted occupancy-probe split candidates for the measured CU " + "count are {split_candidates}. They include non-power-of-two " + "values where useful, are not a whitelist, and must fit the " + "workspace and stage-alignment constraints. Explore outside this " + "set when evidence supports it." +) + +APPEND_GRID_WARNING = ( + "Trusted control-plane warning: current grid has {grid_blocks} " + "blocks for {cu_count} CUs, below the two-blocks-per-CU latency-" + "hiding target. Before micro-optimization, benchmark a finer " + "one-wave zero-barrier grid or multiple legal split-K candidates " + "including combine cost." +) + +#: Shared late-round mandates (the 9th/10th slot of the >=128 and 16..127 +#: portfolios; the <16 portfolio never reaches them). +LATE_ISA_DIAGNOSIS = ( + "Late ISA-diagnosis round. Only if the control-plane plateau gate " + "is open, use one selected ISA Skill and trusted disassembly to " + "shape compiler output through HIP/DUMMA/intrinsics. Raw inline " + "asm remains forbidden. Otherwise continue HIP-only exploration." +) + +LATE_CONDITIONAL_INLINE_ASM = ( + "Final conditional inline-asm round. Raw asm is allowed only when " + "the control plane confirms a HIP plateau and the prior ISA-guided " + "round recorded one concrete compiler limitation plus target " + "instructions. Otherwise make one HIP-only consolidation change." +) + +SMALL: Dict[int, str] = { + 1: "Vectorize contiguous K loads with exact signed-int8 semantics.", + 2: ( + "Increase instruction-level parallelism with independent int32 " + "accumulators or adjacent N outputs; avoid per-K-tile LDS barriers." + ), + 3: ( + "Stage all of tiny A once with at most one barrier, or tune unroll " + "one step if whole-A staging is not cheaper." + ), + 4: ( + "Change one launch variable only: waves per block, N columns per " + "wave, or unroll factor." + ), + 5: ( + "HIP-only memory round: change one vector-load width, contiguous " + "N mapping, or whole-A reuse decision. Raw inline asm is forbidden." + ), + 6: ( + "HIP-only pipeline round: reduce one dependency chain or barrier " + "using ordinary HIP/intrinsics. Raw inline asm is forbidden." + ), + 7: ( + "HIP-only resource round: tune one block size, unroll factor, or " + "live range while preserving coalescing. Raw inline asm is forbidden." + ), + 8: ( + "HIP-only consolidation round: revisit the fastest correct archived " + "mapping and make one final architecture/codegen improvement. Raw " + "inline asm is forbidden." + ), +} + +M16: Dict[int, str] = { + 1: ( + "Establish a minimal 16x16x32 DUMMA tile with the exact API below, " + "one wave per output tile and explicit int32 accumulation. If the " + "seed already has a correct DUMMA kernel, preserve it and instead " + "test the smallest one-wave-per-block, one-N-tile geometry with no " + "cross-wave barrier; do not spend the round reimplementing it." + ), + 2: ( + "Architecture round: measure grid parallelism before polishing. " + "Explore one complete launch geometry among 1/2/4 waves per block " + "and 1/2/4 adjacent N tiles. Prefer enough independent blocks to " + "cover at least all device CUs; report grid_blocks, waves_per_block " + "and estimated_active_cus in proposal.json." + ), + 3: ( + "Architecture round: if the unsplit grid has fewer than two " + "blocks per device CU and K >= 1024, implement and measure " + "split-K=2 plus at least one CU-aligned candidate (which may be " + "non-power-of-two), or test a one-wave zero-barrier geometry that " + "reaches the same parallelism. Write int32 partials into the " + "caller workspace and include the combine+scale kernel in the " + "timed Graph." + ), + 4: ( + "Architecture/pipeline round: explore one of multi-N-tile reuse, " + "A-only staging, or bounded register/LDS prefetch. Retain enough " + "blocks to cover all CUs, state which A/B bytes are reused, and " + "measure whether the change improves normal median/P90." + ), + 5: ( + "HIP-only packed-weight/staging round: compare one packed layout, " + "A-only staging, or B-only staging design. Raw inline asm is forbidden." + ), + 6: ( + "Pipeline round: choose exactly one staging family from direct, " + "A-only LDS, B-only LDS, or A+B LDS using L2/VMEM evidence. Use " + "coalesced 8- or 16-byte cooperative loads and report HBM/LDS byte " + "changes; do not claim asynchronous overlap without evidence." + ), + 7: ( + "Pipeline round: compare single buffering with double buffering " + "only when K>=1024, L2 hit rate is below 70%, and the doubled LDS " + "budget stays below 48 KiB. Count barriers per K step." + ), + 8: ( + "HIP-only resource round: tune one occupancy limiter using actual " + "PMC evidence: waves per block, VGPR live range, LDS footprint, " + "or spill removal. Do not trade repeated HBM reads for occupancy." + ), + 9: LATE_ISA_DIAGNOSIS, + 10: LATE_CONDITIONAL_INLINE_ASM, +} + +LARGE: Dict[int, str] = { + 1: ( + "Establish a correct DUMMA throughput baseline using a 2-D " + "macro-tile. Benchmark 64x64, 64x128, and 128x64 block tiles; " + "record waves per block, VGPRs, LDS bytes, occupancy and TOPS." + ), + 2: ( + "Operand-reuse round: compare direct loads with cooperative " + "A+B LDS staging. Quantify A/B reuse per macro-tile and use " + "vectorized coalesced global loads with a bank-safe LDS layout." + ), + 3: ( + "Pipeline round: compare single and double buffering across K " + "tiles. Retain double buffering only when ISA/PMC evidence shows " + "reduced VMEM stalls without harmful LDS or occupancy growth." + ), + 4: ( + "Tile-shape round: tune M-tile versus N-tile aspect ratio for " + "this exact M/N/K, balancing B reuse, A reuse and enough blocks " + "to occupy every CU. Do not inherit decode launch geometry." + ), + 5: ( + "Packing round: test one weight packing/swizzle that makes each " + "DUMMA B tile vector-loadable and LDS-bank-safe. Include packing " + "outside timing and validate the graph-stable packed layout." + ), + 6: ( + "Epilogue round: fuse per-row and per-column scales, bf16 " + "conversion and the final coalesced store into the compute " + "kernel; remove any unnecessary workspace/combine pass." + ), + 7: ( + "Compute-pipeline round: tune DUMMA issue grouping, prefetch " + "distance and accumulator independence using ISA stall evidence. " + "Raw inline asm remains forbidden." + ), + 8: ( + "Resource round: tune waves per block, VGPR live ranges and LDS " + "footprint from measured occupancy. Recheck the best tile family " + "with normal median/P90 measurements." + ), + 9: LATE_ISA_DIAGNOSIS, + 10: LATE_CONDITIONAL_INLINE_ASM, +} + +BUILTIN_STRATEGY: Dict[str, Any] = { + "repair": { + "faster_wrong": REPAIR_UNKNOWN, + "infrastructure_failure": REPAIR_INFRASTRUCTURE, + "build_failure": REPAIR_BUILD, + }, + "phases": { + "isa_guided_hip": ISA_GUIDED, + "conditional_inline_asm": CONDITIONAL_INLINE_ASM, + }, + "append": { + "split_candidates": APPEND_SPLIT_CANDIDATES, + "grid_warning": APPEND_GRID_WARNING, + }, + "portfolios": { + "small": {str(k): v for k, v in SMALL.items()}, + "m16": {str(k): v for k, v in M16.items()}, + "large": {str(k): v for k, v in LARGE.items()}, + }, +} + +#: Regime keys a harness file may define; a subset falls back per-key. +_PORTFOLIOS = ("small", "m16", "large") + + +# --------------------------------------------------------------------------- # +# Loading +# --------------------------------------------------------------------------- # + +_CACHE_KEY: Optional[Tuple[str, float, float]] = None +_CACHE_VALUE: Optional["RoundStrategy"] = None + + +@dataclass(frozen=True) +class RoundStrategy: + """Round mandates for one harness workspace (built-in or harness file).""" + + repair: Mapping[str, str] + phases: Mapping[str, str] + append: Mapping[str, str] + portfolios: Mapping[str, Mapping[str, str]] + source: str = "builtin" + #: Non-fatal problems found while loading (reported, never raised). + notes: Tuple[str, ...] = field(default=()) + + def text(self, key: str, **fields: Any) -> str: + """Render one template by key path (``repair.build_failure``).""" + node: Any = { + "repair": self.repair, + "phases": self.phases, + "append": self.append, + } + for part in key.split("."): + node = (node or {}).get(part) if isinstance(node, Mapping) else None + if not isinstance(node, str) or not node: + return "" + try: + return node.format(**fields) + except (KeyError, IndexError, ValueError): + return node + + def mandate(self, regime: str, iteration: int) -> Optional[str]: + """The portfolio entry for one round, or None when it has no slot.""" + portfolio = self.portfolios.get(regime) or self.portfolios.get("m16") or {} + entry = portfolio.get(str(int(iteration))) + return entry if isinstance(entry, str) and entry else None + + def slot(self, regime: str, slot: int) -> str: + """A numbered slot, falling back within the portfolio (never raises).""" + portfolio = self.portfolios.get(regime) or {} + for candidate in (str(int(slot)), str(LATE_START_FLOOR)): + entry = portfolio.get(candidate) + if isinstance(entry, str) and entry: + return entry + # Last resort: the highest-numbered slot the portfolio does define. + keys = sorted((k for k in portfolio if str(k).isdigit()), key=int) + return str(portfolio[keys[-1]]) if keys else "" + + +def builtin_strategy() -> Dict[str, Any]: + """A deep-ish copy of the historical defaults (safe to mutate).""" + return { + "repair": dict(BUILTIN_STRATEGY["repair"]), + "phases": dict(BUILTIN_STRATEGY["phases"]), + "append": dict(BUILTIN_STRATEGY["append"]), + "portfolios": { + name: dict(BUILTIN_STRATEGY["portfolios"][name]) + for name in _PORTFOLIOS + }, + } + + +def _component_wired(root: Optional[Path]) -> bool: + """``manifest.yaml`` decides whether the systemprompt component applies.""" + try: + from . import gate_policy as _gates + return _gates.component_wired(COMPONENT, root) + except Exception: # noqa: BLE001 - a missing manifest means "apply it" + return True + + +def _merge(defaults: Dict[str, Any], overlay: Mapping[str, Any], + notes: List[str]) -> Dict[str, Any]: + """Overlay a harness file onto the defaults, key by key. + + A key the file omits keeps its built-in value, so a partial harness file + never blanks out a branch. Wrong types are ignored and reported in + ``notes`` instead of raising: a malformed harness must not stop a run. + """ + out = builtin_strategy() + if not isinstance(overlay, Mapping): + notes.append("strategy file is not a mapping; built-in defaults used") + return out + for section in ("repair", "phases", "append"): + data = overlay.get(section) + if data is None: + continue + if not isinstance(data, Mapping): + notes.append(f"{section}: expected a mapping, ignored") + continue + for key, value in data.items(): + if isinstance(value, str) and value.strip(): + out[section][str(key)] = value + else: + notes.append(f"{section}.{key}: expected a non-empty string, ignored") + portfolios = overlay.get("portfolios") + if portfolios is not None: + if not isinstance(portfolios, Mapping): + notes.append("portfolios: expected a mapping, ignored") + else: + for name, entries in portfolios.items(): + if name not in _PORTFOLIOS: + notes.append(f"portfolios.{name}: unknown regime, ignored") + continue + if not isinstance(entries, Mapping): + notes.append(f"portfolios.{name}: expected a mapping, ignored") + continue + for slot, value in entries.items(): + if isinstance(value, str) and value.strip(): + out["portfolios"][str(name)][str(slot)] = value + else: + notes.append( + f"portfolios.{name}.{slot}: expected a non-empty " + "string, ignored") + return out + + +def load_round_strategy(root: Optional[Path] = None, + *, use_cache: bool = True) -> RoundStrategy: + """Effective round-strategy texts for this harness workspace. + + Order: ``/systemprompt/round_strategy.yaml`` when the component is + wired -> built-in defaults. Cached on (path, file mtime, manifest mtime) so + a mid-run harness swap is picked up, exactly like :mod:`gate_policy`. + """ + global _CACHE_KEY, _CACHE_VALUE + + workspace = Path(root) if root is not None else env_root() + path = workspace / RELATIVE_PATH + manifest = workspace / "manifest.yaml" + try: + key = (str(path), + path.stat().st_mtime if path.is_file() else 0.0, + manifest.stat().st_mtime if manifest.is_file() else 0.0) + except OSError: + key = (str(path), 0.0, 0.0) + + if use_cache and _CACHE_KEY == key and _CACHE_VALUE is not None: + return _CACHE_VALUE + + notes: List[str] = [] + source = "builtin" + data: Dict[str, Any] = builtin_strategy() + if not _component_wired(workspace): + notes.append(f"{COMPONENT} component is wired:false; built-in defaults used") + else: + overlay, error = load_component_yaml(RELATIVE_PATH, root=workspace) + if error: + notes.append(error) + elif overlay: + data = _merge(data, overlay, notes) + source = str(path) + + value = RoundStrategy( + repair=data["repair"], phases=data["phases"], append=data["append"], + portfolios=data["portfolios"], source=source, notes=tuple(notes), + ) + if use_cache: + _CACHE_KEY, _CACHE_VALUE = key, value + return value + + +def reset_cache() -> None: + """Drop the memoised strategy (tests and mid-run swaps).""" + global _CACHE_KEY, _CACHE_VALUE + _CACHE_KEY, _CACHE_VALUE = None, 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_baselines.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_baselines.py index c5f8df79..0aacc03d 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_baselines.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_baselines.py @@ -225,6 +225,35 @@ (8, 4096, 6144, 2048): ("o_proj", 54775.797), (8, 4096, 512, 6144): ("shared_gate_up_proj", 21254.589), (8, 4096, 6144, 256): ("shared_down_proj", 17700.262), + # Model-catalog TP8 M=8 short-decode baselines measured on 2026-09-20 + # (worker29, gfx928) with the same protocol as the 2026-08-12 catalog: + # lmslim int8_utils.matmul_kernel (the W8A8 Triton baseline), GPU events, + # CUDA Graph replay, hot cache, warmups=10, samples=20, launches/sample=5, + # BM16/BN32/BK256 (matmul_int8 default for M<=32). Covers the Hy3 / + # MiniMax M3 / GLM5.2 TP8 operators at M=8 (see + # MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES in the operator API contract). + # Each value is the per-shape median of three independent passes + # (tools/baseline/bench_triton_tp8_m8.py -> tp8_m8_graph_run{1,2,3}.json, + # merged view in tp8_m8_graph_merged.json). The median keeps the estimate + # robust where a single pass was disturbed by a co-tenant on the shared + # card (e.g. minimax shared_down_proj pass 2: 33.760 vs 19.968/19.984). + # Cross-pass spread is <= 1.6% for 12 of 15 shapes, 5.8%/6.8% for + # glm52 q_b_proj / shared_down_proj. + (8, 8, 1280, 4096): ("qkv_proj", 66.432), + (8, 8, 4096, 1024): ("o_proj", 31.472), + (8, 8, 384, 4096): ("shared_gate_up_proj", 57.008), + (8, 8, 4096, 192): ("shared_down_proj", 14.592), + (8, 8, 1280, 6144): ("qkv_proj", 98.865), + (8, 8, 1536, 6144): ("qkv_proj_and_indexer_qk", 99.473), + (8, 8, 6144, 1024): ("o_proj", 32.960), + (8, 8, 768, 6144): ("shared_gate_up_proj", 92.928), + (8, 8, 6144, 384): ("shared_down_proj", 19.984), + (8, 8, 2624, 6144): ("fused_qkv_a_proj", 87.200), + (8, 8, 2048, 2048): ("q_b_proj", 36.704), + (8, 8, 3584, 512): ("kv_b_proj", 15.808), + (8, 8, 6144, 2048): ("o_proj", 58.864), + (8, 8, 512, 6144): ("shared_gate_up_proj", 88.448), + (8, 8, 6144, 256): ("shared_down_proj", 15.136), } 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..ac9068a0 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/w8a8_pipeline.py @@ -10,7 +10,7 @@ import time from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -from typing import Any, Dict +from typing import Any, Dict, Sequence from metainfer.orchestrator.state import StateStore from metainfer.orchestrator.subagent_manager import AgentSpec, SubAgentManager @@ -26,6 +26,326 @@ 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" + +#: Sub-task measurement gate. A DKAO child iterates for hours, so the device it +#: was given can be taken over long after the parent admitted it. Every timed +#: benchmark and every PMC profile re-checks the gate first: a number produced +#: while another workload shares the card is not evidence, it is noise +#: (the same unchanged shape has been observed at 884us -> 3047us). +MEASUREMENT_GATE_VRAM_PERCENT = 90.0 +MEASUREMENT_GATE_UTIL_PERCENT = 0.0 +MEASUREMENT_GATE_WAIT_S = 1800 # re-check every 30 minutes +MEASUREMENT_GATE_MAX_WAITS = 48 # ... for at most 24 hours +#: where the per-check audit trail lands when the task did not pass a state dir +_GATE_LOG_NAME = "measurement_gate.jsonl" + + +class MeasurementGateBlocked(RuntimeError): + """The device did not pass ``VRAM <= 90% and HCU == 0`` in time. + + Raising this instead of measuring is the point: a blocked device yields no + number at all, which the parent treats as an environment failure and + re-measures later (and eventually stops the round) rather than scoring a + polluted reading as a kernel result. + """ + + def __init__(self, gpu: int, checks: int, reason: str = "") -> None: + self.gpu = gpu + self.checks = checks + self.gate_reason = reason + super().__init__( + f"GPU {gpu} did not pass the measurement gate (VRAM <=" + f"{MEASUREMENT_GATE_VRAM_PERCENT:.0f}% and HCU ==" + f"{MEASUREMENT_GATE_UTIL_PERCENT:.0f}%) after {checks} checks" + + (f": {reason}" if reason else "") + ) + + +def _gate_preflight_enabled() -> bool: + """``METAINFER_GPU_PREFLIGHT=off`` is the operator's opt-out of gating. + + Fails *closed*: if the gate module cannot be imported we still gate (all + ``preflight_enabled`` does is read the env/form switch, so a failure here + means something is wrong, not that gating was turned off). + """ + try: + from .gpu_preflight import preflight_enabled + except Exception: # noqa: BLE001 + return True + try: + return bool(preflight_enabled(None)) + except Exception: # noqa: BLE001 + return True + + +def _gate_wait_config(env: Dict[str, str]) -> tuple[int, int]: + def _int(name: str, default: int) -> int: + try: + return int(str(env.get(name) or "").strip() or default) + except (TypeError, ValueError): + return default + + return (_int("METAINFER_GATE_WAIT_SECONDS", MEASUREMENT_GATE_WAIT_S), + _int("METAINFER_GATE_MAX_WAITS", MEASUREMENT_GATE_MAX_WAITS)) + + +def _env_flag_off(env: Dict[str, str]) -> bool: + """True when the operator turned the GPU probe (and so the gate) off.""" + raw = str(env.get("METAINFER_GPU_PREFLIGHT") or "").strip().lower() + return raw in {"0", "false", "no", "off"} + + +def _gate_log(payload: Dict[str, Any], *, state_dir: "Path | None", + env: Dict[str, str]) -> None: + """Best-effort audit row for one gate check (never raises).""" + row = {"ts": time.time(), **payload} + targets = [] + if state_dir is not None: + targets.append(Path(state_dir) / _GATE_LOG_NAME) + raw = str(env.get("METAINFER_CHILD_STATE_DIR") or "").strip() + if raw: + targets.append(Path(raw) / _GATE_LOG_NAME) + for path in targets: + try: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as fh: + fh.write(json.dumps(row, ensure_ascii=False) + "\n") + except OSError: + continue + + +def discard_if_contaminated(*, experiments_path: Path, iteration: int, + iteration_dir: Path, source: Path, + state_dir: Path, + best_commit: str = "") -> Dict[str, Any]: + """Drop a round measured while this worker's device was shared. + + Returns the discard record. When a round is discarded the kernel source is + rolled back to its last good commit, so the round can be re-measured from + the previous state instead of on top of a polluted one. + """ + result = discard_contaminated_round( + experiments_path=experiments_path, iteration=iteration, + state_dir=state_dir, iteration_dir=iteration_dir) + if not result.get("discarded"): + return result + try: + if best_commit: + _run(["git", "reset", "--hard", best_commit], cwd=source) + else: + _run(["git", "checkout", "--", "."], cwd=source) + result["source_restored"] = True + except Exception as exc: # noqa: BLE001 - the discard still stands + result["source_restored"] = False + result["restore_error"] = repr(exc) + return result + + +def ensure_measurement_gate(gpu: int, *, state_dir: "Path | None" = None, + env: Dict[str, str] | None = None, + wait_seconds: int | None = None, + max_waits: int | None = None, + site: str = "benchmark") -> Dict[str, Any]: + """Wait until ``gpu`` passes VRAM <= 90% and HCU == 0 before measuring. + + Returns the passing device check. Raises :class:`MeasurementGateBlocked` + after ``max_waits`` re-checks (default 48 x 30 min = 24 h) so a shared card + can never silently produce a measurement. + """ + env = dict(env or os.environ) + if _env_flag_off(env): + return {"skipped": True, "reason": "METAINFER_GPU_PREFLIGHT disabled"} + from .gpu_preflight import preflight_gpus + + default_wait, default_max = _gate_wait_config(env) + wait_s = int(wait_seconds if wait_seconds is not None else default_wait) + max_waits = int(max_waits if max_waits is not None else default_max) + attempts = 0 + last: Dict[str, Any] = {} + while True: + attempts += 1 + try: + plan = preflight_gpus( + [int(gpu)], samples=1, interval_s=0.0, + vram_limit_percent=MEASUREMENT_GATE_VRAM_PERCENT, + util_tolerance=MEASUREMENT_GATE_UTIL_PERCENT, + ) + check = dict((plan.get("gpus") or {}).get(int(gpu)) or {}) + except Exception as exc: # noqa: BLE001 - fail closed, never measure + check = {"usable": False, "reasons": [f"gate probe failed: {exc!r}"]} + if not check: + check = {"usable": False, "reasons": ["no device state was read"]} + last = check + _gate_log({ + "event": "gate_ok" if check.get("usable") else "gate_blocked", + "site": site, "gpu": int(gpu), "attempt": attempts, + "max_waits": max_waits, "wait_seconds": wait_s, + "usable": bool(check.get("usable")), + "util_percent": check.get("util_percent"), + "vram_percent": check.get("vram_percent"), + "state_source": check.get("state_source"), + "reasons": list(check.get("reasons") or []), + }, state_dir=state_dir, env=env) + if check.get("usable"): + return check + if attempts > max_waits: + reason = "; ".join(str(r) for r in (check.get("reasons") or [])) + raise MeasurementGateBlocked(int(gpu), attempts - 1, reason) + time.sleep(max(0.0, float(wait_s))) + + +def ensure_any_measurement_gate(devices: Sequence[int], *, + state_dir: "Path | None" = None, + env: Dict[str, str] | None = None, + wait_seconds: int | None = None, + max_waits: int | None = None, + site: str = "probe") -> int: + """Wait until *one* of ``devices`` passes the gate; return that device. + + Pinning a helper to a single card is how a task ends up sleeping 24 h while + its siblings sit idle: "GPU 0 is occupied" says nothing about GPU 1. The + candidates keep their given order, so a preferred device still wins when it + is free, and the wait only happens when *no* candidate passes. + """ + candidates = [int(d) for d in devices] + if not candidates: + raise ValueError("ensure_any_measurement_gate needs at least one device") + env = dict(env or os.environ) + if _env_flag_off(env): + return candidates[0] + from .gpu_preflight import preflight_gpus + + default_wait, default_max = _gate_wait_config(env) + wait_s = int(wait_seconds if wait_seconds is not None else default_wait) + max_waits = int(max_waits if max_waits is not None else default_max) + attempts = 0 + while True: + attempts += 1 + try: + plan = preflight_gpus( + candidates, samples=1, interval_s=0.0, + vram_limit_percent=MEASUREMENT_GATE_VRAM_PERCENT, + util_tolerance=MEASUREMENT_GATE_UTIL_PERCENT, + ) + states = dict(plan.get("gpus") or {}) + except Exception as exc: # noqa: BLE001 - fail closed, never measure + states = {d: {"usable": False, + "reasons": [f"gate probe failed: {exc!r}"]} + for d in candidates} + passing = [d for d in candidates if (states.get(d) or {}).get("usable")] + blocked = [d for d in candidates if d not in passing] + _gate_log({ + "event": "gate_ok" if passing else "gate_blocked", + "site": site, "gpu": passing[0] if passing else candidates[0], + "candidates": candidates, "passing": passing, + "attempt": attempts, "max_waits": max_waits, + "wait_seconds": wait_s, "usable": bool(passing), + "reasons": [r for d in blocked + for r in ((states.get(d) or {}).get("reasons") or [])], + }, state_dir=state_dir, env=env) + if passing: + return int(passing[0]) + if attempts > max_waits: + reason = "; ".join( + f"gpu{d}: " + ", ".join( + str(r) for r in ((states.get(d) or {}).get("reasons") or []) + ) for d in blocked) + raise MeasurementGateBlocked(int(candidates[0]), attempts - 1, reason) + time.sleep(max(0.0, float(wait_s))) + + +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, @@ -38,6 +358,7 @@ parse_pmc_csv, ) from .real_pipeline import _last_json, _run, _safe, _status +from .contention import discard_contaminated_round from .result_store import SCHEMA_VERSION, append_jsonl, write_json from .skill_store import generate_merged_skill, generate_worker_skill from .w8a8_baselines import fixed_triton_graph_baseline @@ -108,7 +429,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 +437,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 +470,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 +507,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 +554,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 +579,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 +755,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 +774,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 ) @@ -761,6 +1085,7 @@ def _check_required_files(repo: Path) -> bool: class W8A8Runner: def __init__(self, worker_root: Path, gpu: int) -> None: self.worker_root = worker_root + self.gpu = int(gpu) self.source = worker_root / "source" staged_harness = self.source / "w8a8_bench.py" self.harness = ( @@ -783,6 +1108,27 @@ def __init__(self, worker_root: Path, gpu: int) -> None: ): Path(self.env[key]).mkdir(parents=True, exist_ok=True) + def _ensure_gate(self, site: str) -> Dict[str, Any]: + """Block until this worker's device passes the measurement gate. + + Called immediately before every timed benchmark and PMC profile, so a + device taken over mid-task stops producing numbers instead of producing + polluted ones. ``METAINFER_GPU_PREFLIGHT=off`` keeps the old behaviour. + """ + if not _gate_preflight_enabled(): + return {} + return ensure_measurement_gate( + int(self.gpu), + state_dir=self._gate_state_dir(), + env=self.env, + site=site, + ) + + def _gate_state_dir(self) -> "Path | None": + """Child state dir for gate audit rows, when the task provides one.""" + raw = str(self.env.get("METAINFER_CHILD_STATE_DIR") or "").strip() + return Path(raw) if raw else None + def _prepare_compile_cache(self) -> Dict[str, Any]: """Stage immutable content-addressed compiler inputs for Ninja reuse.""" inputs = [self.source / "csrc" / "bindings.cpp"] @@ -883,6 +1229,7 @@ def benchmark( if not check_correctness: command.append("--skip-correctness") build = self._prepare_compile_cache() + self._ensure_gate("benchmark") result = _run( command, cwd=self.source, env=self.env, timeout=_BENCHMARK_TIMEOUT_S, @@ -952,6 +1299,7 @@ def profile_pmc( ) output_dir.mkdir(parents=True, exist_ok=True) build = self._prepare_compile_cache() + self._ensure_gate("profile_pmc") _run( [ "bash", @@ -1226,7 +1574,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 +2246,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 +2694,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 +2773,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 ) @@ -2488,7 +2852,45 @@ def preflight() -> list[str]: write_json( iteration_dir / "iteration.json", experiment ) - append_jsonl(experiments_path, experiment) + # A round measured while the parent had this device frozen + # (another workload took it over) is not evidence: drop the + # number, quarantine the record and roll the worker back to the + # previous good state so the round is simply re-measured. + contaminated = discard_if_contaminated( + experiments_path=experiments_path, + iteration=iteration, + iteration_dir=iteration_dir, + source=source, + state_dir=self.state_dir, + ) + if contaminated.get("discarded"): + accepted = False + best_commit = _run( + ["git", "rev-parse", "HEAD"], cwd=source + ).stdout.strip() + experiment["contaminated_discarded"] = { + "reason": contaminated.get("reason"), + "quarantined": contaminated.get("quarantined"), + } + self.store.append_timeline( + "worker_round_discarded_contended", + { + "worker_id": assignment.worker_id, + "physical_gpu": assignment.gpu, + "shape_id": shape_id, + "iteration": iteration, + "window": contaminated.get("window"), + "reason": contaminated.get("reason"), + "quarantined": contaminated.get("quarantined"), + "policy": ( + "the device was shared while this round was " + "measured; the number is discarded and the " + "round is re-measured" + ), + }, + ) + else: + append_jsonl(experiments_path, experiment) if candidate_destination is not None: publish_iteration_candidate( iteration_dir, candidate_destination @@ -2711,6 +3113,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 +3127,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 +3236,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 +3296,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 +3464,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_down_proj.sh b/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj.sh new file mode 100644 index 00000000..50003279 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj.sh @@ -0,0 +1,40 @@ +#!/bin/bash +# Verify the recorded 8.67us for hy3_tp8_shared_down_proj_m16 (shape M16 N4096 K192). +# Same protocol as the worker acceptance bench: warmups=100 samples=30 replays=100, +# correctness checked against the cached exact int64 reference. +set -u +WS=/root/zth_agent/MetaInfer/nodes/worker29/workspaces/hy3-dsh-tp8-m16-9-8-0161e718 +GPU=${GPU:-1} +REPEAT=${REPEAT:-2} + +bench () { + local root=$1 label=$2 + 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 192 \ + --reference-cache-dir "$root/cache/references" 2>&1 | tail -3 +} + +for i in $(seq 1 "$REPEAT"); do + echo "########## run $i on GPU$GPU ##########" + for spec in "workers/worker_3:WORKER_3-accepted-build" "final:FINAL-synthesized-build"; do + root="$WS/${spec%%:*}"; label="${spec##*:}" + echo "=== $label (run $i) ===" + bench "$root" "$label" | python3 -c ' +import sys,json +for line in sys.stdin: + line=line.strip() + if not line.startswith("{"): continue + d=json.loads(line) + print("median_us=%.3f min_us=%.3f p90_us=%.3f samples=%s passed=%s ref_cache_hit=%s tops=%.2f bw=%.1f GB/s" % ( + d["median_us"], d["min_us"], d["p90_us"], d.get("samples"), d.get("passed"), + d.get("reference_cache_hit"), d["logical_tops"], d["algorithmic_bandwidth_gb_s"])) +' + done +done diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj_profiler.py b/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj_profiler.py new file mode 100644 index 00000000..3fe7ef2e --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj_profiler.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Independent cross-check for hy3_tp8_shared_down_proj_m16 (M16 N4096 K192). + +Method A: CUDA-event timing over CUDA-graph replay (same as the official harness). +Method B: torch.profiler -> pure GPU kernel time (excludes host launch overhead). + +Run with the worker_3 accepted build (extension cache reused, no rebuild). +""" +from __future__ import annotations + +import importlib.util +import json +import os +import statistics +import sys +from pathlib import Path + +import torch + +WS = Path("/root/zth_agent/MetaInfer/nodes/worker29/workspaces/hy3-dsh-tp8-m16-9-8-0161e718") +ROOT = Path(sys.argv[1]) if len(sys.argv) > 1 else WS / "workers" / "worker_3" +SRC = ROOT / "source" +BUILD = json.loads((ROOT / "cache" / "current_build.json").read_text()) +os.environ["METAINFER_W8A8_COMPILE_SOURCE_DIR"] = BUILD["compile_source_dir"] +os.environ["METAINFER_W8A8_BUILD_KEY"] = BUILD["build_key"] +os.environ["TORCH_EXTENSIONS_DIR"] = str(ROOT / "cache" / "torch") + + +def load_module(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +bench = load_module(SRC / "w8a8_bench.py", "bench_mod") +backend = load_module(SRC / "w8a8_backend.py", "metainfer_w8a8_backend") +backend.load_extension() +api = load_module(SRC / "int8_w8a8_gemm_api.py", "metainfer_w8a8_contract") + +M = int(os.environ.get("SHAPE_M", "16")) +N = int(os.environ.get("SHAPE_N", "4096")) +K = int(os.environ.get("SHAPE_K", "192")) +print(f"== shape M={M} N={N} K={K} root={ROOT.name} ==", flush=True) +a, b, a_scale, b_scale = bench.generate_w8a8_inputs(M, N, K) +out = torch.empty((M, N), dtype=torch.bfloat16, device="cuda") +packed_weight, packed_weight_scale = api.prepare_weight(b, b_scale) +workspace = api.allocate_workspace(M, N, K, a.device) + + +def candidate() -> None: + api.w8a8_gemm_out(a, packed_weight, a_scale, packed_weight_scale, out, workspace) + + +for _ in range(100): + candidate() +torch.cuda.synchronize() + +graph = torch.cuda.CUDAGraph() +with torch.cuda.graph(graph): + candidate() + +# --- Method A: event timing, 30 samples x 100 replays (official protocol) --- +events = [(torch.cuda.Event(True), torch.cuda.Event(True)) for _ in range(30)] +for start, end in events: + start.record() + for _ in range(100): + graph.replay() + end.record() +torch.cuda.synchronize() +per_launch = [s.elapsed_time(e) * 1000.0 / 100.0 for s, e in events] +print("== Method A: CUDA-event graph replay (official protocol) ==") +print(" median_us=%.3f min_us=%.3f p90_us=%.3f n=30" % ( + statistics.median(per_launch), min(per_launch), + sorted(per_launch)[int(0.9 * len(per_launch)) - 1])) + +# --- Method B: profiler kernel-only device time --- +replays = 50 +from torch.profiler import ProfilerActivity, profile # noqa: E402 + +with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(replays): + graph.replay() + torch.cuda.synchronize() + +kernels = [] +for entry in prof.key_averages(): + dev_us = getattr(entry, "self_device_time_total", 0.0) or 0.0 + if dev_us > 0: + kernels.append((entry.key, int(entry.count), dev_us / replays)) +kernels.sort(key=lambda item: -item[2]) +total_per_launch = sum(item[2] for item in kernels) +print("== Method B: torch.profiler GPU kernel time (per graph replay) ==") +for key, count, us in kernels[:6]: + print(" %-58s count=%-4d %8.3f us/replay" % (key[:58], count, us)) +print(" TOTAL kernel time per replay = %.3f us" % total_per_launch) +print() +print("recorded worker accepted median = 8.671 us (min 7.484, p90 9.144); baseline 17.312 us") diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj_sweep.sh b/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj_sweep.sh new file mode 100644 index 00000000..256271a2 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj_sweep.sh @@ -0,0 +1,51 @@ +#!/bin/bash +# Sweep shared_down_proj (hy3_tp8_shared_down_proj_m16, M16 N4096 K192) on idle GPUs. +# Goal: check whether the recorded 8.671us (worker accepted median) is credible. +set -u +WS=/root/zth_agent/MetaInfer/nodes/worker29/workspaces/hy3-dsh-tp8-m16-9-8-0161e718 +OUT=/root/zth_agent/MetaInfer/nodes/worker29/workspaces/hy3-dsh-tp8-m16-9-8-0161e718/diag_down_proj_sweep.log +: > "$OUT" + +bench_once () { # $1=root $2=gpu + local root=$1 gpu=$2 + 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 192 --reference-cache-dir "$root/cache/references" 2>&1 \ + | python3 -c ' +import sys,json +for line in sys.stdin: + line=line.strip() + if line.startswith("{"): + d=json.loads(line) + print("%.3f %.3f %.3f %s" % (d["median_us"], d["min_us"], d["p90_us"], d.get("passed"))) +' +} + +for gpu in 0 1; do + for spec in "workers/worker_3:worker" "final:final"; do + root="$WS/${spec%%:*}"; label="${spec##*:}" + for i in 1 2 3 4; do + res=$(bench_once "$root" "$gpu") + echo "gpu$gpu $label run$i : $res" | tee -a "$OUT" + done + done +done + +echo | tee -a "$OUT" +python3 - "$OUT" <<'EOF' | tee -a "$OUT" +import sys,collections,statistics +groups=collections.defaultdict(list); groupsP=collections.defaultdict(list) +for line in open(sys.argv[1]): + parts=line.split() + if len(parts)<6 or not parts[0].startswith('gpu'): continue + key=(parts[0],parts[1]); groups[key].append(float(parts[3])); groupsP[key].append(float(parts[4])) +print(f"{'group':22s} {'n':>2s} {'median_of_medians':>17s} {'best_median':>12s} {'worst':>8s} {'best_min':>9s}") +for k in sorted(groups): + v=groups[k] + print(f"{k[0]+' '+k[1]:22s} {len(v):>2d} {statistics.median(v):>17.3f} {min(v):>12.3f} {max(v):>8.3f} {min(groupsP[k]):>9.3f}") +EOF +echo "RECORDED worker accepted: median 8.671us / min 7.484us / p90 9.144us ; Triton baseline 17.312us" 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..e9fb3e98 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_make_xlsx.py @@ -0,0 +1,261 @@ +#!/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), + "/".join(sorted({r["validation_state"] for r in rr}))] + ws2.append(vals) + for j, v in enumerate(vals, start=1): + cell = ws2.cell(row=i, column=j) + cell.border = border + if 5 <= j <= 10: + cell.number_format = "0.00" + cell.alignment = Alignment(horizontal="right") + 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, 34], 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 (+TP8 M16 见下条); MiniMax-M3: TP4 M4096, TP8 M16/M4096; GLM5.2: TP8 M4096 —— 共 37 个优化 shape 行，其中 33 行为各任务自身 plan 的最终验收通过算子，4 行为 Hy3 TP8 M16（仅 worker 验收，见下条）；不含默认 42-shape 回归项"], + ["Hy3 TP8 M16 (9.8 任务)", "任务 hy3-dsh-tp8-m16-9-8-0161e718 的 4 个 M=16 shape 目前无成功 final_report：2026-09-09 07:04 最终验收因性能门限失败（o_proj final 11.745 vs worker best 11.044 > 1.05×）。2026-09-10 重跑验收：qkv_proj ✅、o_proj ✅（首测 11.706 超限，第 1 次重测通过）、shared_gate_up_proj 首测 13.1295 vs 门限 13.112（超 0.13%）时按用户要求暂停，shared_down_proj 未测。本表该 4 行数值取 worker 探索期已验收最优 median"], + ["生成时间", "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/_tool_check_decompose.py b/metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_decompose.py new file mode 100644 index 00000000..27754a47 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_decompose.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +"""Decompose the harness metric for hy3_tp8_shared_down_proj_m16. + +per_replay_latency(G) = fixed_overhead + G * kernel_time +Graphs contain G sequential w8a8_gemm_out calls; measure with the official +protocol (samples=30, replays=100) and compare the slope with the profiler +kernel time. +""" +from __future__ import annotations + +import importlib.util +import json +import os +import statistics +from pathlib import Path + +import torch +from torch.profiler import ProfilerActivity, profile + +WS = Path("/root/zth_agent/MetaInfer/nodes/worker29/workspaces/hy3-dsh-tp8-m16-9-8-0161e718") +ROOT = WS / "workers" / "worker_3" +SRC = ROOT / "source" +BUILD = json.loads((ROOT / "cache" / "current_build.json").read_text()) +os.environ["METAINFER_W8A8_COMPILE_SOURCE_DIR"] = BUILD["compile_source_dir"] +os.environ["METAINFER_W8A8_BUILD_KEY"] = BUILD["build_key"] +os.environ["TORCH_EXTENSIONS_DIR"] = str(ROOT / "cache" / "torch") + + +def load_module(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def event_timing(graph, samples=30, replays=100): + evs = [(torch.cuda.Event(True), torch.cuda.Event(True)) for _ in range(samples)] + for s, e in evs: + s.record() + for _ in range(replays): + graph.replay() + e.record() + torch.cuda.synchronize() + per = sorted(s.elapsed_time(e) * 1000.0 / replays for s, e in evs) + return statistics.median(per), per[0] + + +def profiler_kernel_us(graph, replays=50): + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(replays): + graph.replay() + torch.cuda.synchronize() + return sum(float(getattr(e, "self_device_time_total", 0.0) or 0.0) + for e in prof.key_averages()) / replays + + +bench = load_module(SRC / "w8a8_bench.py", "bench_mod") +backend = load_module(SRC / "w8a8_backend.py", "metainfer_w8a8_backend") +backend.load_extension() +api = load_module(SRC / "int8_w8a8_gemm_api.py", "metainfer_w8a8_contract") +M, N, K = 16, 4096, 192 +a8, b8, a_scale, b_scale = bench.generate_w8a8_inputs(M, N, K) +out8 = torch.empty((M, N), dtype=torch.bfloat16, device="cuda") +pw, pws = api.prepare_weight(b8, b_scale) +ws = api.allocate_workspace(M, N, K, a8.device) + + +def gemm(): + api.w8a8_gemm_out(a8, pw, a_scale, pws, out8, ws) + + +# sustained warmup so clocks/state are stable before any measurement +for _ in range(3000): + gemm() +torch.cuda.synchronize() + +results = [] +for g_count in (1, 2, 4, 8): + def body(): + for _ in range(g_count): + gemm() + for _ in range(200): + body() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + body() + med, lo = event_timing(graph) + ker = profiler_kernel_us(graph) + results.append((g_count, med, lo, ker)) + print("G=%-2d event median %8.3f us (min %8.3f) | profiler kernel total %8.3f us -> per kernel %.3f us" + % (g_count, med, lo, ker, ker / g_count), flush=True) + +# linear regression over G +xs = [r[0] for r in results] +ys = [r[1] for r in results] +n = len(xs) +mx, my = sum(xs) / n, sum(ys) / n +slope = sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / sum((x - mx) ** 2 for x in xs) +intercept = my - slope * mx +print() +print("fit: per_replay(G) = %.3f us + G * %.3f us" % (intercept, slope)) +print(" -> fixed per-replay overhead (graph launch/event boundary) = %.3f us" % intercept) +print(" -> marginal cost per gemm kernel = %.3f us" % slope) +print(" -> profiler per-kernel (avg) = %.3f us" % ( + sum(r[3] / r[0] for r in results) / n)) +print() +print("recorded worker accepted median = 8.671 us ; baseline (Triton) = 17.312 us") diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_floor.py b/metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_floor.py new file mode 100644 index 00000000..ab6bc0c5 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_floor.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Measure the fixed per-replay overhead the harness metric includes. + +If (empty-graph floor) + (profiler kernel time) ~= harness per-replay number, +the harness is measuring "per graph-replay call latency" = launch overhead + kernel, +which is the explanation for 8.67us vs ~4.5us kernel-only. +""" +from __future__ import annotations + +import importlib.util +import json +import os +import statistics +from pathlib import Path + +import torch +from torch.profiler import ProfilerActivity, profile + +WS = Path("/root/zth_agent/MetaInfer/nodes/worker29/workspaces/hy3-dsh-tp8-m16-9-8-0161e718") +ROOT = WS / "workers" / "worker_3" +SRC = ROOT / "source" +BUILD = json.loads((ROOT / "cache" / "current_build.json").read_text()) +os.environ["METAINFER_W8A8_COMPILE_SOURCE_DIR"] = BUILD["compile_source_dir"] +os.environ["METAINFER_W8A8_BUILD_KEY"] = BUILD["build_key"] +os.environ["TORCH_EXTENSIONS_DIR"] = str(ROOT / "cache" / "torch") + + +def load_module(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def event_timing(graph, samples=30, replays=100): + evs = [(torch.cuda.Event(True), torch.cuda.Event(True)) for _ in range(samples)] + for s, e in evs: + s.record() + for _ in range(replays): + graph.replay() + e.record() + torch.cuda.synchronize() + per = sorted(s.elapsed_time(e) * 1000.0 / replays for s, e in evs) + return statistics.median(per), per[0], per[int(0.9 * len(per)) - 1] + + +def profiler_kernel_us(graph, replays=50): + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(replays): + graph.replay() + torch.cuda.synchronize() + return sum(float(getattr(e, "self_device_time_total", 0.0) or 0.0) + for e in prof.key_averages()) / replays + + +# ---- 1. empty-ish graph: trivial kernel only (floor of the metric) ---- +tiny_in = torch.zeros(8, dtype=torch.float32, device="cuda") +tiny_out = torch.empty_like(tiny_in) +for _ in range(50): + tiny_out.copy_(tiny_in) +torch.cuda.synchronize() +g_tiny = torch.cuda.CUDAGraph() +with torch.cuda.graph(g_tiny): + tiny_out.copy_(tiny_in) +med, lo, p90 = event_timing(g_tiny) +tiny_kernel = profiler_kernel_us(g_tiny) +print("== trivial-kernel graph (metric floor) ==") +print(" event per replay : median %.3f us min %.3f p90 %.3f" % (med, lo, p90)) +print(" profiler kernel : %.3f us" % tiny_kernel) + +# ---- 2. the real shape, same protocol ---- +bench = load_module(SRC / "w8a8_bench.py", "bench_mod") +backend = load_module(SRC / "w8a8_backend.py", "metainfer_w8a8_backend") +backend.load_extension() +api = load_module(SRC / "int8_w8a8_gemm_api.py", "metainfer_w8a8_contract") +M, N, K = 16, 4096, 192 +a8, b8, a_scale, b_scale = bench.generate_w8a8_inputs(M, N, K) +out8 = torch.empty((M, N), dtype=torch.bfloat16, device="cuda") +pw, pws = api.prepare_weight(b8, b_scale) +ws = api.allocate_workspace(M, N, K, a8.device) + + +def gemm(): + api.w8a8_gemm_out(a8, pw, a_scale, pws, out8, ws) + + +for _ in range(200): + gemm() +torch.cuda.synchronize() +g8 = torch.cuda.CUDAGraph() +with torch.cuda.graph(g8): + gemm() +med8, lo8, p908 = event_timing(g8) +ker8 = profiler_kernel_us(g8) +print() +print("== hy3_tp8_shared_down_proj_m16 (official protocol: samples=30, replays=100) ==") +print(" event per replay : median %.3f us min %.3f p90 %.3f" % (med8, lo8, p908)) +print(" profiler kernel : %.3f us" % ker8) +print(" implied fixed overhead (median - kernel) : %.3f us" % (med8 - ker8)) +print(" floor - tiny kernel : %.3f us" % (med - tiny_kernel)) +print() +print(" recorded worker accepted: median 8.671 us | baseline 17.312 us") diff --git a/metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_timing.py b/metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_timing.py new file mode 100644 index 00000000..b26ae80f --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_timing.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Validate the measurement methodology of the DKAO W8A8 harness. + +Part 1: Large bf16 GEMM -> event timing (same pattern as w8a8_bench.py) vs + torch.profiler GPU kernel time. For a ms-scale kernel the fixed + per-replay overhead is negligible, so both must agree: this validates + the event primitive + report formula on this platform. +Part 2: shared_down_proj (M16 N4096 K192) -> sweep replays_per_sample to expose + the fixed per-replay overhead, and compare with the profiler kernel time. +""" +from __future__ import annotations + +import importlib.util +import json +import os +import statistics +from pathlib import Path + +import torch +from torch.profiler import ProfilerActivity, profile + +WS = Path("/root/zth_agent/MetaInfer/nodes/worker29/workspaces/hy3-dsh-tp8-m16-9-8-0161e718") +ROOT = WS / "workers" / "worker_3" +SRC = ROOT / "source" +BUILD = json.loads((ROOT / "cache" / "current_build.json").read_text()) +os.environ["METAINFER_W8A8_COMPILE_SOURCE_DIR"] = BUILD["compile_source_dir"] +os.environ["METAINFER_W8A8_BUILD_KEY"] = BUILD["build_key"] +os.environ["TORCH_EXTENSIONS_DIR"] = str(ROOT / "cache" / "torch") + + +def load_module(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def event_timing(graph, samples: int, replays: int) -> tuple[float, float, float]: + evs = [(torch.cuda.Event(True), torch.cuda.Event(True)) for _ in range(samples)] + for s, e in evs: + s.record() + for _ in range(replays): + graph.replay() + e.record() + torch.cuda.synchronize() + per = [s.elapsed_time(e) * 1000.0 / replays for s, e in evs] + per.sort() + return statistics.median(per), per[0], per[int(0.9 * len(per)) - 1] + + +def profiler_kernel_us(graph, replays: int) -> float: + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(replays): + graph.replay() + torch.cuda.synchronize() + total = 0.0 + for entry in prof.key_averages(): + total += float(getattr(entry, "self_device_time_total", 0.0) or 0.0) + return total / replays + + +print("=" * 78) +print("PART 1: large bf16 GEMM 4096x4096x4096 (kernel >> launch overhead)") +print("=" * 78) +a = torch.randn(4096, 4096, dtype=torch.bfloat16, device="cuda") +b = torch.randn(4096, 4096, dtype=torch.bfloat16, device="cuda") +out = torch.empty(4096, 4096, dtype=torch.bfloat16, device="cuda") + + +def big() -> None: + torch.matmul(a, b, out=out) + + +for _ in range(10): + big() +torch.cuda.synchronize() +gbig = torch.cuda.CUDAGraph() +with torch.cuda.graph(gbig): + big() +med, lo, p90 = event_timing(gbig, samples=5, replays=10) +ker = profiler_kernel_us(gbig, replays=10) +flops = 2.0 * 4096 ** 3 +print(" event timing : median %9.2f us (min %.2f, p90 %.2f)" % (med, lo, p90)) +print(" profiler : kernel %9.2f us" % ker) +print(" difference : %+0.2f%% -> effective %.1f TFLOPS (bf16)" % ( + (med - ker) / ker * 100.0, flops / (med * 1e-6) / 1e12)) + +print() +print("=" * 78) +print("PART 2: hy3_tp8_shared_down_proj_m16 (M16 N4096 K192), worker_3 accepted build") +print("=" * 78) +bench = load_module(SRC / "w8a8_bench.py", "bench_mod") +backend = load_module(SRC / "w8a8_backend.py", "metainfer_w8a8_backend") +backend.load_extension() +api = load_module(SRC / "int8_w8a8_gemm_api.py", "metainfer_w8a8_contract") + +M, N, K = 16, 4096, 192 +a8, b8, a_scale, b_scale = bench.generate_w8a8_inputs(M, N, K) +out8 = torch.empty((M, N), dtype=torch.bfloat16, device="cuda") +pw, pws = api.prepare_weight(b8, b_scale) +ws = api.allocate_workspace(M, N, K, a8.device) + + +def gemm() -> None: + api.w8a8_gemm_out(a8, pw, a_scale, pws, out8, ws) + + +for _ in range(200): + gemm() +torch.cuda.synchronize() +g8 = torch.cuda.CUDAGraph() +with torch.cuda.graph(g8): + gemm() + +print(" replays/sample sweep (samples=30, official formula = elapsed/replays):") +for replays in (1, 5, 20, 100, 500): + med, lo, p90 = event_timing(g8, samples=30, replays=replays) + print(" replays=%-4d median %7.3f us min %7.3f us p90 %7.3f us" % ( + replays, med, lo, p90)) +ker8 = profiler_kernel_us(g8, replays=50) +print(" profiler pure kernel time: %.3f us (bw %.1f GB/s of int8 weights)" % ( + ker8, (K * N) / (ker8 * 1e-6) / 1e9)) +print(" recorded worker accepted median = 8.671 us / min 7.484 us / baseline 17.312 us") 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..14e8cab0 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.csv @@ -0,0 +1,38 @@ +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 +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,最终验收通过 +hy3-dsh-tp8-m16-9-8-0161e718,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M16-9.8,8,qkv_proj,16,1280,4096,INT8 W8A8,74.782,19.29679036140442,worker_0 accepted/ (n=30 samples),19.29679036140442,19.323190450668335,3.875359507950698,287.53595079506977,,,8.694303915720704,,仅worker验收(未过最终验收) +hy3-dsh-tp8-m16-9-8-0161e718,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M16-9.8,8,o_proj,16,4096,1024,INT8 W8A8,32.137,11.043989658355713,worker_1 accepted/ (n=30 samples),11.043989658355713,11.16318941116333,2.9099085560701914,190.99085560701914,,,12.153011017938878,,仅worker验收(未过最终验收) +hy3-dsh-tp8-m16-9-8-0161e718,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M16-9.8,8,shared_gate_up_proj,16,384,4096,INT8 W8A8,66.527,12.48797059059143,worker_2 accepted/ (n=30 samples),12.48797059059143,12.534370422363281,5.327286729048045,432.7286729048045,,,4.030410516655156,,仅worker验收(未过最终验收) +hy3-dsh-tp8-m16-9-8-0161e718,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M16-9.8,8,shared_down_proj,16,4096,192,INT8 W8A8,17.312,8.670825064182281,worker_3 accepted/ (n=30 samples),8.670825064182281,9.144449830055237,1.9965804720836726,99.65804720836726,,,2.9023563286907703,,仅worker验收(未过最终验收) +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..989c0d54 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.json @@ -0,0 +1,853 @@ +[ + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "task_id": "hy3-dsh-tp8-m16-9-8-0161e718", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP8-M16-9.8", + "tp": 8, + "operator": "qkv_proj", + "M": 16, + "N": 1280, + "K": 4096, + "dtype": "INT8 W8A8", + "baseline_us": 74.782, + "worker_best_us": 19.29679036140442, + "best_variant_source": "worker_0 accepted/ (n=30 samples)", + "final_us": 19.29679036140442, + "final_p90_us": 19.323190450668335, + "speedup_x": 3.875359507950698, + "improvement_pct": 287.53595079506977, + "target_met": null, + "passed": null, + "logical_tops": 8.694303915720704, + "finished_at": "", + "validation_state": "仅worker验收(未过最终验收)" + }, + { + "task_id": "hy3-dsh-tp8-m16-9-8-0161e718", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP8-M16-9.8", + "tp": 8, + "operator": "o_proj", + "M": 16, + "N": 4096, + "K": 1024, + "dtype": "INT8 W8A8", + "baseline_us": 32.137, + "worker_best_us": 11.043989658355713, + "best_variant_source": "worker_1 accepted/ (n=30 samples)", + "final_us": 11.043989658355713, + "final_p90_us": 11.16318941116333, + "speedup_x": 2.9099085560701914, + "improvement_pct": 190.99085560701914, + "target_met": null, + "passed": null, + "logical_tops": 12.153011017938878, + "finished_at": "", + "validation_state": "仅worker验收(未过最终验收)" + }, + { + "task_id": "hy3-dsh-tp8-m16-9-8-0161e718", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP8-M16-9.8", + "tp": 8, + "operator": "shared_gate_up_proj", + "M": 16, + "N": 384, + "K": 4096, + "dtype": "INT8 W8A8", + "baseline_us": 66.527, + "worker_best_us": 12.48797059059143, + "best_variant_source": "worker_2 accepted/ (n=30 samples)", + "final_us": 12.48797059059143, + "final_p90_us": 12.534370422363281, + "speedup_x": 5.327286729048045, + "improvement_pct": 432.7286729048045, + "target_met": null, + "passed": null, + "logical_tops": 4.030410516655156, + "finished_at": "", + "validation_state": "仅worker验收(未过最终验收)" + }, + { + "task_id": "hy3-dsh-tp8-m16-9-8-0161e718", + "model": "Hy3 (gfx928 INT8 W8A8)", + "kernel_repo": "hy3-dsh-TP8-M16-9.8", + "tp": 8, + "operator": "shared_down_proj", + "M": 16, + "N": 4096, + "K": 192, + "dtype": "INT8 W8A8", + "baseline_us": 17.312, + "worker_best_us": 8.670825064182281, + "best_variant_source": "worker_3 accepted/ (n=30 samples)", + "final_us": 8.670825064182281, + "final_p90_us": 9.144449830055237, + "speedup_x": 1.9965804720836726, + "improvement_pct": 99.65804720836726, + "target_met": null, + "passed": null, + "logical_tops": 2.9023563286907703, + "finished_at": "", + "validation_state": "仅worker验收(未过最终验收)" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + }, + { + "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", + "validation_state": "最终验收通过" + } +] \ 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 00000000..dde8c8a4 Binary files /dev/null and b/metainfer/tasks/dcu_kernel_auto_opt/reports/dkao_optimized_operators_bestvariant.xlsx differ diff --git a/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py b/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py index 6f3b845e..fb56ba0d 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" @@ -70,6 +71,132 @@ def _read_jsonl(path: Path) -> list[Dict[str, Any]]: return rows +def _execution_profile(state_dir: Path) -> str: + """``production`` or ``default`` — the queue rank this task runs at.""" + requirements = read_requirements(state_dir) or {} + answers = requirements.get("answers") + raw = "" + if isinstance(answers, dict): + raw = str(answers.get("execution_profile") or "") + raw = raw.strip().lower() + return "production" if raw in {"production", "prod"} else "default" + + +# --------------------------------------------------------------------------- # +# Embedding a harness_evolve child (AHE iteration -> DKAO view) +# --------------------------------------------------------------------------- # +# Every harness_evolve question is a real DKAO run, but it is spawned straight +# from the orchestrator CLI and is therefore NOT in the WebUI task registry -- +# so ``task_or_404`` cannot find it and the DKAO page could not render it. The +# AHE page passes the child's artifact directory instead. Only genuine AHE +# children are accepted, and they are exposed read-only: the roster of a round +# belongs to the HE orchestrator, not to whoever happens to be looking at it. +AHE_CHILD_PREFIX = "ahe-child:" + + +def _ahe_child_entry(state_dir_raw: str): + """Build a read-only entry for a harness_evolve child's artifact dir. + + Raises 404 unless ``state_dir`` is the ``state/`` of a question inside a + *registered* harness_evolve task's ``children/iteration_NNN//`` + tree. The registry is what makes the caller trustworthy: without it any + path on the box could be read through these routes. + """ + from metainfer.server import tasks as _tasks + + state_dir = Path(state_dir_raw) + parts = state_dir.parts + # .../children/iteration_NNN//state + # parts[-1] = "state", parts[-2] = , parts[-3] = iteration_NNN + if state_dir.name != "state" or len(parts) < 3 \ + or not re.fullmatch(r"iteration_\d+", parts[-3]): + raise HTTPException(404, f"not an AHE child directory: {state_dir}") + question_dir = state_dir.parent + iteration_dir = question_dir.parent + children_dir = iteration_dir.parent + if children_dir.name != "children": + raise HTTPException(404, f"not an AHE child directory: {state_dir}") + workspace_dir = children_dir.parent + owner = next( + (e for e in _tasks.list_tasks() + if e.type == "harness-evolve" + and Path(e.workspace_dir or "") == workspace_dir), + None, + ) + if owner is None: + raise HTTPException( + 404, f"no harness-evolve task owns {workspace_dir}") + return _tasks.TaskEntry( + id=f"{AHE_CHILD_PREFIX}{state_dir}", + type=PLUGIN_TYPE, + label=f"{owner.label or owner.id} · {question_dir.name}", + state_dir=str(state_dir), + workspace_dir=str(question_dir / "workspace"), + created_at=owner.created_at, + launcher=owner.launcher, + ) + + +def _entry_for(request: Request, task_id: str): + """``task_or_404``, plus the embed path used by the AHE page.""" + state_dir_raw = str(request.query_params.get("state_dir") or "").strip() + if state_dir_raw: + return _ahe_child_entry(state_dir_raw) + return task_or_404(task_id) + + +def _is_mirrored_child(entry) -> bool: + """True for the read-only AHE-child view (see ``_ahe_child_entry``).""" + return str(getattr(entry, "id", "")).startswith(AHE_CHILD_PREFIX) + + +def _require_writable(entry) -> None: + """Refuse writes through the embedded mirror. + + An AHE round's workers are steered by the HE orchestrator that spawned + them; letting the drill-down mutate them would let a read-only view change + a run's roster mid-round. The UI hides these controls too -- this is the + half that a hand-made request cannot bypass. + """ + if _is_mirrored_child(entry): + raise HTTPException( + 403, "this DKAO run is embedded read-only from a harness_evolve " + "iteration; steer it from its own task page") + + + + +def _latest_gpu_lease(state_dir: Path) -> Dict[str, Any]: + """The last ``gpu_leases_acquired`` row, or ``{}`` when nothing is leased.""" + latest: Dict[str, Any] = {} + for row in _read_jsonl(state_dir / "timeline.jsonl"): + if row.get("type") == "gpu_leases_acquired": + payload = row.get("payload") + if isinstance(payload, dict): + latest = payload + return latest + + +def _gate_status(state_dir: Path, limit: int = 200) -> Dict[str, Any]: + """Compact view of the admission-gate audit trail. + + Answers "why is nothing happening": a task waiting for VRAM<=90%/HCU==0 on + every candidate device leaves the blocked checks and their reasons here. + """ + rows = _read_jsonl(state_dir / "measurement_gate.jsonl")[-limit:] + blocked = [r for r in rows if r.get("event") == "gate_blocked"] + last = rows[-1] if rows else {} + return { + "checks": len(rows), + "blocked": len(blocked), + "last_event": last.get("event"), + "last_ts": last.get("ts"), + "last_site": last.get("site"), + "last_gpu": last.get("gpu"), + "last_reasons": list(last.get("reasons") or []), + } + + def _running_pid(path: Path) -> int | None: record = _load(path, {}) or {} try: @@ -472,8 +599,8 @@ def build_router(plugin) -> APIRouter: router = APIRouter() @router.get("/summary") - def summary(task_id: str) -> Dict[str, Any]: - entry = task_or_404(task_id) + def summary(task_id: str, request: Request) -> Dict[str, Any]: + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) state_dir = state_dir_for(entry) workspace_dir = workspace_dir_for(entry) @@ -487,11 +614,16 @@ def summary(task_id: str) -> Dict[str, Any]: "plan": _load(workspace_dir / "plan.json", None), "workers": workers, "report": _load(workspace_dir / "final_report.json", None), + # production mode: which queue rank this task took and which devices + # it holds, so the operator can tell a real run from a harness round + "execution_profile": _execution_profile(state_dir), + "gpu_leases": _latest_gpu_lease(state_dir), + "gate": _gate_status(state_dir), } @router.get("/iterations") - def iterations(task_id: str) -> Dict[str, Any]: - entry = task_or_404(task_id) + def iterations(task_id: str, request: Request) -> Dict[str, Any]: + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) return read_worker_lanes(workspace_dir_for(entry), state_dir_for(entry)) @@ -499,8 +631,9 @@ def iterations(task_id: str) -> Dict[str, Any]: async def submit_guidance( task_id: str, worker_id: str, request: Request ) -> Dict[str, Any]: - entry = task_or_404(task_id) + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) + _require_writable(entry) plan = _load(workspace_dir_for(entry) / "plan.json", {}) or {} assigned = { str(item.get("worker_id")) @@ -520,9 +653,10 @@ async def submit_guidance( return {"guidance": guidance} @router.post("/workers/{worker_id}/restart") - def restart_worker(task_id: str, worker_id: str) -> Dict[str, Any]: - entry = task_or_404(task_id) + def restart_worker(task_id: str, worker_id: str, request: Request) -> Dict[str, Any]: + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) + _require_writable(entry) if not re.fullmatch(r"worker_[0-3]", worker_id): raise HTTPException(status_code=400, detail="invalid worker id") state_dir = state_dir_for(entry) @@ -603,8 +737,9 @@ async def rename_repo(task_id: str, request: Request) -> Dict[str, Any]: Body: ``{"new_name": ""}``. Refuses while the task's orchestrator is still running. """ - entry = task_or_404(task_id) + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) + _require_writable(entry) state_dir = state_dir_for(entry) workspace_dir = workspace_dir_for(entry) try: @@ -624,8 +759,8 @@ async def rename_repo(task_id: str, request: Request) -> Dict[str, Any]: raise HTTPException(status_code=400, detail=str(exc)) from exc @router.get("/state-graph") - def state_graph(task_id: str) -> Dict[str, Any]: - entry = task_or_404(task_id) + def state_graph(task_id: str, request: Request) -> Dict[str, Any]: + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) state_dir = state_dir_for(entry) run = read_run(state_dir) @@ -642,8 +777,8 @@ def state_graph(task_id: str) -> Dict[str, Any]: ) @router.get("/skills") - def skills(task_id: str) -> Dict[str, Any]: - entry = task_or_404(task_id) + def skills(task_id: str, request: Request) -> Dict[str, Any]: + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) workspace_dir = workspace_dir_for(entry) library = list_skill_library(workspace_dir) @@ -659,9 +794,10 @@ def skills(task_id: str) -> Dict[str, Any]: return library @router.post("/skills/{skill_name}/publish") - def publish(task_id: str, skill_name: str) -> Dict[str, Any]: - entry = task_or_404(task_id) + def publish(task_id: str, skill_name: str, request: Request) -> Dict[str, Any]: + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) + _require_writable(entry) workspace_dir = workspace_dir_for(entry) plan = _load(workspace_dir / "plan.json", {}) or {} if plan.get("execution_mode") in {LEGACY_SMOKE_MODE, SMOKE_MODE}: @@ -684,10 +820,11 @@ def publish(task_id: str, skill_name: str) -> Dict[str, Any]: raise HTTPException(status_code=409, detail=str(exc)) from exc @router.post("/skills/sync") - def sync_skills(task_id: str) -> Dict[str, Any]: + def sync_skills(task_id: str, request: Request) -> Dict[str, Any]: """Manually mirror the authoritative dsh library into the ccb library.""" - entry = task_or_404(task_id) + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) + _require_writable(entry) workspace_dir = workspace_dir_for(entry) return {"sync": sync_skill_libraries(workspace_dir=workspace_dir)} @@ -696,8 +833,9 @@ async def fuse_skill_route(task_id: str, request: Request) -> Dict[str, Any]: """Trigger the main-agent fusion of one pending skill into the dsh library (mirrored to ccb afterwards). Runs in the background; poll ``GET /skills`` ``fuse_status`` for progress.""" - entry = task_or_404(task_id) + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) + _require_writable(entry) workspace_dir = workspace_dir_for(entry) state_dir = state_dir_for(entry) plan = _load(workspace_dir / "plan.json", {}) or {} @@ -751,11 +889,12 @@ def _run_fusion() -> None: } @router.post("/skills/{skill_name}/rollback") - def rollback(task_id: str, skill_name: str) -> Dict[str, Any]: + def rollback(task_id: str, skill_name: str, request: Request) -> Dict[str, Any]: """Restore the latest SKILL.md backup of a fused skill in the dsh library and re-mirror to ccb.""" - entry = task_or_404(task_id) + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) + _require_writable(entry) workspace_dir = workspace_dir_for(entry) try: return {"result": rollback_skill(skill_name, workspace_dir=workspace_dir)} @@ -765,10 +904,10 @@ def rollback(task_id: str, skill_name: str) -> Dict[str, Any]: raise HTTPException(status_code=404, detail=str(exc)) from exc @router.get("/variants") - def variants_index(task_id: str) -> Dict[str, Any]: + def variants_index(task_id: str, request: Request) -> Dict[str, Any]: """Return the fine-grained variant index (which shapes are already captured in the shared variant library).""" - entry = task_or_404(task_id) + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) return {"variants": list_variant_index()} @@ -777,8 +916,9 @@ async def add_variant_route(task_id: str, request: Request) -> Dict[str, Any]: """Add one optimized shape's accepted kernel into the shared variant library, organized by operator type / model / TP / specific operator. Replaces an existing section for the same shape (with a file backup).""" - entry = task_or_404(task_id) + entry = _entry_for(request, task_id) require_task_type(entry, PLUGIN_TYPE) + _require_writable(entry) state_dir = state_dir_for(entry) workspace_dir = workspace_dir_for(entry) try: @@ -810,19 +950,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-agent-fields.js b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-agent-fields.js index 00929b32..9a3f4788 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-agent-fields.js +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-agent-fields.js @@ -4,10 +4,11 @@ // globalThis.__metainferOverrides. The agent_framework field is a plain // select; this component reads its current value from allValues and shows // only the models valid for that framework (ccb: Sonnet/Opus; dsh: -// deepseek-v4-flash), auto-correcting the value when the framework changes. +// deepseek-flash-4.1 / deepseek-v4-flash), auto-correcting the value when the +// framework changes. // // The label→model mapping must stay in sync with -// orchestrator/config.py (AGENT_FRAMEWORKS). +// orchestrator/config.py (CLAUDE_MODELS / DSH_MODEL_IDS). import { html } from "htm/preact"; import { useEffect } from "preact/hooks"; @@ -15,11 +16,12 @@ import { useEffect } from "preact/hooks"; // Framework label → allowed model labels (mirror of config.py). const FRAMEWORK_MODELS = { ccb: ["Opus", "Sonnet"], - dsh: ["deepseek-v4-flash"], + // deepseek-flash-4.1 first: it is the framework default. + dsh: ["deepseek-flash-4.1", "deepseek-v4-flash"], }; const DEFAULT_MODEL = { ccb: "Opus", - dsh: "deepseek-v4-flash", + dsh: "deepseek-flash-4.1", }; function AgentModelField({ field, value, onChange, allValues }) { 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..49e378f7 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js @@ -20,6 +20,31 @@ async function postJson(url, body) { return payload; } +// --------------------------------------------------------------------------- +// Embedding context +// --------------------------------------------------------------------------- +// This view renders a DKAO run. Usually that is a task of this type, resolved +// by task_id. The AHE page also embeds it, pointing at one question of one +// iteration: those runs are spawned by the HE orchestrator and are not in the +// task registry, so the caller passes their artifact directory instead and the +// routes accept it (see routes.py -> _ahe_child_entry). Such an embedded view +// is read-only: the round's workers belong to the HE orchestrator. +// +// Every component below takes `ctx = { taskId, stateDir, readOnly }` so both +// the standalone page and the embedded one share one code path. +export function apiBase({ taskId, stateDir }) { + const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId || "")}`; + return stateDir ? `${base}?state_dir=${encodeURIComponent(stateDir)}` : base; +} + +function ctxOf(props) { + return { + taskId: props.taskId, + stateDir: props.stateDir || "", + readOnly: Boolean(props.readOnly), + }; +} + function StateMachine({ graph }) { const nodes = graph?.nodes || []; if (!nodes.length) return html`

Waiting for the control plane.

`; @@ -324,7 +349,7 @@ function GuidanceCard({ item }) { `; } -function WorkerLane({ lane, taskId, onSaved }) { +function WorkerLane({ lane, ctx, onSaved }) { const [text, setText] = useState(""); const [saving, setSaving] = useState(false); const [restarting, setRestarting] = useState(false); @@ -335,7 +360,7 @@ function WorkerLane({ lane, taskId, onSaved }) { setSaving(true); setMessage(""); try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); await postJson( `${base}/workers/${encodeURIComponent(lane.worker_id)}/guidance`, { text: text.trim() }, @@ -357,7 +382,7 @@ function WorkerLane({ lane, taskId, onSaved }) { setRestarting(true); setMessage(""); try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); await postJson( `${base}/workers/${encodeURIComponent(lane.worker_id)}/restart`, {}, ); @@ -412,7 +437,7 @@ function WorkerLane({ lane, taskId, onSaved }) { ${lane.state} ${lane.assigned_shapes?.join(", ") || "No shapes assigned"} - ${["failed", "timed_out"].includes(lane.state) ? html` + ${!ctx.readOnly && ["failed", "timed_out"].includes(lane.state) ? html` - + `} ${message ? html`${message}` : null} @@ -523,7 +549,7 @@ function SkillFuseStatus({ status, onRollback }) { ${synced && (synced.added?.length || synced.updated?.length) ? ` · 已同步到 ccb 库${synced.added?.length ? `（新增 ${synced.added.length}）` : ""}${synced.updated?.length ? `（更新 ${synced.updated.length}）` : ""}` : ""} - ${status.action === "merge" ? html` + ${status.action === "merge" && onRollback ? html` ` : null} @@ -534,19 +560,19 @@ function SkillFuseStatus({ status, onRollback }) { `; } -function SkillLibrary({ taskId, onClose }) { +function SkillLibrary({ ctx, onClose }) { const [library, setLibrary] = useState(null); const [error, setError] = useState(""); const [busy, setBusy] = useState(""); const load = useCallback(async () => { try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); setLibrary(await fetchJson(`${base}/skills`)); setError(""); } catch (err) { setError(String(err)); } - }, [taskId]); + }, [ctx.taskId, ctx.stateDir]); useEffect(() => { load(); }, [load]); // While a fusion job is running, poll the library so the status flips to // done/error without a manual refresh. @@ -559,7 +585,7 @@ function SkillLibrary({ taskId, onClose }) { if (!window.confirm(`Add "${skill.name}" to the existing skill library (dsh, mirrored to Claude)?`)) return; setBusy(`publish-${skill.name}`); try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); await postJson( `${base}/skills/${encodeURIComponent(skill.name)}/publish`, {}, ); @@ -576,7 +602,7 @@ function SkillLibrary({ taskId, onClose }) { )) return; setBusy(`fuse-${skill.name}`); try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); await postJson(`${base}/skills/fuse`, { skill_name: skill.name }); await load(); } catch (err) { @@ -588,7 +614,7 @@ function SkillLibrary({ taskId, onClose }) { const syncNow = async () => { setBusy("sync"); try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); await postJson(`${base}/skills/sync`, {}); await load(); } catch (err) { @@ -601,7 +627,7 @@ function SkillLibrary({ taskId, onClose }) { if (!window.confirm(`Restore the last backup of skill "${name}"?`)) return; setBusy(`rollback-${name}`); try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); await postJson(`${base}/skills/${encodeURIComponent(name)}/rollback`, {}); await load(); } catch (err) { @@ -617,11 +643,18 @@ function SkillLibrary({ taskId, onClose }) {

Skill library

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

+ ${ctx.readOnly ? null : html` + onClick=${syncNow}>${busy === "sync" ? "Syncing…" : "同步 Skill 库"}`} ${error ? html`
${error}
` : null} + ${readOnly ? html` +
+ 只读镜像 —— 该 DKAO 运行是 harness_evolve 的一轮考题，由 HE 编排器调度； + 本页不提供引导 / 重启 / 发布等写操作。 +
+ ` : null} ${library?.publish_disabled_reason ? html`
${library.publish_disabled_reason} @@ -629,7 +662,8 @@ function SkillLibrary({ taskId, onClose }) { ? ` (${library.quarantined_count} pending files quarantined)` : ""}
` : null} - <${SkillFuseStatus} status=${library?.fuse_status} onRollback=${rollback} /> + <${SkillFuseStatus} status=${library?.fuse_status} + onRollback=${ctx.readOnly ? null : rollback} /> ${!library ? html`

Loading skills…

` : html`
@@ -648,8 +682,8 @@ function SkillLibrary({ taskId, onClose }) { ${(library.pending || []).map((skill) => html` <${SkillFile} skill=${skill} - canPublish=${true} - canFuse=${true} + canPublish=${!ctx.readOnly} + canFuse=${!ctx.readOnly} busy=${busy} onPublish=${publish} onFuse=${fuse} @@ -665,14 +699,17 @@ function SkillLibrary({ taskId, onClose }) { `; } -export default function DcuKernelAutoOptDetail({ taskId, data }) { +export default function DcuKernelAutoOptDetail(props) { + const { taskId, data } = props; + const ctx = ctxOf(props); + const readOnly = ctx.readOnly; const [runtime, setRuntime] = useState({ summary: null, graph: null, lanes: [] }); const [error, setError] = useState(null); const [showSkills, setShowSkills] = useState(false); const refresh = useCallback(async () => { if (!taskId) return; try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); const [summary, graph, iterations] = await Promise.all([ fetchJson(`${base}/summary`), fetchJson(`${base}/state-graph`), @@ -683,7 +720,7 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) { } catch (err) { setError(String(err)); } - }, [taskId]); + }, [ctx.taskId, ctx.stateDir]); useEffect(() => { refresh(); }, [refresh]); useEffect(() => { @@ -698,13 +735,13 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) { const loadVariants = useCallback(async () => { if (!taskId) return; try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); const data = await fetchJson(`${base}/variants`); // Map shape -> variant record (incl. median_us/speedup/source) so the // update flow can compare the existing variant against this task. setVariantIndex(new Map((data.variants || []).map((v) => [v.shape, v]))); } catch (err) { /* non-fatal */ } - }, [taskId]); + }, [ctx.taskId, ctx.stateDir]); useEffect(() => { loadVariants(); }, [loadVariants]); const addVariant = async (shape) => { const existing = variantIndex ? variantIndex.get(shape) : null; @@ -729,7 +766,7 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) { setVariantBusy(shape); setVariantMsg(""); try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); const res = await postJson(`${base}/variants`, { shape_id: shape }); const meta = res.meta || {}; const tp = meta.tp != null ? `tp${meta.tp}` : "tp?"; @@ -766,7 +803,7 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) { setRenameBusy(true); setRenameMsg(""); try { - const base = `/api/dcu-kernel-auto-opt/${encodeURIComponent(taskId)}`; + const base = apiBase(ctx); const res = await postJson(`${base}/rename-repo`, { new_name: target.trim(), }); @@ -787,6 +824,12 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) { return html`
${error ? html`
${error}
` : null} + ${readOnly ? html` +
+ 只读镜像 —— 该 DKAO 运行是 harness_evolve 的一轮考题，由 HE 编排器调度； + 本页不提供引导 / 重启 / 发布等写操作。 +
+ ` : null} ${smokeOnly ? html`
Infrastructure smoke only: this task did not execute the selected @@ -807,6 +850,7 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) { Kernel repository — ${repoName} + ${readOnly ? null : html`
${renameMsg ? html`${renameMsg}` : null} -
+
`} ` : null} +
+

+ 执行画像 + + — ${runtime.summary?.execution_profile === "production" + ? "production（优先级 100，排在 harness_evolve 之前；仍过 VRAM≤90%/HCU==0 闸门）" + : "default（调控/占据模式按 gpu_mode 决定）"} + +

+ ${(() => { + const lease = runtime.summary?.gpu_leases || {}; + const gpus = lease.gpus || []; + const gate = runtime.summary?.gate || {}; + return html` + ${gpus.length ? html` +
+ 已租设备：${gpus.map((g) => `GPU${g}`).join(", ")} + · 模式 ${lease.mode || "-"} + · 优先级 ${lease.priority ?? "-"} + · holder ${lease.holder || "-"} +
+ ` : html`
尚未取得设备租约
`} + ${gate.blocked ? html` +
+ 准入闸门：共 ${gate.checks} 次检查，其中 ${gate.blocked} 次因设备被占用而等待 + — 最近一次 ${gate.last_event || "-"}@${gate.last_site || "-"} + GPU${gate.last_gpu ?? "-"}：${(gate.last_reasons || []).join("; ") || "-"} +
+ ` : html`
准入闸门：暂无阻塞记录（VRAM≤90% 且 HCU==0）
`} + `; + })()} +
+

State machine @@ -841,7 +918,7 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) {

Iterations four isolated worker lanes

${runtime.lanes.map((lane) => html` - <${WorkerLane} lane=${lane} taskId=${taskId} onSaved=${refresh} /> + <${WorkerLane} lane=${lane} ctx=${ctx} onSaved=${refresh} /> `)}
@@ -868,6 +945,7 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) { ${metric(item.metrics.logical_tops ?? item.metrics.tflops, 3)} ${metric(item.metrics.algorithmic_bandwidth_gb_s ?? item.metrics.bandwidth_gb_s, 2)} GB/s + ${readOnly ? null : html` + `} + ${(() => { + 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`
当前: —
`; + })()} `)} @@ -927,10 +1018,11 @@ export default function DcuKernelAutoOptDetail({ taskId, data }) { + ${readOnly ? html`

以只读方式嵌入：不提供发布/同步操作。

` : null} ${showSkills ? html` - <${SkillLibrary} taskId=${taskId} onClose=${() => setShowSkills(false)} /> + <${SkillLibrary} ctx=${ctx} onClose=${() => setShowSkills(false)} /> ` : null}
`; diff --git a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-shape-input.js b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao-shape-input.js index c40b4da7..a0240b49 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 @@ -24,9 +24,11 @@ function parseCsvInts(raw) { var W8A8_M_VALUES = [2, 16, 3072]; // Large-prefill boundary added on 2026-08-06 for TP4; extended to the // Hy3 / MiniMax M3 / GLM5.2 TP8 catalogs on 2026-08-27 via per-topology -// mValues. DeepSeek TP1/TP8 keep the three original M values. The backend -// validates shapes against the contract. +// mValues, plus the M=8 short-decode boundary on 2026-09-20. DeepSeek TP1/TP8 +// keep the three original M values. The backend validates shapes against the +// contract (MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES = (8, 4096)). var W8A8_TP4_M_VALUES = [2, 16, 3072, 4096]; +var W8A8_MODEL_TP8_M_VALUES = [2, 8, 16, 3072, 4096]; // Model label → W8A8 GEMM workload. (M, K) @ (K, N) per weight name and TP // size, mirroring the model cards in the operator docs. DeepSeek keeps the @@ -94,7 +96,7 @@ var MODEL_WORKLOADS = { }, { tp_size: 8, - mValues: [2, 16, 3072, 4096], + mValues: W8A8_MODEL_TP8_M_VALUES, operators: [ { operator: "qkv_proj", K: 4096, N: 1280 }, { operator: "o_proj", K: 1024, N: 4096 }, @@ -129,7 +131,7 @@ var MODEL_WORKLOADS = { }, { tp_size: 8, - mValues: [2, 16, 3072, 4096], + mValues: W8A8_MODEL_TP8_M_VALUES, operators: [ { operator: "qkv_proj", K: 6144, N: 1280 }, { operator: "qkv_proj_and_indexer_qk", K: 6144, N: 1536 }, @@ -167,7 +169,7 @@ var MODEL_WORKLOADS = { }, { tp_size: 8, - mValues: [2, 16, 3072, 4096], + mValues: W8A8_MODEL_TP8_M_VALUES, operators: [ { operator: "fused_qkv_a_proj", K: 6144, N: 2624 }, { operator: "q_b_proj", K: 2048, N: 2048 }, @@ -186,7 +188,8 @@ function modelCatalog(modelLabel) { var shapes = []; (model.topologies || []).forEach(function (topology) { // Per-topology M values override the defaults; TP4 and the Hy3/MiniMax/ - // GLM TP8 catalogs cover the M=4096 large-prefill boundary. + // GLM TP8 catalogs cover the M=8 short-decode and M=4096 large-prefill + // boundaries. var mValues = topology.mValues || (topology.tp_size === 4 ? W8A8_TP4_M_VALUES : W8A8_M_VALUES); topology.operators.forEach(function (item) { @@ -439,6 +442,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 @@ -450,8 +471,10 @@ function ApiDefaultsPreview(_a) {
${catalog.length} logical shapes. Workloads are grouped by TP size and kept as separate task identities. - Each operator is optimized at M=2, M=16 and the large-prefill boundary - M=3072; TP4 additionally covers M=4096. + Every operator covers the decode and prefill boundaries of its own + catalog (per TP size above), e.g. DeepSeek M=2/16/3072 with TP4 adding + M=4096, and the Hy3 / MiniMax M3 / GLM5.2 TP8 catalogs adding the short + decode boundary M=8 and M=4096.
${tpSizes.map(function (tp) { @@ -470,7 +493,16 @@ function ApiDefaultsPreview(_a) { }); } }); - var mValues = tp === 4 ? W8A8_TP4_M_VALUES : W8A8_M_VALUES; + // Derive the M list from the catalog itself so model-specific + // catalogs (Hy3 / MiniMax M3 / GLM5.2 TP8) show their real M + // values (including the M=8 and M=4096 boundaries) instead of the + // DeepSeek defaults. + var mValues = topoShapes + .map(function (shape) { return shape.M; }) + .filter(function (value, index, all) { + return all.indexOf(value) === index; + }) + .sort(function (a, b) { return a - b; }); return html`
@@ -502,7 +534,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 +549,10 @@ function AiShapeSubset(_a) { useEffect(function () { emit(selected); }, []); + var count = catalog.filter(function (shape) { + return selected[shape.id]; + }).length; + return html`
@@ -523,6 +561,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 +620,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 +668,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/conftest.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/conftest.py index 31681ab9..2b693a7b 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/conftest.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/conftest.py @@ -21,6 +21,20 @@ def client(app): return TestClient(app) +@pytest.fixture(autouse=True) +def _no_measurement_gate(monkeypatch): + """Offline tests never wait on a real device. + + Every timed benchmark and PMC profile now passes the admission gate + (VRAM <= 90% and HCU == 0), which on a busy machine means a 30-minute wait + per check. Unit tests exercise command construction and parsing, not + admission, so the gate is off unless a test turns it back on (see + ``test_gpu_preflight.py``). + """ + monkeypatch.setenv("METAINFER_GPU_PREFLIGHT", "0") + yield + + def register_dkao_task( state_dir, workspace_dir, task_id: str = "dkao-1" ) -> TaskEntry: diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_ahe_child_embed.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_ahe_child_embed.py new file mode 100644 index 00000000..fca34b5c --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_ahe_child_embed.py @@ -0,0 +1,175 @@ +"""The AHE page embeds this view for one question of one iteration. + +Every harness_evolve question *is* a real DKAO run, but it is spawned by the HE +orchestrator's CLI and is therefore absent from the WebUI task registry -- so +``task_or_404`` cannot resolve it and the DKAO page could not render it. The AHE +page passes the child's artifact directory instead, which these routes accept +under two conditions, both checked here: + + * the directory really is a question of a *registered* harness_evolve task + (otherwise any path on the box could be read through these routes), and + * it is read-only: a round's roster belongs to the orchestrator that spawned + it, not to whoever is watching it. +""" + +from __future__ import annotations + +import json + +from metainfer.server import tasks as _tasks +from metainfer.server.tasks import TaskEntry + + +def _he_task_with_child(tmp_path, *, register_owner: bool = True, + task_id: str = "he-embed", question: str = "m16_wqkv_a"): + """One registered HE task with one evaluated question on disk.""" + state_dir = tmp_path / "he-state" + workspace = tmp_path / "he-workspace" + child_state = (workspace / "children" / "iteration_003" / question + / "state") + child_ws = child_state.parent / "workspace" + child_state.mkdir(parents=True) + (child_ws / "workers" / "worker_1").mkdir(parents=True) + (child_state / "requirements.json").write_text(json.dumps({ + "task_id": "ahe-it003-wqkv", "task_type": "dcu-kernel-auto-opt", + "answers": {"shapes": []}, + }), encoding="utf-8") + (child_state / "run.json").write_text(json.dumps({ + "task_id": "ahe-it003-wqkv", "task_type": "dcu-kernel-auto-opt", + "current_phase": "parallel_explore", "current_iteration": 2, + "last_update": 123.0, + }), encoding="utf-8") + (child_ws / "plan.json").write_text(json.dumps({ + "shapes": [{"id": "m16_wqkv_a"}], "execution_mode": "gen-and-opt", + }), encoding="utf-8") + (child_ws / "workers" / "worker_1" / "status.json").write_text( + json.dumps({"worker_id": "worker_1", "gpu": 1, "state": "optimizing"}), + encoding="utf-8") + + state_dir.mkdir(parents=True, exist_ok=True) + if register_owner: + _tasks.add_task(TaskEntry( + id=task_id, type="harness-evolve", label="AHE", + state_dir=str(state_dir), workspace_dir=str(workspace), + created_at=0.0)) + return child_state, child_ws, workspace + + +def test_a_registered_he_child_is_served_through_the_dkao_routes( + tmp_path, client): + """The embedded view gets the same payload the standalone page gets.""" + child_state, _, _ = _he_task_with_child(tmp_path) + + resp = client.get( + "/api/dcu-kernel-auto-opt/ignored/summary", + params={"state_dir": str(child_state)}) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["run"]["current_phase"] == "parallel_explore" + assert body["plan"]["shapes"][0]["id"] == "m16_wqkv_a" + assert body["workers"][0]["worker_id"] == "worker_1" + assert body["gate"] is not None + + lanes = client.get( + "/api/dcu-kernel-auto-opt/ignored/iterations", + params={"state_dir": str(child_state)}) + assert lanes.status_code == 200, lanes.text + # one lane per device, and the one this child wrote shows its real state + by_id = {w["worker_id"]: w for w in lanes.json()["workers"]} + assert "worker_1" in by_id + assert by_id["worker_1"]["gpu"] == 1 + assert by_id["worker_1"]["state"] == "optimizing" + + +def test_the_task_id_in_the_path_is_irrelevant_for_an_embedded_child( + tmp_path, client): + """The AHE page keys the URL by question; only state_dir resolves it.""" + child_state, _, _ = _he_task_with_child(tmp_path) + other = _he_task_with_child(tmp_path / "second", task_id="he-embed-2", + question="m16_other") + first = client.get("/api/dcu-kernel-auto-opt/aaa/summary", + params={"state_dir": str(child_state)}) + second = client.get("/api/dcu-kernel-auto-opt/aaa/summary", + params={"state_dir": str(other[0])}) + assert first.status_code == second.status_code == 200 + assert first.json() == second.json() # same content, different paths + + +def test_an_unregistered_directory_is_refused(tmp_path, client): + """Without a registered HE owner, any path could be read: refuse it.""" + child_state, _, _ = _he_task_with_child(tmp_path, register_owner=False) + resp = client.get("/api/dcu-kernel-auto-opt/x/summary", + params={"state_dir": str(child_state)}) + assert resp.status_code == 404 + assert "harness-evolve" in resp.json()["detail"] + + +def test_arbitrary_paths_are_refused(tmp_path, client): + """Not a child of an iteration -> 404, even if the task exists.""" + _, _, workspace = _he_task_with_child(tmp_path) + probes = [ + tmp_path / "he-state", # the HE task's own state + workspace / "children", # too shallow + workspace / "children" / "iteration_003", # the question's parent + workspace / "children" / "iteration_003" / "m16_wqkv_a", # not "state" + tmp_path / "elsewhere" / "state", # unrelated path + ] + for probe in probes: + probe.mkdir(parents=True, exist_ok=True) + resp = client.get("/api/dcu-kernel-auto-opt/x/summary", + params={"state_dir": str(probe)}) + assert resp.status_code == 404, f"{probe} was accepted" + + +def test_without_a_state_dir_the_routes_behave_exactly_as_before( + tmp_path, client): + """No query parameter -> the ordinary registry lookup.""" + resp = client.get("/api/dcu-kernel-auto-opt/nope/summary") + assert resp.status_code == 404 + assert resp.json()["detail"] == "no such task: nope" + + +def test_the_embedded_view_cannot_write(tmp_path, client): + """Read-only for real: the UI hides these, the server refuses them.""" + child_state, _, _ = _he_task_with_child(tmp_path) + writes = [ + ("post", "/workers/worker_1/guidance", {"text": "steer me"}), + ("post", "/workers/worker_1/restart", {}), + ("post", "/skills/some-skill/publish", {}), + ("post", "/skills/sync", {}), + ("post", "/skills/fuse", {"skill_name": "s"}), + ("post", "/skills/some-skill/rollback", {}), + ("post", "/variants", {"shape_id": "m16"}), + ("post", "/rename-repo", {"new_name": "nope"}), + ] + for method, path, payload in writes: + resp = getattr(client, method)( + f"/api/dcu-kernel-auto-opt/ignored{path}", + params={"state_dir": str(child_state)}, json=payload) + assert resp.status_code == 403, f"{path} -> {resp.status_code}" + assert "read-only" in resp.json()["detail"] + + +def test_reads_are_not_blocked_by_the_write_guard(tmp_path, client): + """The guard must not leak onto GETs, or the mirror would be blank.""" + child_state, _, _ = _he_task_with_child(tmp_path) + for path in ("/summary", "/iterations", "/state-graph", "/skills", + "/variants"): + resp = client.get(f"/api/dcu-kernel-auto-opt/ignored{path}", + params={"state_dir": str(child_state)}) + assert resp.status_code == 200, f"{path} -> {resp.status_code}" + + +def test_a_registered_task_of_another_type_is_still_rejected(tmp_path, client): + """The embed path must not become a way around the type check.""" + _he_task_with_child(tmp_path) # HE task exists + other_state = tmp_path / "dkao-state" + other_ws = tmp_path / "dkao-ws" + other_state.mkdir(parents=True) + other_ws.mkdir(parents=True) + _tasks.add_task(TaskEntry( + id="calc-1", type="calc-theoretical-value", label="calc", + state_dir=str(other_state), workspace_dir=str(other_ws), created_at=0.0)) + + resp = client.get("/api/dcu-kernel-auto-opt/calc-1/summary") + assert resp.status_code == 409 # wrong type, as before diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_config.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_config.py index c7ccec16..46bf00d6 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_config.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_config.py @@ -81,7 +81,26 @@ def test_dsh_framework_resolves_wrapper_model(monkeypatch): assert cfg.claude_model == "deepseek/deepseek-v4-flash-0731" -def test_dsh_default_model_when_agent_model_missing(monkeypatch): +def test_dsh_flash_41_label_resolves_to_flash_model(monkeypatch): + """The new frontend label maps to deepseek/deepseek-flash (4.1).""" + monkeypatch.delenv("DSH_AGENT_MODEL", raising=False) + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["agent_framework"] = "dsh" + req["answers"]["agent_model"] = "deepseek-flash-4.1" + + assert load_config(req).claude_model == "deepseek/deepseek-flash" + # Un-pinned, the older label still resolves to its own pinned build. + req["answers"]["agent_model"] = "deepseek-v4-flash" + assert load_config(req).claude_model == "deepseek/deepseek-v4-flash-0731" + + +def test_dsh_model_pin_collapses_catalog(monkeypatch): + """DSH_AGENT_MODEL pins a mixed-version host to one build for every dsh label.""" monkeypatch.setenv("DSH_AGENT_MODEL", "deepseek/deepseek-v4-flash-0731") req = _req(""" shapes: @@ -90,8 +109,24 @@ def test_dsh_default_model_when_agent_model_missing(monkeypatch): worker_0: {gpu: 0, shapes: [m2]} """) req["answers"]["agent_framework"] = "dsh" + req["answers"]["agent_model"] = "deepseek-flash-4.1" - assert load_config(req).claude_model == DSH_DEFAULT_MODEL_ID + assert load_config(req).claude_model == "deepseek/deepseek-v4-flash-0731" + + +def test_dsh_default_model_when_agent_model_missing(monkeypatch): + monkeypatch.delenv("DSH_AGENT_MODEL", raising=False) + req = _req(""" +shapes: + - {id: m2, M: 2, N: 16, K: 32} +assignments: + worker_0: {gpu: 0, shapes: [m2]} +""") + req["answers"]["agent_framework"] = "dsh" + + cfg = load_config(req) + assert cfg.claude_model == DSH_DEFAULT_MODEL_ID + assert cfg.claude_model == "deepseek/deepseek-flash" def test_dsh_rejects_ccb_only_model(monkeypatch): diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contention.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contention.py new file mode 100644 index 00000000..a58b2fc3 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contention.py @@ -0,0 +1,91 @@ +"""Discarding a worker round that was measured on a shared device.""" + +from __future__ import annotations + +import json +import os +import time +from pathlib import Path + +from ..orchestrator.contention import ( + discard_contaminated_round, iteration_is_contaminated, read_contention, +) + + +def _history(path: Path, rounds: int) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as fh: + for i in range(1, rounds + 1): + fh.write(json.dumps({ + "schema_version": 1, "worker_id": "worker_0", "iteration": i, + "shape_id": "op_a", "accepted": i % 2 == 1, + "metrics": {"median_us": 100.0 - i, "p90_us": 100.0 - i, + "passed": True, "mismatch_count": 0}, + }) + "\n") + + +def _window(path: Path, **fields) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(fields), encoding="utf-8") + + +def test_a_round_inside_the_window_is_discarded_and_quarantined(tmp_path): + state = tmp_path / "state" + history = tmp_path / "runs" / "op_a" / "experiments.jsonl" + iteration_dir = tmp_path / "runs" / "op_a" / "iteration3" + iteration_dir.mkdir(parents=True) + _history(history, 3) + now = time.time() + frozen_from, recovered_at = now - 60, now - 30 + # the round's work happened while the device was shared + stamp = frozen_from + 10 + os.utime(iteration_dir, (stamp, stamp)) + _window(state / "contention.json", question="op_a", gpu=0, active=True, + frozen_from=frozen_from, recovered_at=recovered_at) + + result = discard_contaminated_round( + experiments_path=history, iteration=3, state_dir=state, + iteration_dir=iteration_dir) + + assert result["discarded"] is True + assert result["rounds_removed"] == 1 + kept = [json.loads(l) for l in history.read_text().splitlines()] + assert [r["iteration"] for r in kept] == [1, 2] # round 3 removed + quarantined = json.loads(Path(result["quarantined"][0]).read_text()) + assert quarantined["iteration"] == 3 + assert quarantined["metrics"]["median_us"] == 97.0 # evidence preserved + assert quarantined["reason"] == "device_shared_during_measurement" + + +def test_a_round_outside_the_window_is_kept(tmp_path): + state = tmp_path / "state" + history = tmp_path / "runs" / "op_a" / "experiments.jsonl" + iteration_dir = tmp_path / "runs" / "op_a" / "iteration1" + iteration_dir.mkdir(parents=True) + _history(history, 2) + now = time.time() + # the device was taken over AFTER this round finished + _window(state / "contention.json", frozen_from=now + 30, recovered_at=None) + + result = discard_contaminated_round( + experiments_path=history, iteration=1, state_dir=state, + iteration_dir=iteration_dir) + assert result["discarded"] is False + assert len(history.read_text().splitlines()) == 2 + + +def test_no_contention_record_means_nothing_is_discarded(tmp_path): + history = tmp_path / "runs" / "op_a" / "experiments.jsonl" + _history(history, 2) + result = discard_contaminated_round(experiments_path=history, iteration=1, + state_dir=tmp_path / "state") + assert result["discarded"] is False + assert "no contention" in result["reason"] + + +def test_an_open_window_contaminates_the_running_round(tmp_path): + """Contention still active: the round that was running is not trusted.""" + window = {"frozen_from": time.time() - 10, "active": True} + contaminated, why = iteration_is_contaminated(1, window) + assert contaminated is True and "taken over" in why + assert read_contention(tmp_path / "missing") == {} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contract_snapshot.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contract_snapshot.py index ca271898..9a0324ef 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contract_snapshot.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_contract_snapshot.py @@ -148,20 +148,90 @@ def test_live_contract_accepts_model_catalog_tp8_m4096_shapes(): def test_default_optimization_shapes_unchanged_by_model_tp8_boundary(): - # Adding the model-catalog TP8 M=4096 boundary must not change the + # Adding the model-catalog TP8 M=8/M=4096 boundaries must not change the # DeepSeek-only default workload or its serial-validation fallback scope. from ..api.int8w8a8gemm.int8_w8a8_gemm_api import ( DEFAULT_OPTIMIZATION_SHAPES, MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES, ) - assert MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES == (4096,) + assert MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES == (8, 4096) assert len(DEFAULT_OPTIMIZATION_SHAPES) == 42 assert all( str(shape["id"]).startswith("tp") for shape in DEFAULT_OPTIMIZATION_SHAPES ) +def test_live_contract_accepts_model_catalog_tp8_m8_shapes(): + # Model-catalog TP8 short-decode boundary (2026-09-20): Hy3 / MiniMax M3 / + # GLM5.2 TP8 operators are optimizable at M=8 ("Selected shapes only"); + # DeepSeek TP8 defaults keep the original three M values. + contract = resolve_operator_api("Quantized GEMM", "INT8 W8A8") + shapes = { + "hy3_tp8_qkv_proj_m8": { + "tp_size": 8, "operator": "qkv_proj", + "M": 8, "N": 1280, "K": 4096, + }, + "hy3_tp8_o_proj_m8": { + "tp_size": 8, "operator": "o_proj", + "M": 8, "N": 4096, "K": 1024, + }, + "hy3_tp8_shared_gate_up_proj_m8": { + "tp_size": 8, "operator": "shared_gate_up_proj", + "M": 8, "N": 384, "K": 4096, + }, + "hy3_tp8_shared_down_proj_m8": { + "tp_size": 8, "operator": "shared_down_proj", + "M": 8, "N": 4096, "K": 192, + }, + "minimax_tp8_qkv_proj_m8": { + "tp_size": 8, "operator": "qkv_proj", + "M": 8, "N": 1280, "K": 6144, + }, + "minimax_tp8_qkv_proj_and_indexer_qk_m8": { + "tp_size": 8, "operator": "qkv_proj_and_indexer_qk", + "M": 8, "N": 1536, "K": 6144, + }, + "minimax_tp8_o_proj_m8": { + "tp_size": 8, "operator": "o_proj", + "M": 8, "N": 6144, "K": 1024, + }, + "minimax_tp8_shared_gate_up_proj_m8": { + "tp_size": 8, "operator": "shared_gate_up_proj", + "M": 8, "N": 768, "K": 6144, + }, + "minimax_tp8_shared_down_proj_m8": { + "tp_size": 8, "operator": "shared_down_proj", + "M": 8, "N": 6144, "K": 384, + }, + "glm52_tp8_fused_qkv_a_proj_m8": { + "tp_size": 8, "operator": "fused_qkv_a_proj", + "M": 8, "N": 2624, "K": 6144, + }, + "glm52_tp8_q_b_proj_m8": { + "tp_size": 8, "operator": "q_b_proj", + "M": 8, "N": 2048, "K": 2048, + }, + "glm52_tp8_kv_b_proj_m8": { + "tp_size": 8, "operator": "kv_b_proj", + "M": 8, "N": 3584, "K": 512, + }, + "glm52_tp8_o_proj_m8": { + "tp_size": 8, "operator": "o_proj", + "M": 8, "N": 6144, "K": 2048, + }, + "glm52_tp8_shared_gate_up_proj_m8": { + "tp_size": 8, "operator": "shared_gate_up_proj", + "M": 8, "N": 512, "K": 6144, + }, + "glm52_tp8_shared_down_proj_m8": { + "tp_size": 8, "operator": "shared_down_proj", + "M": 8, "N": 6144, "K": 256, + }, + } + validate_contract_shapes(contract, shapes) + + def test_live_contract_accepts_minimax_m3_tp4_shapes(): # Regression for minimaxm3-dsh-tp4-m4096-1-0c2f84a9: the MiniMax M3 TP4 # M=4096 workload stopped in prepare because TP=4 qkv_proj (K, N) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_execution_profile.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_execution_profile.py new file mode 100644 index 00000000..dd9996c1 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_execution_profile.py @@ -0,0 +1,266 @@ +"""Production profile: priority in the device queue, but never a gate bypass. + +The operator's decision (2026-09-16, option 1): a production DKAO task outranks +the evolving harness when devices are handed out, but it still only measures on +a device that passes VRAM<=90% / HCU==0 — a shared card yields stale numbers no +matter who asked for the run. +""" + +from __future__ import annotations + +import json + +from ..orchestrator import config as cfg + +import importlib.util + +import pytest +_GPU_BROKER_AVAILABLE = ( + importlib.util.find_spec("metainfer.orchestrator.gpu_broker") is not None +) +requires_gpu_broker = pytest.mark.skipif( + not _GPU_BROKER_AVAILABLE, + reason="metainfer.orchestrator.gpu_broker is not available in this tree", +) + + +def _lease(monkeypatch, *, grants, profile, answers=None): + """Run lease_gpus_for_task against a fake broker and report the call.""" + from metainfer.orchestrator import gpu_broker + + calls = [] + + class _FakeBroker: + def acquire(self, holder, **kwargs): + calls.append(kwargs) + return [{"gpu": g, "holder": holder} for g in grants] + + monkeypatch.setattr(gpu_broker, "GpuBroker", _FakeBroker) + assignments = [cfg.WorkerAssignment(worker_id="worker_0", gpu=0, + shape_ids=["s1"])] + merged = {"task_id": "t1"} + merged.update(answers or {}) + rebuilt, info = cfg.lease_gpus_for_task( + assignments, merged, production=(profile == "production")) + return calls, rebuilt, info + + +@requires_gpu_broker +def test_production_leases_at_the_highest_priority(monkeypatch): + calls, rebuilt, info = _lease(monkeypatch, grants=[2], profile="production") + assert calls and calls[0]["priority"] == 100 # PRIORITY_PRODUCTION + assert calls[0]["meta"]["profile"] == "production" + assert info["profile"] == "production" and info["gpus"] == [2] + assert rebuilt[0].gpu == 2 + + +@requires_gpu_broker +def test_default_profile_still_uses_production_priority_on_the_broker(monkeypatch): + """DKAO has always asked at priority 100; the profile makes it explicit.""" + calls, _, info = _lease(monkeypatch, grants=[0], profile="default") + assert calls[0]["priority"] == 100 + assert calls[0]["meta"]["profile"] == "default" + assert info["profile"] == "default" + + +@requires_gpu_broker +def test_production_waits_longer_for_a_device_that_passes_the_gate(monkeypatch): + """Nothing passing the gate: production keeps waiting instead of giving up.""" + calls, _, _ = _lease(monkeypatch, grants=[0], profile="production", + answers={"gpu_lease_max_waits": "2", + "production_max_waits": "48", + "gpu_lease_wait_seconds": "0"}) + assert calls[0]["priority"] == 100 + + +@requires_gpu_broker +def test_no_device_after_the_budget_reports_the_gate_honestly(monkeypatch): + import pytest + + class _EmptyBroker: + def acquire(self, holder, **kwargs): + return [] + + from metainfer.orchestrator import gpu_broker + monkeypatch.setattr(gpu_broker, "GpuBroker", _EmptyBroker) + assignments = [cfg.WorkerAssignment(worker_id="worker_0", gpu=0, + shape_ids=["s1"])] + with pytest.raises(RuntimeError) as excinfo: + cfg.lease_gpus_for_task(assignments, {"task_id": "t1", + "gpu_lease_max_waits": "2", + "gpu_lease_wait_seconds": "0"}, + production=True) + message = str(excinfo.value) + assert "production" in message + # and it must say the gate was NOT skipped + assert "not skipped" in message and "shared card" in message + + +@requires_gpu_broker +def test_load_config_accepts_the_profile_and_forces_the_broker(monkeypatch): + """``production`` implies broker-mediated scheduling (that is the queue).""" + from metainfer.orchestrator import gpu_broker + + seen = {} + + class _FakeBroker: + def acquire(self, holder, **kwargs): + seen.update(kwargs) + return [{"gpu": 1, "holder": holder}] + + monkeypatch.setattr(gpu_broker, "GpuBroker", _FakeBroker) + req = { + "task_id": "prod-1", + "answers": { + "execution_mode": "Generate & optimize (auto-create kernel repo)", + "target_repo_path": "11111", + "shape_config": ("model: m\nassignment_mode: ai\nshape_scope: subset\n" + "shapes:\n - {id: s1, M: 16, N: 64, K: 32, tp_size: 8," + " operator: o_proj}\n"), + "max_iterations": "1", + "execution_profile": "production", + # deliberately ask for occupy: production must override it + "gpu_mode": "occupy", + }, + } + config = cfg.load_config(req) + assert config.execution_profile == "production" + assert config.gpu_mode == "scheduled" # broker carries the priority + assert seen["priority"] == 100 + assert config.gpu_leases.get("profile") == "production" + + +def test_an_unknown_profile_is_rejected(): + import pytest + + req = { + "task_id": "bad-1", + "answers": { + "execution_mode": "Generate & optimize (auto-create kernel repo)", + "target_repo_path": "11111", + "shape_config": ("model: m\nassignment_mode: ai\nshape_scope: subset\n" + "shapes:\n - {id: s1, M: 16, N: 64, K: 32, tp_size: 8," + " operator: o_proj}\n"), + "max_iterations": "1", + "execution_profile": "fastest", + }, + } + with pytest.raises(ValueError): + cfg.load_config(req) + + +@requires_gpu_broker +def test_replace_assignments_keeps_the_profile(monkeypatch): + """Re-planning the workers must not silently downgrade a production task.""" + from metainfer.orchestrator import gpu_broker + + class _FakeBroker: + def acquire(self, holder, **kwargs): + return [{"gpu": 3, "holder": holder}] + + monkeypatch.setattr(gpu_broker, "GpuBroker", _FakeBroker) + req = { + "task_id": "prod-2", + "answers": { + "execution_mode": "Generate & optimize (auto-create kernel repo)", + "target_repo_path": "11111", + "shape_config": ("model: m\nassignment_mode: ai\nshape_scope: subset\n" + "shapes:\n - {id: s1, M: 16, N: 64, K: 32, tp_size: 8," + " operator: o_proj}\n"), + "max_iterations": "1", + "execution_profile": "production", + }, + } + config = cfg.load_config(req) + replanned = cfg.replace_assignments(config, config.assignments) + assert replanned.execution_profile == "production" + assert replanned.gpu_leases.get("priority") == 100 + + +def test_gpu_leases_acquired_event_names_the_profile(): + """The timeline event carries the rank, so the GPU view can label a holder.""" + from ..orchestrator.gen_and_opt_pipeline import GenAndOptPipeline + + events = [] + + class _Store: + task_dir = None + + def append_timeline(self, name, payload): + events.append((name, payload)) + + pipeline = GenAndOptPipeline.__new__(GenAndOptPipeline) + pipeline.store = _Store() + pipeline._register_gpu_leases(type("C", (), { + "gpu_leases": {"holder": "dkao:prod-2", "gpus": [3], "mode": "scheduled", + "profile": "production", "priority": 100}, + })()) + acquired = dict(events)["gpu_leases_acquired"] + assert acquired["profile"] == "production" + assert acquired["priority"] == 100 + + +def test_summary_reports_the_profile_and_the_held_devices(tmp_path): + """The task view shows production mode + which devices the task leased.""" + from ..server.routes import _execution_profile, _latest_gpu_lease + + state_dir = tmp_path / "state" + state_dir.mkdir() + (state_dir / "requirements.json").write_text(json.dumps({ + "task_id": "prod-3", + "answers": {"execution_profile": "production"}, + }), encoding="utf-8") + (state_dir / "timeline.jsonl").write_text("\n".join([ + json.dumps({"ts": 1.0, "type": "phase_start", + "payload": {"phase": "prepare"}}), + json.dumps({"ts": 2.0, "type": "gpu_leases_acquired", "payload": { + "holder": "dkao:prod-3", "gpus": [1, 2], "mode": "scheduled", + "profile": "production", "priority": 100}}), + json.dumps({"ts": 3.0, "type": "gpu_leases_released", + "payload": {"holder": "dkao:prod-3", "gpus": [1, 2]}}), + ]), encoding="utf-8") + + assert _execution_profile(state_dir) == "production" + leases = _latest_gpu_lease(state_dir) + assert leases["gpus"] == [1, 2] and leases["priority"] == 100 + + +def test_summary_defaults_to_the_default_profile(tmp_path): + from ..server.routes import _execution_profile, _latest_gpu_lease + + state_dir = tmp_path / "state" + state_dir.mkdir() + (state_dir / "requirements.json").write_text(json.dumps({ + "task_id": "plain-1", "answers": {}}), encoding="utf-8") + assert _execution_profile(state_dir) == "default" + assert _latest_gpu_lease(state_dir) == {} + + +def test_gate_status_summarises_the_blocked_checks(tmp_path): + """A task parked on the gate must say why, in its own state dir.""" + from ..server.routes import _gate_status + + state_dir = tmp_path / "state" + state_dir.mkdir() + rows = [ + {"ts": 1.0, "event": "gate_blocked", "site": "generate_probe", + "gpu": 0, "attempt": 1, "max_waits": 48, + "reasons": ["device busy: HCU 98.0% > 0.0%"]}, + {"ts": 1801.0, "event": "gate_ok", "site": "generate_probe", "gpu": 1, + "passing": [1], "reasons": []}, + ] + (state_dir / "measurement_gate.jsonl").write_text( + "\n".join(json.dumps(r) for r in rows), encoding="utf-8") + + gate = _gate_status(state_dir) + assert gate["checks"] == 2 and gate["blocked"] == 1 + assert gate["last_event"] == "gate_ok" and gate["last_gpu"] == 1 + assert gate["last_reasons"] == [] + + +def test_gate_status_is_empty_without_an_audit_trail(tmp_path): + from ..server.routes import _gate_status + + state_dir = tmp_path / "state" + state_dir.mkdir() + assert _gate_status(state_dir)["checks"] == 0 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_mode.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_mode.py new file mode 100644 index 00000000..213135c2 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_mode.py @@ -0,0 +1,108 @@ +"""DKAO GPU modes: occupy (default, pins devices) vs scheduled (broker leases).""" + +from __future__ import annotations + +import pytest + +import importlib.util + +_GPU_BROKER_AVAILABLE = ( + importlib.util.find_spec("metainfer.orchestrator.gpu_broker") is not None +) +requires_gpu_broker = pytest.mark.skipif( + not _GPU_BROKER_AVAILABLE, + reason="metainfer.orchestrator.gpu_broker is not available in this tree", +) + + +def _assignments(spec): + from ..orchestrator.config import WorkerAssignment + return [WorkerAssignment(worker_id=w, gpu=g, shape_ids=list(s)) + for w, g, s in spec] + + +def test_rebalance_merges_workers_when_leases_are_fewer(): + """2 leases + 4 planned workers -> 2 lanes, every shape kept.""" + from ..orchestrator.config import rebalance_for_gpus + planned = _assignments([ + ("worker_0", 0, ["a"]), ("worker_1", 1, ["b"]), + ("worker_2", 2, ["c"]), ("worker_3", 3, ["d"]), + ]) + lanes = rebalance_for_gpus(planned, [2, 3]) + assert len(lanes) == 2 # one lane per device + assert sorted(l.gpu for l in lanes) == [2, 3] + assert {l.worker_id for l in lanes} == {"worker_2", "worker_3"} + assert sorted(s for l in lanes for s in l.shape_ids) == ["a", "b", "c", "d"] + + +def test_rebalance_keeps_one_lane_per_worker_when_leases_are_plentiful(): + from ..orchestrator.config import rebalance_for_gpus + planned = _assignments([("worker_0", 0, ["a"]), ("worker_1", 1, ["b"])]) + lanes = rebalance_for_gpus(planned, [0, 1, 2, 3]) + assert len(lanes) == 2 # no empty lanes + assert sorted(l.gpu for l in lanes) == [0, 1] + + +@requires_gpu_broker +def test_lease_gpus_for_task_layouts_workers_on_leases(monkeypatch): + from ..orchestrator import config as C + from metainfer.orchestrator import gpu_broker + + class FakeBroker: + def __init__(self, *a, **k): + self.calls = [] + + def acquire(self, holder, **kwargs): + self.calls.append((holder, kwargs)) + return [{"gpu": 2}, {"gpu": 3}] + + monkeypatch.setattr(gpu_broker, "GpuBroker", FakeBroker) + planned = _assignments([ + ("worker_0", 0, ["a"]), ("worker_1", 1, ["b"]), + ("worker_2", 2, ["c"]), ("worker_3", 3, ["d"]), + ]) + lanes, info = C.lease_gpus_for_task(planned, {"task_id": "t-1"}) + assert info["mode"] == "scheduled" and info["gpus"] == [2, 3] + assert info["holder"] == "dkao:t-1" + assert info["requested_workers"] == 4 + assert info["mapping"] == {"worker_2": 2, "worker_3": 3} + assert len(lanes) == 2 + + +@requires_gpu_broker +def test_lease_gpus_for_task_fails_loudly_when_nothing_is_free(monkeypatch): + from ..orchestrator import config as C + from metainfer.orchestrator import gpu_broker + + class EmptyBroker: + def __init__(self, *a, **k): + pass + + def acquire(self, holder, **kwargs): + return [] + + monkeypatch.setattr(gpu_broker, "GpuBroker", EmptyBroker) + with pytest.raises(RuntimeError, match="no GPU lease"): + C.lease_gpus_for_task(_assignments([("worker_0", 0, ["a"])]), + {"task_id": "t-2", "gpu_lease_max_waits": 2}) + + +def test_scheduled_mode_skips_manual_pin_validation(): + """manual (worker_N -> GPU N) is only enforced in occupy mode.""" + from ..orchestrator.config import OptimizerConfig, WorkerAssignment + import dataclasses + base = dict(operator="Quantized GEMM", dtype="INT8 W8A8", hardware="dcu", + kernel_language="HIP C++", claude_model="m", + execution_mode="gen_and_opt", target_repo_path=None, + shapes={}, assignments=[WorkerAssignment("worker_2", 3, ["a"])], + shape_scope="all", mock_iterations=1, + minimum_improvement_percent=1.0) + occupy = OptimizerConfig(assignment_mode="manual", **base) + scheduled = OptimizerConfig(assignment_mode="scheduled", gpu_mode="scheduled", + **base) + assert occupy.gpu_mode == "occupy" + assert scheduled.gpu_mode == "scheduled" + # the pipeline only validates worker_N -> GPU N for occupy + assert (occupy.assignment_mode == "manual" and occupy.gpu_mode == "occupy") + assert not (scheduled.assignment_mode == "manual" + and scheduled.gpu_mode == "occupy") 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..faa21706 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py @@ -0,0 +1,476 @@ +"""GPU admission pre-check: usage rule + pollution marking.""" + +from __future__ import annotations + +import json + +import pytest + +from ..orchestrator import gpu_preflight as gp + +import importlib.util +from pathlib import Path + +_GPU_BROKER_AVAILABLE = ( + importlib.util.find_spec("metainfer.orchestrator.gpu_broker") is not None +) +requires_gpu_broker = pytest.mark.skipif( + not _GPU_BROKER_AVAILABLE, + reason="metainfer.orchestrator.gpu_broker is not available in this tree", +) +_DTK_HIPPROF_AVAILABLE = Path("/opt/dtk/bin/hipprof").is_file() +requires_dtk_hipprof = pytest.mark.skipif( + not _DTK_HIPPROF_AVAILABLE, + reason="/opt/dtk/bin/hipprof is required (DCU toolchain is absent)", +) + +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 ========================================= +""" + +#: What hy-smi actually prints inside a container that cannot open the +#: management interface: the HCU% column (and every other utilisation value) +#: is a placeholder, NOT a zero. +SMI_UNAVAILABLE = """ +================================= System Management Interface ================================== +================================================================================================ +HCU Temp AvgPwr Perf PwrCap VRAM% HCU% Mode +0 52.0C 146.0W high 300.0W 0% N/A Normal +1 50.0C 149.0W high 300.0W 0% N/A Normal +2 50.0C 148.0W high 300.0W 0% N/A Normal +3 51.0C 148.0W high 300.0W 0% N/A 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_na_utilisation_is_unknown_not_zero(): + """``N/A`` must not become ``0.0``: that is what admitted busy cards.""" + parsed = gp.parse_hy_smi(SMI_UNAVAILABLE) + assert set(parsed) == {0, 1, 2, 3} + assert "util_percent" not in parsed[0] # unknown, not idle + assert parsed[0]["power_w"] == 146.0 # real fields are kept + assert gp.smi_is_incomplete(SMI_UNAVAILABLE) is True + assert gp.smi_is_incomplete(SMI_SAMPLE) is False + + +def test_sample_falls_back_to_sysfs_and_gates_busy_cards(): + """hy-smi cannot see HCU%: sysfs decides, and a busy card stays blocked.""" + def sysfs(): + return {0: {"util_percent": 0.0, "vram_percent": 2.0}, + 1: {"util_percent": 0.0, "vram_percent": 2.0}, + 2: {"util_percent": 74.0, "vram_percent": 20.0}, + 3: {"util_percent": 0.0, "vram_percent": 95.0}} + + states = gp.sample_gpu_state([0, 1, 2, 3], samples=1, interval_s=0, + smi_reader=lambda: SMI_UNAVAILABLE, + sysfs_reader=sysfs) + assert states[2]["util_percent"] == 74.0 + assert states[3]["vram_percent"] == 95.0 + assert not any(row.get("unavailable") for row in states.values()) + checks = {g: gp.check_gpu(g, states[g], total_gb=64.0) for g in range(4)} + assert [g for g in range(4) if checks[g]["usable"]] == [0, 1] + assert checks[2]["usable"] is False and checks[3]["usable"] is False + + +def test_no_source_at_all_means_unavailable(monkeypatch): + states = gp.sample_gpu_state([0, 1], samples=1, interval_s=0, + smi_reader=lambda: "", + sysfs_reader=lambda: {}) + assert states[0]["unavailable"] == 1.0 + assert gp.check_gpu(0, states[0])["usable"] is False + + +def test_usability_requires_vram_limit_and_idle_device(): + """Operator rule: VRAM <= 90% *and* HCU == 0, nothing else counts.""" + ok = gp.check_gpu(0, {"vram_percent": 89.0, "util_percent": 0.0}, + total_gb=64.0) + assert ok["usable"] is True and ok["measurement_suspect"] is False + + too_full = gp.check_gpu(1, {"vram_percent": 91.0, "util_percent": 0.0}, + total_gb=64.0) + assert too_full["usable"] is False + assert any("VRAM 91%" in r for r in too_full["reasons"]) + + busy = gp.check_gpu(2, {"vram_percent": 10.0, "util_percent": 3.0}, + total_gb=64.0) + assert busy["usable"] is False # HCU must be 0 + assert any("HCU 3.0%" in r for r in busy["reasons"]) + + assert gp.DEFAULT_VRAM_LIMIT_PERCENT == 90.0 + assert gp.DEFAULT_UTIL_TOLERANCE == 0.0 + + +def test_busy_device_is_unusable_and_suspect(): + busy = gp.check_gpu(2, {"vram_percent": 40.0, "util_percent": 97.0, + "power_w": 251.0}, foreign_pids=[12345]) + assert busy["usable"] is False # a busy device is never used + assert busy["measurement_suspect"] is True + assert any("busy" in r for r in busy["reasons"]) + # power draw is recorded for diagnosis but is not part of the gate + assert busy["power_w"] == 251.0 + assert not 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_an_idle_card_passes_even_at_high_board_power(): + """The gate is exactly "VRAM <= 90% and HCU == 0" — nothing else blocks.""" + check = gp.check_gpu(1, {"vram_percent": 5.0, "util_percent": 0.0, + "power_w": 250.0}) + assert check["usable"] is True + assert check["reasons"] == [] + + +def test_preflight_gate_lists_all_and_only_passing_devices(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": 0.0, + "power_w": 250.0}, + 2: {"vram_percent": 30.0, "util_percent": 4.0}, + 3: {"vram_percent": 97.0, "util_percent": 0.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, 1] + assert plan["suspect_ids"] == [] + assert plan["blocked_ids"] == [2, 3] + assert plan["over_limit_ids"] == [3] + # only gate-passing devices are ever offered + assert plan["preferred"] == [0, 1] + assert "HCU 4.0%" in plan["gate_reasons"]["2"][0] + + +def test_blocked_devices_are_never_offered_for_a_measurement(monkeypatch): + """Both cards full: the gate offers nothing so the caller waits and re-checks. + + Waiting 30 min and retrying (up to 48 times) is the operator rule; silently + measuring on a device that failed the gate is what this prevents. + """ + 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"] == [] + assert plan["clean_ids"] == [] + assert plan["blocked_ids"] == [0, 1] + assert "VRAM 97% > 90% limit" in plan["gate_reasons"]["0"] + + +def test_unreadable_device_state_is_not_usable(monkeypatch): + """No HCU/VRAM reading must never be treated as "idle".""" + monkeypatch.setattr(gp, "foreign_kfd_pids", lambda *a, **k: []) + check = gp.check_gpu(0, {"unavailable": 1.0, + "missing_fields": ["vram_percent", "util_percent"]}) + assert check["usable"] is False + assert check["unavailable"] is True + assert any("state unavailable" in r for r in check["reasons"]) + + +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"] == {} + + +# ---------------------------------------------------------------- sub-task gate +# A DKAO child iterates for hours, so the device can be taken over long after +# the parent admitted it. These tests pin the *child-side* rule: no timed +# benchmark and no PMC profile may start on a device that fails the gate. + +def test_measurement_gate_waits_then_admits(monkeypatch, tmp_path): + from ..orchestrator import w8a8_pipeline as wp + + blocked = {"gpus": {2: {"usable": False, "util_percent": 91.0, + "vram_percent": 10.0, + "reasons": ["device busy: HCU 91.0% > 0.0%"]}}} + ok = {"gpus": {2: {"usable": True, "util_percent": 0.0, + "vram_percent": 5.0, "state_source": "hy-smi"}}} + plans = [blocked, ok] + calls = [] + + def fake_preflight(ids, **kw): + calls.append((list(ids), kw.get("vram_limit_percent"), + kw.get("util_tolerance"))) + return plans.pop(0) + + sleeps = [] + monkeypatch.setattr("metainfer.tasks.dcu_kernel_auto_opt.orchestrator" + ".gpu_preflight.preflight_gpus", fake_preflight) + monkeypatch.setattr(wp, "_env_flag_off", lambda env: False) + monkeypatch.setattr(wp.time, "sleep", lambda s: sleeps.append(s)) + + check = wp.ensure_measurement_gate(2, state_dir=tmp_path, wait_seconds=1800, + max_waits=48, site="benchmark") + assert check["usable"] is True + assert sleeps == [1800] # exactly one 30-minute re-check + assert calls[0][0] == [2] # only the assigned device + assert calls[0][1] == 90.0 and calls[0][2] == 0.0 # the operator's gate + + rows = [json.loads(l) for l in + (tmp_path / "measurement_gate.jsonl").read_text().splitlines()] + assert [r["event"] for r in rows] == ["gate_blocked", "gate_ok"] + assert rows[0]["site"] == "benchmark" and rows[0]["gpu"] == 2 + assert rows[1]["util_percent"] == 0.0 + + +def test_measurement_gate_gives_up_after_max_waits(monkeypatch, tmp_path): + from ..orchestrator import w8a8_pipeline as wp + + always = {"gpus": {3: {"usable": False, "util_percent": 0.0, + "vram_percent": 96.0, + "reasons": ["VRAM 96% > 90% limit"]}}} + monkeypatch.setattr("metainfer.tasks.dcu_kernel_auto_opt.orchestrator" + ".gpu_preflight.preflight_gpus", + lambda ids, **kw: always) + monkeypatch.setattr(wp, "_env_flag_off", lambda env: False) + sleeps = [] + monkeypatch.setattr(wp.time, "sleep", lambda s: sleeps.append(s)) + + with pytest.raises(wp.MeasurementGateBlocked) as excinfo: + wp.ensure_measurement_gate(3, state_dir=tmp_path, wait_seconds=1800, + max_waits=3) + assert excinfo.value.gpu == 3 + assert excinfo.value.checks == 3 # 3 re-checks, then give up + assert "VRAM 96%" in excinfo.value.gate_reason + assert sleeps == [1800, 1800, 1800] + rows = (tmp_path / "measurement_gate.jsonl").read_text().splitlines() + assert len(rows) == 4 # 3 blocked + the final blocked + + +def test_measurement_gate_fails_closed_when_the_probe_breaks(monkeypatch, tmp_path): + from ..orchestrator import w8a8_pipeline as wp + + def boom(*_a, **_k): + raise RuntimeError("hy-smi exploded") + + monkeypatch.setattr("metainfer.tasks.dcu_kernel_auto_opt.orchestrator" + ".gpu_preflight.preflight_gpus", boom) + monkeypatch.setattr(wp, "_env_flag_off", lambda env: False) + monkeypatch.setattr(wp.time, "sleep", lambda s: None) + with pytest.raises(wp.MeasurementGateBlocked): + wp.ensure_measurement_gate(0, state_dir=tmp_path, wait_seconds=0, + max_waits=1) + + +def test_runner_gates_every_measurement(monkeypatch, tmp_path): + """benchmark() and profile_pmc() both pass the gate before launching.""" + from ..orchestrator import w8a8_pipeline as wp + + gated = [] + monkeypatch.setattr(wp.W8A8Runner, "_ensure_gate", + lambda self, site: gated.append((self.gpu, site))) + monkeypatch.setattr(wp.time, "sleep", lambda s: None) # no real gate waits + + runner = object.__new__(wp.W8A8Runner) + runner.worker_root = tmp_path + runner.source = tmp_path + runner.harness = tmp_path / "w8a8_bench.py" + runner.env = {"METAINFER_GPU_PREFLIGHT": "1"} + runner.gpu = 1 + monkeypatch.setattr(wp.W8A8Runner, "_prepare_compile_cache", + lambda self: {"build_key": "k"}) + monkeypatch.setattr(wp.W8A8Runner, "_ensure_reference_cache", + lambda self, m, n, k: None) + monkeypatch.setattr(wp, "_run", + lambda *a, **k: type("R", (), {"stdout": "{}"})()) + monkeypatch.setattr(wp, "_last_json", lambda out: {}) + + runner.benchmark({"M": 16, "N": 64, "K": 32}) + assert gated == [(1, "benchmark")] + + (tmp_path / "profile_pmc.sh").write_text("#!/bin/bash\n", encoding="utf-8") + monkeypatch.setattr(wp, "parse_pmc_csv", + lambda path: {"available": False, "counters": {}}) + runner.profile_pmc({"M": 16, "N": 64, "K": 32}, tmp_path / "out") + assert gated == [(1, "benchmark"), (1, "profile_pmc")] + + +def test_any_measurement_gate_picks_an_idle_sibling(monkeypatch, tmp_path): + """A busy GPU 0 must not park the task while GPU 1 is free.""" + from ..orchestrator import w8a8_pipeline as wp + + plan = {"gpus": { + 0: {"usable": False, "util_percent": 98.0, "vram_percent": 97.0, + "reasons": ["device busy: HCU 98.0% > 0.0%"]}, + 1: {"usable": True, "util_percent": 0.0, "vram_percent": 5.0, + "state_source": "sysfs"}, + }} + calls = [] + monkeypatch.setattr( + "metainfer.tasks.dcu_kernel_auto_opt.orchestrator.gpu_preflight" + ".preflight_gpus", + lambda ids, **kw: (calls.append(list(ids)), plan)[1]) + monkeypatch.setattr(wp, "_env_flag_off", lambda env: False) + sleeps = [] + monkeypatch.setattr(wp.time, "sleep", lambda s: sleeps.append(s)) + + chosen = wp.ensure_any_measurement_gate([0, 1, 2, 3], state_dir=tmp_path, + wait_seconds=1800, site="probe") + assert chosen == 1 # an idle sibling, not the busy GPU 0 + assert sleeps == [] # and no waiting at all + assert calls == [[0, 1, 2, 3]] # the whole candidate set is checked + rows = [json.loads(l) for l in + (tmp_path / "measurement_gate.jsonl").read_text().splitlines()] + assert rows[0]["event"] == "gate_ok" and rows[0]["passing"] == [1] + + +def test_any_measurement_gate_waits_when_no_candidate_passes(monkeypatch, tmp_path): + from ..orchestrator import w8a8_pipeline as wp + + busy = {"gpus": {d: {"usable": False, "util_percent": 99.0, + "vram_percent": 96.0, + "reasons": [f"gpu{d} VRAM 96% > 90% limit"]} + for d in (0, 1)}} + monkeypatch.setattr( + "metainfer.tasks.dcu_kernel_auto_opt.orchestrator.gpu_preflight" + ".preflight_gpus", lambda ids, **kw: busy) + monkeypatch.setattr(wp, "_env_flag_off", lambda env: False) + sleeps = [] + monkeypatch.setattr(wp.time, "sleep", lambda s: sleeps.append(s)) + + with pytest.raises(wp.MeasurementGateBlocked) as excinfo: + wp.ensure_any_measurement_gate([0, 1], state_dir=tmp_path, + wait_seconds=1800, max_waits=2) + assert sleeps == [1800, 1800] # 30-minute re-checks, then give up + assert "gpu0" in excinfo.value.gate_reason + assert "gpu1" in excinfo.value.gate_reason + + +def test_any_measurement_gate_prefers_the_tasks_own_device(monkeypatch, tmp_path): + from ..orchestrator import w8a8_pipeline as wp + + free = {"gpus": {d: {"usable": True, "util_percent": 0.0, + "vram_percent": 1.0} for d in (0, 3)}} + monkeypatch.setattr( + "metainfer.tasks.dcu_kernel_auto_opt.orchestrator.gpu_preflight" + ".preflight_gpus", lambda ids, **kw: free) + monkeypatch.setattr(wp, "_env_flag_off", lambda env: False) + # 3 is the task's own card and is free: the probe stays on it + assert wp.ensure_any_measurement_gate([3, 0], env={"METAINFER_GPU_PREFLIGHT": "1"}) == 3 + + +def _fake_scaffold(root): + """A scaffold that satisfies the trusted Generate preflight file checks.""" + from ..orchestrator import gen_and_opt_pipeline as gp + + root.mkdir(parents=True, exist_ok=True) + for name in (gp._HARNESS_FILENAME, gp.W8A8_API_FILENAME, + gp._GIT_ATTRIBUTES_FILE, "scaffold_manifest.json", + *gp._SCAFFOLD_FILES): + path = root / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("", encoding="utf-8") + (root / "profile_pmc.sh").write_text( + "#!/bin/bash\n/opt/dtk/bin/hipprof --pmc --pmc-type 3 " + '"$source_dir" "$output_dir"\n', encoding="utf-8") + return root + + +@requires_dtk_hipprof +def test_generate_probe_runs_on_a_device_that_passed_the_gate(monkeypatch, tmp_path): + """The probe is gated, but not pinned to GPU 0 (a busy card must not park it).""" + from ..orchestrator import gen_and_opt_pipeline as gp + + source = _fake_scaffold(tmp_path / "scaffold") + contract = gp.file_digest(source / gp.W8A8_API_FILENAME) + seen = {} + + def fake_run(cmd, **kwargs): + if cmd[-1] == "--self-test": + return type("R", (), {"stdout": '{"passed": true}'})() + if cmd[-1] == "--probe": + env = kwargs.get("env") or {} + seen["probe_gpu"] = env.get("HIP_VISIBLE_DEVICES") + # DTK filters twice when both variables name the same non-zero + # index, which hides the device ("No HIP GPUs are available" on + # GPU 1 while GPU 0 survived): binding must set HIP only. + seen["probe_env"] = dict(env) + return type("R", (), {"stdout": json.dumps({ + "visible_devices": 1, "cudagraph_available": True})})() + return type("R", (), {"stdout": ""})() + + monkeypatch.setattr(gp, "_run", fake_run) + monkeypatch.setattr(gp, "_gate_preflight_enabled", lambda: True) + + calls = [] + + def fake_gate(devices, **kwargs): + calls.append((list(devices), kwargs)) + return 1 + + monkeypatch.setattr(gp, "ensure_any_measurement_gate", fake_gate) + out = gp._validate_generate_scaffold( + source, contract_sha256=contract, shapes={"s1": {"M": 16}}, + probe_devices=[0, 1, 2, 3], state_dir=tmp_path / "state") + + assert calls[0][0] == [0, 1, 2, 3] # every candidate may answer + # the wait is auditable in the task's own state dir, not silently invisible + assert calls[0][1]["state_dir"] == tmp_path / "state" + assert calls[0][1]["site"] == "generate_probe" + assert seen["probe_gpu"] == "1" # and the probe honours the choice + assert "ROCR_VISIBLE_DEVICES" not in seen["probe_env"] + assert out["gpu_probe"]["visible_devices"] == 1 + + +@requires_gpu_broker +def test_probe_candidates_put_the_tasks_own_devices_first(): + from ..orchestrator import gen_and_opt_pipeline as gp + + class _Cfg: + assignments = [type("A", (), {"gpu": 2})(), type("A", (), {"gpu": 0})()] + + devices = gp._probe_gate_devices(_Cfg()) + assert devices[:2] == [2, 0] # the task's own cards, in order + assert 1 in devices and 3 in devices # then the idle siblings + assert len(devices) == len(set(devices)) 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_plugin.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plugin.py index 176eb25d..db897b5b 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plugin.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plugin.py @@ -49,7 +49,7 @@ def test_form_schema_is_available(): assert model["default"] == "Opus" assert model["override_component"] == "agent-model" assert [option["label"] for option in model["options"]] == [ - "Opus", "Sonnet", "deepseek-v4-flash" + "Opus", "Sonnet", "deepseek-flash-4.1", "deepseek-v4-flash" ] 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_retest_tool.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_retest_tool.py new file mode 100644 index 00000000..18ed4186 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_retest_tool.py @@ -0,0 +1,52 @@ +"""A/B retest tool: variant lookup, header parsing, verdict and environment.""" + +from __future__ import annotations + +from pathlib import Path + +from ..tools.retest_variant import ( + BENCH_ENV, bench_env, decide, find_variant, parse_variant_header_median, +) + + +def test_find_variant_tolerates_model_label_drift(tmp_path): + """glm should find the shipped glm52 directory.""" + root = tmp_path / "variant" / "int8w8a8-gemm" + glm = root / "glm52" / "TP8" / "M4096" + glm.mkdir(parents=True) + (glm / "shared_down_proj.hip").write_text("// kernel\n", encoding="utf-8") + assert find_variant("glm", 8, 4096, "shared_down_proj", root) == ( + glm / "shared_down_proj.hip") + assert find_variant("glm", 8, 16, "shared_down_proj", root) is None + assert find_variant("hy3", 4, 4096, "o_proj", root) is None + + +def test_parse_variant_header_median_reads_the_record(): + text = ("// @@variant shape=hy3_tp4_o_proj_m4096 commit=abc added=2026-09-14\n" + "// median_us=453.2 p90_us=454.2 speedup=43.87 baseline_us=1.988e+04\n" + "// code follows\n") + path = Path("/tmp/_variant_header_test.hip") + path.write_text(text, encoding="utf-8") + try: + assert parse_variant_header_median(path) == 453.2 + finally: + path.unlink(missing_ok=True) + assert parse_variant_header_median(Path("/tmp/does-not-exist.hip")) is None + + +def test_bench_env_pins_gfx928_and_the_device(): + env = bench_env(2, base={}) + assert env["PYTORCH_ROCM_ARCH"] == BENCH_ENV["PYTORCH_ROCM_ARCH"] == "gfx928" + assert env["HIP_VISIBLE_DEVICES"] == "2" + assert env["CUDA_VISIBLE_DEVICES"] == "2" + + +def test_decide_requires_margin_and_reproducibility(): + # a clear, reproducible win + assert decide(140.0, 280.0, recorded_us=140.0)["verdict"] == "promote" + # too small a margin + assert decide(275.0, 280.0, recorded_us=275.0)["verdict"] == "keep" + # fast but the re-measurement disagrees with the record -> review + assert decide(140.0, 280.0, recorded_us=100.0)["verdict"] == "needs_review" + # missing measurement + assert decide(None, 280.0)["verdict"] == "measurement_failed" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_round_strategy.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_round_strategy.py new file mode 100644 index 00000000..0c3cd43d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_round_strategy.py @@ -0,0 +1,490 @@ +"""systemprompt/round_strategy.yaml is wired: the worker-facing mandate is data. + +What these tests pin down: + +* the built-in texts are byte-identical to the historical Python literals, so an + un-evolved harness tells the worker exactly what it always did; +* the seed file reproduces those defaults (drift guard, like gates.yaml); +* editing the seed changes the runtime mandate, including inside the pipeline's + own call path; +* ``wired: false`` in the manifest makes the component inert; +* a malformed / partial harness file degrades to built-in text instead of + raising (a broken harness is a performance problem, not a crash). +""" + +from __future__ import annotations + +import shutil + +import pytest +import yaml + +from ..orchestrator import round_strategy as rs +from ..orchestrator.harness_io import default_harness_dir +from ..orchestrator.prompts import w8a8_round_strategy + +import importlib.util + +import pytest +_HARNESS_EVOLVE_AVAILABLE = ( + importlib.util.find_spec( + "metainfer.tasks.harness_evolve.orchestrator.pipeline" + ) + is not None +) +requires_harness_evolve = pytest.mark.skipif( + not _HARNESS_EVOLVE_AVAILABLE, + reason="metainfer.tasks.harness_evolve is not available in this tree", +) + +M16 = {"model": "hy3", "M": 16, "N": 2048, "K": 256} +SMALL = {"model": "hy3", "M": 2, "N": 2048, "K": 256} +PREFILL = {"model": "hy3", "M": 3072, "N": 2048, "K": 256} + + +@pytest.fixture(autouse=True) +def _clear_cache(): + rs.reset_cache() + yield + rs.reset_cache() + + +def _workspace(tmp_path, *, 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"]["systemprompt"]["wired"] = wired + (root / "manifest.yaml").write_text( + yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8") + return root + + +def _strategy_file(root): + return root / rs.RELATIVE_PATH + + +def _edit(root, mutate): + data = yaml.safe_load(_strategy_file(root).read_text(encoding="utf-8")) + mutate(data) + _strategy_file(root).write_text( + yaml.safe_dump(data, sort_keys=False, allow_unicode=True), + encoding="utf-8") + + +# --------------------------------------------------------------------------- # +# Built-in defaults and the seed file agree +# --------------------------------------------------------------------------- # + +def test_seed_matches_builtin_defaults(): + """The seed file mirrors the built-in texts (so nothing drifts silently).""" + seed = yaml.safe_load(_strategy_file(default_harness_dir()).read_text( + encoding="utf-8")) + defaults = rs.builtin_strategy() + for section in ("repair", "phases", "append"): + for key, value in defaults[section].items(): + assert seed[section][key].strip() == value, f"{section}.{key}" + for regime, portfolio in defaults["portfolios"].items(): + for slot, value in portfolio.items(): + assert seed["portfolios"][regime][slot].strip() == value, ( + f"portfolios.{regime}.{slot}") + + +def test_seed_round_trips_through_the_loader(tmp_path, monkeypatch): + """Loading the seed yields exactly the built-in texts.""" + root = _workspace(tmp_path) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + loaded = rs.load_round_strategy() + assert loaded.source.endswith("round_strategy.yaml") + assert loaded.repair == rs.builtin_strategy()["repair"] + assert loaded.phases == rs.builtin_strategy()["phases"] + assert loaded.append == rs.builtin_strategy()["append"] + assert loaded.portfolios == rs.builtin_strategy()["portfolios"] + + +def test_no_harness_root_uses_the_plugin_seed(monkeypatch): + """Without METAINFER_HARNESS_ROOT the plugin's own seed is the workspace. + + Resolution falls back to ``harness_default/`` (not to an in-code sentinel), + so the seed file is what production tasks read — and the drift test above + proves it carries the built-in texts. + """ + monkeypatch.delenv("METAINFER_HARNESS_ROOT", raising=False) + rs.reset_cache() + loaded = rs.load_round_strategy() + assert loaded.source.endswith("harness_default/systemprompt/round_strategy.yaml") + assert loaded.portfolios["m16"]["2"].startswith("Architecture round") + assert loaded.portfolios == rs.builtin_strategy()["portfolios"] + + +# --------------------------------------------------------------------------- # +# Editing the file changes the runtime mandate +# --------------------------------------------------------------------------- # + +def test_editing_the_seed_changes_runtime_mandate(tmp_path, monkeypatch): + root = _workspace(tmp_path) + + def mutate(data): + data["portfolios"]["m16"]["2"] = ( + "EVOLVED: benchmark 3-wave and 4-wave geometries before anything else.") + + _edit(root, mutate) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + + rendered = w8a8_round_strategy(M16, 2, [], {}, 10, None) + assert rendered.startswith("EVOLVED: benchmark 3-wave") + # Untouched branches keep their historical text. + assert w8a8_round_strategy(M16, 1, [], {}, 10, None).startswith( + "Establish a minimal 16x16x32 DUMMA tile") + + +def test_edited_text_is_seen_by_the_pipeline_call_path(tmp_path, monkeypatch): + """w8a8_pipeline's non-planner path renders the harness text too.""" + root = _workspace(tmp_path) + + def mutate(data): + data["repair"]["build_failure"] = "EVOLVED BUILD REPAIR for {iteration}" + + _edit(root, mutate) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + monkeypatch.delenv("METAINFER_PLANNER", raising=False) + rs.reset_cache() + + from ..orchestrator import w8a8_pipeline as W + + history = [{"iteration": 4, "build_success": False, + "failure_reason": "hipcc: undefined template", + "artifact_dir": "/tmp/a4"}] + text = W.w8a8_round_strategy(M16, 4, history, {}, max_iterations=10, + isa_policy={}) + assert text == "EVOLVED BUILD REPAIR for 4" + + +def test_repeated_scale_can_be_edited(tmp_path, monkeypatch): + root = _workspace(tmp_path) + + def mutate(data): + data["phases"]["isa_guided_hip"] = "EVOLVED ISA round {completed} of many" + + _edit(root, mutate) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + text = w8a8_round_strategy( + M16, 3, [], {}, 10, + {"phase": "isa_guided_hip", "valid_isa_guided_rounds": 1}) + assert text == "EVOLVED ISA round 2 of many" + + +def test_append_suffixes_can_be_edited(tmp_path, monkeypatch): + root = _workspace(tmp_path) + + def mutate(data): + data["append"]["split_candidates"] = " SUSPECTS={split_candidates}." + + _edit(root, mutate) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + text = w8a8_round_strategy( + M16, 3, [], {"grid_blocks": 500, "device_cu_count": 60}, 10, None) + assert "SUSPECTS=" in text + assert "Trusted occupancy-probe split candidates" not in text + + +# --------------------------------------------------------------------------- # +# The manifest decides whether the file applies at all +# --------------------------------------------------------------------------- # + +def test_unwired_component_falls_back_to_builtin(tmp_path, monkeypatch): + root = _workspace(tmp_path, wired=False) + + def mutate(data): + data["portfolios"]["m16"]["2"] = "EVOLVED: should never be used" + + _edit(root, mutate) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + loaded = rs.load_round_strategy() + assert loaded.source == "builtin" + assert loaded.portfolios["m16"]["2"].startswith("Architecture round") + assert any("wired:false" in note for note in loaded.notes) + + +def test_missing_component_entry_still_applies(tmp_path, monkeypatch): + """A manifest without the component key must not disable the file.""" + root = _workspace(tmp_path) + manifest = yaml.safe_load((root / "manifest.yaml").read_text(encoding="utf-8")) + manifest["components"].pop("systemprompt") + (root / "manifest.yaml").write_text( + yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8") + + def mutate(data): + data["portfolios"]["m16"]["2"] = "EVOLVED without a manifest entry" + + _edit(root, mutate) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + assert w8a8_round_strategy(M16, 2, [], {}, 10, None).startswith("EVOLVED without") + + +# --------------------------------------------------------------------------- # +# Degradation: a broken harness must not break a round +# --------------------------------------------------------------------------- # + +def test_missing_file_uses_builtin(tmp_path, monkeypatch): + root = _workspace(tmp_path) + _strategy_file(root).unlink() + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + loaded = rs.load_round_strategy() + assert loaded.source == "builtin" + assert loaded.repair == rs.builtin_strategy()["repair"] + + +def test_malformed_yaml_is_reported_not_raised(tmp_path, monkeypatch): + root = _workspace(tmp_path) + _strategy_file(root).write_text("repair: [unclosed\n", encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + loaded = rs.load_round_strategy() + assert loaded.source == "builtin" + assert loaded.notes and "round_strategy.yaml" in loaded.notes[0] + assert w8a8_round_strategy(M16, 2, [], {}, 10, None).startswith("Architecture") + + +def test_partial_file_keeps_unlisted_branches(tmp_path, monkeypatch): + """Omitting a key keeps its built-in text instead of blanking the branch.""" + root = _workspace(tmp_path) + # A file that only carries a reformatted large-M round 3. + _strategy_file(root).write_text( + "portfolios:\n large:\n \"3\": EVOLVED large round three\n", + encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + loaded = rs.load_round_strategy() + assert loaded.portfolios["large"]["3"] == "EVOLVED large round three" + assert loaded.portfolios["large"]["2"] == ( + rs.builtin_strategy()["portfolios"]["large"]["2"]) + assert loaded.portfolios["m16"]["2"].startswith("Architecture round") + + +def test_unknown_fields_degrade_to_builtin_text(tmp_path, monkeypatch): + """A renamed format field must not raise mid-round.""" + root = _workspace(tmp_path) + + def mutate(data): + data["append"]["grid_warning"] = " warning for {cus} CUs" + + _edit(root, mutate) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + text = w8a8_round_strategy( + M16, 3, [], {"grid_blocks": 4, "device_cu_count": 60}, 10, None) + # The suffix falls back to the built-in wording rather than raising. + assert "Trusted control-plane warning" in text + + +def test_unknown_regime_is_ignored(tmp_path, monkeypatch): + root = _workspace(tmp_path) + + def mutate(data): + data["portfolios"]["medium"] = {"2": "nonsense"} + + _edit(root, mutate) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + loaded = rs.load_round_strategy() + assert "medium" not in loaded.portfolios + assert any("unknown regime" in note for note in loaded.notes) + + +def test_empty_string_is_ignored(tmp_path, monkeypatch): + """An empty value keeps the built-in text (a mandate is never blank).""" + root = _workspace(tmp_path) + + def mutate(data): + data["portfolios"]["m16"]["2"] = " " + + _edit(root, mutate) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + loaded = rs.load_round_strategy() + assert loaded.portfolios["m16"]["2"].startswith("Architecture round") + + +# --------------------------------------------------------------------------- # +# Cache invalidation: a mid-run harness swap is picked up +# --------------------------------------------------------------------------- # + +def test_cache_is_invalidated_by_file_change(tmp_path, monkeypatch): + root = _workspace(tmp_path) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + assert w8a8_round_strategy(M16, 2, [], {}, 10, None).startswith("Architecture") + + _edit(root, lambda data: data["portfolios"]["m16"].__setitem__( + "2", "EVOLVED after a mid-run swap")) + import os + stat = _strategy_file(root).stat() + os.utime(_strategy_file(root), (stat.st_atime, stat.st_mtime + 10)) + assert w8a8_round_strategy(M16, 2, [], {}, 10, None) == "EVOLVED after a mid-run swap" + + +def test_cache_is_invalidated_by_manifest_change(tmp_path, monkeypatch): + root = _workspace(tmp_path) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + rs.reset_cache() + assert rs.load_round_strategy().source.endswith("round_strategy.yaml") + + manifest = yaml.safe_load((root / "manifest.yaml").read_text(encoding="utf-8")) + manifest["components"]["systemprompt"]["wired"] = False + (root / "manifest.yaml").write_text( + yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8") + import os + stat = (root / "manifest.yaml").stat() + os.utime(root / "manifest.yaml", (stat.st_atime, stat.st_mtime + 10)) + assert rs.load_round_strategy().source == "builtin" + + +# --------------------------------------------------------------------------- # +# The small-M late-round crash is fixed +# --------------------------------------------------------------------------- # + +@pytest.mark.parametrize("max_iterations", [9, 10, 11, 20]) +@pytest.mark.parametrize("iteration", [8, 9, 10, 11, 12]) +def test_small_m_late_rounds_never_raise(max_iterations, iteration): + """M<16 has no slot 9/10; it used to raise KeyError for max_iterations 9..10.""" + text = w8a8_round_strategy(SMALL, iteration, [], {}, max_iterations, None) + assert text + assert isinstance(text, str) + + +def test_small_m_keeps_its_consolidation_round_late(): + text = w8a8_round_strategy(SMALL, 10, [], {}, 20, None) + assert text.startswith("HIP-only consolidation round") + + +def test_normal_regimes_are_unaffected_by_the_fix(): + # 16..127 and >=128 still reach the late slots. + assert w8a8_round_strategy(M16, 9, [], {}, 10, None).startswith( + "Late ISA-diagnosis round") + assert w8a8_round_strategy(M16, 10, [], {}, 10, None).startswith( + "Final conditional inline-asm round") + assert w8a8_round_strategy(PREFILL, 20, [], {}, 20, None).startswith( + "Final conditional inline-asm round") + + +# --------------------------------------------------------------------------- # +# End-to-end: seeding, and the AHE write-scope check reaches the new component +# --------------------------------------------------------------------------- # + +def test_seed_workspace_carries_the_systemprompt_component(tmp_path): + """HE seeds the experiment workspace from harness_default/, so the file ships.""" + from ..orchestrator.harness_io import seed_workspace + + dst = tmp_path / "seeded" + seed_workspace(dst) + assert (dst / rs.RELATIVE_PATH).is_file() + manifest = yaml.safe_load((dst / "manifest.yaml").read_text(encoding="utf-8")) + entry = manifest["components"]["systemprompt"] + assert entry["path"] == rs.RELATIVE_PATH + assert entry["wired"] is True + + +#: Rendered outputs of a 5400-case matrix (M regime x iteration x history x PMC +#: x ISA phase), frozen after the externalisation refactor. Any edit to a +#: mandate — in the built-in defaults *or* in the seed file — changes this +#: digest, which makes an accidental behaviour change impossible to merge +#: unnoticed. Regenerate deliberately when a wording change is intended. +_RENDERED_DIGEST = "8a3bb533a7779a45e6ad2a7b45128d0cc1b406155f83d04f1fc74fcb15255fea" + +_SHAPES = { + "m16": {"model": "hy3", "M": 16, "N": 2048, "K": 256}, + "m64": {"model": "hy3", "M": 64, "N": 4096, "K": 512}, + "large": {"model": "hy3", "M": 3072, "N": 2048, "K": 256}, + "prefill": {"model": "hy3", "M": 4096, "N": 1024, "K": 256}, + "small": {"model": "hy3", "M": 2, "N": 2048, "K": 256}, + "tiny": {"model": "hy3", "M": 1, "N": 2048, "K": 256}, +} +_HISTORIES = { + "empty": [], + "faster_wrong": [ + {"iteration": 2, "correctness_passed": False, "speedup": 1.4, + "artifact_dir": "/tmp/arch/iter2", "failure_reason": "mismatch_count 12"}, + {"iteration": 4, "correctness_passed": True, "speedup": 0.9, + "artifact_dir": "/tmp/arch/iter4", "failure_reason": ""}, + ], + "infra": [{"iteration": 3, "correctness_passed": True, "build_success": True, + "failure_reason": "agent timeout after 900s", + "artifact_dir": "/tmp/i3"}], + "build": [{"iteration": 3, "correctness_passed": True, "build_success": False, + "failure_reason": "hipcc: undefined template", + "artifact_dir": "/tmp/i3"}], + "assert": [{"iteration": 3, "correctness_passed": True, "build_success": True, + "failure_reason": "mismatch_count 3", "artifact_dir": "/tmp/i3"}], +} +_PMCS = [ + {}, + {"grid_blocks": 4, "device_cu_count": 60, "l2_hit_rate": 0.55}, + {"grid_blocks": 400, "device_cu_count": 60, "l2_hit_rate": 0.9}, +] +_ISA_POLICIES = [ + {}, + {"phase": "hip_only"}, + {"phase": "isa_guided_hip", "valid_isa_guided_rounds": 0}, + {"phase": "isa_guided_hip", "valid_isa_guided_rounds": 3}, + {"phase": "conditional_inline_asm"}, +] + + +def test_rendered_mandates_match_the_frozen_digest(): + """The worker-facing mandates are frozen: 5400 rendered strings, one hash.""" + import hashlib + import json + + payload = {} + for regime, shape in _SHAPES.items(): + for iteration in range(1, 13): + for hlabel, history in _HISTORIES.items(): + for pidx, pmc in enumerate(_PMCS): + for iidx, isa in enumerate(_ISA_POLICIES): + payload[ + f"{regime}-it{iteration:02d}-{hlabel}" + f"-pmc{pidx}-isa{iidx}" + ] = w8a8_round_strategy(shape, iteration, history, pmc, 10, isa) + blob = json.dumps(payload, ensure_ascii=False, sort_keys=True).encode("utf-8") + assert len(payload) == 5400 + assert hashlib.sha256(blob).hexdigest() == _RENDERED_DIGEST, ( + "a round mandate changed; if that was intended, regenerate the digest " + "intentionally and describe the new wording in the change manifest") + + +@requires_harness_evolve +def test_ahe_wired_scope_check_covers_the_systemprompt_component(tmp_path): + """An edit to the component while it is wired:false is reported by AHE. + + This is what makes the wider write scope auditable: harness_evolve compares + each component's declared path, so a new component is covered as soon as it + is registered in the manifest. + """ + from metainfer.tasks.harness_evolve.orchestrator.pipeline import ( + _wired_scope_violations, + ) + + before = tmp_path / "before" + after = tmp_path / "after" + shutil.copytree(default_harness_dir(), before) + shutil.copytree(default_harness_dir(), after) + manifest = yaml.safe_load((after / "manifest.yaml").read_text(encoding="utf-8")) + manifest["components"]["systemprompt"]["wired"] = False + (after / "manifest.yaml").write_text( + yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8") + # Unchanged file: not a violation. + assert _wired_scope_violations(before, after) == [] + # Edited file while declared unwired: a violation, with the component named. + text = (after / rs.RELATIVE_PATH).read_text(encoding="utf-8") + (after / rs.RELATIVE_PATH).write_text( + text + "\n# evolved by hand\n", encoding="utf-8") + violations = _wired_scope_violations(before, after) + assert violations == [{"component": "systemprompt", + "path": rs.RELATIVE_PATH}] diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_shape_catalog_panel.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_shape_catalog_panel.py new file mode 100644 index 00000000..10be54fb --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_shape_catalog_panel.py @@ -0,0 +1,126 @@ +"""Panel <-> contract <-> baseline integration for the model shape catalog. + +``static/dkao-shape-input.js`` is what the DKAO "New Task" panel offers the +operator, so it is the source of truth for which shapes a task can submit. +Every shape it can submit must + +* satisfy the frozen operator contract (``validate_optimization_shape``), and +* have a fixed Triton Graph baseline (``fixed_triton_graph_baseline``), + +otherwise the run dies in the ``baseline`` phase with a ValueError long after +the task was created. These tests read the *real* frontend catalog (by +evaluating its self-contained leading section under node, so the assertions +cannot drift from the shipped UI) and check that the three layers agree. +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +from pathlib import Path + +import pytest + +from ..api.int8w8a8gemm.int8_w8a8_gemm_api import ( + MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES, + validate_optimization_shape, +) +from ..orchestrator.w8a8_baselines import fixed_triton_graph_baseline + + +STATIC_JS = Path(__file__).resolve().parent.parent / "static" / "dkao-shape-input.js" +MODEL_CATALOG_MARKERS = ( + "var W8A8_M_VALUES", + 'var DEFAULT_W8A8_SHAPES = modelCatalog("DeepSeek V4 Flash");', +) +# Models whose TP8 catalog carries the M=8 short-decode boundary. +M8_MODELS = ("Hy3 (Hunyuan 3)", "MiniMax M3", "GLM5.2") + + +def _frontend_catalogs() -> dict[str, list[dict[str, object]]]: + """Return {model label: [shape, ...]} from the shipped panel catalog.""" + node = shutil.which("node") + if node is None: + pytest.skip("node is required to evaluate the panel shape catalog") + text = STATIC_JS.read_text(encoding="utf-8") + start, end = MODEL_CATALOG_MARKERS + head = text.index(start) + tail = text.index(end) + len(end) + harness = ( + text[head:tail] + + "\nvar out = {};\n" + + "Object.keys(MODEL_WORKLOADS).forEach(function (label) {\n" + + " out[label] = modelCatalog(label);\n" + + "});\n" + + "console.log(JSON.stringify(out));\n" + ) + proc = subprocess.run( + [node, "-e", harness], + capture_output=True, + text=True, + check=False, + ) + assert proc.returncode == 0, proc.stderr + return json.loads(proc.stdout) + + +@pytest.fixture(scope="module") +def catalogs() -> dict[str, list[dict[str, object]]]: + return _frontend_catalogs() + + +def test_panel_catalog_exposes_m8_for_model_tp8(catalogs): + # The M=8 short-decode boundary added on 2026-09-20 must be offered for + # Hy3 / MiniMax M3 / GLM5.2 TP8 operators. + for model in M8_MODELS: + tp8 = [s for s in catalogs[model] if s["tp_size"] == 8] + assert tp8, model + assert all(s["M"] in (8, 2, 16, 3072, 4096) for s in tp8), model + assert any(s["M"] == 8 for s in tp8), model + assert sorted({int(s["M"]) for s in tp8}) == [2, 8, 16, 3072, 4096] + + +def test_panel_catalog_keeps_deepseek_tp8_defaults(catalogs): + # DeepSeek TP8 intentionally keeps the three original M values, so the + # DeepSeek-only default workload and its validation scope are unchanged. + tp8 = [s for s in catalogs["DeepSeek V4 Flash"] if s["tp_size"] == 8] + assert sorted({int(s["M"]) for s in tp8}) == [2, 16, 3072] + + +def test_panel_m8_shapes_match_the_extra_m_value_contract(catalogs): + assert 8 in MODEL_TP8_EXTRA_OPTIMIZATION_M_VALUES + for model in M8_MODELS: + m8 = [s for s in catalogs[model] if s["tp_size"] == 8 and s["M"] == 8] + # One M=8 shape per catalog operator. + assert len(m8) == len( + {s["operator"] for s in catalogs[model] if s["tp_size"] == 8} + ) + for shape in m8: + validate_optimization_shape(shape) + + +def test_every_panel_tp4_tp8_shape_is_submittable(catalogs): + # A shape the panel offers must both pass the frozen contract and resolve a + # fixed baseline; otherwise the task fails in the baseline phase. + checked = 0 + for model, shapes in catalogs.items(): + for shape in shapes: + if shape["tp_size"] not in (4, 8): + continue + validate_optimization_shape(shape) + record = fixed_triton_graph_baseline(str(shape["id"]), shape) + assert record["median_us"] > 0, (model, shape) + checked += 1 + assert checked >= 70 + + +def test_panel_m8_shapes_resolve_decode_baselines(catalogs): + for model in M8_MODELS: + for shape in catalogs[model]: + if shape["tp_size"] != 8 or shape["M"] != 8: + continue + record = fixed_triton_graph_baseline(str(shape["id"]), shape) + assert record["timing_scope"] == "decode_graph_replay" + assert record["baseline_kind"] == "triton_graph" + assert record["case"], shape 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/tests/test_w8a8_baselines.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_baselines.py index 0eb9d65a..b3a56910 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_baselines.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_w8a8_baselines.py @@ -86,6 +86,37 @@ def test_model_catalog_tp8_m4096_baselines(shape, expected): assert record["timing_scope"] == "prefill_graph_replay" +@pytest.mark.parametrize( + ("shape", "expected"), + [ + # Hy3 TP8 M=8 (measured 2026-09-20, CUDA-graph replay). + ({"tp_size": 8, "M": 8, "N": 1280, "K": 4096}, 66.432), + ({"tp_size": 8, "M": 8, "N": 4096, "K": 1024}, 31.472), + ({"tp_size": 8, "M": 8, "N": 384, "K": 4096}, 57.008), + ({"tp_size": 8, "M": 8, "N": 4096, "K": 192}, 14.592), + # MiniMax M3 TP8 M=8. + ({"tp_size": 8, "M": 8, "N": 1280, "K": 6144}, 98.865), + ({"tp_size": 8, "M": 8, "N": 1536, "K": 6144}, 99.473), + ({"tp_size": 8, "M": 8, "N": 6144, "K": 1024}, 32.960), + ({"tp_size": 8, "M": 8, "N": 768, "K": 6144}, 92.928), + ({"tp_size": 8, "M": 8, "N": 6144, "K": 384}, 19.984), + # GLM5.2 TP8 M=8. + ({"tp_size": 8, "M": 8, "N": 2624, "K": 6144}, 87.200), + ({"tp_size": 8, "M": 8, "N": 2048, "K": 2048}, 36.704), + ({"tp_size": 8, "M": 8, "N": 3584, "K": 512}, 15.808), + ({"tp_size": 8, "M": 8, "N": 6144, "K": 2048}, 58.864), + ({"tp_size": 8, "M": 8, "N": 512, "K": 6144}, 88.448), + ({"tp_size": 8, "M": 8, "N": 6144, "K": 256}, 15.136), + ], +) +def test_model_catalog_tp8_m8_baselines(shape, expected): + # M=8 is the model-catalog short-decode boundary added on 2026-09-20 and + # resolves as a decode baseline (M <= 16). + record = fixed_triton_graph_baseline("shape", shape) + assert record["median_us"] == expected + assert record["timing_scope"] == "decode_graph_replay" + + def test_bootstrap_metrics_are_kept_separate(): bootstrap = {"passed": True, "median_us": 123.0} record = fixed_triton_graph_baseline( diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/__init__.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/__init__.py new file mode 100644 index 00000000..e69de29b 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/bench_triton_tp8_m8.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m8.py new file mode 100644 index 00000000..57de33b6 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m8.py @@ -0,0 +1,238 @@ +#!/usr/bin/env python3 +"""Fair Triton W8A8 GEMM baselines for the model-catalog TP8 shapes at M=8. + +Same protocol as bench_triton_tp8_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 <= 32: + 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 + +``--m`` defaults to 8 (the new decode boundary). Passing ``--m 16`` re-measures +the already-published M=16 catalog entries as a protocol cross-check. +""" + +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") + +# (shape_id stem, operator, K, N) -- Hy3 / MiniMax M3 / GLM5.2 TP8 catalog. +TP8_CASES = [ + ("hy3_tp8_qkv_proj", "qkv_proj", 4096, 1280), + ("hy3_tp8_o_proj", "o_proj", 1024, 4096), + ("hy3_tp8_shared_gate_up_proj", "shared_gate_up_proj", 4096, 384), + ("hy3_tp8_shared_down_proj", "shared_down_proj", 192, 4096), + ("minimax_tp8_qkv_proj", "qkv_proj", 6144, 1280), + ("minimax_tp8_qkv_proj_and_indexer_qk", "qkv_proj_and_indexer_qk", 6144, 1536), + ("minimax_tp8_o_proj", "o_proj", 1024, 6144), + ("minimax_tp8_shared_gate_up_proj", "shared_gate_up_proj", 6144, 768), + ("minimax_tp8_shared_down_proj", "shared_down_proj", 384, 6144), + ("glm52_tp8_fused_qkv_a_proj", "fused_qkv_a_proj", 6144, 2624), + ("glm52_tp8_q_b_proj", "q_b_proj", 2048, 2048), + ("glm52_tp8_kv_b_proj", "kv_b_proj", 512, 3584), + ("glm52_tp8_o_proj", "o_proj", 2048, 6144), + ("glm52_tp8_shared_gate_up_proj", "shared_gate_up_proj", 6144, 512), + ("glm52_tp8_shared_down_proj", "shared_down_proj", 256, 6144), +] + +# Exact default config used by matmul_int8 when 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, +} + + +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_stem, operator, k, n, m, 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) + + def fn(): + 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": f"{case_stem}_m{m}", + "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("--m", type=int, default=8) + 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="") + parser.add_argument("--only", type=str, default="", + help="comma-separated case-stem substrings to keep") + args = parser.parse_args() + m = args.m + + utils = load_int8_utils(args.utils) + cases = TP8_CASES + if args.only: + needles = [s.strip() for s in args.only.split(",") if s.strip()] + cases = [c for c in cases if any(x in c[0] for x in needles)] + + 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':44s} {'N':5s} {'K':5s} {'graph_us':>10s} " + f"{'p90_us':>10s} {'eager_us':>10s}", + flush=True, + ) + results = [] + for case_stem, operator, k, n in cases: + res = run_case( + utils, + case_stem, + operator, + k, + n, + m, + args.warmups, + args.samples, + args.launches_per_sample, + args.mode, + ) + results.append(res) + graph_us = (res.get("graph") or {}).get("median_us") + p90_us = (res.get("graph") or {}).get("p90_us") + eager_us = (res.get("eager") or {}).get("median_us") + print( + f"{res['shape_id']:44s} {n:5d} {k:5d} " + f"{(f'{graph_us:.3f}' if graph_us else '—'):>10s} " + f"{(f'{p90_us:.3f}' if p90_us else '—'):>10s} " + f"{(f'{eager_us:.3f}' if eager_us else '—'):>10s}", + 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/tp8_m8_graph_merged.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_merged.json new file mode 100644 index 00000000..e5104ff8 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_merged.json @@ -0,0 +1,197 @@ +[ + { + "shape_id": "hy3_tp8_qkv_proj_m8", + "M": 8, + "N": 1280, + "K": 4096, + "operator": "qkv_proj", + "pass_medians_us": [ + 66.432, + 66.272, + 67.216 + ], + "published_median_us": 66.432 + }, + { + "shape_id": "hy3_tp8_o_proj_m8", + "M": 8, + "N": 4096, + "K": 1024, + "operator": "o_proj", + "pass_medians_us": [ + 31.472, + 31.488, + 31.424 + ], + "published_median_us": 31.472 + }, + { + "shape_id": "hy3_tp8_shared_gate_up_proj_m8", + "M": 8, + "N": 384, + "K": 4096, + "operator": "shared_gate_up_proj", + "pass_medians_us": [ + 56.896, + 57.008, + 57.792 + ], + "published_median_us": 57.008 + }, + { + "shape_id": "hy3_tp8_shared_down_proj_m8", + "M": 8, + "N": 4096, + "K": 192, + "operator": "shared_down_proj", + "pass_medians_us": [ + 14.592, + 14.592, + 14.592 + ], + "published_median_us": 14.592 + }, + { + "shape_id": "minimax_tp8_qkv_proj_m8", + "M": 8, + "N": 1280, + "K": 6144, + "operator": "qkv_proj", + "pass_medians_us": [ + 98.401, + 98.865, + 99.36 + ], + "published_median_us": 98.865 + }, + { + "shape_id": "minimax_tp8_qkv_proj_and_indexer_qk_m8", + "M": 8, + "N": 1536, + "K": 6144, + "operator": "qkv_proj_and_indexer_qk", + "pass_medians_us": [ + 99.456, + 99.473, + 99.968 + ], + "published_median_us": 99.473 + }, + { + "shape_id": "minimax_tp8_o_proj_m8", + "M": 8, + "N": 6144, + "K": 1024, + "operator": "o_proj", + "pass_medians_us": [ + 32.96, + 33.008, + 32.768 + ], + "published_median_us": 32.96 + }, + { + "shape_id": "minimax_tp8_shared_gate_up_proj_m8", + "M": 8, + "N": 768, + "K": 6144, + "operator": "shared_gate_up_proj", + "pass_medians_us": [ + 92.544, + 93.936, + 92.928 + ], + "published_median_us": 92.928 + }, + { + "shape_id": "minimax_tp8_shared_down_proj_m8", + "M": 8, + "N": 6144, + "K": 384, + "operator": "shared_down_proj", + "pass_medians_us": [ + 19.968, + 33.76, + 19.984 + ], + "published_median_us": 19.984 + }, + { + "shape_id": "glm52_tp8_fused_qkv_a_proj_m8", + "M": 8, + "N": 2624, + "K": 6144, + "operator": "fused_qkv_a_proj", + "pass_medians_us": [ + 87.2, + 88.032, + 86.976 + ], + "published_median_us": 87.2 + }, + { + "shape_id": "glm52_tp8_q_b_proj_m8", + "M": 8, + "N": 2048, + "K": 2048, + "operator": "q_b_proj", + "pass_medians_us": [ + 37.92, + 36.704, + 35.792 + ], + "published_median_us": 36.704 + }, + { + "shape_id": "glm52_tp8_kv_b_proj_m8", + "M": 8, + "N": 3584, + "K": 512, + "operator": "kv_b_proj", + "pass_medians_us": [ + 16.0, + 15.808, + 15.744 + ], + "published_median_us": 15.808 + }, + { + "shape_id": "glm52_tp8_o_proj_m8", + "M": 8, + "N": 6144, + "K": 2048, + "operator": "o_proj", + "pass_medians_us": [ + 58.864, + 58.864, + 58.064 + ], + "published_median_us": 58.864 + }, + { + "shape_id": "glm52_tp8_shared_gate_up_proj_m8", + "M": 8, + "N": 512, + "K": 6144, + "operator": "shared_gate_up_proj", + "pass_medians_us": [ + 89.216, + 88.336, + 88.448 + ], + "published_median_us": 88.448 + }, + { + "shape_id": "glm52_tp8_shared_down_proj_m8", + "M": 8, + "N": 6144, + "K": 256, + "operator": "shared_down_proj", + "pass_medians_us": [ + 16.096, + 15.136, + 15.072 + ], + "published_median_us": 15.136 + } +] \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run1.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run1.json new file mode 100644 index 00000000..a5ec1aa7 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run1.json @@ -0,0 +1,662 @@ +[ + { + "shape_id": "hy3_tp8_qkv_proj_m8", + "operator": "qkv_proj", + "M": 8, + "N": 1280, + "K": 4096, + "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 + }, + "graph": { + "median_us": 66.4322018623352, + "p90_us": 67.2324001789093, + "min_us": 66.20839834213257, + "max_us": 67.4884021282196, + "samples_us": [ + 67.488, + 67.232, + 67.2, + 66.56, + 66.432, + 66.368, + 66.432, + 66.304, + 66.368, + 67.456, + 66.304, + 66.208, + 66.464, + 66.272, + 66.272, + 66.432, + 66.336, + 66.368, + 66.464, + 67.232 + ] + } + }, + { + "shape_id": "hy3_tp8_o_proj_m8", + "operator": "o_proj", + "M": 8, + "N": 4096, + "K": 1024, + "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 + }, + "graph": { + "median_us": 31.47220015525818, + "p90_us": 31.584200263023376, + "min_us": 31.296199560165405, + "max_us": 32.41620063781738, + "samples_us": [ + 32.416, + 31.36, + 31.52, + 31.552, + 31.456, + 31.584, + 31.616, + 31.392, + 31.456, + 31.52, + 31.392, + 31.36, + 31.552, + 31.296, + 31.488, + 31.584, + 31.424, + 31.584, + 31.392, + 31.328 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_gate_up_proj_m8", + "operator": "shared_gate_up_proj", + "M": 8, + "N": 384, + "K": 4096, + "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 + }, + "graph": { + "median_us": 56.89629912376404, + "p90_us": 57.69619941711426, + "min_us": 56.64039850234985, + "max_us": 57.95239806175232, + "samples_us": [ + 56.672, + 57.696, + 56.864, + 56.96, + 56.8, + 57.184, + 56.704, + 56.928, + 57.728, + 57.056, + 56.8, + 56.992, + 56.704, + 57.12, + 56.64, + 56.96, + 56.704, + 56.864, + 56.768, + 57.952 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_down_proj_m8", + "operator": "shared_down_proj", + "M": 8, + "N": 4096, + "K": 192, + "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 + }, + "graph": { + "median_us": 14.592200517654419, + "p90_us": 14.688000082969666, + "min_us": 14.528200030326843, + "max_us": 17.952199280261993, + "samples_us": [ + 14.656, + 14.528, + 14.624, + 14.592, + 14.592, + 14.56, + 14.624, + 14.592, + 14.624, + 14.688, + 14.528, + 14.592, + 14.656, + 14.528, + 14.56, + 14.592, + 14.592, + 14.688, + 17.952, + 14.624 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_m8", + "operator": "qkv_proj", + "M": 8, + "N": 1280, + "K": 6144, + "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 + }, + "graph": { + "median_us": 98.40050041675568, + "p90_us": 98.68860244750977, + "min_us": 98.27240109443665, + "max_us": 99.52059984207153, + "samples_us": [ + 98.497, + 98.336, + 99.361, + 98.337, + 98.337, + 98.401, + 98.401, + 98.337, + 98.465, + 98.272, + 98.432, + 98.689, + 98.273, + 98.273, + 98.561, + 98.529, + 98.305, + 99.521, + 98.336, + 98.4 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_and_indexer_qk_m8", + "operator": "qkv_proj_and_indexer_qk", + "M": 8, + "N": 1536, + "K": 6144, + "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 + }, + "graph": { + "median_us": 99.45639967918396, + "p90_us": 100.32060146331787, + "min_us": 99.2326021194458, + "max_us": 100.41660070419312, + "samples_us": [ + 100.321, + 100.353, + 99.233, + 99.265, + 99.553, + 99.233, + 99.456, + 99.328, + 99.617, + 99.233, + 99.361, + 99.361, + 99.361, + 100.224, + 99.456, + 100.417, + 99.553, + 99.457, + 99.457, + 99.297 + ] + } + }, + { + "shape_id": "minimax_tp8_o_proj_m8", + "operator": "o_proj", + "M": 8, + "N": 6144, + "K": 1024, + "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 + }, + "graph": { + "median_us": 32.96020030975342, + "p90_us": 33.024200797080994, + "min_us": 32.864201068878174, + "max_us": 33.05619955062866, + "samples_us": [ + 32.96, + 32.992, + 32.96, + 33.056, + 32.96, + 33.056, + 32.96, + 32.864, + 32.96, + 32.992, + 32.864, + 32.928, + 33.024, + 32.992, + 32.96, + 33.024, + 32.992, + 32.96, + 32.96, + 32.928 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_gate_up_proj_m8", + "operator": "shared_gate_up_proj", + "M": 8, + "N": 768, + "K": 6144, + "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 + }, + "graph": { + "median_us": 92.54440069198608, + "p90_us": 93.34459900856018, + "min_us": 92.06439852714539, + "max_us": 93.4723973274231, + "samples_us": [ + 93.345, + 92.545, + 92.129, + 92.577, + 92.385, + 92.481, + 92.193, + 92.417, + 92.449, + 93.377, + 92.288, + 92.544, + 93.152, + 92.576, + 92.064, + 93.472, + 92.256, + 92.544, + 92.896, + 92.609 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_down_proj_m8", + "operator": "shared_down_proj", + "M": 8, + "N": 6144, + "K": 384, + "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 + }, + "graph": { + "median_us": 19.968199729919434, + "p90_us": 20.03200054168701, + "min_us": 19.904200732707977, + "max_us": 20.064200460910797, + "samples_us": [ + 19.904, + 20.032, + 20.0, + 19.968, + 19.968, + 19.968, + 19.968, + 19.968, + 20.0, + 20.0, + 20.064, + 19.968, + 19.968, + 19.904, + 19.968, + 19.936, + 20.0, + 19.968, + 20.064, + 20.0 + ] + } + }, + { + "shape_id": "glm52_tp8_fused_qkv_a_proj_m8", + "operator": "fused_qkv_a_proj", + "M": 8, + "N": 2624, + "K": 6144, + "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 + }, + "graph": { + "median_us": 87.20049858093262, + "p90_us": 89.47259783744812, + "min_us": 86.17640137672424, + "max_us": 97.60059714317322, + "samples_us": [ + 87.329, + 87.265, + 86.208, + 86.24, + 94.241, + 87.136, + 86.177, + 86.241, + 86.336, + 86.176, + 86.369, + 86.177, + 87.136, + 97.601, + 89.473, + 89.408, + 89.152, + 89.248, + 89.152, + 89.313 + ] + } + }, + { + "shape_id": "glm52_tp8_q_b_proj_m8", + "operator": "q_b_proj", + "M": 8, + "N": 2048, + "K": 2048, + "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 + }, + "graph": { + "median_us": 37.92020082473755, + "p90_us": 38.24020028114319, + "min_us": 36.80019974708557, + "max_us": 38.624200224876404, + "samples_us": [ + 38.624, + 38.336, + 37.888, + 37.824, + 38.048, + 37.92, + 37.856, + 37.856, + 37.92, + 37.952, + 37.792, + 37.984, + 37.888, + 38.08, + 38.24, + 37.92, + 38.112, + 38.08, + 37.92, + 36.8 + ] + } + }, + { + "shape_id": "glm52_tp8_kv_b_proj_m8", + "operator": "kv_b_proj", + "M": 8, + "N": 3584, + "K": 512, + "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 + }, + "graph": { + "median_us": 16.00010022521019, + "p90_us": 24.86419975757599, + "min_us": 15.71200042963028, + "max_us": 26.464200019836426, + "samples_us": [ + 26.464, + 23.808, + 24.448, + 23.84, + 24.864, + 24.896, + 24.832, + 17.024, + 15.776, + 15.744, + 15.84, + 15.872, + 16.032, + 15.872, + 15.808, + 15.872, + 15.968, + 15.936, + 15.712, + 16.736 + ] + } + }, + { + "shape_id": "glm52_tp8_o_proj_m8", + "operator": "o_proj", + "M": 8, + "N": 6144, + "K": 2048, + "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 + }, + "graph": { + "median_us": 58.864399790763855, + "p90_us": 59.29639935493469, + "min_us": 57.79240131378174, + "max_us": 59.52039957046509, + "samples_us": [ + 59.136, + 59.104, + 58.88, + 59.2, + 59.52, + 58.432, + 58.848, + 58.432, + 59.136, + 58.848, + 59.296, + 57.792, + 59.04, + 57.984, + 58.848, + 59.456, + 58.016, + 58.304, + 58.912, + 57.92 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_gate_up_proj_m8", + "operator": "shared_gate_up_proj", + "M": 8, + "N": 512, + "K": 6144, + "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 + }, + "graph": { + "median_us": 89.21639919281006, + "p90_us": 89.60059881210327, + "min_us": 88.22460174560547, + "max_us": 89.60059881210327, + "samples_us": [ + 89.249, + 89.601, + 88.225, + 89.28, + 88.512, + 89.152, + 89.345, + 88.609, + 88.225, + 89.536, + 88.576, + 89.12, + 89.313, + 89.601, + 89.184, + 89.472, + 88.608, + 89.249, + 89.057, + 89.601 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_down_proj_m8", + "operator": "shared_down_proj", + "M": 8, + "N": 6144, + "K": 256, + "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 + }, + "graph": { + "median_us": 16.096200048923492, + "p90_us": 16.223999857902527, + "min_us": 15.072199702262878, + "max_us": 18.592000007629395, + "samples_us": [ + 16.096, + 16.096, + 16.096, + 15.104, + 16.096, + 16.192, + 16.096, + 16.064, + 16.128, + 15.072, + 16.224, + 15.2, + 16.192, + 18.592, + 16.128, + 16.064, + 15.136, + 16.224, + 16.192, + 16.224 + ] + } + } +] \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run2.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run2.json new file mode 100644 index 00000000..2970f303 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run2.json @@ -0,0 +1,662 @@ +[ + { + "shape_id": "hy3_tp8_qkv_proj_m8", + "operator": "qkv_proj", + "M": 8, + "N": 1280, + "K": 4096, + "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 + }, + "graph": { + "median_us": 66.27230048179626, + "p90_us": 67.26440191268921, + "min_us": 66.11220240592957, + "max_us": 74.72040057182312, + "samples_us": [ + 67.264, + 66.24, + 66.272, + 66.304, + 66.24, + 74.72, + 66.336, + 67.328, + 66.112, + 66.368, + 66.208, + 66.24, + 66.208, + 66.112, + 66.304, + 67.168, + 66.112, + 66.304, + 66.272, + 66.24 + ] + } + }, + { + "shape_id": "hy3_tp8_o_proj_m8", + "operator": "o_proj", + "M": 8, + "N": 4096, + "K": 1024, + "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 + }, + "graph": { + "median_us": 31.488101184368134, + "p90_us": 31.58400058746338, + "min_us": 31.391999125480652, + "max_us": 32.41620063781738, + "samples_us": [ + 32.416, + 31.456, + 31.424, + 31.488, + 31.584, + 31.392, + 31.552, + 31.424, + 31.488, + 31.424, + 31.616, + 31.52, + 31.552, + 31.424, + 31.488, + 31.52, + 31.552, + 31.456, + 31.456, + 31.456 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_gate_up_proj_m8", + "operator": "shared_gate_up_proj", + "M": 8, + "N": 384, + "K": 4096, + "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 + }, + "graph": { + "median_us": 57.00819790363312, + "p90_us": 57.856398820877075, + "min_us": 56.83220028877258, + "max_us": 57.98439979553223, + "samples_us": [ + 57.984, + 57.856, + 56.832, + 57.152, + 57.024, + 57.92, + 56.992, + 56.896, + 56.928, + 56.992, + 56.832, + 57.088, + 56.96, + 57.088, + 56.928, + 57.12, + 56.96, + 57.056, + 56.896, + 57.088 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_down_proj_m8", + "operator": "shared_down_proj", + "M": 8, + "N": 4096, + "K": 192, + "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 + }, + "graph": { + "median_us": 14.591999351978302, + "p90_us": 14.655999839305878, + "min_us": 14.496199786663055, + "max_us": 15.96820056438446, + "samples_us": [ + 14.624, + 14.56, + 14.592, + 14.592, + 14.56, + 14.592, + 14.56, + 15.968, + 14.56, + 14.496, + 14.592, + 14.56, + 14.592, + 14.592, + 14.56, + 14.656, + 14.592, + 14.592, + 14.592, + 14.656 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_m8", + "operator": "qkv_proj", + "M": 8, + "N": 1280, + "K": 6144, + "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 + }, + "graph": { + "median_us": 98.8646000623703, + "p90_us": 99.64860081672668, + "min_us": 98.56060147285461, + "max_us": 100.03260374069214, + "samples_us": [ + 99.617, + 99.745, + 98.561, + 99.649, + 98.913, + 98.656, + 98.624, + 98.976, + 98.72, + 98.593, + 100.033, + 98.625, + 98.721, + 98.817, + 99.52, + 98.72, + 98.752, + 99.552, + 99.616, + 99.041 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_and_indexer_qk_m8", + "operator": "qkv_proj_and_indexer_qk", + "M": 8, + "N": 1536, + "K": 6144, + "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 + }, + "graph": { + "median_us": 99.47250187397003, + "p90_us": 100.35239458084106, + "min_us": 99.29659962654114, + "max_us": 100.41639804840088, + "samples_us": [ + 100.352, + 100.416, + 99.393, + 99.521, + 100.321, + 99.393, + 99.457, + 99.297, + 100.416, + 99.328, + 99.648, + 99.328, + 99.488, + 99.361, + 99.553, + 99.489, + 99.457, + 99.329, + 99.361, + 99.52 + ] + } + }, + { + "shape_id": "minimax_tp8_o_proj_m8", + "operator": "o_proj", + "M": 8, + "N": 6144, + "K": 1024, + "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 + }, + "graph": { + "median_us": 33.00819993019104, + "p90_us": 33.05619955062866, + "min_us": 32.92819857597351, + "max_us": 34.04819965362549, + "samples_us": [ + 34.048, + 32.992, + 32.928, + 33.024, + 33.024, + 33.056, + 32.928, + 33.056, + 32.992, + 32.992, + 32.992, + 32.96, + 32.992, + 32.992, + 33.056, + 32.992, + 33.024, + 33.088, + 33.024, + 33.024 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_gate_up_proj_m8", + "operator": "shared_gate_up_proj", + "M": 8, + "N": 768, + "K": 6144, + "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 + }, + "graph": { + "median_us": 93.93649995326996, + "p90_us": 95.87259888648987, + "min_us": 92.83260107040405, + "max_us": 99.07259941101074, + "samples_us": [ + 93.792, + 93.952, + 93.024, + 92.833, + 93.921, + 93.888, + 93.024, + 92.833, + 93.025, + 93.024, + 93.024, + 99.073, + 95.873, + 95.585, + 96.001, + 95.873, + 95.648, + 95.776, + 95.84, + 95.744 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_down_proj_m8", + "operator": "shared_down_proj", + "M": 8, + "N": 6144, + "K": 384, + "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 + }, + "graph": { + "median_us": 33.76020044088364, + "p90_us": 37.216201424598694, + "min_us": 21.312199532985687, + "max_us": 39.61620032787323, + "samples_us": [ + 39.04, + 34.944, + 33.824, + 34.112, + 33.792, + 33.6, + 37.216, + 34.016, + 32.992, + 33.248, + 33.088, + 33.728, + 39.616, + 32.96, + 35.424, + 33.216, + 33.792, + 22.08, + 21.6, + 21.312 + ] + } + }, + { + "shape_id": "glm52_tp8_fused_qkv_a_proj_m8", + "operator": "fused_qkv_a_proj", + "M": 8, + "N": 2624, + "K": 6144, + "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 + }, + "graph": { + "median_us": 88.03249895572662, + "p90_us": 88.64039778709412, + "min_us": 87.07240223884583, + "max_us": 89.47259783744812, + "samples_us": [ + 89.473, + 87.841, + 88.32, + 88.064, + 88.096, + 88.64, + 87.52, + 87.84, + 88.225, + 88.832, + 87.36, + 87.808, + 88.032, + 88.256, + 87.232, + 88.384, + 88.033, + 88.0, + 87.648, + 87.072 + ] + } + }, + { + "shape_id": "glm52_tp8_q_b_proj_m8", + "operator": "q_b_proj", + "M": 8, + "N": 2048, + "K": 2048, + "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 + }, + "graph": { + "median_us": 36.70420050621033, + "p90_us": 37.88819909095764, + "min_us": 35.84020137786865, + "max_us": 37.95219957828522, + "samples_us": [ + 37.568, + 37.888, + 35.904, + 36.576, + 37.952, + 35.84, + 36.608, + 37.888, + 35.968, + 36.576, + 37.952, + 35.968, + 36.768, + 37.088, + 35.904, + 36.64, + 37.056, + 36.032, + 37.344, + 37.056 + ] + } + }, + { + "shape_id": "glm52_tp8_kv_b_proj_m8", + "operator": "kv_b_proj", + "M": 8, + "N": 3584, + "K": 512, + "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 + }, + "graph": { + "median_us": 15.808099508285522, + "p90_us": 17.728200554847717, + "min_us": 15.71200042963028, + "max_us": 17.85600036382675, + "samples_us": [ + 16.48, + 17.728, + 16.96, + 15.808, + 15.712, + 16.448, + 17.728, + 17.856, + 15.776, + 15.744, + 16.576, + 17.728, + 16.64, + 15.744, + 15.744, + 15.744, + 15.712, + 15.712, + 15.808, + 15.808 + ] + } + }, + { + "shape_id": "glm52_tp8_o_proj_m8", + "operator": "o_proj", + "M": 8, + "N": 6144, + "K": 2048, + "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 + }, + "graph": { + "median_us": 58.86429846286774, + "p90_us": 59.456199407577515, + "min_us": 57.888197898864746, + "max_us": 64.06440138816833, + "samples_us": [ + 59.232, + 58.592, + 58.752, + 58.592, + 58.464, + 58.464, + 58.432, + 57.888, + 58.144, + 58.144, + 58.016, + 64.064, + 59.456, + 59.2, + 59.392, + 59.264, + 58.976, + 59.136, + 59.456, + 59.04 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_gate_up_proj_m8", + "operator": "shared_gate_up_proj", + "M": 8, + "N": 512, + "K": 6144, + "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 + }, + "graph": { + "median_us": 88.33639919757843, + "p90_us": 303.5536050796509, + "min_us": 87.90439963340759, + "max_us": 472.0344066619873, + "samples_us": [ + 89.088, + 88.256, + 88.448, + 88.16, + 88.128, + 88.48, + 88.129, + 88.128, + 88.352, + 88.128, + 88.128, + 88.416, + 87.904, + 88.289, + 88.32, + 162.305, + 303.554, + 450.85, + 295.938, + 472.034 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_down_proj_m8", + "operator": "shared_down_proj", + "M": 8, + "N": 6144, + "K": 256, + "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 + }, + "graph": { + "median_us": 15.136100351810455, + "p90_us": 15.168200433254242, + "min_us": 15.07200002670288, + "max_us": 16.128000617027283, + "samples_us": [ + 16.0, + 15.136, + 15.104, + 15.136, + 15.072, + 15.104, + 15.168, + 15.136, + 15.136, + 15.104, + 15.136, + 15.136, + 15.136, + 15.104, + 16.128, + 15.168, + 15.168, + 15.168, + 15.104, + 15.168 + ] + } + } +] \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run3.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run3.json new file mode 100644 index 00000000..c7586f86 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run3.json @@ -0,0 +1,662 @@ +[ + { + "shape_id": "hy3_tp8_qkv_proj_m8", + "operator": "qkv_proj", + "M": 8, + "N": 1280, + "K": 4096, + "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 + }, + "graph": { + "median_us": 67.21619963645935, + "p90_us": 67.36019849777222, + "min_us": 66.27219915390015, + "max_us": 67.42439866065979, + "samples_us": [ + 67.072, + 67.104, + 67.072, + 67.232, + 67.36, + 67.168, + 67.264, + 67.232, + 67.424, + 67.04, + 67.296, + 66.272, + 67.136, + 67.232, + 67.168, + 67.392, + 66.272, + 67.264, + 67.2, + 67.296 + ] + } + }, + { + "shape_id": "hy3_tp8_o_proj_m8", + "operator": "o_proj", + "M": 8, + "N": 4096, + "K": 1024, + "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 + }, + "graph": { + "median_us": 31.424200534820557, + "p90_us": 31.488201022148132, + "min_us": 31.327998638153076, + "max_us": 32.28819966316223, + "samples_us": [ + 32.288, + 31.424, + 31.488, + 31.424, + 31.584, + 31.424, + 31.424, + 31.488, + 31.392, + 31.36, + 31.456, + 31.424, + 31.456, + 31.424, + 31.392, + 31.328, + 31.392, + 31.456, + 31.328, + 31.424 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_gate_up_proj_m8", + "operator": "shared_gate_up_proj", + "M": 8, + "N": 384, + "K": 4096, + "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 + }, + "graph": { + "median_us": 57.79229998588562, + "p90_us": 57.95220136642456, + "min_us": 56.83220028877258, + "max_us": 57.95239806175232, + "samples_us": [ + 57.792, + 57.856, + 56.832, + 56.928, + 57.792, + 57.888, + 57.792, + 57.952, + 57.824, + 57.952, + 57.632, + 57.888, + 57.728, + 57.792, + 57.76, + 57.888, + 57.728, + 57.952, + 57.792, + 56.96 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_down_proj_m8", + "operator": "shared_down_proj", + "M": 8, + "N": 4096, + "K": 192, + "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 + }, + "graph": { + "median_us": 14.591999351978302, + "p90_us": 14.62399959564209, + "min_us": 14.496000111103058, + "max_us": 14.624199271202087, + "samples_us": [ + 14.56, + 14.528, + 14.624, + 14.496, + 14.56, + 14.592, + 14.56, + 14.592, + 14.592, + 14.56, + 14.592, + 14.528, + 14.528, + 14.592, + 14.624, + 14.592, + 14.624, + 14.624, + 14.56, + 14.592 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_m8", + "operator": "qkv_proj", + "M": 8, + "N": 1280, + "K": 6144, + "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 + }, + "graph": { + "median_us": 99.36040043830872, + "p90_us": 99.5523989200592, + "min_us": 98.43239784240723, + "max_us": 99.61640238761902, + "samples_us": [ + 99.392, + 99.552, + 99.52, + 99.424, + 98.561, + 99.488, + 98.432, + 99.489, + 98.592, + 99.425, + 98.784, + 99.552, + 99.297, + 98.688, + 98.528, + 99.328, + 99.616, + 99.457, + 99.232, + 98.56 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_and_indexer_qk_m8", + "operator": "qkv_proj_and_indexer_qk", + "M": 8, + "N": 1536, + "K": 6144, + "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 + }, + "graph": { + "median_us": 99.9684989452362, + "p90_us": 100.28860569000244, + "min_us": 99.16840195655823, + "max_us": 100.67239999771118, + "samples_us": [ + 99.872, + 100.385, + 99.456, + 100.224, + 99.553, + 100.672, + 100.289, + 100.16, + 100.256, + 99.425, + 100.192, + 99.328, + 100.065, + 99.328, + 100.192, + 99.52, + 100.256, + 99.489, + 99.168, + 99.328 + ] + } + }, + { + "shape_id": "minimax_tp8_o_proj_m8", + "operator": "o_proj", + "M": 8, + "N": 6144, + "K": 1024, + "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 + }, + "graph": { + "median_us": 32.76819884777069, + "p90_us": 32.96020030975342, + "min_us": 32.70399868488312, + "max_us": 33.82419943809509, + "samples_us": [ + 33.824, + 32.864, + 32.896, + 32.704, + 32.768, + 32.896, + 32.736, + 32.8, + 32.832, + 32.736, + 32.704, + 33.024, + 32.736, + 32.736, + 32.896, + 32.736, + 32.736, + 32.768, + 32.96, + 32.736 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_gate_up_proj_m8", + "operator": "shared_gate_up_proj", + "M": 8, + "N": 768, + "K": 6144, + "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 + }, + "graph": { + "median_us": 92.92849898338318, + "p90_us": 93.15239787101746, + "min_us": 92.00040102005005, + "max_us": 93.34459900856018, + "samples_us": [ + 93.216, + 92.992, + 92.192, + 93.12, + 92.225, + 92.16, + 92.897, + 92.224, + 93.345, + 93.056, + 92.0, + 93.088, + 93.152, + 92.96, + 92.128, + 93.024, + 92.224, + 92.192, + 93.056, + 92.192 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_down_proj_m8", + "operator": "shared_down_proj", + "M": 8, + "N": 6144, + "K": 384, + "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 + }, + "graph": { + "median_us": 19.98410001397133, + "p90_us": 20.03220021724701, + "min_us": 19.904200732707977, + "max_us": 20.064200460910797, + "samples_us": [ + 20.0, + 19.968, + 19.968, + 19.936, + 20.032, + 19.968, + 19.904, + 20.032, + 20.064, + 19.968, + 20.0, + 19.968, + 20.032, + 19.968, + 20.032, + 20.032, + 20.064, + 19.936, + 19.968, + 20.032 + ] + } + }, + { + "shape_id": "glm52_tp8_fused_qkv_a_proj_m8", + "operator": "fused_qkv_a_proj", + "M": 8, + "N": 2624, + "K": 6144, + "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 + }, + "graph": { + "median_us": 86.9763970375061, + "p90_us": 87.13639974594116, + "min_us": 86.01620197296143, + "max_us": 87.13639974594116, + "samples_us": [ + 87.04, + 86.944, + 86.144, + 87.136, + 86.176, + 87.136, + 87.04, + 86.912, + 86.976, + 87.04, + 86.912, + 87.072, + 86.016, + 87.136, + 87.008, + 86.24, + 86.976, + 87.136, + 86.112, + 86.976 + ] + } + }, + { + "shape_id": "glm52_tp8_q_b_proj_m8", + "operator": "q_b_proj", + "M": 8, + "N": 2048, + "K": 2048, + "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 + }, + "graph": { + "median_us": 35.79210042953491, + "p90_us": 35.90399920940399, + "min_us": 35.74399948120117, + "max_us": 36.7682009935379, + "samples_us": [ + 36.768, + 35.808, + 35.808, + 35.808, + 35.776, + 35.744, + 35.776, + 35.808, + 35.808, + 35.904, + 35.744, + 35.872, + 36.704, + 35.776, + 35.776, + 35.744, + 35.776, + 35.744, + 35.84, + 35.776 + ] + } + }, + { + "shape_id": "glm52_tp8_kv_b_proj_m8", + "operator": "kv_b_proj", + "M": 8, + "N": 3584, + "K": 512, + "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 + }, + "graph": { + "median_us": 15.744000673294067, + "p90_us": 15.776200592517853, + "min_us": 15.680000185966492, + "max_us": 21.376200020313263, + "samples_us": [ + 15.712, + 15.744, + 15.744, + 15.744, + 15.68, + 21.376, + 15.712, + 15.744, + 15.744, + 15.712, + 15.776, + 15.776, + 15.68, + 15.744, + 15.808, + 15.68, + 15.712, + 15.712, + 15.744, + 15.776 + ] + } + }, + { + "shape_id": "glm52_tp8_o_proj_m8", + "operator": "o_proj", + "M": 8, + "N": 6144, + "K": 2048, + "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 + }, + "graph": { + "median_us": 58.06429982185364, + "p90_us": 58.56040120124817, + "min_us": 57.53620266914368, + "max_us": 59.296202659606934, + "samples_us": [ + 58.56, + 59.296, + 57.824, + 58.112, + 58.208, + 57.6, + 57.792, + 58.528, + 57.76, + 57.536, + 58.208, + 58.24, + 58.848, + 57.888, + 58.24, + 57.696, + 58.208, + 58.016, + 57.856, + 57.824 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_gate_up_proj_m8", + "operator": "shared_gate_up_proj", + "M": 8, + "N": 512, + "K": 6144, + "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 + }, + "graph": { + "median_us": 88.44839930534363, + "p90_us": 88.70440125465393, + "min_us": 87.52040266990662, + "max_us": 88.73639702796936, + "samples_us": [ + 88.736, + 88.448, + 88.608, + 88.64, + 88.384, + 88.416, + 87.712, + 87.744, + 87.584, + 88.64, + 87.52, + 88.224, + 87.68, + 88.704, + 88.448, + 88.609, + 87.52, + 88.576, + 88.672, + 88.704 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_down_proj_m8", + "operator": "shared_down_proj", + "M": 8, + "N": 6144, + "K": 256, + "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 + }, + "graph": { + "median_us": 15.07200002670288, + "p90_us": 15.136000514030457, + "min_us": 15.007999539375305, + "max_us": 15.136200189590454, + "samples_us": [ + 15.136, + 15.072, + 15.04, + 15.072, + 15.008, + 15.008, + 15.072, + 15.072, + 15.072, + 15.04, + 15.136, + 15.104, + 15.136, + 15.104, + 15.072, + 15.072, + 15.104, + 15.104, + 15.104, + 15.072 + ] + } + } +] \ 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/tools/retest_variant.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/retest_variant.py new file mode 100644 index 00000000..744f0207 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/retest_variant.py @@ -0,0 +1,331 @@ +"""A/B retest tool: compare a candidate kernel against the shipped variant. + +Why this exists +--------------- +DKAO records each candidate's median in its ``accepted//manifest.json`` +and each shipped variant's median in its header comment. Those are *records*, +not measurements: if the round that produced them ran while another workload +shared the device, the recorded number can be suppressed (we saw 884us -> 2043 +-> 3047us for one unchanged shape). Before promoting a kernel — especially one +copied out of an old workspace — re-measure both sides under identical +conditions. + +What "A/B" means here +--------------------- +* **A** = the *candidate* kernel (what we might promote). +* **B** = the *variant* currently shipped in DKAO (the incumbent). +* Both are compiled from the same repository skeleton with only the kernel + source swapped, benchmarked on the **same idle device**, with the **same + sampling parameters**, and measured **alternately** (A, B, A, B...). + +Alternating matters: GPU clock/thermal state and any background load drift over +time, so "measure A fully, then B fully" would mix that drift into the +difference. Interleaving makes both sides absorb the same conditions, so the +remaining difference is attributable to the kernel itself. Two rounds are run +so a single fluke cannot decide the outcome; medians are compared and a +candidate must win by ``--threshold`` percent (default 3, matching DKAO's +promotion rule) *and* reproduce its own recorded median (``--recorded-tol``) +before it is written. + +Usage +----- + python -m metainfer.tasks.dcu_kernel_auto_opt.tools.retest_variant \ + --candidate-workspace \ + --shape hy3_tp4_o_proj_m4096 --rounds 2 [--promote] + + # or straight from an accepted kernel file + python ...retest_variant --shape ... --candidate-kernel path/kernel.hip \ + --kernel-root +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import statistics +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +PLUGIN_DIR = Path(__file__).resolve().parents[1] +REPO_ROOT = Path(__file__).resolve().parents[4] +BENCH = PLUGIN_DIR / "assets" / "w8a8_bench.py" +VARIANT_ROOT = PLUGIN_DIR / "variant" / "int8w8a8-gemm" + +#: DKAO compiles for gfx928 only. The int8 DUMMA fragment specialisations in +#: /opt/dtk/include/du_mma.h sit behind ``#if defined(__gfx928__)``, so a +#: default multi-arch torch build fails with "undefined template DUFragment +#: <... signed char ...>". w8a8_pipeline sets this for the same reason. +BENCH_ENV = {"PYTORCH_ROCM_ARCH": "gfx928"} + + +# --------------------------------------------------------------- helpers ---- + +def find_variant(model: str, tp: int, m: int, operator: str, + root: Optional[Path] = None) -> Optional[Path]: + """Shipped variant for a shape, tolerating model-label drift (glm -> glm52).""" + root = Path(root or VARIANT_ROOT) + exact = root / model / f"TP{tp}" / f"M{m}" / f"{operator}.hip" + if exact.is_file(): + return exact + for model_dir in sorted(root.iterdir()) if root.is_dir() else []: + if not model_dir.name.startswith(model): + continue + candidate = model_dir / f"TP{tp}" / f"M{m}" / f"{operator}.hip" + if candidate.is_file(): + return candidate + hits = sorted(root.glob(f"*/TP{tp}/M{m}/{operator}.hip")) if root.is_dir() else [] + return hits[0] if hits else None + + +def parse_variant_header_median(path: Path) -> Optional[float]: + """Median recorded in a variant's header comment (recorded, not measured).""" + try: + for line in path.read_text(encoding="utf-8", + errors="replace").splitlines()[:8]: + if "median_us=" in line: + return float(line.split("median_us=")[1].split()[0]) + except (OSError, ValueError, IndexError): + return None + return None + + +def bench_env(device: int, base: Optional[Dict[str, str]] = None) -> Dict[str, str]: + """Environment for an independent measurement on one device.""" + env = dict(base if base is not None else os.environ) + env["HIP_VISIBLE_DEVICES"] = str(device) + env["CUDA_VISIBLE_DEVICES"] = str(device) + env.update(BENCH_ENV) + return env + + +def decide(candidate_us: Optional[float], variant_us: Optional[float], *, + threshold_percent: float = 3.0, + recorded_us: Optional[float] = None, + recorded_tol_percent: float = 5.0) -> Dict[str, Any]: + """Promotion verdict from measured medians (+ reproducibility check).""" + if not candidate_us or not variant_us: + return {"verdict": "measurement_failed", + "reason": "one side produced no median"} + improvement = (variant_us - candidate_us) / variant_us * 100.0 + recorded_delta = None + if recorded_us: + recorded_delta = (candidate_us - float(recorded_us)) / float(recorded_us) * 100.0 + if improvement < float(threshold_percent): + return {"verdict": "keep", "improvement_percent": improvement, + "recorded_delta_percent": recorded_delta, + "reason": (f"improvement {improvement:.2f}% < required " + f"{float(threshold_percent):.2f}%")} + if recorded_delta is not None and abs(recorded_delta) > float(recorded_tol_percent): + return {"verdict": "needs_review", "improvement_percent": improvement, + "recorded_delta_percent": recorded_delta, + "reason": (f"win {improvement:.2f}% but the re-measurement " + f"differs from the record by {recorded_delta:.2f}%")} + return {"verdict": "promote", "improvement_percent": improvement, + "recorded_delta_percent": recorded_delta, + "reason": f"candidate faster by {improvement:.2f}%"} + + +def accepted_kernel_for(workspace: Path, shape_id: str) -> Optional[Path]: + hits = sorted(Path(workspace).glob( + f"workers/*/accepted/{shape_id}/kernel.hip")) + return hits[0] if hits else None + + +def recorded_median_for(kernel: Path) -> Optional[float]: + try: + manifest = json.loads((kernel.parent / "manifest.json") + .read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + value = (manifest.get("metrics") or {}).get("median_us") + try: + return float(value) + except (TypeError, ValueError): + return None + + +def idle_device(gpu_ids: Optional[List[int]] = None) -> Optional[int]: + """First device satisfying the operating gate (VRAM<=90 and HCU==0).""" + try: + sys.path.insert(0, str(REPO_ROOT)) + from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.gpu_preflight import ( # noqa: E501 + preflight_gpus, + ) + except Exception: # noqa: BLE001 + return None + plan = preflight_gpus(gpu_ids or [0, 1, 2, 3], samples=3, interval_s=2) + clean = plan.get("clean_ids") or [] + return int(clean[0]) if clean else None + + +def prepare_repo(kernel_root: Path, kernel: Path, target: Path) -> Path: + """Copy a kernel repo skeleton and swap in one kernel source.""" + if target.exists(): + shutil.rmtree(target) + shutil.copytree(kernel_root, target) + (target / "csrc" / "w8a8_gemm_hip.hip").write_text( + Path(kernel).read_text(encoding="utf-8", errors="replace"), + encoding="utf-8") + return target + + +def measure(repo: Path, shape: Tuple[int, int, int], device: int, *, + warmups: int = 30, samples: int = 30, replays: int = 50, + reference_cache: Optional[Path] = None, + timeout_s: int = 3600) -> Optional[float]: + m, n, k = shape + cmd = [sys.executable, str(BENCH), "--source", str(repo), + "--m", str(m), "--n", str(n), "--k", str(k), + "--warmups", str(warmups), "--samples", str(samples), + "--replays-per-sample", str(replays), "--skip-correctness"] + if reference_cache is not None and Path(reference_cache).is_dir(): + cmd += ["--reference-cache-dir", str(reference_cache)] + proc = subprocess.run(cmd, capture_output=True, text=True, cwd=str(repo), + env=bench_env(device), timeout=timeout_s) + for line in reversed((proc.stdout or "").splitlines()): + line = line.strip() + if line.startswith("{") and line.endswith("}"): + try: + payload = json.loads(line) + except ValueError: + continue + value = payload.get("median_us") + return float(value) if value else None + return None + + +def ab_retest(*, shape_id: str, candidate_kernel: Path, kernel_root: Path, + model: str, tp: int, m: int, operator: str, + shape: Tuple[int, int, int], rounds: int = 2, + device: Optional[int] = None, warmups: int = 30, samples: int = 30, + replays: int = 50, threshold: float = 3.0, + recorded_tol: float = 5.0, + reference_cache: Optional[Path] = None, + build_root: Path = Path("/tmp/retest_build"), + variant_path: Optional[Path] = None) -> Dict[str, Any]: + """Interleaved A/B measurement of one candidate against the variant.""" + device = device if device is not None else idle_device() + if device is None: + return {"shape": shape_id, "verdict": "no_idle_device"} + variant = variant_path or find_variant(model, tp, m, operator) + if variant is None: + return {"shape": shape_id, "verdict": "no_variant", + "reason": f"no variant for {model}/TP{tp}/M{m}/{operator}"} + cand_repo = prepare_repo(kernel_root, candidate_kernel, + Path(build_root) / shape_id / "cand") + var_repo = prepare_repo(kernel_root, variant, + Path(build_root) / shape_id / "var") + measured: Dict[str, List[float]] = {"cand": [], "var": []} + for _ in range(max(1, int(rounds))): + for label, repo in (("cand", cand_repo), ("var", var_repo)): + value = measure(repo, shape, device, warmups=warmups, + samples=samples, replays=replays, + reference_cache=reference_cache) + if value: + measured[label].append(value) + cand_median = statistics.median(measured["cand"]) if measured["cand"] else None + var_median = statistics.median(measured["var"]) if measured["var"] else None + verdict = decide(cand_median, var_median, threshold_percent=threshold, + recorded_us=recorded_median_for(candidate_kernel), + recorded_tol_percent=recorded_tol) + return {"shape": shape_id, "device": device, + "candidate_medians_us": measured["cand"], + "variant_medians_us": measured["var"], + "candidate_median_us": cand_median, "variant_median_us": var_median, + "variant_path": str(variant), + "recorded_median_us": recorded_median_for(candidate_kernel), + **verdict} + + +# ------------------------------------------------------------------- CLI ---- + +def main(argv: Optional[List[str]] = None) -> int: + ap = argparse.ArgumentParser( + prog="retest-variant", + description="A/B re-measure a candidate kernel against the shipped variant") + ap.add_argument("--shape", required=True, help="shape id, e.g. hy3_tp4_o_proj_m4096") + ap.add_argument("--candidate-workspace", type=Path, + help="AHE child workspace (accepted kernel is looked up)") + ap.add_argument("--candidate-kernel", type=Path, help="explicit kernel.hip") + ap.add_argument("--kernel-root", type=Path, + help="repo skeleton containing csrc/ (defaults to the " + "candidate's worker source)") + ap.add_argument("--model", help="model label (default: parsed from the shape id)") + ap.add_argument("--tp", type=int, help="tensor-parallel size (default: parsed)") + ap.add_argument("--m", type=int, help="M (default: from the accepted manifest)") + ap.add_argument("--operator", help="operator name (default: parsed from the shape id)") + ap.add_argument("--shape-mnk", help="M,N,K of the problem, e.g. 4096,4096,2048") + ap.add_argument("--rounds", type=int, default=2) + ap.add_argument("--warmups", type=int, default=30) + ap.add_argument("--samples", type=int, default=30) + ap.add_argument("--replays", type=int, default=50) + ap.add_argument("--threshold", type=float, default=3.0) + ap.add_argument("--recorded-tol", type=float, default=5.0) + ap.add_argument("--device", type=int, help="GPU id (default: pick an idle one)") + ap.add_argument("--promote", action="store_true", + help="write the variant when the verdict is 'promote'") + ap.add_argument("--json", type=Path, help="write the report to this path") + args = ap.parse_args(argv) + + from metainfer.tasks.harness_evolve.orchestrator.pool import Pool + pool_file = Path("/root/zth_agent/ahe-kernel-repos/registered_pool.yaml") + pool = Pool.from_yaml(pool_file) if pool_file.is_file() else None + inst = pool.get(args.shape) if pool is not None else None + + candidate = args.candidate_kernel + if candidate is None: + if not args.candidate_workspace: + ap.error("--candidate-workspace or --candidate-kernel is required") + candidate = accepted_kernel_for(args.candidate_workspace, args.shape) + if candidate is None or not Path(candidate).is_file(): + ap.error(f"candidate kernel not found for {args.shape}") + + kernel_root = args.kernel_root or Path(candidate).parents[2] / "source" + model = args.model or (inst.model if inst else args.shape.split("_tp")[0]) + tp = args.tp or (inst.tp_size if inst else 4) + m = args.m or (inst.M if inst else 4096) + operator = args.operator or args.shape.split("_m")[0].split("_", 2)[-1] + if args.shape_mnk: + mnk = tuple(int(x) for x in args.shape_mnk.split(",")) + elif inst is not None: + mnk = (inst.M, inst.N, inst.K) + else: + ap.error("--shape-mnk is required when the shape is not in the pool") + + report = ab_retest( + shape_id=args.shape, candidate_kernel=Path(candidate), + kernel_root=Path(kernel_root), model=model, tp=int(tp), m=int(m), + operator=operator, shape=mnk, rounds=args.rounds, device=args.device, + warmups=args.warmups, samples=args.samples, replays=args.replays, + threshold=args.threshold, recorded_tol=args.recorded_tol, + ) + print(json.dumps(report, ensure_ascii=False, indent=2)) + if args.json: + Path(args.json).write_text(json.dumps(report, ensure_ascii=False, indent=2), + encoding="utf-8") + + if args.promote and report.get("verdict") == "promote": + from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.variant_promote import ( # noqa: E501 + model_label_for, promote_variant, + ) + workspace = args.candidate_workspace or Path(candidate).parents[4] + outcome = promote_variant( + workspace_dir=Path(workspace), + answers={"operator": "Quantized GEMM", "dtype": "INT8 W8A8", + "model": model, "tp_size": tp}, + shape_id=args.shape, + source_task=f"{args.shape} (A/B retested)", + correctness_ok=True, min_improvement_percent=args.threshold, + tp=int(tp), m=int(m), + model_label=(model_label_for(args.shape, str(model)) or None), + ) + print(json.dumps({"promotion": outcome}, ensure_ascii=False, indent=2)) + 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..b8032fe2 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,266 +1,78 @@ -// @@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=7cf8d67fe9e0c38f2b4c3c67aded1bf8cc894342 added=2026-09-14 +// median_us=693.8 p90_us=695.3 speedup=98.96 baseline_us=6.866e+04 +// source=test-9-11-1/iteration_002/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 +// (GLM5.2 TP=8 fused_qkv_a_proj, logical (K, N) = (6144, 2624)) // -// Operator contract (fixed by the control plane): +// Operator contract (fixed by csrc/bindings.cpp and int8_w8a8_gemm_api.py): // out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) -// * x_scale[m] * weight_scale[n] ) +// * x_scale[m, 0] * weight_scale[n, 0] ) // -// 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). +// Mandatory decision this round: the tile-shape (M-tile vs N-tile aspect ratio) +// round. The packed_weight architecture, the exact-shape packed-B swizzle, the +// double-buffered 64-K stage pipeline, the bank-safe A LDS stride, the +// conflict-free swizzled B fragment reads and the fused coalesced bf16 epilogue +// are all RETAINED unchanged; only the macro-tile aspect ratio and its derived +// wave split are retuned for M=4096, N=2624, K=6144. No split-K, no workspace +// use, no allocation/compilation/packing/synchronization inside the timed API. // -// 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. +// (1) launch_pack_w8a8_weight packs the exact (k, n) == (6144, 2624) weight +// ONCE, outside the timed region and out of Graph capture, into a +// fragment-interleaved swizzle of 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], +// n16g = n>>4 = 164. Byte-exact bijection 192 x 164 x 512 = k*n = +// 16,121,856: same byte count and same buffer allocation, so every +// captured address is graph-stable. Every other (K, N) keeps the identity +// device-to-device copy. +// (2) the exact-shape guard dispatches the new +// w8a8_dumma_prefill_packedb_db_kernel<256, 128, 64, 512, true>: +// grid = ((N + 127)/128) x (M/256) = 21 x 16 = 336 blocks, +// block = 512 threads = 8 gfx928 wavefronts, +// a 64 x 64 quadrant of sixteen m16n16k32 int32 accumulator fragments +// per wave (4 wave rows x 2 wave columns, the balanced square panel), +// i.e. a 256 x 128 M-major operand-reuse macro-tile, 64-K +// DOUBLE-buffered stages with ONE __syncthreads per stage and the +// stage-(s+1) global loads issued into staging VGPRs before the stage-s +// consume, k0-outer over 96 ascending 64-K stages with kk-inner +// ascending m16n16k32 MMAs (exact int32 accumulation order of the +// reference), and a fused branch-free per-row/column scale + bf16 +// epilogue emitting one coalesced 8-byte store per lane. +// The A LDS tile uses the bank-safe 72-byte row stride (StageK + 8: an +// odd 8-byte count, so the 16 rows of each 16-lane ds_read_b64 phase land +// on 16 distinct bank pairs); the swizzled B tile is a flat BN x 64-byte +// copy of the packed stage, so every col_major B fragment read is one +// conflict-free ds_read_b64 (load_frag8_swz). +// LDS = 2 x (A[256, 72] + B[128, 64]) = 2 x 18,432 + 2 x 8,192 = +// 53,248 B per block, so ONE resident block/CU = 8 waves/CU (the +// two-operand global traffic falls from M*K*21 + N*K*32 = 1,044 MB to +// M*K*21 + N*K*16 = 786 MB, -24.7%, at unchanged wave occupancy). +// (3) the scalar int8/int32 fallback decodes the swizzled layout element-wise +// in ascending kk for (k, n) == (6144, 2624), so the paired M=2 / M=16 +// API shapes with the same (N, K) stay byte-exact; the generic +// single-buffered <64, 64, 128> identity arm covers every other large-M +// shape with n % 64 == 0 and k % 128 == 0. The fallback is never replaced +// by an unconditional shape-specialized launch. // -// 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). +// Correctness: the same weight elements reach the same fragment slots in the +// same ascending k-chunk sequence (k0-outer over 96 ascending 64-K stages, +// kk-inner ascending over two m16n16k32 steps, i.e. the identical (k0 + kk) +// k-chunk sequence 0, 32, 64, 96, ...), and the epilogue applies the exact +// reference multiply order float(dot) * x_scale[row] * weight_scale[col], so +// every bf16 output byte is bit-identical to the scalar reference. Widening BM +// only changes WHICH wave computes which output element; it never reorders the +// K stages or the int32 accumulation. // -// Iteration 3 (mandated pipeline round: single vs double buffering across K -// tiles; retain double buffering only on ISA/PMC evidence of reduced VMEM -// stalls without harmful LDS or occupancy growth): -// * Bottleneck evidence from the accepted iteration-2 source (fresh PMC -// digest 20f3bc9a..., profiles/.../iteration3/pmc.json: profiled -// 2041.095 us, vmem_read 1,574,400, lds_instructions 7,892,992 with -// lds_wait 4,571,574, lds_bank_conflicts 0, 26,112 B LDS, 2 blocks/CU): -// the exact-source ISA of w8a8_dumma_prefill_packedb_kernel<128,64,128> -// shows the stage loop is [global_load_dwordx4 ... -> s_waitcnt -// vmcnt(0) -> ds_write2_b64 -> s_barrier -> 24x ds_read_b64 -> 32x -// v_mmac -> s_barrier] -- the FULL global-load round trip (issue, vmcnt -// wait, LDS flush) is serialized in front of every consume with ZERO -// compute overlap, and the 14 vmcnt + 57 lgkmcnt waits in the symbol sit -// on that critical path. The K=6144 grid (1312 blocks) and 2-block -// co-residency do not hide this per-stage latency. -// * Focused change (one mechanism, all inside csrc/w8a8_gemm_hip.hip): -// the exact-shape guard now dispatches a new -// w8a8_dumma_prefill_packedb_db_kernel<128,64,64> -- SAME macro-tile -// (grid 41x32 = 1312 blocks, 256 threads = 4 waves, 64x32 quadrant per -// wave, eight m16n16k32 int32 accumulators), SAME load_frag8 fragment -// reads, SAME coalesced epilogue -- with the K stage changed from a -// single-buffered 128-K stage (2 __syncthreads, full VMEM latency -// exposed) to a DOUBLE-buffered 64-K stage with ONE __syncthreads per -// stage: the stage-(s+1) global loads are issued into 3 int4 staging -// VGPR per thread BEFORE the stage-s consume, and the prefetched vectors -// are flushed into the idle LDS buffer after the consume (the validated -// M=3072 prefetch order global_load_dwordx4 -> v_mmac -> vmcnt(0) -> -// ds_write2_b64, here in pure HIP: the vmcnt wait lands on the data -// dependency before the ds_write, i.e. AFTER the MMAC burst). -// * Resources: LDS = 2 x (A[128,72] + B[64,72]) = 2 x 13,824 = 27,648 -// B/block (26,112 -> 27,648 B, +5.9%); 2 x 27,648 = 55,296 <= 65,536 B -// keeps TWO resident blocks/CU = 8 waves/CU, occupancy UNCHANGED (the -// StageK=128 double-buffer alternative needs 52,224 B/block -> 1 -// block/CU = 4 waves/CU and is rejected up front as harmful occupancy -// loss). kLdsStride = 72 = 64 + 8: 8-aligned with an odd 8-byte count -// (72/8 = 9) so the 16 rows of each 16-lane ds_read_b64 phase land on 16 -// distinct bank pairs (18r mod 32, r = 0..15) -- the same conflict-free -// 4-phase floor as iteration 2 -- and 72 % 16 == 8 keeps the two-half -// ds_write2_b64 staging form. VGPR grows by the 12 staging registers -// (~72 -> ~84), far below the 128/thread budget for 2 blocks/CU; 0 -// spills expected. Barriers per block: 1 (prologue) + 95 (loop) = 96, -// the same as iteration 2's 2 x 48 (the final stage is consumed without -// a trailing barrier). -// * Correctness: k0-outer over 96 ASCENDING 64-K stages with kk-inner -// ascending over two m16n16k32 steps gives the identical (k0+kk) -// k-chunk sequence 0, 32, 64, 96, ... as iteration 2's 48 x 128-K -// stages, and the element-to-slot fragment mapping is unchanged, so -// every int32 accumulator is bit-identical (mismatch 0 / max_abs_error -// 0.0 expected). -// * Falsifiable retention: double buffering is retained only if the exact -// shape shows median_us < 2276.930923461914 AND p90_us < -// 2540.756378173828 with mismatch_count == 0, max_abs_error == 0.0, -// graph_capture_passed == true, the exact code object shows 0 spills / -// LDS == 27,648 B (2 blocks/CU still fits: 55,296 <= 65,536) and the -// stage ISA shows the prefetch loads issued BEFORE the ds_read_b64 / -// v_mmac burst with the vmcnt wait after the MMACs (load -> mmac -> -// vmcnt(0) -> ds_write2_b64), i.e. reduced exposed VMEM stalls; any -// median/p90 >= current best, any mismatch, any spill, any occupancy -// loss, or any ISA that still serializes the vmcnt wait in front of the -// consume falsifies the round and the control plane restores the -// iteration-2 source (single buffering retained). -// -// Iteration 5 (mandated packing round: test one weight packing/swizzle that -// makes each DUMMA B tile vector-loadable and LDS-bank-safe): -// * Baseline: the accepted iteration-3 kernel -// w8a8_dumma_prefill_packedb_db_kernel<128,64,64> (official median -// 1070.23 us = 123.4 logical TOPS; fresh PMC digest 47542b38..., -// profiles/.../iteration4/pmc.json of the current source: profiled -// 1101.96 us, grid 1312, 96 arch_vgpr / 27,648 B LDS, lds_instructions -// 4,901,632 = 9.7 per wave-stage (6 merged ds_read2_b64 + 3 -// ds_write2_b64), lds_wait 3,701,906, lds_bank_conflicts 12,091,392 = -// 6 x 4 x 4 x 96 x 1312 EXACTLY -- one conflict per 16-lane phase of -// every merged ds_read2_b64, the inherent 2-cycle alias of the +16-row -// (16 x 72 = 1152 B = 9 x 128 B) second-fragment offset, NOT a -// layout defect (the iteration-2 unmerged ds_read_b64 control had 0 -// conflicts at the same 72/136-B strides), vmem_read 1,574,400, L2 hit -// 79.3% with 5,155,404 misses, TCC read 269.3 MB/replay). The kernel is -// LDS-latency-bound, not HBM-bound; the global B operand is the -// remaining layout lever. -// * Mechanism: replace the iteration-2 n-major pack -// packed[n*K + kk] (each 64-K B tile = 64 column chunks of 64 B at -// 6144-B stride -> 64 x 128-B lines touched per tile, 50% line -// utilization, scattered L2 footprint) with a B-PANEL pack -// packed[(n>>6)*(k*64) + (kk>>4)*1024 + (n&63)*16 + (kk&15)] = -// raw[kk*n + n] (41 panels of [K, 64] int8, 16-byte k-runs per column, -// byte-exact bijection verified by simulation): every [64,64] B tile of -// a 64-K stage is then FOUR CONTIGUOUS 1-KiB slabs (one per 16-k-run), -// each staging vector is a 16-byte-aligned contiguous k-run, and a wave -// reads 4 x 256 B of one panel region per stage -- the tile is -// vector-loadable as one coalesced stream with 8x fewer 128-B lines -// touched per tile (32 vs 256) and full line utilization, so L2 tag -// pressure and B HBM fetch traffic drop without any change in LDS -// layout or instruction count. -// * Kernel-side changes are ONLY the two B global-load sites (prologue + -// steady-state prefetch) of w8a8_dumma_prefill_packedb_db_kernel -// (dispatched solely by the exact-shape guard; static_assert BN == 64): -// vector (n_row, k16) is loaded from byte -// (n0>>6)*(k*64) + ((k0>>4)+k16)*1024 + n_row*16 instead of -// (n0+n_row)*k + k0 + k16*16. LDS staging layout (b_tile[64][72], -// two-half ds_write2_b64), the col_major load_frag8 consumes, the -// k0-outer/kk-inner int32 accumulation order, the coalesced epilogue, -// the grid (41x32 = 1312 blocks), threads (256 = 4 waves), occupancy -// (2 blocks/CU, LDS 27,648 B) and VGPR budget are ALL unchanged, so the -// round is a pure packing/swizzle A/B. -// * Correctness: the same weight elements reach the same fragment slots -// in the same ascending k-chunk sequence, so every int32 accumulator is -// bit-identical (mismatch 0 / max_abs_error 0.0 expected); the scalar -// fallback decodes the panel layout when (k,n) == (6144,2624) (ascending -// 16-byte k-runs keep the exact int32 order for the paired M=2/M=16 API -// shapes); the generic identity arms and sibling kernels are untouched. -// * Packing stays outside the timed region and out of Graph capture -// (launch_pack_w8a8_weight, same byte count k*n = 16,121,856 and same -// buffer -> graph-stable captured addresses unchanged; pack relation -// verified by a byte-exact index simulation before dispatch). -// * Predicted deltas: vmem_read_instructions unchanged (~1,574,400), LDS -// counters unchanged (conflicts stay at the 12,091,392 merged-read alias -// floor), L2 misses < 5,155,404 with hit rate UP, TCC read bytes < -// 269,329,728 (B line fetches halve), median ~1000-1070 us. Falsifiable: -// retained only if median_us < 1070.2328491210938 AND p90_us < -// 1081.9308471679688 with mismatch_count == 0, max_abs_error == 0.0, -// graph_capture_passed == true, the exact code object shows 0 spills / -// LDS == 27,648 B / VGPR <= 128 (2 blocks/CU), and PMC confirms L2 -// misses and TCC read bytes strictly below the iteration-4 values with -// vmem_read/lds/conflicts not growing; any median/p90 >= current best, -// any mismatch, any spill, any occupancy loss, or any L2/HBM regression -// falsifies the pack and the control plane restores the iteration-3 -// source. +// Reported launch geometry for this round (iteration 4, tile shape): +// grid_blocks = 21 x 16 = 336 +// waves_per_block = 512 / 64 = 8 +// estimated_active_cus = 120 (K500SM_AI has 120 CUs; 336 blocks with +// 1 resident block/CU = 2.8 blocks/CU occupy every CU) // // Include order is fixed by the control plane: hip_runtime, hip_bfloat16, // then du_mma (this DTK's du_mma.h is not self-contained otherwise). @@ -273,12 +85,18 @@ 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 // 4 wavefronts; must remain a multiple of the gfx928 wavefront size (64). constexpr int kThreadsPerBlock = 256; +// Exact-shape packed-B prefill workgroup: 8 wavefronts = 512 threads, with the +// M-major operand-reuse macro-tile (BM=256 x BN=128) selected by the +// iteration-4 tile-shape round. +constexpr int kPackedThreads = 512; +constexpr int kPackedBM = 256; +constexpr int kPackedBN = 128; // Exact assigned shape (K, N) after TP=8 partitioning. constexpr int kExactN = 2624; @@ -286,21 +104,39 @@ constexpr int kExactK = 6144; using namespace du::dumma; -// Coalesced accumulator epilogue for gfx928 INT8 m16n16k32 (verified against -// du_store_matrix_sync in the accepted worker-29 lineage): the accumulator -// lane mapping is row = lane & 15, col_mod4 = lane >> 4, frag.x[i] maps to -// columns col_mod4 + 4*i. A register 4x4 transpose (two 2x2 steps with -// shfl_xor 16 then 32, one v_cndmask per element per step) re-routes the -// int32 values so lane (r, c4) owns the four CONTIGUOUS columns -// base_col + 4*c4 .. +3; they are scaled (float(dot) * x_scale[row] * -// weight_scale[col], the exact reference multiply order), converted to bf16, -// packed and stored with ONE 8-byte store per lane. The row >= m guard is -// wavefront-uniform (lane & 15 cycles the same 16 rows in every 16-lane -// group), so the shuffles never mix active and inactive lanes. base_col is a -// multiple of 64 for the dispatched tiles, so the float4 weight_scale load -// (col0 % 4 == 0) and the 8-byte store (n even -> out offset a multiple of -// 8) are aligned. -template +// Branch-free round-to-nearest-even fp32 -> bf16 high-16 payload: +// bf16 = (b + 0x7fff + ((b >> 16) & 1)) >> 16, b = __float_as_uint(v) +// For every FINITE fp32 this is exactly the normal path of __float2bfloat16, +// while dropping the ~13-instruction EXEC-bracketed Inf/NaN sequence this DTK +// emits for the library call (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, max-float still rounds to +/-Inf and +/-0.0/subnormal payloads are +// preserved). Everything converted here is finite by construction: |int32 dot| +// < 2^31 and the dequantization scales are finite. +__device__ __forceinline__ unsigned int bf16_rne_high16(float v) { + const unsigned int b = __float_as_uint(v); + return (b + 0x7fffu + ((b >> 16) & 1u)) >> 16; +} + +// Fused per-row/per-column dequantization + bf16 epilogue for the m16n16k32 +// int32 accumulator of gfx928 INT8 DUMMA. Accumulator lane mapping: +// row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> column col_mod4 + 4*i. +// A register 4x4 transpose (two 2x2 steps via shfl_xor 16 then 32) re-routes +// the values so lane (r, c4) owns the four CONTIGUOUS columns base_col + +// 4*c4 .. +3, which are then scaled in the exact reference order +// float(dot) * x_scale[row] * weight_scale[col], rounded to bf16 and stored +// with ONE 8-byte store per lane. 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 are both aligned. The row >= m guard is wavefront-uniform +// (lane & 15 cycles the same 16 rows in each 16-lane group), so the shuffles +// never mix active and inactive lanes. +// +// kGuarded == false is the exact-shape full-tile specialization (m % BM == 0): +// every row of every block tile is strictly inside the output, so the dead +// `row >= m` EXEC bracket around each of the eight unrolled stores is removed +// at compile time. The default kGuarded == true keeps the exact guarded +// semantics for every other m. +template __device__ __forceinline__ void store_prefill_fragment_coalesced( const AccFragment& frag, const float* __restrict__ x_scale, @@ -312,8 +148,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]; @@ -341,8 +179,8 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced( const int f2 = hi ? t0 : a2; const int f3 = hi ? t1 : a3; - // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte - // aligned: base_col is a multiple of 64, n*2 is a multiple of 8). + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte aligned: + // base_col is a multiple of 64 and n*2 is a multiple of 8). const int col0 = base_col + 4 * c4; const float xs = x_scale[row]; const float4 ws = *reinterpret_cast(weight_scale + col0); @@ -352,53 +190,126 @@ __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) | + bf16_rne_high16(v0) | + (bf16_rne_high16(v1) << 16)) | (static_cast( - static_cast(__float2bfloat16(v2))) - << 32) | - (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 -// 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)). +// Explicit 8-byte int8 fragment loader 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; +} + +// Fragment-interleaved B swizzle helpers. // -// K is staged cooperatively in a SINGLE LDS buffer at StageK granularity -// with TWO __syncthreads per stage (one after the staging stores, one after -// all waves consumed the buffer): -// * A[BM, StageK] is staged row-major with flat 16-byte vectors (each -// thread strides over kAVectors = BM*StageK/16 vectors); rows past M are -// zero-filled and masked in the epilogue. -// * B[StageK, BN] is staged row-major from the logical [K, N] weight with -// flat 16-byte vectors (each thread strides over kBVectors = -// StageK*BN/16 vectors). The dispatch guarantees n % BN == 0 and -// k % StageK == 0, so every global vector is aligned and in-bounds. -// * Fragment loads use the library du_load_matrix_sync (row_major for both -// operands) and accumulate with du_mma_sync in ascending k order -- -// k0-outer over stages, kk-inner over kTileK steps -- preserving the -// exact int32 accumulation order of the scalar reference. -// * The epilogue is the coalesced fragment store. +// 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 library col_major loader consumes. +__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); +} + +// 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; +} + +// Tile-shape round: balanced per-wave quadrant. Pick kWaveRows (hence +// kWaveCols = kWaves / kWaveRows) so that the wave's m16n16k32 panel +// (kWaveM16 x kWaveN16 fragments) is as square as possible. A square panel +// minimises the LDS fragment reads per v_mmac, (kWaveM16 + kWaveN16) reads for +// kWaveM16 * kWaveN16 MMAs, and keeps the live operand fragment set -- and +// therefore the VGPR bill of the 16-accumulator M-major macro-tile -- small. +// The search is restricted to exact divisor splits of both the tile fragment +// counts and the wave count so the kernel's quadrant arithmetic stays valid; +// for the retained 128 x 128 / 8-wave geometry it reproduces the historical +// 2 wave rows x 4 wave columns split exactly. +__host__ __device__ constexpr int balanced_wave_rows(int bm_frags, + int bn_frags, + int waves) { + int best_rows = 1; + int best_diff = 1 << 30; + for (int rows = 1; rows <= waves; ++rows) { + if (waves % rows != 0) continue; + if (bm_frags % rows != 0) continue; + const int cols = waves / rows; + if (bn_frags % cols != 0) continue; + const int m16 = bm_frags / rows; + const int n16 = bn_frags / cols; + const int diff = (m16 > n16) ? (m16 - n16) : (n16 - m16); + if (diff < best_diff) { + best_diff = diff; + best_rows = rows; + } + } + return best_rows; +} + +// 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)). // -// Bootstrap keeps this direct/single-buffered by design (no double buffering, -// no split-K, no raw asm); later rounds tune tile, stage, and packing. +// K is staged cooperatively in a SINGLE LDS buffer at StageK granularity with +// TWO __syncthreads per stage. 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 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,269 +420,147 @@ 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). -// Both fragment loaders are the explicit load_frag8 above; the staged bytes -// are the same matrix elements in the same fragment slots as the library -// loaders produced from the identity layout, so every v_mmac input and the -// exact int32 accumulation (k0-outer over 48 stages, kk-inner ascending -// over four m16n16k32 steps) are bit-identical to the accepted iteration-1 -// kernel and to the scalar reference. -template -__global__ __launch_bounds__(kThreadsPerBlock) void -w8a8_dumma_prefill_packedb_kernel( - const int8_t* __restrict__ x_q, // [M, K] row-major - const int8_t* __restrict__ packed_w, // [N, K] n-major (exact shape) - const float* __restrict__ x_scale, // [M] - const float* __restrict__ weight_scale, // [N] - hip_bfloat16* __restrict__ out, // [M, N] +// Exact-shape A-staging load. For the dispatched M=4096 every 256-row block +// tile lies fully inside the output, so the `(g_row < m)` row guard is dead: +// the exact code object would otherwise lower it to unconditional zero-fills +// plus EXEC save/restore uniform-branch pairs per K stage, i.e. dead +// instructions and an EXEC serialization in front of the very +// global_load_dwordx4s 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 for every m % BM != 0. +template +__device__ __forceinline__ int4 load_a_stage_vec( + const int8_t* __restrict__ p, int m, - int n, - int k) { - constexpr int kWaveM16 = BM / 32; - constexpr int kWaveN16 = BN / 32; - static_assert(BM % 32 == 0 && BN % 32 == 0, - "block tile must be a multiple of the wave quadrant"); - static_assert(kThreadsPerBlock % kWaveSize == 0, - "blockDim must be a multiple of the gfx928 wavefront size"); - // Bank-safe LDS row stride: StageK + 8 bytes = 34 dwords. 8-aligned with - // an odd 8-byte count -> every lane's 8-byte fragment is ds_read_b64 at - // the 4-phase conflict-free floor (see the kernel comment above). - constexpr int kLdsStride = StageK + 8; - static_assert(kLdsStride % 8 == 0 && (kLdsStride / 8) % 2 == 1, - "LDS row stride must be 8-aligned with an odd 8-byte count"); - static_assert(kLdsStride % 16 == 8, - "LDS row stride must be 8 mod 16 for the two-half staging"); - - const int tid = static_cast(threadIdx.x); - const int wave = tid / kWaveSize; - const int lane = tid % kWaveSize; - const int wave_row = wave >> 1; // 0 or 1 - const int wave_col = wave & 1; // 0 or 1 - const int m0 = static_cast(blockIdx.y) * BM; - const int n0 = static_cast(blockIdx.x) * BN; - - __shared__ __align__(16) int8_t a_tile[BM * kLdsStride]; - __shared__ __align__(16) int8_t b_tile[BN * kLdsStride]; - - DUFragment - a_frag[kWaveM16]; - DUFragment - b_frag[kWaveN16]; - DUFragment - acc[kWaveM16][kWaveN16]; -#pragma unroll - for (int j = 0; j < kWaveM16; ++j) { -#pragma unroll - for (int i = 0; i < kWaveN16; ++i) { - du_fill_fragment(acc[j][i], 0); - } - } - - // Flat 16-byte staging vectors per stage (dispatch guarantees - // k % StageK == 0 and n % BN == 0). 8 vectors per 128-K row. - constexpr int kAVectors = BM * StageK / static_cast(sizeof(int4)); - constexpr int kBVectors = StageK * BN / static_cast(sizeof(int4)); - constexpr int kVecsPerRow = StageK / static_cast(sizeof(int4)); - - for (int k0 = 0; k0 < k; k0 += StageK) { - // Stage A[BM, StageK] row-major with row stride kLdsStride: thread - // `vec` owns the 16-byte chunk at (local_row, k16); rows past M are - // zero-filled and masked in the epilogue. - for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { - const int local_row = vec / kVecsPerRow; - const int k16 = vec - local_row * kVecsPerRow; - const int g_row = m0 + local_row; - int4 v{0, 0, 0, 0}; - if (g_row < m) { - v = *reinterpret_cast( - x_q + static_cast(g_row) * k + k0 + k16 * 16); - } - // Two 8-byte halves (stride % 16 == 8 -> odd rows are not 16-byte - // aligned, so a single int4 store is illegal; the int64 pair lowers - // to one ds_write2_b64 per thread). - const int64_t* src64 = reinterpret_cast(&v); - int64_t* dst64 = reinterpret_cast( - a_tile + local_row * kLdsStride + k16 * 16); - dst64[0] = src64[0]; - dst64[1] = src64[1]; - } - // Stage B[BN, StageK] n-major from the packed [N, K] weight: thread - // `vec` owns the 16-byte chunk at (n_row, k16); the K axis is always - // in range. - for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { - const int n_row = vec / kVecsPerRow; - const int k16 = vec - n_row * kVecsPerRow; - const int4 v = *reinterpret_cast( - packed_w + static_cast(n0 + n_row) * k + k0 + k16 * 16); - const int64_t* src64 = reinterpret_cast(&v); - int64_t* dst64 = reinterpret_cast( - b_tile + n_row * kLdsStride + k16 * 16); - dst64[0] = src64[0]; - dst64[1] = src64[1]; - } - // Make the staged tiles visible to every wavefront. - __syncthreads(); - - // Consume the stage: each wave accumulates its (BM/2) x (BN/2) quadrant - // with kk-ascending m16n16k32 MMAs; one ds_read_b64 per fragment. - const int local_row = wave_row * (BM / 2); - const int local_col = wave_col * (BN / 2); -#pragma unroll - for (int kk = 0; kk < StageK; kk += kTileK) { -#pragma unroll - for (int j = 0; j < kWaveM16; ++j) { - load_frag8( - a_frag[j], - a_tile + (local_row + j * kTileM) * kLdsStride + kk, - kLdsStride); - } -#pragma unroll - for (int i = 0; i < kWaveN16; ++i) { - load_frag8( - b_frag[i], - b_tile + (local_col + i * kTileN) * kLdsStride + kk, - kLdsStride); - } -#pragma unroll - for (int j = 0; j < kWaveM16; ++j) { -#pragma unroll - for (int i = 0; i < kWaveN16; ++i) { - du_mma_sync(acc[j][i], a_frag[j], b_frag[i], acc[j][i]); - } - } - } - // Single buffer: every wavefront must finish reading LDS before the - // next stage's stores overwrite it. - __syncthreads(); + int g_row) { + if (kFullTile) { + return *reinterpret_cast(p); } + return (g_row < m) ? *reinterpret_cast(p) : int4{0, 0, 0, 0}; +} - const int base_row = m0 + wave_row * (BM / 2); - const int base_col = n0 + wave_col * (BN / 2); -#pragma unroll - for (int j = 0; j < kWaveM16; ++j) { -#pragma unroll - for (int i = 0; i < kWaveN16; ++i) { - store_prefill_fragment_coalesced( - acc[j][i], x_scale, weight_scale, out, m, n, - base_row + j * kTileM, base_col + i * kTileN, lane); - } - } +// Exact-shape A-staging load through the GLOBAL (non-generic) address form: +// the caller passes the kernel's `x_q` base plus the per-stage byte offset, so +// the load visits exactly the same bytes as the pointer form but through the +// global address space (global_load_dwordx4) instead of the generic +// flat_load_dwordx4 path the pointer form lowers to. Used only by the +// exact-shape full-tile instantiation. +__device__ __forceinline__ int4 load_a_stage_vec_global( + const int8_t* __restrict__ x_q, + int64_t offset) { + return *reinterpret_cast(x_q + offset); } -// Iteration-3 double-buffered packed-B prefill kernel (dispatched only for -// (m >= 128, n == 2624, k == 6144)): SAME macro-tile as the accepted -// iteration-2 w8a8_dumma_prefill_packedb_kernel<128,64,128> (grid (N/64) x -// (M/128) = 41x32 = 1312 blocks, 256 threads = 4 wavefronts, 64x32 quadrant -// per wave, eight m16n16k32 int32 accumulators, same load_frag8 fragment -// reads, same bank-safe LDS row strides, same coalesced epilogue). The ONLY -// change is the K pipeline (see the file-header iteration-3 section); -// iteration 5 additionally reads the packed B operand from the B-panel -// layout (see the file-header iteration-5 section and -// w8a8_pack_panel_i8_kernel) -- the LDS staging layout, fragment consumes, -// and accumulation order are untouched: -// * The single-buffered 128-K stage (two __syncthreads; the exact-source -// ISA serializes global_load_dwordx4 -> s_waitcnt vmcnt(0) -> -// ds_write2_b64 -> s_barrier -> consume with no compute overlapping the -// VMEM round trip) becomes a 64-K DOUBLE-buffered stage with ONE -// __syncthreads per stage: the stage-(s+1) global loads are issued into -// 3 int4 staging VGPR per thread BEFORE the stage-s consume, and the -// prefetched vectors are flushed into the idle buffer after the consume, -// so the VMEM latency of stage s+1 overlaps the LDS/MMAC work of stage s -// (validated M=3072 prefetch order, in pure HIP: the vmcnt wait lands on -// the register data dependency right before the ds_write2_b64, i.e. -// after the v_mmac burst). -// * LDS = 2 x (A[128,72] + B[64,72]) = 2 x 13,824 = 27,648 B/block -// (26,112 -> 27,648 B, +5.9%); 2 x 27,648 = 55,296 <= 65,536 B keeps -// TWO resident blocks/CU = 8 waves/CU, occupancy unchanged. kLdsStride = -// 72 = StageK + 8 (8-aligned, odd 8-byte count 9 -> the 16 rows of each -// 16-lane ds_read_b64 phase land on 16 distinct bank pairs, 18r mod 32, -// r = 0..15 -> the same 4-phase conflict-free floor; 72 % 16 == 8 keeps -// the two-half ds_write2_b64 staging form). 12 extra VGPR for the -// staging registers (~72 -> ~84, still <= 128 for 2 blocks/CU). -// * k0-outer over 96 ASCENDING 64-K stages with kk-inner ascending over -// two m16n16k32 steps: the (k0+kk) k-chunk sequence 0, 32, 64, 96, ... -// is identical to iteration 2 (48 x 128-K stages), so every int32 -// accumulator is bit-identical (mismatch 0 / max_abs_error 0.0). -// * Barriers per block: 1 (prologue) + 95 (loop) = 96, the same as -// iteration 2's 2 x 48; the final stage is consumed without a trailing -// barrier. -template -__global__ __launch_bounds__(kThreadsPerBlock) void +// Double-buffered packed-B prefill kernel, dispatched ONLY for +// (m >= 128, n == 2624, k == 6144). Geometry: grid ((N+127)/128) x (M/256) = +// 21 x 16 = 336 blocks, 512 threads = 8 wavefronts, and each wave owns a +// kRowsPerWave x kColsPerWave = 64 x 64 quadrant of sixteen m16n16k32 int32 +// accumulator fragments (wave_row = wave / kWaveCols, wave_col = wave % +// kWaveCols, with the balanced 4 x 2 row/column split chosen by +// balanced_wave_rows), so the block macro-tile is BM x BN = 256 x 128. +// +// OPERAND REUSE (the mandated comparison). Per macro-tile every staged A +// element is consumed by the TWO wave columns (A reuse 2) and every staged B +// element by the FOUR wave rows (B reuse 4). Direct per-wave global fragment +// loads cannot replace the LDS stage: a row_major matrix_a fragment phase is +// lane (row = lane & 15, kq = lane >> 4), so 16 lanes touch 16 rows that are +// K = 6,144 B apart -- 16 x 32 B sectors for 128 useful bytes, a 4x fetch +// amplification -- and the swizzled col_major matrix_b fragments would lose +// their on-chip reuse and multiply their fetch. Global REUSE is the lever, so +// the M-major macro-tile is widened to BM = 256: A is read once per 128-column +// tile (21 re-reads) and B once per 256-row tile (16 re-reads), for a total +// global traffic of M*K*21 + N*K*16 = 786 MB against the 128 x 128 tile's +// 1,044 MB (-24.7%), at the same 8 waves/CU. The staging loads stay +// vectorized/coalesced 16-byte-per-thread and the A LDS stride stays the +// bank-safe 72 B with the flat swizzled B tile. The odd N = 20*128 + 64 is +// covered by clamping the column tile origin to n - BN, so the last block +// recomputes 64 overlapping columns (identical int32 accumulation and identical +// bf16 bytes) instead of predicating a partial tile; every staged B chunk +// therefore addresses a valid n16 group of the packed weight. +// +// The K pipeline is a 64-K DOUBLE-buffered stage: the stage-(s+1) global loads +// are issued into 2 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 runs 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, ... as any other staging, 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. +// +// The packed B operand is the fragment-interleaved swizzle (see +// w8a8_pack_bswz_i8_kernel): the LDS b_tile is the SAME [k32][n16][k8][n][8] +// layout as the pack (2 x (BN/16) 512-byte fragment blocks 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. +// +// kFullTile 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 and the epilogue drops the per-fragment row mask. +template +__global__ __launch_bounds__(kBlockThreads) 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 + 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) { - constexpr int kWaveM16 = BM / 32; - constexpr int kWaveN16 = BN / 32; - static_assert(BM % 32 == 0 && BN % 32 == 0, + constexpr int kWaves = kBlockThreads / kWaveSize; // 8 + // Balanced, tile-aware wave split: 4 wave rows x 2 wave columns for the + // 256 x 128 macro-tile, i.e. a square 64 x 64 per-wave quadrant. + constexpr int kWaveRows = + balanced_wave_rows(BM / kTileM, BN / kTileN, kWaves); // 4 + constexpr int kWaveCols = kWaves / kWaveRows; // 2 + constexpr int kRowsPerWave = BM / kWaveRows; // 64 + constexpr int kColsPerWave = BN / kWaveCols; // 64 + constexpr int kWaveM16 = kRowsPerWave / kTileM; // 4 + constexpr int kWaveN16 = kColsPerWave / kTileN; // 4 + static_assert(kWaveRows * kWaveCols == kWaves, + "wave split must tile the block exactly"); + static_assert(BM % kWaveRows == 0 && BN % kWaveCols == 0, "block tile must be a multiple of the wave quadrant"); - static_assert(kThreadsPerBlock % kWaveSize == 0, + static_assert(kBlockThreads % 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"); + // The packed B weight is the fragment-interleaved swizzle: one 32-k x 16-n + // fragment block per 512 B; a 64-K stage holds 2 k32 blocks x (BN/16) n16 + // groups x 4 k8 groups x 8 n-pairs = BN*4 16-byte chunks. With 512 threads + // and BN = 128 that is exactly one chunk per thread. + static_assert(StageK == 64, "swizzled B pack requires StageK == 64"); + static_assert(BN % 16 == 0, "swizzled B pack requires BN % 16 == 0"); const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; const int lane = tid % kWaveSize; - const int wave_row = wave >> 1; // 0 or 1 - const int wave_col = wave & 1; // 0 or 1 + const int wave_row = wave / kWaveCols; // 0 .. kWaveRows-1 + const int wave_col = wave % kWaveCols; // 0 .. kWaveCols-1 const int m0 = static_cast(blockIdx.y) * BM; - const int n0 = static_cast(blockIdx.x) * BN; + // Clamp the column origin so the odd (n % BN) columns are covered by an + // overlapping tile: every staged B chunk then addresses a valid n16 group of + // the packed weight, and the duplicated columns recompute byte-identically. + const int n0i = static_cast(blockIdx.x) * BN; + const int n0 = (n0i + BN <= n) ? n0i : (n - 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]; + // Swizzled B tile: flat [k32][n16][k8][n_pair][16 B], byte-identical to the + // packed global layout of that stage: 2 x (BN/16) x 512 = BN x 64 B. + __shared__ __align__(16) int8_t b_tile[2][BN * StageK]; DUFragment a_frag[kWaveM16]; @@ -787,55 +576,87 @@ w8a8_dumma_prefill_packedb_db_kernel( } } - // 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. + // Flat 16-byte staging vectors per stage. For the dispatched + // <256, 128, 64, 512> instantiation the 1024 A vectors (2 per thread) and the + // 512 B vectors (1 per thread) map exactly onto the 512 threads, so the + // prefetch payload is 3 int4 staging registers held live across the consume. constexpr int kVecsPerRow = StageK / static_cast(sizeof(int4)); constexpr int kAVectors = BM * StageK / static_cast(sizeof(int4)); constexpr int kBVectors = StageK * BN / static_cast(sizeof(int4)); - constexpr int kAVecPerThread = kAVectors / kThreadsPerBlock; - constexpr int kBVecPerThread = kBVectors / kThreadsPerBlock; - static_assert(kAVectors % kThreadsPerBlock == 0 && - kBVectors % kThreadsPerBlock == 0, + constexpr int kAVecPerThread = kAVectors / kBlockThreads; + constexpr int kBVecPerThread = kBVectors / kBlockThreads; + static_assert(kAVectors % kBlockThreads == 0 && + kBVectors % kBlockThreads == 0, "staging vectors must divide evenly across the block"); + // 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 == kBlockThreads, + "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. + // Per-thread GLOBAL staging byte pointers, held as loop-carried induction + // variables: the prologue loads stage 0 from them, they are then advanced to + // stage 1, and every steady-state iteration issues its prefetch from them and + // walks them by the compile-time stage stride -- StageK bytes for A and one + // 64-K swizzle band `(StageK>>5)*n16g*512` for B. + const int8_t* a_src[kAVecPerThread]; + int a_grow[kAVecPerThread]; +#pragma unroll + for (int v = 0; v < kAVecPerThread; ++v) { + const int vec = tid + v * kBlockThreads; + const int local_row = vec / kVecsPerRow; + const int k16 = vec - local_row * kVecsPerRow; + a_grow[v] = m0 + local_row; + a_src[v] = x_q + static_cast(a_grow[v]) * k + k16 * 16; + } + const int8_t* b_src[kBVecPerThread]; +#pragma unroll + for (int v = 0; v < kBVecPerThread; ++v) { + // Chunk (k32, n16, k8, n_pair) of the stage: BN*2 chunks per k32 block + // (4 k8 groups x 8 n-pairs per n16), 32 chunks per n16 group. + const int vec = tid + v * kBlockThreads; + const int k32 = vec / (BN * 2); + const int rem = vec - k32 * (BN * 2); + b_src[v] = packed_w + + packed_b_swz_offset( + k32, (n0 >> 4) + ((rem >> 5) & (BN / 16 - 1)), n16g, + (rem >> 3) & 3, rem & 7); + } + + // ---- Prologue: issue the stage-0 global loads, flush them into buffer 0 and + // make the tile visible to every wavefront. { - const int k0 = 0; #pragma unroll for (int v = 0; v < kAVecPerThread; ++v) { - const int vec = tid + v * kThreadsPerBlock; + const int vec = tid + v * kBlockThreads; 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}; + if (kFullTile) { + // Exact-shape full-tile A staging through the global (non-generic) + // address form -- the same bytes as the pointer form, but through the + // load path that emits global_load_dwordx4. + a_pre[v] = load_a_stage_vec_global( + x_q, static_cast(g_row) * k + k16 * 16); + } else { + a_pre[v] = load_a_stage_vec(a_src[v], m, a_grow[v]); + } } #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); + // The 16-byte chunk (k32 = vec/(BN*2), n16 = (vec>>5)&(BN/16-1), + // 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; the + // per-thread pointer already carries the k32/n16/k8/n_pair part for + // k0 == 0. 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(b_src[v]); } #pragma unroll for (int v = 0; v < kAVecPerThread; ++v) { - const int vec = tid + v * kThreadsPerBlock; + const int vec = tid + v * kBlockThreads; const int local_row = vec / kVecsPerRow; const int k16 = vec - local_row * kVecsPerRow; const int64_t* src64 = reinterpret_cast(&a_pre[v]); @@ -846,118 +667,157 @@ 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]; + const int vec = tid + v * kBlockThreads; + // 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. - for (int s = 0; s < kStages - 1; ++s) { - // Prefetch stage s+1: issue the global loads (kept in VGPR while the - // stage-s consume runs). - const int k0 = (s + 1) * StageK; + // Advance the staging pointers to stage 1. From here on every steady-state + // half prefetches stage s+1 from these pointers and then walks them by the + // compile-time stage stride, so both global staging addresses are pure + // pointer increments instead of a per-stage address expression. + const int64_t b_step = static_cast(StageK >> 5) * n16g * 512; #pragma unroll - for (int v = 0; v < kAVecPerThread; ++v) { - const int vec = tid + v * kThreadsPerBlock; - const int local_row = vec / kVecsPerRow; - const int k16 = vec - local_row * kVecsPerRow; - const int g_row = m0 + local_row; - a_pre[v] = (g_row < m) - ? *reinterpret_cast( - x_q + static_cast(g_row) * k + k0 + - k16 * 16) - : int4{0, 0, 0, 0}; - } + for (int v = 0; v < kAVecPerThread; ++v) { + a_src[v] += StageK; + } #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); - } + for (int v = 0; v < kBVecPerThread; ++v) { + b_src[v] += b_step; + } - // 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); + // ---- Steady state (double-buffer parity specialization). 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. + // + // Walking the stages in compile-time PAIRS makes each half's parity a + // constant (s = 2*sp + h with h unrolled, so s & 1 == h and (s+1) & 1 == h^1), + // which folds the double-buffer parity offsets into the per-half LDS base + // registers / instruction immediates. The odd leftover for the dispatched + // k = 6144 (kStages = 96 -> kSteady = 95, one unpaired stage) is guarded by ONE + // uniform compare in the second half of the last pair, so correctness does not + // depend on the unroll succeeding. + const int kSteady = kStages - 1; + const int kPairIters = (kSteady + 1) >> 1; + for (int sp = 0; sp < kPairIters; ++sp) { #pragma unroll - for (int kk = 0; kk < StageK; kk += kTileK) { + for (int h = 0; h < 2; ++h) { + const int s = 2 * sp + h; + if (h == 0 || s < kSteady) { + // Compile-time buffer parity: cur == s & 1, nxt == (s + 1) & 1. + const int cur = h; + const int nxt = h ^ 1; + // Prefetch stage s+1: issue the global loads (held in VGPRs while the + // stage-s consume runs) from the loop-carried per-thread stage + // pointers, which already sit at the stage-(s+1) addresses. The + // exact-shape full-tile instantiation forms the two A staging addresses + // in the `x_q + ` global form (so the backend + // keeps the non-generic global_load path), while B keeps the + // loop-carried staging pointer. + const int k0 = (s + 1) * StageK; #pragma unroll - for (int j = 0; j < kWaveM16; ++j) { - load_frag8( - a_frag[j], - a_tile[s & 1] + (local_row + j * kTileM) * kLdsStride + kk, - kLdsStride); - } + for (int v = 0; v < kAVecPerThread; ++v) { + const int vec = tid + v * kBlockThreads; + const int local_row = vec / kVecsPerRow; + const int k16 = vec - local_row * kVecsPerRow; + const int g_row = m0 + local_row; + if (kFullTile) { + a_pre[v] = load_a_stage_vec_global( + x_q, static_cast(g_row) * k + k0 + k16 * 16); + } else { + a_pre[v] = load_a_stage_vec(a_src[v], m, a_grow[v]); + } + } #pragma unroll - for (int i = 0; i < kWaveN16; ++i) { - load_frag8( - b_frag[i], - b_tile[s & 1] + (local_col + i * kTileN) * kLdsStride + kk, - kLdsStride); - } + for (int v = 0; v < kBVecPerThread; ++v) { + // One plain 16-byte staging load per thread from the stage-(s+1) + // swizzle chunk; 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(b_src[v]); + } + + // Consume stage s from buffer cur: each wave accumulates its + // kRowsPerWave x kColsPerWave quadrant with kk-ascending m16n16k32 MMAs; + // one ds_read_b64 per fragment. + const int local_row = wave_row * kRowsPerWave; #pragma unroll - for (int j = 0; j < kWaveM16; ++j) { + for (int kk = 0; kk < StageK; kk += kTileK) { #pragma unroll - for (int i = 0; i < kWaveN16; ++i) { - du_mma_sync(acc[j][i], a_frag[j], b_frag[i], acc[j][i]); + for (int j = 0; j < kWaveM16; ++j) { + load_frag8( + a_frag[j], + a_tile[cur] + (local_row + j * kTileM) * kLdsStride + kk, + static_cast(kLdsStride), + lane); + } +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + // 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[cur] + + ((kk >> 5) * (BN / 16) + wave_col * kWaveN16 + i) * 512, + lane); + } +#pragma unroll + for (int j = 0; j < kWaveM16; ++j) { +#pragma unroll + for (int i = 0; i < kWaveN16; ++i) { + du_mma_sync(acc[j][i], a_frag[j], b_frag[i], acc[j][i]); + } + } } - } - } - // Flush the prefetched stage-(s+1) vectors into the idle buffer. - { - const int buf = (s + 1) & 1; + // Flush the prefetched stage-(s+1) vectors into the idle buffer nxt. + { #pragma unroll - for (int v = 0; v < kAVecPerThread; ++v) { - const int vec = tid + v * kThreadsPerBlock; - const int local_row = vec / kVecsPerRow; - const int k16 = vec - local_row * kVecsPerRow; - const int64_t* src64 = reinterpret_cast(&a_pre[v]); - int64_t* dst64 = reinterpret_cast( - a_tile[buf] + local_row * kLdsStride + k16 * 16); - dst64[0] = src64[0]; - dst64[1] = src64[1]; - } + for (int v = 0; v < kAVecPerThread; ++v) { + const int vec = tid + v * kBlockThreads; + const int local_row = vec / kVecsPerRow; + const int k16 = vec - local_row * kVecsPerRow; + const int64_t* src64 = reinterpret_cast(&a_pre[v]); + int64_t* dst64 = reinterpret_cast( + a_tile[nxt] + local_row * kLdsStride + k16 * 16); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } +#pragma unroll + for (int v = 0; v < kBVecPerThread; ++v) { + const int vec = tid + v * kBlockThreads; + // Flat 16-byte swizzled-tile commit (ds_write_b128). + *reinterpret_cast(b_tile[nxt] + vec * 16) = b_pre[v]; + } + } + // Walk the loop-carried staging pointers to the next stage. StageK is a + // compile-time constant, so each A pointer advances with a single 64-bit + // add pair; b_step is the loop-invariant swizzle band. #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]; + for (int v = 0; v < kAVecPerThread; ++v) { + a_src[v] += StageK; + } +#pragma unroll + for (int v = 0; v < kBVecPerThread; ++v) { + b_src[v] += b_step; + } + __syncthreads(); } } - __syncthreads(); } // ---- Tail: consume the last stage (its data was flushed before the last // barrier; no trailing barrier needed). { const int s = kStages - 1; - const int local_row = wave_row * (BM / 2); - const int local_col = wave_col * (BN / 2); + const int local_row = wave_row * kRowsPerWave; #pragma unroll for (int kk = 0; kk < StageK; kk += kTileK) { #pragma unroll @@ -965,14 +825,16 @@ 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( + 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) { @@ -984,13 +846,15 @@ w8a8_dumma_prefill_packedb_db_kernel( } } - const int base_row = m0 + wave_row * (BM / 2); - const int base_col = n0 + wave_col * (BN / 2); + const int base_row = m0 + wave_row * kRowsPerWave; + const int base_col = n0 + wave_col * kColsPerWave; #pragma unroll for (int j = 0; j < kWaveM16; ++j) { #pragma unroll for (int i = 0; i < kWaveN16; ++i) { - store_prefill_fragment_coalesced( + // kGuarded == !kFullTile -- the full-tile instantiation skips the dead + // per-fragment `row >= m` EXEC bracket. + store_prefill_fragment_coalesced( acc[j][i], x_scale, weight_scale, out, m, n, base_row + j * kTileM, base_col + i * kTileN, lane); } @@ -998,24 +862,33 @@ 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. +// cases (including the paired M=2 and M=16 shapes with the same (N, K) -- 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. +// +// For the exact (k, n) == (6144, 2624) the packed_weight buffer holds the +// fragment-interleaved swizzle (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): +// 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, so the paired M=2 / M=16 API shapes stay byte-exact. __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 +896,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) { + // 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,61 +937,49 @@ __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.) -__global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( - const int8_t* __restrict__ raw, - int8_t* __restrict__ packed, - int k, - int n) { - const int64_t total = static_cast(k) * n; - const int64_t stride = static_cast(gridDim.x) * blockDim.x; - for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + - threadIdx.x; - idx < total; - idx += stride) { - const int col = static_cast(idx / k); - const int kk = static_cast(idx - static_cast(col) * k); - packed[idx] = raw[static_cast(kk) * n + col]; - } -} - -// Exact-shape B-panel pack (iteration 5): for (k, n) == (6144, 2624) the +// 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 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( +// 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 +// pack, 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). +__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]; + } } } @@ -1148,8 +1000,10 @@ __global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( } // namespace // Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the -// caller-provided stream only; never allocates, synchronizes, packs, or -// touches the default stream. +// caller-provided stream only; never allocates, synchronizes, packs, or touches +// the default stream. The workspace is unused for this shape (no split-K): +// m*n*4 = 42,990,080 B exceeds the 16 MiB workspace budget, so +// workspace_split_k_capacity(m, n, k) is zero. extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -1170,96 +1024,60 @@ 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. M = 16*256 is + // exact and N = 2624 = 20*128 + 64 is covered by clamping the last column + // tile origin to n - 128, so every block still stages a full 128-column B + // macro-tile and no B chunk ever leaves the packed weight. The compile-time + // kFullTile specialization is selected when m is a multiple of BM = 256, and + // every other m with this (n, k) keeps the guarded instantiation. 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); + static_cast((n + kPackedBN - 1) / kPackedBN), + static_cast((m + kPackedBM - 1) / kPackedBM)); + const dim3 block(static_cast(kPackedThreads)); + if ((m % kPackedBM) == 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_packedb_db_kernel< + kPackedBM, kPackedBN, 64, kPackedThreads, 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< + kPackedBM, kPackedBN, 64, kPackedThreads, 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 +1100,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 +1129,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. 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 +1148,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; + // 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 +1159,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 +1207,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/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/o_proj.hip index 49710bb2..2a553a70 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/o_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/o_proj.hip @@ -1,76 +1,147 @@ -// @@variant shape=glm_tp8_o_proj_m4096 commit=486f3ec14b6d88c045aa35a5c85d1d44816c8c2d added=2026-08-31 -// median_us=759.9 p90_us=760.8 speedup=72.08 baseline_us=5.478e+04 -// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 -// @@variant shape=glm_tp8_o_proj_m4096 +// @@variant shape=glm_tp8_o_proj_m4096 commit=41d2ef9da654ecb15154a974fa11d84df5e4fc5d added=2026-09-14 +// median_us=651.5 p90_us=654.2 speedup=84.07 baseline_us=5.478e+04 +// source=test-9-11-1/iteration_002/glm_tp8_o_proj_m4096 // INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. // -// Worker: worker_2 (physical GPU 2), assigned shape: +// Worker: worker_0 (physical GPU 0), assigned shape: // glm_tp8_o_proj_m4096 : M=4096, N=6144, K=2048 // -// Bootstrap strategy (iteration 1, correctness-first but a usable profiling -// baseline - a large-Prefill scalar K loop would not be): -// * Large-prefill path (exact (m, n, k) == (4096, 6144, 2048)): -// native INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 -// output tile per block; four wavefronts (256 threads); each wave owns -// a 64x32 quadrant built from eight m16n16k32 accumulator fragments; -// the block cooperatively vector-loads A[128,64] and B[64,64] into one -// single-buffered 64-K LDS stage (15,360 B total: 128*80 + 64*80 with -// 16 B padding per row for bank skew); two __syncthreads per stage; -// fused dot * x_scale[m] * weight_scale[n] epilogue stored directly as -// bf16 from the accumulator fragments. Grid dim3(96, 32) = 3072 blocks -// dwarfs the 120 CUs, so no split-K is needed. -// The weight is read in the identity [K, N] int8 layout (bootstrap: -// launch_pack_w8a8_weight is an identity device-to-device copy for -// every (k, n)). A is staged m-major into a_tile[m][k] and its eight -// fragment k-values per lane are contiguous, so A fragments use the -// direct 8-byte LDS fill (load_fragment8, the exact pattern validated -// on the worker-29 TP4 gate_up lineage iteration 13 and the TP8 hy3 -// o_proj lineage iteration 12). B is staged k-major into -// b_tile[k][n] and its fragments use the library du_load_matrix_sync -// row_major loader (8 per-byte LDS reads strided by ldm per lane - -// correct, direct, and simple; the packed n-major layout + 8-byte B -// fragment fill is a later-round optimization once the packed layout -// is introduced). -// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), -// including all small-M API cases (M=2, M=16) and M=3072 with the same -// (K, N) = (2048, 6144); the fallback reads the identity [K, N] layout -// for every (k, n). -// * launch_pack_w8a8_weight: identity device-to-device copy of the raw -// [K, N] weight and the [N] weight_scale for every (k, n) (bootstrap; -// later rounds may change only the pack HIP implementation and the -// matching GEMM interpretation). Packing never happens inside the -// timed GEMM. -// Iteration 5 (packing round): for the exact (k, n) == (2048, 6144) -// launch_pack_w8a8_weight now produces the n-major [N, K] layout -// packed[n*K + kk] = raw[kk*N + n] (same byte count and buffer -> -// graph-stable), and the exact-shape dispatch uses the new -// w8a8_dumma_256x64x64_packedb_kernel, which stages B n-major into -// bank-safe 72-byte LDS rows and fills each col_major B fragment with -// one contiguous 8-byte ds_read_b64 (zero per-byte reads, zero -// reassembly VALU); the scalar fallback decodes the n-major layout -// when (k, n) matches so the paired M=2/M=16 and M=3072 API shapes -// stay byte-exact. Every other (k, n) keeps the identity copy. -// Iteration 6 (epilogue round): the fused scale->bf16->coalesced -// store epilogue of w8a8_dumma_256x64x64_packedb_kernel is -// restructured to register-batch the per-row and per-column scales: -// each lane loads its 8 x_scale rows and its 2 weight_scale float4 -// columns ONCE (issued on the last K stage, before the final -// protective barrier) and the 16 per-fragment coalesced stores reuse -// them; the previous per-call epilogue re-loaded the scales inside -// each store (PMC vmem_read surplus ~153,600 wavefront loads above -// the pure staging floor). Output bits are identical (same transpose, -// same float(dot) * x_scale[row] * weight_scale[col] order, same bf16 -// rounding). No workspace/combine pass exists: the workspace -// argument stays unused and the epilogue is fully fused in-kernel. +// Round 1 (establish_arch) established ONE complete, correct DUMMA +// architecture family and measured median 753.83 us: +// +// * family: packed_weight + staged A/B, 256x64 output macro-tile, native +// INT8 DUMMA m16n16k32 with int32 accumulation, four 64-lane wavefronts +// (256 threads); +// * weight packing (outside the timed region, before Graph capture, via the +// ABI-preserving launch_pack_w8a8_weight): for the exact (k, n) == +// (2048, 6144) the raw logical [K, N] int8 weight is transposed once into +// the n-major [N, K] layout packed[col * K + kk] = raw[kk * N + col]. +// Same byte count and same caller-owned buffer as the identity layout, so +// every captured address is graph-stable. Every other (k, n) keeps the +// identity device-to-device copy. +// * fragment fills: the m16n16k32 matrix_a slot for lane (r = lane & 15, +// kq = lane >> 4) is eight consecutive k bytes at fixed row, and the +// matrix_b (col_major) slot is eight consecutive k bytes at fixed n on the +// n-major tile, so BOTH operands use one contiguous 8-byte ds_read_b64 per +// fragment (load_fragment8) with no per-byte reads and no mask/OR +// reassembly VALU. +// * LDS bank safety: both padded strides are odd multiples of 8 bytes +// (A = 88 = 22 dwords, B = 72 = 18 dwords). 22*r mod 32 and 18*r mod 32 +// over r = 0..15 are permutations of the 16 even dword starts, so each +// 16-lane ds_read_b64 fragment phase lands on 16 distinct bank pairs. +// Both strides are 8 mod 16, so the 16-byte cooperative staging stores are +// emitted as two int64 halves (the validated ds_write2_b64 pattern). +// * fused coalesced epilogue: the 4x4 shfl_xor register transpose keeps the +// exact per-element multiply order float(dot) * x_scale[row] * +// weight_scale[col] and bf16 rounding, then one 8-byte store per lane per +// fragment (100% store-sector efficiency). Each lane's x_scale rows and +// weight_scale float4 columns are loaded ONCE into registers on the last +// K stage, issued before the final protective barrier, so the epilogue +// reads no scale inside the store chains. The epilogue is fully fused: +// no workspace pass, no separate combine kernel. +// +// Round 2 (mandatory decision: operand-reuse round - compare direct loads +// with cooperative A+B LDS staging, quantify A/B reuse per macro-tile, use +// vectorized coalesced global loads with a bank-safe LDS layout): +// +// FALSIFIABLE BOTTLENECK HYPOTHESIS. The round-1 kernel is bound by the +// L2/global operand traffic its macro-tile reuse ratio forces, not by the +// LDS fragment pipe, not by MMAC issue, and not by HBM. With a +// (tileM, tileN) = (256, 64) macro-tile every A byte is re-read N/tileN = +// 96x and every B byte M/tileM = 16x by the 1,536-block grid, i.e. +// A 805.3 MB + B 201.3 MB = 1,006.6 MB of global reads per operator replay. +// The exact-source PMC record (digest acb11003..., grid 1,536 x 256 +// threads, 151 VGPR, 27,136 B LDS, 957,746 grbm_gui_active cycles at +// ~1.29 GHz) shows only 99.3 MB of HBM reads and 50.3 MB of HBM writes +// (198.5 GB/s counter-derived HBM) against 1,044,480 vmem_read_instructions +// of 16-B-per-lane coalesced staging, i.e. ~1.07 GB of L2 read traffic in +// 740.5 us (~1.4 TB/s). HBM is far from its ceiling, so the traffic that +// matters is the L2-resident operand re-read, and the only structural +// lever on it is the macro-tile reuse ratio 1/tileM + 1/tileN. +// +// Direct fragment loads vs cooperative A+B LDS staging. Direct global +// fragment fills were rejected without measuring them: the m16n16k32 +// fragment owns the k dimension and lane index, so lane (r, kq) needs +// 8 consecutive k bytes at one row; in a 64-lane wave that is 16 different +// rows 2,048 B apart per 16-lane group (8 B useful out of a 32-B sector), +// and each wave reads the same A fragment once per n-group. Per block that +// is >= 4x the L2 sectors of the 16-B-per-thread coalesced cooperative +// staging (which touches every operand byte exactly once per block), before +// counting the duplicated A reads. Cooperative A+B staging therefore +// stays, and this round moves the reuse lever: DOUBLE THE N MACRO-TILE. +// +// Change (ONE bounded factor: the macro-tile reuse geometry). The output +// macro-tile becomes 256x128 at the same 256 threads / 4 wavefronts; the +// 2(m) x 2(n) wave split gives every wave a 128x64 quadrant = 8 m-groups x +// 4 n-groups = 32 m16n16k32 int32 accumulators (the A quadrant height and +// the 88-byte A stride are unchanged, so the eight A fragment fills per kk +// keep the exact round-1 addresses; the four B fragment fills per kk are the +// only new ones). Global staging stays per-thread 16-byte coalesced int4 +// loads, 6 per thread per stage (A rows [0,64),[64,128),[128,192),[192,256) +// at k-chunk (tid&3)*16 -- unchanged from round 1 -- plus TWO B bands now: +// n = b_row0 and b_row0 + 64 at the same k-chunk). The single-buffered +// 64-K LDS stage becomes A[256,88] = 22,528 B + B[128,72] = 9,216 B = +// 31,744 B/block, both strides unchanged and still bank-safe, and the two +// int64-half ds_write2_b64 staging stores are unchanged. +// +// Quantified reuse per macro-tile (bytes of global operand read per +// operator replay, M*K*(N/tileN) + N*K*(M/tileM)): +// (256, 64) -> A 805.3 MB + B 201.3 MB = 1,006.6 MB (round 1) +// (256, 128) -> A 402.7 MB + B 201.3 MB = 604.0 MB (this round, -40%) +// (128, 128) -> A 402.7 MB + B 402.7 MB = 805.3 MB (rejected: the +// halved tileM double-charges B, only -20% total) +// (512, 64) -> A 805.3 MB + B 100.7 MB = 906.0 MB (rejected: the +// unchanged tileN leaves the dominant A term untouched) +// LDS fragment reads per MMAC also fall from (8+2)/16 = 0.625 to +// (8+4)/32 = 0.375 per kk, and the block count falls 1,536 -> 768 (still +// 6.4 blocks/CU on the 120-CU part, so split-K remains unnecessary and the +// caller workspace stays unused; the trusted occupancy-probe split +// candidates [2,3,4,5] are irrelevant to a grid this wide). +// +// Expected effect: if the L2 operand-traffic term is the limiter, median +// moves from ~754 us toward ~500-620 us; if the kernel is instead at the +// MMAC/LDS-issue floor, the larger tile stays within a few percent (the +// extra 64 accumulator VGPRs and the 768-block tail are the only costs). +// +// Round 8 (mandatory decision: fix_build - repair the previous round's +// compile/API failure with the smallest source-owned change; no architecture +// change, no redesign). Rounds 6 and 7 produced no built candidate because the +// worker process was terminated by an infrastructure rate limit, and the +// workspace source was restored to the accepted round-2 kernel. Round 6's +// intended epilogue change is re-applied with the compile-safety correction +// demanded by this round. Preserved byte-for-byte: the 256x128 macro-tile, +// 256 threads / 4 wavefronts, 88-byte A and 72-byte B padded strides, the +// single-buffered 64-K LDS stage (A 22,528 B + B 9,216 B = 31,744 B), the +// one-ds_read_b64-per-fragment fills, the 32 int32 m16n16k32 accumulators and +// their exact accumulation order, the two __syncthreads per stage, the exact +// (m,n,k) dispatch guard, the packed n-major weight path and the scalar +// generic fallback. The only edit is the epilogue conversion. hip's +// __float2bfloat16 lowers on this DTK to a NaN-guarded exec-masked sequence +// per element (~12 extra SALU/VALU plus a divergent branch region for each of +// the 128 conversions per lane); it is replaced by the branchless +// finite-domain RNE helper bf16_rne_finite_bits, which stores the identical +// 16-bit payload for every finite value this epilogue can produce and can +// differ only for NaN/Inf. The 4x4 shfl_xor transpose, the per-element +// multiply order float(dot) * x_scale[row] * weight_scale[col], the tail-row +// guard, the 8-byte coalesced store and the dispatch are unchanged. +// +// 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): it always +// reproduces the exact int32 dot over K in ascending k order, then the fused +// float scale and bf16 store. For (k, n) == (2048, 6144) it decodes the +// n-major packed layout produced above; every other (k, n) reads the identity +// [K, N] layout. // // Rules honored: headers in the known-good order (hip/hip_runtime.h, -// hip/hip_bfloat16.h, du_mma.h); wavefront = 64 and blockDim a multiple of -// 64; int32 accumulation with float conversion only for +// hip/hip_bfloat16.h, du_mma.h); wavefront = 64 and blockDim a multiple of 64; +// int32 accumulation with float conversion only for // dot * x_scale[m] * weight_scale[n]; hipLaunchKernelGGL on the -// caller-provided PyTorch stream; no stream-0 / synchronization / temporary -// allocation in gemm_out; no NVIDIA wmma/mma.sync/PTX/warp-32 masks; all -// barriers on paths reached by every thread in the block; explicit dispatch -// for the assigned shape with a scalar generic fallback for everything else. +// caller-provided PyTorch stream; no allocation / packing / synchronization / +// default-stream launch inside the timed API; no NVIDIA wmma/mma.sync/PTX/ +// FP8/INT4 and no raw inline asm; all __syncthreads() on paths reached by +// every block thread; coalesced lanes follow contiguous N addresses; single +// 31,744 B LDS buffer (no double buffering); explicit exact-shape dispatch for +// the assigned shape with a preserved scalar generic fallback. #include #include @@ -89,62 +160,37 @@ constexpr int kWaveSize = 64; // gfx928 native wavefront constexpr int kTargetM = 4096; constexpr int kTargetN = 6144; constexpr int kTargetK = 2048; -constexpr int kBlockM = 128; -constexpr int kBlockN = 64; -constexpr int kStageK = 64; -constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) -constexpr int kBStride = kBlockN + kBPad; -constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride -constexpr int kBlockThreads = 4 * kWaveSize; - -// Iteration-4 tile-shape constants: 256x64 output tile (M-major doubling of -// the accepted 128x64 tile) at 256 threads = 4 wavefronts, each wave owning a -// 128x32 quadrant (16 m16n16k32 int32 accumulators). -constexpr int kBlockM256 = 2 * kBlockM; - -// Iteration-5 packed-B constants: for the exact (k, n) == (2048, 6144) the -// weight is packed once (outside the timed region) into the n-major [N, K] -// layout packed[n*K + kk] = raw[kk*N + n]; B is then staged n-major into LDS -// rows of 72 B (64 + 8). 72 is an odd 8-byte count, so the 16 lanes of each -// ds_read_b64 fragment-load phase land on 16 distinct bank pairs (18r mod 32, -// r = 0..15 -> all 16 even dword starts: the conflict-free 4-phase floor), -// and 72 % 16 == 8, so staging stores use two int64 halves (ds_write2_b64) -// instead of one misaligned int4. -constexpr int kBPadPacked = 8; // 64 -> 72-byte n-major B LDS row stride -constexpr int kBStridePacked = kStageK + kBPadPacked; -// Iteration-9 (A-stride round) constant: packedb A LDS row stride 64 -> 88 -// bytes, replacing the shared 80-byte kAStride for w8a8_dumma_256x64x64_ -// packedb_kernel only (the identity-layout kernels keep kAStride = 80). -// 88 = 22 dwords: 22r mod 32 is a permutation of the 16 even dword starts -// over r = 0..15, so the 16 lanes of each ds_read_b64 fragment-load phase -// land on 16 distinct bank pairs (the same conflict-free 4-phase floor the -// packed-B stride 72 already uses); the current 80-byte (20-dword) stride -// has 20r mod 32 with period 8, so rows r and r+8 alias onto one bank phase -// and EVERY A fragment read pays a 2-way bank conflict (PMC -// lds_bank_conflicts 20,447,232, deterministic across iterations 5..8). -// 88 % 16 == 8, so the 16-byte staging stores go through two int64 halves -// (ds_write2_b64, the exact accepted packed-B staging pattern) - the same -// instruction count as the current ds_write_b128. A[256,88] = 22,528 B + -// B[64,72] = 4,608 B = 27,136 B/block; 2 blocks/CU = 54,272 <= 65,536, so -// co-residency, VGPR, barrier count and the global-load pattern are all -// unchanged; only the physical LDS addresses move (same bytes stored, same -// bytes read, fragment values bit-identical). -constexpr int kAStridePacked = kStageK + 24; // 64 -> 88-byte A LDS row stride +// Round 2: 256x128 macro-tile, four wavefronts, single-buffered 64-K stage. +// Doubling the N macro-tile halves the A re-read term M*K*(N/tileN) while the +// unchanged tileM keeps the B term N*K*(M/tileM); total global operand reads +// fall 1,006.6 -> 604.0 MB per replay (-40%). +constexpr int kBlockM = 256; +constexpr int kBlockN = 128; +constexpr int kStageK = 64; +constexpr int kBlockThreads = 4 * kWaveSize; // 256 + +// A[256, 88]: m-major, 88 = 22 dwords (odd 8-byte count). 22*r mod 32 over +// r = 0..15 is a permutation of the 16 even dword starts, so every 16-lane +// ds_read_b64 A fragment phase is conflict-free; 88 % 16 == 8 forces the +// two-int64-half (ds_write2_b64) staging store. +constexpr int kAStride = kStageK + 24; // 64 -> 88-byte A LDS row stride +// B[128, 72]: n-major (row index = n), 72 = 18 dwords. 18*r mod 32 over +// r = 0..15 is the same conflict-free permutation; 72 % 16 == 8 gives the same +// two-int64-half staging store. +constexpr int kBStride = kStageK + 8; // 64 -> 72-byte B LDS row stride using bf16_t = hip_bfloat16; -// Direct 8-byte LDS fragment fill for matrix_a on the m-major A tile. -// The m16n16k32 matrix_a row_major loader assigns x[i] = -// p[row*ldm + col_group*8 + i] (row = lane & 15, col_group = lane >> 4), -// i.e. eight consecutive k-values per lane; the library path makes the -// compiler emit per-byte ds_read_u8 loads plus a mask/OR reassembly chain -// per dword that is an arithmetic identity on the v_mmac A operand. Writing -// the same 8 bytes straight into the fragment storage keeps the operand bit -// pattern and the int32 accumulation identical and lets the compiler feed a -// single 8-byte LDS read to the v_mmac. This is the exact pattern validated -// on the worker-29 TP4 gate_up lineage (accepted iteration 13) and the TP8 -// hy3 o_proj lineage (accepted iteration 12). +// Direct 8-byte LDS fragment fill. The m16n16k32 fragment layout assigns +// x[i] = p[(row++) * ldm + col_group*8 + i] (row = lane & 15, +// col_group = lane >> 4), i.e. eight consecutive 8-bit elements per lane on a +// row-major tile. On the m-major A tile these are eight consecutive k bytes of +// one row; on the n-major B tile (row = n) they are eight consecutive k bytes +// of one column. Writing those bytes straight into the fragment storage keeps +// the operand bit pattern, the v_mmac inputs and the int32 accumulation order +// identical to the library loader while emitting one ds_read_b64 instead of +// eight per-byte ds_read_u8 plus a mask/OR reassembly chain per dword. template __device__ __forceinline__ void load_fragment8( Frag& frag, @@ -156,103 +202,46 @@ __device__ __forceinline__ void load_fragment8( *reinterpret_cast(base + off); } -// Coalesced fragment store. The m16n16k32 accumulator lane mapping -// (row = lane & 15, column group c4 = lane >> 4, frag.x[i] -> column -// c4 + 4*i) gives each lane four elements strided by 4 columns, so a direct -// store would be four 2-byte scalar stores per lane whose wavefront -// addresses touch each 32-B sector at 25% utilization. This epilogue -// transposes the 4-element groups within each 4-lane column group (lanes r, -// r+16, r+32, r+48 -- a 4x4 transpose, two 2x2 steps with shfl_xor 16 then -// 32), so lane (r, c4) ends up holding the four CONTIGUOUS columns -// 4*c4 .. 4*c4+3, converts them to bf16, packs 4 bf16 (8 B), and writes ONE -// 8-byte store per lane (100% store sector efficiency). This DTK lowers -// __shfl_xor to ds_bpermute at the block tail where the LDS pipe is idle. -// Only the int32 values are re-routed between lanes; the per-element scale -// multiply order (float(dot) * x_scale[row] * weight_scale[col]) and the -// bf16 rounding are unchanged, so the stored bits are identical to the -// direct per-element store. The row >= m guard is wavefront-uniform (all -// 64 lanes of a wave share the same 16-row window), so the shuffles never -// mix active and inactive lanes; base_col is a multiple of 16 and n*2 a -// multiple of 8, so the float4 weight_scale load and the 8-byte store are -// aligned. -template -__device__ __forceinline__ void store_prefill_fragment_coalesced( - const AccFragment& frag, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - bf16_t* __restrict__ out, - int base_row, - int base_col, - int m, - int n, - int lane) { - const int row = base_row + (lane & 15); - if (row >= m) { - return; // tail-M masking: padded rows never write (wavefront-uniform) - } - const int c4 = lane >> 4; // 0..3 - const int x0 = frag.x[0]; - const int x1 = frag.x[1]; - const int x2 = frag.x[2]; - const int x3 = frag.x[3]; - // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). - const int s0 = __shfl_xor(x0, 16, kWaveSize); - const int s1 = __shfl_xor(x1, 16, kWaveSize); - const int s2 = __shfl_xor(x2, 16, kWaveSize); - const int s3 = __shfl_xor(x3, 16, kWaveSize); - const bool lo = (c4 & 1) == 0; - const int a0 = lo ? x0 : s1; - const int a1 = lo ? s0 : x1; - const int a2 = lo ? x2 : s3; - const int a3 = lo ? s2 : x3; - // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). - const int t0 = __shfl_xor(a0, 32, kWaveSize); - const int t1 = __shfl_xor(a1, 32, kWaveSize); - const int t2 = __shfl_xor(a2, 32, kWaveSize); - const int t3 = __shfl_xor(a3, 32, kWaveSize); - const bool hi = (c4 & 2) == 0; - const int f0 = hi ? a0 : t2; - const int f1 = hi ? a1 : t3; - const int f2 = hi ? t0 : a2; - const int f3 = hi ? t1 : a3; - - // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte - // aligned: base_col is a multiple of 16, n*2 is a multiple of 8). - const int col0 = base_col + 4 * c4; - const float xs = x_scale[row]; - const float4 ws = *reinterpret_cast(weight_scale + col0); - const float v0 = static_cast(f0) * xs * ws.x; - const float v1 = static_cast(f1) * xs * ws.y; - const float v2 = static_cast(f2) * xs * ws.z; - const float v3 = static_cast(f3) * xs * ws.w; - const uint64_t packed = - static_cast( - static_cast(__float2bfloat16(v0))) | - (static_cast( - static_cast(__float2bfloat16(v1))) - << 16) | - (static_cast( - static_cast(__float2bfloat16(v2))) - << 32) | - (static_cast( - static_cast(__float2bfloat16(v3))) - << 48); - *reinterpret_cast(out + static_cast(row) * n + col0) = - packed; +// Round 8 (fix_build / epilogue): plain round-to-nearest-even fp32 -> bf16 +// conversion without hip's __float2bfloat16 NaN/Inf signaling-bit fixup. 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 or +// 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). The helper +// is spelled with a union rather than __float_as_uint so it depends on no +// optional device intrinsic; this is the validated formulation from the +// accepted hy3/TP8/M4096 o_proj lineage, not a new architecture. +__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); } -// Register-batched variant of the coalesced fragment store (iteration 6, -// epilogue round). Identical to store_prefill_fragment_coalesced in every -// output bit: same 4x4 shuffle transpose, same per-element multiply order -// float(dot) * x_scale[row] * weight_scale[col], same __float2bfloat16 -// rounding, same single 8-byte coalesced store. The only difference is that -// the per-row x_scale and per-column weight_scale values are passed in from -// registers: the caller loads each lane's 8 x_scale rows and 2 weight_scale -// float4 columns exactly once per wave (the per-call epilogue re-loaded the -// scales inside each of the 16 fragment stores), so the stored bits are -// unchanged while the epilogue vmem read count per lane per wave drops from -// up to 8 x_scale + 16 float4 loads to exactly 8 + 2 and the loads are -// independent of (and issue before) the shuffle/convert/store chains. +// Fused coalesced fragment store with register-batched scales. The m16n16k32 +// accumulator lane mapping (row = lane & 15, column group c4 = lane >> 4, +// frag.x[i] -> column c4 + 4*i) gives each lane four elements strided by 4 +// columns, so a direct store would be four 2-byte scalar stores per lane whose +// wavefront addresses touch each 32-B sector at 25% utilization. This +// epilogue transposes the 4-element groups within each 4-lane column group +// (lanes r, r+16, r+32, r+48 - a 4x4 transpose, two 2x2 steps with shfl_xor 16 +// then 32), so lane (r, c4) ends up holding the four CONTIGUOUS columns +// 4*c4 .. 4*c4+3, converts them to bf16, packs 4 bf16 (8 B), and writes ONE +// 8-byte store per lane (100% store sector efficiency). Only the int32 values +// are re-routed between lanes; the per-element multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are +// unchanged, so the stored bits are identical to a direct per-element store. +// xs is this lane's x_scale[row] and ws the float4 at +// weight_scale + base_col + 4*c4, both preloaded once by the caller. The +// row >= m guard is wavefront-uniform (all 64 lanes 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 scale load and the +// 8-byte store are aligned. template __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( const AccFragment& frag, @@ -294,612 +283,42 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( const int f2 = hi ? t0 : a2; const int f3 = hi ? t1 : a3; - // Lane (r, c4) owns columns base_col + 4*c4 .. +3 (8 B, 8-byte aligned: - // base_col is a multiple of 16, n*2 is a multiple of 8); ws is the float4 - // at weight_scale + base_col + 4*c4 preloaded by the caller. + // 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*2 is a multiple of 8). const int col0 = base_col + 4 * c4; const float v0 = static_cast(f0) * xs * ws.x; const float v1 = static_cast(f1) * xs * ws.y; const float v2 = static_cast(f2) * xs * ws.z; const float v3 = static_cast(f3) * xs * ws.w; const uint64_t packed = - static_cast( - static_cast(__float2bfloat16(v0))) | - (static_cast( - static_cast(__float2bfloat16(v1))) - << 16) | - (static_cast( - static_cast(__float2bfloat16(v2))) - << 32) | - (static_cast( - static_cast(__float2bfloat16(v3))) - << 48); + 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; } // --------------------------------------------------------------------------- -// Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. -// Each wave owns a 64x32 quadrant (eight m16n16k32 int32 accumulators); the -// block cooperatively stages A[128,64] from x_q (row-major, stride k) and -// B[64,64] from the identity [K, N] weight into a single-buffered LDS -// buffer. Two barriers per stage: one after the cooperative load, one before -// the next stage overwrites LDS. Direct fragment epilogue (coalesced store). +// Exact-shape GEMM: 256x128 output tile per block, four wavefronts. // -// Stage layout (bootstrap, identity [K, N] weight): -// A[128, 80]: m-major, 512 int4s; thread t loads int4 #0 at -// a_tile[(t/4)*80 + (t%4)*16] from x_q[(m0 + t/4)*k + k0 + -// (t%4)*16] and int4 #1 at row +64 (rows 64..127). -// B[64, 80]: k-major, 256 int4s; thread t loads int4 at -// b_tile[(t/4)*80 + (t%4)*16] from -// weight[(k0 + t/4)*N + n0 + (t%4)*16] (16 consecutive n at -// fixed k; one int4 per thread). B fragments are loaded with -// the library row_major loader (ldm = 80): lane (r, c4) holds -// x[i] = b_tile[(kk + r + i)*80 + local_col + c4], i.e. eight -// k-rows strided by ldm - correct for the identity [K, N] -// layout; later rounds can switch to an n-major packed B with -// the 8-byte fragment fill. -// --------------------------------------------------------------------------- -__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - bf16_t* __restrict__ out, - int m, - int n, - int k) { - const int tid = static_cast(threadIdx.x); - const int wave = tid / kWaveSize; - const int lane = tid % kWaveSize; - const int wave_row = wave >> 1; - const int wave_col = wave & 1; - const int m0 = static_cast(blockIdx.y) * kBlockM; - const int n0 = static_cast(blockIdx.x) * kBlockN; - const int local_row = wave_row * 64; - const int local_col = wave_col * 32; - - // Single-buffered 64-K stage: A[128, 80] + B[64, 80] = 15,360 B/block - // (4 blocks/CU fit the 64 KiB LDS budget; the padded 80-byte strides are - // 16-byte-aligned and break the 64-byte LDS bank periodicity). - __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; - __shared__ __align__(16) int8_t b_tile[kBlockN * kBStride]; - - DUFragment - a_frag0, a_frag1; - DUFragment - b_frag0, b_frag1; - DUFragment - acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; - du_fill_fragment(acc00, 0); - du_fill_fragment(acc01, 0); - du_fill_fragment(acc10, 0); - du_fill_fragment(acc11, 0); - du_fill_fragment(acc20, 0); - du_fill_fragment(acc21, 0); - du_fill_fragment(acc30, 0); - du_fill_fragment(acc31, 0); - - // Cooperative staging: each thread owns one 16-byte vector in A rows - // [0,64), one in A rows [64,128), and one 16-byte vector of B. - // A[128,64] is 512 int4s; B[64,64] is 256 int4s. - const int vector_byte_offset = tid * static_cast(sizeof(int4)); - const int stage_row = vector_byte_offset / kStageK; // tid / 4, 0..63 - const int stage_col = vector_byte_offset - stage_row * kStageK; // (tid%4)*16 - const int b_row = tid >> 2; // k row of this thread's B int4 - const int b_n16 = (tid & 3) * 16; // n offset (stage-local) of B int4 - - for (int k0 = 0; k0 < k; k0 += kStageK) { - *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = - *reinterpret_cast( - x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); - *reinterpret_cast(a_tile + - (stage_row + kBlockM / 2) * kAStride + stage_col) = - *reinterpret_cast( - x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + - k0 + stage_col); - *reinterpret_cast(b_tile + b_row * kBStride + b_n16) = - *reinterpret_cast( - weight + (static_cast(k0 + b_row) * n + n0 + b_n16)); - __syncthreads(); - - // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. -#pragma unroll - for (int kk = 0; kk < kStageK; kk += kTileK) { - // A fragments: direct 8-byte LDS fill (m-major A tile; each lane's - // eight fragment k-values are contiguous). - load_fragment8(a_frag0, a_tile + local_row * kAStride + kk, kAStride, - lane); - load_fragment8(a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, - kAStride, lane); - // B fragments: library row_major loader on the k-major [K, N] LDS - // tile (bootstrap, identity layout; ldm = kBStride). - du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, - kBStride); - du_load_matrix_sync(b_frag1, - b_tile + kk * kBStride + local_col + kTileN, - kBStride); - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); - load_fragment8( - a_frag0, a_tile + (local_row + 2 * kTileM) * kAStride + kk, - kAStride, lane); - load_fragment8( - a_frag1, a_tile + (local_row + 3 * kTileM) * kAStride + kk, - kAStride, lane); - du_mma_sync(acc20, a_frag0, b_frag0, acc20); - du_mma_sync(acc21, a_frag0, b_frag1, acc21); - du_mma_sync(acc30, a_frag1, b_frag0, acc30); - du_mma_sync(acc31, a_frag1, b_frag1, acc31); - } - - // Protect the LDS buffer from the next stage's cooperative overwrite. - __syncthreads(); - } - - const int base_row = m0 + local_row; - const int base_col = n0 + local_col; - // Coalesced epilogue store (one 8-byte store per lane per fragment; scales - // + bf16 conversion fused in-kernel, no workspace pass). - store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, - base_row, base_col, m, n, lane); - store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, - base_row, base_col + kTileN, m, n, lane); - store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, - base_row + kTileM, base_col, m, n, lane); - store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, - base_row + kTileM, base_col + kTileN, m, n, - lane); - store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col + kTileN, - m, n, lane); - store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col + kTileN, - m, n, lane); -} - -// --------------------------------------------------------------------------- -// Iteration-3 pipeline-control variant: double-buffered K-stage of the same -// 128x64 tile / 64-K stage / fragment loads / epilogue as the accepted -// single-buffered kernel. Stage s+1's three global loads (2 A int4s + 1 B -// int4 per thread) are issued BEFORE the stage s MMAC burst and the loaded -// vectors stay in VGPR; the publish (ds_write into the idle buffer) happens -// AFTER the burst, so the compiler's vmcnt(0) wait sits behind 16 v_mmac + -// ~70 LDS reads instead of directly behind the load issue (the accepted -// single-buffered kernel waits immediately after 3 loads per wave per -// stage). One __syncthreads per stage (32/block vs 64/block: the barrier -// doubles as read-done for stage s+1 and write-done for stage s; the last -// stage needs no fence before the register epilogue). Costs: LDS -// 15,360 -> 30,720 B/block (A[2][128,80] + B[2][64,80]), residency -// 4 -> 2 blocks/CU = 8 waves/CU; VGPR cap at 2 blocks/CU is 128, so the -// ~12-VGPR loop-carried prefetch payload cannot change residency. Retained -// only if ISA/PMC evidence shows the vmcnt wait moved off the load->ds_write -// path (covered by the MMAC burst) and the measured median/p90 beat the -// 1162.15/1164.03 guard. -// --------------------------------------------------------------------------- -__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_db_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - bf16_t* __restrict__ out, - int m, - int n, - int k) { - const int tid = static_cast(threadIdx.x); - const int wave = tid / kWaveSize; - const int lane = tid % kWaveSize; - const int wave_row = wave >> 1; - const int wave_col = wave & 1; - const int m0 = static_cast(blockIdx.y) * kBlockM; - const int n0 = static_cast(blockIdx.x) * kBlockN; - const int local_row = wave_row * 64; - const int local_col = wave_col * 32; - - // Double-buffered 64-K stage: A[2][128,80] + B[2][64,80] = 30,720 B/block - // (2 blocks/CU fit the 64 KiB LDS budget; padded 80-byte strides and bank - // phases unchanged from the accepted kernel). - __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; - __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBStride]; - - DUFragment - a_frag0, a_frag1; - DUFragment - b_frag0, b_frag1; - DUFragment - acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; - du_fill_fragment(acc00, 0); - du_fill_fragment(acc01, 0); - du_fill_fragment(acc10, 0); - du_fill_fragment(acc11, 0); - du_fill_fragment(acc20, 0); - du_fill_fragment(acc21, 0); - du_fill_fragment(acc30, 0); - du_fill_fragment(acc31, 0); - - // Cooperative staging indexing is identical to the accepted kernel: each - // thread owns two 16-byte A vectors (rows [0,64) and [64,128)) and one - // 16-byte B vector. - const int vector_byte_offset = tid * static_cast(sizeof(int4)); - const int stage_row = vector_byte_offset / kStageK; // tid / 4, 0..63 - const int stage_col = vector_byte_offset - stage_row * kStageK; // (tid%4)*16 - const int b_row = tid >> 2; // k row of this thread's B int4 - const int b_n16 = (tid & 3) * 16; // n offset (stage-local) of B int4 - - const int k_stages = k / kStageK; - - // Loop-carried prefetch payload (stage s+1 global vectors in VGPR). - int4 pa0, pa1, pb0; - - // Prologue: fetch stage 0 into registers, publish into buffer 0, fence. - { - const int k0 = 0; - pa0 = *reinterpret_cast( - x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); - pa1 = *reinterpret_cast( - x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + k0 + - stage_col); - pb0 = *reinterpret_cast( - weight + (static_cast(k0 + b_row) * n + n0 + b_n16)); - *reinterpret_cast(a_tile[0] + stage_row * kAStride + stage_col) = - pa0; - *reinterpret_cast(a_tile[0] + - (stage_row + kBlockM / 2) * kAStride + - stage_col) = pa1; - *reinterpret_cast(b_tile[0] + b_row * kBStride + b_n16) = pb0; - __syncthreads(); - } - - for (int s = 0; s < k_stages; ++s) { - const int cur = s & 1; - const int nxt = cur ^ 1; - - // Prefetch stage s+1: issue the global loads BEFORE the stage s MMAC - // burst; the loaded vectors stay in VGPR and the compiler's vmcnt(0) wait - // is inserted only at the publish below (covered by the burst). - if (s + 1 < k_stages) { - const int k0 = (s + 1) * kStageK; - pa0 = *reinterpret_cast( - x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); - pa1 = *reinterpret_cast( - x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + k0 + - stage_col); - pb0 = *reinterpret_cast( - weight + (static_cast(k0 + b_row) * n + n0 + b_n16)); - } - - // Consume stage s from the current buffer: bit-identical fragment values, - // MMAC sequence, int32 accumulation order and lane mapping to the - // accepted kernel (only the buffer base changes). -#pragma unroll - for (int kk = 0; kk < kStageK; kk += kTileK) { - load_fragment8(a_frag0, a_tile[cur] + local_row * kAStride + kk, - kAStride, lane); - load_fragment8(a_frag1, a_tile[cur] + (local_row + kTileM) * kAStride + - kk, - kAStride, lane); - du_load_matrix_sync(b_frag0, - b_tile[cur] + kk * kBStride + local_col, kBStride); - du_load_matrix_sync(b_frag1, - b_tile[cur] + kk * kBStride + local_col + kTileN, - kBStride); - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); - load_fragment8( - a_frag0, a_tile[cur] + (local_row + 2 * kTileM) * kAStride + kk, - kAStride, lane); - load_fragment8( - a_frag1, a_tile[cur] + (local_row + 3 * kTileM) * kAStride + kk, - kAStride, lane); - du_mma_sync(acc20, a_frag0, b_frag0, acc20); - du_mma_sync(acc21, a_frag0, b_frag1, acc21); - du_mma_sync(acc30, a_frag1, b_frag0, acc30); - du_mma_sync(acc31, a_frag1, b_frag1, acc31); - } - - // Publish stage s+1 into the idle buffer after the burst (vmcnt wait - // covered by the burst), then one fence: read-done for s+1 and - // write-done for s (buffer cur is free for the s+2 publish). - if (s + 1 < k_stages) { - *reinterpret_cast(a_tile[nxt] + stage_row * kAStride + - stage_col) = pa0; - *reinterpret_cast(a_tile[nxt] + - (stage_row + kBlockM / 2) * kAStride + - stage_col) = pa1; - *reinterpret_cast(b_tile[nxt] + b_row * kBStride + b_n16) = pb0; - __syncthreads(); - } - } - - const int base_row = m0 + local_row; - const int base_col = n0 + local_col; - store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, - base_row, base_col, m, n, lane); - store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, - base_row, base_col + kTileN, m, n, lane); - store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, - base_row + kTileM, base_col, m, n, lane); - store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, - base_row + kTileM, base_col + kTileN, m, n, - lane); - store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col + kTileN, - m, n, lane); - store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col + kTileN, - m, n, lane); -} - -// --------------------------------------------------------------------------- -// Iteration-4 tile-shape variant: 256x64 output tile (M-major doubling of the -// accepted 128x64 tile) at 256 threads = 4 wavefronts, each wave owning a -// 128x32 quadrant (sixteen m16n16k32 int32 accumulators). This is the -// measured tile axis pushed one step further: per-FLOP B-fragment LDS reads -// scale as 1/(quadrant m-height) (the 64x128 flip's quadrants were 32x64 and -// lost +8% with a doubled per-FLOP B cost), so doubling the quadrant height -// to 128 HALVES the per-FLOP B ds_read_u8 + reassembly VALU - the measured -// LDS-read/wait bottleneck (PMC lds_wait ~= lds_instructions) - while keeping -// the same 64-column quadrant width (B library row_major loader ldm = 80). -// B reuse doubles (each staged B element is consumed by all 256 block rows), -// so the grid-level B re-read drops from M/128 = 32x to M/256 = 16x (staged B -// traffic 402.7 -> 201.3 MB/replay; A re-read stays N/64 = 96x = 805.3 MB; -// total 1.208 -> 1.007 GB, -17%). Single-buffered 64-K stage: A[256,80] + -// B[64,80] = 25,600 B/block -> 2 blocks/CU = 8 waves/CU, the SAME -// co-residency structure as the accepted double-buffered 128x64 kernel -// (2 blocks x 4 waves), so the comparison is not residency-confounded. -// Two __syncthreads per stage, direct 8-byte A fragment fills, fused -// coalesced scale->bf16 epilogue, row>=m guard. Grid dim3(96,16) = 1536 -// blocks = 12.8 blocks/CU, so the MxN output-tile grid still dwarfs the -// 120 CUs. Fragment values, MMAC sequence, int32 accumulation order -// (k0-outer 64, kk-inner 32, hardware k-order in m16n16k32), -// element-to-lane mapping, barrier count and the fused epilogue are -// bit-identical to the accepted kernels (only the tile height, the quadrant -// split and the staging vectorization change), so exact bf16 equality -// (mismatch 0 / max_abs_error 0.0) is expected without tolerance debate. -// --------------------------------------------------------------------------- -__global__ __launch_bounds__(kBlockThreads) - void w8a8_dumma_256x64x64_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - bf16_t* __restrict__ out, - int m, - int n, - int k) { - const int tid = static_cast(threadIdx.x); - const int wave = tid / kWaveSize; - const int lane = tid % kWaveSize; - // 4 waves: wave_row = (wave>>1)*128 covers rows 0..255 in two 128-row bands - // (waves 2w, 2w+1 share the band), wave_col picks the 32-col half. - const int wave_row = (wave >> 1) * (kBlockM256 / 2); - const int wave_col = wave & 1; - const int m0 = static_cast(blockIdx.y) * kBlockM256; - const int n0 = static_cast(blockIdx.x) * kBlockN; - const int local_row = wave_row; - const int local_col = wave_col * 32; - - // Single-buffered 64-K stage: A[256, 80] + B[64, 80] = 25,600 B/block - // (2 blocks/CU fit the 64 KiB LDS budget; padded 80-byte strides and bank - // phases unchanged from the accepted kernels). - __shared__ __align__(16) int8_t a_tile[kBlockM256 * kAStride]; - __shared__ __align__(16) int8_t b_tile[kBlockN * kBStride]; - - DUFragment - a_frag0, a_frag1, a_frag2, a_frag3, a_frag4, a_frag5, a_frag6, a_frag7; - DUFragment - b_frag0, b_frag1; - // Sixteen m16n16k32 int32 accumulators (8 m-groups x 2 n-groups per wave). - DUFragment - acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31, - acc40, acc41, acc50, acc51, acc60, acc61, acc70, acc71; - du_fill_fragment(acc00, 0); - du_fill_fragment(acc01, 0); - du_fill_fragment(acc10, 0); - du_fill_fragment(acc11, 0); - du_fill_fragment(acc20, 0); - du_fill_fragment(acc21, 0); - du_fill_fragment(acc30, 0); - du_fill_fragment(acc31, 0); - du_fill_fragment(acc40, 0); - du_fill_fragment(acc41, 0); - du_fill_fragment(acc50, 0); - du_fill_fragment(acc51, 0); - du_fill_fragment(acc60, 0); - du_fill_fragment(acc61, 0); - du_fill_fragment(acc70, 0); - du_fill_fragment(acc71, 0); - - // Cooperative staging: A[256,64] is 1,024 int4s (4 per thread: rows - // [0,64), [64,128), [128,192), [192,256)) and B[64,64] is 256 int4s (one - // per thread), 5 loads per thread per stage. - const int a_row0 = tid >> 2; // 0..63 - const int a_col = (tid & 3) * 16; // 0,16,32,48 - - for (int k0 = 0; k0 < k; k0 += kStageK) { - *reinterpret_cast(a_tile + a_row0 * kAStride + a_col) = - *reinterpret_cast( - x_q + static_cast(m0 + a_row0) * k + k0 + a_col); - *reinterpret_cast(a_tile + - (a_row0 + kBlockM256 / 4) * kAStride + a_col) = - *reinterpret_cast( - x_q + static_cast(m0 + a_row0 + kBlockM256 / 4) * k + - k0 + a_col); - *reinterpret_cast(a_tile + - (a_row0 + kBlockM256 / 2) * kAStride + a_col) = - *reinterpret_cast( - x_q + static_cast(m0 + a_row0 + kBlockM256 / 2) * k + - k0 + a_col); - *reinterpret_cast(a_tile + - (a_row0 + 3 * kBlockM256 / 4) * kAStride + - a_col) = - *reinterpret_cast( - x_q + static_cast(m0 + a_row0 + 3 * kBlockM256 / 4) * k + - k0 + a_col); - const int b_row = tid >> 2; // k row of this thread's B int4, 0..63 - const int b_n16 = (tid & 3) * 16; - *reinterpret_cast(b_tile + b_row * kBStride + b_n16) = - *reinterpret_cast( - weight + (static_cast(k0 + b_row) * n + n0 + b_n16)); - __syncthreads(); - - // Each wave consumes its 128x32 quadrant: eight m-groups x two n-groups, - // 16 m16n16k32 MMACs per kk. The A fragment loads and the MMAC sequence - // keep the exact int32 accumulation order (k0-outer 64, kk-inner 32, - // hardware k-order in m16n16k32) and the element-to-lane mapping of the - // accepted kernels; only the number of m-groups per wave changes. -#pragma unroll - for (int kk = 0; kk < kStageK; kk += kTileK) { - load_fragment8(a_frag0, a_tile + local_row * kAStride + kk, kAStride, - lane); - load_fragment8(a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, - kAStride, lane); - load_fragment8(a_frag2, a_tile + (local_row + 2 * kTileM) * kAStride + - kk, - kAStride, lane); - load_fragment8(a_frag3, a_tile + (local_row + 3 * kTileM) * kAStride + - kk, - kAStride, lane); - du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, - kBStride); - du_load_matrix_sync(b_frag1, - b_tile + kk * kBStride + local_col + kTileN, - kBStride); - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); - du_mma_sync(acc20, a_frag2, b_frag0, acc20); - du_mma_sync(acc21, a_frag2, b_frag1, acc21); - du_mma_sync(acc30, a_frag3, b_frag0, acc30); - du_mma_sync(acc31, a_frag3, b_frag1, acc31); - load_fragment8(a_frag4, a_tile + (local_row + 4 * kTileM) * kAStride + - kk, - kAStride, lane); - load_fragment8(a_frag5, a_tile + (local_row + 5 * kTileM) * kAStride + - kk, - kAStride, lane); - load_fragment8(a_frag6, a_tile + (local_row + 6 * kTileM) * kAStride + - kk, - kAStride, lane); - load_fragment8(a_frag7, a_tile + (local_row + 7 * kTileM) * kAStride + - kk, - kAStride, lane); - du_mma_sync(acc40, a_frag4, b_frag0, acc40); - du_mma_sync(acc41, a_frag4, b_frag1, acc41); - du_mma_sync(acc50, a_frag5, b_frag0, acc50); - du_mma_sync(acc51, a_frag5, b_frag1, acc51); - du_mma_sync(acc60, a_frag6, b_frag0, acc60); - du_mma_sync(acc61, a_frag6, b_frag1, acc61); - du_mma_sync(acc70, a_frag7, b_frag0, acc70); - du_mma_sync(acc71, a_frag7, b_frag1, acc71); - } - - // Protect the LDS buffer from the next stage's cooperative overwrite. - __syncthreads(); - } - - const int base_row = m0 + local_row; - const int base_col = n0 + local_col; - store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, - base_row, base_col, m, n, lane); - store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, - base_row, base_col + kTileN, m, n, lane); - store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, - base_row + kTileM, base_col, m, n, lane); - store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, - base_row + kTileM, base_col + kTileN, m, n, - lane); - store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col + kTileN, - m, n, lane); - store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col + kTileN, - m, n, lane); - store_prefill_fragment_coalesced(acc40, x_scale, weight_scale, out, - base_row + 4 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc41, x_scale, weight_scale, out, - base_row + 4 * kTileM, base_col + kTileN, - m, n, lane); - store_prefill_fragment_coalesced(acc50, x_scale, weight_scale, out, - base_row + 5 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc51, x_scale, weight_scale, out, - base_row + 5 * kTileM, base_col + kTileN, - m, n, lane); - store_prefill_fragment_coalesced(acc60, x_scale, weight_scale, out, - base_row + 6 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc61, x_scale, weight_scale, out, - base_row + 6 * kTileM, base_col + kTileN, - m, n, lane); - store_prefill_fragment_coalesced(acc70, x_scale, weight_scale, out, - base_row + 7 * kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced(acc71, x_scale, weight_scale, out, - base_row + 7 * kTileM, base_col + kTileN, - m, n, lane); -} - -// --------------------------------------------------------------------------- -// Iteration-5 packed-B variant of the accepted 256x64 kernel: the B operand -// switches to the one-time n-major [N, K] pack (packed[n*K + kk] = -// raw[kk*N + n], produced by launch_pack_w8a8_weight outside the timed -// region and Graph capture; same byte count and same buffer, so captured -// addresses are unchanged) and the B fragment path switches from the -// library row_major loader (8 per-byte ds_read_u8 strided by ldm per lane -// plus mask/OR reassembly VALU per dword - the measured LDS-read axis) to -// the lineage-validated 8-byte fragment fill (load_fragment8, one -// ds_read_b64 per col_major fragment, zero reassembly VALU) on the n-major -// staged tile. The m16n16k32 matrix_b hardware slot for lane -// (r = lane & 15, kq = lane >> 4) is B[k = kk + kq*8 + i][n = local_col + -// i*16 + r] - eight consecutive k bytes at fixed n - so staging B n-major -// (b_tile[n][k]) makes each lane's eight fragment bytes CONTIGUOUS in LDS: -// the fill reads the same elements into the same fragment slots with the -// same v_mmac inputs and the exact int32 accumulation order (k0-outer 64, -// kk-inner 32, hardware k-order), so exact bf16 equality (mismatch 0 / -// max_abs_error 0.0) is expected without tolerance debate. LDS bank -// safety: the n-major B tile uses the 72-byte row stride (odd 8-byte -// count) so the 16 lanes of each ds_read_b64 phase read 18r mod 32 -// (r = 0..15) distinct bank pairs - conflict-free; 72 % 16 == 8 forces the -// two-int64-half staging stores (ds_write2_b64). Iteration 9 gives the A -// tile the same conflict-free floor: the A row stride is 88 bytes (22 -// dwords; 22r mod 32 over r = 0..15 is a permutation of the 16 even dword -// starts), replacing the 80-byte stride's 2-way class (20r mod 32, period -// 8, rows r and r+8 alias), and A staging uses the same two-int64-half -// ds_write2_b64 pattern (88 % 16 == 8). Everything else is bit-identical -// to the accepted kernel: same bytes staged/read (only the physical LDS -// addresses move), same 128x32 quadrant per wave, same MMAC sequence, two -// __syncthreads per stage, fused coalesced scale->bf16 epilogue, row>=m -// guard, grid dim3(96,16) = 1536 blocks. LDS 22,528 (A) + 4,608 (B) = -// 27,136 B/block -> 2 blocks/CU = 8 waves/CU, the same co-residency -// structure as the accepted kernel. +// 4 waves: wave_row = (wave >> 1) * 128 covers rows 0..255 in two 128-row +// bands (waves 2w and 2w+1 share the band), wave_col = (wave & 1) * 64 selects +// the 64-column half. Each wave consumes a 128x64 quadrant: eight m-groups x +// four n-groups = thirty-two m16n16k32 int32 accumulators. The eight A +// fragment addresses per kk and the 88-byte A stride are byte-identical to +// round 1; only the n-major B tile and its four fragment fills are new. +// +// Staging (6 int4 global loads per thread per stage): +// A[256, 88]: thread t owns 16-byte A chunks at (row = t >> 2, +// col = (t & 3) * 16) in each of the four 64-row bands; each +// store is two int64 halves because 88 % 16 == 8. +// B[128, 72]: thread t owns the 16-byte k-chunk at (n_row = t >> 2, +// k16 = (t & 3) * 16) and at n_row + 64 of the packed n-major +// weight; each store is two int64 halves because 72 % 16 == 8. // --------------------------------------------------------------------------- __global__ __launch_bounds__(kBlockThreads) - void w8a8_dumma_256x64x64_packedb_kernel( + void w8a8_dumma_256x128x64_packedb_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ packed_w, // [N, K] n-major (exact shape) const float* __restrict__ x_scale, @@ -911,100 +330,105 @@ __global__ __launch_bounds__(kBlockThreads) const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; const int lane = tid % kWaveSize; - // 4 waves: wave_row = (wave>>1)*128 covers rows 0..255 in two 128-row bands - // (waves 2w, 2w+1 share the band), wave_col picks the 32-col half. - const int wave_row = (wave >> 1) * (kBlockM256 / 2); - const int wave_col = wave & 1; - const int m0 = static_cast(blockIdx.y) * kBlockM256; + const int wave_row = (wave >> 1) * (kBlockM / 2); + const int wave_col = (wave & 1) * (kBlockN / 2); + const int m0 = static_cast(blockIdx.y) * kBlockM; const int n0 = static_cast(blockIdx.x) * kBlockN; const int local_row = wave_row; - const int local_col = wave_col * 32; + const int local_col = wave_col; - // Single-buffered 64-K stage: A[256, 88] (iteration-9 stride) + B[64, 72] - // n-major = 27,136 B/block (2 blocks/CU fit the 64 KiB LDS budget). - __shared__ __align__(16) int8_t a_tile[kBlockM256 * kAStridePacked]; - __shared__ __align__(16) int8_t b_tile[kBlockN * kBStridePacked]; + // Single-buffered 64-K stage: A[256, 88] + B[128, 72] = 31,744 B/block. + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kBStride]; DUFragment a_frag0, a_frag1, a_frag2, a_frag3, a_frag4, a_frag5, a_frag6, a_frag7; DUFragment - b_frag0, b_frag1; - // Sixteen m16n16k32 int32 accumulators (8 m-groups x 2 n-groups per wave). + b_frag0, b_frag1, b_frag2, b_frag3; + // Thirty-two m16n16k32 int32 accumulators (8 m-groups x 4 n-groups per wave). DUFragment - acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31, - acc40, acc41, acc50, acc51, acc60, acc61, acc70, acc71; + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13, + acc20, acc21, acc22, acc23, acc30, acc31, acc32, acc33, + acc40, acc41, acc42, acc43, acc50, acc51, acc52, acc53, + acc60, acc61, acc62, acc63, acc70, acc71, acc72, acc73; du_fill_fragment(acc00, 0); du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); du_fill_fragment(acc10, 0); du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); du_fill_fragment(acc20, 0); du_fill_fragment(acc21, 0); + du_fill_fragment(acc22, 0); + du_fill_fragment(acc23, 0); du_fill_fragment(acc30, 0); du_fill_fragment(acc31, 0); + du_fill_fragment(acc32, 0); + du_fill_fragment(acc33, 0); du_fill_fragment(acc40, 0); du_fill_fragment(acc41, 0); + du_fill_fragment(acc42, 0); + du_fill_fragment(acc43, 0); du_fill_fragment(acc50, 0); du_fill_fragment(acc51, 0); + du_fill_fragment(acc52, 0); + du_fill_fragment(acc53, 0); du_fill_fragment(acc60, 0); du_fill_fragment(acc61, 0); + du_fill_fragment(acc62, 0); + du_fill_fragment(acc63, 0); du_fill_fragment(acc70, 0); du_fill_fragment(acc71, 0); + du_fill_fragment(acc72, 0); + du_fill_fragment(acc73, 0); - // Cooperative staging: A[256,64] is 1,024 int4s (4 per thread: rows - // [0,64), [64,128), [128,192), [192,256)) and B[64,64] is 256 int4s (one - // per thread), 5 loads per thread per stage. const int a_row0 = tid >> 2; // 0..63 - const int a_col = (tid & 3) * 16; // 0,16,32,48 - // B staging is n-major: thread owns the 16-byte k-chunk at - // (n_row = tid >> 2, k16 = (tid & 3) * 16). + const int a_col = (tid & 3) * 16; // 0, 16, 32, 48 const int b_n_row = tid >> 2; // 0..63 - const int b_k16 = (tid & 3) * 16; // 0,16,32,48 + const int b_k16 = (tid & 3) * 16; // 0, 16, 32, 48 const int base_row = m0 + local_row; const int base_col = n0 + local_col; - // Iteration-6 (epilogue round) register-batched scales: each lane's 16 - // fragment stores cover the 8 m-groups (rows base_row + r + 16*i, - // r = lane & 15) x 2 n-groups (cols base_col + 4*c4 and - // base_col + 16 + 4*c4, c4 = lane >> 4), so each lane needs exactly 8 - // x_scale values and 2 weight_scale float4s. They are loaded ONCE on the - // last K stage, issued BEFORE the final protective barrier (uniform - // branch: k0 is block-uniform, so the barrier below is reached by every - // thread), so the vmem latency overlaps the barrier and the 16 - // shuffle/convert/store chains; the previous per-call epilogue re-loaded - // the scales inside each of the 16 fragment stores (PMC vmem_read - // surplus ~24-25 wavefront loads/wave above the 5/thread/stage staging - // floor matches ~8 x_scale + ~16 float4 re-loads per lane per wave). + + // Epilogue scale registers: each lane's 32 fragment stores cover the 8 + // m-groups (rows base_row + 16*i + r, r = lane & 15) x 4 n-groups + // (cols base_col + 16*j + 4*c4, c4 = lane >> 4), so each lane needs exactly + // 8 x_scale values and 4 weight_scale float4s. They are loaded once on the + // last K stage (uniform branch: k0 is block-uniform, so the following + // barrier is reached by every thread) and reused by all thirty-two coalesced + // stores; xs_m is zeroed for out-of-range tail rows. float xs_m[8]; - float4 ws0, ws1; + float4 ws0, ws1, ws2, ws3; for (int k0 = 0; k0 < k; k0 += kStageK) { const int4 av0 = *reinterpret_cast( x_q + static_cast(m0 + a_row0) * k + k0 + a_col); const int4 av1 = *reinterpret_cast( - x_q + static_cast(m0 + a_row0 + kBlockM256 / 4) * k + k0 + + x_q + static_cast(m0 + a_row0 + kBlockM / 4) * k + k0 + a_col); const int4 av2 = *reinterpret_cast( - x_q + static_cast(m0 + a_row0 + kBlockM256 / 2) * k + k0 + + x_q + static_cast(m0 + a_row0 + kBlockM / 2) * k + k0 + a_col); const int4 av3 = *reinterpret_cast( - x_q + static_cast(m0 + a_row0 + 3 * kBlockM256 / 4) * k + - k0 + a_col); + x_q + static_cast(m0 + a_row0 + 3 * kBlockM / 4) * k + k0 + + a_col); // A rows are 88 B apart (8 mod 16), so each 16-byte staging store is two - // int64 halves (ds_write2_b64, one instruction) - the exact accepted - // packed-B staging pattern below; a single int4 store would be - // misaligned on odd rows. + // int64 halves (ds_write2_b64); a single int4 store would be misaligned on + // odd rows. const int64_t* asrc0 = reinterpret_cast(&av0); const int64_t* asrc1 = reinterpret_cast(&av1); const int64_t* asrc2 = reinterpret_cast(&av2); const int64_t* asrc3 = reinterpret_cast(&av3); int64_t* adst0 = reinterpret_cast( - a_tile + a_row0 * kAStridePacked + a_col); + a_tile + a_row0 * kAStride + a_col); int64_t* adst1 = reinterpret_cast( - a_tile + (a_row0 + kBlockM256 / 4) * kAStridePacked + a_col); + a_tile + (a_row0 + kBlockM / 4) * kAStride + a_col); int64_t* adst2 = reinterpret_cast( - a_tile + (a_row0 + kBlockM256 / 2) * kAStridePacked + a_col); + a_tile + (a_row0 + kBlockM / 2) * kAStride + a_col); int64_t* adst3 = reinterpret_cast( - a_tile + (a_row0 + 3 * kBlockM256 / 4) * kAStridePacked + a_col); + a_tile + (a_row0 + 3 * kBlockM / 4) * kAStride + a_col); adst0[0] = asrc0[0]; adst0[1] = asrc0[1]; adst1[0] = asrc1[0]; @@ -1013,78 +437,95 @@ __global__ __launch_bounds__(kBlockThreads) adst2[1] = asrc2[1]; adst3[0] = asrc3[0]; adst3[1] = asrc3[1]; - // Stage B[64, 64] n-major from the packed [N, K] weight: 16 consecutive - // k bytes per thread at fixed n (one coalesced int4 global load). Row - // stride 72 is 8 mod 16, so the LDS store is two int64 halves - // (ds_write2_b64) - a single int4 store would be misaligned on odd rows. - const int4 bv = *reinterpret_cast( - packed_w + - (static_cast(n0 + b_n_row) * k + k0 + b_k16)); - const int64_t* bsrc64 = reinterpret_cast(&bv); - int64_t* bdst64 = reinterpret_cast( - b_tile + b_n_row * kBStridePacked + b_k16); - bdst64[0] = bsrc64[0]; - bdst64[1] = bsrc64[1]; + // B staging is n-major: 16 consecutive k bytes at fixed n (one coalesced + // int4 global load from the packed [N, K] weight) in two 64-row bands; + // row stride 72 is 8 mod 16, so each LDS store is two int64 halves. + const int4 bv0 = *reinterpret_cast( + packed_w + static_cast(n0 + b_n_row) * k + k0 + b_k16); + const int4 bv1 = *reinterpret_cast( + packed_w + static_cast(n0 + b_n_row + kBlockN / 2) * k + k0 + + b_k16); + const int64_t* bsrc0 = reinterpret_cast(&bv0); + const int64_t* bsrc1 = reinterpret_cast(&bv1); + int64_t* bdst0 = reinterpret_cast( + b_tile + b_n_row * kBStride + b_k16); + int64_t* bdst1 = reinterpret_cast( + b_tile + (b_n_row + kBlockN / 2) * kBStride + b_k16); + bdst0[0] = bsrc0[0]; + bdst0[1] = bsrc0[1]; + bdst1[0] = bsrc1[0]; + bdst1[1] = bsrc1[1]; __syncthreads(); - // Each wave consumes its 128x32 quadrant: eight m-groups x two n-groups, - // 16 m16n16k32 MMACs per kk. The A fragment loads and the MMAC sequence - // keep the exact int32 accumulation order (k0-outer 64, kk-inner 32, - // hardware k-order in m16n16k32) and the element-to-lane mapping of the - // accepted kernels; only the B fragment fill (one 8-byte contiguous LDS - // read per col_major fragment on the n-major tile) changes. + // Each wave consumes its 128x64 quadrant: eight m-groups x four n-groups, + // 32 m16n16k32 MMACs per kk, in the exact int32 accumulation order + // (k0-outer 64, kk-inner 32, hardware k-order inside m16n16k32). The B + // fragments are loaded once per kk and reused by all eight m-groups; the A + // fragments are consumed in two batches of four to bound their live range + // to eight VGPRs at a time. #pragma unroll for (int kk = 0; kk < kStageK; kk += kTileK) { - load_fragment8(a_frag0, a_tile + local_row * kAStridePacked + kk, - kAStridePacked, lane); - load_fragment8(a_frag1, - a_tile + (local_row + kTileM) * kAStridePacked + kk, - kAStridePacked, lane); - load_fragment8(a_frag2, - a_tile + (local_row + 2 * kTileM) * kAStridePacked + kk, - kAStridePacked, lane); - load_fragment8(a_frag3, - a_tile + (local_row + 3 * kTileM) * kAStridePacked + kk, - kAStridePacked, lane); - load_fragment8(b_frag0, b_tile + local_col * kBStridePacked + kk, - kBStridePacked, lane); - load_fragment8(b_frag1, - b_tile + (local_col + kTileN) * kBStridePacked + kk, - kBStridePacked, lane); + load_fragment8(b_frag0, b_tile + local_col * kBStride + kk, kBStride, + lane); + load_fragment8(b_frag1, b_tile + (local_col + kTileN) * kBStride + kk, + kBStride, lane); + load_fragment8(b_frag2, b_tile + (local_col + 2 * kTileN) * kBStride + kk, + kBStride, lane); + load_fragment8(b_frag3, b_tile + (local_col + 3 * kTileN) * kBStride + kk, + kBStride, lane); + load_fragment8(a_frag0, a_tile + local_row * kAStride + kk, kAStride, + lane); + load_fragment8(a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8(a_frag2, a_tile + (local_row + 2 * kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8(a_frag3, a_tile + (local_row + 3 * kTileM) * kAStride + kk, + kAStride, lane); du_mma_sync(acc00, a_frag0, b_frag0, acc00); du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); du_mma_sync(acc10, a_frag1, b_frag0, acc10); du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); du_mma_sync(acc20, a_frag2, b_frag0, acc20); du_mma_sync(acc21, a_frag2, b_frag1, acc21); + du_mma_sync(acc22, a_frag2, b_frag2, acc22); + du_mma_sync(acc23, a_frag2, b_frag3, acc23); du_mma_sync(acc30, a_frag3, b_frag0, acc30); du_mma_sync(acc31, a_frag3, b_frag1, acc31); - load_fragment8(a_frag4, - a_tile + (local_row + 4 * kTileM) * kAStridePacked + kk, - kAStridePacked, lane); - load_fragment8(a_frag5, - a_tile + (local_row + 5 * kTileM) * kAStridePacked + kk, - kAStridePacked, lane); - load_fragment8(a_frag6, - a_tile + (local_row + 6 * kTileM) * kAStridePacked + kk, - kAStridePacked, lane); - load_fragment8(a_frag7, - a_tile + (local_row + 7 * kTileM) * kAStridePacked + kk, - kAStridePacked, lane); + du_mma_sync(acc32, a_frag3, b_frag2, acc32); + du_mma_sync(acc33, a_frag3, b_frag3, acc33); + load_fragment8(a_frag4, a_tile + (local_row + 4 * kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8(a_frag5, a_tile + (local_row + 5 * kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8(a_frag6, a_tile + (local_row + 6 * kTileM) * kAStride + kk, + kAStride, lane); + load_fragment8(a_frag7, a_tile + (local_row + 7 * kTileM) * kAStride + kk, + kAStride, lane); du_mma_sync(acc40, a_frag4, b_frag0, acc40); du_mma_sync(acc41, a_frag4, b_frag1, acc41); + du_mma_sync(acc42, a_frag4, b_frag2, acc42); + du_mma_sync(acc43, a_frag4, b_frag3, acc43); du_mma_sync(acc50, a_frag5, b_frag0, acc50); du_mma_sync(acc51, a_frag5, b_frag1, acc51); + du_mma_sync(acc52, a_frag5, b_frag2, acc52); + du_mma_sync(acc53, a_frag5, b_frag3, acc53); du_mma_sync(acc60, a_frag6, b_frag0, acc60); du_mma_sync(acc61, a_frag6, b_frag1, acc61); + du_mma_sync(acc62, a_frag6, b_frag2, acc62); + du_mma_sync(acc63, a_frag6, b_frag3, acc63); du_mma_sync(acc70, a_frag7, b_frag0, acc70); du_mma_sync(acc71, a_frag7, b_frag1, acc71); + du_mma_sync(acc72, a_frag7, b_frag2, acc72); + du_mma_sync(acc73, a_frag7, b_frag3, acc73); } - // 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). + // Last stage only: prefetch the epilogue's per-row x_scale and per-column + // weight_scale values into registers (uniform branch; 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; @@ -1097,7 +538,12 @@ __global__ __launch_bounds__(kBlockThreads) 4 * c4); ws1 = *reinterpret_cast(weight_scale + base_col + kTileN + 4 * c4); + ws2 = *reinterpret_cast(weight_scale + base_col + + 2 * kTileN + 4 * c4); + ws3 = *reinterpret_cast(weight_scale + base_col + + 3 * kTileN + 4 * c4); } + // Protect the LDS buffer from the next stage's cooperative overwrite. __syncthreads(); } @@ -1106,58 +552,105 @@ __global__ __launch_bounds__(kBlockThreads) base_col, m, n, lane); store_prefill_fragment_coalesced_scaled(acc01, xs_m[0], ws1, out, base_row, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc02, xs_m[0], ws2, out, base_row, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc03, xs_m[0], ws3, out, base_row, + base_col + 3 * kTileN, m, n, lane); store_prefill_fragment_coalesced_scaled(acc10, xs_m[1], ws0, out, base_row + kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced_scaled(acc11, xs_m[1], ws1, out, base_row + kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc12, xs_m[1], ws2, out, + base_row + kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc13, xs_m[1], ws3, out, + base_row + kTileM, + base_col + 3 * kTileN, m, n, lane); store_prefill_fragment_coalesced_scaled(acc20, xs_m[2], ws0, out, base_row + 2 * kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced_scaled(acc21, xs_m[2], ws1, out, base_row + 2 * kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc22, xs_m[2], ws2, out, + base_row + 2 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc23, xs_m[2], ws3, out, + base_row + 2 * kTileM, + base_col + 3 * kTileN, m, n, lane); store_prefill_fragment_coalesced_scaled(acc30, xs_m[3], ws0, out, base_row + 3 * kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced_scaled(acc31, xs_m[3], ws1, out, base_row + 3 * kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc32, xs_m[3], ws2, out, + base_row + 3 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc33, xs_m[3], ws3, out, + base_row + 3 * kTileM, + base_col + 3 * kTileN, m, n, lane); store_prefill_fragment_coalesced_scaled(acc40, xs_m[4], ws0, out, base_row + 4 * kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced_scaled(acc41, xs_m[4], ws1, out, base_row + 4 * kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc42, xs_m[4], ws2, out, + base_row + 4 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc43, xs_m[4], ws3, out, + base_row + 4 * kTileM, + base_col + 3 * kTileN, m, n, lane); store_prefill_fragment_coalesced_scaled(acc50, xs_m[5], ws0, out, base_row + 5 * kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced_scaled(acc51, xs_m[5], ws1, out, base_row + 5 * kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc52, xs_m[5], ws2, out, + base_row + 5 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc53, xs_m[5], ws3, out, + base_row + 5 * kTileM, + base_col + 3 * kTileN, m, n, lane); store_prefill_fragment_coalesced_scaled(acc60, xs_m[6], ws0, out, base_row + 6 * kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced_scaled(acc61, xs_m[6], ws1, out, base_row + 6 * kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc62, xs_m[6], ws2, out, + base_row + 6 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc63, xs_m[6], ws3, out, + base_row + 6 * kTileM, + base_col + 3 * kTileN, m, n, lane); store_prefill_fragment_coalesced_scaled(acc70, xs_m[7], ws0, out, base_row + 7 * kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced_scaled(acc71, xs_m[7], ws1, out, base_row + 7 * kTileM, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc72, xs_m[7], ws2, out, + base_row + 7 * kTileM, + base_col + 2 * kTileN, m, n, lane); + store_prefill_fragment_coalesced_scaled(acc73, xs_m[7], ws3, out, + base_row + 7 * kTileM, + base_col + 3 * kTileN, m, n, lane); } // --------------------------------------------------------------------------- -// Generic scalar fallback: one thread per output element. Exact int32 dot -// over K, then fused x_scale * weight_scale, then bf16 store. Handles every -// unmatched (m, n, k), including all small-M API cases (M=2, M=16) and any -// M with the same (K, N). For the exact (k, n) == (2048, 6144) the weight is -// the n-major packed [N, K] layout (produced once out of the timed region), -// which the fallback decodes; every other (k, n) keeps the identity [K, N] -// layout. +// Generic scalar fallback: one thread per output element. Exact int32 dot +// over K in ascending k order, then fused x_scale * weight_scale, then bf16 +// store. Handles every unmatched (m, n, k), including all small-M API cases +// (M=2, M=16) and any M with the same (K, N). For the exact +// (k, n) == (2048, 6144) the packed_weight buffer holds the n-major transpose +// packed[col * K + kk] = raw[kk * N + col] (produced once outside the timed +// region), which this fallback decodes; every other (k, n) keeps the identity +// [K, N] row-major layout. // --------------------------------------------------------------------------- __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( const int8_t* __restrict__ x_q, @@ -1178,12 +671,6 @@ __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; - // Iteration 5: for the exact (k, n) == (2048, 6144) the packed_weight - // buffer holds the n-major transpose packed[col*k + kk] = raw[kk*n + col] - // (produced once out of the timed region), so the fallback decodes that - // layout there (keeps the paired M=2/M=16 and M=3072 API shapes with the - // same (K, N) byte-exact); every other (k, n) keeps the identity [K, N] - // row-major copy. const bool packed_nmajor = (k == kTargetK && n == kTargetN); const int8_t* b_ptr = packed_nmajor ? weight + static_cast(col) * k : weight + col; @@ -1198,16 +685,13 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( } // --------------------------------------------------------------------------- -// Identity packing (bootstrap; outside the timed region and outside Graph -// capture; the packed buffer keeps the same byte count K*N and the same -// allocated address, so the layout is graph-stable). launch_pack_w8a8_weight -// is an identity device-to-device copy for every (k, n): the packed weight -// is byte-identical to the raw [K, N] weight and the packed scale is -// byte-identical to the raw [N] scale. Later Parallel explore rounds may -// replace these kernels with a real layout transform and must then update -// the matching GEMM interpretation; the identity D2D copy is the generic -// fallback for unmatched (k, n). +// Weight packing, executed once by the trusted binding before Graph capture +// and never inside launch_w8a8_gemm. The packed buffer keeps the same byte +// count as the raw [K, N] weight, so captured addresses are graph-stable. // --------------------------------------------------------------------------- + +// Generic identity device-to-device copy for every (k, n) that has no +// specialized packing. __global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, @@ -1219,14 +703,8 @@ __global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( } } -// Exact-shape n-major pack (iteration 5): for (k, n) == (2048, 6144) the -// weight is transposed once, outside the timed region and out of Graph -// capture, into packed[col * k + kk] = raw[kk * n + col] so every col_major -// B fragment is 8 contiguous K bytes in the packed buffer (and in the -// staged n-major LDS tile). Same byte count and same buffer as the -// identity pack, so captured addresses are unchanged; correctness is -// checked against the raw logical [K, N] weight. One thread per output -// byte; runs once during weight prep. +// Exact-shape n-major transpose for (k, n) == (2048, 6144): +// packed[col * k + kk] = raw[kk * n + col]. One thread per output byte. __global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, @@ -1244,6 +722,7 @@ __global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( } } +// The [N] weight scale needs no layout change. __global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( const float* __restrict__ raw, float* __restrict__ packed, @@ -1278,19 +757,17 @@ extern "C" void launch_w8a8_gemm( (void)workspace_bytes; auto* out_bf16 = reinterpret_cast(out); - // Explicit dispatch. The assigned shape (M=4096, N=6144, K=2048) takes the - // iteration-5 packed-B DUMMA 256x64 path (w8a8_dumma_256x64x64_packedb_kernel - // consumes the n-major packed [N, K] weight produced by - // launch_pack_w8a8_weight; the identity-layout kernels stay compiled for - // re-dispatch only together with an identity pack revert); every other + // Explicit exact-shape dispatch. The assigned shape (M=4096, N=6144, + // K=2048) takes the 256x128 packed-B DUMMA path, which consumes the n-major + // packed [N, K] weight produced by launch_pack_w8a8_weight. Every other // (m, n, k) - including small-M API cases (M=2, M=16) and M=3072 with the - // same (K, N) - takes the scalar fallback, which decodes the n-major - // packed layout when (k, n) == (2048, 6144) and the identity [K, N] - // layout otherwise. + // same (K, N) - takes the scalar fallback, which decodes the n-major packed + // layout when (k, n) == (2048, 6144) and the identity [K, N] layout + // otherwise. if (m == kTargetM && n == kTargetN && k == kTargetK) { - const dim3 grid(kTargetN / kBlockN, kTargetM / kBlockM256); + const dim3 grid(kTargetN / kBlockN, kTargetM / kBlockM); const dim3 block(kBlockThreads); - hipLaunchKernelGGL(w8a8_dumma_256x64x64_packedb_kernel, + hipLaunchKernelGGL(w8a8_dumma_256x128x64_packedb_kernel, grid, block, 0, stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); } else { @@ -1313,13 +790,12 @@ extern "C" void launch_pack_w8a8_weight( int k, int n, hipStream_t stream) { - // Iteration 5: the exact (k, n) == (2048, 6144) weight is packed once, - // outside the timed region and out of Graph capture, into the n-major - // [N, K] layout packed[col*k + kk] = raw[kk*n + col] (same byte count and - // same buffer -> graph-stable addresses) so every col_major B fragment is - // 8 contiguous K bytes; every other (k, n) keeps the identity - // device-to-device copy (the scalar fallback decodes the n-major layout - // when (k, n) matches and the identity layout otherwise). + // For the exact (k, n) == (2048, 6144) the logical [K, N] weight is + // transposed once, outside the timed region and outside Graph capture, into + // the n-major [N, K] layout packed[col * k + kk] = raw[kk * n + col] (same + // byte count and same caller-owned buffer -> graph-stable addresses), so + // every col_major B fragment is 8 contiguous K bytes; every other (k, n) + // keeps the identity device-to-device copy. constexpr int kBlock = 256; const dim3 block(kBlock); 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/glm52/TP8/M4096/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_down_proj.hip index a1b49be9..e07786c9 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_down_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_down_proj.hip @@ -1,313 +1,266 @@ -// @@variant shape=glm_tp8_shared_down_proj_m4096 commit=9b56135e1fec933006c477066840f9c84f497cef added=2026-08-31 -// median_us=278.4 p90_us=279.2 speedup=63.57 baseline_us=1.77e+04 -// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// @@variant shape=glm_tp8_shared_down_proj_m4096 commit=44018b50c8c0557effd9aa58807ccc6a69dead1f added=2026-09-14 +// median_us=146.7 p90_us=147.2 speedup=120.7 baseline_us=1.77e+04 +// source=test-9-11-1/iteration_002/glm_tp8_shared_down_proj_m4096 (retested A/B) // MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). -// worker_3 (iteration 3), assigned shapes: -// glm_tp8_shared_gate_up_proj_m4096 : (M, N, K) = (4096, 512, 6144) -// glm_tp8_shared_down_proj_m4096 : (M, N, K) = (4096, 6144, 256) // -// This file provides the two stable host launch symbols consumed by -// csrc/bindings.cpp: -// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's -// HIP stream; no allocation, packing, -// autotuning, or synchronization. Only the -// caller-provided out and workspace are -// touched (workspace is unused this round). -// * launch_pack_w8a8_weight(...) - out-of-timed-region packing. Identity -// device-to-device copy for every (k, n) -// pair EXCEPT the assigned gate_up weight -// (k, n) == (6144, 512), which is packed -// once as a [K, N] -> [N, K] transpose -// (n-major rows, k contiguous) so the timed -// GEMM consumes B fragments as -// k-contiguous bytes. The pack stays -// out-of-timed, out-of-Graph, same byte -// count, same graph-stable buffer. -// -// Bootstrap strategy (correctness-first, usable profiling baseline): -// * Native INT8 DUMMA m16n16k32 tiled kernel with int32 accumulation for -// every assigned shape with M >= 128 (both assigned shapes qualify). -// * 2-D macro-tile family (64x64, 64x128, 128x64) with one 32x32 quadrant -// per wavefront (four m16n16k32 int32 accumulators kept resident over the -// whole K loop). Single-buffered K loop (K stage 128, two barriers per -// stage): no split-K, no raw asm, no speculative double-buffer pipeline. +// worker_0 (iteration 2, operand-reuse round). Assigned shape: +// glm_tp8_shared_down_proj_m4096 : (M, N, K) = (4096, 6144, 256) // -// Round 1 (baseline): the 64x128 tile measured 517.4 us median / 49.81 TOPS; -// routing the exact gate_up shape to the 64x64 tile measured 432.5 us -// (accepted; 19.6% improvement). +// This file provides the two stable host launch symbols consumed by the +// trusted csrc/bindings.cpp: +// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's +// HIP stream. No allocation, no packing, no +// autotuning, no host/device synchronization +// and no default-stream launch. Only the +// caller-provided out is written; the +// caller-provided workspace is unused +// because there is no split-K. +// * launch_pack_w8a8_weight(...) - out-of-timed-region packing. The exact +// down_proj (k, n) == (256, 6144) pair is +// packed once into the round-5 panel layout +// [N/128][K/64][128][64] (n-major rows of +// 64 k-bytes per panel) so the timed GEMM +// consumes each 128x64 B stage tile as ONE +// contiguous 8 KiB region and each lane's +// fragment stays k-contiguous; every other +// (k, n) stays the identity device-to-device +// copy of the logical [K, N] weight. // -// Round 2 (operand-reuse, mandated decision): cooperative A+B LDS staging is -// kept and direct per-wave global loads are rejected (A is reused 64x per -// macro-tile via kBlockN and B 64x via kBlockM; direct loads would refetch -// both operands once per wavefront, 4x the 28.3 MB staged global traffic -> -// ~1.73M vmem_read instructions vs 432K, and the TP4 gate_up direct-A arm -// measured 7.9x slower). The focused change makes the staged B layout -// bank-safe for the fragment consumers: the gate_up weight is packed once -// [K,N] -> [N,K] (n-major, k contiguous) and B fragments are loaded -// col_major from a 144-B (128+16) row-strided LDS tile, replacing the -// byte-scattered row_major B reads on the identity [K,N] layout that are the -// 17.3M PMC bank-conflict source. Global loads stay 16-byte vectorized and -// coalesced on both operands. LDS stays 19,456 B/block -> 3 blocks/CU -// (58,368 B, unchanged; round-1's "4 blocks/CU" estimate was impossible on -// both LDS (4x19,456 = 77,824 B > 64 KiB) and VGPR (4x256x72 = 73,728 > -// 65,536)). down_proj keeps the identity 64x128 arm; every unmatched shape -// keeps the generic arms and the scalar fallback (which decodes the packed -// [N,K] layout when (k, n) == (6144, 512) so the paired M=2/M=16 gate_up -// tails stay correct). -// * A[kBlockM, kStageK] and B staged in bank-skewed LDS (row strides -// 144/80/144 = 16-byte-aligned non-power-of-two), consumed with the -// library du_load_matrix_sync loaders: A row_major (unchanged, load- -// bearing on every shape), B row_major for the identity [K, N] layout -// (down_proj + unmatched shapes) and B col_major for the gate_up -// [N, K] packed layout (round 2). +// Round 1 (bootstrap, accepted at 313.66 us median): native INT8 DUMMA +// m16n16k32 tiled kernel with exact int32 accumulation for every M >= 128. +// * 64x128 macro-tile: 8 wavefronts, one 32x32 quadrant per wavefront, four +// m16n16k32 int32 accumulators kept resident across the whole K loop. A and +// B are staged cooperatively into bank-skewed LDS and consumed with the +// library du_load_matrix_sync fragment loaders. +// * SINGLE-buffered K loop, K stage 128 (A[64,128] + B[128,128] = 27,648 +// B/block -> 2 blocks/CU): two barriers per stage, no split-K, no raw asm, +// no speculative double-buffer pipeline. K=256 is two stages. // * Fused direct fragment -> x_scale -> weight_scale -> bf16 epilogue using // the verified gfx928 accumulator ownership (row = lane & 15, -// col_mod4 = lane >> 4, x[i] -> columns col_mod4 + 4*i). The int32 dot is -// exact for every assigned shape (max |dot| = 6144*127*127 << 2^31), the -// int32 -> float conversion is exact, and the single fp32 multiply -// followed by one bfloat16 rounding matches the reference bit-for-bit. -// * Tail M rows are zero-filled on load and masked on store, so any -// M >= 128 is supported by the tiled path. -// * All other (m, n, k) - including every M < 128 API case and the paired -// M=2/M=16 shapes with the same (N, K) - go to a scalar int8/int32 -// fallback that decodes the [N, K] packed layout when (k, n) == -// (6144, 512) and the identity [K, N] layout otherwise. -// -// Round 3 (pipeline, mandated decision): compare single and double buffering -// across K tiles. The exact (4096, 512, 6144) guard now routes to -// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true>: a software- -// pipelined 64-K double-buffered K loop (prologue + one barrier per stage, -// 97 barriers for 96 stages) that issues stage s+1's global loads into the -// idle LDS buffer before the stage-s MMAC burst, overlapping the -// global-load latency that the round-2 single buffer serializes in front of -// every stage (ISA-verified order: global_load_dwordx4 -> ds_write_b128 -> -// s_waitcnt -> s_barrier -> ds_read2 x16 -> v_mmac x16 -> s_barrier, with -// the next stage's loads gated behind the previous stage's barrier). The -// 128-K double-buffered variant is arithmetically impossible at 3 blocks/CU -// (2 x 19,456 = 38,912 B/block x 3 = 116,736 > 64 KiB -> 1 block/CU, a -// harmful occupancy loss), so the pipelined comparison uses the 64-K -// double buffer: LDS 20,480 B/block x 3 = 61,440 <= 64 KiB and -// VGPR 64-72 x 256 x 3 <= 65,536, i.e. 3 blocks/CU preserved (no occupancy -// confound). Everything else is byte-identical: same int4 staging pattern -// (same 28.3 MB global traffic and ~432K vmem_read instructions), same -// k0-outer/kk-inner int32 order (stage size changes only regroup exact -// integer sums, so results stay bit-identical), same col_major B fragments -// (stride 80 = 64+16 bank skew), same fused epilogue, same pack, same -// generic single-buffered arms and scalar fallback. Decision rule (mandate): -// retain double buffering only if PMC shows reduced VMEM/LDS stall evidence -// (lds_wait_instructions drops below ~1.2M, lds_instructions stays ~1.97M, -// occupancy stays 3 blocks/CU) AND the median improves >2%; otherwise the -// round's evidence retains the single-buffered round-2 structure. -// -// Round 10 (fragment-load-form round, restarted from the accepted round-3 -// source after the round-9 infra kill): replace the library -// du_load_matrix_sync fragment loads of the exact-shape gate_up kernel with -// the lineage-validated explicit load_frag8 8-byte ds_read_b64 loader. The -// exact (4096, 512, 6144) guard now routes to -// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true, true> (new -// kDirectFrag template arm). The round-3 ISA shows every fragment load as -// ds_read2_b32 + per-byte v_and/v_or_b32_sdwa/v_or3 mask-OR "identity" -// reassembly (~9.4M of the 15.43M VALU instructions per replay); loading the -// same 8 bytes per lane as one int64 (ds_read_b64) removes that VALU and -// halves the fragment bank-conflict factor (32-lane dword phases at stride -// 80 -> 16-lane 8-byte phases). Element-to-slot fragment mapping, v_mmac -// operands, exact k0-outer/kk-inner int32 accumulation, staging/global -// traffic, barriers, tile, grid, LDS bytes, occupancy (3 blocks/CU) and the -// fused epilogue are unchanged; results stay bit-identical. The generic arms -// (<64,128,128>, <128,64,128>, <64,64,128>, all kBNMajor=false) keep the -// library loaders and byte-identical codegen/resources (if constexpr). -// -// Round 11 (LDS bank-skew round, HIP-only; ISA policy keeps raw asm -// disallowed because the HIP plateau rule - 8 valid HIP rounds within +/-2% -// of the then-current best - is unmet). The accepted round-10 kernel still -// shows PMC lds_bank_conflicts = 9,437,184, EXACTLY unchanged from the -// round-3 code object (the -9.07% round-10 win came from VALU removal, not -// from conflicts). 9,437,184 = 192 x 512 blocks x 96 stages: every -// block-stage pays 192 conflict cycles = 128 from the 16 ds_read2_b64 -// fragment reads + 64 from the 8 ds_write_b128 staging stores (the 16-B/lane -// store floor). The read half is a stride defect: at stride 80 = 20 dwords, -// 20*r mod 32 has period 8, so every 16-lane read phase groups rows r and -// r+8 onto the same bank pair - a 2-way conflict on ALL 4 sub-phases of -// every fragment read. The exact (4096, 512, 6144) guard now routes to -// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true, true, true> (new -// kSkew8 template arm) which stages both operands at the down_proj-validated -// 8-mod-16 row strides A = 88 (64+24) and B = 72 (64+8) instead of 80: -// 22*r and 18*r mod 32 are distinct over r = 0..15 and gcd(22,32) = -// gcd(18,32) = 2 does not divide 1, so every 16-lane read phase lands on 16 -// DISTINCT bank pairs - the fragment reads become conflict-free (128 of the -// 192 cycles removed; only the 2-way 16-B/lane staging-store floor of 64 -// cycles remains -> lds_bank_conflicts expected ~3.15M). The 8-mod-16 row -// starts forbid 16-byte staging stores, so each 16-byte int4 is staged as -// two 8-byte halves (ds_write2_b64), the validated down_proj staging form -// (its B tile is exactly stride 72; its A tile exactly stride 88). Global -// loads (int4, same 432,128 vmem_read), pack layout, barrier structure -// (97/block), tile/grid (512 blocks), LDS bytes (A[64,88] + B[64,72] x 2 = -// 20,480 B/block -> 3 blocks/CU unchanged), VGPR, the k0-outer/kk-inner -// int32 order, the element-to-slot fragment mapping and the fused epilogue -// are untouched; only the LDS row strides and the staging store granularity -// change, so results stay bit-identical. Generic arms (<64,128,128>, -// <128,64,128>, <64,64,128>) keep kSkew8 = false and byte-identical -// codegen/resources (if constexpr). +// col_mod4 = lane >> 4, frag.x[i] -> columns col_mod4 + 4*i). The int32 dot +// is exact for every allowed K and the fp32 left-associated multiply chain +// matches the mathematical reference. +// * A second 64x64 arm covers large-M shapes whose N is a multiple of 64 but +// not of 128. Every other (m, n, k) - including every M < 128 API case and +// the paired M=2/M=16 shapes with the same (N, K) - goes to the scalar +// int8/int32 fallback. // -// Round 16 (lead-2 loop-carried payload round, HIP-only; raw asm stays -// DISALLOWED - plateau=false, the recent-valid-improvement window -// [-34.44%, +2.82%, +9.97%] does not meet three valid HIP rounds within -// [-2%, +2%) of the then-current best, and no prior ISA-guided round -// recorded a compiler limitation with target instructions). The accepted -// round-11 exact code object (digest 9889820e, median 293.30 us / p90 -// 293.545) shows the software-pipeline intent is DEFEATED by the compiler's -// wait placement: per stage the steady-state loop is -// global_load_dwordx4 (A prefetch s+1) -> global_load_dwordx4 (B prefetch -// s+1) -> s_waitcnt vmcnt(1) -> s_waitcnt vmcnt(0) <- waits IMMEDIATELY -// -> 2 x ds_write2_b64 (staging stores) -> 4 x ds_read2_b64 (fragment -// reads, stage s) -> 8 x v_mmac CONTIGUOUS (stage s) -> s_barrier -> loop. -// i.e. the stage-(s+1) global loads are waited ~12-16 instructions after -// issue, BEFORE the stage-s MMAC burst, so the full global-load latency -// (~600-1000+ cycles under 120-CU L2 contention) sits on the per-stage -// critical path with zero overlap (the round-12 analysis' -// "prefetch global load -> ds_write2_b64 -> s_barrier -> ds_read2_b64 -> -// v_mmac" chain; PMC: 6.0M VALU + 1.18M LDS + 465K VMEM wave-instructions -// over 293 us x 120 CUs, i.e. ~95-99% issue-stalled). Iteration history -// rules out the other levers on this exact shape: tile aspect (iter 4 -// 32x64: 749.9 us; iter 8 64x128 w8: 395.4 us), 2-wave quadrants (iter 6: -// 629.0 us - wave count per CU is load-bearing), LDS-side fragment -// prefetch (iter 7: 337.97 us, flat/worse), stage 32 / 6 blocks per CU -// (iter 12: 447.4 us - per-stage fixed cost, not occupancy, is the lever). -// The one untested mechanism that directly removes the exposed vmcnt wait: -// a LEAD-2 LOOP-CARRIED REGISTER PAYLOAD (the TP4 o_proj lineage's -// validated "stage-top publish" form). The exact (4096, 512, 6144) guard -// now routes to w8a8_dumma_prefill_tiled_kernel<64,64,64,true,true,true, -// true,true> (new kPayload template arm): each thread keeps the one A int4 -// + one B int4 of stage s+2 in VGPR across the stage-(s+1) barrier, issues -// those global loads at the TOP of iteration s (no wait), and PUBLISHES -// the payload carried from iteration s-1 into the idle LDS buffer at the -// TOP of iteration s (before the stage-s MMAC burst). The first use of the -// payload registers is the next iteration's publish, and the s_barrier -// fences outstanding vmem, so the compiler can no longer place -// s_waitcnt vmcnt(0) immediately after the loads: the wait lands after the -// stage-s fragment reads + MMAC burst (immediately before the barrier), and -// the global-load latency overlaps the whole compute segment instead of -// serializing in front of it. One barrier per 64-K stage is preserved -// (prologue + 96 = 97/block); the publish targets the idle buffer (WAR -// retired by the previous barrier, RAW retired by the stage barrier), so no -// second barrier is introduced. Everything else is byte-identical: 64x64 -// tile / 256 threads / 4 waves / 32x32 quadrant per wave, grid (8,64) = -// 512 blocks, kSkew8 strides 88/72 with ds_write2_b64 staging (8 per -// block-stage), same 16 ds_read2_b64 fragment reads per block-stage (same -// lds_instructions ~1,179,648), same 20,480 B LDS/block -> 3 blocks/CU -// (payload costs 8 VGPRs: ~54 total, 54 x 256 x 3 = 41,472 <= 65,536; no -// spill risk at 3 blocks/CU), same int4 global loads (same ~432,128 -// vmem_read / 32,768 vmem_write), same k0-outer/kk-inner int32 order and -// the same element-to-slot fragment mapping - the payload changes only WHEN -// global bytes land in LDS, never which bytes or the accumulation order, so -// results stay bit-identical (mismatch 0 / max_abs_error 0.0 expected). -// Generic arms (<64,128,128>, <128,64,128>, <64,64,128>, all kBNMajor= -// false) keep kPayload = false and byte-identical codegen/resources -// (if constexpr). +// Round 2 (this session, mandated operand-reuse decision): compare direct +// per-wave global loads with cooperative A+B LDS staging on the exact +// (4096, 6144, 256) shape. +// * Reuse accounting for the 64x128 macro-tile with one 32x32 quadrant per +// wave: every A element participates in kBlockN = 128 output MACs and every +// B element in kBlockM = 64 MACs, so cooperative staging keeps ONE global +// read and ONE LDS write per element while direct per-wave fragment loads +// would re-read the A tile kBlockN/32 = 4 times and the B tile +// kBlockM/32 = 2 times from global (per 64-K stage: 12,288 B staged vs +// 32,768 B of fragment reads, a 2.67x global-traffic amplification on an +// already L2-resident operand stream). Direct loads are rejected on that +// amplification and on the sibling lineage's measured 7.9x direct-arm +// regressions. +// * The bootstrap exact-shape arm reads the identity [K, N] B layout with +// byte-scattered row_major library fragment loads: its exact code object +// and PMC show 3,686,400 LDS instructions, 7,864,320 LDS bank conflicts, +// 5,131,888 LDS waits and 17.4M VALU per operator replay at 313.66 us +// median (11% issue utilization). +// * Focused change: the exact (m, n, k) == (4096, 6144, 256) guard now routes +// to w8a8_dumma_proj_operand_reuse_kernel<64, 128, 64>, which keeps +// cooperative A+B LDS staging, keeps 16-byte vectorized coalesced global +// staging loads on both operands, and fixes the LDS consumer side: +// - the one-time pack transposes B to the n-major [N, K] layout, so the +// staged B tile is [kBlockN, kStageK] with k contiguous and each +// matrix_b col_major fragment lane owns 8 contiguous bytes (the +// identity-layout row_major B reads are the bank-conflict source); +// round 5 keeps the identical staged LDS bytes but re-orders the +// packed buffer into the [N/128][K/64][128][64] panel layout so each +// block-stage B tile is one contiguous global region (see below); +// - A and B fragments are read with the explicit 8-byte load_frag8 +// loader (ONE 8-byte LDS read per fragment) instead of the library +// per-byte loaders plus reassembly VALU; the element-to-slot mapping +// is identical, so the MMAC operands and accumulation are unchanged; +// - LDS row strides are the bank-safe 8-mod-16 values A = 88 (64+24) and +// B = 72 (64+8): 22*r and 18*r mod 32 are 16 distinct values over +// r = 0..15, so every 16-lane ds_read_b64 fragment phase lands on 16 +// distinct bank pairs. The 8-mod-16 row starts forbid 16-byte stores, +// so each staged int4 is written as two 8-byte halves; +// - 64-K STAGE single buffer: 2 barriers/stage (8 per block for K=256), +// LDS A[64,88] + B[128,72] = 14,848 B/block -> 2 blocks/CU (occupancy +// unchanged from the bootstrap's 27,648 B / 2 blocks/CU), 64x128 tile, +// 512 threads / 8 waves, grid (48, 64) = 3072 blocks. +// The k0-outer / kk-inner int32 accumulation order, the element-to-slot +// fragment mapping and the fused epilogue are unchanged, so results stay +// bit-identical (mismatch 0 / max_abs_error 0.0 expected). +// * Expected measured effect: lds_instructions 3,686,400 -> ~540,672, +// lds_bank_conflicts 7,864,320 -> ~1,179,648, lds_wait 5,131,888 -> +// ~162,664, VALU 17.4M -> ~6M, median 313.66 us -> ~275-290 us. +// * Every other shape keeps the bootstrap arms. The scalar fallback decodes +// the packed [N, K] layout for (k, n) == (256, 6144) so the paired M=2/M=16 +// API shapes stay correct, and the packed pair never reaches an +// identity-layout tiled arm. // -// Round 2 (this session, glm_tp8_shared_down_proj_m4096 exact-shape arm): -// the generic <64,128,128> identity-layout arm that served down_proj -// (4096, 6144, 256) measured 314.72 us median with PMC -// lds_bank_conflicts = 7,864,320, ~600 LDS instructions and heavy per-byte -// fragment-reassembly VALU per block-stage (library row_major loaders on -// the identity [K, N] B layout). The exact (4096, 6144, 256) guard now -// routes to w8a8_dumma_prefill_tiled_kernel<64, 128, 64, true, true, true, -// true> - the file's own lineage-validated gate_up machinery transplanted -// to down_proj: B packed once [K,N] -> [N,K] n-major (kBNMajor, same -// transpose kernel as gate_up, same byte count / graph-stable buffer, -// out-of-timed/out-of-Graph), 64-K double-buffered K loop (kDoubleBuffer, -// prologue + 4 stages for K=256, one barrier per stage), explicit 8-byte -// load_frag8 fragment reads (kDirectFrag, one ds_read_b64 per fragment, -// zero reassembly VALU) and the 8-mod-16 bank-skewed LDS strides A=88 / -// B=72 (kSkew8, down_proj-validated, conflict-free 16-lane fragment -// phases). LDS A[64,88] + B[128,72] x 2 = 29,696 B/block -> 2 blocks/CU -// (occupancy unchanged vs the 27,648-B single buffer); grid (48, 64) = -// 3072 blocks; the stage size 128->64 only regroups exact integer sums -// (k0-outer / kk-inner int32 order over k = 0,32,64,96,128,160,192,224 is -// unchanged) and load_frag8 fills the same x[0..7] slots with the same -// little-endian bytes, so results stay bit-identical (mismatch 0 / -// max_abs_error 0.0 expected). +// Round 5 (this session, mandated PACKING round): test one weight +// packing/swizzle that makes each DUMMA B tile vector-loadable and +// LDS-bank-safe, with the pack outside timing and the graph-stable packed +// layout validated. +// * The accepted round-2 arm stages its 128x64 B tile from the plain +// n-major [N, K] packed weight: per block-stage the 512 int4 loads are 16 +// n-rows x 4 int4 per 64-lane warp at a 256-B row stride, i.e. every warp +// instruction touches 16 DISTINCT 128-B lines with only 64 B used in each +// (the same line is completed by another stage of the same block). Still +// 16-byte vector loads, but scattered at line granularity. +// * Focused change: the one-time pack for (k, n) == (256, 6144) ONLY becomes +// the panel layout [N/128][K/64][128][64] (n-major 64-k-byte rows per +// panel, panels ordered n-tile outer / k-tile inner). Byte-exact index +// relation: packed[((nt*(k/64)+kt)*128+nn)*64+kk64] == +// raw[(nt*128+nn)*k + kt*64+kk64]. Each block's per-stage B tile +// (n0 == 128*blockIdx.x, k0 == 64*stage) is now ONE contiguous 8 KiB +// region, so a 64-lane warp instruction reads 1 KiB contiguous (full +// 128-B line utilization) instead of 16 half-used lines; the per-thread +// address stays one base + one local_col*kStageK + v*16 offset. +// * Everything else is byte-identical: same 64-K single-buffered stage with +// two barriers, same 64x128 tile / 512 threads / 8 waves / 32x32 quadrant +// per wave, same grid (48, 64) = 3072 blocks, same LDS A[64,88] + +// B[128,72] = 14,848 B/block -> 2 blocks/CU, same ds_write2_b64 8-byte-half +// staging stores and same conflict-free load_frag8 ds_read_b64 fragment +// reads (the LDS-bank-safe strides 88/72 and the staged LDS bytes are +// untouched - only the GLOBAL byte order of the packed buffer changes), +// same k0-outer/kk-inner int32 order, same fused epilogue -> mismatch 0 / +// max_abs_error 0.0 expected. The scalar fallback decodes the panel layout +// for (k, n) == (256, 6144) so the paired M=2/M=16 tails stay correct; +// every other (k, n) keeps the identity copy. Packing runs +// out-of-timed/out-of-Graph in launch_pack_w8a8_weight, same byte count +// (256*6144), same grid (n*k/16 threads), same graph-stable buffer. +// * Falsifiable gate: median_us < 273.2886219024658 AND p90_us < +// 275.1925849914551 (and mismatch 0 / graph capture true) -> the +// line-scattered B staging was binding and the panel pack wins; otherwise +// the plain n-major pack is retained and this round records that the B +// stream is L2-resident (90.4% hit, same 32 sectors/warp either way) and +// NOT binding on the latency-stalled critical path. // -// Round 3 (this session, glm_tp8_shared_down_proj_m4096 pipeline round): -// mandated single vs double buffering comparison across K tiles. The -// accepted round-2 arm (<64,128,64,true,true,true,true>, 286.53 us median / -// 286.64 us p90, digest cd9791d0) is the 64-K DOUBLE-buffered pipeline; its -// exact code object shows the software pipeline already works on this shape -// (steady-state loop: ds_write2_b64 publish s+1 -> global_load_dwordx4 x2 -// for s+2 -> ds_read2_b64 x4 (fragment reads s+1) -> v_mmac x8 -> s_waitcnt -// vmcnt(0) lgkmcnt(0) -> s_barrier: the stage-(s+2) loads stay in flight -// across the fragment reads + MMAC burst and are waited only after it, i.e. -// reduced VMEM stalls), and PMC shows no harmful LDS or occupancy growth vs -// the pre-round-2 single buffer (29,696 vs 27,648 B/block, 2 blocks/CU -// pinned; lds_wait 162,664 / lds_instructions 540,672 / conflicts -// 1,179,648). The (4096, 6144, 256) guard now routes to the SINGLE-buffered -// form <64,128,64,true,false,true,true>: identical stage size 64, identical -// kBNMajor/kDirectFrag/kSkew8/tile/grid/epilogue, kDoubleBuffer = false, so -// each stage serializes its global-load round trip (load -> wait -> -// ds_write2_b64) in front of the MMAC burst with two barriers per stage -// (8/block vs 5/block) and LDS halves to 14,848 B/block. The single-buffer -// branch now threads kSkew8 and kDirectFrag through its staging/MMAC calls -// (previously hardcoded to the library-loader defaults; generic arms keep -// byte-identical codegen via if constexpr). Falsifiable gate: if the -// single-buffered median is NOT strictly below 286.53368949890137 (p90 -// guard < 286.6385078430176), double buffering is retained with the ISA/PMC -// evidence above; if it IS below, the single buffer wins on this shape -// (LDS halved; VGPR may drop <= 42 and open 3 blocks/CU, the trusted -// occupancy-probe set [2, 3, 4]). Results stay bit-identical either way -// (same stage size, same k0-outer/kk-inner int32 order, same bytes). +// Round 6 (this session, mandated EPILOGUE round): fuse per-row x per-column +// scaling, bf16 conversion and a fully coalesced final store into the compute +// kernel and remove any unnecessary workspace/combine pass. +// * The accepted arm already fuses the scaling, conversion and store into +// w8a8_dumma_proj_operand_reuse_kernel, and the operator never allocates or +// reads the caller workspace (there is no split-K partial and no combine +// kernel), so there is no pass to remove. What the round-2..5 epilogue did +// NOT do is store coalesced: the gfx928 fragment ownership (row = lane & 15, +// col = (lane >> 4) + 4*i) makes each 64-lane *bf16* store instruction +// carry 16 rows x 4 columns, i.e. 16 distinct 128-B global lines with only +// 8 useful bytes in each. The exact code object confirms it: 16 +// global_store_short_d16_hi per thread, interleaved with ~20 scalar +// x_scale/weight_scale loads (the 21 static global loads) and s_waitcnt +// vmcnt chains, and the counter-derived write traffic is 57,966,336 B for a +// 50,331,648 B bf16 output (1.152x sector overhead) out of 61,997,440 B +// total per operator replay. +// * Focused change: the exact-arm epilogue now (a) cooperatively stages the +// 64 x_scale and 128 weight_scale values into the retired A LDS tile with +// one coalesced load per thread, (b) writes each thread's 16 scaled bf16 +// results into a 64-row x 64-column half of the retired B LDS tile with the +// bank-safe 8-mod-16 element stride 68 (136 B), and (c) after one barrier +// writes the half back with 16-byte stores - 8 threads cover one full +// 128-B output line, so a 64-lane store instruction moves 8 complete lines +// instead of 16 partial ones. Two column halves plus one scale-publish +// barrier (5 extra barriers per block); the staged bytes are the same +// values in the same logical [M, N] order, and the scale multiply order +// ((float)acc * x_scale * weight_scale) and the bf16 conversion are +// byte-identical, so results stay bit-exact (mismatch 0 / max_abs_error 0.0 +// expected). +// * The A/B staging buffers are dead after the K loop, so the epilogue reuses +// them: LDS stays 14,848 B/block -> 2 blocks/CU, the grid stays (48, 64) = +// 3072 blocks and the tile stays 64x128 / 512 threads / 8 waves. Every +// other shape keeps the bootstrap arms and the scalar fallback (including +// the shared pair M=2/M=16 down_proj tails) untouched. +// * Falsifiable gate: median_us < 270.556 us (>= 1% under the 273.2886219 us +// official best) AND p90_us <= 275.1925849914551, with the coalesced store +// cutting write_bytes_per_operator_replay from 57,966,336 toward the +// 50,331,648 B theoretical minimum and replacing 393,216 vmem_write +// instructions with ~98,304 (16 bytes each). If the median does not clear +// the 1% bar, the direct 2-byte-scatter epilogue is restored. // -// Round 4 (this session, glm_tp8_shared_down_proj_m4096 tile-shape round): -// the balanced 1:1 aspect probe <64,64,64,true,false,true,true> (256 thr = -// 4 waves, 6144 blocks, LDS 10,240 B/block -> 4 blocks/CU, doubled A -// staging) measured 297.15 us median / 297.97 us p90 and was REJECTED -// (guard restored to the round-3 single-buffered <64,128,64,true,false, -// true,true> winner, 281.65 us / 282.52 us): co-residency was already -// saturated at 2 blocks/CU and the +33% A-side staging cost was not repaid. +// Round 15 (final consolidate / HIP-only polish on the frozen winning arm): +// move the exact arm's scale delivery from the retired-A LDS slab to +// per-thread registers. +// * Frozen winning architecture (accepted best 158.0229902267456 us median / +// 158.49021911621094 us p90, source digest c548257a): 64x128 macro-tile, +// 512 threads / 8 waves with a 32x32 quadrant per wave, 64-K +// single-buffered two-barrier stages, round-5 panel pack +// [N/128][K/64][128][64], grid (48,64) = 3072 blocks, A[64,88] + B[128,72] +// = 14,848 B LDS -> 2 blocks/CU (arch_vgpr 48), round-6 LDS-staged +// coalesced fused epilogue. +// * Evidence this round acts on: the exact-source PMC (120 CU, 3072 x 512, +// 147.532 us profiled) reports lds_instructions 1,631,232 with +// lds_wait_instructions 1,486,046 at ~11% SIMD issue utilization, i.e. ~66 +// LDS instructions per wave of which ~60 stall; the epilogue's scale +// read-back alone is ~10 of those LDS reads per wave (x_scale[row + r] and +// the four weight_scale columns of each of the two column groups, once per +// accumulator fragment), all of them on the post-K-loop critical path. The +// per-block tail is otherwise serialized behind 5 barriers (scale publish + +// publish/WAR per 64-column half). +// * Focused change (one bounded mechanism, exact `(m,n,k) == (4096,6144,256)` +// arm only): each thread loads its OWN two row scales (rows +// local_row + (lane & 15) and local_row + kTileM + (lane & 15)) and eight +// column scales (columns local_col + (lane >> 4) + 4*i for the two kTileN +// column groups) into registers ONCE before the K loop, and +// stage_epilogue_fragment consumes those registers instead of reading +// xs_tile/ws_tile from the retired A tile. Consequences: (a) no +// x_scale/weight_scale is staged into or read back from LDS anywhere in the +// exact arm, (b) the per-block scale-publish barrier disappears (block +// barriers 13 -> 12: 8 K-loop + 4 epilogue), (c) the 10 global scale loads +// per thread are issued before the K loop, so their latency hides behind the +// four 64-K stages instead of sitting in front of the staging multiplies. +// * Byte-identical data path: the values are the same x_scale[m0 + row] / +// weight_scale[n0 + col], the multiply order +// ((float)acc * x_scale * weight_scale) and the bf16 conversion are +// unchanged, and every staged LDS address (stride 68) and every coalesced +// 16-byte global store address is unchanged, so mismatch 0 / max_abs_error +// 0.0 are expected. The guard launches exactly grid(n/128, m/64) with +// m == 4096 and n == 6144, so the registered rows/columns are always in +// range (m0 + local_row + kTileM + 15 <= 4095, n0 + local_col + kTileN + +// 3 + 12 <= 6139). +// * Everything else stays byte-identical: tile, grid, waves, K stage, +// single buffering, panel pack, A/B strides and vectorized staging, LDS +// 14,848 B, kernel name / launch ABI, all shape guards, the generic +// prefill arms and the scalar fallback. +// * Falsifiable prediction: median_us <= 158.0229902267456 (expected +// ~155.8-158.0 us) AND p90_us <= 158.49021911621094, with lds_instructions +// per wave down by ~10 and lds_wait down by a similar amount while +// arch_vgpr stays <= 64 (no spills) and the LDS footprint stays 14,848 B. +// If the median does not improve, the register scales were not on the +// critical path (the LDS broadcast reads were already cheap) and the +// round-6 LDS slab is restored. // -// Round 5 (this session, glm_tp8_shared_down_proj_m4096 PACKING round): test -// one weight packing/swizzle that makes each DUMMA B tile vector-loadable -// and LDS-bank-safe, with the packing outside timing and the graph-stable -// packed layout validated. The accepted round-3 arm stages its 128x64 B -// tile from the plain n-major [N,K] packed weight: per block-stage the 512 -// int4 loads are 16 rows x 4 int4 per warp at 256-B row stride, i.e. every -// warp instruction touches 16 DISTINCT 64-B segments (50% of each 128-B L2 -// line) - vector loads, but scattered at L2 granularity. The round changes -// the one-time pack for (k,n) == (256,6144) ONLY to the panel layout -// [N/128][K/64][128][64] (n-major 64-k-byte rows per panel; byte-exact index -// relation packed[((nt*(k/64)+kt)*128+nn)*64+kk64] == raw[(nt*128+nn)*k + -// kt*64+kk64], verified by simulation): each block's per-stage B tile is now -// ONE contiguous 8 KiB region, so every warp instruction reads 1 KiB -// contiguous (full sector/line utilization; per block-stage the B staging -// collapses from 128 half-used 128-B lines to 64 fully-used lines). The new -// exact-shape arm w8a8_dumma_prefill_tiled_kernel<64,128,64,true,false, -// true,true,false,true> (new kBPanel template flag, default false) is -// byte-identical in EVERYTHING else: same 64-K single-buffered stage, same -// kBNMajor/kDirectFrag/kSkew8, same 64x128 tile / 512 thr / 8 waves / 32x32 -// quadrant, same grid (48,64) = 3072 blocks, same LDS A[64,88]+B[128,72] = -// 14,848 B/block @ 2 blocks/CU, same ds_write2_b64 staging stores and same -// conflict-free load_frag8 ds_read2_b64 fragment reads (LDS-bank-safe -// preserved - only the GLOBAL byte order of the packed buffer changes, -// never the staged LDS bytes), same k0-outer/kk-inner int32 order, same -// fused epilogue -> mismatch 0 / max_abs_error 0.0 expected. The scalar -// fallback decodes the panel layout for (k,n) == (256,6144) so M<128 tails -// stay correct; gate_up (6144,512) keeps its plain [N,K] transpose pack and -// its own arm; generic arms are untouched (kBPanel=false). Packing runs -// out-of-timed/out-of-Graph in launch_pack_w8a8_weight, same byte count -// (256*6144), same graph-stable buffer. Expected effect: B global staging -// becomes perfectly coalesced 1 KiB-per-warp streams with half the L2 line -// count; PMC vmem_read_instructions unchanged (~614,400 - same 512 int4 -// loads/block-stage), lds_* identical (LDS layout untouched), l2_hits/l2_ -// misses shift slightly in the direction of fewer miss sectors, valu ~flat -// (address math per thread unchanged: one base + offset). FALSIFIABLE GATE: -// median_us < 281.6507148742676 AND p90_us < 282.5235176086426 -> B-global -// coalescing was binding on this shape and the panel pack wins; otherwise -// the plain n-major pack is retained and this round records that B reads -// are L2-resident (90.4% hit) and NOT binding on the LDS-bound critical -// path (consistent with the lineage rule that pack layout stops paying once -// fragment loads/strides are fixed). +// Round 17 (fix_build, final round): restore a compiling candidate on the +// frozen round-15 winning arm after round 16 produced no valid proposal (the +// worker process exited on a rate-limit crash before any build ran). The only +// landed correction is the round-16 dead-code cleanup of the exact arm's A +// staging helper, re-applied with a consistent API so the definition and its +// single call site agree: +// * stage_proj_tiles is instantiated ONLY by the exact arm, whose host guard +// requires m == kExactM and launches exactly grid(n/kBlockN, m/kBlockM), so +// every staged row m0 + local_row is in range and the per-int4 +// `global_row < m` compare/select/zero-fill was dead. The helper drops its +// now-unused `m` parameter (the call site drops the argument) and gains a +// static_assert that the exact arm covers whole kBlockM row tiles; every +// staged A byte is unchanged, so results stay bit-exact. +// * Deliberately NOT included: the round-13/16 final-half post-store barrier +// skip. That is a runtime change whose measured p90 regressed at the 158-us +// baseline; this round is a build fix, not a latency change. +// * Everything else is byte-identical to the accepted round-15 best: 64x128 +// macro-tile, 512 threads / 8 waves with a 32x32 quadrant per wave, 64-K +// single-buffered two-barrier stages, round-5 panel pack +// [N/128][K/64][128][64], round-15 register scales, round-6 LDS-staged +// coalesced epilogue (both epilogue barriers kept), A[64,88] + B[128,72] = +// 14,848 B LDS, grid (48,64) = 3072 blocks, 2 blocks/CU, kernel name / +// launch ABI, every shape guard and the generic scalar fallback. +// * Falsifiable prediction: build_success true, correctness exact (mismatch 0 +// / max_abs_error 0.0), graph capture true, median_us flat at the accepted +// best 148.41579914093018 us (the removed compare/select was not on the +// critical path), p90_us <= 148.8046169281006, arch_vgpr <= 64, +// scratch_bytes 0 and LDS 14,848 B. // -// 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. +// Header order is fixed by the control plane: hip_runtime first (du_mma.h is +// not self-contained before it), hip_bfloat16 second, du_mma.h last. #include #include @@ -325,12 +278,32 @@ constexpr int kTileK = 32; // gfx928 wavefront is 64 lanes; every blockDim must be a multiple of 64. constexpr int kWaveSize = 64; -// Large-prefill dispatch geometry. Each macro-tile config uses a 32x32 -// quadrant per wave, so waves per block = (BlockM/32) * (BlockN/32). +// Large-prefill dispatch geometry. Each macro-tile config owns a 32x32 +// quadrant per wavefront, so waves per block = (BlockM/32) * (BlockN/32). constexpr int kPrefillMinM = 128; -constexpr int kPrefillStageK = 128; constexpr int kDefaultBlockM = 64; -constexpr int kDefaultBlockN = 64; +constexpr int kDefaultBlockN = 128; +constexpr int kDefaultStageK = 128; + +// Round 2 exact-shape arm: the assigned (4096, 6144, 256) shape only. +constexpr int kExactM = 4096; +constexpr int kExactN = 6144; +constexpr int kExactK = 256; +constexpr int kProjBlockM = 64; +constexpr int kProjBlockN = 128; +constexpr int kProjStageK = 64; +// The exact-shape arm uses the bank-safe 8-mod-16 LDS row strides A = 88 +// (kProjStageK + 24) and B = 72 (kProjStageK + 8); see the round-2 header +// comment. The stride scheme assumes kStageK == 64. + +// Round 6 (epilogue round) LDS staging geometry for the fused epilogue: the +// 64x128 bf16 output tile is staged one 64-column half at a time (kEpiCols) +// into the retired B staging buffer with the bank-safe 8-mod-16 element row +// stride kEpiStride = 68 bf16 = 136 B (16 row starts land on 16 distinct LDS +// banks and the stride stays 8-byte aligned for the int64 read-back). The +// retired A staging buffer holds the 64 x_scale and 128 weight_scale values. +constexpr int kEpiCols = 64; +constexpr int kEpiStride = 68; constexpr int kScalarThreads = 256; @@ -338,8 +311,9 @@ using namespace du::dumma; // Verified gfx928 INT8 DUMMA accumulator ownership: // row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> columns -// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store -// bf16 directly from the fragment (no accumulator LDS round trip). +// col_mod4 + 4*i. +// Scale by x_scale[row] and weight_scale[col] and store bf16 directly from the +// fragment (no accumulator LDS round trip). Tail rows (row >= m) are masked. template __device__ __forceinline__ void store_prefill_fragment( const AccFragment& frag, @@ -360,61 +334,37 @@ __device__ __forceinline__ void store_prefill_fragment( #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 scaled = + static_cast(frag.x[i]) * xs * weight_scale[col]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); } } -// Round 10 (fragment-load-form round): explicit 8-byte fragment loader -// replacing the library du_load_matrix_sync for the exact-shape gate_up -// kernel. Validated on the sibling TP8 down_proj accepted kernel (iter 8, -// -25.9%: 230.91 -> 171.19 us) and the TP4 gate_up lineage (iter 13 -// load_fragment8, -21%: 280.05 -> 220.71 us). For BOTH the library row_major -// matrix_a and col_major matrix_b m16n16k32 int8 loaders, lane l -> row -// (l & 15), k-quarter (l >> 4), and the loader fetches the same eight -// contiguous bytes at p[row*ldm + (l>>4)*8 .. +7] and stores them into the -// same x[0..7] slots in the same little-endian order (the round-3 ISA shows -// the compiler lowering this to ds_read2_b32 followed by a per-byte -// v_and/v_or_b32_sdwa/v_or3 mask-OR "identity" reassembly - ~9.4M dead-ish -// VALU instructions per replay, ~61% of the 15.43M VALU total). Loading the -// same 8 bytes as one int64 forces the 8-byte ds_read_b64 form instead: -// (a) the mask/OR reassembly VALU disappears, and (b) the LDS access -// changes from ds_read2_b32 (32-lane dword phases - at the stride-80 row -// pattern every phase groups rows 0-15 of k-quarter q with rows 0-15 of -// k-quarter q+1 onto the same 16 bank phases, a 4-way conflict per phase) -// to ds_read_b64 (16-lane phases: rows 0-15 land on 16 bank-pairs -// {20*r mod 32} = 8 distinct pairs hit twice - a 2-way conflict per phase, -// half of the current fragment-conflict factor). The element-to-slot -// mapping is unchanged (f.x[0..7] = the same 8 bytes in the same -// little-endian order), the v_mmac operands and the exact k0-outer/kk-inner -// int32 accumulation sequence are untouched, and the staging stores, LDS -// strides, barriers, tile, grid, occupancy and epilogue are byte-identical. +// Round 2 exact-shape fragment loader. The library row_major matrix_a and +// col_major matrix_b m16n16k32 int8 loaders both map lane l -> row (l & 15), +// k-quarter (l >> 4), and fetch the eight contiguous bytes at +// p[row*ldm + (l >> 4)*8 .. +7] into the fragment's x[0..7] slots in the same +// little-endian order. Loading those 8 bytes as one int64 emits a single +// 8-byte LDS read instead of the library's byte-wise loads plus per-byte +// reassembly VALU; the element-to-slot mapping is identical, so the MMAC +// operands and the exact int32 accumulation 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; +__device__ __forceinline__ void load_frag8(Frag& f, + const int8_t* p, + unsigned ldm, + int lane) { + const unsigned row = static_cast(lane) & 0xf; + const unsigned kq = static_cast(lane) >> 4; const int64_t v = *reinterpret_cast(p + row * ldm + (kq << 3)); reinterpret_cast(f.x)[0] = v; } -// Stage one K tile (k0) of A[kBlockM, kStageK] and B into the caller-offset -// LDS buffers (bank-skewed row strides; int4 vectorized coalesced global -// loads, zero-filled tail M rows). kBNMajor == true stages B from the packed -// [N, K] n-major gate_up layout (n-major rows, k contiguous); otherwise from -// the identity [K, N] layout. Shared by the single-buffered and -// double-buffered K loops; the buffer offset (0 or kABytes/kBMaxBytes) is -// applied by the caller so one body serves both pipeline forms. -// -// kSkew8 (round 11, exact-shape arm only): A and B row strides become the -// down_proj-validated 8-mod-16 values 88 (kStageK + 24) and 72 (kStageK + 8) -// so every 16-lane fragment-read phase lands on 16 distinct bank pairs (see -// the round-11 header comment). The 8-mod-16 row starts forbid 16-byte -// stores, so each 16-byte int4 chunk is staged as two 8-byte halves -// (ds_write2_b64); the global int4 loads are byte-identical. -template +// Stage one K tile (k0) of A[kBlockM, kStageK] (row-major, zero-filled tail M +// rows) and B[kStageK, kBlockN] (identity [K, N] layout) into LDS with +// 16-byte-vectorized coalesced global loads. The B row stride kBlockN + 16 is +// the non-power-of-two bank skew that spreads the row_major fragment reads. +template __device__ __forceinline__ void stage_prefill_tiles( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, @@ -428,20 +378,18 @@ __device__ __forceinline__ void stage_prefill_tiles( int k, int k0) { constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; - // 16-byte bank skew: keeps int4 staging stores aligned while spreading the - // fragment rows across distinct LDS banks. - constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); - constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows - constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); - static_assert(!kSkew8 || kStageK == 64, - "kSkew8 strides 88/72 assume kStageK == 64"); + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kBlockN + sizeof(int4); constexpr int kAInt4PerRow = kStageK / sizeof(int4); constexpr int kAInt4 = kBlockM * kAInt4PerRow; constexpr int kBInt4PerRow = kBlockN / sizeof(int4); constexpr int kBInt4 = kStageK * kBInt4PerRow; - constexpr int kBNInt4PerRow = kStageK / sizeof(int4); - constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; - // Stage A[kBlockM, kStageK] into LDS (zero-filled tail M rows). + static_assert(kAStride % sizeof(int4) == 0, + "A LDS row stride must stay int4 aligned"); + static_assert(kBStride % sizeof(int4) == 0, + "B LDS row stride must stay int4 aligned"); + + // Stage A[kBlockM, kStageK]; rows beyond m contribute zeros. for (int vec = tid; vec < kAInt4; vec += kThreads) { const int local_row = vec / kAInt4PerRow; const int v = vec - local_row * kAInt4PerRow; @@ -452,186 +400,163 @@ __device__ __forceinline__ void stage_prefill_tiles( x_q + static_cast(global_row) * k + k0 + v * static_cast(sizeof(int4))) : int4{0, 0, 0, 0}; - if constexpr (kSkew8) { - // 8-mod-16 row starts: write the 16 bytes as two 8-byte halves so every - // store stays 8-byte aligned (ds_write2_b64 staging; same LDS bytes). - int8_t* dst = - a_tile + local_row * kAStride + v * static_cast(sizeof(int4)); - reinterpret_cast(dst)[0] = - reinterpret_cast(&val)[0]; - reinterpret_cast(dst)[1] = - reinterpret_cast(&val)[1]; - } else { - reinterpret_cast(a_tile + local_row * kAStride)[v] = val; - } + reinterpret_cast(a_tile + local_row * kAStride)[v] = val; } - if constexpr (kBNMajor) { - // Stage B[kBlockN, kStageK] from the packed [N, K] weight (n-major: - // row = n_local, k contiguous). Same vectorized coalesced int4 pattern - // as the identity arm, one per 16 consecutive k bytes. - // - // kBPanel (round 5, down_proj exact-shape arm only): the one-time pack - // for (k,n) == (256,6144) is the panel layout [N/128][K/64][128][64] - - // each block's per-stage B tile is one contiguous 8 KiB region, so a - // 64-lane warp instruction reads 1 KiB contiguous (full L2 line/sector - // utilization) instead of 16 rows x 64 B scattered at 256-B stride - // (50% of each 128-B line). The per-thread address is still one base + - // offset: packed byte offset = panel_base + local_col*64 + v*16 where - // panel_base = (n0/128)*(k/64)*8192 + (k0/64)*8192. The LDS stores are - // untouched (same stride-72 ds_write2_b64 halves -> LDS-bank-safe - // staging and conflict-free fragment reads preserved); only the global - // byte order of the loaded int4 changes, so the staged LDS bytes are - // identical to the plain n-major arm (bit-identical results). - for (int vec = tid; vec < kBNInt4; vec += kThreads) { - const int local_col = vec / kBNInt4PerRow; - const int v = vec - local_col * kBNInt4PerRow; - int4 val; - if constexpr (kBPanel) { - const int64_t panel_base = - (static_cast(n0) / kBlockN) * (k / kStageK) * - (kBlockN * kStageK) + - (k0 / kStageK) * (kBlockN * kStageK); - val = *reinterpret_cast( - weight + panel_base + - static_cast(local_col) * kStageK + - v * static_cast(sizeof(int4))); - } else { - val = *reinterpret_cast( - weight + static_cast(n0 + local_col) * k + k0 + + + // Stage B[kStageK, kBlockN] from the identity [K, N] weight (n%kBlockN == 0 + // is guaranteed by dispatch, so there is no N tail). + for (int vec = tid; vec < kBInt4; vec += kThreads) { + const int kk = vec / kBInt4PerRow; + const int v = vec - kk * kBInt4PerRow; + reinterpret_cast(b_tile + kk * kBStride)[v] = + *reinterpret_cast( + weight + static_cast(k0 + kk) * n + n0 + v * static_cast(sizeof(int4))); - } - if constexpr (kSkew8) { - int8_t* dst = - b_tile + local_col * kBNStride + - v * static_cast(sizeof(int4)); - reinterpret_cast(dst)[0] = - reinterpret_cast(&val)[0]; - reinterpret_cast(dst)[1] = - reinterpret_cast(&val)[1]; - } else { - reinterpret_cast(b_tile + local_col * kBNStride)[v] = val; - } - } - } else { - // Stage B[kStageK, kBlockN] into LDS from the identity [K, N] weight. - for (int vec = tid; vec < kBInt4; vec += kThreads) { - const int kk = vec / kBInt4PerRow; - const int v = vec - kk * kBInt4PerRow; - reinterpret_cast(b_tile + kk * kBStride)[v] = - *reinterpret_cast( - weight + static_cast(k0 + kk) * n + n0 + - v * static_cast(sizeof(int4))); - } } } -// Round 16 (lead-2 loop-carried payload, exact-shape arm only): the register -// half of the lead-2 software pipeline. Each thread prefetches exactly one A -// int4 and one B int4 of a later K stage into thread-local registers and -// returns WITHOUT waiting (the vmcnt wait is emitted by the compiler at the -// first use, one iteration later). The <64,64,64> arm stages kStageK * 64 -// bytes per operand per stage, i.e. kAInt4 == kBNInt4 == kThreads == 256: -// one int4 per thread per operand, so the whole payload costs 8 VGPRs. The -// (row, v) mapping is the same as stage_prefill_tiles, so -// publish_stage_payload writes the exact same LDS addresses the staging -// loop would (bit-identical staging bytes, just loaded one stage earlier). -template -__device__ __forceinline__ void prefetch_stage_payload( +// One wavefront's 32x32-quadrant MMAC burst for the K tile staged at +// a_tile/b_tile. k0-outer / kk-inner int32 accumulation order; the int8 dot is +// exact in int32 for every allowed K. +template +__device__ __forceinline__ void mmac_prefill_stage( + DUFragment& + a_frag0, + DUFragment& + a_frag1, + DUFragment& acc00, + DUFragment& acc01, + DUFragment& acc10, + DUFragment& acc11, + int8_t* a_tile, + int8_t* b_tile, + int local_row, + int local_col) { + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kBlockN + sizeof(int4); + DUFragment + b_frag0, b_frag1; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync( + a_frag0, a_tile + local_row * kAStride + kk, kAStride); + du_load_matrix_sync( + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + du_load_matrix_sync( + b_frag0, b_tile + kk * kBStride + local_col, kBStride); + du_load_matrix_sync( + b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + } +} + +// Round 2 exact-shape staging: A[kBlockM, kStageK] from the row-major [M, K] +// activation (zero-filled tail rows) and B[kBlockN, kStageK] from the packed +// weight, both with 16-byte vectorized coalesced global int4 loads. The LDS +// row strides are the bank-safe 8-mod-16 values A = 88 / B = 72; 8-mod-16 row +// starts forbid 16-byte stores, so each int4 is written as two 8-byte halves +// (ds_write_b64 pair, same LDS bytes as a single 16-byte store would produce). +// +// kBPanel (round 5, exact-shape arm only, default true): B is read from the +// panel layout [N/128][K/64][128][64] instead of the plain n-major [N, K] +// transpose. Within a panel (nt = n0/128, kt = k0/64) the B rows are +// kStageK = 64 bytes apart with k contiguous, so local_col*kStageK + v*16 == +// vec*16 and the whole 64-lane warp instruction reads kBlockN*kStageK/8 = 1 KiB +// of consecutive bytes. The staged LDS bytes are byte-identical to the plain +// form (the same int4 of the same logical [K, N] element block lands at the +// same LDS address), so the LDS-bank-safe fragment side is untouched. +template +__device__ __forceinline__ void stage_proj_tiles( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, + int8_t* a_tile, + int8_t* b_tile, int tid, int m0, int n0, - int m, int k, - int k0, - int4& a_payload, - int4& b_payload) { + int k0) { constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; - constexpr int kAInt4PerRow = kStageK / static_cast(sizeof(int4)); + constexpr int kAStride = kStageK + 24; + constexpr int kBStride = kStageK + 8; + constexpr int kAInt4PerRow = kStageK / sizeof(int4); constexpr int kAInt4 = kBlockM * kAInt4PerRow; - constexpr int kBNInt4PerRow = kStageK / static_cast(sizeof(int4)); - constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; - static_assert(kAInt4 == kThreads && kBNInt4 == kThreads, - "payload arm requires exactly one int4 per thread per " - "operand (exact-shape <64,64,64> arm)"); - static_assert(kBNMajor, "payload arm stages the packed n-major B layout"); - // Same vec -> (local_row, v) mapping as stage_prefill_tiles, so the - // publish half stores to the same LDS addresses. - const int a_local_row = tid / kAInt4PerRow; - const int a_v = tid - a_local_row * kAInt4PerRow; - const int global_row = m0 + a_local_row; - // Tail M rows are zero-filled exactly as in stage_prefill_tiles. - a_payload = global_row < m - ? *reinterpret_cast( - x_q + static_cast(global_row) * k + k0 + - a_v * static_cast(sizeof(int4))) - : int4{0, 0, 0, 0}; - const int b_local_col = tid / kBNInt4PerRow; - const int b_v = tid - b_local_col * kBNInt4PerRow; - b_payload = *reinterpret_cast( - weight + static_cast(n0 + b_local_col) * k + k0 + - b_v * static_cast(sizeof(int4))); -} + constexpr int kBInt4PerRow = kStageK / sizeof(int4); + constexpr int kBInt4 = kBlockN * kBInt4PerRow; + static_assert(kAStride % 8 == 0 && kBStride % 8 == 0, + "8-mod-16 strides must stay 8-byte aligned"); + static_assert(kAInt4 <= kThreads && kBInt4 <= kThreads, + "one staging pass per thread is enough for this arm"); + static_assert(!kBPanel || (kBlockN == kProjBlockN && kStageK == kProjStageK), + "panel pack is specialized to the 128x64 B stage tile"); + static_assert(kExactM % kBlockM == 0, + "exact-arm guard launches whole kBlockM row tiles, so no M " + "tail can reach this staging helper"); -// Round 16: the LDS half of the lead-2 payload pipeline. Writes the payload -// carried in registers (loaded by prefetch_stage_payload one stage earlier) -// into the caller-offset LDS tile with the exact same 8-byte-half -// ds_write2_b64 stores (kSkew8 strides 88/72) that stage_prefill_tiles -// emits, so each thread republishes exactly the bytes it loaded. -template -__device__ __forceinline__ void publish_stage_payload( - int4 a_payload, - int4 b_payload, - int8_t* a_tile, - int8_t* b_tile, - int tid) { - constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; - constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); - constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); - constexpr int kAInt4PerRow = kStageK / static_cast(sizeof(int4)); - constexpr int kAInt4 = kBlockM * kAInt4PerRow; - constexpr int kBNInt4PerRow = kStageK / static_cast(sizeof(int4)); - constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; - static_assert(kAInt4 == kThreads && kBNInt4 == kThreads, - "payload arm requires exactly one int4 per thread per " - "operand (exact-shape <64,64,64> arm)"); - static_assert(kBNMajor, "payload arm stages the packed n-major B layout"); - const int a_local_row = tid / kAInt4PerRow; - const int a_v = tid - a_local_row * kAInt4PerRow; - int8_t* a_dst = - a_tile + a_local_row * kAStride + a_v * static_cast(sizeof(int4)); - if constexpr (kSkew8) { - // 8-mod-16 row starts: write the 16 bytes as two 8-byte halves - // (ds_write2_b64 staging, identical to stage_prefill_tiles). - reinterpret_cast(a_dst)[0] = - reinterpret_cast(&a_payload)[0]; - reinterpret_cast(a_dst)[1] = - reinterpret_cast(&a_payload)[1]; - } else { - reinterpret_cast(a_dst)[0] = a_payload; + // Stage A[kBlockM, kStageK]. The exact-arm host guard requires m == kExactM + // and launches exactly grid(n/kBlockN, m/kBlockM), so every staged row + // m0 + local_row is in range and the former `global_row < m` + // compare/select/zero-fill was dead code (round 17 fix_build). Every staged + // A byte is identical to the zero-filled form. + for (int vec = tid; vec < kAInt4; vec += kThreads) { + const int local_row = vec / kAInt4PerRow; + const int v = vec - local_row * kAInt4PerRow; + const int4 val = *reinterpret_cast( + x_q + static_cast(m0 + local_row) * k + k0 + + v * static_cast(sizeof(int4))); + int8_t* dst = + a_tile + local_row * kAStride + v * static_cast(sizeof(int4)); + reinterpret_cast(dst)[0] = + reinterpret_cast(&val)[0]; + reinterpret_cast(dst)[1] = + reinterpret_cast(&val)[1]; } - const int b_local_col = tid / kBNInt4PerRow; - const int b_v = tid - b_local_col * kBNInt4PerRow; - int8_t* b_dst = - b_tile + b_local_col * kBNStride + b_v * static_cast(sizeof(int4)); - if constexpr (kSkew8) { - reinterpret_cast(b_dst)[0] = - reinterpret_cast(&b_payload)[0]; - reinterpret_cast(b_dst)[1] = - reinterpret_cast(&b_payload)[1]; - } else { - reinterpret_cast(b_dst)[0] = b_payload; + + // Stage B[kBlockN, kStageK] from the packed weight: the round-5 panel + // layout [N/128][K/64][128][64] (one contiguous kBlockN*kStageK = 8 KiB + // region per timed block-stage, fully coalesced 1 KiB per 64-lane warp) or + // the round-2 plain n-major [N, K] transpose. Both forms place the same + // bytes at the same LDS address, so only the global read side changes. + for (int vec = tid; vec < kBInt4; vec += kThreads) { + const int local_col = vec / kBInt4PerRow; + const int v = vec - local_col * kBInt4PerRow; + int4 val; + if constexpr (kBPanel) { + // Panel (nt = n0/kBlockN, kt = k0/kStageK) base in bytes. Within a + // panel the local B rows are kStageK bytes apart and k is contiguous, + // so local_col*kStageK + v*16 == vec*16. + const int64_t panel_base = + (static_cast(n0) / kBlockN) * (k / kStageK) * + (static_cast(kBlockN) * kStageK) + + (static_cast(k0) / kStageK) * + (static_cast(kBlockN) * kStageK); + val = *reinterpret_cast( + weight + panel_base + static_cast(local_col) * kStageK + + v * static_cast(sizeof(int4))); + } else { + // Plain n-major [N, K]: packed[col * k + kk] == raw[kk * n + col]. + val = *reinterpret_cast( + weight + static_cast(n0 + local_col) * k + k0 + + v * static_cast(sizeof(int4))); + } + int8_t* dst = + b_tile + local_col * kBStride + v * static_cast(sizeof(int4)); + reinterpret_cast(dst)[0] = + reinterpret_cast(&val)[0]; + reinterpret_cast(dst)[1] = + reinterpret_cast(&val)[1]; } } -// One wave's 32x32-quadrant MMAC burst for the K tile staged at a_tile/b_tile -// (caller-offset to the target buffer). k0-outer / kk-inner int32 accumulation -// order; stage-size changes only regroup exact integer sums (max |dot| = -// 6144*127*127 << 2^31), so results stay bit-identical. -template -__device__ __forceinline__ void mmac_prefill_stage( +// Round 2 exact-shape MMAC burst: eight explicit 8-byte load_frag8 fragment +// reads (A row_major / B col_major, bank-safe 88/72 strides) and the same +// eight m16n16k32 MMACs in the same k0-outer / kk-inner int32 order the +// library-loader bootstrap used, so the accumulation is bit-identical. +template +__device__ __forceinline__ void mmac_proj_stage( DUFragment& a_frag0, DUFragment& @@ -643,127 +568,257 @@ __device__ __forceinline__ void mmac_prefill_stage( int8_t* a_tile, int8_t* b_tile, int local_row, - int local_col) { - // kSkew8 (round 11, exact-shape arm only): A 88 / B 72 row strides; the - // load_frag8 ldm arguments below pick them up unchanged (same 8-byte - // aligned fragment reads, conflict-free 16-lane bank phases). - constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); - constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows - constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); - if constexpr (kBNMajor) { - if constexpr (kDirectFrag) { - // Round 10 exact-shape path: explicit load_frag8 8-byte ds_read_b64 - // fragment reads replacing the library loaders (see load_frag8 - // comment). Both steps of the 64-K stage are unrolled into - // independent fragment register sets so all eight reads of a stage - // issue before the MMACs; the kk = 0 MMAC group runs first, then the - // kk = kTileK group - the exact k0-outer/kk-inner int32 order and the - // element-to-slot fragment mapping of the library path are preserved - // (f.x[0..7] = the same 8 bytes in the same little-endian order), so - // results stay bit-identical. - static_assert(kStageK == 2 * kTileK, - "kDirectFrag two-step unroll assumes kStageK == 2*kTileK"); - DUFragment - a_frag0_1, a_frag1_1; - DUFragment - b_frag0, b_frag1, b_frag0_1, b_frag1_1; - load_frag8(a_frag0, a_tile + local_row * kAStride, kAStride); - load_frag8(a_frag1, a_tile + (local_row + kTileM) * kAStride, kAStride); - load_frag8(b_frag0, b_tile + local_col * kBNStride, kBNStride); - load_frag8(b_frag1, b_tile + (local_col + kTileN) * kBNStride, - kBNStride); - load_frag8(a_frag0_1, a_tile + local_row * kAStride + kTileK, - kAStride); - load_frag8(a_frag1_1, - a_tile + (local_row + kTileM) * kAStride + kTileK, kAStride); - load_frag8(b_frag0_1, b_tile + local_col * kBNStride + kTileK, - kBNStride); - load_frag8(b_frag1_1, - b_tile + (local_col + kTileN) * kBNStride + kTileK, - kBNStride); - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); - du_mma_sync(acc00, a_frag0_1, b_frag0_1, acc00); - du_mma_sync(acc01, a_frag0_1, b_frag1_1, acc01); - du_mma_sync(acc10, a_frag1_1, b_frag0_1, acc10); - du_mma_sync(acc11, a_frag1_1, b_frag1_1, acc11); - } else { - DUFragment - b_frag0, b_frag1; + int local_col, + int lane) { + constexpr int kAStride = kStageK + 24; + constexpr int kBStride = kStageK + 8; + static_assert(kStageK == 2 * kTileK, + "exact-shape arm unrolls the two 32-K steps of a 64-K stage"); + DUFragment + a_frag0_1, a_frag1_1; + DUFragment + b_frag0, b_frag1, b_frag0_1, b_frag1_1; + load_frag8(a_frag0, a_tile + local_row * kAStride, kAStride, lane); + load_frag8(a_frag1, a_tile + (local_row + kTileM) * kAStride, kAStride, lane); + load_frag8(b_frag0, b_tile + local_col * kBStride, kBStride, lane); + load_frag8(b_frag1, b_tile + (local_col + kTileN) * kBStride, kBStride, lane); + load_frag8(a_frag0_1, a_tile + local_row * kAStride + kTileK, kAStride, lane); + load_frag8(a_frag1_1, a_tile + (local_row + kTileM) * kAStride + kTileK, + kAStride, lane); + load_frag8(b_frag0_1, b_tile + local_col * kBStride + kTileK, kBStride, lane); + load_frag8(b_frag1_1, b_tile + (local_col + kTileN) * kBStride + kTileK, + kBStride, lane); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + 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); +} + +// Round 6 (EPILOGUE round) fused-epilogue helpers. +// +// Rounds 2-5 converted each accumulator fragment straight to bf16 and issued 16 +// scattered 2-byte global stores per thread. One 64-lane store instruction +// carried 16 rows x 4 columns, i.e. 16 distinct 128-B lines with only 8 useful +// bytes in each, and the exact gfx928 code object shows the epilogue +// serialized behind 16 global_store_short plus ~20 scalar x_scale / +// weight_scale loads per thread. Round 6 keeps the per-row x per-column scale +// multiply and the bf16 conversion inside the compute kernel but routes the +// result through a shared-memory staging tile and writes the tile back with +// fully coalesced 16-byte stores: 8 threads cover one full 128-B row segment, +// so a 64-lane store instruction moves 8 complete lines instead of 16 +// fragments of lines. The per-row / per-column scales are staged cooperatively +// into LDS once per block (one coalesced global load per thread), which removes +// the scalar global scale loads entirely; round 15 later replaces that LDS slab +// with per-thread scale registers (see below). +// +// The staged half tile is written as plain bf16 elements; the 8-mod-16 row +// stride (136 B) keeps the 16 row starts on distinct LDS banks and stays +// 8-byte aligned for the int64 read-back. +// +// Round 15 (final consolidate polish): the two row scales and the four +// per-column scales this wave needs for one kTileN column group are passed in +// as registers instead of being read back from an LDS scale slab. The wave's +// rows are local_row + (lane & 15) and local_row + kTileM + (lane & 15), and +// its columns are col + (lane >> 4) + 4*i, so every thread owns exactly two row +// scales and eight column scales (four per column group). They are loaded once +// per thread before the K loop, their global-load latency hides behind the four +// 64-K stages, and the epilogue therefore issues no scale LDS reads at all. +// The multiply order ((float)acc * x_scale * weight_scale), the bf16 conversion +// and every staged LDS / stored global address are unchanged, so results stay +// bit-exact. +template +__device__ __forceinline__ void stage_epilogue_fragment( + const AccFragment& frag, + float xs, + const float (&ws)[4], + hip_bfloat16* __restrict__ epi, + int row, + int col, + int col_base, + int lane) { + const int r = lane & 15; + const int q = lane >> 4; #pragma unroll - for (int kk = 0; kk < kStageK; kk += kTileK) { - du_load_matrix_sync( - a_frag0, a_tile + local_row * kAStride + kk, kAStride); - du_load_matrix_sync( - a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); - du_load_matrix_sync( - b_frag0, b_tile + local_col * kBNStride + kk, kBNStride); - du_load_matrix_sync( - b_frag1, b_tile + (local_col + kTileN) * kBNStride + kk, - kBNStride); - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); - } - } - } else { - DUFragment - b_frag0, b_frag1; + for (int i = 0; i < 4; ++i) { + const float scaled = static_cast(frag.x[i]) * xs * ws[i]; + epi[(row + r) * kStride + (col - col_base) + q + 4 * i] = + __float2bfloat16(scaled); + } +} + +// Coalesced write-back of one staged 64-column bf16 half: thread `tid` owns a +// 16-byte segment of one output row (8 threads per 128-B line), so the whole +// block writes each half with one 16-byte store per thread and fully-used +// 128-B lines, in the same logical [M, N] bf16 order the scalar fallback +// produces. The staged LDS row stride (136 B) is only 8-byte aligned, so the +// LDS side reads two int64 halves and the global side re-forms the aligned +// 16-byte vector. +struct alignas(16) EpiQuad { + int64_t lo; + int64_t hi; +}; + +__device__ __forceinline__ void store_epilogue_half( + const hip_bfloat16* __restrict__ epi, + hip_bfloat16* __restrict__ out, + int m0, + int n0, + int col_base, + int n, + int tid) { + const int row = tid >> 3; + const int seg = tid & 7; + const int64_t* src = + reinterpret_cast(epi + row * kEpiStride + seg * 8); + EpiQuad q; + q.lo = src[0]; + q.hi = src[1]; + hip_bfloat16* dst = out + static_cast(m0 + row) * n + n0 + + col_base + seg * 8; + *reinterpret_cast(dst) = q; +} + +// Round 2 exact-shape kernel: one block computes a kBlockM x kBlockN output +// tile of the assigned (4096, 6144, 256) shape. (kBlockM/32) * (kBlockN/32) +// wavefronts each own a 32x32 quadrant (four m16n16k32 DUMMA accumulators), +// while the block cooperatively stages A[kBlockM, kStageK] from the +// row-major [M, K] activation and B[kBlockN, kStageK] from the round-5 panel +// packed weight. Single-buffered 64-K stages (two barriers per stage). The +// round-6 fused epilogue stages the scaled bf16 output through LDS and writes +// it back coalesced; there is no workspace use and no separate convert or +// combine pass anywhere in this operator. +template +__global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize) void +w8a8_dumma_proj_operand_reuse_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + hip_bfloat16* __restrict__ out, + int m, + int n, + int k) { + constexpr int kWavesN = kBlockN / 32; + constexpr int kAStride = kStageK + 24; + constexpr int kBStride = kStageK + 8; + constexpr int kABytes = kBlockM * kAStride; + constexpr int kBBytes = kBlockN * kBStride; + static_assert(kStageK == 2 * kTileK, + "exact-shape arm unrolls the two 32-K steps of a 64-K stage"); + static_assert(kAStride % 8 == 0 && kBStride % 8 == 0, + "8-mod-16 strides must stay 8-byte aligned"); + static_assert(kBlockM % 32 == 0 && kBlockN % 32 == 0, + "macro-tile must be a whole number of 32x32 quadrants"); + static_assert(kBlockM % kTileM == 0 && kBlockN % kTileN == 0, + "macro-tile must be a whole number of DUMMA tiles"); + + 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; + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + + __shared__ __align__(16) int8_t a_tile[kABytes]; + __shared__ __align__(16) int8_t b_tile[kBBytes]; + + DUFragment + a_frag0, a_frag1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + // Round 15 register scales: the two row scales and the eight column scales + // this wave's epilogue needs are loaded ONCE per thread here, before the K + // loop, so their global-load latency is hidden behind the four 64-K stages + // and the epilogue neither stages nor reads any scale through LDS. local_row + // / local_col are wave_row*32 / wave_col*32; the exact arm always has + // m0 == 64*blockIdx.y so m0 + local_row + kTileM + 15 <= 4095 < m, and + // n0 + local_col + kTileN + 3 + 12 <= 6139 < n, so no tail masking is needed. + const int epi_r = lane & 15; + const int epi_q = lane >> 4; + const float xs0 = x_scale[m0 + local_row + epi_r]; + const float xs1 = x_scale[m0 + local_row + kTileM + epi_r]; + float ws_a[kTileN / 4]; + float ws_b[kTileN / 4]; +#pragma unroll + for (int i = 0; i < kTileN / 4; ++i) { + ws_a[i] = weight_scale[n0 + local_col + epi_q + 4 * i]; + ws_b[i] = weight_scale[n0 + local_col + kTileN + epi_q + 4 * i]; + } + + for (int k0 = 0; k0 < k; k0 += kStageK) { + stage_proj_tiles( + x_q, weight, a_tile, b_tile, tid, m0, n0, k, k0); + __syncthreads(); + + mmac_proj_stage( + a_frag0, a_frag1, acc00, acc01, acc10, acc11, a_tile, b_tile, + local_row, local_col, lane); + // Retire this stage's fragment reads before the next stage's staging + // writes overwrite the single LDS buffer (WAR hazard across waves). + __syncthreads(); + } + + // ---- Round 6 fused coalesced epilogue (round-15 register scales) ------ + // The A and B staging buffers are dead after the K loop (the final + // __syncthreads retires the last fragment reads), so the B tile becomes the + // bf16 output staging area: same 14,848 B footprint, same 2 blocks/CU + // occupancy. The scale factors live in the per-thread registers loaded before + // the K loop, so the retired A tile is not touched here and no scale-publish + // barrier is needed. Rows are always in range for this arm because the host + // guard launches exactly grid(n / kBlockN, m / kBlockM). + constexpr int kEpiHalfBytes = kBlockM * kEpiStride * sizeof(hip_bfloat16); + static_assert(kEpiHalfBytes <= kBBytes, + "one 64-column bf16 half must fit the retired B tile"); + static_assert(kBlockN % kEpiCols == 0, "epilogue halves must tile kBlockN"); + auto* epi = reinterpret_cast(b_tile); + #pragma unroll - for (int kk = 0; kk < kStageK; kk += kTileK) { - du_load_matrix_sync( - a_frag0, a_tile + local_row * kAStride + kk, kAStride); - du_load_matrix_sync( - a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); - du_load_matrix_sync( - b_frag0, b_tile + kk * kBStride + local_col, kBStride); - du_load_matrix_sync( - b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); + for (int half = 0; half < kBlockN / kEpiCols; ++half) { + const int col_base = half * kEpiCols; + if (local_col >= col_base && local_col < col_base + kEpiCols) { + stage_epilogue_fragment( + acc00, xs0, ws_a, epi, local_row, local_col, col_base, lane); + stage_epilogue_fragment( + acc01, xs0, ws_b, epi, local_row, local_col + kTileN, col_base, + lane); + stage_epilogue_fragment( + acc10, xs1, ws_a, epi, local_row + kTileM, local_col, col_base, + lane); + stage_epilogue_fragment( + acc11, xs1, ws_b, epi, local_row + kTileM, local_col + kTileN, + col_base, lane); } + __syncthreads(); + store_epilogue_half(epi, out, m0, n0, col_base, n, tid); + // Retire this half's staging reads before the next half overwrites the + // single staging buffer (WAR hazard across waves). + __syncthreads(); } } // Large-prefill path: one block computes a kBlockM x kBlockN output tile. -// (kBlockM/32) x (kBlockN/32) wavefronts each own a 32x32 quadrant (four +// (kBlockM/32) * (kBlockN/32) wavefronts each own a 32x32 quadrant (four // m16n16k32 DUMMA accumulators), while the block cooperatively stages -// A[kBlockM, kStageK] and B in bank-skewed LDS. +// A[kBlockM, kStageK] and B[kStageK, kBlockN] in bank-skewed LDS. // -// kDoubleBuffer == false: single-buffered K loop (two barriers per stage); -// stage s+1's global loads cannot start until the stage-s barrier, so each -// stage serializes one global-load round trip in front of the compute -// (bootstrap structure; co-resident block streams hide it at the CU level). -// -// kDoubleBuffer == true: software-pipelined K loop (round 3). Stage s+1's -// global loads are issued into the idle LDS buffer before the stage-s MMAC -// burst, overlapping the serialized global-load latency with compute. One -// barrier per stage (plus one prologue): it both retires stage-s fragment -// reads (WAR on the buffer just read) and makes stage s+1's ds_writes -// visible (RAW on the prefetched buffer). The gate_up shape uses -// <64, 64, 64, true, true>: two 64-K tiles = 20,480 B/block x 3 blocks/CU = -// 61,440 <= 64 KiB (occupancy preserved; a 128-K double buffer would need -// 38,912 B/block -> 1 block/CU, a harmful occupancy loss). -// -// kBNMajor selects the B operand layout: -// * false (identity): B staged [K, N] row-major, library row_major B -// fragments (byte-scattered LDS reads; correct for the raw [K, N] weight -// and every non-gate_up shape). -// * true (gate_up only): B staged [N, K] n-major with k contiguous (row -// stride 144 = 128+16, or 80 = 64+16 for the 64-K double buffer), -// library col_major B fragments -> each lane's elements are k-contiguous -// bytes instead of byte-scattered columns. This is the bank-safe consumer -// layout for the B side (the 17.3M PMC bank-conflict source in round 1). -// The one-time pack in launch_pack_w8a8_weight provides the [N, K] buffer -// for (k, n) == (6144, 512) only. -template +// Single-buffered K loop (two barriers per stage): stage s+1's global loads +// cannot start until the stage-s barrier, so each stage serializes one +// global-load round trip in front of its MMAC burst. Co-resident blocks hide +// that latency at the CU level, and the bootstrap deliberately avoids the +// copy/pipeline complexity that later rounds may add. +template __global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize) void w8a8_dumma_prefill_tiled_kernel( const int8_t* __restrict__ x_q, @@ -775,64 +830,22 @@ w8a8_dumma_prefill_tiled_kernel( int n, int k) { constexpr int kWavesN = kBlockN / 32; - constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; - // 16-byte bank skew: keeps int4 staging stores aligned while spreading the - // fragment rows across distinct LDS banks. The staging loops live in - // stage_prefill_tiles below; the strides remain here as the static-asserted - // alignment contract of every instantiation. kSkew8 (round 11, exact-shape - // arm only) switches A/B to the down_proj-validated 8-mod-16 strides - // 88/72; LDS bytes stay A[64,88] + B[64,72] x 2 = 20,480 B/block so the - // 3-blocks/CU occupancy of the accepted kernel is preserved exactly. - constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); - constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows - constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); - constexpr int kBNInt4 = kBlockN * (kStageK / sizeof(int4)); - // Double buffering uses kBuffers = 2 LDS tile copies; the gate_up shape - // pairs it with kStageK = 64 so 3 blocks/CU is preserved (see kernel - // comment above). - constexpr int kBuffers = kDoubleBuffer ? 2 : 1; + constexpr int kAStride = kStageK + sizeof(int4); + constexpr int kBStride = kBlockN + sizeof(int4); constexpr int kABytes = kBlockM * kAStride; - // One shared B tile sized for the larger of the two layouts so an inactive - // arm never allocates dead LDS (keeps <64,64,128> at 19,456 B total and - // 3 blocks/CU). kSkew8 implies kBNMajor, so only the n-major arm is live. - constexpr int kBMaxBytes = - kSkew8 ? kBlockN * kBNStride - : (kStageK * kBStride > kBlockN * kBNStride - ? kStageK * kBStride - : kBlockN * kBNStride); - static_assert(kSkew8 ? (kAStride % 8 == 0) : (kAStride % sizeof(int4) == 0), - "A LDS row stride must stay int4 aligned (or 8-byte aligned " - "under kSkew8)"); - static_assert(kBStride % static_cast(sizeof(int4)) == 0, - "B LDS row stride must stay int4 aligned"); - static_assert( - kSkew8 ? (kBNStride % 8 == 0) : (kBNStride % sizeof(int4) == 0), - "n-major B LDS row stride must stay int4 aligned (or 8-byte aligned " - "under kSkew8)"); - static_assert(!kSkew8 || kBNMajor, - "kSkew8 is only defined for the packed n-major B layout"); - static_assert(!kSkew8 || kStageK == 64, - "kSkew8 strides 88/72 assume kStageK == 64"); + constexpr int kBBytes = kStageK * kBStride; static_assert(kStageK % kTileK == 0, "K stage must be a multiple of the DUMMA K unit"); static_assert(kStageK % static_cast(sizeof(int4)) == 0, "K stage must keep int4 staging aligned"); - static_assert(kBNInt4 % kThreads == 0, - "n-major B staging must divide evenly across threads"); - static_assert(!kPayload || (kBNMajor && kDoubleBuffer && kDirectFrag && - kSkew8 && kStageK == 64 && kBlockM == 64 && - kBlockN == 64), - "payload arm is specialized to the exact-shape <64,64,64> " - "double-buffered load_frag8/kSkew8 arm (one int4 per thread " - "per operand)"); - static_assert( - !kBPanel || (kBNMajor && kStageK == 64 && kBlockN == 128 && !kPayload), - "kBPanel is specialized to the down_proj panel pack (n-major 128x64 " - "per-stage panels, 64-K stages, non-payload pipeline)"); + static_assert(kBlockM % 32 == 0 && kBlockN % 32 == 0, + "macro-tile must be a whole number of 32x32 quadrants"); + static_assert(kBlockM % kTileM == 0 && kBlockN % kTileN == 0, + "macro-tile must be a whole number of DUMMA tiles"); 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; @@ -840,8 +853,8 @@ w8a8_dumma_prefill_tiled_kernel( const int local_row = wave_row * 32; const int local_col = wave_col * 32; - __shared__ __align__(16) int8_t a_tile[kBuffers * kABytes]; - __shared__ __align__(16) int8_t b_tile[kBuffers * kBMaxBytes]; + __shared__ __align__(16) int8_t a_tile[kABytes]; + __shared__ __align__(16) int8_t b_tile[kBBytes]; DUFragment a_frag0, a_frag1; @@ -852,144 +865,44 @@ w8a8_dumma_prefill_tiled_kernel( du_fill_fragment(acc10, 0); du_fill_fragment(acc11, 0); - if constexpr (kDoubleBuffer) { - if constexpr (kPayload) { - // Round 16 lead-2 loop-carried payload pipeline (exact-shape arm - // only): stage s+1's global loads are issued a FULL iteration early - // (top of iteration s-1) into a two-int4-per-thread register payload - // and stay outstanding across the stage-(s-1) fragment reads + MMAC - // burst and the stage barrier; the vmcnt wait lands at the top of - // iteration s (first use of the payload registers), immediately - // before the ds_write2 publish into the idle LDS buffer. The - // s_barrier fences outstanding vmem, so the wait can no longer sit - // right after the load issue: the global-load latency overlaps the - // whole compute segment instead of serializing in front of every - // stage (the round-11 ISA defect). Same 2 buffers, same one barrier - // per stage (prologue + 96 = 97/block), same 16 ds_read2_b64 + - // 8 ds_write2_b64 per block-stage, same staging bytes/addresses - // (bit-identical int32 order). - int4 a_payload, b_payload; - // Prologue: issue the stage-1 global loads first (so they ride behind - // the stage-0 inline staging), stage k0 = 0 into buffer 0 inline - // (same load+wait+write path as the accepted kernel), then the first - // barrier. - prefetch_stage_payload( - x_q, weight, tid, m0, n0, m, k, kStageK, a_payload, b_payload); - stage_prefill_tiles( - x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, 0); - __syncthreads(); - for (int k0 = 0; k0 < k; k0 += kStageK) { - const int buf = (k0 / kStageK) & 1; - if (k0 + kStageK < k) { - // Publish the payload carried from the previous iteration (stage - // k0+kStageK) into the idle buffer: the vmcnt wait for the loads - // issued one iteration earlier is emitted here at the first use, - // after the previous MMAC burst + barrier. - publish_stage_payload( - a_payload, b_payload, a_tile + (buf ^ 1) * kABytes, - b_tile + (buf ^ 1) * kBMaxBytes, tid); - } - if (k0 + 2 * kStageK < k) { - // Issue stage k0+2*kStageK's global loads into the payload - // registers (no wait); they stay in flight across the MMAC burst - // and the barrier below and are waited at the next iteration's - // publish (RAW). - prefetch_stage_payload( - x_q, weight, tid, m0, n0, m, k, k0 + 2 * kStageK, a_payload, - b_payload); - } - // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. - // k0-outer / kk-inner int32 accumulation order. - mmac_prefill_stage( - a_frag0, a_frag1, acc00, acc01, acc10, acc11, - a_tile + buf * kABytes, b_tile + buf * kBMaxBytes, - local_row, local_col); - // One barrier per stage: (a) WAR - all stage-s fragment reads of the - // buffer just read are retired before it becomes the publish target - // two stages later; (b) RAW - every thread's stage-(s+1) ds_writes - // are visible before the next iteration reads them. - __syncthreads(); - } - } else { - // Prologue: stage k0 = 0 into buffer 0, then the pipelined loop. - stage_prefill_tiles( - x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, 0); - __syncthreads(); - for (int k0 = 0; k0 < k; k0 += kStageK) { - const int buf = (k0 / kStageK) & 1; - if (k0 + kStageK < k) { - // Prefetch stage s+1 into the idle buffer: global loads issue here, - // before the stage-s MMAC burst, and complete (RAW) behind the - // single barrier below. - stage_prefill_tiles( - x_q, weight, a_tile + (buf ^ 1) * kABytes, - b_tile + (buf ^ 1) * kBMaxBytes, tid, m0, n0, m, n, k, - k0 + kStageK); - } - // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. - // k0-outer / kk-inner int32 accumulation order. - mmac_prefill_stage( - a_frag0, a_frag1, acc00, acc01, acc10, acc11, - a_tile + buf * kABytes, b_tile + buf * kBMaxBytes, - local_row, local_col); - // One barrier per stage: (a) WAR - all stage-s fragment reads of the - // buffer just read are retired before it becomes the prefetch target - // two stages later; (b) RAW - every thread's stage-(s+1) ds_writes are - // visible before the next iteration reads them. - __syncthreads(); - } - } - } else { - for (int k0 = 0; k0 < k; k0 += kStageK) { - // Stage A[kBlockM, kStageK] and B into the single LDS buffer - // (kSkew8 8-mod-16 row strides 88/72 with ds_write2_b64 staging). - stage_prefill_tiles( - x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, k0); - __syncthreads(); + for (int k0 = 0; k0 < k; k0 += kStageK) { + stage_prefill_tiles( + x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, k0); + __syncthreads(); - // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. - // k0-outer / kk-inner int32 accumulation order (kDirectFrag: eight - // load_frag8 ds_read2_b64 fragment reads, zero reassembly VALU). - mmac_prefill_stage( - a_frag0, a_frag1, acc00, acc01, acc10, acc11, - a_tile, b_tile, local_row, local_col); - // Retire this stage's reads before the next stage's staging writes - // overwrite the single LDS buffer (WAR hazard across waves). - __syncthreads(); - } + mmac_prefill_stage( + a_frag0, a_frag1, acc00, acc01, acc10, acc11, a_tile, b_tile, + local_row, local_col); + // Retire this stage's fragment reads before the next stage's staging + // writes overwrite the single LDS buffer (WAR hazard across waves). + __syncthreads(); } const int base_row = m0 + wave_row * 32; const int base_col = n0 + wave_col * 32; store_prefill_fragment( - acc00, x_scale, weight_scale, out, m, n, - base_row, base_col, lane); + acc00, x_scale, weight_scale, out, m, n, base_row, base_col, lane); store_prefill_fragment( - acc01, x_scale, weight_scale, out, m, n, - base_row, base_col + kTileN, lane); + acc01, x_scale, weight_scale, out, m, n, base_row, base_col + kTileN, + lane); store_prefill_fragment( - acc10, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col, lane); + acc10, x_scale, weight_scale, out, m, n, base_row + kTileM, base_col, + lane); store_prefill_fragment( - acc11, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col + kTileN, lane); + acc11, x_scale, weight_scale, out, m, n, base_row + kTileM, + base_col + kTileN, lane); } // Generic scalar fallback: one output element per grid-stride iteration with // exact int32 accumulation. Correct for any (m, n, k), including the small-M -// API shapes (M=2/16) and any unmatched geometry. The weight is the identity -// [K, N] row-major layout for every (k, n) EXCEPT the gate_up weight -// (k, n) == (6144, 512), whose buffer is packed [N, K] n-major; this kernel -// decodes that layout so the paired gate_up tails stay correct. -__global__ __launch_bounds__(kScalarThreads) void -w8a8_gemm_scalar_kernel( +// API shapes (M=2/16) and any unmatched geometry. The weight buffer is the +// identity [K, N] row-major layout for every (k, n) EXCEPT the packed +// down_proj pair (k, n) == (256, 6144), whose buffer is the round-5 panel +// layout [N/128][K/64][128][64] +// (packed[((nt*4+kt)*128+nn)*64+kk64] == raw[(nt*128+nn)*k + kt*64+kk64] +// with nt = col/128, nn = col%128, kt = kk/64, kk64 = kk%64); this kernel +// decodes that layout so the paired M=2/M=16 down_proj tails stay correct. +__global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, const float* __restrict__ x_scale, @@ -998,15 +911,7 @@ w8a8_gemm_scalar_kernel( int m, int n, int k) { - // Decode the packed layout when this weight was packed that way: - // gate_up (k, n) == (6144, 512) and, since round 2 of this session, - // down_proj (k, n) == (256, 6144) - since round 5 in the PANEL layout - // [N/128][K/64][128][64] (n-major 64-k-byte rows per panel; byte-exact - // relation packed[((nt*4+kt)*128+nn)*64+kk64] == raw[(nt*128+nn)*k + - // kt*64+kk64] with nt = col/128, nn = col%128, kt = kk/64, kk64 = kk%64); - // identity [K, N] otherwise. - const bool b_panel = (k == 256 && n == 6144); - const bool b_nmajor = (k == 6144 && n == 512) || b_panel; + const bool b_panel = (k == kExactK && n == kExactN); const int64_t total = static_cast(m) * n; const int64_t stride = static_cast(gridDim.x) * blockDim.x; for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + @@ -1014,85 +919,56 @@ w8a8_gemm_scalar_kernel( idx < total; idx += stride) { const int row = static_cast(idx / n); const int col = static_cast(idx - static_cast(row) * n); - int32_t acc = 0; const int8_t* a_row = x_q + static_cast(row) * k; - for (int kk = 0; kk < k; ++kk) { - int64_t b_off; - if (b_panel) { - const int nt = col / 128; - const int nn = col - nt * 128; - const int kt = kk / 64; - const int kk64 = kk - kt * 64; - b_off = (static_cast(nt * 4 + kt) * 128 + nn) * 64 + kk64; - } else { - b_off = b_nmajor ? static_cast(col) * k + kk - : static_cast(kk) * n + col; + int32_t acc = 0; + if (b_panel) { + // Round-5 panel decode: the logical [K, N] element (row kk, column + // col) lives at ((nt*(k/64)+kt)*128 + nn)*64 + kk64. + const int nt = col / kProjBlockN; + const int nn = col - nt * kProjBlockN; + const int64_t panel_base = + static_cast(nt) * (k / kProjStageK) * + (kProjBlockN * kProjStageK) + + static_cast(nn) * kProjStageK; + for (int kk = 0; kk < k; ++kk) { + const int kt = kk / kProjStageK; + const int kk64 = kk - kt * kProjStageK; + acc += static_cast(a_row[kk]) * + static_cast( + weight[panel_base + + static_cast(kt) * + (kProjBlockN * kProjStageK) + + kk64]); + } + } else { + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast( + weight[static_cast(kk) * n + col]); } - acc += static_cast(a_row[kk]) * - static_cast(weight[b_off]); } - const float scaled = static_cast(acc) * x_scale[row] * - weight_scale[col]; + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; out[idx] = __float2bfloat16(scaled); } } -// Bootstrap pack: identity device-to-device copy for every (k, n) pair -// (raw [K, N] layout preserved). Runs outside the timed/CUDA-Graph region. -__global__ __launch_bounds__(kScalarThreads) void -w8a8_pack_identity_kernel( +// Bootstrap pack: identity device-to-device copy of the logical [K, N] weight +// and the [N] weight scale. Runs outside the timed/CUDA-Graph region; the +// destination buffers are the graph-stable ones returned by pack_weight. +__global__ __launch_bounds__(kScalarThreads) void w8a8_pack_identity_kernel( const int8_t* __restrict__ raw_weight, const float* __restrict__ weight_scale, int8_t* __restrict__ packed_weight, float* __restrict__ packed_weight_scale, int64_t weight_elems, int n) { - const int64_t idx = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (idx < weight_elems) { - packed_weight[idx] = raw_weight[idx]; - } - if (idx < n) { - packed_weight_scale[idx] = weight_scale[idx]; - } -} - -// Shape-specific gate_up pack: transpose raw [K, N] -> packed [N, K] -// (n-major rows, k contiguous) for (k, n) == (6144, 512) only. Byte-exact -// index relation: packed[cc * k + kk] == raw[kk * n + cc]. Each thread -// gathers the 16 bytes of one aligned int4 k-block from 16 raw rows (stride -// n) and stores one aligned int4, keeping the packed write side fully -// coalesced (same gather + aligned int4-store pattern as the validated -// o_proj transpose repair). Runs outside the timed/CUDA-Graph region; same -// byte count, same graph-stable buffer. -__global__ __launch_bounds__(kScalarThreads) void -w8a8_pack_gateup_transpose_kernel( - const int8_t* __restrict__ raw_weight, - const float* __restrict__ weight_scale, - int8_t* __restrict__ packed_weight, - float* __restrict__ packed_weight_scale, - int k, - int n) { - const int k4 = k / static_cast(sizeof(int4)); - const int64_t total4 = static_cast(n) * k4; - const int64_t stride4 = static_cast(gridDim.x) * blockDim.x; - for (int64_t idx4 = static_cast(blockIdx.x) * blockDim.x + + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - idx4 < total4; idx4 += stride4) { - const int cc = static_cast(idx4 / k4); - const int kblk = static_cast(idx4 - static_cast(cc) * k4); - const int kk = kblk * static_cast(sizeof(int4)); - const int8_t* src = raw_weight + static_cast(kk) * n + cc; - int4 v; - char* vp = reinterpret_cast(&v); -#pragma unroll - for (int i = 0; i < static_cast(sizeof(int4)); ++i) { - vp[i] = src[static_cast(i) * n]; - } - reinterpret_cast(packed_weight + static_cast(cc) * k + - kk)[0] = v; + idx < weight_elems; idx += stride) { + packed_weight[idx] = raw_weight[idx]; } - const int64_t stride = static_cast(gridDim.x) * blockDim.x; for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < n; idx += stride) { @@ -1100,25 +976,23 @@ w8a8_pack_gateup_transpose_kernel( } } -// Round 5 (this session): shape-specific down_proj panel pack for -// (k, n) == (256, 6144) only. Raw [K, N] -> panel layout -// [N/128][K/64][128][64]: n-major rows of 64 k-bytes per panel, panels -// ordered (n-tile outer, k-tile inner). Byte-exact index relation (verified -// by simulation): +// Round 5 (PACKING round) down_proj pack: raw [K, N] -> panel layout +// [N/128][K/64][128][64] for (k, n) == (256, 6144) only. Byte-exact index +// relation (verified by the int4 decode below): // packed[((nt * (k/64) + kt) * 128 + nn) * 64 + kk64] == // raw[(nt * 128 + nn) * k + (kt * 64 + kk64)] // with nt in [0, n/128), kt in [0, k/64), nn in [0, 128), kk64 in [0, 64). // Each packed int4 (idx4) decodes as v = idx4 % 4, row_p = idx4 / 4, // nn = row_p % 128, kt = (row_p / 128) % (k/64), nt = (row_p / 128) / (k/64), -// kk = kt * 64 + v * 16, cc = nt * 128 + nn, and -// packed[idx4*16 .. +16] == raw[(kk + i) * n + cc] for i in 0..15 (the same -// 16-row gather at column cc as the gate_up transpose, stored into a -// different packed offset; each thread stores one aligned int4, so the -// packed write side stays fully coalesced). The panel layout makes every -// per-block B stage tile (128 n x 64 k) ONE contiguous 8 KiB region, so the -// timed kernel's B staging reads 1 KiB per warp instruction contiguously. -// Runs outside the timed/CUDA-Graph region; same byte count, same -// graph-stable buffer, byte-exact vs the raw logical [K, N] weight. +// kk = kt * 64 + v * 16, cc = nt * 128 + nn; the 16 stored bytes are the same +// 16-row gather at column cc as the round-2 transpose, so every thread still +// stores one aligned int4 and the packed write side stays fully coalesced. +// The panel layout makes each timed block's per-stage B tile (128 n x 64 k = +// 8 KiB) ONE contiguous region, so a 64-lane warp's B staging load reads 1 KiB +// contiguous instead of 16 rows x 64 B at a 256-B row stride (16 distinct +// 128-B lines per warp, half of each line used). Runs outside the +// timed/CUDA-Graph region; same byte count, same grid, same graph-stable +// buffer, byte-exact vs the raw logical [K, N] weight. __global__ __launch_bounds__(kScalarThreads) void w8a8_pack_downproj_panel_kernel( const int8_t* __restrict__ raw_weight, @@ -1127,21 +1001,20 @@ w8a8_pack_downproj_panel_kernel( float* __restrict__ packed_weight_scale, int k, int n) { - constexpr int kBlockN = 128; - constexpr int kStageK = 64; - const int64_t total4 = (static_cast(n) * k) / 16; + const int64_t total4 = + (static_cast(n) * k) / static_cast(sizeof(int4)); const int64_t stride4 = static_cast(gridDim.x) * blockDim.x; for (int64_t idx4 = static_cast(blockIdx.x) * blockDim.x + - threadIdx.x; + threadIdx.x; idx4 < total4; idx4 += stride4) { const int v = static_cast(idx4 & 3); const int64_t row_p = idx4 >> 2; - const int nn = static_cast(row_p % kBlockN); - const int64_t kt_tmp = row_p / kBlockN; // nt * (k/kStageK) + kt - const int kt = static_cast(kt_tmp % (k / kStageK)); - const int nt = static_cast(kt_tmp / (k / kStageK)); - const int kk = kt * kStageK + v * 16; - const int cc = nt * kBlockN + nn; + const int nn = static_cast(row_p % kProjBlockN); + const int64_t kt_tmp = row_p / kProjBlockN; // nt * (k/64) + kt + const int kt = static_cast(kt_tmp % (k / kProjStageK)); + const int nt = static_cast(kt_tmp / (k / kProjStageK)); + const int kk = kt * kProjStageK + v * static_cast(sizeof(int4)); + const int cc = nt * kProjBlockN + nn; const int8_t* src = raw_weight + static_cast(kk) * n + cc; int4 val; char* vp = reinterpret_cast(&val); @@ -1149,6 +1022,8 @@ w8a8_pack_downproj_panel_kernel( for (int i = 0; i < static_cast(sizeof(int4)); ++i) { vp[i] = src[static_cast(i) * n]; } + // The panel offset of idx4 is exactly idx4 * 16 bytes: + // ((nt*(k/64)+kt)*128 + nn)*64 + v*16. reinterpret_cast(packed_weight + idx4 * 16)[0] = val; } const int64_t stride = static_cast(gridDim.x) * blockDim.x; @@ -1173,141 +1048,73 @@ extern "C" void launch_w8a8_gemm( int n, int k, hipStream_t stream) { - // No split-K: the provided workspace is not touched. The timed operator - // performs no allocation, synchronization, packing, or default-stream - // launch; only the caller-provided out is written. + // No split-K: the caller-provided workspace is intentionally unused, and the + // timed operator performs no allocation, synchronization, packing, or + // default-stream launch. (void)workspace; (void)workspace_bytes; auto* out_bf16 = static_cast(out); - // Native INT8 DUMMA m16n16k32 tiled path for every large-M shape. Both - // assigned shapes qualify: gate_up (4096, 512, 6144) and down_proj - // (4096, 6144, 256) both have M >= 128, N % 128 == 0 and K % 128 == 0. - // Explicit dispatch by launch geometry; unmatched shapes fall through to - // the scalar path (no shape is specialized away). - if (m >= kPrefillMinM && (n % kDefaultBlockN == 0) && - (k % kPrefillStageK == 0)) { - // Exact-shape guard for the assigned gate_up shape (round-3 pipeline - // decision, round-10 load_frag8, round-11 kSkew8, round-16 kPayload): - // cooperative A+B staging with the n-major [N, K] packed B operand - // (kBNMajor) consumed as col_major fragments, software-pipelined across - // K tiles (kDoubleBuffer) with a 64-K double buffer: 64x64 tile, 512 - // blocks, 20,480 B LDS/block, 3 blocks/CU (61,440 B) - occupancy - // preserved vs the accepted round-2 single buffer (19,456 B x 3 = - // 58,368 B). One barrier per 64-K stage + prologue (97 for 96 stages). - // Round 11 stages A/B at the 8-mod-16 strides 88/72 (kSkew8) so the - // fragment reads are bank-conflict-free (LDS bytes and occupancy - // unchanged). Round 16 (kPayload) issues the stage-(s+1) global loads a - // full iteration early into a two-int4-per-thread register payload and - // publishes them into the idle LDS buffer at the top of iteration s, so - // the vmcnt wait lands after the stage-s MMAC burst instead of in front - // of it (LDS bytes, barriers, traffic, occupancy and results unchanged). - if (m == 4096 && n == 512 && k == 6144) { - const dim3 grid(n / 64, (m + 63) / 64); - const dim3 block(4 * kWaveSize); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel< - 64, 64, 64, true, true, true, true, true>), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - return; - } - // Exact-shape guard for the assigned down_proj shape (round 2 of this - // session): the generic <64,128,128> identity-layout arm measured - // 314.72 us median / 315.18 us p90 with PMC lds_bank_conflicts = - // 7,864,320, ~600 LDS instructions and heavy per-byte fragment - // reassembly VALU per block-stage (library row_major loaders on the - // identity [K,N] B layout). The round-2 guard transplanted the file's - // lineage-validated gate_up machinery: n-major packed B (kBNMajor, pack - // in launch_pack_w8a8_weight for (k,n) == (256,6144)), 64-K - // double-buffered K loop (kDoubleBuffer), explicit 8-byte load_frag8 - // fragment reads (kDirectFrag), 8-mod-16 bank-skewed LDS strides - // A=88/B=72 (kSkew8); accepted at 286.53 us median / 286.64 us p90 with - // PMC lds_instructions 540,672, conflicts 1,179,648, lds_wait 162,664. - // - // Round 3 (this session, pipeline round): compare single and double - // buffering across K tiles. The (4096, 6144, 256) guard now routes to - // w8a8_dumma_prefill_tiled_kernel<64, 128, 64, true, false, true, true> - // - the SINGLE-buffered form of the accepted round-2 arm: same 64-K - // stage, same kBNMajor/kDirectFrag/kSkew8, same tile/grid/epilogue, but - // kDoubleBuffer = false, so each stage serializes one global-load round - // trip (load -> wait -> ds_write2_b64) in front of its MMAC burst with - // two barriers per stage (8/block vs the accepted 5/block), and the LDS - // halves from 29,696 to 14,848 B/block (A[64,88] + B[128,72] x 1). The - // accepted kernel's exact code object shows the double-buffered loop - // already overlaps stage-(s+1) loads with the stage-s MMAC burst - // (global_load x2 -> ds_read2_b64 x4 -> v_mmac x8 -> s_waitcnt - // vmcnt(0) lgkmcnt(0) -> s_barrier), so this round measures whether that - // overlap is worth the extra 14,848 B of LDS: if the single-buffered - // median is not strictly below the accepted 286.53 us (with the p90 - // noise guard), double buffering is retained per the pipeline mandate. - // Results stay bit-identical: stage size, k0-outer/kk-inner int32 order - // and the element-to-slot fragment mapping are unchanged (only the - // buffer count and barrier count differ), so mismatch 0 is expected. - // - // Round 5 (this session, PACKING round): the (4096, 6144, 256) guard - // now routes to w8a8_dumma_prefill_tiled_kernel<64, 128, 64, true, - // false, true, true, false, true> (new kBPanel template flag): the - // one-time pack for (k, n) == (256, 6144) becomes the panel layout - // [N/128][K/64][128][64] (n-major 64-k-byte rows per panel, byte-exact - // vs the raw [K,N] weight, same byte count / graph-stable buffer, - // out-of-timed/out-of-Graph) so each block's per-stage B tile is one - // contiguous 8 KiB region and every warp's int4 loads read 1 KiB - // contiguous (full L2 line utilization) instead of 16 rows x 64 B at - // 256-B stride (half-used lines). Everything else is byte-identical to - // the round-3 winner: same single-buffered 64-K stage, same - // kBNMajor/kDirectFrag/kSkew8, same 64x128 tile / 512 threads / 8 waves, - // same grid (48,64) = 3072 blocks, same LDS A[64,88] + B[128,72] = - // 14,848 B/block @ 2 blocks/CU, same ds_write2_b64 staging and - // load_frag8 fragment reads (LDS bytes and bank behavior untouched), - // same k0-outer/kk-inner int32 order and element-to-slot fragment - // mapping -> mismatch 0 / max_abs_error 0.0 expected. - if (m == 4096 && n == 6144 && k == 256) { - const dim3 grid(n / 128, (m + 63) / 64); + // Native INT8 DUMMA m16n16k32 tiled path for M >= 128. Explicit dispatch by + // launch geometry; every unmatched shape falls through to the scalar path. + if (m >= kPrefillMinM) { + if (m == kExactM && n == kExactN && k == kExactK) { + // Round 2 operand-reuse arm + round 5 panel pack + round 6 coalesced + // fused epilogue + round 15 register scales: cooperative A+B staging with + // the panel-packed B operand consumed as col_major fragments, explicit + // 8-byte load_frag8 reads and bank-safe 8-mod-16 LDS strides (A=88, + // B=72), 64-K single-buffered stages, then an LDS-staged bf16 epilogue + // that applies the per-thread register x_scale/weight_scale, converts and + // writes 16-byte coalesced stores (no workspace, no separate + // convert/combine pass, no LDS scale slab). 64x128 tile, 8 waves + // (32x32 quadrant per wave), grid (6144/128, 4096/64) = (48, 64) = 3072 + // blocks, LDS A[64,88] + B[128,72] = 14,848 B/block (B reused by the + // epilogue) -> 2 blocks/CU, K=256 = four 64-K stages. Weight layout: the + // round-5 panel [N/128][K/64][128][64] (see launch_pack_w8a8_weight). + const dim3 grid(n / kProjBlockN, m / kProjBlockM); const dim3 block(8 * kWaveSize); hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel< - 64, 128, 64, true, false, true, true, false, true>), + HIP_KERNEL_NAME(w8a8_dumma_proj_operand_reuse_kernel< + kProjBlockM, kProjBlockN, kProjStageK>), grid, block, 0, stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); return; } - if (n % 128 == 0) { - // 64x128 tile, 8 waves (32x32 quadrant per wave). - // down_proj: grid (6144/128, 4096/64) = (48, 64) = 3072 blocks. - const dim3 grid(n / 128, (m + 63) / 64); - const dim3 block(8 * kWaveSize); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 128, 128>), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - return; - } - if (m % 128 == 0) { - // 128x64 tile, 8 waves. - const dim3 grid(n / 64, m / 128); - const dim3 block(8 * kWaveSize); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<128, 64, 128>), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - return; - } - { - // 64x64 tile, 4 waves (generic default). - const dim3 grid(n / 64, (m + 63) / 64); - const dim3 block(4 * kWaveSize); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 64, 128>), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - return; + if (!(k == kExactK && n == kExactN)) { + if (n % kDefaultBlockN == 0 && k % kDefaultStageK == 0) { + // 64x128 tile, 8 waves (32x32 quadrant per wave): grid (48, 64) = + // 3072 blocks for the assigned shape, 27,648 B LDS/block -> 2 + // blocks/CU, K=256 = two 128-K stages. + const dim3 grid(n / kDefaultBlockN, (m + kDefaultBlockM - 1) / + kDefaultBlockM); + const dim3 block(8 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel< + kDefaultBlockM, kDefaultBlockN, kDefaultStageK>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } + if (n % 64 == 0 && k % 64 == 0) { + // 64x64 tile, 4 waves: covers large-M shapes whose N is a multiple of + // 64 but not of 128, grid (n/64, ceil(m/64)), 10,240 B LDS/block. + const dim3 grid(n / 64, (m + 63) / 64); + const dim3 block(4 * kWaveSize); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 64, 64>), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); + return; + } } + // The packed (k, n) == (256, 6144) buffer must never reach an + // identity-layout tiled arm: any non-exact large-M shape with that pair + // drops to the layout-decoding scalar fallback below. } - // Generic scalar fallback for every unmatched (m, n, k) and M < 128 - // (including the paired M=2/M=16 API shapes with the same (N, K), which - // decode the packed [N, K] gate_up layout when (k, n) == (6144, 512)). + // Generic scalar fallback for every unmatched (m, n, k), every M < 128, and + // the paired M=2/M=16 API shapes with the same (N, K). It decodes the packed + // panel [N/128][K/64][128][64] layout when (k, n) == (256, 6144). const int64_t total = static_cast(m) * n; const unsigned blocks = static_cast( (total + kScalarThreads - 1) / kScalarThreads); @@ -1325,29 +1132,15 @@ extern "C" void launch_pack_w8a8_weight( int k, int n, hipStream_t stream) { - // Assigned gate_up weight (K=6144, N=512): one-time [K,N] -> [N,K] - // transpose pack (n-major, k contiguous) so the timed GEMM reads B - // fragments as k-contiguous bytes from the staged LDS tile. Runs outside - // the timed/CUDA-Graph region; same byte count, same graph-stable buffer, - // byte-exact vs the raw logical [K,N] weight. - if (k == 6144 && n == 512) { + // Round 5: the assigned down_proj weight (k, n) == (256, 6144) is packed + // once into the panel layout [N/128][K/64][128][64] (n-major 64-k-byte rows + // per panel) so each timed block-stage reads its 128x64 B tile as ONE + // contiguous 8 KiB region while the staged LDS bytes stay identical to the + // round-2 n-major [N, K] transpose. One-time, out-of-timed/out-of-Graph, + // same byte count, same graph-stable buffer, byte-exact vs the raw logical + // [K, N] weight. + if (k == kExactK && n == kExactN) { const int64_t total4 = static_cast(n) * (k / 16); - const unsigned blocks = static_cast( - (total4 + kScalarThreads - 1) / kScalarThreads); - hipLaunchKernelGGL( - w8a8_pack_gateup_transpose_kernel, - dim3(blocks), dim3(kScalarThreads), 0, stream, - raw_weight, weight_scale, packed_weight, packed_weight_scale, k, n); - return; - } - // Round 5 (this session): the assigned down_proj weight (K=256, N=6144) - // gets the PANEL pack [N/128][K/64][128][64] (n-major 64-k-byte rows per - // panel; each per-block B stage tile is one contiguous 8 KiB region so - // the timed kernel's B staging reads 1 KiB per warp instruction - // contiguously). One-time, out-of-timed/out-of-Graph, same byte count, - // same graph-stable buffer, byte-exact vs the raw [K,N] weight. - if (k == 256 && n == 6144) { - const int64_t total4 = (static_cast(n) * k) / 16; const unsigned blocks = static_cast( (total4 + kScalarThreads - 1) / kScalarThreads); hipLaunchKernelGGL( @@ -1357,7 +1150,8 @@ extern "C" void launch_pack_w8a8_weight( return; } // Identity device-to-device copy for every other (k, n) pair, including - // unmatched (K, N). Runs outside the timed/CUDA-Graph region. + // unmatched (K, N). Runs outside the timed/CUDA-Graph region; same byte + // count and same graph-stable destination buffers. const int64_t weight_elems = static_cast(k) * n; const int64_t total = weight_elems > n ? weight_elems : n; const unsigned blocks = static_cast( diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_gate_up_proj.hip index fec33f67..8a57f4bf 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_gate_up_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/shared_gate_up_proj.hip @@ -1,212 +1,103 @@ -// @@variant shape=glm_tp8_shared_gate_up_proj_m4096 commit=df2c88db09d3836df8ff7a4e99e19584fe6d4163 added=2026-08-31 -// median_us=273.8 p90_us=275.3 speedup=77.62 baseline_us=2.125e+04 -// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// @@variant shape=glm_tp8_shared_gate_up_proj_m4096 commit=1aa11051ede5b81487b51f5949af8c5177ffd4a4 added=2026-09-14 +// median_us=249.9 p90_us=251.9 speedup=85.06 baseline_us=2.125e+04 +// source=test-9-11-1/iteration_002/glm_tp8_shared_gate_up_proj_m4096 // MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). -// worker_3 (iteration 3), assigned shapes: +// worker_1 (iteration 1, establish_arch round), assigned shape: // glm_tp8_shared_gate_up_proj_m4096 : (M, N, K) = (4096, 512, 6144) -// glm_tp8_shared_down_proj_m4096 : (M, N, K) = (4096, 6144, 256) // // This file provides the two stable host launch symbols consumed by // csrc/bindings.cpp: // * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's // HIP stream; no allocation, packing, // autotuning, or synchronization. Only the -// caller-provided out and workspace are -// touched (workspace is unused this round). -// * launch_pack_w8a8_weight(...) - out-of-timed-region packing. Identity -// device-to-device copy for every (k, n) -// pair EXCEPT the assigned gate_up weight -// (k, n) == (6144, 512), which is packed -// once as a [K, N] -> [N, K] transpose -// (n-major rows, k contiguous) so the timed -// GEMM consumes B fragments as -// k-contiguous bytes. The pack stays -// out-of-timed, out-of-Graph, same byte -// count, same graph-stable buffer. +// caller-provided out is touched (no +// split-K, so the workspace is untouched). +// * launch_pack_w8a8_weight(...) - out-of-timed-region packing. The assigned +// gate_up weight (k, n) == (6144, 512) is +// packed once as a [K, N] -> [N, K] +// transpose (n-major rows, k contiguous); +// every other (k, n) pair keeps the +// byte-identical identity copy. // -// Bootstrap strategy (correctness-first, usable profiling baseline): -// * Native INT8 DUMMA m16n16k32 tiled kernel with int32 accumulation for -// every assigned shape with M >= 128 (both assigned shapes qualify). -// * 2-D macro-tile family (64x64, 64x128, 128x64) with one 32x32 quadrant -// per wavefront (four m16n16k32 int32 accumulators kept resident over the -// whole K loop). Single-buffered K loop (K stage 128, two barriers per -// stage): no split-K, no raw asm, no speculative double-buffer pipeline. +// Round 1 (establish_arch; architecture experiment = packed_weight family): +// the accepted bootstrap is a correct native INT8 DUMMA m16n16k32 64x64x128 +// tiled kernel (431.615 us median / 432.025 us p90, 59.71 TOPS, 75.35 GB/s, +// passed, graph_capture_passed, mismatch_count 0). Its fresh PMC bundle +// exposes the dominant steady-state defect: +// * lds_bank_conflicts = 17,301,504 over 512 blocks x 48 K-stages = 704 +// conflict cycles per block-stage, sourced by the identity [K, N] B tile: +// the matrix_b row_major loader must gather its 8 elements per lane from +// 8 different k-rows at the kBlockN+16 = 80 B row stride, so the LDS +// reads are byte-scattered and heavily aliased across the wavefront; +// * lds_wait_instructions = 11,954,696 (~486 stall cycles per block-stage); +// * lds_instructions = 7,471,104. +// This round establishes the packed-weight architecture geometry for the +// exact gate_up (K, N) pair: +// * launch_pack_w8a8_weight transposes raw [K, N] into packed [N, K] +// (n-major rows, k contiguous) for (k, n) == (6144, 512) ONLY. Packing +// runs before graph capture, uses the same byte count / graph-stable +// buffer, and is byte-exact: packed[cc * k + kk] == raw[kk * n + cc]. +// * The prefill kernel stages B[kBlockN, kStageK] from that packed layout +// with 16-byte vectorized coalesced int4 global loads into a bank-skewed +// LDS tile (row stride kStageK + 16 B; 80 B at the round-11 64-K packed +// stage) and consumes it with the +// matrix_b col_major fragment loader: each lane's 8 elements are now 8 +// CONSECUTIVE k-bytes of one packed row instead of 8 byte-scattered +// columns, which removes the byte aliasing of the row_major path. +// * A stays row_major from the identity [M, K] activation with the same +// bank-skewed kStageK + 16 B LDS row stride; A staging / global traffic / +// epilogue are unchanged. // -// Round 1 (baseline): the 64x128 tile measured 517.4 us median / 49.81 TOPS; -// routing the exact gate_up shape to the 64x64 tile measured 432.5 us -// (accepted; 19.6% improvement). -// -// Round 2 (operand-reuse, mandated decision): cooperative A+B LDS staging is -// kept and direct per-wave global loads are rejected (A is reused 64x per -// macro-tile via kBlockN and B 64x via kBlockM; direct loads would refetch -// both operands once per wavefront, 4x the 28.3 MB staged global traffic -> -// ~1.73M vmem_read instructions vs 432K, and the TP4 gate_up direct-A arm -// measured 7.9x slower). The focused change makes the staged B layout -// bank-safe for the fragment consumers: the gate_up weight is packed once -// [K,N] -> [N,K] (n-major, k contiguous) and B fragments are loaded -// col_major from a 144-B (128+16) row-strided LDS tile, replacing the -// byte-scattered row_major B reads on the identity [K,N] layout that are the -// 17.3M PMC bank-conflict source. Global loads stay 16-byte vectorized and -// coalesced on both operands. LDS stays 19,456 B/block -> 3 blocks/CU -// (58,368 B, unchanged; round-1's "4 blocks/CU" estimate was impossible on -// both LDS (4x19,456 = 77,824 B > 64 KiB) and VGPR (4x256x72 = 73,728 > -// 65,536)). down_proj keeps the identity 64x128 arm; every unmatched shape -// keeps the generic arms and the scalar fallback (which decodes the packed -// [N,K] layout when (k, n) == (6144, 512) so the paired M=2/M=16 gate_up -// tails stay correct). -// * A[kBlockM, kStageK] and B staged in bank-skewed LDS (row strides -// 144/80/144 = 16-byte-aligned non-power-of-two), consumed with the -// library du_load_matrix_sync loaders: A row_major (unchanged, load- -// bearing on every shape), B row_major for the identity [K, N] layout -// (down_proj + unmatched shapes) and B col_major for the gate_up -// [N, K] packed layout (round 2). +// Geometry (unchanged from the accepted bootstrap, so this round isolates the +// operand-layout factor): +// * 64x64 macro-tile, 4 wavefronts per block (256 threads), each wavefront +// owns one 32x32 quadrant = four resident m16n16k32 int32 accumulators +// held across the whole K loop. +// * SINGLE-buffered K loop: cooperative vectorized (int4) coalesced A/B +// global loads into bank-skewed LDS, then the MMAC burst, then two +// barriers per stage. No split-K, no raw asm. Round 11 (final consolidate +// polish) halves the PACKED arm's K stage from 128 to 64 (kPackedStageK); +// the identity [K, N] arm keeps the 128-K stage byte-identically. +// * Identity [K, N] arm: LDS = A[64, 144] (9,216 B) + B[128, 80] +// (10,240 B) = 19,456 B/block -> 3 blocks/CU. Packed [N, K] arm (round 11 +// geometry): 64-K stage -> LDS = A[64, 80] (5,120 B) + B[64, 80] +// (5,120 B) = 10,240 B/block, so it is no longer LDS-limited and reaches +// the 4-block/CU VGPR ceiling (4 x 256 threads x 64 VGPR = 65,536, the +// whole per-CU VGPR file; 4 x 10,240 = 40,960 B <= 64 KiB LDS), i.e. 16 +// resident waves/CU instead of 12. +// * Grid = (N/64, ceil(M/64)) = (8, 64) = 512 blocks. // * Fused direct fragment -> x_scale -> weight_scale -> bf16 epilogue using -// the verified gfx928 accumulator ownership (row = lane & 15, -// col_mod4 = lane >> 4, x[i] -> columns col_mod4 + 4*i). The int32 dot is -// exact for every assigned shape (max |dot| = 6144*127*127 << 2^31), the -// int32 -> float conversion is exact, and the single fp32 multiply -// followed by one bfloat16 rounding matches the reference bit-for-bit. -// * Tail M rows are zero-filled on load and masked on store, so any -// M >= 128 is supported by the tiled path. -// * All other (m, n, k) - including every M < 128 API case and the paired -// M=2/M=16 shapes with the same (N, K) - go to a scalar int8/int32 +// the gfx928 int8 DUMMA accumulator ownership (row = lane & 15, +// col_mod4 = lane >> 4, frag.x[i] -> columns col_mod4 + 4*i). The int32 +// dot is exact for every assigned K (max |dot| = 6144*127*127 << 2^31), +// and the exact fp32 order (dot * x_scale[m]) * weight_scale[n] then one +// bfloat16 rounding matches the trusted int64 reference bit-for-bit. +// * Tail M rows are zero-filled on load and masked on store, so any M >= 128 +// works with the tiled path. +// * Every other (m, n, k) - including every M < 128 API case and the paired +// M=2/M=16 shapes with the same (N, K) - goes to a scalar int8/int32 // fallback that decodes the [N, K] packed layout when (k, n) == // (6144, 512) and the identity [K, N] layout otherwise. // -// Round 3 (pipeline, mandated decision): compare single and double buffering -// across K tiles. The exact (4096, 512, 6144) guard now routes to -// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true>: a software- -// pipelined 64-K double-buffered K loop (prologue + one barrier per stage, -// 97 barriers for 96 stages) that issues stage s+1's global loads into the -// idle LDS buffer before the stage-s MMAC burst, overlapping the -// global-load latency that the round-2 single buffer serializes in front of -// every stage (ISA-verified order: global_load_dwordx4 -> ds_write_b128 -> -// s_waitcnt -> s_barrier -> ds_read2 x16 -> v_mmac x16 -> s_barrier, with -// the next stage's loads gated behind the previous stage's barrier). The -// 128-K double-buffered variant is arithmetically impossible at 3 blocks/CU -// (2 x 19,456 = 38,912 B/block x 3 = 116,736 > 64 KiB -> 1 block/CU, a -// harmful occupancy loss), so the pipelined comparison uses the 64-K -// double buffer: LDS 20,480 B/block x 3 = 61,440 <= 64 KiB and -// VGPR 64-72 x 256 x 3 <= 65,536, i.e. 3 blocks/CU preserved (no occupancy -// confound). Everything else is byte-identical: same int4 staging pattern -// (same 28.3 MB global traffic and ~432K vmem_read instructions), same -// k0-outer/kk-inner int32 order (stage size changes only regroup exact -// integer sums, so results stay bit-identical), same col_major B fragments -// (stride 80 = 64+16 bank skew), same fused epilogue, same pack, same -// generic single-buffered arms and scalar fallback. Decision rule (mandate): -// retain double buffering only if PMC shows reduced VMEM/LDS stall evidence -// (lds_wait_instructions drops below ~1.2M, lds_instructions stays ~1.97M, -// occupancy stays 3 blocks/CU) AND the median improves >2%; otherwise the -// round's evidence retains the single-buffered round-2 structure. -// -// Round 10 (fragment-load-form round, restarted from the accepted round-3 -// source after the round-9 infra kill): replace the library -// du_load_matrix_sync fragment loads of the exact-shape gate_up kernel with -// the lineage-validated explicit load_frag8 8-byte ds_read_b64 loader. The -// exact (4096, 512, 6144) guard now routes to -// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true, true> (new -// kDirectFrag template arm). The round-3 ISA shows every fragment load as -// ds_read2_b32 + per-byte v_and/v_or_b32_sdwa/v_or3 mask-OR "identity" -// reassembly (~9.4M of the 15.43M VALU instructions per replay); loading the -// same 8 bytes per lane as one int64 (ds_read_b64) removes that VALU and -// halves the fragment bank-conflict factor (32-lane dword phases at stride -// 80 -> 16-lane 8-byte phases). Element-to-slot fragment mapping, v_mmac -// operands, exact k0-outer/kk-inner int32 accumulation, staging/global -// traffic, barriers, tile, grid, LDS bytes, occupancy (3 blocks/CU) and the -// fused epilogue are unchanged; results stay bit-identical. The generic arms -// (<64,128,128>, <128,64,128>, <64,64,128>, all kBNMajor=false) keep the -// library loaders and byte-identical codegen/resources (if constexpr). -// -// Round 11 (LDS bank-skew round, HIP-only; ISA policy keeps raw asm -// disallowed because the HIP plateau rule - 8 valid HIP rounds within +/-2% -// of the then-current best - is unmet). The accepted round-10 kernel still -// shows PMC lds_bank_conflicts = 9,437,184, EXACTLY unchanged from the -// round-3 code object (the -9.07% round-10 win came from VALU removal, not -// from conflicts). 9,437,184 = 192 x 512 blocks x 96 stages: every -// block-stage pays 192 conflict cycles = 128 from the 16 ds_read2_b64 -// fragment reads + 64 from the 8 ds_write_b128 staging stores (the 16-B/lane -// store floor). The read half is a stride defect: at stride 80 = 20 dwords, -// 20*r mod 32 has period 8, so every 16-lane read phase groups rows r and -// r+8 onto the same bank pair - a 2-way conflict on ALL 4 sub-phases of -// every fragment read. The exact (4096, 512, 6144) guard now routes to -// w8a8_dumma_prefill_tiled_kernel<64, 64, 64, true, true, true, true> (new -// kSkew8 template arm) which stages both operands at the down_proj-validated -// 8-mod-16 row strides A = 88 (64+24) and B = 72 (64+8) instead of 80: -// 22*r and 18*r mod 32 are distinct over r = 0..15 and gcd(22,32) = -// gcd(18,32) = 2 does not divide 1, so every 16-lane read phase lands on 16 -// DISTINCT bank pairs - the fragment reads become conflict-free (128 of the -// 192 cycles removed; only the 2-way 16-B/lane staging-store floor of 64 -// cycles remains -> lds_bank_conflicts expected ~3.15M). The 8-mod-16 row -// starts forbid 16-byte staging stores, so each 16-byte int4 is staged as -// two 8-byte halves (ds_write2_b64), the validated down_proj staging form -// (its B tile is exactly stride 72; its A tile exactly stride 88). Global -// loads (int4, same 432,128 vmem_read), pack layout, barrier structure -// (97/block), tile/grid (512 blocks), LDS bytes (A[64,88] + B[64,72] x 2 = -// 20,480 B/block -> 3 blocks/CU unchanged), VGPR, the k0-outer/kk-inner -// int32 order, the element-to-slot fragment mapping and the fused epilogue -// are untouched; only the LDS row strides and the staging store granularity -// change, so results stay bit-identical. Generic arms (<64,128,128>, -// <128,64,128>, <64,64,128>) keep kSkew8 = false and byte-identical -// codegen/resources (if constexpr). +// Iteration 14 (consolidate; final HIP-only polish of the fixed packed_weight +// architecture): the accepted packed arm's four fragment loads per MMAC step +// are switched from the library du_load_matrix_sync lowering (ds_read2_b32 +// plus per-byte v_and / v_or_b32_sdwa / v_or3 mask-OR "identity" reassembly) +// to an explicit 8-byte load_frag8 int64 (ds_read_b64) loader, for the packed +// n-major arm ONLY (new kDirectFrag template flag, set true for the exact +// (K, N) = (6144, 512) dispatch; the identity [K, N] arm keeps the library +// loaders byte-identically). Lane -> (row = lane & 15, k-quarter = lane >> 4) +// and the eight contiguous bytes at p[row * ldm + (lane >> 4) * 8] land in +// the same x[0..7] slots in the same little-endian order, so the v_mmac +// operands, the exact k0-outer / kk-inner int32 accumulation order and every +// result are unchanged. LDS strides (A = kStageK + 16, packed B = kStageK + +// 16), int4 staging, global traffic, single buffer, two barriers per stage, +// tile, grid (512 blocks), 4 waves / 256 threads, 10,240 B LDS, occupancy and +// the fused epilogue are byte-identical; both pack kernels, the host shape +// guards and the scalar fallback are untouched. // -// Round 16 (lead-2 loop-carried payload round, HIP-only; raw asm stays -// DISALLOWED - plateau=false, the recent-valid-improvement window -// [-34.44%, +2.82%, +9.97%] does not meet three valid HIP rounds within -// [-2%, +2%) of the then-current best, and no prior ISA-guided round -// recorded a compiler limitation with target instructions). The accepted -// round-11 exact code object (digest 9889820e, median 293.30 us / p90 -// 293.545) shows the software-pipeline intent is DEFEATED by the compiler's -// wait placement: per stage the steady-state loop is -// global_load_dwordx4 (A prefetch s+1) -> global_load_dwordx4 (B prefetch -// s+1) -> s_waitcnt vmcnt(1) -> s_waitcnt vmcnt(0) <- waits IMMEDIATELY -// -> 2 x ds_write2_b64 (staging stores) -> 4 x ds_read2_b64 (fragment -// reads, stage s) -> 8 x v_mmac CONTIGUOUS (stage s) -> s_barrier -> loop. -// i.e. the stage-(s+1) global loads are waited ~12-16 instructions after -// issue, BEFORE the stage-s MMAC burst, so the full global-load latency -// (~600-1000+ cycles under 120-CU L2 contention) sits on the per-stage -// critical path with zero overlap (the round-12 analysis' -// "prefetch global load -> ds_write2_b64 -> s_barrier -> ds_read2_b64 -> -// v_mmac" chain; PMC: 6.0M VALU + 1.18M LDS + 465K VMEM wave-instructions -// over 293 us x 120 CUs, i.e. ~95-99% issue-stalled). Iteration history -// rules out the other levers on this exact shape: tile aspect (iter 4 -// 32x64: 749.9 us; iter 8 64x128 w8: 395.4 us), 2-wave quadrants (iter 6: -// 629.0 us - wave count per CU is load-bearing), LDS-side fragment -// prefetch (iter 7: 337.97 us, flat/worse), stage 32 / 6 blocks per CU -// (iter 12: 447.4 us - per-stage fixed cost, not occupancy, is the lever). -// The one untested mechanism that directly removes the exposed vmcnt wait: -// a LEAD-2 LOOP-CARRIED REGISTER PAYLOAD (the TP4 o_proj lineage's -// validated "stage-top publish" form). The exact (4096, 512, 6144) guard -// now routes to w8a8_dumma_prefill_tiled_kernel<64,64,64,true,true,true, -// true,true> (new kPayload template arm): each thread keeps the one A int4 -// + one B int4 of stage s+2 in VGPR across the stage-(s+1) barrier, issues -// those global loads at the TOP of iteration s (no wait), and PUBLISHES -// the payload carried from iteration s-1 into the idle LDS buffer at the -// TOP of iteration s (before the stage-s MMAC burst). The first use of the -// payload registers is the next iteration's publish, and the s_barrier -// fences outstanding vmem, so the compiler can no longer place -// s_waitcnt vmcnt(0) immediately after the loads: the wait lands after the -// stage-s fragment reads + MMAC burst (immediately before the barrier), and -// the global-load latency overlaps the whole compute segment instead of -// serializing in front of it. One barrier per 64-K stage is preserved -// (prologue + 96 = 97/block); the publish targets the idle buffer (WAR -// retired by the previous barrier, RAW retired by the stage barrier), so no -// second barrier is introduced. Everything else is byte-identical: 64x64 -// tile / 256 threads / 4 waves / 32x32 quadrant per wave, grid (8,64) = -// 512 blocks, kSkew8 strides 88/72 with ds_write2_b64 staging (8 per -// block-stage), same 16 ds_read2_b64 fragment reads per block-stage (same -// lds_instructions ~1,179,648), same 20,480 B LDS/block -> 3 blocks/CU -// (payload costs 8 VGPRs: ~54 total, 54 x 256 x 3 = 41,472 <= 65,536; no -// spill risk at 3 blocks/CU), same int4 global loads (same ~432,128 -// vmem_read / 32,768 vmem_write), same k0-outer/kk-inner int32 order and -// the same element-to-slot fragment mapping - the payload changes only WHEN -// global bytes land in LDS, never which bytes or the accumulation order, so -// results stay bit-identical (mismatch 0 / max_abs_error 0.0 expected). -// Generic arms (<64,128,128>, <128,64,128>, <64,64,128>, all kBNMajor= -// false) keep kPayload = false and byte-identical codegen/resources -// (if constexpr). -// -// Header order is fixed by the control plane: hip_runtime first (du_mma.h -// is not self-contained before it), hip_bfloat16 second, du_mma.h last. +// Header order is fixed by the control plane: hip_runtime first (du_mma.h is +// not self-contained before it), hip_bfloat16 second, du_mma.h last. #include #include @@ -224,21 +115,51 @@ constexpr int kTileK = 32; // gfx928 wavefront is 64 lanes; every blockDim must be a multiple of 64. constexpr int kWaveSize = 64; -// Large-prefill dispatch geometry. Each macro-tile config uses a 32x32 -// quadrant per wave, so waves per block = (BlockM/32) * (BlockN/32). +// Large-prefill dispatch geometry: one 32x32 DUMMA quadrant per wavefront. constexpr int kPrefillMinM = 128; constexpr int kPrefillStageK = 128; constexpr int kDefaultBlockM = 64; constexpr int kDefaultBlockN = 64; +// Exact assigned shape: the only (K, N) pair whose packed weight is stored +// n-major (transposed). Every other pair keeps the identity [K, N] layout. +constexpr int kGateUpK = 6144; +constexpr int kGateUpN = 512; + +// Round 11 (consolidate): the fixed winning packed_weight architecture keeps +// its 64x64 macro-tile, 4-wave/256-thread block, single-buffered stage +// structure and 16-byte bank skew; only the PACKED arm's K stage depth is +// polished from the generic 128 down to 64. The packed arm's LDS footprint +// drops from A[64,144] + B[64,144] = 18,432 B to A[64,80] + B[64,80] = +// 10,240 B/block, so the 64 KiB per-CU LDS part no longer caps it at 3 +// blocks/CU and it reaches the 4-blocks/CU VGPR ceiling (4 x 256 x 64 = the +// whole 65,536-dword per-CU VGPR file; 4 x 10,240 = 40,960 B). That is 16 +// resident waves/CU instead of 12, giving +33% concurrent coverage of the +// single-buffered global-load -> barrier -> MMAC -> barrier dependency that +// dominates each stage. The current PMC confirms the stage is latency- rather +// than bandwidth-limited: 97.6 GB/s counter-derived HBM bandwidth against +// 88.8 GB/s algorithmic bandwidth at a 91.8% L2 hit rate, so the staged +// operand traffic is far from saturating the memory system. +// +// k = 6144 = 96 x 64, kPackedStageK % kTileK == 0 and the kk inner order +// still advances in 32-K DUMMA steps, so the k0-outer / kk-inner int32 +// accumulation visits the exact same K sequence and every result stays +// bit-identical to the accepted 128-K stage. The identity [K, N] arm, the +// host shape guards, both pack kernels and the scalar fallback are untouched. +constexpr int kPackedStageK = 64; +static_assert(kPackedStageK % kTileK == 0, + "packed K stage must be a multiple of the DUMMA K unit"); +static_assert(kGateUpK % kPackedStageK == 0, + "packed K stage must tile the assigned gate_up K exactly"); + constexpr int kScalarThreads = 256; using namespace du::dumma; // Verified gfx928 INT8 DUMMA accumulator ownership: // row = lane & 15, col_mod4 = lane >> 4, frag.x[i] -> columns -// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store -// bf16 directly from the fragment (no accumulator LDS round trip). +// col_mod4 + 4*i. Scale by x_scale[row] and weight_scale[col] and store bf16 +// directly from the fragment (no accumulator LDS round trip). template __device__ __forceinline__ void store_prefill_fragment( const AccFragment& frag, @@ -259,38 +180,30 @@ __device__ __forceinline__ void store_prefill_fragment( #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]; + const float scaled = + static_cast(frag.x[i]) * xs * weight_scale[col]; out[row * n + col] = __float2bfloat16(scaled); } } -// Round 10 (fragment-load-form round): explicit 8-byte fragment loader -// replacing the library du_load_matrix_sync for the exact-shape gate_up -// kernel. Validated on the sibling TP8 down_proj accepted kernel (iter 8, -// -25.9%: 230.91 -> 171.19 us) and the TP4 gate_up lineage (iter 13 -// load_fragment8, -21%: 280.05 -> 220.71 us). For BOTH the library row_major -// matrix_a and col_major matrix_b m16n16k32 int8 loaders, lane l -> row -// (l & 15), k-quarter (l >> 4), and the loader fetches the same eight -// contiguous bytes at p[row*ldm + (l>>4)*8 .. +7] and stores them into the -// same x[0..7] slots in the same little-endian order (the round-3 ISA shows -// the compiler lowering this to ds_read2_b32 followed by a per-byte -// v_and/v_or_b32_sdwa/v_or3 mask-OR "identity" reassembly - ~9.4M dead-ish -// VALU instructions per replay, ~61% of the 15.43M VALU total). Loading the -// same 8 bytes as one int64 forces the 8-byte ds_read_b64 form instead: -// (a) the mask/OR reassembly VALU disappears, and (b) the LDS access -// changes from ds_read2_b32 (32-lane dword phases - at the stride-80 row -// pattern every phase groups rows 0-15 of k-quarter q with rows 0-15 of -// k-quarter q+1 onto the same 16 bank phases, a 4-way conflict per phase) -// to ds_read_b64 (16-lane phases: rows 0-15 land on 16 bank-pairs -// {20*r mod 32} = 8 distinct pairs hit twice - a 2-way conflict per phase, -// half of the current fragment-conflict factor). The element-to-slot -// mapping is unchanged (f.x[0..7] = the same 8 bytes in the same -// little-endian order), the v_mmac operands and the exact k0-outer/kk-inner -// int32 accumulation sequence are untouched, and the staging stores, LDS -// strides, barriers, tile, grid, occupancy and epilogue are byte-identical. +// Iteration 14 (consolidate polish, packed arm only): explicit 8-byte fragment +// loader replacing the library du_load_matrix_sync calls in the packed +// n-major arm. For BOTH the library row_major matrix_a and col_major +// matrix_b m16n16k32 int8 loaders, lane l -> row (l & 15), k-quarter +// (l >> 4), and the loader fetches the same eight contiguous bytes at +// p[row * ldm + (l >> 4) * 8 .. +7] and stores them into the same x[0..7] +// slots in the same little-endian order (the accepted code object lowers the +// library calls to ds_read2_b32 plus a per-byte v_and / v_or_b32_sdwa / +// v_or3 mask-OR "identity" reassembly). Loading the same 8 bytes as one +// int64 forces the 8-byte ds_read_b64 form: the reassembly VALU disappears +// and the LDS access moves from 32-lane dword phases to 16-lane 8-byte +// phases. The element-to-slot mapping, the v_mmac operands, the +// k0-outer/kk-inner int32 accumulation order, the LDS strides (A = kStageK + +// 16, packed B = kStageK + 16), staging, barriers, tile/grid/occupancy and +// the fused epilogue are unchanged, so results stay bit-identical. template -__device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, - unsigned ldm) { +__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 = @@ -298,22 +211,21 @@ __device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, reinterpret_cast(f.x)[0] = v; } -// Stage one K tile (k0) of A[kBlockM, kStageK] and B into the caller-offset -// LDS buffers (bank-skewed row strides; int4 vectorized coalesced global -// loads, zero-filled tail M rows). kBNMajor == true stages B from the packed -// [N, K] n-major gate_up layout (n-major rows, k contiguous); otherwise from -// the identity [K, N] layout. Shared by the single-buffered and -// double-buffered K loops; the buffer offset (0 or kABytes/kBMaxBytes) is -// applied by the caller so one body serves both pipeline forms. +// Stage one K tile (k0) of A[kBlockM, kStageK] and B into the caller LDS +// buffers with int4-vectorized coalesced global loads. Tail M rows are +// zero-filled so the MACs contribute nothing and the guarded store never +// writes them. // -// kSkew8 (round 11, exact-shape arm only): A and B row strides become the -// down_proj-validated 8-mod-16 values 88 (kStageK + 24) and 72 (kStageK + 8) -// so every 16-lane fragment-read phase lands on 16 distinct bank pairs (see -// the round-11 header comment). The 8-mod-16 row starts forbid 16-byte -// stores, so each 16-byte int4 chunk is staged as two 8-byte halves -// (ds_write2_b64); the global int4 loads are byte-identical. -template +// kBNMajor == true stages B[kBlockN, kStageK] from the packed [N, K] layout +// (n-major rows, k contiguous) that launch_pack_w8a8_weight produces for the +// exact gate_up (k, n) pair: row = local N, columns = k, row stride +// kStageK + 16 B (80 B at the round-11 64-K packed stage). The B fragment +// consumer then reads 8 consecutive k bytes per lane (col_major), replacing +// the byte-scattered identity-layout row_major reads that source the measured +// 17.3M bank conflicts. +// kBNMajor == false stages B[kStageK, kBlockN] from the identity [K, N] +// weight with row stride kBlockN + 16 = 80 B (unchanged bootstrap path). +template __device__ __forceinline__ void stage_prefill_tiles( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, @@ -327,18 +239,17 @@ __device__ __forceinline__ void stage_prefill_tiles( int k, int k0) { constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + constexpr int kInt4 = static_cast(sizeof(int4)); // 16-byte bank skew: keeps int4 staging stores aligned while spreading the // fragment rows across distinct LDS banks. - constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); - constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows - constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); - static_assert(!kSkew8 || kStageK == 64, - "kSkew8 strides 88/72 assume kStageK == 64"); - constexpr int kAInt4PerRow = kStageK / sizeof(int4); + constexpr int kAStride = kStageK + kInt4; + constexpr int kBStride = kBlockN + kInt4; // identity [K, N] rows + constexpr int kBNStride = kStageK + kInt4; // packed [N, K] rows + constexpr int kAInt4PerRow = kStageK / kInt4; constexpr int kAInt4 = kBlockM * kAInt4PerRow; - constexpr int kBInt4PerRow = kBlockN / sizeof(int4); + constexpr int kBInt4PerRow = kBlockN / kInt4; constexpr int kBInt4 = kStageK * kBInt4PerRow; - constexpr int kBNInt4PerRow = kStageK / sizeof(int4); + constexpr int kBNInt4PerRow = kStageK / kInt4; constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; // Stage A[kBlockM, kStageK] into LDS (zero-filled tail M rows). for (int vec = tid; vec < kAInt4; vec += kThreads) { @@ -349,162 +260,52 @@ __device__ __forceinline__ void stage_prefill_tiles( global_row < m ? *reinterpret_cast( x_q + static_cast(global_row) * k + k0 + - v * static_cast(sizeof(int4))) + v * kInt4) : int4{0, 0, 0, 0}; - if constexpr (kSkew8) { - // 8-mod-16 row starts: write the 16 bytes as two 8-byte halves so every - // store stays 8-byte aligned (ds_write2_b64 staging; same LDS bytes). - int8_t* dst = - a_tile + local_row * kAStride + v * static_cast(sizeof(int4)); - reinterpret_cast(dst)[0] = - reinterpret_cast(&val)[0]; - reinterpret_cast(dst)[1] = - reinterpret_cast(&val)[1]; - } else { - reinterpret_cast(a_tile + local_row * kAStride)[v] = val; - } + reinterpret_cast(a_tile + local_row * kAStride)[v] = val; } if constexpr (kBNMajor) { - // Stage B[kBlockN, kStageK] from the packed [N, K] weight (n-major: - // row = n_local, k contiguous). Same vectorized coalesced int4 pattern - // as the identity arm, one per 16 consecutive k bytes. + // Stage B[kBlockN, kStageK] from the packed [N, K] weight (n-major rows, + // k contiguous): row = local N, 16 consecutive k bytes per int4. + static_assert(kBNInt4 % kThreads == 0, + "packed n-major B staging must divide evenly"); for (int vec = tid; vec < kBNInt4; vec += kThreads) { const int local_col = vec / kBNInt4PerRow; const int v = vec - local_col * kBNInt4PerRow; - const int4 val = *reinterpret_cast( - weight + static_cast(n0 + local_col) * k + k0 + - v * static_cast(sizeof(int4))); - if constexpr (kSkew8) { - int8_t* dst = - b_tile + local_col * kBNStride + - v * static_cast(sizeof(int4)); - reinterpret_cast(dst)[0] = - reinterpret_cast(&val)[0]; - reinterpret_cast(dst)[1] = - reinterpret_cast(&val)[1]; - } else { - reinterpret_cast(b_tile + local_col * kBNStride)[v] = val; - } + reinterpret_cast(b_tile + local_col * kBNStride)[v] = + *reinterpret_cast( + weight + static_cast(n0 + local_col) * k + k0 + + v * kInt4); } } else { - // Stage B[kStageK, kBlockN] into LDS from the identity [K, N] weight. + // Stage B[kStageK, kBlockN] from the identity [K, N] weight. for (int vec = tid; vec < kBInt4; vec += kThreads) { const int kk = vec / kBInt4PerRow; const int v = vec - kk * kBInt4PerRow; reinterpret_cast(b_tile + kk * kBStride)[v] = *reinterpret_cast( - weight + static_cast(k0 + kk) * n + n0 + - v * static_cast(sizeof(int4))); + weight + static_cast(k0 + kk) * n + n0 + v * kInt4); } } } -// Round 16 (lead-2 loop-carried payload, exact-shape arm only): the register -// half of the lead-2 software pipeline. Each thread prefetches exactly one A -// int4 and one B int4 of a later K stage into thread-local registers and -// returns WITHOUT waiting (the vmcnt wait is emitted by the compiler at the -// first use, one iteration later). The <64,64,64> arm stages kStageK * 64 -// bytes per operand per stage, i.e. kAInt4 == kBNInt4 == kThreads == 256: -// one int4 per thread per operand, so the whole payload costs 8 VGPRs. The -// (row, v) mapping is the same as stage_prefill_tiles, so -// publish_stage_payload writes the exact same LDS addresses the staging -// loop would (bit-identical staging bytes, just loaded one stage earlier). -template -__device__ __forceinline__ void prefetch_stage_payload( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, - int tid, - int m0, - int n0, - int m, - int k, - int k0, - int4& a_payload, - int4& b_payload) { - constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; - constexpr int kAInt4PerRow = kStageK / static_cast(sizeof(int4)); - constexpr int kAInt4 = kBlockM * kAInt4PerRow; - constexpr int kBNInt4PerRow = kStageK / static_cast(sizeof(int4)); - constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; - static_assert(kAInt4 == kThreads && kBNInt4 == kThreads, - "payload arm requires exactly one int4 per thread per " - "operand (exact-shape <64,64,64> arm)"); - static_assert(kBNMajor, "payload arm stages the packed n-major B layout"); - // Same vec -> (local_row, v) mapping as stage_prefill_tiles, so the - // publish half stores to the same LDS addresses. - const int a_local_row = tid / kAInt4PerRow; - const int a_v = tid - a_local_row * kAInt4PerRow; - const int global_row = m0 + a_local_row; - // Tail M rows are zero-filled exactly as in stage_prefill_tiles. - a_payload = global_row < m - ? *reinterpret_cast( - x_q + static_cast(global_row) * k + k0 + - a_v * static_cast(sizeof(int4))) - : int4{0, 0, 0, 0}; - const int b_local_col = tid / kBNInt4PerRow; - const int b_v = tid - b_local_col * kBNInt4PerRow; - b_payload = *reinterpret_cast( - weight + static_cast(n0 + b_local_col) * k + k0 + - b_v * static_cast(sizeof(int4))); -} - -// Round 16: the LDS half of the lead-2 payload pipeline. Writes the payload -// carried in registers (loaded by prefetch_stage_payload one stage earlier) -// into the caller-offset LDS tile with the exact same 8-byte-half -// ds_write2_b64 stores (kSkew8 strides 88/72) that stage_prefill_tiles -// emits, so each thread republishes exactly the bytes it loaded. -template -__device__ __forceinline__ void publish_stage_payload( - int4 a_payload, - int4 b_payload, - int8_t* a_tile, - int8_t* b_tile, - int tid) { - constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; - constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); - constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); - constexpr int kAInt4PerRow = kStageK / static_cast(sizeof(int4)); - constexpr int kAInt4 = kBlockM * kAInt4PerRow; - constexpr int kBNInt4PerRow = kStageK / static_cast(sizeof(int4)); - constexpr int kBNInt4 = kBlockN * kBNInt4PerRow; - static_assert(kAInt4 == kThreads && kBNInt4 == kThreads, - "payload arm requires exactly one int4 per thread per " - "operand (exact-shape <64,64,64> arm)"); - static_assert(kBNMajor, "payload arm stages the packed n-major B layout"); - const int a_local_row = tid / kAInt4PerRow; - const int a_v = tid - a_local_row * kAInt4PerRow; - int8_t* a_dst = - a_tile + a_local_row * kAStride + a_v * static_cast(sizeof(int4)); - if constexpr (kSkew8) { - // 8-mod-16 row starts: write the 16 bytes as two 8-byte halves - // (ds_write2_b64 staging, identical to stage_prefill_tiles). - reinterpret_cast(a_dst)[0] = - reinterpret_cast(&a_payload)[0]; - reinterpret_cast(a_dst)[1] = - reinterpret_cast(&a_payload)[1]; - } else { - reinterpret_cast(a_dst)[0] = a_payload; - } - const int b_local_col = tid / kBNInt4PerRow; - const int b_v = tid - b_local_col * kBNInt4PerRow; - int8_t* b_dst = - b_tile + b_local_col * kBNStride + b_v * static_cast(sizeof(int4)); - if constexpr (kSkew8) { - reinterpret_cast(b_dst)[0] = - reinterpret_cast(&b_payload)[0]; - reinterpret_cast(b_dst)[1] = - reinterpret_cast(&b_payload)[1]; - } else { - reinterpret_cast(b_dst)[0] = b_payload; - } -} - -// One wave's 32x32-quadrant MMAC burst for the K tile staged at a_tile/b_tile -// (caller-offset to the target buffer). k0-outer / kk-inner int32 accumulation -// order; stage-size changes only regroup exact integer sums (max |dot| = -// 6144*127*127 << 2^31), so results stay bit-identical. +// One wave's 32x32-quadrant MMAC burst for the K tile staged at a_tile/b_tile. +// k0-outer / kk-inner int32 accumulation order; stage-size changes only +// regroup exact integer sums, so results stay bit-identical to the scalar +// fallback and to the int64 reference. +// +// kBNMajor == true consumes B as matrix_b col_major fragments (8 consecutive +// k bytes per lane, row stride kStageK + 16 B) from the packed [N, K] tile; the +// element-to-slot mapping and the v_mmac operand bit patterns are the same as +// the identity row_major path, so the int32 accumulation stays bit-identical. +// +// kDirectFrag == true (iteration 14, packed arm only) replaces the four +// library du_load_matrix_sync fragment loads with the explicit 8-byte +// load_frag8 int64 loader, which halves the LDS read phase count and removes +// the per-byte mask-OR reassembly VALU; the loaded bytes and fragment slots +// are identical. template + bool kDirectFrag = false> __device__ __forceinline__ void mmac_prefill_stage( DUFragment& a_frag0, @@ -518,56 +319,26 @@ __device__ __forceinline__ void mmac_prefill_stage( int8_t* b_tile, int local_row, int local_col) { - // kSkew8 (round 11, exact-shape arm only): A 88 / B 72 row strides; the - // load_frag8 ldm arguments below pick them up unchanged (same 8-byte - // aligned fragment reads, conflict-free 16-lane bank phases). - constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); - constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows - constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); + constexpr int kInt4 = static_cast(sizeof(int4)); + constexpr int kAStride = kStageK + kInt4; + constexpr int kBStride = kBlockN + kInt4; // identity [K, N] rows + constexpr int kBNStride = kStageK + kInt4; // packed [N, K] rows if constexpr (kBNMajor) { - if constexpr (kDirectFrag) { - // Round 10 exact-shape path: explicit load_frag8 8-byte ds_read_b64 - // fragment reads replacing the library loaders (see load_frag8 - // comment). Both steps of the 64-K stage are unrolled into - // independent fragment register sets so all eight reads of a stage - // issue before the MMACs; the kk = 0 MMAC group runs first, then the - // kk = kTileK group - the exact k0-outer/kk-inner int32 order and the - // element-to-slot fragment mapping of the library path are preserved - // (f.x[0..7] = the same 8 bytes in the same little-endian order), so - // results stay bit-identical. - static_assert(kStageK == 2 * kTileK, - "kDirectFrag two-step unroll assumes kStageK == 2*kTileK"); - DUFragment - a_frag0_1, a_frag1_1; - DUFragment - b_frag0, b_frag1, b_frag0_1, b_frag1_1; - load_frag8(a_frag0, a_tile + local_row * kAStride, kAStride); - load_frag8(a_frag1, a_tile + (local_row + kTileM) * kAStride, kAStride); - load_frag8(b_frag0, b_tile + local_col * kBNStride, kBNStride); - load_frag8(b_frag1, b_tile + (local_col + kTileN) * kBNStride, - kBNStride); - load_frag8(a_frag0_1, a_tile + local_row * kAStride + kTileK, - kAStride); - load_frag8(a_frag1_1, - a_tile + (local_row + kTileM) * kAStride + kTileK, kAStride); - load_frag8(b_frag0_1, b_tile + local_col * kBNStride + kTileK, - kBNStride); - load_frag8(b_frag1_1, - b_tile + (local_col + kTileN) * kBNStride + kTileK, - kBNStride); - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); - du_mma_sync(acc00, a_frag0_1, b_frag0_1, acc00); - du_mma_sync(acc01, a_frag0_1, b_frag1_1, acc01); - du_mma_sync(acc10, a_frag1_1, b_frag0_1, acc10); - du_mma_sync(acc11, a_frag1_1, b_frag1_1, acc11); - } else { - DUFragment - b_frag0, b_frag1; + DUFragment + b_frag0, b_frag1; #pragma unroll - for (int kk = 0; kk < kStageK; kk += kTileK) { + for (int kk = 0; kk < kStageK; kk += kTileK) { + if constexpr (kDirectFrag) { + // Iteration 14: one int64 ds_read_b64 per fragment (same 8 bytes per + // lane into the same x[0..7] slots) instead of the library + // ds_read2_b32 + mask/OR reassembly form. + load_frag8(a_frag0, a_tile + local_row * kAStride + kk, kAStride); + load_frag8(a_frag1, + a_tile + (local_row + kTileM) * kAStride + kk, kAStride); + load_frag8(b_frag0, b_tile + local_col * kBNStride + kk, kBNStride); + load_frag8(b_frag1, + b_tile + (local_col + kTileN) * kBNStride + kk, kBNStride); + } else { du_load_matrix_sync( a_frag0, a_tile + local_row * kAStride + kk, kAStride); du_load_matrix_sync( @@ -577,11 +348,11 @@ __device__ __forceinline__ void mmac_prefill_stage( du_load_matrix_sync( b_frag1, b_tile + (local_col + kTileN) * kBNStride + kk, kBNStride); - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); } + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); } } else { DUFragment @@ -592,8 +363,7 @@ __device__ __forceinline__ void mmac_prefill_stage( a_frag0, a_tile + local_row * kAStride + kk, kAStride); du_load_matrix_sync( a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); - du_load_matrix_sync( - b_frag0, b_tile + kk * kBStride + local_col, kBStride); + du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, kBStride); du_load_matrix_sync( b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); du_mma_sync(acc00, a_frag0, b_frag0, acc00); @@ -604,40 +374,23 @@ __device__ __forceinline__ void mmac_prefill_stage( } } -// Large-prefill path: one block computes a kBlockM x kBlockN output tile. -// (kBlockM/32) x (kBlockN/32) wavefronts each own a 32x32 quadrant (four -// m16n16k32 DUMMA accumulators), while the block cooperatively stages -// A[kBlockM, kStageK] and B in bank-skewed LDS. +// Single-buffered large-prefill path: one block computes a kBlockM x kBlockN +// output tile. (kBlockM/32) x (kBlockN/32) wavefronts each own a 32x32 +// quadrant (four m16n16k32 DUMMA accumulators) while the block cooperatively +// stages A[kBlockM, kStageK] and B in bank-skewed LDS. // -// kDoubleBuffer == false: single-buffered K loop (two barriers per stage); -// stage s+1's global loads cannot start until the stage-s barrier, so each -// stage serializes one global-load round trip in front of the compute -// (bootstrap structure; co-resident block streams hide it at the CU level). +// kBNMajor selects the B operand layout: false is the identity [K, N] staging +// with row_major fragments (any (k, n)); true is the packed [N, K] staging +// with col_major fragments (exact gate_up (k, n) only). // -// kDoubleBuffer == true: software-pipelined K loop (round 3). Stage s+1's -// global loads are issued into the idle LDS buffer before the stage-s MMAC -// burst, overlapping the serialized global-load latency with compute. One -// barrier per stage (plus one prologue): it both retires stage-s fragment -// reads (WAR on the buffer just read) and makes stage s+1's ds_writes -// visible (RAW on the prefetched buffer). The gate_up shape uses -// <64, 64, 64, true, true>: two 64-K tiles = 20,480 B/block x 3 blocks/CU = -// 61,440 <= 64 KiB (occupancy preserved; a 128-K double buffer would need -// 38,912 B/block -> 1 block/CU, a harmful occupancy loss). +// Grid covers N in whole tiles (dispatch guarantees N % kBlockN == 0) and M +// in ceil-divided tiles (row tail handled by the zero-fill + masked store). // -// kBNMajor selects the B operand layout: -// * false (identity): B staged [K, N] row-major, library row_major B -// fragments (byte-scattered LDS reads; correct for the raw [K, N] weight -// and every non-gate_up shape). -// * true (gate_up only): B staged [N, K] n-major with k contiguous (row -// stride 144 = 128+16, or 80 = 64+16 for the 64-K double buffer), -// library col_major B fragments -> each lane's elements are k-contiguous -// bytes instead of byte-scattered columns. This is the bank-safe consumer -// layout for the B side (the 17.3M PMC bank-conflict source in round 1). -// The one-time pack in launch_pack_w8a8_weight provides the [N, K] buffer -// for (k, n) == (6144, 512) only. +// kDirectFrag (iteration 14) selects the explicit 8-byte load_frag8 fragment +// loader for the packed n-major arm; the identity arm and every generic +// geometry keep the library du_load_matrix_sync loaders. template + bool kDirectFrag = false> __global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize) void w8a8_dumma_prefill_tiled_kernel( const int8_t* __restrict__ x_q, @@ -649,56 +402,45 @@ w8a8_dumma_prefill_tiled_kernel( int n, int k) { constexpr int kWavesN = kBlockN / 32; - constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; - // 16-byte bank skew: keeps int4 staging stores aligned while spreading the - // fragment rows across distinct LDS banks. The staging loops live in - // stage_prefill_tiles below; the strides remain here as the static-asserted - // alignment contract of every instantiation. kSkew8 (round 11, exact-shape - // arm only) switches A/B to the down_proj-validated 8-mod-16 strides - // 88/72; LDS bytes stay A[64,88] + B[64,72] x 2 = 20,480 B/block so the - // 3-blocks/CU occupancy of the accepted kernel is preserved exactly. - constexpr int kAStride = kSkew8 ? kStageK + 24 : kStageK + sizeof(int4); - constexpr int kBStride = kBlockN + sizeof(int4); // identity [K,N] rows - constexpr int kBNStride = kSkew8 ? kStageK + 8 : kStageK + sizeof(int4); - constexpr int kBNInt4 = kBlockN * (kStageK / sizeof(int4)); - // Double buffering uses kBuffers = 2 LDS tile copies; the gate_up shape - // pairs it with kStageK = 64 so 3 blocks/CU is preserved (see kernel - // comment above). - constexpr int kBuffers = kDoubleBuffer ? 2 : 1; + constexpr int kInt4 = static_cast(sizeof(int4)); + constexpr int kAStride = kStageK + kInt4; + constexpr int kBStride = kBlockN + kInt4; // identity [K, N] rows + constexpr int kBNStride = kStageK + kInt4; // packed [N, K] rows constexpr int kABytes = kBlockM * kAStride; - // One shared B tile sized for the larger of the two layouts so an inactive - // arm never allocates dead LDS (keeps <64,64,128> at 19,456 B total and - // 3 blocks/CU). kSkew8 implies kBNMajor, so only the n-major arm is live. - constexpr int kBMaxBytes = - kSkew8 ? kBlockN * kBNStride - : (kStageK * kBStride > kBlockN * kBNStride - ? kStageK * kBStride - : kBlockN * kBNStride); - static_assert(kSkew8 ? (kAStride % 8 == 0) : (kAStride % sizeof(int4) == 0), - "A LDS row stride must stay int4 aligned (or 8-byte aligned " - "under kSkew8)"); - static_assert(kBStride % static_cast(sizeof(int4)) == 0, - "B LDS row stride must stay int4 aligned"); - static_assert( - kSkew8 ? (kBNStride % 8 == 0) : (kBNStride % sizeof(int4) == 0), - "n-major B LDS row stride must stay int4 aligned (or 8-byte aligned " - "under kSkew8)"); - static_assert(!kSkew8 || kBNMajor, - "kSkew8 is only defined for the packed n-major B layout"); - static_assert(!kSkew8 || kStageK == 64, - "kSkew8 strides 88/72 assume kStageK == 64"); + constexpr int kBBytes = kStageK * kBStride; + constexpr int kBNBytes = kBlockN * kBNStride; + // The B tile is sized exactly for the layout the instantiation uses, so an + // inactive arm never allocates dead LDS: identity (kStageK=128) 128 x 80 = + // 10,240 B, packed (kStageK=64) 64 x 80 = 5,120 B. With A[64, 80] = + // 5,120 B the packed arm is 10,240 B/block, so the packed instantiation now + // reaches the 4-blocks/CU VGPR ceiling (4 x 256 x 64 = 65,536 VGPR; + // 4 x 10,240 = 40,960 <= 64 KiB LDS) instead of the 3 blocks/CU the 18,432 B + // packed arm allowed. The identity arm keeps its 19,456 B / 3 blocks/CU. + constexpr int kBAlloc = kBNMajor ? kBNBytes : kBBytes; + // Consolidation guard: the packed arm must keep the 4-blocks/CU LDS budget + // that its 64-K stage buys; this fires the build if a future stage/tile edit + // would silently drop the packed arm back to 3 blocks/CU. + static_assert(!kBNMajor || 4 * (kABytes + kBAlloc) <= 64 * 1024, + "packed arm must keep 4 blocks/CU on the 64 KiB LDS part"); static_assert(kStageK % kTileK == 0, "K stage must be a multiple of the DUMMA K unit"); - static_assert(kStageK % static_cast(sizeof(int4)) == 0, + static_assert(kStageK % kInt4 == 0, "K stage must keep int4 staging aligned"); - static_assert(kBNInt4 % kThreads == 0, - "n-major B staging must divide evenly across threads"); - static_assert(!kPayload || (kBNMajor && kDoubleBuffer && kDirectFrag && - kSkew8 && kStageK == 64 && kBlockM == 64 && - kBlockN == 64), - "payload arm is specialized to the exact-shape <64,64,64> " - "double-buffered load_frag8/kSkew8 arm (one int4 per thread " - "per operand)"); + static_assert(kAStride % kInt4 == 0, + "A LDS row stride must stay int4 aligned"); + static_assert(kBStride % kInt4 == 0, + "B LDS row stride must stay int4 aligned"); + static_assert(kBNStride % kInt4 == 0, + "packed B LDS row stride must stay int4 aligned"); + static_assert(kBlockM % 32 == 0 && kBlockN % 32 == 0, + "each wavefront owns a 32x32 DUMMA quadrant"); + static_assert(!kDirectFrag || kBNMajor, + "direct 8-byte fragment loader is defined for the packed " + "n-major arm only"); + static_assert(!kBNMajor || (kBlockN * (kStageK / kInt4)) % + (kBlockM * kBlockN / 1024 * kWaveSize) == + 0, + "packed n-major B staging must divide evenly across threads"); const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; @@ -710,8 +452,8 @@ w8a8_dumma_prefill_tiled_kernel( const int local_row = wave_row * 32; const int local_col = wave_col * 32; - __shared__ __align__(16) int8_t a_tile[kBuffers * kABytes]; - __shared__ __align__(16) int8_t b_tile[kBuffers * kBMaxBytes]; + __shared__ __align__(16) int8_t a_tile[kABytes]; + __shared__ __align__(16) int8_t b_tile[kBAlloc]; DUFragment a_frag0, a_frag1; @@ -722,111 +464,20 @@ w8a8_dumma_prefill_tiled_kernel( du_fill_fragment(acc10, 0); du_fill_fragment(acc11, 0); - if constexpr (kDoubleBuffer) { - if constexpr (kPayload) { - // Round 16 lead-2 loop-carried payload pipeline (exact-shape arm - // only): stage s+1's global loads are issued a FULL iteration early - // (top of iteration s-1) into a two-int4-per-thread register payload - // and stay outstanding across the stage-(s-1) fragment reads + MMAC - // burst and the stage barrier; the vmcnt wait lands at the top of - // iteration s (first use of the payload registers), immediately - // before the ds_write2 publish into the idle LDS buffer. The - // s_barrier fences outstanding vmem, so the wait can no longer sit - // right after the load issue: the global-load latency overlaps the - // whole compute segment instead of serializing in front of every - // stage (the round-11 ISA defect). Same 2 buffers, same one barrier - // per stage (prologue + 96 = 97/block), same 16 ds_read2_b64 + - // 8 ds_write2_b64 per block-stage, same staging bytes/addresses - // (bit-identical int32 order). - int4 a_payload, b_payload; - // Prologue: issue the stage-1 global loads first (so they ride behind - // the stage-0 inline staging), stage k0 = 0 into buffer 0 inline - // (same load+wait+write path as the accepted kernel), then the first - // barrier. - prefetch_stage_payload( - x_q, weight, tid, m0, n0, m, k, kStageK, a_payload, b_payload); - stage_prefill_tiles( - x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, 0); - __syncthreads(); - for (int k0 = 0; k0 < k; k0 += kStageK) { - const int buf = (k0 / kStageK) & 1; - if (k0 + kStageK < k) { - // Publish the payload carried from the previous iteration (stage - // k0+kStageK) into the idle buffer: the vmcnt wait for the loads - // issued one iteration earlier is emitted here at the first use, - // after the previous MMAC burst + barrier. - publish_stage_payload( - a_payload, b_payload, a_tile + (buf ^ 1) * kABytes, - b_tile + (buf ^ 1) * kBMaxBytes, tid); - } - if (k0 + 2 * kStageK < k) { - // Issue stage k0+2*kStageK's global loads into the payload - // registers (no wait); they stay in flight across the MMAC burst - // and the barrier below and are waited at the next iteration's - // publish (RAW). - prefetch_stage_payload( - x_q, weight, tid, m0, n0, m, k, k0 + 2 * kStageK, a_payload, - b_payload); - } - // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. - // k0-outer / kk-inner int32 accumulation order. - mmac_prefill_stage( - a_frag0, a_frag1, acc00, acc01, acc10, acc11, - a_tile + buf * kABytes, b_tile + buf * kBMaxBytes, - local_row, local_col); - // One barrier per stage: (a) WAR - all stage-s fragment reads of the - // buffer just read are retired before it becomes the publish target - // two stages later; (b) RAW - every thread's stage-(s+1) ds_writes - // are visible before the next iteration reads them. - __syncthreads(); - } - } else { - // Prologue: stage k0 = 0 into buffer 0, then the pipelined loop. - stage_prefill_tiles( - x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, 0); - __syncthreads(); - for (int k0 = 0; k0 < k; k0 += kStageK) { - const int buf = (k0 / kStageK) & 1; - if (k0 + kStageK < k) { - // Prefetch stage s+1 into the idle buffer: global loads issue here, - // before the stage-s MMAC burst, and complete (RAW) behind the - // single barrier below. - stage_prefill_tiles( - x_q, weight, a_tile + (buf ^ 1) * kABytes, - b_tile + (buf ^ 1) * kBMaxBytes, tid, m0, n0, m, n, k, - k0 + kStageK); - } - // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. - // k0-outer / kk-inner int32 accumulation order. - mmac_prefill_stage( - a_frag0, a_frag1, acc00, acc01, acc10, acc11, - a_tile + buf * kABytes, b_tile + buf * kBMaxBytes, - local_row, local_col); - // One barrier per stage: (a) WAR - all stage-s fragment reads of the - // buffer just read are retired before it becomes the prefetch target - // two stages later; (b) RAW - every thread's stage-(s+1) ds_writes are - // visible before the next iteration reads them. - __syncthreads(); - } - } - } else { - for (int k0 = 0; k0 < k; k0 += kStageK) { - // Stage A[kBlockM, kStageK] and B into the single LDS buffer. - stage_prefill_tiles( - x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, k0); - __syncthreads(); - - // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. - // k0-outer / kk-inner int32 accumulation order. - mmac_prefill_stage( - a_frag0, a_frag1, acc00, acc01, acc10, acc11, - a_tile, b_tile, local_row, local_col); - // Retire this stage's reads before the next stage's staging writes - // overwrite the single LDS buffer (WAR hazard across waves). - __syncthreads(); - } + for (int k0 = 0; k0 < k; k0 += kStageK) { + // Stage A[kBlockM, kStageK] and B into the single LDS buffer, then + // publish them to every wavefront. + stage_prefill_tiles( + x_q, weight, a_tile, b_tile, tid, m0, n0, m, n, k, k0); + __syncthreads(); + // Each wave computes its 32x32 quadrant: four m16n16k32 DUMMA tiles. + // k0-outer / kk-inner int32 accumulation order. + mmac_prefill_stage( + a_frag0, a_frag1, acc00, acc01, acc10, acc11, + a_tile, b_tile, local_row, local_col); + // Retire this stage's reads before the next stage's staging writes + // overwrite the single LDS buffer (WAR hazard across waves). + __syncthreads(); } const int base_row = m0 + wave_row * 32; @@ -847,10 +498,10 @@ w8a8_dumma_prefill_tiled_kernel( // Generic scalar fallback: one output element per grid-stride iteration with // exact int32 accumulation. Correct for any (m, n, k), including the small-M -// API shapes (M=2/16) and any unmatched geometry. The weight is the identity -// [K, N] row-major layout for every (k, n) EXCEPT the gate_up weight -// (k, n) == (6144, 512), whose buffer is packed [N, K] n-major; this kernel -// decodes that layout so the paired gate_up tails stay correct. +// API shapes (M=2/16). The weight is the identity [K, N] row-major layout for +// every (k, n) EXCEPT the assigned gate_up pair (k, n) == (6144, 512), whose +// buffer is packed [N, K] n-major; this kernel decodes that layout so the +// paired gate_up M=2/M=16 tails stay correct. __global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( const int8_t* __restrict__ x_q, @@ -861,7 +512,7 @@ w8a8_gemm_scalar_kernel( int m, int n, int k) { - const bool b_nmajor = (k == 6144 && n == 512); + const bool b_nmajor = (k == kGateUpK && n == kGateUpN); const int64_t total = static_cast(m) * n; const int64_t stride = static_cast(gridDim.x) * blockDim.x; for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + @@ -878,14 +529,15 @@ w8a8_gemm_scalar_kernel( acc += static_cast(a_row[kk]) * static_cast(b_base[kk * b_stride]); } - const float scaled = static_cast(acc) * x_scale[row] * - weight_scale[col]; + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; out[idx] = __float2bfloat16(scaled); } } -// Bootstrap pack: identity device-to-device copy for every (k, n) pair -// (raw [K, N] layout preserved). Runs outside the timed/CUDA-Graph region. +// Bootstrap pack: identity device-to-device copy of the raw [K, N] weight and +// the [N] fp32 scale for every (k, n) pair that is not the assigned gate_up +// pair. Runs outside the timed/CUDA-Graph region. __global__ __launch_bounds__(kScalarThreads) void w8a8_pack_identity_kernel( const int8_t* __restrict__ raw_weight, @@ -909,9 +561,8 @@ w8a8_pack_identity_kernel( // index relation: packed[cc * k + kk] == raw[kk * n + cc]. Each thread // gathers the 16 bytes of one aligned int4 k-block from 16 raw rows (stride // n) and stores one aligned int4, keeping the packed write side fully -// coalesced (same gather + aligned int4-store pattern as the validated -// o_proj transpose repair). Runs outside the timed/CUDA-Graph region; same -// byte count, same graph-stable buffer. +// coalesced. Runs outside the timed/CUDA-Graph region; same byte count, same +// graph-stable buffer. __global__ __launch_bounds__(kScalarThreads) void w8a8_pack_gateup_transpose_kernel( const int8_t* __restrict__ raw_weight, @@ -924,7 +575,7 @@ w8a8_pack_gateup_transpose_kernel( const int64_t total4 = static_cast(n) * k4; const int64_t stride4 = static_cast(gridDim.x) * blockDim.x; for (int64_t idx4 = static_cast(blockIdx.x) * blockDim.x + - threadIdx.x; + threadIdx.x; idx4 < total4; idx4 += stride4) { const int cc = static_cast(idx4 / k4); const int kblk = static_cast(idx4 - static_cast(cc) * k4); @@ -936,8 +587,8 @@ w8a8_pack_gateup_transpose_kernel( for (int i = 0; i < static_cast(sizeof(int4)); ++i) { vp[i] = src[static_cast(i) * n]; } - reinterpret_cast(packed_weight + static_cast(cc) * k + - kk)[0] = v; + reinterpret_cast( + packed_weight + static_cast(cc) * k + kk)[0] = v; } const int64_t stride = static_cast(gridDim.x) * blockDim.x; for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + @@ -962,80 +613,43 @@ extern "C" void launch_w8a8_gemm( int k, hipStream_t stream) { // No split-K: the provided workspace is not touched. The timed operator - // performs no allocation, synchronization, packing, or default-stream - // launch; only the caller-provided out is written. + // performs no allocation, synchronization, packing, autotuning, or + // default-stream launch; only the caller-provided out is written. (void)workspace; (void)workspace_bytes; auto* out_bf16 = static_cast(out); - // Native INT8 DUMMA m16n16k32 tiled path for every large-M shape. Both - // assigned shapes qualify: gate_up (4096, 512, 6144) and down_proj - // (4096, 6144, 256) both have M >= 128, N % 128 == 0 and K % 128 == 0. - // Explicit dispatch by launch geometry; unmatched shapes fall through to - // the scalar path (no shape is specialized away). + // Native INT8 DUMMA m16n16k32 tiled path for every shape whose geometry + // matches the staged 64x64 tiles. The assigned shape (4096, 512, 6144) + // takes the packed-weight (n-major B) arm with the round-11 64-K stage + // (kPackedStageK, 10,240 B LDS -> 4 blocks/CU) and the iteration-14 direct + // 8-byte load_frag8 fragment loader (kDirectFrag); any other matching + // geometry keeps the identity [K, N] row_major B arm on the 128-K stage + // with the library loaders. Unmatched geometry and M < 128 fall through to + // the scalar fallback (no shape is specialized away). if (m >= kPrefillMinM && (n % kDefaultBlockN == 0) && (k % kPrefillStageK == 0)) { - // Exact-shape guard for the assigned gate_up shape (round-3 pipeline - // decision, round-10 load_frag8, round-11 kSkew8, round-16 kPayload): - // cooperative A+B staging with the n-major [N, K] packed B operand - // (kBNMajor) consumed as col_major fragments, software-pipelined across - // K tiles (kDoubleBuffer) with a 64-K double buffer: 64x64 tile, 512 - // blocks, 20,480 B LDS/block, 3 blocks/CU (61,440 B) - occupancy - // preserved vs the accepted round-2 single buffer (19,456 B x 3 = - // 58,368 B). One barrier per 64-K stage + prologue (97 for 96 stages). - // Round 11 stages A/B at the 8-mod-16 strides 88/72 (kSkew8) so the - // fragment reads are bank-conflict-free (LDS bytes and occupancy - // unchanged). Round 16 (kPayload) issues the stage-(s+1) global loads a - // full iteration early into a two-int4-per-thread register payload and - // publishes them into the idle LDS buffer at the top of iteration s, so - // the vmcnt wait lands after the stage-s MMAC burst instead of in front - // of it (LDS bytes, barriers, traffic, occupancy and results unchanged). - if (m == 4096 && n == 512 && k == 6144) { - const dim3 grid(n / 64, (m + 63) / 64); - const dim3 block(4 * kWaveSize); + const dim3 grid(n / kDefaultBlockN, + (m + kDefaultBlockM - 1) / kDefaultBlockM); + const dim3 block(kDefaultBlockM * kDefaultBlockN / 1024 * kWaveSize); + if (k == kGateUpK && n == kGateUpN) { hipLaunchKernelGGL( HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel< - 64, 64, 64, true, true, true, true, true>), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - return; - } - if (n % 128 == 0) { - // 64x128 tile, 8 waves (32x32 quadrant per wave). - // down_proj: grid (6144/128, 4096/64) = (48, 64) = 3072 blocks. - const dim3 grid(n / 128, (m + 63) / 64); - const dim3 block(8 * kWaveSize); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 128, 128>), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - return; - } - if (m % 128 == 0) { - // 128x64 tile, 8 waves. - const dim3 grid(n / 64, m / 128); - const dim3 block(8 * kWaveSize); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<128, 64, 128>), + kDefaultBlockM, kDefaultBlockN, kPackedStageK, + true, true>), grid, block, 0, stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); - return; - } - { - // 64x64 tile, 4 waves (generic default). - const dim3 grid(n / 64, (m + 63) / 64); - const dim3 block(4 * kWaveSize); + } else { hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 64, 128>), + HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel< + kDefaultBlockM, kDefaultBlockN, kPrefillStageK, + false>), grid, block, 0, stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); - return; } + return; } - // Generic scalar fallback for every unmatched (m, n, k) and M < 128 - // (including the paired M=2/M=16 API shapes with the same (N, K), which - // decode the packed [N, K] gate_up layout when (k, n) == (6144, 512)). const int64_t total = static_cast(m) * n; const unsigned blocks = static_cast( (total + kScalarThreads - 1) / kScalarThreads); @@ -1053,13 +667,14 @@ extern "C" void launch_pack_w8a8_weight( int k, int n, hipStream_t stream) { - // Assigned gate_up weight (K=6144, N=512): one-time [K,N] -> [N,K] + // Assigned gate_up weight (K=6144, N=512): one-time [K, N] -> [N, K] // transpose pack (n-major, k contiguous) so the timed GEMM reads B // fragments as k-contiguous bytes from the staged LDS tile. Runs outside // the timed/CUDA-Graph region; same byte count, same graph-stable buffer, - // byte-exact vs the raw logical [K,N] weight. - if (k == 6144 && n == 512) { - const int64_t total4 = static_cast(n) * (k / 16); + // byte-exact vs the raw logical [K, N] weight. + if (k == kGateUpK && n == kGateUpN) { + const int64_t total4 = + static_cast(n) * (k / static_cast(sizeof(int4))); const unsigned blocks = static_cast( (total4 + kScalarThreads - 1) / kScalarThreads); hipLaunchKernelGGL( @@ -1069,7 +684,8 @@ extern "C" void launch_pack_w8a8_weight( return; } // Identity device-to-device copy for every other (k, n) pair, including - // unmatched (K, N). Runs outside the timed/CUDA-Graph region. + // unmatched (K, N) and the paired M=2/M=16 API tails with the assigned + // (K, N). Runs outside the timed/CUDA-Graph region. const int64_t weight_elems = static_cast(k) * n; const int64_t total = weight_elems > n ? weight_elems : n; const unsigned blocks = static_cast( 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/TP4/M4096/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip index 94a1fee0..83102b88 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip @@ -1,71 +1,212 @@ -// @@variant shape=hy3_tp4_o_proj_m4096 commit=7c36244e756baf1e8268807cb6481eb118a30133 added=2026-08-23 -// median_us=803.9 p90_us=805.1 tops=85.48 bandwidth_gb_s=62.65 speedup=24.73 baseline_us=1.988e+04 -// source=hy3-tp4-m4096-dsh-test1-bfdf3002 +// @@variant shape=hy3_tp4_o_proj_m4096 commit=a030cfe61386a23a42e172497185fff364c58a85 added=2026-09-14 +// median_us=453.2 p90_us=454.2 speedup=43.87 baseline_us=1.988e+04 +// source=test-9-11-1/iteration_002/hy3_tp4_o_proj_m4096 (retested A/B) // INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. // -// Worker: worker_1 (physical GPU 1), assigned shape: -// hy3_tp4_o_proj_m4096 : M=4096, N=4096, K=2048 +// Worker: worker_0 (physical GPU 0). Assigned shape: +// hy3_tp4_o_proj_m4096 : M=4096, N=4096, K=2048 (logical (K, N) = (2048, 4096)) +// +// Iteration 1 - architecture establishment (one bounded DUMMA geometry round). +// +// * Exact assigned-shape path (m, n, k) == (4096, 4096, 2048): native gfx928 +// INT8 DUMMA m16n16k32 with int32 accumulation and a 128x64 output macro- +// tile per block. +// - Four wavefronts (256 threads); each wave owns a 64x32 quadrant built +// from eight m16n16k32 accumulator fragments (acc00..acc31). +// - Grid is (N/64) x (M/128) = 64 x 32 = 2048 blocks. With /64 KiB LDS +// and this kernel's 30,720 B/block the resource budget allows two +// resident blocks per CU, so the 2048-block grid covers all 120 CUs +// (~8.5 waves of 240 resident blocks). +// - 64-deep K stages, double buffered: the ping-pong LDS pair holds +// A[128,64] and packed-B[64,64] panels (2 x 20 KiB + 2 x 5 KiB +// buffers, 30,720 B total). The loop-carried int4 payload for stage +// k0 + 64 is published into the idle buffer at the top of stage k0 +// (before the MMAC burst) and the global prefetch for stage k0 + 128 +// has a full stage of lead, so global load latency retires under the +// burst and only one __syncthreads per K stage is needed. +// - Both LDS tiles use an 80-byte row stride (five 16-byte bank phases) +// instead of the power-of-two 64-byte stride, which breaks the +// periodicity that made every DUMMA fragment load collide on a handful +// of LDS banks; all cooperative int4 staging stores stay 16-byte +// aligned. +// - The epilogue fuses float(dot) * x_scale[m] * weight_scale[n] and +// stores bf16 directly from the accumulator fragments (no extra LDS +// round trip, no separate epilogue kernel): verified gfx928 int8 +// m16n16k32 ownership is lane&15 -> row, lane>>4 -> col%4, and x[i] -> +// columns col%4 + 4*i. +// - No split-K: the 2048-block MxN grid already exposes ample +// parallelism, so the workspace is never touched and no combine pass +// is paid inside the timed region. +// +// Iteration 2 - B-operand LDS transpose for vectorized fragment loads: +// PMC/ISA audit of the iteration-1 source (11,272,192 LDS instructions, +// 31,457,280 bank conflicts) showed the operand split exactly. The A +// operand's row-major loader emits one 8-byte ds_read2_b32 per m16n16k32 +// fragment (its per-lane bytes are contiguous in K), but the B operand's +// row-major loader walks the 80-byte LDS row stride, so each B fragment +// costs eight scalar ds_read_u8: 32 of the 40 loop LDS reads per 64-K +// stage. This round quantifies operand reuse per 128x64 macro-tile and +// keeps cooperative A+B LDS staging (direct global operand loads would +// re-request ~1 GB of A plus ~0.5 GB of B per replay and lean on L2). The +// focused change is the B staging layout: the tile is now n-major +// ([64 n][64 k] with the same 80-byte bank-phase stride) and is consumed +// with the col_major loader, whose per-lane bytes are eight contiguous K +// bytes. Each B fragment therefore collapses from eight ds_read_u8 to a +// single ds_read2_b32, cutting the loop LDS reads from ~40 to ~12 per +// stage. The global B panel is repacked to [N/64, 64, K] (same byte count) +// so the cooperative int4 stage load stays one coalesced 64-byte segment +// per n row and no transpose is paid inside the timed GEMM. A/B reuse per +// macro-tile is unchanged (each staged A byte feeds 64 output columns, each +// staged B byte 128 output rows), the LDS footprint stays 30,720 B, and two +// blocks/CU remain resident. +// +// Iteration 3 - single-buffered K pipeline (round's mandated single-vs-double +// buffering comparison): +// PMC/ISA of the accepted iteration-2 source (3,932,160 LDS instructions, +// 18,874,368 bank conflicts = 4.80 conflict cycles/LDS, +// lds_wait_instructions 3,680,076 = 93.6% of LDS instructions, 1,089,536 VMEM +// read instructions, 30,720 B LDS, 80 VGPR) shows the kernel is latency +// bound, not LDS-instruction bound: after the iteration-2 vectorization the +// LDS stream is only 12 ds_read2_b32 per 64-K stage, yet almost every one +// still registers a wait, and the only resource that keeps residency at 2 +// blocks/CU is LDS (30,720 * 3 > 65,536) while the 80-VGPR budget already +// allows 3 (61,440 <= 65,536). The round therefore drops the ping-pong pair +// for a single 15,360 B A+B tile and re-publishes the next stage after the +// burst instead of at the top of the stage. This trades one extra barrier per +// K stage and the publish-under-burst overlap for a third resident block +// (12 waves/CU instead of 8). Falsifiable: if the kernel is truly stall-bound, +// the extra block hides enough LDS/VMEM latency to move median below +// 736.46 us; if the exposed LDS store latency and doubled barrier count +// dominate, median regresses and double buffering is retained per the round's +// decision rule. Fragment loads, MMAC schedule, staging mapping, global +// access pattern, packed layout, and the generic fallback are unchanged. +// +// Iteration 4 - square macro-tile (M-tile : N-tile aspect-ratio round): +// PMC/ISA of the accepted iteration-3 source (3,932,160 LDS instructions, +// 18,874,368 bank conflicts = 4.80 conflict cycles/LDS, 8,923,656 +// lds_wait_instructions, 1,089,536 VMEM reads, 15,360 B LDS at 88 arch VGPR, +// 2048 blocks, 594 us of the 655 us median) shows the residual cost is +// per-MAC issue overhead, not LDS capacity: with the 2:1 128x64 macro-tile +// the block still issues 60 LDS wave-instructions (48 fragment reads + 12 +// cooperative stage stores), 12 VMEM loads and 2 barriers for every 524,288 +// MACs, and each 64x32 wave quadrant pays six 8-byte fragment loads for only +// eight m16n16k32 MMACs (0.75 LDS loads/MMAC). This round retunes the +// M-tile : N-tile aspect ratio from 2:1 to 1:1 with the same K depth and the +// same single-buffered pipeline: the macro-tile becomes 128x128 (grid +// 32x32 = 1024 blocks, still 8.5 blocks per each of the 120 CUs) and each of +// the four wavefronts owns a 64x64 quadrant built from sixteen m16n16k32 +// accumulators. A square quadrant is the fragment-shape optimum: eight +// 8-byte fragment loads feed sixteen MMACs (0.5 LDS loads/MMAC, -33%), a +// stage store grows to four int4 per thread so the block pays 4 VMEM loads +// instead of 3 for twice the MACs (-33% VMEM loads/MAC), and the two +// barriers are amortized over 2x the work (-50% barriers/MAC). A 1:1 tile +// also balances the reconvergent operand traffic exactly -- each staged A +// byte feeds 128 output columns and each staged B byte feeds 128 output +// rows, so per-replay L2 reads fall from ~768 MB (A 512 MB + B 256 MB) to +// ~512 MB (A 256 MB + B 256 MB). LDS is 20,480 B (A[128,80] + B[128,80]), +// so two blocks stay resident (40,960 B); VGPR grows by the eight extra int32 +// accumulator fragments (~115 expected, under the 128-VGPR two-block cap). +// Falsifiable: if the residual time is per-MAC instruction/latency overhead +// that a square quadrant and half the barriers can hide, median_us drops +// below the current 654.9467849731445; if the kernel instead needs the 12 +// resident waves of the 128x64 tile (or 16 accumulators push VGPR past 128), +// median rises and the 2:1 tile is retained. K depth, the single-buffer +// publish schedule, the 80-byte five-bank-phase strides, the exact int32 +// accumulation order, the packed layout, the shape guards and the scalar +// fallback are unchanged. +// +// Iteration 5 - packed-B fragment-interleaved swizzle (round's mandated +// packing/swizzle round): +// PMC/ISA of the accepted iteration-4 source (2,621,440 LDS instructions, +// 12,582,912 LDS bank conflicts = 4.80 conflict cycles per LDS instruction, +// 4,869,824 lds_wait_instructions, 827,392 VMEM reads, 20,480 B LDS at 128 +// arch VGPR, 1024 blocks, 500.3 us profiled against the 544.8 us median) +// shows the B operand is already one 8-byte vector LDS read per fragment but +// still not bank-safe: the n-major B tile's 80-byte (20-word) row stride +// makes the 16-row sweep of every col_major B fragment alias the same bank +// pairs, so the conflicts survive the iteration-2 vectorization. This round +// replaces the [N/64, 64, K] panel layout with ONE fragment-interleaved +// swizzle, still a byte-count-preserving permutation of the logical [K, N] +// weight. For each 64-K stage and each 128-column block the packed data is +// sixteen 512-byte fragment blocks ordered [k32][n16], each block stored as +// [k8 0..3][n 0..15][8 B]: +// packed[((k>>6)*(N/128) + (n>>7))*16 + ((k>>5)&1)*8 + ((n>>4)&7)]*512 +// + (((k>>3)&3)*16 + (n&15))*8 + (k&7) +// An explicit col_major 8-byte loader (lane n = lane&15, k8 = lane>>4 reads +// p + k8*128 + n*8) then puts each 16-lane phase on exactly one 128-byte k8 +// block -- all 32 banks exactly once, conflict-free by construction. The +// block's 64-K B stage is one contiguous 8 KiB region, so the cooperative +// stage store becomes two flat 16-byte copies per thread (chunk tid and +// chunk tid+256) instead of two strided n-row vectors, and the global read +// is a fully coalesced 8 KiB stream. The explicit loader's byte placement is +// identical to du_load_matrix_sync, so every v_mmac operand, the +// k0-outer/kk-inner accumulation order and every int32 accumulator are +// bit-identical. A keeps its row-major 80-byte stride and the single- +// buffered 64-K pipeline with its publish schedule is unchanged; LDS/block +// falls from 20,480 B to 18,432 B (10,240 A + 8,192 B), so two blocks still +// stay resident at the 128-VGPR launch bound. Falsifiable: if the residual +// LDS cost is bank-conflict serialization, lds_bank_conflicts fall sharply +// and median_us drops below the current 544.7953414916992; if the conflicts +// were already covered by the two resident blocks and the explicit loader +// only adds address ALU, median rises and the panel layout is retained. +// +// Iteration 6 - fused coalesced epilogue (round's mandated epilogue round): +// PMC/ISA of the accepted iteration-5 source (2,097,152 LDS instructions, +// 6,291,456 LDS bank conflicts, 4,884,093 lds_wait_instructions, 827,392 +// VMEM reads, 262,144 VMEM writes, 18,432 B LDS at 120 arch VGPR, 1024 +// blocks, 435.5 us profiled against the 479.2 us median) shows the residual +// non-GEMM cost is the epilogue store, not the compute loop: the direct +// accumulator store issues four 2-byte scalar global stores per lane per +// fragment, and because the m16n16k32 ownership gives each lane four columns +// strided by 4 (row = lane&15, c4 = lane>>4, x[i] -> c4 + 4*i) a store +// instruction scatters over 16 rows x 8 bytes and touches each 32-byte +// sector at 25% utilization, so the epilogue pays 262,144 wave-level store +// instructions hauling 4x the output sectors it needs. This round fuses the +// per-row/per-column scales and the bf16 conversion into ONE register-only +// transpose store: a 4x4 in-wave transpose (two 2x2 steps, __shfl_xor 16 +// then 32; this DTK lowers it to ds_bpermute at the block tail where the LDS +// pipe is idle) re-routes the four int32 values so lane (r, c4) ends up +// holding the four CONTIGUOUS columns 4*c4..4*c4+3, which are then scaled, +// converted to bf16 and written as ONE 8-byte store per lane. Every 32-byte +// sector is 100% utilized and the epilogue issues 4x fewer store +// instructions (16 per wave instead of 64); the per-column weight scale also +// collapses from four scalar loads to one float4 load per fragment. Only the +// int32 values are re-routed between lanes, and the per-element multiply +// order float(dot) * x_scale[row] * weight_scale[col] plus +// __float2bfloat16 are unchanged, so every stored bf16 bit is identical. The +// kernel already performs no split-K and never touches the workspace (the +// timed operator is a single compute dispatch), so no workspace/combine pass +// exists to remove; LDS stays 18,432 B, and the single-buffered 64-K +// pipeline, the fragment loaders, the exact int32 accumulation order, the +// shape guards and the scalar fallback are unchanged. Falsifiable: if the +// residual time is store-sector over-fetch and the shuffle ALU is hidden, +// vmem_write instructions fall ~4x and median_us drops below the current +// 479.2323112487793; if the transpose ds_bpermute/v_cndmask chain contends +// with the LDS pipe or the extra live scale registers spill, median rises +// and the direct store is retained. +// +// * launch_pack_w8a8_weight packs exactly (k, n) == (2048, 4096) into the +// fragment-interleaved swizzle above (byte-count preserving, same +// allocation and same graph-stable captured addresses). Every other +// (k, n) keeps the identity device-to-device copy, and +// packed_weight_scale is always the identity copy of weight_scale. Packing +// happens before Graph capture and never inside the timed GEMM. // -// Iteration 1 - DUMMA throughput baseline with a 2-D macro-tile (128x64). -// * Large-prefill path (exact (m, n, k) == (4096, 4096, 2048)): -// native INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 -// output tile per block; four wavefronts (256 threads); each wave owns -// a 64x32 quadrant built from eight m16n16k32 fragments; block -// cooperatively vector-loads A[128,64] and packed B panels into 26 KiB -// of LDS with a padded 80-byte B row stride; 64-deep K double buffering -// with a register prefetch of the next stage (one loop-carried barrier -// per 64-K stage); fused dot * x_scale[m] * weight_scale[n] epilogue -// stored directly as bf16 from the accumulator fragments. -// Rationale vs the bootstrap 64x64x128 single-buffer tile (verified on -// the exact shape in the trusted ledger): 128x64 doubles B reuse across -// rows, halves the grid to 2048 blocks, stages each B stage as one -// coalesced 4 KiB stream from pre-packed [N/64, K, 64] panels, breaks -// the 64-byte LDS bank periodicity with an 80-byte stride, and hides -// global->LDS latency with double buffering. -// * launch_pack_w8a8_weight packs exactly (k, n) == (2048, 4096) into -// contiguous [N/64, K, 64] int8 panels; every other (k, n) keeps the -// identity copy. Packing never happens inside the timed GEMM. // * Generic scalar int8/int32 fallback for every unmatched (m, n, k), -// including the paired M=2 decode shape with the same (K, N); the -// fallback reads the packed panel layout whenever (k, n) == (2048, 4096) -// and the identity layout otherwise. +// including all M < 128 API cases and the paired small-M shape with the +// same (N, K) == (4096, 2048). The fallback decodes the fragment- +// interleaved swizzle for that exact (k, n) pair and the identity layout +// otherwise, so it is correct for every packed/unpacked weight. // -// Iteration 3 - A-tile LDS bank-conflict fix: -// PMC on the accepted iteration-1 source: 77,594,624 LDS bank conflicts -// (~6.88 per LDS instruction). Root cause in the fragment loads: the A -// tile still uses an unpadded 64-byte row stride (16 banks/row), so the -// bank pattern repeats every two rows and an m16n16k32 A-fragment load -// (64 lanes x 8 B over 16 rows) puts eight lanes on the same bank pair -// (8-way conflicts), while the already padded 80-byte B rows (20 -// banks/row, period 8) only reach 4-way over their 32-row span - matching -// the PMC mix (8 A-loads x 8-way + 4 B-loads x 4-way per 64-K stage). -// Focused change: pad the A LDS rows 64 -> 80 bytes (five bank phases, -// same as B) so each A-fragment load drops from 8-way to 2-way bank -// conflicts. LDS grows 26,624 -> 30,720 B and stays at 2 blocks/CU -// (61,440 B of the 64 KiB LDS budget); 80 is a multiple of 16 so every -// cooperative int4 stage store stays aligned; the packed [N/64, K, 64] -// global layout, exact int32 accumulation, grid, and fallback are -// untouched. +// 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 the caller-provided +// PyTorch HIP stream, so it is CUDA-Graph safe. This round does not use the +// workspace at all. // -// Iteration 10 - stage-boundary restructure (early publish, one full stage of -// write->read distance): -// Rounds 7 (in-loop A-fragment prefetch, 845.0 us) and 8 (occupancy probe, -// 866.6 us) falsified the in-kk LDS latency and wave-count theories and -// indicted the structural stage-start chain: [last MMACs] -> [s_waitcnt -// vmcnt for the prefetch loads] -> [3x ds_write_b128] -> [s_waitcnt -// lgkmcnt(0)] -> [s_barrier] -> [stage-start fragment loads] (ISA 0x3BC0- -// 0x3BF4 on the accepted source). The next-stage LDS publish currently -// happens AFTER the MMAC burst, so the store issue+retire latency and the -// fresh-write visibility sit directly between the barrier and the next -// stage's first reads. Focused change: the loop-carried int4 payload for -// stage k0+kStageK is published into the idle ping-pong buffer at the TOP -// of stage k0 (before the burst) and the global prefetch for stage -// k0+2*kStageK now has a full extra stage of lead, so the store latency -// retires under the burst and the stage-end barrier publishes data that has -// been resident in LDS for a full stage before the next stage reads it. -// One barrier per stage, 30,720 B LDS (2 blocks/CU), grid, packed layout, -// exact int32 accumulation order, shape guards, and the scalar fallback are -// unchanged. +// Include 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 before the HIP +// runtime headers; du_mma_common.h is not an installed public header). #include #include @@ -79,29 +220,65 @@ using namespace du::dumma; constexpr int kTileM = 16; constexpr int kTileN = 16; -constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 +constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 constexpr int kWaveSize = 64; // gfx928 native wavefront + constexpr int kTargetM = 4096; constexpr int kTargetN = 4096; constexpr int kTargetK = 2048; -constexpr int kBlockM = 128; -constexpr int kBlockN = 64; -constexpr int kStageK = 64; -constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) -constexpr int kBStride = kBlockN + kBPad; -constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride + +constexpr int kBlockM = 128; // output tile rows per block +constexpr int kBlockN = 128; // output tile columns per block +constexpr int kStageK = 64; // K depth staged per barrier +constexpr int kBPad = 16; // 64 -> 80-byte A LDS row stride (five bank phases) +constexpr int kAStride = kStageK + kBPad; constexpr int kBlockThreads = 4 * kWaveSize; +// Iteration 5 - fragment-interleaved packed-B swizzle. Each 32-k x 16-n DUMMA +// B fragment block is 512 B stored as [k8 0..3][n 0..15][8 B]; a 64-K stage of +// a 128-column block is (kStageK/kTileK) * (kBlockN/kTileN) = 16 such blocks, +// i.e. 8,192 B held in exactly the packed global order. The cooperative stage +// store is therefore two flat 16-byte copies per thread (chunk tid and chunk +// tid + kBlockThreads) and the col_major fragment load is one 8-byte read per +// lane at p + k8*128 + n*8. +constexpr int kSwzBlock = 512; // bytes per fragment block +constexpr int kKGroup = 8; // int8 k bytes owned by one col_major lane +constexpr int kBSwzBlocks = (kStageK / kTileK) * (kBlockN / kTileN); +constexpr int kBSwzBytes = kBSwzBlocks * kSwzBlock; +constexpr int kBChunks = kBSwzBytes / static_cast(sizeof(int4)); +static_assert(kBChunks == 2 * kBlockThreads, + "each thread stages exactly two int4 chunks of packed B"); +static_assert(kBSwzBytes == kBlockN * kStageK, + "the swizzle is byte-count preserving for the B stage"); + using bf16_t = hip_bfloat16; // --------------------------------------------------------------------------- -// Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. -// Each wave owns a 64x32 quadrant (eight m16n16k32 accumulators); the block -// cooperatively stages A[128,64] and packed-B [64,64] panels in two ping-pong -// buffers of 64-K stages (26 KiB total LDS). One barrier per stage. +// Iteration 6 - fused, fully coalesced epilogue. The per-row x_scale, the +// per-column weight_scale, the bf16 conversion and the store all happen here +// inside the compute kernel (no separate epilogue/combine pass and no +// workspace traffic). +// +// Verified gfx928 int8 m16n16k32 accumulator ownership (cross-checked against +// du_store_matrix_sync): lane & 15 selects the fragment row, lane >> 4 selects +// the column group c4, and x[i] maps to columns c4 + 4*i. A direct store is +// therefore four 2-byte scalar stores per lane whose wavefront addresses touch +// each 32-byte sector at 25% utilization. This helper transposes the 4-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; this DTK 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, and +// writes them as ONE 8-byte store (100% sector efficiency). The per-column +// weight scale collapses from four scalar loads to one float4 load. Only the +// int32 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 to the direct +// per-element store. The row >= m guard is wavefront-uniform (base_row is a +// multiple of 16 and m is a multiple of 16 on every launched path), so the +// shuffles never mix active and inactive lanes. // --------------------------------------------------------------------------- template -__device__ __forceinline__ void store_prefill_fragment( +__device__ __forceinline__ void store_prefill_fragment_coalesced( const AccFragment& frag, const float* __restrict__ x_scale, const float* __restrict__ weight_scale, @@ -111,25 +288,103 @@ __device__ __forceinline__ void store_prefill_fragment( 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); + return; // tail-M masking: padded rows never write (wavefront-uniform) } + const int c4 = lane >> 4; // 0..3 + const int x0 = frag.x[0]; + const int x1 = frag.x[1]; + const int x2 = frag.x[2]; + const int x3 = frag.x[3]; + // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). + const int s0 = __shfl_xor(x0, 16, kWaveSize); + const int s1 = __shfl_xor(x1, 16, kWaveSize); + const int s2 = __shfl_xor(x2, 16, kWaveSize); + const int s3 = __shfl_xor(x3, 16, kWaveSize); + const bool lo = (c4 & 1) == 0; + const int a0 = lo ? x0 : s1; + const int a1 = lo ? s0 : x1; + const int a2 = lo ? x2 : s3; + const int a3 = lo ? s2 : x3; + // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). + const int t0 = __shfl_xor(a0, 32, kWaveSize); + const int t1 = __shfl_xor(a1, 32, kWaveSize); + const int t2 = __shfl_xor(a2, 32, kWaveSize); + const int t3 = __shfl_xor(a3, 32, kWaveSize); + const bool hi = (c4 & 2) == 0; + const int f0 = hi ? a0 : t2; + const int f1 = hi ? a1 : t3; + const int f2 = hi ? t0 : a2; + const int f3 = hi ? t1 : a3; + + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte aligned: + // base_col is a multiple of 16, n*2 is a multiple of 8). + const int col0 = base_col + 4 * c4; + const float xs = x_scale[row]; + const float4 ws = *reinterpret_cast(weight_scale + col0); + const float v0 = static_cast(f0) * xs * ws.x; + const float v1 = static_cast(f1) * xs * ws.y; + const float v2 = static_cast(f2) * xs * ws.z; + const float v3 = static_cast(f3) * xs * ws.w; + const uint64_t packed = + static_cast( + static_cast(__float2bfloat16(v0))) | + (static_cast( + static_cast(__float2bfloat16(v1))) + << 16) | + (static_cast( + static_cast(__float2bfloat16(v2))) + << 32) | + (static_cast( + static_cast(__float2bfloat16(v3))) + << 48); + *reinterpret_cast(out + static_cast(row) * n + col0) = + packed; } -__global__ __launch_bounds__(kBlockThreads) void -w8a8_dumma_128x64x64_packed_kernel( +// --------------------------------------------------------------------------- +// Iteration 5 - explicit col_major B fragment loader on the fragment- +// interleaved swizzled tile. The gfx928 int8 m16n16k32 col_major matrix_b +// mapping is n = lane & 15, k8 = lane >> 4 (each lane owns eight contiguous k +// bytes), so on a [k8 0..3][n 0..15][8 B] block the lane's data sits at +// p + k8*128 + n*8. One 8-byte vector LDS read per lane, and each 16-lane +// phase spans exactly one 128-byte k8 block -- all 32 banks exactly once, so +// the fragment load is conflict-free by construction. The byte placement is +// identical to du_load_matrix_sync, so the +// v_mmac_i32_16x16x32_i8 operands and the int32 accumulation are bit-identical +// to the previous n-major panel layout. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void load_b_frag8_swizzled( + DUFragment& frag, + 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(&frag.x[0]) = v; +} + +// --------------------------------------------------------------------------- +// Exact-shape prefill: 128x128 output tile per block, four wavefronts of 64 +// lanes. Each wave owns a 64x64 quadrant (sixteen m16n16k32 accumulators); the +// block cooperatively stages A[128,64] (row-major, k-major per lane) and the +// fragment-interleaved packed-B 64-K stage (16 x 512 B = 8,192 B, [k32][n16] +// block order with each block stored [k8][n][8 B]) in ONE 64-K buffer (18,432 B +// total) with two barriers per stage (iteration-3 pipeline). The swizzled B +// tile plus the explicit col_major 8-byte loader makes every B fragment one +// conflict-free ds_read_b64 (each 16-lane phase spans exactly one 128-byte k8 +// block), and the square 64x64 quadrant pays eight fragment loads per sixteen +// MMACs instead of six per eight. +// --------------------------------------------------------------------------- +// The 2 in the launch bound asks the compiler to keep two blocks (eight waves) +// resident per CU, which caps the VGPR budget at 128 per thread: with sixteen +// int32 accumulator fragments (64 VGPRs) plus the four-int4 pipeline payload +// the natural allocation lands under that cap, and the bound prevents a +// compiler choice that would silently drop residency to one block (four +// waves) per CU. +__global__ __launch_bounds__(kBlockThreads, 2) void +w8a8_dumma_128x128x64_packed_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, const float* __restrict__ x_scale, @@ -145,164 +400,237 @@ w8a8_dumma_128x64x64_packed_kernel( const int wave_col = wave & 1; const int m0 = static_cast(blockIdx.y) * kBlockM; const int n0 = static_cast(blockIdx.x) * kBlockN; - // Packed layout: panel p = B[:, p*64:(p+1)*64] stored as [K][64] rows. - const int b_panel_offset = - static_cast(blockIdx.x) * kTargetK * kBlockN; + // Iteration-5 packed layout: for each 64-K stage and each 128-column block the + // weight is stored as 16 fragment blocks of 512 B in [k32][n16] order, so this + // block's whole 64-K B stage is ONE contiguous kBSwzBytes region at + // ((k0/64) * (N/128) + blockIdx.x) * kBSwzBytes. The staging is therefore a + // flat 16-byte copy per chunk (no transpose, no padding). + const int b_colblocks = kTargetN / kBlockN; const int local_row = wave_row * 64; - const int local_col = wave_col * 32; + const int local_col = wave_col * 64; - __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; - __shared__ __align__(16) int8_t b_tile[2][kStageK * kBStride]; + // Iteration 5 - square macro-tile with the swizzled B stage. A[128,64] rows + // keep the 80-byte five-bank-phase stride (128 * 80 = 10,240 B) and the + // packed-B 64-K stage is a dense 8,192 B swizzled region, so the block uses + // 18,432 B. Two blocks stay resident (36,864 B of the 64 KiB LDS budget) and + // the 16 accumulator fragments land under the 128-VGPR two-block cap. + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBSwzBytes]; DUFragment - a_frag0, a_frag1; - DUFragment - b_frag0, b_frag1; + a_frag0, a_frag1, a_frag2, a_frag3; + DUFragment + b_frag0, b_frag1, b_frag2, b_frag3; DUFragment - acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; + acc00, acc01, acc02, acc03, + acc10, acc11, acc12, acc13, + acc20, acc21, acc22, acc23, + acc30, acc31, acc32, acc33; du_fill_fragment(acc00, 0); du_fill_fragment(acc01, 0); + du_fill_fragment(acc02, 0); + du_fill_fragment(acc03, 0); du_fill_fragment(acc10, 0); du_fill_fragment(acc11, 0); + du_fill_fragment(acc12, 0); + du_fill_fragment(acc13, 0); du_fill_fragment(acc20, 0); du_fill_fragment(acc21, 0); + du_fill_fragment(acc22, 0); + du_fill_fragment(acc23, 0); du_fill_fragment(acc30, 0); du_fill_fragment(acc31, 0); + du_fill_fragment(acc32, 0); + du_fill_fragment(acc33, 0); - // Prologue: cooperative vectorized staging of stage 0. Each thread owns one - // int4 in A rows [0,64) and one in A rows [64,128), plus one int4 of packed - // B. A[128,64] is 512 int4s; B[64,64] is 256 int4s. - const int vector_byte_offset = tid * sizeof(int4); + // Prologue: cooperative vectorized staging of stage 0 straight into the one + // LDS tile. Each thread owns one int4 in A rows [0,64), one in A rows + // [64,128) (tid/4 is the row, (tid%4)*16 the K chunk) and two int4 of the + // packed-B stage (chunk tid and chunk tid + 256). A's tiles keep their + // 80-byte row stride; the packed-B stage is contiguous in global memory and + // in LDS, so both of its 16-byte copies are flat and fully coalesced. + constexpr int kInt4 = static_cast(sizeof(int4)); + const int vector_byte_offset = tid * kInt4; const int stage_row = vector_byte_offset / kStageK; const int stage_col = vector_byte_offset - stage_row * kStageK; - *reinterpret_cast(a_tile[0] + stage_row * kAStride + stage_col) = + const int b_stage_base0 = static_cast(blockIdx.x) * kBSwzBytes; + *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = *reinterpret_cast( x_q + static_cast(m0 + stage_row) * k + stage_col); - *reinterpret_cast(a_tile[0] + + *reinterpret_cast(a_tile + (stage_row + kBlockM / 2) * kAStride + stage_col) = *reinterpret_cast( x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + stage_col); - *reinterpret_cast(b_tile[0] + stage_row * kBStride + stage_col) = - *reinterpret_cast( - weight + b_panel_offset + stage_row * kBlockN + stage_col); + *reinterpret_cast(b_tile + tid * kInt4) = + *reinterpret_cast(weight + b_stage_base0 + tid * kInt4); + *reinterpret_cast(b_tile + (tid + kBlockThreads) * kInt4) = + *reinterpret_cast(weight + b_stage_base0 + + (tid + kBlockThreads) * kInt4); __syncthreads(); - // Loop-carried register payload for the stage after the one being computed - // (stage k0 + kStageK): loaded one full stage ahead and published into the - // idle ping-pong buffer at the TOP of the following stage. The publish's - // LDS store latency then retires under the MMAC burst, and the stage-end - // barrier publishes it a full stage before the next stage's first fragment - // loads - so the structural LDS write -> barrier -> read chain leaves the - // stage boundary (rounds 7/8 falsified in-loop A-fragment prefetch and - // occupancy; this is the prescribed stage-boundary restructure). + // Loop-carried register payload for stage kStageK, the first stage that has + // to be republished into the single buffer. It is consumed by the stage-0 + // burst first and only published after that burst retires, so it keeps a + // one-full-stage lead from load to LDS store, exactly as before. int4 cur_a0{}; int4 cur_a1{}; - int4 cur_b{}; + int4 cur_b0{}; + int4 cur_b1{}; if (k > kStageK) { cur_a0 = *reinterpret_cast( x_q + static_cast(m0 + stage_row) * k + kStageK + stage_col); cur_a1 = *reinterpret_cast( x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + kStageK + stage_col); - cur_b = *reinterpret_cast( - weight + b_panel_offset + (kStageK + stage_row) * kBlockN + stage_col); + cur_b0 = *reinterpret_cast( + weight + (b_colblocks + static_cast(blockIdx.x)) * kBSwzBytes + + tid * kInt4); + cur_b1 = *reinterpret_cast( + weight + (b_colblocks + static_cast(blockIdx.x)) * kBSwzBytes + + (tid + kBlockThreads) * kInt4); } - int current = 0; for (int k0 = 0; k0 < k; k0 += kStageK) { - const int next_k = k0 + kStageK; - const int next = current ^ 1; - // Publish the loop-carried payload (stage next_k) into the idle buffer - // BEFORE the MMAC burst; the stage-end barrier below publishes it a full - // stage before it is read, hiding the store latency under the burst. - if (next_k < k) { - *reinterpret_cast(a_tile[next] + stage_row * kAStride + - stage_col) = cur_a0; - *reinterpret_cast(a_tile[next] + - (stage_row + kBlockM / 2) * kAStride + - stage_col) = cur_a1; - *reinterpret_cast( - b_tile[next] + stage_row * kBStride + stage_col) = cur_b; - } - // Prefetch stage k0 + 2*kStageK into the payload registers (a full stage - // of extra global->LDS lead) so the publish never waits on the loads. - const int load_k = k0 + 2 * kStageK; - if (load_k < k) { - cur_a0 = *reinterpret_cast( - x_q + static_cast(m0 + stage_row) * k + load_k + stage_col); - cur_a1 = *reinterpret_cast( - x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + - load_k + stage_col); - cur_b = *reinterpret_cast( - weight + b_panel_offset + (load_k + stage_row) * kBlockN + - stage_col); - } - - // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. + // Each wave consumes its 64x64 quadrant from the single LDS tile: sixteen + // m16n16k32 MMACs per kk from four A and four B fragments (0.5 LDS loads + // per MMAC instead of the 0.75 of the 128x64 tile). #pragma unroll for (int kk = 0; kk < kStageK; kk += kTileK) { + // Swizzled B block index of this wave's four 16-n fragments: k32 half + // (kk/32) selects the [k32] outer index, and this wave's 64-column + // quadrant starts at n16 group wave_col*4. + const int b_swz = ((kk >> 5) << 3) + (wave_col << 2); + load_b_frag8_swizzled(b_frag0, b_tile + (b_swz + 0) * kSwzBlock, lane); + load_b_frag8_swizzled(b_frag1, b_tile + (b_swz + 1) * kSwzBlock, lane); du_load_matrix_sync( - b_frag0, b_tile[current] + kk * kBStride + local_col, kBStride); - du_load_matrix_sync( - b_frag1, b_tile[current] + kk * kBStride + local_col + kTileN, - kBStride); - du_load_matrix_sync( - a_frag0, a_tile[current] + local_row * kAStride + kk, kAStride); + a_frag0, a_tile + (local_row + 0 * kTileM) * kAStride + kk, + kAStride); du_load_matrix_sync( - a_frag1, a_tile[current] + (local_row + kTileM) * kAStride + kk, + a_frag1, a_tile + (local_row + 1 * kTileM) * kAStride + kk, kAStride); du_mma_sync(acc00, a_frag0, b_frag0, acc00); du_mma_sync(acc01, a_frag0, b_frag1, acc01); du_mma_sync(acc10, a_frag1, b_frag0, acc10); du_mma_sync(acc11, a_frag1, b_frag1, acc11); + load_b_frag8_swizzled(b_frag2, b_tile + (b_swz + 2) * kSwzBlock, lane); + load_b_frag8_swizzled(b_frag3, b_tile + (b_swz + 3) * kSwzBlock, lane); du_load_matrix_sync( - a_frag0, a_tile[current] + (local_row + 2 * kTileM) * kAStride + kk, + a_frag2, a_tile + (local_row + 2 * kTileM) * kAStride + kk, kAStride); du_load_matrix_sync( - a_frag1, a_tile[current] + (local_row + 3 * kTileM) * kAStride + kk, + a_frag3, a_tile + (local_row + 3 * kTileM) * kAStride + kk, kAStride); - du_mma_sync(acc20, a_frag0, b_frag0, acc20); - du_mma_sync(acc21, a_frag0, b_frag1, acc21); - du_mma_sync(acc30, a_frag1, b_frag0, acc30); - du_mma_sync(acc31, a_frag1, b_frag1, acc31); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + du_mma_sync(acc22, a_frag2, b_frag2, acc22); + du_mma_sync(acc23, a_frag2, b_frag3, acc23); + du_mma_sync(acc32, a_frag3, b_frag2, acc32); + du_mma_sync(acc33, a_frag3, b_frag3, acc33); + du_mma_sync(acc20, a_frag2, b_frag0, acc20); + du_mma_sync(acc21, a_frag2, b_frag1, acc21); + du_mma_sync(acc30, a_frag3, b_frag0, acc30); + du_mma_sync(acc31, a_frag3, b_frag1, acc31); } + // Single-buffer hazard: the tile must not be overwritten while any wave is + // still reading it. This is the extra barrier the single-buffered pipeline + // pays per K stage (the ping-pong version published into the idle buffer + // at the top of the stage and needed only one barrier). + __syncthreads(); + const int next_k = k0 + kStageK; if (next_k < k) { - // The stage-end barrier publishes the payload written at the top of - // this stage (a full stage before the next stage reads it) and protects - // buffer `current` from the publish at the top of the next stage. + // Publish stage next_k into the now-idle single tile. The store latency + // is no longer hidden under a burst, so the second barrier below is + // required before the next stage's first fragment loads. + *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = + cur_a0; + *reinterpret_cast(a_tile + + (stage_row + kBlockM / 2) * kAStride + + stage_col) = cur_a1; + *reinterpret_cast(b_tile + tid * kInt4) = cur_b0; + *reinterpret_cast(b_tile + (tid + kBlockThreads) * kInt4) = + cur_b1; __syncthreads(); - current = next; + // Refill the register payload one stage further out; the lead from this + // load to its publish at the end of the next iteration is still one full + // stage, so the global->register latency is covered as before. + const int load_k = k0 + 2 * kStageK; + if (load_k < k) { + cur_a0 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row) * k + load_k + + stage_col); + cur_a1 = *reinterpret_cast( + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + + load_k + stage_col); + const int b_load_base = + ((load_k >> 6) * b_colblocks + static_cast(blockIdx.x)) * + kBSwzBytes; + cur_b0 = *reinterpret_cast(weight + b_load_base + + tid * kInt4); + cur_b1 = *reinterpret_cast( + weight + b_load_base + (tid + kBlockThreads) * kInt4); + } } } 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_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, + base_row, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, + base_row, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced(acc02, x_scale, weight_scale, out, + base_row, base_col + 2 * kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc03, x_scale, weight_scale, out, + base_row, base_col + 3 * kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, + lane); + store_prefill_fragment_coalesced(acc12, x_scale, weight_scale, out, + base_row + kTileM, base_col + 2 * kTileN, m, + n, lane); + store_prefill_fragment_coalesced(acc13, x_scale, weight_scale, out, + base_row + kTileM, base_col + 3 * kTileN, m, + n, lane); + store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, m, + n, lane); + store_prefill_fragment_coalesced(acc22, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + 2 * kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc23, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + 3 * kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, + lane); + store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, m, + n, lane); + store_prefill_fragment_coalesced(acc32, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + 2 * kTileN, + m, n, lane); + store_prefill_fragment_coalesced(acc33, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + 3 * kTileN, + m, n, lane); } // --------------------------------------------------------------------------- -// Generic scalar fallback: one thread per output element. Exact int32 dot -// over K, then fused x_scale * weight_scale, then bf16 store. Handles every -// unmatched (m, n, k), including all small-M API cases. When (k, n) matches -// the packed pair (2048, 4096) the weight is read from the [N/64, K, 64] -// panel layout; otherwise from the identity layout. +// 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 M < 128. When (k, n) matches the packed +// pair (2048, 4096) the weight is read from the iteration-5 fragment- +// interleaved swizzle; otherwise from the identity layout. // --------------------------------------------------------------------------- __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( const int8_t* __restrict__ x_q, @@ -323,15 +651,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; if (n == kTargetN && k == kTargetK) { - const int panel = col / kBlockN; - const int panel_col = col - panel * kBlockN; - const int64_t panel_offset = - static_cast(panel) * kTargetK * kBlockN; + // Decode the same swizzle the packer writes (ascending kk, exact order). + const int colblock = col / kBlockN; + const int n16 = (col % kBlockN) / kTileN; + const int r = col % kTileN; for (int kk = 0; kk < k; ++kk) { + const int ks = kk % kStageK; + const int block = ((kk / kStageK) * (kTargetN / kBlockN) + colblock) * + kBSwzBlocks + + (ks / kTileK) * (kBlockN / kTileN) + n16; + const int64_t idx = + static_cast(block) * kSwzBlock + + static_cast((ks % kTileK) / kKGroup) * (kTileN * kKGroup) + + static_cast(r) * kKGroup + (kk % kKGroup); acc += static_cast(x_q[static_cast(row) * k + kk]) * - static_cast(weight[panel_offset + - static_cast(kk) * kBlockN + - panel_col]); + static_cast(weight[idx]); } } else { const int8_t* a_row = x_q + static_cast(row) * k; @@ -341,36 +675,55 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( static_cast(b_col[static_cast(kk) * n]); } } - const float scaled = static_cast(acc) * - x_scale[row] * weight_scale[col]; + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; out[linear] = __float2bfloat16(scaled); } // --------------------------------------------------------------------------- -// Weight packing: exact (k, n) == (2048, 4096) becomes contiguous -// [N/64, K, 64] int8 panels (each int4 stays within one 64-column row, so -// loads are coalesced 4 KiB streams per 64-K stage); identity copy for any -// other (k, n). Scales are always copied identity. +// Weight packing: the exact (k, n) == (2048, 4096) weight is permuted once, +// outside the timed region and outside Graph capture, into the iteration-5 +// fragment-interleaved swizzle. For each 64-K stage and each 128-column block +// there are kBSwzBlocks = 16 fragment blocks of 512 B ordered [k32][n16], each +// stored as [k8 0..3][n 0..15][8 B]. Every packed 16-byte chunk is therefore a +// flat 16-byte store of two 8-byte k-runs (n rows 2*n_pair and 2*n_pair+1 of +// one k8 group), which is exactly the chunk the cooperative stage store later +// copies into LDS. The byte count k*n, the allocation and the graph-stable +// captured addresses are unchanged; every other (k, n) keeps the identity +// device-to-device copy, and the scales are always copied identity. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(256) void w8a8_pack_64n_panel_kernel( +__global__ __launch_bounds__(256) void w8a8_pack_bswz_i8_kernel( const int8_t* __restrict__ raw_weight, int8_t* __restrict__ packed_weight) { - const int vec = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - constexpr int kVectors = + const int chunk = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + constexpr int kChunks = kTargetK * kTargetN / static_cast(sizeof(int4)); - if (vec >= kVectors) { + if (chunk >= kChunks) { return; } - constexpr int kPanelBytes = kTargetK * kBlockN; - const int byte_offset = vec * sizeof(int4); - const int panel = byte_offset / kPanelBytes; - const int panel_offset = byte_offset - panel * kPanelBytes; - const int kk = panel_offset / kBlockN; - const int col = panel_offset - kk * kBlockN; - reinterpret_cast(packed_weight)[vec] = - *reinterpret_cast( - raw_weight + static_cast(kk) * kTargetN + - panel * kBlockN + col); + constexpr int kChunksPerBlock = kSwzBlock / static_cast(sizeof(int4)); + constexpr int kPairsPerK8 = kTileN / 2; // 16-byte chunks per k8 group + const int within = chunk % kChunksPerBlock; + const int block = chunk / kChunksPerBlock; + const int kk8 = within / kPairsPerK8; + const int n_pair = within % kPairsPerK8; + const int n16 = block % (kBlockN / kTileN); + const int k32 = (block / (kBlockN / kTileN)) % (kStageK / kTileK); + const int colblock = (block / kBSwzBlocks) % (kTargetN / kBlockN); + const int stage = block / (kBSwzBlocks * (kTargetN / kBlockN)); + const int kk_base = stage * kStageK + k32 * kTileK + kk8 * kKGroup; + const int col0 = colblock * kBlockN + n16 * kTileN + n_pair * 2; + // Transposed gather: one 16-byte chunk = two 8-byte k-runs of adjacent n. + int4 value; + int8_t* bytes = reinterpret_cast(&value); +#pragma unroll + for (int i = 0; i < kKGroup; ++i) { + bytes[i] = raw_weight[static_cast(kk_base + i) * kTargetN + col0]; + bytes[kKGroup + i] = + raw_weight[static_cast(kk_base + i) * kTargetN + col0 + 1]; + } + *reinterpret_cast(packed_weight + + chunk * static_cast(sizeof(int4))) = value; } __global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( @@ -388,7 +741,7 @@ __global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( const float* __restrict__ raw, float* __restrict__ packed, int count) { - const int linear = blockIdx.x * blockDim.x + threadIdx.x; + const int linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (linear < count) { packed[linear] = raw[linear]; } @@ -397,10 +750,10 @@ __global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( } // namespace // --------------------------------------------------------------------------- -// Stable host launch symbols consumed by csrc/bindings.cpp. -// Both launchers are pure dispatch: no allocation, no packing, no -// synchronization, no default-stream launch; they run on the caller-provided -// PyTorch stream and are CUDA-Graph safe. +// Stable host launch symbols consumed by csrc/bindings.cpp. Both launchers are +// pure dispatch: no allocation, no packing, no synchronization, no +// default-stream launch; they run on the caller-provided PyTorch stream and +// are CUDA-Graph safe. // --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, @@ -414,28 +767,28 @@ 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; // no split-K: the GEMM never touches the workspace (void)workspace_bytes; auto* out_bf16 = reinterpret_cast(out); // Explicit dispatch. The assigned shape (M=4096, N=4096, K=2048) takes the - // packed-panel DUMMA 128x64 path; every other (m, n, k) - including - // small-M API cases and the paired M=2 shape with the same (K, N) - takes - // the scalar fallback, which understands the packed panel layout. + // swizzled-packed DUMMA 128x128 path; every other (m, n, k) - including every + // M < 128 API case and the paired small-M shape with the same (N, K) - takes + // the scalar fallback, which understands the fragment-interleaved swizzle. if (m == kTargetM && n == kTargetN && k == kTargetK) { - const dim3 grid(kTargetN / kBlockN, kTargetM / kBlockM); + const dim3 grid(static_cast(kTargetN / kBlockN), + static_cast(kTargetM / kBlockM)); const dim3 block(kBlockThreads); - hipLaunchKernelGGL(w8a8_dumma_128x64x64_packed_kernel, - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); + hipLaunchKernelGGL(w8a8_dumma_128x128x64_packed_kernel, grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); } else { constexpr int kBlock = 256; const int64_t total = static_cast(m) * n; const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); const dim3 block(kBlock); - hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, m, n, + k); } (void)hipGetLastError(); } @@ -455,26 +808,25 @@ extern "C" void launch_pack_w8a8_weight( const dim3 weight_grid( static_cast((kPackVectors + kBlock - 1) / kBlock)); const dim3 block(kBlock); - hipLaunchKernelGGL(w8a8_pack_64n_panel_kernel, - weight_grid, block, 0, stream, - raw_weight, packed_weight); + hipLaunchKernelGGL(w8a8_pack_bswz_i8_kernel, weight_grid, block, 0, + stream, raw_weight, packed_weight); } else { const int64_t weight_count = static_cast(k) * n; - const dim3 weight_grid( - static_cast((weight_count + kBlock - 1) / kBlock)); + if (weight_count > 0) { + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL(w8a8_pack_identity_kernel, weight_grid, block, 0, + stream, raw_weight, packed_weight, weight_count); + } + } + if (n > 0) { + const dim3 scale_grid(static_cast( + (static_cast(n) + kBlock - 1) / kBlock)); const dim3 block(kBlock); - hipLaunchKernelGGL(w8a8_pack_identity_kernel, - weight_grid, block, 0, stream, - raw_weight, packed_weight, weight_count); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, scale_grid, block, 0, + stream, weight_scale, packed_weight_scale, n); } - - const int64_t scale_count = n; - const dim3 scale_grid( - static_cast((scale_count + kBlock - 1) / kBlock)); - const dim3 block(kBlock); - hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, - scale_grid, block, 0, stream, - weight_scale, packed_weight_scale, n); (void)hipGetLastError(); } // @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/qkv_proj.hip index 919b94e1..443ac4bc 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/qkv_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M4096/qkv_proj.hip @@ -1,122 +1,54 @@ -// @@variant shape=hy3_tp4_qkv_proj_m4096 commit=61a0f079ab9622cc6c83a9c44af6ecd626ef0640 added=2026-08-23 -// median_us=756.7 p90_us=760.2 tops=113.5 bandwidth_gb_s=63.78 speedup=31.78 baseline_us=2.405e+04 -// source=hy3-tp4-m4096-dsh-test1-bfdf3002 -// MetaInfer W8A8 INT8 GEMM bootstrap for gfx928 (K500SM_AI). +// @@variant shape=hy3_tp4_qkv_proj_m4096 commit=2eebb654a8b34a0ec8bdca109fa8daef07be1e91 added=2026-09-14 +// median_us=625.3 p90_us=626.8 speedup=38.45 baseline_us=2.405e+04 +// source=test-9-11-1/iteration_002/hy3_tp4_qkv_proj_m4096 (retested A/B) +// MetaInfer W8A8 INT8 GEMM for gfx928 (K500SM_AI), CHILD worker_0. // -// Worker: worker_0 (physical GPU 0) -// Assigned shape: hy3_tp4_qkv_proj_m4096 (M=4096, N=2560, K=4096) +// Assigned shape (iteration 1 / establish_arch round, packed_weight family): +// hy3_tp4_qkv_proj_m4096 : M=4096, N=2560, K=4096 // -// Operator contract (fixed by the control plane): -// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) +// Frozen logical contract (validated in int8_w8a8_gemm_api.py): +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) // * x_scale[m] * weight_scale[n] ) -// with the packed weight layout equal to the logical [K, N] row-major int8 -// weight for the bootstrap (launch_pack_w8a8_weight is an identity D2D copy). +// The bootstrap packed weight was the logical [K, N] row-major int8 weight +// (identity copy). This session keeps that layout for every shape except the +// exact assigned pair (K=4096, N=2560), which launch_pack_w8a8_weight packs +// once (outside the timed region and outside Graph capture) into the +// 16-byte-chunk layout packed16[(k/16) * N + n] with the same k*n byte count. // -// Iteration 3 (harness) strategy: single-buffered 64x128 macro-tile with -// cooperative A+B LDS staging -- the matched control arm of the mandated -// pipeline comparison (single vs double buffering across K tiles). -// * Large-M path (M >= 128, N % 128 == 0, K % 128 == 0): native INT8 DUMMA -// m16n16k32 tiled kernel with a 64x128 output tile per block and 256 -// threads (4 wavefronts). Each wave owns a 32x64 quadrant (eight -// m16n16k32 int32 accumulator fragments). K is staged cooperatively in a -// SINGLE LDS buffer at 64-K granularity: the next stage's 3 x int4 -// (1 A + 2 B vectors) are loaded from global only after the current -// stage's 16-MMAC DUMMA burst and its all-consumption barrier, waited -// on, and committed as int32 LDS stores into the same buffer, followed -// by a store-visibility barrier -- TWO __syncthreads per 64-K stage -// (128 over K=4096). The global-load latency is fully exposed (no -// compute overlap), which isolates the value of the double-buffer -// overlap at the same 64-K stage size, odd-word LDS strides, and -// 2-blocks/CU residency as the accepted iteration-2 double-buffered -// kernel (1108.58 us median; iteration-1 single-buffered 128-K stage: -// 1296.50 us). LDS/block = 64x68 + 64x132 = 12,800 B -> 2 resident -// blocks/CU (VGPR-bound). The grid (N/128) x (ceil(M/64)) supplies 1280 -// blocks for the assigned shape. -// * Comparison evidence (fact ledger): double-64 vs single-128 improved -// the benchmark median 1296.50 -> 1108.58 us, but PMC moved against the -// 'reduced VMEM stalls' reading: vmem_read_instructions stayed flat at -// 1,177,600, lds_wait_instructions rose 19.48M -> 20.69M, bank conflicts -// rose 5.24M -> 7.86M, L2 hit fell 89.46% -> 88.25%. This round re-runs -// the single-buffered pipeline at the SAME 64-K stage to decide -// retention on matched evidence (benchmark median/P90 is authoritative -// per the interpretation guard; PMC counters are diagnostic). -// * Operand-reuse accounting per 64x128 macro-tile (A[64,64] + B[64,128] -// per 64-K stage): cooperative staging loads each A element from global -// once per stage and reuses it for 8 MMAs (512 MMA-uses per element over -// K=4096), and each B element for 2 MMAs per stage (128 MMA-uses over -// K=4096). Without staging, the two waves sharing each A row / B column -// segment would refetch everything (2x global traffic) with per-lane -// 8-byte fragment loads (4x more global-load instructions), so the -// cooperative A+B staging path is kept. -// * LDS row strides stay padded to odd word counts (A: 68 B = 17 words, -// B: 132 B = 33 words). The 64x64 bootstrap used unpadded 128/64-byte -// rows, so every INT8 m16n16k32 fragment load collapsed onto a handful -// of LDS banks (16-way conflicts; 199.2M conflict events over 24.9M LDS -// instructions in PMC). The odd word stride makes the lane row/k-group -// index contribute to the bank, spreading fragment loads across banks -// (A and B: 16-way -> 2-4-way per LDS instruction) at 2 blocks/CU. -// * The epilogue applies float(dot) * x_scale[m] * weight_scale[n] and -// stores bf16 directly from the accumulator fragments using the verified -// gfx928 int8 m16n16k32 lane mapping (row = lane & 15, col_mod4 = -// lane >> 4, frag.x[i] -> columns col_mod4 + 4*i). -// * Fallbacks are preserved unchanged: the 64x64 kernel for large-M shapes -// with N % 64 == 0 but N % 128 != 0, and the scalar int8/int32 -// grid-stride fallback for every other (m, n, k) including the paired -// M=2 API shape. launch_pack_w8a8_weight stays an identity -// device-to-device copy valid for every (K, N). +// Strategy: native INT8 DUMMA m16n16k32 tiled kernel with int32 accumulation. +// * Large-M generic path (m >= 128, n % 128 == 0, k % 64 == 0): one 64x128 +// output macro-tile per block, 256 threads = 4 gfx928 wavefronts of 64 +// lanes. Each wave owns a 32x64 quadrant = eight m16n16k32 int32 +// accumulator fragments (2 row halves x 4 column quarters). K is staged +// cooperatively in a SINGLE LDS buffer at 64-K granularity, consumed by a +// 32-K unrolled fragment/MMA loop. Exactly two __syncthreads per stage; +// no split-K, no double buffering, no raw asm. +// * Exact-shape path (m >= 128, n == 2560, k == 4096): the SAME 64x128 +// macro-tile / 4-wavefront / single-buffered-64-K geometry, but B is +// packed 16-byte-chunk and consumed with one ds_read2_b64 per m16n16k32 +// fragment instead of eight ds_read_u8. Measured motivation (PMC on the +// bootstrap digest): lds_wait_instructions 42,668,821 = 1.76x +// lds_instructions (24,248,320), lds_bank_conflicts 7,864,320 = 32% of +// LDS instructions -> the LDS pipe is the dominant stall class at +// 99.6 TOPS, and the byte-granular row-major B fragment loads are the +// primary contributor. +// * Grid for the assigned shape: (2560 / 128) x (4096 / 64) = 20 x 64 = +// 1280 blocks, 4 waves/block, 13,568 B LDS/block -> 3 resident +// blocks/CU on the 120-CU K500SM_AI. +// * Scalar int8/int32 grid-stride fallback for every unmatched (m, n, k) +// and for small-M API cases (m < 128), which keeps the paired M=2 shape +// on the fallback path; that fallback decodes the packed layout for the +// exact (k, n) pair and the logical [K, N] layout otherwise. // -// Iteration 5 (packing round) strategy: keep the accepted iteration-3 -// pipeline (single-buffered 64-K stage, 64x128 macro-tile, 3 blocks/CU) and -// add ONE weight packing/swizzle for the exact assigned shape (K=4096, -// N=2560): launch_pack_w8a8_weight transposes the logical [K, N] weight to -// an n-major packed[n][k] layout once, outside the timed region and outside -// Graph capture (same byte count k*n, same allocation, graph-stable packed -// layout). The exact shape routes (before the generic 64x128 path) to a new -// w8a8_dumma_prefill_64x128_packedb_kernel that stages B into an n-major LDS -// tile with a 16-byte-aligned 80-byte row stride (20 words, five bank -// phases) and consumes it with col_major m16n16k32 B fragments via an -// explicit 8-byte loader (load_b_frag8): each lane's 8 fragment elements at -// p[row*ldm + col + i] are one contiguous 8-byte-aligned chunk (one -// ds_read2_b64 per B fragment instead of 8 byte-granular ds_read_u8 at -// stride 132 in the accepted kernel's ISA), and each 16-byte staging vector -// commits with one ds_write_b128 instead of four ds_write_b32. A stays -// row-major (already vectorized ds_read2_b32 in the accepted ISA); the int32 -// accumulation order (k0-outer, kk-inner) and the fused scale/bf16 epilogue -// are unchanged, so results stay bit-identical. All other shapes keep the -// identity pack and the accepted iteration-3 kernel byte-for-byte; the 64x64 -// kernel, the scalar fallback (which decodes packed[n][k] when (k, n) -// matches), M tails, and the exact shape guards are preserved. +// The timed operator (launch_w8a8_gemm) only dispatches kernels on the +// caller-provided HIP stream: it never allocates, compiles, autotunes, packs +// weights, synchronizes the host or device, or launches on the default stream. +// launch_pack_w8a8_weight performs the one-time exact-shape chunk packing +// and an identity copy for every other (K, N); both are outside Graph capture. // -// Iteration 6 (epilogue round) strategy: keep the accepted iteration-5 -// packedb kernel's pipeline byte-identical (single-buffered 64-K stage, -// n-major B with 80-B rows, load_b_frag8, 3 blocks/CU) and replace ONLY its -// final store. The scales and bf16 conversion are already fused in the -// kernel and launch_w8a8_gemm has no workspace/combine pass (workspace is -// unused), so the remaining epilogue inefficiency is the store pattern: the -// m16n16k32 accumulator lane mapping (row = lane & 15, column group -// c4 = lane >> 4, frag.x[i] -> column c4 + 4*i) makes the direct store four -// 2-byte scalar stores per lane whose wavefront addresses touch every 32-B -// sector at 25% utilization. store_prefill_fragment_coalesced transposes the -// 4-element groups within each 4-lane column group (lanes r, r+16, r+32, -// r+48) with two 2x2 shuffle steps (8 shfl_xor + 8 v_cndmask; this DTK -// lowers each shfl_xor to one ds_bpermute -- +64 LDS permutes per wavefront -// at the block tail, where the LDS pipe is otherwise idle), so lane (r, c4) -// owns the four CONTIGUOUS columns 4*c4 .. 4*c4+3, then converts to bf16, -// packs 4 bf16 (8 B), and issues ONE 8-byte store per lane: 64 lanes x 8 B = -// 512 B per fragment per wavefront in 16 fully-used 32-B sectors (100% -// store sector efficiency; the assigned shape's vmem_write_instructions drop -// 163,840 -> 40,960). Only the int32 values are re-routed between lanes; the -// per-element float scale/bf16 math is unchanged, so output bits are -// identical. The generic 64x128/64x64 kernels and the scalar fallback are -// untouched (their store_prefill_fragment epilogue is preserved -// byte-for-byte). -// -// The timed operator (launch_w8a8_gemm) performs no allocation, compilation, -// autotuning, packing, host/device synchronization, or default-stream launch: -// it only dispatches kernels on the caller-provided HIP stream. -// -// Include order is fixed by the control plane: hip_runtime, hip_bfloat16, -// then du_mma (this DTK's du_mma.h is not self-contained otherwise). +// Include order is fixed by the control plane: hip_runtime, then +// hip_bfloat16, then du_mma (this DTK's du_mma.h is not self-contained when +// included before the HIP runtime headers). Do not include du_mma_common.h. #include #include @@ -126,57 +58,99 @@ namespace { +using namespace du::dumma; + +// gfx928 native wavefront is 64 lanes; blockDim must be a multiple of 64. constexpr int kWaveSize = 64; + +// Native INT8 DUMMA shape. constexpr int kTileM = 16; constexpr int kTileN = 16; constexpr int kTileK = 32; + +// 64x128 macro-tile, single-buffered 64-K LDS stage, 4 wavefronts. constexpr int kBlockM = 64; -constexpr int kBlockN = 64; -constexpr int kStageK = 128; -// 4 wavefronts; must remain a multiple of the gfx928 wavefront size (64). +constexpr int kBlockN = 128; +constexpr int kStageK = 64; constexpr int kThreadsPerBlock = 256; -// Iteration 3 (pipeline round): 64x128 macro-tile, SINGLE-buffered 64-K K -// stage (matched control for the single-vs-double buffering comparison). -// LDS row strides stay padded to odd word counts (A: 68 B = 17 words, -// B: 132 B = 33 words) so the INT8 m16n16k32 fragment loads spread across -// banks instead of collapsing 16-way onto one bank (128-byte rows are 32 -// words, i.e. 0 mod 32 for every row -> single-bank collisions). -constexpr int kBlockM128 = 64; -constexpr int kBlockN128 = 128; -// Single-buffered K stage: one 64-K buffer = 12,800 B/block (the accepted -// iteration-2 double buffer was 2 x 64 = 25,600 B; the iteration-1 -// single-buffered 128-K stage was 25,344 B). Occupancy stays 2 blocks/CU -// (VGPR-bound at arch_vgpr ~112), identical to both prior kernels. -constexpr int kStageK128 = 64; -constexpr int kAStride128 = kStageK128 + 4; // 68 bytes per A row (17 words) -constexpr int kBStride128 = kBlockN128 + 4; // 132 bytes per B row (33 words) - -// Iteration 5 (packing round): packed n-major B layout for the exact -// assigned shape (K=4096, N=2560). B is transposed once, outside timing, to -// packed[n][k]; the packed kernel stages B into an n-major LDS tile with a -// 16-byte-aligned non-power-of-two row stride (80 B = 20 words, five bank -// phases) so every DUMMA m16n16k32 B fragment load (col_major) reads its -// lane's 8 bytes contiguously (ds_read2_b64) and every 16-byte staging -// vector commits with one ds_write_b128 (b_n*80 + b_k16 is 0 mod 16). -constexpr int kPackedBStride = 80; // 64 data + 16 pad bytes (20 words) +// Odd-word (4-byte-padded) LDS row strides: 17 and 33 words. +constexpr int kAStride = kStageK + 4; // 68 bytes per A row +constexpr int kBStride = kBlockN + 4; // 132 bytes per B row + +// Iteration 2 (operand-reuse round) strategy: keep the accepted iteration-1 +// cooperative A+B LDS staging decision and fix ONLY the packed B LDS layout to +// be bank-safe. Operand-reuse accounting per 64x128 macro-tile per 64-K stage +// (A[64,64] + B[64,128] = 12,288 staged bytes): the 2x2 wave grid reads every +// staged A element 2x and every staged B element 2x (48 m16n16k32 fragment +// loads per block stage, contiguous 512-byte fragments), i.e. the cooperative +// path moves 12,288 global bytes but only 2x that through LDS. Loading the +// DUMMA fragments directly from global would instead refetch both operands +// twice with 4,096-byte-strided (>32x uncoalesced) per-lane 8-byte addresses, +// so cooperative A+B staging is retained and only its LDS layout changes: +// * the weight is now packed once (untimed, outside Graph capture) into a +// 16-byte-chunk layout packed16[(k/16) * N + n] holding +// weight[16*(c/N) .. +15][c%N], same k*n byte count, so the GEMM's B +// staging reads become 4 coalesced 256-byte runs per 64-lane wavefront +// (consecutive lanes -> consecutive 16-byte chunks) at unchanged volume; +// * the B LDS tile uses an 8-byte-aligned 72-byte (18-word) row stride and +// the staging map lane = kq*16 + row4 (row4 = lane & 15 = n-row inside a +// 16-row group, kq = (lane >> 4) & 3 = 16-byte k-chunk), so both the +// col_major m16n16k32 fragment read (n = lane & 15, k = 8*(lane>>4)+i; bank +// = 18*row4 + 2*kq) and the staged commit at n_local*72 + 16*kq spread the +// 16-lane issue phase over 16 distinct even banks plus their odd +// neighbours. The previous 80-byte (20-word) row stride left only 8 +// distinct banks per 16-row phase; the fresh PMC on the iteration-1 digest +// records lds_bank_conflicts 23,592,960 over lds_instructions 3,932,160 +// (6.0 events per LDS read instruction), the largest LDS defect. +// * the chunk commit is two alignment-provable 8-byte qword stores (the tile +// is declared as int64_t), the fragment registers receive byte-identical +// values, and the int32 k0-outer/kk-inner accumulation order is unchanged. +// LDS/block = 64*68 + 128*72 = 13,568 B; the 64x128 geometry, 256 threads / +// 4 wavefronts, single-buffered 64-K stage, two barriers per stage, 3 +// blocks/CU occupancy, fused scale/bf16 epilogue, exact shape guards, the +// generic 64x128 fallback, the scalar fallback (which decodes the chunked +// layout for the exact (k, n)) and every other shape's path are unchanged. +// +// Iteration 3 (fix_build round): the iteration-2 candidate above failed to +// compile for exactly one reason -- launch_pack_w8a8_weight still referenced +// the removed w8a8_pack_nmajor_b_kernel symbol after the pack kernel was +// rewritten as w8a8_pack_k16chunk_b_kernel. The only change this round is +// rebinding that host call to the defined 16-byte-chunk pack kernel (plus +// comment text); the kernel bodies, geometry, guards and layouts are byte-for- +// byte the iteration-2 candidate. +// +// Iteration 1 (establish_arch round, packed_weight family) for the exact +// assigned shape (K=4096, N=2560): launch_pack_w8a8_weight transposed the +// logical [K, N] row-major weight once, outside the timed region and outside +// Graph capture, into an n-major layout with the SAME byte count (k*n), so +// allocation and graph-stable addresses are unchanged. Every DUMMA m16n16k32 B +// fragment collapsed to ONE contiguous 8-byte lane chunk (one ds_read2_b64) +// instead of eight byte-granular ds_read_u8 at the 132-byte row_major stride. +// +// Iteration 2 (operand-reuse round): the pack is now a 16-byte-chunk layout +// packed16[(k/16) * N + n] (chunk c = weight[16*(c/N) .. +15][c%N]) so the +// staging global reads are 4 coalesced 256-byte runs per wavefront, and the B +// LDS tile row stride becomes 72 bytes (18 words, 8-byte aligned) with the +// staging map lane = kq*16 + row4. A 16-row phase then touches 16 distinct +// even banks (18*row4) plus their odd neighbours instead of the 8 banks the +// 20-word stride produced, removing the measured 6.0 bank-conflict events per +// LDS read instruction (23,592,960 conflicts / 3,932,160 LDS instructions). +constexpr int kPackedBStride = 72; // 64 data + 8 pad bytes (18 words) +constexpr int kPackedBWords = kPackedBStride / 8; // 9 qwords per n-row constexpr int kPackedBK = 4096; // exact K of the packed assigned shape constexpr int kPackedBN = 2560; // exact N of the packed assigned shape -// Iteration 18 (fragment-pipeline round): the packedb compute loop splits the -// 64-K stage into two explicit 32-K steps (kk = 0, kk = 32) and hoists ALL -// fragment loads for both steps before the first v_mmac. This requires -// exactly two kTileK steps per stage. -static_assert(kStageK128 == 2 * kTileK, - "pipelined kk loop requires exactly two 32-K steps per stage"); +// Fallback launch geometry. +constexpr int kFallbackThreads = 256; +constexpr int kPackThreads = 256; -using namespace du::dumma; - -// Direct accumulator epilogue for gfx928 INT8 m16n16k32 (verified against -// du_store_matrix_sync): row = lane & 15, col_mod4 = lane >> 4, -// frag.x[i] maps to columns col_mod4 + 4*i. Stores +// Direct accumulator epilogue for the gfx928 INT8 m16n16k32 DUMMA lane +// mapping (verified against du_store_matrix_sync): row = lane & 15, +// col_mod4 = lane >> 4, and frag.x[i] maps to columns col_mod4 + 4*i. Stores // bf16(float(dot) * x_scale[row] * weight_scale[col]); out-of-range rows are -// masked (tail-M handling). +// masked so a partial M tail is handled safely. Adjacent lanes stay on +// adjacent columns in the fastest-changing N dimension. template __device__ __forceinline__ void store_prefill_fragment( const AccFragment& frag, @@ -202,101 +176,10 @@ __device__ __forceinline__ void store_prefill_fragment( } } -// Iteration 6 (epilogue round): coalesced fragment store. The m16n16k32 -// accumulator lane mapping (row = lane & 15, column group c4 = lane >> 4, -// frag.x[i] -> column c4 + 4*i) gives each lane four elements strided by 4 -// columns, so the direct store (store_prefill_fragment) is four 2-byte -// scalar stores per lane whose wavefront addresses touch each 32-B sector at -// 25% utilization (8 B used of 32 B; 4 store instructions per fragment per -// wave). This epilogue transposes the 4-element groups within each 4-lane -// column group (lanes r, r+16, r+32, r+48 -- a 4x4 transpose, two 2x2 steps -// with shfl_xor 16 then 32, one v_cndmask per element per step; no staging -// tile round trip), so lane (r, c4) ends up holding the four CONTIGUOUS -// columns 4*c4 .. 4*c4+3, converts them to bf16, packs 4 bf16 (8 B), and -// writes ONE 8-byte store per lane: 64 lanes x 8 B = 512 B per fragment per -// wavefront in 16 fully-used 32-B sectors (100% store sector efficiency; -// vmem_write_instructions for the assigned shape drops 163,840 -> 40,960). -// Note: this DTK's __shfl_xor lowers to ds_bpermute (one LDS permute per -// shuffle, i.e. +64 LDS instructions per wavefront per block -- accepted -// because it runs at the block tail where the LDS pipe is otherwise idle; -// there is no global round trip and no staging tile). Only the int32 values -// are re-routed between lanes -- the per-element scale multiply order -// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are -// unchanged, so the stored bits are identical to store_prefill_fragment. -// The row>=m guard is wavefront-uniform (lane & 15 cycles the same 16 rows -// in every 16-lane group), so the shuffles never mix active and inactive -// lanes; base_col is a multiple of 64 and n*2 a multiple of 8, so the float4 -// weight_scale load (col0 % 4 == 0) and the 8-byte store are aligned. -template -__device__ __forceinline__ void store_prefill_fragment_coalesced( - const AccFragment& frag, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - hip_bfloat16* __restrict__ out, - int m, - int n, - int base_row, - int base_col, - int lane) { - const int row = base_row + (lane & 15); - if (row >= m) { - return; - } - const int c4 = lane >> 4; // 0..3 - const int x0 = frag.x[0]; - const int x1 = frag.x[1]; - const int x2 = frag.x[2]; - const int x3 = frag.x[3]; - // Step 1: transpose the 2x2 in-block pairs (c4 bit0 <-> element bit0). - const int s0 = __shfl_xor(x0, 16, kWaveSize); - const int s1 = __shfl_xor(x1, 16, kWaveSize); - const int s2 = __shfl_xor(x2, 16, kWaveSize); - const int s3 = __shfl_xor(x3, 16, kWaveSize); - const bool lo = (c4 & 1) == 0; - const int a0 = lo ? x0 : s1; - const int a1 = lo ? s0 : x1; - const int a2 = lo ? x2 : s3; - const int a3 = lo ? s2 : x3; - // Step 2: transpose the 2x2 blocks (c4 bit1 <-> element bit1). - const int t0 = __shfl_xor(a0, 32, kWaveSize); - const int t1 = __shfl_xor(a1, 32, kWaveSize); - const int t2 = __shfl_xor(a2, 32, kWaveSize); - const int t3 = __shfl_xor(a3, 32, kWaveSize); - const bool hi = (c4 & 2) == 0; - const int f0 = hi ? a0 : t2; - const int f1 = hi ? a1 : t3; - const int f2 = hi ? t0 : a2; - const int f3 = hi ? t1 : a3; - - // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte - // aligned: base_col is a multiple of 64, n*2 is a multiple of 8). - const int col0 = base_col + 4 * c4; - const float xs = x_scale[row]; - const float4 ws = *reinterpret_cast(weight_scale + col0); - const float v0 = static_cast(f0) * xs * ws.x; - const float v1 = static_cast(f1) * xs * ws.y; - const float v2 = static_cast(f2) * xs * ws.z; - const float v3 = static_cast(f3) * xs * ws.w; - const uint64_t packed = - static_cast( - static_cast(__float2bfloat16(v0))) | - (static_cast( - static_cast(__float2bfloat16(v1))) - << 16) | - (static_cast( - static_cast(__float2bfloat16(v2))) - << 32) | - (static_cast( - static_cast(__float2bfloat16(v3))) - << 48); - *reinterpret_cast(out + row * n + col0) = packed; -} - -// Large-M prefill kernel: 64x64 output tile per block, K staged in LDS -// (single buffer), four waves each owning a 32x32 quadrant = four -// m16n16k32 int8->int32 DUMMA accumulators. +// Large-M prefill kernel: 64x128 output tile per block, single-buffered 64-K +// K stage, 256 threads / 4 wavefronts. __global__ __launch_bounds__(kThreadsPerBlock) void -w8a8_dumma_prefill_64x64_kernel( +w8a8_dumma_prefill_64x128_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, const float* __restrict__ x_scale, @@ -313,120 +196,8 @@ w8a8_dumma_prefill_64x64_kernel( const int m0 = static_cast(blockIdx.y) * kBlockM; const int n0 = static_cast(blockIdx.x) * kBlockN; - __shared__ __align__(16) int8_t a_tile[kBlockM * kStageK]; - __shared__ __align__(16) int8_t b_tile[kStageK * kBlockN]; - - DUFragment - a_frag0, a_frag1; - DUFragment - b_frag0, b_frag1; - DUFragment - acc00, acc01, acc10, acc11; - du_fill_fragment(acc00, 0); - du_fill_fragment(acc01, 0); - du_fill_fragment(acc10, 0); - du_fill_fragment(acc11, 0); - - for (int k0 = 0; k0 < k; k0 += kStageK) { - // Stage A[64,128] (row-major, zero-filled past M). - constexpr int kAVectors = kBlockM * kStageK / static_cast(sizeof(int4)); - for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { - const int byte_offset = vec * static_cast(sizeof(int4)); - const int local_row = byte_offset / kStageK; - const int kk = byte_offset - local_row * kStageK; - const int global_row = m0 + local_row; - reinterpret_cast(a_tile)[vec] = - global_row < m - ? *reinterpret_cast( - x_q + global_row * k + k0 + kk) - : int4{0, 0, 0, 0}; - } - - // Stage B[128,64] (row-major [K, N] logical weight layout). - constexpr int kBVectors = kStageK * kBlockN / static_cast(sizeof(int4)); - for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { - const int byte_offset = vec * static_cast(sizeof(int4)); - const int kk = byte_offset / kBlockN; - const int col = byte_offset - kk * kBlockN; - reinterpret_cast(b_tile)[vec] = - *reinterpret_cast( - weight + (k0 + kk) * n + n0 + col); - } - __syncthreads(); - - // Consume the stage: each wave does four m16n16k32 MMAs per 32-K step. - const int local_row = wave_row * 32; - const int local_col = wave_col * 32; -#pragma unroll - for (int kk = 0; kk < kStageK; kk += kTileK) { - du_load_matrix_sync( - a_frag0, a_tile + local_row * kStageK + kk, kStageK); - du_load_matrix_sync( - a_frag1, a_tile + (local_row + kTileM) * kStageK + kk, kStageK); - du_load_matrix_sync( - b_frag0, b_tile + kk * kBlockN + local_col, kBlockN); - du_load_matrix_sync( - b_frag1, b_tile + kk * kBlockN + local_col + kTileN, kBlockN); - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); - } - // Single buffer: make sure every wave finished reading LDS before the - // next stage overwrites it. - __syncthreads(); - } - - const int base_row = m0 + wave_row * 32; - const int base_col = n0 + wave_col * 32; - store_prefill_fragment( - acc00, x_scale, weight_scale, out, m, n, - base_row, base_col, lane); - store_prefill_fragment( - acc01, x_scale, weight_scale, out, m, n, - base_row, base_col + kTileN, lane); - store_prefill_fragment( - acc10, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col, lane); - store_prefill_fragment( - acc11, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col + kTileN, lane); -} - -// Large-M prefill kernel #2 (iteration 3, pipeline round): 64x128 output -// tile per block, 256 threads (4 wavefronts). Each wave owns a 32x64 -// quadrant = eight m16n16k32 int32 accumulators (2 row halves x 4 column -// quarters). K is staged cooperatively in a SINGLE LDS buffer at 64-K -// granularity as the matched control for the mandated single-vs-double -// buffering comparison: it shares the iteration-2 double-buffered kernel's -// 64-K stage size, odd-word LDS strides (A: 68 B = 17 words, B: 132 B = 33 -// words), 16-MMAC DUMMA burst, and 2-blocks/CU residency, but uses ONE -// 12,800-B buffer and TWO __syncthreads per stage (128 over K): the next -// stage's 3 x int4 (1 A + 2 B) are loaded from global only after the -// current burst and its all-consumption barrier, vmcnt-waited, and -// committed into the same buffer, then a store-visibility barrier. The -// global-load latency is fully exposed (no compute overlap), isolating the -// value of the double-buffer overlap at matched stage size and occupancy. -__global__ __launch_bounds__(kThreadsPerBlock) void -w8a8_dumma_prefill_64x128_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - hip_bfloat16* __restrict__ out, - int m, - int n, - int k) { - const int tid = static_cast(threadIdx.x); - const int wave = tid / kWaveSize; - const int lane = tid % kWaveSize; - const int wave_row = wave >> 1; // 0 or 1 - const int wave_col = wave & 1; // 0 or 1 - const int m0 = static_cast(blockIdx.y) * kBlockM128; - const int n0 = static_cast(blockIdx.x) * kBlockN128; - - __shared__ __align__(16) int8_t a_tile[kBlockM128 * kAStride128]; - __shared__ __align__(16) int8_t b_tile[kStageK128 * kBStride128]; + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; DUFragment a_frag0, a_frag1; @@ -443,22 +214,21 @@ w8a8_dumma_prefill_64x128_kernel( du_fill_fragment(acc12, 0); du_fill_fragment(acc13, 0); - const int num_stages = k / kStageK128; + const int num_stages = k / kStageK; // Cooperative staging mapping (fixed per thread, reused every stage): - // A[64,64] -> 256 int4 vectors; thread `tid` owns vector `tid` - // (row = tid/4, 16-B column group = (tid%4)*16), so each - // wavefront covers 16 rows x 64 B contiguous per row. - // B[64,128] -> 512 int4 vectors; thread `tid` owns vectors `tid` and - // `tid+256` (kk = tid/8, column group = (tid%8)*16), so - // each wavefront covers 8 kk rows x 128 B contiguous. + // A[64, 64] -> 256 int4 vectors; thread `tid` owns vector `tid` + // (row = tid/4, 16-B column group = (tid%4)*16). + // B[64, 128] -> 512 int4 vectors; thread `tid` owns vectors `tid` and + // `tid+256` (kk = tid/8, 16-B column group = (tid%8)*16), + // i.e. kk rows 0..31 and 32..63. const int a_row = tid >> 2; const int a_col16 = (tid & 3) << 4; const int b_kk0 = tid >> 3; const int b_col16 = (tid & 7) << 4; - // VGPR payload for one full 64-K stage (3 x int4 = 12 VGPR), loaded just - // before it is committed (no one-stage-ahead overlap in this control). + // One-stage register payload (3 x int4 = 12 VGPR), loaded just before it is + // committed; there is no stage-ahead overlap in this bootstrap. int4 vA, vB0, vB1; // Prologue: load stage 0 into registers, commit it to the single buffer, @@ -471,22 +241,24 @@ w8a8_dumma_prefill_64x128_kernel( vB0 = *reinterpret_cast( weight + b_kk0 * n + n0 + b_col16); vB1 = *reinterpret_cast( - weight + (b_kk0 + kStageK128 / 2) * n + n0 + b_col16); - // Commit as int32 stores (skips the 4-byte pad column; 4 x ds_write_b32 - // per 16-byte vector because the odd strides are 4 B mod 16). + weight + (b_kk0 + kStageK / 2) * n + n0 + b_col16); + // Commit as int32 stores: the odd-word strides leave the rows 4-byte but + // not 16-byte aligned, so four ds_write_b32 per 16-byte vector. int32_t* adst = reinterpret_cast( - a_tile + a_row * kAStride128 + a_col16); + a_tile + a_row * kAStride + a_col16); adst[0] = vA.x; adst[1] = vA.y; adst[2] = vA.z; adst[3] = vA.w; int32_t* bdst0 = reinterpret_cast( - b_tile + b_kk0 * kBStride128 + b_col16); + b_tile + b_kk0 * kBStride + b_col16); bdst0[0] = vB0.x; bdst0[1] = vB0.y; bdst0[2] = vB0.z; bdst0[3] = vB0.w; - int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + // kBStride is a whole number of int32 words, so the second half-row is + // exactly (kStageK/2) rows away. + int32_t* bdst1 = bdst0 + (kStageK / 2) * (kBStride / 4); bdst1[0] = vB1.x; bdst1[1] = vB1.y; bdst1[2] = vB1.z; @@ -496,27 +268,23 @@ w8a8_dumma_prefill_64x128_kernel( for (int s = 0; s < num_stages; ++s) { // Consume the single buffer: each wave does eight m16n16k32 MMAs per - // 32-K step over its 32x64 quadrant. + // 32-K step over its own 32x64 quadrant. const int local_row = wave_row * 32; const int local_col = wave_col * 64; #pragma unroll - for (int kk = 0; kk < kStageK128; kk += kTileK) { + for (int kk = 0; kk < kStageK; kk += kTileK) { du_load_matrix_sync( - a_frag0, a_tile + local_row * kAStride128 + kk, kAStride128); + a_frag0, a_tile + local_row * kAStride + kk, kAStride); du_load_matrix_sync( - a_frag1, a_tile + (local_row + kTileM) * kAStride128 + kk, - kAStride128); + a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, kAStride); du_load_matrix_sync( - b_frag0, b_tile + kk * kBStride128 + local_col, kBStride128); + b_frag0, b_tile + kk * kBStride + local_col, kBStride); du_load_matrix_sync( - b_frag1, b_tile + kk * kBStride128 + local_col + kTileN, - kBStride128); + b_frag1, b_tile + kk * kBStride + local_col + kTileN, kBStride); du_load_matrix_sync( - b_frag2, b_tile + kk * kBStride128 + local_col + 2 * kTileN, - kBStride128); + b_frag2, b_tile + kk * kBStride + local_col + 2 * kTileN, kBStride); du_load_matrix_sync( - b_frag3, b_tile + kk * kBStride128 + local_col + 3 * kTileN, - kBStride128); + b_frag3, b_tile + kk * kBStride + local_col + 3 * kTileN, kBStride); 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); @@ -526,18 +294,15 @@ w8a8_dumma_prefill_64x128_kernel( du_mma_sync(acc12, a_frag1, b_frag2, acc12); du_mma_sync(acc13, a_frag1, b_frag3, acc13); } - // All consumption of the single buffer is complete only after every - // wavefront passes this barrier; the stage s+1 stores below overwrite - // that same buffer, so they must wait for it. + // Every wavefront must finish reading the single buffer before the next + // stage overwrites it. __syncthreads(); - // Load stage s+1 and commit it into the single buffer. The global-load - // latency is fully exposed here: the loads are issued after the burst - // (nothing to overlap) and the compiler's vmcnt wait before the DS - // stores is on the critical path. `s + 1 < num_stages` is block-uniform, - // so the branch and its barrier are free of divergence. + // Load stage s+1 from global and commit it into the same buffer. The + // `s + 1 < num_stages` guard is block-uniform, so the barrier below is + // reached by every thread in the block. if (s + 1 < num_stages) { - const int s1 = (s + 1) * kStageK128; + const int s1 = (s + 1) * kStageK; const int g_row = m0 + a_row; vA = (g_row < m) ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) @@ -545,26 +310,25 @@ w8a8_dumma_prefill_64x128_kernel( vB0 = *reinterpret_cast( weight + (s1 + b_kk0) * n + n0 + b_col16); vB1 = *reinterpret_cast( - weight + (s1 + b_kk0 + kStageK128 / 2) * n + n0 + b_col16); + weight + (s1 + b_kk0 + kStageK / 2) * n + n0 + b_col16); int32_t* adst = reinterpret_cast( - a_tile + a_row * kAStride128 + a_col16); + a_tile + a_row * kAStride + a_col16); adst[0] = vA.x; adst[1] = vA.y; adst[2] = vA.z; adst[3] = vA.w; int32_t* bdst0 = reinterpret_cast( - b_tile + b_kk0 * kBStride128 + b_col16); + b_tile + b_kk0 * kBStride + b_col16); bdst0[0] = vB0.x; bdst0[1] = vB0.y; bdst0[2] = vB0.z; bdst0[3] = vB0.w; - int32_t* bdst1 = bdst0 + (kStageK128 / 2) * (kBStride128 / 4); + int32_t* bdst1 = bdst0 + (kStageK / 2) * (kBStride / 4); bdst1[0] = vB1.x; bdst1[1] = vB1.y; bdst1[2] = vB1.z; bdst1[3] = vB1.w; - // Make the stage s+1 stores visible to every wavefront before the - // next burst. + // Make the stage s+1 stores visible before the next burst. __syncthreads(); } } @@ -572,96 +336,80 @@ w8a8_dumma_prefill_64x128_kernel( const int base_row = m0 + wave_row * 32; const int base_col = n0 + wave_col * 64; store_prefill_fragment( - acc00, x_scale, weight_scale, out, m, n, - base_row, base_col, lane); + acc00, x_scale, weight_scale, out, m, n, base_row, base_col, lane); store_prefill_fragment( - acc01, x_scale, weight_scale, out, m, n, - base_row, base_col + kTileN, lane); + acc01, x_scale, weight_scale, out, m, n, base_row, base_col + kTileN, + lane); store_prefill_fragment( - acc02, x_scale, weight_scale, out, m, n, - base_row, base_col + 2 * kTileN, lane); + acc02, x_scale, weight_scale, out, m, n, base_row, + base_col + 2 * kTileN, lane); store_prefill_fragment( - acc03, x_scale, weight_scale, out, m, n, - base_row, base_col + 3 * kTileN, lane); + acc03, x_scale, weight_scale, out, m, n, base_row, + base_col + 3 * kTileN, lane); store_prefill_fragment( - acc10, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col, lane); + acc10, x_scale, weight_scale, out, m, n, base_row + kTileM, base_col, + lane); store_prefill_fragment( - acc11, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col + kTileN, lane); + acc11, x_scale, weight_scale, out, m, n, base_row + kTileM, + base_col + kTileN, lane); store_prefill_fragment( - acc12, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col + 2 * kTileN, lane); + acc12, x_scale, weight_scale, out, m, n, base_row + kTileM, + base_col + 2 * kTileN, lane); store_prefill_fragment( - acc13, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col + 3 * kTileN, lane); + acc13, x_scale, weight_scale, out, m, n, base_row + kTileM, + base_col + 3 * kTileN, lane); } -// Iteration 5 (packing round): explicit 8-byte loader for the col_major -// m16n16k32 B fragment (the gate_up-lineage-validated load_fragment8 -// pattern). The col_major fragment mapping is n = lane&15, -// k = 8*(lane>>4) + i; with the n-major LDS tile at row stride 80 the -// lane's 8 elements are contiguous (p[row*ldm + col .. +7], 8-byte aligned -// because ldm=80 and col are multiples of 8), so this compiles to ONE -// ds_read2_b64 instead of 8 ds_read_u8 + mask/OR reassembly. The byte -// placement is identical to du_load_matrix_sync, so the -// v_mmac_i32_16x16x32_i8 fragment registers receive the same values. +// Iteration 1/2 (packed_weight family): explicit 8-byte loader for the +// col_major m16n16k32 B fragment. With the n-major LDS tile at row stride +// ldm (72 bytes in iteration 2, 80 in iteration 1), the col_major fragment +// mapping (n = lane & 15, k = 8 * (lane >> 4) + i) places each lane's 8 +// elements at p[(lane & 15) * ldm + ((lane >> 4) << 3) .. +7]: one contiguous, +// 8-byte-aligned chunk (ldm and every k offset are multiples of 8), so this is +// ONE ds_read2_b64 instead of eight ds_read_u8 + mask/OR reassembly. The bytes +// placed in the fragment registers are exactly what +// du_load_matrix_sync would produce, so the +// v_mmac_i32_16x16x32_i8 operands and the int32 accumulation are unchanged. __device__ __forceinline__ void load_b_frag8( DUFragment& f, const int8_t* __restrict__ p, - int ldm) { - const unsigned row = static_cast(__lane_id()) & 0xfu; - const unsigned col = (static_cast(__lane_id()) >> 4) << 3; + int ldm, + int lane) { + const int row = lane & 15; + const int col = (lane >> 4) << 3; const int64_t v = *reinterpret_cast(p + row * ldm + col); *reinterpret_cast(&f.x[0]) = v; } -// Large-M prefill kernel #3 (iteration 5, packing round): 64x128 macro-tile -// for the exact assigned shape (K=4096, N=2560) with the weight packed -// n-major. The pipeline is byte-identical to the accepted iteration-3 kernel -// (single-buffered 64-K stage, TWO __syncthreads per stage with fully -// exposed global-load latency, 3 x int4 VGPR payload, A staged row-major -// with the 68-byte odd-word stride); ONLY the B operand path changes: -// * pack: the logical [K, N] weight is transposed once (outside the timed -// region and Graph capture) to packed[n][k] by launch_pack_w8a8_weight. -// * staging: thread tid owns rows n_local = tid>>2 and n_local+64, each -// with the 16-byte k-run (tid&3)*16; two global_load_dwordx4 are -// committed with ONE ds_write_b128 each (16-byte-aligned LDS address: -// row stride 80 and the k16 offset are both 0 mod 16), vs four -// ds_write_b32 per vector in the accepted kernel. -// * fragments: load_b_frag8 (col_major mapping n = local_col + lane&15, -// k = kk + 8*(lane>>4) + i) reads each lane's 8 elements as one -// 8-byte-aligned contiguous chunk -> one ds_read2_b64 per B fragment, -// vs 8 byte-granular ds_read_u8 at stride 132 in the accepted kernel's -// ISA. The mapping is the transpose of the accepted row_major mapping, -// so the v_mmac_i32_16x16x32_i8 accumulation order is unchanged and the -// int32 results are bit-identical. Bank safety: the 20-word row stride -// spreads each k-group's 16 n-lanes over 8 distinct banks with five -// bank phases. -// * A path, epilogue, tail masking, and occupancy are unchanged. -// LDS/block = 64x68 + 128x80 = 14,592 B -> 3 resident blocks/CU -// (VGPR-bound at arch_vgpr ~80), same as the accepted kernel. -// -// Iteration 18 (fragment-pipeline round) strategy: keep the accepted -// iteration-6 packedb kernel byte-identical in EVERYTHING except the inner -// fragment schedule of the 64-K-stage compute loop. The accepted kernel's -// exact gfx928 ISA (iteration 7 evidence) shows the compiler already -// software-pipelines the kk=32 A fragment loads under the kk=0 MMAC burst, -// but its four kk=32 B-fragment ds_read_b64 loads are issued only AFTER -// MMACs 0-7, with s_waitcnt lgkmcnt(4..1) interleaved into MMACs 8-11 (16+ -// lgkmcnt waits per stage; PMC lds_wait 6,560,572 = 1.18x lds_instructions, -// the dominant stall class of a SIMD ~36%-issue-busy kernel). This round -// hoists ALL twelve ds_read fragment loads of both 32-K steps before the -// first v_mmac of the stage (explicit two-step unroll of the kk loop with -// separate kk=32 fragment variables), so the LDS latency concentrates in one -// wait and the 16-MMAC stream has no lgkmcnt waits. The MMAC sequence per -// (m,n) output tile and the int32 k0-outer/kk-inner accumulation order are -// unchanged, so output bits stay identical. Occupancy target is retained: -// arch_vgpr 72 -> <= 85 at 3 blocks/CU (65,280 <= 65,536 VGPRs), LDS 14,592 -// B, grid 1280 x 256, TWO __syncthreads per stage, staging mapping, packed -// B layout, coalesced epilogue, and the exact shape guard (m >= 128 && -// n == 2560 && k == 4096) all unchanged; generic 64x128/64x64 kernels, the -// pack kernel, and the scalar fallback are untouched. +// Iteration 2 (operand-reuse round): assemble one 8-byte qword from the two +// low 32-bit halves of a packed 16-byte staging vector. The chunk commit uses +// two of these per 16-byte vector (72-byte LDS rows are only 8-byte aligned), +// which keeps the store side of the B tile bank-safe at identical byte volume. +__device__ __forceinline__ int64_t pack_lo32_pair(int lo, int hi) { + return static_cast(static_cast(lo)) | + (static_cast(static_cast(hi)) << 32); +} + +// Iteration 1/2 (packed_weight family): packed B variant of the 64x128 +// macro-tile for the exact assigned shape (K=4096, N=2560). The geometry, +// K staging depth (64), barrier count (two __syncthreads per stage), A path +// (row-major, 68-byte odd-word stride, 4 x ds_write_b32 per 16-byte vector), +// 8-MMA-per-32-K burst, fused scale/bf16 epilogue and 3-blocks/CU occupancy +// target are the same as the generic kernel. ONLY the B operand path changes: +// * staging (iteration 2): lane = b_kq*16 + b_row4 with b_row4 = lane & 15 +// (n-row inside a 16-row group) and b_kq = (lane >> 4) & 3 (16-byte +// k-chunk of the 64-K stage); wave w owns n-rows w*32 + b_row4 and +// w*32 + b_row4 + 16. Consecutive lanes read consecutive 16-byte chunks of +// packed16[(k/16)*N + n] (4 fully coalesced 256-byte runs per wavefront), +// and each chunk commits as two alignment-provable 8-byte qword stores into +// the bank-safe 72-byte-stride LDS tile at n_local*72 + 16*b_kq. +// * consume: load_b_frag8 reads each lane's 8 B bytes as ONE ds_read2_b64 +// (vs 8 ds_read_u8 at stride 132); with the 18-word row stride the +// 16-lane phase spreads over 16 distinct even banks + their odd +// neighbours (the 20-word stride left 8), removing the measured 6.0 bank +// conflicts per LDS read instruction. +// LDS/block = 64x68 + 128x72 = 13,568 B (fits the single-buffer 48 KiB +// headroom rule); grid = (2560 / 128) x (4096 / 64) = 20 x 64 = 1280 blocks. __global__ __launch_bounds__(kThreadsPerBlock) void w8a8_dumma_prefill_64x128_packedb_kernel( const int8_t* __restrict__ x_q, @@ -677,12 +425,13 @@ w8a8_dumma_prefill_64x128_packedb_kernel( const int lane = tid % kWaveSize; const int wave_row = wave >> 1; // 0 or 1 const int wave_col = wave & 1; // 0 or 1 - const int m0 = static_cast(blockIdx.y) * kBlockM128; - const int n0 = static_cast(blockIdx.x) * kBlockN128; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; - __shared__ __align__(16) int8_t a_tile[kBlockM128 * kAStride128]; - // n-major B tile: 128 n-rows of 80 bytes (64 data + 16 pad). - __shared__ __align__(16) int8_t b_tile[kBlockN128 * kPackedBStride]; + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + // Iteration 2: n-major B tile, 128 n-rows of 72 bytes (64 data + 8 pad), + // declared as qwords so each 8-byte chunk commit is alignment-provable. + __shared__ __align__(16) int64_t b_tile[kBlockN * kPackedBWords]; DUFragment a_frag0, a_frag1; @@ -699,173 +448,167 @@ w8a8_dumma_prefill_64x128_packedb_kernel( du_fill_fragment(acc12, 0); du_fill_fragment(acc13, 0); - const int num_stages = k / kStageK128; + const int num_stages = k / kStageK; - // A staging mapping unchanged (row-major x_q, odd-word stride 68). + // A staging mapping unchanged. Iteration 2 B staging mapping for the chunked + // packed layout: lane = b_kq*16 + b_row4 with b_row4 = lane & 15 (n-row + // inside a 16-row group) and b_kq = (lane >> 4) & 3 (16-byte k-chunk of the + // 64-K stage); wave w owns n-rows w*32 + b_row4 and w*32 + b_row4 + 16. const int a_row = tid >> 2; const int a_col16 = (tid & 3) << 4; - // B staging mapping for packed[n][k]: thread tid owns the 16-byte k-run - // (b_k16..b_k16+15) of n-rows b_n and b_n + 64. - const int b_n = tid >> 2; - const int b_k16 = (tid & 3) << 4; + const int b_row4 = lane & 15; + const int b_kq = (lane >> 4) & 3; + const int b_n0 = wave * 32 + b_row4; + const int b_n1 = b_n0 + 16; + const int b_kqo = b_kq << 1; // qword offset of this 16-byte chunk + + // Iteration 18 (consolidate round, final HIP-only polish): every staging + // address inside the 64-stage loop is a loop-invariant base plus a constant + // per-stage byte stride, and every per-stage offset fits in 32 bits. + // * A: the row offset g_row*k (int64) is formed once (iteration 16); the + // stage-dependent part is a 32-bit running offset advanced by kStageK. + // * B: ((kg + b_kq)*n + n0 + b_n) << 4 for stage s is + // ((b_kq*n + n0 + b_n) << 4) + s*(kStageK/16)*n*16. The base is at most + // (3*2560 + 2559 + 127) << 4 = 165,856 and the running offset at most + // 63*163,840 + 165,856 = 10,487,776, both inside int32, so the compiler + // no longer emits a 64-bit v_mad_i64_i32 multiply per B operand per + // stage; only a 32-bit add. The two B pointers differ by exactly + // (b_n1 - b_n0) << 4 = 256 bytes, so one running offset serves both. + // * the LDS commit destinations are loop-invariant and hoisted as well. + // Same bytes from the same addresses into the same registers and LDS slots: + // every fragment byte and every int32 accumulator is bit-identical. + const int g_row = m0 + a_row; + const bool a_row_ok = (g_row < m); + const int64_t a_row_off = static_cast(g_row) * k; + int a_stage_off = a_col16; // s1 + a_col16, advanced by kStageK + const int a_step = kStageK; + int b_off = (b_kq * n + n0 + b_n0) << 4; // 32-bit packed-B byte offset + const int b_off_n1 = b_off + (16 << 4); + const int b_step = ((kStageK >> 4) * n) << 4; + int32_t* const adst = reinterpret_cast( + a_tile + a_row * kAStride + a_col16); + int64_t* const bdst0 = b_tile + b_n0 * kPackedBWords + b_kqo; + int64_t* const bdst1 = b_tile + b_n1 * kPackedBWords + b_kqo; - // VGPR payload for one full 64-K stage (3 x int4 = 12 VGPR). int4 vA, vB0, vB1; // Prologue: load stage 0 into registers, commit it to the single buffers, // and make it visible before the first burst. { - const int g_row = m0 + a_row; - vA = (g_row < m) - ? *reinterpret_cast(x_q + g_row * k + a_col16) - : int4{0, 0, 0, 0}; - vB0 = *reinterpret_cast( - packed_b + (n0 + b_n) * k + b_k16); - vB1 = *reinterpret_cast( - packed_b + (n0 + b_n + kBlockN128 / 2) * k + b_k16); + vA = a_row_ok ? *reinterpret_cast(x_q + a_row_off + a_stage_off) + : int4{0, 0, 0, 0}; + // Chunk c = (k/16)*n + n_idx holds packed_b[c*16 .. +15] = + // weight[16*(c/n) .. 16*(c/n)+15][c%n]; stage 0 uses k-chunks 0..3. + vB0 = *reinterpret_cast(packed_b + b_off); + vB1 = *reinterpret_cast(packed_b + b_off_n1); // Commit A as int32 stores (68-byte rows are 4 mod 16 -> no b128). - int32_t* adst = reinterpret_cast( - a_tile + a_row * kAStride128 + a_col16); adst[0] = vA.x; adst[1] = vA.y; adst[2] = vA.z; adst[3] = vA.w; - // Commit B as 16-byte vector stores: b_n*80 + b_k16 is 0 mod 16. - *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB0; - *reinterpret_cast( - b_tile + (b_n + kBlockN128 / 2) * kPackedBStride + b_k16) = vB1; + // Commit B as two 8-byte qword stores per 16-byte chunk (72-byte rows are + // 8 mod 16, so no single b128); n_local*72 + 16*b_kq is a multiple of 8. + bdst0[0] = pack_lo32_pair(vB0.x, vB0.y); + bdst0[1] = pack_lo32_pair(vB0.z, vB0.w); + bdst1[0] = pack_lo32_pair(vB1.x, vB1.y); + bdst1[1] = pack_lo32_pair(vB1.z, vB1.w); __syncthreads(); } for (int s = 0; s < num_stages; ++s) { // Consume the single buffer: each wave does eight m16n16k32 MMAs per - // 32-K step over its 32x64 quadrant. + // 32-K step over its own 32x64 quadrant. const int local_row = wave_row * 32; const int local_col = wave_col * 64; - const int8_t* abase = a_tile + local_row * kAStride128; - const int8_t* bbase = b_tile + local_col * kPackedBStride; - // Iteration 18 (fragment-pipeline round): software-pipeline the kk=32 - // fragment loads UNDER the kk=0 MMAC burst. The accepted kernel's exact - // ISA already prefetches the kk=32 A fragments this way, but its four - // kk=32 B ds_read_b64 loads are issued only AFTER MMACs 0-7, with - // s_waitcnt lgkmcnt(4..1) interleaved into MMACs 8-11 (16+ lgkmcnt waits - // per stage; PMC lds_wait 6,560,572 = 1.18x lds_instructions). Loading - // both 32-K steps' operands before the first v_mmac concentrates the LDS - // latency into one wait and leaves the 16-MMAC stream free of lgkmcnt - // waits. Same 12 ds_read instructions per stage, same fragment values, - // same int32 k0-outer/kk-inner accumulation order (the MMAC sequence per - // output tile is unchanged) -> bit-identical results. - du_load_matrix_sync(a_frag0, abase, kAStride128); - du_load_matrix_sync(a_frag1, abase + kTileM * kAStride128, kAStride128); - load_b_frag8(b_frag0, bbase, kPackedBStride); - load_b_frag8(b_frag1, bbase + kTileN * kPackedBStride, kPackedBStride); - load_b_frag8(b_frag2, bbase + 2 * kTileN * kPackedBStride, kPackedBStride); - load_b_frag8(b_frag3, bbase + 3 * kTileN * kPackedBStride, kPackedBStride); - DUFragment - a2_frag0, a2_frag1; - DUFragment - b2_frag0, b2_frag1, b2_frag2, b2_frag3; - du_load_matrix_sync(a2_frag0, abase + kTileK, kAStride128); - du_load_matrix_sync( - a2_frag1, abase + kTileM * kAStride128 + kTileK, kAStride128); - load_b_frag8(b2_frag0, bbase + kTileK, kPackedBStride); - load_b_frag8( - b2_frag1, bbase + kTileN * kPackedBStride + kTileK, kPackedBStride); - load_b_frag8( - b2_frag2, bbase + 2 * kTileN * kPackedBStride + kTileK, - kPackedBStride); - load_b_frag8( - b2_frag3, bbase + 3 * kTileN * kPackedBStride + kTileK, - kPackedBStride); - // kk = 0 burst (eight m16n16k32 MMAs). - du_mma_sync(acc00, a_frag0, b_frag0, acc00); - du_mma_sync(acc01, a_frag0, b_frag1, acc01); - du_mma_sync(acc02, a_frag0, b_frag2, acc02); - du_mma_sync(acc03, a_frag0, b_frag3, acc03); - du_mma_sync(acc10, a_frag1, b_frag0, acc10); - du_mma_sync(acc11, a_frag1, b_frag1, acc11); - du_mma_sync(acc12, a_frag1, b_frag2, acc12); - du_mma_sync(acc13, a_frag1, b_frag3, acc13); - // kk = 32 burst (same accumulators, same MMAC order as the accepted - // two-iteration kk loop). - du_mma_sync(acc00, a2_frag0, b2_frag0, acc00); - du_mma_sync(acc01, a2_frag0, b2_frag1, acc01); - du_mma_sync(acc02, a2_frag0, b2_frag2, acc02); - du_mma_sync(acc03, a2_frag0, b2_frag3, acc03); - du_mma_sync(acc10, a2_frag1, b2_frag0, acc10); - du_mma_sync(acc11, a2_frag1, b2_frag1, acc11); - du_mma_sync(acc12, a2_frag1, b2_frag2, acc12); - du_mma_sync(acc13, a2_frag1, b2_frag3, acc13); - // All consumption of the single buffer is complete only after every - // wavefront passes this barrier; the stage s+1 stores below overwrite - // that same buffer, so they must wait for it. + const int8_t* abase = a_tile + local_row * kAStride; + const int8_t* bbase = + reinterpret_cast(b_tile) + local_col * kPackedBStride; +#pragma unroll + for (int kk = 0; kk < kStageK; kk += kTileK) { + du_load_matrix_sync(a_frag0, abase + kk, kAStride); + du_load_matrix_sync(a_frag1, abase + kTileM * kAStride + kk, kAStride); + load_b_frag8(b_frag0, bbase + kk, kPackedBStride, lane); + load_b_frag8( + b_frag1, bbase + kTileN * kPackedBStride + kk, kPackedBStride, lane); + load_b_frag8( + b_frag2, bbase + 2 * kTileN * kPackedBStride + kk, kPackedBStride, + lane); + load_b_frag8( + b_frag3, bbase + 3 * kTileN * kPackedBStride + kk, kPackedBStride, + lane); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc02, a_frag0, b_frag2, acc02); + du_mma_sync(acc03, a_frag0, b_frag3, acc03); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc12, a_frag1, b_frag2, acc12); + du_mma_sync(acc13, a_frag1, b_frag3, acc13); + } + // Every wavefront must finish reading the single buffer before the next + // stage overwrites it. __syncthreads(); - // Load stage s+1 and commit it into the single buffer (fully exposed - // global-load latency, exactly as in the accepted iteration-3 kernel). + // Load stage s+1 from global and commit it into the same buffers. The + // `s + 1 < num_stages` guard is block-uniform, so the barrier below is + // reached by every thread in the block. if (s + 1 < num_stages) { - const int s1 = (s + 1) * kStageK128; - const int g_row = m0 + a_row; - vA = (g_row < m) - ? *reinterpret_cast(x_q + g_row * k + s1 + a_col16) + // Iteration 18: advance the hoisted 32-bit staging offsets by the + // compile-time per-stage strides instead of re-forming the full row / + // chunk index (and its 64-bit multiply) in every one of the 64 stages. + a_stage_off += a_step; + b_off += b_step; + vA = a_row_ok + ? *reinterpret_cast(x_q + a_row_off + a_stage_off) : int4{0, 0, 0, 0}; - vB0 = *reinterpret_cast( - packed_b + (n0 + b_n) * k + s1 + b_k16); - vB1 = *reinterpret_cast( - packed_b + (n0 + b_n + kBlockN128 / 2) * k + s1 + b_k16); - int32_t* adst = reinterpret_cast( - a_tile + a_row * kAStride128 + a_col16); + vB0 = *reinterpret_cast(packed_b + b_off); + vB1 = *reinterpret_cast(packed_b + b_off + (16 << 4)); adst[0] = vA.x; adst[1] = vA.y; adst[2] = vA.z; adst[3] = vA.w; - *reinterpret_cast(b_tile + b_n * kPackedBStride + b_k16) = vB0; - *reinterpret_cast( - b_tile + (b_n + kBlockN128 / 2) * kPackedBStride + b_k16) = vB1; - // Make the stage s+1 stores visible to every wavefront before the - // next burst. + bdst0[0] = pack_lo32_pair(vB0.x, vB0.y); + bdst0[1] = pack_lo32_pair(vB0.z, vB0.w); + bdst1[0] = pack_lo32_pair(vB1.x, vB1.y); + bdst1[1] = pack_lo32_pair(vB1.z, vB1.w); + // Make the stage s+1 stores visible before the next burst. __syncthreads(); } } const int base_row = m0 + wave_row * 32; const int base_col = n0 + wave_col * 64; - // Iteration 6 (epilogue round): coalesced stores -- register 4x4 - // transpose (8 shfl_xor -> ds_bpermute + 8 v_cndmask, no staging tile) - // then ONE 8-byte store per lane per fragment (4x fewer VMEM store - // instructions, 100% store sector utilization vs 25% for the scalar - // 2-byte stores). - store_prefill_fragment_coalesced( - acc00, x_scale, weight_scale, out, m, n, - base_row, base_col, lane); - store_prefill_fragment_coalesced( - acc01, x_scale, weight_scale, out, m, n, - base_row, base_col + kTileN, lane); - store_prefill_fragment_coalesced( - acc02, x_scale, weight_scale, out, m, n, - base_row, base_col + 2 * kTileN, lane); - store_prefill_fragment_coalesced( - acc03, x_scale, weight_scale, out, m, n, - base_row, base_col + 3 * kTileN, lane); - store_prefill_fragment_coalesced( - acc10, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col, lane); - store_prefill_fragment_coalesced( - acc11, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col + kTileN, lane); - store_prefill_fragment_coalesced( - acc12, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col + 2 * kTileN, lane); - store_prefill_fragment_coalesced( - acc13, x_scale, weight_scale, out, m, n, - base_row + kTileM, base_col + 3 * kTileN, lane); + store_prefill_fragment( + acc00, x_scale, weight_scale, out, m, n, base_row, base_col, lane); + store_prefill_fragment( + acc01, x_scale, weight_scale, out, m, n, base_row, base_col + kTileN, + lane); + store_prefill_fragment( + acc02, x_scale, weight_scale, out, m, n, base_row, + base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc03, x_scale, weight_scale, out, m, n, base_row, + base_col + 3 * kTileN, lane); + store_prefill_fragment( + acc10, x_scale, weight_scale, out, m, n, base_row + kTileM, base_col, + lane); + store_prefill_fragment( + acc11, x_scale, weight_scale, out, m, n, base_row + kTileM, + base_col + kTileN, lane); + store_prefill_fragment( + acc12, x_scale, weight_scale, out, m, n, base_row + kTileM, + base_col + 2 * kTileN, lane); + store_prefill_fragment( + acc13, x_scale, weight_scale, out, m, n, base_row + kTileM, + base_col + 3 * kTileN, lane); } -// Generic scalar int8/int32 fallback for unmatched shapes and small-M API -// cases. One output element per grid-stride step; exact int32 accumulation -// (K <= 4096 keeps the int8 dot well within int32 range), then the fused -// float scale and bf16 store. -__global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( +// Generic scalar int8/int32 fallback for every unmatched (m, n, k) and for +// small-M API cases (m < 128, including the paired M=2 shape). One output +// element per grid-stride step; exact int32 accumulation (the assigned K +// keeps the int8 dot well inside int32 range), then the fused float scale and +// bf16 store. +__global__ __launch_bounds__(kFallbackThreads) void w8a8_scalar_gemm_kernel( const int8_t* __restrict__ a, const int8_t* __restrict__ b, const float* __restrict__ x_scale, @@ -883,19 +626,22 @@ __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 (packing round): the exact assigned shape (K=4096, - // N=2560) uses the packed n-major layout packed[n][k] (set up once by - // launch_pack_w8a8_weight); every other (k, n) keeps the logical - // [K, N] row-major layout. The branch is grid-uniform per launch. + // Iteration 2 (operand-reuse round): the exact assigned shape + // (K=4096, N=2560) is packed by launch_pack_w8a8_weight into the chunked + // packed16[(k/16) * N + n] layout, so this fallback must decode that layout + // for the paired M=2 API shape: element (kk, col) lives at + // packed16[((kk >> 4) * n + col) * 16 + (kk & 15)]. Every other (k, n) + // keeps the logical [K, N] row-major layout. Grid-uniform per launch. const bool packed_b = (k == kPackedBK && n == kPackedBN); - const int8_t* b_col = - packed_b ? b + static_cast(col) * k : b + col; + const int8_t* b_col = b + col; int32_t acc = 0; for (int kk = 0; kk < k; ++kk) { - const int32_t bw = packed_b - ? static_cast(b_col[kk]) - : static_cast( - b_col[static_cast(kk) * n]); + const int32_t bw = + packed_b + ? static_cast( + b[(static_cast(kk >> 4) * n + col) * 16 + + (kk & 15)]) + : static_cast(b_col[static_cast(kk) * n]); acc += static_cast(a_row[kk]) * bw; } const float scaled = static_cast(acc) * @@ -906,7 +652,7 @@ __global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( // Identity device-to-device byte copy (pack_weight bootstrap, valid for any // (K, N)). -__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( +__global__ __launch_bounds__(kPackThreads) void w8a8_pack_identity_i8_kernel( const int8_t* __restrict__ src, int8_t* __restrict__ dst, int64_t count) { @@ -920,7 +666,7 @@ __global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( } // Identity device-to-device float copy for the weight scales. -__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( +__global__ __launch_bounds__(kPackThreads) void w8a8_pack_identity_f32_kernel( const float* __restrict__ src, float* __restrict__ dst, int64_t count) { @@ -933,45 +679,54 @@ __global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( } } -// Iteration 5 (packing round): one-time n-major transpose of the logical -// [K, N] row-major weight into packed[n][k] for the exact assigned shape -// (K=4096, N=2560). Runs only from launch_pack_w8a8_weight, outside the -// timed region and outside Graph capture. The packed buffer keeps the same -// byte count (k*n) as the identity pack, so allocations and graph-stable -// addresses are unchanged. Element (k, n) of the logical weight lands at -// packed[n * K + k]; each thread copies one 16-byte k-run (coalesced read -// side; the strided write side is off the critical path). The guard -// guarantees n % 16 == 0, so every 16-byte chunk lies inside one logical -// row and the transpose is exact. -__global__ __launch_bounds__(256) void w8a8_pack_nmajor_b_kernel( +// Iteration 2 (operand-reuse round): one-time 16-byte-chunk packing of the +// logical [K, N] row-major weight into packed16[(k/16) * N + n], where chunk +// c holds packed16[c*16 + i] = weight[(c / N) * 16 + i][c % N] for i in 0..15. +// Runs only from launch_pack_w8a8_weight, outside the timed region and outside +// Graph capture. The packed buffer keeps the same byte count (k*n) as the +// identity pack, so allocations and graph-stable addresses are unchanged. Each +// thread owns one 16-byte chunk: the 16 strided byte gathers (stride N) are +// fully coalesced across the wavefront's threads and the 16-byte store is +// contiguous, so the pack stays bandwidth-clean. The guard requires +// k % 16 == 0, true for the exact assigned shape. +__global__ __launch_bounds__(kPackThreads) void w8a8_pack_k16chunk_b_kernel( const int8_t* __restrict__ src, int8_t* __restrict__ dst, int k, int n) { - const int64_t chunks = (static_cast(k) * n) >> 4; + const int64_t chunks = (static_cast(k) >> 4) * n; const int64_t stride = static_cast(gridDim.x) * blockDim.x; for (int64_t c = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; c < chunks; c += stride) { - const int64_t off = c << 4; - const int kk = static_cast(off / n); - const int col = static_cast(off - static_cast(kk) * n); - const int4 v = *reinterpret_cast(src + off); - const int8_t* bytes = reinterpret_cast(&v); + const int64_t kg = c / n; + const int n_idx = static_cast(c - kg * n); + const int8_t* scol = src + (kg << 4) * n + n_idx; + int8_t* d = dst + (c << 4); #pragma unroll for (int i = 0; i < 16; ++i) { - // Logical element (k = kk, n = col + i) lands at packed[(col+i)*K+kk]. - dst[static_cast(col + i) * k + kk] = bytes[i]; + d[i] = scol[static_cast(i) * n]; } } } +inline int64_t capped_blocks(int64_t elements, int threads) { + int64_t blocks = (elements + threads - 1) / threads; + if (blocks > 4096) { + blocks = 4096; + } + if (blocks < 1) { + blocks = 1; + } + return blocks; +} + } // namespace // Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the -// caller-provided stream only; never allocates, synchronizes, or touches the -// default stream. +// caller-provided stream only; never allocates, synchronizes, packs, or +// touches the default stream. extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -991,16 +746,18 @@ extern "C" void launch_w8a8_gemm( } auto* out_ptr = static_cast(out); - // Iteration 5 (packing round): exact assigned shape (K=4096, N=2560) - // routes to the packed n-major B variant of the 64x128 kernel. The guard - // is exact (m >= 128, n == 2560, k == 4096) and sits BEFORE the generic - // 64x128 path; every other shape keeps its existing path (generic 64x128 - // for n % 128 == 0, 64x64 for n % 64 == 0, scalar fallback otherwise) and - // the identity-packed layout. + // Iteration 1/2 (packed_weight family): the exact assigned shape + // (K=4096, N=2560) routes to the packed 16-byte-chunk B variant of the + // 64x128 kernel, whose B operand comes from the one-time packed16[(k/16)*N+n] + // layout prepared by launch_pack_w8a8_weight. The guard is exact (m >= 128, + // n == 2560, k == 4096) and sits BEFORE the generic path; every other + // (m, n, k) keeps its existing path and the identity-packed layout. The + // m >= 128 guard keeps the paired M=2 API shape on the scalar fallback, + // which decodes the chunked layout for this exact (k, n). if (m >= 128 && n == kPackedBN && k == kPackedBK) { const dim3 grid( - static_cast(n / kBlockN128), - static_cast((m + kBlockM128 - 1) / kBlockM128)); + static_cast(n / kBlockN), + static_cast((m + kBlockM - 1) / kBlockM)); const dim3 block(static_cast(kThreadsPerBlock)); hipLaunchKernelGGL( w8a8_dumma_prefill_64x128_packedb_kernel, @@ -1019,44 +776,17 @@ extern "C" void launch_w8a8_gemm( return; } - // Native INT8 DUMMA m16n16k32 prefill path (2-D macro-tile 64x128, - // single-buffered 64-K K stage; pipeline-round control arm) for every - // large-M shape with compatible geometry (covers the assigned M=4096, - // N=2560, K=4096 shape: 2560 % 128 == 0). The 128-K guard is written - // explicitly so the dispatch condition is bit-identical to the accepted - // iteration-1 kernel even though the internal K stage is now 64 - // (single-buffered). - if (m >= 128 && (n % kBlockN128) == 0 && (k % 128) == 0) { - const dim3 grid( - static_cast(n / kBlockN128), - static_cast((m + kBlockM128 - 1) / kBlockM128)); - const dim3 block(static_cast(kThreadsPerBlock)); - hipLaunchKernelGGL( - w8a8_dumma_prefill_64x128_kernel, - grid, - block, - 0, - stream, - a, - b, - x_scale, - weight_scale, - out_ptr, - m, - n, - k); - return; - } - - // Native INT8 DUMMA m16n16k32 prefill path (64x64 tile) for large-M shapes - // whose N is not a multiple of 128 but is a multiple of 64. + // Native INT8 DUMMA m16n16k32 prefill path for every large-M shape with + // compatible geometry (covers the assigned M=4096, N=2560, K=4096 shape: + // 2560 % 128 == 0 and 4096 % 64 == 0). The explicit m >= 128 guard keeps + // the small-M API shapes on the scalar fallback. if (m >= 128 && (n % kBlockN) == 0 && (k % kStageK) == 0) { const dim3 grid( static_cast(n / kBlockN), static_cast((m + kBlockM - 1) / kBlockM)); const dim3 block(static_cast(kThreadsPerBlock)); hipLaunchKernelGGL( - w8a8_dumma_prefill_64x64_kernel, + w8a8_dumma_prefill_64x128_kernel, grid, block, 0, @@ -1074,15 +804,7 @@ extern "C" void launch_w8a8_gemm( // Scalar int8/int32 fallback for unmatched (m, n, k) and small-M API cases. const int64_t total = static_cast(m) * n; - constexpr int kFallbackThreads = 256; - int64_t blocks = (total + kFallbackThreads - 1) / kFallbackThreads; - if (blocks > 4096) { - blocks = 4096; - } - if (blocks < 1) { - blocks = 1; - } - const dim3 grid(static_cast(blocks)); + const dim3 grid(static_cast(capped_blocks(total, kFallbackThreads))); const dim3 block(kFallbackThreads); hipLaunchKernelGGL( w8a8_scalar_gemm_kernel, @@ -1100,12 +822,19 @@ extern "C" void launch_w8a8_gemm( k); } -// Stable host launch symbol consumed by csrc/bindings.cpp. For the exact -// assigned shape (K=4096, N=2560) it performs the one-time n-major B pack -// (logical [K, N] -> packed[n][k]) outside the timed region and outside -// Graph capture; every other (K, N) keeps the identity device-to-device -// copy, valid for every (K, N). The packed buffer size is k*n in both -// cases, so the allocation and graph-stable packed layout are unchanged. +// Stable host launch symbol consumed by csrc/bindings.cpp. Iteration 2 +// (operand-reuse round) behaviour: for the exact assigned pair (K=4096, +// N=2560) it performs the one-time 16-byte-chunk packing +// packed16[(k/16) * N + n] (outside the timed region and Graph capture) and +// for every other (K, N) the identity device-to-device copy of the logical +// [K, N] int8 weight. The N fp32 weight scales are copied identically in both +// cases. Both layouts keep the same byte count as the raw weight, so +// allocation and graph-stable addresses are unchanged. +// +// Iteration 3 (fix_build round) correction: this host call previously +// referenced the removed w8a8_pack_nmajor_b_kernel symbol, which was the only +// compile error of the iteration-2 candidate. It now binds the defined +// w8a8_pack_k16chunk_b_kernel; nothing else in this file changed. extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -1114,19 +843,22 @@ extern "C" void launch_pack_w8a8_weight( int k, int n, hipStream_t stream) { - constexpr int kPackThreads = 256; const int64_t weight_count = static_cast(k) * n; if (weight_count > 0) { if (k == kPackedBK && n == kPackedBN) { + // Iteration 2 (operand-reuse round): exact assigned shape + // (K=4096, N=2560) -> one-time 16-byte-chunk packing + // packed16[(k/16) * N + n]. Same k*n byte count as the identity pack, so + // the allocation and the graph-stable packed layout are unchanged. + // Outside the timed region. const int64_t chunks = weight_count >> 4; - int64_t blocks = (chunks + kPackThreads - 1) / kPackThreads; - if (blocks > 4096) { - blocks = 4096; - } + const dim3 grid( + static_cast(capped_blocks(chunks, kPackThreads))); + const dim3 block(kPackThreads); hipLaunchKernelGGL( - w8a8_pack_nmajor_b_kernel, - dim3(static_cast(blocks)), - dim3(kPackThreads), + w8a8_pack_k16chunk_b_kernel, + grid, + block, 0, stream, raw_weight, @@ -1134,14 +866,13 @@ extern "C" void launch_pack_w8a8_weight( k, n); } else { - int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; - if (blocks > 4096) { - blocks = 4096; - } + const dim3 grid( + static_cast(capped_blocks(weight_count, kPackThreads))); + const dim3 block(kPackThreads); hipLaunchKernelGGL( w8a8_pack_identity_i8_kernel, - dim3(static_cast(blocks)), - dim3(kPackThreads), + grid, + block, 0, stream, raw_weight, @@ -1150,20 +881,18 @@ extern "C" void launch_pack_w8a8_weight( } } if (n > 0) { - int64_t blocks = (static_cast(n) + kPackThreads - 1) / - kPackThreads; - if (blocks > 4096) { - blocks = 4096; - } + const dim3 grid( + static_cast(capped_blocks(n, kPackThreads))); + const dim3 block(kPackThreads); hipLaunchKernelGGL( w8a8_pack_identity_f32_kernel, - dim3(static_cast(blocks)), - dim3(kPackThreads), + grid, + block, 0, stream, weight_scale, packed_weight_scale, - n); + static_cast(n)); } } // @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/o_proj.hip 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..3c5fe22e 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,758 +1,1099 @@ -// @@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=a1271e87052a802cf1a620ad5f58ba9d526103cd added=2026-09-17 +// median_us=14.43 p90_us=14.44 speedup=5.183 baseline_us=74.78 +// source=hy3-tp8-m16-test9-16-619177f8 +// Initial correctness-first INT8 W8A8 GEMM for gfx928 (K500SM_AI). // -// 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. +// Bootstrap strategy (iteration 0): one plain scalar kernel computes the +// complete int32 dot product for every output element. There is no DUMMA, no +// split-K, no LDS staging, no inline assembly, no software pipelining and no +// autotuning. The only goals are: +// * compile fast and portably for gfx928, +// * launch on PyTorch's current HIP stream with no host/device sync, +// * be exact against bf16((A.float() @ B.float()) * x_scale * w_scale.T). // -// 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. +// Layout contract for this iteration (identity pack): +// a : int8 [M, K] row-major (x_q) +// b : int8 [K, N] row-major (packed_weight) +// x_scale : fp32 [M] +// weight_scale : fp32 [N] +// out : bf16 [M, N] row-major // -// 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): +// launch_pack_w8a8_weight is an identity device-to-device copy of both the +// int8 weight and the fp32 weight scale. Later Parallel explore rounds may +// change that packing and the matching GEMM interpretation. // -// 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 1 established the minimal native gfx928 DUMMA INT8 m16n16k32 +// tile: ONE 64-lane wavefront per block, one 16x16 output tile per block, +// grid = N/16 = 80 one-wave blocks, blockDim 64, A/B fragments loaded +// directly from row-major global at every 32-K step, explicit int32 +// accumulation, and a single-wave LDS epilogue. It reached 119.72 us +// (official best), well behind the 74.78 us Triton graph baseline. // -// Both kernels launch on the caller's stream inside the timed Graph, so the -// combine cost is included in the measured operator latency. +// Iteration 2 (architecture: launch-geometry exploration). +// PMC/ISA evidence for that 80-block grid: +// * grid_blocks = 80 for device_cu_count = 120 -> 40 CUs idle and at most +// one resident wave per busy CU, far below the two-blocks-per-CU +// latency-hiding target; +// * the disassembled K-loop is a serial chain -- 16 x global_load_ubyte, +// a run of s_waitcnt vmcnt(N...) drains, then exactly one v_mmac -- so +// each wave stalls on global latency every 32-K step with no other wave +// on the CU to cover it; +// * counter-derived HBM bandwidth is ~45 GB/s, i.e. the device is nowhere +// near memory-bandwidth bound: the kernel is latency/underfill bound. +// The N dimension only has N/16 = 80 tiles and M is already the full 16, so +// the single geometry that can raise independent block count is split-K. +// This round therefore keeps the complete one-wave / one-N-tile geometry +// (waves_per_block = 1, adjacent_n_tiles = 1, no cross-wave barrier, direct +// global fragment loads, one v_mmac per 32-K step) and adds a 3-way split of +// the 128 stage-aligned (k/32) K-steps: +// * grid = (N/16) * 3 = 240 one-wave blocks for 120 CUs => exactly two +// independent blocks per CU, the mandated latency-hiding target; +// * split s owns stages [s*T/S, (s+1)*T/S) with T = k/32 = 128 and S = 3, +// i.e. 42/43/43 stages of 32 K each -- every range stays 32-aligned, so +// the stage-alignment constraint and the workspace constraint (S <= 16 +// planes) both hold; +// * each block accumulates its K sub-range exactly in int32 and stores the +// 16x16 partial into plane s of the caller-provided workspace +// (int32 [S][M][N]); a separate 128-thread combine kernel sums the S +// int32 planes (exact, order-free integer addition), applies +// static_cast(acc) * x_scale[row] * weight_scale[col] in the same +// order as the scalar fallback, and writes bf16 -- so the output stays +// bit-identical while the combine is a cheap L2-resident second launch. // -// 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 3 (architecture: split-K occupancy sweep). +// The iteration-2 falsifiable claim held: the 240 one-wave blocks (2 blocks +// per CU) cut the wall to 47.81 us median, 1.564x over the 74.78 us Triton +// Graph baseline, so the 80-block underfill was real. The current-best +// PMC still reports grid_blocks = 240 for device_cu_count = 120 -- only TWO +// of the four SIMDs of every CU ever hold a resident wave -- and the +// counter-derived HBM rate is still just 123.8 GB/s, far below device +// bandwidth, so the kernel remains latency/occupancy bound rather than +// bandwidth bound. The trusted CU-aligned probe set for this shape is +// {2,3,4,5,6,8,9,12}; with N/16 = 80 one-wave N tiles the grid is 80*S, so +// the CU-aligned splits are the multiples of 3 (80*S must be a multiple of +// 120). split-K=2 yields only 160 blocks = 1.33 blocks/CU -- still below +// the mandated two-blocks-per-CU target -- so it is rejected, and the next +// CU-aligned candidate is split-K=6: +// * grid = (N/16) * 6 = 480 one-wave blocks for 120 CUs => exactly FOUR +// independent one-wave blocks per CU, i.e. one resident wave on each of +// the four SIMDs, instead of 2 of 4 in iteration 2; +// * the 128 32-aligned K-stages are partitioned 21/21/22/21/21/22 (rows +// 672/672/704/672/672/704) -- disjoint, ascending, exact cover of +// [0, 4096), every boundary a multiple of 32 -- so the stage-alignment +// and exact-int32 contracts are unchanged; +// * workspace need = 6 * M * N * 4 = 491,520 B, inside the contract's +// 16-plane 1,310,720 B allocation; the partial kernel, its plane +// indexing and the combine+scale kernel are already generic in split_k, +// so the only source change is the launcher geometry constant +// kSplitKTarget 3 -> 6 (the combine stays inside the timed Graph). +// Falsifiable claim: if the two idle SIMDs per CU are the limiter, the +// per-block serial 32-K chain halves (43 -> ~21 steps) while the number of +// resident waves doubles, so the official Graph median/P90 must drop +// materially below the 47.81/47.85 us iteration-2 best (expected ~34-40 us, +// with the extra 2 partial planes and the slightly longer combine cost +// included). If it stays flat or regresses, the wall is the per-step +// load->wait->reassemble->mma latency or per-CU issue rather than SIMD +// occupancy, and the next round must attack the 16 direct global byte loads +// per stage instead of the split factor. // -// 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 4 (accepted, architecture/pipeline: multi-N-tile reuse). +// The iteration-3 result falsified its own prediction: moving kSplitKTarget +// 3 -> 6 cut the wall only 47.81 -> 45.08 us (-5.7%) instead of the +// predicted ~34-40 us, so doubling the resident one-wave blocks per CU did +// NOT remove the limiter. The reason is an invariance the split factor +// cannot break: the total m16n16k32 mma count is fixed at (N/16)*(K/32) = +// 80*128 = 10,240, so per device SIMD (480 of them) the number of serial +// load->wait->reassemble->mma steps is 10,240/480 = 21.3 for ANY split +// factor -- finer splits only shorten each block while multiplying the +// block count by the same amount. The measured PMC still shows the poison +// per-step pattern: 16 global_load_ubyte (8 for the A fragment + 8 for the +// single B fragment) followed by a ladder of scattered s_waitcnt vmcnt and +// exactly ONE v_mmac per 32-K step, with no cross-iteration load overlap. +// This round therefore keeps the split-K partial+combine architecture and +// the 480-block / 4-blocks-per-CU co-residency but widens every one-wave +// block to TWO adjacent 16-wide N tiles (32 N columns): +// * per 32-K step the wave loads the A fragment (16 m-rows x 32 K bytes = +// 512 B of x_q) EXACTLY ONCE and issues TWO independent du_mma_sync +// calls -- one per 16-wide N tile -- so those A bytes are REUSED across +// both output tiles (the per-output-tile A global loads and their +// wait/reassembly cost halve); +// * every B byte (k x 32 N bytes per block slice) is still read exactly +// once -- B has no reuse (each N column belongs to exactly one block); +// * the 128 32-aligned K stages are split 12 ways and each block owns TWO +// output tiles, so the grid is (N/(16*2)) * 12 = 40 * 12 = 480 one-wave +// blocks for 120 CUs = exactly 4 blocks/CU, i.e. one resident wave per +// SIMD, identical co-residency to iteration 3 (enough blocks to cover +// every CU); +// * each block's serial chain drops from ~21 single-mma steps to <= 11 +// two-mma steps (the 128 stages partition 8 x 11 + 4 x 10 = rows +// 352 x 8 + 320 x 4, disjoint, ascending, exact 32-aligned cover of +// [0, 4096)), so if the wall is the per-step global latency then the +// latency paid per m16n16k32 mma is halved; +// * workspace need = 12 * M * N * 4 = 983,040 B, inside the contract's +// 16-plane 1,310,720 B allocation; the partial kernel body, its plane +// indexing and the combine+scale kernel stay generic in split_k (the +// runtime clamp still shrinks split_k to the supplied workspace), and +// the combine stays inside the timed Graph. +// Falsifiable claim: if the per-step direct-global load->wait->mma latency +// is the limiter, the per-output A-load cost halves and the serial chain +// per block halves, so the official Graph median/P90 must drop materially +// below the 45.08/45.17 us iteration-3 best (expected ~33-39 us, with the +// six extra partial planes and the longer 12-way combine included). If it +// stays flat or regresses, the limiter is L2/HBM-side byte traffic or the +// per-step issue of the byte-granular loads themselves rather than latency +// amortization, and the next round must stage the operands in LDS. // -// 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 5 (measured and NOT accepted: B-only LDS staging). +// Staging only the B operand into a per-block LDS mirror (16-B int4 +// cooperative copy + one same-wave barrier, then du_load_matrix_sync from LDS +// at row stride 32) left the loop time unchanged and paid the serial staging +// prologue on top: the official Graph median/P90 moved 30.374/30.420 -> +// 31.309/31.350 us (+3.1%), so the round was rejected and the source was +// reverted to the iteration-4 best recorded here. The falsification is +// informative: removing two thirds of the per-step global byte loads (the 16 +// B loads of the 24) did NOT shorten the per-step chain, so the residual +// serialized operand fetch on the MMA critical path is the A half, and the +// B-only design additionally left a mixed dual-latency step (LDS B + global +// A) whose LDS reads replay the same bank 4x at the iteration-5 row pitch. // -// 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 6 (pipeline round, this round): A+B whole-slice LDS staging. +// Mandated staging-family decision, made from the L2/VMEM evidence: +// * VMEM: the exact-source PMC of the iteration-4 code object (digest +// d5784ae4...) shows the rolled K loop issuing 24 byte-granular +// global_load_ubyte per 32-K step (vmem_read 122,880 = 24 x 5,120 +// block-steps) behind an 18-entry scattered s_waitcnt vmcnt ladder, all +// on the MMA path with no cross-iteration load overlap. +// * L2: the same replay reports 61.237% L2 hits, and the per-kernel TCC/EA +// traffic ledger shows the partial kernel reading only 5,359,488 B while +// its load instructions request ~7.9 MB (A 2.62 MB + B 5.24 MB). The +// 40x-reused x_q A operand is therefore ~95% L2-resident (A HBM misses +// ~= 5,359,488 - 5,242,880 = 116,608 B) and the streaming B operand is +// the compulsory HBM traffic read exactly once. No staging family can +// change the HBM byte count: direct, A-only, B-only and A+B all fetch the +// same 5,242,880 B of B and the same ~116,608 B of A misses. The choice +// is therefore only about WHERE the operand fetch latency sits. +// * B-only was already measured (iteration 5) and regressed, so between the +// remaining families the evidence selects A+B: it is the ONLY one that +// removes the A-side global wait (which the B-only result proves is the +// serialization point) AND the B-side ladder at the same time, replacing +// both with one batched cooperative prologue of 16-B loads. +// Design (one-wave geometry, blockDim 64, unchanged): +// * Each 64-lane block stages BOTH operands of its whole K slice once: +// stage_a = the 16 m-rows x 256 slice bytes of x_q at row stride 272 +// (256 data + 16 pad for bank skew, 4,352 B) and stage_b = the +// 256 k-rows x 32 B B-strip at row stride 36 (32 data + 4 pad, 9,216 B). +// Both are EXACT byte mirrors of the global regions the direct loader +// read, filled by 16-B int4 cooperative copies (A: rows/16 = 16 chunks +// per m-row -> 256 chunks / 64 lanes = 4 chunks per lane, one 4-deep +// batched group; B: 2 chunks per k-row -> 512 chunks / 64 lanes = 8 +// chunks per lane, two 4-deep batched groups), one same-wave +// __syncthreads(), and then a K loop with ZERO per-step global loads and +// ZERO vmcnt waits: A from LDS at stride 272, B0/B1 from LDS at stride 36. +// * Pad choice is deliberate: the row_major B fragment loader reads each +// lane's bytes at (8*(lane>>4) + i)*ldm, so a 32-B ldm puts the four lane +// groups 256 B (== 0 mod 128) apart and replays one LDS bank 4x, while +// 36 B puts them 288 B (== 32 mod 128) apart -> banks b, b+8, b+16, b+24; +// the A pad 272 B (== 16 mod 128) likewise skews its 16 row reads. The +// pads are never written and never read. +// * Geometry moves 12 -> 16 splits because a whole-slice A stage does not +// fit at split 12 with 4 resident blocks/CU (A 16x368 + B 352x36 + 2,048 +// = 20,608 B -> floor(65,536/20,608) = 3 resident waves/CU, below the +// measured 4). At S = 16 the slices are uniform 8 x 32 stages = 256 rows +// (128 == 16*8), per-block LDS is 4,352 + 9,216 + 2,048 = 15,616 B -> +// floor(65,536/15,616) = 4 resident blocks/CU (62,464 B <= 65,536), and +// the grid is (N/32) x 16 = 40 x 16 = 640 one-wave blocks = 5.33/CU over +// 120 CUs, i.e. the measured 4-blocks/CU co-residency is preserved. +// The workspace need is 16 int32 planes = 1,310,720 B = the caller's full +// fixed allocation for this shape (the existing clamp still shrinks +// split_k for any smaller buffer, and a slice larger than the staging +// capacity keeps the byte-for-byte direct-global loop below). +// * Every partial is still an exact int32 ascending 32-aligned sub-sum over +// the SAME fragment bytes (the staged mirrors are exact copies), the +// combine kernel is unchanged, and no cross-wave synchronization or +// cross-block state is introduced. +// Byte ledger: HBM read/write bytes UNCHANGED (B 5,242,880 B read once + A +// misses ~116,608 B + 1,310,720 B of partial planes written/read); LDS added +// per block = 4,352 (stage_a) + 9,216 (stage_b) = 13,568 B written and +// re-read 8 times as fragments (8 x 1,536 B = 12,288 B), with the bank-skew +// pads chosen so the fragment reads are conflict-free. +// Falsifiable claim: if the per-step serialized operand fetch (not the +// compulsory HBM bytes) is the wall, the LDS-only K loop plus the batched +// 16-B staging prologue must drop the official Graph median/P90 materially +// below 30.374/30.420 us (expected ~21-27 us, combine's +4 planes included). +// If it instead stays flat or regresses, the wall at 4 waves/CU is the +// DRAM/LDS service rate or per-CU issue rather than the per-step wait ladder, +// and the next round must attack the operand stream's DRAM granularity (a +// packed [N/32][K][32] weight layout) rather than the staging family. // -// 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 7 (measured and NOT accepted: depth-1 register double buffering). +// Giving the rolled LDS-only K loop two fragment register sets moved the +// official Graph median 21.412 -> 20.938 us but the P90 21.427 -> 25.333 us +// (+18%, max 26.557), so the round was rejected on the tail and the source was +// reverted to the iteration-6 best recorded here. The lesson: at ONE resident +// wave per SIMD there is no sibling wave to absorb the tail, so the fix must +// add wave-level parallelism rather than reorder a single wave's own chain. // -// 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 8 (occupancy/resource round, this round): raise the waves-per-block +// occupancy limiter from 1 to 2. The exact-source PMC of the accepted +// iteration-6 code object (digest 4fe713e6..., grid 640 = 40 N groups x 16 +// splits, workgroup 64, arch_vgpr 80, sgpr 64, scratch 0, LDS 15,616 B/block +// -> floor(65,536/15,616) = 4 resident blocks/CU, i.e. exactly ONE resident +// wave per SIMD) shows the residual wall is per-SIMD single-wave latency +// exposure, not issue or HBM traffic: 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 +// (lds_wait_instructions 11,842, lds_bank_conflicts 133,120, vmem_read 7,680 = +// only the one-time staging loads). This round makes each block TWO 64-lane +// wavefronts (blockDim 64 -> 128, __launch_bounds__(128)) and splits the block +// 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 m16n16k32 accumulator chain, +// one du_mma per step per wave (same total MMA count). LDS per block is +// UNCHANGED (stage_a 4,352 + stage_b 9,216 + 2 x 1,024 acc tiles = 15,616 B), +// so 4 resident blocks/CU still fit but now carry 8 resident waves/CU = 2 per +// SIMD: one wave's LDS->MMA stall is covered by its co-resident sibling. The +// one-time staging is PARTITIONED, not duplicated: wave 0 copies the shared A +// slice, wave 1 copies the B strip (the iteration-6 cooperative loops +// unchanged), so exactly the same global bytes are read once per block -- no +// repeated HBM reads are traded for occupancy -- and the staging barrier +// becomes a cross-wave __syncthreads(). The only deliberately added cost is +// LDS re-reads: the shared A fragment is now loaded from the staged mirror by +// both waves (LDS bandwidth is not the wall at this issue utilization). Every +// partial int32 is the identical ascending sum over the identical staged bytes +// for its own 16-column tile, so the output stays bit-identical and the combine +// kernel, the 16-plane workspace contract, the exact (m,n,k) == (16,1280,4096) +// guard, the workspace-less direct-global path and the scalar fallback for +// every unmatched shape (paired M=2 included) are all unchanged. +// Falsifiable claim: if the per-SIMD single-wave latency exposure is the wall, +// the official Graph median/P90 must drop below 21.411954/21.427149 us +// (expected ~19.5-20.9 us median) with mismatch_count 0; if it stays flat or +// regresses, the wall at 8 waves/CU is the LDS service rate or per-CU issue +// rather than wave-level coverage. // -// 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). +// Iteration 9 (combine-restructure round, this round): attack the OTHER timed +// dispatch of the operator -- the split-K combine. Every previous round +// optimized only the 640-block partial kernel, but the operator replay runs +// TWO kernels on the same stream and the combine is a real part of the wall: +// its exact-source PMC (source digest e2a23702..., 160 x 128-thread +// one-element-per-thread blocks) profiles at 3.52 us against 19.52 us for the +// partial kernel, and the static code object's rolled plane loop is a +// DEPENDENT load -> s_waitcnt vmcnt -> int32 add chain of 16 DRAM round trips +// per output element (one 4-B global load per plane iteration) with only 160 +// blocks (~1.33 warps/SIMD) to hide it. This round restructures ONLY the +// combine kernel and its launch: +// * the guard launcher switches the combine from 160 linear 128-thread +// blocks to a 2-D grid of (n / 64) column groups x m rows = 20 x 16 = 320 +// ONE-wave 64-lane blocks, and the kernel's split_k == 16 specialization +// gives every lane exactly ONE output element (row = blockIdx.y, +// col = blockIdx.x * 64 + lane): all sixteen plane loads are issued as +// INDEPENDENT reads into a 16-int register array (one vmcnt drain, +// 16-way memory-level parallelism per lane) and only then summed in the +// unchanged ascending split order s = 0..15; +// * each per-plane wavefront load now reads 64 lanes x 4 B = 256 B +// CONTIGUOUS (four consecutive 64-B lines of one m-row) instead of four +// runs at the 5,120-B m-row stride. +// The partial kernel, the 640-block / 16-split geometry, the staged byte +// mirrors, the k-ascending int32 partials, the 16-plane 1,310,720 B workspace +// contract, the workspace-less fallback and the scalar fallback are +// byte-for-byte unchanged. Integer addition is exact and order-free +// (|acc| <= K*127*127 ~= 6.6e7 << 2^31), the fp32 scale order and the bf16 +// conversion are element-for-element identical, so every output byte is +// bit-identical. A clamped split_k < 16 still takes the unchanged generic +// combine path with its original 128-thread launch, so an undersized +// workspace cannot change behavior. +// Falsifiable claim: if the combine's per-element dependent plane-load chain +// (and its underfilled 160-block grid) is a material part of the 19.51 us +// operator wall, the official Graph median/P90 must drop materially below +// 19.510394/19.535989 us (expected ~17.3-18.9 us); if it stays flat, the +// combine was already fully hidden behind the partial kernel's drain and the +// remaining wall is entirely partial-kernel side. // -// 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 11 (DS-side staging-latency round, this round): shorten the +// one-time staging join of the accepted iteration-9/10 partial kernel. The +// exact-source PMC of that object (source digest bd963ec6..., grid 640 = 40 N +// groups x 16 splits, workgroup 128, LDS 15,616 B, arch_vgpr 80, scratch 0, +// profiled 20.159 us) shows the kernel is deeply latency-bound, not +// issue/bandwidth-bound: only 486,400 VALU + 112,640 LDS wave-instructions for +// 1,280 resident wavefronts x 8 K-steps, while vmem_read_instructions is just +// 7,680 (= the one-time staging loads, 12 int4 loads per block) and the +// block's critical path therefore still contains the operand staging prologue. +// The staging is PARTITIONED (wave 0 copies the 4-chunk-per-lane A slice, +// wave 1 copies the 8-chunk-per-lane B strip) and the cross-wave +// __syncthreads() joins both, so the join costs the SLOWER wave. Wave 0 +// already issues its whole A copy as one 4-deep batch, but wave 1's 512-chunk +// B strip (rows*2 = 256*2 = 512 chunks == 8 x 64 lanes) is emitted as two +// serialized 4-deep iterations: iteration k+1's four 16-B global loads cannot +// issue until iteration k's four dependent LDS stores have consumed their +// registers, so every block pays TWO dependent global round trips on the +// staging join instead of one. Change (one focused HIP edit, +// csrc/w8a8_gemm_hip.hip only): when the B strip is exactly the guard's +// 8-chunk-per-lane shape (b_chunks == 8 * kDummaWaveSize == 512, which the +// exact (m,n,k) == (16,1280,4096) guard always produces because rows == +// kStageMaxRows == 256), wave 1 now runs ONE straight-line 8-deep batched +// group: all eight 16-B source loads are issued before the first LDS store, +// so a single vmcnt drain covers the whole strip and the second global round +// trip leaves every block's staging join. The generic rolled 4-deep loop +// stays verbatim as the else path for any non-guard chunk count. The staged +// bytes, their LDS addresses (the same 36-B row pitch and store decomposition) +// and therefore every fragment byte, every k-ascending int32 partial and the +// combine output stay bit-identical; the 16-plane 1,310,720 B workspace +// contract, the workspace-less direct-global fallback, the exact shape guard, +// the scalar fallback for every unmatched shape (paired M=2 included), the +// combine kernel and its 320-block one-row-per-block launch, the pack identity +// and every stream/Graph property are unchanged. +// Falsifiable claim: if the block's staging join is a material part of the +// 18.980 us wall, removing the second global round trip from wave 1's B copy +// must move the official Graph median/P90 below 18.980430/19.183630 us +// (expected ~18.3-18.9 us median) with mismatch_count 0; if it stays flat or +// regresses, the staging join is already fully covered by the co-resident +// blocks and the residual wall is the LDS-only K loop. // -// 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 13 (HIP-only consolidation round, this round): halve the split-K +// partial-plane round trip, the only non-compulsory component of the +// operator's memory traffic. The exact-source PMC of the accepted iteration-11 +// object (source digest 210fbfcfde28bc26d9c1c93a8a5f91353db51c5426853e85b5701dd44fa81fa8, +// primary kernel 640 blocks x 128 threads, LDS 15,616 B, arch_vgpr 80, sgpr 48, +// scratch 0, profiled 20.0 us; combine 320 one-wave blocks, profiled 2.72 us; +// official unprofiled Graph median/P90 18.145450/18.166250 us) reports, per +// operator replay, 6,716,736 B read + 1,351,680 B written = 8,068,416 B, and +// the memory_traffic ledger splits it into 5,368,320 B for the partial kernel +// (B 5,242,880 compulsory + A + the 1,310,720 B of int32 partial planes it +// writes) and 1,348,416 B read + 40,960 B written for the combine (the same +// 1,310,720 B of planes re-read). The combine's own L2 hit rate is just +// 6.849% (l2_hits 1,597 / l2_misses 21,721) against the partial kernel's +// 54.392%, i.e. the 16 int32 planes are written through L2, evicted by the +// streaming B strip, and pulled back from DRAM by the immediately following +// combine: 2,621,440 B of the 8,068,416 B total (32.5%) is pure split-K +// bookkeeping against a compulsory A 65,536 B + B 5,242,880 B + out 40,960 B +// = 5,349,376 B. Change (one focused HIP edit, csrc/w8a8_gemm_hip.hip only): +// move kSplitKTarget 16 -> 8 so the split-K bookkeeping halves to 8 planes +// (655,360 B written + 655,360 B re-read), removing 1,310,720 B (16.2%) of the +// per-replay traffic. +// * Geometry: grid = (N/(16*T)) * 8 = 40 * 8 = 320 blocks of 128 threads; +// the 128 32-K stages partition into uniform 16-stage (512-row) slices +// (128 == 8 x 16), every boundary still a multiple of 32, so the +// stage-alignment and exact-int32 contracts are unchanged. The workspace +// need is 8 x 80 KiB = 655,360 B, well inside the contract's fixed +// 1,310,720 B / 16-plane allocation, and the existing clamp still shrinks +// split_k for a smaller buffer. +// * Staging: a 512-row slice is twice the unchanged 256-row staging +// capacity, so the kernel now stages and consumes its slice in ascending +// 256-row PHASES (exactly two at the guard). Each phase is byte-for-byte +// the accepted iteration-11 geometry: wave 0 stages the 256-chunk A phase +// as one 4-deep batch and wave 1 stages the 512-chunk B strip as ONE +// straight-line 8-deep batch behind a single vmcnt drain, then the phase's +// 8 K-steps accumulate into the SAME acc_frag in ascending k order, and an +// end-of-phase __syncthreads() releases the LDS for the next phase. The +// staged bytes, their bank-skewed LDS addresses, the 36-B B row pitch, the +// 8-deep batch and therefore every fragment byte are unchanged, so every +// k-ascending int32 partial and the combine output stay bit-identical. +// * LDS per block is unchanged at 15,616 B (stage_a 4,352 + stage_b 9,216 + +// 2 x 1,024 acc tiles), so four resident blocks/CU still fit by capacity +// (2 waves/block = 8 waves/CU = 2 per SIMD by capacity); the 320-block +// grid is a single resident generation. +// * The combine guard split and its 2-D one-wave launch move to +// kCombineSplits = 8 (the same 20 x 16 = 320 one-wave blocks, eight +// independent plane loads per lane in registers); a clamped split_k < 8 +// keeps the unchanged generic 128-thread combine path. +// The exact shape guard, the workspace-less fall-through, the scalar fallback +// for every unmatched shape (paired M=2 included), the pack identity, the +// stream/Graph contract and every fallback path are unchanged; no raw inline +// asm and no ISA skill is used (phase hip_only, plateau false, +// raw_inline_asm_allowed false, skill_allowed false). +// Falsifiable claim: if the 2,621,440 B split-K plane round trip is a material +// part of the 18.145 us wall, halving it must move the official Graph +// median/P90 below 18.145450/18.166250 us (expected ~16.5-17.5 us median, +// ~4-9% faster, from the 1,310,720 B traffic reduction minus the reduced grid +// parallelism) with mismatch_count 0 and unchanged Graph capture; if it stays +// flat or regresses, the plane traffic was hidden behind the operand stream +// and the residual wall is the compulsory 5.35 MB operand/out stream. // -// 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. +// The scalar kernel below stays byte-for-byte as the generic fallback for +// every unmatched shape, including the paired M=2 API shape with the same +// (N, K); the (m, n, k) == (16, 1280, 4096) guard is the only call site of +// the split-K DUMMA path, and if the caller supplies no usable workspace the +// exact shape safely falls through to that scalar fallback. // -// Graph-safety: this file only launches kernels on the caller-provided stream -// (PyTorch's current HIP stream). It performs no allocation, compilation, -// autotuning, packing, host synchronization, or device synchronization, and -// touches no buffer other than the caller-provided out/workspace pointers. -// The split-K choice is a host-side static dispatch decision (env override or -// default 16) made once per launch call, never inside the Graph replay, so -// capture/replay determinism holds. -// -// Header order is deliberate: hip_runtime.h, then hip_bfloat16.h, then -// du_mma.h (du_mma.h on this DTK is not self-contained before the HIP runtime -// headers). +// Iteration 17 (HIP-only consolidation round, this round): complete the B +// memory path by packing the cold once-read weight into the col_major +// fragment layout, the lineage rule that the sibling TP4/TP8 M=16 shapes +// validated ("one ds_read_b64 per step, was 8 ds_read_u8"). Evidence for the +// accepted iteration-13 object (source digest 920506893c362e1bcefa5c36ffb0776fbfee00a9dd004943b66fc55d386ae572, +// primary kernel 320 blocks x 128 threads, LDS 15,616 B, arch_vgpr 80, sgpr 48, +// scratch 0, profiled 18.24 us; combine 320 one-wave blocks, profiled 2.72 us; +// official unprofiled Graph median/P90 16.614190340042114/16.642990112304688 +// us): the exact-source code object's rolled LDS-only K loop issues ONE +// ds_read2_b32 for the A fragment, but EIGHT scatter ds_read_u8 per lane for +// the B fragment (immediates i*36 on the k-major 36-B row pitch) plus the +// byte-mask/shift/OR reassembly that packs those eight bytes into the v_mmac +// operand -- ~10 of the ~11 LDS wave-instructions and ~2/3 of the ~41 VALU +// wave-instructions per wave-step, with lds_wait_instructions 15,137 and +// lds_bank_conflicts 117,760 against lds_instructions 108,800. The B +// fragment's 8 bytes per lane are contiguous in memory if and only if the +// staged B tile is n-major, because the library col_major matrix_b loader +// reads p[(lane&15)*ldm + 8*(lane>>4) + i]. Change (one focused HIP edit, +// csrc/w8a8_gemm_hip.hip only): +// * for the exact (K, N) == (4096, 1280) weight pair, launch_pack_w8a8_weight +// now writes the n-major transpose packed[n*K + k] = raw[k*N + n] instead +// of the identity copy (same buffer, same k*n bytes, outside the timed +// region/Graph; the packed layout is opaque per the API contract). Every +// other (K, N) keeps the identity copy. +// * the guard kernel stages its 32-column B strip n-major: +// stage_b[n * kStageBColStride + k] = packed[n0 + n][phase_begin + k] at +// the 16-B-aligned 272-B column stride, so each 16-B cooperative chunk is +// one ds_write_b128 and each column's phase is one contiguous run. +// * the B fragment is declared col_major and loaded from the n-major mirror +// (ldm = 272), so the K loop's B fragment becomes ONE contiguous 8-byte +// LDS read per step. +// * the scalar fallback (paired M=2 and every unmatched shape) and the +// DUMMA kernel's direct-global safety net decode the same permutation for +// this (K, N); all other shapes are untouched. +// The same logical (k, n) element lands in the same v_mmac fragment register, +// every staged byte is a copy of the same packed byte, and the k-ascending +// int32 accumulation is unchanged, so every partial and every output byte stay +// bit-identical; LDS per block drops 15,616 -> 15,104 B (still 4 resident +// blocks/CU), and the grid, plane geometry, combine kernel, workspace +// contract, stream/Graph properties and fallbacks are unchanged. No raw +// inline asm and no ISA skill is used (phase hip_only, plateau false, +// raw_inline_asm_allowed false, skill_allowed false). +// Falsifiable claim: if the B fragment's 8-scatter-byte LDS read plus its +// reassembly VALU is a material part of the 16.614 us wall, replacing it with +// one 8-byte LDS read (and ~4x fewer LDS wave-instructions) must move the +// official Graph median/P90 below 16.614190/16.642990 us (expected +// ~15.9-16.3 us median, 2-4% faster) with mismatch_count 0 and unchanged +// Graph capture; if it stays flat or regresses, the LDS instruction count is +// not the wall and the residual wall is the compulsory operand byte stream. #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; +// 128 threads = 2 waves of the native gfx928 64-lane wavefront, and a +// multiple of 64 as required by this DTK. +constexpr int kBlockThreads = 128; -// Minimal DUMMA tile constants (gfx928 INT8 support is m16n16k32). -constexpr int kDummaM = 16; -constexpr int kDummaN = 16; -constexpr int kDummaK = 32; -constexpr int kDummaThreads = 64; // one wavefront (64 lanes) per block -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}; +// Iteration 17: the exact-guard weight pair (K, N) == (4096, 1280) whose +// packed weight buffer is the n-major transpose packed[n*K + k] = +// raw[k*N + n] instead of the identity copy. launch_pack_w8a8_weight writes +// that permutation (outside the timed region / Graph), the split-K DUMMA +// kernel consumes it through its n-major LDS stage with the library col_major +// matrix_b fragment loader, and every other consumer of the packed weight +// (the scalar fallback that serves the paired M=2 shape and every unmatched +// (m, n, k), plus the DUMMA kernel's direct-global safety net) decodes the +// same permutation. Every other (K, N) keeps the byte-for-byte identity pack. +constexpr int kGuardK = 4096; +constexpr int kGuardN = 1280; -// --------------------------------------------------------------------------- -// Scalar GEMM kernel: one thread computes one output element (generic -// fallback for every unmatched shape, including the paired M=2 shape). -// --------------------------------------------------------------------------- -// Adjacent lanes own adjacent N columns (the fastest-changing dimension), so -// output stores are coalesced across each wavefront. -__global__ __launch_bounds__(kScalarBlockThreads) void -w8a8_scalar_gemm_kernel( +// One thread -> one output element. Consecutive threads stay adjacent in N, +// the fastest-changing dimension, so weight loads are coalesced. +__global__ __launch_bounds__(kBlockThreads) void w8a8_scalar_kernel( const int8_t* __restrict__ a, const int8_t* __restrict__ b, const float* __restrict__ x_scale, const float* __restrict__ weight_scale, - hip_bfloat16* __restrict__ out, + __hip_bfloat16* __restrict__ out, int m, int n, int k) { - const int64_t linear = - static_cast(blockIdx.x) * kScalarBlockThreads + threadIdx.x; - const int64_t total = static_cast(m) * n; - if (linear >= total) { + const int64_t total = + static_cast(m) * static_cast(n); + const int64_t idx = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { return; } - const int row = static_cast(linear / n); - const int col = static_cast(linear - static_cast(row) * n); - - // Complete K loop, k ascending, accumulated exactly in int32. The maximum - // API K (<= 6144) keeps the int8 dot well inside int32 range. - const int8_t* a_row = a + static_cast(row) * k; - // Iteration 11: the exact (k,n)==(4096,1280) pair is packed [N,K] n-major - // by launch_pack_w8a8_weight; the fallback decodes that layout so it stays - // correct against the packed buffer (paired M=2 validation shape included). - // Every other (K,N) keeps the raw [K,N] row-major decode. - const bool packed = (k == kPackedK && n == kPackedN); - const int8_t* b_col = - b + (packed ? static_cast(col) * k : col); - const int64_t b_stride = packed ? 1 : static_cast(n); + + const int row = static_cast(idx / n); + const int col = static_cast(idx - static_cast(row) * n); + + const int8_t* a_ptr = 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[static_cast(kk) * b_stride]); + if (k == kGuardK && n == kGuardN) { + // Iteration 17: the exact-guard weight buffer holds the n-major + // transpose packed[n*K + k] = raw[k*N + n] (a pure byte permutation), + // so the fallback decodes that layout here. The k-ascending int32 sum + // and therefore every output byte are unchanged. + const int8_t* b_col = b + static_cast(col) * k; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_ptr[kk]) * + static_cast(b_col[kk]); + } + } else { + const int8_t* b_ptr = b + col; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_ptr[kk]) * static_cast(*b_ptr); + b_ptr += n; + } } - // Same float32 evaluation order as the exact reference: - // (float(dot) * x_scale[m]) * weight_scale[n], then bf16 (RN). + // Integer dot stays exact in int32; the two fp32 scales are applied once at + // the very end and the result is stored as bf16. 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); + out[idx] = __float2bfloat16(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 split-K 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-lane wavefront per N tile (two wavefronts per block +// since iteration 8). Iteration 4 gives every block kTilesPerWave = 2 +// adjacent 16x16 output tiles, and the 128 stage-aligned (k/32) K-steps are +// split kSplitKTarget ways so that the grid becomes (N/(16*T))*S blocks +// instead of N/16. Every K step is an explicit int32 du_mma_sync over a +// 32-aligned ascending k chunk, so each partial is an exact int32 sub-sum and +// the combine below is bit-identical to the scalar kernel's total (exact +// integer addition, order-free; |total| <= K*127*127 ~= 6.6e7 << 2^31). +// +// Iteration 6 adds whole-slice A+B LDS staging: each block first mirrors its +// whole A slice and its whole 32-column B strip into LDS with 16-B cooperative +// copies (exact byte mirrors, bank-skewed row pitches), and the K loop then +// reads every fragment from LDS with no global load and no vmcnt wait on the +// MMA path. See the staging block below for the capacity, bank-skew and +// bit-exactness arguments. +// +// Iteration 8 raises waves-per-block from 1 to 2: wave w owns N tile w of the +// block's 32 columns and runs its own independent accumulator chain, the +// one-time staging is partitioned across the two waves (no operand byte is +// fetched twice), and LDS per block is unchanged, so 4 resident blocks/CU +// become 8 resident waves/CU = 2 per SIMD. +// +// The gfx928 int8 m16n16k32 accelerator fragment ownership is opaque, so the +// epilogue materializes each tile with du_store_matrix_sync into a per-wave +// LDS staging array (16x16 int32 = 1,024 B each, library-defined mem_row_major +// placement) and from there stores the raw int32 partials into plane `split` +// of the caller workspace. blockDim is two wavefronts, so the staging +// __syncthreads() is a genuine cross-wave barrier inside one block (wave 0's A +// slice is read by wave 1 and vice versa for the B strip); there is no +// cross-block synchronization anywhere. // --------------------------------------------------------------------------- -// 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; + +// Iteration 4: two adjacent 16-wide N tiles per block (32 N columns). +// Iteration 8: one 64-lane wave owns each of those two tiles, so the block is +// two wavefronts and the A fragment of a 32-K step is shared by both waves +// through LDS; B has no reuse (every B byte belongs to exactly one block). +constexpr int kTilesPerWave = 2; + +// Iteration 13 target split of the K dimension: 8. With kTilesPerWave = 2 the +// grid is (N / (16*2)) * S = 40 * S blocks, so S = 8 gives 320 one-generation +// blocks over the 120-device-CU K500SM_AI (2.67 blocks/CU by grid, all resident +// because the unchanged 15,616 B/block LDS footprint still fits four blocks/CU +// = 8 resident waves/CU = 2 per SIMD by capacity). The move 16 -> 8 halves the +// ONLY non-compulsory component of the operator's memory traffic: the split-K +// partial planes (8 x 80 KiB = 655,360 B written and re-read instead of +// 16 x 80 KiB = 1,310,720 B), i.e. 1,310,720 B of the measured 8,068,416 B +// per-operator-replay traffic. Each split now owns 16 x 32 K stages = 512 +// rows, which is twice the unchanged 256-row staging capacity, so the kernel +// stages and consumes its slice in ascending 256-row PHASES (exactly two at +// the guard); see the phase loop below. The slices stay UNIFORM +// (128 == 8 x 16) and every boundary stays a multiple of 32, so the +// stage-alignment and exact-int32 contracts are untouched. The workspace +// holds 16 int32 planes for this shape, so 8 planes fit with room to spare and +// the runtime clamp still covers smaller buffers. +constexpr int kSplitKTarget = 8; + +// Iteration 6 staging capacity, tied to the guard launch (k == 4096): 8 +// stages x 32 K = 256 K rows of the slice are staged per phase. Kept +// deliberately independent of kSplitKTarget since iteration 13: the 512-row +// slice of the 8-way split is staged as two 256-row phases, so the staging +// capacity, the A bank-skew pad and the 8-deep B batch are the accepted +// iteration-11 geometry (iteration 17 only re-lays the B mirror n-major). +constexpr int kStageMaxSteps = 8; +constexpr int kStageMaxRows = kStageMaxSteps * kDummaK; // 256 rows +// stage_a mirrors the block's A slice: 16 m-rows x rows bytes of x_q stored +// m-major at kStageARowStride = rows + 16 pad bytes. 272 B (== 16 mod 128 B) +// skews the row_major A fragment loader's 16 row reads instead of replaying +// one bank, and 272 is a multiple of 16 so the cooperative 16-B stores stay +// 16-B aligned. +constexpr int kStageARowStride = kStageMaxRows + 16; // 272 B +// Iteration 17: stage_b is an n-MAJOR mirror of the block's 32-column B strip. +// Element B[k][n0 + n] lives at stage_b[n * kStageBColStride + k], i.e. each +// N column owns a contiguous run of kStageBColStride bytes (the phase's k rows +// followed by 16 pad bytes), and the stride is 16-B aligned so every 16-B +// staged chunk is one aligned ds_write_b128. The library col_major matrix_b +// fragment loader reads lane l's 8 fragment bytes as +// p[(lane & 15) * ldm + 8 * (lane >> 4) + i] (i = 0..7) -- EIGHT CONTIGUOUS +// bytes -- so on the n-major tile the K loop's B fragment becomes ONE 8-byte +// LDS read per step instead of the iteration-13 row_major loader's 8 scatter +// ds_read_u8 plus the byte-reassembly VALU that fed them. The same logical +// (k, n) element lands in the same v_mmac fragment register in both layouts, +// so every k-ascending int32 partial and the whole output stay bit-identical. +constexpr int kStageBColStride = kStageMaxRows + 16; // 272 B +constexpr int kStageABytes = kDummaM * kStageARowStride; // 4,352 B +constexpr int kStageBBytes = + kDummaN * kTilesPerWave * kStageBColStride; // 8,704 B + +__global__ __launch_bounds__(kDummaWaveSize * kTilesPerWave) +void w8a8_dumma_m16_splitk_kernel( + const int8_t* __restrict__ a, // x_q [M, K] row-major + const int8_t* __restrict__ b, // packed_weight [N, K] n-major (guard) + int32_t* __restrict__ partials, // workspace [S][M, 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; + int k, + int split_k, + int groups /* n / (kDummaN * kTilesPerWave) */) { + // Iteration 8: blockDim is TWO 64-lane wavefronts. Wave w owns N tile w of + // the block's 2 x 16 columns (cols [n0 + 16*w, n0 + 16*w + 16)) and runs its + // own independent m16n16k32 accumulator chain for that tile; the staging of + // each ascending 256-row phase is partitioned across the waves, not + // duplicated (wave 0 copies the shared A phase, wave 1 the B strip phase), + // so every global byte is still read exactly once per block -- no repeated + // HBM reads are traded for occupancy. + const int tid = static_cast(threadIdx.x); // blockDim == 128 + const int w = tid / kDummaWaveSize; // wave id inside block: 0/1 + const int lane = tid - w * kDummaWaveSize; // lane id within the wave + const int group = static_cast(blockIdx.x) % groups; + const int split = static_cast(blockIdx.x) / groups; + const int n0 = group * (kDummaN * kTilesPerWave); - // 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; + // Two LDS staging tiles, one per 16-wide N tile of this block, plus the + // iteration-6 whole-slice operand mirrors: stage_a (the shared A slice) and + // stage_b (the block's 32-column n-major B strip). Static LDS = 2 x 1,024 + // + 4,352 + 8,704 = 15,104 B/block, so floor(65,536 / 15,104) = 4 resident + // blocks/CU = 8 resident waves/CU (2 per SIMD) since iteration 8. + __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[kStageABytes]; + __shared__ __align__(16) int8_t stage_b[kStageBBytes]; 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). + // Iteration 17: the B fragment is col_major over the n-major stage_b, so its + // 8 bytes per lane are contiguous in LDS (one 8-byte read per step). du::dumma::DUFragment - b_frag; + b_frag; // this wave's own 16-wide B half of the block's 32-column strip du::dumma::DUFragment - acc_frag; + acc_frag; // this wave's accumulator for its n16 tile du::dumma::du_fill_fragment(acc_frag, 0); - for (int i = 0; i < steps; ++i) { - const int kk = k0 + i * kDummaK; - du::dumma::du_load_matrix_sync(a_frag, a + kk, k); - // Packed n-major row n starts at b[n*k + k0]; the library col_major - // loader addresses p[(lane&15)*ldm + ((lane>>4)<<3) + i], i.e. n-major - // rows of k bytes, so the tile's 16 columns live at n0..n0+15 n-rows - // with row stride k. - du::dumma::du_load_matrix_sync( - b_frag, b + static_cast(n0) * k + kk, k); - du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } + // Stage-aligned split: split s owns stages [s*T/S, (s+1)*T/S) of the + // T = k/kDummaK 32-K stages, so every k range is a multiple of 32 and the + // ranges tile [0, k) exactly with no gap and no overlap. + const int total_stages = k / kDummaK; + const int stage_begin = (split * total_stages) / split_k; + const int stage_end = ((split + 1) * total_stages) / split_k; + const int k_begin = stage_begin * kDummaK; + const int k_end = stage_end * kDummaK; - // 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); -} + // Iteration 13: split_k == 8 makes each split's slice 16 x 32-K stages = 512 + // rows, twice the unchanged 256-row staging capacity, so the slice is staged + // and consumed in ascending 256-row PHASES (exactly two at the exact-shape + // guard). Each phase reproduces the accepted iteration-11 staging geometry + // byte-for-byte -- A: 256 16-B chunks = 4 per lane on wave 0 (one 4-deep + // batch); B: 512 chunks = 8 per lane on wave 1 (ONE straight-line 8-deep + // batch behind a single vmcnt drain) -- and the phase's 8 K-steps accumulate + // into the SAME acc_frag in ascending k order. The staged bytes, their LDS + // addresses and therefore every fragment byte are identical to what the + // 16-way split staged per 32-K step; only the number of int32 partial planes + // written and re-read by the split-K round trip halves. + for (int phase_begin = k_begin; phase_begin < k_end; + phase_begin += kStageMaxRows) { + const int remain = k_end - phase_begin; + const int phase_rows = (remain < kStageMaxRows) ? remain : kStageMaxRows; -// Template launch helper for the partial kernel (one wavefront per block, -// grid = (num_tiles, SPLIT_K)). -template -void launch_splitk_partial( - const int8_t* a, - const int8_t* b, - int32_t* partials, - int n, - int k, - hipStream_t stream, - int num_tiles) { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_m16_splitk_partial_kernel), - dim3(static_cast(num_tiles), static_cast(SPLIT_K)), - dim3(static_cast(kDummaThreads)), - 0, - stream, - a, - b, - partials, - n, - k); -} - -// Iteration 11 packed B-fragment loader (validated on the TP4 gate_up -// (4096,768) lineage): with n-major [N,K] storage and a col_major fragment, -// du_mma.h's matrix_b col_major address rule is x[i] = p[(lane&15)*ldm + -// ((lane>>4)<<3) + i] -- lane l's 8 fragment bytes are consecutive in -// memory, so one 8-byte LDS read per fragment replaces the 8 ds_read_u8 at -// stride-64 addresses that the row-major loader emits. du_mma_sync consumes -// the fragment as one 64-bit value per lane, so writing the packed bytes -// straight into x[0..7] is exactly the library's own storage convention. -// p must be 8 B aligned (guaranteed: b_s is __align__(256), kBStride % 8 == -// 0, step offsets are multiples of 32, and (lane&15)*kBStride % 8 == 0). -__device__ __forceinline__ void load_b_frag_packed( - du::dumma::DUFragment& b_frag, - const int8_t* p, int ldm, int lane) { - const int row = lane & 15; - const int col = (lane >> 4) << 3; - const uint64_t v = *reinterpret_cast( - p + static_cast(row) * ldm + col); - *reinterpret_cast(b_frag.x) = v; -} - -// --------------------------------------------------------------------------- -// Whole-K-slice LDS staging partial kernel (M == 16, iteration 7): -// partial[tile, slice][m, n] = int32_dot(x_q[m, k0:k1], weight[k0:k1, n]) -// --------------------------------------------------------------------------- -// Same 32-aligned non-uniform K slices and the same k-ascending int32 -// du_mma accumulation as w8a8_dumma_m16_splitk_partial_kernel, so the -// combine+scale kernel below sums bit-identical partials and the outputs are -// bit-identical to the accepted iteration-3 kernel. Differences: -// - A (16 x slice) is staged ONCE into LDS with coalesced dword global -// loads (4 B/lane, 256 B per wavefront instruction); B (16 x slice) is -// staged as 16-B int4 loads straight from the packed [N,K] n-major -// weight (iteration 11) into an n-major LDS tile. -// - One __syncthreads() after staging, then a zero-barrier LDS-only K -// loop: du_load_matrix_sync reads the A fragment from LDS and -// load_b_frag_packed reads each B fragment with one ds_read_b64 (~30-cycle -// LDS latency) instead of a global round trip per K step. -// Layouts (all contiguous): -// x_q [16, K] int8 row-major (ldm = K) -// weight [K, N] int8 row-major for every (K,N) except the exact pair -// (k,n)==(4096,1280), which is packed [N,K] n-major (element -// (k,n) at P[n*K+k]; iteration 11) -// partials [(SPLIT_K * num_tiles) * 256] int32, plane = slice*num_tiles + -// tile, row-major 16x16 tile per plane (caller workspace) -// LDS: a_s[16][kMaxKlen + 4] (compile-time row stride kMaxKlen+4, an odd -// dword count -> the library's matrix_a row-major lane pattern (row = -// lane&15, 8 consecutive bytes per lane) is bank-conflict-free); -// b_s[16][kMaxKlen + 16] (iteration 11: n-major, row stride kMaxKlen+16 = -// 16-B aligned for the int4 staging stores, 8-B aligned for the col_major -// fragment b64 reads, not a multiple of 128 B -> the 16 staged n-rows do not -// alias onto one LDS bank phase; each lane's 8 fragment bytes are -// consecutive, so every B fragment loads with ONE ds_read_b64 instead of 8 -// ds_read_u8). Per-template LDS fits: S=6 22,848 B, S=8 16,704 B, S=9 -// 15,680 B, S=12 11,584 B, S=16 8,512 B (all <= 64 KiB; resident blocks/CU -// by LDS: 2/3/4/5/7). Only launched -// under the exact (m,n,k) == (16,1280,4096) guard, so total_steps = k/32 = -// 128 is exact for the compile-time LDS sizing below. -template -__global__ __launch_bounds__(kDummaThreads) void -w8a8_dumma_m16_slicestage_partial_kernel( - const int8_t* __restrict__ a, - const int8_t* __restrict__ b, - int32_t* __restrict__ partials, - int n, - int k) { - const int tile = static_cast(blockIdx.x); - const int slice = static_cast(blockIdx.y); - const int n0 = tile * kDummaN; - - // 32-aligned non-uniform slice bounds (identical formula to the direct - // kernel: slices tile [0, K) exactly once in ascending order). - const int total_steps = k / kDummaK; // 4096 / 32 = 128 under the guard - const int base = total_steps / SPLIT_K; - const int rem = total_steps - base * SPLIT_K; - const int steps = base + (slice < rem ? 1 : 0); - const int k0 = (slice * base + (slice < rem ? slice : rem)) * kDummaK; - const int klen = steps * kDummaK; - - // Compile-time LDS sizing for the longest slice this instantiation can - // receive (base+1 steps, or base when SPLIT_K divides 128 evenly). - constexpr int kMaxSteps = 128 / SPLIT_K + (128 % SPLIT_K == 0 ? 0 : 1); - constexpr int kMaxKlen = kMaxSteps * kDummaK; - __shared__ int8_t a_s[16][kMaxKlen + 4]; - // Iteration 11: B is staged n-major (16 columns x kMaxKlen+16 bytes, row - // stride kMaxKlen+16) so col_major B fragments read 8 consecutive bytes - // per lane (one ds_read_b64). kMaxKlen is a multiple of 32, so the row - // stride is 16-B aligned (int4 staging stores stay aligned), 8-B aligned - // (fragment b64 reads stay aligned) and not a multiple of 128 B (no LDS - // bank-phase aliasing across the 16 staged n-rows). - __shared__ __align__(256) int8_t b_s[16][kMaxKlen + 16]; - - // Row strides are compile-time constants equal to the array row strides - // (a_s: kMaxKlen+4 bytes = odd dword count; b_s: kMaxKlen+16 bytes), so - // the staging stores and the K-loop fragment reads always agree even for - // short slices (klen < kMaxKlen). The A stride stays an odd dword count - // (bank-conflict-free row-major lane pattern); the B n-major stride keeps - // every fragment b64 read 8 B aligned and off the single bank phase. - constexpr int kAStride = kMaxKlen + 4; - constexpr int kBStride = kMaxKlen + 16; - - // Stage A: 16 rows x klen bytes = 4*klen dwords; consecutive lanes copy - // consecutive dwords of one row (coalesced 256 B per wavefront load). - // Iteration 12 (staging MLP, alignment-safe): the fresh exact-source ISA - // (iteration-12 profile, digest ab7fc108...) shows both staging loops - // compile to a fully serialized load -> s_waitcnt vmcnt(0) -> store chain - // (A: 16 sequential global round trips per block, B: 4), and the accepted - // iteration-11 round already validated this fix on the B side (16 - // serialized dword loads -> 4 aligned int4 loads). Both loops below batch - // 4-deep in source: 4 independent loads are issued before their 4 stores, - // so one vmcnt drain covers 4 in-flight round trips (A 16 -> 4, B 4 -> 1). - // A stays dword-wide: 4-B loads AND stores are aligned at every a_s row, - // while kAStride = kMaxKlen+4 is not 16-B aligned, so int4 A staging would - // repeat round 9's misaligned-store regression (58.56 us) -- that round - // falsified misaligned int4 staging, not the MLP idea. B stays int4 (its - // global rows are 16-B aligned and kBStride % 16 == 0). Fragment reads, - // LDS strides, the K loop and the int32 accumulation order are untouched, - // so outputs stay bit-identical (0 mismatches) and Graph capture/replay - // with changed contents must pass. klen % 32 == 0 => a_dwords % 128 == 0, - // so the batch loop is exact whenever klen % 128 == 0 (S=16: klen = 256 -> - // 1024 dwords = 4 batches exactly; S=6/9/12's odd slices fall through to - // the scalar tail, which preserves correctness for every S in the sweep). - const int a_dwords = 4 * klen; - const int a_dw_per_row = klen / 4; - constexpr int kABatch = 4 * kDummaThreads; - { - const auto ld_a = [&](int e) { - const int row = e / a_dw_per_row; - return *reinterpret_cast( - a + static_cast(row) * k + k0 + (e - row * a_dw_per_row) * 4); - }; - const auto st_a = [&](int e, int32_t v) { - const int row = e / a_dw_per_row; - *reinterpret_cast(&a_s[row][(e - row * a_dw_per_row) * 4]) = v; - }; - int e = static_cast(threadIdx.x); - for (; e + 3 * kDummaThreads < a_dwords; e += kABatch) { - const int32_t v0 = ld_a(e); - const int32_t v1 = ld_a(e + kDummaThreads); - const int32_t v2 = ld_a(e + 2 * kDummaThreads); - const int32_t v3 = ld_a(e + 3 * kDummaThreads); - st_a(e, v0); - st_a(e + kDummaThreads, v1); - st_a(e + 2 * kDummaThreads, v2); - st_a(e + 3 * kDummaThreads, v3); - } - for (; e < a_dwords; e += kDummaThreads) { - st_a(e, ld_a(e)); + // Iteration 6 A+B phase staging, PARTITIONED across the two wavefronts + // since iteration 8 (B mirror is n-major since iteration 17). Every + // global operand byte of this phase is loaded exactly once, in 16-B int4 + // granules, into an EXACT byte mirror of the operand region the fragment + // loaders read: stage_a holds a[m][phase_begin .. phase_begin+phase_rows) + // at row stride kStageARowStride, and stage_b holds the block's 32 + // B-columns' phase as b_lds[n][k], i.e. element + // packed[n0 + n][phase_begin + k] at n * kStageBColStride + k (the 16 pad + // bytes per column are never written and never read). Wave 0 copies the A + // phase and wave 1 copies the B strip -- the union reads exactly the same + // global bytes as the single-wave iteration-6 prologue, so no operand byte + // is fetched twice. Each copy uses 64 lanes x 16 B per wavefront + // instruction; the A phase is one 4-deep batch and the exact-guard B strip + // is one 8-deep batch, so one vmcnt drain covers each wave's whole phase + // and the two waves' round trips overlap across the SIMDs. n0 and + // phase_begin are 32-aligned under the exact-shape guard and k % 16 == 0, + // n % 16 == 0, so every 16-B source load is 16-B aligned, and every staged + // B chunk lands at a 16-B-aligned n-major address (one ds_write_b128). + const bool staged_fits = phase_rows <= kStageMaxRows; + if (staged_fits) { + if (w == 0) { + // A slice phase: 16 m-rows x phase_rows bytes = phase_rows int4 + // chunks, per_row = phase_rows/16 chunks per m-row (phase_rows is a + // multiple of 32, hence of 16). + const int per_row = phase_rows / 16; + const int a_chunks = kDummaM * per_row; // == phase_rows + const int4* __restrict__ a4 = reinterpret_cast(a); + int4* __restrict__ ald4 = reinterpret_cast(stage_a); + const int kb4 = phase_begin >> 4; // phase start in global int4 columns + const int a4_stride = k >> 4; // global A int4 columns per m-row + const auto ld_a = [&](int e) { + const int m = e / per_row; + const int t = e - m * per_row; + return a4[static_cast(m) * a4_stride + kb4 + t]; + }; + const auto st_a = [&](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_a(e); + const int4 v1 = ld_a(e + kDummaWaveSize); + const int4 v2 = ld_a(e + 2 * kDummaWaveSize); + const int4 v3 = ld_a(e + 3 * kDummaWaveSize); + st_a(e, v0); + st_a(e + kDummaWaveSize, v1); + st_a(e + 2 * kDummaWaveSize, v2); + st_a(e + 3 * kDummaWaveSize, v3); + } + for (; e < a_chunks; e += kDummaWaveSize) { + st_a(e, ld_a(e)); + } + } else { + // Iteration 17 B strip: the n-major packed weight stores each N + // column's K axis contiguously, so a phase is + // chunks_per_n = phase_rows / 16 int4 chunks PER N column (16 at the + // guard) and consecutive lanes cover consecutive 16-B chunks of one + // column's 256-B phase run. Each chunk is written CONTIGUOUSLY into + // the n-major LDS tile at stage_b[nl * kStageBColStride + t * 16], i.e. + // one aligned ds_write_b128 per chunk, exactly the bytes the col_major + // fragment loader reads back. + const int chunks_per_n = phase_rows / 16; // 16 at the guard + const int b_chunks = (kDummaN * kTilesPerWave) * chunks_per_n; + const int4* __restrict__ b4 = reinterpret_cast(b); + int4* __restrict__ bld4 = reinterpret_cast(stage_b); + const int b4_per_n = k >> 4; // packed int4 chunks per N column + const auto ld_b = [&](int e) { + const int nl = e / chunks_per_n; // local N column 0..31 + const int t = e - nl * chunks_per_n; // k-chunk in this phase + return b4[static_cast(n0 + nl) * b4_per_n + + (phase_begin >> 4) + t]; + }; + // 272-B n-row stride and 16-B chunks keep every store 16-B aligned. + const auto st_b = [&](int e, const int4& v) { + const int nl = e / chunks_per_n; + const int t = e - nl * chunks_per_n; + bld4[nl * (kStageBColStride >> 4) + t] = v; + }; + int e = lane; + // Iteration 11: the exact-guard B strip is exactly 8 chunks per lane + // (phase_rows == kStageMaxRows == 256 -> b_chunks == 512 == 8 x 64 + // lanes), so that case 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, and the second dependent global + // round trip leaves the block's staging join. Any other chunk count + // keeps the unchanged rolled 4-deep loop below. + if (b_chunks == 8 * kDummaWaveSize) { + const int4 v0 = ld_b(e + 0 * kDummaWaveSize); + const int4 v1 = ld_b(e + 1 * kDummaWaveSize); + const int4 v2 = ld_b(e + 2 * kDummaWaveSize); + const int4 v3 = ld_b(e + 3 * kDummaWaveSize); + const int4 v4 = ld_b(e + 4 * kDummaWaveSize); + const int4 v5 = ld_b(e + 5 * kDummaWaveSize); + const int4 v6 = ld_b(e + 6 * kDummaWaveSize); + const int4 v7 = ld_b(e + 7 * kDummaWaveSize); + st_b(e + 0 * kDummaWaveSize, v0); + st_b(e + 1 * kDummaWaveSize, v1); + st_b(e + 2 * kDummaWaveSize, v2); + st_b(e + 3 * kDummaWaveSize, v3); + st_b(e + 4 * kDummaWaveSize, v4); + st_b(e + 5 * kDummaWaveSize, v5); + st_b(e + 6 * kDummaWaveSize, v6); + st_b(e + 7 * kDummaWaveSize, v7); + e += 8 * kDummaWaveSize; + } + for (; e + 3 * kDummaWaveSize < b_chunks; e += 4 * kDummaWaveSize) { + const int4 v0 = ld_b(e); + const int4 v1 = ld_b(e + kDummaWaveSize); + const int4 v2 = ld_b(e + 2 * kDummaWaveSize); + const int4 v3 = ld_b(e + 3 * kDummaWaveSize); + st_b(e, v0); + st_b(e + kDummaWaveSize, v1); + st_b(e + 2 * kDummaWaveSize, v2); + st_b(e + 3 * kDummaWaveSize, v3); + } + for (; e < b_chunks; e += kDummaWaveSize) { + st_b(e, ld_b(e)); + } + } } - } + // Cross-wave barrier: wave 0's staged A phase is read by wave 1 (and vice + // versa for the B strip), so both tiles are ready for both waves here, and + // it also orders this phase's staging after every lane's previous-phase + // LDS reads. + __syncthreads(); // stage_a / stage_b ready for every lane of the block - // 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)); + // Iteration 8 LDS-only K loop, one independent chain per wavefront. Each + // 32-K step reads the shared A fragment from the staged A mirror at row + // stride kStageARowStride and this wave's own 16-wide B half from the + // iteration-17 n-major B mirror at column stride kStageBColStride, where + // the library col_major loader takes the lane's 8 fragment bytes as ONE + // contiguous 8-byte read -- no global load and no vmcnt wait on the MMA + // path. The ascending k0 order keeps each partial's int32 accumulation + // exact, and the staged tiles are exact byte copies of the global bytes the + // iteration-4 path read, so this branch is bit-identical to the direct-global + // loop below. + if (staged_fits) { + for (int k0 = phase_begin; k0 < phase_begin + phase_rows; + k0 += kDummaK) { + const int local_row = k0 - phase_begin; + du::dumma::du_load_matrix_sync(a_frag, stage_a + local_row, + kStageARowStride); + du::dumma::du_load_matrix_sync( + b_frag, stage_b + local_row + w * kDummaN * kStageBColStride, + kStageBColStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + } else { + // Safety net kept for a runtime-clamped split whose phase would exceed + // the staging capacity (an undersized workspace clamps split_k < 8): the + // iteration-4 direct-global path, here consuming the same n-major packed + // weight through the col_major loader at column stride K, so a clamped + // launch stays bit-exact and can never overrun stage_a / stage_b. This + // branch never executes at the exact-shape guard launch, where every + // phase is exactly kStageMaxRows rows. + for (int k0 = phase_begin; k0 < phase_begin + phase_rows; + 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 + w * kDummaN) * k + k0, k); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } } + // End-of-phase barrier: no lane may overwrite stage_a / stage_b for the + // next phase until every lane has consumed this phase's staged fragments. + __syncthreads(); } + // Each wave materializes ONLY its own n16 tile (its private acc_tile slot) + // and publishes that tile's 16 columns to the plane; the two waves' plane + // ranges are disjoint, so no cross-wave dependency exists in the epilogue + // (the barrier below only orders each wave's own LDS store/read-back). + int32_t* my_tile = (w == 0) ? acc_tile0 : acc_tile1; + du::dumma::du_store_matrix_sync(my_tile, acc_frag, kDummaN, + du::dumma::mem_row_major); __syncthreads(); - du::dumma::DUFragment - a_frag; - // Iteration 11: col_major B fragments match the n-major staged tile, so - // each lane's 8 fragment bytes are consecutive (one ds_read_b64 per - // fragment in the K loop instead of 8 ds_read_u8). - du::dumma::DUFragment - b_frag; - du::dumma::DUFragment - acc_frag; - du::dumma::du_fill_fragment(acc_frag, 0); - - // Zero-barrier LDS-only K loop: every fragment read comes from LDS, so no - // global latency sits on the MMA critical path after the single staging - // barrier above. - const int lane = static_cast(threadIdx.x) & 63; - for (int i = 0; i < steps; ++i) { - du::dumma::du_load_matrix_sync(a_frag, &a_s[0][i * kDummaK], kAStride); - // One 8-B ds_read_b64 per B fragment (col_major: each lane's 8 fragment - // bytes are consecutive in the n-major tile) instead of 8 ds_read_u8 at - // the 20-B k-major pitch. - load_b_frag_packed(b_frag, &b_s[0][i * kDummaK], kBStride, lane); - du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + // 256 int32 elements / 64 lanes: each lane writes four partials of its own + // tile, kept row-major within the plane so the combine below reads them + // coalesced. Wave 1's tile lands kDummaN columns to the right of wave 0's. + int32_t* plane = + partials + static_cast(split) * kDummaM * n; +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; + linear += kDummaWaveSize) { + const int row = linear / kDummaN; + const int col = linear - row * kDummaN; + const int64_t dst = + static_cast(row) * n + n0 + w * kDummaN + col; + plane[dst] = my_tile[linear]; } - - // Private int32 partial plane: same slice-major layout as the direct - // kernel (plane = slice*num_tiles + tile). - const int num_tiles = static_cast(gridDim.x); - du::dumma::du_store_matrix_sync( - partials + (slice * num_tiles + tile) * (kDummaM * kDummaN), - acc_frag, - kDummaN, - du::dumma::mem_row_major); } -// Template launch helper for the staged partial kernel (one wavefront per -// block, grid = (num_tiles, SPLIT_K)); only instantiated for the SPLIT_K -// values whose longest slice fits the 64 KiB LDS (6, 8, 9, 12, 16). -template -void launch_slicestage_partial( - const int8_t* a, - const int8_t* b, - int32_t* partials, +// Combine pass: sum the S exact int32 partial planes (order-free integer +// addition), then apply the two fp32 scales in the same order as the scalar +// fallback and store bf16. Every plane is M*N int32 = 80 KiB for this shape, +// so the whole combine is L2-resident. +// +// Iteration 9 (combine restructure, HIP-only): the guard launch +// (split_k == kCombineSplits == 8 since iteration 13) now runs a 2-D grid of +// one-wave 64-lane blocks -- (n / kCombineColsPerBlock) column groups x m rows +// -- with exactly one output element per lane (row = blockIdx.y, +// col = blockIdx.x * 64 + lane). The eight plane loads are issued as +// INDEPENDENT reads into a register array (one vmcnt drain, 8-way +// memory-level parallelism per lane) and summed afterwards in the unchanged +// ascending split order s = 0..7, instead of the rolled dependent load -> +// wait -> int32 add chain of 8 DRAM round trips per element. Each per-plane +// wavefront load also reads 64 lanes x 4 B = 256 B contiguous (four +// consecutive 64-B lines of one m-row). Every output element's int32 sum, +// fp32 scale order and bf16 conversion are unchanged, so the output is +// bit-identical; the generic path below (any clamped split_k < 8) keeps the +// original 128-thread linear mapping and the original dependent loop. +constexpr int kCombineThreads = 128; +constexpr int kCombineSplits = 8; // == kSplitKTarget of the guard +constexpr int kCombineWaveSize = 64; // one gfx928 wavefront per block +constexpr int kCombineColsPerBlock = 64; // n-columns per one-wave block + +__global__ __launch_bounds__(kCombineThreads) void w8a8_splitk_combine_kernel( + const int32_t* __restrict__ partials, // [S][M, N] int32 + const float* __restrict__ x_scale, // [M] + const float* __restrict__ weight_scale, // [N] + __hip_bfloat16* __restrict__ out, // [M, N] bf16 + int total /* m * n */, 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); -} + int split_k) { + const int lane = static_cast(threadIdx.x); -// --------------------------------------------------------------------------- -// Combine + scale kernel (M == 16): one 16x16 tile per block, one thread per -// output element. Exact int32 sum of the SPLIT_K workspace partials -// (ascending split order), then x_scale/weight_scale scaling and the bf16 -// (RN-even) store. No barrier is needed: every thread reads only its own -// 16x16 element's S partial values (written by the previous kernel on the -// same stream). -// --------------------------------------------------------------------------- -__global__ __launch_bounds__(kCombineThreads) void -w8a8_dumma_m16_combine_scale_kernel( - const int32_t* __restrict__ partials, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - hip_bfloat16* __restrict__ out, - int split_k, - int num_tiles, - int n) { - const int tile = static_cast(blockIdx.x); - const int linear = static_cast(threadIdx.x); // 0..255 - const int row = linear >> 4; - const int col = linear & 15; - const int out_col = tile * kDummaN + col; + // Iteration 9 guard specialization: one output element per lane, all sixteen + // plane reads issued as independent loads before the ascending sum. The + // launcher pairs this branch with the 2-D one-wave grid (blockDim == 64, + // grid == (n/64, m)); the blockDim test keeps any other launch shape on the + // generic path. + if (split_k == kCombineSplits && + static_cast(blockDim.x) == kCombineColsPerBlock) { + const int row = static_cast(blockIdx.y); + const int col = + static_cast(blockIdx.x) * kCombineColsPerBlock + lane; + const int64_t elem = static_cast(row) * n + col; + int32_t plane[kCombineSplits]; +#pragma unroll + for (int s = 0; s < kCombineSplits; ++s) { + plane[s] = partials[static_cast(s) * total + elem]; + } + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kCombineSplits; ++s) { + acc += plane[s]; // ascending split order: identical exact int32 sum + } + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[elem] = __float2bfloat16(scaled); + return; + } + + const int i = static_cast(blockIdx.x) * kCombineThreads + lane; + if (i >= total) { + return; + } - int32_t sum = 0; + int32_t acc = 0; for (int s = 0; s < split_k; ++s) { - sum += partials[(s * num_tiles + tile) * (kDummaM * kDummaN) + linear]; + acc += partials[static_cast(s) * total + i]; } + const int row = i / n; + const int col = i - row * n; const float scaled = - static_cast(sum) * x_scale[row] * weight_scale[out_col]; - out[row * n + out_col] = __float2bfloat16(scaled); + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[i] = __float2bfloat16(scaled); } -// --------------------------------------------------------------------------- -// Weight pack op (outside the timed region). -// --------------------------------------------------------------------------- -// For every (K,N) except the exact pair (k,n)==(4096,1280) the packed weight -// stays a contiguous [K, N] int8 buffer with an [N, 1] fp32 scale, -// byte-identical to the raw inputs. For (k,n)==(4096,1280) (iteration 11) -// launch_pack_w8a8_weight transposes the raw [K,N] weight to an [N,K] -// n-major packed layout (element (kk,nn) -> dst[nn*k + kk]) so DUMMA -// col_major B fragments read 8 consecutive bytes per lane (one ds_read_b64 -// instead of 8 ds_read_u8); the staged/direct DUMMA kernels and the scalar -// fallback (for the paired M=2 validation shape) all decode that layout. -__global__ __launch_bounds__(kCopyBlockThreads) void -w8a8_identity_copy_i8_kernel( +// Identity pack: raw int8 weight -> packed int8 weight. +__global__ __launch_bounds__(kBlockThreads) void identity_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 i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { + dst[i] = src[i]; } } -__global__ __launch_bounds__(kCopyBlockThreads) void -w8a8_identity_copy_f32_kernel( +// Identity pack: fp32 weight scale -> packed fp32 weight scale. +__global__ __launch_bounds__(kBlockThreads) void identity_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]; + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { + dst[i] = src[i]; } } -// 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, +// Iteration 17 exact-guard n-major transpose pack: +// packed[n * K + k] = raw[k * N + n] ((K, N) == (4096, 1280) only) +// The pack runs once in launch_pack_w8a8_weight, outside the timed operator +// region and outside Graph capture, so a plain 32 x 32 LDS-tiled transpose is +// more than enough. Same byte count and the same packed_weight buffer as the +// identity pack; the packed layout is opaque per the API contract and every +// consumer below decodes this permutation for the guard pair. +__global__ __launch_bounds__(256) void pack_nmajor_i8_kernel( + const int8_t* __restrict__ src, // raw weight [K, N] row-major + int8_t* __restrict__ dst, // packed n-major [N, K] row-major 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) { + __shared__ int8_t tile[32][33]; + const int tx = static_cast(threadIdx.x); // 0..31 + const int ty = static_cast(threadIdx.y); // 0..7 + const int col0 = static_cast(blockIdx.x) * 32; // N (output row) + const int row0 = static_cast(blockIdx.y) * 32; // K (input row) + // Load tile[k_local][n_local] = src[(row0 + k_local) * n + col0 + n_local]. + for (int j = ty; j < 32; j += 8) { + const int row = row0 + j; + if (row < k && col0 + tx < n) { + tile[j][tx] = src[static_cast(row) * n + (col0 + tx)]; + } + } + __syncthreads(); + // Store its transpose: dst[(col0 + n_local) * k + row0 + k_local]. + for (int j = ty; j < 32; j += 8) { + const int col = col0 + j; + if (col < n && row0 + tx < k) { + dst[static_cast(col) * k + (row0 + tx)] = tile[tx][j]; + } + } +} + +void launch_identity_i8( + const int8_t* src, + int8_t* dst, + int64_t count, + hipStream_t stream) { + if (count <= 0) { return; } - const int kk = static_cast(linear / n); - const int nn = static_cast(linear - static_cast(kk) * n); - dst[static_cast(nn) * k + kk] = src[linear]; + const int64_t blocks = (count + kBlockThreads - 1) / kBlockThreads; + hipLaunchKernelGGL( + identity_i8_kernel, + dim3(static_cast(blocks)), + dim3(kBlockThreads), + 0, + stream, + src, + dst, + count); +} + +void launch_pack_nmajor_i8( + const int8_t* src, + int8_t* dst, + int k, + int n, + hipStream_t stream) { + if (k <= 0 || n <= 0) { + return; + } + const dim3 block(32, 8); + const dim3 grid( + static_cast((n + 31) / 32), + static_cast((k + 31) / 32)); + hipLaunchKernelGGL( + pack_nmajor_i8_kernel, grid, block, 0, stream, src, dst, k, n); +} + +void launch_identity_f32( + const float* src, + float* dst, + int64_t count, + hipStream_t stream) { + if (count <= 0) { + return; + } + const int64_t blocks = (count + kBlockThreads - 1) / kBlockThreads; + hipLaunchKernelGGL( + identity_f32_kernel, + dim3(static_cast(blocks)), + dim3(kBlockThreads), + 0, + stream, + src, + dst, + count); } } // namespace -// --------------------------------------------------------------------------- -// Stable host launch symbols (declared in csrc/bindings.cpp). -// --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -765,126 +1106,145 @@ extern "C" void launch_w8a8_gemm( int n, int k, hipStream_t stream) { - auto* out_bf16 = static_cast(out); + // The timed op may only use the caller-provided out/workspace on the + // caller stream; no host/device sync and no allocation happen here. + if (a == nullptr || b == nullptr || x_scale == nullptr || + weight_scale == nullptr || out == nullptr) { + return; + } + if (m <= 0 || n <= 0 || k <= 0) { + 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)). - if (m == 16 && n == 1280 && k == 4096) { - const int num_tiles = n / kDummaN; // 1280 / 16 = 80 + __hip_bfloat16* out_bf16 = reinterpret_cast<__hip_bfloat16*>(out); + + // Exact-shape DUMMA split-K specialization (iteration 8): the only call + // site of w8a8_dumma_m16_splitk_kernel. Multi-N-tile reuse plus whole-slice + // A+B LDS staging, with two 64-lane wavefronts per block (iteration 8): + // wave w owns one of the block's kTilesPerWave = 2 adjacent 16-wide N tiles + // and runs its own accumulator chain, while the block stages its A slice and + // its whole 32-column B strip into LDS (partitioned across the two waves) + // and runs an LDS-only K loop; since iteration 13 a 512-row slice is staged + // as two ascending 256-row phases. The K dimension is split S ways, so + // grid = (N/(16*T))*S blocks of 128 threads: the 8-way split gives + // 40*8 = 320 blocks for the 120 device CUs with uniform 16-stage (512-row) + // slices, each staged as two 256-row phases in the iteration-17 + // 15,104 B/block LDS footprint (four resident blocks/CU by capacity). S is + // clamped to the workspace capacity (the API guarantees 16 int32 planes) and + // to the number of 32-K stages; a clamped split whose phase exceeded the + // staging capacity would use the kernel's byte-for-byte direct-global loop + // instead (unreachable at the guard, where every phase is exactly 256 rows). + // Guarded by the exact (m, n, k) + // triple so every other shape -- notably the paired M=2 API shape with the + // same (N, K) -- keeps the scalar fallback below; if no usable workspace is + // supplied the exact shape also falls through. Iteration 9 changes only the + // second timed dispatch: at the unclamped guard split the 8-plane combine + // runs as a 2-D one-wave grid (n/64 x m blocks of 64 lanes, one output + // element per lane, all 8 plane loads independent in registers) while any + // runtime-clamped split_k < 8 keeps the original generic combine launch. + if (m == 16 && n == 1280 && k == 4096 && workspace != nullptr) { + const int64_t plane_elems = static_cast(m) * n; 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; - } - } + plane_elems * static_cast(sizeof(int32_t)); + const int total_stages = k / kDummaK; + int split_k = kSplitKTarget; + if (plane_bytes > 0) { + const int64_t max_split = workspace_bytes / plane_bytes; + if (split_k > max_split) { + split_k = static_cast(max_split < 0 ? 0 : max_split); } + } else { + split_k = 0; + } + if (split_k > total_stages) { + split_k = total_stages; } - // 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 (split_k >= 1) { + int32_t* partials = reinterpret_cast(workspace); + // 32-column groups: n / 32 = 40 for the guarded exact shape. Both + // factors are exact divisors of the guard's (n, k), so no tail path is + // needed. + const int groups = n / (kDummaN * kTilesPerWave); + const dim3 grid(static_cast(groups * split_k)); + const dim3 block(static_cast(kDummaWaveSize * kTilesPerWave)); + hipLaunchKernelGGL( + w8a8_dumma_m16_splitk_kernel, + grid, + block, + 0, + stream, + a, + b, + partials, + n, + k, + split_k, + groups); + + // Iteration 9: the unclamped guard split (split_k == 8 since iteration + // 13) launches the combine on the 2-D one-wave grid (20 column groups x + // 16 rows = 320 blocks of 64 lanes) that gives every lane one output + // element and all eight plane loads as independent register reads; a + // runtime-clamped split_k < 8 keeps the original linear 128-thread launch + // and the kernel's generic dependent-plane loop. + const int total = m * n; + if (split_k == kCombineSplits) { + hipLaunchKernelGGL( + w8a8_splitk_combine_kernel, + dim3(static_cast(n / kCombineColsPerBlock), + static_cast(m)), + dim3(static_cast(kCombineWaveSize)), + 0, + stream, + partials, + x_scale, + weight_scale, + out_bf16, + total, + n, + split_k); + } else { + const int combine_blocks = + (total + kCombineThreads - 1) / kCombineThreads; + hipLaunchKernelGGL( + w8a8_splitk_combine_kernel, + dim3(static_cast(combine_blocks)), + dim3(static_cast(kCombineThreads)), + 0, + stream, + partials, + x_scale, + weight_scale, + out_bf16, + total, + n, + split_k); } - } - 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); + const int64_t total = + static_cast(m) * static_cast(n); + const int64_t blocks = (total + kBlockThreads - 1) / kBlockThreads; + + // Generic scalar fallback for every unmatched (m, n, k), including the + // paired M=2 API shape with the same (N, K). + hipLaunchKernelGGL( + w8a8_scalar_kernel, + dim3(static_cast(blocks)), + dim3(kBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_bf16, + m, + n, + k); } extern "C" void launch_pack_w8a8_weight( @@ -895,47 +1255,28 @@ 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); + if (raw_weight == nullptr || weight_scale == nullptr || + packed_weight == nullptr || packed_weight_scale == nullptr) { + return; + } + if (k <= 0 || n <= 0) { + return; } - 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 weight_count = + static_cast(k) * static_cast(n); + if (k == kGuardK && n == kGuardN) { + // Iteration 17: the exact-guard weight pair is transposed into the + // n-major layout packed[n*K + k] = raw[k*N + n] consumed by the guard + // DUMMA kernel's n-major LDS stage and by every fallback that decodes the + // packed buffer for this (K, N). Same element count and buffer as the + // identity pack. + launch_pack_nmajor_i8(raw_weight, packed_weight, k, n, stream); + } else { + // Generic fallback for every unmatched (K, N): identity copy. + launch_identity_i8(raw_weight, packed_weight, weight_count, stream); + } + launch_identity_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/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..b4247bad 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,392 +1,206 @@ -// @@variant shape=hy3_tp8_shared_gate_up_proj_m16 commit=2321c42b8618aae23700bc2cf3532b59c8f6c8d8 added=2026-08-27 -// median_us=10.38 p90_us=10.42 -// source=hy3-dsh-tp8-m16-2-dc53295a +// @@variant shape=hy3_tp8_shared_gate_up_proj_m16 commit=d36c2aa242f6657f2fbb136f2edc7e0bf394900a added=2026-09-17 +// median_us=9.17 p90_us=9.427 speedup=7.255 baseline_us=66.53 +// source=hy3-tp8-m16-test9-16-619177f8 // MetaInfer W8A8 INT8 GEMM - HIP implementation (gfx928 / K500SM_AI). // -// worker_2 iteration 8: split-K=15 one-wave LDS-staged partial GEMM + -// split-15 combine for the assigned shape hy3_tp8_shared_gate_up_proj_m16: -// M = 16, N = 384, K = 4096 (M*N = 6144 output elements) +// worker_2 iterations 3-5 for hy3_tp8_shared_gate_up_proj_m16: +// M = 16, N = 384, K = 4096. Iteration 3 = mandated architecture round +// (accepted). Iteration 5 = mandated packed-weight/staging round: +// the DUMMA-fragment-order weight pack + register-direct B transport +// (the B LDS stage of iteration 3 is removed). // -// Operator contract: -// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// Contract: +// out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) // * x_scale[m, 0] * weight_scale[n, 0] ) // -// Iteration 1 = minimal gfx928 DUMMA INT8 bootstrap (exact mandatory tile): -// * one 64-lane wavefront per block, one 16x16 output tile per wave, -// explicit int32 accumulation, grid = N/16 = 24 blocks, direct global -// fragment loads, LDS-drain epilogue; measured 117.91 us median with -// only 24 of 120 CUs active (24 waves, 1 per CU, no co-residency). +// State after iteration 2: the 16-way uniform split-K launch measured +// 11.350 us median / 11.389 us p90 with grid = (N/16) x 16 = 384 one-wave +// blocks (3.2 blocks/CU), 9,216 B LDS/block, and a SEPARATE 24-block combine +// kernel (two dispatches per replay). Fresh PMC on that exact source reports +// lds_bank_conflicts = 254,976 (664 per block) and lds_wait_instructions = +// 6,112 concentrated in the B fragment path, because B is staged k-major +// ([k][16], leading dimension 16) and the row_major fragment loader reads each +// lane's 8 elements as 8 ds_read_u8 at stride 16 -> a structural 8-way bank +// conflict per B fragment read (verified in the exact code object: +// ds_read_u8 v3, v38 offset:4096 / 4112 / 4128 / ... stride 16). The A +// fragment already reads one ds_read2_b32 because A is staged row-major with +// the 8 elements consecutive. The same PMC block shows B staging issues 16-B +// loads at a 384-B stride (half of every 32-B sector wasted) and the second +// dispatch pays a full kernel-launch tail after the partial grid drains. // -// Iteration 4 = accepted split-K=10 grid-parallelism probe (19.84 us): -// * K=4096 split into 10 non-uniform 32-aligned slices, grid = 24 N-tiles -// x 10 splits = 240 one-wave blocks (2 per CU); direct global fragment -// loads in the K loop; separate 24-block combine kernel sums the 10 -// int32 planes in ascending split order (bit-identical int32 order); -// both launches inside the timed Graph region. -// * PMC: 49,152 vmem_read instructions for the partial kernel (~205 per -// block = byte-level fragment loads + reassembly, one vmcnt wait per -// load before each mmac), lds_instructions = 0 -> the K loop is -// serialized on global load latency (issue/latency-bound). +// This round is the mandated architecture change (one-wave, ZERO-barrier, +// CU-aligned, non-power-of-two split; trusted probe set [2,5,10,15]) plus ONE +// consolidation of the two measured B/combine bottlenecks, using the exact +// recipe previously validated on this shape: // -// Iteration 6 = accepted split-K=16 LDS-staged kernel (18.73 us median / -// 19.04 us p90): block tile 16x32 (2 waves x 1 N-tile of 16 columns), 128 -// thr, grid = 12 N-groups x split-K=16 = 192 blocks = 1.6 blocks/CU = 3.2 -// waves/CU; whole uniform 256-row slice staged once (A 16x256 stride 272, -// B 32x256 stride 48), one __syncthreads, zero-barrier LDS-only K loop -// (8 steps x 1 v_mmac per wave); workspace = 16 int32 planes = 393,216 B = -// the caller's workspace_split_k_capacity(16,384,4096) budget; combine = -// 24 blocks x 128 thr, ascending split sum; PMC: 192 blocks < 2/CU target. +// * K is split kSplitK = 15 ways with 32-aligned non-uniform slices +// (8 x 288 = 9 DUMMA steps + 7 x 256 = 8 steps; 8*288 + 7*256 = 4096). +// Grid = 24 N-groups x 15 splits = 360 one-wave blocks = EXACTLY 3 +// blocks/CU on the 120 measured device CUs - integer balanced, no tail +// wave, and a shrink of the per-block slice (288/256 vs 256) that keeps +// 3 co-resident blocks inside the LDS budget (3 x 4,868 B << 64 KiB). +// * [iteration 3 - pack superseded by the iteration-5 block below] the +// weight is packed once, OUTSIDE the timed region and outside the +// Graph, by launch_pack_w8a8_weight into the per-tile slice-major-chunk +// layout for (k,n)==(4096,384) ONLY: every 16-B chunk +// holds 16 consecutive k of one n column and the 4,096 chunks of a tile +// are stored in ascending (slice, n, k-chunk) order. Every other (K,N) +// keeps the identity copy. +// * [iteration 3 - B transport superseded by the iteration-5 block below] +// the B slice is staged into an n-major LDS layout b_lds[n][k] (element +// (k,n) at n*304 + k) from the contiguous packed region, so each staged +// wave-load covers 1,024 contiguous bytes (full 32-B sectors, sequential +// DRAM pages) and each lane's 8 B fragment elements are 8 CONSECUTIVE +// bytes -> exactly ONE ds_read2_b32 per DUMMA step (was 8 bank-conflicted +// ds_read_u8). The col_major loader reads the same elements as the old +// row_major loader (x[i] = element (k0+col+i, n_base+row)), so every +// fragment register - and therefore the exact int32 accumulation order - +// is unchanged. +// * A is NOT staged through LDS: each lane loads its own 8-B fragment slice +// for every DUMMA step into a per-lane register array (int2 global loads +// issued back-to-back) and the K loop consumes it straight from +// registers; A is 64 KiB total and L2-resident, and removing its LDS +// round trip drops the per-block LDS footprint to 4,868 B. +// * the 24-block combine is FUSED into the partial kernel as a +// last-arrival tail: every block publishes its int32 plane tile, lane 0 +// releases with __threadfence() + atomicAdd on a per-tile counter, the +// 15th arriver acquires, sums the 15 planes in ascending split order +// (exact integer order), applies the reference float order +// (float(acc) * x_scale[row]) * weight_scale[col], stores bf16 and resets +// its counter. Exactly ONE dispatch per replay. // -// Iteration 7 = measured split-K=10 one-wave grid (this file's round-7 -// candidate, REJECTED by the p90 guard): block tile 16x16 (1 wavefront, 64 -// thr), grid = 24 N-groups x split-K=10 = 240 one-wave blocks = exactly -// 2 blocks/CU; K=4096 -> 10 non-uniform 32-aligned slices (8 x 416 + 2 x -// 384); whole slice staged once (A stride 432, B stride 48; 26,880 B/block -// = 2 x 26,880 = 53,760 B/CU at 2 blocks, fits 64 KiB), 0 s_barrier (the -// single-wave __syncthreads is eliminated; lgkmcnt waits cover the ds -// hazard, verified in the exact gfx928 code object); 10 planes = 245,760 B; -// combine = 24 blocks x 128 thr. Measured median 15.432 us (-21% vs -// iteration 6) but p90 19.434 > best 19.040 -> p90_guard_passed = false; -// the 30 samples were bimodal (4 slow windows 19.4-23.2 us then 26 fast -// windows 15.4-15.9 us): a machine-state artifact, not kernel behavior. +// --------------------------------------------------------------------------- +// Iteration 5 (this revision) - mandated packed-weight/staging round: the +// DUMMA-fragment-order weight pack + register-direct B transport. +// +// The accepted iteration-3 source (digest 77a1ea90..., 9.674 us median / +// 9.751 us p90 / 9.631 us min, 360 one-wave blocks = exactly 3 blocks/CU, +// fused last-arrival tail, ONE dispatch) carried B through LDS: 5 batched +// 16-B global loads -> `s_waitcnt vmcnt(0)` (every B byte had to arrive +// before ANY ds_write) -> 5 ds_write_b128 with per-chunk index math -> 2 full +// `s_waitcnt lgkmcnt(0)` drains -> 9 ds_read2_b32 back into the fragment +// registers. The exact iteration-4 code object shows that whole path +// (load_wait0_ds_write_windows = 5, 15 lgkmcnt(0) drains, lds_instructions +// 5,424 = 15/block, lds_bank_conflicts 14,592, lds_bytes 4,868) while the PMC +// shows only ~199 VALU instructions per 64-lane wave in ~42,829 cycles: with +// just 3 co-resident waves per CU there is no other work to hide a +// serialization stage, so every turn of the B path is exposed latency. // -// Iteration 8 (this file's active M=16 path) = mandatory HIP-only -// occupancy round: tune one occupancy limiter (LDS footprint -> blocks/CU) -// with PMC evidence. The split-K=10 kernel's 26,880 B LDS/block caps it at -// exactly 2 blocks/CU (3 x 26,880 = 80,640 > 64 KiB). split-K=15 (trusted -// probe set [2,5,10,15]) shortens every slice (288/256 vs 416/384 rows), -// shrinking the per-block LDS footprint to 18,688 B -> 3 x 18,688 = 56,064 -// B/CU <= 64 KiB, so 3 one-wave blocks/CU become legal, and the grid = -// 24 N-groups x 15 splits = 360 one-wave blocks = EXACTLY 3 blocks/CU on -// 120 CUs (above the 2-blocks/CU latency-hiding target, integer-balanced, -// no tail wave). No repeated HBM reads: B stays once-read (1,572,864 B), -// A stays L2-hot, every block still stages its own slice exactly once. -// * block tile 16x16 (1 wavefront x 1 N-tile), 64 thr; grid = 24 x 15 = -// 360 one-wave blocks = exactly 3 blocks/CU; -// * K=4096 -> 15 non-uniform 32-aligned slices: 8 x 288 (9 DUMMA steps) -// + 7 x 256 (8 steps); k0-ascending within each slice and ascending -// split sum in the combine preserve the bit-identical int32 order; -// * whole slice staged once into bank-skewed LDS (A 16x288 stride 304, -// B 288x16 stride 48; 18,688 B/block), batched 16-B vector loads, one -// __syncthreads (0 s_barrier for the 1-wave block), then a zero-barrier -// LDS-only K loop (9/8 steps x 1 A-frag + 1 B-frag ds_read + 1 v_mmac); -// * workspace = 15 int32 planes [15][16][384] = 368,640 B (<= the -// caller's 393,216 B capacity; guard updated, unsplit DUMMA fallback -// when workspace < 15 planes); combine = 24 blocks x 128 thr, all 15 -// plane loads issued into a register array before the ascending sum; -// * total replay traffic ~2.32 MB (B 1.57 MB + planes RW 737,280 B + -// out 12,288 B) vs 2.47 MB for the accepted S=16 kernel (-6%); vs S=10 -// the plane RW grows +245,760 B (+50%) - the combine cost of more -// splits being benchmarked against the occupancy gain. +// This round changes ONE thing (the weight pack), and one stage disappears +// with it: launch_pack_w8a8_weight (outside the timed region and outside the +// Graph, same 1,572,864-byte buffer, (k,n)==(4096,384) ONLY) now emits the +// weight in the exact m16n16k32 matrix_b fragment order of the 64-lane +// wavefront, so a block's whole B slice becomes a straight sequential +// register stream and the B LDS stage is deleted: // -// Iteration 9 (this file's active M=16 path) = HIP-only two-launch attack. -// Round 8 measured 15.529 us median / 16.024 us p90 and its falsifiable -// branch ("if median stays >= 15.4, the next round attacks the two-launch -// structure - fused/single-kernel combine - instead of geometry") fired. -// Change: fuse the 24-block combine INTO the partial kernel as a -// last-arrival tail using the validated TP4 qkv lineage protocol - every -// block writes its 16x16 int32 plane tile, __threadfence() (release), then -// atomicAdd(&counters[tile], 1); the block whose atomicAdd returns -// kSplitK-1 (the 15th arriver of its N-group) __threadfence()s (acquire), -// sums the 15 planes in ascending split order (exact int32 order preserved -// -> outputs bit-identical), applies both scales, stores bf16, and resets -// its counter to 0 (atomicExch) so the next replay restarts from 0. The 24 -// counters live in the workspace tail (96 B right after the 15 planes; -// guard updated to 368,736 B) and are zeroed once by the host before the -// first launch (a static-flag hipMemsetAsync on the caller stream - issued -// in the eager precheck, outside the timed Graph in the normal flow). -// Result: ONE kernel launch per replay (grid 360 one-wave blocks, 3/CU, -// LDS unchanged 18,688 B); the combine no longer pays a second dispatch -// and each N-group's combine tail overlaps the still-running partials of -// other groups instead of waiting for full grid completion + relaunch. +// * the 65,536 B of one 16-column tile are the 15 split slices in +// ascending split order; inside a slice the bytes are the fragment slots +// of the wave, TWO k32 steps per 1,024-B pair slot: byte = +// slice_base(s) + (step>>1)*1024 + lane*16 + (step&1)*8 + i with +// lane = ((k - k_base)>>3 & 3)*16 + (col&15), i = (k - k_base)&7; +// * a 9-step (288-k) slice is 4 pair slots plus one 512-B single slot for +// the odd 9th step (8 B/lane at lane*8); an 8-step (256-k) slice is 4 +// pair slots; 8*4,608 + 7*4,096 = 65,536 B per tile exactly, so the byte +// count and the graph-stable buffer are unchanged; +// * the timed kernel then fetches B fragment registers with 4 x int4 (one +// per step pair, 16 B/lane, 1,024 CONTIGUOUS bytes per wave-load) plus, +// for the 9-step slices, one int2 (8 B/lane, 512 contiguous bytes) - the +// SAME 4/5 load instructions and the same bytes per lane as the accepted +// staged path, but each load address is lane*16/lane*8 inside a +// contiguous slot: no chunk-index math, no ds_write, no __syncthreads, no +// lgkmcnt drain and no ds_read is needed, and the K loop can consume each +// step pair as soon as it arrives instead of after the whole slice; +// * b_slice[s].x[i] is the same element the accepted col_major LDS loader +// produced, x[i] = element(k0 + (lane>>4)*8 + i, n_base + (lane&15)), so +// every matrix_b fragment register - and therefore every v_mmac input and +// the exact int32 accumulation order - is unchanged; +// * A is unchanged (register-only, 8/9 batched int2 loads, zero LDS), the +// grid is unchanged (24 N-groups x 15 splits = 360 one-wave blocks = +// exactly 3 blocks/CU), and the partial-plane store, the fused +// last-arrival combine tail, the workspace guard (15 planes + 24 counters +// = 368,736 B <= 393,216 B) and both fallbacks are unchanged; the generic +// scalar and unsplit M=16 DUMMA fallbacks decode the new pack elementwise +// through the updated packed384_byte_offset, so they stay bit-exact +// (including the paired M=2 shape). // -// Repair 1/4 (correctness): the round-9 draft let ALL 64 lanes of every -// block execute atomicAdd(&counters[tile], 1) (64 increments per block = -// 960 per tile per replay), so the kSplitK-1 condition fired on the 15th -// increment inside the FIRST-arriving block - before its 14 siblings had -// stored their planes - and the unguarded atomicExch reset made several -// blocks combine per replay. Exact check failed: 6048/6144 mismatched, -// first_mismatch (0,0) actual=0.0 (premature combiner summed unwritten, -// zero planes). Fix: exactly one arrival per block (lane 0 fence + -// atomicAdd, LDS s_is_last flag, barriers) + lane-0-only atomicExch reset, -// mirroring the validated TP4 qkv fused tail. GEMM/plane-store body and the -// ascending 15-plane sum are unchanged (bit-identical int32 order). +// Predicted signature if the mechanism is real: LDS per block drops from +// 4,868 B to the 4-B s_is_last flag; lds_instructions 5,424 -> ~0; +// lds_bank_conflicts 14,592 -> 0; lds_wait_instructions 461 -> 0; +// group_segment_fixed_size ~4,868 -> ~4 B; waitcnt drops by ~20 (all +// lgkmcnt drains and the staging vmcnt(0) pair vanish); vmem_read_instructions +// stays ~5,832 (same B bytes and the same 4/5 B load instructions per block); +// memory_traffic.total_bytes_per_operator_replay stays ~2.05 MB. If the +// median does not improve, the B LDS round trip was fully hidden by the 3 +// co-resident waves and the kernel sits at its launch/streaming floor. +// --------------------------------------------------------------------------- // -// Iteration 10 (this file's active M=16 path) = final HIP-only -// consolidation round (raw inline asm not allowed by the ISA policy: -// plateau=false, valid HIP rounds < 8). Round 9 measured 14.516 us median / -// 14.544 us p90 and its falsifiable branch ("if median stays >= 14.5 ... the -// single-wave combine tail (60 loads/lane) is the new critical path - the -// next round vectorizes the combiner's plane loads (int4) or shrinks the -// plane count instead") fired (14.516 >= 14.5). Change: vectorize ONLY the -// fused combine tail's plane loads - each lane now owns 4 CONSECUTIVE -// columns (row = lane>>2, col0 = (lane&3)*4) and reads each of the 15 -// planes with one aligned int4 (16-B) load instead of 4 scalar 4-B loads -// (60 -> 15 vmem loads per lane, 4x fewer tail load instructions, same -// bytes, same L2-request footprint); weight_scale is read once as float4 -// and the 4 adjacent bf16 stores merge into one 8-B store. Per-element -// int32 accumulation order is unchanged (ascending split sum per output -// element) -> outputs remain bit-identical (0 mismatches). Everything else -// is untouched: grid = 24 N-groups x 15 splits = 360 one-wave blocks = 3 -// blocks/CU, per-block LDS 18,688 B, zero-barrier LDS-only K loop, -// workspace = 15 planes + 24 counters = 368,736 B, one launch per replay, -// exact (m,n,k) guard, unsplit DUMMA fallback, generic scalar fallback. -// Falsifiable: (1) median_us < 14.516 (target ~13.6-14.2) with p90 <= 14.544 -// (guard vs the accepted best); (2) 0 mismatches + graph_capture_passed; -// (3) PMC: grid 360 / workgroup 64 / LDS 18,688 / s_barrier 0 unchanged, -// arch_vgpr stays ~40-60 with 0 spills, the tail portion of vmem_read drops -// ~4x (staging ~3,264 unchanged; tail 60 -> 15 loads/lane); (4) replay -// traffic unchanged ~2.07 MB. If median stays >= 14.516 with p90 <= 14.544, -// the tail was not the critical path and the kernel sits at its -// streaming/launch floor -> the round documents the plateau evidence. +// Iteration 7 (this revision) - mandated pipeline/buffering decision plus the +// removal of the only avoidable term left in the HBM byte ledger. // -// Iteration 11 (this file's active M=16 path) = final conditional inline-asm -// round -> mandatory HIP-only consolidation (raw asm is NOT allowed: isa_policy -// says phase=hip_only, plateau=false, raw_inline_asm_allowed=false, -// valid_hip_rounds=7 < 8 - the control plane has not confirmed a HIP plateau). -// Round 10 measured 13.925 us median / 13.972 us p90 (< 14.516 target), so the -// int4-vectorized tail branch did NOT fire and the tail is not the critical -// path anymore. The accepted kernel's fresh exact gfx928 code object exposes -// the next concrete compiler limitation: the slice staging is SERIALIZED - -// every global_load_dwordx4 is followed by s_waitcnt vmcnt(0) before its -// ds_write_b128 (load_wait0_ds_write_windows = 18 across the two slice -// instantiations, 10 long + 8 short), so each block pays ~10/8 full -// global-memory latencies one after another before the K loop's first -// ds_read (the last write is followed by lgkmcnt(0) waits). Change: in -// split15_slice_body, issue ALL staging global loads into one register array -// (int4 s_buf[10/8]) BEFORE any ds_write - the compiler is then forced to put -// the loads in flight back-to-back and move the vmcnt waits to the write -// side, collapsing the staging critical path to one memory latency. Same -// addresses, same per-lane order, same LDS layout, same __syncthreads, same -// zero-barrier LDS-only K loop, same fused last-arrival tail: the int32 -// accumulation order is untouched (outputs bit-identical, 0 mismatches -// expected). Everything else is untouched: grid = 24 N-groups x 15 splits = -// 360 one-wave blocks = exactly 3 blocks/CU (trusted probe set member 15), -// per-block LDS 18,688 B (still the occupancy binder: 3 x 18,688 = 56,064 B -// <= 64 KiB), workspace = 15 int32 planes + 24 arrival counters = 368,736 B -// (<= 393,216 B budget; unsplit DUMMA fallback when smaller), one launch per -// replay, fence/atomic/barrier protocol unchanged, exact (m,n,k) guard, -// generic scalar fallback for every unmatched shape including the paired M=2. -// Falsifiable predictions: (1) median_us < 13.925 (target ~11.5-12.8) with -// p90_us <= 13.972 (guard vs the accepted best); (2) correctness 0 mismatches -// + graph_capture_passed (single-launch structure, write-in-place planes, -// self-resetting counters, no clear); (3) ISA: the first staging vmcnt(0) -// now covers ~7 outstanding global_load_dwordx4 (old kernel: exactly 1) - -// memory-level parallelism per block rises 1 -> ~7-10 and the staging -// critical path collapses from ~10 serialized latencies to ~1-2; vmcnt(0) -// 20 -> ~19, arch_vgpr ~69-75 (locally verified 69 VGPR / 40 SGPR) with 0 -// spills, grid 360 / LDS 18,692 / s_barrier 0 / v_mmac 17 / ds_write_b128 18 -// / ds_read2_b32 17 / ds_read_u8 136 unchanged; (4) PMC vmem_read ~4,392 and -// replay traffic ~2.06 MB unchanged (same loads, same bytes). If -// median stays >= 13.925 with p90 <= 13.972, the 3-blocks/CU overlap already -// hid the staging latency and the kernel sits at its streaming/launch floor - -// the round then documents the plateau evidence for this shape. +// MANDATED DECISION (single vs double buffering). The gate is evaluated on the +// exact accepted source: K = 4096 >= 1024 (true), l2_hit_rate_percent 51.521 +// < 70 (true), doubled LDS budget 2 x 4 B = 8 B << 48 KiB (true). The round +// nevertheless keeps SINGLE buffering, because the kernel has no LDS pipeline +// stage to buffer: LDS is 4 B/block, the fused tail's s_is_last flag +// (lds_instructions 720 = 2/block, lds_bank_conflicts 0, lds_wait 8). Both +// operands travel global -> fragment registers directly; the exact gfx928 +// code object of the accepted source shows 9 x global_load_dwordx2 (A) and +// 4 x global_load_dwordx4 + 1 x global_load_dwordx2 (B) issued back-to-back +// ahead of the chain, with the compiler already software-pipelining the +// consumption (s_waitcnt vmcnt(7) -> first v_mmac ... vmcnt(0) -> last +// v_mmac). BARRIERS PER K STEP = 0 for the current kernel (zero s_barrier in +// the K loop). The double-buffered alternative is the iterations 1-3 LDS +// stage, which needed 2 s_barrier and 2 full s_waitcnt lgkmcnt(0) per stage +// and whose deletion in iteration 5 measured +3.2% (9.674 -> 9.365 us). The +// single/double comparison therefore rejects double buffering on measured +// evidence, not on the gate. // -// Iteration 15 (this file's active M=16 path) = final conditional inline-asm -// round -> mandatory HIP-only consolidation (raw asm is NOT allowed: isa_policy -// says phase=hip_only, plateau=false, raw_inline_asm_allowed=false - the -// control plane has NOT confirmed a HIP plateau: the three recent valid -// improvements were -28.62% / -0.033% / +3.86% (iterations 12/13/14), the -// -28.62% being the rejected round-12 regression and +3.86% the round-14 -// packed-B median gain, and no prior ISA-guided round recorded a compiler -// limitation plus target instructions). Round 14 (n-major-packed B staging -// loads only) measured median 12.512 us / p90 13.064 us - best median ever, -// rejected only by the p90 guard (13.064 > 13.052) - and the source was -// restored to the accepted iteration-11 digest. The accepted kernel's exact -// gfx928 code object exposes the remaining half of the B path: the K loop -// still reads each lane's 8 B-fragment elements as 8 ds_read_u8 from the -// k-major stride-48 LDS (4-way bank conflicted; the persistent PMC signature -// since iteration 9), while the A fragment already reads its 8 consecutive -// bytes as ONE ds_read2_b32. The TP4 gate_up/down_proj lineage validated the -// fix for exactly this ("one ds_read_b64 per step, was 8 ds_read_u8"; rule 3: -// 'Pack the cold once-read weight (B) ... into the fragment's layout'). -// Change (one consolidation): complete the round-14 B-pack into a per-tile -// n-major transpose packed[tile][n][kk] = raw[kk*384 + tile*16 + n] -// ((k,n)==(4096,384) ONLY; every other (k,n) keeps the identity pack), stage -// B into an n-major LDS layout b_lds[n][k] (element (k,n) at n*304 + k; -// column stride 304 = 288 + 16 bank skew, 16-B aligned so the staging writes -// stay ds_write_b128), and load the B fragment with the library col_major -// loader: a.x[i] = p[row*304 + col + i] reads the lane's 8 elements as 8 -// CONSECUTIVE bytes -> one ds_read2_b32 per step (was 8 ds_read_u8 at stride -// 48). Same elements, same x[i] register assignment (x[i] = element -// (k0+col+i, n_base+row) in both layouts), same v_mmac, so the int32 -// accumulation order is untouched (k-ascending within each slice, ascending -// split sum - outputs bit-identical, 0 mismatches expected). The B staging -// loads also get the round-14 locality win: each lane holds one 16-B -// row-chunk of 16 consecutive k values of one n column (chunk c -> n = c&15, -// k-chunk = c>>4) - 64 lanes cover 16 x 64-B contiguous sectors per wave-load -// (was 64 scattered 16-B requests), sequential HBM pages, no straddles. LDS -// per block drops 18,688 -> 9,728 B (A 16x304 = 4,864 + B 16x304 = 4,864), -// so 3 x 9,728 = 29,184 B/CU << 64 KiB (LDS is no longer the occupancy -// binder, but the grid stays 360 = exactly 3 blocks/CU, so residency is -// unchanged); per-block LDS instructions drop ~4x (136 ds_read_u8 + 17 -// ds_read2_b32 -> 34 ds_read2_b32), attacking the per-CU LDS-pipe pressure -// that round 13's evidence identified as the shared per-CU resource. -// Everything else is untouched: grid = 24 N-groups x 15 splits = 360 -// one-wave blocks = exactly 3 blocks/CU (trusted probe-set member 15), -// workspace = 15 int32 planes + 24 arrival counters = 368,736 B (<= the -// caller's 393,216 B budget; unsplit DUMMA fallback when smaller), one launch -// per replay, fence/atomic/barrier protocol unchanged (lane-0 atomicAdd, LDS -// s_is_last flag, lane-0 atomicExch reset), exact (m,n,k) guard, generic -// scalar fallback for every unmatched shape including the paired M=2 (both -// fallbacks decode the packed layout for (k,n)==(4096,384): the scalar path -// reads element (kk, col) at (col>>4)*65536 + (col&15)*4096 + kk; the unsplit -// DUMMA path loads B with the row_major loader at ldm = K from -// packed + tile*65536 + k0, giving p[(col+i)*4096 + row] = element -// (k0+col+i, n_base+row) - identical values to the raw layout). -// Falsifiable predictions: (1) median_us < 12.995 (target ~12.0-12.7; the -// round-14 pack alone measured 12.512 and this also cuts the per-block LDS -// read instructions 8x and the per-CU LDS-pipe work ~4x) with p90_us <= -// 13.052 (guard vs the accepted best p90); (2) correctness 0 mismatches -// (bit-identical int32 order) and graph_capture_passed (single-launch -// structure unchanged, write-in-place planes, self-resetting counters, pack -// out-of-timed/out-of-Graph, no clear); (3) PMC/ISA: grid_blocks 360 / -// workgroup_size 64 / s_barrier 0 / v_mmac 17 / ds_write_b128 18 unchanged, -// ds_read_u8 136 -> 0, ds_read2_b32 17 -> 34 (A + B), group_segment_fixed_size -// 18,692 -> ~9,732, arch_vgpr ~69 +/- a few with 0 spills, occupancy 3 -// blocks/CU (now grid-bound), lds_instructions ~31,632 -> ~20k, -// lds_bank_conflicts ~63,744 -> ~15-25k (the B u8 4-way reads disappear), -// vmem_read_instructions ~4,392 unchanged (same loads, same bytes); -// (4) memory_traffic.total_bytes_per_operator_replay unchanged ~2.06 MB (no -// HBM delta). If median stays >= 12.995 with p90 <= 13.052, the LDS pipe was -// not on the critical path either (the kernel is at its launch/memory floor) -// and this final round documents the plateau evidence for this shape instead -// of claiming a gain. +// The round's ONE change attacks the partial planes. The accepted source's +// memory_traffic is read 1,665,856 B (= weight 1,572,864 + A 65,536 + 27,456) +// and write 381,024 B (= partials 368,640 + out 12,288 + counters 96), total +// 2,046,880 B per replay at 218.56 GB/s counter-derived; the 368,640 B of +// partial-plane write-back is 18% of all HBM bytes for values that are read +// back exactly once by the fused tail. Those 15 int32 planes are replaced by +// ONE hot per-tile int32 accumulator (24 tiles x 256 int32 = 24,576 B, +// L2-resident) updated with global atomicAdd, and the fused last-arrival tail +// reads that single accumulator, applies the reference float order +// (float(acc) * x_scale[row]) * weight_scale[col], stores bf16 and resets the +// tile to zero. Integer addition is exact and order-independent, so the summed +// int32 - and therefore every output bf16 - is bit-identical to the 15-plane +// ascending-sum chain. Unchanged: grid (24 N-groups x 15 splits = 360 one-wave +// blocks = exactly 3.0 blocks/CU on 120 CUs), slices (8 x 288 + 7 x 256, +// 32-aligned), weight pack, A register path, counter protocol, fallbacks. The +// workspace requirement drops from 368,736 B to 24,672 B, which frees the +// 393,216-B partial-plane cap that held the split at 15 - the hand-off for a +// later blocks/CU probe. // -// Iteration 17 (this file's active M=16 path) = HIP-only B-pack ORDER -// consolidation (isa_policy: phase=hip_only, plateau=false, -// raw_inline_asm_allowed=false, skill_allowed=false - the three recent valid -// HIP rounds are NOT all within [-2%, +2%): -0.033% / +3.86% / +8.35%, so the -// plateau is not proven and one more HIP-only change is mandated; the -// round-16 attempt died in the agent infrastructure before a valid proposal, -// so this round returns to the accepted iteration-15 source (digest -// b39a985c...) and starts ONE new bounded experiment). Round 15 (n-major B -// pack + n-major LDS) measured median 11.994 us / p90 12.293 us - accepted - -// and left the remaining half of the B memory path on the table: the n-major -// pack places chunk (n, k-chunk) at n*4096 + k-chunk*16, so each staged B -// wave-load touches 16 x 64-B sectors at 4,096-B stride (16 scattered L2 -// requests / up to 16 DRAM pages) - the round-14 contiguous-load locality -// was traded away to buy the one-ds_read2_b32 B fragment read. Change (one -// consolidation, pack ORDER only): the pack kernel now emits the SAME 16-B -// chunks (16 consecutive k of one n column) in ascending (slice, n, k-chunk) -// order - chunk index = slice_chunk_base(s) + n*cps(s) + k-chunk, byte -// address = tile*65536 + chunk*16 - so each staged wave-load (64 consecutive -// chunks) covers 1,024 CONTIGUOUS bytes (sequential L2 sectors and -// sequential DRAM pages; was 16 scattered 64-B sectors at 4,096-B stride). -// Chunk content, the n-major LDS destination b_lds[n][k] (n*304 + k), the -// one-ds_read2_b32 B fragment read, the v_mmac count, the zero-barrier K -// loop and the fused last-arrival tail are all UNTOUCHED: every staged byte -// and every fragment register is identical to iteration 15, so the int32 -// accumulation order is bit-identical (0 mismatches expected). The -// long-slice staging write index changes from (c&15, c>>4) to (c/18, c%18) -// (the short 256-row instantiation keeps (c>>4, c&15); divisions resolved -// by the compiler, staging path only); the generic scalar fallback and the -// unsplit DUMMA fallback decode the new pack order through a shared -// packed384_byte_offset helper (per-element gather in the unsplit path, -// fallback-only: x[i] = element (k0 + (lane>>4)*8 + i, n_base + (lane&15)), -// the same fragment registers as the library row_major loader). -// Everything else is untouched: grid = 24 N-groups x 15 splits = 360 -// one-wave blocks = exactly 3 blocks/CU (trusted probe-set member 15), -// per-block LDS 9,732 B, workspace = 15 int32 planes + 24 arrival counters -// = 368,736 B (<= 393,216 B budget; unsplit DUMMA fallback when smaller), -// one launch per replay, fence/atomic/barrier protocol unchanged, exact -// (m,n,k) guard, generic scalar fallback for every unmatched shape -// including the paired M=2. -// Falsifiable predictions: (1) median_us < 11.994 (target ~11.2-11.8; the B -// staging wave-loads become fully sequential, removing the 16-scattered- -// sector DRAM page pattern) with p90_us <= 12.293 (guard vs the accepted -// best p90); (2) correctness 0 mismatches (bit-identical int32 order) and -// graph_capture_passed (single-launch structure unchanged, write-in-place -// planes, self-resetting counters, pack out-of-timed/out-of-Graph, no -// clear); (3) PMC/ISA: grid_blocks 360 / workgroup_size 64 / lds_bytes -// 9,732 / s_barrier 0 / v_mmac 17 / ds_read2_b32 34 / ds_write_b128 18 -// unchanged, arch_vgpr ~69-72 with 0 spills, vmem_read ~4,392 unchanged -// (same loads, same bytes), the B staging address stream becomes sequential -// ((k_base + c)*16 per chunk) so l2_misses ~32,321 may drop (fewer sector -// straddles / page effects) while lds_bank_conflicts may shift (the -// long-slice staging write bank pattern changes; instruction count and -// bytes identical); (4) memory_traffic.total_bytes_per_operator_replay -// unchanged ~2.06 MB (same bytes, better DRAM locality). If median stays -// >= 11.994 with p90 <= 12.293, the B staging DRAM pattern was not on the -// critical path either (the kernel is at its launch/latency floor) and this -// round documents the next plateau datum for this shape. +// Predicted signature if the mechanism is real: write_bytes_per_operator_replay +// 381,024 -> ~37,000 (out 12,288 + counters 96 + accumulator write-back +// <= 24,576), read_bytes 1,665,856 -> ~1,638,400, total 2,046,880 -> ~1,675,000 +// (-18%); the kernel body's 4 global_store_dword per lane become 4 no-return +// global_atomic_add; lds_bytes stays 4 and ds_read/ds_write stay 1/1 (the +// flag); vmem_read_instructions unchanged; the last arriver's 15 int4 plane +// reads collapse to one int4 accumulator read per lane. If the median does not +// move, the 368,640 B of partial write-back was absorbed by L2 and the kernel +// is latency-bound at 3 warps/CU - which is what the freed 24,672-B workspace +// then lets the next round attack with a >= 4-6 blocks/CU split. +// --------------------------------------------------------------------------- // -// Iteration 21 (this file's active M=16 path) = HIP-only, ONE consolidation: -// A register-only transport (the validated TP4 o_proj lineage recipe -// "register-only K-loop transport with depth-1 prefetch"; raw asm is NOT -// allowed: isa_policy says phase=hip_only, plateau=false, -// raw_inline_asm_allowed=false, skill_allowed=false - the three recent valid -// HIP improvements +3.86% / +8.35% / +5.35% are not within [-2%, +2%), so the -// HIP plateau is not proven and one more HIP-only change is mandated; rounds -// 18/19/20 died in the agent infrastructure before valid proposals, so this -// round returns to the accepted iteration-17 source (digest -// 88ea95a832fbc79e4d625b2aa22df3a0d533d6eb2c6aff08529c7e60a92942c4, verified -// sha256-identical before editing) and starts ONE new bounded experiment - no -// round-18/19/20 partial candidate was repaired or replayed). The accepted -// kernel (11.385 us median / 11.954 us p90) still stages the small L2-hot A -// slice (16 x 288/256 int8) through LDS: per block that is 5/4 16-B global -// loads + 5/4 ds_write_b128 + 9/8 ds_read2_b32 (2-way bank conflicted, -// structural for any 16-B-aligned stride) that the K loop then reads back. -// Every block reads its own A rows exactly once, so the LDS round trip only -// adds instructions and a serialization point; the A values are 64 KiB total -// and L2-resident (re-read by all 360 blocks), so direct global fragment -// loads are cheap. Change (one consolidation): (1) split15_slice_body no -// longer stages A; each lane preloads its own 8-B fragment slice for EVERY -// DUMMA step into a per-lane register array a_slice[kStepsT] (9/8 -// global_load_dwordx2 per lane, issued back-to-back BEFORE any ds_write - -// the iteration-11 register-array batching pattern - so the A slice costs -// ONE memory latency and zero LDS instructions); (2) the K loop consumes the -// A fragments from registers (du_mma_sync(acc, a_slice[s], b_frag, acc)) - -// same elements per lane per step as the LDS row_major loader (verified -// against du_mma.hpp: x[i] = p[row*ldm + col + i], row = lane&15, col = -// (lane>>4)*8, ldm = k), so the int32 accumulation order is untouched -// (k-ascending per split, ascending split sum -> outputs bit-identical, 0 -// mismatches expected); (3) the shared a_lds array is deleted -> per-block -// LDS drops 9,732 -> 4,868 B (B 4,864 + s_is_last 4). B staging, the B -// col_major LDS reads, the zero-barrier structure (one __syncthreads before -// the B reads, unchanged), the slice-major pack, the fused last-arrival tail -// and every fallback are UNTOUCHED. Everything else is unchanged: grid = 24 -// N-groups x 15 splits = 360 one-wave blocks = exactly 3 blocks/CU (trusted -// probe-set member 15; LDS is no longer the binder at 3 x 4,868 = 14,604 -// B/CU), workspace = 15 int32 planes + 24 arrival counters = 368,736 B (<= -// the caller's 393,216 B budget; unsplit DUMMA fallback when smaller), one -// launch per replay, fence/atomic/barrier protocol unchanged (lane-0 -// atomicAdd, LDS s_is_last flag, lane-0 atomicExch reset), exact (m,n,k) -// guard, generic scalar fallback for every unmatched shape including the -// paired M=2 (the packed384_byte_offset decode is untouched). -// Falsifiable predictions: (1) median_us < 11.385 (target ~10.8-11.3; the -// per-block staging loses the A LDS round trip and the K loop loses 9/8 -// LDS reads) with p90_us <= 11.954 (guard vs the accepted best p90); -// (2) correctness 0 mismatches (bit-identical int32 order; M=2 paired -// fallback exact via the packed decode) and graph_capture_passed (single- -// launch structure unchanged, write-in-place planes, self-resetting -// counters, no clear); (3) PMC/ISA: grid_blocks 360 / workgroup_size 64 / -// s_barrier 0 / v_mmac 17 unchanged; group_segment_fixed_size 9,732 -> -// ~4,868 B; ds_write_b128 18 -> ~9/8 (B staging only); ds_read2_b32 34 -> -// ~17 (B only); ds_read_u8 stays 0; the A slice appears as 9/8 batched -// global_load_dwordx2 before the first B ds_write (explicit int2 loads; if -// instead the ISA shows the A loads sunk into the K loop or scalarized to -// global_load_ubyte, the compiler refused the batching/vectorization and -// this round documents that finding); arch_vgpr ~69 -> ~85-95 (a_slice -// holds 18/16 extra VGPR) with 0 spills (per-wave ~90 x 64 = 5,760 VGPR << -// 65,536/CU, occupancy stays 3 blocks/CU = grid-bound); lds_instructions -// ~10,128 -> ~5-6k and lds_bank_conflicts ~29,184 -> ~15k (the A 2-way -// fragment-read conflicts disappear); vmem_read_instructions ~4,392 -> -// ~10,464 (A loads become 8-B per step instead of 16-B staged chunks - SAME -// bytes, more instructions, all issued in one batch); (4) -// memory_traffic.total_bytes_per_operator_replay unchanged ~2.06 MB (same -// A+B bytes; A stays L2-hot). If median stays >= 11.385 with p90 <= 11.954, -// the A LDS round trip was not on the critical path either (the kernel is -// at its launch/latency floor with 3 waves/CU) and this round documents the -// next plateau datum for this shape. +// The int32 accumulation is exact (integer addition is order independent), so +// the split cannot change the result bit-for-bit: k-ascending inside each +// slice, ascending split sum at the tail = the scalar chain. // -// The generic scalar path remains the fallback for every unmatched -// (m,n,k), including the paired M=2 shape with the same (N,K); the M=16 -// DUMMA path is guarded by the exact (m,n,k) triple. +// The generic scalar kernel remains the fallback for every unmatched (m,n,k) +// including the paired M=2 shape with the same (N,K) = (384,4096); it decodes +// the (k,n)==(4096,384) pack so that fallback stays exact. The unsplit M=16 +// DUMMA kernel is the fallback when the caller workspace cannot hold the +// per-tile int32 accumulator + 24 counters (24,672 B as of iteration 7; it was +// 15 partial planes + counters = 368,736 B through iteration 6). The M=16 +// split-K path is guarded by the full (m,n,k) triple. // // Host symbols (matched against csrc/bindings.cpp): -// launch_w8a8_gemm(...) -// launch_pack_w8a8_weight(...) +// launch_w8a8_gemm(...) -- assigned-shape split-K / fallbacks +// launch_pack_w8a8_weight(...) -- fragment-order pack / identity // // Known-good include order for this DTK: HIP runtime first, then the HIP // bfloat16 header, then du_mma.h. @@ -399,59 +213,86 @@ namespace { +// --------------------------------------------------------------------------- +// Layout constants of the exact (k,n)==(4096,384) DUMMA-fragment-order weight +// pack (iteration 5). One 16-column tile is 65,536 B = 15 split slices in +// ascending split order; a slice is the concatenation of the wavefront's +// matrix_b fragment slots, TWO k32 steps per 1,024-B pair slot: +// offset_in_slice(step, lane, i) = +// (step>>1)*1024 + lane*16 + (step&1)*8 + i +// with lane = ((k - k_base)>>3 & 3)*16 + n_local, i = (k - k_base)&7, +// n_local = col&15. The 9-step (288-k) slices have 4 pair slots plus one +// 512-B single slot for the odd 9th step (8 B/lane at lane*8); the 8-step +// (256-k) slices have 4 pair slots. 8*4,608 + 7*4,096 = 65,536 B per tile, so +// the buffer size, the pack launch and the graph-stable address are unchanged. +// --------------------------------------------------------------------------- +constexpr int kFragPackNumLong = 8; // 8 x 288-k + 7 x 256-k +constexpr int kFragPackSliceLong = 288; +constexpr int kFragPackSliceShort = 256; +constexpr int kFragPackStepK = 32; // kDummaK +constexpr int kFragPackPairBytes = 1024; // 64 lanes * 16 B (2 steps) +constexpr int kFragPackLongBytes = 4608; // 4*1024 + 512 (9 steps) +constexpr int kFragPackShortBytes = 4096; // 4*1024 (8 steps) +constexpr int kFragPackLongRegion = 8 * 4608; // 36,864 B of long slices +constexpr int kFragPackTileBytes = 65536; // 4096 * 16 + // Byte offset of element (k, col) inside the exact (k,n)==(4096,384) -// slice-major-chunk weight pack (iteration 17; see -// w8a8_pack_weight_gateup_slicemajor_kernel): each 16-B chunk holds 16 -// CONSECUTIVE k values of one n column, and chunks are stored in ascending -// (slice, n, k-chunk) order. With tile = col>>4, n_local = col&15 and -// kchunk = k>>4: -// kchunk < 144 (the 8 long 288-k slices): s = kchunk/18, kk' = kchunk%18, -// chunk base = s*288, chunks per n = 18; -// kchunk >= 144 (the 7 short 256-k slices): c2 = kchunk-144, -// s = 8 + c2/16, kk' = c2%16, chunk base = 2304 + (c2/16)*256, -// chunks per n = 16; -// offset = tile*65536 + (chunk_base + n_local*cps + kk')*16 + (k & 15). -// Fallback paths only (the timed M=16 staging path uses the same mapping -// implicitly through the pack's chunk order). 144 = 8 long slices * 18 -// k-chunks per 288-k slice. +// DUMMA-fragment-order weight pack (see the constants above and the pack +// kernel below): byte = tile*65536 + slice_base(s) + (step>>1)*1024 + +// lane*16 + (step&1)*8 + i, with s/step/i/lane derived from (k, col) exactly +// as the pack kernel derives them from the byte index. Fallback paths only +// (generic scalar and unsplit M=16 DUMMA); the timed split-K kernel reads the +// same slots directly through its own template-indexed addresses. __device__ __forceinline__ int64_t packed384_byte_offset(int k, int col) { const int64_t tile = static_cast(col) >> 4; // col / 16 const int64_t n_local = static_cast(col) & 15; - const int kchunk = k >> 4; - int64_t base; - int kk_p, cps; - if (kchunk < 144) { - const int s = kchunk / 18; - base = static_cast(s) * 288; - kk_p = kchunk % 18; - cps = 18; + const int long_end = kFragPackNumLong * kFragPackSliceLong; // 2304 + int s, k_base; + if (k < long_end) { + s = k / kFragPackSliceLong; + k_base = s * kFragPackSliceLong; } else { - const int c2 = kchunk - 144; - const int s = 8 + c2 / 16; - base = 2304 + static_cast(c2 / 16) * 256; - kk_p = c2 % 16; - cps = 16; + s = kFragPackNumLong + (k - long_end) / kFragPackSliceShort; + k_base = long_end + (s - kFragPackNumLong) * kFragPackSliceShort; } - return tile * 65536 + (base + n_local * cps + kk_p) * 16 + (k & 15); + const int kk = k - k_base; + const int step = kk >> 5; // kk / kFragPackStepK + const int within = kk & 31; // kk % kFragPackStepK + const int lane = (within >> 3) * 16 + static_cast(n_local); + const int i = within & 7; + const int64_t base = + (s < kFragPackNumLong) + ? static_cast(s) * kFragPackLongBytes + : static_cast(kFragPackLongRegion) + + static_cast(s - kFragPackNumLong) * + kFragPackShortBytes; + const int64_t slot = + (s < kFragPackNumLong && step == kFragPackSliceLong / kFragPackStepK - 1) + ? 4 * kFragPackPairBytes + static_cast(lane) * 8 + : static_cast(step >> 1) * kFragPackPairBytes + + static_cast(lane) * 16 + (step & 1) * 8; + return tile * kFragPackTileBytes + base + slot + i; } -// Scalar W8A8 GEMM: one thread computes one output element (m, n). -// -// Adjacent threads map to adjacent N columns (the fastest-changing dimension -// of both `out` and the [K, N] weight), so global reads coalesce in N. -// The K loop is executed exactly once per output element in ascending k and -// accumulated in int32; every assigned K keeps the exact int8 dot within -// int32 range (max |dot| = 128*128*4096 = 67,108,864 < 2^31). The result is -// converted to float only for `dot * x_scale[m] * weight_scale[n]`, then -// stored as bfloat16, matching the contract reference exactly. + +// --------------------------------------------------------------------------- +// Generic scalar fallback: one thread per output element. // -// gfx928: wavefront = 64, blockDim must be a multiple of 64 (128 here). -__global__ __launch_bounds__(128) void w8a8_gemm_scalar_kernel( - const int8_t* __restrict__ a, // x_q [M, K] row-major - const int8_t* __restrict__ b, // weight [K, N] row-major - const float* __restrict__ x_scale, // [M, 1] contiguous - const float* __restrict__ weight_scale, // [N, 1] contiguous - hip_bfloat16* __restrict__ out, // [M, N] row-major +// The linear index is row-major over [M, N], so adjacent lanes write adjacent +// N columns in `out` and read adjacent N columns of the weight in the raw +// layout: both streams coalesce in the fastest-changing N dimension. For +// (k,n)==(4096,384) the weight buffer instead holds the DUMMA-fragment-order +// pack, decoded per element via packed384_byte_offset so this fallback stays +// exact for the paired M=2 shape and every unmatched m. +// --------------------------------------------------------------------------- +constexpr int kScalarBlockThreads = 128; + +__global__ __launch_bounds__(kScalarBlockThreads) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, // x_q [M, K] row-major + const int8_t* __restrict__ b, // weight [K, N] row-major / pack + const float* __restrict__ x_scale, // [M, 1] contiguous + const float* __restrict__ weight_scale, // [N, 1] contiguous + hip_bfloat16* __restrict__ out, // [M, N] row-major int m, int n, int k) { @@ -466,17 +307,14 @@ __global__ __launch_bounds__(128) void w8a8_gemm_scalar_kernel( const int col = static_cast(idx - static_cast(row) * n); const int8_t* __restrict__ a_row = a + static_cast(row) * k; - // For (k,n)==(4096,384) the weight buffer holds the per-tile - // slice-major-chunk pack (iteration 17; see packed384_byte_offset); every - // other (k,n) keeps the raw [K, N] row-major layout (identity pack). - // Decode so this fallback stays exact for the paired M=2 shape and all - // unmatched m. const bool packed_384 = (k == 4096 && n == 384); const int64_t b_row_stride = packed_384 ? 0 : static_cast(n); int32_t acc = 0; #pragma unroll 8 for (int kk = 0; kk < k; ++kk) { + // Exact int32 accumulation: max |dot| = 128 * 128 * 4096 = 67,108,864 + // which is < 2^31 for every assigned K. const int8_t b_val = packed_384 ? b[packed384_byte_offset(kk, col)] : b[static_cast(col) + @@ -485,39 +323,30 @@ __global__ __launch_bounds__(128) void w8a8_gemm_scalar_kernel( } // Reference order: (float(dot) * x_scale[m]) * weight_scale[n], then bf16. - // hip_bfloat16's float constructor performs the supported RNE conversion. - const float scaled = static_cast(acc) * x_scale[row] * - weight_scale[col]; + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; out[idx] = hip_bfloat16(scaled); } // --------------------------------------------------------------------------- -// Minimal gfx928 DUMMA INT8 path (assigned shape M=16, N=384, K=4096). -// Tile: 16x16x32, int32 accumulation, one 64-lane wavefront per block, -// one 16x16 output tile per wave (grid = N/16 = 24 blocks). -// --------------------------------------------------------------------------- - // gfx928 DUMMA INT8 m16n16k32 tile constants. +// --------------------------------------------------------------------------- constexpr int kDummaM = 16; constexpr int kDummaN = 16; constexpr int kDummaK = 32; -constexpr int kDummaBlockThreads = 64; // one wavefront per block +constexpr int kDummaBlockThreads = 64; // one 64-lane wavefront per block -// M=16 W8A8 GEMM with the minimal 16x16x32 DUMMA tile: one wavefront owns one -// independent 16x16 output tile. A (x_q [16, K] row-major) and B (weight -// [K, N] row-major) fragments are loaded directly from global memory every -// DUMMA step with the library loaders; the int32 accumulator is explicit and -// drained through LDS in a coalesced epilogue that applies x_scale and -// weight_scale before the bf16 store. No LDS staging, no double buffering, -// and no cross-wave barrier in the K loop (a single wavefront per block; -// the one __syncthreads in the epilogue is reached by every block thread). +// Unsplit M=16 DUMMA fallback (used only when the caller workspace cannot hold +// the per-tile int32 accumulator + counters). One wavefront owns one 16x16 tile +// with direct library fragment loads and an LDS-drained epilogue. For +// (k,n)==(4096,384) the B fragment is gathered per element from the pack. __global__ __launch_bounds__(kDummaBlockThreads) void w8a8_gemm_m16_dumma_kernel( - const int8_t* __restrict__ a, // x_q [16, K] row-major - const int8_t* __restrict__ b, // weight [K, N] row-major - const float* __restrict__ x_scale, // [16, 1] contiguous - const float* __restrict__ weight_scale, // [N, 1] contiguous - hip_bfloat16* __restrict__ out, // [16, N] row-major + const int8_t* __restrict__ a, // x_q [16, K] row-major + const int8_t* __restrict__ b, // weight [K, N] row-major / pack + const float* __restrict__ x_scale, // [16, 1] contiguous + const float* __restrict__ weight_scale, // [N, 1] contiguous + hip_bfloat16* __restrict__ out, // [16, N] row-major int n, int k) { const int lane = static_cast(threadIdx.x); // 0..63 @@ -535,31 +364,20 @@ w8a8_gemm_m16_dumma_kernel( acc_frag; du::dumma::du_fill_fragment(acc_frag, 0); - // K is divisible by kDummaK for every assigned shape. A fragment rows are - // x_q rows (ldm = k); B fragment rows are weight k-rows (ldm = n), with the - // tile's 16 columns at n_base. The explicit int32 accumulation is - // k0-ascending over exact 32-element dot products, so the final int32 total - // is bit-identical to the scalar chain (integer arithmetic is exact). - // Iteration 17: this fallback is only reachable for (m,n,k)==(16,384,4096), - // whose weight buffer holds the per-tile slice-major-chunk pack. The - // library row_major loader (a.x[i] = p[(col+i)*ldm + row], row = lane&15, - // col = (lane>>4)*8) cannot decode that non-linear order, so the B - // fragment is filled with per-element gathers from the pack: x[i] = - // element (k0 + (lane>>4)*8 + i, n_base + (lane&15)) - the SAME fragment - // registers the raw-layout loader would produce (the round-15 n-major - // decode asserted the same mapping). Fallback path only - never on the - // timed M=16 split-K kernel. const bool packed_384 = (k == 4096 && n == 384); #pragma unroll 4 for (int k0 = 0; k0 < k; k0 += kDummaK) { du::dumma::du_load_matrix_sync(a_frag, a + k0, k); if (packed_384) { + // The library row_major loader (x[i] = p[(col+i)*ldm + row], row = + // lane&15, col = (lane>>4)*8) cannot decode the pack, so gather the + // SAME fragment registers per element: x[i] = element + // (k0 + (lane>>4)*8 + i, n_base + (lane&15)). const int n_local = lane & 15; const int k_off = (lane >> 4) * 8; #pragma unroll for (int i = 0; i < 8; ++i) { - b_frag.x[i] = - b[packed384_byte_offset(k0 + k_off + i, n_base + n_local)]; + b_frag.x[i] = b[packed384_byte_offset(k0 + k_off + i, n_base + n_local)]; } } else { du::dumma::du_load_matrix_sync( @@ -586,45 +404,32 @@ w8a8_gemm_m16_dumma_kernel( } // --------------------------------------------------------------------------- -// Iteration 8 (HIP round): split-K=15 one-wave LDS-staged partial GEMM + -// split-15 combine (mandatory occupancy round - one limiter: LDS footprint -// -> blocks/CU; trusted probe set [2,5,10,15]). +// Split-K=15 one-wave, zero-barrier, CU-aligned GEMM + fused last-arrival tail +// (iteration 3) with register-direct A/B transport (iteration 5) and one hot +// per-tile int32 accumulator (iteration 7). // -// Round 7 measured the split-K=10 one-wave grid (240 blocks = exactly -// 2 blocks/CU, 26,880 B LDS/block = 53,760 B/CU) at 15.43 us median but was -// rejected by the p90 guard (19.43 > 19.04; bimodal samples - machine -// state, not kernel behavior). Its 26,880 B/block LDS caps occupancy at -// exactly 2 blocks/CU (3 x 26,880 = 80,640 B > 64 KiB). split-K=15 shortens -// every slice (288/256 vs 416/384 rows), shrinking the per-block LDS -// footprint to 18,688 B -> 3 x 18,688 = 56,064 B/CU <= 64 KiB, so 3 one-wave -// blocks/CU become legal, and the grid = 24 N-groups x 15 splits = 360 -// one-wave blocks = EXACTLY 3 blocks/CU on 120 CUs (above the 2-blocks/CU -// latency-hiding target, integer-balanced, no tail wave). B stays once-read -// from HBM and A stays L2-hot (no repeated global reads - each block stages -// its own slice exactly once). -// * block tile 16x16 (1 wavefront x 1 N-tile), 64 thr; grid = 24 x 15 = -// 360 one-wave blocks = exactly 3 blocks/CU; -// * K=4096 -> 15 non-uniform 32-aligned slices: 8 x 288 (9 DUMMA steps) -// + 7 x 256 (8 steps); k0-ascending within each slice, ascending split -// sum -> bit-identical int32 order (0 mismatches expected); -// * whole slice staged once into bank-skewed LDS (A 16x288 stride 304, -// B 288x16 stride 48; 18,688 B/block), batched 16-B vector loads, one -// __syncthreads (emitted as 0 s_barrier for the 1-wave block), then a -// zero-barrier LDS-only K loop (9/8 steps x 1 A-frag + 1 B-frag ds_read -// + 1 v_mmac, 9/8 v_mmac per block); -// * workspace = 15 int32 planes [15][16][384] = 368,640 B (<= the -// caller's 393,216 B workspace_split_k_capacity budget; every launch -// overwrites every partial tile - no workspace clear); -// * combine kernel: 24 blocks x 128 thr, all 15 plane loads issued into a -// register array before the ascending split sum (exact int32 order), -// then (float(acc) * x_scale[row]) * weight_scale[col] -> bf16; -// * replay traffic ~2.32 MB (B 1.57 MB + planes RW 737,280 B + out 12,288 -// B) vs 2.47 MB for the accepted S=16 kernel (-6%); vs S=10 the plane -// RW grows +245,760 B (+50%) - the combine cost of more splits is the -// explicit trade being benchmarked against the occupancy gain. +// * block tile 16x16 (1 wavefront x 1 N-tile), 64 thr; grid = 24 N-groups x +// 15 splits = 360 one-wave blocks = EXACTLY 3 blocks/CU on 120 CUs; +// * K=4096 -> 15 non-uniform 32-aligned slices: 8 x 288 (9 DUMMA steps) + +// 7 x 256 (8 steps); k0-ascending within each slice, exact order-free +// int32 reduction at the tail -> bit-identical int32 sum; +// * A is register-only (per-lane int2 fragment slice per step, no LDS); +// * B (the once-read, cold weight) is packed DUMMA-fragment-order outside +// the timed region (iteration 5) and read straight into its fragment +// registers with 4 x int4 (+ 1 x int2 for the 9-step slices) per lane - +// NO LDS stage, no barrier, no lgkmcnt drain; +// * no __syncthreads in the block body (0 s_barrier for a 1-wave block) and +// a zero-barrier register-only K loop (0 barriers per K step); +// * workspace = ONE int32 accumulator per tile [24][16][16] = 24,576 B plus +// 24 arrival counters = 24,672 B (<= the caller's 393,216 B budget); each +// slice adds its 256 int32 with global atomicAdd into its tile (hot in L2, +// no per-replay write-back of 368,640 B of planes); +// * fused tail: exactly ONE arrival per block (lane 0 fence + atomicAdd, +// acquisition fence, LDS s_is_last flag, barriers); the last arriver reads +// the single accumulated tile, applies both scales, stores bf16, resets +// the tile accumulator and the counter. One dispatch per replay. // --------------------------------------------------------------------------- - -constexpr int kSplitK = 15; // 15 planes; workspace 368,640 B +constexpr int kSplitK = 15; // 15 slices; grid = 24 x 15 = 360 constexpr int kNumNGroups = 24; // N=384 / 16 (partial block tiles) constexpr int kNumLongSplits = 8; // 8 x 288-row slices + 7 x 256-row constexpr int kSliceLong = 288; // 9 x kDummaK steps (non-uniform) @@ -632,70 +437,61 @@ constexpr int kSliceShort = 256; // 8 x kDummaK steps constexpr int kStepsLong = kSliceLong / kDummaK; // 9 constexpr int kStepsShort = kSliceShort / kDummaK; // 8 constexpr int kPartialBlockThreads = 64; // one wavefront per block -constexpr int kCounterInts = kNumNGroups; // one arrival counter per N-group -constexpr int kStagedAStride = 304; // 288 + 16 bank skew (16-B aligned) -// Iteration 15: B is staged n-major (element (k,n) at n*kStagedBStride + k, -// the transpose of the A staging convention), so the column stride must hold -// the longest k-slice (288) plus bank skew; 304 = 288 + 16 keeps every -// staging ds_write_b128 and every 8-byte fragment read 16-/8-B aligned. -constexpr int kStagedBStride = 304; // 288 + 16 bank skew (16-B aligned) - -// Partial GEMM: block (tile, split) computes one 16x16 output tile for the -// split's contiguous K slice and stores the int32 partial to -// partials[split][row][n]. The whole slice is staged into LDS once with -// batched 16-B vector loads, one barrier, then a zero-barrier LDS-only K -// loop. k0-ascending within the slice; the combine kernel sums slices in -// ascending split order, so the total int32 accumulation order is -// bit-identical to the scalar chain. The slice body is a compile-time -// template (kSliceKT, kStepsT) so the staging loops and the K loop are -// fully unrolled; the grid-uniform `split < kNumLongSplits` branch picks the -// 288- or 256-row instantiation for each block (one launch total). +constexpr int kCounterInts = kNumNGroups; // one arrival counter per N-group + +// Iteration 7 workspace: ONE int32 accumulator per 16-column tile instead of +// 15 int32 partial planes. 24 tiles x 256 int32 = 24,576 B total, hot in L2, +// so the 368,640 B of per-replay partial-plane write-back disappears. The +// 24 arrival counters stay after the accumulator. Every accumulator element is +// zeroed once by the host (like the counters) and is reset by the tile's fused +// last-arrival tail at the end of every replay, so no clear is needed inside +// the timed region. +constexpr int kAccumIntsPerTile = kDummaM * kDummaN; // 256 +constexpr int64_t kAccumInts = + static_cast(kNumNGroups) * kAccumIntsPerTile; // 6,144 +constexpr int64_t kAccumBytes = + kAccumInts * static_cast(sizeof(int32_t)); // 24,576 B + +// One 16x16 tile for one 32-aligned K slice. The slice body is a compile-time +// template (kSliceKT, kStepsT) so the load loops and the K loop are fully +// unrolled; the grid-uniform `split < kNumLongSplits` branch picks the 288- or +// 256-row instantiation for each block (one launch total). +// +// Iteration 5: both operands go straight into their fragment registers and NO +// LDS is used at all, so the body is just "load registers, then kStepsT +// v_mmac". A comes from the raw row-major activation; B comes from the +// DUMMA-fragment-order pack, whose per-tile slice region is the wavefront's +// matrix_b fragment slots in k-ascending order (two k32 steps per 1,024-B +// pair slot, one 512-B single slot for an odd 9th step). +// Iteration 7: the slice's 16x16 int32 result is added straight into the +// tile's single accumulator with four per-lane global atomicAdd (no plane, no +// per-split store base). Integer addition is exact and order-independent, so +// the reduction is bit-identical to the ascending 15-plane sum. template __device__ __forceinline__ void split15_slice_body( const int8_t* __restrict__ a, // x_q [16, K] row-major - const int8_t* __restrict__ b, // weight [K, N] row-major - int32_t* __restrict__ partials, // [kSplitK][16][N] int32 - int n, + const int8_t* __restrict__ b, // weight fragment-order pack + int32_t* __restrict__ accum, // [24][16][16] int32 int k, - int split, int tile, - int n_base, - int k_base, - int8_t (*b_lds)[kStagedBStride]) { + int split, + int k_base) { const int tid = static_cast(threadIdx.x); - // Iteration 21 (HIP-only, A register-only transport; the validated TP4 - // o_proj lineage recipe "register-only K-loop transport"). The A slice is - // small (16 x kSliceKT = 4,608 / 4,096 B, L2-hot 64-KiB activation) and - // every block reads its own rows exactly once, so staging it through LDS - // only buys the LDS round trip: per block it costs kALoads 16-B global - // loads + kALoads ds_write_b128 + kStepsT ds_read2_b32 (2-way bank - // conflicted). This round removes the A LDS stage entirely: each lane - // loads its own 8-B fragment slice for EVERY DUMMA step into a per-lane - // register array a_slice[kStepsT] (the library matrix_a row_major loader - // with ldm = k reads x[i] = p[row*k + k_base + s*kDummaK + col + i], row = - // lane&15, col = (lane>>4)*8 - the SAME 8 elements per lane per step the - // LDS version read, so every fragment register is identical and the int32 - // accumulation order is bit-identical). All kStepsT A loads are issued in - // one register-array block BEFORE any ds_write (the iteration-11 batching - // pattern that forced the staging loads back-to-back in the accepted code - // object), so the A slice costs one memory latency and zero LDS - // instructions; the K loop then consumes the fragments straight from - // registers. Guarded lanes are impossible (kStepsT * 64 lanes * 8 B = - // kSliceKT * 16 B exactly, no tail). + // Register-only A transport (iteration 3, unchanged). The A slice is small + // (16 x kSliceKT = 4,608/4,096 B, L2-hot 64-KiB activation) and every block + // reads its own rows exactly once, so staging it through LDS only buys the + // LDS round trip. Each lane loads its own 8-B fragment slice for EVERY DUMMA + // step into a per-lane register array: the library matrix_a row_major loader + // with ldm=k reads x[i] = p[row*k + k_base + s*kDummaK + col + i], row = + // lane&15, col = (lane>>4)*8. The 8 elements are 8 CONSECUTIVE bytes of one + // row (8-B aligned), so one explicit int2 load per lane per step and a + // register-to-register memcpy fill the same fragment registers the LDS + // loader produced. All kStepsT loads are issued back-to-back, ahead of the B + // loads, so the A slice costs one memory latency and zero LDS ops. du::dumma::DUFragment a_slice[kStepsT]; - // Explicit vector loads: the library loader assigns the 8 bytes - // individually, which LLVM scalarizes to 8 global_load_ubyte per step per - // lane from GLOBAL memory (verified in the round-21 draft code object: - // 136 ubyte loads = the old "direct fragment byte-load" poison). The - // fragment's 8 elements are 8 CONSECUTIVE bytes of one row (m = lane&15) - // at k = k_base + s*kDummaK + (lane>>4)*8, 8-B aligned (k_base, kDummaK - // and the col offset are all multiples of 8; x_q is tensor-aligned), so - // one explicit int2 load per lane per step forces ONE global_load_dwordx2 - // and the __builtin_memcpy (8 bytes, register-to-register) fills the same - // fragment registers the LDS row_major loader produced. const int a_row = tid & 15; // m index (library loader row) const int a_col = (tid >> 4) * 8; // k offset (library loader col) const int8_t* a_base = @@ -706,114 +502,91 @@ __device__ __forceinline__ void split15_slice_body( __builtin_memcpy(a_slice[s].x, &raw, 8); } - // Stage the B slice (kSliceKT rows x 16 cols) with coalesced 16-B vector - // loads, then one barrier (B only - the A stage is gone). B chunks - // (iteration 17) come from the per-tile slice-major-chunk pack (see - // w8a8_pack_weight_gateup_slicemajor_kernel): chunk c holds 16 CONSECUTIVE - // k values of one n column (n = c/18 long, c/16 short; k-chunk = c%18 / - // c%16) and the chunks are stored in ascending c order, so the slice's - // chunks occupy the CONTIGUOUS byte range [(k_base + c)*16] of the tile - - // 64 lanes cover 1,024 CONTIGUOUS bytes per wave-load (the round-14 - // locality win). Same chunk content, same n-major LDS destination, same - // per-lane order -> the fragment reads and the int32 accumulation stay - // bit-identical. - // - // Iteration 11: issue ALL staging global loads into one register array - // BEFORE any ds_write. The accepted kernel's exact gfx928 code object - // serialized the slice staging: every global_load_dwordx4 was followed by - // s_waitcnt vmcnt(0) before its ds_write_b128 (load_wait0_ds_write_windows - // = 18 across the two slice instantiations, 10 long + 8 short), so each - // block's staging paid ~10 (long) / 8 (short) full global-memory latencies - // one after another before the K loop could issue its first ds_read - // (lgkmcnt(0)). Holding the 16-B chunks in a register array first forces - // the compiler to issue the loads back-to-back (all in flight - // concurrently) and move the vmcnt waits to the write side, collapsing the - // staging critical path to one memory latency. Guarded lanes (the unrolled - // tail beyond the chunk count) write a zero placeholder so no - // uninitialized register is ever consumed. - constexpr int kChunksPerBRow = kSliceKT / 16; // 18 / 16 (per n column) - constexpr int kBLoads = (kSliceKT + 63) / 64; // 5 / 4 - int4 s_buf[kBLoads]; + // Register-direct B transport (iteration 5). This block's split owns a + // CONTIGUOUS region of its tile: slice_base(s) + 4 x 1,024-B pair slots, + // followed (9-step slices only) by one 512-B single slot for the odd last + // step. One int4 load per lane per pair slot covers 1,024 CONTIGUOUS bytes + // per wave-load and lands BOTH matrix_b fragments of that pair in their + // final registers: its low 8 bytes are x[i] of step 2m and its high 8 bytes + // x[i] of step 2m+1, i.e. the same x[i] = element(k0 + (lane>>4)*8 + i, + // n_base + (lane&15)) the accepted col_major LDS loader produced. The loads + // are issued back-to-back with no ds_write, no __syncthreads and no lgkmcnt + // drain in between, so the K loop can consume each step pair as soon as it + // arrives instead of after the whole slice. + du::dumma::DUFragment + b_slice[kStepsT]; + constexpr int kPairSlots = kStepsT / 2; // 4 + const int8_t* b_split = + b + static_cast(tile) * (kDummaN * k) + + ((split < kNumLongSplits) + ? static_cast(split) * kFragPackLongBytes + : static_cast(kFragPackLongRegion) + + static_cast(split - kNumLongSplits) * + kFragPackShortBytes); #pragma unroll - for (int i = 0; i < kBLoads; ++i) { - const int c = tid + 64 * i; - // Iteration 17: slice-major-chunk pack - chunk c (n = c/kChunksPerBRow, - // k-chunk = c%kChunksPerBRow) lives at byte (k_base + c)*16 of the - // tile, so a wave-load reads 1,024 contiguous bytes. - s_buf[i] = (c < kSliceKT) - ? *reinterpret_cast( - b + static_cast(tile) * (kDummaN * k) + - (static_cast(k_base) + c) * 16) - : int4{0, 0, 0, 0}; + for (int m = 0; m < kPairSlots; ++m) { + const int4 raw = *reinterpret_cast( + b_split + static_cast(m) * kFragPackPairBytes + tid * 16); + __builtin_memcpy(b_slice[2 * m].x, &raw, 8); + __builtin_memcpy(b_slice[2 * m + 1].x, + reinterpret_cast(&raw) + 8, 8); } -#pragma unroll - for (int i = 0; i < kBLoads; ++i) { - const int c = tid + 64 * i; - if (c < kSliceKT) { - // n-major LDS destination (iteration 15, UNCHANGED): element (k,n) at - // n*kStagedBStride + k. Chunk c maps to (n = c/kChunksPerBRow, - // k-chunk = c%kChunksPerBRow): for the 256-row instantiation this is - // the same (c&15, c>>4) as iteration 15; for the 288-row - // instantiation the chunk order is n-major (18 chunks per column) and - // the division is resolved by the compiler (staging path only). Still - // one ds_write_b128 per 16-B chunk. - const int b_row = c / kChunksPerBRow; - *reinterpret_cast( - &b_lds[b_row][(c - b_row * kChunksPerBRow) * 16]) = s_buf[i]; - } + if (kStepsT % 2 != 0) { + // 9-step slices: the odd last step is a 512-B single slot (8 B/lane). + const int2 raw = *reinterpret_cast( + b_split + kPairSlots * kFragPackPairBytes + tid * 8); + __builtin_memcpy(b_slice[kStepsT - 1].x, &raw, 8); } - __syncthreads(); - - // The block's one wavefront owns the whole 16x16 tile. Zero-barrier K - // loop: kStepsT steps of kDummaK=32; the A fragment comes from the - // preloaded register array a_slice[s] (iteration 21, no LDS), the B - // fragment from b_lds (iteration 15: col_major on the n-major layout, ldm - // 304), one v_mmac per step (9 or 8 v_mmac per block). The col_major - // loader computes b.x[i] = p[row*304 + col + i] (row = n_local = lane&15, - // col = k offset (lane>>4)*8), i.e. the lane's 8 elements B[k0+col+i] - // [n_base+row] are 8 CONSECUTIVE bytes of one n column -> ONE ds_read2_b32 - // per step (was 8 ds_read_u8 at stride 48). x[i] gets the same element as - // the old row_major loader, so the fragment registers (and the int32 - // accumulation order) are bit-identical. - du::dumma::DUFragment - b_frag; + // The block's one wavefront owns the whole 16x16 tile. Zero-barrier K loop: + // kStepsT steps of kDummaK=32, A and B both straight from the preloaded + // register arrays, one v_mmac per step (9 or 8 v_mmac per block), zero + // s_barrier per K step. The int32 accumulation inside the slice is untouched + // (k-ascending); across slices the addition is exact and order-free. du::dumma::DUFragment acc_frag; du::dumma::du_fill_fragment(acc_frag, 0); #pragma unroll for (int s = 0; s < kStepsT; ++s) { - const int k0 = s * kDummaK; - du::dumma::du_load_matrix_sync(b_frag, &b_lds[0][k0], kStagedBStride); - du::dumma::du_mma_sync(acc_frag, a_slice[s], b_frag, acc_frag); + du::dumma::du_mma_sync(acc_frag, a_slice[s], b_slice[s], acc_frag); } - // Direct int32 partial store: partials[split][row][n_base .. n_base+15]. - // Every launch overwrites every partial tile, so no workspace clear is - // needed before capture or between replays. - int32_t* partial_base = - partials + static_cast(split) * kDummaM * n + n_base; - du::dumma::du_store_matrix_sync(partial_base, acc_frag, n, - du::dumma::mem_row_major); + // Direct int32 accumulation into the per-tile accumulator (iteration 7). + // The library mem_row_major store owns one element per lane at row = tid&15, + // column = (tid>>4), then x[1], x[2], x[3] at +4, +8, +12 columns (the exact + // mapping the iteration-6 candidate verified against the exact code object: + // one 64-bit base with 0/16/32/48-byte offsets). Four no-return atomicAdd + // publish exactly those values into accum[tile][row][col], which stays hot + // in L2 (24,576 B total) - the 368,640-B partial-plane round trip of + // iterations 3-6 is gone. Integer addition is exact and order-independent, + // so the tile sum, and therefore every bf16 output, is unchanged. + int32_t* accum_lane = + accum + static_cast(tile) * kAccumIntsPerTile + + (tid & 15) * kDummaN + (tid >> 4); +#pragma unroll + for (int i = 0; i < 4; ++i) { + atomicAdd(accum_lane + 4 * i, acc_frag.x[i]); + } } -// Fused last-arrival combine tail (iteration 9): removes the separate -// 24-block combine kernel -> exactly ONE kernel launch per replay. +// Fused last-arrival combine tail: removes the separate 24-block combine +// kernel -> exactly ONE kernel launch per replay. // -// Protocol (the validated TP4 qkv lineage pattern): -// * every block writes its 16x16 int32 plane tile, then __threadfence() -// (release) and atomicAdd(&counters[tile], 1); -// * the block whose atomicAdd returns kSplitK-1 is the 15th (last) -// arriver of its N-group: it __threadfence()s (acquire), sums the 15 -// planes in ascending split order (exact int32 order), applies both -// scales, stores bf16, and resets its counter to 0 (atomicExch) so the -// next replay restarts from 0. The host zeroes the 24 counters once -// before the first launch; the self-reset keeps every later replay -// deterministic (stream-ordered, no racing increments). +// Protocol (iteration 7): every block adds its 16x16 int32 result into the +// tile's single accumulator with global atomicAdd, then lane 0 +// __threadfence()es (release) and atomicAdd(&counters[tile], 1); the block +// whose atomicAdd returns kSplitK-1 is the 15th (last) arriver of its +// N-group: it acquires, reads the summed accumulator, applies both scales, +// stores bf16, resets the tile's accumulator to zero and resets its counter to +// 0 (atomicExch) so the next replay restarts from zero. The host zeroes the +// whole scratch region once before the first launch; the tail's self-reset +// keeps every later replay deterministic. Only lane 0 ever increments the +// counter (the round-9 lineage lost correctness by letting all 64 lanes +// increment). __device__ __forceinline__ void split15_combine_tail( - const int32_t* __restrict__ partials, // [kSplitK][16][N] int32 + int32_t* __restrict__ accum, // [24][16][16] int32 (hot) const float* __restrict__ x_scale, // [16, 1] contiguous const float* __restrict__ weight_scale, // [N, 1] contiguous hip_bfloat16* __restrict__ out, // [16, N] row-major @@ -822,24 +595,13 @@ __device__ __forceinline__ void split15_combine_tail( int n) { const int lane = static_cast(threadIdx.x); // 0..63 (one wavefront) const int n_base = tile * kDummaN; - const int plane = kDummaM * n; // int32 stride per split - // Repair 1/4 (iteration 9): exactly ONE arrival per block. Only lane 0 - // fences, atomically increments the tile counter and publishes the result - // through LDS; the other 63 lanes must NOT atomicAdd. The round-9 draft had - // all 64 lanes incrementing -> 64 x 15 = 960 increments per tile per - // replay, so the 15th increment (arrived == kSplitK-1) fired inside the - // FIRST-arriving block, before its 14 siblings had stored their planes: - // the premature combiner summed mostly-zero planes (first_mismatch - // actual=0.0 at (0,0)) and the unguarded atomicExch reset made several - // blocks combine per replay (6048/6144 mismatched). This mirrors the - // validated TP4 qkv fused tail: tid-0 fence + atomicAdd, LDS flag, barriers. __syncthreads(); __shared__ int s_is_last; if (lane == 0) { - __threadfence(); // release plane stores + __threadfence(); // release this slice const int arrived = atomicAdd(&counters[tile], 1); - __threadfence(); // acquire other planes + __threadfence(); // acquire other slices s_is_last = (arrived == kSplitK - 1) ? 1 : 0; } __syncthreads(); @@ -847,93 +609,84 @@ __device__ __forceinline__ void split15_combine_tail( return; // not the last arriver } - // 256 outputs / 64 lanes = 4 CONSECUTIVE columns per lane (row = lane>>2, - // col0 = (lane & 3) * 4): one aligned int4 (16-B) plane load per split - // replaces the round-9 loop's 4 scalar 4-B loads per element, cutting the - // tail's vmem load instruction count 4x (60 -> 15 per lane) with no byte - // or L2-request increase (each 16-B load is fully consumed; the 4 lanes - // of one row cover a contiguous 64-B segment). The per-element int32 - // accumulation order is UNCHANGED: for each output element the 15 plane - // values are still summed in ascending split order, so the outputs stay - // bit-identical to the scalar chain. weight_scale is read once as a - // float4 (n_base + col0 is 16-B aligned) and the 4 adjacent bf16 stores - // merge into one 8-B store. Fence/atomic/barrier protocol untouched. + // 256 accumulators / 64 lanes = 4 CONSECUTIVE columns per lane (row = lane>>2, + // col0 = (lane & 3) * 4): ONE aligned int4 (16-B) accumulator load per lane + // replaces the 15 dependent int4 plane loads (15 -> 1 per lane), and the + // accumulator is only 24,576 B (L2-resident) instead of 368,640 B of planes. + // The value read is the exact int32 sum of all 15 slices (integer addition + // is order-independent), so the reference float order + // (float(acc) * x_scale[row]) * weight_scale[col] produces bit-identical + // bf16 outputs. weight_scale is read once as a float4 (16-B aligned) and the + // 4 adjacent bf16 stores merge into one 8-B store. const int row = lane >> 2; // 0..15 (4 lanes per row) const int col0 = (lane & 3) * 4; // 0,4,8,12 (16-B aligned) - const int32_t* p = partials + row * n + n_base + col0; + int32_t* p = + accum + static_cast(tile) * kAccumIntsPerTile + + row * kDummaN + col0; + const int4 v = *reinterpret_cast(p); const float4 ws = *reinterpret_cast(weight_scale + n_base + col0); const float xs = x_scale[row]; - int32_t acc0 = 0, acc1 = 0, acc2 = 0, acc3 = 0; -#pragma unroll - for (int s = 0; s < kSplitK; ++s) { - const int4 v = *reinterpret_cast( - p + static_cast(s) * plane); - acc0 += v.x; - acc1 += v.y; - acc2 += v.z; - acc3 += v.w; - } hip_bfloat16* o = out + static_cast(row) * n + n_base + col0; - o[0] = hip_bfloat16(static_cast(acc0) * xs * ws.x); - o[1] = hip_bfloat16(static_cast(acc1) * xs * ws.y); - o[2] = hip_bfloat16(static_cast(acc2) * xs * ws.z); - o[3] = hip_bfloat16(static_cast(acc3) * xs * ws.w); - - // All kSplitK arrivals of this replay have happened; restart for the next - // replay (the kernel's completion on the stream orders this reset before - // the next replay's first atomicAdd). One lane is enough; all 15 blocks - // have already passed their atomicAdd, so no racing increment survives. + o[0] = hip_bfloat16(static_cast(v.x) * xs * ws.x); + o[1] = hip_bfloat16(static_cast(v.y) * xs * ws.y); + o[2] = hip_bfloat16(static_cast(v.z) * xs * ws.z); + o[3] = hip_bfloat16(static_cast(v.w) * xs * ws.w); + + // Every slice of this tile has already arrived (that is what made this block + // the last arriver), so the reset cannot race with another atomicAdd: the + // accumulator restarts at zero for the next replay, exactly like the + // counter, and the next replay's kernel is stream-ordered after this one. + // The store is to the same address just read, so the compiler keeps it after + // the load that produced v. + *reinterpret_cast(p) = make_int4(0, 0, 0, 0); if (lane == 0) { atomicExch(&counters[tile], 0); } } +// Iteration 5: no LDS slice stage at all - both operands are fetched straight +// into fragment registers (A from the raw activation, B from the +// DUMMA-fragment-order pack), so the only shared storage left in this kernel is +// the 4-B s_is_last flag of the fused combine tail. +// Iteration 7: the per-slice result is accumulated into accum[tile] with +// global atomicAdd (no per-split planes); split/n_base are only needed to pick +// the slice and its K base. __global__ __launch_bounds__(kPartialBlockThreads) void w8a8_gemm_m16_dumma_splitk15_lds_staged_partial_kernel( const int8_t* __restrict__ a, // x_q [16, K] row-major - const int8_t* __restrict__ b, // weight [K, N] row-major - int32_t* __restrict__ partials, // [kSplitK][16][N] int32 + const int8_t* __restrict__ b, // weight fragment-order pack + int32_t* __restrict__ accum, // [24][16][16] int32 const float* __restrict__ x_scale, // [16, 1] contiguous const float* __restrict__ weight_scale, // [N, 1] contiguous hip_bfloat16* __restrict__ out, // [16, N] row-major int32_t* __restrict__ counters, // [kNumNGroups] monotonic int n, int k) { - const int tile = static_cast(blockIdx.x) / kSplitK; // 0..kNumNGroups-1 - const int split = static_cast(blockIdx.x) % kSplitK; // 0..kSplitK-1 - const int n_base = tile * kDummaN; + const int tile = static_cast(blockIdx.x) / kSplitK; // 0..23 + const int split = static_cast(blockIdx.x) % kSplitK; // 0..14 const int k_base = (split < kNumLongSplits) ? split * kSliceLong : kNumLongSplits * kSliceLong + (split - kNumLongSplits) * kSliceShort; - // Iteration 15: n-major B stage, element (k,n) at n*kStagedBStride + k - // (16 n columns x 304 B = 4,864 B; the k-slice length is carried by the - // template instantiation, not by the shared array). Iteration 21: the A - // slice is no longer staged in LDS (register-only transport), so the - // per-block LDS is B + the combine flag only. - __shared__ __align__(16) int8_t b_lds[kDummaN][kStagedBStride]; - if (split < kNumLongSplits) { - split15_slice_body(a, b, partials, n, k, split, - tile, n_base, k_base, b_lds); + split15_slice_body(a, b, accum, k, tile, split, + k_base); } else { - split15_slice_body(a, b, partials, n, k, split, - tile, n_base, k_base, b_lds); + split15_slice_body(a, b, accum, k, tile, split, + k_base); } - // Fused combine tail (iteration 9): one launch per replay; the last - // arriver of each N-group sums the 15 planes and stores bf16. - split15_combine_tail(partials, x_scale, weight_scale, out, counters, tile, + split15_combine_tail(accum, x_scale, weight_scale, out, counters, tile, n); } -// Identity device-to-device weight pack (iteration 1: the DUMMA kernel reads -// B with the raw [K, N] row-major layout, so identity packing is exact). -// Kept valid for every (K, N) including unmatched shapes; the exact -// (k,n)==(4096,384) shape uses the per-tile slice-major-chunk pack below -// instead (iteration 17). +// --------------------------------------------------------------------------- +// Weight packing (outside the timed region). +// --------------------------------------------------------------------------- +constexpr int64_t kGateUpTotal = 4096 * 384; // 1,572,864 bytes + __global__ void w8a8_pack_weight_identity_kernel( const int8_t* __restrict__ src, int8_t* __restrict__ dst, @@ -945,60 +698,73 @@ __global__ void w8a8_pack_weight_identity_kernel( } } -// Per-tile slice-major-chunk weight pack for the exact (k,n)==(4096,384) -// shape (iteration 17): each 16-B chunk holds 16 CONSECUTIVE k values of one -// n column, and the 4,096 chunks of a tile are stored in ascending (slice, -// n, k-chunk) order: -// chunk index = slice_chunk_base(s) + n_local*chunks_per_slice(s) + kk' -// with the 8 long 288-k slices (s < 8: base s*288, 18 chunks per n) followed -// by the 7 short 256-k slices (s >= 8: base 8*288 + (s-8)*256, 16 chunks per -// n); byte address = tile*65536 + chunk_index*16, chunk content = -// raw[k*384 + tile*16 + n_local] for k = k_base(s) + kk'*16 .. +15 (k_base -// numerically equals the chunk base: 288 k and 288 chunks per long slice). -// This is the iteration-15 n-major pack with the same chunk CONTENT but -// reordered memory (n*4096 + kk*16 -> slice-major chunk order), so a staged -// wave-load (64 consecutive chunks) covers 1,024 CONTIGUOUS bytes -// (sequential L2 sectors and sequential DRAM pages; was 16 x 64-B sectors -// at 4,096-B stride) while the n-major LDS destination, the -// one-ds_read2_b32 B fragment read and the int32 accumulation order are all -// untouched (outputs bit-identical to iteration 15). One-time, -// out-of-timed-region, out-of-Graph, same byte count and same buffer -// (graph-stable addresses unchanged). Every other (K, N) keeps the identity -// pack. Fallbacks decode this order via packed384_byte_offset. +// Per-tile DUMMA-fragment-order weight pack for the exact (k,n)==(4096,384) +// shape (iteration 5): byte idx inside a tile is the matrix_b fragment byte +// that the timed split-K kernel loads, so each block's B slice is a straight +// sequential register stream (4 x 16-B pair slots + one 8-B odd slot per lane) +// with no LDS stage in between. The inverse mapping (element -> byte) is +// packed384_byte_offset; the layout constants above define the geometry: +// 15 split slices in ascending split order; +// slice s byte base: s*4608 (8 long 288-k slices, 9 steps) +// 8*4608 + (s-8)*4096 (7 short 256-k slices, 8 steps); +// inside a slice: (step>>1)*1024 + lane*16 + (step&1)*8 + i, with +// lane = ((k-k_base)>>3 & 3)*16 + n_local, i = (k-k_base)&7. The odd 9th +// step of a long slice is a single 512-B slot: 4*1024 + lane*8 + i. +// One-time, out-of-timed-region, out-of-Graph; same byte count and the same +// graph-stable buffer. Every other (K, N) keeps the identity pack. Fallbacks +// decode this order via packed384_byte_offset. __global__ void w8a8_pack_weight_gateup_slicemajor_kernel( const int8_t* __restrict__ src, int8_t* __restrict__ dst) { const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - constexpr int64_t kTotal = 4096 * 384; // 1,572,864 - if (idx >= kTotal) { + if (idx >= kGateUpTotal) { return; } const int64_t tile = idx >> 16; // / 65536 const int64_t rem = idx - (tile << 16); // % 65536 (byte in tile) - const int64_t chunk = rem >> 4; // 16-B chunk index in tile - const int64_t in_chunk = rem & 15; // byte within the chunk - int64_t base; // chunk index of the slice start - int64_t n_local, kk_p; - if (chunk < 8 * 288) { // long slice (18 chunks per n) - const int64_t r = chunk % 288; - base = (chunk / 288) * 288; - n_local = r / 18; - kk_p = r % 18; - } else { // short slice (16 chunks per n) - const int64_t c2 = chunk - 8 * 288; - const int64_t r = c2 % 256; - base = 8 * 288 + (c2 / 256) * 256; - n_local = r / 16; - kk_p = r % 16; + int s; // split index + int step; // DUMMA k32 step inside the slice + int lane; // wavefront lane of the fragment byte + int i; // byte inside the lane's 8-B fragment + if (rem < kFragPackLongRegion) { // long 288-k slices + s = static_cast(rem / kFragPackLongBytes); + const int64_t r = rem - static_cast(s) * kFragPackLongBytes; + if (r < 4 * kFragPackPairBytes) { // 4 pair slots (2 steps each) + const int m = static_cast(r >> 10); // r / 1024 + const int rr = static_cast(r & 1023); + lane = rr >> 4; + const int pos = rr & 15; + step = 2 * m + (pos >> 3); + i = pos & 7; + } else { // odd 9th step: 512-B single slot + const int rr = static_cast(r - 4 * kFragPackPairBytes); + lane = rr >> 3; + i = rr & 7; + step = 8; + } + } else { // short 256-k slices + const int64_t c2 = rem - kFragPackLongRegion; + const int64_t sidx = c2 / kFragPackShortBytes; + s = kFragPackNumLong + static_cast(sidx); + const int64_t r = c2 - sidx * kFragPackShortBytes; + const int m = static_cast(r >> 10); + const int rr = static_cast(r & 1023); + lane = rr >> 4; + const int pos = rr & 15; + step = 2 * m + (pos >> 3); + i = pos & 7; } - // base == k_base (k index of the slice start) numerically: 288 k / 288 - // chunks per long slice, 256 / 256 per short slice. - const int64_t k = base + kk_p * 16 + in_chunk; + const int n_local = lane & 15; + const int k_base = (s < kFragPackNumLong) + ? s * kFragPackSliceLong + : kFragPackNumLong * kFragPackSliceLong + + (s - kFragPackNumLong) * kFragPackSliceShort; + const int64_t k = static_cast(k_base) + step * kFragPackStepK + + (lane >> 4) * 8 + i; dst[idx] = src[k * 384 + tile * 16 + n_local]; } -// Identity device-to-device scale pack (bootstrap). __global__ void w8a8_pack_scale_identity_kernel( const float* __restrict__ src, float* __restrict__ dst, @@ -1011,20 +777,19 @@ __global__ void w8a8_pack_scale_identity_kernel( } // namespace -// Timed GEMM entry point (inside the CUDA/HIP Graph region). -// -// Graph-safety contract: no allocation, no compilation, no autotuning, no -// weight packing, no host/device synchronization, no default-stream launch. -// Launches only on the caller-provided stream and uses only caller-provided -// tensors. The iteration-9 split-K=15 path uses the caller's workspace as the -// int32 partial planes (write-in-place every launch, no clear needed) plus 24 -// arrival counters in the workspace tail and issues exactly ONE kernel per -// replay (LDS-staged partial GEMM with a fused last-arrival combine tail); -// the iteration-1 unsplit kernel is the workspace-too-small fallback. +// --------------------------------------------------------------------------- +// launch_w8a8_gemm // -// Exact-shape specializations are guarded HERE, before the generic scalar -// fallback, so the paired M=2 API shape with the same (N, K) and every other -// shape still reach the scalar path below. +// Timed operator: launches on the caller-provided (PyTorch current) stream. +// No allocation, no compilation, no autotuning, no packing, no host or device +// synchronization and no default-stream launch. The assigned-shape path uses +// the caller-provided workspace for ONE int32 accumulator per 16-column tile +// (24 x 256 x 4 = 24,576 B) plus 24 arrival counters (96 B) and issues exactly +// ONE kernel per replay (split-K=15 one-wave zero-barrier GEMM with global +// atomicAdd accumulation and a fused last-arrival tail, register-direct A/B +// transport, no LDS slice stage); every accumulator element is reset by its +// tile's last arriver, so no clear is needed between replays. +// --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -1037,53 +802,41 @@ extern "C" void launch_w8a8_gemm( int n, int k, hipStream_t stream) { - (void)workspace; - (void)workspace_bytes; - - if (m <= 0 || n <= 0 || k <= 0) { - return; - } - const int64_t total = static_cast(m) * n; - if (total <= 0) { + if (m <= 0 || n <= 0 || k <= 0 || a == nullptr || b == nullptr || + out == nullptr) { return; } - // Exact-shape M=16 specialization (iteration 9): split-K=15 LDS-staged - // partial GEMM with a FUSED last-arrival combine tail -> exactly one - // launch per replay (the round-8 falsifiable branch fired: median - // 15.53 >= 15.4, so the two-launch structure is attacked instead of - // geometry). grid = 24 N-groups (16x16 tiles) x 15 K-slices = 360 one-wave - // blocks = exactly 3 blocks/CU on 120 CUs; K=4096 -> 15 non-uniform - // 32-aligned slices (8 x 288 = 9 DUMMA steps + 7 x 256 = 8 steps); each - // block stages its whole slice in LDS once (18,688 B/block -> 3 x 18,688 = - // 56,064 B/CU <= 64 KiB), runs a zero-barrier LDS-only K loop, stores its - // int32 plane, then the 15th arriver of its N-group sums the 15 planes in - // ascending k order, applies both scales and stores bf16. The 24 arrival - // counters live in the 96 B right after the 15 planes and are zeroed once - // before the first launch (outside the timed Graph in the normal - // eager-first flow). The iteration-1 unsplit DUMMA kernel remains the - // fallback when the caller's workspace cannot hold 15 partial planes plus - // the counters (368,640 + 96 B <= the workspace_split_k_capacity budget of - // 393,216 B); the paired M=2 API shape with the same (N,K) and every - // other shape still reach the scalar fallback below. + hip_bfloat16* out_bf16 = reinterpret_cast(out); + + // Exact-shape M=16 dispatch (hy3_tp8_shared_gate_up_proj_m16): split-K=15 + // one-wave zero-barrier partial GEMM + fused last-arrival combine tail. + // grid = 24 N-groups x 15 K-slices = 360 one-wave blocks = exactly 3 + // blocks/CU on 120 CUs; K=4096 -> 8 x 288 + 7 x 256 (32-aligned). + // Any other (m, n, k) - including the paired M=2 shape with the same + // (N, K) = (384, 4096), and every other M=16 (N, K) - falls through to the + // generic scalar fallback (which decodes the (4096,384) pack). if (m == 16 && n == 384 && k == 4096) { - constexpr int64_t kPartialBytes = - static_cast(kSplitK) * kDummaM * 384 * 4; // 368,640 B constexpr int64_t kCounterBytes = static_cast(kCounterInts) * sizeof(int32_t); // 96 B - if (workspace != nullptr && - workspace_bytes >= kPartialBytes + kCounterBytes) { - auto* partials = reinterpret_cast(workspace); - int32_t* counters = partials + kSplitK * kDummaM * 384; - // Zero the 24 arrival counters once before the first launch. Issued on + // Workspace is the per-tile int32 accumulator (24,576 B) followed by the + // 24 arrival counters. The caller's validated workspace is >= 393,216 B + // (16 partial planes), so this guard is strictly more permissive than the + // pre-iteration-7 requirement and the same exact-shape caller keeps the + // split-K path. + constexpr int64_t kScratchBytes = kAccumBytes + kCounterBytes; // 24,672 B + if (workspace != nullptr && workspace_bytes >= kScratchBytes) { + int32_t* accum = reinterpret_cast(workspace); + int32_t* counters = accum + kAccumInts; + // Zero the whole scratch region once before the first launch. Issued on // the caller stream in the eager precheck, so the captured Graph - // contains no memset node; if the first call is captured, the memset - // becomes a replayable Graph node (still exact). The kernel's - // self-reset keeps every later replay deterministic. - static bool counters_zeroed = false; - if (!counters_zeroed) { - counters_zeroed = true; - hipMemsetAsync(counters, 0, kCounterBytes, stream); + // contains no memset node in the normal eager-first flow; the kernel's + // fused tail self-resets every accumulator tile and counter, so every + // later replay restarts from zero. + static bool scratch_zeroed = false; + if (!scratch_zeroed) { + scratch_zeroed = true; + hipMemsetAsync(accum, 0, kScratchBytes, stream); } hipLaunchKernelGGL( w8a8_gemm_m16_dumma_splitk15_lds_staged_partial_kernel, @@ -1093,50 +846,63 @@ extern "C" void launch_w8a8_gemm( stream, a, b, - partials, + accum, x_scale, weight_scale, - reinterpret_cast(out), + out_bf16, counters, n, k); - } else { - hipLaunchKernelGGL(w8a8_gemm_m16_dumma_kernel, dim3(n / kDummaN), - dim3(kDummaBlockThreads), 0, stream, a, b, x_scale, - weight_scale, reinterpret_cast(out), - n, k); + return; } + // Workspace too small: unsplit M=16 DUMMA fallback. + hipLaunchKernelGGL( + w8a8_gemm_m16_dumma_kernel, + dim3(n / kDummaN), + dim3(kDummaBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_bf16, + n, + k); return; } - constexpr int kBlockThreads = 128; // multiple of the gfx928 wavefront (64) - const int grid = static_cast((total + kBlockThreads - 1) / kBlockThreads); - + const int64_t total = static_cast(m) * n; + const int64_t blocks = + (total + kScalarBlockThreads - 1) / kScalarBlockThreads; + const dim3 grid(static_cast(blocks)); + const dim3 block(static_cast(kScalarBlockThreads)); hipLaunchKernelGGL( w8a8_gemm_scalar_kernel, - dim3(grid), - dim3(kBlockThreads), + grid, + block, 0, stream, a, b, x_scale, weight_scale, - reinterpret_cast(out), + out_bf16, m, n, k); } -// Optional out-of-timed-region weight packing. Iteration 1: identity -// device-to-device copy of raw_weight -> packed_weight and weight_scale -> -// packed_weight_scale, valid for any (K, N) and any unmatched shape (the -// DUMMA kernel reads the raw [K, N] row-major layout). Iteration 17: the -// exact (k,n)==(4096,384) shape is packed per-tile slice-major-chunk -// instead (see w8a8_pack_weight_gateup_slicemajor_kernel: same 16-B chunks -// as iteration 15, chunk order (slice, n, k-chunk) so staged wave-loads are -// 1,024-B contiguous); every other (K, N) keeps the identity copy. Same -// byte count, same buffer, one-time and outside the timed Graph region. +// --------------------------------------------------------------------------- +// launch_pack_w8a8_weight +// +// Optional out-of-timed-region weight packing. Identity device-to-device copy +// of raw_weight -> packed_weight and weight_scale -> packed_weight_scale valid +// for any (K, N); the exact (k,n)==(4096,384) shape is packed per-tile +// DUMMA-fragment-order instead (same byte count, same buffer, one-time and +// outside the timed Graph region). The packed copies are asynchronous on the +// caller stream. +// --------------------------------------------------------------------------- extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -1146,7 +912,7 @@ extern "C" void launch_pack_w8a8_weight( int n, hipStream_t stream) { const int64_t total = static_cast(k) * n; - if (total > 0) { + if (raw_weight != nullptr && packed_weight != nullptr && total > 0) { constexpr int kPackBlock = 256; const int grid = static_cast((total + kPackBlock - 1) / kPackBlock); if (k == 4096 && n == 384) { @@ -1165,7 +931,7 @@ extern "C" void launch_pack_w8a8_weight( total); } } - if (n > 0) { + if (weight_scale != nullptr && packed_weight_scale != nullptr && n > 0) { constexpr int kScaleBlock = 64; const int grid = (n + kScaleBlock - 1) / kScaleBlock; hipLaunchKernelGGL( 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. diff --git a/metainfer/tasks/evolve_kernel/orchestrator/orchestrator.py b/metainfer/tasks/evolve_kernel/orchestrator/orchestrator.py index e2f0bcbb..ab8bf92f 100644 --- a/metainfer/tasks/evolve_kernel/orchestrator/orchestrator.py +++ b/metainfer/tasks/evolve_kernel/orchestrator/orchestrator.py @@ -109,6 +109,7 @@ def run_with_requirements( effort=effort, extra_add_dirs=[repo_root, logs_root], snapshot_file=paths["agents_file"], + state_dir=state_dir, ) orch = Orchestrator(req=req, store=store, cfg=cfg, manager=manager) diff --git a/metainfer/tasks/evolve_kernel/server/_state_readers.py b/metainfer/tasks/evolve_kernel/server/_state_readers.py index 27effcb6..17f28164 100644 --- a/metainfer/tasks/evolve_kernel/server/_state_readers.py +++ b/metainfer/tasks/evolve_kernel/server/_state_readers.py @@ -5,6 +5,7 @@ from __future__ import annotations +import difflib import json from pathlib import Path from typing import Any, Dict, List, Optional @@ -213,3 +214,94 @@ def read_reference_kernel(workspace_dir: Path) -> Dict[str, Any]: "path": str(path), "lines": len(code.splitlines()), } + + +# --------------------------------------------------------------------------- # +# Kernel source & diff +# --------------------------------------------------------------------------- # + +# Metadata carried alongside a single kernel's source (the fields the diff +# view needs to explain a change: how fast it was, which iteration added it, +# and which library kernel it was derived from). +_KERNEL_META_KEYS = ( + "exec_time_ms", + "complexity_score", + "combined_score", + "iteration_added", + "parent_id", +) + + +def read_kernel_source(workspace_dir: Path, kernel_id: str) -> Dict[str, Any]: + """Return the full source + metadata for one library kernel, by id. + + ``read_kernel_library`` ships every kernel's code in one payload; this + resolves a single kernel so callers that already know the id (the diff + view) don't have to scan the list, and so the lookup has one definition. + """ + for k in read_kernel_library(workspace_dir)["kernels"]: + if str(k.get("id")) == str(kernel_id): + code = k.get("code") or "" + return { + "exists": True, + "id": k.get("id"), + "code": code, + "lines": len(code.splitlines()), + "meta": {key: k.get(key) for key in _KERNEL_META_KEYS}, + } + return {"exists": False, "id": kernel_id, "code": "", "lines": 0, "meta": {}} + + +def read_kernel_diff(workspace_dir: Path, kernel_id: str, + base: str = "reference") -> Dict[str, Any]: + """Unified diff of a library kernel against a baseline. + + ``base`` is ``"reference"`` (the original kernel the task started from) + or ``"parent"`` (the library kernel this one was derived from). Returns + the diff text plus added/removed line counts; a missing kernel, a root + kernel with no parent, or a parent that is no longer in the library + (the library evicts past ``MAX_LIBRARY_SIZE``) yields ``exists: False`` + with an ``error`` reason rather than raising. + """ + src = read_kernel_source(workspace_dir, kernel_id) + if not src["exists"]: + return _diff_missing(base, f"unknown kernel {kernel_id!r}") + + if base == "parent": + parent_id = src["meta"].get("parent_id") + if not parent_id: + return _diff_missing(base, "kernel has no parent") + base_src = read_kernel_source(workspace_dir, parent_id) + if not base_src["exists"]: + return _diff_missing( + base, f"parent {str(parent_id)[:8]} is no longer in the library") + base_code = base_src["code"] + base_label = f"kernel {str(parent_id)[:8]}" + else: + ref = read_reference_kernel(workspace_dir) + base_code = ref["code"] if ref["exists"] else "" + base_label = "reference" + + diff_lines = list(difflib.unified_diff( + base_code.splitlines(), (src["code"] or "").splitlines(), + fromfile=f"a/{base_label}", + tofile=f"b/kernel {str(kernel_id)[:8]}", + lineterm="", + )) + added = sum(1 for ln in diff_lines + if ln.startswith("+") and not ln.startswith("+++")) + removed = sum(1 for ln in diff_lines + if ln.startswith("-") and not ln.startswith("---")) + return { + "exists": True, + "base": base, + "base_label": base_label, + "diff": "\n".join(diff_lines), + "added": added, + "removed": removed, + } + + +def _diff_missing(base: str, error: str) -> Dict[str, Any]: + return {"exists": False, "base": base, "base_label": base, + "diff": "", "added": 0, "removed": 0, "error": error} diff --git a/metainfer/tasks/evolve_kernel/server/routes.py b/metainfer/tasks/evolve_kernel/server/routes.py index 82b91d58..07cf05a9 100644 --- a/metainfer/tasks/evolve_kernel/server/routes.py +++ b/metainfer/tasks/evolve_kernel/server/routes.py @@ -95,6 +95,16 @@ def ok_reference_kernel(task_id: str) -> Dict[str, Any]: require_task_type(entry, PLUGIN_TYPE) return _state_readers.read_reference_kernel(workspace_dir_for(entry)) + # ---- Kernel diff ---- + + @router.get("/kernels/{kernel_id}/diff") + def ok_kernel_diff(task_id: str, kernel_id: str, + base: str = "reference") -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return _state_readers.read_kernel_diff( + workspace_dir_for(entry), kernel_id, base) + # ---- QA ---- register_qa_routes(router, plugin, prefix="/qa") diff --git a/metainfer/tasks/evolve_kernel/static/ok-detail.js b/metainfer/tasks/evolve_kernel/static/ok-detail.js index 9e316db0..b6c3b015 100644 --- a/metainfer/tasks/evolve_kernel/static/ok-detail.js +++ b/metainfer/tasks/evolve_kernel/static/ok-detail.js @@ -9,13 +9,14 @@ import { html } from "htm/preact"; import { useCallback, useEffect, useState } from "preact/hooks"; import { StateGraph } from "app/ok-state-graph"; -import { KernelLibrary } from "app/ok-kernel-library"; +import { KernelLibrary, formatMs } from "app/ok-kernel-library"; import { AgentsPanel } from "app/agents-panel"; import { Timeline } from "app/timeline"; +import { fmtDur } from "app/utils"; import { getIterations, getCharts, getStateGraph, getKernelLibrary, getCorrectnessHarness, getPerfHarness, - getReferenceKernel, + getReferenceKernel, getKernelDiff, } from "app/ok-runtime-api"; const withTimeout = (p, ms = 8000) => @@ -128,6 +129,76 @@ function MetricCards({ library, iterations }) {
`; } +// ---- Iteration progress ---- + +// One row per optimization iteration. The metric cards only surface the *last* +// iteration's speedup and the library's current best, so a run that is stuck — +// or that regressed after a good iteration — looks identical to one that is +// climbing. This table is the per-iteration trail: what each iteration set out +// to do (`goal`), how it ended, and what it measured. Rows that set a new best +// exec time are marked, which is the "is it actually improving?" signal. +// `status` is written by IterationRecord as one of running|success|failed. +const ITER_STATUS_CLASS = { + success: "green", + failed: "red", + running: "yellow", // in flight — no result yet, but nothing broke either +}; + +function IterationProgress({ iterations }) { + const rows = iterations || []; + if (rows.length === 0) { + return html`

No iterations recorded yet.

`; + } + + // Best-so-far marking. Iterations that failed or are still running carry no + // `perf.exec_time_ms`, so they simply never take the mark. + const bestIterations = new Set(); + let best = Infinity; + for (const it of rows) { + const ms = it.perf && it.perf.exec_time_ms; + if (typeof ms === "number" && ms > 0 && ms < best) { + best = ms; + bestIterations.add(it.iteration); + } + } + + return html`
+ + + + + + + + + ${rows.map((it) => { + const perf = it.perf || {}; + const improved = bestIterations.has(it.iteration); + // These numbers come out of an LLM-generated benchmark harness, so + // coerce rather than trusting the type (a string "1.02" is fine). + const speedup = perf.speedup == null ? null : Number(perf.speedup); + return html` + + + + + + + + + `; + })} + +
#StatusExec TimeSpeedupLibraryDurationGoal
${it.iteration} + ${it.status || "—"} + ${perf.exec_time_ms != null ? formatMs(perf.exec_time_ms) : "—"} + ${improved ? html`▲` : null} + ${Number.isFinite(speedup) ? speedup.toFixed(3) + "×" : "—"}${it.kernel_library_size ?? "—"}${it.duration_s ? fmtDur(it.duration_s) : "—"} + ${it.goal || (it.failure_reason ? "failed: " + it.failure_reason : "—")} +
+
`; +} + // ---- Harness status badges ---- function HarnessStatus({ harness, label }) { @@ -140,6 +211,101 @@ function HarnessStatus({ harness, label }) { return html`${label}: not generated`; } +// ---- Kernel source / diff inspector ---- + +// Line-numbered source rendering: the raw string in a 
 gives no way to
+// cite a line, so emit one 
 per line with a gutter number.
+function SourceLines({ code }) {
+  const lines = (code || "").split("\n");
+  return html`

+    ${lines.map((ln, i) => html`
+      

+        ${i + 1}
+        ${ln}
+      

+    `)}
+  
`;
+}
+
+// Unified-diff rendering: classify each line by prefix so additions and
+// deletions stand out; hunk headers (@@) get their own muted styling.
+function DiffLines({ diff }) {
+  const lines = (diff || "").split("\n");
+  if (!lines.some((l) => l.length > 0)) {
+    return html`
No differences.
`;
+  }
+  return html`

+    ${lines.map((ln, i) => {
+      let cls = "ok-diff-ctx";
+      if (ln.startsWith("+++") || ln.startsWith("---")) cls = "ok-diff-meta";
+      else if (ln.startsWith("@@")) cls = "ok-diff-hunk";
+      else if (ln.startsWith("+")) cls = "ok-diff-add";
+      else if (ln.startsWith("-")) cls = "ok-diff-del";
+      return html`
${ln || " "}
`;
+    })}
+  
`;
+}
+
+const DIFF_MODES = [
+  { key: "source", label: "Source" },
+  { key: "reference", label: "Diff vs reference" },
+  { key: "parent", label: "Diff vs parent" },
+];
+
+function KernelInspector({ taskId, kernel, onClose }) {
+  const [mode, setMode] = useState("source");
+  const [diff, setDiff] = useState(null);
+  const [diffErr, setDiffErr] = useState(null);
+
+  useEffect(() => {
+    if (mode === "source" || !taskId || !kernel) return;
+    let cancelled = false;
+    setDiff(null);
+    setDiffErr(null);
+    withTimeout(getKernelDiff(taskId, kernel.id, mode))
+      .then((d) => { if (!cancelled) setDiff(d); })
+      .catch((e) => { if (!cancelled) setDiffErr(String(e.message || e)); });
+    return () => { cancelled = true; };
+  }, [taskId, kernel && kernel.id, mode]);
+
+  const code = kernel.code || kernel.code_preview || "";
+  const parentId = kernel.parent_id ? String(kernel.parent_id).slice(0, 8) : null;
+
+  return html`

+    
Kernel: ${String(kernel.id).slice(0, 8)}…
+      
+    

+    

+      

+        ${DIFF_MODES.map((m) => html`
+          
+        `)}
+      

+      ${mode !== "source" && diff && diff.exists ? html`
+        
+          +${diff.added}
+          −${diff.removed}
+          vs ${diff.base_label}
+        
+      ` : null}
+    

+
+    ${mode === "source"
+      ? html`<${SourceLines} code=${code} />`
+      : diffErr
+        ? html`
Failed to load diff: ${diffErr}
`
+        : !diff
+          ? html`
Loading diff…
`
+          : !diff.exists
+            ? html`
${diff.error || "No diff available."}
`
+            : html`<${DiffLines} diff=${diff.diff} />`}
+  
`;
+}
+
 // ---- Main view ----
 
 export default function OptKernelDetailView({
@@ -167,6 +333,15 @@ export default function OptKernelDetailView({
 
     <${MetricCards} library=${rt.library} iterations=${rt.iterations} />
 
+    

+      
Iteration Progress
+        
+          (▲ = new best exec time)
+        
+      

+      <${IterationProgress} iterations=${rt.iterations} />
+    

+
     

       

         
State Machine

@@ -185,12 +360,10 @@ export default function OptKernelDetailView({
     

 
     ${selectedKernel ? html`
-      

-        
Kernel: ${selectedKernel.id.slice(0, 8)}…
-          
-        

-        
${selectedKernel.code || selectedKernel.code_preview || "Code not available"}

-      

+      <${KernelInspector}
+        taskId=${taskId}
+        kernel=${selectedKernel}
+        onClose=${() => setSelectedKernel(null)} />
     ` : null}
 
     

diff --git a/metainfer/tasks/evolve_kernel/static/ok-kernel-library.js b/metainfer/tasks/evolve_kernel/static/ok-kernel-library.js
index 6bfc58ac..31dc6576 100644
--- a/metainfer/tasks/evolve_kernel/static/ok-kernel-library.js
+++ b/metainfer/tasks/evolve_kernel/static/ok-kernel-library.js
@@ -10,7 +10,7 @@ function scoreClass(val) {
   return "ok-score-bad";
 }
 
-function formatMs(ms) {
+export function formatMs(ms) {
   if (ms == null || ms === 0) return "—";
   if (ms < 1) return (ms * 1000).toFixed(1) + " μs";
   if (ms < 1000) return ms.toFixed(3) + " ms";
diff --git a/metainfer/tasks/evolve_kernel/static/ok-runtime-api.js b/metainfer/tasks/evolve_kernel/static/ok-runtime-api.js
index 8f572de0..f2ef898a 100644
--- a/metainfer/tasks/evolve_kernel/static/ok-runtime-api.js
+++ b/metainfer/tasks/evolve_kernel/static/ok-runtime-api.js
@@ -51,6 +51,14 @@ export async function getReferenceKernel(taskId) {
   return res.json();
 }
 
+export async function getKernelDiff(taskId, kernelId, base = "reference") {
+  const res = await fetch(
+    `${BASE(taskId)}/kernels/${encodeURIComponent(kernelId)}/diff` +
+    `?base=${encodeURIComponent(base)}`);
+  if (!res.ok) throw new Error(`kernel diff: ${res.status}`);
+  return res.json();
+}
+
 export async function getRetrospective(taskId, n) {
   const res = await fetch(`${BASE(taskId)}/iterations/${n}/retrospective`);
   if (!res.ok) throw new Error(`retrospective ${n}: ${res.status}`);
diff --git a/metainfer/tasks/evolve_kernel/static/ok.css b/metainfer/tasks/evolve_kernel/static/ok.css
index 8b4988c2..47a17ae7 100644
--- a/metainfer/tasks/evolve_kernel/static/ok.css
+++ b/metainfer/tasks/evolve_kernel/static/ok.css
@@ -115,6 +115,33 @@
   color: #d29922;
 }
 
+.ok-badge.red {
+  background: rgba(248, 81, 73, 0.15);
+  color: #f85149;
+}
+
+/* ---- Iteration progress table ---- */
+
+.ok-best-tag {
+  color: #3fb950;
+  margin-left: 0.25rem;
+}
+
+/* Goals are free-form sentences; keep rows scannable and let the full text
+   show on hover (the cell carries it as a title). */
+.ok-goal {
+  max-width: 26rem;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+/* A long run logs dozens of iterations — keep the panel bounded. */
+.ok-iter-scroll {
+  max-height: 22rem;
+  overflow: auto;
+}
+
 .ok-flow-indicator {
   display: flex;
   align-items: center;
@@ -162,3 +189,107 @@
   color: var(--muted, #8b949e);
   margin-top: 0.25rem;
 }
+
+/* ---- Kernel source / diff inspector ---- */
+
+.ok-inspector-toolbar {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 0.75rem;
+  margin-bottom: 0.75rem;
+  flex-wrap: wrap;
+}
+
+.ok-inspector-modes {
+  display: flex;
+  gap: 0.4rem;
+}
+
+.ok-inspector-modes .btn.active {
+  background: var(--accent-bg, rgba(88, 166, 255, 0.15));
+  color: var(--accent, #58a6ff);
+  border-color: var(--accent, #58a6ff);
+}
+
+.ok-diff-stat {
+  display: inline-flex;
+  align-items: center;
+  gap: 0.4rem;
+  font-size: 0.8rem;
+  font-family: monospace;
+}
+
+.ok-diff-stat-add {
+  color: #3fb950;
+  font-weight: 600;
+}
+
+.ok-diff-stat-del {
+  color: #f85149;
+  font-weight: 600;
+}
+
+.ok-source,
+.ok-diff {
+  max-height: 600px;
+  overflow: auto;
+  background: var(--code-bg, #0d1117);
+  border: 1px solid var(--border-color, #30363d);
+  border-radius: 6px;
+  font-family: monospace;
+  font-size: 0.8rem;
+  line-height: 1.45;
+}
+
+.ok-diff.empty {
+  padding: 1rem;
+}
+
+.ok-diff.empty.error {
+  color: #f85149;
+}
+
+.ok-source-row,
+.ok-diff-row {
+  display: flex;
+  white-space: pre;
+}
+
+.ok-source-num {
+  flex: 0 0 auto;
+  width: 3.5rem;
+  padding: 0 0.6rem;
+  text-align: right;
+  color: var(--muted, #8b949e);
+  user-select: none;
+  border-right: 1px solid var(--border-color, #30363d);
+}
+
+.ok-source-line {
+  flex: 1 1 auto;
+  padding: 0 0.6rem;
+}
+
+.ok-diff-row {
+  padding: 0 0.75rem;
+}
+
+.ok-diff-add {
+  background: rgba(63, 185, 80, 0.15);
+  color: #3fb950;
+}
+
+.ok-diff-del {
+  background: rgba(248, 81, 73, 0.15);
+  color: #f85149;
+}
+
+.ok-diff-hunk {
+  color: var(--accent, #58a6ff);
+  background: var(--accent-bg, rgba(88, 166, 255, 0.08));
+}
+
+.ok-diff-meta {
+  color: var(--muted, #8b949e);
+}
diff --git a/metainfer/tasks/evolve_kernel/tests/test_iteration_progress.py b/metainfer/tasks/evolve_kernel/tests/test_iteration_progress.py
new file mode 100644
index 00000000..fe6df0a3
--- /dev/null
+++ b/metainfer/tasks/evolve_kernel/tests/test_iteration_progress.py
@@ -0,0 +1,63 @@
+"""Contract for the fields the WebUI "Iteration Progress" table renders.
+
+``static/ok-detail.js`` builds one row per iteration straight from
+``read_iterations``: ``status``, ``goal``, ``duration_s``,
+``kernel_library_size`` and ``perf.{exec_time_ms,speedup}``. Iterations that
+never got as far as measuring (failed / still running) must still come back as
+rows — the table falls back to "—" for the metrics it doesn't have — so the
+reader must not require a complete record, and an iteration that was only
+partially written must not take the whole list down.
+"""
+
+from __future__ import annotations
+
+import json
+
+from metainfer.tasks.evolve_kernel.server import _state_readers as sr
+
+
+def _write_iteration(state_dir, n, **fields):
+    # The writer dumps `IterationRecord.to_dict()` (== asdict), so every field is
+    # always on disk — an unmeasured iteration carries `perf: {}`, not a missing
+    # key. Fixtures pass the fields explicitly to mirror that shape.
+    iters = state_dir / "iterations"
+    iters.mkdir(parents=True, exist_ok=True)
+    (iters / f"{n:03d}.json").write_text(
+        json.dumps({"iteration": n, **fields}), encoding="utf-8")
+
+
+def test_read_iterations_keeps_the_fields_the_table_renders(tmp_path):
+    _write_iteration(tmp_path, 2, status="success", goal="wider tiles + double buffer",
+                     duration_s=41.5, kernel_library_size=3,
+                     perf={"exec_time_ms": 0.121, "speedup": 1.02,
+                           "combined_score": 58.1})
+    _write_iteration(tmp_path, 1, status="failed", duration_s=2.0,
+                     failure_reason="harness crashed", perf={})
+
+    rows = sr.read_iterations(tmp_path)
+    assert [r["iteration"] for r in rows] == [1, 2]
+
+    failed, measured = rows
+    # A failed iteration is still a row; its empty `perf` leaves the metric
+    # columns at "—" in the table.
+    assert failed["status"] == "failed"
+    assert failed["failure_reason"] == "harness crashed"
+    assert failed["perf"] == {}
+
+    assert measured["perf"]["exec_time_ms"] == 0.121
+    assert measured["perf"]["speedup"] == 1.02
+    assert measured["goal"] == "wider tiles + double buffer"
+    assert measured["duration_s"] == 41.5
+    assert measured["kernel_library_size"] == 3
+
+
+def test_read_iterations_skips_an_unreadable_record(tmp_path):
+    _write_iteration(tmp_path, 1, status="success", perf={"exec_time_ms": 0.2})
+    (tmp_path / "iterations" / "002.json").write_text("{not json", encoding="utf-8")
+
+    rows = sr.read_iterations(tmp_path)
+    assert [r["iteration"] for r in rows] == [1]
+
+
+def test_read_iterations_empty_when_no_directory(tmp_path):
+    assert sr.read_iterations(tmp_path) == []
diff --git a/metainfer/tasks/evolve_kernel/tests/test_kernel_diff.py b/metainfer/tasks/evolve_kernel/tests/test_kernel_diff.py
new file mode 100644
index 00000000..ed235404
--- /dev/null
+++ b/metainfer/tasks/evolve_kernel/tests/test_kernel_diff.py
@@ -0,0 +1,189 @@
+"""Unit tests for the kernel source/diff readers in server/_state_readers.py.
+
+Covers ``read_kernel_source`` (resolve one library kernel by id) and
+``read_kernel_diff`` (unified diff against the reference or a parent kernel),
+including the non-raising failure paths the diff view depends on.
+"""
+
+from __future__ import annotations
+
+import json
+import tempfile
+from pathlib import Path
+
+from metainfer.tasks.evolve_kernel.server._state_readers import (
+    read_kernel_diff,
+    read_kernel_source,
+)
+
+REF_CODE = "\n".join([
+    "def kernel(x):",
+    "    return x",
+])
+
+PARENT_CODE = "\n".join([
+    "def kernel(x):",
+    "    # parent tweak",
+    "    return x",
+])
+
+CHILD_CODE = "\n".join([
+    "def kernel(x):",
+    "    # child tweak",
+    "    return x * 2",
+])
+
+
+def _write_library(ws: Path, kernels):
+    (ws / "kernel_library.json").write_text(
+        json.dumps(kernels, indent=2), encoding="utf-8")
+
+
+def _write_reference(ws: Path, code: str):
+    ref_dir = ws / "reference"
+    ref_dir.mkdir(parents=True, exist_ok=True)
+    (ref_dir / "original_kernel.py").write_text(code, encoding="utf-8")
+
+
+def _kernel(kid, code, parent_id=None, **meta):
+    entry = {
+        "id": kid,
+        "code": code,
+        "parent_id": parent_id,
+        "exec_time_ms": meta.get("exec_time_ms", 1.0),
+        "complexity_score": meta.get("complexity_score", 0.0),
+        "combined_score": meta.get("combined_score", 7.0),
+        "iteration_added": meta.get("iteration_added", 0),
+    }
+    return entry
+
+
+# --------------------------------------------------------------------------- #
+# read_kernel_source
+# --------------------------------------------------------------------------- #
+
+
+class TestReadKernelSource:
+    def test_finds_kernel_by_id(self):
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_library(ws, [
+                _kernel("aaa", REF_CODE),
+                _kernel("bbb", CHILD_CODE, parent_id="aaa",
+                        exec_time_ms=2.5, iteration_added=7),
+            ])
+            src = read_kernel_source(ws, "bbb")
+            assert src["exists"] is True
+            assert src["id"] == "bbb"
+            assert src["code"] == CHILD_CODE
+            assert src["lines"] == len(CHILD_CODE.splitlines())
+            assert src["meta"]["parent_id"] == "aaa"
+            assert src["meta"]["exec_time_ms"] == 2.5
+            assert src["meta"]["iteration_added"] == 7
+
+    def test_matches_by_string_coercion(self):
+        """ids may be written as non-str (e.g. int); lookup must coerce."""
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_library(ws, [{"id": 42, "code": REF_CODE}])
+            src = read_kernel_source(ws, "42")
+            assert src["exists"] is True
+            assert src["code"] == REF_CODE
+
+    def test_missing_kernel(self):
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_library(ws, [_kernel("aaa", REF_CODE)])
+            src = read_kernel_source(ws, "nope")
+            assert src["exists"] is False
+            assert src["code"] == ""
+            assert src["meta"] == {}
+
+    def test_no_library_file(self):
+        with tempfile.TemporaryDirectory() as d:
+            src = read_kernel_source(Path(d), "anything")
+            assert src["exists"] is False
+
+
+# --------------------------------------------------------------------------- #
+# read_kernel_diff
+# --------------------------------------------------------------------------- #
+
+
+class TestReadKernelDiff:
+    def test_diff_vs_reference(self):
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_reference(ws, REF_CODE)
+            _write_library(ws, [_kernel("child", CHILD_CODE)])
+            diff = read_kernel_diff(ws, "child", base="reference")
+            assert diff["exists"] is True
+            assert diff["base"] == "reference"
+            assert diff["base_label"] == "reference"
+            assert diff["added"] == 2  # "# child tweak" and "return x * 2"
+            assert diff["removed"] == 1  # "return x"
+            assert "a/reference" in diff["diff"]
+            assert "b/kernel child" in diff["diff"]
+
+    def test_diff_vs_parent(self):
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_reference(ws, REF_CODE)
+            _write_library(ws, [
+                _kernel("parent1234", PARENT_CODE),
+                _kernel("child5678", CHILD_CODE, parent_id="parent1234"),
+            ])
+            diff = read_kernel_diff(ws, "child5678", base="parent")
+            assert diff["exists"] is True
+            assert diff["base_label"] == "kernel parent12"
+            assert "a/kernel parent12" in diff["diff"]
+
+    def test_root_kernel_has_no_parent(self):
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_reference(ws, REF_CODE)
+            _write_library(ws, [_kernel("root", REF_CODE)])
+            diff = read_kernel_diff(ws, "root", base="parent")
+            assert diff["exists"] is False
+            assert diff["error"] == "kernel has no parent"
+
+    def test_parent_evicted_from_library(self):
+        """A parent no longer present (evicted) must not raise."""
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_reference(ws, REF_CODE)
+            _write_library(ws, [
+                _kernel("child5678", CHILD_CODE, parent_id="gone0000"),
+            ])
+            diff = read_kernel_diff(ws, "child5678", base="parent")
+            assert diff["exists"] is False
+            assert "no longer in the library" in diff["error"]
+
+    def test_unknown_kernel(self):
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_reference(ws, REF_CODE)
+            _write_library(ws, [_kernel("aaa", REF_CODE)])
+            diff = read_kernel_diff(ws, "ghost")
+            assert diff["exists"] is False
+            assert "unknown kernel" in diff["error"]
+
+    def test_missing_reference_is_full_addition(self):
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_library(ws, [_kernel("child", CHILD_CODE)])
+            diff = read_kernel_diff(ws, "child", base="reference")
+            assert diff["exists"] is True
+            assert diff["removed"] == 0
+            assert diff["added"] == len(CHILD_CODE.splitlines())
+
+    def test_identical_code_yields_empty_body(self):
+        with tempfile.TemporaryDirectory() as d:
+            ws = Path(d)
+            _write_reference(ws, REF_CODE)
+            _write_library(ws, [_kernel("same", REF_CODE)])
+            diff = read_kernel_diff(ws, "same", base="reference")
+            assert diff["exists"] is True
+            assert diff["added"] == 0
+            assert diff["removed"] == 0
+            assert diff["diff"] == ""
diff --git a/metainfer/tasks/find_low_hanging_kernel/orchestrator/orchestrator.py b/metainfer/tasks/find_low_hanging_kernel/orchestrator/orchestrator.py
index 8350f887..1735ac58 100644
--- a/metainfer/tasks/find_low_hanging_kernel/orchestrator/orchestrator.py
+++ b/metainfer/tasks/find_low_hanging_kernel/orchestrator/orchestrator.py
@@ -131,6 +131,7 @@ def run_with_requirements(
         effort=effort,
         extra_add_dirs=[repo_root, logs_root, workspace_dir, *user_paths],
         snapshot_file=paths["agents_file"],
+        state_dir=state_dir,
     )
     pipeline = Pipeline(req=req, store=store, cfg=cfg, manager=manager)
 
diff --git a/metainfer/tasks/gen_cpp_infer_framework/orchestrator/orchestrator.py b/metainfer/tasks/gen_cpp_infer_framework/orchestrator/orchestrator.py
index 1cbf1570..ce9c3bf4 100644
--- a/metainfer/tasks/gen_cpp_infer_framework/orchestrator/orchestrator.py
+++ b/metainfer/tasks/gen_cpp_infer_framework/orchestrator/orchestrator.py
@@ -225,6 +225,7 @@ def run_with_requirements(
         extra_add_dirs=[notebooks_dir, repo_root, workspace_dir, logs_root, state_dir],
         snapshot_file=paths["agents_file"],
         budget=budget,
+        state_dir=state_dir,
     )
     # Wire the hard-exhausted callback NOW that the manager exists.
     # When the hard threshold is crossed, every in-flight agent gets
diff --git a/metainfer/tasks/gen_infer_framework/orchestrator/orchestrator.py b/metainfer/tasks/gen_infer_framework/orchestrator/orchestrator.py
index 4ee576aa..9b070f30 100644
--- a/metainfer/tasks/gen_infer_framework/orchestrator/orchestrator.py
+++ b/metainfer/tasks/gen_infer_framework/orchestrator/orchestrator.py
@@ -213,6 +213,7 @@ def run_with_requirements(
         extra_add_dirs=[notebooks_dir, repo_root, workspace_dir, logs_root],
         snapshot_file=paths["agents_file"],
         budget=budget,
+        state_dir=state_dir,
     )
     # Wire the hard-exhausted callback NOW that the manager exists.
     # When the hard threshold is crossed, every in-flight agent gets
diff --git a/metainfer/tasks/gen_infer_framework_cpp/orchestrator/orchestrator.py b/metainfer/tasks/gen_infer_framework_cpp/orchestrator/orchestrator.py
index da8c8f63..739dae69 100644
--- a/metainfer/tasks/gen_infer_framework_cpp/orchestrator/orchestrator.py
+++ b/metainfer/tasks/gen_infer_framework_cpp/orchestrator/orchestrator.py
@@ -262,6 +262,7 @@ def run_with_requirements(
         ],
         snapshot_file=paths["agents_file"],
         budget=budget,
+        state_dir=state_dir,
     )
     # Wire the hard-exhausted callback NOW that the manager exists.
     # When the hard threshold is crossed, every in-flight agent gets
diff --git a/metainfer/tasks/harness_evolve/FLOW.md b/metainfer/tasks/harness_evolve/FLOW.md
new file mode 100644
index 00000000..383df151
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/FLOW.md
@@ -0,0 +1,252 @@
+# HE 流程（固定协议）
+
+> **这是 harness_evolve（AHE 外层循环）的规范文本，不是设计笔记。**
+> 代码必须与本文一致；改协议 = 改本文 + 改代码 + 改契约测试
+> （`tests/test_flow_contract.py` 会检查本文存在且关键约束仍在）。
+> 实现位置见文末「代码对应」。
+
+---
+
+## 0. 边界：HE 只发起 DKAO 任务
+
+**HE 的唯一职责：把这一轮的算子清单交给 DKAO，等它交回测量结果，据此判门、决定下一步。**
+
+| 事项 | 归属 |
+|---|---|
+| 取题 / 组卷（这一轮测哪些算子） | **HE** |
+| 把每个算子起成一个 DKAO 子任务、传 harness 快照、等它跑完、收测量结果 | **HE** |
+| 判性能门 / 泛化门、决定 analyze→evolve 还是待批 | **HE** |
+| 每轮题数 = 4 的倍数（一波铺满 4 张 HCU） | **HE**（只做这个尺寸决定） |
+| 卡干不干净（VRAM ≤ 90% 且 HCU == 0）、等卡、闸门超时 | **DKAO**（`ensure_measurement_gate`） |
+| 卡被抢导致的污染测量作废 / 回滚 | **DKAO**（contamination 处理） |
+| 内核编译、基准、正确性、PMC、最终验收 | **DKAO** |
+| 子任务崩溃/工具链缺失等"环境欠一次测量"的重试 | **HE**（重发一次 DKAO，见 §5） |
+
+HE **不做**的事（照这条检查任何新代码）：
+
+- 不租卡（不调用 GPU broker）、不读卡状态、不因为"卡忙"等待或否决；
+- 不 `SIGSTOP`/冻结子进程（那是早前为"卡被抢"设计的机制，已废弃：卡的事是 DKAO 的事）；
+- 不判 `measurement_suspect`、不写污染窗口；
+- 不因为自身对环境的判断而终止一轮——能终止一轮的只有"子任务没交回可用测量"这件事本身。
+
+HE 为每道题只做一个**布局决定**：设备 `= index % 4`（第 i 道题绑第 i%4 张卡），
+它只保证"一波不超过 4 道题、一题一张卡"，不代表 HE 认为那张卡空闲。
+把每道题**立刻**发给 DKAO；卡不空时，是 DKAO 的子进程在闸门那里等。
+
+---
+
+## 1. 取题与组卷
+
+### 1.1 每轮题数固定为 4 的倍数
+
+- 机器 4 张 HCU（`GPU_COUNT = 4`），**一波 = 4 题**；一轮题数只能是 4 的倍数。
+- 表单字段 `per_round_budget`（Round budget）只提供 4 的倍数选项；CLI/接口传非 4 倍数时**向上取整**
+  （3→4、5→8、7→8），合法值原样保留（4→4、12→12）；非法/空 → 默认 4（`OPERATORS_PER_GATE`）。
+- 题数**只按配置抽**，不按空闲卡数缩水：卡忙是 DKAO 的事，缩题数会让一轮考不出结论。
+  （历史上那个 `fit_questions_to_devices` 反向开关已移除。）
+
+### 1.2 抽题规则
+
+从注册池（`pool_source`，真实实测池）中抽，**随机**，且：
+
+- 只抽**可判**算子：既有 variant 又有可用基线的才抽（否则算不出 win/loss，只会稀释比例）；
+  可判算子不足一波（4 个）时**直接停任务**（`final_status=no_question_pool`，事件
+  `variant_pool_too_small`）——缩到 3 题会改变门的含义（"≥75% 赢"在 3 题上等于"必须全胜"）；
+- **保证同时含 decode 与 prefill**：至少 1 个 decode（`M ≤ 32`）和 1 个 prefill（`M > 32`）；
+  其余名额随机填满——一轮只考一类 shape 会把更难、更有价值的那一半漏掉；
+- 第 1 轮（baseline）：随机抽 4 题 = **A1**，这 4 题**同时固化为本轮的泛化考卷**（只写一次，之后永不替换）；
+- 第 2 轮起（performance）：**重新随机抽 4 题 = A_i**，并**排除考卷**（`A_i ≠ A1`，也不同于上一轮的题）；
+- **抽题可复现**：随机种子由 `(task_id, iteration, stage)` 派生（表单 `round_seed` 可显式钉住），
+  并写进该轮的 `variant_round_selection` 事件（`seed` / `count` / `judgeable_in_pool`），
+  这样"为什么抽到这 4 题"事后能查、那一轮也能重跑。
+
+---
+
+## 2. 轮次与阶段
+
+一次"通过"叫一个 stage，记录在 `runs/iteration_NNN/input/stage.json`：
+
+| stage | 轮次编号 | 做什么 | 判哪门 | 消耗轮次 |
+|---|---|---|---|---|
+| **baseline** | iteration 1 | 抽 A1（4 题）交 DKAO 实测；这 4 题固化为考卷 | 都不判 | 0 |
+| **performance** | iteration 2 起 | 抽新题 A_i（排除考卷）交 DKAO 实测 | 只判性能门 | 1.0 |
+| **generalization** | 紧随其后的那一趟（"第 2.5 轮"） | 同一版 harness **再发一次 DKAO**，题目固定为 A1 | 只判泛化门 | 0.5 |
+
+> "紧随其后的那一趟"= 性能门通过后**立刻**用同一版 harness 再发一次评测（占下一趟 pass，
+> 有自己的 `runs/iteration_NNN` 目录），不是在同一趟里重测；批准时记录的是**赢下性能门的那一轮编号**。
+
+- **第 1 轮不判门**：它只负责"量出基线"。它测出来的更快 kernel **直接写进 variant 池**
+  （阈值 `promote_min_improvement_percent`，默认 3%），不需要人工批准——"kernel 是事实，harness 才要批"。
+  若第 1 轮有任一算子没拿到可用测量 → 整个 HE 任务立刻停（`baseline_round_failed`），不留半截证据。
+- **一轮的完整生命周期**（第 2 轮起）：
+  1. 抽新题 A_i（4 的倍数、含 decode+prefill、排除考卷）→ 一题一个 DKAO 子任务；
+  2. 判**性能门**（§3）；不过 → 回退到上一好 harness → analyze → evolve 产出新 harness → 下一轮重新随机抽题；
+  3. 过了性能门 → 用**同一版 harness**再发一次 DKAO，题目 = 固化考卷 A1 → 判**泛化门**（§4）；
+  4. 两门都过 → 写 `pending_promotion.json`，任务停在 `awaiting_approval`，**等人工批准**才接入生产 DKAO；
+     任何一门不过 → 同上走 analyze→evolve→下一轮。
+
+---
+
+## 3. 性能门（performance stage）
+
+对手是**每个算子当前的 variant**（该算子目前的最优解：成绩 + 内核）。
+
+- **赢** = 严格快于 variant（含 2% 噪声带；平局不算赢）；
+- 通过条件：**≥75% 的题赢**（4 题即至少 3 题），**且没赢的那些不得低于 variant 的 80%**。
+
+## 4. 泛化门（generalization stage）
+
+**仅在性能门通过后开**。不抽新题：**重考第 1 轮那份固化考卷 A1**，同一版 harness 再发一次 DKAO。
+
+- 通过条件：**≥50% 的题赢**，同样 **80% 地板**（任一题低于 variant 的 80% → 不过）；
+- **平局不算输**：子任务以 `best_known` 热身起跑，考卷上的算子正是第 1 轮刚写进池的 kernel，
+  重考大概率复现同一个 kernel；把"复现自己"判成失败等于要求搜索策略每轮都超越自己。
+  泛化门要的是"无败场 + 真退步就否决"。
+
+### 门的三种结论：PASS / FAIL / **INCOMPLETE**
+
+门是对 **harness** 的断言，所以只有在这轮真的产出了证据时才允许下结论：
+
+| 结论 | 何时 | 循环怎么走 |
+|---|---|---|
+| **PASS** | 全部题都测到，且满足比例 + 地板 | 性能门过 → 冻结重考；泛化门过 → 写 `pending_promotion`，停 `awaiting_approval` 等批准 |
+| **FAIL** | 已测到的部分**已经决定输**：有题低于 80% 地板，或"已赢 + 未测"仍达不到所需胜场数 | 回退上一好 harness → analyze → evolve → 下一轮重新抽题 |
+| **INCOMPLETE** | 有题没测到，而结果**仍然开放** | **本轮不下结论**：记 `round_incomplete`，不写候选、不改判 harness；下一趟重新测量（泛化门缺测时**保留**冻结的重考请求，考卷仍欠一个判决） |
+
+这条分界很重要：过去任何缺测都直接判 FAIL，等于把"机器没给数"记成"harness 退步"，
+既让循环去修一个没坏的 harness，也会在泛化重考时把候选白白花掉。
+
+## 5. 测量没拿到怎么办（唯一允许的重发）
+
+- 子任务**崩溃/工具链缺失/无 report** → HE 认为"环境欠一次测量"，**重发一次 DKAO**
+  （重发是全新的一次 DKAO、独立子任务目录；上限 `env_retry_attempts`，默认 3 次尝试）。
+- 子任务**被 DKAO 自己的闸门拒测**（exit 75 / `gate_blocked`）→ **不重发**：DKAO 已经在它那张卡上
+  等满自己的预算（30 分钟一次 × 48 = 24 小时）才给的否决，HE 没有可换的卡，重发等于无视它的结论。
+  该轮判 `round_incomplete`（cause `gpu_gate_blocked`）→ **终态**，等人工决定（不会被自动 resume）。
+- 某一轮**有题缺测但结果仍开放**（见 §3 的 INCOMPLETE）→ 该轮不下结论，事件 `round_incomplete`
+  （`gate` / `unmeasured` / `measured`），循环按 §3 的表继续。
+- 每次尝试都写独立目录（`children/iteration_NNN`、`..._a2`、`..._generalization`），
+  失败那次的日志原样留档；DKAO 会把子任务工作区的 `main` 绑定到本次 kernel repo，共用目录会让重试秒退。
+
+---
+
+## 6. 判据与产物
+
+| 文件 | 内容 |
+|---|---|
+| `variant_table.json` | 逐算子 variant + 固化考卷（`groups.generalization_ids`） |
+| `runs/iteration_NNN/input/stage.json` | 这次是哪个 stage + `round_cost` |
+| `.../baseline_round.json` | 第 1 轮：判定 + 入池结果 |
+| `.../performance_gate.json` / `.../generalization_gate.json` | 两门的逐题判定（wins/needed/floor） |
+| `.../decision.json` | 本轮 system decision |
+| `pending_generalization.json` | 冻结的重考请求 |
+| `pending_promotion.json` / `promoted_harness.json` | 待批候选 / 已接入版本 |
+| `children/<尝试目录>/<算子>/` | 每个算子一次 DKAO 运行（独立目录） |
+
+批准后的动作：`approve` 把候选 harness 接入 DKAO（备份上一版）、用**写泛化门那次重考**的成绩更新 variant、
+把该轮算子内核按"最好者胜"写进生产池；任务结束（`promoted:<版本>`）。**只有被接受那一轮的优化进入生产。**
+
+---
+
+## 7. 代码对应（改协议时一并改）
+
+| 协议条目 | 实现 |
+|---|---|
+| 题数 = 4 的倍数、一波 4 题 | `orchestrator/rounds.py`: `GPU_COUNT` / `OPERATORS_PER_GATE` / `normalized_per_gate` |
+| 取题（随机 / 可判 / decode+prefill / 排除考卷） | `orchestrator/rounds.py`: `sample_operators` / `regime_of` |
+| 轮次与阶段、考卷固化 | `orchestrator/rounds.py`: `select_round_questions` / `record_paper` / `round_stage` |
+| 两门判据（75% / 50% + 80% 地板 + 2% 噪声带） | `orchestrator/decision_engine.py`（`gate_verdict`）+ `rounds.py: evaluate_gate` |
+| 每轮把算子清单发给 DKAO、设备 = index % 4 | `orchestrator/adapters/eval.py`: `DkaoCliEvaluator.evaluate` / `_run_wave` |
+| 环境相关的重发与终态 | `adapters/eval.py`: `_env_failed` / `_stop_on_incomplete_round` |
+| 表单题数只给 4 的倍数 | `form.yaml`: `per_round_budget` |
+| 契约测试 | `tests/test_flow_contract.py` |
+
+**环境相关的判据不在本文件**：卡干净与否、等多久、污染怎么作废，全部写在 DKAO 侧
+（`dcu_kernel_auto_opt/orchestrator/gpu_preflight.py`、`w8a8_pipeline.py` 的测量闸门、
+`contention.py`；DKAO 侧的报错归档见 `dcu_kernel_auto_opt/error_catalog/`）。
+
+---
+
+## 8. 旧路径（suite / legacy）：保留，但必须自报家门
+
+新任务**一律**用本文件的固定协议（配一个实测池文件）。旧的 suite/legacy 路径仍然保留
+（离线套件、2026-09-15 之前的老实验），但它**不固化考卷、不判两门、不按 4 的倍数取题**——
+它的轮次不可能产出可接入 DKAO 的 harness。所以从 2026-09-22 起，它必须自己说出来：
+
+| 说给谁看 | 形式 |
+|---|---|
+| 盯着进程的人 | orchestrator 日志一条 `WARNING: running the LEGACY protocol (…). New tasks must use a measured pool file — see FLOW.md` |
+| 任务的审计记录 | timeline 事件 `legacy_protocol_selected`（含 `reason` / `pool_source` / `pool_file` / `policy`） |
+| 页面上的人 | 实验目录写 `protocol.json`，任务详情页据此在顶部显示黄色横幅「⚠ 本任务未按固定协议运行」 |
+
+**触发条件**（任一即判定为旧路径）：`round_questions: legacy|suite|off`；没配池文件（只用内联 suite）；
+池文件路径不存在；用了 `builtin` 假池。
+
+> 结论性的取舍（要不要彻底删掉这条路径）见 `README.md` 的「现状与边界」；在删除之前，
+> 这条告警是防止它被误用的唯一保障。
+
+---
+
+## 9. 停止与续跑（operator 操作说明）
+
+**停止 = 暂停，不是杀进程。** 这两件事在本系统里必须分得很清：一个进程"只是没了"和"崩了"在判据上
+完全一样，而 supervisor / janitor / `should_launch_resume` 的存在就是为了把崩掉的跑起来——所以
+"只杀进程的停止"会在几秒到 5 分钟内被自动复活。
+
+### 9.1 点按钮（推荐）
+
+任务详情页（HE 页面）状态机下方那一行：**「停止任务」**（琥珀色边框，仅在 `running` 时出现）。
+
+一次点击按固定顺序做完六步（代码：`server/routes.py: stop`）：
+
+| # | 动作 | 为什么必须在这个位置 |
+|---|---|---|
+| 1 | 写 `workspace/operator_stop.json` | **标记先落地**：janitor 一旦先看到"没人在跑"，就会自动拉起它 |
+| 2 | 停 supervisor（写 `supervisor.stop` + SIGTERM） | 它存在的意义就是"把死掉的实验拉起来"，必须先停它 |
+| 3 | 杀掉本轮所有 DKAO 子任务（按 cmdline 匹配 workspace 路径） | 子任务各自 `start_new_session`，父进程死了它们还活着 |
+| 4 | 杀掉 HE 编排器（`_run_alive`：`orchestrator.pid` → `resume.pid` → cmdline 扫描） | 最后杀，避免它在子任务写到一半时被打断 |
+| 5 | **再扫一次子任务** | 编排器把"子任务消失"当环境失败，会在第 3~4 步之间发一波新的（`env_retry_attempts`）；漏掉它，页面写着"已停止"而新子任务还在卡上跑几小时 |
+| 6 | `run.json` 写 `finished: true` / `final_status: stopped_by_operator` | 让所有自动路径收手；`finished` 同时是 `/resume` 的前置条件 |
+
+**关键点**：运行中的编排器**自己不读** `operator_stop.json`——只有 janitor（`_operator_stopped`）和
+supervisor 读它，CLI 只在"刻意重启"时删它。所以标记只防"自动复活"，**不会**让活着的进程停下来，
+第 4 步不能省。
+
+### 9.2 停止后一定不会被自动拉起
+
+janitor 每 300 s 扫一遍，命中标记就跳过；`stopped_by_operator` 也在 `TERMINAL_STATUSES` 里，
+`should_launch_resume` 直接返回 False。想确认，看 WebUI 日志（`/root/zth_agent/webui-8765.log`）：
+
+```
+[he-janitor] {'task': '', 'action': 'skip', 'reason': 'stopped by operator'}
+```
+
+### 9.3 续跑
+
+页面上「续跑到该轮次上限」，或 `POST /api/harness-evolve//resume {"target": N}`。
+`_launch_resume` 会**先删掉 operator-stop 标记**再起进程，所以续跑不会被自己的标记挡住。
+
+被停掉的那一轮**重跑**：抽题种子由 `(task_id, iteration, stage)` 派生 → 抽到的还是**同样那批题**；
+那一轮没有判完，所以它的成绩**不会**写回 variant（写回只发生在"轮次判完"或"harness 被人工批准"时）。
+已经跑完的子任务证据留在磁盘上，但同样不计入。
+
+### 9.4 不要做什么
+
+| 做法 | 后果 |
+|---|---|
+| `kill -9 <编排器 pid>`（任何"只杀进程"） | 看起来像崩溃 → janitor 在 ≤5 分钟内 `supervisor+resume` 自动拉起；子任务在各自 session 里**不会**被杀，变成孤儿继续占卡 |
+| shell 任务页的 **Kill** 按钮 | 它会按路径杀整棵树（比裸 kill 干净），但写的是 `final_status: killed`，**不在终态集合里**；没有 operator-stop 标记时，janitor 仍会按"未达目标轮数"把它拉起来 |
+| shell 任务页的 **Reset** | 那是真删除：日志 / 迭代产物 / 运行记录全清，只留 `requirements.json`，不可恢复 |
+| 手删 `operator_stop.json` | 等于解除暂停保护：下一次 janitor tick 就会把它当崩溃恢复 |
+
+### 9.5 代码对应
+
+| 事项 | 位置 |
+|---|---|
+| 停止的六步顺序 | `server/routes.py: stop` |
+| 宽口径活性判断（orchestrator.pid → resume.pid → cmdline） | `server/routes.py: _run_alive`（窄口径 `_resume_running` 只认 `resume.pid`，New Task 起的任务没有它） |
+| 子任务定位（各自独立 session，靠 cmdline 认领） | `server/routes.py: _live_dkao_children` |
+| 信号发送（先 `killpg` 再 `kill`，僵尸算已停） | `server/routes.py: _terminate_pid` |
+| 终态集合 / 自动复活判据 | `server/routes.py: TERMINAL_STATUSES` / `should_launch_resume` |
+| 续跑并清除标记 | `server/routes.py: _launch_resume`、`orchestrator/cli.py: _clear_operator_stop` |
+| 契约测试 | `tests/test_operator_stop.py` |
diff --git a/metainfer/tasks/harness_evolve/README.md b/metainfer/tasks/harness_evolve/README.md
new file mode 100644
index 00000000..85b645e0
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/README.md
@@ -0,0 +1,197 @@
+# harness_evolve（v0）— AHE 外层循环（headless）
+
+把 dcu_kernel_auto_opt 任务实例当评测单元跑 AHE evaluate→analyze→evolve 的
+"第一版可用"实现。设计依据：
+
+- **协议（固定，先读这个）：[`FLOW.md`](FLOW.md)** —— 每轮题数、取题规则、两门判据、
+  HE 与 DKAO 的边界；代码必须与它一致（`tests/test_flow_contract.py` 守着）
+- `MetaInfer/docs/dkao_harness_eval_protocol.md`（评测协议 v2）
+- `MetaInfer/docs/ahe_dkao_integration_plan.md`（接入架构）
+- 官方参考：`/root/zth_agent/ahe-ref`（MIT）
+
+## 变体轮次协议（固定，规范文本见 [`FLOW.md`](FLOW.md)）
+
+**协议以 [`FLOW.md`](FLOW.md) 为准**：每轮题数固定为 4 的倍数（4 张 HCU 一波铺满）、
+第 1 轮随机抽 A1 并**固化为泛化考卷**、性能门 75% + 80% 地板、泛化门（仅性能门通过后开）
+50% + 80% 地板、过关后等人工批准；**HE 只发起 DKAO 任务，不管理环境**。
+本节只留实现要点；与 `FLOW.md` 不一致时以 `FLOW.md` + 契约测试为准。
+
+**停止 / 续跑怎么操作**（按钮做了什么、为什么"裸杀进程"会被自动复活、哪些做法不可逆）：
+见 [`FLOW.md`](FLOW.md) §9「停止与续跑（operator 操作说明）」。一句话：**页面的「停止任务」= 暂停**，
+产物保留、不会被自动拉起，之后用「续跑到该轮次上限」接着跑；**直接 kill 进程或 shell 的 Kill 按钮
+都不等于暂停**。
+
+语义（先统一）：
+
+- **baseline**：每个算子固定的 Triton 基线（`baseline_us`），永不变。
+- **variant**：每个算子**当前的最优解**（成绩 + 内核），是候选必须打败的对手；
+  **只在人工批准接入某一版 harness 时更新**（`variant_table.json`）。
+- **harness**：搜索策略（planner/gates）。它只影响 DKAO 怎么找内核，不参与测量。
+
+
+**阶段（stage）**：一次通过循环叫一个 stage，`runs/iteration_NNN/input/stage.json` 记录它是哪一个：
+
+| stage | 轮次编号 | 做什么 | 判哪门 | 预算 |
+|---|---|---|---|---|
+| **baseline** | iteration 1 | 随机抽 4 题（= `A1`），交给 DKAO 实测；这 4 题**固化为泛化考卷** | **都不判** | 0 |
+| **performance** | iteration 2 起 | 重新随机抽 4 题（`A2`，排除考卷），交给 DKAO 实测 | **只判性能门** | 1.0 |
+| **generalization** | 与上手**同一编号** | 过关后**再用同一版 harness 发动一次 DKAO**，题目固定为 `A1` | **只判泛化门** | 0.5 |
+
+- **第 1 轮是基线轮**：它不判任何门、不产生 `pending_promotion.json`，但**它测出来的更快 kernel
+  直接写入 variant 池**（`promote_kernels_for_round`，阈值与生产一致：`promote_min_improvement_percent`，
+  默认 3%；写池前自动备份池文件并追加 history）。这是"kernel 结果是事实，harness 才需要人批准"的分离。
+  第 1 轮若**没能对每个算子都拿到可用测量**（子任务崩溃 / DKAO 自己的 GPU 门拒测 / 无 report）
+  → 整个 HE 任务**立即停止**（`final_status=baseline_round_failed`），不留半截证据。
+- **性能门**：随机抽的那 4 题**必须真正快于**各自 variant（严格优于，含 2% 噪声带）——≥75% WIN
+  （4 题即至少 3 题），且没赢的不得低于 variant 的 80%。平局**不算赢**。
+- **泛化门**：重考考卷那 4 题，≥50% WIN，同样 80% 地板。**平局不算输**：DKAO 子任务是
+  `warm_start: best_known` 起跑，考卷上的算子正是基线轮刚写进池的 kernel，重考大概率复现同一个
+  kernel（实测 delta = 0.0%）；把它判成失败等于要求"搜索策略每轮都比自己上一轮更快"。
+  所以泛化门要求的是**无败场 + 真的退步就否决**（任一算子低于 variant 的 80% → 不过）。
+- **过关后**：写 `pending_promotion.json`，任务停在 `awaiting_approval`，**等你人工批准**才把 harness
+  接进生产 DKAO。任何一门不过 → 回退到"当前已发布/上一好 harness"、analyze→evolve、进入下一轮
+  重新抽题（上限 `max_rounds`，默认 10；泛化轮算 0.5 轮，所以一个"过关再重考"的完整循环消耗 1.5 轮）。
+- **harness 与 DKAO 的边界**：HE 不执行任何评测本身，它只对每个算子**起一个 DKAO 子任务**
+  （`dcu_kernel_auto_opt.orchestrator.cli run`），把当前 harness 快照（`METAINFER_HARNESS_ROOT`）和
+  本轮算子清单交过去；测量、GPU 门、contamination 全在 DKAO 内。泛化轮就是**再发一次**这样的任务，
+  只是题目固定为考卷 —— 代码上只是再调一次 `evaluator.evaluate()`，DKAO 一行不用改。
+  每一个"评估尝试"都有独立的子任务目录（`children/iteration_NNN`、重试 `..._a2`、重考 `..._generalization`），
+  保证重考是"重新发动一次 DKAO"、重试也是新的一次 DKAO，而不是接着上一批继续跑。
+  这一点是必须的：DKAO 会把子任务工作区的 `main` 符号链接绑死到它本次要用的 kernel repo，
+  复用同一个目录会让重试**秒退**（`workspace main already points to ...`），
+  于是环境一抖就烧完整个重试预算、看起来像永久失败。重试目录独立后，
+  失败那次尝试的日志与状态也原样留在旁边作为证据。
+
+**环境（卡）不是 HE 的事**：卡被占用、闸门怎么等、污染测量怎么作废，全部由 DKAO 的子进程
+自己处理（`ensure_measurement_gate` / contamination），判据文本见 `FLOW.md` §0 与 §5。
+HE 过去在父层做过的三件事都已移除：复查卡状态、`SIGSTOP` 冻结子任务、
+按空闲卡数缩小题数（`fit_questions_to_devices`）——冻结并不会把卡让给别人，
+缩题数只会让一轮考不出结论。
+
+**机制门**（宽松）：每轮对上一轮 `change_manifest` 声明的机制做 grading，结果写进
+`decision.json.mechanism_gate` 与时间线（`mechanism_gate` 事件）；它**不否决**已通过两门的候选，
+但"晋升了却说不清为什么"会留痕。第 1 轮（baseline）没有门可 grade，其 `mechanism_gate` 为空。
+
+**批准后的动作**：
+- `approve`：把候选 harness 发布进 DKAO（备份上一版 + `promoted_harness.json`）、
+  用**被接受那一轮**（写泛化门的那次重考）的成绩更新 variant、把该轮算子的内核写入生产池
+  （逐算子"最好者胜"：低于池内 `best_known_us` 才写入，先备份池文件），任务结束
+  （`final_status=promoted:<版本>`）。**只有被接受那一轮的优化进入生产，其余轮次丢弃。**
+- `deny`：丢弃候选，恢复迭代。
+- `rollback-harness --to last-good|h-N|seed`：把生产版本回退。
+
+命令行：
+
+```bash
+python -m metainfer.tasks.harness_evolve.orchestrator.cli pending  --state-dir  --workspace-dir 
+python -m metainfer.tasks.harness_evolve.orchestrator.cli approve  --state-dir  --workspace-dir  [--by NAME]
+python -m metainfer.tasks.harness_evolve.orchestrator.cli deny     --state-dir  --workspace-dir  [--reason TEXT]
+python -m metainfer.tasks.harness_evolve.orchestrator.cli rollback-harness --state-dir  --workspace-dir  --to last-good
+```
+
+WebUI：`GET /promotion` 看待批准候选与生产版本，`POST /promotion/approve|deny` 做决定。
+
+产物：`variant_table.json`（逐算子 variant + 泛化考卷）、`variant_history.jsonl`（谁把哪个
+算子刷成了多少）、`runs/iteration_NNN/input/stage.json`（这次是哪个 stage）、
+`.../baseline_round.json`（第 1 轮：判定 + 入池结果）、`.../performance_gate.json`、
+`.../generalization_gate.json`、`.../decision.json`、`pending_generalization.json`（冻结的重考请求）、
+`pending_promotion.json`、`promoted_harness.json`、
+`promotion_log.jsonl`；内核入库会更新池文件的 `best_known_us`/`history`，池文件备份在
+`/promoted/.<时间戳>.yaml`。
+DKAO 子任务目录：`children/<尝试目录>/<算子>`，其中尝试目录为 `iteration_NNN`（首次）、
+`iteration_NNN_a2`（环境失败后的第 2 次尝试，最多 `env_retry_attempts` 次）、
+`iteration_NNN_generalization`（泛化重考）；任何两个尝试都不共用目录。
+
+回滚修复（本轮）：`workspace_revision` 现在记录**本轮自己的** commit（此前读的是上一轮
+manifest 的字段，恒为空）；`champion_override.json` 缺失时不再抛异常。
+
+## HE 不管理 GPU（2026-09-22 起，operator 规则）
+
+**HE 只发起 DKAO 任务，卡的事归 DKAO。** 判据文本见 [`FLOW.md`](FLOW.md) §0。
+
+- HE 不租卡、不读卡状态、不因为"卡忙"等待或否决、不冻结子进程。
+- 卡干不干净由 **DKAO 子进程**在测量闸门里判（VRAM ≤ 90% 且 HCU == 0；30 分钟一次、
+  最长 24 小时），写在子任务的 `state/measurement_gate.jsonl`——那是唯一存在闸门判定的地方。
+- HE 侧只剩每轮的**设备布局**记录：`runs/iteration_NNN/input/benchmark/gpu_preflight.json`
+  （哪道题绑了哪张卡、`index % 4`）。它是布局快照，不是占用检查。
+- 子任务被 DKAO 闸门拒测（exit 75）→ HE **不重发**，本轮判 `round_incomplete`
+  （cause `gpu_gate_blocked`，终态，等人工决定）；子任务崩溃 / 工具链缺失 → HE 重发一次
+  （`env_retry_attempts`）。
+- 逃生阀：`METAINFER_GPU_PREFLIGHT=0` 关的是 **DKAO 侧**的闸门。
+
+## 运行（dry-run，离线可跑通闭环）
+
+```bash
+# 1) 写 requirements.json（answers 含 suite_yaml / max_iterations / execution_mode）
+#    answers.execution_mode: dry-run（默认，确定性 fixture）| dkao-cli（需 worker29）
+# 2) headless 运行（框架契约）
+python3 -m metainfer.tasks.harness_evolve.orchestrator.cli run \
+  requirements.json --state-dir  --workspace-dir 
+```
+
+## 目录布局（每轮双代际，同 ahe-ref 语义）
+
+```
+/                       # = AHE experiment root
+├── workspace/                         # 可演化 harness（从 dcu harness_default 播种）
+├── runs/iteration_NNN/
+│   ├── input/workspace/               # 本轮评测快照
+│   ├── input/benchmark/results.json
+│   ├── input/diff.json                # vs 上轮（flipped/regressed）
+│   ├── input/change_evaluation.json   # 上轮 change_manifest 的归因（verdicts）
+│   ├── input/analysis/overview.md
+│   └── evolve/change_manifest.json    # 本轮 Evolve 输出
+├── best_ever.json                     # 超过则更新；低于则自动回滚 workspace
+├── iteration_scores.jsonl
+└── report.md
+```
+
+## Pool v2（考题协议 v2，2026-09-09）
+
+- 新增 `orchestrator/pool.py`（注册池 + τ_family 自动判过 + 分层抽样/heldout）与
+  `orchestrator/round_plan.py`（单角色出卷 + 三条护栏：池内 / 与上轮重叠 ≥60% / budget）。
+- 表单新增 `pool_source`（`builtin` | 池 yaml 路径）与 `per_round_budget`；置空 = 旧 suite 模式（行为不变）。
+- pool 模式下每轮 evaluate 只跑 round_plan 选中的题，`evolve/round_plan.json` 输出下轮考卷，
+  决策（best_ever/回滚/归因）用重叠对口径；详见 `MetaInfer/docs/dkao_suite_protocol_v2.md`。
+
+## 真实接入 v1（2026-09-09）
+
+- `decision_engine.py`：BASELINE / PROMOTE / PROMOTE_UNEXPLAINED / CONFIRM_REQUIRED / NO_SIGNAL / SPECIALIZE / REJECT / REJECT_OVERFIT；性能门优先、机制门解释、held-out 可专用化。
+- `DkaoCliEvaluator`：每题一个真实 DKAO 子任务，并发<=4，强制 `METAINFER_KERNEL_REPOS=/root/zth_agent/ahe-kernel-repos`；解析 final_report 与 child/repo/性能元数据。
+- `AgentEvolver`：单角色 DSH Agent 同时修改 harness、写完整 change_manifest、出下一轮 round_plan；确认轮保持候选字节不变。
+- 前端：每轮显示 system decision、child task/repo、median/p90、change attribution。
+
+- 实测注册池构建：`python3 -m metainfer.tasks.harness_evolve.orchestrator.pool_builder`
+  只读扫描历史 final_report → `/root/zth_agent/ahe-kernel-repos/registered_pool.yaml`（34 实例）。
+- dkao-cli 模式拒绝 builtin 假池；必须使用实测池文件。
+- planner 策略已数据化：`dcu_kernel_auto_opt/harness_default/planner_policy.yaml`（wired:true），
+  AgentEvolver 改它即可改变下一轮状态→方案选择。
+
+## 现状与边界（2026-09-22）
+
+运行方式与判据：**协议见 [`FLOW.md`](FLOW.md)**（固定文本，契约测试 `tests/test_flow_contract.py` 守着）。
+
+- ✅ **变体轮次协议是默认路径**：每轮 4 的倍数题、第 1 轮固化泛化考卷、性能门 75% + 泛化门 50%
+  （各带 80% 地板）、过关后 `awaiting_approval` 等人工批准。实测状态：5 个 HE 任务里
+  1 个（`test-9-17-1-07bb5cd8`）两门通过并写出 `pending_promotion.json`，**至今未被批准**；
+  全库尚无 `promoted_harness.json` —— 循环还没真正闭合过一次。
+- ✅ `dkao-cli` evaluator 是生产路径（不再是骨架）：每题一个真实 DKAO 子任务、并发 ≤4、
+  `METAINFER_KERNEL_REPOS=/root/zth_agent/ahe-kernel-repos`、harness 快照经
+  `METAINFER_HARNESS_ROOT` 传给子任务。
+- ✅ `AgentEvolver` 已实现（`evolve_mode: agent`，DSH Agent 改 harness 并写 change_manifest /
+  round_plan）；`evolve_mode: dry-run` 仍可用于离线闭环。
+- ✅ **HE 不管理 GPU**：不租卡、不读卡、不冻结子进程、不按空闲卡缩题（细节见 `FLOW.md` §0 与
+  下面「HE 不管理 GPU」节）。
+- ✅ 门有三种结论 **PASS / FAIL / INCOMPLETE**：缺测且结果仍开放时不给 harness 判词（`FLOW.md` §3）。
+- ⚠️ **测量噪声未标定**：门用 2% 噪声带、子任务用 `bench_profile=quick`，但"同一版 harness
+  反复重测同一批题"的噪声基线还没有实测数据；判据基准（variant/pool 的 `best_known_us`）
+  直接来自 DKAO 的测量，因此噪声会直接影响门的结论。
+- ⚠️ **待批候选没有提醒机制**：`awaiting_approval` 是终态，不批就一直停着（上面那个候选已躺 5 天）。
+- ⚠️ **旧路径（suite / legacy）保留但会自报家门**：它不判两门、不固化考卷，因此不可能产出可接入的
+  harness；现在它会在日志里打 WARNING、在 timeline 写 `legacy_protocol_selected`、并在任务页面顶部
+  显示黄色横幅（细节见 `FLOW.md` §8）。彻底删掉这条路径仍是**未决**：删掉会一起失去 suite 模式
+  （离线 dry-run 闭环）与"Evolve Agent 出下一轮考卷"（`round_plan.json`，默认协议下这份 plan 目前
+  只写不用）；在删除之前，这条告警就是防止它被误用的保障。
+- ⚠️ `harness_source` 默认指向 dcu_kernel_auto_opt 的 `harness_default/` 种子；
+  指向 AHE 演化出的 workspace（git 仓库）即可继续"harness 自动进化"。
+- ⚠️ 历史文档 `docs/ahe_dkao_integration_plan.md` 等描述的是**父层管卡**的旧设计，已作废。
diff --git a/metainfer/tasks/harness_evolve/__init__.py b/metainfer/tasks/harness_evolve/__init__.py
new file mode 100644
index 00000000..e711c2a3
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/__init__.py
@@ -0,0 +1,13 @@
+"""harness_evolve — AHE outer-loop task type for MetaInfer.
+
+Drives dcu_kernel_auto_opt task instances as evaluation units, per
+``MetaInfer/docs/dkao_harness_eval_protocol.md`` (evaluation protocol v2) and
+``MetaInfer/docs/ahe_dkao_integration_plan.md``.
+
+This is the "first usable version": a headless orchestrator running the AHE
+iteration layout with pluggable evaluators/evolvers (dry-run default; DKAO-CLI
+and AgentEvolver reserved for worker29 validation), plus read-only WebUI routes.
+"""
+
+from .orchestrator import plugin as _task_plugin  # noqa: F401
+from .server import plugin as _web_plugin  # noqa: F401
diff --git a/metainfer/tasks/harness_evolve/form.yaml b/metainfer/tasks/harness_evolve/form.yaml
new file mode 100644
index 00000000..726b92b7
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/form.yaml
@@ -0,0 +1,209 @@
+# harness_evolve（AHE outer loop）——真实测评任务表单。
+#
+# 考卷不由人工填写：每轮考题由单角色 Evolve Agent 的 round_plan 从注册池选择，
+# 控制面强制池内、重叠>=60%、预算封顶。评测只用真实 DKAO（dkao-cli），
+# 无 dry-run 虚拟测试选项。
+
+- key: pool_source
+  question: "考题注册池 YAML（已接入、有实测基线的 GEMM 算子；dkao-cli 禁止使用内置假池）"
+  header: "Question pool"
+  required: false
+  form: file
+  default: "/root/zth_agent/ahe-kernel-repos/registered_pool.yaml"
+
+# 协议见 FLOW.md §1.1：题数固定为 4 的倍数（4 张 HCU 一波铺满），所以这里只给
+# 4 的倍数选项；CLI/接口传非倍数时由 normalized_per_gate 向上取整。
+- key: per_round_budget
+  question: "每轮考题数（固定为 4 的倍数：4 张 HCU 一波铺满）"
+  header: "Round budget"
+  required: false
+  form: select
+  default: "4"
+  options:
+    - label: "4"
+      description: "默认：一波 4 题，正好铺满 4 张卡"
+    - label: "8"
+      description: "两波共 8 题"
+    - label: "12"
+      description: "三波共 12 题"
+    - label: "16"
+      description: "四波共 16 题"
+    - label: "20"
+      description: "五波共 20 题（池子不够时按池子大小取）"
+
+# NOTE on the two numbers below — they are different layers and used to read
+# alike, which is exactly how a run once ended up doing rounds nobody asked for:
+#   default_rounds = one question, inside DKAO: how many optimization rounds its
+#                    worker may try (this is DKAO's own "Max iters").
+#   rounds         = the whole HE task: how many iterations it may run, one
+#                    iteration = one round.
+# The select field's options are the on-screen text (the renderer shows neither
+# the question text nor per-option descriptions for a select), so the option
+# labels stay plain numbers and the meaning lives in the question text.
+- key: default_rounds
+  question: "【单题内部】每道考题交给 DKAO 后，它的 worker 最多做几轮内核优化——这是单题预算，不是任务轮次（下面 Iterations 才是）。含义：1=只做一次基线验证；3=快速试 3 轮；5=默认，与生产 DKAO 的 Max iters 一致；10/20=更充分但更慢。该值原样传给 DKAO 的 Max iters。"
+  header: "Max iters"
+  required: false
+  form: select
+  default: "5"
+  options:
+    - label: "1"
+      description: "只做一次基线验证（单题不迭代）"
+    - label: "3"
+      description: "快速试 3 轮"
+    - label: "5"
+      description: "默认：单题最多 5 轮，与生产 DKAO 一致"
+    - label: "10"
+      description: "单题最多 10 轮，更充分但更慢"
+    - label: "20"
+      description: "单题最多 20 轮，最慢"
+
+- key: rounds
+  question: "【任务上限，唯一的轮次字段】最多跑几轮，1 轮 = 1 次迭代。第 1 轮是基线轮（只测量、不判门）；之后每一轮 = 重抽一组新题判性能门，过关则用同一版 harness 重考考卷判泛化门（重考不额外占一轮）。任何一轮两门全过就立刻停下、等你批准接入生产；比如 Rounds=10 而第 5 轮就过了，任务就在第 5 轮结束。轮次用尽仍未过关则以 rounds_exhausted 结束。"
+  header: "Rounds"
+  required: false
+  form: number
+  default: "10"
+
+- key: harness_source
+  question: "Harness 起点：dcu_default = 内置种子；或指向已演化 workspace 的绝对路径"
+  header: "Harness src"
+  required: false
+  form: text
+  default: "dcu_default"
+
+- key: execution_mode
+  question: "评测方式（真实 DKAO 专用）"
+  header: "Eval mode"
+  required: false
+  form: select
+  default: "dkao-cli"
+  options:
+    - label: "dkao-cli"
+      description: "真实 DKAO GPU 评测（worker29 容器 + dsh agent）"
+
+- key: ahe_repo_root
+  question: "AHE 专用 kernel 仓库根（与人工 /root/zth_agent/kernel-repos 隔离）"
+  header: "AHE repos"
+  required: false
+  form: file
+  default: "/root/zth_agent/ahe-kernel-repos"
+
+- key: child_timeout_minutes
+  question: "单个 DKAO 考题子任务超时（分钟；0=不限）"
+  header: "Child timeout"
+  required: false
+  form: number
+  default: "0"
+
+- key: planner_enabled
+  question: "DKAO 子任务是否启用 planner（METAINFER_PLANNER=1）：只有开启后，被测 harness 的 planner_policy 才真正生效"
+  header: "Planner"
+  required: false
+  form: select
+  default: "true"
+  options:
+    - label: "true"
+      description: "注入 METAINFER_PLANNER=1：轮次指令由 harness 的 planner_policy/catalog 渲染（HE 演化 planner 的前提）"
+    - label: "false"
+      description: "使用 DKAO 默认菜单，planner_policy 改动不生效（仅用于对照/回退）"
+
+- key: validate_scope
+  question: "子任务最终验证范围：task = 只验证本题 shape（快）；api = 全量 API shape 回归（慢，生产默认）"
+  header: "Validate scope"
+  required: false
+  form: select
+  default: "task"
+  options:
+    - label: "task"
+      description: "只做本题 shape 的串行验证（HE 每题只有 1 个 shape，省去 42 个回归 shape 的固定开销）"
+    - label: "api"
+      description: "全量 API shape 验证（含未优化 shape 的回归）"
+
+- key: bench_profile
+  question: "benchmark 采样密度：quick = warmups30/samples30/replays50（稳定档）；full = 资产默认(100/30/100)"
+  header: "Bench profile"
+  required: false
+  form: select
+  default: "quick"
+  options:
+    - label: "quick"
+      description: "较快且稳定（warmups30/samples30/replays50），兼顾可比性"
+    - label: "full"
+      description: "资产默认的高密度采样"
+
+- key: mechanism_policy
+  question: "机制门严格度：lenient = 性能显著提升即可晋级（未验证记 PROMOTE_UNEXPLAINED，仅「矛盾证据」才要求复测）；strict = 无机制证据即要求复测；off = 不做机制判定"
+  header: "Mechanism gate"
+  required: false
+  form: select
+  default: "lenient"
+  options:
+    - label: "lenient"
+      description: "宽松：性能优先，机制命中/部分命中/未观测都可晋级；只有明确矛盾（如 harness 未生效、planner 被关闭）才复测一轮"
+    - label: "strict"
+      description: "保守：机制未验证就要求复测一轮"
+    - label: "off"
+      description: "关闭机制判定（只看性能）"
+
+- key: promote_kernels
+  question: "本轮评测出更快且正确的算子时，自动替换 DKAO 共享 variant（写前备份、可回滚）"
+  header: "Kernel promotion"
+  required: false
+  form: select
+  default: "true"
+  options:
+    - label: "true"
+      description: "自动把比现有 variant 快 ≥ 阈值的成品写回 DKAO variant 树"
+    - label: "false"
+      description: "只判定不写入（promotion.json 仍记录每题结果）"
+
+- key: promote_min_improvement_percent
+  question: "自动替换要求的最小提升幅度（相对现有 variant 的 median，%）"
+  header: "Promote threshold"
+  required: false
+  form: number
+  default: "3"
+
+- key: evolve_mode
+  question: "Evolve 模式：agent = 真实单角色 Evolve Agent 改 harness 并出下一轮考卷；dry-run 仅供调试占位"
+  header: "Evolve mode"
+  required: false
+  form: select
+  default: "agent"
+  options:
+    - label: "agent"
+      description: "真实 DSH Evolve Agent（改 harness + round_plan）"
+    - label: "dry-run"
+      description: "调试占位：只写 manifest、不改 harness"
+
+- key: evolve_model
+  question: "真实 Evolve Agent 模型 id（dsh 情况下可选；deepseek-flash-4.1 = 当前默认，deepseek-v4-flash = 旧 pinned 构建）"
+  header: "Evolve model"
+  required: false
+  form: select
+  default: "deepseek/deepseek-flash"
+  options:
+    - label: "deepseek/deepseek-flash"
+      description: "DeepSeek Flash 4.1（表单标签 deepseek-flash-4.1）——当前默认"
+    - label: "deepseek/deepseek-v4-flash-0731"
+      description: "DeepSeek V4 Flash 旧 pinned 构建（复现历史实验时选它）"
+
+- key: evolve_timeout_seconds
+  question: "真实 Evolve Agent 单轮超时（秒）"
+  header: "Evolve timeout"
+  required: false
+  form: number
+  default: "3600"
+
+- key: agent_framework
+  question: "DKAO 子任务的 Agent framework"
+  header: "Agent fw"
+  required: false
+  form: select
+  default: "dsh"
+  options:
+    - label: "dsh"
+      description: "DeepSeek Harness"
+    - label: "ccb"
+      description: "Claude Code"
diff --git a/metainfer/tasks/harness_evolve/orchestrator/__init__.py b/metainfer/tasks/harness_evolve/orchestrator/__init__.py
new file mode 100644
index 00000000..e219a487
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/__init__.py
@@ -0,0 +1,7 @@
+"""harness_evolve orchestrator package."""
+
+from metainfer.orchestrator.tasks import register
+
+from .plugin import PLUGIN
+
+register(PLUGIN)
diff --git a/metainfer/tasks/harness_evolve/orchestrator/adapters/__init__.py b/metainfer/tasks/harness_evolve/orchestrator/adapters/__init__.py
new file mode 100644
index 00000000..8eb218d7
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/adapters/__init__.py
@@ -0,0 +1 @@
+"""Evaluation and evolve adapters for harness_evolve (v0)."""
diff --git a/metainfer/tasks/harness_evolve/orchestrator/adapters/eval.py b/metainfer/tasks/harness_evolve/orchestrator/adapters/eval.py
new file mode 100644
index 00000000..9270459f
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/adapters/eval.py
@@ -0,0 +1,834 @@
+"""Evaluation adapters for harness_evolve.
+
+Real evaluator contract:
+- every question becomes one real dcu_kernel_auto_opt child task;
+- artificial/production kernel repos are READ-ONLY seeds;
+- all child output repos are forced under the isolated AHE repo root
+  (default /root/zth_agent/ahe-kernel-repos);
+- per-child final_report / worker rounds / token budget are normalized for the
+  parent AHE benchmark and frontend.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import json
+import os
+import re
+import subprocess
+import sys
+import time
+import uuid
+
+from pathlib import Path
+from typing import Any, Dict, List
+
+import yaml
+
+from ..config import ExperimentConfig, InstanceSpec
+from ..process_control import ChildProcess, ProcessRegistry, spawn
+from ..rounds import (
+    CHILDREN_DIR, GPU_COUNT, STAGE_GENERALIZATION, STAGE_PERFORMANCE,
+    device_for_index, stage_dir_name,
+)
+
+
+def _round_stage(cfg: Any) -> str:
+    """Which stage this evaluation is: the pipeline sets ``answers[round_stage]``.
+
+    Missing means "an old/legacy run": the performance stage is the shape the
+    adapter had before the stages existed, so that is the fallback.
+    """
+    answers = getattr(cfg, "answers", None) or {}
+    try:
+        return str(answers.get("round_stage") or STAGE_PERFORMANCE)
+    except AttributeError:
+        return STAGE_PERFORMANCE
+
+#: The exit code DKAO's in-harness measurement gate uses when it has waited its
+#: whole budget (``METAINFER_GATE_MAX_WAITS``, 48 x 30 min = 24 h) without ever
+#: seeing a clean device (VRAM <= 90% and HCU == 0). Mirrors
+#: ``metainfer/tasks/dcu_kernel_auto_opt/assets/w8a8_bench.py::GATE_EXIT_CODE``;
+#: duplicated on purpose so the HE adapter does not import the child's asset
+#: module just for a number, and asserted against it in the tests.
+GATE_EXIT_CODE = 75
+
+
+class Evaluator:
+    name = "base"
+
+    def evaluate(
+        self,
+        cfg: ExperimentConfig,
+        workspace: Path,
+        iteration: int,
+        instances: "List[InstanceSpec] | None" = None,
+    ) -> Dict[str, Dict[str, Any]]:
+        raise NotImplementedError
+
+
+class DryRunEvaluator(Evaluator):
+    """Deterministic offline evaluator (tests / CI / protocol dry runs)."""
+
+    name = "dry-run"
+
+    def evaluate(self, cfg, workspace, iteration, instances=None):
+        targets = list(instances) if instances else list(cfg.suite)
+        fixture: Dict[str, Any] = dict(cfg.answers.get("dry_fixture") or {})
+        # ``dry_fixture_all`` sets the default for questions the fixture does
+        # not name, so an offline run can be steered wholesale (used by the
+        # variant-round end-to-end test).
+        default_pass = cfg.answers.get("dry_fixture_all")
+        results: Dict[str, Dict[str, Any]] = {}
+        for idx, inst in enumerate(targets):
+            override = fixture.get(inst.id)
+            if override is None:
+                override = default_pass
+            passed = bool(override) if override is not None else (idx != 0)
+            threshold = float(inst.pass_median_us_le or 1.0)
+            scale = 0.8 if passed else 1.4
+            median = round(threshold * scale, 3)
+            results[inst.id] = {
+                "status": "completed",
+                "passed": passed,
+                "correctness_ok": True,
+                "p90_ok": True,
+                "median_us": median,
+                "p90_us": median * 1.01,
+                "rounds_used": inst.budget_rounds,
+                "tokens": 0,
+                "reason": "dry-run fixture",
+                "child_task_id": f"dry-{iteration}-{inst.id}",
+                "repo_path": None,
+                "harness_revision": _harness_revision(workspace),
+            }
+        return results
+
+
+class DkaoCliEvaluator(Evaluator):
+    """Real DKAO evaluator (one isolated child task per question)."""
+
+    name = "dkao-cli"
+
+    def __init__(self) -> None:
+        # live child handles, so a question can be terminated while it runs
+        # (see process_control.py). Freezing a child because "its" device was
+        # taken over is gone with the rest of HE's GPU management: the child's
+        # own gate never measures on a busy card in the first place.
+        self.children = ProcessRegistry()
+
+    def evaluate(self, cfg, workspace, iteration, instances=None):
+        targets = list(instances) if instances else list(cfg.suite)
+        if not targets:
+            return {}
+        stage = _round_stage(cfg)
+        repo_root = Path(
+            str(cfg.answers.get("ahe_repo_root") or
+                "/root/zth_agent/ahe-kernel-repos")
+        ).expanduser().resolve()
+        repo_root.mkdir(parents=True, exist_ok=True)
+        # One DKAO child per question, under this *attempt's* own directory.
+        # Each attempt is a different DKAO run, so none of them may share a
+        # child workspace: DKAO binds that workspace to the exact kernel repo it
+        # was launched with, and a second run pointed at the same directory dies
+        # on its own ownership check before it does any work.
+        children_root = cfg.exp_dir / CHILDREN_DIR
+        children_root.mkdir(parents=True, exist_ok=True)
+        # HE does not manage GPU occupancy. Every question this round is handed
+        # to its own DKAO child immediately, pinned to a device by the fixed
+        # ``index % GPU_COUNT`` rotation; whether that device is free is DKAO's
+        # business, not ours. The child's measurement gate
+        # (``ensure_measurement_gate``: VRAM <= 90% and HCU == 0, 30 min per
+        # check, at most 24 h) waits for the card, and the child's own
+        # contamination handling discards a round measured on a shared device.
+        # HE therefore never takes a lease, never pre-checks the cards and never
+        # freezes a child for a device it does not own — it only records the
+        # round's device layout (``_save_preflight``) for audit.
+        #
+        # The wave exists for one reason only: not to launch more children than
+        # there are devices (a wave is at most GPU_COUNT questions), so N
+        # questions land as N / 4 sequential waves with one question per device
+        # in each.
+        layout = self._device_layout(cfg, iteration, targets)
+        self._save_preflight(cfg, iteration, layout)
+        self._append_timeline(cfg, "gpu_layout_fixed", {
+            "iteration": iteration,
+            "gpu_count": GPU_COUNT,
+            "devices": layout["devices"],
+            "managed_by": layout["managed_by"],
+        })
+
+        results: Dict[str, Dict[str, Any]] = {}
+        pending = list(targets)
+        max_attempts = int(cfg.answers.get("env_retry_attempts") or 3)
+        attempt = 0
+        gate_blocked: List[str] = []
+        while pending:
+            wave_size = min(GPU_COUNT, len(pending))
+            wave = pending[:wave_size]
+            attempt += 1
+            child_root = (children_root
+                          / stage_dir_name(stage, iteration, attempt=attempt))
+            child_root.mkdir(parents=True, exist_ok=True)
+            if attempt > 1:
+                # The retry owns a fresh directory, so the failed attempt's
+                # evidence -- its logs above all -- stays where it is.
+                self._append_timeline(cfg, "env_retry_attempt", {
+                    "iteration": iteration, "stage": stage, "attempt": attempt,
+                    "dir": child_root.name,
+                    "questions": [inst.id for inst in wave],
+                })
+            wave_results = self._run_wave(cfg, workspace, iteration, wave,
+                                          repo_root, child_root,
+                                          wave_size)
+            results.update(wave_results)
+            # A child that produced no usable number (build/toolchain/crash) is
+            # retried: the environment owes us a measurement. Two things are NOT
+            # retried:
+            #  - a device the child's own gate refused to measure on (exit 75):
+            #    DKAO already waited its whole 24 h budget on that card. HE has
+            #    no cards to reassign, so relaunching would only ignore a verdict
+            #    it has no standing to overrule. The round stops instead.
+            #  - a contention window: gone entirely — HE no longer detects or
+            #    reacts to contention, DKAO's gate and contamination handling do.
+            gave_up = [inst for inst in pending
+                       if (wave_results.get(inst.id) or {}).get("gate_blocked")]
+            gate_blocked.extend(inst.id for inst in gave_up)
+            retry = [inst for inst in pending
+                     if inst not in gave_up
+                     and self._env_failed(wave_results.get(inst.id) or {})]
+            done = [inst for inst in pending if inst not in retry]
+            pending = [inst for inst in pending if inst not in done]
+            if not retry:
+                break
+            if attempt >= max_attempts:
+                # The budget counts attempts, not retries: with
+                # ``env_retry_attempts: 2`` a question runs at most twice.
+                self._append_timeline(cfg, "env_retry_exhausted", {
+                    "iteration": iteration, "attempts": attempt,
+                    "pending": [inst.id for inst in retry],
+                })
+                break
+            pending = retry
+        self._stop_on_incomplete_round(cfg, iteration, targets, results,
+                                       reason=("gpu_gate_blocked" if gate_blocked
+                                               else "environment_failed"))
+        return results
+
+    def _save_preflight(self, cfg: ExperimentConfig, iteration: int,
+                        layout: Dict[str, Any]) -> None:
+        """Record the round's device *layout* next to the iteration's evidence.
+
+        This is a layout record, not an occupancy check: it says which device
+        each question was handed to and that HE took no lease. Whether that
+        device was free at the time is answered by the child's own
+        ``measurement_gate.jsonl`` — the only place a gate verdict exists now.
+        """
+        stage = _round_stage(cfg)
+        name = "gpu_preflight.json" if stage != STAGE_GENERALIZATION else \
+            "gpu_preflight_generalization.json"
+        try:
+            bench_dir = (cfg.exp_dir / "runs" / f"iteration_{iteration:03d}"
+                         / "input" / "benchmark")
+            bench_dir.mkdir(parents=True, exist_ok=True)
+            (bench_dir / name).write_text(
+                json.dumps(layout, ensure_ascii=False, indent=2),
+                encoding="utf-8")
+        except OSError:
+            pass
+
+    @classmethod
+    def _device_layout(cls, cfg: ExperimentConfig, iteration: int,
+                       targets: List[Any]) -> Dict[str, Any]:
+        """The fixed ``index % GPU_COUNT`` device assignment for one round."""
+        devices = [device_for_index(idx) for idx, _ in enumerate(targets)]
+        return {
+            "schema": "he-gpu-layout/1",
+            "enabled": False,
+            "managed_by": "dcu_kernel_auto_opt",
+            "reason": ("harness_evolve does not manage GPU occupancy; DKAO's "
+                       "measurement gate (VRAM<=90% and HCU==0) decides when a "
+                       "device may be measured on"),
+            "leases": [],
+            "gpu_count": GPU_COUNT,
+            "iteration": iteration,
+            "devices": {
+                str(getattr(inst, "id", f"question_{idx}")): gpu
+                for idx, (inst, gpu) in enumerate(zip(targets, devices))
+            },
+            "assigned_devices": sorted(set(devices)),
+        }
+
+    def _run_wave(self, cfg: ExperimentConfig, workspace: Path, iteration: int,
+                  wave: List[Any], repo_root: Path,
+                  child_root: Path,
+                  wave_size: int) -> Dict[str, Dict[str, Any]]:
+        """Run one wave: one DKAO child per question, one question per device.
+
+        The device is chosen by the fixed ``index % GPU_COUNT`` rotation rather
+        than by a reading of the cards: HE hands the task over, DKAO's own
+        measurement gate decides when the card is clean enough to time a
+        kernel. Nothing here watches the devices — a round that DKAO's gate
+        holds up is visible in the child's ``measurement_gate.jsonl``, not in a
+        verdict HE invents.
+        """
+        results: Dict[str, Dict[str, Any]] = {}
+        if not wave:
+            return results
+        with concurrent.futures.ThreadPoolExecutor(max_workers=wave_size) as ex:
+            futures = {
+                ex.submit(
+                    self._run_one, cfg, workspace, iteration, inst,
+                    device_for_index(idx),
+                    repo_root, child_root, None,
+                ): inst.id
+                for idx, inst in enumerate(wave)
+            }
+            for fut in concurrent.futures.as_completed(futures):
+                iid = futures[fut]
+                try:
+                    results[iid] = fut.result()
+                except Exception as exc:  # noqa: BLE001
+                    import traceback as _traceback
+                    tb_text = _traceback.format_exc(limit=8)
+                    results[iid] = {
+                        "status": "exception",
+                        "passed": False,
+                        "correctness_ok": False,
+                        "p90_ok": False,
+                        "median_us": None,
+                        "p90_us": None,
+                        "reason": (
+                            f"evaluator exception: {exc}\n{tb_text}"
+                        ),
+                        "child_task_id": None,
+                        "repo_path": None,
+                        "harness_revision": _harness_revision(workspace),
+                    }
+        return results
+
+    #: statuses that mean "no usable measurement because the environment broke".
+    #: ``partial_worker_result`` is deliberately absent: the child may have
+    #: produced a valid measurement whose only problem is that final synthesis
+    #: did not land. Treating that as an environment failure requeued questions
+    #: that had already measured, and the retry then collided with the child's
+    #: own kernel repo — which is how a round ended up with zero medians.
+    ENV_FAILURE_STATUSES = ("missing_report", "exception", "timeout")
+
+    #: statuses that mean "the workers measured, the round did not finish"
+    PARTIAL_STATUSES = ("partial_worker_result",)
+
+    #: reason fragments that mean "the toolchain/environment broke", i.e. the
+    #: question deserves a retry rather than being counted as a bad kernel
+    ENV_FAILURE_HINTS = ("compile", "build failed", "ninja", "hipcc",
+                         "undefined template", "no space left",
+                         "outofmemory", "out of memory", "timed out",
+                         "timeout", "killed")
+
+    def _env_failed(self, result: Dict[str, Any],
+                    *, retry: bool = False) -> bool:
+        """A failure caused by the environment (not by the candidate kernel).
+
+        The rule is "no usable number => the environment owes us a measurement":
+        a child that could not be built, timed out or crashed is retried (30
+        minutes of rest, up to 3 attempts). A child that *did* measure is not
+        retried — its number is evidence even when the run ended early (for
+        example ``partial_worker_result``: the workers measured, final synthesis
+        did not land).
+
+        ``retry=True`` marks a re-measurement of a question that already ran
+        once, where even a measured partial is an environment problem: its
+        device was taken over mid-run, so the number describes a shared card.
+        """
+        if not isinstance(result, dict):
+            return False
+        if result.get("passed") is True:
+            return False
+        status = str(result.get("status") or "").lower()
+        measured = result.get("median_us") not in (None, "")
+        if status in self.ENV_FAILURE_STATUSES:
+            return True
+        if measured:
+            # a real number: only a retry treats a partial finish as an
+            # environment problem, and only because the device was shared
+            return bool(retry and status in self.PARTIAL_STATUSES)
+        reason = str(result.get("reason") or "").lower()
+        if any(hint in reason for hint in self.ENV_FAILURE_HINTS):
+            return True
+        return True
+
+    def _request_stop(self, cfg: ExperimentConfig, reason: str,
+                      **fields: Any) -> None:
+        """Ask the outer loop to stop the whole experiment (terminal state)."""
+        try:
+            (cfg.exp_dir / "stop_requested.json").write_text(
+                json.dumps({"reason": reason, **fields}, ensure_ascii=False,
+                           indent=2), encoding="utf-8")
+        except OSError:
+            pass
+
+    #: statuses that mean "the round ended without a measurement"
+    INCOMPLETE_STATUSES = ENV_FAILURE_STATUSES + ("rejected", "failed")
+
+    def _incomplete_ids(self, targets: List[Any],
+                        results: Dict[str, Dict[str, Any]]) -> List[str]:
+        """Questions the round never measured (no number, or a broken run).
+
+        A round that ran out of retries used to be recorded as "these
+        questions just failed", and the decisions downstream then compared
+        against whatever numbers happened to exist. That silently weakens the
+        evidence for a promotion, so the caller stops the run instead.
+        """
+        out: List[str] = []
+        for inst in targets:
+            iid = getattr(inst, "id", str(inst))
+            res = results.get(iid)
+            if not isinstance(res, dict) or res.get("passed") is True:
+                continue
+            if not self._is_measured(res):
+                out.append(iid)
+        return out
+
+    def _is_measured(self, result: Dict[str, Any]) -> bool:
+        """True when the result carries a usable number of its own."""
+        if not isinstance(result, dict):
+            return False
+        median = result.get("median_us")
+        if median in (None, ""):
+            return False
+        status = str(result.get("status") or "").lower()
+        return status not in self.ENV_FAILURE_STATUSES
+
+    def _stop_on_incomplete_round(self, cfg: ExperimentConfig, iteration: int,
+                                  targets: List[Any],
+                                  results: Dict[str, Dict[str, Any]],
+                                  *, reason: str) -> bool:
+        """True when the round must stop instead of deciding on partial data."""
+        missing = self._incomplete_ids(list(targets), results or {})
+        if not missing:
+            return False
+        unmeasured = sorted(set(missing))
+        self._append_timeline(cfg, "round_incomplete", {
+            "iteration": iteration, "reason": reason,
+            "unmeasured": unmeasured,
+            "measured": sorted(set(results or {}) - set(unmeasured)),
+        })
+        self._request_stop(cfg, "round_incomplete", iteration=iteration,
+                           cause=reason, unmeasured=unmeasured)
+        return True
+
+    def _append_timeline(self, cfg: ExperimentConfig, event: str,
+                         payload: Dict[str, Any]) -> None:
+        try:
+            from metainfer.tasks.harness_evolve.orchestrator.state import (
+                append_timeline,
+            )
+            append_timeline(cfg.state_dir, event, payload)
+        except Exception:  # noqa: BLE001 - progress records are best effort
+            pass
+
+    def _run_one(
+        self, cfg: ExperimentConfig, workspace: Path, iteration: int,
+        inst: InstanceSpec, index: int, repo_root: Path, child_root: Path,
+        preflight: Optional[Dict[str, Any]] = None,
+        gpu_override: "int | None" = None,
+    ) -> Dict[str, Any]:
+        # Device layout: the machine exposes 4 GPUs; spread the concurrent
+        # children over 0..3 in submission order and repeat cyclically, so a
+        # 4-multiple question set keeps every card busy exactly once per wave.
+        # (worker id == GPU id is enforced by DKAO manual assignments, hence
+        # _build_dkao_requirements names the worker worker_.)
+        gpu = (int(gpu_override) if gpu_override is not None
+               else device_for_index(index))
+        suffix = uuid.uuid4().hex[:6]
+        child_id = _slug(
+            f"ahe-{cfg.task_id}-it{iteration:03d}-candidate-{inst.id}-{suffix}"
+        )
+        repo_name = child_id
+        state_dir = child_root / inst.id / "state"
+        work_dir = child_root / inst.id / "workspace"
+        state_dir.mkdir(parents=True, exist_ok=True)
+        work_dir.mkdir(parents=True, exist_ok=True)
+        req = self._build_dkao_requirements(
+            cfg, inst, iteration, child_id, repo_name, gpu, workspace
+        )
+        req_path = child_root / inst.id / "requirements.json"
+        req_path.parent.mkdir(parents=True, exist_ok=True)
+        req_path.write_text(json.dumps(req, ensure_ascii=False, indent=2),
+                            encoding="utf-8")
+        log_path = child_root / inst.id / "orchestrator-external.log"
+
+        cmd = [
+            _python_executable(), "-m",
+            "metainfer.tasks.dcu_kernel_auto_opt.orchestrator.cli", "run",
+            str(req_path), "--state-dir", str(state_dir),
+            "--workspace-dir", str(work_dir),
+        ]
+        env = _child_env(cfg, repo_root, workspace,
+                         child_state_dir=state_dir)
+        timeout_min = float(cfg.answers.get("child_timeout_minutes") or 0)
+        timeout = timeout_min * 60 if timeout_min > 0 else None
+        started = time.time()
+        # Spawn (not run): the parent keeps a handle so it can terminate the
+        # child on a timeout. Freezing a child for "its" device is gone with the
+        # rest of HE's GPU management — that device is DKAO's to arbitrate.
+        child = spawn(cmd, cwd=str(Path(__file__).resolve().parents[5]),
+                      env=env, log_path=log_path)
+        self.children.add(inst.id, child)
+        rc, timed_out = 0, False
+        try:
+            if timeout:
+                rc = child.proc.wait(timeout=timeout)
+            else:
+                rc = child.proc.wait()
+        except subprocess.TimeoutExpired:
+            child.terminate(reason=f"child timeout after {timeout:.0f}s")
+            rc, timed_out = 124, True
+        finally:
+            self.children.drop(inst.id)
+            try:
+                child.proc.wait(timeout=30)
+            except Exception:  # noqa: BLE001 - already dead or unkillable
+                pass
+        terminated_by = child.terminate_reason
+        duration = time.time() - started
+        report_path = work_dir / "final_report.json"
+        report = _load_json(report_path) or {}
+        normalized = (
+            _normalize_report(inst, report)
+            if report else _normalize_workspace_result(inst, work_dir)
+        )
+        normalized.update({
+            "child_task_id": child_id,
+            "parent_ahe_task_id": cfg.task_id,
+            "ahe_iteration": iteration,
+            "candidate_role": "candidate",
+            "repo_path": str(repo_root / repo_name),
+            "repo_root": str(repo_root),
+            "harness_revision": _harness_revision(workspace),
+            "child_state_dir": str(state_dir),
+            "child_workspace_dir": str(work_dir),
+            "log_path": str(log_path),
+            "duration_s": round(duration, 3),
+            "returncode": rc,
+            "timed_out": timed_out,
+            "gpu": gpu,
+            "terminated_by": terminated_by or None,
+        })
+        # No ``gpu_preflight`` / ``measurement_suspect`` stamping here any more:
+        # HE takes no reading of the device it handed out, so it has no verdict
+        # to attach. The child records every admission check it made in
+        # ``/measurement_gate.jsonl``, which is what the HE
+        # detail page reads when it shows why a question is waiting.
+        #
+        # ``GATE_EXIT_CODE`` (75) is the child's *own* gate saying "I waited 48
+        # x 30 min and no card ever became clean". That is an authoritative
+        # verdict, not a flaky environment: HE does not repeal it by launching
+        # the same child again, so it is marked and the round stops.
+        if rc == GATE_EXIT_CODE:
+            normalized["gate_blocked"] = True
+            normalized.setdefault("status", "environment_failed")
+            normalized["reason"] = (
+                f"DKAO measurement gate gave up on device "
+                f"{gpu} after its full wait budget"
+            )
+        #
+        # final_report.json does not carry a canonical rounds field (its shape
+        # changed over DKAO revisions), so fall back to the authoritative
+        # per-attempt records the workers wrote.
+        counted = _worker_rounds(work_dir)
+        if counted and int(normalized.get("rounds_used") or 0) < counted:
+            normalized["rounds_used"] = counted
+        if rc != 0 and normalized.get("status") == "missing_report":
+            normalized["reason"] = _tail(log_path)
+        return normalized
+
+    def _build_dkao_requirements(
+        self, cfg: ExperimentConfig, inst: InstanceSpec, iteration: int,
+        child_id: str, repo_name: str, gpu: int, workspace: Path,
+    ) -> Dict[str, Any]:
+        s = inst.shape
+        shape_yaml = {
+            "model": str(s.get("model", "AHE registered GEMM")),
+            "assignment_mode": "manual",
+            "shape_scope": "subset",
+            "shapes": [{
+                "id": inst.id,
+                "M": int(s["M"]), "N": int(s["N"]), "K": int(s["K"]),
+                "tp_size": int(s.get("tp_size", 8)),
+                "operator": str(s["operator"]),
+            }],
+            "assignments": {
+                # DKAO's manual mode requires worker_N to map to GPU N
+                # (gen_and_opt_pipeline validates worker_id == gpu). Name the
+                # single worker after the assigned GPU so parallel children
+                # spread across devices 0..3 without tripping the check.
+                f"worker_{gpu}": {"gpu": gpu, "shapes": [inst.id]},
+            },
+        }
+        return {
+            "task_id": child_id,
+            "task_type": "dcu-kernel-auto-opt",
+            "label": child_id,
+            "operator": "Quantized GEMM",
+            "kernel_language": "HIP C++",
+            "target_hardware": "K500SM_AI / gfx928",
+            "dtype": "INT8 W8A8",
+            "agent_framework": cfg.agent_framework,
+            # The child DKAO task runs on the same DSH model generation as the
+            # AHE evolve agent (agent_model is a DKAO form label, resolved to an
+            # id by DKAO's own orchestrator/config.py).
+            "agent_model": "deepseek-flash-4.1",
+            "execution_mode": "Generate & optimize (auto-create kernel repo)",
+            "target_repo_path": repo_name,
+            "model": str(s.get("model", "AHE registered GEMM")),
+            "shape_assignment_mode": "Manual by GPU",
+            "shape_scope": "Selected shapes only",
+            "shape_config": yaml.safe_dump(shape_yaml, sort_keys=False),
+            "correctness_ref": "PyTorch eager",
+            "perf_target": "AHE automatic baseline-relative criterion",
+            "max_iterations": str(inst.budget_rounds),
+            "minimum_improvement_percent": 0,
+            # Parent/trace metadata: ignored by old DKAO code, preserved in req.
+            "parent_ahe_task_id": cfg.task_id,
+            "ahe_iteration": iteration,
+            "harness_revision": _harness_revision(workspace),
+            "candidate_role": "candidate",
+        }
+
+
+def _normalize_workspace_result(inst: InstanceSpec, work_dir: Path) -> Dict[str, Any]:
+    """Best-effort evidence when final synthesis/serial validate fails.
+
+    Partial worker metrics are visible in the AHE frontend but never pass the
+    system gate (no final validation = candidate cannot be promoted).
+    """
+    for path in sorted((work_dir / "workers").glob("*/result.json")):
+        data = _load_json(path) or {}
+        shape = (data.get("shapes") or {}).get(inst.id) or {}
+        metrics = shape.get("metrics") or {}
+        if not metrics:
+            continue
+        median, p90 = metrics.get("median_us"), metrics.get("p90_us")
+        correctness = bool(
+            metrics.get("graph_capture_passed") is True
+            and metrics.get("correctness_passed_in_precheck") is True
+        )
+        rounds_path = path.parent / "runs" / inst.id / "experiments.jsonl"
+        rounds = 0
+        if rounds_path.is_file():
+            rounds = len(rounds_path.read_text(
+                encoding="utf-8", errors="replace"
+            ).splitlines())
+        return {
+            "status": "partial_worker_result",
+            "passed": False,
+            "correctness_ok": correctness,
+            "p90_ok": bool(p90 is not None),
+            "median_us": float(median) if median is not None else None,
+            "p90_us": float(p90) if p90 is not None else None,
+            "rounds_used": rounds,
+            "tokens": None,
+            "new_best_known": False,
+            "reason": "final report missing; worker best preserved as evidence",
+        }
+    return {
+        "status": "missing_report", "passed": False,
+        "correctness_ok": False, "p90_ok": False,
+        "median_us": None, "p90_us": None, "rounds_used": 0,
+        "tokens": None, "reason": "final_report.json and worker result missing",
+    }
+
+
+def _normalize_report(inst: InstanceSpec, report: Dict[str, Any]) -> Dict[str, Any]:
+    if not report:
+        return {
+            "status": "missing_report", "passed": False,
+            "correctness_ok": False, "p90_ok": False,
+            "median_us": None, "p90_us": None, "rounds_used": 0,
+            "tokens": None, "reason": "final_report.json missing",
+        }
+    final = (report.get("final_validation") or {}).get(inst.id) or {}
+    worker = (report.get("worker_validation") or {}).get(inst.id) or {}
+    metrics = final or (worker.get("metrics") or {})
+    correctness = bool(
+        metrics.get("passed") is True
+        or (
+            metrics.get("graph_capture_passed") is True
+            and (metrics.get("mismatch_count") in (None, 0))
+            and report.get("status") == "success"
+        )
+    )
+    median = metrics.get("median_us")
+    p90 = metrics.get("p90_us")
+    return {
+        "status": str(report.get("status") or "unknown"),
+        "passed": bool(report.get("status") == "success" and correctness),
+        "correctness_ok": correctness,
+        "p90_ok": bool(p90 is not None),
+        "median_us": float(median) if median is not None else None,
+        "p90_us": float(p90) if p90 is not None else None,
+        "rounds_used": _count_rounds(report, inst.id),
+        "tokens": None,
+        "new_best_known": False,
+        "reason": "final_report parsed",
+        "final_report_status": report.get("status"),
+    }
+
+
+def _count_rounds(report: Dict[str, Any], shape_id: str) -> int:
+    # final report does not currently expose a canonical rounds_used field;
+    # worker result/details remain available through child_workspace_dir.
+    workers = report.get("workers") or {}
+    for worker in workers.values():
+        shape = (worker.get("shapes") or {}).get(shape_id) if isinstance(worker, dict) else None
+        if isinstance(shape, dict):
+            try:
+                return int(shape.get("iterations") or shape.get("rounds_used") or 0)
+            except (TypeError, ValueError):
+                pass
+    return 0
+
+
+def _harness_revision(workspace: Path) -> str:
+    manifest = _load_json(workspace / "manifest.json")
+    if not manifest and (workspace / "manifest.yaml").is_file():
+        try:
+            manifest = yaml.safe_load(
+                (workspace / "manifest.yaml").read_text(encoding="utf-8")
+            ) or {}
+        except (OSError, ValueError):
+            manifest = {}
+    return str((manifest or {}).get("revision") or "unknown")
+
+
+def _load_json(path: Path) -> Dict[str, Any] | None:
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    return value if isinstance(value, dict) else None
+
+
+def _slug(value: str) -> str:
+    value = re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-")
+    return value[:120] or f"ahe-child-{uuid.uuid4().hex[:8]}"
+
+
+def _tail(path: Path, size: int = 2000) -> str:
+    try:
+        text = path.read_text(encoding="utf-8", errors="replace")
+    except OSError:
+        return "child failed; log unavailable"
+    return text[-size:]
+
+
+def _truthy(value: Any, default: bool = False) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, bool):
+        return value
+    return str(value).strip().lower() not in {"0", "false", "no", "off", ""}
+
+
+def _child_env(cfg: ExperimentConfig, repo_root: Path,
+               workspace: Path,
+               child_state_dir: Optional[Path] = None) -> Dict[str, str]:
+    """Environment for one DKAO child task.
+
+    ``METAINFER_PLANNER=1`` is what makes the evaluated harness' planner
+    actually drive DKAO's round instructions; without it the harness snapshot
+    is inert for planning decisions. The form switch (default on) keeps that
+    auditable.
+
+    ``METAINFER_CHILD_STATE_DIR`` points the child's own measurement gate at a
+    state directory, so every "gate blocked / gate ok" check it makes while
+    iterating is auditable next to the round's other evidence.
+    """
+    env = os.environ.copy()
+    env["METAINFER_KERNEL_REPOS"] = str(repo_root)
+    env["METAINFER_HARNESS_ROOT"] = str(workspace)
+    if child_state_dir is not None:
+        env["METAINFER_CHILD_STATE_DIR"] = str(child_state_dir)
+    if _truthy(cfg.answers.get("planner_enabled"), default=True):
+        env["METAINFER_PLANNER"] = "1"
+    else:
+        env.pop("METAINFER_PLANNER", None)
+
+    # Validation budget: one question = one shape, so AHE defaults to the
+    # task-scoped serial validation with the quick bench profile (production
+    # DKAO tasks keep the full API sweep unless their form says otherwise).
+    answers = cfg.answers
+    scope = str(answers.get("validate_scope") or "task").strip().lower()
+    env["METAINFER_VALIDATE_SCOPE"] = scope if scope in {"api", "task"} else "task"
+    profile = str(answers.get("bench_profile") or "quick").strip().lower()
+    quick = {}
+    if profile in {"quick", "fast"}:
+        try:
+            from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.validation_budget import (  # noqa: E501
+                QUICK_BENCH,
+            )
+            quick = dict(QUICK_BENCH)
+        except Exception:  # noqa: BLE001
+            quick = {"warmups": 10, "samples": 20, "replays_per_sample": 20}
+    explicit = {
+        "METAINFER_BENCH_WARMUPS": answers.get("bench_warmups"),
+        "METAINFER_BENCH_SAMPLES": answers.get("bench_samples"),
+        "METAINFER_BENCH_REPLAYS": answers.get("bench_replays"),
+    }
+    if quick:
+        env["METAINFER_BENCH_WARMUPS"] = str(quick["warmups"])
+        env["METAINFER_BENCH_SAMPLES"] = str(quick["samples"])
+        env["METAINFER_BENCH_REPLAYS"] = str(quick["replays_per_sample"])
+    for name, value in explicit.items():
+        if value not in (None, ""):
+            try:
+                env[name] = str(int(value))
+            except (TypeError, ValueError):
+                pass
+    return env
+
+
+def _worker_rounds(work_dir: Path) -> int:
+    """Rounds actually attempted, from each worker's experiments.jsonl."""
+    total = 0
+    workers = work_dir / "workers"
+    if not workers.is_dir():
+        return 0
+    for exp_file in workers.glob("*/runs/*/experiments.jsonl"):
+        try:
+            total += len(exp_file.read_text(
+                encoding="utf-8", errors="replace").splitlines())
+        except OSError:
+            continue
+    return total
+
+
+def _python_executable() -> str:
+    """A reliable python binary for spawning DKAO child orchestrators.
+
+    ``sys.executable`` can be ``''`` when the WebUI launcher renames argv[0]
+    to a friendly process name (``metainfer-orchestrator``) via Popen's
+    ``executable=`` split: CPython then cannot resolve its own path, and
+    spawning a child with the empty string fails with
+    ``PermissionError [Errno 13] Permission denied: ''``. Prefer the explicit
+    ``METAINFER_PYTHON`` override, then a valid ``sys.executable``, then
+    ``python3`` on PATH.
+    """
+    import shutil
+
+    env_py = str(os.environ.get("METAINFER_PYTHON") or "").strip()
+    if env_py:
+        return env_py
+    exe = str(sys.executable or "").strip()
+    if exe and os.path.isfile(exe) and os.access(exe, os.X_OK):
+        return exe
+    for cand in ("python3", "python"):
+        found = shutil.which(cand)
+        if found:
+            return found
+    return "/usr/bin/python3"
+
+
+def build_evaluator(mode: str) -> Evaluator:
+    return DkaoCliEvaluator() if mode == "dkao-cli" else DryRunEvaluator()
diff --git a/metainfer/tasks/harness_evolve/orchestrator/attribution.py b/metainfer/tasks/harness_evolve/orchestrator/attribution.py
new file mode 100644
index 00000000..440ff0c3
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/attribution.py
@@ -0,0 +1,134 @@
+"""Attribution / flip / rollback helpers (ported semantics from ahe-ref)."""
+
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+from typing import Any, Dict, List
+
+
+def pass_rate(results: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
+    items = list(results.values())
+    n = len(items)
+    n_pass = sum(1 for r in items if r.get("passed") is True)
+    return {
+        "n_total": n,
+        "n_pass": n_pass,
+        "pass_rate": (n_pass / n) if n else 0.0,
+    }
+
+
+def diff_results(
+    prev: Dict[str, Dict[str, Any]] | None,
+    curr: Dict[str, Dict[str, Any]],
+) -> Dict[str, Any]:
+    """Cross-iteration diff (bool-level; three-state noise lives in real eval).
+
+    flipped: fail -> pass   regressed: pass -> fail   unchanged otherwise.
+    """
+    prev = prev or {}
+    flipped: List[str] = []
+    regressed: List[str] = []
+    for inst_id, r in curr.items():
+        p_ok = bool((prev.get(inst_id) or {}).get("passed") is True)
+        c_ok = bool(r.get("passed") is True)
+        if not p_ok and c_ok:
+            flipped.append(inst_id)
+        elif p_ok and not c_ok:
+            regressed.append(inst_id)
+    return {
+        "flipped": sorted(flipped),
+        "regressed": sorted(regressed),
+        "net": len(flipped) - len(regressed),
+    }
+
+
+def evaluate_changes(manifest: Dict[str, Any], diff: Dict[str, Any]) -> Dict[str, Any]:
+    """change_manifest -> per-change verdicts (schema mirrors ahe-ref)."""
+    flipped_set = set(diff.get("flipped", []))
+    regressed_set = set(diff.get("regressed", []))
+    all_predicted = set()
+    all_risk = set()
+    evaluations: List[Dict[str, Any]] = []
+
+    for chg in manifest.get("changes", []):
+        predicted = list(chg.get("predicted_fixes", []) or [])
+        risks = list(chg.get("risk_tasks", []) or [])
+        all_predicted.update(predicted)
+        all_risk.update(risks)
+        actually_fixed = [t for t in predicted if t in flipped_set]
+        still_failed = [t for t in predicted if t not in flipped_set]
+        risk_realized = [t for t in risks if t in regressed_set]
+        n_fixed = len(actually_fixed)
+        n_pred = len(predicted)
+        n_risk = len(risk_realized)
+
+        if n_risk > 0 and n_fixed == 0:
+            verdict = "HARMFUL"
+        elif n_risk > 0 and n_fixed > 0:
+            verdict = "MIXED"
+        elif n_pred and n_fixed == n_pred:
+            verdict = "EFFECTIVE"
+        elif n_fixed > 0:
+            verdict = "PARTIALLY_EFFECTIVE"
+        else:
+            verdict = "INEFFECTIVE"
+
+        evaluations.append({
+            "change_id": chg.get("id", "unknown"),
+            "description": chg.get("description", ""),
+            "files": list(chg.get("files", []) or []),
+            "predicted_fixes": predicted,
+            "actually_fixed": actually_fixed,
+            "still_failed": still_failed,
+            "predicted_risks": risks,
+            "risk_realized": risk_realized,
+            "hit_rate": f"{n_fixed}/{n_pred}" if n_pred else "0/0",
+            "verdict": verdict,
+        })
+
+    unattributed = sorted(
+        t for t in regressed_set if t not in (all_predicted | all_risk)
+    )
+    return {
+        "evaluating_iteration": int(manifest.get("iteration", 0)),
+        "change_evaluations": evaluations,
+        "unattributed_regressions": unattributed,
+        "summary": ", ".join(
+            f"{e['change_id']}: {e['verdict']}" for e in evaluations
+        ),
+    }
+
+
+def copy_tree(src: Path, dst: Path) -> None:
+    dst.mkdir(parents=True, exist_ok=True)
+    for item in src.iterdir():
+        if item.name == ".git":
+            continue
+        target = dst / item.name
+        if item.is_dir():
+            shutil.copytree(item, target, dirs_exist_ok=True)
+        else:
+            shutil.copy2(item, target)
+
+
+def replace_tree(src: Path, dst: Path) -> None:
+    """Replace dst's non-git content with src (true rollback semantics)."""
+    dst.mkdir(parents=True, exist_ok=True)
+    for item in list(dst.iterdir()):
+        if item.name == ".git":
+            continue
+        if item.is_dir():
+            shutil.rmtree(item)
+        else:
+            item.unlink()
+    copy_tree(src, dst)
+
+
+def save_json(path: Path, payload: Dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2),
+                   encoding="utf-8")
+    tmp.replace(path)
diff --git a/metainfer/tasks/harness_evolve/orchestrator/champion_selection.py b/metainfer/tasks/harness_evolve/orchestrator/champion_selection.py
new file mode 100644
index 00000000..da5e72f2
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/champion_selection.py
@@ -0,0 +1,211 @@
+"""Generation-level champion selection for harness_evolve.
+
+The decision engine compares the candidate harness against a *champion*. Using
+"the last generation that happened to be PROMOTEd" is wrong: a generation can
+hold the best measured result and still not be promoted in its own round (its
+change needed confirmation, or only one question won). If the next generation
+is worse but does get promoted, that worse harness would take the champion
+slot.
+
+This module keeps one record per evaluated generation and picks the champion by
+*measured merit*:
+
+1. generations whose verdict was REJECT / REJECT_OVERFIT are disqualified
+   (they were proven not to generalise or to regress);
+2. more passing questions wins (absolute count on each generation's own set);
+3. otherwise, on the questions the two generations share, more WINs wins;
+4. otherwise the generation with the better cross-shape median wins;
+5. ties keep the earlier generation (stability over novelty).
+
+The champion is therefore "the best harness we have measured so far", which is
+what every later round must beat — even if that generation was never promoted
+at the time.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional
+
+GENERATIONS_FILE = "generations.jsonl"
+_REJECTED_VERDICTS = {"REJECT", "REJECT_OVERFIT"}
+_PENDING_VERDICTS = {"CONFIRM_REQUIRED"}
+
+
+@dataclass
+class Generation:
+    iteration: int
+    verdict: str = ""
+    results: Dict[str, Dict[str, Any]] = field(default_factory=dict)
+    snapshot_dir: str = ""
+    workspace_revision: str = ""
+    pass_count: int = 0
+    rejected: bool = False
+    pending: bool = False
+
+    def __post_init__(self) -> None:
+        # the verdict decides qualification; keep the two flags in sync
+        if self.verdict in _REJECTED_VERDICTS:
+            self.rejected = True
+        if self.verdict in _PENDING_VERDICTS:
+            self.pending = True
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "iteration": self.iteration,
+            "verdict": self.verdict,
+            "pass_count": self.pass_count,
+            "rejected": self.rejected,
+            "pending": self.pending,
+            "snapshot_dir": self.snapshot_dir,
+            "workspace_revision": self.workspace_revision,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any],
+                  results: Optional[Mapping[str, Any]] = None) -> "Generation":
+        verdict = str(data.get("verdict") or "")
+        return cls(
+            iteration=int(data.get("iteration") or 0),
+            verdict=verdict,
+            results=dict(results or {}),
+            snapshot_dir=str(data.get("snapshot_dir") or ""),
+            workspace_revision=str(data.get("workspace_revision") or ""),
+            pass_count=int(data.get("pass_count") or 0),
+            rejected=bool(data.get("rejected", verdict in _REJECTED_VERDICTS)),
+            pending=bool(data.get("pending", verdict in _PENDING_VERDICTS)),
+        )
+
+
+def _median(row: Mapping[str, Any]) -> Optional[float]:
+    try:
+        value = float(row.get("median_us"))
+    except (TypeError, ValueError):
+        return None
+    return value if value > 0 else None
+
+
+def load_generations(exp_dir: Path) -> List[Generation]:
+    """Every evaluated generation, oldest first, with its measured results."""
+    exp = Path(exp_dir)
+    path = exp / GENERATIONS_FILE
+    records: List[Dict[str, Any]] = []
+    if path.is_file():
+        for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+            try:
+                row = json.loads(line)
+            except ValueError:
+                continue
+            if isinstance(row, dict):
+                records.append(row)
+    out: List[Generation] = []
+    for row in sorted(records, key=lambda r: int(r.get("iteration") or 0)):
+        iteration = int(row.get("iteration") or 0)
+        results_path = (exp / "runs" / f"iteration_{iteration:03d}" / "input"
+                        / "benchmark" / "results.json")
+        results: Dict[str, Any] = {}
+        if results_path.is_file():
+            try:
+                results = json.loads(results_path.read_text(encoding="utf-8")) or {}
+                results = results.get("results", results)
+            except (OSError, ValueError):
+                results = {}
+        out.append(Generation.from_dict(row, results))
+    return out
+
+
+def record_generation(exp_dir: Path, generation: Generation) -> None:
+    """Append (or replace) one generation record."""
+    exp = Path(exp_dir)
+    path = exp / GENERATIONS_FILE
+    rows: List[Dict[str, Any]] = []
+    if path.is_file():
+        for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+            try:
+                row = json.loads(line)
+            except ValueError:
+                continue
+            if isinstance(row, dict) and int(row.get("iteration") or 0) != generation.iteration:
+                rows.append(row)
+    rows.append(generation.to_dict())
+    rows.sort(key=lambda r: int(r.get("iteration") or 0))
+    path.write_text(
+        "\n".join(json.dumps(r, ensure_ascii=False) for r in rows) + "\n",
+        encoding="utf-8")
+
+
+def _pass_count(results: Mapping[str, Mapping[str, Any]]) -> int:
+    return sum(1 for row in results.values() if row.get("passed") is True)
+
+
+def _wins(a: Mapping[str, Mapping[str, Any]],
+          b: Mapping[str, Mapping[str, Any]], noise_percent: float) -> int:
+    """How many questions generation ``a`` beats generation ``b`` on."""
+    wins = 0
+    for iid in set(a) & set(b):
+        ma, mb = _median(a[iid]), _median(b[iid])
+        if ma is None or mb is None:
+            continue
+        delta = (mb - ma) / mb * 100.0        # positive => a is faster
+        if delta > noise_percent:
+            wins += 1
+    return wins
+
+
+def compare_generations(a: Generation, b: Generation, *,
+                        noise_percent: float = 2.0) -> str:
+    """Return ``\"a\"``, ``\"b\"`` or ``\"tie\"`` (a wins / b wins / no winner)."""
+    if a.rejected and not b.rejected:
+        return "b"
+    if b.rejected and not a.rejected:
+        return "a"
+    a_pass, b_pass = _pass_count(a.results), _pass_count(b.results)
+    if a_pass != b_pass:
+        return "a" if a_pass > b_pass else "b"
+    a_wins = _wins(a.results, b.results, noise_percent)
+    b_wins = _wins(b.results, a.results, noise_percent)
+    if a_wins != b_wins:
+        return "a" if a_wins > b_wins else "b"
+    # shared-question medians: compare the geometric mean of ratios
+    shared = [iid for iid in set(a.results) & set(b.results)
+              if _median(a.results[iid]) and _median(b.results[iid])]
+    if shared:
+        ratio = 1.0
+        for iid in shared:
+            ratio *= _median(a.results[iid]) / _median(b.results[iid])
+        ratio **= 1.0 / len(shared)
+        if ratio < 1.0 - noise_percent / 100.0:
+            return "a"
+        if ratio > 1.0 + noise_percent / 100.0:
+            return "b"
+    return "tie"
+
+
+def best_generation(generations: List[Generation], *,
+                    noise_percent: float = 2.0) -> Optional[Generation]:
+    """The champion: best measured generation among non-rejected ones."""
+    candidates = [g for g in generations if not g.rejected and g.results]
+    if not candidates:
+        return None
+    best = candidates[0]
+    for gen in candidates[1:]:
+        winner = compare_generations(gen, best, noise_percent=noise_percent)
+        if winner == "a":
+            best = gen
+    return best
+
+
+def summarize(generations: List[Generation]) -> List[Dict[str, Any]]:
+    """Cross-round review table: one row per generation."""
+    champion = best_generation(generations)
+    rows: List[Dict[str, Any]] = []
+    for gen in generations:
+        rows.append({
+            **gen.to_dict(),
+            "champion": bool(champion and champion.iteration == gen.iteration),
+            "medians": {iid: _median(row)
+                        for iid, row in sorted(gen.results.items())},
+        })
+    return rows
diff --git a/metainfer/tasks/harness_evolve/orchestrator/cli.py b/metainfer/tasks/harness_evolve/orchestrator/cli.py
new file mode 100644
index 00000000..d7cdf1a0
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/cli.py
@@ -0,0 +1,462 @@
+"""CLI entry point for harness_evolve orchestrator.
+
+Framework contract::
+
+    python -m metainfer.tasks.harness_evolve.orchestrator.cli run \\
+         --state-dir DIR --workspace-dir DIR
+
+Explicit publish gate::
+
+    python -m metainfer.tasks.harness_evolve.orchestrator.cli publish \\
+        --exp-dir DIR [--harness] [--harness-target PATH] \\
+        [--kernels] [--pool PATH] [--reason TEXT]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import signal
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Iterator, Optional
+
+from metainfer.orchestrator._bootstrap import clear_pid_file, write_pid_file
+
+from .config import load_experiment_config
+from .pipeline import completed_iterations, run_experiment
+from .state import read_target_iterations, write_target_iterations
+from .publish import check_publishable, publish_harness, publish_kernels
+
+_DEFAULT_POOL = Path("/root/zth_agent/ahe-kernel-repos/registered_pool.yaml")
+
+
+@contextmanager
+def _orchestrator_pid(state_dir: Path, task_id: str) -> Iterator[Path]:
+    """Own the task's ``orchestrator.pid`` for as long as this process runs.
+
+    The WebUI decides whether a task is *running* (green dot, live kill
+    button) from this file and nothing else: ``Launcher.status()`` reads
+    ``{"pid", "started_at"}`` and checks that the pid is still the same
+    process. An orchestrator that never writes it looks stopped for its whole
+    life, which is exactly how harness_evolve used to behave while a run was
+    in progress.
+
+    Written before any work starts and cleared on every exit path — normal
+    return, exception, or SIGTERM/SIGINT (the WebUI's Stop button) — so the
+    file never claims a dead pid. Signals use ``os._exit`` after clearing, so
+    an ``atexit`` handler would not be enough on its own.
+
+    Covers both entry points the server uses: ``run`` (started from the New
+    Task form) and ``resume`` (started by the supervisor to continue a run).
+    """
+    pid_file = state_dir / "orchestrator.pid"
+
+    def _on_signal(signum, _frame):
+        try:
+            clear_pid_file(pid_file)
+        except Exception:  # noqa: BLE001 - best effort while dying
+            pass
+        os._exit(143 if signum == signal.SIGTERM else 130)
+
+    previous = (signal.getsignal(signal.SIGTERM), signal.getsignal(signal.SIGINT))
+    write_pid_file(pid_file, task_id)
+    signal.signal(signal.SIGTERM, _on_signal)
+    signal.signal(signal.SIGINT, _on_signal)
+    try:
+        yield pid_file
+    finally:
+        signal.signal(signal.SIGTERM, previous[0])
+        signal.signal(signal.SIGINT, previous[1])
+        try:
+            clear_pid_file(pid_file)
+        except Exception:  # noqa: BLE001 - never mask the real failure
+            pass
+
+
+def clear_stop_request(exp_dir: Path, *, state_dir: Path | None = None) -> "list[str]":
+    """Drop a previous stop signal so a resumed round can actually run.
+
+    A run that stopped on purpose (``round_incomplete`` — e.g. because a child's
+    DKAO gate never got a clean device) leaves ``stop_requested.json`` behind,
+    and the pipeline treats it as "stop now" after the next evaluate. A human
+    asking for another round has decided the condition is over, so the signal is
+    recorded in the timeline and then removed: without this, every resume would
+    stop again immediately.
+    """
+    import time as _time
+
+    cleared: list = []
+    stop_path = Path(exp_dir) / "stop_requested.json"
+    if not stop_path.is_file():
+        return cleared
+    try:
+        data = json.loads(stop_path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        data = {}
+    try:
+        from .state import append_timeline
+        append_timeline(Path(state_dir or exp_dir), "stop_request_cleared",
+                        {"reason": (data or {}).get("reason"),
+                         "cleared_at": _time.time()})
+    except Exception:  # noqa: BLE001 - the audit row is best effort
+        pass
+    try:
+        stop_path.unlink()
+        cleared.append("stop_requested.json")
+    except OSError:
+        pass
+    return cleared
+
+
+#: The round target is a small number of judged passes; the WebUI clamps its own
+#: "target" input to this range, so the form's value is clamped the same way.
+_TARGET_MIN, _TARGET_MAX = 1, 50
+
+
+def _configured_rounds(cfg: Any) -> int:
+    """The round limit the New Task form asked for (``rounds``).
+
+    One round is one iteration: the baseline round is round 1, every later round
+    is a performance pass plus — when it wins — the generalization retake inside
+    that same iteration. ``max_iterations`` is the field's old name, still read
+    so runs created before the rename keep their limit.
+    """
+    for key in ("rounds", "max_iterations"):
+        raw = cfg.answers.get(key)
+        if raw in (None, ""):
+            continue
+        try:
+            value = int(raw)
+        except (TypeError, ValueError):
+            continue
+        return max(_TARGET_MIN, min(_TARGET_MAX, value))
+    return _TARGET_MAX // 5          # 10, the form's default
+
+
+def _seed_rounds_from_form(cfg: Any, workspace_dir: Path) -> int:
+    """Pin the form's round limit where the loop and the WebUI both read it."""
+    rounds = _configured_rounds(cfg)
+    write_target_iterations(Path(workspace_dir), rounds)
+    print(f"[harness_evolve] rounds = {rounds} "
+          f"(from the form's Rounds field)", flush=True)
+    return rounds
+
+
+def run_with_requirements(
+    requirements_path: Path,
+    *,
+    state_dir: Path,
+    workspace_dir: Path,
+) -> int:
+    state_dir.mkdir(parents=True, exist_ok=True)
+    task_id = _task_id_of(requirements_path, fallback=state_dir.parent.name)
+    with _orchestrator_pid(state_dir, task_id):
+        cfg = load_experiment_config(requirements_path, state_dir, workspace_dir)
+        _seed_rounds_from_form(cfg, workspace_dir)
+        report = run_experiment(cfg)
+    print(f"[harness_evolve] done -> {report}")
+    return 0
+
+
+def _clear_operator_stop(workspace_dir: Path) -> "list[str]":
+    """Drop the operator-stop marker (see the WebUI's Stop button).
+
+    ``operator_stop.json`` is what keeps a deliberately stopped run from being
+    quietly restarted by the supervisor, the janitor sweep or a plain resume.
+    A resume is a human asking for the run to continue, so the marker must not
+    outlive it -- otherwise the run we are starting now would be treated as
+    already stopped. Idempotent: no marker, nothing to report.
+    """
+    path = Path(workspace_dir) / "operator_stop.json"
+    try:
+        path.unlink()
+        return ["operator_stop.json"]
+    except OSError:
+        return []
+
+
+def _task_id_of(requirements_path: Path, *, fallback: str) -> str:
+    """The task id to stamp in the PID file, read from requirements.json.
+
+    Best effort: the PID file must exist before anything else can fail, so a
+    missing or malformed requirements file falls back to the directory name
+    rather than raising.
+    """
+    try:
+        data = json.loads(requirements_path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return fallback
+    task_id = str(data.get("task_id") or "").strip()
+    return task_id or fallback
+
+
+def resume_command(args: argparse.Namespace) -> int:
+    """Continue an existing experiment with one (or more) further rounds.
+
+    The already-measured champion and previous round come back from disk, so
+    only the new generation of the harness is evaluated.
+    """
+    state_dir = Path(args.state_dir)
+    workspace_dir = Path(args.workspace_dir)
+    req_path = (Path(args.requirements) if args.requirements
+                else state_dir / "requirements.json")
+    if not req_path.is_file():
+        print(f"[resume] requirements not found: {req_path}", flush=True)
+        return 2
+    cfg = load_experiment_config(req_path, state_dir, workspace_dir)
+    # A resumed run is just as "live" to the WebUI as a freshly started one, so
+    # it owns the same PID file (see _orchestrator_pid).
+    with _orchestrator_pid(state_dir, cfg.task_id):
+        def _another_driver() -> Optional[int]:
+            """PID of a *different* live orchestrator on this experiment.
+
+            Two drivers on one experiment is how a round ends up measured twice
+            on the same four cards: the janitor's resume used to open iteration
+            N+1 while iteration N was still being evaluated. Nothing here counts
+            ourselves — the pid file this process just wrote is its own.
+            """
+            try:
+                from .supervisor import task_process_alive
+                return task_process_alive(state_dir)
+            except Exception:  # noqa: BLE001 - never block on our own bookkeeping
+                return None
+
+        done = completed_iterations(cfg.exp_dir)
+        if done < 1:
+            # Two cases: the task is already running (started by the server) and
+            # we should simply stay out of the way, or nothing is running and we
+            # are the ones to start round 1.
+            live = _another_driver()
+            if live:
+                print(f"[resume] task already running (pid {live}); nothing to do",
+                      flush=True)
+                return 0
+            print("[resume] no completed iteration yet; starting iteration 1",
+                  flush=True)
+            # ``target`` is optional on this path (programmatic callers build a
+            # bare Namespace), so read it defensively.
+            first_target = getattr(args, "target", None)
+            if first_target is not None:
+                write_target_iterations(cfg.exp_dir, max(1, int(first_target)))
+            elif read_target_iterations(cfg.exp_dir) is None:
+                _seed_rounds_from_form(cfg, cfg.exp_dir)
+            report = run_experiment(cfg, start_iteration=1, iterations_to_run=1)
+            print(f"[harness_evolve] started -> {report}")
+            return 0
+        # Same guard on the "there is history" path: a completed pass on disk
+        # says nothing about whether the run that produced it is *still going*.
+        running = _another_driver()
+        if running:
+            print(f"[resume] task already running (pid {running}); nothing to do",
+                  flush=True)
+            return 0
+        explicit_target = getattr(args, "target", None)
+        if explicit_target is not None:
+            write_target_iterations(cfg.exp_dir, max(1, int(explicit_target)))
+            print(f"[resume] target iterations set to {int(explicit_target)}",
+                  flush=True)
+        elif read_target_iterations(cfg.exp_dir) is None:
+            # No target was ever stored (a run created before the form's
+            # Iterations field was wired): fall back to it instead of running
+            # unbounded until the round budget trips.
+            _seed_rounds_from_form(cfg, cfg.exp_dir)
+        start = done + 1
+        rounds = max(1, int(args.iterations or 1))
+        cleared = clear_stop_request(cfg.exp_dir, state_dir=state_dir)
+        if cleared:
+            print(f"[resume] cleared previous stop signal: {cleared}", flush=True)
+        cleared += _clear_operator_stop(workspace_dir)
+        print(f"[resume] completed iterations: {done}; running "
+              f"iteration {start}..{start + rounds - 1}", flush=True)
+        champion_iteration = getattr(args, "champion_iteration", None)
+        if champion_iteration:
+            print(f"[resume] comparison baseline pinned to iteration "
+                  f"{int(champion_iteration)}", flush=True)
+        report = run_experiment(cfg, start_iteration=start,
+                                iterations_to_run=rounds,
+                                champion_iteration=champion_iteration)
+    print(f"[harness_evolve] resumed -> {report}")
+    return 0
+
+
+def _dcu_harness_default() -> Path:
+    return (
+        Path(__file__).resolve().parents[2]
+        / "dcu_kernel_auto_opt" / "harness_default"
+    )
+
+
+def _promotion_context(args: argparse.Namespace):
+    """Load the experiment config for the promotion commands."""
+    state_dir = Path(args.state_dir)
+    workspace_dir = Path(args.workspace_dir)
+    req_path = state_dir / "requirements.json"
+    if not req_path.is_file():
+        raise SystemExit(f"[promote] requirements not found: {req_path}")
+    return load_experiment_config(req_path, state_dir, workspace_dir)
+
+
+def pending_command(args: argparse.Namespace) -> int:
+    from .promotion import promoted_harness, read_pending
+
+    cfg = _promotion_context(args)
+    pending = read_pending(cfg.exp_dir)
+    promoted = promoted_harness(cfg.exp_dir)
+    print(json.dumps({"pending": pending, "promoted": promoted},
+                     ensure_ascii=False, indent=2))
+    return 0 if pending else 1
+
+
+def approve_command(args: argparse.Namespace) -> int:
+    from .promotion import approve_promotion
+
+    cfg = _promotion_context(args)
+    result = approve_promotion(cfg, cfg.exp_dir, approved_by=str(args.by))
+    if not result.get("ok"):
+        print(f"[approve] refused: {result.get('errors')}")
+        return 2
+    print(f"[approve] promoted {result['version']} -> "
+          f"{result['record']['production_dir']}")
+    print(f"[approve] variants updated: {result['variant_update']['updated']}")
+    print(f"[approve] previous version backed up at "
+          f"{result['record']['previous_backup']}")
+    return 0
+
+
+def deny_command(args: argparse.Namespace) -> int:
+    from .promotion import deny_promotion
+
+    cfg = _promotion_context(args)
+    result = deny_promotion(cfg, cfg.exp_dir, reason=str(args.reason))
+    if not result.get("ok"):
+        print(f"[deny] refused: {result.get('errors')}")
+        return 2
+    print(f"[deny] discarded candidate {result['denied']}; "
+          "resume the experiment to keep iterating")
+    return 0
+
+
+def rollback_command(args: argparse.Namespace) -> int:
+    from .promotion import rollback_harness
+
+    cfg = _promotion_context(args)
+    result = rollback_harness(cfg.exp_dir, to=str(args.to))
+    if not result.get("ok"):
+        print(f"[rollback] refused: {result.get('errors')}")
+        return 2
+    print(f"[rollback] restored from {result['restored_from']}")
+    print(f"[rollback] replaced tree backed up at "
+          f"{result['record']['replaced_backup']}")
+    return 0
+
+
+def publish_command(args: argparse.Namespace) -> int:
+    exp = Path(args.exp_dir).resolve()
+    if not (exp / "config_snapshot.json").is_file():
+        print(f"[publish] not a harness_evolve experiment: {exp}", flush=True)
+        return 2
+    did = False
+    if args.harness or args.all:
+        target = Path(args.harness_target or _dcu_harness_default())
+        result = publish_harness(exp, target=target, reason=args.reason)
+        print(f"[publish][harness] ok={result['ok']} "
+              f"{result.get('published')} {result.get('errors') or ''}",
+              flush=True)
+        if not result.get("ok"):
+            return 3
+        did = True
+    if args.kernels or args.all:
+        pool = Path(args.pool or _DEFAULT_POOL)
+        result = publish_kernels(exp, pool_path=pool, reason=args.reason)
+        print(f"[publish][kernels] ok={result['ok']} "
+              f"published={result.get('published')} "
+              f"updated={result.get('updated', [])} "
+              f"{result.get('errors') or ''}", flush=True)
+        if not result.get("ok"):
+            return 3
+        did = True
+    if not did:
+        gate = check_publishable(exp)
+        print(f"[publish] check only: ok={gate.get('ok')} "
+              f"iteration={gate.get('iteration')} "
+              f"verdict={gate.get('verdict')} errors={gate.get('errors') or []}",
+              flush=True)
+        return 0 if gate.get("ok") else 3
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="metainfer-harness-evolve")
+    sub = parser.add_subparsers(dest="command", required=True)
+    run_p = sub.add_parser("run")
+    run_p.add_argument("requirements", type=Path)
+    run_p.add_argument("--state-dir", type=Path, required=True)
+    run_p.add_argument("--workspace-dir", type=Path, required=True)
+
+    res = sub.add_parser("resume", help="continue an existing experiment")
+    res.add_argument("--state-dir", type=Path, required=True)
+    res.add_argument("--workspace-dir", type=Path, required=True)
+    res.add_argument("--requirements", type=Path, default=None,
+                     help="defaults to /requirements.json")
+    res.add_argument("--iterations", type=int, default=1,
+                     help="how many further rounds to run (default 1)")
+    res.add_argument("--target", type=int, default=None,
+                     help="total iterations to reach (dynamic; re-read every round)")
+    res.add_argument("--champion-iteration", type=int, default=None,
+                     help="pin the comparison baseline to a finished iteration")
+    res.set_defaults(func=resume_command)
+
+    pub = sub.add_parser("publish", help="explicit publish gate")
+    pub.add_argument("--exp-dir", type=Path, required=True)
+    pub.add_argument("--harness", action="store_true",
+                     help="publish champion harness to production default seed")
+    pub.add_argument("--harness-target", type=Path, default=None,
+                     help="override harness target (default: dcu harness_default)")
+    pub.add_argument("--kernels", action="store_true",
+                     help="publish new best-known kernels into the registered pool")
+    pub.add_argument("--all", action="store_true")
+    pub.add_argument("--pool", type=Path, default=None)
+    pub.add_argument("--reason", type=str, default="explicit publish")
+    pub.set_defaults(func=publish_command)
+
+    # Human decision on a candidate harness that passed both gates. Nothing is
+    # written to production DKAO until `approve` runs.
+    appr = sub.add_parser("approve", help="approve a pending harness promotion")
+    appr.add_argument("--state-dir", type=Path, required=True)
+    appr.add_argument("--workspace-dir", type=Path, required=True)
+    appr.add_argument("--by", type=str, default="operator")
+    appr.set_defaults(func=approve_command)
+
+    deny = sub.add_parser("deny", help="reject a pending harness promotion")
+    deny.add_argument("--state-dir", type=Path, required=True)
+    deny.add_argument("--workspace-dir", type=Path, required=True)
+    deny.add_argument("--reason", type=str, default="operator")
+    deny.set_defaults(func=deny_command)
+
+    roll = sub.add_parser("rollback-harness",
+                          help="restore a previously promoted harness")
+    roll.add_argument("--state-dir", type=Path, required=True)
+    roll.add_argument("--workspace-dir", type=Path, required=True)
+    roll.add_argument("--to", type=str, default="last-good",
+                      help="version label (h-2), last-good, or seed")
+    roll.set_defaults(func=rollback_command)
+
+    pend = sub.add_parser("pending", help="show the pending promotion, if any")
+    pend.add_argument("--state-dir", type=Path, required=True)
+    pend.add_argument("--workspace-dir", type=Path, required=True)
+    pend.set_defaults(func=pending_command)
+
+    args = parser.parse_args(argv)
+    if args.command == "run":
+        return run_with_requirements(
+            args.requirements,
+            state_dir=args.state_dir,
+            workspace_dir=args.workspace_dir,
+        )
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/metainfer/tasks/harness_evolve/orchestrator/config.py b/metainfer/tasks/harness_evolve/orchestrator/config.py
new file mode 100644
index 00000000..aac57434
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/config.py
@@ -0,0 +1,170 @@
+"""Suite and experiment configuration for harness_evolve."""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List
+
+import yaml
+
+from .rounds import normalized_per_gate
+
+#: default instance used by dry-run when the suite spec is empty
+DRY_DEFAULT_INSTANCES = [
+    {
+        "id": "hy3_tp8_qkv_m4096_w4",
+        "shape": {"model": "hy3", "tp_size": 8, "operator": "qkv_proj",
+                  "M": 4096, "N": 1280, "K": 4096},
+        "warm_start_from": "best_known",
+        "budget_rounds": 4,
+        "pass_median_us_le": 500.0,
+    },
+    {
+        "id": "glm52_tp4_o_m16_w4",
+        "shape": {"model": "glm52", "tp_size": 4, "operator": "o_proj",
+                  "M": 16, "N": 6144, "K": 2048},
+        "warm_start_from": "best_known",
+        "budget_rounds": 4,
+        "pass_median_us_le": 60.0,
+    },
+]
+
+
+@dataclass
+class InstanceSpec:
+    id: str
+    shape: Dict[str, Any]
+    budget_rounds: int = 4
+    warm_start_from: str = "best_known"
+    pass_median_us_le: float = 0.0
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "InstanceSpec":
+        return cls(
+            id=str(d["id"]),
+            shape=dict(d.get("shape") or {}),
+            budget_rounds=int(d.get("budget_rounds") or 4),
+            warm_start_from=str(d.get("warm_start_from") or "best_known"),
+            pass_median_us_le=float(d.get("pass_median_us_le") or 0.0),
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "id": self.id,
+            "shape": self.shape,
+            "budget_rounds": self.budget_rounds,
+            "warm_start_from": self.warm_start_from,
+            "pass_median_us_le": self.pass_median_us_le,
+        }
+
+
+@dataclass
+class ExperimentConfig:
+    task_id: str
+    state_dir: Path
+    workspace_dir: Path
+    suite: List[InstanceSpec]
+    max_iterations: int = 3
+    harness_source: str = "dcu_default"   # dir path or "dcu_default"
+    execution_mode: str = "dry-run"       # dry-run | dkao-cli
+    agent_framework: str = "dsh"
+    pool_source: str = ""                 # "" = legacy suite mode; builtin|path = pool v2
+    per_round_budget: int = 0             # 0 = unlimited (legacy); pool budget cap
+    answers: Dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def exp_dir(self) -> Path:
+        return self.workspace_dir
+
+    def dump_snapshot(self, path: Path) -> None:
+        path.write_text(json.dumps({
+            "schema_version": 1,
+            "task_id": self.task_id,
+            "max_iterations": self.max_iterations,
+            "harness_source": self.harness_source,
+            "execution_mode": self.execution_mode,
+            "pool_source": self.pool_source,
+            "per_round_budget": self.per_round_budget,
+            "suite": [s.to_dict() for s in self.suite],
+        }, ensure_ascii=False, indent=2), encoding="utf-8")
+
+    @property
+    def pool_mode(self) -> bool:
+        return bool(self.pool_source) and self.pool_source != "suite"
+
+
+def load_experiment_config(
+    requirements_path: Path, state_dir: Path, workspace_dir: Path
+) -> ExperimentConfig:
+    req = json.loads(requirements_path.read_text(encoding="utf-8"))
+    # The shell/server frontend flattens form answers onto the requirements
+    # top level (same contract dcu_kernel_auto_opt consumes via its
+    # ``_answers`` fallback). Accept either an ``answers`` dict (CLI/tests) or
+    # the flattened top level so a WebUI-launched experiment honors the form.
+    req_answers = req.get("answers")
+    answers = req_answers if isinstance(req_answers, dict) else req
+    task_id = str(req.get("task_id", "task"))
+
+    suite_yaml = str(answers.get("suite_yaml") or "")
+    raw_instances: List[Dict[str, Any]] = []
+    if suite_yaml.strip():
+        data = yaml.safe_load(suite_yaml)
+        if isinstance(data, dict):
+            raw_instances = list(data.get("instances") or [])
+        elif isinstance(data, list):
+            raw_instances = list(data)
+    if not raw_instances:
+        raw_instances = DRY_DEFAULT_INSTANCES
+
+    return ExperimentConfig(
+        task_id=task_id,
+        state_dir=Path(state_dir),
+        workspace_dir=Path(workspace_dir),
+        suite=[InstanceSpec.from_dict(d) for d in raw_instances],
+        max_iterations=int(answers.get("max_iterations") or 3),
+        harness_source=str(answers.get("harness_source") or "dcu_default"),
+        execution_mode=str(answers.get("execution_mode") or "dry-run"),
+        agent_framework=str(answers.get("agent_framework") or "dsh"),
+        pool_source=str(answers.get("pool_source") or ""),
+        per_round_budget=_normalized_budget(answers.get("per_round_budget")),
+        answers=dict(answers),
+    )
+
+
+def _normalized_budget(value: Any) -> int:
+    """This round's question count: always a multiple of 4 (FLOW.md §1.1).
+
+    The machine has four HCUs, so a round is sized in whole waves; a request of
+    5 or 7 questions would leave devices idle in the second wave, so it is
+    rounded up (never down) here, at the single point where the form/CLI answer
+    enters the run. ``0``/empty keeps the legacy meaning "no cap, use the whole
+    pool", which is why this cannot simply be ``normalized_per_gate`` (that one
+    maps 0 to the 4-question default).
+    """
+    try:
+        count = int(value)
+    except (TypeError, ValueError):
+        return 0
+    if count <= 0:
+        return 0
+    return normalized_per_gate(count)
+
+
+def harness_seed_dir(harness_source: str) -> Path:
+    """Resolve the harness seed directory.
+
+    ``dcu_default`` maps to the dcu_kernel_auto_opt plugin's harness_default/.
+    An explicit path (absolute or relative to the experiment workspace) is used
+    as-is. Later this becomes the "champion workspace" an AHE experiment forks.
+    """
+    if harness_source and harness_source != "dcu_default":
+        p = Path(harness_source).expanduser()
+        return p.resolve()
+    dcu_plugin = (
+        Path(__file__).resolve().parents[2]
+        / "dcu_kernel_auto_opt" / "harness_default"
+    )
+    return dcu_plugin.resolve()
diff --git a/metainfer/tasks/harness_evolve/orchestrator/decision_engine.py b/metainfer/tasks/harness_evolve/orchestrator/decision_engine.py
new file mode 100644
index 00000000..ce4d4b73
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/decision_engine.py
@@ -0,0 +1,588 @@
+"""System-level harness decision engine.
+
+Separates the performance gate (did the agent system get better?) from the
+mechanism gate (did it improve for the reason declared in change_manifest?).
+True performance is primary: an unexplained but reproducible win may be
+promoted; a mechanism miss alone never turns a real stable win into REJECT.
+
+Verdicts:
+  BASELINE               first measured version, becomes champion
+  PROMOTE                clear performance win + mechanism evidence hit
+  CONFIRM_REQUIRED       clear performance win, mechanism absent/missed; repeat
+  PROMOTE_UNEXPLAINED    repeated clear win but mechanism still missed/unknown
+  NO_SIGNAL              inside noise; champion unchanged, candidate archived
+  SPECIALIZE             target family wins but held-out regresses; routeable
+  REJECT                 correctness or primary-suite significant regression
+  REJECT_OVERFIT         main improves but held-out regresses and not routeable
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, List, Mapping, Optional
+
+
+@dataclass(frozen=True)
+class DecisionPolicy:
+    noise_percent: float = 2.0
+    min_independent_wins: int = 2
+    max_loss_ratio: float = 0.10
+    max_heldout_losses: int = 0
+    #: The performance gate is a *question-count* rule: the candidate passes
+    #: when at least this fraction of the counted questions (probe questions
+    #: excluded) are measurably faster than the champion. With the default 0.75
+    #: that is "3 of 4 questions beat the current variant", regardless of how
+    #: far behind the remaining question is.
+    min_win_ratio: float = 0.75
+    #: One hard red line inside that rule: a single question that is slower by
+    #: more than this percentage — or that fails correctness — still rejects the
+    #: candidate, so "three small wins plus one collapse" cannot be promoted.
+    hard_regression_percent: float = 30.0
+    #: A round whose paired set is mostly unmeasurable (crashed children,
+    #: suppressed measurements) cannot support a verdict: it is recorded as
+    #: INCONCLUSIVE instead of REJECT.
+    min_valid_pair_ratio: float = 0.5
+    #: How strictly the mechanism gate is enforced when a real performance win
+    #: has no (or only partial) mechanism evidence:
+    #:   ``lenient`` (default) — unverified/partial evidence may still promote
+    #:       (recorded as PROMOTE_UNEXPLAINED); only a *contradiction* forces
+    #:       one confirmation round;
+    #:   ``strict`` — anything short of a verified hit requires confirmation;
+    #:   ``off`` — skip the mechanism gate entirely.
+    mechanism_policy: str = "lenient"
+
+
+def _median(r: Mapping[str, Any]) -> Optional[float]:
+    try:
+        v = float(r.get("median_us"))
+        return v if v > 0 else None
+    except (TypeError, ValueError):
+        return None
+
+
+_MISSING_STATUSES = {"missing_report", "exception", "timeout", "partial_worker_result"}
+
+
+def _correctness_hard_fail(r: Mapping[str, Any]) -> bool:
+    return r.get("correctness_ok") is False or r.get("correctness_passed") is False
+
+
+def _measurement_missing(r: Mapping[str, Any]) -> bool:
+    """True when there is no usable measurement at all.
+
+    A crashed child (``missing_report``) or a timeout is an *environment*
+    failure: treating it as HARD_FAIL/LOSS turns infrastructure trouble into a
+    fake harness regression, which is how a run can be "rejected" five times
+    in a row while nothing was actually measured.
+    """
+    if r.get("median_us") in (None, ""):
+        return True
+    return str(r.get("status") or "").strip().lower() in _MISSING_STATUSES
+
+
+def compare_instances(
+    champion: Mapping[str, Mapping[str, Any]],
+    candidate: Mapping[str, Mapping[str, Any]],
+    *,
+    noise_percent: float,
+) -> Dict[str, Any]:
+    """Paired comparison on ids present in both result sets."""
+    ids = sorted(set(champion) & set(candidate))
+    rows: Dict[str, Any] = {}
+    counts = {"WIN": 0, "LOSS": 0, "NO_SIGNAL": 0, "HARD_FAIL": 0,
+              "INVALID": 0}
+    pass_before = pass_after = 0
+
+    for iid in ids:
+        before, after = champion[iid], candidate[iid]
+        pass_before += int(before.get("passed") is True)
+        pass_after += int(after.get("passed") is True)
+        b, a = _median(before), _median(after)
+        if _measurement_missing(after):
+            verdict, delta = "INVALID", None
+        elif _correctness_hard_fail(after):
+            verdict, delta = "HARD_FAIL", None
+        elif b is None or a is None:
+            verdict, delta = "INVALID", None
+        else:
+            delta = (b - a) / b * 100.0  # positive = candidate faster
+            if delta > noise_percent:
+                verdict = "WIN"
+            elif delta < -noise_percent:
+                verdict = "LOSS"
+            else:
+                verdict = "NO_SIGNAL"
+        counts[verdict] += 1
+        rows[iid] = {
+            "champion_median_us": b,
+            "candidate_median_us": a,
+            "delta_percent": round(delta, 4) if delta is not None else None,
+            "verdict": verdict,
+            "champion_passed": before.get("passed") is True,
+            "candidate_passed": after.get("passed") is True,
+        }
+
+    return {
+        "overlap_ids": ids,
+        "per_instance": rows,
+        "counts": counts,
+        "pass_before": pass_before,
+        "pass_after": pass_after,
+        "delta_pass": pass_after - pass_before,
+    }
+
+
+def _manifest_changes(manifest: Optional[Mapping[str, Any]]) -> List[Mapping[str, Any]]:
+    if not manifest:
+        return []
+    return [x for x in manifest.get("changes", []) if isinstance(x, Mapping)]
+
+
+def _mechanism_status(
+    manifest: Optional[Mapping[str, Any]],
+    mechanism_checks: Optional[Mapping[str, bool]],
+) -> Dict[str, Any]:
+    changes = _manifest_changes(manifest)
+    declared = [c for c in changes if c.get("mechanism_signature")]
+    if not declared:
+        return {"status": "NOT_DECLARED", "hits": [], "misses": []}
+    if not mechanism_checks:
+        return {
+            "status": "NOT_CHECKED",
+            "hits": [],
+            "misses": [str(c.get("id", "unknown")) for c in declared],
+            "unobserved": [],
+            "contradicted": [],
+        }
+    hits, misses, unobserved, contradicted = [], [], [], []
+    for c in declared:
+        cid = str(c.get("id", "unknown"))
+        raw = mechanism_checks.get(cid)
+        if raw is True or raw == "hit":
+            hits.append(cid)
+        elif raw == "partial":
+            hits.append(cid)          # lenient: partial evidence counts
+        elif raw == "unobserved":
+            unobserved.append(cid)
+        elif raw == "contradicted":
+            contradicted.append(cid)
+            misses.append(cid)
+        else:
+            misses.append(cid)
+    if contradicted:
+        status = "MISS"
+    elif hits and not misses:
+        status = "HIT"
+    elif hits and misses:
+        status = "PARTIAL"
+    elif unobserved and not misses:
+        status = "UNOBSERVED"
+    else:
+        status = "MISS"
+    return {"status": status, "hits": hits, "misses": misses,
+            "unobserved": unobserved, "contradicted": contradicted}
+
+
+def _collapse(row: Mapping[str, Any], limit_percent: float) -> bool:
+    """One question that regressed past the red line (or broke correctness)."""
+    delta = row.get("delta_percent")
+    try:
+        if delta is not None and float(delta) < -abs(float(limit_percent)):
+            return True
+    except (TypeError, ValueError):
+        pass
+    return row.get("verdict") == "HARD_FAIL"
+
+
+def _routeable(manifest: Optional[Mapping[str, Any]]) -> bool:
+    for c in _manifest_changes(manifest):
+        scope = c.get("scope") or {}
+        if isinstance(scope, Mapping) and scope.get("routeable") is True:
+            return True
+    return False
+
+
+def decide(
+    champion: Optional[Mapping[str, Mapping[str, Any]]],
+    candidate: Mapping[str, Mapping[str, Any]],
+    *,
+    manifest: Optional[Mapping[str, Any]] = None,
+    heldout_champion: Optional[Mapping[str, Mapping[str, Any]]] = None,
+    heldout_candidate: Optional[Mapping[str, Mapping[str, Any]]] = None,
+    mechanism_checks: Optional[Mapping[str, bool]] = None,
+    confirmation_reproduced: bool = False,
+    purposes: Optional[Mapping[str, str]] = None,
+    policy: DecisionPolicy = DecisionPolicy(),
+) -> Dict[str, Any]:
+    """Return a complete, auditable system-level decision record.
+
+    ``purposes`` labels each question (regression / repair / probe). Probe
+    questions are *generalisation* evidence only: they never enter the
+    pass-rate or win/loss pair that decides promotion, but a probe that
+    regresses blocks promotion, because it means the change does not
+    generalise beyond the family it was tuned on.
+    """
+    if champion is None:
+        return {
+            "verdict": "BASELINE",
+            "performance_gate": {"status": "BASELINE"},
+            "mechanism_gate": {"status": "NOT_APPLICABLE"},
+            "heldout_gate": {"status": "NOT_APPLICABLE"},
+            "action": "SET_CHAMPION",
+        }
+
+    purposes = purposes or {}
+    probe_set = {i for i, p in purposes.items() if p == "probe"}
+    main_ids = [i for i in candidate if i not in probe_set]
+    if not main_ids:
+        main_ids = list(candidate)
+    champion_main = {k: v for k, v in champion.items() if k in main_ids}
+    candidate_main = {k: v for k, v in candidate.items() if k in main_ids}
+
+    perf = compare_instances(champion_main, candidate_main,
+                             noise_percent=policy.noise_percent)
+    counts = perf["counts"]
+    #: correctness/missing-measurement failures are always fatal: a comparison
+    #: missing one side is not evidence that the candidate is faster.
+    hard_fail = counts["HARD_FAIL"] > 0 or counts["INVALID"] > 0
+    measured = counts["WIN"] + counts["LOSS"] + counts["NO_SIGNAL"]
+    # Question-count performance gate: >= min_win_ratio of the measured
+    # questions must be measurably faster than the champion. NO_SIGNAL (inside
+    # the noise band) is not a win, so it counts against the ratio.
+    win_ratio = (counts["WIN"] / measured) if measured else 0.0
+    enough_wins = measured > 0 and win_ratio >= float(policy.min_win_ratio)
+    # Hard red line: one question that fell off a cliff still rejects.
+    collapses = [
+        iid for iid, row in (perf.get("per_instance") or {}).items()
+        if _collapse(row, policy.hard_regression_percent)
+    ]
+    clear_improvement = bool(enough_wins and not hard_fail and not collapses)
+
+    held = None
+    heldout_regression = False
+    if heldout_champion is not None and heldout_candidate is not None:
+        held = compare_instances(
+            heldout_champion, heldout_candidate,
+            noise_percent=policy.noise_percent,
+        )
+        heldout_regression = (
+            held["counts"]["HARD_FAIL"] > 0
+            or held["counts"]["LOSS"] > policy.max_heldout_losses
+            or held["delta_pass"] < 0
+        )
+
+    mechanism = _mechanism_status(manifest, mechanism_checks)
+
+    #: Counted questions that produced a usable pairing (WIN/LOSS/NO_SIGNAL).
+    measurable = counts["WIN"] + counts["LOSS"] + counts["NO_SIGNAL"]
+    #: All counted questions inside the noise band: "no measurable difference"
+    #: rather than "worse". Drives both the gate status and the verdict.
+    indistinguishable = bool(measured) and counts["WIN"] == 0 and counts["LOSS"] == 0
+    paired_total = len(perf["overlap_ids"])
+    coverage = (measured / paired_total) if paired_total else 1.0
+    inconclusive = bool(paired_total) and coverage < float(
+        policy.min_valid_pair_ratio)
+
+    generalization: Dict[str, Any] = {"status": "NOT_RUN", "comparison": None,
+                                      "ids": sorted(probe_set)}
+    if probe_set:
+        probe_champion = {k: v for k, v in champion.items() if k in probe_set}
+        probe_candidate = {k: v for k, v in candidate.items() if k in probe_set}
+        if set(probe_champion) & set(probe_candidate):
+            probe_perf = compare_instances(
+                probe_champion, probe_candidate,
+                noise_percent=policy.noise_percent)
+            probe_counts = probe_perf["counts"]
+            generalization = {
+                "status": ("FAIL" if (probe_counts["HARD_FAIL"] > 0
+                                      or probe_counts["LOSS"] > 0) else "PASS"),
+                "comparison": probe_perf,
+                "ids": sorted(probe_set),
+            }
+
+    if inconclusive:
+        # Not a result: keep the champion, do not count it as a failure.
+        verdict, action = "INCONCLUSIVE", "KEEP_CHAMPION_ARCHIVE_CANDIDATE"
+    elif hard_fail or collapses:
+        verdict, action = "REJECT", "RESTORE_CHAMPION"
+    elif clear_improvement and heldout_regression:
+        if _routeable(manifest):
+            verdict, action = "SPECIALIZE", "ARCHIVE_SPECIALIZED"
+        else:
+            verdict, action = "REJECT_OVERFIT", "RESTORE_CHAMPION"
+    elif clear_improvement and generalization["status"] == "FAIL":
+        # A win inside the tuned families that breaks an untouched family is
+        # over-fitting, not progress: require one confirmation round.
+        if confirmation_reproduced:
+            verdict, action = "PROMOTE_UNEXPLAINED", "SET_CHAMPION"
+        else:
+            verdict, action = "CONFIRM_REQUIRED", "ARCHIVE_AND_CONFIRM"
+    elif clear_improvement:
+        mstatus = mechanism["status"]
+        policy_mode = str(policy.mechanism_policy or "lenient").lower()
+        if mstatus in {"HIT", "PARTIAL"}:
+            verdict, action = "PROMOTE", "SET_CHAMPION"
+        elif policy_mode == "off":
+            verdict, action = "PROMOTE_UNEXPLAINED", "SET_CHAMPION"
+        elif policy_mode == "lenient" and mstatus in {"UNOBSERVED",
+                                                      "NOT_CHECKED",
+                                                      "NOT_DECLARED"}:
+            # Performance is primary: an unverified win may still take the
+            # champion slot, but it is recorded as unexplained.
+            verdict, action = "PROMOTE_UNEXPLAINED", "SET_CHAMPION"
+        elif confirmation_reproduced:
+            verdict, action = "PROMOTE_UNEXPLAINED", "SET_CHAMPION"
+        else:
+            verdict, action = "CONFIRM_REQUIRED", "ARCHIVE_AND_CONFIRM"
+    else:
+        # Below the win ratio. Distinguish "we measured no difference" (an
+        # honest no-signal: keep the champion, try again) from "the candidate
+        # is measurably worse" (a real reject that counts toward the circuit
+        # breaker). A round where every counted question sat inside the noise
+        # band is the former; anything with LOSSes is the latter.
+        if indistinguishable:
+            verdict, action = "NO_SIGNAL", "KEEP_CHAMPION_ARCHIVE_CANDIDATE"
+        else:
+            verdict, action = "REJECT", "RESTORE_CHAMPION"
+
+    performance_status = (
+        "INCONCLUSIVE" if inconclusive else
+        "FAIL" if hard_fail or collapses else
+        "PASS" if clear_improvement else
+        "NO_SIGNAL" if indistinguishable else "FAIL"
+    )
+    heldout_status = (
+        "NOT_RUN" if held is None else
+        "FAIL" if heldout_regression else "PASS"
+    )
+    record = {
+        "verdict": verdict,
+        "action": action,
+        "policy": {
+            "noise_percent": policy.noise_percent,
+            "min_win_ratio": policy.min_win_ratio,
+            "hard_regression_percent": policy.hard_regression_percent,
+            "min_independent_wins": policy.min_independent_wins,
+            "max_loss_ratio": policy.max_loss_ratio,
+            "max_heldout_losses": policy.max_heldout_losses,
+            "mechanism_policy": policy.mechanism_policy,
+        },
+        "performance_gate": {"status": performance_status,
+                             "measurable_pairs": measurable,
+                             "paired_total": paired_total,
+                             "coverage": round(coverage, 3),
+                             # the question-count rule, in one line:
+                             "win_ratio": round(win_ratio, 3),
+                             "min_win_ratio": policy.min_win_ratio,
+                             "wins_needed": _wins_needed(measured, policy),
+                             "collapses": sorted(collapses),
+                             **perf},
+        "mechanism_gate": mechanism,
+        "generalization_gate": generalization,
+        "purpose_breakdown": {
+            iid: {"purpose": purposes.get(iid, "regression"),
+                  "counted": iid not in probe_set}
+            for iid in sorted(set(champion) | set(candidate))
+        },
+        "paired_ids": sorted(set(champion_main) & set(candidate_main)),
+        "probe_ids": sorted(probe_set),
+        "heldout_gate": {
+            "status": heldout_status,
+            "comparison": held,
+        },
+        "reason": _reason(verdict, perf, mechanism, heldout_regression,
+                          win_ratio=win_ratio,
+                          collapses=collapses,
+                          policy=policy),
+    }
+    if generalization["status"] == "FAIL":
+        tail = ("promoted after reproduction" if verdict == "PROMOTE_UNEXPLAINED"
+                else ("rejected" if action == "RESTORE_CHAMPION"
+                      else "repeat once"))
+        record["reason"] = (
+            "performance improved inside the tuned families but an untouched "
+            f"family regressed (over-fitting): {tail}"
+        )
+    return record
+
+
+def _needed_wins(measured: int, ratio: float) -> int:
+    """Ceiling of ``measured * ratio``: 4 operators at 75% need 3 wins."""
+    if measured <= 0:
+        return 0
+    import math
+    return max(1, min(measured, int(math.ceil(measured * float(ratio)))))
+
+
+def _wins_needed(measured: int, policy: DecisionPolicy) -> int:
+    """How many of the counted questions must win (round-level policy)."""
+    return _needed_wins(measured, policy.min_win_ratio)
+
+
+#: The two harness gates, on the operator-count rule the operator specified.
+#: Both compare the candidate against the *current variant* of each operator and
+#: both refuse a candidate that lets any single operator fall below a floor, so
+#: "many small wins plus one collapse" cannot pass.
+PERFORMANCE_GATE_WIN_RATIO = 0.75      # 4 operators -> at least 3 must win
+PERFORMANCE_GATE_FLOOR_PERCENT = 0.80  # every non-winning operator >= 80% of the variant
+GENERALIZATION_GATE_WIN_RATIO = 0.50   # 4 operators -> at least 2 must win
+GENERALIZATION_GATE_FLOOR_PERCENT = 0.80
+
+#: How a tie (a difference inside the comparison's noise band) is scored.
+#: ``"loss"``: only a strict win counts — the performance gate's rule, where
+#: fresh operators must genuinely come out ahead.
+#: ``"neutral"``: a tie is evidence of *holding the line*, not of failure — the
+#: generalization gate's rule. The retake re-measures the paper on the same
+#: harness, and DKAO warm-starts from the kernel the baseline round just put in
+#: the pool, so re-finding that same kernel is the expected outcome; scoring it
+#: as a loss would demand that a search policy out-optimize its own best result
+#: every time. A real regression (below the floor) still fails the gate.
+TIE_AS_LOSS = "loss"
+TIE_AS_NEUTRAL = "neutral"
+
+
+def gate_verdict(reference: Mapping[str, Mapping[str, Any]],
+                 candidate: Mapping[str, Mapping[str, Any]],
+                 *,
+                 kind: str = "performance",
+                 noise_percent: float = 2.0,
+                 win_ratio: Optional[float] = None,
+                 floor_percent: Optional[float] = None,
+                 ties: Optional[str] = None,
+                 ) -> Dict[str, Any]:
+    """Judge one round of a gate on the operator-count rule.
+
+    ``kind`` is ``"performance"`` (>= 75% of operators must beat their variant)
+    or ``"generalization"`` (>= 50%). In both cases every operator that did not
+    win must still be at least ``floor_percent`` of its variant, and an operator
+    that produced no usable measurement fails the gate outright rather than
+    being counted as a loss.
+
+    ``ties`` decides how a difference inside ``noise_percent`` is scored; it
+    defaults to the gate's own rule (performance: a tie is not a win;
+    generalization: a tie holds the line). Every operator that is genuinely
+    slower than its variant (and every one below the floor) blocks the gate.
+    """
+    if win_ratio is None:
+        win_ratio = (PERFORMANCE_GATE_WIN_RATIO if kind == "performance"
+                     else GENERALIZATION_GATE_WIN_RATIO)
+    if floor_percent is None:
+        floor_percent = (PERFORMANCE_GATE_FLOOR_PERCENT if kind == "performance"
+                         else GENERALIZATION_GATE_FLOOR_PERCENT)
+    if ties is None:
+        ties = TIE_AS_LOSS if kind == "performance" else TIE_AS_NEUTRAL
+    comparison = compare_instances(reference, candidate,
+                                   noise_percent=noise_percent)
+    rows = comparison.get("per_instance") or {}
+    judged = ["WIN", "LOSS", "NO_SIGNAL"]
+    all_judged = [iid for iid, row in rows.items()
+                  if row.get("verdict") in judged]
+    ties_ids = [iid for iid in all_judged
+                if rows[iid]["verdict"] == "NO_SIGNAL"]
+    losses = [iid for iid in all_judged if rows[iid]["verdict"] == "LOSS"]
+    # Neutral ties leave the denominator: they are neither a win to count nor a
+    # loss to answer for (the floor is what catches a real regression).
+    counted = ([iid for iid in all_judged if iid not in set(ties_ids)]
+               if ties == TIE_AS_NEUTRAL else all_judged)
+    wins = [iid for iid in counted if rows[iid]["verdict"] == "WIN"]
+    unmeasured = sorted(
+        iid for iid, row in rows.items()
+        if row.get("verdict") in {"INVALID", "HARD_FAIL"})
+    too_slow = sorted(
+        iid for iid in all_judged
+        if _below_floor(rows[iid], floor_percent))
+    measured = len(counted)
+    ratio = (len(wins) / measured) if measured else 0.0
+    needed = _needed_wins(measured, win_ratio)
+    reasons: List[str] = []
+    if unmeasured:
+        reasons.append("no usable measurement: " + ", ".join(unmeasured))
+    if too_slow:
+        reasons.append(
+            f"below {floor_percent:.0%} of the variant: " + ", ".join(too_slow))
+    if measured and len(wins) < needed:
+        reasons.append(f"{len(wins)}/{measured} operators beat the variant "
+                       f"(needs {needed} = {win_ratio:.0%})")
+    all_tied = bool(all_judged) and len(ties_ids) == len(all_judged)
+    if all_tied:
+        # Every operator tied. Scored as holding the line (the generalization
+        # gate's rule) that is a pass on the floor requirement alone; scored as
+        # a loss it means the candidate proved nothing.
+        if ties == TIE_AS_LOSS:
+            reasons.append("no operator beat its variant (all tied)")
+    elif not measured and not unmeasured:
+        reasons.append("nothing comparable to the variant")
+    record = {
+        "kind": kind,
+        "status": "PASS" if not reasons else "FAIL",
+        "win_ratio": round(ratio, 3),
+        "required_win_ratio": win_ratio,
+        "floor_percent_of_variant": floor_percent,
+        "tie_policy": ties,
+        "operators": sorted(rows),
+        "wins": sorted(wins),
+        "losses": sorted(losses),
+        "ties": sorted(ties_ids),
+        "counted": measured,
+        "wins_needed": needed,
+        "unmeasured": unmeasured,
+        "below_floor": too_slow,
+        "reasons": reasons,
+        "comparison": comparison,
+    }
+    return record
+
+
+def _below_floor(row: Mapping[str, Any], floor_percent: float) -> bool:
+    """True when a measured operator is worse than ``floor_percent`` of its variant."""
+    delta = row.get("delta_percent")
+    try:
+        if delta is None:
+            return False
+        return float(delta) < (float(floor_percent) - 1.0) * 100.0
+    except (TypeError, ValueError):
+        return False
+
+
+
+def _reason(verdict: str, perf: Mapping[str, Any], mechanism: Mapping[str, Any],
+            heldout_regression: bool, *, win_ratio: float = 0.0,
+            collapses: Optional[List[str]] = None,
+            policy: Optional[DecisionPolicy] = None) -> str:
+    policy = policy or DecisionPolicy()
+    counts = perf.get("counts") or {}
+    measured = (int(counts.get("WIN", 0)) + int(counts.get("LOSS", 0))
+                + int(counts.get("NO_SIGNAL", 0)))
+    summary = (f"{counts.get('WIN', 0)}/{measured} counted questions faster "
+               f"than the champion (needs {_wins_needed(measured, policy)}, "
+               f"{policy.min_win_ratio:.0%})")
+    if verdict == "INCONCLUSIVE":
+        return ("too few usable measurements this round (crashed children or "
+                "suppressed readings); no verdict")
+    if verdict == "PROMOTE":
+        return f"clear performance improvement ({summary}); declared mechanism verified"
+    if verdict == "PROMOTE_UNEXPLAINED":
+        return (f"reproduced performance improvement ({summary}); mechanism "
+                "remains unverified")
+    if verdict == "CONFIRM_REQUIRED":
+        return (f"performance improved ({summary}) but mechanism is "
+                "absent/unverified; repeat once")
+    if verdict == "SPECIALIZE":
+        return ("target suite improved but held-out regressed; preserve as "
+                "routeable variant")
+    if verdict == "REJECT_OVERFIT":
+        return "target suite improved but held-out regression exceeded boundary"
+    if verdict == "REJECT":
+        if collapses:
+            return (f"hard red line crossed on {', '.join(sorted(collapses))} "
+                    f"(> {policy.hard_regression_percent:.0f}% slower or "
+                    f"incorrect); {summary}")
+        if int(counts.get("HARD_FAIL", 0)) or int(counts.get("INVALID", 0)):
+            return ("correctness failure or an unusable measurement in the "
+                    "comparison; not a result we can trust")
+        return (f"performance gate failed: {summary} "
+                f"(needed {policy.min_win_ratio:.0%} of counted questions)")
+    if verdict == "NO_SIGNAL":
+        return (f"candidate is inside the noise band ({summary}); keep "
+                "champion and archive candidate")
+    return "first measured version establishes the champion baseline"
diff --git a/metainfer/tasks/harness_evolve/orchestrator/evolve.py b/metainfer/tasks/harness_evolve/orchestrator/evolve.py
new file mode 100644
index 00000000..5c11bd0f
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/evolve.py
@@ -0,0 +1,272 @@
+"""Single-role Evolve adapters.
+
+The real AgentEvolver modifies the harness workspace and, in the same turn,
+produces BOTH:
+  - _ahe_change_manifest.json (what changed / system prediction / mechanism)
+  - _ahe_round_plan.json      (which registered questions to run next + why)
+The pipeline moves these to runs/iteration_NNN/evolve/ after validation.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+#: Model id the DSH Evolve Agent runs on by default. The form's `evolve_model`
+#: field picks a label (see form.yaml) and the pipeline passes the id down to
+#: ``bridge/dsh/dsh_agent.py --model``; this constant is the fallback used by
+#: callers that were given no answer at all.
+#: ``deepseek/deepseek-flash`` is the 4.1 Flash revision (form label
+#: ``deepseek-flash-4.1``). ``deepseek/deepseek-v4-flash-0731`` is the older
+#: pinned build, kept selectable for reproducing earlier experiments.
+DEFAULT_EVOLVE_MODEL = "deepseek/deepseek-flash"
+LEGACY_EVOLVE_MODEL = "deepseek/deepseek-v4-flash-0731"
+
+
+def _empty_manifest(iteration: int) -> Dict[str, Any]:
+    return {
+        "schema_version": 1,
+        "iteration": iteration,
+        "changes": [],
+        "verification": {
+            "status": "pending", "fixes": [], "regressions": [],
+            "false_predictions": [], "verdict": None,
+        },
+        "note": "dry-run evolve: no automatic changes",
+    }
+
+
+class Evolver:
+    name = "base"
+
+    def apply(self, workspace: Path, iteration: int,
+              query_context: Dict[str, Any]) -> Dict[str, Any]:
+        raise NotImplementedError
+
+
+class DryRunEvolver(Evolver):
+    name = "dry-run"
+
+    def apply(self, workspace, iteration, query_context):
+        fake = os.environ.get("HARNESS_EVOLVE_FAKE_CHANGE", "") in {"1", "true"}
+        manifest = _empty_manifest(iteration)
+        changed = False
+        summary_lines = [f"# Evolve summary — iteration {iteration} (dry-run)", "",
+                         "No automatic harness change (dry-run evolver)."]
+        if fake:
+            manifest = {
+                "schema_version": 1, "iteration": iteration,
+                "changes": [{
+                    "id": f"chg-{iteration}", "type": "new",
+                    "component": "gates",
+                    "description": "synthetic change for attribution testing",
+                    "files": ["gates.yaml"],
+                    "failure_evidence": ["dry-run synthetic failure"],
+                    "root_cause": "synthetic", "targeted_fix": "synthetic gates edit",
+                    "failure_pattern": "synthetic", "predicted_fixes": [],
+                    "risk_tasks": [],
+                    "scope": {"expected_improve": ["synthetic"],
+                              "expected_unchanged": [], "at_risk": [],
+                              "routeable": False},
+                    "system_prediction": {"direction": "improve",
+                                          "expected_effect": ["pass-rate increase"]},
+                    "mechanism_signature": ["synthetic mechanism fires"],
+                }],
+                "verification": {"status": "pending", "fixes": [],
+                                 "regressions": [], "false_predictions": [],
+                                 "verdict": None},
+            }
+            changed = True
+            summary_lines.append("Emitted one synthetic change (fake mode).")
+        return {"summary": "\n".join(summary_lines), "manifest": manifest,
+                "changed": changed, "round_plan": None}
+
+
+class AgentEvolver(Evolver):
+    """Real single-role DSH evolve agent."""
+
+    name = "agent"
+    MANIFEST_NAME = "_ahe_change_manifest.json"
+    PLAN_NAME = "_ahe_round_plan.json"
+
+    def __init__(self, *, model: str = DEFAULT_EVOLVE_MODEL,
+                 timeout_seconds: float = 3600):
+        self.model = model
+        self.timeout_seconds = timeout_seconds
+
+    def apply(self, workspace, iteration, query_context):
+        manifest_path = workspace / self.MANIFEST_NAME
+        plan_path = workspace / self.PLAN_NAME
+        manifest_path.unlink(missing_ok=True)
+        plan_path.unlink(missing_ok=True)
+        prompt = self._prompt(workspace, iteration, query_context)
+        wrapper = (
+            Path(__file__).resolve().parents[2] / "dcu_kernel_auto_opt"
+            / "bridge" / "dsh" / "dsh_agent.py"
+        )
+        cmd = [
+            str(wrapper), "-p", "--output-format", "stream-json",
+            "--input-format", "text", "--verbose",
+            "--permission-mode", "bypassPermissions",
+            "--add-dir", str(workspace), "--model", self.model,
+            "--effort", "max", "--tools", "Read,Glob,Grep,Write",
+            "--disallowedTools", "Bash,Skill,WebFetch,WebSearch",
+        ]
+        try:
+            proc = subprocess.run(
+                cmd, input=prompt, text=True, capture_output=True,
+                timeout=self.timeout_seconds, env=os.environ.copy(),
+            )
+        except subprocess.TimeoutExpired:
+            # A slow/failed evolve turn must not abort the whole experiment:
+            # the iteration's evaluation value is already recorded, so fall
+            # back to a no-change placeholder and let the pipeline finish
+            # (champion kept, report + finished still written).
+            return _evolve_failed(
+                iteration,
+                f"evolve agent timed out after {self.timeout_seconds:.0f}s",
+            )
+        except Exception as exc:  # noqa: BLE001 - never crash the outer loop
+            return _evolve_failed(
+                iteration, f"evolve agent crashed: {exc!r}"
+            )
+        if proc.returncode != 0:
+            return _evolve_failed(
+                iteration,
+                f"evolve agent failed rc={proc.returncode}: "
+                f"{(proc.stderr or '')[-800:]}",
+            )
+        manifest = _read_json(manifest_path)
+        plan = _read_json(plan_path)
+        missing = [
+            name for name, value in
+            ((self.MANIFEST_NAME, manifest), (self.PLAN_NAME, plan))
+            if value is None
+        ]
+        if missing:
+            return _evolve_failed(
+                iteration,
+                "evolve agent did not write " + ", ".join(missing),
+            )
+        manifest_path.unlink(missing_ok=True)
+        plan_path.unlink(missing_ok=True)
+        summary = _last_result(proc.stdout) or "Evolve Agent completed."
+        return {"summary": summary, "manifest": manifest, "round_plan": plan,
+                "changed": bool(manifest.get("changes"))}
+
+    def _prompt(self, workspace: Path, iteration: int,
+                context: Dict[str, Any]) -> str:
+        pool_ids = list(context.get("pool_ids") or [])
+        selected = list(context.get("selected") or [])
+        return f"""You are the SINGLE AHE Evolve Agent for DKAO INT8 W8A8 GEMM.
+
+Iteration: {iteration}
+Harness workspace (the ONLY directory you may modify): {workspace}
+Analysis overview: {context.get('overview')}
+Previous results/diff are embedded below.
+
+Current selected questions: {json.dumps(selected)}
+Registered pool ids (you may select ONLY these): {json.dumps(pool_ids)}
+Results: {json.dumps(context.get('results') or {}, ensure_ascii=False)[:16000]}
+Diff: {json.dumps(context.get('diff') or {}, ensure_ascii=False)}
+Decision: {json.dumps(context.get('decision') or {}, ensure_ascii=False)[:8000]}
+
+Your two responsibilities in this ONE role:
+1) Evidence-driven harness evolution: inspect workspace, make at most ONE logical
+   change to planner/gates/prompt/skill/memory; do not touch model config. Every
+   change needs failure_evidence, root_cause, targeted_fix, system_prediction,
+   scope (expected_improve/unchanged/at_risk/routeable), mechanism_signature.
+2) Choose the NEXT round exam from registered pool ids. Give strategy
+   explore|exploit|balanced, rationale and selected ids. The machine assigns
+   each question a role (2 regression anchors / 1 repair / 1 cross-family
+   probe) and enforces pool membership, budget, family quota and >=60%
+   overlap: your picks are used where they fit a role and replaced where they
+   do not (the correction is recorded in the round rationale). Weights are
+   ignored — priority comes from the role, and the probe never counts toward
+   the pass-rate decision.
+
+Mechanism evidence (machine-checkable, not prose): every change MUST also
+carry a structured `mechanism_check` describing how a robot can verify that the
+mechanism actually ran. Allowed kinds:
+  {{"kind": "planner_plan_ids", "expect_any": ["memory_layout", "epilogue_fusion"],
+   "expect_order": ["architecture_explore", "memory_layout"]}}
+  {{"kind": "harness_revision"}}   # candidate revision must be the one children load
+  {{"kind": "gate_events", "expect_any": ["gate_blocked"]}}
+  {{"kind": "gate_values", "expect_any": ["round_acceptance_improvement_percent"]}}
+Use planner_plan_ids when the change edits planner_policy/planner_catalog,
+gate_values when it edits gates.yaml (the runtime reads that file now), and
+gate_events for gating behaviour that shows up in the child timeline. A change without mechanism_check stays
+"unobserved" (allowed, but it cannot be verified as effective).
+
+Write strict JSON files INSIDE workspace:
+- {manifest_path_for_prompt(workspace, self.MANIFEST_NAME)}
+  {{"schema_version":1,"iteration":{iteration},"changes":[...],
+    "verification":{{"status":"pending","fixes":[],"regressions":[],
+    "false_predictions":[],"verdict":null}}}}
+- {manifest_path_for_prompt(workspace, self.PLAN_NAME)}
+  {{"iteration":{iteration},"strategy":"balanced","rationale":"...",
+    "selected":[...],"weights":{{}}}}
+
+Use file tools only. Do not compile, benchmark, run DKAO, edit outside workspace,
+or invent pool ids. Finish after rereading both JSON files from disk.
+"""
+
+
+def manifest_path_for_prompt(workspace: Path, name: str) -> str:
+    return str(workspace / name)
+
+
+def _evolve_failed(iteration: int, reason: str) -> Dict[str, Any]:
+    """No-change evolve placeholder used when the evolve turn fails.
+
+    The outer loop keeps the iteration's evaluation result (champion
+    unchanged) and still writes evolve artifacts + report, instead of
+    crashing mid-experiment and losing the whole round.
+    """
+    manifest = {
+        "schema_version": 1,
+        "iteration": iteration,
+        "changes": [],
+        "verification": {
+            "status": "evolve_failed", "fixes": [], "regressions": [],
+            "false_predictions": [], "verdict": None,
+        },
+        "note": reason,
+    }
+    summary = (
+        f"# Evolve summary — iteration {iteration} (failed)\n\n"
+        f"{reason}\n\nNo harness change was applied; "
+        "the champion is unchanged."
+    )
+    return {"summary": summary, "manifest": manifest, "round_plan": None,
+            "changed": False}
+
+
+def _read_json(path: Path) -> Optional[Dict[str, Any]]:
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    return value if isinstance(value, dict) else None
+
+
+def _last_result(stdout: str) -> str:
+    result = ""
+    for line in stdout.splitlines():
+        try:
+            event = json.loads(line)
+        except ValueError:
+            continue
+        if event.get("type") == "result":
+            result = str(event.get("result") or "")
+    return result
+
+
+def build_evolver(mode: str, *, model: str = DEFAULT_EVOLVE_MODEL,
+                  timeout_seconds: float = 3600) -> Evolver:
+    if mode == "agent":
+        return AgentEvolver(model=model, timeout_seconds=timeout_seconds)
+    return DryRunEvolver()
diff --git a/metainfer/tasks/harness_evolve/orchestrator/mechanism_checks.py b/metainfer/tasks/harness_evolve/orchestrator/mechanism_checks.py
new file mode 100644
index 00000000..22976fb1
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/mechanism_checks.py
@@ -0,0 +1,352 @@
+"""Mechanism gate: check whether a harness change was actually exercised.
+
+The Evolve Agent declares, per change, *why* it should work (a
+``mechanism_signature``) and — since this module landed — how to verify that
+the mechanism really ran (a structured ``mechanism_check``). This module
+collects the candidate iteration's evidence and grades each declared change:
+
+  ``hit``          the declared mechanism was observed
+  ``partial``      planner/agent activity observed, but not the declared plan
+  ``unobserved``   nothing to observe (no record) — *not* a contradiction
+  ``contradicted`` the change demonstrably did not take effect (e.g. the
+                   candidate harness revision was never loaded, or the
+                   planner produced no plan records although the change
+                   requires it)
+
+Grades feed ``decision_engine``: under the default lenient policy a real
+performance win may promote with ``hit``/``partial``/``unobserved`` evidence
+(the latter recorded as PROMOTE_UNEXPLAINED), while a *contradiction* forces
+one confirmation round so an implausible explanation cannot silently win.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional, Tuple
+
+from .rounds import child_dir_map
+
+_PLAN_ID_KEYS = ("plan_id", "planner_plan_id", "plan")
+_GATE_EVENT_HINTS = ("gate", "guard", "threshold", "reject", "blocked")
+
+
+def _load_json(path: Path) -> Dict[str, Any]:
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return {}
+    return value if isinstance(value, dict) else {}
+
+
+def _jsonl(path: Path) -> List[Dict[str, Any]]:
+    out: List[Dict[str, Any]] = []
+    if not path.is_file():
+        return out
+    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+        try:
+            row = json.loads(line)
+        except ValueError:
+            continue
+        if isinstance(row, dict):
+            out.append(row)
+    return out
+
+
+def _candidate_gates(candidate_workspace: Optional[Path]) -> Dict[str, Any]:
+    """Gate values the candidate harness declares (its gates.yaml + wiring)."""
+    if candidate_workspace is None:
+        return {}
+    root = Path(candidate_workspace)
+    try:
+        import yaml
+        manifest = yaml.safe_load(
+            (root / "manifest.yaml").read_text(encoding="utf-8")) or {}
+        data = yaml.safe_load(
+            (root / "gates.yaml").read_text(encoding="utf-8")) or {}
+    except (OSError, ValueError):
+        return {}
+    components = manifest.get("components") or {}
+    gates_entry = components.get("gates") if isinstance(components, dict) else {}
+    return {
+        "gates": data.get("gates") if isinstance(data, dict) else None,
+        "wired": bool((gates_entry or {}).get("wired", True)),
+    }
+
+
+def collect_candidate_evidence(exp: Path, iteration: int,
+                               planner_enabled: Optional[bool] = None,
+                               candidate_workspace: Optional[Path] = None,
+                               ) -> Dict[str, Any]:
+    """Gather observable evidence from one iteration's DKAO children."""
+    exp = Path(exp)
+    children_root = exp / "children"
+    snapshot_manifest = (
+        exp / "runs" / f"iteration_{iteration:03d}" / "input" / "workspace"
+        / "manifest.yaml"
+    )
+    candidate_revision = ""
+    try:
+        import yaml
+        candidate_revision = str(
+            (yaml.safe_load(snapshot_manifest.read_text(encoding="utf-8"))
+             or {}).get("revision") or ""
+        )
+    except (OSError, ValueError):
+        candidate_revision = ""
+
+    observed_gate_snapshots: List[Dict[str, Any]] = []
+    agent_plan_ids: List[str] = []
+    planner_plan_ids: List[str] = []
+    harness_revisions: List[str] = []
+    gate_events: List[str] = []
+    children: Dict[str, Any] = {}
+    # An iteration can own several attempt directories (a first run, its
+    # environment retries, and the generalization retake): evidence is gathered
+    # from all of them, because a mechanism that fired in any attempt fired.
+    for child in (child_dir_map(children_root, iteration).values()
+                  if children_root.is_dir() else []):
+        workspace = child / "workspace"
+        child_plans: List[str] = []
+        for exp_file in (workspace / "workers").glob(
+                "*/runs/*/experiments.jsonl"):
+            for row in _jsonl(exp_file):
+                for key in _PLAN_ID_KEYS:
+                    value = row.get(key)
+                    if isinstance(value, str) and value.strip():
+                        child_plans.append(value.strip())
+                        break
+        agent_plan_ids.extend(child_plans)
+
+        # Hard evidence: the plans the DKAO planner itself selected for
+        # this worker's rounds (written by _round_strategy_text).
+        for plans_file in (workspace / "workers").glob(
+                "*/planner_plans.jsonl"):
+            for row in _jsonl(plans_file):
+                value = row.get("plan_id")
+                if isinstance(value, str) and value.strip():
+                    planner_plan_ids.append(value.strip())
+
+        scaffold = _load_json(workspace / "main" / "scaffold_manifest.json")
+        harness = scaffold.get("harness") if isinstance(scaffold, dict) else {}
+        revision = ""
+        if isinstance(harness, dict):
+            revision = str(harness.get("revision") or "")
+        if revision:
+            harness_revisions.append(revision)
+
+        gate_snapshot = _load_json(child / "state" / "gates_effective.json")
+        if gate_snapshot:
+            observed_gate_snapshots.append(gate_snapshot)
+
+        for event in _jsonl(child / "state" / "timeline.jsonl"):
+            etype = str(event.get("type") or "")
+            if any(hint in etype for hint in _GATE_EVENT_HINTS):
+                gate_events.append(etype)
+
+        children[child.name] = {
+            "agent_plan_ids": child_plans,
+            "harness_revision": revision,
+        }
+
+    return {
+        "iteration": iteration,
+        "planner_enabled": planner_enabled,
+        "candidate_harness_revision": candidate_revision,
+        "observed_harness_revisions": sorted(set(harness_revisions)),
+        # ``plan_ids`` stays the agents' self-reported ids for backwards
+        # compatibility; ``planner_plan_ids`` is the authoritative record.
+        "plan_ids": agent_plan_ids,
+        "agent_plan_ids": agent_plan_ids,
+        "planner_plan_ids": planner_plan_ids,
+        "plan_id_source": ("planner" if planner_plan_ids
+                           else ("agent" if agent_plan_ids else "none")),
+        "candidate_gates": _candidate_gates(candidate_workspace),
+        "observed_gates": observed_gate_snapshots,
+        "gate_events": sorted(set(gate_events)),
+        "children": children,
+    }
+
+
+def _grade_harness_revision(change: Mapping[str, Any],
+                            evidence: Mapping[str, Any]) -> Tuple[str, str]:
+    observed = list(evidence.get("observed_harness_revisions") or [])
+    expected = str(evidence.get("candidate_harness_revision") or "")
+    if not observed:
+        return "unobserved", "no child recorded which harness revision it loaded"
+    if expected and all(rev == expected for rev in observed):
+        return "hit", f"candidate harness revision {expected} was loaded"
+    if expected and all(rev != expected for rev in observed):
+        return ("contradicted",
+                f"children loaded {observed} but the candidate revision is "
+                f"{expected}: the change never took effect")
+    return "partial", f"harness revisions observed: {observed}"
+
+
+def _grade_plan_ids(change: Mapping[str, Any],
+                    evidence: Mapping[str, Any]) -> Tuple[str, str]:
+    check = change.get("mechanism_check") or {}
+    planner_ids = list(evidence.get("planner_plan_ids") or [])
+    agent_ids = list(evidence.get("agent_plan_ids")
+                     or evidence.get("plan_ids") or [])
+    observed = planner_ids or agent_ids
+    expect_any = [str(x) for x in (check.get("expect_any") or [])]
+    expect_order = [str(x) for x in (check.get("expect_order") or [])]
+    if not expect_any and not expect_order:
+        return "unobserved", "mechanism_check declares no expected plan ids"
+    # A disabled planner is hard counter-evidence: a planner_policy change
+    # cannot have been consumed, whatever plan ids the agents reported for
+    # themselves (those come from the workers' own proposals).
+    if evidence.get("planner_enabled") is False:
+        return ("contradicted",
+                "the planner was disabled for this iteration, so a "
+                "planner_policy change could not have taken effect"
+                + (f" (agent-reported plan ids: {sorted(set(observed))})"
+                   if observed else ""))
+    if not observed:
+        return ("unobserved",
+                "no planner plan records were produced this iteration")
+    if not planner_ids:
+        # Only the workers' own proposals mention plans; that cannot prove the
+        # planner drove the mandate, so it is soft evidence at best.
+        hits_any = [pid for pid in agent_ids
+                    if pid in expect_any or pid in expect_order]
+        if hits_any:
+            return ("partial",
+                    "plan ids match only agent self-reports (no planner "
+                    f"record): {sorted(set(hits_any))}")
+        return ("partial",
+                f"agent-reported plan ids {sorted(set(agent_ids))} do not "
+                f"match the declared ones")
+    if expect_order:
+        idx = 0
+        for pid in observed:
+            if idx < len(expect_order) and pid == expect_order[idx]:
+                idx += 1
+        if idx == len(expect_order):
+            return "hit", f"observed expected plan order {expect_order}"
+    hit_any = [pid for pid in observed if pid in expect_any]
+    if hit_any:
+        return "hit", f"observed declared plan(s) {sorted(set(hit_any))}"
+    return ("partial",
+            f"planner produced plans {sorted(set(observed))} but none of the "
+            f"declared ones {sorted(set(expect_any + expect_order))}")
+
+
+def _grade_gate_events(change: Mapping[str, Any],
+                       evidence: Mapping[str, Any]) -> Tuple[str, str]:
+    check = change.get("mechanism_check") or {}
+    expected = [str(x) for x in (check.get("expect_any") or [])]
+    observed = list(evidence.get("gate_events") or [])
+    if not expected:
+        return "unobserved", "mechanism_check declares no expected gate events"
+    hits = [event for event in observed if event in expected]
+    if hits:
+        return "hit", f"observed declared gate event(s) {hits}"
+    return "unobserved", f"declared gate events absent (observed {observed})"
+
+
+def _grade_gate_values(change: Mapping[str, Any],
+                       evidence: Mapping[str, Any]) -> Tuple[str, str]:
+    """Did children actually run with the gate values the harness declares?"""
+    candidate = evidence.get("candidate_gates") or {}
+    declared = candidate.get("gates")
+    observed = list(evidence.get("observed_gates") or [])
+    if not candidate:
+        return "unobserved", "candidate harness gates could not be read"
+    if not candidate.get("wired", True):
+        return ("contradicted",
+                "gates.yaml is not wired in the candidate manifest, so its "
+                "values cannot reach the runtime")
+    if not observed:
+        return ("unobserved",
+                "children recorded no effective-gate snapshot for this run")
+    check = change.get("mechanism_check") or {}
+    paths = [str(p) for p in (check.get("expect_any") or [])]
+    for snapshot in observed:
+        seen = snapshot.get("gates")
+        if not isinstance(seen, Mapping) or declared is None:
+            continue
+        if paths:
+            mismatched = []
+            for path in paths:
+                node_seen: Any = seen
+                node_declared: Any = declared
+                for part in path.split("."):
+                    node_seen = node_seen.get(part) if isinstance(node_seen, Mapping) else None
+                    node_declared = (node_declared.get(part)
+                                     if isinstance(node_declared, Mapping) else None)
+                if node_seen != node_declared:
+                    mismatched.append(path)
+            if mismatched:
+                return ("contradicted",
+                        f"children ran with different values for {mismatched}")
+            return "hit", f"children ran with the declared values for {paths}"
+        if dict(seen) == dict(declared):
+            return "hit", "children ran with the candidate gate values"
+        return ("contradicted",
+                "children ran with gate values that differ from the candidate")
+    return "unobserved", "no comparable gate snapshot"
+
+
+_GRADERS = {
+    "harness_revision": _grade_harness_revision,
+    "planner_plan_ids": _grade_plan_ids,
+    "planner_plan_id_sequence": _grade_plan_ids,
+    "gate_events": _grade_gate_events,
+    "gate_values": _grade_gate_values,
+}
+
+
+def evaluate_changes(
+    manifest: Optional[Mapping[str, Any]],
+    evidence: Mapping[str, Any],
+) -> Tuple[Dict[str, str], Dict[str, Any]]:
+    """Grade every declared change; returns (checks, detail)."""
+    changes = [c for c in ((manifest or {}).get("changes") or [])
+               if isinstance(c, Mapping)]
+    checks: Dict[str, str] = {}
+    detail: Dict[str, Any] = {"changes": []}
+    for change in changes:
+        cid = str(change.get("id") or "unknown")
+        check = change.get("mechanism_check")
+        declared = bool(change.get("mechanism_signature"))
+        if not declared:
+            checks[cid] = "unobserved"
+            detail["changes"].append({
+                "id": cid, "grade": "unobserved",
+                "reason": "change declares no mechanism_signature",
+            })
+            continue
+        if not isinstance(check, Mapping):
+            checks[cid] = "unobserved"
+            detail["changes"].append({
+                "id": cid, "grade": "unobserved",
+                "reason": ("no structured mechanism_check declared; "
+                           "signature kept as documentation only"),
+            })
+            continue
+        kind = str(check.get("kind") or "").strip()
+        grader = _GRADERS.get(kind)
+        if grader is None:
+            checks[cid] = "unobserved"
+            detail["changes"].append({
+                "id": cid, "grade": "unobserved",
+                "reason": f"unknown mechanism_check kind {kind!r}",
+            })
+            continue
+        grade, reason = grader(change, evidence)
+        checks[cid] = grade
+        detail["changes"].append({
+            "id": cid, "grade": grade, "kind": kind, "reason": reason,
+            "expect": {k: v for k, v in check.items() if k != "kind"},
+        })
+    detail["planner_plan_ids"] = list(evidence.get("planner_plan_ids") or [])
+    detail["agent_plan_ids"] = list(evidence.get("agent_plan_ids") or [])
+    detail["summary"] = {
+        "hit": [c for c, g in checks.items() if g == "hit"],
+        "partial": [c for c, g in checks.items() if g == "partial"],
+        "unobserved": [c for c, g in checks.items() if g == "unobserved"],
+        "contradicted": [c for c, g in checks.items() if g == "contradicted"],
+    }
+    return checks, detail
diff --git a/metainfer/tasks/harness_evolve/orchestrator/phases.py b/metainfer/tasks/harness_evolve/orchestrator/phases.py
new file mode 100644
index 00000000..8a5de058
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/phases.py
@@ -0,0 +1,40 @@
+"""harness_evolve phases (state graph) for the shell."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+PREPARE = "prepare"
+EVALUATE = "evaluate"
+ANALYZE = "analyze"
+EVOLVE = "evolve"
+REPORT = "report"
+FINISHED = "finished"
+
+_ORDER = [PREPARE, EVALUATE, ANALYZE, EVOLVE, REPORT, FINISHED]
+_LABELS = {
+    PREPARE: "Prepare",
+    EVALUATE: "Evaluate",
+    ANALYZE: "Analyze",
+    EVOLVE: "Evolve",
+    REPORT: "Report",
+    FINISHED: "Finished",
+}
+
+
+def graph_payload(current: str, **_: Any) -> Dict[str, Any]:
+    nodes = [
+        {"id": phase, "label": _LABELS[phase], "is_terminal": phase == FINISHED}
+        for phase in _ORDER
+    ]
+    edges = [
+        {"from": _ORDER[i], "to": _ORDER[i + 1], "label": "ok",
+         "outcomes": ["ok"]}
+        for i in range(len(_ORDER) - 1)
+    ]
+    return {
+        "current": current,
+        "nodes": nodes,
+        "edges": edges,
+        "terminal_nodes": [FINISHED],
+    }
diff --git a/metainfer/tasks/harness_evolve/orchestrator/pipeline.py b/metainfer/tasks/harness_evolve/orchestrator/pipeline.py
new file mode 100644
index 00000000..7fa6894d
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/pipeline.py
@@ -0,0 +1,2422 @@
+"""harness_evolve main pipeline (AHE outer loop).
+
+Iteration layout (double generation, mirrors ahe-ref semantics):
+    runs/iteration_NNN/
+      input/workspace/          snapshot of the workspace evaluated this loop
+      input/benchmark/results.json
+      input/benchmark/scored_ids.json      (pool mode: ids scored this round)
+      input/diff.json           vs previous loop's eval (overlap ids)
+      input/analysis/overview.md
+      evolve/change_manifest.json
+      evolve/round_plan.json               (pool mode: next round's question set)
+      evolve/evolve_summary.md
+      evolve/workspace/         (only when the evolver reported a change)
+
+Root artifacts: best_ever.json, iteration_scores.jsonl, report.md.
+Auto rollback: when an iteration's decision pass rate falls below best-ever, the
+best snapshot is restored over the live workspace (workspace/).
+
+Two suite modes:
+- legacy (pool_source=""): evaluate cfg.suite every round (v0 behaviour).
+- pool v2 (pool_source="builtin"|): questions come from the registered
+  pool; the single Evolve role's round_plan picks next-round sets; machine
+  guardrails enforce in-pool + >=60% overlap + budget; decisions (best-ever /
+  rollback / attribution) are computed on the overlap pairs.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+import time
+import traceback
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from .adapters.eval import build_evaluator
+from .champion_selection import (
+    Generation, best_generation, load_generations, record_generation,
+)
+from .attribution import (
+    copy_tree, diff_results, evaluate_changes, pass_rate, replace_tree, save_json,
+)
+from .config import ExperimentConfig, InstanceSpec, harness_seed_dir
+from .decision_engine import DecisionPolicy, decide
+from .evolve import DEFAULT_EVOLVE_MODEL, build_evolver
+from .phases import (
+    ANALYZE, EVALUATE, EVOLVE, FINISHED, PREPARE, REPORT,
+)
+from .pool import Pool, builtin_pool
+from .promotion import (
+    restore_promoted_harness, write_pending_promotion,
+)
+from .rounds import (
+    GENERALIZATION_ROUND_COST, OPERATORS_PER_GATE, PERFORMANCE_ROUND_COST,
+    STAGE_BASELINE, STAGE_GENERALIZATION, STAGE_PERFORMANCE,
+    baseline_round_verdict, child_dir_map, clear_pending_generalization,
+    evaluate_gate, group_state, normalized_per_gate,
+    read_pending_generalization, record_paper, resolve_child_dir, round_stage,
+    select_round_questions, write_pending_generalization,
+)
+from .variant import (
+    load_variant_table, pool_instances, record_variants, save_variant_table,
+)
+from .question_roles import select_questions
+from .round_plan import (
+    RoundPlan, enforce_guardrails, load_rotation, load_round_plan,
+    mark_selected,
+)
+from .state import (
+    append_timeline, harness_revision, loaded_harness_revision,
+    read_target_iterations, set_run, write_harness_revision,
+)
+
+BEST_EVER_FILE = "best_ever.json"
+SCORES_FILE = "iteration_scores.jsonl"
+#: This run's question plan (performance set + the fixed generalization paper).
+ROUND_QUESTIONS_FILE = "round_questions.json"
+REPORT_FILE = "report.md"
+LAST_SCORED_FILE = "last_scored_ids.json"
+
+
+def _uses_variant_rounds(cfg: ExperimentConfig) -> bool:
+    """Whether this run measures operators against the variant table.
+
+    This is the **default protocol** (operator decision, 2026-09-15): one HE
+    run chases one promotable harness, every round draws fresh random operators
+    judged against their current variant, and round 1 fixes the generalization
+    paper. Pool mode without a real pool file, and legacy suite mode, keep the
+    old path; ``round_questions: legacy`` restores it explicitly.
+    """
+    raw = str(cfg.answers.get("round_questions") or "").strip().lower()
+    if raw in {"legacy", "suite", "off", "0", "false"}:
+        return False
+    if raw in {"variant", "on", "1", "true", "yes"}:
+        return True
+    if not cfg.pool_mode:
+        return False
+    source = str(cfg.answers.get("question_pool") or cfg.pool_source or "").strip()
+    if not source or source == "builtin":
+        return False          # nothing to draw operators from: stay on the old path
+    return Path(source).expanduser().is_file()
+
+
+def _question_pool_path(cfg: ExperimentConfig) -> Optional[Path]:
+    """The pool file holding the operators a round may draw from."""
+    raw = str(cfg.answers.get("question_pool") or cfg.pool_source or "").strip()
+    if not raw or raw == "builtin":
+        return None
+    path = Path(raw).expanduser()
+    return path if path.is_file() else None
+
+
+def _round_rng(cfg: ExperimentConfig, iteration: int,
+               stage: str = "") -> "tuple[Any, str]":
+    """This round's draw RNG, plus the seed it came from.
+
+    The draw is what the round is judged on, so "why these four operators?" has
+    to be answerable later. ``round_seed`` in the form pins it explicitly; with
+    no pin the seed is derived from the task id, the iteration and the stage,
+    which is still a fresh draw per round (A_i ≠ A_{i-1}) but reproducible when
+    somebody has to re-run one pass.
+
+    Returns ``(rng, seed_label)``.
+    """
+    import random
+
+    pinned = cfg.answers.get("round_seed")
+    if pinned in (None, ""):
+        seed: Any = f"{cfg.task_id}:{iteration}:{stage or 'round'}"
+        return random.Random(str(seed)), str(seed)
+    try:
+        return random.Random(int(pinned) + int(iteration)), f"round_seed={pinned}+{iteration}"
+    except (TypeError, ValueError):
+        seed = f"{pinned}:{iteration}"
+        return random.Random(seed), seed
+
+
+def _legacy_protocol_reason(cfg: ExperimentConfig) -> str:
+    """Why this run is not on the fixed protocol, in one sentence.
+
+    ``FLOW.md`` describes the only protocol a new HE task may use (variant
+    rounds: a fixed generalization paper, two gates, rounds sized in waves of 4).
+    The older paths still exist for offline suite runs and for experiments
+    started before 2026-09-15 — but "opt-in" was *invisible*: nothing said the
+    run had left the protocol, so a pass that judges neither gate looked like any
+    other. This sentence is what the operator reads instead.
+    """
+    raw = str(cfg.answers.get("round_questions") or "").strip().lower()
+    if raw in {"legacy", "suite", "off", "0", "false"}:
+        return f"round_questions={raw!r} was requested explicitly"
+    source = str(cfg.answers.get("question_pool") or cfg.pool_source or "").strip()
+    if not source:
+        return ("no question pool is configured, so the questions come from the "
+                "inline suite and neither gate is judged")
+    if source == "builtin":
+        return "the built-in demo pool is in use instead of a measured pool file"
+    path = Path(source).expanduser()
+    if not path.is_file():
+        return f"the configured pool file does not exist: {source}"
+    return "this run is not on the variant-round protocol"
+
+
+def _announce_legacy_protocol(cfg: ExperimentConfig, exp: Path) -> None:
+    """Say out loud that this run bypasses the fixed protocol. Never raises.
+
+    Three places, because the three audiences are different: the orchestrator log
+    (whoever is watching the process), the task timeline (the run's own record,
+    and what the WebUI lists), and ``protocol.json`` (what the task page reads to
+    put a banner on screen instead of a run that merely looks quiet).
+    """
+    reason = _legacy_protocol_reason(cfg)
+    payload = {
+        "protocol": "legacy",
+        "mode": "legacy",
+        "reason": reason,
+        "round_questions": cfg.answers.get("round_questions") or None,
+        "pool_source": cfg.pool_source or None,
+        "pool_file": (str(_question_pool_path(cfg)) if _question_pool_path(cfg)
+                      else None),
+        "since": time.time(),
+        "policy": ("FLOW.md §1–§4 is the fixed protocol for new tasks: a frozen "
+                   "generalization paper, the performance/generalization gates, "
+                   "rounds sized in waves of 4. This run judges neither gate, so "
+                   "its rounds cannot promote a harness."),
+    }
+    try:
+        append_timeline(cfg.state_dir, "legacy_protocol_selected", payload)
+    except Exception:  # noqa: BLE001 - progress records are best effort
+        pass
+    try:
+        from .attribution import save_json
+
+        save_json(Path(exp) / "protocol.json", payload)
+    except Exception:  # noqa: BLE001
+        pass
+    try:
+        print(f"[harness_evolve] WARNING: running the LEGACY protocol ({reason}). "
+              f"New tasks must use a measured pool file — see FLOW.md",
+              file=sys.stderr, flush=True)
+    except OSError:
+        pass
+
+
+def _judgeable_ids(instances: Dict[str, Dict[str, Any]],
+                   table: Dict[str, Any]) -> List[str]:
+    """Pool operators that can produce a win/loss (they have a variant)."""
+    from .variant import resolve_variant
+
+    out: List[str] = []
+    for iid in instances:
+        try:
+            if resolve_variant(table or {}, iid, instances=instances) is not None:
+                out.append(str(iid))
+        except Exception:  # noqa: BLE001 - a malformed row is simply not judgeable
+            continue
+    return sorted(out)
+
+
+def _mechanism_gate_for_round(cfg: ExperimentConfig, exp: Path,
+                              iteration: int) -> Dict[str, Any]:
+    """Grade the previous round's declared mechanism against what was observed.
+
+    Round N evaluates the harness change that round N-1's evolve produced, so
+    the manifest to grade is the previous round's one. Round 1 tests the seed,
+    which declares nothing: the gate is then ``NOT_DECLARED``, which the lenient
+    policy allows.
+    """
+    manifest = None
+    if iteration > 1:
+        manifest = _load_json(
+            exp / "runs" / f"iteration_{iteration - 1:03d}" / "evolve"
+            / "change_manifest.json")
+    if not manifest:
+        return {"status": "NOT_DECLARED", "hits": [], "misses": [],
+                "reason": "no previous-round change manifest to verify"}
+    try:
+        from .mechanism_checks import (
+            collect_candidate_evidence, evaluate_changes,
+        )
+        from .decision_engine import _mechanism_status
+
+        evidence = collect_candidate_evidence(
+            exp, iteration,
+            planner_enabled=_truthy(cfg.answers.get("planner_enabled"),
+                                    default=True),
+            candidate_workspace=(exp / "runs" / f"iteration_{iteration:03d}"
+                                 / "input" / "workspace"))
+        checks, detail = evaluate_changes(manifest, evidence)
+        status = _mechanism_status(manifest, checks)
+        status["detail"] = detail
+        append_timeline(cfg.state_dir, "mechanism_gate", {
+            "iteration": iteration, "status": status.get("status"),
+            "hits": status.get("hits"), "misses": status.get("misses"),
+            "contradicted": status.get("contradicted"),
+        })
+        return status
+    except Exception as exc:  # noqa: BLE001 - the gate is advisory here
+        append_timeline(cfg.state_dir, "mechanism_check_failed",
+                        {"iteration": iteration, "error": repr(exc)})
+        return {"status": "NOT_CHECKED", "hits": [], "misses": [],
+                "error": repr(exc)}
+
+
+def _seed_workspace(cfg: ExperimentConfig) -> Path:
+    workspace = cfg.exp_dir / "workspace"
+    if not (workspace / "manifest.yaml").exists():
+        copy_tree(harness_seed_dir(cfg.harness_source), workspace)
+    if not (workspace / ".git").exists():
+        subprocess.run(["git", "init"], cwd=workspace, check=True,
+                       stdout=subprocess.DEVNULL)
+        subprocess.run(["git", "config", "user.name", "MetaInfer AHE"],
+                       cwd=workspace, check=True)
+        subprocess.run(["git", "config", "user.email", "ahe@metainfer.local"],
+                       cwd=workspace, check=True)
+        subprocess.run(["git", "add", "-A"], cwd=workspace, check=True)
+        subprocess.run(["git", "commit", "-m", "baseline: seed harness"],
+                       cwd=workspace, check=True, stdout=subprocess.DEVNULL)
+    return workspace
+
+
+def _head_revision(workspace: Path) -> str:
+    return subprocess.check_output(
+        ["git", "rev-parse", "HEAD"], cwd=workspace, text=True
+    ).strip()
+
+
+def _current_revision(workspace: Path) -> str:
+    """HEAD of the live harness workspace, or ``""`` when it is not a repo.
+
+    Used to stamp each round with *its own* revision: reading it from the
+    previous round's manifest left ``workspace_revision`` empty in every
+    generation record, which is what made the rollback machinery unable to
+    restore a precise revision.
+    """
+    try:
+        return _head_revision(workspace)
+    except (OSError, subprocess.SubprocessError):
+        return ""
+
+
+def _commit_workspace(workspace: Path, iteration: int) -> str:
+    """Commit one evolved harness generation; return the exact revision."""
+    subprocess.run(["git", "add", "-A"], cwd=workspace, check=True)
+    diff = subprocess.run(
+        ["git", "diff", "--cached", "--quiet"], cwd=workspace
+    )
+    if diff.returncode != 0:
+        subprocess.run(
+            ["git", "commit", "-m", f"evolve: iteration {iteration}"],
+            cwd=workspace, check=True, stdout=subprocess.DEVNULL,
+        )
+    return subprocess.check_output(
+        ["git", "rev-parse", "HEAD"], cwd=workspace, text=True
+    ).strip()
+
+
+def _load_json(path: Path) -> Optional[Dict[str, Any]]:
+    if not path.is_file():
+        return None
+    try:
+        return json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+
+
+def _load_jsonl(path: Path) -> List[Dict[str, Any]]:
+    out: List[Dict[str, Any]] = []
+    if not path.is_file():
+        return out
+    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+        try:
+            row = json.loads(line)
+        except ValueError:
+            continue
+        if isinstance(row, dict):
+            out.append(row)
+    return out
+
+
+def _load_ids(path: Path) -> List[str]:
+    data = _load_json(path)
+    return list(data) if isinstance(data, list) else []
+
+
+def _save_scores(cfg: ExperimentConfig, entry: Dict[str, Any]) -> None:
+    with (cfg.exp_dir / SCORES_FILE).open("a", encoding="utf-8") as fh:
+        fh.write(json.dumps(entry, ensure_ascii=False) + "\n")
+
+
+def _update_best_ever(cfg: ExperimentConfig, iteration: int, rate: float,
+                      snapshot: Path) -> Dict[str, Any]:
+    best = _load_json(cfg.exp_dir / BEST_EVER_FILE) or {}
+    if not best or rate > float(best.get("pass_rate", -1.0)):
+        best = {
+            "iteration": iteration,
+            "pass_rate": rate,
+            "snapshot_dir": str(snapshot),
+        }
+        save_json(cfg.exp_dir / BEST_EVER_FILE, best)
+    return best
+
+
+def _write_overview(input_dir: Path, iteration: int,
+                    results: Dict[str, Dict[str, Any]],
+                    diff: Optional[Dict[str, Any]]) -> None:
+    lines = [
+        f"# Analysis overview — iteration {iteration}",
+        "",
+        "| instance | passed | median_us | reason |",
+        "|---|---:|---:|---|",
+    ]
+    for inst_id, r in sorted(results.items()):
+        lines.append(
+            f"| {inst_id} | {r.get('passed')} | {r.get('median_us')} | "
+            f"{r.get('reason')} |"
+        )
+    if diff:
+        lines.append("")
+        lines.append(f"- flipped: {diff.get('flipped')}")
+        lines.append(f"- regressed: {diff.get('regressed')}")
+    analysis = input_dir / "analysis"
+    analysis.mkdir(parents=True, exist_ok=True)
+    (analysis / "overview.md").write_text("\n".join(lines) + "\n",
+                                          encoding="utf-8")
+
+
+def _variant_instance_spec(entry: Dict[str, Any], *, budget_rounds: int) -> InstanceSpec:
+    """Build an :class:`InstanceSpec` straight from a pool entry (no Pool object)."""
+    contract = dict(entry.get("contract") or {})
+    variant = entry.get("best_known_us")
+    baseline = entry.get("baseline_us")
+    threshold = variant if variant is not None else (baseline or 0.0)
+    return InstanceSpec.from_dict({
+        "id": str(contract.get("id") or entry.get("id") or ""),
+        "shape": {
+            "model": str(contract.get("model") or ""),
+            "tp_size": int(contract.get("tp_size") or 8),
+            "operator": str(contract.get("operator") or ""),
+            "M": int(contract.get("M") or 0),
+            "N": int(contract.get("N") or 0),
+            "K": int(contract.get("K") or 0),
+        },
+        "budget_rounds": int(budget_rounds),
+        "warm_start_from": "best_known",
+        "pass_median_us_le": float(threshold or 0.0),
+    })
+
+
+def _pool_instance_spec(inst: Any, budget_rounds: int = 4) -> InstanceSpec:
+    def _g(name: str, default: Any = "") -> Any:
+        value = getattr(inst, name, default)
+        return default if value is None else value
+
+    return InstanceSpec.from_dict({
+        "id": inst.id,
+        "shape": {
+            "model": _g("model"),
+            "tp_size": _g("tp_size", 8),
+            "operator": _g("operator"),
+            "M": _g("M", 0),
+            "N": _g("N", 0),
+            "K": _g("K", 0),
+        },
+        "budget_rounds": budget_rounds,
+        "warm_start_from": "best_known",
+        # Dry evaluator uses this to synthesize median; choose the registered
+        # best/target scale so pool.auto_pass still yields first-fail/rest-pass.
+        "pass_median_us_le": (
+            getattr(inst, "best_known_us", None)
+            or 0.5 * float(getattr(inst, "baseline_us", 1.0))
+        ),
+    })
+
+
+def _harness_validate_error(workspace: Path) -> Optional[str]:
+    """Malformed harness component (bad YAML) must never be benchmarked.
+
+    Returns the first offending file's error text, or None when valid.
+    """
+    import yaml as _yaml
+    for path in sorted(workspace.rglob("*.yaml")) + sorted(workspace.rglob("*.yml")):
+        try:
+            _yaml.safe_load(path.read_text(encoding="utf-8"))
+        except Exception as exc:  # noqa: BLE001
+            return f"{path.name}: {exc}"
+    return None
+
+
+def _wired_scope_violations(before: Path, after: Path) -> List[Dict[str, str]]:
+    """Components a harness change touched although the manifest says unwired.
+
+    A component with ``wired: false`` is not read by the runtime, so changing
+    it cannot show any effect — the round would be spent for nothing. The
+    manifest is therefore load-bearing: it declares which components apply.
+    """
+    import yaml as _yaml
+
+    before = Path(before)
+    after = Path(after)
+    manifest_path = after / "manifest.yaml"
+    if not manifest_path.is_file():
+        return []
+    try:
+        manifest = _yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {}
+    except (OSError, ValueError):
+        return []
+    components = manifest.get("components")
+    if not isinstance(components, dict):
+        return []
+    out: List[Dict[str, str]] = []
+    for name, entry in components.items():
+        if not isinstance(entry, dict) or entry.get("wired", True):
+            continue
+        relative = str(entry.get("path") or "")
+        if not relative:
+            continue
+        old_file = before / relative
+        new_file = after / relative
+        try:
+            old_text = old_file.read_text(encoding="utf-8") if old_file.is_file() else None
+            new_text = new_file.read_text(encoding="utf-8") if new_file.is_file() else None
+        except OSError:
+            continue
+        if old_text != new_text:
+            out.append({"component": str(name), "path": relative})
+    return out
+
+
+def _load_pool(cfg: ExperimentConfig) -> Optional[Pool]:
+    if not cfg.pool_mode:
+        return None
+    if cfg.pool_source == "builtin":
+        if cfg.execution_mode == "dkao-cli":
+            raise ValueError(
+                "dkao-cli requires a measured registered pool YAML; "
+                "builtin contains dry-run fixture baselines only"
+            )
+        return builtin_pool()
+    return Pool.from_yaml(Path(cfg.pool_source).expanduser())
+
+
+def _record_state(cfg: ExperimentConfig, phase: str,
+                  iteration: Optional[int] = None,
+                  event: Optional[str] = None,
+                  **payload: Any) -> None:
+    """Mirror the outer-loop position into run.json + timeline (WebUI)."""
+    fields: Dict[str, Any] = {"current_phase": phase}
+    if phase == FINISHED:
+        fields["finished"] = True
+        fields["final_status"] = "success"
+    if iteration is not None:
+        fields["current_iteration"] = iteration
+    set_run(cfg.state_dir, **fields)
+    if event:
+        append_timeline(cfg.state_dir, event,
+                        {"iteration": iteration, **payload})
+
+
+def completed_iterations(exp: Path) -> int:
+    """Highest iteration number that already has *end-of-pass* artifacts.
+
+    "Done" is a judged pass: the variant protocol leaves ``baseline_round.json``
+    (round 1), ``performance_gate.json`` / ``generalization_gate.json``, or a
+    ``decision.json`` behind — reading only ``results.json`` made this return 0
+    for a run that had just finished three real rounds, and ``resume`` then
+    refused to continue (it believes nothing has been measured yet).
+
+    ``input/stage.json`` is deliberately **not** in the list: it is written when
+    a pass *starts* (it records which stage this pass is). Counting it made a
+    round that had just begun indistinguishable from a finished one, so the
+    resume the sweep launches computed ``start = done + 1`` and opened the next
+    round in parallel with the one being measured.
+    """
+    runs = exp / "runs"
+    if not runs.is_dir():
+        return 0
+    done = 0
+    for d in sorted(runs.glob("iteration_*")):
+        if not d.is_dir():
+            continue
+        input_dir = d / "input"
+        judged = any(
+            (input_dir / name).is_file()
+            for name in ("decision.json", "baseline_round.json",
+                         "performance_gate.json", "generalization_gate.json")
+        ) or (input_dir / "benchmark" / "results.json").is_file()
+        if not judged:
+            continue
+        try:
+            done = max(done, int(d.name.split("_")[-1]))
+        except ValueError:
+            continue
+    return done
+
+
+def _iteration_results(exp: Path, iteration: int) -> Optional[Dict[str, Any]]:
+    """Benchmark results recorded for one iteration (None when absent)."""
+    data = _load_json(
+        exp / "runs" / f"iteration_{iteration:03d}"
+        / "input" / "benchmark" / "results.json"
+    )
+    if not data:
+        return None
+    results = data.get("results")
+    return dict(results) if isinstance(results, dict) else None
+
+
+def _record_champion_generation(cfg: ExperimentConfig, exp: Path, *,
+                                iteration: int, decision: Dict[str, Any],
+                                decision_rate: Any, snapshot: Path,
+                                live_workspace: Path,
+                                results: Dict[str, Any],
+                                current: Dict[str, Any]) -> Dict[str, Any]:
+    """Record this generation and update the champion. **Never raises.**
+
+    Which generation the champion points at is what every later round compares
+    against, so losing this bookkeeping silently is not acceptable. It *used* to
+    be swallowed into a timeline row and nothing else (26 ``champion_record_failed
+    / AttributeError`` rows in one run on 2026-09-11), which left the run
+    comparing against a champion chain nobody could vouch for.
+
+    A failure is now loud (stderr, i.e. the orchestrator log), carries its
+    traceback, and is followed by a rebuild of ``best_ever.json`` from the
+    generations already on disk so the chain survives the error.
+
+    Returns ``{"best_ever", "recorded", "recovered", ...}``; the caller adopts
+    ``best_ever``.
+    """
+    best = dict(current or {})
+    outcome: Dict[str, Any] = {"iteration": iteration, "recorded": False,
+                               "recovered": False, "best_ever": best}
+    try:
+        record_generation(exp, Generation(
+            iteration=iteration,
+            verdict=str(decision.get("verdict") or ""),
+            results=dict(results),
+            snapshot_dir=str(snapshot),
+            # This generation's own revision (not the previous round's
+            # manifest field, which was always empty here).
+            workspace_revision=_current_revision(live_workspace),
+            pass_count=sum(1 for r in results.values()
+                           if isinstance(r, dict) and r.get("passed") is True),
+        ))
+        pinned = _load_json(exp / "champion_override.json") or {}
+        pinned_iter = int(pinned.get("champion_iteration") or 0)
+        measured_now = None
+        if pinned_iter:
+            # An explicit baseline (resume --champion-iteration) stays in force
+            # until it is changed: automation must not silently override a
+            # human's comparison choice.
+            pinned_results = _iteration_results(exp, pinned_iter)
+            best = {
+                "iteration": pinned_iter,
+                "pass_rate": (decision_rate
+                              if pinned_iter == iteration else None),
+                "snapshot_dir": str(
+                    exp / "runs" / f"iteration_{pinned_iter:03d}"
+                    / "input" / "workspace"),
+                "workspace_revision": _current_revision(live_workspace),
+                "verdict": "MANUAL_BASELINE",
+            }
+            save_json(exp / BEST_EVER_FILE, best)
+            append_timeline(cfg.state_dir, "champion_pinned_held", {
+                "iteration": iteration, "champion_iteration": pinned_iter,
+                "pinned_results_present": bool(pinned_results),
+            })
+        else:
+            measured_now = best_generation(load_generations(exp))
+        if measured_now is not None and measured_now.results:
+            best = {
+                "iteration": measured_now.iteration,
+                "pass_rate": (decision_rate
+                              if measured_now.iteration == iteration else None),
+                "snapshot_dir": measured_now.snapshot_dir or str(snapshot),
+                "workspace_revision": measured_now.workspace_revision,
+                "verdict": ("BEST_MEASURED"
+                            if measured_now.iteration != iteration
+                            else str(decision.get("verdict") or "")),
+            }
+            save_json(exp / BEST_EVER_FILE, best)
+            append_timeline(cfg.state_dir, "champion_best_measured", {
+                "iteration": iteration,
+                "champion_iteration": measured_now.iteration,
+                "verdict": best["verdict"],
+            })
+        outcome.update({"recorded": True, "best_ever": best})
+        return outcome
+    except Exception as exc:  # noqa: BLE001 - never break the loop
+        append_timeline(cfg.state_dir, "champion_record_failed", {
+            "iteration": iteration, "error": repr(exc),
+            "traceback": traceback.format_exc()[-1500:],
+        })
+        try:
+            print(f"[harness_evolve] champion bookkeeping failed at iteration "
+                  f"{iteration}: {exc!r} — rebuilding best_ever from the "
+                  f"generations on disk", file=sys.stderr, flush=True)
+        except OSError:
+            pass
+        outcome["error"] = repr(exc)
+        recovered = None
+        try:
+            recovered = best_generation(load_generations(exp))
+        except Exception as exc2:  # noqa: BLE001 - the fallback is best effort
+            append_timeline(cfg.state_dir, "champion_rebuild_failed",
+                            {"iteration": iteration, "error": repr(exc2)})
+        if recovered is not None and recovered.results:
+            best = {
+                "iteration": recovered.iteration,
+                "pass_rate": (decision_rate
+                              if recovered.iteration == iteration else None),
+                "snapshot_dir": recovered.snapshot_dir or str(snapshot),
+                "workspace_revision": recovered.workspace_revision,
+                "verdict": "REBUILT_AFTER_ERROR",
+            }
+            try:
+                save_json(exp / BEST_EVER_FILE, best)
+                outcome["recovered"] = True
+                append_timeline(cfg.state_dir, "champion_rebuilt", {
+                    "iteration": iteration,
+                    "champion_iteration": recovered.iteration,
+                })
+            except OSError:
+                pass
+        outcome["best_ever"] = best
+        return outcome
+
+
+def _truthy(value: Any, default: bool = True) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, bool):
+        return value
+    return str(value).strip().lower() not in {"0", "false", "no", "off", ""}
+
+
+def _promote_iteration(cfg: ExperimentConfig, exp: Path, iteration: int,
+                       results: Dict[str, Any],
+                       instances: List[InstanceSpec],
+                       *, freeze: bool = False,
+                       ) -> Optional[Dict[str, Any]]:
+    """Promote this iteration's better-than-variant kernels into DKAO.
+
+    A question that beats the existing variant for its shape by at least
+    ``promote_min_improvement_percent`` (default 3%) and did not fail
+    correctness becomes the new shared variant — i.e. DKAO gains a better
+    operator. Writes are backed up by the variant store and both the outcome
+    table and an append-only log are recorded for audit/rollback.
+    """
+    if not _truthy(cfg.answers.get("promote_kernels"), default=True):
+        return None
+    children = exp / "children"
+    if not children.is_dir():
+        return None
+    try:
+        from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.variant_promote import (  # noqa: E501
+            model_label_for, promote_variant,
+        )
+    except Exception as exc:  # noqa: BLE001 - promotion must never break the loop
+        return {"iteration": iteration, "enabled": True, "error": repr(exc),
+                "results": []}
+    threshold = float(cfg.answers.get("promote_min_improvement_percent") or 3.0)
+    outcomes: List[Dict[str, Any]] = []
+    for inst in instances:
+        result = results.get(inst.id) or {}
+        # The question's newest attempt owns the workspace that produced the
+        # result (a retry runs in its own directory).
+        # The question's newest attempt owns the workspace that produced the
+        # result (a retry runs in its own directory).
+        child_dir = resolve_child_dir(children, iteration, inst.id)
+        workspace = (child_dir / "workspace") if child_dir else Path("")
+        if not workspace.is_dir():
+            outcomes.append({"shape": inst.id, "action": "skipped",
+                             "reason": "no child workspace for this question"})
+            continue
+        shape = inst.shape or {}
+        answers = {
+            "operator": "Quantized GEMM",
+            "dtype": "INT8 W8A8",
+            "model": str(shape.get("model") or ""),
+            "tp_size": shape.get("tp_size"),
+        }
+        outcome = promote_variant(
+            workspace_dir=workspace,
+            answers=answers,
+            shape_id=inst.id,
+            source_task=f"{cfg.task_id}/iteration_{iteration:03d}/{inst.id}",
+            correctness_ok=result.get("correctness_ok"),
+            min_improvement_percent=threshold,
+            tp=shape.get("tp_size"),
+            m=shape.get("M"),
+            model_label=(model_label_for(inst.id, str(shape.get("model") or ""))
+                         or None),
+            dry_run=freeze,
+        )
+        outcome["ahe_passed"] = result.get("passed")
+        outcome["ahe_median_us"] = result.get("median_us")
+        outcomes.append(outcome)
+        if freeze:
+            with (exp / "promotion_candidates.jsonl").open(
+                    "a", encoding="utf-8") as fh:
+                fh.write(json.dumps({
+                    "ts": time.time(), "iteration": iteration,
+                    "shape": inst.id, "workspace": str(workspace),
+                    "source_task": f"{cfg.task_id}/iteration_{iteration:03d}/{inst.id}",
+                    "answers": answers, "tp": shape.get("tp_size"),
+                    "m": shape.get("M"),
+                    "model_label": (model_label_for(
+                        inst.id, str(shape.get("model") or "")) or None),
+                    "action": outcome.get("action"),
+                    "ahe_passed": result.get("passed"),
+                    "ahe_median_us": result.get("median_us"),
+                    "variant_median_us": outcome.get("old_median_us"),
+                }, ensure_ascii=False) + "\n")
+
+    promoted = [o["shape"] for o in outcomes
+                if o.get("action") in {"added", "updated"}]
+    payload = {
+        "iteration": iteration,
+        "enabled": True,
+        "frozen": bool(freeze),
+        "min_improvement_percent": threshold,
+        "promoted": promoted,
+        "results": outcomes,
+        "note": ("frozen: candidates recorded, promotion happens once at the end"
+                 if freeze else None),
+    }
+    save_json(exp / "runs" / f"iteration_{iteration:03d}" / "promotion.json",
+              payload)
+    with (exp / "promotion_log.jsonl").open("a", encoding="utf-8") as fh:
+        for outcome in outcomes:
+            fh.write(json.dumps({
+                "ts": time.time(), "iteration": iteration,
+                "task_id": cfg.task_id, **outcome,
+            }, ensure_ascii=False) + "\n")
+    append_timeline(cfg.state_dir, "kernel_promotion", {
+        "iteration": iteration, "promoted": promoted,
+        "min_improvement_percent": threshold, "frozen": bool(freeze),
+    })
+    return payload
+
+
+def _promote_baseline_kernels(cfg: ExperimentConfig, exp: Path,
+                              variant_pool: Dict[str, Any],
+                              results: Dict[str, Any]) -> Dict[str, Any]:
+    """Round 1's faster kernels go into the variant pool — as the new bar.
+
+    A kernel result is a fact about the hardware, not a verdict about a
+    harness: the operator's rule is that the baseline round's improvements enter
+    the variant pool immediately, so every later gate (and DKAO's own
+    ``warm_start: best_known``) compares against what the baseline round
+    actually achieved. The *harness* still waits for approval — this only moves
+    kernels.
+
+    Only operators the round actually measured are touched, only a real
+    improvement (>= ``promote_min_improvement_percent`` faster than the recorded
+    best) is written, and the pool file is backed up before being rewritten.
+    """
+    if not results:
+        return {"ok": True, "updated": [], "promoted": [],
+                "skipped": "no measurements"}
+    pool_path = _question_pool_path(cfg)
+    if pool_path is None:
+        return {"ok": True, "updated": [], "promoted": [],
+                "skipped": "no question pool file"}
+    try:
+        from .promotion import promote_kernels_for_round
+    except Exception as exc:  # noqa: BLE001 - promotion must never break the loop
+        return {"ok": False, "updated": [], "errors": [repr(exc)]}
+    measured = {
+        iid: row for iid, row in (results or {}).items()
+        if isinstance(row, dict) and row.get("correctness_ok") is not False
+        and row.get("median_us") is not None
+    }
+    # Same improvement rule as every other kernel write (production default 3%),
+    # so a measurement inside the noise band cannot churn the pool.
+    threshold = float(cfg.answers.get("promote_min_improvement_percent") or 3.0)
+    kept: Dict[str, Any] = {}
+    for iid, row in measured.items():
+        old = (variant_pool.get(iid) or {}).get("best_known_us")
+        try:
+            new = float(row["median_us"])
+            old_value = float(old) if old is not None else None
+        except (TypeError, ValueError):
+            continue
+        if new <= 0:
+            continue
+        if old_value is not None and \
+                (old_value - new) / old_value * 100.0 < threshold:
+            continue
+        kept[iid] = row
+    if not kept:
+        append_timeline(cfg.state_dir, "baseline_kernels_promoted", {
+            "iteration": 1, "operators": [], "pool": str(pool_path),
+            "reason": f"no operator improved by >= {threshold:.2f}%",
+        })
+        return {"ok": True, "updated": [], "promoted": [],
+                "pool": str(pool_path),
+                "reason": f"no operator improved by >= {threshold:.2f}%"}
+    version = "baseline:it1"
+    outcome = promote_kernels_for_round(
+        exp, pool_path, kept, version=version, iteration=1,
+        reason="baseline round: faster kernels enter the variant pool")
+    updated = outcome.get("updated") or []
+    outcome["promoted"] = sorted(str(u["operator_id"]) for u in updated)
+    append_timeline(cfg.state_dir, "baseline_kernels_promoted", {
+        "iteration": 1, "operators": [u["operator_id"] for u in updated],
+        "pool": str(pool_path), "backup": outcome.get("backup"),
+        "min_improvement_percent": threshold,
+    })
+    for change in updated:
+        # Keep the experiment's own view in step with the pool it just moved.
+        iid = str(change["operator_id"])
+        entry = variant_pool.get(iid)
+        if isinstance(entry, dict):
+            entry["best_known_us"] = change["median_us"]
+    return outcome
+
+
+_FAILING_VERDICTS = {"REJECT", "REJECT_OVERFIT", "INCONCLUSIVE"}
+_RECOVERING_VERDICTS = {"PROMOTE", "PROMOTE_UNEXPLAINED", "NO_SIGNAL"}
+
+
+def _next_fail_streak(verdict: str, streak: int) -> int:
+    """Consecutive failing rounds.
+
+    INCONCLUSIVE counts: rounds where nothing could be measured are exactly the
+    situation that used to spin forever (a run reached 513 rounds with 510
+    REJECTs while a shared GPU made every measurement unusable).
+    """
+    verdict = str(verdict or "")
+    if verdict in _FAILING_VERDICTS:
+        return int(streak) + 1
+    if verdict in _RECOVERING_VERDICTS:
+        return 0
+    return int(streak)
+
+
+def _stop_requested(exp: Path) -> Optional[Dict[str, Any]]:
+    """A request (from the evaluator) to stop the whole experiment.
+
+    Raised when no device satisfied the gate for the configured number of
+    hourly checks: waiting forever would burn wall-clock time, so the task
+    stops with an auditable reason instead.
+    """
+    data = _load_json(Path(exp) / "stop_requested.json")
+    return data or None
+
+
+def _variant_policy(cfg: ExperimentConfig) -> str:
+    """``freeze`` (default): collect candidates, promote once at the end.
+
+    Freezing keeps the reference implementation identical across rounds, so
+    every round compares under the same conditions. ``per_round`` restores the
+    old behaviour (write the variant as soon as a round beats it).
+
+    Under the variant-round protocol this knob is bypassed on purpose: the
+    baseline round's faster kernels go into the variant pool immediately (the
+    operator's rule — a kernel result is a fact, while the *harness* is what
+    waits for approval).
+    """
+    raw = str(cfg.answers.get("variant_policy") or "freeze").strip().lower()
+    return "per_round" if raw in {"per_round", "always", "each", "immediate"} \
+        else "freeze"
+
+
+def _seed_variant_entry(operators: Dict[str, Any], iid: str,
+                        entry: Dict[str, Any]) -> None:
+    """Mirror one pool ``best_known`` into the experiment's variant table."""
+    operators[iid] = {
+        "median_us": entry.get("best_known_us"),
+        "p90_us": None,
+        "correctness_ok": True,
+        "kernel_source": "pool:best_known",
+        "harness_version": None,
+        "source_iteration": None,
+        "family": entry.get("family") or "",
+        "contract": entry.get("contract") or {},
+        "conditions": {"origin": "variant_pool_refresh"},
+        "recorded_at": None,
+    }
+
+
+def _reload_variant_baseline(cfg: ExperimentConfig,
+                             variant_pool: Dict[str, Any],
+                             variant_table: Dict[str, Any]) -> List[str]:
+    """Pull externally-updated ``best_known_us`` back into this run's view.
+
+    The baseline round writes the kernels it beat the variant with straight
+    into the question pool. Every later round must therefore compare against —
+    and DKAO must warm-start from — that *new* number, not the one this process
+    read at startup. Returns the operators whose reference moved, so the gates
+    can say out loud that the bar was raised.
+    """
+    path = _question_pool_path(cfg)
+    if path is None:
+        return []
+    try:
+        fresh = pool_instances(path)
+    except (OSError, ValueError):
+        return []
+    moved: List[str] = []
+    operators = variant_table.setdefault("operators", {})
+    for iid, entry in fresh.items():
+        variant_pool[iid] = entry
+        best = entry.get("best_known_us")
+        recorded = (operators.get(iid) or {}).get("median_us")
+        try:
+            better = best is not None and (recorded is None
+                                           or float(best) < float(recorded))
+        except (TypeError, ValueError):
+            better = False
+        if not better:
+            continue
+        moved.append(str(iid))
+        _seed_variant_entry(operators, str(iid), entry)
+    if moved:
+        save_variant_table(cfg.exp_dir, variant_table)
+        append_timeline(cfg.state_dir, "variant_reference_refreshed", {
+            "operators": sorted(moved)})
+    return sorted(moved)
+
+
+def _baseline_stop(cfg: ExperimentConfig, exp: Path, *, iteration: int,
+                   verdict: Dict[str, Any]) -> None:
+    """Round 1 could not measure the paper: the run stops right there.
+
+    Round 1 judges nothing, so it has no gate to fail — but it owes one usable
+    measurement per operator: those kernels seed the pool every later round has
+    to beat, and the paper is the run's exam. If a child crashed, was blocked by
+    DKAO's own measurement gate or reported nothing usable, there is no honest
+    way to continue, so the task ends with an auditable reason instead of
+    inventing evidence.
+    """
+    reason = "; ".join(verdict.get("reasons") or []) or "baseline round failed"
+    save_json(Path(exp) / "stop_requested.json", {
+        "iteration": iteration, "reason": "baseline_round_failed",
+        "detail": {"failed": verdict.get("failed"),
+                   "measured": verdict.get("measured"),
+                   "message": reason},
+    })
+    append_timeline(cfg.state_dir, "baseline_round_failed", {
+        "iteration": iteration, "failed": verdict.get("failed"),
+        "measured": verdict.get("measured"), "reason": reason,
+    })
+    _record_state(cfg, FINISHED, iteration, event="baseline_round_failed",
+                  failed=list(verdict.get("failed") or []))
+    set_run(cfg.state_dir, final_status="baseline_round_failed")
+
+
+
+def _track_best_known(cfg: ExperimentConfig, exp: Path, iteration: int,
+                      results: Dict[str, Any]) -> Dict[str, Any]:
+    """Per-shape best-known curve + explicit regression detection.
+
+    Before this, a round whose kernel was slower than the known best only
+    silently failed to promote. Recording the curve makes "this round got
+    worse starting point" visible in the record and the UI.
+    """
+    path = exp / "best_known.json"
+    data = _load_json(path) or {}
+    regressions: List[Dict[str, Any]] = []
+    tolerance = 1.02
+    for iid, row in (results or {}).items():
+        if not isinstance(row, dict):
+            continue
+        median = row.get("median_us")
+        entry = data.setdefault(iid, {"best_median_us": None,
+                                      "best_iteration": None, "history": []})
+        entry.setdefault("history", []).append({
+            "iteration": iteration, "median_us": median,
+            "passed": row.get("passed"),
+        })
+        entry["history"] = entry["history"][-40:]
+        try:
+            value = float(median)
+        except (TypeError, ValueError):
+            continue
+        if value <= 0:
+            continue
+        best = entry.get("best_median_us")
+        if best is None or value < float(best):
+            entry["best_median_us"] = value
+            entry["best_iteration"] = iteration
+        elif value > float(best) * tolerance:
+            regressions.append({
+                "shape": iid, "iteration": iteration, "median_us": value,
+                "best_median_us": best, "best_iteration": entry.get("best_iteration"),
+                "delta_percent": round((value - float(best)) / float(best) * 100, 2),
+            })
+    save_json(path, data)
+    if regressions:
+        with (exp / "regressions.jsonl").open("a", encoding="utf-8") as fh:
+            for row in regressions:
+                fh.write(json.dumps({"ts": time.time(), **row},
+                                    ensure_ascii=False) + "\n")
+        append_timeline(cfg.state_dir, "kernel_regression", {
+            "iteration": iteration, "regressions": regressions,
+        })
+    return {"regressions": regressions}
+
+
+def _finalize_promotion(cfg: ExperimentConfig, exp: Path) -> Optional[Dict[str, Any]]:
+    """Promote the best kernel per shape once, after a frozen experiment."""
+    if not _truthy(cfg.answers.get("promote_kernels"), default=True):
+        return None
+    candidates_path = exp / "promotion_candidates.jsonl"
+    if not candidates_path.is_file():
+        return None
+    best: Dict[str, Dict[str, Any]] = {}
+    for line in candidates_path.read_text(encoding="utf-8",
+                                          errors="replace").splitlines():
+        try:
+            row = json.loads(line)
+        except ValueError:
+            continue
+        if not isinstance(row, dict) or row.get("ahe_passed") is not True:
+            continue
+        if row.get("action") not in {"would_add", "would_update", "added",
+                                    "updated"}:
+            continue
+        current = best.get(str(row.get("shape")))
+        if current is None or (row.get("ahe_median_us") or float("inf")) < (
+                current.get("ahe_median_us") or float("inf")):
+            best[str(row.get("shape"))] = row
+    if not best:
+        payload = {"enabled": True, "frozen": True, "promoted": [],
+                   "results": [], "note": "no eligible candidate"}
+        save_json(exp / "final_promotion.json", payload)
+        return payload
+    try:
+        from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.variant_promote import (  # noqa: E501
+            promote_variant,
+        )
+    except Exception as exc:  # noqa: BLE001
+        return {"enabled": True, "error": repr(exc), "results": []}
+    threshold = float(cfg.answers.get("promote_min_improvement_percent") or 3.0)
+    outcomes: List[Dict[str, Any]] = []
+    for shape_id, row in sorted(best.items()):
+        workspace = Path(str(row.get("workspace") or ""))
+        if not workspace.is_dir():
+            outcomes.append({"shape": shape_id, "action": "skipped",
+                             "reason": "candidate workspace is gone"})
+            continue
+        outcome = promote_variant(
+            workspace_dir=workspace,
+            answers=row.get("answers") or {},
+            shape_id=shape_id,
+            source_task=str(row.get("source_task") or ""),
+            correctness_ok=True,
+            min_improvement_percent=threshold,
+            tp=row.get("tp"), m=row.get("m"),
+            model_label=row.get("model_label"),
+        )
+        outcome["ahe_median_us"] = row.get("ahe_median_us")
+        outcome["candidate_iteration"] = row.get("iteration")
+        outcomes.append(outcome)
+    promoted = [o["shape"] for o in outcomes
+                if o.get("action") in {"added", "updated"}]
+    payload = {
+        "enabled": True, "frozen": True, "min_improvement_percent": threshold,
+        "promoted": promoted, "results": outcomes,
+        "candidates_considered": sorted(best),
+    }
+    save_json(exp / "final_promotion.json", payload)
+    with (exp / "promotion_log.jsonl").open("a", encoding="utf-8") as fh:
+        for outcome in outcomes:
+            fh.write(json.dumps({"ts": time.time(), "iteration": -1,
+                                 "task_id": cfg.task_id, "final": True,
+                                 **outcome}, ensure_ascii=False) + "\n")
+    append_timeline(cfg.state_dir, "final_kernel_promotion", {
+        "promoted": promoted, "candidates": sorted(best),
+    })
+    return payload
+
+
+def _write_round_review(cfg: ExperimentConfig, exp: Path) -> Dict[str, Any]:
+    """Cross-round review: every generation, medians, champion and regressions."""
+    from .champion_selection import best_generation, load_generations, summarize
+
+    generations = load_generations(exp)
+    rows = summarize(generations)
+    champion = best_generation(generations)
+    review = {
+        "task_id": cfg.task_id,
+        "generations": rows,
+        "champion_iteration": champion.iteration if champion else None,
+        "champion_verdict": champion.verdict if champion else None,
+        "final_promotion": _load_json(exp / "final_promotion.json"),
+        "best_known": _load_json(exp / "best_known.json"),
+        "variant_policy": _variant_policy(cfg),
+        "pinned_baseline": _load_json(exp / "champion_override.json"),
+    }
+    save_json(exp / "round_review.json", review)
+    lines = ["# Cross-round review (AHE)", "",
+             f"- variant policy: **{review['variant_policy']}**",
+             f"- champion generation: **iteration {review['champion_iteration']}**"
+             f" ({review['champion_verdict']})", ""]
+    if rows:
+        shapes = sorted({iid for row in rows for iid in row.get("medians", {})})
+        header = "| gen | verdict | pass | " + " | ".join(shapes) + " |"
+        lines += [header, "|" + "---|" * (3 + len(shapes))]
+        for row in rows:
+            medians = row.get("medians") or {}
+            cells = []
+            for shape in shapes:
+                value = medians.get(shape)
+                cells.append(f"{value:.1f}" if isinstance(value, (int, float))
+                             else "-")
+            lines.append(
+                f"| {row['iteration']}{' ⭐' if row.get('champion') else ''} "
+                f"| {row.get('verdict')} | {row.get('pass_count')} | "
+                + " | ".join(cells) + " |")
+        lines.append("")
+    best_known = review["best_known"] or {}
+    if best_known:
+        lines.append("## Best known per shape")
+        for shape, entry in sorted(best_known.items()):
+            lines.append(
+                f"- `{shape}`: {entry.get('best_median_us')} us "
+                f"(iteration {entry.get('best_iteration')})")
+        lines.append("")
+    promotion = review["final_promotion"] or {}
+    if promotion:
+        lines.append("## Final promotion")
+        for outcome in promotion.get("results") or []:
+            lines.append(
+                f"- `{outcome.get('shape')}`: {outcome.get('action')} "
+                f"({outcome.get('reason') or outcome.get('improvement_percent') or ''})")
+    (exp / "round_review.md").write_text("\n".join(lines) + "\n",
+                                         encoding="utf-8")
+    append_timeline(cfg.state_dir, "round_review_written", {
+        "champion_iteration": review["champion_iteration"],
+        "generations": len(rows),
+    })
+    return review
+
+
+def _code_reload_pending(cfg: ExperimentConfig,
+                         loaded: str) -> Optional[Dict[str, Any]]:
+    """Whether the harness code on disk changed since this process started.
+
+    A round runs for hours, so an operator who fixes the loop (a new GPU gate,
+    a new retry policy) cannot expect the live process to honour it — the
+    modules are already imported. Rather than running one more round of stale
+    logic, the loop stops here: the recorded revision marker makes the restart
+    idempotent (a fresh process records its own revision and carries on), and
+    the supervisor/server resumes at the next iteration with the new code.
+    """
+    current = harness_revision()
+    if not loaded or current == loaded:
+        return None
+    return {
+        "loaded_revision": loaded,
+        "on_disk_revision": current,
+        "resumes_at_iteration": None,  # filled by the caller
+    }
+
+
+def _rounds_used_on_disk(exp: Path) -> float:
+    """How many rounds this experiment has already spent.
+
+    ``rounds_used`` lives in ``run.json`` (the WebUI reads it). When it is
+    missing — a run created before the field — it is derived from the round
+    artifacts, so a resumed run does not silently start with a fresh budget.
+    """
+    stored = _load_json(Path(exp) / "run.json") or {}
+    raw = stored.get("rounds_used")
+    try:
+        if raw is not None:
+            return max(0.0, float(raw))
+    except (TypeError, ValueError):
+        pass
+    used = 0.0
+    for stage_file in (Path(exp) / "runs").glob("iteration_*/input/stage.json"):
+        data = _load_json(stage_file) or {}
+        try:
+            used += float(data.get("round_cost") or 0.0)
+        except (TypeError, ValueError):
+            continue
+    return used
+
+
+#: A round is an iteration: the baseline round is round 1, and each later round
+#: is one performance pass plus (when it wins) the generalization retake inside
+#: the same iteration, so a retake never consumes a round of its own.
+ROUNDS_DEFAULT = 10
+ROUNDS_MAX = 50
+
+#: Cost guardrail, deliberately far above any real run: a judged pass costs 1.0
+#: and a retake 0.5, so ROUNDS_MAX rounds of steady retakes still cost well under
+#: this. It exists only so a pathological loop cannot spend forever.
+COST_GUARDRAIL = 100.0
+
+
+def _rounds_limit(cfg: ExperimentConfig) -> int:
+    """How many rounds (iterations) this task may run — the form's "Rounds".
+
+    ``max_iterations`` is still honoured for runs whose requirements were built
+    before the form field was renamed; the default is ROUNDS_DEFAULT.
+    """
+    for key in ("rounds", "max_iterations"):
+        raw = cfg.answers.get(key)
+        if raw in (None, ""):
+            continue
+        try:
+            return max(1, min(ROUNDS_MAX, int(raw)))
+        except (TypeError, ValueError):
+            continue
+    return ROUNDS_DEFAULT
+
+
+def _rounds_used_on_disk(exp: Path) -> float:
+    """The judged-pass cost this experiment has already spent (1.0 / 0.5).
+
+    Only used by the cost guardrail; the round *limit* is an iteration number.
+    """
+    stored = _load_json(Path(exp) / "run.json") or {}
+    raw = stored.get("rounds_used")
+    try:
+        if raw is not None:
+            return max(0.0, float(raw))
+    except (TypeError, ValueError):
+        pass
+    used = 0.0
+    for stage_file in (Path(exp) / "runs").glob("iteration_*/input/stage.json"):
+        data = _load_json(stage_file) or {}
+        try:
+            used += float(data.get("round_cost") or 0.0)
+        except (TypeError, ValueError):
+            continue
+    return used
+
+
+def run_experiment(cfg: ExperimentConfig, *,
+                   start_iteration: int = 1,
+                   iterations_to_run: Optional[int] = None,
+                   champion_iteration: Optional[int] = None) -> Path:
+    """Run (or continue) the AHE outer loop.
+
+    ``start_iteration`` > 1 resumes an existing experiment: the champion and
+    previous-round state are rebuilt from disk (best_ever + the iteration
+    artifacts) so a new generation of the harness can be evaluated against
+    the already-measured champion without re-running earlier rounds.
+    """
+    exp = cfg.exp_dir
+    (exp / "runs").mkdir(parents=True, exist_ok=True)
+    cfg.dump_snapshot(exp / "config_snapshot.json")
+    # Record the revision of the code this process loaded, so a later round can
+    # tell "the operator changed the harness" from "nothing happened".
+    loaded_revision = write_harness_revision(cfg.state_dir)
+    _record_state(
+        cfg, PREPARE,
+        event="experiment_start" if start_iteration <= 1 else "experiment_resume",
+        start_iteration=start_iteration,
+        max_iterations=cfg.max_iterations,
+        execution_mode=cfg.execution_mode,
+        pool_source=cfg.pool_source,
+        harness_revision=loaded_revision,
+    )
+
+    live_workspace = _seed_workspace(cfg)
+    evaluator = build_evaluator(cfg.execution_mode)
+    evolver = build_evolver(
+        str(cfg.answers.get("evolve_mode") or "dry-run"),
+        model=str(cfg.answers.get("evolve_model") or DEFAULT_EVOLVE_MODEL),
+        timeout_seconds=float(cfg.answers.get("evolve_timeout_seconds") or 3600),
+    )
+
+    pool = _load_pool(cfg)
+    # FLOW.md §1.1: a round is sized in whole waves of 4 questions. The
+    # configured count is already normalized on load (config.py); this is the
+    # second call site, where "no cap" means "the whole pool" and the pool size
+    # itself is the only remaining source of a non-multiple count.
+    budget = cfg.per_round_budget or (len(pool.all_ids()) if pool else 0)
+    if budget > 0:
+        budget = normalized_per_gate(budget)
+    purposes: Dict[str, str] = {}
+    selection_notes: List[str] = []
+    prev_selected: List[str] = _load_ids(exp / LAST_SCORED_FILE)
+
+    prev_results: Optional[Dict[str, Any]] = None
+    champion_results: Optional[Dict[str, Any]] = None
+    pending_confirmation = False
+    best_ever: Dict[str, Any] = _load_json(exp / BEST_EVER_FILE) or {}
+
+    if start_iteration > 1:
+        # Resume: the champion is whichever generation best_ever recorded; its
+        # measured results and the previous round's set come back from disk.
+        # ``champion_iteration`` lets the caller pin the comparison baseline
+        # explicitly (e.g. "compare iteration 3 against iteration 2").
+        prev_results = _iteration_results(exp, start_iteration - 1)
+        # The champion is the best *measured* generation among the ones that
+        # were not rejected: a generation can hold the best result and still
+        # not have been promoted in its own round (it needed confirmation, or
+        # too few questions won). Comparing against "the last promoted
+        # generation" would let a worse harness take the champion slot.
+        measured = best_generation(load_generations(exp))
+        champ_iter = int(champion_iteration
+                         or (measured.iteration if measured else 0)
+                         or best_ever.get("iteration")
+                         or start_iteration - 1)
+        if champion_iteration is None and measured is not None and measured.results:
+            champion_results = dict(measured.results)
+            best_ever = {
+                "iteration": measured.iteration,
+                "pass_rate": None,
+                "snapshot_dir": measured.snapshot_dir or str(
+                    exp / "runs" / f"iteration_{measured.iteration:03d}"
+                    / "input" / "workspace"),
+                "workspace_revision": measured.workspace_revision,
+                "verdict": "BEST_MEASURED",
+            }
+            save_json(exp / BEST_EVER_FILE, best_ever)
+        else:
+            champion_results = _iteration_results(exp, champ_iter)
+        if champion_iteration is not None:
+            snapshot_dir = (exp / "runs" / f"iteration_{champ_iter:03d}"
+                            / "input" / "workspace")
+            scores = [row for row in _load_jsonl(exp / SCORES_FILE)
+                      if int(row.get("iteration") or 0) == champ_iter]
+            revision = ""
+            manifest = _load_json(
+                exp / "runs" / f"iteration_{champ_iter:03d}" / "evolve"
+                / "change_manifest.json")
+            if manifest:
+                revision = str(manifest.get("workspace_revision") or "")
+            best_ever = {
+                "iteration": champ_iter,
+                "pass_rate": (scores[-1].get("pass_rate") if scores else None),
+                "snapshot_dir": str(snapshot_dir),
+                "workspace_revision": revision,
+                "verdict": "MANUAL_BASELINE",
+            }
+            save_json(exp / BEST_EVER_FILE, best_ever)
+            save_json(exp / "champion_override.json",
+                      {"champion_iteration": champ_iter,
+                       "pinned_by": "resume --champion-iteration"})
+            append_timeline(cfg.state_dir, "champion_baseline_pinned",
+                            {"iteration": champ_iter})
+        if not prev_selected:
+            prev_selected = _load_ids(exp / LAST_SCORED_FILE)
+        set_run(cfg.state_dir, finished=False, final_status=None)
+
+    # Stop condition is re-read every round, so the WebUI can raise or lower
+    # the target while the loop is running ("run to iteration N"). The hard
+    # cap only bounds a runaway target.
+    HARD_ITERATION_CAP = 50
+    #: One HE run chases one promotable harness and stops when it gets one.
+    #: The round budget keeps a run that never succeeds from going forever.
+    max_rounds = max(1.0, COST_GUARDRAIL)
+
+    # Variant table: what every operator is measured against. Seeded from the
+    # existing pool (each operator's best known result) and only updated when a
+    # harness is accepted and promoted, so the reference never drifts.
+    variant_pool: Dict[str, Any] = {}
+    variant_table: Dict[str, Any] = {"operators": {}, "groups": {}}
+    if not _uses_variant_rounds(cfg):
+        _announce_legacy_protocol(cfg, exp)
+    if _uses_variant_rounds(cfg):
+        pool_path = _question_pool_path(cfg)
+        variant_pool = pool_instances(pool_path) if pool_path else {}
+        variant_table = load_variant_table(exp, pool_path=pool_path)
+        per_gate = normalized_per_gate(
+            cfg.answers.get("questions_per_round")
+            or cfg.answers.get("per_round_budget")
+            or OPERATORS_PER_GATE)
+        append_timeline(cfg.state_dir, "variant_table_ready", {
+            "operators": len(variant_table.get("operators") or {}),
+            "pool_operators": len(variant_pool),
+            "harness_version": variant_table.get("harness_version"),
+            "questions_per_gate": per_gate,
+        })
+        if not variant_pool:
+            # Nothing to measure: fail loudly instead of spinning through the
+            # round budget. This happens when the pool file is missing/empty.
+            append_timeline(cfg.state_dir, "variant_pool_empty", {
+                "pool": str(_question_pool_path(cfg) or cfg.pool_source)})
+            _record_state(cfg, FINISHED, start_iteration,
+                          event="variant_pool_empty")
+            set_run(cfg.state_dir, final_status="no_question_pool")
+            return _final_report(cfg)
+        judgeable = _judgeable_ids(variant_pool, variant_table)
+        if len(judgeable) < per_gate:
+            # A round is a whole wave of 4 questions (FLOW.md §1.1). Drawing
+            # fewer would silently change what the gate means: with 3 measured
+            # operators "at least 75% win" becomes "all three must win". Say so
+            # and stop instead of judging a round the pool cannot fill.
+            append_timeline(cfg.state_dir, "variant_pool_too_small", {
+                "pool": str(_question_pool_path(cfg) or cfg.pool_source),
+                "judgeable": judgeable,
+                "judgeable_count": len(judgeable),
+                "questions_per_round": per_gate,
+                "policy": ("a round is sized in whole waves of 4; add judgeable "
+                           "operators (variant + baseline) to the pool or lower "
+                           "the round size"),
+            })
+            _record_state(cfg, FINISHED, start_iteration,
+                          event="variant_pool_too_small")
+            set_run(cfg.state_dir, final_status="no_question_pool")
+            return _final_report(cfg)
+    # ``iterations_to_run`` is the *current call's* batch (the server resumes one
+    # round at a time), not the run's bound: the bound is the form's "Rounds".
+    fallback_end = start_iteration + (iterations_to_run or cfg.max_iterations) - 1
+    while_end = max(fallback_end, start_iteration) + HARD_ITERATION_CAP
+    iteration = start_iteration
+    #: Round budget. A pass that judges the performance gate costs a full
+    #: round; the generalization retake (the same harness handed to DKAO again
+    #: with the frozen paper) costs half a round — the operator calls it
+    #: "iteration 2.5". ``iteration`` therefore stays an integer (it names the
+    #: ``runs/iteration_NNN`` directory) while this counter is fractional.
+    # Reported cost of the judged passes so far (1.0 per performance pass, 0.5
+    # per retake). It is bookkeeping for the guardrail below, not the run limit:
+    # the limit is an iteration number.
+    rounds_used = _rounds_used_on_disk(exp)
+    rounds_limit = _rounds_limit(cfg)
+    #: Round 1 could not measure the paper, so the run has no foundation and
+    #: already ended: the tail below must not overwrite that verdict.
+    baseline_stopped = False
+    while iteration <= while_end:
+        target = read_target_iterations(exp)
+        set_run(cfg.state_dir, target_iterations=target, max_rounds=max_rounds,
+                rounds_used=rounds_used)
+        # The form's "Rounds" is the limit; the WebUI target can raise or lower
+        # it for a live run. Either way the unit is a round = an iteration, so
+        # round 10 means the run may start iteration 10 and no further.
+        rounds_limit_now = max(1, int(target)) if target is not None else rounds_limit
+        # The variant protocol counts rounds here; the legacy/pool path keeps its
+        # historical "requested iterations" bound (its callers pass an explicit
+        # batch and its rounds are not the same unit).
+        if not _uses_variant_rounds(cfg) and iteration > fallback_end:
+            break
+        if iteration > rounds_limit_now:
+            append_timeline(cfg.state_dir, "rounds_reached", {
+                "iteration": iteration - 1, "rounds": rounds_limit_now,
+                "target": target,
+                "stage": "variant" if _uses_variant_rounds(cfg) else "legacy"})
+            break
+        if _uses_variant_rounds(cfg) and rounds_used >= max_rounds:
+            # The cost guardrail tripped: far above any real run, it only stops
+            # a pathological loop. The run limit is the form's "Rounds".
+            append_timeline(cfg.state_dir, "round_budget_exhausted", {
+                "rounds_used": rounds_used, "max_rounds": max_rounds,
+                "rounds_limit": rounds_limit_now,
+                "reason": "cost guardrail tripped"})
+            set_run(cfg.state_dir, final_status="rounds_exhausted",
+                    rounds_used=rounds_used)
+            break
+        # A pending retake means the previous performance stage passed: this
+        # pass re-measures the frozen paper with the same harness bytes and
+        # judges the generalization gate. Nothing else happens in it.
+        pending_retake = (read_pending_generalization(exp)
+                          if _uses_variant_rounds(cfg) else None)
+        stage = round_stage(iteration=iteration, pending=pending_retake)
+        #: Set when this pass produced its own gate decision (baseline and the
+        #: generalization retake do): the frozen-champion decision engine must
+        #: not run, and the loop steps on by itself.
+        own_decision = False
+        hist_index = len(load_generations(exp))
+        iter_dir = exp / "runs" / f"iteration_{iteration:03d}"
+        input_dir = iter_dir / "input"
+        evolve_dir = iter_dir / "evolve"
+
+        # --- select this round's instances ---------------------------------
+        # Variant-driven rounds: the operators come from the question pool
+        # (which already records each operator's current best), and they are
+        # judged against the variant table instead of a frozen champion. This
+        # is what makes "one HE run chases one promotable harness" work.
+        if _uses_variant_rounds(cfg):
+            # The round draws exactly the configured number of operators (a
+            # multiple of the 4 HCUs): the count comes from the New Task form
+            # and is never reduced to fit the cards. Whether a card is free is
+            # DKAO's business — the child's engine waits on its own admission
+            # gate — so a round is always the size the operator asked for.
+            fit = {"requested": per_gate, "count": per_gate,
+                   "free_devices": None, "reduced": False}
+            draw_count = fit["count"] or per_gate
+            reference_moved: List[str] = []
+            if stage == STAGE_PERFORMANCE:
+                # The baseline round may have raised the bar (its kernels go
+                # straight into the pool); every later round has to see it.
+                reference_moved = _reload_variant_baseline(
+                    cfg, variant_pool, variant_table)
+            round_questions = select_round_questions(
+                iteration=iteration, table=variant_table,
+                instances=variant_pool, state={"groups": group_state(variant_table)},
+                count=draw_count, rng=_round_rng(cfg, iteration, stage)[0],
+                pending=pending_retake)
+            selected = list(round_questions["performance_ids"])
+            purposes = {iid: "performance" for iid in selected}
+            for iid in round_questions.get("generalization_ids") or []:
+                purposes.setdefault(iid, "generalization")
+            if stage == STAGE_GENERALIZATION:
+                purposes = {iid: "generalization" for iid in selected}
+            selection_notes = [
+                {
+                    STAGE_BASELINE: ("baseline round: measures the new run's "
+                                     "generalization paper; judges no gate"),
+                    STAGE_PERFORMANCE: ("performance gate: fresh random "
+                                        "operators, judged against each "
+                                        "operator's variant"),
+                    STAGE_GENERALIZATION: ("generalization retake: the same "
+                                           "harness handed to DKAO again with "
+                                           "the frozen paper"),
+                }.get(stage, "variant-driven round"),
+                f"questions={len(selected)} (requested {fit['requested']})"]
+            if reference_moved:
+                selection_notes.append(
+                    "variant reference refreshed from the baseline round's "
+                    "kernels: " + ", ".join(reference_moved))
+            # Fresh pool data also decides what DKAO warm-starts from, so the
+            # retake and the performance stage always launch children against
+            # the current best-known kernel.
+            inst_budget = int(cfg.answers.get("default_rounds") or 5)
+            instances = [
+                _variant_instance_spec({"id": i, **dict(variant_pool[i])},
+                                       budget_rounds=inst_budget)
+                for i in selected if i in variant_pool]
+            save_json(input_dir / "benchmark" / "scored_ids.json", selected)
+            save_json(input_dir / "benchmark" / "purposes.json", {
+                "iteration": iteration, "stage": stage,
+                "purposes": purposes,
+                "families": [], "rationale": "variant-driven selection",
+                "notes": selection_notes,
+                "generalization_ids": round_questions.get("generalization_ids"),
+                "defines_paper": round_questions.get("defines_paper"),
+            })
+            save_json(exp / ROUND_QUESTIONS_FILE, {
+                "iteration": iteration, "stage": stage,
+                "performance_ids": selected,
+                "generalization_ids": round_questions.get("generalization_ids"),
+                "defines_paper": round_questions.get("defines_paper"),
+            })
+            save_json(input_dir / "stage.json", {
+                "iteration": iteration, "stage": stage,
+                "round_cost": {
+                    # The baseline round judges nothing, so it costs nothing:
+                    # the budget is spent on judged passes.
+                    STAGE_BASELINE: 0.0,
+                    STAGE_GENERALIZATION: GENERALIZATION_ROUND_COST,
+                }.get(stage, PERFORMANCE_ROUND_COST),
+                "rounds_used_before": rounds_used,
+                "retake_of_iteration": (pending_retake or {}).get("iteration"),
+            })
+            prev_selected = list(selected)
+            if round_questions.get("defines_paper"):
+                # Persist the run's fixed generalization paper immediately: it
+                # is the exam for the rest of the run, so it must survive even
+                # if the baseline round is the last thing that happens.
+                variant_table["groups"] = record_paper(
+                    variant_table, iteration=iteration,
+                    generalization_ids=round_questions["generalization_ids"],
+                    performance_ids=selected)
+                save_variant_table(exp, variant_table)
+            append_timeline(cfg.state_dir, "variant_round_selection", {
+                "iteration": iteration, "stage": stage,
+                "performance_ids": selected,
+                "generalization_ids": round_questions.get("generalization_ids"),
+                "defines_paper": round_questions.get("defines_paper"),
+                # Which draw this was: the seed is what makes "why these four
+                # operators?" reproducible after the fact.
+                "seed": _round_rng(cfg, iteration, stage)[1],
+                "count": draw_count,
+                "judgeable_in_pool": len(_judgeable_ids(variant_pool, variant_table)),
+                "rounds_used": rounds_used,
+            })
+        elif pool is not None:
+            plan = None
+            if iteration > 1:
+                prev_dir = iter_dir.parent / f"iteration_{iteration - 1:03d}"
+                plan = load_round_plan(prev_dir / "evolve" / "round_plan.json")
+            # Role-based selection: regression anchors + repair + a
+            # cross-family generalisation probe, with family quota, overlap
+            # and least-recently-used rotation. The evolve agent's picks are
+            # honoured where they fit a role.
+            rotation = load_rotation(exp)
+            selection = select_questions(
+                pool=pool,
+                budget=budget or len(pool.all_ids()),
+                prev_selected=prev_selected,
+                prev_results=prev_results,
+                plan=plan,
+                rotation=rotation,
+            )
+            selected = selection.selected
+            purposes = dict(selection.purposes)
+            selection_notes = list(selection.notes)
+            inst_budget = int(cfg.answers.get("default_rounds") or 5)
+            instances = [
+                _pool_instance_spec(pool.get(i), budget_rounds=inst_budget)
+                for i in selected
+            ]
+            save_json(input_dir / "benchmark" / "scored_ids.json", selected)
+            save_json(input_dir / "benchmark" / "purposes.json", {
+                "iteration": iteration,
+                "purposes": purposes,
+                "families": selection.families,
+                "rationale": selection.rationale,
+                "notes": selection_notes,
+            })
+            save_json(exp / LAST_SCORED_FILE, selected)
+            mark_selected(exp, selected, iteration)
+            append_timeline(cfg.state_dir, "question_selection", {
+                "iteration": iteration,
+                "selected": selected,
+                "purposes": purposes,
+                "families": selection.families,
+                "notes": selection_notes,
+            })
+            prev_selected = list(selected)
+        else:
+            instances = list(cfg.suite)
+
+        # --- safety net: a malformed evolved harness rolls back immediately ---
+        malformed = _harness_validate_error(live_workspace)
+        if malformed and iteration > 1:
+            best_dir = Path(str(best_ever.get("snapshot_dir", "")))
+            if best_dir.is_dir():
+                replace_tree(best_dir, live_workspace)
+                best_rev = best_ever.get("workspace_revision")
+                if best_rev:
+                    subprocess.run(
+                        ["git", "reset", "--hard", str(best_rev)],
+                        cwd=live_workspace, check=True,
+                        stdout=subprocess.DEVNULL,
+                    )
+            save_json(iter_dir / "rollback.json", {
+                "iteration": iteration,
+                "verdict": "REJECT",
+                "action": "RESTORE_CHAMPION",
+                "reason": f"harness malformed: {malformed}",
+                "restored_from": str(best_dir),
+            })
+            save_json(input_dir / "decision.json", {
+                "iteration": iteration,
+                "verdict": "REJECT",
+                "action": "RESTORE_CHAMPION",
+                "reason": f"harness malformed: {malformed}",
+            })
+            # end the experiment early: champion restored and recorded
+            break
+
+        # --- snapshot + evaluate the live workspace ------------------------
+        snapshot = input_dir / "workspace"
+        copy_tree(live_workspace, snapshot)
+        _record_state(cfg, EVALUATE, iteration, event="iteration_start",
+                      stage=stage,
+                      selected=list(selected) if pool is not None else
+                      [i.id for i in instances])
+
+        # The evaluator hands every question to its own DKAO child. Which stage
+        # this is decides only where those children live on disk — the retake
+        # gets its own directory so it is a fresh DKAO job, not a continuation
+        # of the performance stage's optimization. DKAO itself is untouched.
+        cfg.answers["round_stage"] = stage
+        results = evaluator.evaluate(cfg, snapshot, iteration, instances=instances)
+        stop = _stop_requested(exp)
+        if stage == STAGE_BASELINE and _uses_variant_rounds(cfg):
+            # The baseline round is the run's foundation: it judges no gate, but
+            # it owes one usable measurement per operator (those kernels set the
+            # bar every later round has to clear, and the paper is the run's
+            # exam). Nothing — a crashed child, DKAO's own measurement gate
+            # refusing the cards, a missing report — may be papered over.
+            verdict = baseline_round_verdict(results, selected)
+            if stop or verdict["status"] != "PASS":
+                reason = (str(stop.get("reason")) if stop
+                          else "; ".join(verdict.get("reasons") or []))
+                append_timeline(cfg.state_dir, "round_failed", {
+                    "iteration": iteration, "stage": stage,
+                    "reason": reason or "baseline round failed",
+                    "failed": verdict.get("failed"),
+                })
+                save_json(input_dir / "benchmark" / "results.json", {
+                    "iteration": iteration, "results": results,
+                    "round_failed": reason or "baseline round failed",
+                })
+                _record_state(cfg, ANALYZE, iteration,
+                              event="baseline_round_failed")
+                _baseline_stop(cfg, exp, iteration=iteration, verdict=verdict)
+                baseline_stopped = True
+                break
+        if stop:
+            append_timeline(cfg.state_dir, "round_failed", {
+                "iteration": iteration, "stage": stage,
+                "reason": str(stop.get("reason") or "stop requested"),
+                "detail": {k: v for k, v in stop.items() if k != "iteration"},
+            })
+            try:
+                (exp / "stop_requested.json").unlink()
+            except OSError:
+                pass
+            save_json(input_dir / "benchmark" / "results.json", {
+                "iteration": iteration, "results": results,
+                "round_failed": str(stop.get("reason") or "stop requested"),
+            })
+            if _uses_variant_rounds(cfg):
+                if stage == STAGE_GENERALIZATION:
+                    # The retake is spent: a failed retake is a failed
+                    # candidate, so the frozen request is cleared and the next
+                    # pass is a normal round again. Its budget was already spent
+                    # when it was launched.
+                    clear_pending_generalization(exp)
+                else:
+                    rounds_used += PERFORMANCE_ROUND_COST
+                _record_state(cfg, ANALYZE, iteration,
+                              event="round_failed_retry",
+                              rounds_used=rounds_used)
+                iteration += 1
+                continue
+            _record_state(cfg, FINISHED, iteration, event="stop_requested",
+                          **{k: v for k, v in stop.items() if k != "iteration"})
+            set_run(cfg.state_dir,
+                    final_status=str(stop.get("reason") or "stopped"))
+            break
+
+        # --- variant stages: baseline, then performance, then generalization -
+        # Every comparison is against each operator's *current variant* (the
+        # best kernel the variant pool holds). Which gates run depends on the
+        # stage:
+        #   baseline       — judges nothing; its faster kernels go into the
+        #                    variant pool so later rounds have a real bar;
+        #   performance    — fresh random operators, performance gate only;
+        #   generalization — the *same* harness handed to DKAO a second time
+        #                    with the frozen paper, generalization gate only.
+        # Only a generalization pass that wins parks the candidate for a human.
+        if _uses_variant_rounds(cfg):
+            perf_ids = list(selected)
+            perf_gate: Optional[Dict[str, Any]] = None
+            general_gate: Optional[Dict[str, Any]] = None
+            mechanism: Dict[str, Any] = {}
+            baseline_promotion: Optional[Dict[str, Any]] = None
+            reference_moved: List[str] = []
+
+            if stage == STAGE_BASELINE:
+                # The baseline round owes evidence, not a verdict.
+                baseline_verdict = baseline_round_verdict(results, perf_ids)
+                if baseline_verdict["status"] != "PASS":
+                    _baseline_stop(cfg, exp, iteration=iteration,
+                                   verdict=baseline_verdict)
+                    _record_state(cfg, ANALYZE, iteration,
+                                  event="baseline_round_failed")
+                    break
+                # Faster kernels are a kernel fact, not a harness verdict: they
+                # enter the variant pool immediately, so every later gate — and
+                # DKAO's own warm start — compares against what this round
+                # achieved. Nothing about the *harness* moves here.
+                baseline_promotion = _promote_baseline_kernels(
+                    cfg, exp, variant_pool, results)
+                record_variants(
+                    exp, variant_table, results,
+                    version=str(variant_table.get("harness_version") or "seed"),
+                    iteration=iteration,
+                    groups=variant_table.get("groups"))
+                reference_moved = sorted(
+                    f"{u['operator_id']} ({u.get('previous_best_us')} -> "
+                    f"{u.get('median_us')} us)"
+                    for u in ((baseline_promotion or {}).get("updated") or []))
+                save_json(input_dir / "baseline_round.json", {
+                    "iteration": iteration, "stage": stage,
+                    "verdict": baseline_verdict,
+                    "pool_promotion": baseline_promotion,
+                    "reference_moved": reference_moved,
+                })
+                append_timeline(cfg.state_dir, "baseline_round_measured", {
+                    "iteration": iteration,
+                    "operators": perf_ids,
+                    "pool_promoted": reference_moved,
+                    "paper": perf_ids,
+                })
+
+            if stage != STAGE_BASELINE:
+                # One gate per stage. The performance stage judges the fresh
+                # draw; the generalization stage judges the paper it was
+                # launched to re-measure.
+                kind = ("generalization" if stage == STAGE_GENERALIZATION
+                        else "performance")
+                gate = evaluate_gate(kind=kind, operator_ids=perf_ids,
+                                     results=results, table=variant_table,
+                                     instances=variant_pool)
+                gate["iteration"] = iteration
+                gate["stage"] = stage
+                if kind == "performance":
+                    perf_gate = gate
+                    save_json(input_dir / "performance_gate.json", gate)
+                    # The round's own results, where every reader of this run
+                    # expects them (the WebUI, resume, and the round review).
+                    save_json(input_dir / "benchmark" / "results.json",
+                              {"iteration": iteration, "stage": stage,
+                               "results": results})
+                else:
+                    general_gate = gate
+                    save_json(input_dir / "generalization_gate.json", gate)
+                    save_json(input_dir / "benchmark"
+                              / "generalization_results.json",
+                              {"iteration": iteration, "results": results})
+                    save_json(input_dir / "benchmark" / "results.json",
+                              {"iteration": iteration, "stage": stage,
+                               "results": results})
+                append_timeline(cfg.state_dir, f"{kind}_gate", {
+                    "iteration": iteration, "stage": stage,
+                    "status": gate["status"], "wins": gate["wins"],
+                    "wins_needed": gate["wins_needed"],
+                    "below_floor": gate["below_floor"],
+                    "unmeasured": gate["unmeasured"],
+                    "operators": gate["operators"],
+                    **({"paper": perf_ids}
+                       if kind == "generalization" else {}),
+                })
+
+            # Mechanism gate (observability, lenient): a harness change must say
+            # *why* it should help and produce evidence for it. A contradicted
+            # mechanism does not veto an accepted win — the counter gates do
+            # that — but it is recorded, so "we promoted it and cannot explain
+            # it" is visible instead of silent.
+            if stage != STAGE_BASELINE:
+                mechanism = _mechanism_gate_for_round(cfg, exp, iteration)
+
+            if stage != STAGE_BASELINE and gate.get("status") == "INCOMPLETE":
+                # The round drew questions but did not measure all of them, and
+                # the missing ones still leave the outcome open. There is no
+                # verdict to reach: judging here would blame the harness for a
+                # number the machine never produced — and on the retake it would
+                # spend the candidate on it. Record it as an incomplete round and
+                # try again; a frozen retake stays frozen, because the paper is
+                # still owed a judgement.
+                append_timeline(cfg.state_dir, "round_incomplete", {
+                    "iteration": iteration,
+                    "stage": stage,
+                    "gate": ("generalization" if stage == STAGE_GENERALIZATION
+                             else "performance"),
+                    "unmeasured": gate["unmeasured"],
+                    "measured": sorted(set(results) - set(gate["unmeasured"])),
+                    "reasons": gate["reasons"],
+                    "rounds_used": rounds_used,
+                    "policy": ("no harness verdict this round; the missing "
+                               "questions are measured in a later pass"),
+                })
+                if stage != STAGE_GENERALIZATION:
+                    # A fresh round was spent, verdict or not.
+                    rounds_used += PERFORMANCE_ROUND_COST
+                restore_promoted_harness(cfg, exp, live_workspace)
+                _record_state(cfg, ANALYZE, iteration,
+                              event="round_incomplete",
+                              stage=stage,
+                              unmeasured=gate["unmeasured"],
+                              rounds_used=rounds_used)
+                iteration += 1
+                continue
+
+            both_passed = bool(
+                stage == STAGE_GENERALIZATION and general_gate
+                and general_gate["status"] == "PASS")
+            if stage == STAGE_BASELINE:
+                verdict, action = "BASELINE_MEASURED", "ITERATE"
+                reason = ("baseline round: the paper was measured and the "
+                          "faster kernels entered the variant pool; no gate is "
+                          "judged this round")
+            elif both_passed:
+                verdict, action = "PROMOTE_CANDIDATE", "AWAIT_APPROVAL"
+                reason = ("performance and generalization gates passed; "
+                          "waiting for human approval before publishing the "
+                          "harness to DKAO")
+            elif stage == STAGE_GENERALIZATION:
+                verdict, action = "REJECT", "RESTORE_CHAMPION"
+                reason = ("generalization gate "
+                          f"{(general_gate or {}).get('status')}")
+            else:
+                verdict, action = "REJECT", "RESTORE_CHAMPION"
+                reason = f"performance gate {(perf_gate or {}).get('status')}"
+            decision = {
+                "iteration": iteration,
+                "stage": stage,
+                "round_cost": {
+                    STAGE_BASELINE: 0.0,
+                    STAGE_GENERALIZATION: GENERALIZATION_ROUND_COST,
+                }.get(stage, PERFORMANCE_ROUND_COST),
+                "verdict": verdict,
+                "action": action,
+                "reason": reason,
+                "performance_gate": perf_gate,
+                "generalization_gate": general_gate,
+                "mechanism_gate": mechanism,
+                "baseline_promotion": baseline_promotion,
+                "variant_reference": {
+                    iid: {"median_us": row.get("median_us"),
+                          "source": row.get("variant_source")}
+                    for iid, row in ((perf_gate or general_gate or {})
+                                     .get("comparison", {})
+                                     .get("per_instance") or {}).items()
+                },
+            }
+            save_json(input_dir / "decision.json", decision)
+            # The baseline round answers to no gate, so the legacy decision
+            # engine (which compares against a frozen champion) has nothing to
+            # say about it and must not overwrite this verdict.
+            own_decision = stage == STAGE_BASELINE
+            score_row = pass_rate(results)
+            with (exp / SCORES_FILE).open("a", encoding="utf-8") as fh:
+                fh.write(json.dumps({
+                    **score_row, "iteration": iteration, "stage": stage,
+                    "rounds_used": rounds_used,
+                    "decision_pass_rate": score_row.get("pass_rate"),
+                    "ts": time.time(),
+                }, ensure_ascii=False) + "\n")
+            _record_state(cfg, ANALYZE, iteration, event="iteration_decision",
+                          stage=stage,
+                          verdict=decision["verdict"],
+                          action=decision["action"],
+                          rounds_used=rounds_used,
+                          performance_gate=(perf_gate or {}).get("status"),
+                          generalization_gate=(general_gate or {}).get("status"),
+                          mechanism_gate=mechanism.get("status"))
+            if both_passed:
+                # The harness that just proved itself is the one parked for
+                # approval: the revision frozen when the retake was scheduled,
+                # and the round number that *won* the performance gate — the
+                # retake is half a round, not a round of its own.
+                write_pending_promotion(
+                    exp, iteration=int(pending_retake.get("iteration")
+                                       or iteration),
+                    revision=(pending_retake.get("harness_revision")
+                              or _head_revision(live_workspace)),
+                    snapshot_dir=str(snapshot),
+                    performance_gate=perf_gate or general_gate or {},
+                    generalization_gate=general_gate,
+                    results=results)
+                set_run(cfg.state_dir, final_status="awaiting_approval",
+                        rounds_used=rounds_used)
+                append_timeline(cfg.state_dir, "awaiting_approval", {
+                    "iteration": iteration,
+                    "performance_wins": (pending_retake.get("performance_wins")
+                                         or []),
+                    "generalization_wins": general_gate["wins"],
+                })
+                return _final_report(cfg)
+            if stage == STAGE_PERFORMANCE:
+                # The performance gate passed: do NOT re-measure the paper in
+                # this pass. Freeze the retake — same harness bytes, the paper's
+                # operators — and let the next pass hand that job to DKAO.
+                if (perf_gate or {}).get("status") == "PASS":
+                    pending = write_pending_generalization(
+                        exp, iteration=iteration,
+                        harness_revision=_head_revision(live_workspace),
+                        snapshot_dir=str(snapshot),
+                        operator_ids=round_questions.get("generalization_ids")
+                        or perf_ids,
+                        performance_gate=perf_gate or {})
+                    # Both halves are spent the moment the retake is launched:
+                    # the empty half is a launch in which the whole frozen
+                    # document is retested. Accounting for it here (rather than
+                    # when the retake is judged) also keeps the budget honest on
+                    # the run's last pass, which ends inside the retake.
+                    rounds_used += (PERFORMANCE_ROUND_COST
+                                    + GENERALIZATION_ROUND_COST)
+                    append_timeline(cfg.state_dir, "generalization_retake_frozen", {
+                        "iteration": iteration,
+                        "operators": pending["operator_ids"],
+                        "harness_revision": pending["harness_revision"],
+                        "rounds_used": rounds_used,
+                    })
+                    # This pass is spent: the retake owns the next one. It runs
+                    # as its own stage, with its own DKAO launch of the frozen
+                    # paper on the very same harness bytes.
+                    iteration += 1
+                    continue
+                rounds_used += PERFORMANCE_ROUND_COST
+                restore_promoted_harness(cfg, exp, live_workspace)
+                iteration += 1
+                continue
+            if stage == STAGE_GENERALIZATION:
+                # The retake did not win: the candidate is spent, so the frozen
+                # request is cleared and the next pass is a normal round again.
+                # Its budget was already spent when it was launched.
+                clear_pending_generalization(exp)
+                restore_promoted_harness(cfg, exp, live_workspace)
+                iteration += 1
+                continue
+            # The baseline round judges nothing and keeps the iteration: the
+            # harness it just measured still has to evolve before the next
+            # round can test it.
+            iteration += 1
+            continue
+        # Pool mode owns the pass criterion: turn real/dry measurements into an
+        # automatic family-baseline pass decision (no user threshold required).
+        if pool is not None:
+            for iid, result in results.items():
+                median = result.get("median_us")
+                if iid in pool.instances and isinstance(median, (int, float)):
+                    auto = pool.auto_pass(
+                        iid, float(median),
+                        correctness_ok=bool(result.get("correctness_ok")),
+                        p90_ok=bool(result.get("p90_ok")),
+                    )
+                    result["raw_evaluator_passed"] = result.get("passed")
+                    result["passed"] = auto["passed"]
+                    result["automatic_pass"] = auto
+        save_json(input_dir / "benchmark" / "results.json", {
+            "iteration": iteration,
+            "results": results,
+        })
+        stats = pass_rate(results)
+
+        # overlap decision stats (pool mode keeps decisions comparable)
+        overlap_stats = None
+        if prev_results is not None:
+            overlap = {
+                k: v for k, v in results.items() if k in prev_results
+            }
+            if overlap:
+                overlap_stats = pass_rate(overlap)
+
+        decision_rate = (
+            float(overlap_stats["pass_rate"])
+            if pool is not None and overlap_stats is not None
+            else float(stats["pass_rate"])
+        )
+        _save_scores(cfg, {
+            "iteration": iteration,
+            **stats,
+            "overlap_pass_rate": (
+                overlap_stats["pass_rate"] if overlap_stats is not None else None
+            ),
+            "decision_pass_rate": decision_rate,
+            "ts": time.time(),
+        })
+
+        diff = diff_results(prev_results, results)
+        save_json(input_dir / "diff.json", diff)
+        _write_overview(input_dir, iteration, results, diff)
+
+        # --- change attribution for the PREVIOUS loop's manifest -----------
+        prev_manifest = None
+        if iteration > 1:
+            prev_manifest = _load_json(
+                (iter_dir.parent / f"iteration_{iteration - 1:03d}"
+                 / "evolve" / "change_manifest.json")
+            )
+            if prev_manifest:
+                save_json(
+                    input_dir / "change_evaluation.json",
+                    evaluate_changes(prev_manifest, diff),
+                )
+
+        # --- system-level decision engine ----------------------------------
+        # Performance is primary; mechanism evidence explains a win but cannot
+        # veto a repeated real improvement. Held-out inputs land in a later
+        # adapter iteration; current v1 writes NOT_RUN when absent.
+        # Mechanism gate: grade the previous loop's declared changes against
+        # what the candidate iteration actually did (planner plan ids, the
+        # harness revision children loaded, gate events).
+        mechanism_checks_map = None
+        if prev_manifest is not None:
+            try:
+                from .mechanism_checks import (
+                    collect_candidate_evidence,
+                    evaluate_changes as evaluate_mechanism_checks,
+                )
+                evidence = collect_candidate_evidence(
+                    exp, iteration,
+                    planner_enabled=_truthy(
+                        cfg.answers.get("planner_enabled"), default=True),
+                    candidate_workspace=snapshot,
+                )
+                mechanism_checks_map, mechanism_detail = (
+                    evaluate_mechanism_checks(prev_manifest, evidence))
+                save_json(input_dir / "mechanism_evidence.json", {
+                    "iteration": iteration,
+                    "evidence": evidence,
+                    "checks": mechanism_checks_map,
+                    "detail": mechanism_detail,
+                })
+            except Exception as exc:  # noqa: BLE001 - never break the loop
+                append_timeline(cfg.state_dir, "mechanism_check_failed",
+                                {"iteration": iteration, "error": repr(exc)})
+
+        wired_violations: List[Dict[str, str]] = []
+        try:
+            previous_workspace = (
+                exp / "runs" / f"iteration_{iteration - 1:03d}"
+                / "input" / "workspace"
+            )
+            if iteration > 1 and previous_workspace.is_dir():
+                wired_violations = _wired_scope_violations(
+                    previous_workspace, live_workspace)
+                if wired_violations:
+                    append_timeline(cfg.state_dir,
+                                    "harness_wired_scope_violation",
+                                    {"iteration": iteration,
+                                     "components": wired_violations})
+        except Exception:  # noqa: BLE001 - advisory only
+            wired_violations = []
+
+        # No pairing is dropped for a shared device any more: HE does not watch
+        # the cards, and DKAO's admission gate refuses to time a kernel on a
+        # busy one in the first place (a blocked round produces no measurement
+        # at all, which the round handles as an environment failure).
+        comparison_candidate = results
+        comparison_champion = champion_results
+
+        if not own_decision:
+            decision = decide(
+                comparison_champion,
+                comparison_candidate,
+                manifest=prev_manifest,
+                mechanism_checks=mechanism_checks_map,
+                confirmation_reproduced=pending_confirmation,
+                purposes=purposes,
+                policy=DecisionPolicy(
+                    mechanism_policy=str(
+                        cfg.answers.get("mechanism_policy") or "lenient"),
+                ),
+            )
+            decision.update({
+                "iteration": iteration,
+                "candidate_snapshot": str(snapshot),
+                "decision_pass_rate": decision_rate,
+                "wired_scope_violations": wired_violations,
+            })
+            save_json(input_dir / "decision.json", decision)
+            champion_outcome = _record_champion_generation(
+                cfg, exp, iteration=iteration, decision=decision,
+                decision_rate=decision_rate, snapshot=snapshot,
+                live_workspace=live_workspace, results=results,
+                current=best_ever)
+            best_ever = champion_outcome.get("best_ever", best_ever)
+            _record_state(cfg, ANALYZE, iteration, event="iteration_decision",
+                          verdict=decision["verdict"], action=decision["action"],
+                          decision_pass_rate=decision_rate)
+        try:
+            if not _uses_variant_rounds(cfg):
+                # Legacy path only: promote this round's kernels immediately.
+                # Under the variant protocol nothing reaches production until a
+                # harness is accepted and approved, so a rejected round must not
+                # leave its kernels behind (operator rule: only the accepted
+                # round's optimizations survive).
+                _promote_iteration(cfg, exp, iteration, results, instances,
+                                   freeze=(_variant_policy(cfg) == "freeze"))
+            _track_best_known(cfg, exp, iteration, results)
+        except Exception as exc:  # noqa: BLE001 - never break the outer loop
+            append_timeline(cfg.state_dir, "kernel_promotion_failed",
+                            {"iteration": iteration, "error": repr(exc)})
+
+        # No circuit breaker: under the variant protocol a round that does not
+        # beat the current variants is a normal outcome, and the run's budget
+        # (max_rounds) is what bounds it. Stopping early on "three rejections"
+        # would kill a run that was still making progress on its own terms.
+        verdict = str(decision.get("verdict") or "")
+
+        if own_decision:
+            # The variant-round protocol decides for itself (see the gates
+            # above): there is no frozen-champion action to take here. The
+            # baseline round measured the paper and moved the variant pool; a
+            # passing performance stage re-launches as the generalization
+            # retake; a failing gate already restored the promoted harness.
+            iteration += 1
+            continue
+
+        action = decision["action"]
+        if action == "SET_CHAMPION":
+            best_ever = {
+                "iteration": iteration,
+                "pass_rate": decision_rate,
+                "snapshot_dir": str(snapshot),
+                "workspace_revision": _head_revision(live_workspace),
+                "verdict": decision["verdict"],
+            }
+            save_json(exp / BEST_EVER_FILE, best_ever)
+            champion_results = dict(results)
+            pending_confirmation = False
+        elif action in {"RESTORE_CHAMPION", "KEEP_CHAMPION_ARCHIVE_CANDIDATE",
+                        "ARCHIVE_SPECIALIZED"}:
+            if action == "ARCHIVE_SPECIALIZED":
+                copy_tree(snapshot, exp / "specialized" / f"iteration_{iteration:03d}")
+            best_dir = Path(str(best_ever.get("snapshot_dir", "")))
+            if best_dir.is_dir():
+                replace_tree(best_dir, live_workspace)
+                best_rev = best_ever.get("workspace_revision")
+                if best_rev:
+                    subprocess.run(
+                        ["git", "reset", "--hard", str(best_rev)],
+                        cwd=live_workspace, check=True,
+                        stdout=subprocess.DEVNULL,
+                    )
+            disposition = {
+                "iteration": iteration,
+                "verdict": decision["verdict"],
+                "action": action,
+                "champion_snapshot": str(best_dir),
+                "candidate_snapshot": str(snapshot),
+                "reason": decision["reason"],
+            }
+            save_json(iter_dir / "candidate_disposition.json", disposition)
+            # rollback.json is reserved for true reject/overfit, not no-signal.
+            if decision["verdict"] in {"REJECT", "REJECT_OVERFIT"}:
+                save_json(iter_dir / "rollback.json", disposition)
+            pending_confirmation = False
+        elif action == "ARCHIVE_AND_CONFIRM":
+            # Leave the candidate live for one confirmation round; no champion
+            # switch yet. A reproduced win can become PROMOTE_UNEXPLAINED.
+            save_json(iter_dir / "candidate_disposition.json", {
+                "iteration": iteration,
+                "verdict": decision["verdict"],
+                "action": action,
+                "candidate_snapshot": str(snapshot),
+                "reason": decision["reason"],
+            })
+            pending_confirmation = True
+
+        # --- analyze -> evolve ---------------------------------------------
+        # A confirmation round must retest the SAME candidate; do not let the
+        # Evolve Agent mutate it between the first unexplained win and repeat.
+        final_round = (
+            (target is not None and iteration >= int(target))
+            or (target is None and iteration >= fallback_end)
+        )
+        if final_round:
+            # No next round exists to evaluate a new generation: evolving here
+            # would leave an unverified harness behind.
+            evolve_out = {
+                "summary": (
+                    f"# Final round — iteration {iteration}\n\n"
+                    "Evolve skipped: the target round was reached, so no "
+                    "unverified next generation is produced."
+                ),
+                "manifest": {
+                    "schema_version": 1,
+                    "iteration": iteration,
+                    "changes": [],
+                    "verification": {"status": "target_reached_skip_evolve"},
+                },
+                "round_plan": None,
+                "changed": False,
+            }
+            append_timeline(cfg.state_dir, "evolve_skipped_final_round", {
+                "iteration": iteration, "target": target,
+                "fallback_end": fallback_end,
+            })
+        elif action == "ARCHIVE_AND_CONFIRM":
+            evolve_out = {
+                "summary": (
+                    f"# Confirmation scheduled — iteration {iteration}\n\n"
+                    "Candidate kept byte-identical for one repeat evaluation."
+                ),
+                "manifest": {
+                    "schema_version": 1,
+                    "iteration": iteration,
+                    "changes": [],
+                    "verification": {"status": "confirmation_pending"},
+                },
+                "round_plan": None,
+                "changed": False,
+            }
+        else:
+            evolve_out = evolver.apply(
+                live_workspace, iteration,
+                query_context={
+                    "overview": str(input_dir / "analysis" / "overview.md"),
+                    "results": results,
+                    "diff": diff,
+                    "decision": decision,
+                    "selected": (list(selected) if pool is not None else
+                                 [i.id for i in instances]),
+                    "pool_ids": (pool.all_ids() if pool is not None else
+                                 [i.id for i in instances]),
+                },
+            )
+        _record_state(cfg, EVOLVE, iteration, event="evolve_start",
+                      evolve_mode=cfg.answers.get("evolve_mode"))
+        evolve_dir.mkdir(parents=True, exist_ok=True)
+        revision = _commit_workspace(live_workspace, iteration)
+        evolve_out["manifest"]["workspace_revision"] = revision
+        save_json(evolve_dir / "change_manifest.json", evolve_out["manifest"])
+        (evolve_dir / "evolve_summary.md").write_text(
+            evolve_out["summary"] + "\n", encoding="utf-8"
+        )
+        if pool is not None:
+            agent_plan = evolve_out.get("round_plan")
+            if isinstance(agent_plan, dict):
+                plan_obj = RoundPlan.from_dict(agent_plan)
+                # Validate agent output immediately. The next iteration applies
+                # the overlap/budget guardrail and can auto-fill if needed.
+                pool.assert_in_pool(plan_obj.selected)
+            else:
+                plan_obj = RoundPlan(
+                    iteration=iteration,
+                    strategy="balanced",
+                    rationale="dry-run default: keep same scored set",
+                    selected=list(selected),
+                )
+            save_json(evolve_dir / "round_plan.json", plan_obj.to_dict())
+        if evolve_out.get("changed"):
+            copy_tree(live_workspace, evolve_dir / "workspace")
+
+        prev_results = results
+        iteration += 1
+
+        # Round boundary: apply a code change instead of running another round
+        # of stale logic. The supervisor/server restarts the process, which
+        # records its own revision and resumes at this iteration.
+        reload_reason = _code_reload_pending(cfg, loaded_revision)
+        if reload_reason is not None:
+            reload_reason["resumes_at_iteration"] = iteration
+            reload_reason["finished_iteration"] = iteration - 1
+            append_timeline(cfg.state_dir, "code_reload_deferred", reload_reason)
+            _record_state(
+                cfg, PREPARE, iteration - 1,
+                event="code_reload_deferred",
+                loaded_revision=reload_reason["loaded_revision"],
+                on_disk_revision=reload_reason["on_disk_revision"],
+                resumes_at_iteration=iteration,
+            )
+            set_run(cfg.state_dir,
+                    final_status=f"code_reload_pending:{iteration}")
+            return _final_report(cfg)
+
+    if baseline_stopped:
+        # The task already ended on the baseline round: report and stop without
+        # touching final_status (which ``_baseline_stop`` owns).
+        _record_state(cfg, REPORT, event="experiment_report")
+        return _final_report(cfg)
+    try:
+        # Variant-round protocol promotes kernels through the pool (the baseline
+        # round) and the approved round, never through the frozen-candidate
+        # sweep: running it here would promote kernels from rejected rounds.
+        if _variant_policy(cfg) == "freeze" and not _uses_variant_rounds(cfg):
+            _finalize_promotion(cfg, exp)
+    except Exception as exc:  # noqa: BLE001 - reporting must still happen
+        append_timeline(cfg.state_dir, "final_promotion_failed",
+                        {"error": repr(exc)})
+    try:
+        _write_round_review(cfg, exp)
+    except Exception as exc:  # noqa: BLE001
+        append_timeline(cfg.state_dir, "round_review_failed",
+                        {"error": repr(exc)})
+    _record_state(cfg, REPORT, event="experiment_report")
+    report = _final_report(cfg)
+    _record_state(cfg, FINISHED, event="experiment_finished")
+    return report
+
+
+def _final_report(cfg: ExperimentConfig) -> Path:
+    exp = cfg.exp_dir
+    lines = [
+        "# harness_evolve report (AHE outer loop)",
+        "",
+        f"- task: {cfg.task_id}",
+        f"- iterations: {cfg.max_iterations}",
+        f"- execution_mode: {cfg.execution_mode}",
+        f"- suite size: {len(cfg.suite)}",
+        f"- pool_mode: {cfg.pool_mode}",
+    ]
+    best = _load_json(exp / BEST_EVER_FILE) or {}
+    if best:
+        lines.append(
+            f"- best-ever: iteration {best.get('iteration')} "
+            f"pass_rate {best.get('pass_rate')}"
+        )
+    lines.append("")
+    lines.append("See runs/iteration_*/ for per-iteration artifacts.")
+    report = exp / REPORT_FILE
+    report.write_text("\n".join(lines) + "\n", encoding="utf-8")
+    return report
diff --git a/metainfer/tasks/harness_evolve/orchestrator/plugin.py b/metainfer/tasks/harness_evolve/orchestrator/plugin.py
new file mode 100644
index 00000000..094d62bd
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/plugin.py
@@ -0,0 +1,11 @@
+"""Task plugin descriptor for harness_evolve."""
+
+from metainfer.orchestrator.tasks.base import TaskPlugin
+
+
+PLUGIN = TaskPlugin(
+    task_type="harness-evolve",
+    cli_module="metainfer.tasks.harness_evolve.orchestrator.cli",
+    phases_module="metainfer.tasks.harness_evolve.orchestrator.phases",
+    diagnostic_globs=("*.json", "*.jsonl", "*.md", "*.log"),
+)
diff --git a/metainfer/tasks/harness_evolve/orchestrator/pool.py b/metainfer/tasks/harness_evolve/orchestrator/pool.py
new file mode 100644
index 00000000..132cbec0
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/pool.py
@@ -0,0 +1,241 @@
+"""Registered-instance pool + automatic pass criteria (suite protocol v2).
+
+A "question" can only come from the pool: instances the user registered with a
+frozen Triton baseline and a valid contract (they are the operators actually
+integrated & measured, e.g. DeepSeek INT8 W8A8 GEMM family). AHE decides pass
+automatically (``tau_family`` derived from each family's accepted history) and
+picks per-round question sets via round_plan + machine guardrails (see
+``docs/dkao_suite_protocol_v2.md``).
+"""
+
+from __future__ import annotations
+
+import json
+import statistics
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set
+
+DEFAULT_TAU = 0.6  # cold-start default: must beat baseline by >= 1/DEFAULT_TAU? see auto_pass
+
+
+def family_of(M: int, operator: str) -> str:
+    regime = "decode" if M <= 32 else "prefill"
+    return f"{regime}__{operator}"
+
+
+@dataclass
+class PoolInstance:
+    id: str
+    model: str
+    tp_size: int
+    operator: str
+    M: int
+    N: int
+    K: int
+    baseline_us: float
+    family: str = ""
+    best_known_us: Optional[float] = None
+    history: List[Dict[str, Any]] = field(default_factory=list)
+    last_selected_round: Optional[int] = None
+
+    @property
+    def key(self) -> str:
+        return self.id
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "id": self.id,
+            "contract": {
+                "model": self.model, "tp_size": self.tp_size,
+                "operator": self.operator, "M": self.M, "N": self.N,
+                "K": self.K,
+            },
+            "family": self.family or family_of(self.M, self.operator),
+            "baseline_us": self.baseline_us,
+            "best_known_us": self.best_known_us,
+            "history": self.history,
+            "last_selected_round": self.last_selected_round,
+        }
+
+
+def _entry_to_instance(d: Dict[str, Any]) -> PoolInstance:
+    c = d.get("contract") or {}
+    baseline = d.get("baseline_us")
+    if not isinstance(baseline, (int, float)) or baseline <= 0:
+        raise ValueError(f"instance {d.get('id')} needs positive baseline_us")
+    inst = PoolInstance(
+        id=str(d["id"]),
+        model=str(c.get("model", "")),
+        tp_size=int(c.get("tp_size", 8)),
+        operator=str(c.get("operator", "")),
+        M=int(c.get("M", 0)),
+        N=int(c.get("N", 0)),
+        K=int(c.get("K", 0)),
+        baseline_us=float(baseline),
+        family=str(d.get("family") or family_of(int(c.get("M", 0)),
+                                                 str(c.get("operator", "")))),
+        best_known_us=(
+            float(d["best_known_us"]) if d.get("best_known_us") else None
+        ),
+        history=list(d.get("history") or []),
+        last_selected_round=(
+            int(d["last_selected_round"]) if d.get("last_selected_round") is not None
+            else None
+        ),
+    )
+    return inst
+
+
+class Pool:
+    """Registered question pool (file-backed or built-in fixture)."""
+
+    def __init__(self, instances: Optional[Dict[str, PoolInstance]] = None):
+        self.instances: Dict[str, PoolInstance] = instances or {}
+
+    @classmethod
+    def from_yaml(cls, path: Path) -> "Pool":
+        import yaml
+        data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
+        entries = data.get("instances") if isinstance(data, dict) else data
+        pool = cls()
+        for d in entries or []:
+            inst = _entry_to_instance(d)
+            pool.register(inst)
+        return pool
+
+    def register(self, inst: PoolInstance) -> None:
+        if inst.id in self.instances:
+            raise ValueError(f"duplicate pool instance: {inst.id}")
+        self.instances[inst.id] = inst
+
+    def get(self, inst_id: str) -> PoolInstance:
+        return self.instances[inst_id]
+
+    def all_ids(self) -> List[str]:
+        return sorted(self.instances)
+
+    def assert_in_pool(self, ids: List[str]) -> None:
+        unknown = [i for i in ids if i not in self.instances]
+        if unknown:
+            raise ValueError(f"instances not in pool: {unknown}")
+
+    # -- per-family automatic target (tau) --------------------------------
+    def family_ratios(self, family: str) -> List[float]:
+        ratios = []
+        for inst in self.instances.values():
+            if (inst.family or family_of(inst.M, inst.operator)) != family:
+                continue
+            for rec in inst.history:
+                if rec.get("accepted") is True and rec.get("median_us"):
+                    try:
+                        r = float(rec["median_us"]) / float(inst.baseline_us)
+                    except (TypeError, ValueError, ZeroDivisionError):
+                        continue
+                    if 0 < r < 1.0:
+                        ratios.append(r)
+        return ratios
+
+    def tau_family(self, family: str, quantile: float = 0.6,
+                   default: float = DEFAULT_TAU) -> float:
+        ratios = self.family_ratios(family)
+        if not ratios:
+            return default
+        ratios.sort()
+        idx = min(len(ratios) - 1, max(0, int(len(ratios) * quantile)))
+        return max(0.05, ratios[idx])
+
+    # -- auto pass ---------------------------------------------------------
+    def auto_pass(self, inst_id: str, median_us: float, *,
+                  correctness_ok: bool, p90_ok: bool) -> Dict[str, Any]:
+        inst = self.get(inst_id)
+        tau = self.tau_family(inst.family or family_of(inst.M, inst.operator))
+        target_us = tau * float(inst.baseline_us)
+        perf_ok = median_us <= target_us
+        passed = bool(correctness_ok and perf_ok and p90_ok)
+        return {
+            "instance": inst_id,
+            "passed": passed,
+            "tau_family": tau,
+            "target_us": round(target_us, 3),
+            "median_us": round(float(median_us), 3),
+            "correctness_ok": bool(correctness_ok),
+            "p90_ok": bool(p90_ok),
+        }
+
+    # -- coverage / sampling helpers ---------------------------------------
+    def stratified_sample(self, budget: int,
+                          exclude: Optional[Set[str]] = None) -> List[str]:
+        """Deterministic coverage-first sample: at least one per family/M-regime
+        bucket when budget allows, then fill by insertion order."""
+        exclude = exclude or set()
+        by_family: Dict[str, List[str]] = {}
+        for iid in self.all_ids():
+            if iid in exclude:
+                continue
+            inst = self.instances[iid]
+            fam = inst.family or family_of(inst.M, inst.operator)
+            by_family.setdefault(fam, []).append(iid)
+        picked: List[str] = []
+        seen: Set[str] = set()
+        for fam in sorted(by_family):
+            if len(picked) >= budget:
+                break
+            iid = by_family[fam][0]
+            picked.append(iid)
+            seen.add(iid)
+        for iid in self.all_ids():
+            if len(picked) >= budget:
+                break
+            if iid not in seen and iid not in exclude:
+                picked.append(iid)
+                seen.add(iid)
+        return picked
+
+    def historical_ids(self) -> List[str]:
+        """Instances with at least one accepted history record (anchor set)."""
+        return [
+            iid for iid, inst in self.instances.items() if inst.history
+        ]
+
+    def heldout_ids(self, k: int, *,
+                    exclude: Optional[Set[str]] = None) -> List[str]:
+        """Least-recently-selected instances (lightweight held-out probe)."""
+        exclude = exclude or set()
+        candidates = [
+            (inst.last_selected_round if inst.last_selected_round is not None
+             else -10**9, iid)
+            for iid, inst in self.instances.items()
+            if iid not in exclude
+        ]
+        candidates.sort()
+        return [iid for _, iid in candidates[:k]]
+
+
+# ---------------------------------------------------------------------------
+# Built-in fixture pool (dry-run / tests). Real deployment loads a user pool
+# file (registered operators with measured baselines).
+# ---------------------------------------------------------------------------
+def _default_instances() -> Dict[str, PoolInstance]:
+    rows = [
+        ("dsv4_tp8_qkv_proj_m16", "deepseek-v4", 8, "qkv_proj", 16, 1280, 4096, 100.0, 45.0),
+        ("hy3_tp8_qkv_proj_m16", "hy3", 8, "qkv_proj", 16, 1280, 4096, 90.0, 40.0),
+        ("hy3_tp8_o_proj_m16", "hy3", 8, "o_proj", 16, 4096, 1024, 60.0, 22.0),
+        ("glm52_tp4_gate_up_m16", "glm52", 4, "shared_gate_up_proj", 16, 1024, 4096, 70.0, 30.0),
+        ("hy3_tp8_qkv_proj_m4096", "hy3", 8, "qkv_proj", 4096, 1280, 4096, 10000.0, 5000.0),
+        ("hy3_tp8_o_proj_m4096", "hy3", 8, "o_proj", 4096, 4096, 1024, 8000.0, 3000.0),
+        ("hy3_tp8_gate_up_m4096", "hy3", 8, "shared_gate_up_proj", 4096, 384, 4096, 7000.0, 2500.0),
+    ]
+    insts = {}
+    for iid, model, tp, op, m, n, k, base, best in rows:
+        hist = [{"accepted": True, "median_us": best}]
+        insts[iid] = PoolInstance(
+            id=iid, model=model, tp_size=tp, operator=op, M=m, N=n, K=k,
+            baseline_us=base, family=family_of(m, op),
+            best_known_us=best, history=hist, last_selected_round=None,
+        )
+    return insts
+
+
+def builtin_pool() -> Pool:
+    return Pool(_default_instances())
diff --git a/metainfer/tasks/harness_evolve/orchestrator/pool_builder.py b/metainfer/tasks/harness_evolve/orchestrator/pool_builder.py
new file mode 100644
index 00000000..ab41f8b6
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/pool_builder.py
@@ -0,0 +1,105 @@
+"""Build a measured AHE question pool from historical DKAO final reports.
+
+Sources are read-only (MetaInfer nodes/*/workspaces/*/final_report.json).
+Output belongs under the isolated AHE root, e.g.
+/root/zth_agent/ahe-kernel-repos/registered_pool.yaml.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+import yaml
+
+from .pool import family_of
+
+
+def build_pool(meta_root: Path) -> Dict[str, Any]:
+    entries: Dict[str, Dict[str, Any]] = {}
+    for report_path in sorted(meta_root.glob("nodes/*/workspaces/*/final_report.json")):
+        try:
+            report = json.loads(report_path.read_text(encoding="utf-8"))
+        except (OSError, ValueError):
+            continue
+        config = report.get("config") or {}
+        initial = report.get("initial_metrics") or {}
+        final = report.get("final_validation") or {}
+        for shape in config.get("shapes", []) or []:
+            sid = str(shape.get("id") or "")
+            if not sid:
+                continue
+            baseline = (initial.get(sid) or {}).get("baseline_us")
+            if not isinstance(baseline, (int, float)) or baseline <= 0:
+                continue
+            metrics = final.get(sid) or {}
+            median = metrics.get("median_us")
+            current = entries.get(sid)
+            best = (
+                float(median) if isinstance(median, (int, float)) and median > 0
+                else None
+            )
+            history = []
+            if best is not None:
+                history.append({
+                    "accepted": True,
+                    "median_us": best,
+                    "p90_us": metrics.get("p90_us"),
+                    "source_report": str(report_path),
+                })
+            row = {
+                "id": sid,
+                "contract": {
+                    "model": str(config.get("model") or report.get("model") or
+                                 report_path.parent.name),
+                    "tp_size": int(shape.get("tp_size", 8)),
+                    "operator": str(shape.get("operator") or ""),
+                    "M": int(shape.get("M", 0)),
+                    "N": int(shape.get("N", 0)),
+                    "K": int(shape.get("K", 0)),
+                },
+                "family": family_of(int(shape.get("M", 0)),
+                                    str(shape.get("operator") or "")),
+                "baseline_us": float(baseline),
+                "best_known_us": best,
+                "history": history,
+                "last_selected_round": None,
+            }
+            if current is None:
+                entries[sid] = row
+            else:
+                # Keep the fixed baseline from the first valid report; merge
+                # histories and retain the fastest measured best-known.
+                current["history"].extend(history)
+                if best is not None and (
+                    current.get("best_known_us") is None
+                    or best < current["best_known_us"]
+                ):
+                    current["best_known_us"] = best
+                    current["contract"] = row["contract"]
+    return {
+        "schema_version": 1,
+        "source": str(meta_root.resolve()),
+        "instances": [entries[k] for k in sorted(entries)],
+    }
+
+
+def main(argv=None) -> int:
+    p = argparse.ArgumentParser()
+    p.add_argument("--meta-root", type=Path,
+                   default=Path("/root/zth_agent/MetaInfer"))
+    p.add_argument("--out", type=Path,
+                   default=Path("/root/zth_agent/ahe-kernel-repos/registered_pool.yaml"))
+    args = p.parse_args(argv)
+    data = build_pool(args.meta_root)
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    args.out.write_text(yaml.safe_dump(data, sort_keys=False, allow_unicode=True),
+                        encoding="utf-8")
+    print(f"pool: {len(data['instances'])} instances -> {args.out}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/metainfer/tasks/harness_evolve/orchestrator/process_control.py b/metainfer/tasks/harness_evolve/orchestrator/process_control.py
new file mode 100644
index 00000000..abc4d994
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/process_control.py
@@ -0,0 +1,133 @@
+"""Child-process handles: spawn, watch, terminate.
+
+Why this exists
+---------------
+A question runs as a real DKAO child task for hours, and HE has to be able to
+stop that child when the round itself is over — most importantly when the child
+outstays its timeout (``child_timeout_minutes``) and would otherwise hold a
+device forever.
+
+* ``spawn`` starts the child in its own session, so the child and everything it
+  spawned (agent sessions, compilers, benchmark subprocesses) can be signalled
+  as one group;
+* ``terminate`` escalates ``SIGTERM`` → ``SIGKILL`` on that group.
+
+What is deliberately **not** here any more: freezing (``SIGSTOP``/``SIGCONT``).
+It existed to pause a child whose device had been taken over — exactly the
+"HE manages the environment" coupling that FLOW.md §0 removed. Whether a card is
+clean enough to measure on is DKAO's verdict, taken in the child's own
+measurement gate, and the child's contamination handling deals with a round that
+was measured on a shared device. A stopped process keeps its VRAM anyway, so
+pausing never handed the card to anybody.
+"""
+
+from __future__ import annotations
+
+import os
+import signal
+import subprocess
+import threading
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence
+
+
+class ChildProcess:
+    """One child task under process control."""
+
+    def __init__(self, proc: "subprocess.Popen[Any]") -> None:
+        self.proc = proc
+        self.pid = int(proc.pid)
+        # start_new_session=True made the child a session/group leader, so the
+        # whole tree shares this pgid.
+        try:
+            self.pgid = os.getpgid(self.pid)
+        except OSError:
+            self.pgid = self.pid
+        self.terminated: bool = False
+        self.terminate_reason: str = ""
+
+    # ------------------------------------------------------------- queries
+    def poll(self) -> Optional[int]:
+        return self.proc.poll()
+
+    def alive(self) -> bool:
+        return self.proc.poll() is None
+
+    # ------------------------------------------------------------- control
+    def _signal_group(self, sig: int) -> bool:
+        """Signal the whole process group; fall back to the child alone."""
+        try:
+            os.killpg(self.pgid, sig)
+            return True
+        except OSError:
+            pass
+        try:
+            os.kill(self.pid, sig)
+            return True
+        except OSError:
+            return False
+
+    def terminate(self, *, reason: str = "", grace_s: float = 20.0,
+                  now: Optional[float] = None) -> bool:
+        """Stop the child for good (``SIGTERM`` then ``SIGKILL``)."""
+        if self.terminated:
+            return False
+        self.terminate_reason = reason
+        self._signal_group(signal.SIGTERM)
+        deadline = time.time() + max(0.0, float(grace_s))
+        while self.alive() and time.time() < deadline:
+            time.sleep(0.2)
+        if self.alive():
+            self._signal_group(signal.SIGKILL)
+        self.terminated = True
+        return True
+
+    def state(self) -> Dict[str, Any]:
+        return {
+            "pid": self.pid,
+            "pgid": self.pgid,
+            "alive": self.alive(),
+            "terminated": self.terminated,
+            "terminate_reason": self.terminate_reason,
+        }
+
+
+def spawn(command: Sequence[str], *, cwd: str, env: Dict[str, str],
+          log_path: Path) -> ChildProcess:
+    """Start a child task in its own session (so it can be signalled as a tree)."""
+    log = log_path.open("w", encoding="utf-8")
+    proc = subprocess.Popen(
+        list(command), cwd=cwd, env=env, stdout=log,
+        stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL,
+        text=True, start_new_session=True,
+    )
+    log.close()          # the child keeps its own descriptor
+    return ChildProcess(proc)
+
+
+class ProcessRegistry:
+    """Live child handles, keyed by question id (thread-safe enough for us)."""
+
+    def __init__(self) -> None:
+        self._children: Dict[str, ChildProcess] = {}
+        self._lock = threading.Lock()
+
+    def add(self, question_id: str, child: ChildProcess) -> None:
+        with self._lock:
+            self._children[question_id] = child
+
+    def get(self, question_id: str) -> Optional[ChildProcess]:
+        with self._lock:
+            return self._children.get(question_id)
+
+    def drop(self, question_id: str) -> None:
+        with self._lock:
+            self._children.pop(question_id, None)
+
+    def items(self) -> List[Any]:
+        with self._lock:
+            return list(self._children.items())
+
+    def state(self) -> Dict[str, Any]:
+        return {qid: child.state() for qid, child in self.items()}
diff --git a/metainfer/tasks/harness_evolve/orchestrator/promotion.py b/metainfer/tasks/harness_evolve/orchestrator/promotion.py
new file mode 100644
index 00000000..5b80aa00
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/promotion.py
@@ -0,0 +1,416 @@
+"""Human-approved promotion of an evolved harness into production DKAO.
+
+The operator's rule: an HE run exists to produce one promotable harness, and
+nothing is written to DKAO without a human decision. So a candidate that passes
+both gates does **not** publish itself:
+
+1. the round writes ``pending_promotion.json`` (candidate revision, snapshot,
+   both gate records) and the run parks in ``awaiting_approval``;
+2. ``approve`` copies the candidate harness into production (backing up the
+   previous version), records it, updates the variant table and finishes the
+   run;
+3. ``deny`` throws the candidate away and the caller keeps iterating.
+
+``rollback_harness`` puts the previous production version back, which is what
+makes an approval safe to make.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from .attribution import copy_tree, replace_tree, save_json
+from .state import append_timeline
+
+PENDING_FILE = "pending_promotion.json"
+PROMOTED_FILE = "promoted_harness.json"
+PROMOTION_LOG = "promotion_log.jsonl"
+
+
+def _load_json(path: Path) -> Any:
+    try:
+        return json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+
+
+def _ts() -> str:
+    return time.strftime("%Y%m%d-%H%M%S")
+
+
+def _log(exp_dir: Path, **fields: Any) -> None:
+    row = {"ts": time.time(), **fields}
+    try:
+        Path(exp_dir).mkdir(parents=True, exist_ok=True)
+        with (Path(exp_dir) / PROMOTION_LOG).open("a", encoding="utf-8") as fh:
+            fh.write(json.dumps(row, ensure_ascii=False) + "\n")
+    except OSError:
+        pass
+
+
+def version_for(iteration: int) -> str:
+    """The harness version label used in records (h-1 is the seed)."""
+    return f"h-{int(iteration)}"
+
+
+def write_pending_promotion(exp_dir: Path, *, iteration: int, revision: str,
+                            snapshot_dir: str,
+                            performance_gate: Dict[str, Any],
+                            generalization_gate: Dict[str, Any],
+                            results: Dict[str, Any]) -> Path:
+    """Record a candidate that passed both gates and is awaiting approval."""
+    exp_dir = Path(exp_dir)
+    payload = {
+        "schema_version": 1,
+        "created_at": time.time(),
+        "iteration": iteration,
+        "version": version_for(iteration),
+        "revision": revision,
+        "snapshot_dir": snapshot_dir,
+        "performance_gate": performance_gate,
+        "generalization_gate": generalization_gate,
+        "operator_results": {
+            iid: {k: row.get(k) for k in
+                  ("median_us", "p90_us", "passed", "correctness_ok",
+                   "status", "repo_path")}
+            for iid, row in (results or {}).items() if isinstance(row, dict)
+        },
+        "awaiting": "human approval (approve / deny)",
+    }
+    save_json(exp_dir / PENDING_FILE, payload)
+    _log(exp_dir, action="pending_promotion", iteration=iteration,
+         version=payload["version"], revision=revision)
+    return exp_dir / PENDING_FILE
+
+
+def read_pending(exp_dir: Path) -> Optional[Dict[str, Any]]:
+    data = _load_json(Path(exp_dir) / PENDING_FILE)
+    return data if isinstance(data, dict) else None
+
+
+def clear_pending(exp_dir: Path) -> None:
+    try:
+        (Path(exp_dir) / PENDING_FILE).unlink()
+    except OSError:
+        pass
+
+
+def promoted_harness(exp_dir: Path) -> Optional[Dict[str, Any]]:
+    data = _load_json(Path(exp_dir) / PROMOTED_FILE)
+    return data if isinstance(data, dict) else None
+
+
+def _production_dir() -> Path:
+    from .config import harness_seed_dir
+    return harness_seed_dir("dcu_default")
+
+
+def _versioned_backup(exp_dir: Path, version: str) -> Path:
+    """Keep a named copy of what production looked like before a promotion."""
+    target = _production_dir()
+    backup_root = Path(exp_dir) / "promoted"
+    backup_root.mkdir(parents=True, exist_ok=True)
+    backup = backup_root / version
+    if target.is_dir() and not backup.exists():
+        copy_tree(target, backup)
+    return backup
+
+
+def promote_harness(exp_dir: Path, pending: Dict[str, Any], *,
+                    approved_by: str = "operator") -> Dict[str, Any]:
+    """Write the approved candidate harness into production DKAO.
+
+    The previous production tree is kept as a named copy (and a timestamped
+    one) so a promotion is always reversible.
+    """
+    exp_dir = Path(exp_dir)
+    snapshot = Path(str(pending.get("snapshot_dir") or ""))
+    if not snapshot.is_dir():
+        return {"ok": False, "errors": [f"candidate snapshot missing: {snapshot}"]}
+    target = _production_dir()
+    version = str(pending.get("version") or version_for(pending.get("iteration", 0)))
+    previous = promoted_harness(exp_dir) or {}
+    previous_version = previous.get("version") or "h-1"
+    backup = _versioned_backup(exp_dir, previous_version)
+    live_backup = target.parent / f"{target.name}.{_ts()}"
+    if target.is_dir():
+        copy_tree(target, live_backup)
+    replace_tree(snapshot, target)
+
+    record = {
+        "schema_version": 1,
+        "version": version,
+        "iteration": pending.get("iteration"),
+        "revision": pending.get("revision"),
+        "snapshot_dir": str(snapshot),
+        "production_dir": str(target),
+        "backup_dir": str(backup),
+        "previous_version": previous_version,
+        "previous_backup": str(live_backup),
+        "approved_by": approved_by,
+        "approved_at": time.time(),
+        "performance_gate": {
+            "status": (pending.get("performance_gate") or {}).get("status"),
+            "wins": (pending.get("performance_gate") or {}).get("wins"),
+        },
+        "generalization_gate": {
+            "status": (pending.get("generalization_gate") or {}).get("status"),
+            "wins": (pending.get("generalization_gate") or {}).get("wins"),
+            "paper": (pending.get("generalization_gate") or {}).get("operators"),
+        },
+    }
+    save_json(exp_dir / PROMOTED_FILE, record)
+    _log(exp_dir, action="promoted", version=version,
+         iteration=pending.get("iteration"), revision=pending.get("revision"),
+         production_dir=str(target), backup=str(backup),
+         previous_version=previous_version, approved_by=approved_by)
+    return {"ok": True, "record": record, "backup": str(backup),
+            "previous_backup": str(live_backup)}
+
+
+def rollback_harness(exp_dir: Path, *, to: str = "last-good") -> Dict[str, Any]:
+    """Restore a previously promoted harness version into production.
+
+    ``to`` accepts a version label (``h-2``), ``last-good`` (the version the
+    current one replaced) or ``seed`` (the built-in harness).
+    """
+    exp_dir = Path(exp_dir)
+    current = promoted_harness(exp_dir) or {}
+    target = _production_dir()
+    if to == "seed":
+        source = harness_seed_source()
+        label = "seed"
+    elif to in {"last-good", "last_good", ""}:
+        source = Path(str(current.get("backup_dir") or ""))
+        label = str(current.get("previous_version") or "seed")
+        if not source.is_dir():
+            source = harness_seed_source()
+            label = "seed"
+    else:
+        source = exp_dir / "promoted" / str(to)
+        label = str(to)
+    if not Path(source).is_dir():
+        return {"ok": False, "errors": [f"no harness tree to restore: {source}"]}
+    live_backup = target.parent / f"{target.name}.{_ts()}"
+    if target.is_dir():
+        copy_tree(target, live_backup)
+    replace_tree(Path(source), target)
+    record = {
+        "schema_version": 1,
+        "version": f"rollback:{label}",
+        "rolled_back_at": time.time(),
+        "production_dir": str(target),
+        "restored_from": str(source),
+        "replaced_backup": str(live_backup),
+        "replaced_version": current.get("version"),
+    }
+    save_json(exp_dir / PROMOTED_FILE, record)
+    _log(exp_dir, action="rollback", to=label, restored_from=str(source),
+         replaced_backup=str(live_backup))
+    return {"ok": True, "record": record, "restored_from": str(source)}
+
+
+def harness_seed_source() -> Path:
+    """The built-in seed harness that ships with the plugin (h-1)."""
+    here = Path(__file__).resolve().parents[2]
+    return here / "dcu_kernel_auto_opt" / "harness_default"
+
+
+def restore_promoted_harness(cfg: Any, exp_dir: Path,
+                             live_workspace: Path) -> Dict[str, Any]:
+    """Put the experiment workspace back to the promoted (production) harness.
+
+    A rejected round must not leave its own edit live, and it must not fall all
+    the way back to the seed either: the reference for the next round is what
+    production currently runs.
+    """
+    promoted = promoted_harness(exp_dir)
+    source: Optional[Path] = None
+    revision = ""
+    if promoted and promoted.get("snapshot_dir"):
+        candidate = Path(str(promoted["snapshot_dir"]))
+        if candidate.is_dir():
+            source = candidate
+            revision = str(promoted.get("revision") or "")
+    if source is None:
+        source = harness_seed_source()
+    replace_tree(source, live_workspace)
+    if revision:
+        import subprocess
+        try:
+            subprocess.run(["git", "reset", "--hard", revision],
+                           cwd=live_workspace, check=True,
+                           stdout=subprocess.DEVNULL)
+        except (OSError, subprocess.CalledProcessError):
+            pass
+    append_timeline(cfg.state_dir, "workspace_restored_to_promoted", {
+        "from": str(source),
+        "revision": revision or None,
+        "version": (promoted or {}).get("version"),
+    })
+    return {"restored_from": str(source), "revision": revision}
+
+
+def _versioned_pool_backup(exp_dir: Path, pool_path: Path) -> Path:
+    """Copy the pool next to the experiment before rewriting production data."""
+    import shutil
+    root = Path(exp_dir) / "promoted"
+    root.mkdir(parents=True, exist_ok=True)
+    backup = root / f"{pool_path.stem}.{_ts()}{pool_path.suffix}"
+    try:
+        shutil.copy2(pool_path, backup)
+    except OSError:
+        return pool_path
+    return backup
+
+
+def promote_kernels_for_round(exp_dir: Path, pool_path: Path,
+                              results: Dict[str, Any], *,
+                              version: str, iteration: int,
+                              reason: str = "harness promotion",
+                              ) -> Dict[str, Any]:
+    """Write an accepted round's operator kernels into the production pool.
+
+    Only the operators this round actually measured are touched, and only when
+    they beat what the pool already records — "best known wins", applied
+    surgically instead of rescanning every historical child report. The pool is
+    backed up before being rewritten.
+    """
+    import yaml
+    from .variant import _num
+
+    pool_path = Path(pool_path)
+    if not pool_path.is_file():
+        return {"ok": False, "updated": [],
+                "errors": [f"pool file missing: {pool_path}"]}
+    try:
+        data = yaml.safe_load(pool_path.read_text(encoding="utf-8")) or {}
+    except (OSError, ValueError) as exc:  # noqa: BLE001
+        return {"ok": False, "updated": [], "errors": [str(exc)]}
+    instances = {str(i.get("id")): i for i in (data.get("instances") or [])
+                 if isinstance(i, dict)}
+    updated: List[Dict[str, Any]] = []
+    for iid, row in (results or {}).items():
+        if not isinstance(row, dict):
+            continue
+        entry = instances.get(str(iid))
+        if entry is None:
+            continue
+        median = _num(row.get("median_us"))
+        if median is None or row.get("correctness_ok") is False:
+            continue
+        best = _num(entry.get("best_known_us"))
+        if best is not None and median >= best:
+            continue
+        entry["best_known_us"] = float(median)
+        history = entry.setdefault("history", [])
+        if isinstance(history, list):
+            history.append({
+                "accepted": True,
+                "median_us": float(median),
+                "p90_us": row.get("p90_us"),
+                "source_report": row.get("repo_path") or row.get("child_task_id"),
+                "harness_version": version,
+                "iteration": iteration,
+            })
+        updated.append({"operator_id": str(iid), "median_us": float(median),
+                        "previous_best_us": best})
+    if not updated:
+        return {"ok": True, "updated": [], "pool": str(pool_path),
+                "reason": reason}
+    backup = _versioned_pool_backup(exp_dir, pool_path)
+    pool_path.write_text(
+        yaml.safe_dump(data, sort_keys=False, allow_unicode=True),
+        encoding="utf-8")
+    _log(exp_dir, action="kernels_promoted", version=version,
+         iteration=iteration, pool=str(pool_path), updated=updated,
+         backup=str(backup))
+    return {"ok": True, "updated": updated, "pool": str(pool_path),
+            "backup": str(backup), "reason": reason}
+
+
+def resolve_question_pool(exp_dir: Path) -> Optional[Path]:
+    """The pool file this experiment draws its operators from, if any."""
+    exp_dir = Path(exp_dir)
+    for name in ("requirements.json", "config_snapshot.json"):
+        data = _load_json(exp_dir / name)
+        if not isinstance(data, dict):
+            continue
+        answers = (data.get("answers")
+                   if isinstance(data.get("answers"), dict) else data)
+        for key in ("question_pool", "pool_source"):
+            raw = str((answers or {}).get(key) or "").strip()
+            if raw and raw != "builtin":
+                candidate = Path(raw).expanduser()
+                if candidate.is_file():
+                    return candidate
+    return None
+
+
+def approve_promotion(cfg: Any, exp_dir: Path, *,
+                      approved_by: str = "operator") -> Dict[str, Any]:
+    """Approve the pending candidate: promote it, update variants, finish run."""
+    from .variant import load_variant_table, record_variants
+
+    exp_dir = Path(exp_dir)
+    pending = read_pending(exp_dir)
+    if not pending:
+        return {"ok": False, "errors": ["no pending promotion"]}
+    result = promote_harness(exp_dir, pending, approved_by=approved_by)
+    if not result.get("ok"):
+        return result
+    version = result["record"]["version"]
+    table = load_variant_table(exp_dir)
+    # The candidate won on the operators it measured; those become the new
+    # variants (and therefore the reference for the next HE run).
+    variant_update = record_variants(
+        exp_dir, table, pending.get("operator_results") or {},
+        version=version, iteration=int(pending.get("iteration") or 0),
+        groups={"generalization_ids":
+                (pending.get("generalization_gate") or {}).get("operators") or []})
+    # ... and their kernels go into production, so the accepted round is the
+    # only one whose optimizations survive.
+    pool_path = resolve_question_pool(exp_dir)
+    if pool_path is not None:
+        kernel_update = promote_kernels_for_round(
+            exp_dir, pool_path, pending.get("operator_results") or {},
+            version=version, iteration=int(pending.get("iteration") or 0))
+    else:
+        kernel_update = {"ok": True, "updated": [],
+                         "skipped": "no question pool configured"}
+    clear_pending(exp_dir)
+    append_timeline(cfg.state_dir, "harness_promoted", {
+        "version": version,
+        "iteration": pending.get("iteration"),
+        "approved_by": approved_by,
+        "variant_updates": variant_update.get("updated"),
+        "kernels_promoted": [u["operator_id"]
+                             for u in kernel_update.get("updated", [])],
+        "backup": result.get("backup"),
+    })
+    try:
+        from .state import set_run
+        set_run(cfg.state_dir, final_status=f"promoted:{version}")
+    except Exception:  # noqa: BLE001
+        pass
+    return {**result, "variant_update": variant_update,
+            "kernel_update": kernel_update, "version": version}
+
+
+def deny_promotion(cfg: Any, exp_dir: Path, *, reason: str = "operator") -> Dict[str, Any]:
+    """Throw the pending candidate away; the run keeps iterating."""
+    exp_dir = Path(exp_dir)
+    pending = read_pending(exp_dir)
+    if not pending:
+        return {"ok": False, "errors": ["no pending promotion"]}
+    clear_pending(exp_dir)
+    _log(exp_dir, action="denied", iteration=pending.get("iteration"),
+         version=pending.get("version"), reason=reason)
+    append_timeline(cfg.state_dir, "promotion_denied", {
+        "iteration": pending.get("iteration"),
+        "version": pending.get("version"), "reason": reason,
+    })
+    return {"ok": True, "denied": pending.get("version")}
diff --git a/metainfer/tasks/harness_evolve/orchestrator/publish.py b/metainfer/tasks/harness_evolve/orchestrator/publish.py
new file mode 100644
index 00000000..74168872
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/publish.py
@@ -0,0 +1,200 @@
+"""Publish gate: promote an AHE experiment's champion into production DKAO.
+
+Nothing is ever written to the production DKAO automatically. Publish only
+happens through this explicit gate, and only when the requested evidence is on
+disk:
+
+Harness publish (champion -> dcu_kernel_auto_opt/harness_default):
+  required: a promoted iteration (verdict PROMOTE / PROMOTE_UNEXPLAINED) with a
+            held-out gate that did NOT fail; target dir is backed up first and
+            a publish_log.jsonl record is appended for rollback.
+
+Kernel publish (new best-knowns -> registered pool):
+  scans AHE child DKAO runs for final-success shapes whose median beats the pool
+  best_known; the pool YAML is backed up, then best_known/history are updated.
+
+Every publish creates a timestamped backup of the destination before writing.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import yaml
+
+from .attribution import copy_tree, replace_tree, save_json
+
+PUBLISH_LOG = "publish_log.jsonl"
+_PROMOTABLE = {"PROMOTE", "PROMOTE_UNEXPLAINED"}
+
+
+def _ts() -> str:
+    return time.strftime("%Y%m%d-%H%M%S")
+
+
+def _backup(path: Path, suffix: str) -> Path:
+    """Copy destination into a timestamped backup (keep production live)."""
+    if not path.exists():
+        return path  # nothing to back up yet
+    backup = path.parent / f"{path.name}.{suffix}-{_ts()}"
+    if path.is_dir():
+        copy_tree(path, backup)
+    else:
+        backup.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(path, backup)
+    return backup
+
+
+def _decision_at(exp: Path, iteration: int) -> Optional[Dict[str, Any]]:
+    path = exp / "runs" / f"iteration_{iteration:03d}" / "input" / "decision.json"
+    if not path.is_file():
+        return None
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    return data if isinstance(data, dict) else None
+
+
+def _best_ever(exp: Path) -> Optional[Dict[str, Any]]:
+    path = exp / "best_ever.json"
+    if not path.is_file():
+        return None
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    return data if isinstance(data, dict) else None
+
+
+def check_publishable(exp: Path) -> Dict[str, Any]:
+    """Validate the experiment carries promotable, non-overfit evidence."""
+    best = _best_ever(exp)
+    if not best:
+        return {"ok": False, "errors": ["no best_ever.json (nothing promoted)"]}
+    iteration = int(best.get("iteration", 0))
+    decision = _decision_at(exp, iteration)
+    if not decision:
+        return {"ok": False, "errors": [f"no decision.json at iteration {iteration}"]}
+    verdict = decision.get("verdict")
+    errors: List[str] = []
+    if verdict not in _PROMOTABLE:
+        errors.append(f"iteration {iteration} verdict {verdict!r} is not promotable")
+    held = ((decision.get("heldout_gate") or {}).get("status"))
+    if held == "FAIL":
+        errors.append("held-out gate failed (overfit guard)")
+    snapshot = best.get("snapshot_dir")
+    if not snapshot or not Path(snapshot).is_dir():
+        errors.append("best snapshot dir missing")
+    if errors:
+        return {"ok": False, "errors": errors}
+    return {
+        "ok": True,
+        "iteration": iteration,
+        "verdict": verdict,
+        "snapshot_dir": str(snapshot),
+        "workspace_revision": best.get("workspace_revision"),
+    }
+
+
+def publish_harness(exp: Path, *, target: Path, reason: str = "explicit publish"
+                    ) -> Dict[str, Any]:
+    """Copy the champion harness snapshot over the production default seed."""
+    gate = check_publishable(exp)
+    if not gate["ok"]:
+        return {"ok": False, "errors": gate["errors"], "published": False}
+    source = Path(gate["snapshot_dir"])
+    backup = _backup(target, "pre-publish")
+    replace_tree(source, target)
+    record = {
+        "ts": time.time(),
+        "kind": "harness",
+        "iteration": gate["iteration"],
+        "verdict": gate["verdict"],
+        "source_snapshot": str(source),
+        "workspace_revision": gate.get("workspace_revision"),
+        "target": str(target),
+        "backup": str(backup) if backup != target else None,
+        "reason": reason,
+    }
+    _append_log(exp, record)
+    return {"ok": True, "published": True, "record": record}
+
+
+def _child_reports(exp: Path) -> List[Tuple[str, Dict[str, Any]]]:
+    reports: List[Tuple[str, Dict[str, Any]]] = []
+    for path in sorted((exp / "children").glob("*/*/workspace/final_report.json")):
+        try:
+            data = json.loads(path.read_text(encoding="utf-8"))
+        except (OSError, ValueError):
+            continue
+        if isinstance(data, dict):
+            reports.append((str(path), data))
+    return reports
+
+
+def publish_kernels(exp: Path, *, pool_path: Path, reason: str = "explicit publish"
+                    ) -> Dict[str, Any]:
+    """Update pool best_known/history from final-success child DKAO runs."""
+    if not pool_path.is_file():
+        return {"ok": False, "published": False,
+                "errors": [f"pool file missing: {pool_path}"]}
+    try:
+        data = yaml.safe_load(pool_path.read_text(encoding="utf-8")) or {}
+    except (OSError, ValueError) as exc:
+        return {"ok": False, "published": False, "errors": [str(exc)]}
+    instances = {i["id"]: i for i in data.get("instances", [])}
+    updated: List[str] = []
+    for src, report in _child_reports(exp):
+        if report.get("status") != "success":
+            continue
+        final = report.get("final_validation") or {}
+        config = report.get("config") or {}
+        for shape in config.get("shapes", []) or []:
+            sid = str(shape.get("id") or "")
+            entry = instances.get(sid)
+            metrics = final.get(sid) or {}
+            median = metrics.get("median_us")
+            if (
+                entry is None
+                or not isinstance(median, (int, float))
+                or median <= 0
+            ):
+                continue
+            best = entry.get("best_known_us")
+            if best is None or float(median) < float(best):
+                entry["best_known_us"] = float(median)
+                hist = entry.setdefault("history", [])
+                hist.append({
+                    "accepted": True,
+                    "median_us": float(median),
+                    "p90_us": metrics.get("p90_us"),
+                    "source_report": src,
+                })
+                updated.append(sid)
+    if not updated:
+        return {"ok": True, "published": False, "updated": [],
+                "record": {"kind": "kernels", "reason": reason}}
+    backup = _backup(pool_path, "pre-publish")
+    pool_path.write_text(yaml.safe_dump(data, sort_keys=False, allow_unicode=True),
+                         encoding="utf-8")
+    record = {
+        "ts": time.time(),
+        "kind": "kernels",
+        "updated": updated,
+        "target": str(pool_path),
+        "backup": str(backup) if backup != pool_path else None,
+        "reason": reason,
+    }
+    _append_log(exp, record)
+    return {"ok": True, "published": True, "updated": updated, "record": record}
+
+
+def _append_log(exp: Path, record: Dict[str, Any]) -> None:
+    with (exp / PUBLISH_LOG).open("a", encoding="utf-8") as fh:
+        fh.write(json.dumps(record, ensure_ascii=False) + "\n")
diff --git a/metainfer/tasks/harness_evolve/orchestrator/question_roles.py b/metainfer/tasks/harness_evolve/orchestrator/question_roles.py
new file mode 100644
index 00000000..f9eb1c79
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/question_roles.py
@@ -0,0 +1,428 @@
+"""Role-based question selection for AHE rounds.
+
+The previous rule let the Evolve agent pick any in-pool set (guardrails only
+enforced membership, budget and >=60% overlap). That produced the failure we
+actually hit on 9-8-8: three of four questions came from the same operator
+family, an over-generalised policy change looked like a win, and the decision
+promoted it — after which two unrelated shapes regressed.
+
+This module gives every question in a round an explicit **role**, so a round
+measures what it claims to measure:
+
+  regression  2 slots — questions that PASSED last round and are inside the
+              change's declared scope; spread across families. Primary causal
+              evidence that the change did not break what worked.
+  repair      1 slot — questions that LOST last round (or sit in scope.at_risk);
+              when nothing lost, the least-recently-examined question of a
+              different family.
+  probe       1 slot — deliberately a *different* family than everything else
+              (cold-start generalisation probe). It is recorded but carries
+              zero weight in the pass-rate/pairing decision.
+
+Machine constraints on top: at least ``min_families`` distinct families (when
+the pool allows), at most ``max_per_family`` questions per family, >=60%
+overlap with the previous round, and least-recently-used rotation as the
+tie-breaker so the pool is not pinned to the same four shapes forever.
+
+The agent's own ``selected`` list is honoured where it satisfies a role, and
+replaced where it does not — the correction is recorded in the rationale.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Mapping, Optional, Sequence
+
+PURPOSE_REGRESSION = "regression"
+PURPOSE_REPAIR = "repair"
+PURPOSE_PROBE = "probe"
+MAIN_PURPOSES = (PURPOSE_REGRESSION, PURPOSE_REPAIR)
+
+MIN_OVERLAP_RATIO = 0.6
+
+
+@dataclass
+class Selection:
+    selected: List[str] = field(default_factory=list)
+    purposes: Dict[str, str] = field(default_factory=dict)
+    rationale: str = ""
+    families: List[str] = field(default_factory=list)
+    notes: List[str] = field(default_factory=list)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "selected": self.selected,
+            "purposes": self.purposes,
+            "rationale": self.rationale,
+            "families": self.families,
+            "notes": self.notes,
+        }
+
+
+def _family(pool: Any, iid: str) -> str:
+    try:
+        inst = pool.get(iid)
+    except Exception:  # noqa: BLE001
+        return ""
+    return str(getattr(inst, "family", "") or "")
+
+
+def _regime(pool: Any, iid: str) -> str:
+    """decode (M<=32) vs prefill (M>32) — a stronger axis than operator."""
+    try:
+        inst = pool.get(iid)
+    except Exception:  # noqa: BLE001
+        return ""
+    try:
+        m = int(getattr(inst, "M", 0) or 0)
+    except (TypeError, ValueError):
+        return ""
+    if m <= 0:
+        return ""
+    return "decode" if m <= 32 else "prefill"
+
+
+def _rotation_key(rotation: Mapping[str, int], iid: str) -> int:
+    try:
+        return int(rotation.get(iid, 0))
+    except (TypeError, ValueError):
+        return 0
+
+
+def _scope_ids(plan: Optional[Any]) -> List[str]:
+    """Ids the change declares it will improve / must keep unchanged."""
+    ids: List[str] = []
+    changes = getattr(plan, "changes", None)
+    if changes is None and isinstance(plan, Mapping):
+        changes = plan.get("changes")
+    for change in changes or []:
+        scope = change.get("scope") if isinstance(change, Mapping) else None
+        if not isinstance(scope, Mapping):
+            continue
+        for bucket in ("expected_improve", "unchanged", "at_risk"):
+            bucket_ids = scope.get(bucket)
+            if isinstance(bucket_ids, str):
+                bucket_ids = [bucket_ids]
+            for iid in bucket_ids or []:
+                if isinstance(iid, str) and iid not in ids:
+                    ids.append(iid)
+    return ids
+
+
+def _pick_spread(candidates: Sequence[str], pool: Any, used_families: set,
+                 count: int, used_regimes: Optional[set] = None) -> List[str]:
+    """Pick ``count`` ids preferring uncovered regimes, then families.
+
+    Regime first: decode (M<=32) and prefill (M>32) need completely different
+    kernels, so a round that covers four prefill families but no decode
+    question would miss the harder, more valuable half of the space.
+    """
+    used_regimes = used_regimes if used_regimes is not None else set()
+
+    def rank(iid: str) -> tuple:
+        regime = _regime(pool, iid)
+        family = _family(pool, iid)
+        return (
+            0 if (regime and regime not in used_regimes) else 1,
+            0 if (family and family not in used_families) else 1,
+        )
+
+    picked: List[str] = []
+    remaining = list(candidates)
+    # Greedy: re-rank after every pick so the second anchor prefers the regime
+    # the first one did not cover (a pre-sorted pass would keep picking the
+    # same regime).
+    while len(remaining) > 0 and len(picked) < count:
+        remaining.sort(key=lambda i: (rank(i), candidates.index(i)))
+        best = remaining[0]
+        regime = _regime(pool, best)
+        family = _family(pool, best)
+        redundant = bool(family and family in used_families
+                         and regime in used_regimes)
+        if redundant and len(remaining) > (count - len(picked)):
+            remaining.pop(0)
+            continue
+        remaining.pop(0)
+        picked.append(best)
+        if family:
+            used_families.add(family)
+        if regime:
+            used_regimes.add(regime)
+    return picked
+
+
+def select_questions(*, pool: Any, budget: int,
+                     prev_selected: Sequence[str] = (),
+                     prev_results: Optional[Mapping[str, Mapping[str, Any]]] = None,
+                     plan: Optional[Any] = None,
+                     rotation: Optional[Mapping[str, int]] = None,
+                     min_families: int = 3,
+                     max_per_family: int = 2,
+                     anchor_min: int = 2) -> Selection:
+    """Return this round's question set with per-question roles."""
+    budget = int(budget)
+    rotation = rotation or {}
+    prev_results = prev_results or {}
+    notes: List[str] = []
+    purposes: Dict[str, str] = {}
+    stages: Dict[str, str] = {}
+
+    all_ids = list(pool.all_ids())
+    families_available = {_family(pool, i) for i in all_ids}
+    families_available.discard("")
+    target_families = min(int(min_families), max(1, len(families_available)))
+
+    if budget <= 0 or not all_ids:
+        return Selection(selected=[], purposes={}, rationale="empty pool")
+
+    prev = [i for i in prev_selected if i in pool.instances]
+    passed = [i for i in prev if (prev_results.get(i) or {}).get("passed") is True]
+    failed = [i for i in prev
+              if prev_results.get(i) is not None
+              and (prev_results.get(i) or {}).get("passed") is not True]
+    scope_all = [i for i in _scope_ids(plan) if i in pool.instances]
+    scope_improve = [i for i in _scope_ids(plan) if i in pool.instances]
+    scope_risk: List[str] = []
+    changes = getattr(plan, "changes", None)
+    if changes is None and isinstance(plan, Mapping):
+        changes = plan.get("changes")
+    for change in changes or []:
+        scope = change.get("scope") if isinstance(change, Mapping) else None
+        if not isinstance(scope, Mapping):
+            continue
+        for iid in scope.get("at_risk") or []:
+            if isinstance(iid, str) and iid in pool.instances and iid not in scope_risk:
+                scope_risk.append(iid)
+        for iid in scope.get("expected_improve") or []:
+            if isinstance(iid, str) and iid in pool.instances and iid not in scope_improve:
+                scope_improve.append(iid)
+
+    agent_ids: List[str] = []
+    plan_selected = getattr(plan, "selected", None)
+    if plan_selected is None and isinstance(plan, Mapping):
+        plan_selected = plan.get("selected")
+    for iid in plan_selected or []:
+        if isinstance(iid, str) and iid in pool.instances and iid not in agent_ids:
+            agent_ids.append(iid)
+
+    chosen: List[str] = []
+    used_families: set = set()
+    used_regimes: set = set()
+
+    def _take(iid: str, purpose: str, stage: str) -> None:
+        if iid in chosen or len(chosen) >= budget:
+            return
+        chosen.append(iid)
+        purposes[iid] = purpose
+        stages[iid] = stage
+        used_families.add(_family(pool, iid))
+        if _regime(pool, iid):
+            used_regimes.add(_regime(pool, iid))
+
+    def _family_counts() -> Dict[str, int]:
+        counts: Dict[str, int] = {}
+        for iid in chosen:
+            fam = _family(pool, iid)
+            counts[fam] = counts.get(fam, 0) + 1
+        return counts
+
+    def _lru(candidates: Sequence[str]) -> List[str]:
+        return sorted(candidates, key=lambda i: (_rotation_key(rotation, i), i))
+
+    def _drop_one(protect: Sequence[str] = (PURPOSE_REGRESSION,)) -> Optional[str]:
+        """Remove the least important question to make room."""
+        order = ["fill", "agent", "repair", "probe", "regression"]
+        for stage in order:
+            victims = [i for i in chosen if stages.get(i) == stage]
+            for victim in reversed(victims):
+                if purposes.get(victim) in protect:
+                    # keep at least two regression anchors
+                    if purposes.get(victim) == PURPOSE_REGRESSION and \
+                            list(purposes.values()).count(PURPOSE_REGRESSION) > 2:
+                        pass
+                    else:
+                        continue
+                chosen.remove(victim)
+                purposes.pop(victim, None)
+                stages.pop(victim, None)
+                return victim
+        return None
+
+    # 1) regression anchors (2): passed last round, scope-aware, family-spread
+    anchor_pool = [i for i in scope_improve + scope_all if i in passed] or list(passed)
+    if not anchor_pool:
+        anchor_pool = list(pool.historical_ids())
+        if prev:
+            notes.append("no passed question last round; anchors from history")
+    anchor_want = 2 if budget >= 3 else 1
+    for iid in _pick_spread(anchor_pool, pool, used_families, anchor_want,
+                            used_regimes):
+        _take(iid, PURPOSE_REGRESSION, "anchor")
+
+    # 2) agent picks that already fit a role
+    for iid in agent_ids:
+        if len(chosen) >= budget:
+            break
+        if iid in passed and iid not in chosen:
+            _take(iid, PURPOSE_REGRESSION, "agent")
+        elif iid in failed and iid not in chosen:
+            _take(iid, PURPOSE_REPAIR, "agent")
+
+    # 3) repair slot: at-risk scope first, then whatever lost, then LRU
+    if budget >= 2 and PURPOSE_REPAIR not in purposes.values():
+        repair_pool = [i for i in scope_risk + list(failed)
+                       if i not in chosen]
+        if not repair_pool:
+            repair_pool = _lru([i for i in all_ids if i not in chosen])
+            if failed:
+                notes.append("repair slot filled by least-recently-examined question")
+        for iid in _pick_spread(repair_pool, pool, used_families, 1,
+                                used_regimes):
+            _take(iid, PURPOSE_REPAIR, "repair")
+
+    # 4) fill remaining slots: scope-improve first, then least-recently-used
+    if len(chosen) < budget:
+        fill_pool = [i for i in scope_improve if i not in chosen]
+        fill_pool += _lru([i for i in all_ids if i not in chosen
+                           and i not in fill_pool])
+        for iid in _pick_spread(fill_pool, pool, used_families,
+                                budget - len(chosen), used_regimes):
+            _take(iid, PURPOSE_REGRESSION if iid in passed else PURPOSE_REPAIR,
+                  "fill")
+
+    # 5) generalisation probe: prefer a family nothing else covers
+    if budget >= 2:
+        uncovered = _lru([i for i in all_ids if i not in chosen
+                          and _family(pool, i)
+                          and _family(pool, i) not in used_families])
+        if not uncovered:
+            uncovered = _lru([i for i in all_ids if i not in chosen
+                              and _regime(pool, i)
+                              and _regime(pool, i) not in used_regimes])
+        if uncovered:
+            probe_id = uncovered[0]
+            if len(chosen) >= budget:
+                victim = _drop_one()
+                if victim is not None:
+                    notes.append(f"replaced {victim} with cross-family probe "
+                                 f"{probe_id}")
+            _take(probe_id, PURPOSE_PROBE, "probe")
+        else:
+            # No untouched family left: promote the freshest question whose
+            # family appears once in this round (it is the generalisation
+            # signal), instead of swapping the set around.
+            counts = _family_counts()
+            singles = [i for i in chosen
+                       if counts.get(_family(pool, i), 0) == 1
+                       and i not in probe_ids(purposes)]
+            singles.sort(key=lambda i: (
+                purposes.get(i) == PURPOSE_REGRESSION,
+                _rotation_key(rotation, i), i))
+            if singles:
+                probe_id = singles[0]
+                purposes[probe_id] = PURPOSE_PROBE
+                stages[probe_id] = "probe"
+                notes.append(f"probe is the unique-family question "
+                             f"{probe_id} already in the round")
+            else:
+                fallback = _lru([i for i in all_ids if i not in chosen])
+                if fallback:
+                    probe_id = fallback[0]
+                    victim = _drop_one() if len(chosen) >= budget else None
+                    if victim is not None:
+                        notes.append(f"replaced {victim} with probe {probe_id}")
+                    _take(probe_id, PURPOSE_PROBE, "probe")
+                    notes.append("pool exhausted per family; probe is the "
+                                 "freshest remaining question")
+                else:
+                    notes.append("no cross-family probe available in the pool")
+
+    # ---------------------------------------------------- machine guardrails
+    if max_per_family > 0:
+        counts = _family_counts()
+        for iid in list(chosen):
+            fam = _family(pool, iid)
+            if counts.get(fam, 0) <= max_per_family:
+                continue
+            replacement = next(
+                (c for c in _lru([i for i in all_ids if i not in chosen])
+                 if _family(pool, c)
+                 and counts.get(_family(pool, c), 0) < max_per_family
+                 and _family(pool, c) not in used_families),
+                None,
+            )
+            if replacement is None:
+                continue
+            counts[fam] -= 1
+            counts[_family(pool, replacement)] = counts.get(
+                _family(pool, replacement), 0) + 1
+            purposes[replacement] = purposes.get(iid, PURPOSE_REPAIR)
+            stages[replacement] = stages.get(iid, "fill")
+            purposes.pop(iid, None)
+            stages.pop(iid, None)
+            chosen[chosen.index(iid)] = replacement
+            used_families = {_family(pool, c) for c in chosen}
+            notes.append(f"family quota: {iid} -> {replacement}")
+
+    # overlap: keep round-to-round comparability
+    if prev:
+        need = int(MIN_OVERLAP_RATIO * len(prev))
+        overlap = len(set(chosen) & set(prev))
+        for iid in prev:
+            if overlap >= need or not chosen:
+                break
+            if iid in chosen:
+                continue
+            victim = _drop_one(protect=(PURPOSE_REGRESSION, PURPOSE_REPAIR))
+            if victim is None:
+                break
+            notes.append(f"overlap backfill: {victim} -> {iid}")
+            _take(iid, PURPOSE_REGRESSION, "overlap")
+            overlap += 1
+
+    # anchors: at least anchor_min ids with accepted history
+    anchor_set = set(pool.historical_ids())
+    have = sum(1 for i in chosen if i in anchor_set)
+    if have < anchor_min:
+        for iid in _lru([i for i in anchor_set if i not in chosen]):
+            if have >= anchor_min:
+                break
+            victim = _drop_one(protect=(PURPOSE_REGRESSION,))
+            if victim is None:
+                break
+            _take(iid, PURPOSE_REGRESSION, "anchor")
+            have += 1
+            notes.append(f"anchor backfill: added {iid}")
+
+    families = sorted({_family(pool, i) for i in chosen if _family(pool, i)})
+    regimes = sorted({_regime(pool, i) for i in chosen if _regime(pool, i)})
+    if len(regimes) < 2 and len({_regime(pool, i) for i in all_ids
+                                 if _regime(pool, i)}) >= 2:
+        notes.append(f"regime coverage {regimes} covers a single regime")
+    if len(families) < target_families:
+        notes.append(f"family coverage {len(families)} < target {target_families}")
+
+    rationale = (
+        "roles: " + ", ".join(f"{i}={purposes.get(i)}" for i in chosen)
+        + f" | families: {', '.join(families)}"
+        + f" | regimes: {', '.join(regimes)}"
+        + (f" | notes: {'; '.join(notes)}" if notes else "")
+    )
+    return Selection(selected=list(chosen),
+                     purposes={i: purposes[i] for i in chosen if i in purposes},
+                     rationale=rationale, families=families, notes=notes)
+
+
+def overlap_ratio(selected: Sequence[str], prev: Sequence[str]) -> float:
+    if not prev:
+        return 1.0
+    return len(set(selected) & set(prev)) / len(set(prev))
+
+
+def main_purposes(purposes: Mapping[str, str]) -> List[str]:
+    """Ids that take part in the pass-rate / win-loss decision."""
+    return [i for i, p in (purposes or {}).items() if p in MAIN_PURPOSES]
+
+
+def probe_ids(purposes: Mapping[str, str]) -> List[str]:
+    return [i for i, p in (purposes or {}).items() if p == PURPOSE_PROBE]
diff --git a/metainfer/tasks/harness_evolve/orchestrator/repo_cleanup.py b/metainfer/tasks/harness_evolve/orchestrator/repo_cleanup.py
new file mode 100644
index 00000000..083319c2
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/repo_cleanup.py
@@ -0,0 +1,300 @@
+"""任务级 kernel 仓库清理（harness_evolve / AHE）。
+
+背景
+----
+AHE 每轮为每道题物化一个独立的候选 kernel 仓库，命名规则见
+``harness_evolve/orchestrator/adapters/eval.py``::
+
+    child_id = _slug(f"ahe-{cfg.task_id}-it{iteration:03d}-candidate-{inst.id}-{suffix}")
+
+因此 ``ahe--it-...`` 的前缀可以**精确**定位"某任务拥有的仓库"。
+一个跑满 N 轮的 AHE 任务会留下 约 4N 个仓库（4 = 一轮题数/卡数），
+而 ``sys_shell`` 里原有的 ``_purge_owned_kernel_repo`` 只删**单个**仓库，
+且 task_type 门禁写死为 ``dcu-kernel-auto-opt``，对 HE 完全不生效 —— 于是
+任务删除后这些仓库永久残留。
+
+安全设计（三层）
+----------------
+1. **前缀精确匹配**：只处理 ``ahe-{task_id}-it`` 开头的目录，天然不会碰到
+   人工维护的 ``/root/zth_agent/kernel-repos`` 或别的任务的仓库。
+2. **池白名单**：``registered_pool.yaml`` 中 ``source_report`` 引用的仓库
+   一律保留 —— 即便前缀匹配上了也绝不删。
+3. **路径校验**：跳过符号链接、跳过非本目录直接子项、解析后必须仍在 repo_root 内。
+
+默认 ``dry_run=True``：只报告，不删除。
+"""
+
+from __future__ import annotations
+
+import re
+import shutil
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+DEFAULT_AHE_REPO_ROOT = Path("/root/zth_agent/ahe-kernel-repos")
+_REGISTERED_POOL_NAME = "registered_pool.yaml"
+
+# 匹配仓库名本身：形如 ``ahe-{task_id}-it{NNN}-candidate-...``
+# （见 eval.py 的 ``_slug(f"ahe-{task_id}-it{iteration:03d}-candidate-...")``）。
+#
+# 只锚定名称模式、不要求前后是 '/' 或引号：池文件里同一个仓库既可能写成
+# ``.../ahe-xxx-it001-.../final_report.json``（后面跟 '/'），也可能写成
+# ``.../ahe-xxx-it001-...``（行尾）或 JSON 里的 ``"repo_path": ".../ahe-xxx-it001-..."``，
+# 加边界锚点会漏掉其中一种。``-it<数字>`` 这个标记本身已足够特异：
+# 它既不会命中仓库根目录 ``ahe-kernel-repos``，也不会命中 ``kernel-repos`` 下
+# 人工维护的仓库（那些名字里没有 ahe- 前缀）。
+_REPO_REF_RE = re.compile(r"(ahe-[A-Za-z0-9._-]*?-it\d+[A-Za-z0-9._-]*)")
+
+
+def resolve_ahe_repo_root(answers: Optional[Dict[str, Any]] = None) -> Path:
+    """与 eval.py 的解析保持一致：answers['ahe_repo_root'] 优先，否则默认值。"""
+    raw = None
+    if isinstance(answers, dict):
+        raw = answers.get("ahe_repo_root")
+    base = Path(str(raw or DEFAULT_AHE_REPO_ROOT)).expanduser()
+    try:
+        return base.resolve()
+    except OSError:
+        return base
+
+
+def repo_name_prefix(task_id: str) -> str:
+    """该任务拥有的仓库名前缀。结尾的 'it' 是关键：避免 task_id 互为前缀时误伤。"""
+    return f"ahe-{task_id}-it"
+
+
+def load_pool_keep_list(repo_root: Path) -> set[str]:
+    """从 registered_pool.yaml 提取所有被引用的仓库目录名（白名单）。"""
+    keep: set[str] = set()
+    pool = repo_root / _REGISTERED_POOL_NAME
+    if not pool.is_file():
+        return keep
+    try:
+        text = pool.read_text(encoding="utf-8", errors="replace")
+    except OSError:
+        return keep
+    for name in _REPO_REF_RE.findall(text):
+        name = name.strip()
+        if name and name != _REGISTERED_POOL_NAME:
+            keep.add(name)
+    return keep
+
+
+def is_task_registered(task_id: str) -> bool:
+    """该 task_id 是否仍在 WebUI 注册表中（即任务尚未被删除）。
+
+    注册表位置与 :mod:`metainfer.server.tasks` 一致：
+    ``/.metainfer/registry.json``。任一节点目录命中即视为仍注册。
+    """
+    import json as _json
+    root = _meta_root()
+    if root is None:
+        return False
+    try:
+        candidates = sorted(root.glob("nodes/*/.metainfer/registry.json"))
+    except OSError:
+        return False
+    for reg in candidates:
+        try:
+            data = _json.loads(reg.read_text(encoding="utf-8"))
+        except (ValueError, OSError):
+            continue
+        for t in data.get("tasks", []) or []:
+            if t.get("id") == task_id:
+                return True
+    return False
+
+
+def _meta_root() -> Optional[Path]:
+    """定位 MetaInfer 仓库根（含 nodes/ 的那一层）。"""
+    # 本文件位于 /metainfer/tasks/harness_evolve/orchestrator/repo_cleanup.py
+    here = Path(__file__).resolve()
+    for parent in here.parents:
+        if (parent / "metainfer").is_dir() and (parent / "nodes").is_dir():
+            return parent
+    return None
+
+
+def owned_repos(
+    task_id: str,
+    repo_root: Path,
+    keep: Optional[set[str]] = None,
+) -> tuple[List[Path], List[str]]:
+    """返回 (可删除仓库列表, 被保留的原因说明)。
+
+    只做识别与安全检查，不产生任何副作用 —— 便于 dry-run 与测试。
+    """
+    keep = keep if keep is not None else load_pool_keep_list(repo_root)
+    prefix = repo_name_prefix(task_id)
+    deletable: List[Path] = []
+    skipped: List[str] = []
+    if not repo_root.is_dir():
+        return deletable, [f"repo_root 不存在: {repo_root}"]
+    try:
+        entries = sorted(repo_root.iterdir())
+    except OSError as exc:
+        return deletable, [f"无法读取 repo_root: {exc}"]
+
+    for p in entries:
+        if not p.name.startswith(prefix):
+            continue
+        # 保护 1：池引用
+        if p.name in keep:
+            skipped.append(f"{p.name} (被 registered_pool.yaml 引用，保留)")
+            continue
+        # 保护 2：只删真实目录，不碰符号链接
+        if p.is_symlink():
+            skipped.append(f"{p.name} (符号链接，跳过)")
+            continue
+        if not p.is_dir():
+            skipped.append(f"{p.name} (非目录，跳过)")
+            continue
+        # 保护 3：解析后必须仍在 repo_root 内
+        try:
+            resolved = p.resolve()
+        except OSError:
+            skipped.append(f"{p.name} (无法解析路径，跳过)")
+            continue
+        if resolved.parent != repo_root:
+            skipped.append(f"{p.name} (解析后越出 repo_root，跳过)")
+            continue
+        deletable.append(p)
+    return deletable, skipped
+
+
+def purge_task_kernel_repos(
+    task_id: str,
+    repo_root: Path,
+    keep: Optional[set[str]] = None,
+    dry_run: bool = True,
+    force: bool = False,
+) -> Dict[str, Any]:
+    """删除某 HE 任务拥有的候选 kernel 仓库。
+
+    :param dry_run: True 时只统计，不做任何删除（默认）。
+    :param force: 任务**仍在注册表**中时是否照删。默认 False —— 即"任务还没被
+        删除就拒绝清理它的仓库"，避免误伤在跑/待恢复的实验。
+    :returns: 摘要字典，含 matched/deleted/skipped/bytes 等。
+    """
+    deletable, skipped = owned_repos(task_id, repo_root, keep)
+
+    # 保护 0（最重要）：任务仍在注册表 → 其仓库可能被在跑或待恢复的实验使用
+    if not force and is_task_registered(task_id):
+        return {
+            "task_id": task_id,
+            "repo_root": str(repo_root),
+            "prefix": repo_name_prefix(task_id),
+            "dry_run": dry_run,
+            "matched": len(deletable),
+            "deleted": [],
+            "failed": [],
+            "skipped": skipped + [
+                f"任务 {task_id} 仍在注册表中（未被删除）；"
+                "如确认要删，请加 --force / force=True"
+            ],
+            "bytes": 0,
+            "bytes_human": "0B",
+            "refused": "task_still_registered",
+        }
+
+    # 统计体积（删除前算，删后无从得知）
+    bytes_total = 0
+    sizes: Dict[str, int] = {}
+    for p in deletable:
+        sz = 0
+        for sub in p.rglob("*"):
+            try:
+                if sub.is_file() and not sub.is_symlink():
+                    sz += sub.stat().st_size
+            except OSError:
+                continue
+        sizes[p.name] = sz
+        bytes_total += sz
+
+    deleted: List[str] = []
+    failed: List[str] = []
+    if not dry_run:
+        for p in deletable:
+            try:
+                shutil.rmtree(p)
+                deleted.append(p.name)
+            except OSError as exc:
+                failed.append(f"{p.name}: {exc}")
+
+    return {
+        "task_id": task_id,
+        "repo_root": str(repo_root),
+        "prefix": repo_name_prefix(task_id),
+        "dry_run": dry_run,
+        "matched": len(deletable),
+        "deleted": deleted,
+        "failed": failed,
+        "skipped": skipped,
+        "bytes": bytes_total,
+        "bytes_human": _human(bytes_total),
+    }
+
+
+def _human(n: int) -> str:
+    step = float(n)
+    for unit in ("B", "K", "M", "G", "T"):
+        if step < 1024 or unit == "T":
+            return f"{step:.1f}{unit}"
+        step /= 1024
+    return f"{step:.1f}T"
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    import argparse
+    import json
+
+    ap = argparse.ArgumentParser(
+        description="清理某个 harness_evolve 任务拥有的候选 kernel 仓库（默认 dry-run）",
+    )
+    ap.add_argument("task_id", help="HE 任务的 task_id，例如 test-9-8-5-e10567a1")
+    ap.add_argument("--repo-root", default=None,
+                    help=f"AHE 仓库根，默认 {DEFAULT_AHE_REPO_ROOT}")
+    ap.add_argument("--apply", action="store_true",
+                    help="真正执行删除（不加此参数只做 dry-run）")
+    ap.add_argument("--force", action="store_true",
+                    help="任务仍在注册表中时也照删（危险，默认拒绝）")
+    ap.add_argument("--json", action="store_true", help="以 JSON 输出")
+    args = ap.parse_args(argv)
+
+    repo_root = Path(args.repo_root).expanduser() if args.repo_root \
+        else DEFAULT_AHE_REPO_ROOT
+    summary = purge_task_kernel_repos(
+        args.task_id, repo_root, dry_run=not args.apply, force=args.force,
+    )
+
+    if args.json:
+        print(json.dumps(summary, ensure_ascii=False, indent=2))
+        return 0 if not summary["failed"] else 1
+
+    if summary.get("refused"):
+        print(f"[拒绝清理] task_id={summary['task_id']}")
+        print(f"  {summary['skipped'][-1]}")
+        return 2
+
+    mode = "DRY-RUN（未删除任何东西）" if summary["dry_run"] else "已执行删除"
+    print(f"[{mode}] task_id={summary['task_id']}")
+    print(f"  repo_root : {summary['repo_root']}")
+    print(f"  匹配前缀  : {summary['prefix']}")
+    print(f"  命中仓库  : {summary['matched']} 个，合计 {summary['bytes_human']}")
+    if summary["skipped"]:
+        print(f"  受保护跳过: {len(summary['skipped'])} 个")
+        for s in summary["skipped"][:10]:
+            print(f"    - {s}")
+    if summary["deleted"]:
+        print(f"  已删除    : {len(summary['deleted'])} 个")
+    if summary["failed"]:
+        print(f"  删除失败  : {len(summary['failed'])} 个")
+        for f in summary["failed"][:10]:
+            print(f"    - {f}")
+    if summary["dry_run"] and summary["matched"]:
+        print("\n  确认无误后，加 --apply 执行删除。")
+    return 0 if not summary["failed"] else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/metainfer/tasks/harness_evolve/orchestrator/round_plan.py b/metainfer/tasks/harness_evolve/orchestrator/round_plan.py
new file mode 100644
index 00000000..de02566c
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/round_plan.py
@@ -0,0 +1,209 @@
+"""round_plan: per-round question set chosen by the single Evolve role.
+
+The evolve step produces BOTH the harness change and the next round's
+``round_plan.json`` (strategy / rationale / selected / weights). Machine
+guardrails in this module keep agent-chosen sets comparable and in-pool:
+  - all ids must be in the pool;
+  - overlap with the previous round's scored set >= 60% (auto back-fill);
+  - budget cap enforced.
+Decisions (flip / attribution / best-ever / rollback) are made on the overlap
+pairs only (see docs/dkao_suite_protocol_v2.md).
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from .pool import Pool
+
+MIN_OVERLAP_RATIO = 0.6
+VALID_STRATEGIES = ("explore", "exploit", "balanced")
+
+
+@dataclass
+class RoundPlan:
+    iteration: int
+    strategy: str = "balanced"
+    rationale: str = ""
+    selected: List[str] = field(default_factory=list)
+    weights: Dict[str, float] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "iteration": self.iteration,
+            "strategy": self.strategy,
+            "rationale": self.rationale,
+            "selected": self.selected,
+            "weights": self.weights,
+        }
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "RoundPlan":
+        return cls(
+            iteration=int(d.get("iteration", 0)),
+            strategy=str(d.get("strategy") or "balanced"),
+            rationale=str(d.get("rationale") or ""),
+            selected=[str(x) for x in (d.get("selected") or [])],
+            weights={
+                str(k): float(v) for k, v in (d.get("weights") or {}).items()
+            },
+        )
+
+
+def load_round_plan(path: Path) -> Optional[RoundPlan]:
+    if not path.is_file():
+        return None
+    try:
+        return RoundPlan.from_dict(
+            json.loads(path.read_text(encoding="utf-8"))
+        )
+    except (OSError, ValueError):
+        return None
+
+
+def validate_in_pool(plan: RoundPlan, pool: Pool) -> None:
+    if plan.strategy not in VALID_STRATEGIES:
+        raise ValueError(
+            f"strategy {plan.strategy!r} not in {VALID_STRATEGIES}"
+        )
+    pool.assert_in_pool(plan.selected)
+    if len(set(plan.selected)) != len(plan.selected):
+        raise ValueError("round_plan.selected contains duplicates")
+
+
+def _top_up_anchors(
+    picked: List[str],
+    seen: set,
+    pool: Pool,
+    budget: int,
+    anchor_min: int,
+) -> List[str]:
+    """Guarantee at least ``anchor_min`` instances WITH accepted history.
+
+    Cold (baseline-only) shapes are allowed as questions, but every round must
+    carry >= anchor_min historical anchors so family tau stays anchored and
+    results stay comparable. Replaces trailing non-anchor picks if needed.
+    """
+    anchor_set = set(pool.historical_ids())
+    have = sum(1 for i in picked if i in anchor_set)
+    if have >= anchor_min:
+        return picked
+    for aid in pool.historical_ids():
+        if have >= anchor_min:
+            break
+        if aid in seen:
+            continue
+        # make room by dropping the last non-anchor pick (never drop an anchor)
+        drop_idx = next(
+            (i for i in range(len(picked) - 1, -1, -1)
+             if picked[i] not in anchor_set),
+            None,
+        )
+        if drop_idx is not None:
+            old = picked.pop(drop_idx)
+            seen.discard(old)
+        picked.append(aid)
+        seen.add(aid)
+        have += 1
+        if len(picked) > budget:  # safety; budget >= anchor_min is required
+            picked = picked[:budget]
+    if have < anchor_min:
+        raise ValueError(
+            f"pool has fewer than {anchor_min} historical anchors; "
+            "cannot satisfy the anchor rule"
+        )
+    return picked
+
+
+def enforce_guardrails(
+    prev_selected: List[str],
+    plan: Optional[RoundPlan],
+    pool: Pool,
+    budget: int,
+    anchor_min: int = 2,
+) -> List[str]:
+    """Return the final scored set for the next round.
+
+    Guardrails:
+    - budget must be >= anchor_min;
+    - plan None / empty -> stratified default sample;
+    - otherwise start from plan.selected (in pool);
+    - overlap with previous scored set >= MIN_OVERLAP_RATIO (auto back-fill);
+    - >= anchor_min instances must carry accepted history (anchor rule);
+    - every id stays in the pool; total capped at budget.
+    """
+    budget = int(budget)
+    if budget < anchor_min:
+        raise ValueError(f"budget {budget} < anchor_min {anchor_min}")
+    if plan is None or not plan.selected:
+        picked = pool.stratified_sample(budget)
+        return _top_up_anchors(picked, set(picked), pool, budget, anchor_min)[:budget]
+
+    validate_in_pool(plan, pool)
+    picked: List[str] = []
+    seen: set = set()
+    for iid in plan.selected:
+        if iid not in seen:
+            picked.append(iid)
+            seen.add(iid)
+        if len(picked) >= budget:
+            break
+
+    prev = [i for i in prev_selected if i in pool.instances]
+    if prev:
+        overlap = len(set(picked) & set(prev))
+        need = int(MIN_OVERLAP_RATIO * len(prev))
+        for iid in prev:
+            if overlap >= need or len(picked) >= budget:
+                break
+            if iid not in seen:
+                picked.append(iid)
+                seen.add(iid)
+                overlap += 1
+
+    if len(picked) < budget:
+        for iid in pool.stratified_sample(budget):
+            if len(picked) >= budget:
+                break
+            if iid not in seen:
+                picked.append(iid)
+                seen.add(iid)
+    return _top_up_anchors(picked, seen, pool, budget, anchor_min)[:budget]
+
+
+ROTATION_FILE = "pool_rotation.json"
+
+
+def load_rotation(exp_dir: Path) -> Dict[str, int]:
+    """Per-question last-selected round (sidecar; the pool yaml stays read-only)."""
+    path = Path(exp_dir) / ROTATION_FILE
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return {}
+    if not isinstance(data, dict):
+        return {}
+    out: Dict[str, int] = {}
+    for key, value in data.items():
+        try:
+            out[str(key)] = int(value)
+        except (TypeError, ValueError):
+            continue
+    return out
+
+
+def mark_selected(exp_dir: Path, ids: List[str], iteration: int) -> Dict[str, int]:
+    """Record that ``ids`` were examined in ``iteration`` (rotation fairness)."""
+    rotation = load_rotation(exp_dir)
+    for iid in ids:
+        rotation[str(iid)] = int(iteration)
+    try:
+        (Path(exp_dir) / ROTATION_FILE).write_text(
+            json.dumps(rotation, ensure_ascii=False, indent=2, sort_keys=True),
+            encoding="utf-8")
+    except OSError:
+        pass
+    return rotation
diff --git a/metainfer/tasks/harness_evolve/orchestrator/rounds.py b/metainfer/tasks/harness_evolve/orchestrator/rounds.py
new file mode 100644
index 00000000..5df5d13e
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/rounds.py
@@ -0,0 +1,481 @@
+"""Round question selection and gate evaluation for harness_evolve.
+
+One HE run chases one deliverable: a harness version that can be promoted into
+DKAO. The round structure is fixed by the operator:
+
+* **Round 1 is a baseline round.** It draws 4 operators at random from the
+  question pool, has DKAO measure them, and **judges nothing**: no gate is
+  evaluated and nothing is promoted. The kernels that beat their variant go
+  straight into the variant pool (that is a kernel fact, not a harness
+  verdict). These 4 operators become the fixed **generalization paper**
+  (``generalization_ids``) for the rest of the run.
+* **Every later round** draws a fresh random set of 4 operators (never reusing
+  the paper) and judges **only the performance gate**.
+* When the performance gate passes, the *same* harness is handed to DKAO again
+  with the paper's operators — that retake is the **generalization gate**, and
+  it counts as half a round (``GENERALIZATION_ROUND_COST``).
+* Both gates are counter rules against the variant: the performance gate needs
+  >= 75% wins, the generalization gate >= 50% wins, and in both cases a
+  non-winning operator may not fall below 80% of its variant.
+
+Because a rejected candidate never reaches production, each round measures
+against the same reference: "is this harness better than what the variant pool
+currently holds".
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import random
+import time
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Sequence
+
+from .decision_engine import (
+    GENERALIZATION_GATE_WIN_RATIO, PERFORMANCE_GATE_WIN_RATIO,
+    gate_verdict,
+)
+from .variant import pool_instances, reference_results, resolve_variant
+
+#: Operators per gate. The machine exposes 4 HCUs, so every round is sized in
+#: multiples of 4: that keeps a wave able to fill all four devices at once and
+#: keeps the two gates symmetric (the paper is one wave too).
+OPERATORS_PER_GATE = 4
+GPU_COUNT = 4
+
+#: A round is split into two stages. ``performance`` hands a fresh operator set
+#: to DKAO and judges only the performance gate; when that passes, the very same
+#: harness is handed to DKAO a second time with the frozen paper — that retake
+#: is the ``generalization`` stage. Round 1 is neither: it is a *baseline* round
+#: that measures the paper and judges nothing.
+STAGE_BASELINE = "baseline"
+STAGE_PERFORMANCE = "performance"
+STAGE_GENERALIZATION = "generalization"
+
+#: What each stage costs against the round budget. A normal round is 1.0; the
+#: generalization retake is half a round (the operator asked for "iteration
+#: 2.5"), because it re-runs a fixed, already-chosen paper instead of opening a
+#: new round of the search.
+PERFORMANCE_ROUND_COST = 1.0
+GENERALIZATION_ROUND_COST = 0.5
+
+#: The frozen "retake the paper with this harness" request.
+PENDING_GENERALIZATION_FILE = "pending_generalization.json"
+
+#: Directory holding one evaluation attempt's DKAO children
+#: (``children//``). Each attempt -- a first run and
+#: every environment retry of it, and the generalization retake -- gets its own,
+#: so no two DKAO runs ever share a child workspace.
+CHILDREN_DIR = "children"
+
+
+def round_stage(*, iteration: int, pending: Optional[Dict[str, Any]]) -> str:
+    """Which stage this pass through the loop is.
+
+    A pending generalization retake wins over the round number: the loop may
+    come back to the same iteration to retake the paper, and that pass is the
+    ``generalization`` stage even though its directory is ``iteration_002``.
+    """
+    if pending:
+        return STAGE_GENERALIZATION
+    return STAGE_BASELINE if int(iteration) <= 1 else STAGE_PERFORMANCE
+
+
+def iteration_of_dir_name(name: str) -> int:
+    """The iteration a ``children/`` directory belongs to (0 if none).
+
+    Names are ``iteration_NNN`` plus optional suffixes (``_a2`` for a retry,
+    ``_generalization`` for the retake). The iteration is always the second
+    underscore-separated field, so it is read by position: a regex would also
+    match the attempt counter and silently merge two rounds' children.
+    """
+    parts = str(name).split("_")
+    if len(parts) < 2 or parts[0] != "iteration":
+        return 0
+    try:
+        return int(parts[1])
+    except ValueError:
+        return 0
+
+
+def stage_dir_name(stage: str, iteration: int, *, attempt: int = 1) -> str:
+    """Directory under ``/children`` holding one evaluation attempt.
+
+    Every attempt gets its own directory. A retry of one question is a *new*
+    DKAO run, and DKAO binds the child workspace's ``main`` symlink to the exact
+    kernel repo it was told to use: reusing the directory makes the retry die in
+    seconds with "workspace main already points to ..., not requested ...",
+    which turns a transient environment hiccup into an apparently permanent
+    failure that burns the whole retry budget. A fresh directory also keeps each
+    attempt's evidence instead of overwriting the one that failed.
+    """
+    name = f"iteration_{int(iteration):03d}"
+    if int(attempt) > 1:
+        name += f"_a{int(attempt)}"
+    if stage == STAGE_GENERALIZATION:
+        name += "_generalization"
+    return name
+
+
+def attempt_dirs(children_root: Path, iteration: int) -> List[Path]:
+    """Every attempt directory of one iteration, oldest attempt first.
+
+    The retake sorts after the performance pass it belongs to (both carry
+    ``_generalization`` in the name only for the retake's own directory).
+    """
+    root = Path(children_root)
+    if not root.is_dir():
+        return []
+    found = [d for d in root.iterdir()
+             if d.is_dir() and iteration_of_dir_name(d.name) == int(iteration)]
+    return sorted(found, key=lambda d: (d.name.count("_a"),
+                                        "_generalization" in d.name, d.name))
+
+
+def resolve_child_dir(children_root: Path, iteration: int,
+                      question_id: str) -> Optional[Path]:
+    """The directory holding one question's current attempt (newest wins).
+
+    ``None`` means the question has no child in this iteration at all: returning
+    a path that does not exist would make "not measured" look like "measured,
+    directory missing" to every caller.
+    """
+    found = [d / str(question_id) for d in attempt_dirs(children_root, iteration)
+             if (d / str(question_id)).is_dir()]
+    return found[-1] if found else None
+
+
+def child_dir_map(children_root: Path, iteration: int) -> Dict[str, Path]:
+    """``{question_id: directory}`` for one iteration, newest attempt wins."""
+    out: Dict[str, Path] = {}
+    for attempt_dir in attempt_dirs(children_root, iteration):
+        for qid in sorted(p.name for p in attempt_dir.iterdir() if p.is_dir()):
+            out[qid] = attempt_dir / qid
+    return out
+
+
+def stage_seconds(pending: Optional[Dict[str, Any]]) -> int:
+    """Stable suffix making a retake's child ids distinct from the first pass."""
+    if not pending:
+        return 0
+    try:
+        return int(pending.get("frozen_at") or 0)
+    except (TypeError, ValueError):
+        return 0
+
+
+def device_for_index(index: int) -> int:
+    """The device a question at ``index`` in a round is handed to.
+
+    A fixed rotation, not a reading of the cards: ``harness_evolve`` does not
+    manage GPU occupancy (DKAO's own admission gate does), so the assignment
+    only has to be a stable, even layout — N questions land as N / 4 sequential
+    waves with one question per device in each.
+    """
+    return int(index) % GPU_COUNT
+
+
+def normalized_per_gate(value: Any, default: int = OPERATORS_PER_GATE) -> int:
+    """Round the requested question count to a non-zero multiple of 4.
+
+    A count that is not a multiple of 4 would leave devices idle in a wave, so
+    it is rounded up to the next multiple (and anything unusable falls back to
+    the default).
+    """
+    try:
+        count = int(value)
+    except (TypeError, ValueError):
+        return default
+    if count <= 0:
+        return default
+    if count % GPU_COUNT:
+        count = ((count // GPU_COUNT) + 1) * GPU_COUNT
+    return count
+
+
+#: An operator is "decode" at small M and "prefill" at large M, matching the
+#: family naming the pool uses (``decode__o_proj`` / ``prefill__o_proj``).
+DECODE_MAX_M = 32
+REGIME_DECODE = "decode"
+REGIME_PREFILL = "prefill"
+
+
+def regime_of(entry: Dict[str, Any]) -> str:
+    """Which regime an operator belongs to (family first, shape as fallback)."""
+    family = str(entry.get("family") or "")
+    if family.startswith(REGIME_DECODE + "__"):
+        return REGIME_DECODE
+    if family.startswith(REGIME_PREFILL + "__"):
+        return REGIME_PREFILL
+    contract = entry.get("contract") or {}
+    try:
+        m = int(contract.get("M") or 0)
+    except (TypeError, ValueError):
+        m = 0
+    return REGIME_DECODE if 0 < m <= DECODE_MAX_M else REGIME_PREFILL
+
+
+def sample_operators(instances: Dict[str, Dict[str, Any]], *,
+                     table: Optional[Dict[str, Any]] = None,
+                     count: int = OPERATORS_PER_GATE,
+                     exclude: Iterable[str] = (),
+                     rng: Optional[random.Random] = None,
+                     require_judgeable: bool = True,
+                     require_regimes: bool = True,
+                     ) -> List[str]:
+    """A random operator sample, skipping ids that cannot be judged.
+
+    ``require_judgeable`` keeps out operators with neither a variant nor a fixed
+    baseline: those could not produce a win/loss, so they would only dilute the
+    ratio.
+
+    ``require_regimes`` guarantees the sample spans both regimes (at least one
+    decode and at least one prefill) so a round is never judged on a single
+    shape class; the rest of the sample is filled at random from everything
+    that is left.
+    """
+    excluded = {str(x) for x in exclude}
+    available = [iid for iid in instances if iid not in excluded]
+    if require_judgeable:
+        available = [iid for iid in available
+                     if resolve_variant(table or {}, iid,
+                                        instances=instances) is not None]
+    if not available:
+        return []
+    picker = rng or random
+    count = max(1, min(int(count), len(available)))
+    picked: List[str] = []
+    if require_regimes:
+        for regime in (REGIME_DECODE, REGIME_PREFILL):
+            members = [iid for iid in available
+                       if regime_of(instances[iid]) == regime
+                       and iid not in picked]
+            if members and len(picked) < count:
+                picked.append(picker.choice(sorted(members)))
+    rest = [iid for iid in available if iid not in picked]
+    picker.shuffle(rest)
+    picked.extend(rest[:max(0, count - len(picked))])
+    return sorted(picked[:count])
+
+
+def _statement_groups(state: Dict[str, Any]) -> Dict[str, Any]:
+    """The persisted groups, whichever shape the caller handed over.
+
+    The pipeline passes what :func:`group_state` returned (the groups dict
+    itself); older callers wrap it as ``{"groups": ...}``. Reading only one of
+    the two shapes silently loses the frozen paper — and a lost paper means the
+    performance draw may hand the generalization exam's own operators back.
+    """
+    if not isinstance(state, dict):
+        return {}
+    nested = state.get("groups")
+    if isinstance(nested, dict):
+        return dict(nested)
+    if "generalization_ids" in state:
+        return dict(state)
+    return {}
+
+
+def select_round_questions(*, iteration: int,
+                           table: Dict[str, Any],
+                           instances: Dict[str, Dict[str, Any]],
+                           state: Dict[str, Any],
+                           count: int = OPERATORS_PER_GATE,
+                           rng: Optional[random.Random] = None,
+                           pending: Optional[Dict[str, Any]] = None,
+                           ) -> Dict[str, Any]:
+    """Which operators this pass measures, and in which role.
+
+    Returns ``{"performance_ids", "generalization_ids", "defines_paper",
+    "stage"}``.
+
+    * ``pending`` set → this pass is the **generalization retake**: it measures
+      exactly the paper, again, and judges the generalization gate. No new
+      operators are drawn.
+    * otherwise round 1 → the **baseline** round: it draws the paper (which
+      becomes the fixed exam) and measures it, judging nothing.
+    * otherwise → **performance** stage: a fresh random draw excluding the
+      paper; the paper comes back only if this pass wins and the retake runs.
+    """
+    groups = _statement_groups(state)
+    paper = list(groups.get("generalization_ids") or [])
+    if pending:
+        ids = [str(i) for i in (pending.get("operator_ids") or paper)]
+        return {"performance_ids": ids, "generalization_ids": ids,
+                "defines_paper": False, "stage": STAGE_GENERALIZATION}
+    if iteration <= 1 or not paper:
+        # Round 1 defines the generalization paper and measures it, but it is a
+        # baseline round: nothing is judged and nothing is promoted from it.
+        draw = sample_operators(instances, table=table, count=count, rng=rng)
+        return {"performance_ids": draw, "generalization_ids": list(draw),
+                "defines_paper": True, "stage": STAGE_BASELINE}
+    performance = sample_operators(instances, table=table, count=count,
+                                   exclude=paper, rng=rng)
+    return {"performance_ids": performance,
+            "generalization_ids": paper, "defines_paper": False,
+            "stage": STAGE_PERFORMANCE}
+
+
+def _load_json(path: Path) -> Any:
+    try:
+        return json.loads(Path(path).read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+
+
+def read_pending_generalization(exp_dir: Path) -> Optional[Dict[str, Any]]:
+    """The frozen retake request, if the last performance stage passed.
+
+    While this file exists, the next pass through the loop is the
+    generalization stage: the same harness bytes go to DKAO again, with the
+    paper's operators. Nothing else is decided until that retake is judged.
+    """
+    data = _load_json(Path(exp_dir) / PENDING_GENERALIZATION_FILE)
+    return data if isinstance(data, dict) else None
+
+
+def write_pending_generalization(exp_dir: Path, *,
+                                 iteration: int,
+                                 harness_revision: str,
+                                 snapshot_dir: Optional[str],
+                                 operator_ids: Sequence[str],
+                                 performance_gate: Dict[str, Any],
+                                 ) -> Dict[str, Any]:
+    """Freeze "retake the paper with this harness" for the next pass."""
+    exp_dir = Path(exp_dir)
+    payload = {
+        "schema": "he-pending-generalization/1",
+        "iteration": int(iteration),
+        "operator_ids": [str(i) for i in operator_ids],
+        "harness_revision": str(harness_revision or ""),
+        "snapshot_dir": str(snapshot_dir) if snapshot_dir else None,
+        "performance_wins": list(performance_gate.get("wins") or []),
+        "performance_gate_status": performance_gate.get("status"),
+        "frozen_at": time.time(),
+    }
+    exp_dir.mkdir(parents=True, exist_ok=True)
+    path = exp_dir / PENDING_GENERALIZATION_FILE
+    tmp = path.with_suffix(".json.tmp")
+    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2,
+                              sort_keys=True), encoding="utf-8")
+    tmp.replace(path)
+    return payload
+
+
+def clear_pending_generalization(exp_dir: Path) -> None:
+    try:
+        (Path(exp_dir) / PENDING_GENERALIZATION_FILE).unlink()
+    except OSError:
+        pass
+
+
+def evaluate_gate(*, kind: str, operator_ids: Sequence[str],
+                  results: Dict[str, Any],
+                  table: Dict[str, Any],
+                  instances: Dict[str, Dict[str, Any]],
+                  ) -> Dict[str, Any]:
+    """Judge one gate's operator set against the variants they must beat."""
+    ids = [iid for iid in operator_ids]
+    reference = reference_results(table, ids, instances=instances)
+    candidate = {iid: results[iid] for iid in ids if iid in results}
+    verdict = gate_verdict(reference, candidate, kind=kind)
+    verdict["iteration"] = None
+    verdict["requested_ids"] = ids
+    verdict["requested"] = len(ids)
+    missing = [iid for iid in ids if iid not in results]
+    if missing:
+        unmeasured = sorted(set(verdict["unmeasured"]) | set(missing))
+        verdict["unmeasured"] = unmeasured
+        verdict["reasons"] = list(verdict["reasons"]) + [
+            "not measured this round: " + ", ".join(sorted(missing))]
+        # A gate is a statement about the harness, so it may only be made when
+        # the round actually produced the evidence. Two cases are decided here:
+        #
+        #  * the measured part already loses — a question below the floor, or so
+        #    few wins that even if every missing question won, the count would
+        #    not reach the required wins. That is a FAIL, and saying so now
+        #    saves the round a relaunch;
+        #  * otherwise the outcome is open: the harness may be fine and the
+        #    machine may have eaten a measurement. Marking that FAIL blames the
+        #    harness for a missing number (and, on the generalization retake,
+        #    spends the candidate on it), so the round carries no verdict and is
+        #    re-measured instead.
+        needed_total = (max(1, min(len(ids),
+                                 math.ceil(len(ids) * gate_ratio(kind))))
+                        if ids else 0)
+        verdict["wins_needed_over_requested"] = needed_total
+        possible_wins = len(verdict["wins"]) + len(unmeasured)
+        if verdict["below_floor"] or possible_wins < needed_total:
+            verdict["status"] = "FAIL"
+            verdict["reasons"] = list(verdict["reasons"]) + [
+                "gate cannot pass even if every unmeasured question won "
+                f"({possible_wins} possible win(s) vs {needed_total} needed)"]
+        else:
+            verdict["status"] = "INCOMPLETE"
+    return verdict
+
+
+def group_state(table: Dict[str, Any]) -> Dict[str, Any]:
+    """The persisted groups (the fixed generalization paper lives here)."""
+    groups = table.get("groups")
+    return dict(groups) if isinstance(groups, dict) else {}
+
+
+def record_paper(table: Dict[str, Any], *, iteration: int,
+                 generalization_ids: Sequence[str],
+                 performance_ids: Sequence[str]) -> Dict[str, Any]:
+    """Persist this run's fixed generalization paper on the variant table.
+
+    The paper is written once and never replaced: it is the run's frozen exam,
+    so a later round must not be able to move the goalposts.
+    """
+    groups = group_state(table)
+    groups.setdefault("generalization_ids", list(generalization_ids))
+    groups.setdefault("generalization_defined_at", iteration)
+    groups["last_performance_ids"] = list(performance_ids)
+    return groups
+
+
+def baseline_round_verdict(results: Dict[str, Any],
+                          operator_ids: Sequence[str]) -> Dict[str, Any]:
+    """Whether the baseline round produced usable evidence for every operator.
+
+    Round 1 judges no gate, so there is no win ratio to check — but it still
+    owes one usable measurement per operator: the kernels that beat their
+    variant are what seeds the pool for every later round to beat, and a paper
+    that was never really measured cannot serve as the run's exam.
+    """
+    ids = [str(i) for i in operator_ids]
+    unusable: List[str] = []
+    for iid in ids:
+        row = results.get(iid)
+        if not isinstance(row, dict):
+            unusable.append(iid)
+            continue
+        median = row.get("median_us")
+        try:
+            value = float(median)
+        except (TypeError, ValueError):
+            unusable.append(iid)
+            continue
+        if value <= 0 or row.get("correctness_ok") is False:
+            unusable.append(iid)
+    return {
+        "status": "PASS" if not unusable else "FAIL",
+        "operators": ids,
+        "measured": sorted(set(ids) - set(unusable)),
+        "failed": sorted(unusable),
+        "reasons": ([] if not unusable else
+                    ["baseline round produced no usable measurement for: "
+                     + ", ".join(sorted(unusable))]),
+    }
+
+
+def gate_ratio(kind: str) -> float:
+    return (PERFORMANCE_GATE_WIN_RATIO if kind == "performance"
+            else GENERALIZATION_GATE_WIN_RATIO)
+
+
+def paper_path(exp_dir: Path) -> Path:
+    return Path(exp_dir) / "generalization_paper.json"
diff --git a/metainfer/tasks/harness_evolve/orchestrator/state.py b/metainfer/tasks/harness_evolve/orchestrator/state.py
new file mode 100644
index 00000000..3ceb55cf
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/state.py
@@ -0,0 +1,195 @@
+"""Experiment state recording for harness_evolve (file-system-as-database).
+
+The WebUI launcher runs the HE orchestrator as a bare ``cli run`` subprocess,
+so nothing outside the pipeline would otherwise write the per-task
+``run.json`` / ``timeline.jsonl`` that the shell reads for the phase graph
+(``state_reader.read_run`` / ``read_timeline``). This module gives the
+pipeline a small writer that keeps those files in the same shape DKAO writes:
+``run.json`` fields merged with ``state_reader`` defaults, timeline rows as
+``{"ts": float, "type": str, "payload": dict}``.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+_RUN_DEFAULTS: Dict[str, Any] = {
+    "task_id": None,
+    "task_type": None,
+    "created_at": 0,
+    "current_iteration": 0,
+    "current_phase": "idle",
+    "last_update": 0,
+    "finished": False,
+    "final_status": None,
+    "last_outcome": None,
+    "last_transition_label": None,
+    "notes": [],
+}
+
+
+def _load_json(path: Path) -> Optional[Dict[str, Any]]:
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    return value if isinstance(value, dict) else None
+
+
+def read_run(state_dir: Path) -> Dict[str, Any]:
+    data = _load_json(Path(state_dir) / "run.json")
+    if data is None:
+        return dict(_RUN_DEFAULTS)
+    return {**_RUN_DEFAULTS, **data}
+
+
+def set_run(state_dir: Path, **fields: Any) -> Dict[str, Any]:
+    """Merge fields into run.json (thread-safe via atomic replace)."""
+    state_dir = Path(state_dir)
+    state_dir.mkdir(parents=True, exist_ok=True)
+    path = state_dir / "run.json"
+    run = read_run(state_dir)
+    run.update(fields)
+    run["last_update"] = time.time()
+    _atomic_write_json(path, run)
+    return run
+
+
+def append_timeline(state_dir: Path, event_type: str,
+                    payload: Optional[Dict[str, Any]] = None) -> None:
+    state_dir = Path(state_dir)
+    state_dir.mkdir(parents=True, exist_ok=True)
+    path = state_dir / "timeline.jsonl"
+    row = {
+        "ts": time.time(),
+        "type": event_type,
+        "payload": payload or {},
+    }
+    with path.open("a", encoding="utf-8") as fh:
+        fh.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+
+def _atomic_write_json(path: Path, data: Dict[str, Any]) -> None:
+    tmp = path.with_suffix(".json.tmp")
+    tmp.write_text(json.dumps(data, ensure_ascii=False, indent=2),
+                   encoding="utf-8")
+    os.replace(tmp, path)
+
+
+def last_phase_event(state_dir: Path) -> Optional[Dict[str, Any]]:
+    """Most recent timeline entry (used by tests / diagnostics)."""
+    path = Path(state_dir) / "timeline.jsonl"
+    if not path.is_file():
+        return None
+    last: Optional[Dict[str, Any]] = None
+    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+        try:
+            row = json.loads(line)
+        except ValueError:
+            continue
+        if isinstance(row, dict):
+            last = row
+    return last
+
+
+def list_timeline(state_dir: Path) -> List[Dict[str, Any]]:
+    out: List[Dict[str, Any]] = []
+    path = Path(state_dir) / "timeline.jsonl"
+    if not path.is_file():
+        return out
+    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+        try:
+            row = json.loads(line)
+        except ValueError:
+            continue
+        if isinstance(row, dict):
+            out.append(row)
+    return out
+
+
+#: Where a running process records the harness code it actually loaded, so a
+#: later round can notice the files on disk moved on without it.
+REVISION_FILE = "harness_revision.json"
+
+
+def harness_revision() -> str:
+    """Fingerprint of the harness_evolve *code* (not the experiment workspace).
+
+    Rounds run for hours, so an operator who fixes the loop cannot expect the
+    running process to pick the change up. This is how the next round notices
+    that a restart is required instead of silently running stale logic.
+    """
+    import hashlib
+
+    root = Path(__file__).resolve().parent
+    digest = hashlib.sha256()
+    for path in sorted(root.rglob("*.py")):
+        try:
+            data = path.read_bytes()
+        except OSError:
+            continue
+        digest.update(path.relative_to(root).as_posix().encode("utf-8"))
+        digest.update(b"\0")
+        digest.update(data)
+        digest.update(b"\0")
+    return digest.hexdigest()[:16]
+
+
+def write_harness_revision(state_dir: Path,
+                           revision: Optional[str] = None) -> str:
+    """Record the revision this process actually loaded."""
+    value = str(revision or harness_revision())
+    state_dir = Path(state_dir)
+    state_dir.mkdir(parents=True, exist_ok=True)
+    _atomic_write_json(state_dir / REVISION_FILE, {
+        "revision": value, "pid": os.getpid(), "started_at": time.time()})
+    return value
+
+
+def loaded_harness_revision(state_dir: Path) -> Optional[str]:
+    data = _load_json(Path(state_dir) / REVISION_FILE)
+    if not data:
+        return None
+    value = data.get("revision")
+    return str(value) if value else None
+
+
+#: How many iterations the experiment should reach in TOTAL, kept in a file
+#: outside the running process so the target can be adjusted while the outer
+#: loop is running (raise it to continue, lower it to stop earlier).
+TARGET_FILE = "target_iterations.json"
+
+
+def target_iterations_path(exp_dir: Path) -> Path:
+    return Path(exp_dir) / TARGET_FILE
+
+
+def read_target_iterations(exp_dir: Path) -> Optional[int]:
+    """Total iteration count the experiment should reach (None = unset)."""
+    data = _load_json(target_iterations_path(exp_dir))
+    if not data:
+        return None
+    try:
+        value = int(data.get("target"))
+    except (TypeError, ValueError):
+        return None
+    return value if value >= 1 else None
+
+
+def write_target_iterations(exp_dir: Path, target: int) -> Dict[str, Any]:
+    exp_dir = Path(exp_dir)
+    exp_dir.mkdir(parents=True, exist_ok=True)
+    payload = {"target": int(target), "updated_at": time.time()}
+    _atomic_write_json(target_iterations_path(exp_dir), payload)
+    return payload
+
+
+def clear_target_iterations(exp_dir: Path) -> None:
+    try:
+        target_iterations_path(exp_dir).unlink()
+    except OSError:
+        pass
diff --git a/metainfer/tasks/harness_evolve/orchestrator/supervisor.py b/metainfer/tasks/harness_evolve/orchestrator/supervisor.py
new file mode 100644
index 00000000..c961efd7
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/supervisor.py
@@ -0,0 +1,737 @@
+"""Incident supervisor for unattended harness_evolve runs.
+
+The outer loop is long (hours per round) and can die in ways that are not the
+harness' fault (the GPU filled up, an agent transport closed, the host killed
+the process). This supervisor watches the experiment and, per incident class,
+takes the *safe* action automatically:
+
+  resource       the host ran out of memory and killed the run (exit 137 /
+                 OOM markers in the log) -> wait (default 30 min), then restart
+                 the round. HE never inspects the cards: a device verdict is
+                 DKAO's (FLOW.md §0)
+  agent          flaky agent transport/timeout -> exponential backoff restart
+  code           deterministic traceback/config error -> do NOT loop; stop and
+                 report (retrying would replay the same failure)
+  unknown        restart once, then stop
+
+Every decision is a pure function (:func:`classify_incident`,
+:func:`decide_action`) so the behaviour is testable without touching GPUs, and
+every action is appended to ``incidents.jsonl`` for audit. Repeated failures of
+the same class trip a circuit breaker that stops the run instead of burning
+compute unattended.
+
+Optionally (``--diagnose-agent``) it asks a read-only DSH agent to write
+``incident_report.md`` with a root-cause hypothesis; that report is advisory
+only — it never performs restarts itself.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+
+CATEGORY_RESOURCE = "resource"
+CATEGORY_AGENT = "agent"
+CATEGORY_CODE = "code"
+CATEGORY_UNKNOWN = "unknown"
+
+_RESOURCE_MARKERS = (
+    "outofmemory", "out of memory", "hipErrorOutOfMemory", "bad_alloc",
+    "cannot allocate memory", "no space left", "memory exhausted",
+)
+_AGENT_MARKERS = (
+    "transportclosederror", "timed out after", "timeoutexpired",
+    "nonzero exit 1", "agent attempt", "connection reset",
+)
+_CODE_MARKERS = ("traceback (most recent call last)",)
+
+
+def classify_incident(*, log_tail: str = "",
+                      exit_code: Optional[int] = None) -> Tuple[str, List[str]]:
+    """Return ``(category, evidence)`` for one observed failure.
+
+    The verdict comes from *this run's own* evidence (exit status and log text);
+    HE holds no device verdict of its own (FLOW.md §0), so nothing here reads
+    the cards.
+    """
+    text = (log_tail or "").lower()
+    evidence: List[str] = []
+    if exit_code in (137, -9):
+        evidence.append(f"exit code {exit_code} (killed, usually OOM)")
+    for marker in _RESOURCE_MARKERS:
+        if marker.lower() in text:
+            evidence.append(f"log mentions {marker!r}")
+            break
+    if evidence:
+        return CATEGORY_RESOURCE, evidence
+    for marker in _AGENT_MARKERS:
+        if marker.lower() in text:
+            return CATEGORY_AGENT, [f"log mentions {marker!r}"]
+    for marker in _CODE_MARKERS:
+        if marker in text:
+            return CATEGORY_CODE, ["deterministic traceback in the log"]
+    return CATEGORY_UNKNOWN, evidence or ["no recognisable failure signature"]
+
+
+def decide_action(*, category: str, same_class_streak: int,
+                  waits_done: int, restarts_done: int,
+                  running: bool,
+                  max_restarts: int = 3, max_waits: int = 6,
+                  breaker_streak: int = 3,
+                  wait_seconds: int = 1800) -> Dict[str, Any]:
+    """Pure policy: what to do about one incident.
+
+    ``action`` is one of ``restart`` | ``wait`` | ``stop`` | ``run`` (healthy)
+    | ``done`` (target already reached).
+    """
+    if running:
+        return {"action": "run", "sleep_s": 0,
+                "reason": "experiment process is alive"}
+    if same_class_streak >= breaker_streak and category != CATEGORY_RESOURCE:
+        return {"action": "stop", "sleep_s": 0,
+                "reason": (f"circuit breaker: {same_class_streak} consecutive "
+                           f"{category} failures")}
+    if category == CATEGORY_CODE:
+        return {"action": "stop", "sleep_s": 0,
+                "reason": ("deterministic code/config failure; restarting "
+                           "would replay it")}
+    if category == CATEGORY_RESOURCE:
+        if waits_done >= max_waits:
+            return {"action": "stop", "sleep_s": 0,
+                    "reason": f"resource still short after {waits_done} waits"}
+        return {"action": "wait", "sleep_s": int(wait_seconds),
+                "reason": ("memory pressure: wait, re-check free VRAM, then "
+                           "restart")}
+    if restarts_done >= max_restarts:
+        return {"action": "stop", "sleep_s": 0,
+                "reason": f"gave up after {restarts_done} restarts"}
+    backoff = wait_seconds if category == CATEGORY_AGENT else 300
+    return {"action": "restart", "sleep_s": min(backoff, 1800),
+            "reason": f"{category} failure; restart with backoff"}
+
+
+def pid_is_live(pid: Optional[int], *, marker: str = "harness_evolve",
+                proc_root: Path = Path("/proc")) -> bool:
+    """Live (non-zombie) process whose cmdline mentions ``marker``.
+
+    Shared by the supervisor and the server so both agree on what "running"
+    means: a zombie child or a recycled pid must never count as alive.
+
+    ``proc_root`` exists so tests can build a fake process table (the same hook
+    ``live_child_count`` already has); it defaults to the real ``/proc``.
+    """
+    if not pid or pid == os.getpid():
+        return False
+    proc = Path(proc_root) / str(int(pid))
+    if not proc.exists():
+        return False
+    try:
+        state = (proc / "stat").read_text(encoding="utf-8").split(") ", 1)[1].split()[0]
+    except (OSError, IndexError):
+        return False
+    if state == "Z":
+        return False
+    try:
+        cmdline = (proc / "cmdline").read_bytes().replace(b"\0", b" ").decode(
+            "utf-8", "replace")
+    except OSError:
+        return False
+    return marker in cmdline
+
+
+def pid_from_file(path: Path) -> Optional[int]:
+    """The PID recorded in a pid file, whether it is JSON or a bare number.
+
+    Both writers in this package use JSON (``{"pid": …}``); an older process may
+    have left a bare number. Reading only the bare form is what made every
+    current pid file unparseable, which silently pushed every liveness question
+    onto the ``/proc`` cmdline scan — and that scan then matched the *caller
+    itself*, so a resume could never start round 1.
+    """
+    try:
+        raw = path.read_text(encoding="utf-8").strip()
+    except OSError:
+        return None
+    if not raw:
+        return None
+    if raw.startswith("{"):
+        try:
+            data = json.loads(raw)
+        except ValueError:
+            return None
+        try:
+            return int(data.get("pid"))
+        except (TypeError, ValueError):
+            return None
+    try:
+        return int(raw.split()[0])
+    except (IndexError, ValueError):
+        return None
+
+
+def task_process_alive(state_dir: Path, *,
+                       exclude: Optional[int] = None,
+                       proc_root: Path = Path("/proc"),
+                       self_pid: Optional[int] = None) -> Optional[int]:
+    """PID of *another* live orchestrator for this task, whichever path started it.
+
+    Tasks can be started by the server ("run") or by the supervisor
+    ("resume"); ``resume.pid`` alone is not enough, and a server-started task
+    has no resume pid at all. Falling back to a cmdline scan keeps the
+    supervisor from "restarting" a perfectly healthy run.
+
+    The caller's own pid is never returned (and ``exclude`` adds another): a
+    resume asking "is somebody else running this experiment?" must not answer
+    itself, or it will stand down forever and no round will ever start.
+    """
+    state_dir = Path(state_dir)
+    excluded = {int(self_pid) if self_pid is not None else os.getpid()}
+    if exclude:
+        excluded.add(int(exclude))
+    for name in ("orchestrator.pid", "resume.pid"):
+        pid = pid_from_file(state_dir / name)
+        if not pid or pid in excluded:
+            continue
+        if pid_is_live(pid, proc_root=proc_root):
+            return pid
+    marker = str(state_dir)
+    for entry in Path(proc_root).iterdir():
+        if not entry.name.isdigit():
+            continue
+        pid = int(entry.name)
+        if pid in excluded:
+            continue
+        try:
+            cmdline = (entry / "cmdline").read_bytes().replace(b"\0", b" ").decode(
+                "utf-8", "replace")
+        except OSError:
+            continue
+        if "harness_evolve.orchestrator.cli" in cmdline and marker in cmdline:
+            return pid
+    return None
+
+
+class Supervisor:
+    """Watch one experiment and keep it moving toward its target round."""
+
+    def __init__(self, state_dir: Path, workspace_dir: Path, *,
+                 target: Optional[int] = None, interval_s: int = 30,
+                 wait_seconds: int = 1800, max_restarts: int = 3,
+                 max_waits: int = 6, diagnose_agent: bool = False,
+                 stall_seconds: int = 2400,
+                 clock: Callable[[], float] = time.time,
+                 sleeper: Callable[[float], None] = time.sleep,
+                 resume_launcher: Optional[Callable[[], int]] = None,
+                 ) -> None:
+        self.state_dir = Path(state_dir)
+        self.workspace_dir = Path(workspace_dir)
+        self.exp_dir = self.workspace_dir
+        self.target = target
+        self.interval_s = int(interval_s)
+        self.wait_seconds = int(wait_seconds)
+        self.max_restarts = int(max_restarts)
+        self.max_waits = int(max_waits)
+        self.diagnose_agent = bool(diagnose_agent)
+        self.stall_seconds = int(stall_seconds)
+        self._resume_proc = None
+        self.clock = clock
+        self.sleeper = sleeper
+        self._resume_launcher = resume_launcher
+        self.state_path = self.exp_dir / "supervisor_state.json"
+        self.incidents_path = self.exp_dir / "incidents.jsonl"
+        self.stop_flag = self.exp_dir / "supervisor.stop"
+
+    # ---------------------------------------------------------------- state
+    def read_state(self) -> Dict[str, Any]:
+        try:
+            return json.loads(self.state_path.read_text(encoding="utf-8"))
+        except (OSError, ValueError):
+            return {}
+
+    def write_state(self, **fields: Any) -> Dict[str, Any]:
+        state = self.read_state()
+        state.update(fields)
+        state["updated_at"] = self.clock()
+        try:
+            self.state_path.write_text(
+                json.dumps(state, ensure_ascii=False, indent=2),
+                encoding="utf-8")
+        except OSError:
+            pass
+        return state
+
+    def record_incident(self, **fields: Any) -> None:
+        row = {"ts": self.clock(), **fields}
+        try:
+            with self.incidents_path.open("a", encoding="utf-8") as fh:
+                fh.write(json.dumps(row, ensure_ascii=False) + "\n")
+        except OSError:
+            pass
+
+    # ----------------------------------------------------------- inspection
+    def _read_json(self, path: Path) -> Dict[str, Any]:
+        try:
+            return json.loads(path.read_text(encoding="utf-8"))
+        except (OSError, ValueError):
+            return {}
+
+    def run_state(self) -> Dict[str, Any]:
+        return self._read_json(self.state_dir / "run.json")
+
+    def current_target(self) -> Optional[int]:
+        if self.target is not None:
+            return int(self.target)
+        data = self._read_json(self.exp_dir / "target_iterations.json")
+        try:
+            return int(data.get("target"))
+        except (TypeError, ValueError):
+            return None
+
+    # NOTE: there is deliberately no ``free_gb()`` / device query here any more.
+    # HE does not read the cards (FLOW.md §0): whether a device is clean enough
+    # to measure on is DKAO's verdict, taken in the child's own gate, and the
+    # only place it is recorded is the child's ``measurement_gate.jsonl``.
+
+    def resume_pid(self) -> Optional[int]:
+        try:
+            return int((self.state_dir / "resume.pid")
+                       .read_text(encoding="utf-8").split()[0])
+        except (OSError, ValueError, IndexError):
+            return None
+
+    def process_alive(self) -> bool:
+        """True when *any* orchestrator process for this task is alive.
+
+        Covers both launch paths (server "run" and supervisor "resume") and
+        rejects zombies, which is what the stuck-then-restarting failure mode
+        needed.
+        """
+        return task_process_alive(self.state_dir) is not None
+
+    def progress_age(self) -> Optional[float]:
+        """Seconds since *any* part of the run last made progress.
+
+        The outer loop is mostly waiting: during ``evaluate`` it can sit for
+        hours while the DKAO children work, and the outer timeline barely
+        changes. Counting child state as progress prevents the stall guard
+        from killing a perfectly healthy long round (which would turn into a
+        restart storm).
+        """
+        newest = 0.0
+        for name in ("run.json", "timeline.jsonl"):
+            try:
+                newest = max(newest, (self.state_dir / name).stat().st_mtime)
+            except OSError:
+                continue
+        iteration = int(self.run_state().get("current_iteration") or 0)
+        children_root = self.exp_dir / "children"
+        if iteration and children_root.is_dir():
+            # One round can own two child directories: the performance pass
+            # (``iteration_NNN``) and, when that pass won its gate, the
+            # generalization retake (``iteration_NNN_generalization``). Both are
+            # part of "what is running right now".
+            for name in (f"iteration_{iteration:03d}",
+                         f"iteration_{iteration:03d}_generalization"):
+                current = children_root / name
+                if not current.is_dir():
+                    continue
+                for child in current.iterdir():
+                    for relative in ("state/timeline.jsonl", "state/run.json",
+                                     "orchestrator-external.log"):
+                        try:
+                            newest = max(newest,
+                                         (child / relative).stat().st_mtime)
+                        except OSError:
+                            continue
+        if newest <= 0:
+            return None
+        return max(0.0, self.clock() - newest)
+
+    def live_child_count(self, proc_root: Path = Path("/proc")) -> int:
+        """Live DKAO children of the iteration currently being evaluated.
+
+        The stall guard needs *processes*, not only file mtimes: a child that is
+        generating or measuring a kernel can be silent for an hour, and reading
+        that as "the run died" produced a duplicate resume per stall.
+        """
+        iteration = int(self.run_state().get("current_iteration") or 0)
+        if not iteration:
+            return 0
+        # ``iteration_NNN`` also matches ``iteration_NNN_generalization``: both
+        # belong to the round being evaluated (the retake is the same round).
+        needle = f"children/iteration_{iteration:03d}"
+        alive = 0
+        try:
+            entries = list(proc_root.iterdir())
+        except OSError:
+            return 0
+        for entry in entries:
+            if not entry.name.isdigit():
+                continue
+            try:
+                cmdline = (entry / "cmdline").read_bytes().replace(
+                    b"\0", b" ").decode("utf-8", "replace")
+            except OSError:
+                continue
+            if needle not in cmdline or "dcu_kernel_auto_opt" not in cmdline:
+                continue
+            try:
+                state = (entry / "stat").read_text().split(") ", 1)[1].split()[0]
+            except (OSError, IndexError, ValueError):
+                continue
+            if state != "Z":
+                alive += 1
+        return alive
+
+    def stale_run_pids(self) -> List[int]:
+        """Orchestrator processes for this task that are still alive."""
+        marker = str(self.state_dir)
+        out: List[int] = []
+        try:
+            entries = list(Path("/proc").iterdir())
+        except OSError:
+            return out
+        for entry in entries:
+            if not entry.name.isdigit():
+                continue
+            pid = int(entry.name)
+            try:
+                cmdline = (entry / "cmdline").read_bytes().replace(
+                    b"\0", b" ").decode("utf-8", "replace")
+            except OSError:
+                continue
+            if "harness_evolve.orchestrator.cli" in cmdline and marker in cmdline:
+                out.append(pid)
+        return out
+
+    def terminate_stale_run(self, pids: Optional[Sequence[int]] = None
+                            ) -> List[int]:
+        """Kill a stalled run's process groups; return the pids that were signalled.
+
+        Their children are killed with them (each orchestrator starts its own
+        session), so the GPU leases they hold are reaped by the broker instead
+        of blocking the round that replaces them.
+        """
+        if pids is None:
+            pids = self.stale_run_pids()
+        pids = [int(pid) for pid in pids if int(pid) != os.getpid()]
+        for pid in pids:
+            try:
+                os.killpg(os.getpgid(pid), signal.SIGTERM)
+            except OSError:
+                try:
+                    os.kill(pid, signal.SIGTERM)
+                except OSError:
+                    continue
+        deadline = self.clock() + 10.0
+        while self.clock() < deadline:
+            if not any(pid_is_live(pid) for pid in pids):
+                return pids
+            self.sleeper(0.5)
+        for pid in pids:
+            if not pid_is_live(pid):
+                continue
+            try:
+                os.killpg(os.getpgid(pid), signal.SIGKILL)
+            except OSError:
+                try:
+                    os.kill(pid, signal.SIGKILL)
+                except OSError:
+                    continue
+        return pids
+
+    def reap_child(self) -> Optional[int]:
+        """Collect a finished child so it cannot linger as a zombie."""
+        proc = getattr(self, "_resume_proc", None)
+        if proc is None:
+            return None
+        code = proc.poll()
+        if code is None:
+            return None
+        self._resume_proc = None
+        self.record_incident(action="resume_exited", category=CATEGORY_UNKNOWN,
+                             pid=proc.pid, returncode=code)
+        return int(code)
+
+    def _log_tail(self, size: int = 4000) -> str:
+        parts: List[str] = []
+        for name in ("resume.log", "orchestrator.log"):
+            path = self.state_dir / name
+            if path.is_file():
+                try:
+                    parts.append(path.read_text(encoding="utf-8",
+                                                errors="replace")[-size:])
+                except OSError:
+                    continue
+        return "\n".join(parts)
+
+    # ------------------------------------------------------------- actions
+    @staticmethod
+    def project_root() -> Path:
+        """Directory that contains the ``metainfer`` package.
+
+        Getting this wrong makes the restarted process die instantly with
+        ``ModuleNotFoundError: No module named 'metainfer'`` — which is exactly
+        how a supervised run can look "stuck in analyze".
+        """
+        here = Path(__file__).resolve()
+        for parent in here.parents:
+            if (parent / "metainfer" / "tasks").is_dir():
+                return parent
+        return here.parents[4]
+
+    def _launch_resume(self) -> int:
+        if self._resume_launcher is not None:
+            return int(self._resume_launcher())
+        meta_root = self.project_root()
+        log = (self.state_dir / "resume.log").open("a", encoding="utf-8")
+        cmd = [
+            sys.executable, "-m",
+            "metainfer.tasks.harness_evolve.orchestrator.cli", "resume",
+            "--state-dir", str(self.state_dir),
+            "--workspace-dir", str(self.workspace_dir),
+        ]
+        env = dict(os.environ)
+        env.setdefault("PYTHONPATH", str(meta_root))
+        proc = subprocess.Popen(
+            cmd, cwd=str(meta_root), stdout=log, stderr=subprocess.STDOUT,
+            stdin=subprocess.DEVNULL, env=env, start_new_session=True)
+        log.close()
+        # Keep the handle so the child can be reaped: an unreaped child stays
+        # as a zombie, and a zombie still answers ``os.kill(pid, 0)``.
+        self._resume_proc = proc
+        (self.state_dir / "resume.pid").write_text(str(proc.pid),
+                                                   encoding="utf-8")
+        return int(proc.pid)
+
+    def write_incident_report(self, incident: Dict[str, Any]) -> Optional[Path]:
+        """Ask a read-only DSH agent for a root-cause report (advisory)."""
+        if not self.diagnose_agent:
+            return None
+        # Local import: this advisory path must not add an import edge from the
+        # supervisor to the evolve module at load time.
+        from .evolve import DEFAULT_EVOLVE_MODEL
+
+        wrapper = (Path(__file__).resolve().parents[2]
+                   / "dcu_kernel_auto_opt" / "bridge" / "dsh" / "dsh_agent.py")
+        if not wrapper.is_file():
+            return None
+        report = self.exp_dir / "incident_report.md"
+        prompt = (
+            "You are diagnosing a crashed kernel-optimization experiment. "
+            f"Read these logs if useful: {self.state_dir}/resume.log, "
+            f"{self.state_dir}/orchestrator.log, and the incident record at "
+            f"{self.incidents_path}. Latest incident: {json.dumps(incident)[:2000]}\n"
+            f"Write a concise markdown report to {report} with: (1) most likely "
+            "root cause with the log lines that support it, (2) whether a "
+            "restart would help, (3) one concrete prevention suggestion. "
+            "Do not run commands; analysis and the file write only."
+        )
+        try:
+            proc = subprocess.run(
+                [sys.executable, str(wrapper), "-p", "--output-format",
+                 "stream-json", "--input-format", "text",
+                 "--permission-mode", "bypassPermissions",
+                 "--add-dir", str(self.exp_dir), "--model",
+                 DEFAULT_EVOLVE_MODEL, "--effort", "low",
+                 "--tools", "Read,Glob,Grep,Write",
+                 "--disallowedTools", "Bash,Skill,WebFetch,WebSearch"],
+                input=prompt, text=True, capture_output=True, timeout=1200,
+                env=dict(os.environ))
+        except Exception:  # noqa: BLE001 - diagnostics are best effort
+            return None
+        return report if report.is_file() else None
+
+    # ---------------------------------------------------------------- loop
+    def step(self) -> Dict[str, Any]:
+        """One supervision tick. Returns the decision record."""
+        if self.stop_flag.exists():
+            self.write_state(status="stopped", reason="stop flag present")
+            return {"action": "stop", "reason": "stop flag present"}
+
+        self.reap_child()
+        run = self.run_state()
+        target = self.current_target()
+        iteration = int(run.get("current_iteration") or 0)
+        finished = bool(run.get("finished"))
+        alive = self.process_alive()
+        stalled = False
+        if alive and not finished:
+            age = self.progress_age()
+            if age is not None and age > float(self.stall_seconds):
+                live_children = self.live_child_count()
+                if live_children:
+                    # A DKAO child can spend an hour generating a kernel without
+                    # writing its own state. That is work, not a stall, and
+                    # calling it one used to start a *second* resume process
+                    # that then fought the first one for the same four devices.
+                    self.record_incident(action="stall_ignored",
+                                         category=CATEGORY_UNKNOWN,
+                                         idle_seconds=int(age),
+                                         live_children=live_children,
+                                         reason="children of this iteration "
+                                                "are still running")
+                else:
+                    stalled = True
+                    alive = False
+                    self.record_incident(action="stalled",
+                                         category=CATEGORY_UNKNOWN,
+                                         pid=self.resume_pid(),
+                                         idle_seconds=int(age),
+                                         reason="no state written for too long")
+                    # Never leave the stalled process behind: an orphan keeps
+                    # its GPU leases and starves the round that replaces it.
+                    killed = self.terminate_stale_run()
+                    if killed:
+                        self.record_incident(action="stalled_killed",
+                                             category=CATEGORY_UNKNOWN,
+                                             pids=killed)
+
+        if target is not None and iteration >= target and finished:
+            self.write_state(status="done", iteration=iteration,
+                             target=target)
+            return {"action": "done", "reason": f"target {target} reached",
+                    "iteration": iteration, "target": target}
+
+        same_streak = int(self.read_state().get("same_class_streak") or 0)
+        waits = int(self.read_state().get("waits_done") or 0)
+        restarts = int(self.read_state().get("restarts_done") or 0)
+
+        incident: Optional[Dict[str, Any]] = None
+        category = CATEGORY_UNKNOWN
+        if alive:
+            decision = {"action": "run", "sleep_s": 0,
+                        "reason": "experiment process is alive"}
+        else:
+            category, evidence = classify_incident(log_tail=self._log_tail())
+            last_category = str(self.read_state().get("last_category") or "")
+            same_streak = same_streak + 1 if category == last_category else 1
+            incident = {
+                "category": category, "evidence": evidence,
+                "iteration": iteration,
+                "finished": finished, "target": target,
+            }
+            decision = decide_action(
+                category=category, same_class_streak=same_streak,
+                waits_done=waits, restarts_done=restarts, running=False,
+                max_restarts=self.max_restarts, max_waits=self.max_waits,
+                wait_seconds=self.wait_seconds)
+
+        action = decision["action"]
+        if action == "done":
+            self.write_state(status="done")
+            return decision
+
+        if action == "run":
+            self.write_state(status="running", target=target,
+                             iteration=iteration, last_category=None,
+                             same_class_streak=0, waits_done=0,
+                             restarts_done=0)
+            return decision
+
+        if action == "stop":
+            self.write_state(status="stopped", reason=decision["reason"],
+                             last_category=category,
+                             same_class_streak=same_streak)
+            self.record_incident(action="stop", category=category,
+                                 reason=decision["reason"],
+                                 evidence=(incident or {}).get("evidence"))
+            if incident:
+                self.write_incident_report({**incident, **decision})
+            return decision
+
+        if action == "wait":
+            self.record_incident(action="wait", category=category,
+                                 seconds=decision["sleep_s"],
+                                 reason=decision["reason"],
+                                 evidence=(incident or {}).get("evidence"))
+            self.write_state(status="waiting", last_category=category,
+                             same_class_streak=same_streak,
+                             waits_done=waits + 1,
+                             waiting_until=self.clock() + decision["sleep_s"],
+                             reason=decision["reason"])
+            self.sleeper(decision["sleep_s"])
+            # The wait is over: restart the run. Whether a card is clean by then
+            # is DKAO's business (FLOW.md §0) — its gate waits on the child's
+            # behalf, so there is nothing for HE to re-check here.
+            pid = self._launch_resume()
+            self.record_incident(action="restart_after_wait",
+                                 category=category, pid=pid)
+            self.write_state(status="running", pid=pid,
+                             waits_done=waits + 1)
+            return {**decision, "action": "restart", "pid": pid}
+
+        # restart (backoff)
+        if decision.get("sleep_s"):
+            self.sleeper(decision["sleep_s"])
+        pid = self._launch_resume()
+        self.record_incident(action="restart", category=category, pid=pid,
+                             reason=decision["reason"],
+                             evidence=(incident or {}).get("evidence"))
+        self.write_state(status="running", pid=pid, last_category=category,
+                         same_class_streak=same_streak,
+                         restarts_done=restarts + 1)
+        return {**decision, "pid": pid}
+
+    def run_forever(self, max_ticks: Optional[int] = None) -> Dict[str, Any]:
+        self.write_state(status="starting", target=self.current_target(),
+                         pid=os.getpid())
+        ticks = 0
+        last: Dict[str, Any] = {}
+        while max_ticks is None or ticks < max_ticks:
+            last = self.step()
+            ticks += 1
+            if last.get("action") in {"done", "stop"}:
+                break
+            self.sleeper(self.interval_s)
+        return last
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    ap = argparse.ArgumentParser(prog="harness-evolve-supervisor")
+    sub = ap.add_subparsers(dest="command", required=True)
+    watch = sub.add_parser("watch", help="supervise one experiment")
+    watch.add_argument("--state-dir", type=Path, required=True)
+    watch.add_argument("--workspace-dir", type=Path, required=True)
+    watch.add_argument("--target", type=int, default=None)
+    watch.add_argument("--interval", type=int, default=30)
+    watch.add_argument("--wait-seconds", type=int, default=1800)
+    watch.add_argument("--max-restarts", type=int, default=3)
+    watch.add_argument("--max-waits", type=int, default=6)
+    watch.add_argument("--diagnose-agent", action="store_true")
+    watch.add_argument("--stall-seconds", type=int, default=2400)
+    args = ap.parse_args(argv)
+
+    if args.command == "watch":
+        sup = Supervisor(
+            args.state_dir, args.workspace_dir, target=args.target,
+            interval_s=args.interval, wait_seconds=args.wait_seconds,
+            max_restarts=args.max_restarts, max_waits=args.max_waits,
+            diagnose_agent=args.diagnose_agent,
+            stall_seconds=args.stall_seconds,
+        )
+        (args.workspace_dir).mkdir(parents=True, exist_ok=True)
+        (args.state_dir).mkdir(parents=True, exist_ok=True)
+        pid_file = args.workspace_dir / "supervisor.pid"
+        pid_file.write_text(str(os.getpid()), encoding="utf-8")
+        try:
+            final = sup.run_forever()
+        finally:
+            try:
+                pid_file.unlink()
+            except OSError:
+                pass
+        print(json.dumps(final, ensure_ascii=False))
+        return 0
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/metainfer/tasks/harness_evolve/orchestrator/variant.py b/metainfer/tasks/harness_evolve/orchestrator/variant.py
new file mode 100644
index 00000000..db41a7d7
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/orchestrator/variant.py
@@ -0,0 +1,324 @@
+"""Per-operator variant table: the thing a candidate harness is measured against.
+
+Terminology (fixed by the operator):
+
+* **baseline** — the fixed Triton number for an operator (``baseline_us`` in the
+  question pool). It never changes.
+* **variant** — the *current best* result for an operator, plus the kernel that
+  produced it. It is what every candidate is compared with.
+* **harness** — the search policy (planner/gates). It does not measure or
+  validate anything; it only steers which experiments DKAO tries, so "the new
+  harness is better" means "it searches better", which is decided by comparing
+  candidate results against the variants.
+
+Update rule (operator decision, option **b**): the variant table only moves when
+a harness is accepted and promoted into DKAO. A round that measures something
+faster but is then rejected leaves the variant untouched — the comparison basis
+must stay "what production currently is", otherwise the baseline would drift
+while production did not.
+
+The table is seeded from what already exists (the registered question pool's
+``best_known_us``, and the experiment's own ``best_known.json``); operators
+without a recorded variant are simply compared against their fixed baseline
+until they are measured for real.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional
+
+import yaml
+
+VARIANT_TABLE = "variant_table.json"
+VARIANT_HISTORY = "variant_history.jsonl"
+
+#: Fallback chain used when an operator has no recorded variant yet.
+SOURCE_VARIANT = "variant"
+SOURCE_BEST_KNOWN = "best_known"
+SOURCE_BASELINE = "baseline"
+SOURCE_NONE = "none"
+
+
+def _load_json(path: Path) -> Any:
+    try:
+        return json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+
+
+def _num(value: Any) -> Optional[float]:
+    try:
+        number = float(value)
+    except (TypeError, ValueError):
+        return None
+    return number if number > 0 else None
+
+
+def pool_instances(pool_path: Path) -> Dict[str, Dict[str, Any]]:
+    """``{operator_id: {baseline_us, best_known_us, family, contract}}``."""
+    data = _load_json(Path(pool_path))
+    if not isinstance(data, dict):
+        try:
+            data = yaml.safe_load(Path(pool_path).read_text(encoding="utf-8"))
+        except (OSError, ValueError, yaml.YAMLError):
+            return {}
+    out: Dict[str, Dict[str, Any]] = {}
+    for entry in (data or {}).get("instances") or []:
+        if not isinstance(entry, dict):
+            continue
+        iid = str(entry.get("id") or "")
+        if not iid:
+            continue
+        out[iid] = {
+            "baseline_us": _num(entry.get("baseline_us")),
+            "best_known_us": _num(entry.get("best_known_us")),
+            "family": str(entry.get("family") or ""),
+            "contract": dict(entry.get("contract") or {}),
+            "history": list(entry.get("history") or []),
+        }
+    return out
+
+
+def seed_from_pool(table: Dict[str, Any],
+                   instances: Dict[str, Dict[str, Any]]) -> int:
+    """Fill in variants for operators the pool already knows the best for.
+
+    Only operators that have a measured ``best_known_us`` (an existing kernel
+    that beat the baseline) become variants; the rest stay on their baseline
+    until a round measures them. Nothing is re-measured here.
+    """
+    operators = table.setdefault("operators", {})
+    seeded = 0
+    for iid, entry in instances.items():
+        best = entry.get("best_known_us")
+        if best is None or iid in operators:
+            continue
+        operators[iid] = {
+            "median_us": best,
+            "p90_us": None,
+            "correctness_ok": True,
+            "kernel_source": "pool:best_known",
+            "harness_version": table.get("harness_version"),
+            "source_iteration": None,
+            "family": entry.get("family") or "",
+            "contract": entry.get("contract") or {},
+            "conditions": {"origin": "registered_pool"},
+            "recorded_at": None,
+        }
+        seeded += 1
+    return seeded
+
+
+def seed_from_experiment(table: Dict[str, Any], exp_dir: Path) -> int:
+    """Fill in variants from the experiment's own ``best_known.json``."""
+    data = _load_json(Path(exp_dir) / "best_known.json")
+    if not isinstance(data, dict):
+        return 0
+    operators = table.setdefault("operators", {})
+    seeded = 0
+    for iid, entry in data.items():
+        if not isinstance(entry, dict) or iid in operators:
+            continue
+        best = _num(entry.get("best_median_us"))
+        if best is None:
+            continue
+        operators[str(iid)] = {
+            "median_us": best,
+            "p90_us": None,
+            "correctness_ok": True,
+            "kernel_source": f"experiment:iteration_{entry.get('best_iteration')}",
+            "harness_version": table.get("harness_version"),
+            "source_iteration": entry.get("best_iteration"),
+            "family": "",
+            "contract": {},
+            "conditions": {"origin": "experiment_best_known"},
+            "recorded_at": None,
+        }
+        seeded += 1
+    return seeded
+
+
+def load_variant_table(exp_dir: Path, *, pool_path: Optional[Path] = None,
+                       create: bool = True) -> Dict[str, Any]:
+    """Read the table, seeding it from existing data on first use."""
+    exp_dir = Path(exp_dir)
+    path = exp_dir / VARIANT_TABLE
+    table = _load_json(path)
+    if not isinstance(table, dict):
+        table = {
+            "schema_version": 1,
+            "created_at": time.time(),
+            "harness_version": None,
+            "harness_dir": None,
+            "operators": {},
+            "groups": {},
+            "seed_sources": [],
+        }
+    changed = False
+    if pool_path is not None and Path(pool_path).is_file():
+        seeded = seed_from_pool(table, pool_instances(Path(pool_path)))
+        if seeded:
+            table.setdefault("seed_sources", []).append(
+                {"source": str(pool_path), "seeded": seeded, "ts": time.time()})
+            changed = True
+    seeded_exp = seed_from_experiment(table, exp_dir)
+    if seeded_exp:
+        table.setdefault("seed_sources", []).append(
+            {"source": "best_known.json", "seeded": seeded_exp, "ts": time.time()})
+        changed = True
+    if changed or not path.is_file():
+        if create:
+            save_variant_table(exp_dir, table)
+    return table
+
+
+def save_variant_table(exp_dir: Path, table: Dict[str, Any]) -> Path:
+    exp_dir = Path(exp_dir)
+    exp_dir.mkdir(parents=True, exist_ok=True)
+    table["updated_at"] = time.time()
+    path = exp_dir / VARIANT_TABLE
+    tmp = path.with_suffix(".json.tmp")
+    tmp.write_text(json.dumps(table, ensure_ascii=False, indent=2, sort_keys=True),
+                   encoding="utf-8")
+    tmp.replace(path)
+    return path
+
+
+def resolve_variant(table: Dict[str, Any], operator_id: str, *,
+                    instances: Optional[Dict[str, Dict[str, Any]]] = None,
+                    ) -> Optional[Dict[str, Any]]:
+    """The number this operator must beat, and where it came from.
+
+    Order: recorded variant -> pool ``best_known_us`` -> fixed Triton
+    ``baseline_us``. Returns ``None`` when nothing at all is known, in which
+    case the operator cannot be judged this round.
+    """
+    entry = (table.get("operators") or {}).get(operator_id)
+    if isinstance(entry, dict):
+        median = _num(entry.get("median_us"))
+        if median is not None:
+            return {"operator_id": operator_id, "median_us": median,
+                    "source": SOURCE_VARIANT,
+                    "harness_version": entry.get("harness_version"),
+                    "p90_us": entry.get("p90_us"),
+                    "kernel_source": entry.get("kernel_source"),
+                    "recorded_at": entry.get("recorded_at")}
+    pool_entry = (instances or {}).get(operator_id) or {}
+    best = _num(pool_entry.get("best_known_us"))
+    if best is not None:
+        return {"operator_id": operator_id, "median_us": best,
+                "source": SOURCE_BEST_KNOWN, "harness_version": None,
+                "p90_us": None, "kernel_source": "pool:best_known",
+                "recorded_at": None}
+    baseline = _num(pool_entry.get("baseline_us"))
+    if baseline is not None:
+        return {"operator_id": operator_id, "median_us": baseline,
+                "source": SOURCE_BASELINE, "harness_version": None,
+                "p90_us": None, "kernel_source": "triton_baseline",
+                "recorded_at": None}
+    return None
+
+
+def reference_results(table: Dict[str, Any], operator_ids: Iterable[str], *,
+                      instances: Optional[Dict[str, Dict[str, Any]]] = None,
+                      ) -> Dict[str, Dict[str, Any]]:
+    """A ``champion``-shaped result dict for :func:`compare_instances`.
+
+    Each operator is represented by the variant it must beat, so the existing
+    comparison/reporting code works unchanged.
+    """
+    out: Dict[str, Dict[str, Any]] = {}
+    for iid in operator_ids:
+        resolved = resolve_variant(table, iid, instances=instances)
+        if resolved is None:
+            continue
+        out[iid] = {
+            "status": "reference",
+            "passed": True,
+            "correctness_ok": True,
+            "median_us": resolved["median_us"],
+            "p90_us": resolved.get("p90_us"),
+            "variant_source": resolved["source"],
+            "variant_harness_version": resolved.get("harness_version"),
+            "kernel_source": resolved.get("kernel_source"),
+        }
+    return out
+
+
+def unjudgeable(table: Dict[str, Any], operator_ids: Iterable[str], *,
+                instances: Optional[Dict[str, Dict[str, Any]]] = None,
+                ) -> List[str]:
+    """Operators with no variant and no baseline: they cannot be scored."""
+    return [iid for iid in operator_ids
+            if resolve_variant(table, iid, instances=instances) is None]
+
+
+def variant_beats(resolved: Dict[str, Any],
+                  metrics: Dict[str, Any]) -> bool:
+    """Whether a measurement is faster than the variant it was compared with."""
+    candidate = _num((metrics or {}).get("median_us"))
+    if candidate is None:
+        return False
+    return candidate < _num(resolved.get("median_us")) or False
+
+
+def record_variants(exp_dir: Path, table: Dict[str, Any], results: Dict[str, Any],
+                    *, version: str, iteration: int,
+                    groups: Optional[Dict[str, Any]] = None,
+                    ) -> Dict[str, Any]:
+    """Update the table after a harness is accepted and promoted (rule b).
+
+    Only measured, correct results are recorded. Every change is appended to
+    ``variant_history.jsonl`` so it is always clear which harness set which
+    number.
+    """
+    exp_dir = Path(exp_dir)
+    operators = table.setdefault("operators", {})
+    recorded: List[Dict[str, Any]] = []
+    history_path = exp_dir / VARIANT_HISTORY
+    now = time.time()
+    for iid, row in (results or {}).items():
+        if not isinstance(row, dict):
+            continue
+        median = _num(row.get("median_us"))
+        if median is None or row.get("correctness_ok") is False:
+            continue
+        previous = operators.get(iid) or {}
+        previous_median = _num(previous.get("median_us"))
+        if previous_median is not None and median >= previous_median:
+            continue          # only a real improvement moves the variant
+        operators[iid] = {
+            "median_us": median,
+            "p90_us": _num(row.get("p90_us")),
+            "correctness_ok": True,
+            "kernel_source": row.get("repo_path") or row.get("child_task_id"),
+            "harness_version": version,
+            "source_iteration": iteration,
+            "family": row.get("family") or previous.get("family") or "",
+            "contract": previous.get("contract") or {},
+            "conditions": {
+                "bench_profile": row.get("bench_profile"),
+                "gate": "vram<=90,hcu==0",
+            },
+            "recorded_at": now,
+        }
+        change = {"ts": now, "operator_id": iid, "median_us": median,
+                  "previous_median_us": previous_median,
+                  "harness_version": version, "iteration": iteration}
+        recorded.append(change)
+        try:
+            history_path.parent.mkdir(parents=True, exist_ok=True)
+            with history_path.open("a", encoding="utf-8") as fh:
+                fh.write(json.dumps(change, ensure_ascii=False) + "\n")
+        except OSError:
+            pass
+    table["harness_version"] = version
+    if groups:
+        table["groups"] = dict(groups)
+    table["last_promoted_at"] = now
+    save_variant_table(exp_dir, table)
+    return {"updated": sorted(c["operator_id"] for c in recorded),
+            "changes": recorded}
diff --git a/metainfer/tasks/harness_evolve/server/__init__.py b/metainfer/tasks/harness_evolve/server/__init__.py
new file mode 100644
index 00000000..3352dc2e
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/server/__init__.py
@@ -0,0 +1 @@
+"""harness_evolve server package (placeholder; headless usage needs no WebUI)."""
diff --git a/metainfer/tasks/harness_evolve/server/plugin.py b/metainfer/tasks/harness_evolve/server/plugin.py
new file mode 100644
index 00000000..41de1f04
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/server/plugin.py
@@ -0,0 +1,34 @@
+"""harness_evolve Web plugin registration."""
+
+from pathlib import Path
+
+from metainfer.server.registry import WebPlugin, register
+
+from .routes import build_router
+
+PLUGIN_TYPE = "harness-evolve"
+_FRONTEND_DIR = Path(__file__).resolve().parent.parent / "static"
+
+plugin = WebPlugin(
+    type=PLUGIN_TYPE,
+    label="Harness Evolution (AHE)",
+    description=(
+        "AHE outer loop: evolve the DKAO harness across task instances with "
+        "double-generation iterations, change attribution and auto rollback."
+    ),
+    build_router=build_router,
+    detail_view_module="app/he-detail",
+    frontend_dir=_FRONTEND_DIR,
+    extra_stylesheets=["he.css"],
+    # The AHE page embeds the DKAO task view to mirror one question of one
+    # iteration (see LiveDkaoMirror in he-detail.js). Point at the DKAO
+    # plugin's own bundle rather than shipping a second copy, so the embedded
+    # view and the standalone DKAO page cannot drift apart. Listed explicitly
+    # because auto-discovery only yields ``app/`` keys.
+    importmap_entries={
+        "app/dkao-detail":
+            "/static/plugins/dcu-kernel-auto-opt/dkao-detail.js?v=CACHE_BUST",
+    },
+)
+
+register(plugin)
diff --git a/metainfer/tasks/harness_evolve/server/routes.py b/metainfer/tasks/harness_evolve/server/routes.py
new file mode 100644
index 00000000..cefca4a1
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/server/routes.py
@@ -0,0 +1,1553 @@
+"""Read-only Web routes for harness_evolve.
+
+Serves the AHE experiment artifacts (workspace_dir of a task entry):
+- /summary            report.md text + config snapshot + best_ever + scores
+- /iterations         per-iteration benchmark/diff/change_evaluation/evolve/rollback
+- /iterations/{n}/overview   analysis overview markdown (plain text)
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Mapping
+
+from fastapi import APIRouter, HTTPException, Request
+
+from metainfer.server._helpers import (
+    require_task_type,
+    state_dir_for,
+    task_or_404,
+    workspace_dir_for,
+)
+from metainfer.server.state_reader import read_requirements, read_run
+
+from ..orchestrator import phases
+from ..orchestrator.phases import PREPARE
+from ..orchestrator.rounds import (
+    CHILDREN_DIR, GPU_COUNT, child_dir_map, device_for_index,
+    iteration_of_dir_name, resolve_child_dir,
+)
+from ..orchestrator.state import (
+    read_target_iterations, write_target_iterations,
+)
+
+PLUGIN_TYPE = "harness-evolve"
+
+
+def _load(path: Path, default: Any) -> Any:
+    try:
+        return json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return default
+
+
+def _experiment_config(entry: Any, workspace_dir: Path):
+    """Load the experiment config behind one task (for promotion commands)."""
+    from ..orchestrator.config import load_experiment_config
+
+    state_dir = Path(str(getattr(entry, "state_dir", "") or ""))
+    req = state_dir / "requirements.json"
+    if not req.is_file():
+        raise HTTPException(404, f"requirements.json missing in {state_dir}")
+    return load_experiment_config(req, state_dir, workspace_dir)
+
+
+def _approve_with_config(entry: Any, workspace_dir: Path, by: str):
+    from ..orchestrator.promotion import approve_promotion
+
+    cfg = _experiment_config(entry, workspace_dir)
+    return approve_promotion(cfg, cfg.exp_dir, approved_by=by)
+
+
+def _deny_with_config(entry: Any, workspace_dir: Path, reason: str):
+    from ..orchestrator.promotion import deny_promotion
+
+    cfg = _experiment_config(entry, workspace_dir)
+    return deny_promotion(cfg, cfg.exp_dir, reason=reason)
+
+
+def _text(path: Path) -> str | None:
+    if not path.is_file():
+        return None
+    try:
+        return path.read_text(encoding="utf-8", errors="replace")
+    except OSError:
+        return None
+
+
+def _jsonl(path: Path) -> List[Dict[str, Any]]:
+    if not path.is_file():
+        return []
+    rows = []
+    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+        try:
+            item = json.loads(line)
+        except ValueError:
+            continue
+        if isinstance(item, dict):
+            rows.append(item)
+    return rows
+
+
+def _iterations(exp: Path) -> List[Dict[str, Any]]:
+    runs = exp / "runs"
+    if not runs.is_dir():
+        return []
+    out: List[Dict[str, Any]] = []
+    for d in sorted(runs.glob("iteration_*")):
+        if not d.is_dir():
+            continue
+        i_dir = d / "input"
+        e_dir = d / "evolve"
+        entry: Dict[str, Any] = {
+            "name": d.name,
+            # Which gate this round judged, and what it concluded. The page must
+            # show the round's real criterion (FLOW.md §3: the variant
+            # comparison), not the pool's tau_family * baseline annotation.
+            "stage": _load(i_dir / "stage.json", None),
+            "performance_gate": _load(i_dir / "performance_gate.json", None),
+            "generalization_gate": _load(
+                i_dir / "generalization_gate.json", None),
+            "baseline_round": _load(i_dir / "baseline_round.json", None),
+            "benchmark": _load(i_dir / "benchmark" / "results.json", None),
+            "scored_ids": _load(i_dir / "benchmark" / "scored_ids.json", None),
+            "diff": _load(i_dir / "diff.json", None),
+            "decision": _load(i_dir / "decision.json", None),
+            "change_evaluation": _load(
+                i_dir / "change_evaluation.json", None
+            ),
+            "rollback": _load(d / "rollback.json", None),
+            "overview": (
+                _text(i_dir / "analysis" / "overview.md") or ""
+            )[:8000],
+            "evolve_manifest": _load(
+                e_dir / "change_manifest.json", None
+            ),
+            "round_plan": _load(e_dir / "round_plan.json", None),
+            "evolve_summary": _text(e_dir / "evolve_summary.md") or "",
+            "promotion": _load(d / "promotion.json", None),
+            "mechanism_evidence": _load(
+                i_dir / "mechanism_evidence.json", None),
+            "purposes": (_load(i_dir / "benchmark" / "purposes.json", {})
+                         or {}).get("purposes", {}),
+        }
+        out.append(entry)
+    return out
+
+
+def _child_best_median(workspace_dir: Path) -> Dict[str, Any]:
+    """Fast per-child perf peek: best accepted median + completed rounds."""
+    best: Optional[float] = None
+    rounds = 0
+    runs_root = workspace_dir / "workers"
+    if runs_root.is_dir():
+        for w in runs_root.glob("worker_*"):
+            for exp_file in (w / "runs").glob("*/experiments.jsonl"):
+                for row in _jsonl(exp_file):
+                    rounds += 1
+                    m = row.get("metrics") or {}
+                    median = m.get("median_us")
+                    if row.get("accepted") and isinstance(median, (int, float)):
+                        best = median if best is None else min(best, float(median))
+    return {"best_median_us": best, "rounds": rounds}
+
+
+def _child_gate_audit(child_state_dir: Path, limit: int = 200) -> Dict[str, Any]:
+    """One child's own DKAO admission-gate audit trail.
+
+    ``harness_evolve`` does not manage GPU occupancy: it hands each question to
+    a DKAO child and the child's engine decides when a card is clean enough to
+    time a kernel. So the reason a question is "doing nothing" lives in that
+    child's ``measurement_gate.jsonl`` — ``gate_blocked`` rows carry the reason
+    (``VRAM 96% > 90% limit`` / ``device busy: HCU 100.0% > 0.0%``), the attempt
+    number and the wait interval. This reads it read-only and never judges.
+
+    One subtlety, because it decides what the UI may claim: a *passing* row also
+    carries a ``reasons`` list, and those reasons belong to the *other*
+    candidates that failed (``gpu`` on such a row is the device that passed
+    while ``reasons`` describes the one that did not). So the reasons are only
+    surfaced when the last verdict actually blocked, and the device this child
+    is using is taken from the passing row, not from the reasons.
+    """
+    rows = _jsonl(child_state_dir / "measurement_gate.jsonl")[-limit:]
+    # A gate call writes one row per channel, so identical rows double up.
+    seen: set = set()
+    unique: List[Dict[str, Any]] = []
+    for row in rows:
+        key = (round(float(row.get("ts") or 0), 1), row.get("event"),
+               row.get("attempt"), row.get("site"))
+        if key in seen:
+            continue
+        seen.add(key)
+        unique.append(row)
+    blocked = [r for r in unique if r.get("event") == "gate_blocked"]
+    passed = [r for r in unique if r.get("event") == "gate_pass"
+              or r.get("usable") is True]
+    last = unique[-1] if unique else {}
+    last_blocked = bool(last) and last.get("event") == "gate_blocked"
+
+    def _device_of(row: Dict[str, Any]) -> Optional[int]:
+        try:
+            return int(row["gpu"]) if row.get("gpu") is not None else None
+        except (TypeError, ValueError):
+            return None
+
+    # Which card this child is *measuring on* — not which card some gate call
+    # happened to run on. The generate-stage probe may legitimately run on any
+    # passing device ("the task's own card first, otherwise any idle one"), so a
+    # probe row on GPU2 says nothing about a child pinned to GPU0. Rows from the
+    # measuring sites (benchmark / PMC) do name the child's own card, so they are
+    # preferred; only when none exists does a probe (or a block) stand in.
+    own_card = None
+    for row in reversed(passed):
+        if str(row.get("site") or "") in ("benchmark", "profile_pmc"):
+            own_card = _device_of(row)
+            if own_card is not None:
+                break
+    if own_card is None:
+        for row in reversed(passed):
+            own_card = _device_of(row)
+            if own_card is not None:
+                break
+    device = own_card
+    if device is None:
+        for row in reversed(blocked):
+            device = _device_of(row)
+            if device is not None:
+                break
+    probe_device = None
+    for row in reversed(passed):
+        if str(row.get("site") or "") == "generate_probe":
+            probe_device = _device_of(row)
+            if probe_device is not None:
+                break
+    return {
+        "checks": len(unique),
+        "blocked": len(blocked),
+        "waited_seconds": sum(
+            float(r.get("wait_seconds") or 0) for r in blocked),
+        "device": device,
+        # Every device this child's gate ever looked at: the card it measures on
+        # plus any card a probe ran on. Kept apart from ``device`` so the page
+        # can say "pinned to GPU0, probe passed on GPU2" instead of claiming the
+        # child moved.
+        "devices": sorted({d for d in (_device_of(r) for r in unique)
+                           if d is not None}),
+        "probe_device": probe_device,
+        "device_from": ("gate_pass" if device is not None and passed
+                        else ("gate_blocked" if device is not None else None)),
+        "last_event": last.get("event"),
+        "last_ts": last.get("ts"),
+        "last_site": last.get("site"),
+        "last_gpu": last.get("gpu"),
+        "last_usable": last.get("usable"),
+        "last_candidates": list(last.get("candidates") or []),
+        "last_passing": list(last.get("passing") or []),
+        # Only meaningful (and only about *this* device) when the verdict blocked.
+        "last_reasons": (list(last.get("reasons") or []) if last_blocked
+                         else []),
+        "last_attempt": last.get("attempt"),
+        "max_waits": last.get("max_waits"),
+        "wait_seconds": last.get("wait_seconds"),
+        "blocked_reasons": [
+            {"ts": r.get("ts"), "site": r.get("site"), "gpu": r.get("gpu"),
+             "attempt": r.get("attempt"), "reasons": list(r.get("reasons") or [])}
+            for r in blocked[-10:]
+        ],
+    }
+
+
+def _layout_devices(exp: Path, num: int) -> Dict[str, int]:
+    """``question -> device`` exactly as HE recorded it for one round.
+
+    ``gpu_preflight.json`` is a layout record, not an occupancy check: it says
+    which card each question was handed to. That — not the card some later probe
+    used — is what a per-card view must group by.
+    """
+    data = _load(exp / "runs" / f"iteration_{num:03d}" / "input" / "benchmark"
+                 / "gpu_preflight.json", None)
+    raw = data.get("devices") if isinstance(data, dict) else None
+    out: Dict[str, int] = {}
+    if isinstance(raw, dict):
+        for qid, gpu in raw.items():
+            try:
+                out[str(qid)] = int(gpu)
+            except (TypeError, ValueError):
+                continue
+    return out
+
+
+def _assigned_device(requirements: Optional[Dict[str, Any]]) -> Optional[int]:
+    """The card HE pinned for one child, from its ``shape_config``.
+
+    The HE evaluator pins each question with a manual DKAO assignment, so the
+    child's own requirements are a reliable second source when the round's layout
+    record is missing (older rounds, hand-written requirements).
+    """
+    if not isinstance(requirements, dict):
+        return None
+    raw = requirements.get("shape_config")
+    cfg: Any = raw
+    if not isinstance(raw, dict):
+        try:
+            import yaml
+
+            cfg = yaml.safe_load(str(raw or "")) or {}
+        except Exception:  # noqa: BLE001 - a malformed config is simply no answer
+            return None
+    assignments = cfg.get("assignments") if isinstance(cfg, dict) else None
+    if isinstance(assignments, dict):
+        for worker in assignments.values():
+            if isinstance(worker, dict) and worker.get("gpu") is not None:
+                try:
+                    return int(worker["gpu"])
+                except (TypeError, ValueError):
+                    continue
+    return None
+
+
+def _gpu_cards(exp: Path, num: int,
+               answers: Optional[Dict[str, Any]] = None,
+               ) -> List[Dict[str, Any]]:
+    """Devices of one AHE iteration with the questions DKAO is measuring on them.
+
+    One card per physical device, showing how far each child's inner DKAO task
+    is and what it has measured. This replaces the old lease view read-only: HE
+    holds no lease and reads no device state, so the card carries the child's
+    own evidence — the device the child's gate cleared it for, and why it had to
+    wait when it did.
+
+    The device comes from that child's gate records. Only when a child has no
+    gate record at all (it never reached a measurement) does the fixed
+    ``index % GPU_COUNT`` rotation stand in, and then it is flagged as a
+    prediction rather than reported as an observation.
+    """
+    children = _children_live(exp, num, answers)
+    own = child_dir_map(exp / CHILDREN_DIR, num)
+    layout = _layout_devices(exp, num)
+    cards: Dict[int, Dict[str, Any]] = {}
+    for idx, (qid, info) in enumerate(sorted(children.items())):
+        state_dir = Path(str(info.get("child_state_dir") or (base / qid / "state")))
+        gate = _child_gate_audit(state_dir)
+        requirements = _load((own.get(qid) or Path()) / "requirements.json", None)
+        assigned = _assigned_device(requirements)
+        # Which card this question belongs to. HE's own layout for the round is
+        # the answer ("index % 4", one question per card); a child that ran
+        # before the layout was recorded falls back to the device HE pinned in
+        # its requirements. The child's gate device is only a last resort: it can
+        # name a *probe* card, which would put a child pinned to GPU0 under GPU2
+        # and make a four-card view show two.
+        if qid in layout:
+            gpu, source = layout[qid], "layout"
+        elif assigned is not None:
+            gpu, source = assigned, "assignment"
+        elif gate.get("device") is not None:
+            gpu, source = int(gate["device"]), (gate.get("device_from") or "gate")
+        else:
+            gpu, source = device_for_index(idx), "predicted"
+        info.update({
+            "question": qid,
+            "gpu": gpu,
+            "gpu_source": source,
+            "assigned_device": assigned,
+            "gate_device": gate.get("device"),
+            "gate_audit": gate,
+            "requirements": requirements,
+        })
+        cards.setdefault(gpu, {"gpu": gpu, "questions": []})
+        cards[gpu]["questions"].append(info)
+    return [cards[g] for g in sorted(cards)]
+
+
+def _gate_map(answers: Dict[str, Any], ids: List[str]) -> Dict[str, Dict[str, Any]]:
+    """The *pool* criterion per instance: baseline, family tau, target median.
+
+    This is ``pool.auto_pass`` (median <= tau_family * baseline). It is NOT the
+    variant-round gate: since 2026-09-15 the performance/generalization gates
+    compare against each operator's current *variant* (FLOW.md §3, served by
+    :func:`_real_gate_map`), and a variant round never reaches ``auto_pass``
+    (the variant branch ``continue``s before it). Keep this for the legacy / suite
+    path, which really was judged this way, and expose it to the page as
+    ``pool_reference`` so it can never be mistaken for the round's gate.
+    """
+    from ..orchestrator.pool import Pool
+
+    src = str(answers.get("pool_source") or "").strip()
+    if not src:
+        return {}
+    try:
+        pool = Pool.from_yaml(Path(src).expanduser())
+    except Exception:  # noqa: BLE001 - gates are advisory for the UI
+        return {}
+    out: Dict[str, Dict[str, Any]] = {}
+    for iid in ids:
+        inst = pool.instances.get(iid)
+        if inst is None:
+            continue
+        tau = pool.tau_family(inst.family or "")
+        out[iid] = {
+            "baseline_us": inst.baseline_us,
+            "tau_family": tau,
+            "target_us": round(tau * float(inst.baseline_us), 3),
+            "best_known_us": inst.best_known_us,
+            "family": inst.family,
+        }
+    return out
+
+
+#: The gate's own knobs, mirrored from ``orchestrator.decision_engine`` so the
+#: page can describe the criterion the orchestrator really applies (FLOW.md §3)
+#: instead of re-deriving it. One place changes both.
+GATE_WIN_RATIO = {"performance": 0.75, "generalization": 0.50}
+GATE_FLOOR_PERCENT = 0.80
+GATE_NOISE_PERCENT = 2.0
+
+VARIANT_TABLE_FILE = "variant_table.json"
+
+
+def _round_input_dir(exp: Path, num: int) -> Path:
+    return exp / "runs" / f"iteration_{num:03d}" / "input"
+
+
+def _round_stage(exp: Path, num: int) -> Optional[str]:
+    """The stage this round declared (``baseline`` / ``performance`` / ...).
+
+    ``None`` means the round predates the staging (legacy / suite path, or the
+    short-lived transitional build that wrote a gate document without a
+    ``stage.json``).
+    """
+    doc = _load(_round_input_dir(exp, num) / "stage.json", None)
+    name = doc.get("stage") if isinstance(doc, dict) else None
+    return str(name) if name else None
+
+
+def _round_gate_doc(exp: Path, num: int,
+                    stage: Optional[str]) -> Optional[Dict[str, Any]]:
+    """The gate document this round judged, or ``None`` when it judged none.
+
+    The baseline round (FLOW.md §2) judges nothing, so it has no document: its
+    questions still *have* an opponent — the one the round is about to measure
+    and write back — which is what the page shows for it.
+    """
+    i_dir = _round_input_dir(exp, num)
+    if stage == "generalization":
+        return _load(i_dir / "generalization_gate.json", None)
+    if stage == "performance":
+        return _load(i_dir / "performance_gate.json", None)
+    if stage is None:                    # transitional build: document, no stage
+        for name in ("performance_gate.json", "generalization_gate.json"):
+            doc = _load(i_dir / name, None)
+            if isinstance(doc, dict):
+                return doc
+    return None
+
+
+def _pool_source(answers: Mapping[str, Any]) -> str:
+    return str(
+        answers.get("question_pool") or answers.get("pool_source") or ""
+    ).strip()
+
+
+def _variant_sources(exp: Path, answers: Mapping[str, Any]):
+    """``(variant_table, pool_instances)``, exactly as the orchestrator sees them."""
+    from ..orchestrator.variant import pool_instances
+
+    table = _load(exp / VARIANT_TABLE_FILE, None)
+    table = table if isinstance(table, dict) else {}
+    instances: Dict[str, Dict[str, Any]] = {}
+    src = _pool_source(answers)
+    if src:
+        path = Path(src).expanduser()
+        if path.is_file():
+            try:
+                instances = pool_instances(path)
+            except Exception:  # noqa: BLE001 - a bad pool is "no answer", not a 500
+                instances = {}
+    return table, instances
+
+
+def _real_gate_map(exp: Path, num: int, ids: List[str],
+                   answers: Mapping[str, Any],
+                   ) -> Dict[str, Any]:
+    """The criterion this round actually judges, per question.
+
+    FLOW.md §3: the performance gate compares each question against that
+    operator's *current variant* — a strict win for >= 75% of the questions,
+    with every question at least 80% of its variant (the floor). That is a
+    different number from ``tau_family * baseline_us``, so the two are served
+    separately and the page labels them by protocol:
+
+    * ``variant`` (a ``stage.json`` exists, a gate document exists, or the run
+      kept a variant table): the round is judged against the variant;
+    * ``legacy``: the pre-2026-09-15 / suite path, which really did use the pool
+      criterion — the page says so instead of dressing it up as a variant match.
+
+    ``target_us`` keeps its historic name for the round table, but for a variant
+    round it is the variant (``variant_us``), never a family baseline.
+    """
+    from ..orchestrator.variant import resolve_variant
+
+    stage = _round_stage(exp, num)
+    gate = _round_gate_doc(exp, num, stage)
+    table, instances = _variant_sources(exp, answers)
+    on_protocol = (stage is not None or isinstance(gate, dict)
+                   or (exp / VARIANT_TABLE_FILE).is_file())
+    judged = ((gate or {}).get("comparison") or {}).get("per_instance") or {}
+    win_ratio = GATE_WIN_RATIO.get(
+        str(stage or ""), GATE_WIN_RATIO["performance"])
+    targets: Dict[str, Dict[str, Any]] = {}
+    for qid in ids:
+        resolved = resolve_variant(table, qid, instances=instances) or {}
+        variant_us = resolved.get("median_us")
+        row = judged.get(qid) or {}
+        targets[qid] = {
+            "protocol": "variant" if on_protocol else "legacy",
+            "stage": stage,
+            "variant_us": variant_us,
+            "variant_source": resolved.get("source"),
+            "variant_p90_us": resolved.get("p90_us"),
+            "kernel_source": resolved.get("kernel_source"),
+            "target_us": variant_us,
+            "judged": ({
+                "verdict": row.get("verdict"),
+                "candidate_median_us": row.get("candidate_median_us"),
+                "delta_percent": row.get("delta_percent"),
+            } if row else None),
+            "win_ratio": win_ratio,
+            "wins_needed": None,
+            "floor_percent": GATE_FLOOR_PERCENT,
+            "noise_percent": GATE_NOISE_PERCENT,
+        }
+    round_gate: Optional[Dict[str, Any]] = None
+    if isinstance(gate, dict):
+        round_gate = {
+            "kind": gate.get("kind") or stage,
+            "status": gate.get("status"),
+            "wins": gate.get("wins"),
+            "wins_needed": gate.get("wins_needed"),
+            "operators": gate.get("operators"),
+            "counted": gate.get("counted"),
+            "below_floor": gate.get("below_floor"),
+            "unmeasured": gate.get("unmeasured"),
+            "reasons": gate.get("reasons"),
+        }
+        for qid in targets:
+            targets[qid]["wins_needed"] = gate.get("wins_needed")
+    return {
+        "protocol": "variant" if on_protocol else "legacy",
+        "stage": stage,
+        "round": round_gate,
+        "targets": targets,
+    }
+
+
+def _children_live(exp: Path, num: int,
+                   answers: Optional[Dict[str, Any]] = None,
+                   ) -> Dict[str, Dict[str, Any]]:
+    """Live per-question state from each real DKAO child task.
+
+    Every question is one DKAO child whose orchestrator writes its own
+    ``state/run.json``; expose the child's current_phase (Prepare -> ... ->
+    Finished of the *inner* DKAO loop) plus its perf peek and the gate it is
+    judged by so the HE detail page can show the Evaluate step drilling into
+    each child.
+
+    ``gate`` is the round's real criterion (the variant comparison). For rounds
+    that predate the variant protocol it degenerates to the pool criterion, with
+    ``protocol == "legacy"`` on it so the page can say which path produced the
+    verdict. ``pool_reference`` always carries the pool criterion as a clearly
+    named annotation — it is not a gate on the variant protocol.
+    """
+    base = exp / CHILDREN_DIR
+    dirs = child_dir_map(base, num)          # {question: newest attempt dir}
+    out: Dict[str, Dict[str, Any]] = {}
+    if not dirs:
+        return out
+    ids = sorted(dirs)
+    gates = _gate_map(answers or {}, ids)
+    real = _real_gate_map(exp, num, ids, answers or {})
+    for qid, d in sorted(dirs.items()):
+        run = read_run(d / "state")
+        perf = _child_best_median(d / "workspace")
+        gate = real["targets"].get(qid)
+        pool_ref = gates.get(qid)
+        if (gate is not None and gate.get("protocol") == "legacy"
+                and pool_ref is not None):
+            # A legacy round was judged by the family criterion, so that is the
+            # target its page must show — unchanged from how it was decided.
+            gate = dict(gate, target_us=pool_ref.get("target_us"))
+        out[qid] = {
+            "current_phase": run.get("current_phase") or "idle",
+            "current_iteration": run.get("current_iteration"),
+            "finished": bool(run.get("finished")),
+            "final_status": run.get("final_status"),
+            "last_update": run.get("last_update"),
+            "best_median_us": perf["best_median_us"],
+            "rounds": perf["rounds"],
+            "gate": gate,
+            "pool_reference": pool_ref,
+            "gate_protocol": real["protocol"],
+            "gate_stage": real["stage"],
+            "child_dir": str(d),
+            "attempt_dir": d.parent.name,
+            "child_state_dir": str(d / "state"),
+        }
+    return out
+
+
+def _tail_text(path: Path, size: int = 1200) -> str:
+    if not path.is_file():
+        return ""
+    try:
+        text = path.read_text(encoding="utf-8", errors="replace")
+    except OSError:
+        return ""
+    return text[-size:]
+
+
+def _dkao_child_payload(exp: Path, num: int, child: str) -> Dict[str, Any]:
+    """Aggregate one HE question's inner DKAO task into a DKAO-like payload.
+
+    Each question is a real dcu_kernel_auto_opt child that writes the same
+    artifacts a WebUI DKAO task does (state/run.json + timeline + agents.json,
+    workspace/workers/*). This mirrors what the DKAO detail page reads so the
+    HE page can render the child with the same look: 8-step inner state
+    machine + worker lanes + agents.
+    """
+    if not child or child in {".", ".."} or "/" in child or "\\" in child:
+        raise HTTPException(404, "unknown child")
+    # A question can own several attempt directories (a first run, its retries,
+    # the generalization retake): the drill-down always opens the newest one.
+    d = resolve_child_dir(exp / CHILDREN_DIR, num, child)
+    if d is None or not d.is_dir():
+        raise HTTPException(404, "child not found")
+
+    state_dir = d / "state"
+    run = read_run(state_dir)
+    run_phase = run.get("current_phase") or "prepare"
+    # DKAO phases machine (Prepare..Finished, 8 steps incl. Baseline).
+    from metainfer.tasks.dcu_kernel_auto_opt.orchestrator import phases as dkao
+    graph = dkao.graph_payload(run_phase, include_baseline=True)
+
+    agents_raw = [
+        item for item in ((_load(state_dir / "agents.json", None) or {})
+                          .get("agents") or [])
+        if isinstance(item, dict)
+    ]
+
+    workers: List[Dict[str, Any]] = []
+    wdir = d / "workspace" / "workers"
+    if wdir.is_dir():
+        for w in sorted(wdir.iterdir()):
+            if not w.is_dir():
+                continue
+            status = _load(w / "status.json", {}) or {}
+            bootstrap = _load(w / "bootstrap_result.json", {}) or {}
+            result = _load(w / "result.json", {}) or {}
+            shape_id = status.get("shape_id")
+            metrics: Dict[str, Any] = {}
+            if shape_id:
+                metrics = (
+                    ((result.get("shapes") or {}).get(shape_id) or {})
+                    .get("metrics") or {}
+                )
+            # Per-round optimization history (the same source DKAO lanes use).
+            experiments: List[Dict[str, Any]] = []
+            rounds = 0
+            runs_dir = w / "runs"
+            if runs_dir.is_dir():
+                for exp_file in sorted(runs_dir.glob("*/experiments.jsonl")):
+                    if not shape_id:
+                        shape_id = exp_file.parent.name
+                    for row in _jsonl(exp_file):
+                        m = row.get("metrics") or {}
+                        experiments.append({
+                            "iteration": row.get("iteration"),
+                            "median_us": m.get("median_us"),
+                            "p90_us": m.get("p90_us"),
+                            "accepted": bool(row.get("accepted")),
+                            "correctness_passed": bool(
+                                row.get("correctness_passed")
+                            ),
+                            "speedup": row.get("speedup"),
+                            "baseline_us": row.get("baseline_us"),
+                            "timestamp": row.get("timestamp"),
+                        })
+                rounds = len(experiments)
+                experiments.sort(key=lambda e: (
+                    float(e.get("timestamp") or 0),
+                    int(e.get("iteration") or 0),
+                ))
+            accepted_medians = [
+                e["median_us"] for e in experiments
+                if e.get("accepted") and e.get("median_us") is not None
+            ]
+            workers.append({
+                "worker_id": w.name,
+                "state": status.get("state") or "idle",
+                "iteration": status.get("iteration"),
+                "shape_id": shape_id,
+                "physical_gpu": status.get("physical_gpu"),
+                "bootstrap_status": bootstrap.get("status"),
+                "median_us": metrics.get("median_us"),
+                "p90_us": metrics.get("p90_us"),
+                "passed": metrics.get("passed"),
+                "rounds": rounds,
+                "best_median_us": (
+                    min(accepted_medians) if accepted_medians else None
+                ),
+                "experiments": experiments,
+                "last_update": status.get("last_update"),
+            })
+
+    # Agents: annotate worker turns with the matching per-round performance
+    # (agent name is "worker_N--iterK" -> experiments.jsonl iteration K)
+    import re as _re
+    perf_index: Dict[Any, Dict[str, Any]] = {}
+    for w in workers:
+        for e in w.get("experiments") or []:
+            perf_index[(w["worker_id"], e.get("iteration"))] = e
+    agents: List[Dict[str, Any]] = []
+    for item in agents_raw:
+        entry: Dict[str, Any] = {
+            "name": item.get("name"),
+            "role": item.get("role"),
+            "status": item.get("status"),
+            "phase": item.get("phase"),
+            "error": item.get("error"),
+        }
+        match = _re.match(r"^(worker_\d+)-.+-iter(\d+)$",
+                          str(item.get("name") or ""))
+        if match:
+            worker_id, iteration = match.group(1), int(match.group(2))
+            turn = perf_index.get((worker_id, iteration))
+            if turn:
+                entry["iteration"] = iteration
+                entry["median_us"] = turn.get("median_us")
+                entry["p90_us"] = turn.get("p90_us")
+                entry["speedup"] = turn.get("speedup")
+                entry["baseline_us"] = turn.get("baseline_us")
+                entry["accepted"] = turn.get("accepted")
+        agents.append(entry)
+
+    return {
+        "child": child,
+        "run": {
+            "current_phase": run_phase,
+            "current_iteration": run.get("current_iteration"),
+            "finished": bool(run.get("finished")),
+            "final_status": run.get("final_status"),
+        },
+        "graph": graph,
+        "agents": agents,
+        "workers": workers,
+        "repo_path": str(d / "workspace" / "main"),
+        "child_workspace_dir": str(d / "workspace"),
+        "log_tail": _tail_text(d / "orchestrator-external.log"),
+    }
+
+
+def _resume_running(state_dir: Path) -> Optional[int]:
+    """PID of a live harness_evolve run (zombies and recycled pids excluded)."""
+    try:
+        pid = int((Path(state_dir) / "resume.pid")
+                  .read_text(encoding="utf-8").split()[0])
+    except (OSError, ValueError, IndexError):
+        return None
+    try:
+        from ..orchestrator.supervisor import task_process_alive
+    except Exception:  # noqa: BLE001 - fall back to a plain existence check
+        try:
+            os.kill(pid, 0)
+            return pid
+        except OSError:
+            return None
+    return task_process_alive(state_dir)
+
+
+def _run_alive(state_dir: Path) -> Optional[int]:
+    """PID of a live orchestrator for this task, whichever path started it.
+
+    The sweep has to distinguish "nobody is running this experiment" from "this
+    experiment is running fine". ``_resume_running`` answers only the narrower
+    question (is there a *resume*?), and a task created from the New Task form
+    never writes ``resume.pid`` — so a one-second-old, perfectly healthy run
+    looked like an unsupervised corpse, and the sweep started a second driver
+    for it (which then opened the next round while round 1 was still being
+    evaluated). ``task_process_alive`` reads ``orchestrator.pid`` first and
+    falls back to a cmdline scan, which is exactly the question here.
+    """
+    try:
+        from ..orchestrator.supervisor import task_process_alive
+        return task_process_alive(Path(state_dir))
+    except Exception:  # noqa: BLE001 - never break the sweep
+        return _resume_running(state_dir)
+
+
+#: Terminal states that must never be auto-restarted: the loop stopped on
+#: purpose (a round could not measure every question, operator abort, circuit
+#: breaker) and needs a human decision. Restarting them would silently re-enter
+#: the same condition.
+#:
+#: ``gpu_gate_blocked`` is deliberately absent: since HE stopped managing GPU
+#: occupancy it holds no device verdict of its own and never stops a run for
+#: one. What it can do is stop because a *child's* DKAO gate refused to measure
+#: — that arrives as ``round_incomplete`` with ``cause`` ``gpu_gate_blocked``
+#: (see ``DkaoCliEvaluator._stop_on_incomplete_round``).
+TERMINAL_STATUSES = {"round_incomplete",
+                     "aborted_by_operator", "circuit_breaker",
+                     # an explicit Stop press: only a human asking again (or
+                     # the resume paths clearing it) may bring the run back
+                     "stopped_by_operator",
+                     # variant-round protocol: a run parks for a human decision
+                     # and ends once a harness is promoted (or the budget runs
+                     # out) — none of those should be silently restarted.
+                     "awaiting_approval", "rounds_exhausted"}
+
+
+def should_launch_resume(run: Mapping[str, Any], target: Any,
+                         resume_running: bool) -> bool:
+    """Whether an unattended start must (re)launch the experiment process.
+
+    Covers three cases: nothing started yet, the previous process died
+    mid-round (``finished`` is False), and the target round is not reached —
+    but never a run that deliberately stopped for a terminal reason.
+    """
+    if resume_running or not run:
+        return False
+    if str(run.get("final_status") or "") in TERMINAL_STATUSES:
+        return False
+    if not run.get("finished"):
+        return True
+    if target in (None, ""):
+        return False
+    try:
+        return int(run.get("current_iteration") or 0) < int(target)
+    except (TypeError, ValueError):
+        return False
+
+
+def _launch_resume(state_dir: Path, workspace_dir: Path,
+                   iterations: int,
+                   champion_iteration: Any = None) -> Dict[str, Any]:
+    """Spawn the resume CLI in the background and record its pid."""
+    import subprocess
+    import sys
+
+    # Spawning the experiment again is the operator saying "continue": without
+    # this, the marker left by Stop would make the supervisor and the janitor
+    # refuse to supervise the run we are starting right now.
+    _clear_operator_stop(Path(workspace_dir))
+    meta_root = Path(__file__).resolve().parents[4]
+    log_path = state_dir / "resume.log"
+    log = log_path.open("a", encoding="utf-8")
+    log.write(f"\n=== resume requested: +{iterations} round(s) "
+              f"@ {time.strftime('%Y-%m-%d %H:%M:%S')} ===\n")
+    log.flush()
+    cmd = [
+        sys.executable, "-m",
+        "metainfer.tasks.harness_evolve.orchestrator.cli", "resume",
+        "--state-dir", str(state_dir),
+        "--workspace-dir", str(workspace_dir),
+        "--iterations", str(iterations),
+    ]
+    if champion_iteration not in (None, ""):
+        cmd += ["--champion-iteration", str(int(champion_iteration))]
+    env = dict(os.environ)
+    env.setdefault("PYTHONPATH", str(meta_root))
+    proc = subprocess.Popen(
+        cmd, cwd=str(meta_root), stdout=log, stderr=subprocess.STDOUT,
+        stdin=subprocess.DEVNULL, env=env, start_new_session=True,
+    )
+    log.close()
+    (state_dir / "resume.pid").write_text(str(proc.pid), encoding="utf-8")
+    return {"ok": True, "pid": proc.pid, "iterations": iterations,
+            "log": str(log_path)}
+
+
+#: Presence of this file means "the operator stopped this run on purpose".
+#: It is the only thing that distinguishes a deliberate stop from a crash, and
+#: every path that would helpfully restart an unsupervised run must respect it
+#: (the janitor sweep, the supervisor, the resume launcher). Cleared by the
+#: resume paths, i.e. when a human asks for the run to continue again.
+OPERATOR_STOP_FILE = "operator_stop.json"
+
+
+def _operator_stop_path(workspace_dir: Path) -> Path:
+    return Path(workspace_dir) / OPERATOR_STOP_FILE
+
+
+def _run_live(state_dir: Path, workspace_dir: Path) -> bool:
+    """Whether this experiment is doing something right now.
+
+    Deliberately not just ``orchestrator.pid``: children outlive their parent's
+    bookkeeping (a round can be mid-flight while the parent is briefly between
+    steps), and a run started before that file existed has no entry at all.
+    """
+    if _resume_running(Path(state_dir)):
+        return True
+    if _supervisor_running(Path(workspace_dir)):
+        return True
+    return bool(_live_dkao_children(Path(workspace_dir)))
+
+
+def _operator_stopped(workspace_dir: Path) -> Optional[Dict[str, Any]]:
+    """The stop record if the operator stopped this run, else None."""
+    return _load(_operator_stop_path(workspace_dir), None) or None
+
+
+def _clear_operator_stop(workspace_dir: Path) -> bool:
+    """Drop the operator-stop marker so the run may be driven again."""
+    path = _operator_stop_path(workspace_dir)
+    try:
+        path.unlink()
+        return True
+    except OSError:
+        return False
+
+
+def _live_dkao_children(workspace_dir: Path) -> List[Dict[str, Any]]:
+    """Live DKAO child runs belonging to this HE workspace.
+
+    Children are spawned with ``start_new_session=True``, so signalling the HE
+    orchestrator does NOT reach them: they would keep their agents and their
+    place in DKAO's admission queue after a "stop". They are found by their
+    command line, which names this experiment's artifact directory -- the one
+    thing that identifies them without a registry of our own.
+    """
+    marker = f"{Path(workspace_dir)}/children/".encode()
+    found: List[Dict[str, Any]] = []
+    for entry in Path("/proc").iterdir():
+        if not entry.name.isdigit():
+            continue
+        pid = int(entry.name)
+        if pid == os.getpid():
+            continue
+        try:
+            cmdline = (entry / "cmdline").read_bytes().replace(b"\0", b" ")
+        except OSError:
+            continue
+        if marker not in cmdline \
+                or b"dcu_kernel_auto_opt.orchestrator.cli" not in cmdline:
+            continue
+        found.append({"pid": pid, "cmd": cmdline.decode("utf-8", "replace")})
+    return found
+
+
+def _proc_state(pid: int) -> Optional[str]:
+    """``/proc/`` state letter, or None when the process is gone."""
+    try:
+        stat = Path(f"/proc/{int(pid)}/stat").read_text(encoding="utf-8")
+    except OSError:
+        return None
+    try:
+        return stat.split(") ", 1)[1].split()[0]
+    except IndexError:
+        return None
+
+
+def _terminate_pid(pid: int, *, grace_s: float = 20.0) -> bool:
+    """SIGTERM a process group (then the process), and report whether it went.
+
+    ``os.killpg`` first because these processes lead their own session: the
+    DKAO orchestrator spawns agent processes that must not outlive it.
+
+    A zombie counts as stopped -- it executes no more, it just has not been
+    reaped yet (we are usually not its parent). ``os.kill(pid, 0)`` succeeds
+    for exactly that case, which would make a successful stop report itself as
+    a failure.
+    """
+    def _stopped() -> bool:
+        state = _proc_state(pid)
+        return state is None or state == "Z"
+
+    for target in (-pid, pid):
+        try:
+            os.kill(target, 15)
+            break
+        except OSError:
+            continue
+    deadline = time.time() + grace_s
+    while time.time() < deadline:
+        if _stopped():
+            return True
+        time.sleep(0.2)
+    for target in (-pid, pid):
+        try:
+            os.kill(target, 9)
+            break
+        except OSError:
+            continue
+    time.sleep(0.3)
+    return _stopped()
+
+
+def _supervisor_running(workspace_dir: Path) -> Optional[int]:
+    """PID of a live incident supervisor for this experiment (None when idle).
+
+    A zombie must not count: ``os.kill(pid, 0)`` succeeds for a defunct process
+    whose parent has not reaped it, and that stale pid then blocks every future
+    supervision of the task ("supervisor already running" for a process that
+    cannot observe anything).
+    """
+    pid_file = Path(workspace_dir) / "supervisor.pid"
+    try:
+        pid = int(pid_file.read_text(encoding="utf-8").split()[0])
+    except (OSError, ValueError, IndexError):
+        return None
+    try:
+        state = Path(f"/proc/{pid}/stat").read_text().split(") ", 1)[1]
+    except (OSError, IndexError, ValueError):
+        return None
+    if state.split()[0] == "Z":
+        return None
+    try:
+        os.kill(pid, 0)
+    except OSError:
+        return None
+    return pid
+
+
+def _launch_supervisor(state_dir: Path, workspace_dir: Path, target: Any,
+                       diagnose: bool = False) -> Dict[str, Any]:
+    """Start the unattended supervisor (self-healing) for this experiment."""
+    import subprocess
+    import sys
+
+    meta_root = Path(__file__).resolve().parents[4]
+    log_path = Path(workspace_dir) / "supervisor.log"
+    log = log_path.open("a", encoding="utf-8")
+    log.write(f"\n=== supervisor requested: target={target} diagnose={diagnose} "
+              f"@ {time.strftime('%Y-%m-%d %H:%M:%S')} ===\n")
+    log.flush()
+    cmd = [
+        sys.executable, "-m",
+        "metainfer.tasks.harness_evolve.orchestrator.supervisor", "watch",
+        "--state-dir", str(state_dir), "--workspace-dir", str(workspace_dir),
+    ]
+    if target not in (None, ""):
+        cmd += ["--target", str(int(target))]
+    if diagnose:
+        cmd += ["--diagnose-agent"]
+    env = dict(os.environ)
+    env.setdefault("PYTHONPATH", str(meta_root))
+    proc = subprocess.Popen(
+        cmd, cwd=str(meta_root), stdout=log, stderr=subprocess.STDOUT,
+        stdin=subprocess.DEVNULL, env=env, start_new_session=True,
+    )
+    log.close()
+    (Path(workspace_dir) / "supervisor.pid").write_text(str(proc.pid),
+                                                         encoding="utf-8")
+    return {"pid": proc.pid, "log": str(log_path)}
+
+
+def _supervisor_payload(workspace_dir: Path) -> Dict[str, Any]:
+    workspace_dir = Path(workspace_dir)
+    state: Dict[str, Any] = {}
+    try:
+        state = json.loads((workspace_dir / "supervisor_state.json")
+                           .read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        state = {}
+    incidents: List[Dict[str, Any]] = []
+    path = workspace_dir / "incidents.jsonl"
+    if path.is_file():
+        for line in path.read_text(encoding="utf-8", errors="replace").splitlines()[-40:]:
+            try:
+                incidents.append(json.loads(line))
+            except ValueError:
+                continue
+    return {
+        "running": _supervisor_running(workspace_dir) is not None,
+        "pid": _supervisor_running(workspace_dir),
+        "state": state,
+        "incidents": incidents,
+        "has_report": (workspace_dir / "incident_report.md").is_file(),
+    }
+
+
+_JANITOR_STARTED = False
+
+
+def janitor_sweep() -> List[Dict[str, Any]]:
+    """Last-resort recovery: restart experiments nobody is supervising.
+
+    The supervisor normally handles crashes, but if the supervisor itself dies
+    (OOM, kill, container restart) the experiment would sit idle forever — the
+    exact "stuck in analyze" symptom. This sweep runs inside the server and
+    re-launches supervision for any task that is unfinished and has neither a
+    live experiment process nor a live supervisor.
+    """
+    from metainfer.server.tasks import list_tasks
+
+    actions: List[Dict[str, Any]] = []
+    for entry in list_tasks():
+        if getattr(entry, "type", None) != PLUGIN_TYPE:
+            continue
+        try:
+            state_dir = state_dir_for(entry)
+            workspace_dir = workspace_dir_for(entry)
+            stopped = _operator_stopped(workspace_dir)
+            if stopped is not None:
+                # "Last-resort recovery" must not overrule a human: this is the
+                # one sweep that would otherwise restart a task minutes after
+                # the operator stopped it.
+                actions.append({"task": entry.id, "action": "skip",
+                                "reason": "stopped by operator",
+                                "stopped_at": stopped.get("stopped_at")})
+                continue
+            run = read_run(state_dir)
+            if not run:
+                continue
+            target = read_target_iterations(workspace_dir)
+            run_live = _run_alive(state_dir)
+            supervisor_live = _supervisor_running(workspace_dir)
+            if run_live:
+                # Somebody *is* running this experiment — most often the process
+                # the New Task form started a second ago. Starting a second
+                # driver here is what made a fresh task open two rounds at once.
+                actions.append({"task": entry.id, "action": "skip",
+                                "reason": "run is alive", "pid": run_live})
+                continue
+            if supervisor_live:
+                if should_launch_resume(run, target, False):
+                    resume = _launch_resume(state_dir, workspace_dir, 1,
+                                            champion_iteration=None)
+                    actions.append({"task": entry.id, "action": "resume",
+                                    "pid": resume.get("pid")})
+                continue
+            if not should_launch_resume(run, target, False):
+                continue
+            out = _launch_supervisor(state_dir, workspace_dir, target,
+                                     diagnose=False)
+            resume = _launch_resume(state_dir, workspace_dir, 1,
+                                    champion_iteration=None)
+            actions.append({"task": entry.id, "action": "supervisor+resume",
+                            "supervisor_pid": out.get("pid"),
+                            "resume_pid": resume.get("pid")})
+        except Exception as exc:  # noqa: BLE001 - never break the sweep
+            actions.append({"task": getattr(entry, "id", "?"),
+                            "action": "error", "error": repr(exc)})
+    return actions
+
+
+def _start_janitor(interval_s: int = 300) -> None:
+    global _JANITOR_STARTED
+    if _JANITOR_STARTED:
+        return
+    _JANITOR_STARTED = True
+    import threading
+
+    def loop() -> None:
+        while True:
+            time.sleep(max(30, int(interval_s)))
+            try:
+                actions = janitor_sweep()
+                for action in actions:
+                    if action.get("action") != "error":
+                        print(f"[he-janitor] {action}", flush=True)
+            except Exception:  # noqa: BLE001
+                continue
+
+    threading.Thread(target=loop, daemon=True, name="he-janitor").start()
+
+
+def build_router(plugin) -> APIRouter:
+    router = APIRouter()
+    if os.environ.get("METAINFER_DISABLE_JANITOR", "").strip().lower() not in {
+            "1", "true", "yes"}:
+        _start_janitor()
+
+    @router.get("/summary")
+    def summary(task_id: str) -> Dict[str, Any]:
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        exp = workspace_dir_for(entry)
+        state_dir = state_dir_for(entry)
+        return {
+            "config_snapshot": _load(exp / "config_snapshot.json", None),
+            "report": _text(exp / "report.md"),
+            "best_ever": _load(exp / "best_ever.json", None),
+            "scores": _jsonl(exp / "iteration_scores.jsonl"),
+            # Which protocol this run is on. ``protocol.json`` is written by the
+            # orchestrator when it leaves the fixed one; the page puts a banner
+            # on screen from it, because a run that judges neither gate looks
+            # exactly like a run that is merely quiet.
+            "protocol": _load(exp / "protocol.json", None)
+            or {"protocol": "variant", "mode": "variant",
+                "reason": "FLOW.md fixed protocol"},
+            # Whether the operator pressed Stop: the page needs it to explain
+            # why nothing is running, and to offer the way back.
+            "operator_stop": _operator_stopped(exp),
+            # Liveness as this plugin knows it. The shell's green dot reads
+            # ``orchestrator.pid``, which a run started before that file existed
+            # never wrote -- the process is alive and the dot is blue. The page
+            # must not inherit that blind spot, or Stop would be hidden on
+            # exactly the runs that need it.
+            "running": _run_live(state_dir, exp),
+        }
+
+    @router.get("/iterations")
+    def iterations(task_id: str) -> List[Dict[str, Any]]:
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        return _iterations(workspace_dir_for(entry))
+
+    @router.get("/state-graph")
+    def state_graph(task_id: str) -> Dict[str, Any]:
+        """Outer-loop AHE phase machine (Prepare..Finished)."""
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        run = read_run(state_dir_for(entry))
+        return phases.graph_payload(run.get("current_phase") or PREPARE)
+
+    @router.get("/iterations/{num}/gpu")
+    def iteration_gpu(task_id: str, num: int) -> Dict[str, Any]:
+        """Read-only device cards for one AHE iteration.
+
+        Replaces the old lease view: HE holds no lease and takes no reading of
+        the cards, so the page shows the layout HE chose (``index % 4``) plus
+        what each child's *own* DKAO gate says about the device it was handed.
+        """
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        req = read_requirements(state_dir_for(entry)) or {}
+        answers = req.get("answers")
+        if not isinstance(answers, dict):
+            answers = req
+        exp = workspace_dir_for(entry)
+        return {
+            "iteration": num,
+            "gpu_count": GPU_COUNT,
+            "managed_by": "dcu_kernel_auto_opt",
+            "leases": [],
+            "layout": _load(
+                exp / "runs" / f"iteration_{num:03d}" / "input" / "benchmark"
+                / "gpu_preflight.json", None),
+            "cards": _gpu_cards(exp, num, answers),
+        }
+
+    @router.get("/iterations/{num}/children")
+    def iteration_children(task_id: str, num: int) -> Dict[str, Dict[str, Any]]:
+        """Live inner DKAO phase per question of one AHE iteration."""
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        req = read_requirements(state_dir_for(entry)) or {}
+        answers = req.get("answers")
+        if not isinstance(answers, dict):
+            answers = req
+        return _children_live(workspace_dir_for(entry), num, answers)
+
+    @router.post("/resume")
+    async def resume(task_id: str, request: Request) -> Dict[str, Any]:
+        """Continue the experiment with further round(s) from the current state."""
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        state_dir = state_dir_for(entry)
+        workspace_dir = workspace_dir_for(entry)
+        try:
+            body = await request.json()
+        except Exception:  # noqa: BLE001 - body is optional
+            body = {}
+        body = body or {}
+        target = body.get("target")
+        if target is not None:
+            target = max(1, int(target))
+            write_target_iterations(workspace_dir, target)
+
+        live = _resume_running(state_dir)
+        if live:
+            # Already running: a new target is picked up by the next round.
+            return {
+                "ok": True, "running": True, "pid": live,
+                "target": read_target_iterations(workspace_dir),
+                "updated": target is not None,
+                "message": ("target updated; the running loop picks it up "
+                            "at the next round" if target is not None
+                            else "already running"),
+            }
+        run = read_run(state_dir)
+        if not run.get("finished"):
+            raise HTTPException(
+                409, "task is not finished yet; wait for the current round")
+        iterations = max(1, int(body.get("iterations") or 1))
+        champion_iteration = body.get("champion_iteration")
+        out = _launch_resume(state_dir, workspace_dir, iterations,
+                             champion_iteration=champion_iteration)
+        out["target"] = read_target_iterations(workspace_dir)
+        return out
+
+    @router.post("/autopilot")
+    async def autopilot(task_id: str, request: Request) -> Dict[str, Any]:
+        """Unattended mode: run to the target round, self-healing on incidents."""
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        state_dir = state_dir_for(entry)
+        workspace_dir = workspace_dir_for(entry)
+        try:
+            body = await request.json()
+        except Exception:  # noqa: BLE001
+            body = {}
+        body = body or {}
+        target = body.get("target")
+        if target is not None:
+            write_target_iterations(workspace_dir, max(1, int(target)))
+        target = read_target_iterations(workspace_dir)
+        live = _supervisor_running(workspace_dir)
+        if live:
+            return {"ok": True, "supervisor_pid": live, "target": target,
+                    "message": "supervisor already running; target updated"}
+        out = _launch_supervisor(state_dir, workspace_dir, target,
+                                 diagnose=bool(body.get("diagnose_agent")))
+        resume_pid = None
+        run = read_run(state_dir)
+        if should_launch_resume(run, target, bool(_resume_running(state_dir))):
+            # Covers both "not started yet" and "the previous process died
+            # mid-round": the run must never be left unsupervised.
+            champion = body.get("champion_iteration")
+            resume = _launch_resume(state_dir, workspace_dir, 1,
+                                    champion_iteration=champion)
+            resume_pid = resume.get("pid")
+        return {"ok": True, "supervisor_pid": out["pid"], "target": target,
+                "resume_pid": resume_pid, "log": out["log"]}
+
+    @router.get("/review")
+    def review(task_id: str) -> Dict[str, Any]:
+        """Cross-round review: generations, best-known curve, regressions."""
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        ws = workspace_dir_for(entry)
+        regressions: List[Dict[str, Any]] = []
+        path = ws / "regressions.jsonl"
+        if path.is_file():
+            for line in path.read_text(encoding="utf-8",
+                                      errors="replace").splitlines()[-40:]:
+                try:
+                    regressions.append(json.loads(line))
+                except ValueError:
+                    continue
+        return {
+            "review": _load(ws / "round_review.json", None),
+            "best_known": _load(ws / "best_known.json", {}),
+            "regressions": regressions,
+        }
+
+    @router.get("/promotion")
+    def promotion(task_id: str) -> Dict[str, Any]:
+        """Pending harness promotion + what production currently is.
+
+        Nothing reaches DKAO without an explicit approval, so the UI needs to
+        see the candidate, its gate evidence, and the current production
+        version in one place.
+        """
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        ws = workspace_dir_for(entry)
+        from ..orchestrator.promotion import promoted_harness, read_pending
+
+        pending = read_pending(ws)
+        promoted = promoted_harness(ws)
+        table_path = ws / "variant_table.json"
+        table = _load(table_path, {}) or {}
+        return {
+            "pending": pending,
+            "promoted": promoted,
+            "variant_table": {
+                "harness_version": table.get("harness_version"),
+                "operators": len(table.get("operators") or {}),
+                "groups": table.get("groups") or {},
+                "last_promoted_at": table.get("last_promoted_at"),
+            },
+        }
+
+    @router.post("/promotion/approve")
+    async def promotion_approve(task_id: str, request: Request) -> Dict[str, Any]:
+        """Human approval: publish the pending harness into DKAO."""
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        workspace_dir = workspace_dir_for(entry)
+        try:
+            body = await request.json()
+        except Exception:  # noqa: BLE001 - body is optional
+            body = {}
+        by = str((body or {}).get("by") or "operator")
+        result = _approve_with_config(entry, workspace_dir, by)
+        return {"ok": bool(result.get("ok")), "result": result,
+                "approved_by": by}
+
+    @router.post("/promotion/deny")
+    async def promotion_deny(task_id: str, request: Request) -> Dict[str, Any]:
+        """Reject the pending candidate; the run keeps iterating."""
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        workspace_dir = workspace_dir_for(entry)
+        try:
+            body = await request.json()
+        except Exception:  # noqa: BLE001
+            body = {}
+        reason = str((body or {}).get("reason") or "operator")
+        result = _deny_with_config(entry, workspace_dir, reason)
+        return {"ok": bool(result.get("ok")), "result": result}
+
+    @router.get("/supervisor")
+    def supervisor(task_id: str) -> Dict[str, Any]:
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        return _supervisor_payload(workspace_dir_for(entry))
+
+    @router.post("/supervisor/stop")
+    def supervisor_stop(task_id: str) -> Dict[str, Any]:
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        workspace_dir = workspace_dir_for(entry)
+        (workspace_dir / "supervisor.stop").write_text(
+            str(time.time()), encoding="utf-8")
+        pid = _supervisor_running(workspace_dir)
+        if pid:
+            try:
+                os.kill(pid, 15)
+            except OSError:
+                pass
+        return {"ok": True, "stopped_pid": pid}
+
+    @router.post("/stop")
+    def stop(task_id: str) -> Dict[str, Any]:
+        """Stop the whole run, leaving it restartable.
+
+        Stopping is not killing: this is the operator's "pause". Everything
+        that would keep the round moving is brought down -- the incident
+        supervisor (so it cannot restart what we are stopping), the running
+        DKAO children, and the HE orchestrator itself -- and the run is left
+        explicitly *finished*, so no automatic path revives it while the
+        artifacts stay intact for a later resume.
+
+        Killing instead would be indistinguishable from a crash, which is
+        exactly what ``should_launch_resume`` and the janitor sweep exist to
+        recover from: they would restart the task seconds later.
+
+        Order matters, and the order is the one below: marker (1) so no sweep
+        revives the run while it is being torn down; supervisor (2), children
+        (3), orchestrator (4), one more child sweep for the env-retry wave the
+        orchestrator may have launched while dying (5), and only then the
+        terminal state (6). A live orchestrator never reads the marker itself,
+        so skipping (4) leaves a process that answers the loss of its children
+        by re-launching them.
+        """
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        state_dir = state_dir_for(entry)
+        workspace_dir = workspace_dir_for(entry)
+        run_before = read_run(state_dir)
+        iteration = run_before.get("current_iteration")
+
+        # 1. The marker first: it is what makes this stop stick. A concurrent
+        #    janitor tick must see it before it decides to relaunch anything.
+        stopped_at = time.time()
+        _operator_stop_path(workspace_dir).write_text(json.dumps({
+            "stopped_at": stopped_at,
+            "by": "webui",
+            "iteration": iteration,
+            "phase": run_before.get("current_phase"),
+        }, indent=2), encoding="utf-8")
+
+        # 2. The supervisor, before the orchestrator: it exists to restart a
+        #    dead experiment, and would do exactly that to the one we stop.
+        (workspace_dir / "supervisor.stop").write_text(
+            str(stopped_at), encoding="utf-8")
+        supervisor_pid = _supervisor_running(workspace_dir)
+        if supervisor_pid:
+            try:
+                os.kill(supervisor_pid, 15)
+            except OSError:
+                pass
+
+        # 3. The children, before the parent: the parent's own shutdown would
+        #    otherwise leave them (and their agents) running.
+        children = _live_dkao_children(workspace_dir)
+        stopped_children = [c["pid"] for c in children
+                            if _terminate_pid(c["pid"])]
+
+        # 4. The orchestrator itself, now that nothing below it is mid-write.
+        #    ``_run_alive`` and NOT ``_resume_running``: the latter reads only
+        #    ``resume.pid``, which a task created from the New Task form never
+        #    writes -- so Stop answered "stopped" while the orchestrator kept
+        #    running, and immediately re-launched the children it had just lost
+        #    as an env retry. ``_run_alive`` reads ``orchestrator.pid``,
+        #    ``resume.pid`` and finally scans cmdlines; the janitor sweep learned
+        #    this same lesson first (see its docstring).
+        orchestrator_pid = _run_alive(state_dir)
+        stopped_ok = not orchestrator_pid or _terminate_pid(orchestrator_pid)
+
+        # 5. Once more for the retry wave. A question whose child vanished is an
+        #    environment failure to the evaluator, which answers it with a fresh
+        #    child (``env_retry_attempts``) -- and it can launch that wave in the
+        #    seconds between step 3 and step 4. Miss it and the UI says "stopped"
+        #    while a brand-new DKAO child optimizes on a card for hours.
+        already = set(stopped_children)
+        relaunched = [c["pid"] for c in _live_dkao_children(workspace_dir)
+                      if c["pid"] not in already]
+        stopped_children += [pid for pid in relaunched if _terminate_pid(pid)]
+
+        # 6. Make it visible and final: a round cut short must not look like a
+        #    round that finished, and ``finished`` is what tells every
+        #    automatic path to keep its hands off.
+        run = read_run(state_dir)
+        notes = list(run.get("notes") or [])
+        if "Stopped by operator." not in notes:
+            notes.append("Stopped by operator.")
+        run.update({
+            "finished": True,
+            "final_status": "stopped_by_operator",
+            "last_outcome": "stopped_by_operator",
+            "last_transition_label": "operator stop",
+            "last_update": stopped_at,
+            "notes": notes,
+        })
+        tmp = state_dir / ".run.json.stop.tmp"
+        tmp.write_text(json.dumps(run, indent=2), encoding="utf-8")
+        os.replace(tmp, state_dir / "run.json")
+
+        # A stale resume.pid would make /resume answer "already running" about
+        # a process we just killed.
+        try:
+            (state_dir / "resume.pid").unlink()
+        except OSError:
+            pass
+
+        return {
+            "ok": True,
+            "stopped": stopped_ok,
+            "stopped_children": stopped_children,
+            "relaunched_children": relaunched,
+            "stopped_supervisor_pid": supervisor_pid,
+            "stopped_orchestrator_pid": orchestrator_pid,
+            "iteration": iteration,
+            "message": ("已停止；产物保留，点「继续运行」可从当前进度接着跑"
+                        if stopped_ok else
+                        "停止信号已发送，但编排进程未及时退出；可再点一次"),
+        }
+
+    @router.get("/dkao/{num}/{child}")
+    def dkao_child(task_id: str, num: int, child: str) -> Dict[str, Any]:
+        """DKAO-like detail of one HE question (inner 8-step machine etc.)."""
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        return _dkao_child_payload(workspace_dir_for(entry), num, child)
+
+    @router.get("/iterations/{num}/overview")
+    def overview(task_id: str, num: int) -> str:
+        entry = task_or_404(task_id)
+        require_task_type(entry, PLUGIN_TYPE)
+        path = (
+            workspace_dir_for(entry) / "runs"
+            / f"iteration_{num:03d}" / "input" / "analysis" / "overview.md"
+        )
+        text = _text(path)
+        if text is None:
+            raise HTTPException(404, "overview not found")
+        return text
+
+    return router
diff --git a/metainfer/tasks/harness_evolve/static/dkao.css b/metainfer/tasks/harness_evolve/static/dkao.css
new file mode 100644
index 00000000..55f7e00f
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/static/dkao.css
@@ -0,0 +1,909 @@
+.dkao-detail {
+  display: grid;
+  grid-template-columns: minmax(0, 2fr) minmax(320px, 1fr);
+  gap: var(--sp-3);
+  align-items: start;
+}
+
+.dkao-state-panel,
+.dkao-iterations-panel {
+  grid-column: 1 / -1;
+}
+
+.dkao-state-machine {
+  display: flex;
+  align-items: center;
+  gap: 10px;
+  padding: 24px;
+  overflow-x: auto;
+}
+
+.dkao-state-step {
+  min-width: 132px;
+  padding: 13px 16px;
+  border: 1px solid var(--border);
+  border-radius: var(--radius);
+  background: var(--bg);
+  color: var(--fg-dim);
+  display: flex;
+  align-items: center;
+  gap: 9px;
+  white-space: nowrap;
+}
+
+.dkao-state-step.active {
+  color: var(--fg);
+  border-color: var(--accent);
+  background: rgba(88, 166, 255, 0.13);
+  box-shadow: 0 0 0 1px rgba(88, 166, 255, 0.2);
+}
+
+.dkao-state-step.terminal.active {
+  border-color: var(--ok);
+  background: rgba(63, 185, 80, 0.12);
+}
+
+.dkao-state-index {
+  width: 23px;
+  height: 23px;
+  border-radius: 50%;
+  background: var(--panel-2);
+  color: var(--accent-2);
+  display: inline-flex;
+  align-items: center;
+  justify-content: center;
+  font-size: var(--fs-xs);
+  font-weight: 700;
+}
+
+.dkao-state-arrow {
+  color: var(--muted);
+  font-size: 20px;
+}
+
+.dkao-explore-summary {
+  padding: 0 24px 24px;
+}
+
+.dkao-explore-heading {
+  margin: 0 0 10px;
+  color: var(--fg-dim);
+  font-size: var(--fs-xs);
+  letter-spacing: 0.04em;
+  text-transform: uppercase;
+}
+
+.dkao-explore-grid {
+  display: grid;
+  grid-template-columns: repeat(4, minmax(0, 1fr));
+  gap: 10px;
+}
+
+.dkao-explore-gpu {
+  min-width: 0;
+  padding: 12px;
+  border: 1px solid var(--border-soft);
+  border-radius: var(--radius);
+  background: var(--bg);
+}
+
+.dkao-explore-gpu.stale {
+  border-color: var(--warn);
+}
+
+.dkao-explore-gpu > header {
+  display: flex;
+  align-items: flex-start;
+  justify-content: space-between;
+  gap: 8px;
+}
+
+.dkao-explore-gpu > header > div {
+  display: flex;
+  min-width: 0;
+  flex-direction: column;
+  gap: 2px;
+}
+
+.dkao-explore-gpu > header strong {
+  color: var(--accent-2);
+}
+
+.dkao-explore-gpu > header small {
+  overflow: hidden;
+  color: var(--muted);
+  font-family: var(--fs-mono);
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+.dkao-explore-step {
+  margin: 12px 0 8px;
+  color: var(--fg);
+  font-size: var(--fs-sm);
+}
+
+.dkao-explore-gpu dl {
+  display: grid;
+  grid-template-columns: 56px minmax(0, 1fr);
+  gap: 4px 8px;
+  margin: 0;
+  font-size: var(--fs-xs);
+}
+
+.dkao-explore-gpu dt {
+  color: var(--muted);
+}
+
+.dkao-explore-gpu dd {
+  margin: 0;
+  overflow: hidden;
+  color: var(--fg-dim);
+  font-family: var(--fs-mono);
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+.dkao-explore-progress {
+  height: 5px;
+  margin-top: 11px;
+  overflow: hidden;
+  border-radius: 3px;
+  background: var(--panel-2);
+}
+
+.dkao-explore-progress > span {
+  display: block;
+  height: 100%;
+  border-radius: inherit;
+  background: var(--accent);
+  transition: width 180ms ease;
+}
+
+.dkao-explore-warning {
+  margin: 8px 0 0;
+  color: var(--warn);
+  font-size: var(--fs-xxs);
+}
+
+.dkao-worker-lanes {
+  display: flex;
+  flex-direction: column;
+}
+
+.dkao-worker-lane {
+  display: grid;
+  grid-template-columns: 125px minmax(280px, 0.9fr) minmax(340px, 1.1fr);
+  min-height: 132px;
+  border-bottom: 1px solid var(--border-soft);
+}
+
+.dkao-worker-lane:last-child {
+  border-bottom: 0;
+}
+
+.dkao-worker-lane.unassigned {
+  opacity: 0.65;
+}
+
+.dkao-worker-meta {
+  padding: 10px;
+  border-right: 1px solid var(--border-soft);
+  background: var(--bg-elev);
+  display: flex;
+  flex-direction: column;
+  align-items: flex-start;
+  gap: 5px;
+}
+
+.dkao-worker-meta strong {
+  color: var(--accent-2);
+  font-family: var(--fs-mono);
+}
+
+.dkao-worker-meta small {
+  color: var(--muted);
+  word-break: break-word;
+}
+
+.dkao-restart-worker {
+  width: 100%;
+  margin-top: auto;
+  padding: 5px 7px;
+  border-color: var(--warn);
+  color: var(--warn);
+  font-size: var(--fs-xs);
+}
+
+.dkao-iteration-strip {
+  display: flex;
+  gap: 7px;
+  padding: 7px 8px 9px;
+  overflow-x: auto;
+  min-height: 89px;
+  scrollbar-width: thin;
+}
+
+.dkao-iteration-card {
+  flex: 0 0 148px;
+  border: 1px solid var(--border-soft);
+  border-radius: var(--radius-sm);
+  background: var(--bg);
+  padding: 7px 8px;
+}
+
+.dkao-iteration-card.accepted {
+  border-color: rgba(63, 185, 80, 0.55);
+}
+
+.dkao-iteration-card.running {
+  border-color: rgba(88, 166, 255, 0.7);
+  background: rgba(88, 166, 255, 0.06);
+}
+
+.dkao-running-pulse {
+  height: 4px;
+  margin-top: 7px;
+  overflow: hidden;
+  border-radius: 2px;
+  background: var(--panel-2);
+}
+
+.dkao-running-pulse > span {
+  display: block;
+  width: 42%;
+  height: 100%;
+  border-radius: inherit;
+  background: var(--accent);
+  animation: dkao-iteration-running 1.25s ease-in-out infinite alternate;
+}
+
+@keyframes dkao-iteration-running {
+  from { transform: translateX(0); }
+  to { transform: translateX(138%); }
+}
+
+.dkao-iteration-title {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  margin-bottom: 5px;
+}
+
+.dkao-iteration-card dl {
+  display: grid;
+  grid-template-columns: 44px 1fr;
+  gap: 2px 5px;
+  margin: 0;
+  font-size: var(--fs-xs);
+}
+
+.dkao-iteration-card dt {
+  color: var(--muted);
+}
+
+.dkao-iteration-card dd {
+  margin: 0;
+  color: var(--fg-dim);
+  font-family: var(--fs-mono);
+}
+
+.dkao-lane-section {
+  min-width: 0;
+  border-right: 1px solid var(--border-soft);
+  display: flex;
+  flex-direction: column;
+}
+
+.dkao-lane-section:last-child {
+  border-right: 0;
+}
+
+.dkao-lane-heading {
+  min-height: 27px;
+  padding: 5px 9px;
+  border-bottom: 1px solid var(--border-soft);
+  display: flex;
+  justify-content: space-between;
+  color: var(--fg-dim);
+  font-size: var(--fs-xs);
+}
+
+.dkao-lane-heading span {
+  color: var(--muted);
+}
+
+.dkao-plan-strip {
+  display: flex;
+  gap: 7px;
+  padding: 7px 8px 9px;
+  overflow-x: auto;
+  min-height: 89px;
+  scrollbar-width: thin;
+}
+
+.dkao-plan-card {
+  flex: 0 0 190px;
+  border: 1px solid var(--border-soft);
+  border-radius: var(--radius-sm);
+  background: var(--bg);
+  padding: 7px 8px;
+}
+
+.dkao-plan-card.manual,
+.dkao-plan-card.guidance.pending {
+  border-color: rgba(88, 166, 255, 0.55);
+}
+
+.dkao-plan-title {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  margin-bottom: 5px;
+}
+
+.dkao-plan-card p {
+  margin: 0 0 4px;
+  color: var(--fg-dim);
+  font-size: var(--fs-xs);
+  line-height: 1.35;
+  display: -webkit-box;
+  -webkit-box-orient: vertical;
+  -webkit-line-clamp: 2;
+  overflow: hidden;
+}
+
+.dkao-plan-card small {
+  color: var(--muted);
+  display: block;
+  white-space: nowrap;
+  overflow: hidden;
+  text-overflow: ellipsis;
+}
+
+.dkao-guidance-form {
+  display: flex;
+  gap: 6px;
+  padding: 0 8px 7px;
+}
+
+.dkao-guidance-form input {
+  min-width: 0;
+  flex: 1;
+  padding: 6px 8px;
+  font-size: var(--fs-xs);
+}
+
+.dkao-guidance-form button {
+  flex: 0 0 auto;
+  padding: 5px 9px;
+  font-size: var(--fs-xs);
+}
+
+.dkao-guidance-message {
+  color: var(--accent-2);
+  padding: 0 8px 6px;
+}
+
+.dkao-validation-subtitle {
+  display: flex;
+  align-items: baseline;
+  justify-content: space-between;
+  gap: 12px;
+  margin: 0 0 10px;
+  color: var(--fg-dim);
+  font-size: var(--fs-xs);
+  letter-spacing: 0.04em;
+  text-transform: uppercase;
+}
+
+.dkao-validation-subtitle span {
+  color: var(--muted);
+  font-weight: 400;
+  letter-spacing: 0;
+  text-transform: none;
+}
+
+.dkao-events {
+  max-height: 320px;
+  overflow: auto;
+  padding: 8px 14px;
+}
+
+.dkao-events > div {
+  display: grid;
+  grid-template-columns: 150px 1fr;
+  gap: 12px;
+  padding: 7px 0;
+  border-bottom: 1px solid var(--border-soft);
+  font-size: var(--fs-xs);
+}
+
+.dkao-events code {
+  color: var(--accent-2);
+}
+
+.dkao-events span {
+  color: var(--muted);
+  word-break: break-word;
+}
+
+.dkao-skill-entry,
+.dkao-skill-library {
+  grid-column: 1 / -1;
+}
+
+.dkao-skill-entry,
+.dkao-skill-library-title {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 18px;
+  padding: 14px 16px;
+}
+
+.dkao-skill-entry h2,
+.dkao-skill-entry p,
+.dkao-skill-library-title h2,
+.dkao-skill-library-title p {
+  margin: 0;
+}
+
+.dkao-skill-columns {
+  display: grid;
+  grid-template-columns: 1fr 1fr;
+  min-height: 360px;
+  border-top: 1px solid var(--border-soft);
+}
+
+.dkao-skill-column {
+  min-width: 0;
+  padding: 14px;
+}
+
+.dkao-skill-column + .dkao-skill-column {
+  border-left: 1px solid var(--border-soft);
+}
+
+.dkao-skill-column h3 {
+  display: flex;
+  justify-content: space-between;
+  margin: 0;
+}
+
+.dkao-skill-column h3 span {
+  color: var(--muted);
+  font-family: var(--fs-mono);
+}
+
+.dkao-skill-column > p {
+  min-height: 20px;
+  margin: 4px 0 10px;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+.dkao-skill-list {
+  max-height: 560px;
+  overflow: auto;
+  display: flex;
+  flex-direction: column;
+  gap: 7px;
+}
+
+.dkao-skill-file {
+  border: 1px solid var(--border-soft);
+  border-radius: var(--radius-sm);
+  background: var(--bg);
+}
+
+.dkao-skill-file summary {
+  cursor: pointer;
+  padding: 9px 10px;
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 10px;
+}
+
+.dkao-skill-file summary > span:first-child {
+  min-width: 0;
+  display: flex;
+  flex-direction: column;
+}
+
+.dkao-skill-file summary strong,
+.dkao-skill-file summary small {
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+.dkao-skill-file summary small {
+  color: var(--muted);
+}
+
+.dkao-skill-file pre {
+  max-height: 380px;
+  overflow: auto;
+  margin: 0;
+  padding: 12px;
+  border-top: 1px solid var(--border-soft);
+  color: var(--fg-dim);
+  white-space: pre-wrap;
+  word-break: break-word;
+  font-size: var(--fs-xs);
+}
+
+@media (max-width: 1100px) {
+  .dkao-detail {
+    grid-template-columns: 1fr;
+  }
+
+  .dkao-validation-panel,
+  .dkao-events-panel {
+    grid-column: 1;
+  }
+
+  .dkao-worker-lane {
+    grid-template-columns: 115px minmax(260px, 1fr) minmax(320px, 1fr);
+    overflow-x: auto;
+  }
+
+  .dkao-skill-columns {
+    grid-template-columns: 1fr;
+  }
+
+  .dkao-skill-column + .dkao-skill-column {
+    border-left: 0;
+    border-top: 1px solid var(--border-soft);
+  }
+}
+
+/* ---- shape input component ------------------------------------------------ */
+
+.dkao-shape-input {
+  border: 1px solid var(--border);
+  border-radius: var(--radius);
+  overflow: hidden;
+}
+
+.dkao-shape-tabs {
+  display: flex;
+  align-items: center;
+  padding: 6px;
+  background: var(--panel-2);
+  border-bottom: 1px solid var(--border);
+  gap: 4px;
+}
+
+.dkao-shape-tab {
+  padding: 5px 11px;
+  border: 1px solid transparent;
+  border-radius: var(--radius-sm);
+  background: transparent;
+  color: var(--fg-dim);
+  font-size: var(--fs-xs);
+  cursor: pointer;
+  transition: background 0.12s;
+}
+
+.dkao-shape-tab:hover:not(:disabled) {
+  background: var(--bg);
+  color: var(--fg);
+}
+
+.dkao-shape-tab.active {
+  background: var(--bg);
+  color: var(--fg);
+  border-color: var(--border);
+  box-shadow: 0 1px 2px rgba(0,0,0,.06);
+}
+
+.dkao-shape-tab:disabled {
+  opacity: .45;
+  cursor: default;
+}
+
+.dkao-shape-spacer {
+  width: 12px;
+}
+
+.dkao-shape-mode-note {
+  margin-left: auto;
+  padding-right: 6px;
+  color: var(--muted);
+  font-size: var(--fs-xxs);
+}
+
+.dkao-shape-body {
+  padding: 10px;
+}
+
+.dkao-shape-body textarea.input {
+  margin: 0;
+  border: 0;
+  border-radius: 0;
+  resize: vertical;
+}
+
+.dkao-guided-form {
+  display: flex;
+  flex-direction: column;
+  gap: 10px;
+}
+
+.dkao-guided-info {
+  padding: 9px 11px;
+  background: var(--bg-elev);
+  border-radius: var(--radius-sm);
+  border: 1px solid var(--border-soft);
+  color: var(--fg-dim);
+  font-size: var(--fs-xs);
+  line-height: 1.45;
+}
+
+.dkao-guided-info strong {
+  color: var(--accent-2);
+}
+
+.dkao-guided-row {
+  display: flex;
+  align-items: baseline;
+  gap: 8px;
+  flex-wrap: wrap;
+}
+
+.dkao-guided-label {
+  flex: 0 0 130px;
+  min-width: 100px;
+  font-size: var(--fs-xs);
+  color: var(--fg-dim);
+  font-family: var(--fs-mono);
+  text-align: right;
+  padding-top: 6px;
+}
+
+.dkao-guided-input {
+  flex: 1 1 160px;
+  min-width: 120px;
+  font-family: var(--fs-mono);
+}
+
+.dkao-guided-hint {
+  flex: 1 1 200px;
+  min-width: 140px;
+  font-size: var(--fs-xxs);
+  color: var(--muted);
+  padding-top: 6px;
+}
+
+.dkao-guided-check {
+  display: flex;
+  align-items: center;
+  gap: 6px;
+  font-size: var(--fs-xs);
+  color: var(--fg-dim);
+  cursor: pointer;
+  padding-left: 138px;
+}
+
+.dkao-manual-assignment {
+  display: flex;
+  flex-direction: column;
+  gap: 10px;
+}
+
+.dkao-assignment-actions {
+  display: flex;
+  align-items: center;
+  flex-wrap: wrap;
+  gap: 5px;
+  color: var(--muted);
+  font-size: var(--fs-xxs);
+}
+
+.dkao-shape-catalog {
+  overflow: hidden;
+  border: 1px solid var(--border-soft);
+  border-radius: var(--radius-sm);
+}
+
+.dkao-shape-catalog-header {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  padding: 8px 10px;
+  background: var(--panel-2);
+  font-size: var(--fs-xs);
+}
+
+.dkao-shape-catalog-header span {
+  color: var(--muted);
+  font-size: var(--fs-xxs);
+}
+
+.dkao-tp-filter {
+  display: flex;
+  align-items: center;
+  gap: 4px;
+  padding: 6px 10px;
+  border-bottom: 1px solid var(--border-soft);
+  background: var(--panel-2);
+}
+
+.dkao-tp-filter-label {
+  margin-right: 4px;
+  color: var(--muted);
+  font-size: var(--fs-xxs);
+  letter-spacing: 0.05em;
+  text-transform: uppercase;
+}
+
+.dkao-shape-catalog-items {
+  display: grid;
+  grid-template-columns: repeat(2, minmax(0, 1fr));
+  gap: 5px;
+  padding: 7px;
+}
+
+.dkao-shape-choice {
+  display: flex;
+  align-items: center;
+  min-width: 0;
+  gap: 7px;
+  padding: 7px 8px;
+  border: 1px solid var(--border-soft);
+  border-radius: var(--radius-sm);
+  color: var(--fg-dim);
+  cursor: pointer;
+}
+
+.dkao-shape-choice.selected {
+  border-color: var(--accent);
+  background: color-mix(in srgb, var(--accent) 7%, var(--bg));
+  color: var(--fg);
+}
+
+.dkao-shape-choice > span {
+  display: flex;
+  min-width: 0;
+  flex-direction: column;
+  gap: 2px;
+}
+
+.dkao-shape-choice strong,
+.dkao-shape-choice small {
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+.dkao-shape-choice strong {
+  font-family: var(--fs-mono);
+  font-size: var(--fs-xxs);
+}
+
+.dkao-shape-choice small {
+  color: var(--muted);
+  font-size: 10px;
+}
+
+.dkao-gpu-grid {
+  display: grid;
+  grid-template-columns: repeat(2, minmax(0, 1fr));
+  gap: 8px;
+}
+
+.dkao-gpu-card {
+  min-width: 0;
+  overflow: hidden;
+  border: 1px solid var(--border-soft);
+  border-radius: var(--radius-sm);
+  background: var(--bg-elev);
+}
+
+.dkao-gpu-card > header {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  padding: 8px 10px;
+  border-bottom: 1px solid var(--border-soft);
+  background: var(--panel-2);
+  font-size: var(--fs-xs);
+}
+
+.dkao-gpu-card > header span {
+  color: var(--muted);
+  font-size: var(--fs-xxs);
+}
+
+.dkao-gpu-shapes {
+  display: flex;
+  min-height: 54px;
+  flex-direction: column;
+  gap: 6px;
+  padding: 7px;
+}
+
+.dkao-gpu-shape {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 8px;
+  min-width: 0;
+  padding: 7px 8px;
+  border: 1px solid var(--border-soft);
+  border-radius: var(--radius-sm);
+  background: var(--bg);
+}
+
+.dkao-gpu-shape > div {
+  display: flex;
+  min-width: 0;
+  flex-direction: column;
+  gap: 2px;
+}
+
+.dkao-gpu-shape strong,
+.dkao-gpu-shape small {
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+.dkao-gpu-shape strong {
+  font-family: var(--fs-mono);
+  font-size: var(--fs-xxs);
+}
+
+.dkao-gpu-shape small {
+  color: var(--muted);
+  font-size: 10px;
+}
+
+.dkao-gpu-shape select {
+  max-width: 72px;
+  padding: 3px 4px;
+  border: 1px solid var(--border);
+  border-radius: var(--radius-sm);
+  background: var(--panel-2);
+  color: var(--fg);
+  font-size: var(--fs-xxs);
+}
+
+.dkao-gpu-empty {
+  display: grid;
+  min-height: 38px;
+  place-items: center;
+  color: var(--muted);
+  font-size: var(--fs-xxs);
+}
+
+@media (max-width: 760px) {
+  .dkao-shape-catalog-items,
+  .dkao-gpu-grid,
+  .dkao-explore-grid {
+    grid-template-columns: 1fr;
+  }
+}
+
+@media (min-width: 761px) and (max-width: 1180px) {
+  .dkao-explore-grid {
+    grid-template-columns: repeat(2, minmax(0, 1fr));
+  }
+}
+
+/* Kernel repository rename panel */
+.dkao-repo-panel .dkao-repo-actions {
+  display: flex;
+  align-items: center;
+  gap: 12px;
+  margin-top: 6px;
+}
+
+.dkao-repo-panel .dkao-rename-msg {
+  color: var(--muted);
+  font-size: var(--fs-xs);
+}
diff --git a/metainfer/tasks/harness_evolve/static/he-api.js b/metainfer/tasks/harness_evolve/static/he-api.js
new file mode 100644
index 00000000..975f6b55
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/static/he-api.js
@@ -0,0 +1,142 @@
+// Fetch helpers for harness_evolve's task-specific endpoints.
+// Mounted by the shell at /api/harness-evolve//*.
+
+const TASK_BASE = (taskId) =>
+  `/api/harness-evolve/${encodeURIComponent(taskId)}`;
+
+export async function getSummary(taskId) {
+  const r = await fetch(`${TASK_BASE(taskId)}/summary`);
+  if (!r.ok) throw new Error(`summary: ${r.status}`);
+  return r.json();
+}
+
+export async function getIterations(taskId) {
+  const r = await fetch(`${TASK_BASE(taskId)}/iterations`);
+  if (!r.ok) throw new Error(`iterations: ${r.status}`);
+  return r.json();
+}
+
+export async function getStateGraph(taskId) {
+  const r = await fetch(`${TASK_BASE(taskId)}/state-graph`);
+  if (!r.ok) throw new Error(`state-graph: ${r.status}`);
+  return r.json();
+}
+
+export async function getChildren(taskId, n) {
+  const r = await fetch(`${TASK_BASE(taskId)}/iterations/${n}/children`);
+  if (!r.ok) throw new Error(`children ${n}: ${r.status}`);
+  return r.json();
+}
+
+export async function getDkaoChild(taskId, n, child) {
+  const r = await fetch(`${TASK_BASE(taskId)}/dkao/${n}/${encodeURIComponent(child)}`);
+  if (!r.ok) throw new Error(`dkao ${child}: ${r.status}`);
+  return r.json();
+}
+
+// Continue the experiment toward a target round count (dynamic: call again
+// while it runs to raise/lower the target).
+export async function resumeExperiment(taskId, target = null, iterations = null,
+                                      championIteration = null) {
+  const body = {};
+  if (target != null) body.target = Number(target);
+  if (iterations != null) body.iterations = Number(iterations);
+  if (championIteration != null) body.champion_iteration = Number(championIteration);
+  const r = await fetch(`${TASK_BASE(taskId)}/resume`, {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify(body),
+  });
+  const payload = await r.json().catch(() => ({}));
+  if (!r.ok) throw new Error(payload.detail || `resume: ${r.status}`);
+  return payload;
+}
+
+export async function getOverview(taskId, n) {
+  const r = await fetch(`${TASK_BASE(taskId)}/iterations/${n}/overview`);
+  if (!r.ok) throw new Error(`overview ${n}: ${r.status}`);
+  return r.text();
+}
+
+// Stop the run and leave it restartable: brings down the supervisor, the
+// running DKAO children and the orchestrator, and records an operator stop so
+// nothing restarts it automatically. Resume ("启动续跑至目标") clears that.
+export async function stopExperiment(taskId) {
+  const r = await fetch(`${TASK_BASE(taskId)}/stop`, { method: "POST" });
+  const payload = await r.json().catch(() => ({}));
+  if (!r.ok) throw new Error(payload.detail || `stop: ${r.status}`);
+  return payload;
+}
+
+// Unattended mode: run to a target round with automatic incident handling.
+export async function startAutopilot(taskId, target = null,
+                                     championIteration = null,
+                                     diagnose = true) {
+  const body = { diagnose_agent: Boolean(diagnose) };
+  if (target != null) body.target = Number(target);
+  if (championIteration != null) body.champion_iteration = Number(championIteration);
+  const r = await fetch(`${TASK_BASE(taskId)}/autopilot`, {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify(body),
+  });
+  const payload = await r.json().catch(() => ({}));
+  if (!r.ok) throw new Error(payload.detail || `autopilot: ${r.status}`);
+  return payload;
+}
+
+export async function getSupervisor(taskId) {
+  const r = await fetch(`${TASK_BASE(taskId)}/supervisor`);
+  if (!r.ok) throw new Error(`supervisor: ${r.status}`);
+  return r.json();
+}
+
+export async function stopSupervisor(taskId) {
+  const r = await fetch(`${TASK_BASE(taskId)}/supervisor/stop`, { method: "POST" });
+  const payload = await r.json().catch(() => ({}));
+  if (!r.ok) throw new Error(payload.detail || `stop: ${r.status}`);
+  return payload;
+}
+
+export async function getReview(taskId) {
+  const r = await fetch(`${TASK_BASE(taskId)}/review`);
+  if (!r.ok) throw new Error(`review: ${r.status}`);
+  return r.json();
+}
+
+// Read-only device cards for one iteration: which device each question was
+// handed to, plus what that child's own DKAO admission gate says.
+export async function getIterationGpu(taskId, n) {
+  const r = await fetch(`${TASK_BASE(taskId)}/iterations/${n}/gpu`);
+  if (!r.ok) throw new Error(`gpu ${n}: ${r.status}`);
+  return r.json();
+}
+
+// Human approval of a candidate harness (nothing reaches DKAO without it).
+export async function getPromotion(taskId) {
+  const r = await fetch(`${TASK_BASE(taskId)}/promotion`);
+  if (!r.ok) throw new Error(`promotion: ${r.status}`);
+  return r.json();
+}
+
+export async function approvePromotion(taskId, by = "operator") {
+  const r = await fetch(`${TASK_BASE(taskId)}/promotion/approve`, {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ by }),
+  });
+  const payload = await r.json().catch(() => ({}));
+  if (!r.ok) throw new Error(payload.detail || `approve: ${r.status}`);
+  return payload;
+}
+
+export async function denyPromotion(taskId, reason = "operator") {
+  const r = await fetch(`${TASK_BASE(taskId)}/promotion/deny`, {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ reason }),
+  });
+  const payload = await r.json().catch(() => ({}));
+  if (!r.ok) throw new Error(payload.detail || `deny: ${r.status}`);
+  return payload;
+}
diff --git a/metainfer/tasks/harness_evolve/static/he-detail.js b/metainfer/tasks/harness_evolve/static/he-detail.js
new file mode 100644
index 00000000..cede7a84
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/static/he-detail.js
@@ -0,0 +1,1346 @@
+// harness_evolve task detail body (AHE experiment viewer).
+// Rendered by the task-detail shell when detail_view_module === "app/he-detail".
+// Read-only: fetches /api/harness-evolve//* and renders summary + per
+// iteration benchmark / diff / change evaluation / overview / rollback.
+
+import { html } from "htm/preact";
+import { useCallback, useEffect, useMemo, useState } from "preact/hooks";
+import { approvePromotion, denyPromotion, getChildren, getDkaoChild, getIterationGpu, getIterations, getPromotion, getReview, getStateGraph, getSummary, getSupervisor, resumeExperiment, startAutopilot, stopExperiment, stopSupervisor } from "app/he-api";
+// Shared polling hook: a hidden tab must not keep hammering the server, and a
+// round is only polled while it can still change (see usePollWhileVisible).
+import { usePollWhileVisible as pollWhileVisible } from "app/utils";
+// The DKAO task page itself, reused as a read-only mirror of one question.
+// Reused rather than reimplemented so the two views cannot drift: the same
+// component renders a child here and renders a standalone DKAO task.
+import DcuKernelAutoOptDetail from "app/dkao-detail";
+
+const withTimeout = (p, ms = 8000) =>
+  Promise.race([
+    p,
+    new Promise((_, rej) => setTimeout(() => rej(new Error("timeout")), ms)),
+  ]);
+
+// Latency formatting for the round table, which has no local helper of its own.
+const fmtUs = (v) => (v == null ? "-" : (v.toFixed ? v.toFixed(2) : String(v)));
+
+// The variant comparison's verdict (FLOW.md §3), not the round's PASS/FAIL.
+// A tie is deliberately not a win: the performance gate counts only a strict
+// win (outside the 2% noise band), so "平" is a distinct, visible outcome.
+function GateVerdictBadge({ verdict }) {
+  const cls = { WIN: "ok", LOSS: "bad", NO_SIGNAL: "warn" }[verdict] || "warn";
+  const label = { WIN: "赢", LOSS: "输", NO_SIGNAL: "平·不算赢" }[verdict]
+    || verdict || "-";
+  return html`${label}`;
+}
+
+// Where the number a question must beat came from. "variant" = the experiment's
+// own variant table, which round 1 seeds from the pool's best_known before any
+// kernel is measured here — say which, so "已记录的 variant" is never read as
+// "measured by this run" when it was not.
+const VARIANT_SOURCE_LABEL = {
+  best_known: "池子 best_known",
+  baseline: "固定 Triton baseline（尚无 variant，退到基线）",
+};
+
+function variantSourceLabel(gate) {
+  if (!gate) return "-";
+  if (gate.variant_source && gate.variant_source !== "variant") {
+    return VARIANT_SOURCE_LABEL[gate.variant_source] || gate.variant_source;
+  }
+  return String(gate.kernel_source || "").startsWith("pool:")
+    ? "本实验 variant（由池子 best_known 播种）"
+    : "本实验已记录的 variant";
+}
+
+// The round's gate line: which stage this round is, which gate it owed, and what
+// that gate concluded. Read from the round's own gate document, never inferred.
+function RoundGateLine({ iter }) {
+  const stage = iter && iter.stage && iter.stage.stage;
+  const gate = (iter && (iter.generalization_gate || iter.performance_gate))
+    || null;
+  if (!stage) {
+    return html`

+      本轮无 stage.json：legacy / suite 路径，未按固定协议判性能门
+    
`;
+  }
+  const KIND = { baseline: "基线轮（不判门）", performance: "性能门",
+                 generalization: "泛化门" };
+  const cls = !gate ? "run"
+    : (gate.status === "PASS" ? "ok"
+       : (gate.status === "INCOMPLETE" ? "warn" : "bad"));
+  return html`

+    stage ${stage} · ${KIND[stage] || stage}
+    ${gate
+      ? html` · 结论 ${gate.status}
+          ${(gate.wins || []).length}/${(gate.operators || []).length} 赢（需 ${gate.wins_needed}）
+          ${(gate.below_floor || []).length
+            ? html` · 低于地板：${gate.below_floor.join(", ")}` : null}
+          ${(gate.unmeasured || []).length
+            ? html` · 未测到：${gate.unmeasured.join(", ")}` : null}`
+      : html` · 尚未判门（第 1 轮只量基线，不判门）`}
+  
`;
+}
+
+// One question's real gate cell: the variant it must beat, plus the verdict when
+// the round has been judged. The pool criterion (tau_family * baseline) is NOT
+// the variant-round gate, so it is only ever shown as a labelled annotation.
+function QuestionGateCell({ gate, poolRef }) {
+  if (!gate) return html`-`;
+  const poolLine = (poolRef && poolRef.target_us != null) ? html`
+    

+      池子口径参考 ≤${fmtUs(poolRef.target_us)} µs
+      （τ ${fmtUs(poolRef.tau_family)} × baseline；本协议不判门）
+    
` : null;
+  if (gate.protocol === "legacy") {
+    // Rounds before 2026-09-15 really were judged this way: say so plainly.
+    const need = (poolRef && poolRef.baseline_us && poolRef.target_us)
+      ? (poolRef.baseline_us / poolRef.target_us) : null;
+    return html`
+      
≤${fmtUs(gate.target_us)} µs

+      legacy 判据${need ? ` · ≥${need.toFixed(2)}× baseline` : ""}
+    `;
+  }
+  const judged = gate.judged;
+  const src = variantSourceLabel(gate);
+  const title = `性能门对手 = 该算子当前 variant ${fmtUs(gate.variant_us)} µs（来源：${src}）`
+    + ` · 赢 = 严格快 ${gate.noise_percent}% 以上（平局不算赢）`
+    + ` · 地板 = variant 的 ${Math.round((gate.floor_percent || 0.8) * 100)}%`
+    + (judged ? ` · 本轮实测 ${fmtUs(judged.candidate_median_us)} µs` : "")
+    + (poolRef && poolRef.target_us != null
+       ? ` · 池子口径 τ×baseline = ${fmtUs(poolRef.target_us)} µs（本协议不判门）` : "");
+  return html`
+    
≤${fmtUs(gate.variant_us)} µs variant

+    ${judged
+      ? html`<${GateVerdictBadge} verdict=${judged.verdict} />
+             ${judged.delta_percent > 0 ? "+" : ""}${fmtUs(judged.delta_percent)}%`
+      : html`${gate.stage === "baseline"
+            ? "第 1 轮不判门 · 本轮定下的对手"
+            : "本轮尚未判门"}`}
+    ${poolLine}
+  `;
+}
+
+// Outer-loop AHE step machine (Prepare -> Evaluate -> ... -> Finished).
+// Uses the same classes as the DKAO page (.dkao-state-*) for identical looks.
+function StateMachine({ graph }) {
+  const nodes = graph?.nodes || [];
+  if (!nodes.length) return html`
Waiting for the control plane.
`;
+  return html`
+    

+      ${nodes.map((node, index) => html`
+        

+          ${index + 1}
+          ${node.label}
+        

+        ${index < nodes.length - 1 ? html`→` : null}
+      `)}
+    

+  `;
+}
+
+// Inner DKAO phases (each question IS one DKAO task running its own loop).
+const DKAO_PHASE_LABEL = {
+  idle: "…",
+  prepare: "Prepare",
+  generate_kernel_repo: "Generate kernel repo",
+  baseline: "Baseline",
+  parallel_explore: "Parallel explore",
+  skill_synthesis: "Skill synthesis",
+  serial_validate: "Serial validate",
+  report: "Report",
+  finished: "Finished",
+};
+
+function DkaoPhaseChip({ child }) {
+  if (!child) return html`no child yet`;
+  const phase = child.current_phase || "idle";
+  if (child.finished) {
+    const cls = child.final_status === "success" ? "ok" : "bad";
+    return html`finished${child.final_status ? ` · ${child.final_status}` : ""}`;
+  }
+  return html`${DKAO_PHASE_LABEL[phase] || phase}`;
+}
+
+function StateBadge({ state }) {
+  const s = state || "";
+  const bad = /fail|error|reject|killed|timeout/.test(s);
+  const done = /pass|accept|success|complete|finish|done|publish/.test(s);
+  const cls = bad ? "bad" : done ? "ok" : "run";
+  return html`${s || "-"}`;
+}
+
+// Per-question inner DKAO view (renders the same way the DKAO detail page
+// does: the 8-step DKAO state machine + worker lanes + agents + log).
+function ChildDkaoView({ payload }) {
+  if (!payload) return html`
Loading DKAO child…
`;
+  const run = payload.run || {};
+  return html`
+    

+      

+        
DKAO child — ${payload.child}

+        

+          phase: ${DKAO_PHASE_LABEL[run.current_phase] || run.current_phase}
+          · iteration ${run.current_iteration ?? "-"}
+          ${run.finished ? html` · <${StateBadge} state=${run.final_status || "finished"} />` : ""}
+          · repo: ${payload.repo_path || "-"}
+        

+      

+      <${StateMachine} graph=${payload.graph} />
+      

+        
Workers

+        ${payload.workers && payload.workers.length ? html`
+          
+            
+            
+              ${payload.workers.map((w) => html`
+                
+                  
+                  
+                  
+                  
+                  
+                  
+                  
+                `)}
+            
+          
workerstateiterationshapeGPUroundsbest median us
${w.worker_id}<${StateBadge} state=${w.state} />${w.iteration ?? "-"}${w.shape_id || "-"}${w.physical_gpu ?? "-"}${w.rounds ?? 0}${w.best_median_us != null ? w.best_median_us : "-"}
` : html`
no worker lanes yet
`}
+        ${(() => {
+          const rounds = (payload.workers || [])
+            .flatMap((w) => (w.experiments || []).map((e) => ({ worker: w.worker_id, ...e })))
+            .sort((a, b) => (a.iteration || 0) - (b.iteration || 0));
+          if (!rounds.length) return null;
+          return html`
+            
Optimization rounds (accepted median per attempt, speedup vs the frozen Triton baseline)

+            
+              
+              
+                ${rounds.map((e) => html`
+                  
+                    
+                    
+                    
+                    
+                    
+                    
+                    
+                  `)}
+              
+            
workeritermedian usp90 usbaseline usspeedupaccepted
${e.worker}${e.iteration ?? "-"}${e.median_us != null ? e.median_us.toFixed ? e.median_us.toFixed(2) : e.median_us : "—"}${e.p90_us != null ? (e.p90_us.toFixed ? e.p90_us.toFixed(2) : e.p90_us) : "—"}${e.baseline_us != null ? (e.baseline_us.toFixed ? e.baseline_us.toFixed(2) : e.baseline_us) : "—"}${e.speedup != null ? (typeof e.speedup === "number" ? `${e.speedup.toFixed(2)}×` : e.speedup) : "-"}${e.accepted ? html`accept` : html`reject`}
`;
+        })()}
+      

+      

+        
Agents

+        ${payload.agents && payload.agents.length ? html`
+          
+            
+            
+              ${payload.agents.map((a) => html`
+                
+                  
+                  
+                  
+                  
+                  
+                  
+                `)}
+            
+          
agentrolestatusitermedian usspeedup vs baseline
${a.name || "-"}${a.role || "-"}<${StateBadge} state=${a.status} />${a.iteration != null ? a.iteration : "-"}${a.median_us != null ? (a.median_us.toFixed ? a.median_us.toFixed(2) : a.median_us) : "-"}${a.speedup != null
+                        ? html`${typeof a.speedup === "number" ? a.speedup.toFixed(2) + "×" : a.speedup}${a.accepted ? "" : " (rejected)"}`
+                        : "-"}
` : html`
no agents yet
`}
+      

+      ${payload.log_tail ? html`
child orchestrator log
+        
${payload.log_tail}
` : null}
+    
`;
+}
+
+function VerdictBadge({ verdict }) {
+  const cls = { EFFECTIVE: "ok", PARTIALLY_EFFECTIVE: "warn", MIXED: "warn",
+                INEFFECTIVE: "bad", HARMFUL: "bad", BASELINE: "ok",
+                PROMOTE: "ok", PROMOTE_UNEXPLAINED: "ok",
+                CONFIRM_REQUIRED: "warn", NO_SIGNAL: "warn", SPECIALIZE: "warn",
+                REJECT: "bad", REJECT_OVERFIT: "bad" }[verdict] || "warn";
+  return html`${verdict || "-"}`;
+}
+
+function ScoreTable({ scores }) {
+  if (!scores || !scores.length) return null;
+  return html`
+    
+      
+      
+        ${scores.map((s) => html`
+          
+          `)}
+      
+    
iterationpasstotalrate
${s.iteration}${s.n_pass}${s.n_total}${(100 * (s.pass_rate || 0)).toFixed(1)}%
`;
+}
+
+// --------------------------------------------------------------------------- //
+// Live DKAO mirror: one iteration's questions, rendered by the DKAO page itself
+// --------------------------------------------------------------------------- //
+// Each question of an iteration is a real DKAO run. Those runs are spawned by
+// the HE orchestrator's CLI and are absent from the task registry, so the
+// embedded view is addressed by the child's artifact directory instead of a
+// task id -- hence the `stateDir` prop, which the DKAO routes accept read-only.
+// The switcher follows the still-running question by default, because that is
+// the one the operator opened the iteration to watch.
+function LiveDkaoMirror({ taskId, iter, initialQuestion = null }) {
+  const num = useMemo(() => {
+    const m = (iter.name || "").match(/(\d+)\s*$/);
+    return m ? Number(m[1]) : 0;
+  }, [iter.name]);
+
+  const [children, setChildren] = useState(null);
+  const [sel, setSel] = useState(initialQuestion);
+  const [err, setErr] = useState(null);
+  //: How many DKAO pages this round shows at once. "grid" is the default: a
+  //: round is four questions on four cards, and comparing them is the point —
+  //: one-at-a-time behind tabs hid three quarters of the round. Each cell is the
+  //: same read-only DKAO page, pointed at that question's own state dir.
+  const [layout, setLayout] = useState("grid");
+
+  useEffect(() => {
+    if (!taskId || !num) return;
+    let alive = true;
+    withTimeout(getChildren(taskId, num))
+      .then((c) => { if (alive) { setChildren(c || {}); setErr(null); } })
+      .catch((e) => { if (alive) { setChildren({}); setErr(String(e)); } });
+    return () => { alive = false; };
+  }, [taskId, num]);
+
+  const ids = useMemo(() => Object.keys(children || {}).sort(), [children]);
+  // Follow whichever question is still running; fall back to the first, and
+  // never let a stale selection survive a refresh that dropped it.
+  const selected = useMemo(() => {
+    if (sel && ids.includes(sel)) return sel;
+    const running = ids.find((id) => children[id] && !children[id].finished);
+    return running || ids[0] || null;
+  }, [ids, sel, children]);
+
+  if (!num) return null;
+  if (children && !ids.length) {
+    return html`

+      
DKAO 实时视图 — ${iter.name}

+      
这一轮还没有派发 DKAO 子任务（或产物已被清理）。

+    
`;
+  }
+  const child = selected ? children[selected] : null;
+  const childSummary = (c) => html`
+    
+      · phase ${c.current_phase || "-"}
+      · iter ${c.current_iteration ?? "-"}
+      ${c.rounds ? html` · ${c.rounds} rounds` : null}
+      ${c.best_median_us != null ? html` · best ≈${Number(c.best_median_us).toFixed(2)} µs` : null}
+    `;
+  const mirror = (id, c) => html`
+    <${DcuKernelAutoOptDetail} taskId=${id} stateDir=${c.child_state_dir}
+      readOnly=${true} />`;
+  const statusBadge = (c) => (c.finished
+    ? html`${c.final_status || "finished"}`
+    : html`${c.current_phase || "…"} #${c.current_iteration ?? "-"}`);
+  const layouts = [
+    ["grid", "四宫格（四个同时看）"],
+    ["stack", "单列堆叠（四个都展开）"],
+    ["single", "只看一个（切换）"],
+  ];
+
+  return html`
+    

+      
DKAO 实时视图 — ${iter.name}
+        （该轮每道考题的 DKAO 页面，只读；可同时看四个）

+      

+        
+        ${layouts.map(([key, label]) => html`
+          `)}
+        四个子任务是各自独立的 DKAO 运行（一卡一题），所以四个页面互不影响
+      

+      

+        ${ids.map((id) => {
+          const c = children[id] || {};
+          const cls = selected === id ? "he-live-tab active" : "he-live-tab";
+          return html`
+            `;
+        })}
+      

+      ${!children ? html`
Loading DKAO children…
` : null}
+      ${err ? html`
无法读取子任务列表：${err}
` : null}
+      ${layout === "single"
+        ? (child && child.child_state_dir ? html`
+            

+              子任务数据目录
+              ${child.child_state_dir}
+              ${childSummary(child)}
+            

+            ${mirror(selected, child)}`
+          : html`
这一轮没有可显示的子任务。
`)
+        : html`
+          

+            ${ids.map((id) => {
+              const c = children[id] || {};
+              if (!c.child_state_dir) return null;
+              return html`
+                

+                  

+                    
+                    ${statusBadge(c)}
+                    ${childSummary(c)}
+                    
+                    
+                  

+                  
${mirror(id, c)}

+                
`;
+            })}
+          
`}
+    
`;
+}
+
+function IterBlock({ iter, taskId }) {
+  const results = (iter.benchmark && iter.benchmark.results) || {};
+  const rows = Object.entries(results);
+  return html`
+    

+      ${iter.name}
+      

+        ${iter.decision ? html`
+          

+            
System decision

+            <${VerdictBadge} verdict=${iter.decision.verdict} />
+            
${iter.decision.reason || ""}

+            performance: ${iter.decision.performance_gate?.status || "-"}
+            · mechanism: ${iter.decision.mechanism_gate?.status || "-"}
+            · generalization: ${iter.decision.generalization_gate?.status || "-"}
+            · held-out: ${iter.decision.heldout_gate?.status || "-"}
+            ${iter.decision.generalization_gate?.status === "FAIL" ? html`
+              
非本族探针退化 → 判为过拟合（探针不计入 pass rate，但阻止晋级）
` : null}
+            
+              决策配对 ${(iter.decision.paired_ids || []).length} 题
+              ${(iter.decision.probe_ids || []).length
+                ? `· 探针 ${(iter.decision.probe_ids || []).join(", ")} 不计入判定` : ""}
+            
+          
` : null}
+        ${iter.rollback ? html`
auto-rollback: ${iter.rollback.reason}
` : null}
+        ${iter.mechanism_evidence ? html`
+          

+            
Mechanism gate
+              （机制是否真的被执行）

+            ${iter.decision && iter.decision.mechanism_gate
+              ? html`<${VerdictBadge} verdict=${iter.decision.mechanism_gate.status} />` : null}
+            
+              
+              
+                ${((iter.mechanism_evidence.detail || {}).changes || []).map((c) => html`
+                  
+                    
+                    
+                    
+                    
+                  `)}
+              
+            
changegradekindevidence
${c.id}${(() => {
+                      const cls = { hit: "ok", partial: "ok",
+                                    unobserved: "warn", contradicted: "bad" }[c.grade] || "warn";
+                      return html`${c.grade}`;
+                    })()}${c.kind || "-"}${c.reason || ""}

+            
+              planner: ${String((iter.mechanism_evidence.evidence || {}).planner_enabled)}
+              · plan ids observed: ${((iter.mechanism_evidence.evidence || {}).plan_ids || []).join(", ") || "—"}
+              · harness revisions loaded: ${((iter.mechanism_evidence.evidence || {}).observed_harness_revisions || []).join(", ") || "—"}
+            
+          
` : null}
+        ${iter.promotion && (iter.promotion.results || []).length ? html`
+          

+            
Kernel promotion → DKAO variants
+              （阈值 ${iter.promotion.min_improvement_percent}% · 写前备份）

+            
+              
+              
+                ${iter.promotion.results.map((p) => html`
+                  
+                    
+                    
+                    
+                    
+                    
+                  `)}
+              
+            
shapeactionvariant median → newimprovementreason
${p.shape}${(() => {
+                      const cls = { added: "ok", updated: "ok", "would-add": "ok",
+                                    "would-update": "ok" }[p.action]
+                        || (p.action === "rejected" ? "bad" : "warn");
+                      return html`${p.action}`;
+                    })()}${p.old_median_us != null ? (p.old_median_us.toFixed ? p.old_median_us.toFixed(2) : p.old_median_us) : "—"}
+                        → ${p.new_median_us != null ? (p.new_median_us.toFixed ? p.new_median_us.toFixed(2) : p.new_median_us) : "—"}${p.improvement_percent != null ? `${p.improvement_percent.toFixed(2)}%` : "-"}
+                        ${p.speedup != null ? html`(${typeof p.speedup === "number" ? p.speedup.toFixed(2) + "× baseline" : p.speedup})` : null}${p.reason || ""} ${p.backup ? html`backup: ${p.backup}` : null}

+          
` : null}
+        ${rows.length ? html`
+          
+            
+            
+              ${rows.map(([id, r]) => html`
+                
+                
+                
+                
+                
+                `)}
+            
+          
instancerolepassedmedian/p90 uschild/reporeason
${id}${(() => {
+                  const purpose = (iter.purposes || {})[id];
+                  if (!purpose) return html`—`;
+                  const cls = purpose === "probe" ? "warn" : "ok";
+                  return html`${purpose}`;
+                })()}${r.passed ? html`pass`
+                                : html`fail`}${r.median_us != null ? r.median_us : "-"} / ${r.p90_us != null ? r.p90_us : "-"}
${r.child_task_id || "-"}
${r.repo_path || ""}		${r.reason || ""}
` : html`
no benchmark results yet
`}
+        <${LiveDkaoMirror} taskId=${taskId} iter=${iter} />
+        ${iter.diff ? html`
+          
flipped: ${(iter.diff.flipped || []).join(", ") || "-"}
+          · regressed: ${(iter.diff.regressed || []).join(", ") || "-"}
` : null}
+        ${iter.change_evaluation ? html`
+          
Change attribution (previous loop)

+          
+            
+            
+              ${(iter.change_evaluation.change_evaluations || []).map((c) => html`
+                
+                `)}
+            
+          
changehitverdict
${c.change_id}: ${(c.description || "").slice(0, 60)}${c.hit_rate}${html`<${VerdictBadge} verdict=${c.verdict} />`}
` : null}
+        ${iter.overview ? html`
analysis overview
+          
${iter.overview}
` : null}
+      

+    
`;
+}
+
+
+// Human decision on a candidate harness. Nothing is written to DKAO until the
+// operator approves, so this panel is the gate: it shows the candidate, the
+// evidence of both gates, and the current production version.
+// A run that left the fixed protocol (FLOW.md §1–§4) judges neither gate, so it
+// looks exactly like a run that is merely quiet. The orchestrator writes
+// protocol.json + a timeline event when that happens; this puts it on screen.
+function ProtocolBanner({ protocol }) {
+  if (!protocol || protocol.mode !== "legacy") return null;
+  return html`
+    

+      
⚠ 本任务未按固定协议运行 （legacy / suite 路径）

+      

+        原因：${protocol.reason || "unknown"}。
+        这条路径不固化泛化考卷、不判性能门/泛化门、不按 4 的倍数取题，
+        因此它的轮次**不会产出可接入 DKAO 的 harness**。
+        新任务请配置实测池文件（问卷池 Question pool），协议见 FLOW.md。
+      

+    

+  `;
+}
+
+// Four cards, one glance: this round's DKAO children as four columns.
+// The drill-down below still shows one child in full (read-only DKAO mirror);
+// this answers the other question — "what is each card doing right now?" —
+// without switching tabs. Every number is the child's own evidence: phase and
+// iteration come from its run.json, the median from its accepted rounds, the
+// gate line from its measurement_gate.jsonl, and the card membership from HE's
+// layout record for the round (never from the card a probe happened to use).
+function FourCardOverview({ data, onSelect, selected }) {
+  const cards = (data && data.cards) || [];
+  const gpuCount = (data && data.gpu_count) || 4;
+  if (!cards.length) return null;
+  const fmt = (v) => (v == null ? "-" : (v.toFixed ? v.toFixed(2) : v));
+  const byGpu = new Map(cards.map((c) => [Number(c.gpu), c]));
+  const now = Date.now() / 1000;
+  const quiet = (q) => {
+    const ts = Number(q.last_update || 0);
+    if (!ts || q.finished) return null;
+    return Math.max(0, Math.round(now - ts));
+  };
+  const statusChip = (q) => {
+    if (q.finished) {
+      const ok = q.final_status === "success" || q.final_status === "ok";
+      return html`已结束${q.final_status ? ` · ${q.final_status}` : ""}`;
+    }
+    const a = q.gate_audit || {};
+    if (a.last_event === "gate_blocked" && a.last_usable === false) {
+      return html`等卡（闸门未放行）`;
+    }
+    const secs = quiet(q);
+    if (secs != null && secs > 900) {
+      return html`静默 ${Math.round(secs / 60)}min`;
+    }
+    return html`运行中`;
+  };
+  const sourceLabel = (src) => (
+    src === "layout" ? "本轮布局"
+      : src === "assignment" ? "题面绑定"
+        : src === "predicted" ? "预测（无记录）" : "子任务门禁");
+  const cells = [];
+  for (let gpu = 0; gpu < gpuCount; gpu += 1) cells.push(gpu);
+  return html`
+    

+      
四卡一览 （第 ${data && data.iteration} 轮 · 一卡一题 · 数字都来自子任务自己的记录）

+      

+        ${cells.map((gpu) => {
+          const card = byGpu.get(gpu);
+          const questions = (card && card.questions) || [];
+          return html`
+            

+              

+                GPU${gpu}
+                ${questions.length ? `${questions.length} 题` : "本轮未派题"}
+              

+              ${questions.length === 0
+                ? html`
空闲
`
+                : questions.map((q) => {
+                    const a = q.gate_audit || {};
+                    const g = q.gate || {};
+                    const secs = quiet(q);
+                    // The bar is the round's real gate: the operator's current
+                    // variant (legacy rounds kept the family τ×baseline target,
+                    // which is what judged them).
+                    const target = g.target_us;
+                    const onTarget = (q.best_median_us != null && target != null)
+                      ? q.best_median_us <= target : null;
+                    return html`
+                      

+                        

+                          
+                          ${statusChip(q)}
+                        

+                        

+                          
阶段

+                          
${DKAO_PHASE_LABEL[q.current_phase] || q.current_phase || "idle"}

+                          
轮次

+                          
iter ${q.current_iteration ?? "-"} · ${q.rounds || 0} 次测量

+                          
当前

+                          
${q.best_median_us != null
+                                ? html`${fmt(q.best_median_us)} µs
+                                    ${target != null
+                                      ? html`
+                                          / 必须打败 ≤ ${fmt(target)}${onTarget === null ? "" : (onTarget ? " ✓" : " ✗")}`
+                                      : null}`
+                                : html`尚未测出`}

+                          
闸门

+                          
${a.checks == null || a.checks === 0
+                                ? html`暂无记录`
+                                : html`${a.blocked ? html`等待 ${a.blocked} 次` : html`已通过`}
+                                    ${secs != null ? html` · ${secs < 60 ? `${secs}s` : `${Math.round(secs / 60)}min`} 前更新` : null}`}

+                          
来源

+                          
${sourceLabel(q.gpu_source)}

+                        

+                      
`;
+                  })}
+            
`;
+        })}
+      

+    

+  `;
+}
+
+function PromotionPanel({ promo, busy, onApprove, onDeny }) {
+  if (!promo) return null;
+  const pending = promo.pending;
+  const promoted = promo.promoted;
+  const table = promo.variant_table || {};
+  const gateLine = (gate, need) => {
+    if (!gate) return html`—`;
+    const wins = (gate.wins || []).length;
+    const ops = (gate.operators || []).length;
+    const ok = gate.status === "PASS";
+    return html`${gate.status}
+      ${wins}/${ops} win（需 ${need}）`;
+  };
+  return html`
+    

+      
Harness promotion （接入 DKAO 需要人工批准）

+      ${pending ? html`
+        

+          

+            候选版本 ${pending.version}
+            · 来自第 ${pending.iteration} 轮
+            · revision ${(pending.revision || "-").slice(0, 12)}
+          

+          
+            
+              
+              
+              
+                  
+            
+          
性能门${gateLine(pending.performance_gate, 3)}
泛化门${gateLine(pending.generalization_gate, 2)}
机制门（宽松，仅记录）${(pending.mechanism_gate && pending.mechanism_gate.status) || "NOT_RECORDED"}

+          

+            
+            
+          

+        
` : html`
+        
当前没有待批准候选。
`}
+      

+        生产版本：${(promoted && promoted.version) || "h-1（种子）"}
+        ${promoted && promoted.approved_at ? html` · 于 ${new Date(promoted.approved_at * 1000).toLocaleString()}` : null}
+        · variant 表：${table.operators || 0} 个算子 / 版本 ${table.harness_version || "未发布"}
+      

+    

+  `;
+}
+
+// One iteration's device cards, read-only: HE no longer leases cards — each
+// question is a DKAO child and the child's own admission gate decides when the
+// device it was handed to is clean enough to time a kernel. Folded twice on
+// purpose: the section folds (default collapsed) and so does every card and
+// every per-question detail, so a 4-card round stays one scan line until asked.
+function GpuCardsPanel({ data, num }) {
+  const cards = (data && data.cards) || [];
+  const nQuestions = cards.reduce((acc, c) => acc + ((c.questions || []).length), 0);
+  const fmt = (v) => (v == null ? "-" : (v.toFixed ? v.toFixed(2) : v));
+  return html`
+    

+      
+        GPU 设备
+        
+          — 第 ${num} 轮 · ${cards.length} 张卡上有 ${nQuestions} 个 DKAO 子任务
+          · 本任务不持有租约，卡由 DKAO 的准入闸门管理
+        
+      
+      ${!cards.length
+        ? html`
本轮尚未派发任何 DKAO 子任务（无设备卡片）
`
+        : html`

+            ${cards.map((card) => html`
+              

+                
+                  GPU${card.gpu}
+                  
+                    ${(card.questions || []).map((q) => q.question).join(", ") || "空闲"}
+                    · ${(card.questions || []).length} 个问题
+                  
+                
+                ${(card.questions || []).map((q) => {
+                  const a = q.gate_audit || {};
+                  const g = q.gate || {};
+                  const p = q.pool_reference || {};
+                  const req = q.requirements;
+                  // Where this card membership came from, most authoritative
+                  // first: HE's layout record for the round, the device HE
+                  // pinned in the child's requirements, the child's own gate —
+                  // and only as a last resort the index % 4 prediction.
+                  const predicted = q.gpu_source === "predicted";
+                  const SOURCE_LABEL = {
+                    layout: "该卡来自本轮布局记录（HE 派题时选定）",
+                    assignment: "该卡来自子任务 requirements 的绑定",
+                    gate_pass: "该卡来自子任务自己的门禁记录",
+                    gate_blocked: "该卡来自子任务自己的门禁记录（最近一次被阻塞）",
+                    predicted: "布局预测：index % 4，尚未有门禁记录",
+                  };
+                  // A gate_ok row also carries reasons — they describe the
+                  // candidates that failed, not this device — so the wording
+                  // keys off the verdict, never off a non-empty reasons list.
+                  const blockedRow = a.last_event === "gate_blocked";
+                  const passing = (a.last_passing || []).map((c) => `GPU${c}`);
+                  const busy = (a.last_candidates || [])
+                    .filter((c) => !(a.last_passing || []).includes(c))
+                    .map((c) => `GPU${c}`);
+                  const gateLine = a.checks === 0 || a.checks == null
+                    ? `准入闸门：暂无记录（VRAM≤90% 且 HCU==0）`
+                    : blockedRow
+                      ? `准入闸门：共 ${a.checks} 次检查，其中 ${a.blocked ?? 0} 次因设备被占用而等待 — 最近一次 ${a.last_event}@${a.last_site || "-"} GPU${a.last_gpu ?? "-"}：${(a.last_reasons || []).join("; ") || "-"}`
+                      : `准入闸门：共 ${a.checks} 次检查，最近一次检查通过`
+                        + (passing.length ? `（可用卡：${passing.join("、")}`
+                            + (busy.length ? `；检查时 ${busy.join("、")} 被占用` : "") + `）` : "");
+                  return html`
+                    

+                      
+                        ${q.question}
+                        
+                          · ${DKAO_PHASE_LABEL[q.current_phase] || q.current_phase || "idle"}
+                          · ${SOURCE_LABEL[q.gpu_source] || "该卡来源未知"}
+                          · iter ${q.current_iteration ?? "-"}
+                          · ${q.rounds || 0} rounds
+                          · best ${q.best_median_us != null ? `${fmt(q.best_median_us)} us` : "-"}
+                          ${q.finished ? ` · 已结束${q.final_status ? `（${q.final_status}）` : ""}` : ""}
+                        
+                      
+                      

+                        
${gateLine}

+                        

+                          ${a.probe_device != null && a.device != null && a.probe_device !== a.device
+                            ? `探针曾在 GPU${a.probe_device} 上通过（与本题所在卡无关）；` : ""}
+                          ${g.protocol === "legacy"
+                            ? html`legacy 判据：baseline ${p.baseline_us != null ? `${fmt(p.baseline_us)} us` : "-"}
+                                · tau ${p.tau_family != null ? fmt(p.tau_family) : "-"}
+                                · target ≤ ${g.target_us != null ? `${fmt(g.target_us)} us` : "-"}`
+                            : html`性能门对手（variant）：${g.variant_us != null ? `${fmt(g.variant_us)} us` : "-"}
+                                （${variantSourceLabel(g)}）
+                                ${g.judged
+                                  ? html` · 本轮实测 ${fmt(g.judged.candidate_median_us)} us
+                                          <${GateVerdictBadge} verdict=${g.judged.verdict} />`
+                                  : null}`}
+                          ${p.target_us != null
+                            ? html` · 池子口径 τ×baseline = ${fmt(p.target_us)} us（本协议不判门）`
+                            : null}
+                          ${a.waited_seconds && (blockedRow || a.checks == null)
+                            ? ` · 累计等待 ${Math.round(a.waited_seconds / 60)} min` : ""}
+                        

+                      

+                      ${(a.blocked_reasons || []).length ? html`
+                        

+                          阻塞历史（最近 ${(a.blocked_reasons || []).length} 次）
+                          
+                            
+                            
+                              ${(a.blocked_reasons || []).map((b) => html`
+                                
+                                
+                                
+                                
+                                
+                              `)}
+                            
+                          
when站点GPU第几次原因
${b.ts ? new Date(b.ts * 1000).toLocaleTimeString() : "-"}${b.site || "-"}${b.gpu ?? "-"}${b.attempt ?? "-"}${(b.reasons || []).join("; ") || "-"}

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

+                          子任务 requirements
+                          
${JSON.stringify(req, null, 2)}

+                        
` : null}
+                      ${q.child_state_dir ? html`
+                        

+                          child_state_dir
+                          
${q.child_state_dir}

+                        
` : null}
+                    
`;
+                })}
+              
`)}
+          
`}
+    
`;
+}
+
+// --------------------------------------------------------------------------- //
+// Round picker for the DKAO blocks
+// --------------------------------------------------------------------------- //
+// Each round is an independent DKAO launch, so which round you look at is a
+// real switch. Default is "follow the newest round"; picking an older round
+// pins it (and the 5s polling stops, because history does not move).
+function RoundControls({ roundN, latestNum, roundNums, selIterN, setSelIterN,
+                         ids, selChild, childrenMap, onSelectChild, onRefresh }) {
+  return html`
+    

+      
+      
+      
+      ${ids && ids.length ? html`
+        
+        
+        ` : null}
+    

+    ${selIterN == null ? null : html`

+      正在查看历史轮次 iteration ${roundN}（本块与上面的 Live questions 都停在这一轮，不再自动刷新）
`}
+  `;
+}
+
+export default function HeDetailView({ taskId, run, status, data }) {
+  const [summary, setSummary] = useState(null);
+  const [iters, setIters] = useState([]);
+  const [graph, setGraph] = useState(null);
+  const [children, setChildren] = useState(null);
+  const [selChild, setSelChild] = useState(null);
+  const [dkao, setDkao] = useState(null);
+  const [resumeMsg, setResumeMsg] = useState(null);
+  const [sup, setSup] = useState(null);
+  const [review, setReview] = useState(null);
+  const [iterGpu, setIterGpu] = useState(null);
+  const [targetIters, setTargetIters] = useState(null);
+  const [baseIter, setBaseIter] = useState("");        // "" = auto (best_ever)
+  const [promo, setPromo] = useState(null);
+  const [promoBusy, setPromoBusy] = useState(false);
+  //: Which round's DKAO children the drill-down shows. Rounds are independent
+  //: DKAO launches, so this is a real switch, not a filter: picking round 5
+  //: re-reads that round's children (each has its own state dir).
+  //: ``null`` = follow the newest round until the operator picks one.
+  const [selIterN, setSelIterN] = useState(null);
+
+  const latestNum = useMemo(() => {
+    if (!iters.length) return 0;
+    const m = (iters[iters.length - 1].name || "").match(/(\d+)$/);
+    return m ? Number(m[1]) : 0;
+  }, [iters]);
+
+  // Every round the run has artifacts for, oldest first.
+  const roundNums = useMemo(
+    () => iters
+      .map((it) => {
+        const m = (it.name || "").match(/(\d+)$/);
+        return m ? Number(m[1]) : 0;
+      })
+      .filter((n) => n > 0),
+    [iters],
+  );
+  // The round actually on screen. Following means "the newest one", and the
+  // selection is clamped to what exists so a stale choice cannot outlive its
+  // round (Reset, or a fresh run).
+  const roundN = useMemo(() => {
+    if (!selIterN) return latestNum;
+    if (!roundNums.length || roundNums.includes(selIterN)) return selIterN;
+    return latestNum;
+  }, [selIterN, latestNum, roundNums]);
+  const following = roundN === latestNum;
+
+  const refresh = useCallback(async () => {
+    if (!taskId) return;
+    const [s, its, g] = await Promise.all([
+      withTimeout(getSummary(taskId)).catch((e) => { console.warn("he summary:", e); return null; }),
+      withTimeout(getIterations(taskId)).catch((e) => { console.warn("he iterations:", e); return []; }),
+      withTimeout(getStateGraph(taskId)).catch((e) => { console.warn("he graph:", e); return null; }),
+    ]);
+    setSummary(s || null);
+    setIters(its || []);
+    setGraph(g || null);
+  }, [taskId]);
+
+  const refreshPromotion = useCallback(async () => {
+    if (!taskId) return;
+    const p = await withTimeout(getPromotion(taskId))
+      .catch((e) => { console.warn("he promotion:", e); return null; });
+    if (p) setPromo(p);
+  }, [taskId]);
+
+  const approvePromo = useCallback(async () => {
+    if (!taskId || promoBusy) return;
+    setPromoBusy(true);
+    try {
+      const r = await approvePromotion(taskId, "webui");
+      setResumeMsg(r && r.ok ? `已批准并接入：${r.result && r.result.version}`
+                             : "批准被拒绝（见日志）");
+    } catch (e) {
+      setResumeMsg(`批准失败：${e.message}`);
+    } finally {
+      setPromoBusy(false);
+      refreshPromotion();
+      refresh();
+    }
+  }, [taskId, promoBusy, refreshPromotion, refresh]);
+
+  const denyPromo = useCallback(async () => {
+    if (!taskId || promoBusy) return;
+    setPromoBusy(true);
+    try {
+      const r = await denyPromotion(taskId, "webui");
+      setResumeMsg(r && r.ok ? "已拒绝该候选，可继续迭代" : "拒绝失败（见日志）");
+    } catch (e) {
+      setResumeMsg(`拒绝失败：${e.message}`);
+    } finally {
+      setPromoBusy(false);
+      refreshPromotion();
+    }
+  }, [taskId, promoBusy, refreshPromotion]);
+
+  const refreshChildren = useCallback(async () => {
+    if (!taskId) return;
+    if (roundN < 1) { setChildren(null); return; }
+    const c = await withTimeout(getChildren(taskId, roundN))
+      .catch((e) => { console.warn("he children:", e); return null; });
+    if (c) setChildren(c);
+  }, [taskId, roundN]);
+
+  useEffect(() => { refresh(); }, [refresh]);
+  pollWhileVisible(refresh, 6000);
+  useEffect(() => { refreshChildren(); }, [refreshChildren]);
+  pollWhileVisible(refreshPromotion, 10000);
+  // Only the newest round is still moving, and only until its children are
+  // done: a round the operator pinned (or a round whose questions have all
+  // finished) is history and is fetched once, on selection. This also removes
+  // the duplicate: the drill-down below keeps its own children view, and both
+  // used to poll the same URL.
+  const roundMoves = useMemo(() => {
+    const kids = Object.values(children || {});
+    return kids.length === 0 || kids.some((c) => !c.finished);
+  }, [children]);
+  pollWhileVisible(refreshChildren, 5000, {
+    enabled: Boolean(taskId) && following && latestNum >= 1 && roundMoves,
+  });
+
+  const best = summary && summary.best_ever;
+  const lastIter = iters.length ? iters[iters.length - 1] : null;
+  // Benchmark results belong to the round on screen, exactly like its children
+  // do: the two blocks below must never describe different rounds.
+  const selIterObj = useMemo(
+    () => iters.find((it) => {
+      const m = (it.name || "").match(/(\d+)$/);
+      return m && Number(m[1]) === roundN;
+    }) || null,
+    [iters, roundN],
+  );
+  const results = (selIterObj && selIterObj.benchmark
+                   && selIterObj.benchmark.results) || {};
+  // Show every live DKAO child even before its benchmark row exists: child
+  // run.json phases appear as soon as evaluation spawns the subtasks, while
+  // benchmark results only materialize after a child finishes.
+  const childrenMap = children || {};
+  const liveIds = [...new Set([...Object.keys(results), ...Object.keys(childrenMap)])];
+  const liveKey = liveIds.join("|");
+  const liveRows = liveIds.map((id) => ({
+    id,
+    r: results[id] || null,
+    child: childrenMap[id] || null,
+  }));
+
+  const loadChild = useCallback(async (id) => {
+    if (!id || roundN < 1) return;
+    setSelChild(id);
+    const p = await withTimeout(getDkaoChild(taskId, roundN, id))
+      .catch((e) => { console.warn("he dkao child:", e); return null; });
+    setDkao(p || null);
+  }, [taskId, roundN]);
+
+  // Default to a still-running child so the DKAO drill-down has content.
+  useEffect(() => {
+    if (!liveIds.length || (selChild && liveIds.includes(selChild))) return;
+    const running = liveIds.find((id) => childrenMap[id] && !childrenMap[id].finished);
+    loadChild(running || liveIds[0]);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [liveKey, selChild]);
+  // Same rule as the child list: poll only while the shown round is the newest
+  // one and the question it shows can still move. ``loadChild`` also sets the
+  // selection, so the interval refreshes the payload in place.
+  pollWhileVisible(() => loadChild(selChild), 5000, {
+    enabled: Boolean(taskId) && roundN >= 1 && following && Boolean(selChild)
+      && Boolean(children && children[selChild] && !children[selChild].finished),
+  });
+
+  useEffect(() => {
+    if (!taskId) return;
+    let stop = false;
+    const tick = async () => {
+      const payload = await withTimeout(getSupervisor(taskId))
+        .catch(() => null);
+      if (!stop && payload) setSup(payload);
+    };
+    tick();
+    const id = setInterval(tick, 6000);
+    return () => { stop = true; clearInterval(id); };
+  }, [taskId]);
+
+  useEffect(() => {
+    if (!taskId) return;
+    let stop = false;
+    const tick = async () => {
+      const payload = await withTimeout(getReview(taskId)).catch(() => null);
+      if (!stop && payload) setReview(payload);
+    };
+    tick();
+    const id = setInterval(tick, 15000);
+    return () => { stop = true; clearInterval(id); };
+  }, [taskId]);
+
+  // Read-only device cards of the selected iteration (HE holds no lease).
+  useEffect(() => {
+    if (!taskId || latestNum < 1) { setIterGpu(null); return; }
+    let stop = false;
+    const tick = async () => {
+      const payload = await withTimeout(getIterationGpu(taskId, latestNum))
+        .catch(() => null);
+      if (!stop && payload) setIterGpu(payload);
+    };
+    tick();
+    const id = setInterval(tick, 10000);
+    return () => { stop = true; clearInterval(id); };
+  }, [taskId, latestNum]);
+
+  const startAuto = async () => {
+    if (!taskId) return;
+    const want = Math.max(1, Math.min(50,
+      Number(targetIters) || (iters.length + 1)));
+    setResumeMsg("starting unattended mode…");
+    try {
+      const out = await startAutopilot(taskId, want,
+        baseIter === "" ? null : Number(baseIter), true);
+      setResumeMsg(`自动模式已启动（supervisor pid ${out.supervisor_pid}，目标 ${out.target} 轮）`);
+      refresh();
+    } catch (e) {
+      setResumeMsg(`自动模式启动失败：${e.message}`);
+    }
+  };
+  const stopAuto = async () => {
+    if (!taskId) return;
+    try {
+      await stopSupervisor(taskId);
+      setResumeMsg("已请求停止自动模式");
+      const payload = await getSupervisor(taskId).catch(() => null);
+      if (payload) setSup(payload);
+    } catch (e) {
+      setResumeMsg(`停止失败：${e.message}`);
+    }
+  };
+
+  const applyTarget = async () => {
+    if (!taskId) return;
+    const want = Math.max(1, Math.min(50, Number(targetIters) || (iters.length + 1)));
+    setResumeMsg("applying…");
+    try {
+      const out = await resumeExperiment(taskId, want, null,
+                                        baseIter === "" ? null : Number(baseIter));
+      if (out.running) {
+        setResumeMsg(`目标已更新为 ${out.target} 轮（当前运行中，下一轮生效）`);
+      } else {
+        setResumeMsg(`已启动续跑（pid ${out.pid}），目标 ${out.target ?? want} 轮`);
+      }
+      refresh();
+    } catch (e) {
+      setResumeMsg(`操作失败：${e.message}`);
+    }
+  };
+  const isFinished = Boolean(graph && graph.current === "finished");
+  const [stopBusy, setStopBusy] = useState(false);
+  // The orchestrator process is what "running" means here; a stopped run keeps
+  // its artifacts and its phase, so the graph alone cannot tell them apart.
+  // The shell's view comes from orchestrator.pid, which runs started before
+  // that file existed never wrote -- so trust HE's own answer too, or Stop
+  // would be hidden on a run that is very much alive.
+  const isRunning = Boolean((status && status.running) || (summary && summary.running));
+
+  const stopRun = async () => {
+    if (!taskId || stopBusy) return;
+    setStopBusy(true);
+    setResumeMsg("正在停止…");
+    try {
+      const out = await stopExperiment(taskId);
+      setResumeMsg(out.message + (out.stopped_children && out.stopped_children.length
+        ? `（同时停掉 ${out.stopped_children.length} 个 DKAO 子任务）` : ""));
+      refresh();
+      refreshPromotion();
+    } catch (e) {
+      setResumeMsg(`停止失败：${e.message}`);
+    } finally {
+      setStopBusy(false);
+    }
+  };
+
+  return html`
+    

+      <${ProtocolBanner} protocol=${summary && summary.protocol} />
+      <${PromotionPanel} promo=${promo} busy=${promoBusy}
+        onApprove=${approvePromo} onDeny=${denyPromo} />
+      ${(summary && summary.operator_stop && !isRunning) ? html`
+        

+          
⏸ 已被 operator 停止 （暂停，不是结束）

+          

+            停在 iteration ${summary.operator_stop.iteration ?? "-"}（phase ${summary.operator_stop.phase || "-"}）·
+            不会被 janitor / supervisor 自动拉起；本轮没判完，所以这一轮的成绩不计入 variant。
+            继续跑：点左边的「续跑到该轮次上限」（会先清掉停止标记，该轮重跑、抽到的是同一批题）。
+            操作说明见 FLOW.md §9。
+          

+        
` : null}
+      <${StateMachine} graph=${graph} />
+      

+        
+         setTargetIters(e.target.value)} />
+        
+        ${isRunning ? html`
+          
+          停止=暂停：产物保留、不会被自动重启；之后点左边的「续跑到该轮次上限」继续` : null}
+        
+        
+        已完成 ${iters.length} 轮${graph && graph.current === "finished" ? "（已停止）" : "（运行中）"} · 运行中可随时改，下一轮生效
+        ${resumeMsg ? html`${resumeMsg}` : null}
+        ${isFinished
+          ? html`复用已评测的冠军，只评测新一版 harness（不重跑历史迭代）`
+          : null}
+      

+      

+        
+        
+        
+          自动模式：${sup && sup.running ? `运行中（pid ${sup.pid}）` : ((sup && sup.state && sup.state.status) || "未启动")}
+          ${sup && sup.incidents && sup.incidents.length ? ` · 事故 ${sup.incidents.length}` : ""}
+          ${sup && sup.has_report ? " · 诊断报告已生成" : ""}
+        
+      

+      <${FourCardOverview} data=${iterGpu} selected=${selChild}
+        onSelect=${(id) => setSelChild(id)} />
+      <${GpuCardsPanel} data=${iterGpu} num=${latestNum} />
+      ${review && review.review ? html`
+        

+          
跨轮复核 （每代 harness 成绩 · 冠军代 · best-known）

+          

+            冠军代 iteration ${review.review.champion_iteration}
+            variant 策略：${review.review.variant_policy}
+              ${(review.review.final_promotion || {}).promoted
+                ? ` · 收尾晋升：${(review.review.final_promotion.promoted || []).join(", ")}` : ""}
+              ${review.regressions && review.regressions.length
+                ? ` · ⚠ 退步 ${review.regressions.length} 次` : ""}
+            
+          

+          
+            
+            
+              ${(review.review.generations || []).map((g) => html`
+                
+                  
+                  
+                  
+                  
+                `)}
+            
+          
genverdictpassmedians (us)
${g.iteration}${g.champion ? " ⭐" : ""}${g.rejected ? " ✗" : ""}${g.verdict || "-"}${g.pass_count}${Object.entries(g.medians || {}).map(([k, v]) =>
+                    `${k.split("_").slice(-3).join("_")}=${v == null ? "-" : Number(v).toFixed(1)}`).join("  ")}

+          ${review.regressions && review.regressions.length ? html`
+            

+              退步记录（kernel 低于历史最优）
+              
+                
+                
+                  ${review.regressions.slice(-8).reverse().map((r) => html`
+                    
+                    
+                    `)}
+                
+              
iterationshapemedianbestdelta
${r.iteration}${r.shape}${r.median_us}${r.best_median_us} (gen ${r.best_iteration})+${r.delta_percent}%

+            
` : null}
+        
` : null}
+      ${sup && sup.incidents && sup.incidents.length ? html`
+        

+          
Incidents （supervisor 自动处置记录）

+          
+            
+            
+              ${sup.incidents.slice(-8).reverse().map((i) => html`
+                
+                  
+                  
+                  
+                  
+                `)}
+            
+          
whencategoryactiondetail
${new Date((i.ts || 0) * 1000).toLocaleTimeString()}${i.category || "-"}${i.action || "-"}${i.reason || (i.evidence || []).join("; ") || ""}
+                    ${i.seconds ? html`wait ${Math.round(i.seconds / 60)}m` : null}

+        
` : null}
+      

+        
Best ever

+          
${best ? `${(100 * best.pass_rate).toFixed(1)}%` : "-"}

+          ${best ? html`
iteration ${best.iteration}
` : null}
+        

+        
Iterations

+          
${iters.length}

+        
Mode

+          
${(summary && summary.config_snapshot && summary.config_snapshot.execution_mode) || "-"}

+        

+      

+      ${liveRows.length ? html`
+        

+          
Live questions — iteration ${roundN}
+            （每行 = 一个真实 DKAO 子任务）

+          <${RoundGateLine} iter=${selIterObj} />
+          
+            
+              
+              
+              
+              
+              
+              
+            
+            
+              ${liveRows.map(({ id, r, child }) => html`
+                
+                  
+                  
+                  
+                  
+                  
+                  
+                `)}
+            
+          
questionDKAO child phase子任务验收median(best) us必须打败（variant）rounds
${id}<${DkaoPhaseChip} child=${child} />
+                    ${!r ? html`running`
+                      : r.passed ? html`pass`
+                      : r.status === "exception" || r.status === "missing_report" || r.status === "partial_worker_result"
+                        ? html`${r.status}`
+                        : html`${r.status || "pending"}`}
+                  ${(r && r.median_us != null) ? r.median_us
+                        : (child && child.best_median_us != null) ? `≈${child.best_median_us}` : "-"}<${QuestionGateCell}
+                        gate=${child && child.gate}
+                        poolRef=${child && child.pool_reference} />${(r && r.rounds_used != null) ? r.rounds_used : (child ? (child.rounds || 0) : 0)}

+        

+        

+          
DKAO task drill-down
+            （选轮次后看那一轮的 DKAO 子任务；只读）

+          <${RoundControls}
+            roundN=${roundN} latestNum=${latestNum} roundNums=${roundNums}
+            selIterN=${selIterN} setSelIterN=${setSelIterN}
+            ids=${liveIds} selChild=${selChild} childrenMap=${childrenMap}
+            onSelectChild=${loadChild} onRefresh=${loadChild} />
+          <${ChildDkaoView} payload=${dkao} />
+        
` : html`
+        

+          
DKAO task drill-down
+            （选轮次后看那一轮的 DKAO 子任务；只读）

+          <${RoundControls}
+            roundN=${roundN} latestNum=${latestNum} roundNums=${roundNums}
+            selIterN=${selIterN} setSelIterN=${setSelIterN}
+            ids=${null} selChild=${null} childrenMap=${{}}
+            onSelectChild=${loadChild} onRefresh=${loadChild} />
+          
这一轮没有派发 DKAO 子任务（或产物已被清理）。

+        
`}
+      

+        
Pass rate per iteration

+        ${summary && html`<${ScoreTable} scores=${summary.scores} />`}
+      

+      

+        
Iterations

+        ${iters.length ? iters.map((it) => html`<${IterBlock} iter=${it} taskId=${taskId} />`)
+                        : html`
no iterations yet (task not started / dry artifacts absent)
`}
+        ${summary && summary.report ? html`
report.md
+          
${summary.report}
` : null}
+      

+    
`;
+}
diff --git a/metainfer/tasks/harness_evolve/static/he.css b/metainfer/tasks/harness_evolve/static/he.css
new file mode 100644
index 00000000..18632b2d
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/static/he.css
@@ -0,0 +1,166 @@
+/* harness_evolve detail view — mirrors the DKAO look & feel.
+   Reuses the dkao state-machine classes (copied into this plugin dir) and the
+   shell theme variables so the AHE outer loop reads like a DKAO task page. */
+@import url("./dkao.css");
+
+.he-detail {
+  display: flex;
+  flex-direction: column;
+  gap: var(--sp-4, 10px);
+}
+
+/* summary header cards */
+.he-cards { display: flex; flex-wrap: wrap; gap: var(--sp-3, 8px); }
+.he-card {
+  flex: 1 1 160px;
+  min-width: 160px;
+  border: 1px solid var(--border, #ddd);
+  border-radius: var(--radius, 8px);
+  padding: var(--sp-3) var(--sp-4);
+  background: var(--bg, #fff);
+}
+.he-card h4 { margin: 0 0 4px; font-size: var(--fs-xs, 12px);
+  text-transform: uppercase; letter-spacing: .04em; color: var(--fg-dim, #666); }
+.he-card .big { font-size: 26px; font-weight: 700; color: var(--fg, #111); }
+.he-card .sub { color: var(--fg-dim, #666); font-size: var(--fs-xs, 12px); }
+
+/* generic content panels */
+.he-section {
+  border: 1px solid var(--border, #ddd);
+  border-radius: var(--radius, 8px);
+  padding: var(--sp-3) var(--sp-4);
+  background: var(--bg, #fff);
+}
+.he-section > h3, .he-section > h4 { margin: 0 0 var(--sp-3); }
+.he-section h4 small { font-weight: 400; color: var(--fg-dim, #666); }
+
+/* tables */
+table.he-tbl { border-collapse: collapse; width: 100%; font-size: var(--fs-xs, 12px); }
+.he-tbl th { text-align: left; padding: 6px 8px; border-bottom: 1px solid var(--border, #ddd);
+  background: var(--panel-2, #f2f2f2); color: var(--fg-dim, #555); font-weight: 600; }
+.he-tbl td { padding: 6px 8px; border-bottom: 1px solid var(--border-soft, #eee);
+  vertical-align: top; }
+.he-tbl .mono { font-family: var(--fs-mono, monospace); font-size: var(--fs-xs, 12px); }
+
+/* badges (DKAO-colored) */
+.badge { display: inline-block; padding: 2px 8px; border-radius: var(--radius-sm, 6px);
+  font-size: var(--fs-xs, 11px); font-weight: 600; color: #fff; white-space: nowrap; }
+.badge.ok { background: var(--ok, #2e7d32); }
+.badge.bad { background: var(--err, #c62828); }
+.badge.warn { background: var(--warn, #ef6c00); }
+.badge.run { background: var(--accent, #1565c0); }
+
+/* details / pre */
+.he-iter summary { cursor: pointer; font-weight: 600; padding: 6px 0; }
+.he-pre { background: var(--bg-elev, #f6f6f6); border: 1px solid var(--border, #ddd);
+  border-radius: var(--radius, 8px); padding: var(--sp-3) var(--sp-4);
+  white-space: pre-wrap; font-family: var(--fs-mono, monospace);
+  font-size: var(--fs-xs, 12px); max-height: 320px; overflow: auto; }
+
+/* child drill-down row */
+.he-subrow { display: flex; align-items: center; gap: 10px; flex-wrap: wrap;
+  margin: 4px 0 10px; }
+.he-subrow label { font-size: var(--fs-xs, 12px); color: var(--fg-dim, #555); }
+.he-select { padding: 4px 8px; border: 1px solid var(--border, #ddd);
+  border-radius: var(--radius-sm, 6px); background: var(--bg, #fff);
+  color: var(--fg, #111); font-size: var(--fs-xs, 12px); max-width: 640px; }
+.he-btn { padding: 4px 10px; border: 1px solid var(--border, #ddd);
+  border-radius: var(--radius-sm, 6px); background: var(--panel-2, #eee);
+  color: var(--fg, #111); cursor: pointer; font-size: var(--fs-xs, 12px); }
+.he-child-dkao { display: flex; flex-direction: column; gap: var(--sp-4, 10px); }
+
+/* read-only device cards (GPU 设备) — folded section, folded card, folded detail */
+.he-gpu-panel > summary { cursor: pointer; padding: var(--sp-2, 6px) 0; }
+.he-gpu-panel > summary small { font-weight: 400; }
+.he-gpu-grid { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr));
+  gap: var(--sp-3, 8px); margin-top: var(--sp-2, 6px); }
+.he-gpu-card { min-width: 0; border: 1px solid var(--border-soft, #eee);
+  border-radius: var(--radius-sm, 6px); background: var(--bg-elev, #f6f6f6); }
+.he-gpu-card > summary { display: flex; align-items: baseline; justify-content: space-between;
+  gap: 10px; padding: 7px 10px; cursor: pointer; font-size: var(--fs-xs, 12px); }
+.he-gpu-card[open] > summary { border-bottom: 1px solid var(--border-soft, #eee); }
+.he-gpu-card > summary strong { font-family: var(--fs-mono, monospace); }
+.he-gpu-card > summary small { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.he-gpu-q { border-top: 1px solid var(--border-soft, #eee); }
+.he-gpu-card > .he-gpu-q:first-of-type { border-top: none; }
+.he-gpu-q > summary { display: flex; align-items: baseline; flex-wrap: wrap; gap: 6px;
+  padding: 6px 10px; cursor: pointer; font-size: var(--fs-xs, 12px); }
+.he-gpu-q > summary small { font-size: var(--fs-xxs, 11px); }
+.he-gpu-gate { display: flex; flex-direction: column; gap: 2px;
+  padding: 0 10px 6px; font-size: var(--fs-xs, 12px); }
+.he-gpu-sub { margin: 4px 10px 8px; font-size: var(--fs-xs, 12px); }
+.he-gpu-sub > summary { cursor: pointer; color: var(--fg-dim, #666); font-size: var(--fs-xxs, 11px); }
+.he-gpu-sub .mono { word-break: break-all; }
+.he-gpu-empty { display: grid; min-height: 38px; place-items: center;
+  font-size: var(--fs-xs, 12px); }
+
+/* live DKAO mirror inside one iteration — the DKAO page rendered read-only */
+.he-live-dkao { border-top: 1px solid var(--border-soft, #eee); padding-top: var(--sp-3, 8px); }
+.he-live-tabs { flex-wrap: wrap; gap: 6px; }
+.he-live-tab { display: inline-flex; align-items: center; gap: 6px; cursor: pointer;
+  padding: 5px 9px; border: 1px solid var(--border-soft, #ddd);
+  border-radius: var(--radius-sm, 6px); background: var(--bg-elev, #f6f6f6);
+  font-size: var(--fs-xs, 12px); }
+.he-live-tab.active { border-color: var(--accent, #2b6cb0); font-weight: 600; }
+.he-live-tab small { display: inline-flex; }
+.he-live-host { display: flex; align-items: baseline; flex-wrap: wrap; gap: 6px;
+  margin: var(--sp-2, 6px) 0; font-size: var(--fs-xs, 12px); }
+.he-live-host .mono { word-break: break-all; font-size: var(--fs-xxs, 11px); }
+.dkao-readonly-banner { border-left: 3px solid var(--accent, #2b6cb0); }
+
+/* Stop control: stopping is a pause, so it reads as a normal action that
+   happens to be destructive to the *process*, not to the artifacts. */
+.he-btn-stop { border-color: var(--warn, #b7791f); color: var(--warn, #b7791f); }
+.he-btn-stop:hover:not(:disabled) { background: var(--warn, #b7791f); color: #fff; }
+.he-btn-stop:disabled { opacity: 0.6; cursor: default; }
+
+/* four-card overview (四卡一览) — one column per HCU for this round's children.
+   Deliberately not `.he-cards`: that one is the flex summary-card row. */
+.he-four-cards { display: grid; grid-template-columns: repeat(4, minmax(0, 1fr));
+  gap: var(--sp-3, 8px); margin-top: var(--sp-2, 6px); }
+@media (max-width: 1200px) {
+  .he-four-cards { grid-template-columns: repeat(2, minmax(0, 1fr)); }
+}
+.he-card-col { min-width: 0; border: 1px solid var(--border-soft, #eee);
+  border-radius: var(--radius-sm, 6px); background: var(--bg-elev, #f6f6f6);
+  padding: 8px 10px; }
+.he-card-col.idle { opacity: .6; }
+.he-card-col > header { display: flex; align-items: baseline; justify-content: space-between;
+  gap: 8px; margin-bottom: 6px; }
+.he-card-col > header strong { font-family: var(--fs-mono, monospace); color: var(--accent-2, #2b6cb0); }
+.he-card-empty { display: grid; min-height: 34px; place-items: center; font-size: var(--fs-xs, 12px); }
+.he-card-q { border-top: 1px solid var(--border-soft, #eee); padding-top: 6px; margin-top: 6px; }
+.he-card-q:first-of-type { border-top: none; margin-top: 0; padding-top: 0; }
+.he-card-q.selected { outline: 1px solid var(--accent, #2b6cb0); }
+.he-card-q-head { display: flex; align-items: center; justify-content: space-between;
+  gap: 6px; margin-bottom: 4px; }
+.he-card-q-name { padding: 0; border: none; background: none; cursor: pointer;
+  font-family: var(--fs-mono, monospace); font-size: var(--fs-xs, 12px);
+  color: var(--fg, #111); text-align: left; word-break: break-all; }
+.he-card-q-name:hover { text-decoration: underline; }
+.he-card-metrics { display: grid; grid-template-columns: 34px minmax(0, 1fr);
+  gap: 2px 6px; margin: 0; font-size: var(--fs-xxs, 11px); }
+.he-card-metrics dt { color: var(--fg-dim, #666); }
+.he-card-metrics dd { margin: 0; min-width: 0; overflow-wrap: anywhere; }
+
+/* DKAO 实时视图：同时看四个子任务（一卡一题，四个独立的 DKAO 页面）。
+   默认四宫格；窄屏自动退化为单列。每个单元格是一个只读 DKAO 页面的容器。 */
+.he-live-layout { gap: 6px; }
+.he-live-layout .he-btn.active { border-color: var(--accent, #2b6cb0); font-weight: 600; }
+.he-live-grid { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr));
+  gap: var(--sp-3, 8px); align-items: start; margin-top: var(--sp-2, 6px); }
+.he-live-stack { display: grid; grid-template-columns: minmax(0, 1fr);
+  gap: var(--sp-3, 8px); margin-top: var(--sp-2, 6px); }
+@media (max-width: 1500px) { .he-live-grid { grid-template-columns: minmax(0, 1fr); } }
+.he-live-cell { min-width: 0; border: 1px solid var(--border-soft, #eee);
+  border-radius: var(--radius-sm, 6px); background: var(--bg, #fff);
+  padding: 2px 8px 8px; }
+.he-live-cell.selected { border-color: var(--accent, #2b6cb0); }
+.he-live-cell-head { display: flex; align-items: baseline; flex-wrap: wrap; gap: 6px;
+  padding: 5px 0 6px; border-bottom: 1px solid var(--border-soft, #eee); }
+.he-live-cell-title { padding: 0; border: none; background: none; cursor: pointer;
+  font-size: var(--fs-xs, 12px); color: var(--fg, #111); text-align: left;
+  word-break: break-all; }
+.he-live-cell-title:hover { text-decoration: underline; }
+.he-live-cell-grow { flex: 1 1 auto; }
+.he-live-cell-body { min-width: 0; overflow-x: auto; }
diff --git a/metainfer/tasks/harness_evolve/tests/__init__.py b/metainfer/tasks/harness_evolve/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/metainfer/tasks/harness_evolve/tests/conftest.py b/metainfer/tasks/harness_evolve/tests/conftest.py
new file mode 100644
index 00000000..8df577db
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/conftest.py
@@ -0,0 +1,19 @@
+"""Shared fixtures for harness_evolve plugin route tests."""
+
+from __future__ import annotations
+
+import pytest
+from fastapi.testclient import TestClient
+
+from metainfer.testing import isolated_env  # noqa: F401 - re-export as fixture
+from metainfer.server import app as app_module
+
+
+@pytest.fixture
+def app(isolated_env):
+    return app_module.create_app()
+
+
+@pytest.fixture
+def client(app):
+    return TestClient(app)
diff --git a/metainfer/tasks/harness_evolve/tests/test_champion_record_failure.py b/metainfer/tasks/harness_evolve/tests/test_champion_record_failure.py
new file mode 100644
index 00000000..744b66b0
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_champion_record_failure.py
@@ -0,0 +1,135 @@
+"""Losing the champion bookkeeping must be loud, and must not lose the champion.
+
+``best_ever.json`` + the generation log decide what later rounds compare
+against. When that write failed it used to be swallowed into a timeline row
+(``champion_record_failed``) and nothing else — 26 of them in one run on
+2026-09-11, all ``AttributeError`` — so the run carried on against a champion
+chain nobody could vouch for.
+
+Now a failure is visible (stderr, i.e. the orchestrator log), keeps its
+traceback, and is followed by a rebuild of the champion from the generations
+already on disk.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from metainfer.tasks.harness_evolve.orchestrator import pipeline as P
+from metainfer.tasks.harness_evolve.orchestrator.champion_selection import (
+    Generation, record_generation,
+)
+from metainfer.tasks.harness_evolve.orchestrator.config import ExperimentConfig
+
+
+def _cfg(tmp_path: Path) -> ExperimentConfig:
+    cfg = ExperimentConfig(
+        task_id="he-champion", state_dir=tmp_path / "state",
+        workspace_dir=tmp_path / "exp", suite=[], answers={},
+        execution_mode="dry-run", per_round_budget=4, max_iterations=1,
+    )
+    cfg.state_dir.mkdir(parents=True, exist_ok=True)
+    cfg.exp_dir.mkdir(parents=True, exist_ok=True)
+    return cfg
+
+
+def _generation(exp: Path, iteration: int, median_us: float) -> Generation:
+    """A generation on disk, the way the loop leaves one behind.
+
+    ``load_generations`` reads each generation's measurements back from
+    ``runs/iteration_NNN/input/benchmark/results.json``, so a usable champion
+    needs both files — that is what the rebuild fallback replays.
+    """
+    generation = Generation(
+        iteration=iteration,
+        verdict="PROMOTE",
+        results={"op00": {"median_us": median_us, "passed": True}},
+        snapshot_dir=f"/tmp/snap-{iteration}",
+        workspace_revision=f"rev-{iteration}",
+        pass_count=1,
+    )
+    record_generation(exp, generation)
+    results_path = (exp / "runs" / f"iteration_{iteration:03d}" / "input"
+                    / "benchmark" / "results.json")
+    results_path.parent.mkdir(parents=True, exist_ok=True)
+    results_path.write_text(json.dumps(
+        {"iteration": iteration, "results": generation.results}),
+        encoding="utf-8")
+    return generation
+
+
+def _events(cfg: ExperimentConfig) -> list:
+    path = Path(cfg.state_dir) / "timeline.jsonl"
+    if not path.is_file():
+        return []
+    return [json.loads(line)["type"]
+            for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
+
+
+def test_a_healthy_record_writes_the_champion(tmp_path):
+    cfg = _cfg(tmp_path)
+    _generation(cfg.exp_dir, 1, median_us=11.0)
+
+    outcome = P._record_champion_generation(
+        cfg, cfg.exp_dir, iteration=1, decision={"verdict": "PROMOTE"},
+        decision_rate=1.0, snapshot=Path("/tmp/snap-1"),
+        live_workspace=cfg.workspace_dir, results={"op00": {"passed": True}},
+        current={})
+
+    assert outcome["recorded"] is True and outcome["recovered"] is False
+    best = json.loads((cfg.exp_dir / "best_ever.json").read_text())
+    assert best["iteration"] == 1
+    assert "champion_best_measured" in _events(cfg)
+
+
+def test_a_failed_record_is_loud_and_rebuilds_the_champion(tmp_path, monkeypatch,
+                                                           capsys):
+    cfg = _cfg(tmp_path)
+    # a previous, good generation is already on disk
+    _generation(cfg.exp_dir, 1, median_us=11.0)
+
+    def boom(_exp, _generation):
+        raise AttributeError("'NoneType' object has no attribute 'get'")
+
+    monkeypatch.setattr(P, "record_generation", boom)
+
+    outcome = P._record_champion_generation(
+        cfg, cfg.exp_dir, iteration=2, decision={"verdict": "PROMOTE"},
+        decision_rate=0.5, snapshot=Path("/tmp/snap-2"),
+        live_workspace=cfg.workspace_dir, results={"op00": {"passed": True}},
+        current={})
+
+    # never raises, and the error is on the record with its traceback
+    assert outcome["recorded"] is False and outcome["error"]
+    rows = [json.loads(line) for line in
+            (Path(cfg.state_dir) / "timeline.jsonl").read_text(
+                encoding="utf-8").splitlines() if line.strip()]
+    failed = [row for row in rows if row["type"] == "champion_record_failed"]
+    assert failed and "AttributeError" in failed[0]["payload"]["error"]
+    assert "NoneType" in failed[0]["payload"]["traceback"]
+    # ... and it is visible in the orchestrator log, not only in a JSONL row
+    assert "champion bookkeeping failed" in capsys.readouterr().err
+    # the champion chain survives: rebuilt from the generation on disk
+    assert outcome["recovered"] is True
+    best = json.loads((cfg.exp_dir / "best_ever.json").read_text())
+    assert best["iteration"] == 1 and best["verdict"] == "REBUILT_AFTER_ERROR"
+    assert "champion_rebuilt" in _events(cfg)
+
+
+def test_a_failure_with_nothing_to_rebuild_from_still_returns(tmp_path,
+                                                              monkeypatch):
+    """No usable generation on disk: report, keep the old champion, carry on."""
+    cfg = _cfg(tmp_path)
+    monkeypatch.setattr(P, "record_generation",
+                        lambda *a, **k: (_ for _ in ()).throw(RuntimeError("disk full")))
+    previous = {"iteration": 7, "verdict": "BEST_MEASURED"}
+
+    outcome = P._record_champion_generation(
+        cfg, cfg.exp_dir, iteration=8, decision={"verdict": "REJECT"},
+        decision_rate=0.0, snapshot=Path("/tmp/snap-8"),
+        live_workspace=cfg.workspace_dir, results={}, current=previous)
+
+    assert outcome["recorded"] is False and outcome["recovered"] is False
+    assert outcome["best_ever"] == previous
+    assert "champion_record_failed" in _events(cfg)
diff --git a/metainfer/tasks/harness_evolve/tests/test_champion_selection.py b/metainfer/tasks/harness_evolve/tests/test_champion_selection.py
new file mode 100644
index 00000000..9fc0b2ee
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_champion_selection.py
@@ -0,0 +1,74 @@
+"""Champion = best measured generation, not "last promoted generation"."""
+
+from __future__ import annotations
+
+from metainfer.tasks.harness_evolve.orchestrator.champion_selection import (
+    Generation, best_generation, compare_generations, load_generations,
+    record_generation, summarize,
+)
+
+
+def _gen(iteration, medians, *, verdict="PROMOTE", passed=None) -> Generation:
+    results = {}
+    for iid, median in medians.items():
+        ok = True if passed is None else bool(passed.get(iid, True))
+        results[iid] = {"median_us": median, "passed": ok,
+                        "correctness_ok": True}
+    return Generation(iteration=iteration, verdict=verdict, results=results,
+                      pass_count=sum(1 for r in results.values()
+                                     if r["passed"]))
+
+
+def test_better_generation_keeps_the_champion_slot():
+    """The reported scenario: gen1 40us (unpromoted), gen2 60us (promoted)."""
+    gen1 = _gen(1, {"a": 40.0, "b": 40.0}, verdict="CONFIRM_REQUIRED")
+    gen2 = _gen(2, {"a": 60.0, "b": 60.0}, verdict="PROMOTE")
+    assert compare_generations(gen1, gen2) == "a"
+    assert best_generation([gen1, gen2]).iteration == 1
+
+
+def test_rejected_generation_is_disqualified():
+    good = _gen(1, {"a": 40.0})
+    bad = _gen(2, {"a": 20.0}, verdict="REJECT")      # fast but regressed
+    assert best_generation([good, bad]).iteration == 1
+
+
+def test_no_signal_generation_does_not_displace_a_better_one():
+    keeper = _gen(1, {"a": 40.0}, verdict="PROMOTE")
+    flat = _gen(2, {"a": 41.0}, verdict="NO_SIGNAL")
+    assert best_generation([keeper, flat]).iteration == 1
+
+
+def test_more_passing_questions_wins_before_raw_speed():
+    fewer = _gen(1, {"a": 30.0, "b": 100.0}, passed={"a": True, "b": False})
+    more = _gen(2, {"a": 32.0, "b": 60.0}, passed={"a": True, "b": True})
+    assert best_generation([fewer, more]).iteration == 2
+
+
+def test_earlier_generation_wins_a_tie():
+    first = _gen(1, {"a": 50.0})
+    second = _gen(2, {"a": 50.5})                     # within the noise band
+    assert best_generation([first, second]).iteration == 1
+
+
+def test_round_trip_persistence(tmp_path):
+    exp = tmp_path / "exp"
+    (exp / "runs" / "iteration_001" / "input" / "benchmark").mkdir(parents=True)
+    (exp / "runs" / "iteration_001" / "input" / "benchmark"
+     / "results.json").write_text('{"results": {"a": {"median_us": 40.0, "passed": true}}}',
+                                  encoding="utf-8")
+    record_generation(exp, _gen(1, {"a": 40.0}, verdict="CONFIRM_REQUIRED"))
+    gens = load_generations(exp)
+    assert len(gens) == 1
+    assert gens[0].iteration == 1 and gens[0].pending is True
+    assert gens[0].results["a"]["median_us"] == 40.0
+    assert best_generation(gens).iteration == 1
+
+    # recording the same iteration again replaces the row
+    record_generation(exp, _gen(1, {"a": 35.0}, verdict="PROMOTE"))
+    gens = load_generations(exp)
+    assert len(gens) == 1 and gens[0].verdict == "PROMOTE"
+
+    rows = summarize(gens)
+    assert rows[0]["champion"] is True
+    assert rows[0]["medians"] == {"a": 40.0}
diff --git a/metainfer/tasks/harness_evolve/tests/test_config_req_contract.py b/metainfer/tasks/harness_evolve/tests/test_config_req_contract.py
new file mode 100644
index 00000000..5db5a04d
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_config_req_contract.py
@@ -0,0 +1,100 @@
+"""load_experiment_config must honor the WebUI/shell requirements contract.
+
+The shell/server flattens form answers onto the requirements top level (the
+same contract dcu_kernel_auto_opt consumes); the legacy CLI/tests path passes
+them under an ``answers`` dict. Both must yield the same ExperimentConfig so a
+WebUI-launched harness-evolve experiment actually runs dkao-cli against the
+registered pool instead of silently falling back to dry-run defaults.
+"""
+
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from metainfer.tasks.harness_evolve.orchestrator.config import (
+    load_experiment_config,
+)
+
+_FLAT_FIELDS = {
+    "task_id": "he-webui-task",
+    "task_type": "harness-evolve",
+    "label": "webui task",
+    "raw_request": "",
+    "pool_source": "/root/zth_agent/ahe-kernel-repos/registered_pool.yaml",
+    "per_round_budget": 4,
+    "default_rounds": 5,
+    "max_iterations": 1,
+    "harness_source": "dcu_default",
+    "execution_mode": "dkao-cli",
+    "ahe_repo_root": "/root/zth_agent/ahe-kernel-repos",
+    "child_timeout_minutes": "0",
+    "evolve_mode": "agent",
+    "evolve_model": "deepseek/deepseek-flash",
+    "evolve_timeout_seconds": "1800",
+    "agent_framework": "dsh",
+}
+
+
+def _write_req(tmp_path, payload) -> "object":
+    req = tmp_path / "requirements.json"
+    req.write_text(json.dumps(payload), encoding="utf-8")
+    return req
+
+
+def test_flat_top_level_requirements_are_honored(tmp_path):
+    """WebUI form answers flattened onto req top level must not be ignored."""
+    cfg = load_experiment_config(
+        _write_req(tmp_path, _FLAT_FIELDS),
+        tmp_path / "state",
+        tmp_path / "ws",
+    )
+    assert cfg.task_id == "he-webui-task"
+    assert cfg.execution_mode == "dkao-cli"
+    assert cfg.pool_source == _FLAT_FIELDS["pool_source"]
+    assert cfg.pool_mode is True
+    assert cfg.per_round_budget == 4
+    assert cfg.max_iterations == 1
+    assert cfg.harness_source == "dcu_default"
+    assert cfg.agent_framework == "dsh"
+    assert cfg.answers.get("evolve_mode") == "agent"
+    assert cfg.answers.get("default_rounds") == 5
+
+
+def test_answers_dict_requirements_still_honored(tmp_path):
+    """Legacy CLI/tests contract (answers dict) keeps working.
+
+    Identity fields live on the req top level; the form answers (everything
+    under ``answers``) are what load_experiment_config reads for behaviour.
+    """
+    form_fields = {
+        k: v for k, v in _FLAT_FIELDS.items()
+        if k not in {"task_id", "task_type", "label", "raw_request"}
+    }
+    payload = {
+        "task_id": _FLAT_FIELDS["task_id"],
+        "task_type": _FLAT_FIELDS["task_type"],
+        "answers": form_fields,
+    }
+    cfg = load_experiment_config(
+        _write_req(tmp_path, payload),
+        tmp_path / "state",
+        tmp_path / "ws",
+    )
+    assert cfg.task_id == "he-webui-task"
+    assert cfg.execution_mode == "dkao-cli"
+    assert cfg.pool_mode is True
+    assert cfg.per_round_budget == 4
+
+
+def test_empty_requirements_fall_back_to_dry_defaults(tmp_path):
+    """No form fields at all -> documented dry-run defaults (baseline only)."""
+    cfg = load_experiment_config(
+        _write_req(tmp_path, {"task_id": "empty", "task_type": "harness-evolve"}),
+        tmp_path / "state",
+        tmp_path / "ws",
+    )
+    assert cfg.execution_mode == "dry-run"
+    assert cfg.pool_mode is False
+    assert cfg.max_iterations == 3
diff --git a/metainfer/tasks/harness_evolve/tests/test_decision_engine.py b/metainfer/tasks/harness_evolve/tests/test_decision_engine.py
new file mode 100644
index 00000000..06fcc8dc
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_decision_engine.py
@@ -0,0 +1,196 @@
+"""System-level five-state decision engine tests."""
+
+from __future__ import annotations
+
+from ..orchestrator.decision_engine import DecisionPolicy, decide
+
+
+def _results(medians, *, passed=True, correctness=True):
+    return {
+        k: {
+            "median_us": v,
+            "passed": passed,
+            "correctness_ok": correctness,
+        }
+        for k, v in medians.items()
+    }
+
+
+def _manifest(routeable=False):
+    return {
+        "iteration": 1,
+        "changes": [{
+            "id": "chg-1",
+            "mechanism_signature": ["plan sequence changes"],
+            "scope": {"routeable": routeable},
+        }],
+    }
+
+
+def test_first_result_establishes_baseline():
+    d = decide(None, _results({"a": 100.0}))
+    assert d["verdict"] == "BASELINE"
+    assert d["action"] == "SET_CHAMPION"
+
+
+def test_promote_when_performance_and_mechanism_hit():
+    champ = _results({"a": 100.0, "b": 100.0})
+    cand = _results({"a": 90.0, "b": 88.0})
+    d = decide(champ, cand, manifest=_manifest(),
+               mechanism_checks={"chg-1": True})
+    assert d["verdict"] == "PROMOTE"
+    assert d["performance_gate"]["status"] == "PASS"
+    assert d["mechanism_gate"]["status"] == "HIT"
+
+
+def test_lenient_default_promotes_unexplained_win_without_repeat():
+    """Default policy is lenient: an unverified win still takes the champion."""
+    champ = _results({"a": 100.0, "b": 100.0})
+    cand = _results({"a": 90.0, "b": 90.0})
+    d = decide(champ, cand, manifest=_manifest(), mechanism_checks={})
+    assert d["verdict"] == "PROMOTE_UNEXPLAINED"
+    assert d["action"] == "SET_CHAMPION"
+    assert d["policy"]["mechanism_policy"] == "lenient"
+
+
+def test_strict_policy_requires_confirmation_then_promotes():
+    champ = _results({"a": 100.0, "b": 100.0})
+    cand = _results({"a": 90.0, "b": 90.0})
+    strict = DecisionPolicy(mechanism_policy="strict")
+    first = decide(champ, cand, manifest=_manifest(), mechanism_checks={},
+                   policy=strict)
+    assert first["verdict"] == "CONFIRM_REQUIRED"
+    second = decide(champ, cand, manifest=_manifest(), mechanism_checks={},
+                    confirmation_reproduced=True, policy=strict)
+    assert second["verdict"] == "PROMOTE_UNEXPLAINED"
+    assert second["action"] == "SET_CHAMPION"
+
+
+def test_no_signal_keeps_champion_without_rejecting_candidate():
+    champ = _results({"a": 100.0, "b": 100.0})
+    cand = _results({"a": 99.5, "b": 101.0})
+    d = decide(champ, cand, manifest=_manifest())
+    assert d["verdict"] == "NO_SIGNAL"
+    assert d["action"] == "KEEP_CHAMPION_ARCHIVE_CANDIDATE"
+
+
+def test_correctness_or_primary_regression_rejects():
+    champ = _results({"a": 100.0, "b": 100.0})
+    cand = _results({"a": 120.0, "b": 90.0})
+    cand["b"]["correctness_ok"] = False
+    d = decide(champ, cand, manifest=_manifest())
+    assert d["verdict"] == "REJECT"
+    assert d["action"] == "RESTORE_CHAMPION"
+
+
+def test_heldout_regression_specializes_if_routeable_else_rejects_overfit():
+    champ = _results({"a": 100.0, "b": 100.0})
+    cand = _results({"a": 80.0, "b": 80.0})
+    held_champ = _results({"h": 100.0})
+    held_cand = _results({"h": 120.0})
+
+    specialized = decide(
+        champ, cand, manifest=_manifest(routeable=True),
+        mechanism_checks={"chg-1": True},
+        heldout_champion=held_champ, heldout_candidate=held_cand,
+    )
+    assert specialized["verdict"] == "SPECIALIZE"
+    assert specialized["action"] == "ARCHIVE_SPECIALIZED"
+
+    rejected = decide(
+        champ, cand, manifest=_manifest(routeable=False),
+        mechanism_checks={"chg-1": True},
+        heldout_champion=held_champ, heldout_candidate=held_cand,
+    )
+    assert rejected["verdict"] == "REJECT_OVERFIT"
+
+
+def test_harness_validate_error_detects_bad_yaml(tmp_path):
+    import shutil as _sh
+    from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.harness_io import default_harness_dir
+    from ..orchestrator.pipeline import _harness_validate_error
+    wd = tmp_path / "ws"
+    _sh.copytree(default_harness_dir(), wd)
+    assert _harness_validate_error(wd) is None
+    (wd / "planner_policy.yaml").write_text("not: [valid\n", encoding="utf-8")
+    assert _harness_validate_error(wd) is not None
+
+
+# --------------------------------------------------------- the 75% rule
+# The performance gate is a question-count rule: >= 75% of the counted
+# questions must be measurably faster than the champion ("3 of 4 beat the
+# current variant"), with one hard red line for a single question that
+# collapses (> 30% slower) or fails correctness.
+
+def test_three_of_four_question_wins_pass_the_performance_gate():
+    champ = _results({"a": 100.0, "b": 100.0, "c": 100.0, "d": 100.0})
+    cand = _results({"a": 90.0, "b": 90.0, "c": 91.0, "d": 110.0})  # d is 10% slower
+    d = decide(champ, cand, manifest=_manifest(), mechanism_checks={"chg-1": True})
+    assert d["performance_gate"]["status"] == "PASS"
+    assert d["performance_gate"]["win_ratio"] == 0.75
+    assert d["performance_gate"]["wins_needed"] == 3
+    assert d["verdict"] == "PROMOTE"
+
+
+def test_two_of_four_question_wins_fail_the_performance_gate():
+    champ = _results({"a": 100.0, "b": 100.0, "c": 100.0, "d": 100.0})
+    cand = _results({"a": 90.0, "b": 90.0, "c": 130.0, "d": 130.0})
+    d = decide(champ, cand, manifest=_manifest(), mechanism_checks={"chg-1": True})
+    assert d["verdict"] == "REJECT"
+    assert d["performance_gate"]["status"] == "FAIL"
+    assert d["performance_gate"]["win_ratio"] == 0.5
+    assert "2/4 counted questions faster" in d["reason"]
+
+
+def test_one_collapsing_question_rejects_despite_three_wins():
+    """The red line: 3 small wins do not buy one 40%-slower question."""
+    champ = _results({"a": 100.0, "b": 100.0, "c": 100.0, "d": 100.0})
+    cand = _results({"a": 90.0, "b": 90.0, "c": 90.0, "d": 140.0})
+    d = decide(champ, cand, manifest=_manifest(), mechanism_checks={"chg-1": True})
+    assert d["verdict"] == "REJECT"
+    assert d["performance_gate"]["collapses"] == ["d"]
+    assert "hard red line" in d["reason"]
+
+
+def test_three_questions_all_winning_passes():
+    champ = _results({"a": 100.0, "b": 100.0, "c": 100.0})
+    cand = _results({"a": 80.0, "b": 80.0, "c": 80.0})
+    d = decide(champ, cand, manifest=_manifest(), mechanism_checks={"chg-1": True})
+    assert d["policy"]["min_win_ratio"] == 0.75
+    assert d["performance_gate"]["wins_needed"] == 3
+    assert d["performance_gate"]["status"] == "PASS"
+
+
+def test_a_draw_inside_the_noise_band_counts_against_the_ratio():
+    """Only strictly measurable wins count: a tie is not a win."""
+    champ = _results({"a": 100.0, "b": 100.0, "c": 100.0, "d": 100.0})
+    cand = _results({"a": 90.0, "b": 90.0, "c": 100.0, "d": 130.0})
+    d = decide(champ, cand, manifest=_manifest(), mechanism_checks={"chg-1": True})
+    gate = d["performance_gate"]
+    assert gate["counts"]["NO_SIGNAL"] == 1 and gate["counts"]["WIN"] == 2
+    assert gate["win_ratio"] == 0.5
+    assert d["verdict"] == "REJECT"
+
+
+def test_probe_questions_are_excluded_from_the_ratio():
+    """A probe is generalisation evidence, not part of the 75% denominator."""
+    champ = _results({"a": 100.0, "b": 100.0, "c": 100.0, "probe": 100.0})
+    cand = _results({"a": 90.0, "b": 90.0, "c": 90.0, "probe": 200.0})
+    d = decide(champ, cand, manifest=_manifest(), mechanism_checks={"chg-1": True},
+               purposes={"probe": "probe"})
+    # 3/3 counted questions won -> gate passes, but the probe regression blocks
+    assert d["performance_gate"]["win_ratio"] == 1.0
+    assert d["generalization_gate"]["status"] == "FAIL"
+    assert d["verdict"] == "CONFIRM_REQUIRED"
+
+
+def test_unusable_measurement_is_fatal_even_at_three_of_four():
+    """A comparison missing one side is not "the candidate is faster"."""
+    champ = _results({"a": 100.0, "b": 100.0, "c": 100.0, "d": 100.0})
+    cand = _results({"a": 90.0, "b": 90.0, "c": 90.0, "d": 100.0})
+    cand["d"] = {"median_us": None, "passed": False, "correctness_ok": False,
+                 "status": "missing_report"}
+    d = decide(champ, cand, manifest=_manifest(), mechanism_checks={"chg-1": True})
+    assert d["performance_gate"]["counts"]["INVALID"] == 1
+    assert d["verdict"] == "REJECT"
+    assert "unusable measurement" in d["reason"]
diff --git a/metainfer/tasks/harness_evolve/tests/test_evolve_failure.py b/metainfer/tasks/harness_evolve/tests/test_evolve_failure.py
new file mode 100644
index 00000000..84edf825
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_evolve_failure.py
@@ -0,0 +1,87 @@
+"""AgentEvolver must never crash the outer loop.
+
+A timed-out or failed evolve turn keeps the iteration's evaluation value:
+apply() returns a no-change placeholder (champion unchanged) so the pipeline
+still writes evolve artifacts + report + finished.
+"""
+
+from __future__ import annotations
+
+import subprocess
+
+import pytest
+
+from metainfer.tasks.harness_evolve.orchestrator import evolve as ev
+
+
+def _ctx() -> dict:
+    return {
+        "overview": "/tmp/overview.md",
+        "results": {},
+        "diff": {},
+        "decision": {"verdict": "BASELINE"},
+        "selected": ["a"],
+        "pool_ids": ["a", "b"],
+    }
+
+
+def test_timeout_returns_placeholder(tmp_path, monkeypatch):
+    def _boom(*_a, **_kw):
+        raise subprocess.TimeoutExpired(cmd="dsh_agent.py", timeout=1800)
+
+    monkeypatch.setattr(ev.subprocess, "run", _boom)
+    out = ev.AgentEvolver(timeout_seconds=1800).apply(tmp_path, 3, _ctx())
+    assert out["changed"] is False
+    assert out["round_plan"] is None
+    assert out["manifest"]["changes"] == []
+    assert out["manifest"]["verification"]["status"] == "evolve_failed"
+    assert "timed out" in out["manifest"]["note"]
+    assert "failed" in out["summary"]
+
+
+def test_nonzero_exit_returns_placeholder(tmp_path, monkeypatch):
+    class _P:
+        returncode = 1
+        stdout = ""
+        stderr = "boom: agent crashed"
+
+    monkeypatch.setattr(ev.subprocess, "run", lambda *a, **k: _P())
+    out = ev.AgentEvolver().apply(tmp_path, 1, _ctx())
+    assert out["manifest"]["verification"]["status"] == "evolve_failed"
+    assert "rc=1" in out["manifest"]["note"]
+
+
+def test_missing_artifacts_returns_placeholder(tmp_path, monkeypatch):
+    class _P:
+        returncode = 0
+        stdout = ""
+        stderr = ""
+
+    monkeypatch.setattr(ev.subprocess, "run", lambda *a, **k: _P())
+    out = ev.AgentEvolver().apply(tmp_path, 1, _ctx())
+    assert out["manifest"]["verification"]["status"] == "evolve_failed"
+    assert "did not write" in out["manifest"]["note"]
+
+
+def test_successful_evolve_still_returns_artifacts(tmp_path, monkeypatch):
+    class _P:
+        returncode = 0
+        stdout = '{"type": "result", "result": "done"}\n'
+        stderr = ""
+
+    def _run(*_a, **_kw):
+        # The agent writes its artifacts during the run (apply clears them
+        # before spawning), so emulate that side effect here.
+        (tmp_path / ev.AgentEvolver.MANIFEST_NAME).write_text(
+            '{"schema_version": 1, "iteration": 1, "changes": [{"id": "c1"}],'
+            ' "verification": {"status": "pending"}}', encoding="utf-8")
+        (tmp_path / ev.AgentEvolver.PLAN_NAME).write_text(
+            '{"iteration": 1, "strategy": "balanced", "selected": ["a"]}',
+            encoding="utf-8")
+        return _P()
+
+    monkeypatch.setattr(ev.subprocess, "run", _run)
+    out = ev.AgentEvolver().apply(tmp_path, 1, _ctx())
+    assert out["changed"] is True
+    assert out["round_plan"]["selected"] == ["a"]
+    assert out["summary"] == "done"
diff --git a/metainfer/tasks/harness_evolve/tests/test_flow_contract.py b/metainfer/tasks/harness_evolve/tests/test_flow_contract.py
new file mode 100644
index 00000000..af83b0c0
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_flow_contract.py
@@ -0,0 +1,220 @@
+"""The HE flow is a fixed contract: `FLOW.md` states it, the code obeys it.
+
+Three things are pinned here, all of them from the operator's own wording:
+
+1. **The document exists and still says the protocol** — question counts are
+   multiples of 4, round 1 draws the frozen generalization paper from a random
+   set that spans decode *and* prefill, the two gates are 75% / 50% with an 80%
+   floor, and a failed performance gate goes back to analyze→evolve instead of
+   being promoted.
+2. **The multiple of 4 is enforced, not suggested** — the form offers only
+   multiples of 4, and a count that arrives from the CLI is rounded up at the
+   single point where it enters the run (``config._normalized_budget``).
+3. **HE only launches DKAO tasks** — no device reads, no waiting on cards, no
+   freezing a child for a device. Every one of those existed at some point and
+   was removed; this test is what keeps them out.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+HE_PKG = Path(__file__).resolve().parents[1]
+ORCH = HE_PKG / "orchestrator"
+
+
+def _read(path: Path) -> str:
+    return path.read_text(encoding="utf-8")
+
+
+def _code_only(path: Path) -> str:
+    """Source with comments and string literals removed.
+
+    A contract test has to judge *code*: the comments in ``eval.py`` explaining
+    that the device verdict belongs to DKAO are exactly what we want to keep,
+    and they would otherwise read as "HE mentions the measurement gate".
+    """
+    import io
+    import tokenize
+
+    tokens = []
+    with path.open("rb") as fh:
+        for tok in tokenize.tokenize(fh.readline):
+            if tok.type in (tokenize.COMMENT, tokenize.STRING, tokenize.NL,
+                            tokenize.NEWLINE, tokenize.INDENT,
+                            tokenize.DEDENT):
+                continue
+            tokens.append(tok.string)
+    return " ".join(tokens)
+
+
+# --------------------------------------------------------------- 1. the document
+
+def test_the_flow_document_exists_and_is_called_out_by_the_readme():
+    flow = HE_PKG / "FLOW.md"
+    assert flow.is_file(), "the fixed protocol document is missing"
+    assert "FLOW.md" in _read(HE_PKG / "README.md"), (
+        "README must point at FLOW.md instead of restating the protocol")
+
+
+def test_the_flow_document_states_every_rule_the_operator_fixed():
+    doc = _read(HE_PKG / "FLOW.md")
+    required = [
+        "4 的倍数",              # question count is a multiple of 4
+        "decode", "prefill",     # round 1 spans both regimes
+        "M ≤ 32",                # decode definition
+        "A1", "A_i",             # round-1 paper vs later draws
+        "75%", "50%", "80%",     # the two gates and the floor
+        "仅在性能门通过后开",     # generalization only after performance
+        "固化考卷",               # the paper never moves
+        "只发起 DKAO 任务",       # the boundary, in one line
+        "不租卡", "不读卡状态", "不 `SIGSTOP`",   # what HE must not do
+        "analyze", "evolve",     # the failure path
+        "awaiting_approval",     # promotion waits for a human
+    ]
+    missing = [token for token in required if token not in doc]
+    assert not missing, f"FLOW.md no longer states: {missing}"
+
+
+# ------------------------------------------------- 2. the multiple-of-4 rule
+
+def test_the_machine_and_the_wave_are_both_four():
+    from metainfer.tasks.harness_evolve.orchestrator.rounds import (
+        GPU_COUNT, OPERATORS_PER_GATE, normalized_per_gate,
+    )
+
+    assert GPU_COUNT == 4 and OPERATORS_PER_GATE == 4
+    assert normalized_per_gate(4) == 4
+    assert normalized_per_gate(8) == 8
+    assert normalized_per_gate(5) == 8      # up, never down
+    assert normalized_per_gate(3) == 4
+    assert normalized_per_gate(None) == 4
+
+
+def test_the_form_only_offers_multiples_of_four():
+    import metainfer.tasks  # noqa: F401 - registers the plugins
+    from metainfer.server.forms import load_form_schema
+
+    schema = load_form_schema("harness-evolve")
+    field = next(f for f in schema["fields"] if f["key"] == "per_round_budget")
+
+    assert field["type"] == "select", (
+        "a free number field lets the operator ask for 3 or 5 questions, which "
+        "leaves devices idle in the wave")
+    counts = [int(option["label"]) for option in field["options"]]
+    assert counts and all(count % 4 == 0 for count in counts), counts
+    assert int(field["default"]) % 4 == 0
+
+
+def test_a_count_from_the_cli_is_rounded_up_where_it_enters_the_run():
+    from metainfer.tasks.harness_evolve.orchestrator.config import (
+        _normalized_budget,
+    )
+
+    assert _normalized_budget(3) == 4
+    assert _normalized_budget(5) == 8
+    assert _normalized_budget(12) == 12
+    assert _normalized_budget(0) == 0, (
+        "0 keeps the legacy meaning 'no cap / whole pool'")
+    assert _normalized_budget("bad") == 0
+
+
+def test_the_round_is_drawn_from_the_pool_in_whole_waves():
+    """The selection itself honours the count it is given (no silent shrink)."""
+    import random
+
+    from metainfer.tasks.harness_evolve.orchestrator.rounds import (
+        select_round_questions,
+    )
+
+    instances = {}
+    for idx in range(8):
+        m = 16 if idx % 2 == 0 else 4096
+        regime = "decode" if m <= 32 else "prefill"
+        iid = f"op{idx:02d}"
+        instances[iid] = {
+            "baseline_us": 100.0 * (idx + 1),
+            "best_known_us": 40.0 * (idx + 1),        # an existing variant
+            "family": f"{regime}__fam{idx}",
+            "contract": {"M": m, "N": 64, "K": 32},
+        }
+    table = {"operators": {}, "groups": {}, "harness_version": None}
+
+    picked = select_round_questions(iteration=1, table=table,
+                                    instances=instances, state={}, count=8,
+                                    rng=random.Random(7))
+
+    assert len(picked["performance_ids"]) == 8
+    assert picked["generalization_ids"] == picked["performance_ids"], (
+        "round 1's draw *is* the frozen paper")
+    assert picked["defines_paper"] is True
+
+
+# ------------------------------------------- 3. HE does not manage environment
+
+def test_the_supervisor_never_reads_the_cards():
+    from metainfer.tasks.harness_evolve.orchestrator.supervisor import Supervisor
+
+    code = _code_only(ORCH / "supervisor.py")
+    for forbidden in ("gpu_preflight", "preflight_gpus", "free_gb",
+                      "sample_gpu_state", "foreign_kfd_pids", "GpuBroker"):
+        assert forbidden not in code, (
+            f"HE must not consult {forbidden}: a device verdict is DKAO's "
+            "(FLOW.md §0)")
+    assert not hasattr(Supervisor, "free_gb")
+
+
+def test_nothing_freezes_a_child_for_a_device_any_more():
+    src = _code_only(ORCH / "process_control.py")
+    for forbidden in ("signal.SIGSTOP", "signal.SIGCONT", "def freeze",
+                      "def resume("):
+        assert forbidden not in src, forbidden
+    # terminate is what remains, and it must still work as a group signal
+    assert "def terminate" in src and "killpg" in src
+    assert "def freeze" not in src
+
+
+def test_classify_incident_is_driven_by_the_runs_own_evidence():
+    from metainfer.tasks.harness_evolve.orchestrator.supervisor import (
+        CATEGORY_AGENT, CATEGORY_CODE, CATEGORY_RESOURCE, CATEGORY_UNKNOWN,
+        classify_incident,
+    )
+
+    assert classify_incident(exit_code=137)[0] == CATEGORY_RESOURCE
+    assert classify_incident(log_tail="std::bad_alloc")[0] == CATEGORY_RESOURCE
+    assert classify_incident(log_tail="TransportClosedError")[0] == CATEGORY_AGENT
+    assert classify_incident(
+        log_tail="Traceback (most recent call last):")[0] == CATEGORY_CODE
+    assert classify_incident(log_tail="nothing to see")[0] == CATEGORY_UNKNOWN
+    with pytest.raises(TypeError):
+        classify_incident(log_tail="x", free_gb=1.0)      # the knob is gone
+
+
+def test_the_evaluator_only_decides_the_device_layout():
+    src = _code_only(ORCH / "adapters" / "eval.py")
+    assert "device_for_index" in src, (
+        "the device is a pure layout decision (index % 4)")
+    for forbidden in ("ensure_measurement_gate", "MEASUREMENT_GATE_VRAM_PERCENT",
+                      "measurement_suspect", "contention", "gpu_preflight",
+                      "preflight_gpus"):
+        assert forbidden not in src, (
+            "environment verdicts live in DKAO, not in HE")
+
+
+def test_the_gate_numbers_are_the_protocol_numbers():
+    from metainfer.tasks.harness_evolve.orchestrator.decision_engine import (
+        GENERALIZATION_GATE_FLOOR_PERCENT, GENERALIZATION_GATE_WIN_RATIO,
+        PERFORMANCE_GATE_FLOOR_PERCENT, PERFORMANCE_GATE_WIN_RATIO,
+    )
+    from metainfer.tasks.harness_evolve.orchestrator.rounds import (
+        GENERALIZATION_ROUND_COST, PERFORMANCE_ROUND_COST,
+    )
+
+    assert PERFORMANCE_GATE_WIN_RATIO == 0.75
+    assert GENERALIZATION_GATE_WIN_RATIO == 0.50
+    assert PERFORMANCE_GATE_FLOOR_PERCENT == 0.80
+    assert GENERALIZATION_GATE_FLOOR_PERCENT == 0.80
+    assert PERFORMANCE_ROUND_COST == 1.0
+    assert GENERALIZATION_ROUND_COST == 0.5
diff --git a/metainfer/tasks/harness_evolve/tests/test_form_rounds.py b/metainfer/tasks/harness_evolve/tests/test_form_rounds.py
new file mode 100644
index 00000000..190a3b2a
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_form_rounds.py
@@ -0,0 +1,137 @@
+"""The two round-ish numbers in the New Task form must mean what they say.
+
+They used to be easy to confuse, and confusing them had a real cost: the outer
+loop's bound came from a hidden ``max_rounds`` (default 10) plus whatever target
+the WebUI happened to carry, while the form's own answer was only recorded in
+the report. A run therefore kept producing rounds nobody had asked for.
+
+The operator's rule now is one round field:
+
+* **Max iters** (``default_rounds``) — *one question's* optimization budget
+  inside DKAO. It is not a round count and must reach the child task unchanged.
+* **Rounds** (``rounds``) — *the whole task's* limit, counted in iterations: the
+  baseline round is round 1, and every later round is a performance pass plus
+  (when it wins) the generalization retake inside that same iteration. Passing
+  both gates in round 5 of 10 ends the run there, waiting for approval.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from metainfer.tasks.harness_evolve.orchestrator import cli as he_cli
+from metainfer.tasks.harness_evolve.orchestrator.config import (
+    load_experiment_config,
+)
+from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+    COST_GUARDRAIL, ROUNDS_DEFAULT, _rounds_limit,
+)
+from metainfer.tasks.harness_evolve.orchestrator.state import (
+    read_target_iterations,
+)
+
+
+def _write_req(tmp_path: Path, **answers) -> Path:
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "task_id": "form-rounds",
+        "task_type": "harness-evolve",
+        "execution_mode": "dry-run",
+        "evolve_mode": "dry-run",
+        "pool_source": "",
+    }
+    payload.update(answers)
+    req = tmp_path / "requirements.json"
+    req.write_text(json.dumps(payload), encoding="utf-8")
+    return req
+
+
+def _cfg(tmp_path: Path, **answers):
+    return load_experiment_config(
+        _write_req(tmp_path, **answers), tmp_path / "state", tmp_path / "ws")
+
+
+# ---------------------------------------------------------------- Max iters
+
+def test_max_iters_reaches_each_child_task_unchanged(tmp_path):
+    """``default_rounds`` is a per-question DKAO budget, not a round count."""
+    cfg = _cfg(tmp_path, default_rounds="7")
+    assert cfg.answers["default_rounds"] == "7"
+    from metainfer.tasks.harness_evolve.orchestrator.adapters.eval import (
+        DkaoCliEvaluator,
+    )
+    from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+        _variant_instance_spec,
+    )
+
+    spec = _variant_instance_spec(
+        {"id": "op", "baseline_us": 10.0, "best_known_us": 9.0},
+        budget_rounds=int(cfg.answers["default_rounds"]))
+    assert spec.budget_rounds == 7
+    # and it is handed to DKAO as its own "Max iters" field
+    req = DkaoCliEvaluator()._build_dkao_requirements(
+        cfg, spec, 1, "child-id", "repo-name", 0, tmp_path / "ws")
+    assert req["max_iterations"] == "7"
+
+
+# ------------------------------------------------------------------- Rounds
+
+def test_rounds_is_the_only_round_limit(tmp_path):
+    """There is one limit, counted in rounds, and a round is an iteration."""
+    assert _rounds_limit(_cfg(tmp_path, rounds=4)) == 4
+    # the pre-rename field name is still honoured
+    assert _rounds_limit(_cfg(tmp_path / "b", max_iterations=6)) == 6
+    # absent: the form's default
+    assert _rounds_limit(_cfg(tmp_path / "c")) == ROUNDS_DEFAULT
+    # nonsense is clamped into a sane range
+    assert _rounds_limit(_cfg(tmp_path / "d", rounds=9999)) == 50
+    assert _rounds_limit(_cfg(tmp_path / "e", rounds=-3)) == 1
+    # the cost guardrail can never be the thing that stops a real run
+    assert COST_GUARDRAIL >= 4 * ROUNDS_DEFAULT
+
+
+def test_rounds_is_pinned_where_the_loop_and_the_webui_read_it(tmp_path):
+    """The form's answer becomes the run's target before round 1."""
+    cfg = _cfg(tmp_path, rounds=7)
+    assert read_target_iterations(cfg.exp_dir) is None
+    he_cli._seed_rounds_from_form(cfg, cfg.exp_dir)
+    assert read_target_iterations(cfg.exp_dir) == 7
+    # a runaway value is clamped exactly like the WebUI's own target box
+    cfg2 = _cfg(tmp_path / "big", rounds=9999)
+    he_cli._seed_rounds_from_form(cfg2, cfg2.exp_dir)
+    assert read_target_iterations(cfg2.exp_dir) == 50
+
+
+# --------------------------------------------------- the form's own contract
+
+def test_the_form_has_one_round_field_and_it_is_machine_readable():
+    """Loadable *and* parseable: a select submits its label.
+
+    The explanation must not live in an option label (the backend runs ``int()``
+    on some answers), and the outer loop must expose exactly one round field.
+    """
+    import metainfer.tasks  # noqa: F401 - registers the plugins
+    from metainfer.server.forms import load_form_schema, validate_submission
+
+    schema = load_form_schema("harness-evolve")
+    assert schema is not None and schema["fields"]
+    keys = [f["key"] for f in schema["fields"]]
+    by_key = {f["key"]: f for f in schema["fields"]}
+
+    assert "rounds" in keys and "max_rounds" not in keys, (
+        "the outer loop exposes exactly one round field")
+    for field in ("default_rounds", "rounds", "per_round_budget"):
+        assert str(by_key[field]["default"]).isdigit()
+    for option in by_key["default_rounds"]["options"]:
+        int(option["label"])          # raises if a label grew an explanation
+
+    # the two fields name their layer, so "Max iters" cannot be read as rounds
+    assert "单题内部" in by_key["default_rounds"]["help"]
+    assert "任务上限" in by_key["rounds"]["help"]
+
+    answers = {k: str(f["default"]) for k, f in by_key.items()
+               if f.get("default") not in (None, "")}
+    answers["pool_source"] = "/root/zth_agent/ahe-kernel-repos/registered_pool.yaml"
+    result = validate_submission("harness-evolve", answers)
+    assert result["ok"], result.get("errors")
diff --git a/metainfer/tasks/harness_evolve/tests/test_four_card_overview.py b/metainfer/tasks/harness_evolve/tests/test_four_card_overview.py
new file mode 100644
index 00000000..fb60ed78
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_four_card_overview.py
@@ -0,0 +1,123 @@
+"""The four-card overview: what the page shows, and where the numbers come from.
+
+Two halves, and both are pinned because each has been wrong in production:
+
+1. **Membership** — which card a question is shown under. It comes from HE's
+   layout record for the round, *not* from the child's last gate row: the
+   generate-stage probe may run on any passing device, so three children pinned
+   to GPU0/1/2 were all listed under GPU2 and the page showed two cards instead
+   of four. (``tests/test_gpu_cards.py`` pins the precedence in detail; here the
+   payload contract the view reads is pinned.)
+2. **Content** — the view reads only fields the endpoint really provides, and it
+   is wired into the detail page (a component nobody renders is not a feature).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from metainfer.tasks.harness_evolve.server.routes import _gpu_cards
+
+HE = Path(__file__).resolve().parents[1]
+JS = HE / "static" / "he-detail.js"
+CSS = HE / "static" / "he.css"
+
+
+def _seed(exp: Path, num: int, qid: str, *, layout_gpu: int, phase: str = "parallel_explore") -> None:
+    attempt = exp / "children" / f"iteration_{num:03d}" / qid
+    state = attempt / "state"
+    state.mkdir(parents=True, exist_ok=True)
+    (state / "run.json").write_text(json.dumps({
+        "task_id": qid, "task_type": "dcu-kernel-auto-opt",
+        "current_phase": phase, "current_iteration": 2,
+        "finished": False, "last_update": 1790066400.0,
+    }), encoding="utf-8")
+    (attempt / "requirements.json").write_text(json.dumps({
+        "task_id": qid,
+        "shape_config": ("model: m\nassignment_mode: manual\nshapes:\n- id: " + qid
+                         + "\nassignments:\n  worker_0:\n    gpu: " + str(layout_gpu)
+                         + "\n    shapes:\n    - " + qid + "\n"),
+    }), encoding="utf-8")
+
+
+def _layout(exp: Path, num: int, mapping: dict) -> None:
+    bench = exp / "runs" / f"iteration_{num:03d}" / "input" / "benchmark"
+    bench.mkdir(parents=True, exist_ok=True)
+    (bench / "gpu_preflight.json").write_text(json.dumps({
+        "iteration": num, "gpu_count": 4, "managed_by": "dcu_kernel_auto_opt",
+        "devices": mapping,
+    }), encoding="utf-8")
+
+
+def test_a_four_question_round_becomes_four_cards(tmp_path):
+    exp = tmp_path / "exp"
+    questions = {"q0": 0, "q1": 1, "q2": 2, "q3": 3}
+    _layout(exp, 1, questions)
+    for qid, gpu in questions.items():
+        _seed(exp, 1, qid, layout_gpu=gpu)
+
+    cards = _gpu_cards(exp, 1)
+
+    assert [c["gpu"] for c in cards] == [0, 1, 2, 3], (
+        "one card per question — the view must be able to show all four")
+    for card in cards:
+        assert len(card["questions"]) == 1
+        q = card["questions"][0]
+        # every field the overview renders
+        assert {"gpu", "question", "current_phase", "current_iteration",
+                "best_median_us", "rounds", "finished", "last_update",
+                "gate", "gate_audit", "gpu_source"} <= set(q)
+        assert q["gpu_source"] == "layout"
+        assert q["gate_audit"].get("devices") is not None
+
+
+def test_a_grouped_card_still_carries_every_question(tmp_path):
+    """Two questions on one card (a second wave) must both be listed."""
+    exp = tmp_path / "exp"
+    _layout(exp, 2, {"q0": 0, "q1": 0})
+    _seed(exp, 2, "q0", layout_gpu=0)
+    _seed(exp, 2, "q1", layout_gpu=0)
+
+    cards = _gpu_cards(exp, 2)
+
+    assert len(cards) == 1 and cards[0]["gpu"] == 0
+    assert [q["question"] for q in cards[0]["questions"]] == ["q0", "q1"]
+
+
+# --------------------------------------------------------------- the view side
+
+def test_the_four_card_component_exists_and_is_rendered():
+    src = JS.read_text(encoding="utf-8")
+    assert "function FourCardOverview(" in src
+    assert "<${FourCardOverview} data=${iterGpu}" in src, (
+        "the overview must be rendered with the round's card payload")
+    assert 'class="he-four-cards"' in src
+
+
+def test_the_view_reads_only_fields_the_endpoint_provides():
+    src = JS.read_text(encoding="utf-8")
+    component = src.split("function FourCardOverview(", 1)[1].split(
+        "function PromotionPanel(", 1)[0]
+    for field in ("current_phase", "current_iteration", "best_median_us",
+                  "rounds", "finished", "last_update", "gate_audit",
+                  "gpu_source", "question"):
+        assert field in component, f"the card cell should show {field}"
+    # and it must not invent a metric the backend never sends
+    assert "p90_us" not in component
+
+
+def test_the_grid_is_four_columns_with_a_narrow_screen_fallback():
+    css = CSS.read_text(encoding="utf-8")
+    block = css.split(".he-four-cards {", 1)[1].split("}", 1)[0]
+    assert "repeat(4" in block, "four cards side by side"
+    assert "@media" in css.split(".he-four-cards {", 1)[1][:400], (
+        "narrow screens fall back to two columns")
+
+
+def test_the_overview_click_opens_the_child_drilldown():
+    src = JS.read_text(encoding="utf-8")
+    component = src.split("function FourCardOverview(", 1)[1].split(
+        "function PromotionPanel(", 1)[0]
+    assert "onSelect(q.question)" in component
+    assert "onSelect=${(id) => setSelChild(id)}" in src
diff --git a/metainfer/tasks/harness_evolve/tests/test_freeze_and_review.py b/metainfer/tasks/harness_evolve/tests/test_freeze_and_review.py
new file mode 100644
index 00000000..ae908624
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_freeze_and_review.py
@@ -0,0 +1,125 @@
+"""Scheme 1: frozen variants, best-known tracking, final promotion, review."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+    _finalize_promotion, _track_best_known, _variant_policy,
+)
+
+
+def _cfg(tmp_path: Path, **answers) -> SimpleNamespace:
+    state = tmp_path / "st"
+    state.mkdir(parents=True, exist_ok=True)
+    return SimpleNamespace(state_dir=state, task_id="t",
+                           answers={"promote_kernels": "true", **answers})
+
+
+def test_variant_policy_defaults_to_freeze(tmp_path):
+    assert _variant_policy(_cfg(tmp_path)) == "freeze"
+    assert _variant_policy(_cfg(tmp_path, variant_policy="per_round")) == "per_round"
+    assert _variant_policy(_cfg(tmp_path, variant_policy="immediate")) == "per_round"
+
+
+def test_best_known_curve_and_regression_detection(tmp_path):
+    cfg = _cfg(tmp_path)
+    exp = tmp_path / "exp"
+    exp.mkdir()
+
+    _track_best_known(cfg, exp, 1, {"a": {"median_us": 40.0, "passed": True}})
+    first = json.loads((exp / "best_known.json").read_text(encoding="utf-8"))
+    assert first["a"]["best_median_us"] == 40.0
+    assert first["a"]["best_iteration"] == 1
+
+    out = _track_best_known(cfg, exp, 2, {"a": {"median_us": 60.0, "passed": True}})
+    assert out["regressions"] and out["regressions"][0]["shape"] == "a"
+    assert out["regressions"][0]["best_median_us"] == 40.0
+    assert out["regressions"][0]["delta_percent"] == 50.0
+    assert (exp / "regressions.jsonl").is_file()
+
+    data = json.loads((exp / "best_known.json").read_text(encoding="utf-8"))
+    assert data["a"]["best_median_us"] == 40.0        # 60us never becomes best
+    assert data["a"]["best_iteration"] == 1
+    assert [h["median_us"] for h in data["a"]["history"]] == [40.0, 60.0]
+
+    # a better round updates the best
+    _track_best_known(cfg, exp, 3, {"a": {"median_us": 25.0, "passed": True}})
+    data = json.loads((exp / "best_known.json").read_text(encoding="utf-8"))
+    assert data["a"]["best_median_us"] == 25.0 and data["a"]["best_iteration"] == 3
+
+
+def test_finalize_promotion_promotes_the_best_candidate_not_the_newest(
+        tmp_path, monkeypatch):
+    """40us from round 1 must win over 60us from round 2."""
+    cfg = _cfg(tmp_path)
+    exp = tmp_path / "exp"
+    good = exp / "children" / "iteration_001" / "shapeA" / "workspace"
+    worse = exp / "children" / "iteration_002" / "shapeA" / "workspace"
+    good.mkdir(parents=True)
+    worse.mkdir(parents=True)
+    rows = [
+        {"iteration": 1, "shape": "shapeA", "workspace": str(good),
+         "ahe_passed": True, "action": "would_update", "ahe_median_us": 40.0,
+         "answers": {"model": "hy3"}, "source_task": "t/1"},
+        {"iteration": 2, "shape": "shapeA", "workspace": str(worse),
+         "ahe_passed": True, "action": "would_update", "ahe_median_us": 60.0,
+         "answers": {"model": "hy3"}, "source_task": "t/2"},
+        {"iteration": 2, "shape": "shapeB", "workspace": str(worse),
+         "ahe_passed": False, "action": "would_add", "ahe_median_us": 5.0,
+         "answers": {}},
+    ]
+    (exp / "promotion_candidates.jsonl").write_text(
+        "\n".join(json.dumps(r) for r in rows) + "\n", encoding="utf-8")
+
+    calls = []
+    import metainfer.tasks.dcu_kernel_auto_opt.orchestrator.variant_promote as vp
+    monkeypatch.setattr(vp, "promote_variant",
+                        lambda **kw: calls.append(kw) or {
+                            "shape": kw["shape_id"], "action": "updated"})
+
+    payload = _finalize_promotion(cfg, exp)
+    assert payload["promoted"] == ["shapeA"]
+    assert len(calls) == 1
+    assert calls[0]["workspace_dir"] == good          # the 40us candidate
+    assert Path(str(payload["results"][0]["shape"])) == Path("shapeA")
+    assert (exp / "final_promotion.json").is_file()
+    # a failing candidate is never promoted, however fast it looks
+    assert "shapeB" not in payload["promoted"]
+
+
+def test_finalize_promotion_respects_the_switch(tmp_path):
+    cfg = _cfg(tmp_path, promote_kernels="false")
+    exp = tmp_path / "exp"
+    exp.mkdir()
+    (exp / "promotion_candidates.jsonl").write_text("{}\n", encoding="utf-8")
+    assert _finalize_promotion(cfg, exp) is None
+
+
+def test_fail_streak_counts_failures_and_resets_on_progress():
+    from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+        _next_fail_streak,
+    )
+    assert _next_fail_streak("REJECT", 0) == 1
+    assert _next_fail_streak("REJECT_OVERFIT", 1) == 2
+    # unmeasurable rounds count too: that is the runaway-loop signature
+    assert _next_fail_streak("INCONCLUSIVE", 2) == 3
+    assert _next_fail_streak("NO_SIGNAL", 3) == 0
+    assert _next_fail_streak("PROMOTE", 3) == 0
+    assert _next_fail_streak("PROMOTE_UNEXPLAINED", 3) == 0
+
+
+def test_circuit_breaker_trips_after_three_failing_rounds():
+    from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+        _next_fail_streak,
+    )
+    streak = 0
+    tripped_at = None
+    for round_no, verdict in enumerate(["REJECT", "INCONCLUSIVE", "REJECT"], 1):
+        streak = _next_fail_streak(verdict, streak)
+        if streak >= 3:
+            tripped_at = round_no
+            break
+    assert tripped_at == 3
diff --git a/metainfer/tasks/harness_evolve/tests/test_gate_display_contract.py b/metainfer/tasks/harness_evolve/tests/test_gate_display_contract.py
new file mode 100644
index 00000000..211cfc1a
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_gate_display_contract.py
@@ -0,0 +1,431 @@
+"""What the round table's gate column shows, and where those numbers come from.
+
+The page used to print ``tau_family * baseline_us`` — the *pool* criterion — under
+the header "pass gate ≤ us". On the variant-round protocol (FLOW.md §3, the fixed
+protocol for every HE task since 2026-09-15) that number is not the gate: a round
+is judged by pairing each question against that operator's current **variant**,
+needing a strict win for >= 75% of the questions with nothing below 80% of its
+variant. Showing the pool criterion there made a 20.00x requirement look like the
+bar while the real bar was ~163x, i.e. it hid the very thing the round is judged
+on — and on an old run it read as if a kernel 2-4x slower than the variant had
+passed a variant comparison.
+
+Two halves, both pinned because each can be wrong independently:
+
+* the endpoint must serve the round's real criterion (``_real_gate_map``, driven
+  end to end through ``_children_live``);
+* the view must read it, keep the pool criterion only as a labelled annotation,
+  and still render the legacy path — which really *was* judged by that number.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import shutil
+import subprocess
+from pathlib import Path
+
+import pytest
+import yaml
+
+from metainfer.tasks.harness_evolve.server.routes import (
+    _children_live,
+    _iterations,
+)
+
+HE = Path(__file__).resolve().parents[1]
+JS = HE / "static" / "he-detail.js"
+
+#: One pool instance whose *pool* criterion (tau * baseline = 0.5 * 1000) is far
+#: from its variant (200): the two can never be confused in an assertion.
+POOL_BASELINE_US = 1000.0
+POOL_VARIANT_US = 200.0
+POOL_TARGET_US = 500.0
+
+
+def _pool_file(tmp_path: Path) -> Path:
+    path = tmp_path / "registered_pool.yaml"
+    path.write_text(yaml.safe_dump({
+        "schema_version": 1,
+        "instances": [{
+            "id": "q0",
+            "contract": {"model": "m", "tp_size": 8, "operator": "o_proj",
+                         "M": 4096, "N": 64, "K": 32},
+            "family": "prefill__fam",
+            "baseline_us": POOL_BASELINE_US,
+            "best_known_us": POOL_VARIANT_US,
+            # One accepted record with ratio 0.5 pins tau_family at 0.5.
+            "history": [{"accepted": True, "median_us": POOL_TARGET_US}],
+        }],
+    }), encoding="utf-8")
+    return path
+
+
+def _child(exp: Path, num: int, qid: str = "q0") -> None:
+    """A finished-looking DKAO child with one accepted measurement."""
+    d = exp / "children" / f"iteration_{num:03d}" / qid
+    (d / "state").mkdir(parents=True, exist_ok=True)
+    (d / "state" / "run.json").write_text(json.dumps({
+        "task_id": qid, "task_type": "dcu-kernel-auto-opt",
+        "current_phase": "finished", "current_iteration": 3,
+        "finished": True, "final_status": "success",
+    }), encoding="utf-8")
+    run = d / "workspace" / "workers" / "worker_0" / "runs" / qid
+    run.mkdir(parents=True, exist_ok=True)
+    (run / "experiments.jsonl").write_text(json.dumps({
+        "iteration": 1, "accepted": True, "metrics": {"median_us": 150.0},
+    }) + "\n", encoding="utf-8")
+
+
+def _variant_table(exp: Path, median_us: float = POOL_VARIANT_US,
+                   kernel_source: str = "pool:best_known") -> None:
+    (exp / "variant_table.json").write_text(json.dumps({
+        "schema_version": 1, "operators": {
+            "q0": {"median_us": median_us, "kernel_source": kernel_source,
+                   "conditions": {"origin": "registered_pool"}},
+        }, "groups": {},
+    }), encoding="utf-8")
+
+
+def _round(exp: Path, num: int, stage: str | None, *,
+           gate: dict | None = None, gate_name: str = "performance_gate.json",
+           ) -> None:
+    i_dir = exp / "runs" / f"iteration_{num:03d}" / "input"
+    (i_dir / "benchmark").mkdir(parents=True, exist_ok=True)
+    if stage is not None:
+        (i_dir / "stage.json").write_text(json.dumps({
+            "iteration": num, "stage": stage, "round_cost": 1.0,
+        }), encoding="utf-8")
+    if gate is not None:
+        (i_dir / gate_name).write_text(json.dumps(gate), encoding="utf-8")
+
+
+def _gate_doc(status: str = "PASS", verdict: str = "WIN",
+              delta: float = 12.5) -> dict:
+    return {
+        "kind": "performance", "status": status, "win_ratio": 1.0,
+        "required_win_ratio": 0.75, "operators": ["q0"], "wins": ["q0"],
+        "losses": [], "ties": [], "counted": 1, "wins_needed": 1,
+        "unmeasured": [], "below_floor": [], "reasons": [],
+        "comparison": {"per_instance": {"q0": {
+            "champion_median_us": POOL_VARIANT_US,
+            "candidate_median_us": 175.0,
+            "delta_percent": delta,
+            "verdict": verdict,
+        }}},
+    }
+
+
+def _answers(pool: Path) -> dict:
+    return {"pool_source": str(pool)}
+
+
+# ----------------------------------------------------------- what the gate is
+
+
+def test_a_performance_round_is_judged_against_the_variant(tmp_path):
+    exp = tmp_path / "exp"
+    pool = _pool_file(tmp_path)
+    _child(exp, 2)
+    _variant_table(exp)
+    _round(exp, 2, "performance", gate=_gate_doc())
+
+    gate = _children_live(exp, 2, _answers(pool))["q0"]["gate"]
+
+    assert gate["protocol"] == "variant"
+    assert gate["stage"] == "performance"
+    # the number the question must beat is the variant, not tau * baseline
+    assert gate["variant_us"] == POOL_VARIANT_US
+    assert gate["target_us"] == POOL_VARIANT_US
+    assert gate["judged"]["verdict"] == "WIN"
+    assert gate["judged"]["delta_percent"] == 12.5
+    assert gate["win_ratio"] == 0.75
+    assert gate["floor_percent"] == 0.80
+    assert gate["noise_percent"] == 2.0
+
+
+def test_the_pool_criterion_is_served_separately_and_labelled(tmp_path):
+    """It must still be reachable (old runs used it) — but never as the gate."""
+    exp = tmp_path / "exp"
+    pool = _pool_file(tmp_path)
+    _child(exp, 2)
+    _variant_table(exp)
+    _round(exp, 2, "performance", gate=_gate_doc())
+
+    live = _children_live(exp, 2, _answers(pool))["q0"]
+
+    assert live["gate"]["target_us"] != live["pool_reference"]["target_us"], (
+        "the two criteria must not collapse into one number")
+    assert live["pool_reference"]["target_us"] == POOL_TARGET_US
+    assert live["pool_reference"]["tau_family"] == 0.5
+    assert live["pool_reference"]["baseline_us"] == POOL_BASELINE_US
+    assert live["gate_protocol"] == "variant"
+
+
+def test_the_baseline_round_has_an_opponent_but_no_verdict(tmp_path):
+    """Round 1 judges nothing (FLOW.md §2) — but it does fix the opponent."""
+    exp = tmp_path / "exp"
+    pool = _pool_file(tmp_path)
+    _child(exp, 1)
+    _variant_table(exp)
+    _round(exp, 1, "baseline")
+
+    gate = _children_live(exp, 1, _answers(pool))["q0"]["gate"]
+
+    assert gate["protocol"] == "variant"
+    assert gate["stage"] == "baseline"
+    assert gate["judged"] is None
+    assert gate["target_us"] == POOL_VARIANT_US
+
+
+def test_a_legacy_round_keeps_the_criterion_it_was_judged_by(tmp_path):
+    """No stage, no gate document, no variant table: the pre-2026-09-15 path.
+
+    Those rounds were decided by pool.auto_pass, so that is the target their page
+    must keep showing — labelled as legacy, not as a variant comparison.
+    """
+    exp = tmp_path / "exp"
+    pool = _pool_file(tmp_path)
+    _child(exp, 3)
+    _round(exp, 3, None)
+
+    live = _children_live(exp, 3, _answers(pool))["q0"]
+
+    assert live["gate"]["protocol"] == "legacy"
+    assert live["gate"]["judged"] is None
+    assert live["gate"]["target_us"] == POOL_TARGET_US
+    assert live["pool_reference"]["target_us"] == POOL_TARGET_US
+
+
+def test_the_round_payload_carries_its_stage_and_gate(tmp_path):
+    """The page renders the round's own gate document, so serve it per round."""
+    exp = tmp_path / "exp"
+    _child(exp, 2)
+    _round(exp, 2, "performance", gate=_gate_doc(status="FAIL",
+                                                 verdict="LOSS", delta=-18.0))
+
+    entry = _iterations(exp)[0]
+
+    assert entry["stage"]["stage"] == "performance"
+    assert entry["performance_gate"]["status"] == "FAIL"
+    assert entry["performance_gate"]["wins_needed"] == 1
+
+
+def test_a_gate_round_without_a_stage_document_is_still_read(tmp_path):
+    """The 2026-09-18 transitional build wrote the gate but no stage.json."""
+    exp = tmp_path / "exp"
+    pool = _pool_file(tmp_path)
+    _child(exp, 1)
+    _round(exp, 1, None, gate=_gate_doc())
+
+    gate = _children_live(exp, 1, _answers(pool))["q0"]["gate"]
+
+    assert gate["protocol"] == "variant"
+    assert gate["judged"]["verdict"] == "WIN"
+
+
+# --------------------------------------------------------------- the view side
+
+
+def _js() -> str:
+    return JS.read_text(encoding="utf-8")
+
+
+def _component(name: str, end: str = "\nfunction ") -> str:
+    src = _js()
+    assert f"function {name}(" in src, f"{name} is gone"
+    return src.split(f"function {name}(", 1)[1].split(end, 1)[0]
+
+
+def test_the_gate_column_no_longer_advertises_the_pool_criterion():
+    src = _js()
+    assert "pass gate ≤ us" not in src, (
+        "the column must not be titled like a gate it does not show")
+    assert "必须打败（variant）" in src, "the column must name the real opponent"
+    component = _component("QuestionGateCell")
+    assert "gate.variant_us" in component, (
+        "the cell must render the variant, not a family baseline")
+    assert "gate.judged" in component, (
+        "a judged round must show WIN/LOSS, not just a number")
+
+
+def test_the_pool_number_is_only_an_explicitly_labelled_annotation():
+    component = _component("QuestionGateCell")
+    assert "poolRef" in component, (
+        "the pool criterion must arrive as its own argument")
+    assert "本协议不判门" in component, (
+        "showing tau * baseline without saying it is not the gate is the bug")
+    # `≥N× baseline` may only survive inside the legacy branch.
+    if "× baseline" in component:
+        legacy = component.split('gate.protocol === "legacy"', 1)[1]
+        assert "baseline_us / poolRef.target_us" in legacy, (
+            "the × baseline phrasing belongs to the legacy criterion only")
+
+
+def test_the_view_keeps_the_legacy_path_readable():
+    component = _component("QuestionGateCell")
+    assert 'gate.protocol === "legacy"' in component, (
+        "legacy rounds were judged by the family criterion; say so")
+    src = _js()
+    assert "legacy 判据" in src
+
+
+def test_the_round_shows_which_gate_it_owed_and_what_it_concluded():
+    src = _js()
+    assert "function RoundGateLine(" in src
+    assert "<${RoundGateLine} iter=${selIterObj} />" in src, (
+        "the round's gate line must be rendered with the round on screen")
+    line = _component("RoundGateLine")
+    for field in ("performance_gate", "generalization_gate", "wins_needed",
+                  "below_floor"):
+        assert field in line, f"the gate line should show {field}"
+
+
+def test_the_child_result_column_is_not_presented_as_the_gate():
+    src = _js()
+    assert "子任务验收" in src, (
+        "the r.passed column is DKAO's own acceptance, not the performance gate")
+    assert "DKAO 子任务自己的验收" in src
+
+
+def test_the_four_card_target_is_the_rounds_real_bar():
+    """The cards read the same `target_us`; it must be the variant on-protocol."""
+    component = _component("FourCardOverview", "\nfunction PromotionPanel(")
+    assert "必须打败 ≤" in component
+    assert "target = g.target_us" in component
+
+
+#: The repo ships no JS test runner (see tests/test_live_dkao_mirror.py), so the
+#: module is only *parsed*, not executed: a syntax error in a file with no build
+#: step would otherwise surface as a blank page and nothing else.
+@pytest.mark.skipif(shutil.which("node") is None, reason="node not installed")
+def test_the_detail_module_parses(tmp_path):
+    copy = tmp_path / "he-detail.mjs"
+    copy.write_text(_js(), encoding="utf-8")
+    proc = subprocess.run(["node", "--check", str(copy)],
+                          capture_output=True, text=True)
+    assert proc.returncode == 0, proc.stderr
+
+
+# ------------------------------------------------- the cell, actually executed
+
+#: Parsing is not rendering: reading ``g.baseline_us`` off the wrong object
+#: renders "-" instead of raising, which is how the pool criterion came to be
+#: shown as the gate in the first place. So the cell is also executed — with the
+#: repo's own vendored htm/preact (no browser, no bundler, no network): the vnode
+#: tree htm returns is plain data and can be walked for its text.
+VENDOR = HE.parent / "sys_shell" / "static" / "vendor"
+
+_RENDER_DRIVER = """
+import { QuestionGateCell } from "./he-detail.mjs";
+function text(v, out = []) {
+  if (v == null || v === false || v === true) return out;
+  if (typeof v === "string" || typeof v === "number") { out.push(String(v)); return out; }
+  if (Array.isArray(v)) { for (const x of v) text(x, out); return out; }
+  if (typeof v === "object") {
+    const p = v.props || {};
+    if (typeof v.type === "function") return text(v.type(p), out);
+    if (p.children != null) text(p.children, out);
+  }
+  return out;
+}
+const poolRef = { baseline_us: 1000.0, tau_family: 0.5, target_us: 500.0 };
+const base = { variant_source: "variant", kernel_source: "pool:best_known",
+               target_us: 200.0, variant_us: 200.0, win_ratio: 0.75,
+               floor_percent: 0.8, noise_percent: 2.0 };
+const cases = {
+  baseline: { ...base, protocol: "variant", stage: "baseline", judged: null },
+  judged: { ...base, protocol: "variant", stage: "performance",
+            judged: { verdict: "WIN", candidate_median_us: 175.0, delta_percent: 12.5 } },
+  tie: { ...base, protocol: "variant", stage: "performance",
+         judged: { verdict: "NO_SIGNAL", candidate_median_us: 200.0, delta_percent: 0.0 } },
+  legacy: { ...base, protocol: "legacy", stage: null, judged: null,
+            target_us: 500.0 },
+};
+const out = {};
+for (const [k, gate] of Object.entries(cases)) {
+  out[k] = text(QuestionGateCell({ gate, poolRef })).join(" ").replace(/\\s+/g, " ").trim();
+}
+out.none = text(QuestionGateCell({ gate: null, poolRef: null })).join(" ").trim();
+console.log(JSON.stringify(out));
+"""
+
+
+def _imported_names(src: str) -> set:
+    """Every named import he-detail.js takes from the app/* modules.
+
+    Derived from the source so the stub cannot rot: adding an import to the view
+    must not turn a render check into an unrelated failure.
+    """
+    names = set()
+    for match in re.finditer(r'import\s*\{([^}]*)\}\s*from\s*"(app/[^"]+)"', src):
+        for part in match.group(1).split(","):
+            name = part.strip().split(" as ")[0].strip()
+            if name:
+                names.add(name)
+    assert names, "he-detail.js imports nothing from app/*? update this test"
+    return names
+
+
+def _render_gate_cell(tmp_path: Path) -> dict:
+    """Render ``QuestionGateCell`` in node against the vendored htm/preact."""
+    if not VENDOR.is_dir():
+        pytest.skip("vendored htm/preact is missing")
+    work = tmp_path / "render"
+    work.mkdir()
+    # The vendor bundles import each other by bare specifier (the browser reads
+    # them through the shell importmap); a flat dir with relative specifiers is
+    # the node equivalent.
+    for name in ("htm.js", "htm-core.js", "preact.js", "preact-hooks.js"):
+        src = (VENDOR / name).read_text(encoding="utf-8")
+        src = src.replace('from"preact"', 'from"./preact.js"')
+        src = src.replace('from"htm"', 'from"./htm-core.js"')
+        (work / name).write_text(src, encoding="utf-8")
+    (work / "stub.js").write_text(
+        "const noop = async () => ({});\n"
+        + "\n".join(f"export const {n} = noop;"
+                    for n in sorted(_imported_names(_js())))
+        + "\nexport default function Stub() { return null; }\n",
+        encoding="utf-8")
+    src = _js()
+    for spec in ("app/he-api", "app/utils", "app/dkao-detail",
+                 "htm/preact", "preact/hooks"):
+        assert f'from "{spec}"' in src, f"{spec} import moved; update this test"
+    src = src.replace('from "app/he-api"', 'from "./stub.js"')
+    src = src.replace('from "app/utils"', 'from "./stub.js"')
+    src = src.replace('from "app/dkao-detail"', 'from "./stub.js"')
+    src = src.replace('from "htm/preact"', 'from "./htm.js"')
+    src = src.replace('from "preact/hooks"', 'from "./preact-hooks.js"')
+    src += "\nexport { QuestionGateCell };\n"
+    (work / "he-detail.mjs").write_text(src, encoding="utf-8")
+    (work / "driver.mjs").write_text(_RENDER_DRIVER, encoding="utf-8")
+    proc = subprocess.run(["node", str(work / "driver.mjs")],
+                          capture_output=True, text=True)
+    assert proc.returncode == 0, proc.stderr
+    return json.loads(proc.stdout)
+
+
+@pytest.mark.skipif(shutil.which("node") is None, reason="node not installed")
+def test_the_cell_renders_the_variant_and_labels_the_pool_number(tmp_path):
+    rendered = _render_gate_cell(tmp_path)
+
+    # The opponent is the variant; the pool number is never bare.
+    assert "≤ 200.00 µs" in rendered["baseline"]
+    assert "variant" in rendered["baseline"]
+    assert "第 1 轮不判门" in rendered["baseline"]
+    assert "500.00 µs" in rendered["baseline"], (
+        "the pool criterion should still be visible — as an annotation")
+    assert "本协议不判门" in rendered["baseline"]
+
+
+@pytest.mark.skipif(shutil.which("node") is None, reason="node not installed")
+def test_the_cell_renders_the_verdict_and_stays_empty_without_data(tmp_path):
+    rendered = _render_gate_cell(tmp_path)
+
+    assert "赢" in rendered["judged"] and "12.50" in rendered["judged"]
+    assert "平·不算赢" in rendered["tie"], "a tie is not a win for this gate"
+    assert rendered["none"] == "-", "no gate data must not invent a number"
+    # The legacy path keeps the number that judged it, and says so.
+    assert "legacy 判据" in rendered["legacy"]
+    assert "≥2.00× baseline" in rendered["legacy"]
diff --git a/metainfer/tasks/harness_evolve/tests/test_gpu_cards.py b/metainfer/tasks/harness_evolve/tests/test_gpu_cards.py
new file mode 100644
index 00000000..94b766b2
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_gpu_cards.py
@@ -0,0 +1,361 @@
+"""HE hands DKAO its questions and reads the gate verdict back, read-only.
+
+HE used to lease GPUs and police contention itself. That is gone: the device a
+question runs on is a fixed layout decision, and the only gate verdict is the
+one the child's own DKAO engine wrote into ``measurement_gate.jsonl``. These
+tests pin both halves — the layout HE records, and the read-only audit the
+detail page renders — plus the rule that a gate verdict is not retried away.
+"""
+
+from __future__ import annotations
+
+import json
+
+from metainfer.tasks.harness_evolve.orchestrator.adapters.eval import (
+    GATE_EXIT_CODE, DkaoCliEvaluator,
+)
+from metainfer.tasks.harness_evolve.orchestrator.rounds import (
+    GPU_COUNT, device_for_index,
+)
+from metainfer.tasks.harness_evolve.server.routes import (
+    _child_gate_audit, _gpu_cards,
+)
+
+
+def _gate_row(ts: float, event: str, *, attempt: int = 1, site: str = "benchmark",
+              gpu: int = 0, reasons=None, wait_seconds: float = 1800.0,
+              candidates=None, passing=None, usable=None):
+    row = {
+        "ts": ts, "event": event, "site": site, "gpu": gpu, "attempt": attempt,
+        "wait_seconds": wait_seconds, "max_waits": 48,
+        "reasons": list(reasons or []),
+    }
+    if candidates is not None:
+        row["candidates"] = list(candidates)
+        row["passing"] = list(passing or [])
+        row["usable"] = bool(passing) if usable is None else usable
+    return row
+
+
+def _write_gate(child_state_dir, rows):
+    child_state_dir.mkdir(parents=True, exist_ok=True)
+    (child_state_dir / "measurement_gate.jsonl").write_text(
+        "".join(json.dumps(r) + "\n" for r in rows), encoding="utf-8")
+
+
+# --------------------------------------------------------------- the layout
+
+
+def test_the_layout_is_a_fixed_rotation_and_says_he_manages_nothing():
+    """N questions land as N / 4 waves with one question per device."""
+    assert GPU_COUNT == 4
+    assert [device_for_index(i) for i in range(8)] == [0, 1, 2, 3, 0, 1, 2, 3]
+
+    class _Inst:
+        def __init__(self, i):
+            self.id = i
+
+    class _Cfg:
+        pass
+
+    layout = DkaoCliEvaluator._device_layout(
+        _Cfg(), 2, [_Inst(f"q{i}") for i in range(4)])
+    assert layout["schema"] == "he-gpu-layout/1"
+    assert layout["managed_by"] == "dcu_kernel_auto_opt"
+    # HE holds no lease at all — the field exists so a reader can see that
+    assert layout["leases"] == [] and layout["enabled"] is False
+    assert layout["devices"] == {"q0": 0, "q1": 1, "q2": 2, "q3": 3}
+    assert layout["assigned_devices"] == [0, 1, 2, 3]
+    assert "measurement gate" in layout["reason"]
+
+
+def test_the_layout_is_recorded_next_to_the_iteration_evidence(tmp_path):
+    """Audit trail: same file the DKAO side reads, in the run's input dir."""
+    (tmp_path / "exp").mkdir()
+    cfg = type("Cfg", (), {"exp_dir": tmp_path / "exp"})()
+    ev = DkaoCliEvaluator()
+    ev._save_preflight(cfg, 3, {"schema": "he-gpu-layout/1", "leases": []})
+    audit = (tmp_path / "exp" / "runs" / "iteration_003" / "input"
+             / "benchmark" / "gpu_preflight.json")
+    assert audit.is_file()
+    body = json.loads(audit.read_text(encoding="utf-8"))
+    assert body["schema"] == "he-gpu-layout/1" and body["leases"] == []
+
+
+# ------------------------------------------------------- the gate audit read
+
+
+def test_gate_audit_counts_checks_waits_and_the_last_verdict(tmp_path):
+    _write_gate(tmp_path / "state", [
+        _gate_row(1.0, "gate_ok", gpu=0, candidates=[0, 1, 2, 3], passing=[0]),
+        _gate_row(2.0, "gate_blocked", reasons=["VRAM 96.0% > 90.0% limit"]),
+        _gate_row(2.0, "gate_blocked", reasons=["VRAM 96.0% > 90.0% limit"]),
+        _gate_row(3.0, "gate_blocked", attempt=2, site="worker", gpu=2,
+                  reasons=["device busy: HCU 100.0% > 0.0%"]),
+        _gate_row(4.0, "gate_ok", gpu=2, candidates=[0, 2], passing=[2]),
+    ])
+    audit = _child_gate_audit(tmp_path / "state")
+    assert audit["checks"] == 4                # the duplicated channel row folds
+    assert audit["blocked"] == 2
+    assert audit["waited_seconds"] == 3600.0   # two 30-minute waits
+    assert audit["last_event"] == "gate_ok"
+    assert audit["last_site"] == "benchmark"
+    assert audit["max_waits"] == 48
+    assert audit["last_passing"] == [2] and audit["last_candidates"] == [0, 2]
+    # the blocked rows are kept so the page can show *why* it waited
+    reasons = [r["reasons"] for r in audit["blocked_reasons"]]
+    assert ["VRAM 96.0% > 90.0% limit"] in reasons
+    assert ["device busy: HCU 100.0% > 0.0%"] in reasons
+
+
+def test_a_passing_row_never_reports_its_neighbours_as_this_device_reason(
+        tmp_path):
+    """Real DKAO shape: a passing row lists the *other* cards it rejected.
+
+    ``ensure_any_measurement_gate`` puts the failed candidates' reasons on a
+    ``gate_ok`` row while ``gpu`` is the card that passed. Reporting those
+    reasons next to this device would show a block that never happened, so the
+    audit must drop them for a passing verdict — and must still name the device
+    it cleared.
+    """
+    _write_gate(tmp_path / "state", [
+        _gate_row(10.0, "gate_ok", site="generate_probe", gpu=1,
+                  candidates=[0, 1, 2, 3], passing=[1, 2, 3], usable=True,
+                  reasons=["device busy: HCU 97.5% > 0.0%"]),
+    ])
+    audit = _child_gate_audit(tmp_path / "state")
+    assert audit["last_event"] == "gate_ok" and audit["last_usable"] is True
+    assert audit["last_reasons"] == []                 # not this device's block
+    assert audit["blocked"] == 0 and audit["waited_seconds"] == 0
+    assert audit["device"] == 1 and audit["device_from"] == "gate_pass"
+    assert audit["last_passing"] == [1, 2, 3]
+
+
+def test_a_blocked_row_keeps_its_reasons_and_the_device_it_blocked_on(
+        tmp_path):
+    _write_gate(tmp_path / "state", [
+        _gate_row(10.0, "gate_blocked", site="benchmark", gpu=3, attempt=1,
+                  candidates=[3], passing=[], usable=False,
+                  reasons=["VRAM 94% > 90% limit",
+                           "device busy: HCU 90.8% > 0.0%"]),
+    ])
+    audit = _child_gate_audit(tmp_path / "state")
+    assert audit["last_event"] == "gate_blocked"
+    assert audit["last_usable"] is False
+    assert audit["last_reasons"] == ["VRAM 94% > 90% limit",
+                                     "device busy: HCU 90.8% > 0.0%"]
+    assert audit["device"] == 3 and audit["device_from"] == "gate_blocked"
+    assert audit["waited_seconds"] == 1800.0
+
+
+def test_gate_audit_of_a_child_that_never_gated_is_empty(tmp_path):
+    audit = _child_gate_audit(tmp_path / "missing")
+    assert audit["checks"] == 0 and audit["blocked"] == 0
+    assert audit["last_event"] is None and audit["blocked_reasons"] == []
+    assert audit["waited_seconds"] == 0
+    assert audit["device"] is None and audit["device_from"] is None
+    assert audit["last_reasons"] == []
+
+
+def test_child_gate_audit_matches_the_dkao_route_verdict(tmp_path):
+    """HE must read the child's rows the same way DKAO's own page does."""
+    from metainfer.tasks.dcu_kernel_auto_opt.server.routes import _gate_status
+
+    state = tmp_path / "state"
+    _write_gate(state, [
+        _gate_row(1.0, "gate_blocked", reasons=["device busy: HCU 100.0% > 0.0%"]),
+        _gate_row(2.0, "gate_ok", gpu=1, candidates=[1], passing=[1]),
+    ])
+    dkao = _gate_status(state)
+    he = _child_gate_audit(state)
+    assert dkao["blocked"] == he["blocked"] == 1
+    assert dkao["last_event"] == he["last_event"] == "gate_ok"
+
+
+# ------------------------------------------------------------ the card view
+
+
+def _seed_child(exp, num, qid, *, phase="parallel_explore", gate_rows=()):
+    state = exp / "children" / f"iteration_{num:03d}" / qid / "state"
+    state.mkdir(parents=True, exist_ok=True)
+    (state / "run.json").write_text(json.dumps({
+        "task_id": qid, "task_type": "dcu-kernel-auto-opt",
+        "current_phase": phase, "current_iteration": 1,
+    }), encoding="utf-8")
+    if gate_rows:
+        _write_gate(state, gate_rows)
+    return state
+
+
+def test_gpu_cards_group_questions_by_their_own_device(tmp_path):
+    """The card is the device the child's gate recorded, not the enumeration."""
+    exp = tmp_path / "exp"
+    # child A was cleared for device 3, child B for device 0: the cards must
+    # follow the children's own records rather than alphabetical order
+    _seed_child(exp, 2, "m16_a", gate_rows=[
+        _gate_row(1.0, "gate_blocked", gpu=3,
+                  reasons=["device busy: HCU 100.0% > 0.0%"]),
+        _gate_row(2.0, "gate_ok", gpu=3, candidates=[3], passing=[3]),
+    ])
+    _seed_child(exp, 2, "m16_b", gate_rows=[
+        _gate_row(3.0, "gate_ok", gpu=0, candidates=[0], passing=[0]),
+    ])
+    _seed_child(exp, 2, "m16_c")            # no gate record: layout prediction
+
+    cards = _gpu_cards(exp, 2)
+    by_gpu = {c["gpu"]: c for c in cards}
+    assert sorted(by_gpu) == [0, 2, 3]
+    assert [q["question"] for q in by_gpu[3]["questions"]] == ["m16_a"]
+    assert [q["question"] for q in by_gpu[0]["questions"]] == ["m16_b"]
+    # m16_c is the third child alphabetically -> index 2 -> predicted device 2
+    assert [q["question"] for q in by_gpu[2]["questions"]] == ["m16_c"]
+    assert by_gpu[3]["questions"][0]["gate_audit"]["blocked"] == 1
+    assert by_gpu[3]["questions"][0]["gpu_source"] == "gate_pass"
+    # nothing recorded this round's layout and the child has no requirements,
+    # so its device is only a prediction — and it must not look like an
+    # observation or like HE's own layout record
+    assert by_gpu[2]["questions"][0]["gpu_source"] == "predicted"
+
+
+def test_a_child_without_gate_records_is_marked_as_a_prediction(tmp_path):
+    """An observed device and a predicted one must not look the same."""
+    exp = tmp_path / "exp"
+    _seed_child(exp, 1, "no_gate_yet")
+    cards = _gpu_cards(exp, 1)
+    q = cards[0]["questions"][0]
+    assert q["gpu_source"] == "predicted"
+    assert q["gate_audit"]["checks"] == 0
+    assert q["gate_audit"]["device"] is None
+
+
+def test_the_round_layout_beats_a_probe_run_on_another_card(tmp_path):
+    """The real case: three children were listed under GPU2 because their
+    *generate probe* ran there, even though HE had given them GPU0/1/2/3.
+
+    A probe may legitimately run on any passing device ("this task's own card
+    first, otherwise any idle one"), so it must never decide which card a
+    question is shown under — otherwise a four-card view shows two cards.
+    """
+    exp = tmp_path / "exp"
+    (exp / "runs" / "iteration_002" / "input" / "benchmark").mkdir(parents=True)
+    (exp / "runs" / "iteration_002" / "input" / "benchmark"
+     / "gpu_preflight.json").write_text(json.dumps({
+         "iteration": 2, "gpu_count": 4, "managed_by": "dcu_kernel_auto_opt",
+         "devices": {"m16_a": 0, "m16_b": 1},
+     }), encoding="utf-8")
+    # both children's last *passing* row is the probe, on device 2
+    _seed_child(exp, 2, "m16_a", gate_rows=[
+        _gate_row(1.0, "gate_pass", gpu=2, site="generate_probe",
+                  candidates=[0, 1, 2, 3], passing=[2]),
+    ])
+    _seed_child(exp, 2, "m16_b", gate_rows=[
+        _gate_row(2.0, "gate_pass", gpu=2, site="generate_probe",
+                  candidates=[0, 1, 2, 3], passing=[2]),
+    ])
+
+    cards = _gpu_cards(exp, 2)
+    by_gpu = {c["gpu"]: c for c in cards}
+
+    assert sorted(by_gpu) == [0, 1], "the layout decides the card"
+    assert [q["question"] for q in by_gpu[0]["questions"]] == ["m16_a"]
+    assert [q["question"] for q in by_gpu[1]["questions"]] == ["m16_b"]
+    assert by_gpu[0]["questions"][0]["gpu_source"] == "layout"
+    # ... and the probe's card is still reported, as a separate fact
+    assert by_gpu[0]["questions"][0]["gate_audit"]["probe_device"] == 2
+    assert by_gpu[0]["questions"][0]["gate_audit"]["devices"] == [2]
+
+
+def test_a_measuring_site_row_decides_over_a_probe_row(tmp_path):
+    """Without a layout record, the card that *measured* is the child's card."""
+    exp = tmp_path / "exp"
+    _seed_child(exp, 1, "m16_a", gate_rows=[
+        _gate_row(1.0, "gate_pass", gpu=2, site="generate_probe",
+                  candidates=[0, 2], passing=[2]),
+        _gate_row(2.0, "gate_pass", gpu=0, site="benchmark",
+                  candidates=[0], passing=[0]),
+    ])
+    cards = _gpu_cards(exp, 1)
+    q = cards[0]["questions"][0]
+    assert cards[0]["gpu"] == 0 and q["gpu_source"] == "gate_pass"
+    assert q["gate_audit"]["probe_device"] == 2
+
+
+def test_gpu_cards_carry_the_child_phase_and_never_a_lease(tmp_path):
+    exp = tmp_path / "exp"
+    _seed_child(exp, 1, "only", phase="serial_validate")
+    cards = _gpu_cards(exp, 1)
+    assert len(cards) == 1
+    q = cards[0]["questions"][0]
+    assert q["question"] == "only"
+    assert q["current_phase"] == "serial_validate"
+    assert q["gpu"] == 0
+    assert "gate_audit" in q
+    # nothing in this payload is a lease HE took
+    assert "lease" not in q and "holder" not in q
+
+
+def test_gpu_cards_are_empty_before_any_child_is_dispatched(tmp_path):
+    exp = tmp_path / "exp"
+    (exp / "children").mkdir(parents=True)
+    assert _gpu_cards(exp, 4) == []
+
+
+# -------------------------------------------------- the verdict is not ignored
+
+
+def test_gate_exit_code_matches_the_dkao_asset_that_raises_it():
+    """The duplicated constant must not drift from the child's own value."""
+    import re
+    from pathlib import Path
+
+    asset = (Path(__file__).resolve().parents[2] / "dcu_kernel_auto_opt"
+             / "assets" / "w8a8_bench.py")
+    match = re.search(r"^GATE_EXIT_CODE\s*=\s*(\d+)", asset.read_text(
+        encoding="utf-8"), re.MULTILINE)
+    assert match, "w8a8_bench.py no longer defines GATE_EXIT_CODE"
+    assert int(match.group(1)) == GATE_EXIT_CODE
+
+
+# ------------------------------------------------------------- the Web route
+
+
+def test_the_web_route_serves_the_cards_and_offers_no_lease_control(
+        tmp_path, client):
+    """The page's own endpoint: read-only, no lease, no release/drain verbs."""
+    from metainfer.server import tasks as _tasks
+    from metainfer.server.tasks import TaskEntry
+
+    state_dir = tmp_path / "state"
+    workspace = tmp_path / "ws"
+    state_dir.mkdir(parents=True)
+    workspace.mkdir(parents=True)
+    (state_dir / "requirements.json").write_text(json.dumps({
+        "task_id": "he-gpu", "task_type": "harness-evolve",
+        "answers": {"ahe_repo_root": str(tmp_path / "repos")},
+    }), encoding="utf-8")
+    (state_dir / "run.json").write_text(json.dumps({"task_id": "he-gpu"}),
+                                        encoding="utf-8")
+    _seed_child(workspace, 2, "m16_a", gate_rows=[
+        _gate_row(1.0, "gate_blocked", gpu=3,
+                  reasons=["device busy: HCU 100.0% > 0.0%"]),
+    ])
+    _tasks.add_task(TaskEntry(
+        id="he-gpu", type="harness-evolve", label="device cards",
+        state_dir=str(state_dir), workspace_dir=str(workspace),
+        created_at=0.0))
+
+    resp = client.get("/api/harness-evolve/he-gpu/iterations/2/gpu")
+    assert resp.status_code == 200, resp.text
+    body = resp.json()
+    assert body["managed_by"] == "dcu_kernel_auto_opt"
+    assert body["leases"] == []                # HE holds nothing
+    assert body["gpu_count"] == 4
+    assert [c["gpu"] for c in body["cards"]] == [3]
+    q = body["cards"][0]["questions"][0]
+    assert q["question"] == "m16_a"
+    assert q["gate_audit"]["blocked"] == 1
+    assert q["gate_audit"]["last_reasons"] == ["device busy: HCU 100.0% > 0.0%"]
+    # the operator controls are gone with the lease they controlled
+    for verb in ("release", "drain"):
+        assert client.post(f"/api/harness-evolve/he-gpu/gpu/{verb}").status_code == 404
+    assert client.get("/api/harness-evolve/he-gpu/gpu").status_code == 404
diff --git a/metainfer/tasks/harness_evolve/tests/test_legacy_protocol_warning.py b/metainfer/tasks/harness_evolve/tests/test_legacy_protocol_warning.py
new file mode 100644
index 00000000..adc835d7
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_legacy_protocol_warning.py
@@ -0,0 +1,167 @@
+"""A run that leaves the fixed protocol has to say so.
+
+``FLOW.md`` describes the only protocol a new HE task may use: variant rounds, a
+frozen generalization paper, the two gates, rounds sized in waves of 4. Two older
+paths still exist — suite mode (no measured pool file) and an explicit
+``round_questions: legacy`` — and both skip all of that.
+
+The operator's decision (2026-09-22) was **not** to delete them but to make
+leaving the protocol impossible to miss: a WARNING in the orchestrator log, an
+event on the task timeline, ``protocol.json`` for the task page, and a banner in
+the detail view. That is what these tests pin.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import yaml
+
+from metainfer.server import tasks as _tasks
+from metainfer.server.tasks import TaskEntry
+from metainfer.tasks.harness_evolve.orchestrator import pipeline as P
+from metainfer.tasks.harness_evolve.orchestrator.config import (
+    load_experiment_config,
+)
+from metainfer.tasks.harness_evolve.orchestrator.pipeline import run_experiment
+
+
+def _requirements(tmp_path: Path, answers: dict) -> Path:
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    req = tmp_path / "requirements.json"
+    req.write_text(json.dumps({"task_id": "he-legacy", "answers": answers},
+                              ensure_ascii=False), encoding="utf-8")
+    return req
+
+
+def _cfg(tmp_path: Path, answers: dict):
+    return load_experiment_config(_requirements(tmp_path, answers),
+                                 tmp_path / "state", tmp_path / "ws")
+
+
+def _events(cfg) -> list:
+    path = Path(cfg.state_dir) / "timeline.jsonl"
+    if not path.is_file():
+        return []
+    return [json.loads(line) for line in
+            path.read_text(encoding="utf-8").splitlines() if line.strip()]
+
+
+def _pool(tmp_path: Path, operators: int = 8) -> Path:
+    path = tmp_path / "pool.yaml"
+    path.write_text(yaml.safe_dump({
+        "schema_version": 1,
+        "instances": [
+            {"id": f"op{i:02d}", "family": f"decode__fam{i}",
+             "contract": {"M": 16}, "baseline_us": 100.0 + i,
+             "best_known_us": 40.0 + 4 * i, "history": []}
+            for i in range(operators)
+        ],
+    }, sort_keys=False), encoding="utf-8")
+    return path
+
+
+# ------------------------------------------------------------------- the reason
+
+def test_the_reason_names_the_actual_cause(tmp_path):
+    explicit = _cfg(tmp_path / "a", {"execution_mode": "dry-run",
+                                     "round_questions": "legacy"})
+    assert "round_questions" in P._legacy_protocol_reason(explicit)
+
+    suite = _cfg(tmp_path / "b", {"execution_mode": "dry-run"})
+    assert "no question pool" in P._legacy_protocol_reason(suite)
+
+    builtin = _cfg(tmp_path / "c", {"execution_mode": "dry-run",
+                                    "pool_source": "builtin"})
+    assert "demo pool" in P._legacy_protocol_reason(builtin)
+
+    missing = _cfg(tmp_path / "d", {"execution_mode": "dry-run",
+                                    "pool_source": str(tmp_path / "nope.yaml")})
+    assert "does not exist" in P._legacy_protocol_reason(missing)
+
+
+# ------------------------------------------------- what the run leaves behind
+
+def test_a_suite_run_says_it_left_the_protocol(tmp_path, capsys):
+    cfg = _cfg(tmp_path, {"execution_mode": "dry-run", "max_iterations": 1})
+
+    run_experiment(cfg)
+
+    announced = [row for row in _events(cfg)
+                 if row["type"] == "legacy_protocol_selected"]
+    assert announced, "leaving the protocol must be on the record"
+    payload = announced[0]["payload"]
+    assert payload["mode"] == "legacy"
+    assert "no question pool" in payload["reason"]
+    assert "FLOW.md" in payload["policy"]
+
+    # ... in a file the task page reads, so it can warn instead of looking idle
+    written = json.loads((cfg.exp_dir / "protocol.json").read_text())
+    assert written["mode"] == "legacy" and written["reason"] == payload["reason"]
+
+    # ... and in the log, where somebody watching the process will see it
+    assert "WARNING" in capsys.readouterr().err
+
+
+def test_a_variant_run_says_nothing(tmp_path):
+    """The default protocol must not raise the banner (no false positives)."""
+    pool = _pool(tmp_path)
+    cfg = _cfg(tmp_path, {"execution_mode": "dry-run", "max_iterations": 1,
+                          "pool_source": str(pool), "question_pool": str(pool),
+                          "per_round_budget": "4"})
+
+    run_experiment(cfg)
+
+    assert not [row for row in _events(cfg)
+                if row["type"] == "legacy_protocol_selected"]
+    assert not (cfg.exp_dir / "protocol.json").exists()
+
+
+# --------------------------------------------------------------- on the page
+
+def _register(state_dir: Path, workspace_dir: Path, task_id: str):
+    state_dir.mkdir(parents=True, exist_ok=True)
+    workspace_dir.mkdir(parents=True, exist_ok=True)
+    _tasks.add_task(TaskEntry(id=task_id, type="harness-evolve", label="he",
+                              state_dir=str(state_dir),
+                              workspace_dir=str(workspace_dir),
+                              created_at=0.0))
+
+
+def test_the_task_page_can_warn_about_a_legacy_run(client, isolated_env,
+                                                   tmp_path):
+    home = isolated_env["home"]
+    state = home / "tasks" / "he-legacy-view"
+    ws = home / "workspaces" / "he-legacy-view"
+    _register(state, ws, "he-legacy-view")
+    (ws / "protocol.json").write_text(json.dumps({
+        "mode": "legacy", "protocol": "legacy",
+        "reason": "no question pool is configured",
+    }), encoding="utf-8")
+
+    body = client.get("/api/harness-evolve/he-legacy-view/summary").json()
+
+    assert body["protocol"]["mode"] == "legacy"
+    assert "no question pool" in body["protocol"]["reason"]
+
+
+def test_a_variant_task_reports_the_fixed_protocol(client, isolated_env):
+    home = isolated_env["home"]
+    state = home / "tasks" / "he-variant-view"
+    ws = home / "workspaces" / "he-variant-view"
+    _register(state, ws, "he-variant-view")
+
+    body = client.get("/api/harness-evolve/he-variant-view/summary").json()
+
+    assert body["protocol"]["mode"] == "variant"
+    assert "FLOW.md" in body["protocol"]["reason"]
+
+
+def test_the_detail_view_renders_the_banner_only_for_legacy():
+    """The page must read ``summary.protocol`` and stay quiet otherwise."""
+    src = (Path(__file__).resolve().parents[1] / "static"
+           / "he-detail.js").read_text(encoding="utf-8")
+    assert "ProtocolBanner" in src
+    assert 'mode !== "legacy"' in src, "the banner is legacy-only"
+    assert "summary && summary.protocol" in src, "it reads the summary payload"
diff --git a/metainfer/tasks/harness_evolve/tests/test_live_dkao_mirror.py b/metainfer/tasks/harness_evolve/tests/test_live_dkao_mirror.py
new file mode 100644
index 00000000..7b2d3422
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_live_dkao_mirror.py
@@ -0,0 +1,125 @@
+"""The AHE page embeds the DKAO task view inside each iteration.
+
+The two halves of that feature live in different plugins, so nothing about it
+is checked by either plugin's own tests:
+
+  * the AHE detail module imports ``app/dkao-detail`` (the *other* plugin's
+    bundle), which only works if the importmap carries that key -- the JS
+    itself cannot be executed here (no build step, no browser), so the wiring
+    is asserted where it is decidable: the module the server actually serves,
+    the importmap it is served with, and the bundle's own embed interface;
+  * the embedded view is addressed by the child's artifact directory, because
+    HE children are spawned by the orchestrator and are not tasks.
+
+Deliberately narrow: this does not try to be a JS test suite.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+
+HE_DETAIL = ("metainfer/tasks/harness_evolve/static/he-detail.js")
+DKAO_DETAIL = "metainfer/tasks/dcu_kernel_auto_opt/static/dkao-detail.js"
+
+
+def _source(rel: str) -> str:
+    from pathlib import Path
+    return (Path(__file__).resolve().parents[4] / rel).read_text(encoding="utf-8")
+
+
+def _importmap(client) -> dict:
+    resp = client.get("/")
+    assert resp.status_code == 200
+    block = re.search(
+        r']*type="importmap"[^>]*>(.*?)', resp.text,
+        re.DOTALL)
+    assert block, "the shell no longer ships an importmap"
+    return json.loads(block.group(1)).get("imports", {})
+
+
+def test_the_ahe_page_resolves_the_dkao_view_it_imports(client):
+    """`app/dkao-detail` must exist, or the AHE page dies on import."""
+    imports = _importmap(client)
+    assert "app/dkao-detail" in imports, sorted(imports)
+    url = imports["app/dkao-detail"]
+    assert url.startswith("/static/plugins/dcu-kernel-auto-opt/")
+    assert "CACHE_BUST" not in url            # substituted, not leaked
+    asset = client.get(url.split("?")[0])
+    assert asset.status_code == 200, url
+    assert "DcuKernelAutoOptDetail" in asset.text
+
+
+def test_he_detail_actually_imports_and_mounts_the_dkao_view():
+    src = _source(HE_DETAIL)
+    assert re.search(
+        r'import\s+DcuKernelAutoOptDetail\s+from\s+"app/dkao-detail"', src), \
+        "he-detail.js no longer imports the DKAO view"
+    mount = re.search(r"<\$\{DcuKernelAutoOptDetail\}([^/]*)/>", src)
+    assert mount, "he-detail.js never renders the DKAO view"
+    # A mirrored child is read-only and located by its artifact directory;
+    # without these the embedded calls would 404 or offer write controls.
+    assert "stateDir=" in mount.group(1)
+    assert "readOnly=${true}" in mount.group(1)
+
+
+def test_the_dkao_view_supports_the_embed_contract():
+    """The other side of the same contract, in the DKAO bundle."""
+    src = _source(DKAO_DETAIL)
+    assert "export function apiBase(" in src, \
+        "dkao-detail.js no longer exports the embed URL builder"
+    assert "state_dir=" in src, "apiBase no longer passes state_dir"
+    assert "readOnly" in src, "dkao-detail.js dropped its read-only mode"
+    # The write controls must be gated on it, not merely styled.
+    for control in ("onClick=${() => addVariant(shape)}",
+                    "onClick=${renameRepo}",
+                    "onClick=${syncNow}"):
+        idx = src.index(control)
+        window = src[max(0, idx - 900):idx]
+        assert "readOnly" in window, f"{control} is not gated by readOnly"
+
+
+def test_the_iteration_block_hands_the_task_down_to_the_mirror():
+    """Without taskId the mirror cannot fetch the iteration's children."""
+    src = _source(HE_DETAIL)
+    assert re.search(r"<\$\{IterBlock\}\s+iter=\$\{it\}\s+taskId=\$\{taskId\}", src), \
+        "IterBlock is rendered without taskId"
+    assert re.search(r"<\$\{LiveDkaoMirror\}\s+taskId=\$\{taskId\}", src), \
+        "IterBlock does not mount the live DKAO mirror"
+
+
+def test_the_mirror_follows_a_running_question_by_default():
+    src = _source(HE_DETAIL)
+    body = src[src.index("function LiveDkaoMirror("):]
+    body = body[:body.index("\nfunction IterBlock(")]
+    assert "!children[id].finished" in body, \
+        "the switcher no longer auto-selects the running question"
+    assert "getChildren(taskId, num)" in body, \
+        "the mirror no longer reads the iteration's children"
+
+
+# ------------------------------------------------------------------ stop button
+
+
+def test_the_stop_button_is_wired_to_a_running_task():
+    """Stop must be reachable, and only while there is something to stop."""
+    src = _source(HE_DETAIL)
+    # `status` is the shell's liveness view; without it in the props the
+    # button silently never renders (it is passed, but must be destructured).
+    assert re.search(r"export default function HeDetailView\(\{[^}]*\bstatus\b",
+                     src), "HeDetailView does not accept the shell's status prop"
+    assert re.search(r"import\s*\{[^}]*stopExperiment[^}]*\}\s*from\s*\"app/he-api\"",
+                     src), "he-detail.js does not import stopExperiment"
+    assert "stopExperiment(taskId)" in src, "the button calls nothing"
+    # Offered exactly when the orchestrator is alive, not when it merely has
+    # artifacts on disk (a stopped run keeps its current_phase).
+    assert "status && status.running" in src, \
+        "the Stop button is not gated on the orchestrator being alive"
+
+
+def test_the_stop_api_posts_to_the_stop_route():
+    src = _source("metainfer/tasks/harness_evolve/static/he-api.js")
+    body = src[src.index("export async function stopExperiment"):]
+    body = body[:body.index("\n}")]
+    assert "${TASK_BASE(taskId)}/stop" in body
+    assert 'method: "POST"' in body
diff --git a/metainfer/tasks/harness_evolve/tests/test_live_mirror_layouts.py b/metainfer/tasks/harness_evolve/tests/test_live_mirror_layouts.py
new file mode 100644
index 00000000..b858e08a
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_live_mirror_layouts.py
@@ -0,0 +1,64 @@
+"""The DKAO live view shows the round's four questions at once.
+
+A round is four questions on four cards (FLOW.md §1.1), so comparing them is the
+point of the live view. It used to render exactly one: four tabs, one DKAO page,
+three quarters of the round hidden behind a click.
+
+These are source-level checks, the same way the rest of this plugin's front-end
+contracts are pinned (the repo ships no JS test runner): the component must
+render every child by default, keep the single view reachable per child, and the
+grid must collapse on a narrow screen instead of squeezing four pages.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+HE = Path(__file__).resolve().parents[1]
+JS = HE / "static" / "he-detail.js"
+CSS = HE / "static" / "he.css"
+
+
+def _component() -> str:
+    src = JS.read_text(encoding="utf-8")
+    return src.split("function LiveDkaoMirror(", 1)[1].split(
+        "function IterBlock(", 1)[0]
+
+
+def test_every_child_is_rendered_by_default():
+    component = _component()
+    assert 'useState("grid")' in component, "four at once is the default"
+    assert "ids.map((id) =>" in component, "every question gets a cell"
+    assert "he-live-cell" in component
+    assert "mirror(id, c)" in component, (
+        "each cell renders that question's own DKAO page")
+    assert "taskId=${id} stateDir=${c.child_state_dir}" in component, (
+        "…pointed at its own state dir, so the four are independent")
+
+
+def test_the_three_layouts_are_offered():
+    component = _component()
+    for key in ('"grid"', '"stack"', '"single"'):
+        assert key in component, f"missing layout {key}"
+    assert "setLayout(key)" in component
+
+
+def test_the_single_view_is_still_one_click_away_per_child():
+    component = _component()
+    assert 'setLayout("single")' in component, (
+        "drilling into one question must stay possible from the grid")
+
+
+def test_a_question_without_a_state_dir_is_skipped_not_rendered_empty():
+    component = _component()
+    assert "if (!c.child_state_dir) return null;" in component
+
+
+def test_the_grid_is_two_columns_and_collapses_on_narrow_screens():
+    css = CSS.read_text(encoding="utf-8")
+    block = css.split(".he-live-grid {", 1)[1].split("}", 1)[0]
+    assert "repeat(2" in block
+    tail = css.split(".he-live-grid {", 1)[1][:600]
+    assert "@media" in tail, "narrow screens must fall back to one column"
+    assert ".he-live-cell-body" in css and "overflow-x: auto" in css, (
+        "a 50%-wide cell must scroll its tables, not clip them")
diff --git a/metainfer/tasks/harness_evolve/tests/test_liveness.py b/metainfer/tasks/harness_evolve/tests/test_liveness.py
new file mode 100644
index 00000000..cff96bd9
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_liveness.py
@@ -0,0 +1,277 @@
+"""A live HE run must look live in the WebUI (the green dot).
+
+The shell owns the "running" indicator and derives it from one thing:
+``Launcher.status()``, which validates the pid recorded in the task's
+``orchestrator.pid``. harness_evolve used to write no such file, so a run in
+progress was reported as stopped for its whole life. These tests pin the
+contract to the shell's own reader rather than to a private helper, so a
+regression makes the green dot disappear again instead of silently passing.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import pytest
+
+from metainfer.server import launcher as launcher_mod
+from metainfer.server import paths as paths_mod
+from metainfer.server import proc as proc_mod
+from metainfer.tasks.harness_evolve.orchestrator import cli as he_cli
+from metainfer.testing import isolated_env  # noqa: F401 - used as a fixture
+
+
+@pytest.fixture
+def shell(monkeypatch):
+    """The shell's *real* status reader (the fixture swaps in a fake one).
+
+    The whole point of these tests is that a live HE run is visible to the
+    code the WebUI actually runs, so a stand-in for ``LocalLauncher`` would
+    test the stand-in. ``METAINFER_ROOT`` isolation from ``isolated_env`` is
+    still used, which is why the fixture is requested above.
+
+    One measurement is neutralised: ``LocalLauncher.status`` validates a pid by
+    comparing the kernel's process start time -- derived from ``/proc`` ticks
+    plus the *recorded* boot time -- against the wall-clock ``started_at`` in
+    the PID file, with a 2 s tolerance. Both are the same clock in production.
+    This sandbox's recorded boot time is off from its wall clock by a couple of
+    seconds *and the offset moves* (measured: 0.23 s, then 2.4 s three seconds
+    later), so a perfectly healthy process lands past the tolerance and the
+    shell reports ``pid-dead`` -- a property of this host, not of harness_evolve.
+
+    The fixture re-derives that offset on every call, immediately before the
+    comparison, and subtracts it. The real code path, the real tolerance and
+    the real wall-clock ``started_at`` from the PID file are all still used.
+    """
+    real_start_time = proc_mod.pid_start_time
+    def _start_time(pid):
+        value = real_start_time(pid)
+        if value is None:
+            return None
+        own = real_start_time(os.getpid())
+        if own is None:          # /proc unavailable: leave the value alone
+            return value
+        # Re-measured here, in the same instant as the comparison, because this
+        # host's offset drifts.
+        return float(value) + (time.time() - float(own))
+
+    monkeypatch.setattr(proc_mod, "pid_start_time", _start_time)
+
+    def _status(task_id: str):
+        return launcher_mod.LocalLauncher().status(task_id)
+    return _status
+
+
+def _requirements(tmp_path: Path, task_id: str, *, rounds: int = 1) -> Path:
+    req = tmp_path / "requirements.json"
+    req.write_text(json.dumps({
+        "task_id": task_id,
+        "task_type": "harness-evolve",
+        "execution_mode": "dry-run",
+        "max_iterations": rounds,
+        "evolve_mode": "dry-run",
+        "per_round_budget": 2,
+        "pool_source": "",
+    }), encoding="utf-8")
+    return req
+
+
+def _pid_file(task_id: str) -> Path:
+    """Where the shell looks for it (``server.paths.task_dir/...``)."""
+    return paths_mod.task_dir(task_id) / "orchestrator.pid"
+
+
+def test_run_stamps_the_pid_file_the_shell_reads(
+        isolated_env, shell, tmp_path, monkeypatch):
+    """The server's own status reader must see a live run while it runs."""
+    task_id = "he-live"
+    state_dir = _pid_file(task_id).parent
+    cfg = he_cli.load_experiment_config(
+        _requirements(tmp_path, task_id), state_dir, tmp_path / "ws")
+
+    seen: dict = {}
+
+    def fake_run(cfg_, *args, **kwargs):
+        # Exactly what the shell does while the run is in progress.
+        status = shell(task_id).to_dict()
+        seen.update(status)
+        seen["payload"] = json.loads(
+            (state_dir / "orchestrator.pid").read_text(encoding="utf-8"))
+        return tmp_path / "ws" / "report.md"
+
+    monkeypatch.setattr(he_cli, "load_experiment_config", lambda *a, **k: cfg)
+    monkeypatch.setattr(he_cli, "run_experiment", fake_run)
+
+    assert he_cli.run_with_requirements(
+        _requirements(tmp_path, task_id), state_dir=state_dir,
+        workspace_dir=tmp_path / "ws") == 0
+
+    # live: the shell must have seen a running task with a validated pid
+    assert seen["running"] is True, ("SHELLVIEW=%r" % (seen,))
+    assert seen["pid"] == os.getpid()
+    assert seen["exit_hint"] == "pid-alive"
+    assert seen["payload"]["task_id"] == task_id     # not the directory name
+    assert isinstance(seen["payload"]["started_at"], float)
+
+
+def test_the_pid_file_is_cleared_when_the_run_ends(
+        isolated_env, shell, tmp_path, monkeypatch):
+    """A finished run must not claim a pid, and must say when it finished."""
+    task_id = "he-finished"
+    state_dir = _pid_file(task_id).parent
+    cfg = he_cli.load_experiment_config(
+        _requirements(tmp_path, task_id), state_dir, tmp_path / "ws")
+    monkeypatch.setattr(he_cli, "load_experiment_config", lambda *a, **k: cfg)
+    monkeypatch.setattr(he_cli, "run_experiment",
+                        lambda cfg_, *a, **k: tmp_path / "report.md")
+
+    he_cli.run_with_requirements(_requirements(tmp_path, task_id),
+                                 state_dir=state_dir,
+                                 workspace_dir=tmp_path / "ws")
+
+    status = shell(task_id).to_dict()
+    assert status["running"] is False
+    assert status["pid"] is None
+    assert status["exit_hint"] == "pid-file-cleared"
+    assert isinstance(status["finished_at"], float)
+
+
+def test_a_crashing_run_still_clears_its_pid_file(
+        isolated_env, shell, tmp_path, monkeypatch):
+    """An exception must not leave a stale "running" claim behind."""
+    task_id = "he-crash"
+    state_dir = _pid_file(task_id).parent
+    cfg = he_cli.load_experiment_config(
+        _requirements(tmp_path, task_id), state_dir, tmp_path / "ws")
+
+    def boom(cfg_, *a, **k):
+        raise RuntimeError("pipeline exploded")
+
+    monkeypatch.setattr(he_cli, "load_experiment_config", lambda *a, **k: cfg)
+    monkeypatch.setattr(he_cli, "run_experiment", boom)
+
+    try:
+        he_cli.run_with_requirements(_requirements(tmp_path, task_id),
+                                     state_dir=state_dir,
+                                     workspace_dir=tmp_path / "ws")
+    except RuntimeError:
+        pass
+    status = shell(task_id).to_dict()
+    assert status["running"] is False and status["finished_at"] is not None
+
+
+def test_sigterm_clears_the_pid_file_before_exiting(
+        isolated_env, tmp_path, monkeypatch):
+    """The WebUI's Stop button sends SIGTERM: it must mark the task stopped."""
+    task_id = "he-sigterm"
+    state_dir = _pid_file(task_id).parent
+    pid_file = state_dir / "orchestrator.pid"
+    state_dir.mkdir(parents=True, exist_ok=True)
+
+    # The only thing a handler can safely do here is the real one: the test
+    # process would die on os._exit, so drive the handler body directly.
+    exited: list = []
+    monkeypatch.setattr(he_cli.os, "_exit", lambda code: exited.append(code))
+    captured: dict = {}
+    monkeypatch.setattr(he_cli.signal, "signal",
+                        lambda sig, fn: captured.__setitem__(sig, fn))
+
+    with he_cli._orchestrator_pid(state_dir, task_id):
+        assert pid_file.is_file()
+        handler = captured[he_cli.signal.SIGTERM]
+        handler(he_cli.signal.SIGTERM, None)
+
+    assert exited == [143]
+    assert json.loads(pid_file.read_text(encoding="utf-8"))["pid"] is None
+    assert json.loads(pid_file.read_text(encoding="utf-8"))["finished_at"]
+
+
+def test_resume_also_marks_the_task_live(
+        isolated_env, shell, tmp_path, monkeypatch):
+    """A supervisor-started resume is just as live as a fresh run."""
+    task_id = "he-resume"
+    state_dir = _pid_file(task_id).parent
+    state_dir.mkdir(parents=True, exist_ok=True)
+    req = state_dir / "requirements.json"
+    req.write_text(json.dumps({
+        "task_id": task_id, "task_type": "harness-evolve",
+        "execution_mode": "dry-run", "max_iterations": 1,
+        "evolve_mode": "dry-run", "pool_source": "",
+    }), encoding="utf-8")
+    cfg = he_cli.load_experiment_config(req, state_dir, tmp_path / "ws")
+
+    seen: dict = {}
+
+    def fake_run(cfg_, *a, **k):
+        seen.update(shell(task_id).to_dict())
+        return tmp_path / "report.md"
+
+    monkeypatch.setattr(he_cli, "load_experiment_config", lambda *a, **k: cfg)
+    monkeypatch.setattr(he_cli, "completed_iterations", lambda exp: 1)
+    monkeypatch.setattr(he_cli, "run_experiment", fake_run)
+
+    args = type("Args", (), {
+        "state_dir": str(state_dir), "workspace_dir": str(tmp_path / "ws"),
+        "requirements": None, "iterations": 1, "target": None,
+        "champion_iteration": None,
+    })()
+    assert he_cli.resume_command(args) == 0
+    assert seen["running"] is True and seen["exit_hint"] == "pid-alive", ("SHELLVIEW=%r" % (seen,))
+    assert shell(task_id).to_dict()["running"] is False
+
+
+def test_the_real_cli_binary_writes_and_clears_the_pid_file(
+        isolated_env, shell, tmp_path):
+    """End to end, with no mocks: the actual process the server spawns."""
+    task_id = "he-real-cli"
+    state_dir = _pid_file(task_id).parent
+    ws = tmp_path / "ws"
+    req = _requirements(tmp_path, task_id)
+
+    # The server spawns the orchestrator with the package importable (its
+    # launcher prepends ``sys.path`` to PYTHONPATH); do the same so this is a
+    # faithful copy of that spawn rather than a plain pytest environment.
+    repo_root = Path(__file__).resolve().parents[4]   # /metainfer -> 
+    env = dict(os.environ)
+    env["PYTHONPATH"] = os.pathsep.join(
+        [str(repo_root)] + [p for p in env.get("PYTHONPATH", "").split(os.pathsep)
+                            if p])
+    proc = subprocess.Popen(
+        [sys.executable, "-m",
+         "metainfer.tasks.harness_evolve.orchestrator.cli", "run",
+         str(req), "--state-dir", str(state_dir), "--workspace-dir", str(ws)],
+        cwd=str(repo_root), env=env,
+        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    pid_file = state_dir / "orchestrator.pid"
+    try:
+        deadline = time.time() + 60
+        while not pid_file.is_file() and time.time() < deadline:
+            if proc.poll() is not None:
+                break
+            time.sleep(0.01)
+        # The process identity the shell needs is in the file the real binary
+        # wrote: the orchestrator's own pid and the wall-clock start it stamped.
+        # (Whether the *shell* then calls it live is asserted in the tests
+        # above, where the host's clock offset can be cancelled out; doing that
+        # here would mean fighting a sandbox clock that jumps by seconds.)
+        if pid_file.is_file():
+            payload = json.loads(pid_file.read_text(encoding="utf-8"))
+            assert payload["pid"] == proc.pid
+            assert payload["task_id"] == task_id
+            assert float(payload["started_at"]) <= time.time() + 60
+    finally:
+        out, _ = proc.communicate(timeout=120)
+
+    assert proc.returncode == 0, f"rc={proc.returncode} out={out!r}"
+
+    # Once the process is gone the verdict is deterministic -- no clock needed.
+    status = shell(task_id).to_dict()
+    assert status["running"] is False
+    assert status["pid"] is None and status["exit_hint"] == "pid-file-cleared"
+    cleared = json.loads(pid_file.read_text(encoding="utf-8"))
+    assert cleared["pid"] is None and isinstance(cleared["finished_at"], float)
diff --git a/metainfer/tasks/harness_evolve/tests/test_mechanism_gate.py b/metainfer/tasks/harness_evolve/tests/test_mechanism_gate.py
new file mode 100644
index 00000000..9dc136b6
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_mechanism_gate.py
@@ -0,0 +1,293 @@
+"""Mechanism gate: evidence grading + lenient/strict decision behaviour."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from metainfer.tasks.harness_evolve.orchestrator.decision_engine import (
+    DecisionPolicy, decide,
+)
+from metainfer.tasks.harness_evolve.orchestrator.mechanism_checks import (
+    collect_candidate_evidence, evaluate_changes,
+)
+
+WIN = {"a": {"passed": True, "median_us": 100.0, "correctness_ok": True,
+             "new_best_known": True},
+       "b": {"passed": True, "median_us": 100.0, "correctness_ok": True,
+             "new_best_known": True}}
+CAND = {"a": {"passed": True, "median_us": 50.0, "correctness_ok": True},
+        "b": {"passed": True, "median_us": 50.0, "correctness_ok": True}}
+
+
+def _manifest(check=None, change_id="chg-1"):
+    change = {"id": change_id, "component": "planner_policy",
+              "mechanism_signature": ["fallback.m16 contains epilogue_fusion"]}
+    if check is not None:
+        change["mechanism_check"] = check
+    return {"schema_version": 1, "iteration": 1, "changes": [change],
+            "verification": {"status": "pending"}}
+
+
+# ---------------------------------------------------------------- grader ----
+
+def test_plan_ids_hit_partial_and_contradiction():
+    check = {"kind": "planner_plan_ids",
+             "expect_any": ["memory_layout", "epilogue_fusion"]}
+    hit = {"planner_plan_ids": ["architecture_explore", "memory_layout"],
+           "planner_enabled": True}
+    part = {"planner_plan_ids": ["architecture_explore", "grid_splitk"],
+            "planner_enabled": True}
+    contra = {"planner_plan_ids": [], "planner_enabled": False}
+    unob = {"planner_plan_ids": [], "planner_enabled": True}
+
+    assert evaluate_changes(_manifest(check), hit)[0]["chg-1"] == "hit"
+    assert evaluate_changes(_manifest(check), part)[0]["chg-1"] == "partial"
+    assert evaluate_changes(_manifest(check), contra)[0]["chg-1"] == "contradicted"
+    assert evaluate_changes(_manifest(check), unob)[0]["chg-1"] == "unobserved"
+
+
+def test_expected_plan_order_hit():
+    check = {"kind": "planner_plan_ids",
+             "expect_order": ["architecture_explore", "memory_layout"]}
+    ev = {"planner_plan_ids": ["architecture_explore", "grid_splitk",
+                               "memory_layout"],
+          "planner_enabled": True}
+    assert evaluate_changes(_manifest(check), ev)[0]["chg-1"] == "hit"
+
+
+def test_harness_revision_grading():
+    check = {"kind": "harness_revision"}
+    ok = {"candidate_harness_revision": "rev2",
+          "observed_harness_revisions": ["rev2"]}
+    wrong = {"candidate_harness_revision": "rev2",
+             "observed_harness_revisions": ["rev1"]}
+    none = {"candidate_harness_revision": "rev2",
+            "observed_harness_revisions": []}
+    assert evaluate_changes(_manifest(check), ok)[0]["chg-1"] == "hit"
+    assert evaluate_changes(_manifest(check), wrong)[0]["chg-1"] == "contradicted"
+    assert evaluate_changes(_manifest(check), none)[0]["chg-1"] == "unobserved"
+
+
+def test_agent_only_plan_ids_are_soft_evidence():
+    """Without a planner record, matching ids are only partial evidence."""
+    check = {"kind": "planner_plan_ids", "expect_any": ["memory_layout"]}
+    ev = {"agent_plan_ids": ["memory_layout"], "planner_enabled": True}
+    checks, detail = evaluate_changes(_manifest(check), ev)
+    assert checks["chg-1"] == "partial"
+    assert "agent" in detail["changes"][0]["reason"]
+
+
+def test_missing_check_stays_unobserved():
+    checks, detail = evaluate_changes(_manifest(None), {"plan_ids": []})
+    assert checks["chg-1"] == "unobserved"
+    assert "mechanism_check" in detail["changes"][0]["reason"]
+
+
+def test_evidence_collection_reads_children(tmp_path):
+    exp = tmp_path / "exp"
+    ws = exp / "children" / "iteration_001" / "shapeA" / "workspace"
+    runs = ws / "workers" / "worker_0" / "runs" / "shapeA"
+    runs.mkdir(parents=True)
+    (runs / "experiments.jsonl").write_text(
+        json.dumps({"iteration": 1, "plan_id": "memory_layout"}) + "\n"
+        + json.dumps({"iteration": 2, "plan_id": "epilogue_fusion"}) + "\n",
+        encoding="utf-8")
+    (ws / "main").mkdir(parents=True)
+    (ws / "main" / "scaffold_manifest.json").write_text(
+        json.dumps({"harness": {"revision": "rev-cand"}}), encoding="utf-8")
+    tl = exp / "children" / "iteration_001" / "shapeA" / "state"
+    tl.mkdir(parents=True)
+    (tl / "timeline.jsonl").write_text(
+        json.dumps({"ts": 1.0, "type": "gate_blocked", "payload": {}}) + "\n",
+        encoding="utf-8")
+    snap = exp / "runs" / "iteration_001" / "input" / "workspace"
+    snap.mkdir(parents=True)
+    (snap / "manifest.yaml").write_text("revision: rev-cand\n", encoding="utf-8")
+
+    (ws / "workers" / "worker_0" / "planner_plans.jsonl").write_text(
+        json.dumps({"iteration": 1, "plan_id": "memory_layout",
+                    "source": "planner"}) + "\n", encoding="utf-8")
+
+    ev = collect_candidate_evidence(exp, 1, planner_enabled=True)
+    assert ev["agent_plan_ids"] == ["memory_layout", "epilogue_fusion"]
+    assert ev["planner_plan_ids"] == ["memory_layout"]
+    assert ev["plan_id_source"] == "planner"
+    assert ev["observed_harness_revisions"] == ["rev-cand"]
+    assert ev["candidate_harness_revision"] == "rev-cand"
+    assert "gate_blocked" in ev["gate_events"]
+
+
+# --------------------------------------------------------------- decision ----
+
+def test_lenient_promotes_unobserved_win():
+    d = decide(WIN, CAND, manifest=_manifest({"kind": "planner_plan_ids",
+                                              "expect_any": ["memory_layout"]}),
+               mechanism_checks={"chg-1": "unobserved"},
+               policy=DecisionPolicy(mechanism_policy="lenient"))
+    assert d["verdict"] == "PROMOTE_UNEXPLAINED"
+    assert d["action"] == "SET_CHAMPION"
+    assert d["mechanism_gate"]["status"] == "UNOBSERVED"
+
+
+def test_lenient_still_confirms_on_contradiction():
+    d = decide(WIN, CAND, manifest=_manifest({"kind": "planner_plan_ids",
+                                              "expect_any": ["memory_layout"]}),
+               mechanism_checks={"chg-1": "contradicted"},
+               policy=DecisionPolicy(mechanism_policy="lenient"))
+    assert d["verdict"] == "CONFIRM_REQUIRED"
+    assert d["mechanism_gate"]["contradicted"] == ["chg-1"]
+
+
+def test_contradiction_promotes_after_reproduction():
+    d = decide(WIN, CAND, manifest=_manifest({"kind": "harness_revision"}),
+               mechanism_checks={"chg-1": "contradicted"},
+               confirmation_reproduced=True,
+               policy=DecisionPolicy(mechanism_policy="lenient"))
+    assert d["verdict"] == "PROMOTE_UNEXPLAINED"
+
+
+def test_hit_promotes_and_strict_policy_confirms_unobserved():
+    hit = decide(WIN, CAND, manifest=_manifest({"kind": "harness_revision"}),
+                 mechanism_checks={"chg-1": "hit"},
+                 policy=DecisionPolicy(mechanism_policy="lenient"))
+    assert hit["verdict"] == "PROMOTE"
+    strict = decide(WIN, CAND,
+                    manifest=_manifest({"kind": "harness_revision"}),
+                    mechanism_checks={"chg-1": "unobserved"},
+                    policy=DecisionPolicy(mechanism_policy="strict"))
+    assert strict["verdict"] == "CONFIRM_REQUIRED"
+
+
+def test_off_policy_ignores_mechanism():
+    d = decide(WIN, CAND, manifest=_manifest({"kind": "harness_revision"}),
+               mechanism_checks={"chg-1": "contradicted"},
+               policy=DecisionPolicy(mechanism_policy="off"))
+    assert d["verdict"] == "PROMOTE_UNEXPLAINED"
+
+
+# ------------------------------------------------------- generalization ----
+
+def test_probe_regression_blocks_promotion():
+    """A win in the tuned families must not promote if an untouched one broke."""
+    champ = {"a": {"passed": True, "median_us": 100.0, "correctness_ok": True},
+             "b": {"passed": True, "median_us": 100.0, "correctness_ok": True},
+             "p": {"passed": True, "median_us": 100.0, "correctness_ok": True}}
+    cand = {"a": {"passed": True, "median_us": 50.0, "correctness_ok": True},
+            "b": {"passed": True, "median_us": 50.0, "correctness_ok": True},
+            "p": {"passed": True, "median_us": 300.0, "correctness_ok": True}}
+    purposes = {"a": "regression", "b": "repair", "p": "probe"}
+    d = decide(champ, cand, purposes=purposes)
+    assert d["generalization_gate"]["status"] == "FAIL"
+    assert d["verdict"] == "CONFIRM_REQUIRED"
+    assert "over-fitting" in d["reason"]
+    # the probe is not part of the paired decision set
+    assert "p" not in d["paired_ids"] and d["probe_ids"] == ["p"]
+    # reproduces -> allowed through, still recorded
+    d2 = decide(champ, cand, purposes=purposes, confirmation_reproduced=True)
+    assert d2["verdict"] == "PROMOTE_UNEXPLAINED"
+    assert d2["generalization_gate"]["status"] == "FAIL"
+
+
+def test_probe_pass_allows_normal_promotion():
+    champ = {"a": {"passed": True, "median_us": 100.0, "correctness_ok": True},
+             "b": {"passed": True, "median_us": 100.0, "correctness_ok": True},
+             "p": {"passed": True, "median_us": 100.0, "correctness_ok": True}}
+    cand = {"a": {"passed": True, "median_us": 50.0, "correctness_ok": True},
+            "b": {"passed": True, "median_us": 50.0, "correctness_ok": True},
+            "p": {"passed": True, "median_us": 98.0, "correctness_ok": True}}
+    d = decide(champ, cand,
+               purposes={"a": "regression", "b": "repair", "p": "probe"})
+    assert d["generalization_gate"]["status"] == "PASS"
+    assert d["verdict"] == "PROMOTE_UNEXPLAINED"   # lenient default
+    assert d["purpose_breakdown"]["p"]["counted"] is False
+
+
+# ------------------------------------------------------- gate_values ------
+
+def _gate_change(paths=("round_acceptance_improvement_percent",)):
+    return _manifest({"kind": "gate_values", "expect_any": list(paths)})
+
+
+def test_gate_values_hit_when_children_ran_with_declared_values():
+    ev = {
+        "candidate_gates": {"wired": True,
+                            "gates": {"round_acceptance_improvement_percent": 2.5}},
+        "observed_gates": [
+            {"gates": {"round_acceptance_improvement_percent": 2.5}},
+        ],
+    }
+    checks, detail = evaluate_changes(_gate_change(), ev)
+    assert checks["chg-1"] == "hit"
+    assert "declared values" in detail["changes"][0]["reason"]
+
+
+def test_gate_values_contradicted_when_values_differ():
+    ev = {
+        "candidate_gates": {"wired": True,
+                            "gates": {"round_acceptance_improvement_percent": 2.5}},
+        "observed_gates": [
+            {"gates": {"round_acceptance_improvement_percent": 1.0}},
+        ],
+    }
+    checks, _ = evaluate_changes(_gate_change(), ev)
+    assert checks["chg-1"] == "contradicted"
+
+
+def test_gate_values_contradicted_when_component_is_unwired():
+    ev = {"candidate_gates": {"wired": False, "gates": {}},
+          "observed_gates": [{"gates": {}}]}
+    checks, detail = evaluate_changes(_gate_change(), ev)
+    assert checks["chg-1"] == "contradicted"
+    assert "not wired" in detail["changes"][0]["reason"]
+
+
+def test_gate_values_unobserved_without_snapshot():
+    ev = {"candidate_gates": {"wired": True, "gates": {}}, "observed_gates": []}
+    checks, _ = evaluate_changes(_gate_change(), ev)
+    assert checks["chg-1"] == "unobserved"
+
+
+# ------------------------------------------------- invalid / inconclusive ----
+
+def test_missing_measurement_is_invalid_not_hard_fail():
+    """A crashed child must not look like a performance regression."""
+    champ = {"a": {"passed": True, "median_us": 100.0, "correctness_ok": True},
+             "b": {"passed": True, "median_us": 100.0, "correctness_ok": True}}
+    cand = {"a": {"passed": False, "median_us": None, "correctness_ok": False,
+                  "status": "missing_report"},
+            "b": {"passed": True, "median_us": 90.0, "correctness_ok": True}}
+    d = decide(champ, cand)
+    assert d["performance_gate"]["counts"]["INVALID"] == 1
+    assert d["performance_gate"]["counts"]["HARD_FAIL"] == 0
+
+
+def test_mostly_unmeasurable_round_is_inconclusive():
+    champ = {k: {"passed": True, "median_us": 100.0, "correctness_ok": True}
+             for k in "abcd"}
+    cand = {
+        "a": {"passed": False, "median_us": None, "status": "missing_report"},
+        "b": {"passed": False, "median_us": None, "status": "missing_report"},
+        "c": {"passed": False, "median_us": None, "status": "missing_report"},
+        "d": {"passed": True, "median_us": 100.0, "correctness_ok": True},
+    }
+    d = decide(champ, cand)
+    assert d["verdict"] == "INCONCLUSIVE"
+    assert d["action"] == "KEEP_CHAMPION_ARCHIVE_CANDIDATE"
+    assert d["performance_gate"]["status"] == "INCONCLUSIVE"
+    assert "no verdict" in d["reason"]
+
+
+def test_measurable_regression_still_rejects():
+    """Sanity: a genuine, fully measured regression still rejects.
+
+    Two losses are required: ``max_loss_ratio`` deliberately tolerates a single
+    losing question so one noisy shape cannot trigger a rollback.
+    """
+    champ = {k: {"passed": True, "median_us": 100.0, "correctness_ok": True}
+             for k in ("a", "b")}
+    cand = {k: {"passed": True, "median_us": 300.0, "correctness_ok": True}
+            for k in ("a", "b")}
+    d = decide(champ, cand)
+    assert d["verdict"] == "REJECT"
+    assert d["performance_gate"]["counts"]["LOSS"] == 2
diff --git a/metainfer/tasks/harness_evolve/tests/test_one_driver.py b/metainfer/tasks/harness_evolve/tests/test_one_driver.py
new file mode 100644
index 00000000..95c5def1
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_one_driver.py
@@ -0,0 +1,240 @@
+"""One experiment, one driver — and "finished" means finished.
+
+The sweep in the server exists so an experiment whose supervisor died does not
+sit idle forever. It got that wrong in two independent ways, and together they
+produced the reported symptom ("HE opened the next round before DKAO finished"):
+
+* it asked ``_resume_running`` whether the task was running, but a task created
+  from the New Task form never writes ``resume.pid`` — only ``orchestrator.pid``
+  — so a one-second-old, healthy run looked like an unsupervised corpse;
+* the resume it then started computed the next iteration from
+  ``completed_iterations``, which counted ``input/stage.json`` — a file written
+  when a pass *starts* — so "just began" read as "already judged" and the next
+  round opened on top of the running one.
+
+Each piece is pinned here separately, because either one alone is enough to
+double-drive an experiment.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import pytest
+
+from metainfer.tasks.harness_evolve.orchestrator import supervisor as sup
+from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+    completed_iterations,
+)
+from metainfer.tasks.harness_evolve.server import routes as R
+
+
+# ------------------------------------------------------- pid files and liveness
+
+def test_pid_from_file_reads_both_forms(tmp_path):
+    (tmp_path / "json.pid").write_text(
+        json.dumps({"pid": 4242, "task_id": "t"}), encoding="utf-8")
+    (tmp_path / "bare.pid").write_text("4242\n", encoding="utf-8")
+    (tmp_path / "empty.pid").write_text("", encoding="utf-8")
+    (tmp_path / "junk.pid").write_text("not-a-pid", encoding="utf-8")
+
+    assert sup.pid_from_file(tmp_path / "json.pid") == 4242
+    assert sup.pid_from_file(tmp_path / "bare.pid") == 4242
+    assert sup.pid_from_file(tmp_path / "empty.pid") is None
+    assert sup.pid_from_file(tmp_path / "junk.pid") is None
+    assert sup.pid_from_file(tmp_path / "missing.pid") is None
+
+
+def test_a_runner_never_finds_itself(tmp_path):
+    """The regression that made a resume stand down forever.
+
+    In production the resume's own cmdline is
+    ``python3 -m ...harness_evolve.orchestrator.cli resume --state-dir ``
+    — which matches the ``/proc`` scan this function falls back to. Answering
+    "somebody is already running this" with its own pid meant a run with no
+    completed iteration could never be restarted.
+    """
+    state = tmp_path / "state"
+    state.mkdir()
+    proc_root = tmp_path / "proc"
+    mine = 11111
+    (proc_root / str(mine)).mkdir(parents=True)
+    (proc_root / str(mine) / "cmdline").write_bytes(
+        f"python3 -m metainfer.tasks.harness_evolve.orchestrator.cli resume "
+        f"--state-dir {state} --workspace-dir {tmp_path / 'ws'} --iterations 1"
+        .encode())
+
+    # asking as *that* process: nobody else is running the experiment
+    assert sup.task_process_alive(state, proc_root=proc_root,
+                                  self_pid=mine) is None
+    # asking as somebody else: yes, that process is running it
+    assert sup.task_process_alive(state, proc_root=proc_root,
+                                  self_pid=mine + 1) == mine
+    # an explicit exclusion works too (the supervisor skipping its own child)
+    assert sup.task_process_alive(state, proc_root=proc_root, self_pid=mine + 1,
+                                  exclude=mine) is None
+
+
+def test_a_dead_pid_file_is_not_alive(tmp_path):
+    state = tmp_path / "state"
+    state.mkdir()
+    (state / "orchestrator.pid").write_text(
+        json.dumps({"pid": 999999, "task_id": "t"}), encoding="utf-8")
+    assert sup.task_process_alive(state) is None
+
+
+# ------------------------------------------------------------------- the sweep
+
+def _entry(state: Path, ws: Path, task_id: str = "t-he"):
+    class Entry:
+        id = task_id
+        type = R.PLUGIN_TYPE
+        state_dir = str(state)
+        workspace_dir = str(ws)
+
+    return Entry()
+
+
+def _run_json(state: Path, *, finished: bool = False, iteration: int = 1) -> None:
+    (state / "run.json").write_text(
+        json.dumps({"task_id": "t-he", "task_type": "harness-evolve",
+                    "current_iteration": iteration, "finished": finished}),
+        encoding="utf-8")
+
+
+@pytest.fixture()
+def sweep_env(tmp_path, monkeypatch):
+    """A janitor that would launch everything, over a task nobody runs."""
+    state = tmp_path / "state"
+    ws = tmp_path / "ws"
+    state.mkdir()
+    ws.mkdir()
+    launched: list = []
+    monkeypatch.setattr("metainfer.server.tasks.list_tasks",
+                        lambda: [_entry(state, ws)])
+    monkeypatch.setattr(R, "read_target_iterations", lambda _w: 5)
+    monkeypatch.setattr(R, "_launch_supervisor",
+                        lambda *a, **k: launched.append("supervisor") or {"pid": 11})
+    monkeypatch.setattr(R, "_launch_resume",
+                        lambda *a, **k: launched.append("resume") or {"pid": 22})
+    return {"state": state, "ws": ws, "launched": launched}
+
+
+def test_the_sweep_leaves_a_server_started_run_alone(sweep_env, monkeypatch):
+    """The exact production shape: run alive, no resume.pid, no supervisor.pid.
+
+    This is what a task looks like one second after the New Task form starts it
+    — and the old liveness check called that "nobody is supervising".
+    """
+    monkeypatch.setattr(R, "_run_alive", lambda _s: 4242)
+    monkeypatch.setattr(R, "_supervisor_running", lambda _w: None)
+    _run_json(sweep_env["state"], finished=False, iteration=1)
+
+    actions = R.janitor_sweep()
+
+    assert sweep_env["launched"] == [], (
+        "a live run must never get a second driver")
+    assert actions and actions[0]["action"] == "skip"
+    assert actions[0]["pid"] == 4242
+
+
+def test_the_sweep_recovers_a_dead_run(sweep_env, monkeypatch):
+    """Nothing alive and unfinished: that is the case the sweep exists for."""
+    monkeypatch.setattr(R, "_run_alive", lambda _s: None)
+    monkeypatch.setattr(R, "_supervisor_running", lambda _w: None)
+    _run_json(sweep_env["state"], finished=False, iteration=2)
+
+    actions = R.janitor_sweep()
+
+    assert sweep_env["launched"] == ["supervisor", "resume"]
+    assert actions[0]["action"] == "supervisor+resume"
+
+
+def test_the_sweep_restarts_only_the_run_when_a_supervisor_survives(
+        sweep_env, monkeypatch):
+    monkeypatch.setattr(R, "_run_alive", lambda _s: None)
+    monkeypatch.setattr(R, "_supervisor_running", lambda _w: 777)
+    _run_json(sweep_env["state"], finished=False, iteration=2)
+
+    actions = R.janitor_sweep()
+
+    assert sweep_env["launched"] == ["resume"]
+    assert actions[0]["action"] == "resume"
+
+
+def test_the_sweep_never_restarts_a_terminal_run(sweep_env, monkeypatch):
+    monkeypatch.setattr(R, "_run_alive", lambda _s: None)
+    monkeypatch.setattr(R, "_supervisor_running", lambda _w: None)
+    (sweep_env["state"] / "run.json").write_text(
+        json.dumps({"current_iteration": 1, "finished": True,
+                    "final_status": "round_incomplete"}), encoding="utf-8")
+
+    # a run that stopped on purpose gets no action at all: no supervisor, no
+    # resume, and it needs a human decision
+    assert R.janitor_sweep() == []
+    assert sweep_env["launched"] == []
+
+
+def test_the_sweep_skips_a_task_the_operator_stopped(sweep_env, monkeypatch):
+    (sweep_env["ws"] / "operator_stop.json").write_text(
+        json.dumps({"stopped_at": time.time()}), encoding="utf-8")
+    _run_json(sweep_env["state"], finished=False, iteration=1)
+
+    assert R.janitor_sweep()[0]["action"] == "skip"
+    assert sweep_env["launched"] == []
+
+
+# --------------------------------------------------- "finished" means finished
+
+def _pass_dir(exp: Path, iteration: int, *files: str) -> Path:
+    input_dir = exp / "runs" / f"iteration_{iteration:03d}" / "input"
+    input_dir.mkdir(parents=True, exist_ok=True)
+    for name in files:
+        path = input_dir / name
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text("{}", encoding="utf-8")
+    return input_dir
+
+
+def test_a_pass_that_only_started_is_not_a_completed_iteration(tmp_path):
+    """``stage.json`` is written when the pass opens the round, not when it ends.
+
+    Counting it made the resume launched by the sweep compute ``start = 2``
+    while iteration 1 was still being measured — two rounds on four cards.
+    """
+    exp = tmp_path / "exp"
+    _pass_dir(exp, 1, "stage.json")          # began, produced nothing yet
+
+    assert completed_iterations(exp) == 0
+
+
+def test_a_pass_that_produced_a_verdict_is_completed(tmp_path):
+    exp = tmp_path / "exp"
+    _pass_dir(exp, 1, "stage.json", "baseline_round.json")
+    _pass_dir(exp, 2, "stage.json", "performance_gate.json")
+
+    assert completed_iterations(exp) == 2
+
+    _pass_dir(exp, 3, "stage.json")          # round 3 is under way
+    assert completed_iterations(exp) == 2, (
+        "an open round must not advance the counter")
+
+
+def test_a_decided_pass_counts_even_without_a_gate_file(tmp_path):
+    exp = tmp_path / "exp"
+    _pass_dir(exp, 1, "stage.json", "decision.json")
+
+    assert completed_iterations(exp) == 1
+
+
+def test_the_generalization_retake_counts_for_its_iteration(tmp_path):
+    exp = tmp_path / "exp"
+    _pass_dir(exp, 2, "stage.json", "performance_gate.json")
+    _pass_dir(exp, 2, "generalization_gate.json")
+
+    assert completed_iterations(exp) == 2
diff --git a/metainfer/tasks/harness_evolve/tests/test_operator_stop.py b/metainfer/tasks/harness_evolve/tests/test_operator_stop.py
new file mode 100644
index 00000000..f9944867
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_operator_stop.py
@@ -0,0 +1,437 @@
+"""Stopping an AHE run must be a pause, not a death.
+
+The two are easy to confuse in this codebase: an experiment whose process is
+merely gone looks *exactly* like one that crashed, and the supervisor, the
+janitor sweep and ``should_launch_resume`` all exist to bring crashed runs back
+automatically. So a Stop button that only kills the process would restart the
+task seconds later.
+
+These tests pin the difference: Stop brings down everything that would keep the
+round moving (supervisor, DKAO children, orchestrator), records an explicit
+stop that no automatic path overrules, and leaves the artifacts intact so the
+resume paths can pick the run back up.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+
+from metainfer.server import tasks as _tasks
+from metainfer.server.state_reader import read_run
+from metainfer.server.tasks import TaskEntry
+
+from metainfer.tasks.harness_evolve.server.routes import (
+    OPERATOR_STOP_FILE,
+    _operator_stopped,
+    janitor_sweep,
+    should_launch_resume,
+)
+
+
+def _he_task(tmp_path, task_id="he-stop", *, finished=False,
+             final_status="", iteration=2, target=5, resume_pid=None):
+    state_dir = tmp_path / "state"
+    workspace = tmp_path / "ws"
+    state_dir.mkdir(parents=True, exist_ok=True)
+    workspace.mkdir(parents=True, exist_ok=True)
+    (state_dir / "requirements.json").write_text(json.dumps({
+        "task_id": task_id, "task_type": "harness-evolve",
+        "answers": {"per_round_budget": 4},
+    }), encoding="utf-8")
+    run = {"task_id": task_id, "task_type": "harness-evolve",
+           "current_phase": "evaluate", "current_iteration": iteration,
+           "finished": finished, "last_update": 1.0}
+    if final_status:
+        run["final_status"] = final_status
+    (state_dir / "run.json").write_text(json.dumps(run), encoding="utf-8")
+    if target is not None:
+        (workspace / "target_iterations.json").write_text(
+            json.dumps({"target": target}), encoding="utf-8")
+    if resume_pid is not None:
+        (state_dir / "resume.pid").write_text(str(resume_pid), encoding="utf-8")
+    _tasks.add_task(TaskEntry(
+        id=task_id, type="harness-evolve", label="AHE",
+        state_dir=str(state_dir), workspace_dir=str(workspace), created_at=0.0))
+    return state_dir, workspace
+
+
+def test_stop_marks_the_run_finished_and_terminal(tmp_path, client):
+    """A cut-short round must not look like a finished one, and must stay put."""
+    state_dir, workspace = _he_task(tmp_path)
+
+    resp = client.post("/api/harness-evolve/he-stop/stop")
+    assert resp.status_code == 200, resp.text
+    body = resp.json()
+    assert body["ok"] is True
+    assert body["iteration"] == 2
+
+    run = read_run(state_dir)
+    assert run["finished"] is True
+    assert run["final_status"] == "stopped_by_operator"
+    assert "Stopped by operator." in run["notes"]
+
+    # The marker is what stops every automatic path from reviving it.
+    marker = json.loads((workspace / OPERATOR_STOP_FILE).read_text())
+    assert marker["iteration"] == 2
+    assert marker["phase"] == "evaluate"
+    assert marker["by"] == "webui"
+
+
+def test_a_stopped_run_is_terminal_and_not_restarted(tmp_path, client):
+    """should_launch_resume must refuse the exact state Stop produces."""
+    state_dir, _ = _he_task(tmp_path)
+    client.post("/api/harness-evolve/he-stop/stop")
+    assert should_launch_resume(read_run(state_dir), 5, False) is False
+
+
+def test_a_crashed_run_is_still_recovered(tmp_path, client):
+    """The guard must not disarm crash recovery: only Stop is special."""
+    state_dir, _ = _he_task(tmp_path, finished=False)
+    assert should_launch_resume(read_run(state_dir), 5, False) is True
+
+
+def test_the_janitor_never_restarts_a_stopped_run(tmp_path, client,
+                                                 monkeypatch):
+    """The sweep that exists to resurrect dead runs must respect a human."""
+    _he_task(tmp_path)
+    calls = []
+    monkeypatch.setattr(
+        "metainfer.tasks.harness_evolve.server.routes._launch_supervisor",
+        lambda *a, **k: calls.append("supervisor") or {"pid": 1})
+    monkeypatch.setattr(
+        "metainfer.tasks.harness_evolve.server.routes._launch_resume",
+        lambda *a, **k: calls.append("resume") or {"pid": 1})
+
+    assert calls == []
+    client.post("/api/harness-evolve/he-stop/stop")
+    actions = janitor_sweep()
+    assert calls == [], "the janitor restarted a run the operator stopped"
+    assert any(a.get("action") == "skip" for a in actions), actions
+
+
+def test_the_janitor_still_recovers_an_unstopped_run(tmp_path, client,
+                                                     monkeypatch):
+    """Sanity check on the previous test: without Stop, it does relaunch."""
+    _he_task(tmp_path)
+    calls = []
+    monkeypatch.setattr(
+        "metainfer.tasks.harness_evolve.server.routes._launch_supervisor",
+        lambda *a, **k: calls.append("supervisor") or {"pid": 1})
+    monkeypatch.setattr(
+        "metainfer.tasks.harness_evolve.server.routes._launch_resume",
+        lambda *a, **k: calls.append("resume") or {"pid": 1})
+    janitor_sweep()
+    assert calls == ["supervisor", "resume"]
+
+
+def test_stop_clears_a_stale_resume_pid(tmp_path, client):
+    """Otherwise /resume would answer "already running" about a dead process."""
+    state_dir, _ = _he_task(tmp_path, resume_pid=999999)   # not our process
+    resp = client.post("/api/harness-evolve/he-stop/stop")
+    assert resp.status_code == 200
+    assert not (state_dir / "resume.pid").exists()
+    follow_up = client.post("/api/harness-evolve/he-stop/resume", json={})
+    assert follow_up.status_code == 200, follow_up.text
+
+
+def test_resume_clears_the_marker_so_the_run_can_continue(tmp_path, client,
+                                                          monkeypatch):
+    """"Restart later" has to actually work: the marker must not outlive it.
+
+    The marker is cleared by ``_launch_resume`` itself -- the function that
+    decides the run continues -- so its call is inspected at the spawn instead
+    of being replaced: replacing it would assert nothing about clearing.
+    """
+    import subprocess
+
+    state_dir, workspace = _he_task(tmp_path)
+    client.post("/api/harness-evolve/he-stop/stop")
+    assert (workspace / OPERATOR_STOP_FILE).exists()
+
+    spawned = {}
+
+    class _FakeProc:
+        pid = 4242
+
+    def fake_popen(cmd, *a, **k):
+        spawned["cmd"] = cmd
+        # The state the run is launched in: this is what the supervisor and the
+        # janitor will see from the moment the process exists.
+        spawned["marker_gone"] = not (workspace / OPERATOR_STOP_FILE).exists()
+        return _FakeProc()
+
+    monkeypatch.setattr(subprocess, "Popen", fake_popen)
+    resp = client.post("/api/harness-evolve/he-stop/resume", json={})
+    assert resp.status_code == 200, resp.text
+    assert "resume" in spawned.get("cmd", []), "resume refused to relaunch"
+    assert spawned["marker_gone"], \
+        "the stop marker was still there when the new run was spawned"
+    assert _operator_stopped(workspace) is None
+
+
+def test_the_cli_resume_clears_the_marker_too(tmp_path):
+    """The server is not the only way back in; the CLI path must agree."""
+    from metainfer.tasks.harness_evolve.orchestrator.cli import (
+        _clear_operator_stop,
+    )
+    workspace = tmp_path / "ws"
+    workspace.mkdir()
+    (workspace / OPERATOR_STOP_FILE).write_text("{}", encoding="utf-8")
+    assert _clear_operator_stop(workspace) == [OPERATOR_STOP_FILE]
+    assert not (workspace / OPERATOR_STOP_FILE).exists()
+    assert _clear_operator_stop(workspace) == []          # idempotent
+
+
+def test_summary_reports_the_stop_so_the_page_can_explain_it(tmp_path, client):
+    _he_task(tmp_path)
+    before = client.get("/api/harness-evolve/he-stop/summary").json()
+    assert before["operator_stop"] is None
+    client.post("/api/harness-evolve/he-stop/stop")
+    after = client.get("/api/harness-evolve/he-stop/summary").json()
+    assert after["operator_stop"]["iteration"] == 2
+
+
+def test_stop_only_targets_this_task_type(tmp_path, client):
+    """A DKAO task id must not reach the HE stop route."""
+    _he_task(tmp_path)
+    _tasks.add_task(TaskEntry(
+        id="dkao-x", type="dcu-kernel-auto-opt", label="d",
+        state_dir=str(tmp_path / "d"), workspace_dir=str(tmp_path / "dw"),
+        created_at=0.0))
+    assert client.post("/api/harness-evolve/dkao-x/stop").status_code == 409
+    assert client.post("/api/harness-evolve/nope/stop").status_code == 404
+
+
+def test_stop_reports_children_it_brought_down(tmp_path, client, monkeypatch):
+    """Stopping must reach the DKAO children, which outlive the parent."""
+    _he_task(tmp_path)
+    from metainfer.tasks.harness_evolve.server import routes as r
+    killed = []
+    monkeypatch.setattr(r, "_live_dkao_children",
+                        lambda ws: [{"pid": 111, "cmd": "child a"},
+                                    {"pid": 222, "cmd": "child b"}])
+    monkeypatch.setattr(r, "_terminate_pid",
+                        lambda pid, **k: killed.append(pid) or True)
+    body = client.post("/api/harness-evolve/he-stop/stop").json()
+    assert body["stopped_children"] == [111, 222]
+    # the second sweep must not re-kill what the first one already took down
+    assert killed == [111, 222]
+    assert body["relaunched_children"] == []
+
+
+def test_stop_reaches_an_orchestrator_that_never_wrote_resume_pid(
+        tmp_path, client, monkeypatch):
+    """A run started from the New Task form has no ``resume.pid``.
+
+    Stop used to look for the orchestrator through ``_resume_running``, which
+    reads only that file: on such a run it found nothing, answered "stopped",
+    and left the orchestrator alive -- which then re-launched the children Stop
+    had just killed as an env retry. ``_run_alive`` is the wide check
+    (``orchestrator.pid``, then ``resume.pid``, then a cmdline scan).
+    """
+    state_dir, _ = _he_task(tmp_path)          # no resume_pid on purpose
+    from metainfer.tasks.harness_evolve.server import routes as r
+
+    killed = []
+    monkeypatch.setattr(r, "_live_dkao_children", lambda ws: [])
+    monkeypatch.setattr(r, "_run_alive", lambda sd: 4242)
+    monkeypatch.setattr(r, "_terminate_pid",
+                        lambda pid, **k: killed.append(pid) or True)
+
+    body = client.post("/api/harness-evolve/he-stop/stop").json()
+    assert body["stopped_orchestrator_pid"] == 4242
+    assert killed == [4242], "the orchestrator was left running"
+    assert body["stopped"] is True
+    # ... and a task that wrote neither pid file is still "nothing to stop"
+    assert r._resume_running(state_dir) is None
+
+
+def test_stop_sweeps_the_retry_wave_the_orchestrator_launched_while_dying(
+        tmp_path, client, monkeypatch):
+    """The env retry can land between the child sweep and the parent kill.
+
+    That wave is a *new* child, not one of the pids already stopped, so a single
+    sweep leaves it optimizing for hours while the page says "stopped".
+    """
+    _he_task(tmp_path)
+    from metainfer.tasks.harness_evolve.server import routes as r
+
+    waves = [[{"pid": 111, "cmd": "first wave"}],
+             [{"pid": 111, "cmd": "first wave"},
+              {"pid": 333, "cmd": "retry wave"}]]
+    killed = []
+    monkeypatch.setattr(r, "_live_dkao_children", lambda ws: waves.pop(0))
+    monkeypatch.setattr(r, "_run_alive", lambda sd: 4242)
+    monkeypatch.setattr(r, "_terminate_pid",
+                        lambda pid, **k: killed.append(pid) or True)
+
+    body = client.post("/api/harness-evolve/he-stop/stop").json()
+    assert killed == [111, 4242, 333], killed
+    assert body["stopped_children"] == [111, 333]
+    assert body["relaunched_children"] == [333]
+
+
+def test_a_terminate_failure_is_reported_not_hidden(tmp_path, client,
+                                                    monkeypatch):
+    """The operator has to know whether the process actually went away.
+
+    A stale pid is not a live process, so this drives ``_live_dkao_children``
+    instead: it is the only path that yields a pid the server will try to kill.
+    """
+    state_dir, _ = _he_task(tmp_path)
+    from metainfer.tasks.harness_evolve.server import routes as r
+    monkeypatch.setattr(r, "_live_dkao_children",
+                        lambda ws: [{"pid": 111, "cmd": "child"}])
+    monkeypatch.setattr(r, "_terminate_pid", lambda pid, **k: False)
+    monkeypatch.setattr(r, "_run_alive", lambda sd: 999999)
+
+    body = client.post("/api/harness-evolve/he-stop/stop").json()
+    assert body["stopped"] is False
+    assert body["stopped_children"] == []          # it did not go away
+    assert "可再点一次" in body["message"]
+    # ... and the run is still marked stopped, so nothing restarts it blindly
+    assert read_run(state_dir)["final_status"] == "stopped_by_operator"
+
+
+def test_terminate_pid_actually_stops_a_real_process(tmp_path):
+    """The one path mocks cannot prove: does the signal really land?
+
+    Stop is only a pause if the processes are genuinely gone -- a surviving
+    DKAO child would keep its agents and its place in DKAO's admission queue
+    while the UI says "stopped".
+    """
+    import subprocess
+    import sys
+    import time
+
+    from metainfer.tasks.harness_evolve.server.routes import _terminate_pid
+
+    # Own session, like every process the experiment spawns; the SIGTERM goes
+    # to the group first, so this must not be the test runner's group.
+    proc = subprocess.Popen([sys.executable, "-c",
+                             "import time\nwhile True: time.sleep(1)\n"],
+                            start_new_session=True)
+    try:
+        assert proc.poll() is None
+        assert _terminate_pid(proc.pid, grace_s=5.0) is True
+        deadline = time.time() + 5.0
+        while proc.poll() is None and time.time() < deadline:
+            time.sleep(0.1)
+        assert proc.poll() is not None, "the process survived the stop"
+    finally:
+        if proc.poll() is None:
+            proc.kill()
+        proc.wait(timeout=5)
+
+
+def test_terminate_pid_reports_a_process_it_cannot_kill(monkeypatch):
+    """A refusal must come back as False, not as a silent success."""
+    from metainfer.tasks.harness_evolve.server import routes as r
+
+    monkeypatch.setattr(r.os, "kill", lambda *a, **k: None)   # never dies
+    monkeypatch.setattr(r, "_proc_state", lambda pid: "S")    # ... and runs on
+    assert r._terminate_pid(12345, grace_s=0.3) is False
+
+
+def test_a_zombie_counts_as_stopped():
+    """Otherwise every successful stop of a non-child reports failure."""
+    from metainfer.tasks.harness_evolve.server import routes as r
+
+    assert r._proc_state(os.getpid()) in {"R", "S", "D"}
+    assert r._proc_state(99999999) is None          # gone
+
+
+def test_terminate_pid_treats_a_zombie_as_stopped(monkeypatch):
+    from metainfer.tasks.harness_evolve.server import routes as r
+
+    monkeypatch.setattr(r.os, "kill", lambda *a, **k: None)
+    monkeypatch.setattr(r, "_proc_state", lambda pid: "Z")
+    assert r._terminate_pid(12345, grace_s=0.3) is True
+
+
+# ------------------------------------------------- liveness without a pid file
+
+
+def test_summary_reports_a_live_run_the_shell_cannot_see(tmp_path, client,
+                                                        monkeypatch):
+    """Stop must appear on a run whose orchestrator.pid was never written.
+
+    The shell's green dot reads that file; runs started before it existed are
+    alive with a blue dot, and inheriting that blind spot would hide Stop on
+    exactly the runs that need it.
+    """
+    from metainfer.tasks.harness_evolve.server import routes as r
+
+    _he_task(tmp_path, resume_pid=4242)
+    monkeypatch.setattr(r, "_resume_running", lambda sd: 4242)
+    monkeypatch.setattr(r, "_supervisor_running", lambda ws: None)
+    monkeypatch.setattr(r, "_live_dkao_children", lambda ws: [])
+    body = client.get("/api/harness-evolve/he-stop/summary").json()
+    assert body["running"] is True
+    # and the shell's own view stays untouched: this is a UI hint, not a
+    # rewrite of the shared liveness contract
+    shell = client.get("/api/sys-shell/he-stop").json()["status"]
+    assert shell.get("running") is False
+
+
+def test_summary_sees_liveness_through_a_running_child(tmp_path, client,
+                                                       monkeypatch):
+    """A round can be mid-flight while the parent is between steps."""
+    from metainfer.tasks.harness_evolve.server import routes as r
+
+    _he_task(tmp_path)
+    monkeypatch.setattr(r, "_resume_running", lambda sd: None)
+    monkeypatch.setattr(r, "_supervisor_running", lambda ws: None)
+    monkeypatch.setattr(r, "_live_dkao_children",
+                        lambda ws: [{"pid": 7, "cmd": "child"}])
+    assert client.get("/api/harness-evolve/he-stop/summary"
+                      ).json()["running"] is True
+
+
+def test_summary_is_not_running_after_a_stop(tmp_path, client, monkeypatch):
+    """The button must disappear once the stop has actually landed."""
+    from metainfer.tasks.harness_evolve.server import routes as r
+
+    _he_task(tmp_path)
+    monkeypatch.setattr(r, "_live_dkao_children", lambda ws: [])
+    client.post("/api/harness-evolve/he-stop/stop")
+    body = client.get("/api/harness-evolve/he-stop/summary").json()
+    assert body["running"] is False
+    assert body["operator_stop"] is not None
+
+
+# --------------------------------------------- what the operator is told to do
+
+
+def _read(rel: str) -> str:
+    from pathlib import Path
+    return (Path(__file__).resolve().parents[1] / rel).read_text(encoding="utf-8")
+
+
+def test_the_page_calls_it_a_pause_and_offers_the_way_back():
+    """The button is where most operators will ever read about this."""
+    src = _read("static/he-detail.js")
+    assert "停止任务（暂停）" in src, "the button no longer says it is a pause"
+    assert "不会被 janitor / supervisor 自动拉起" in src, (
+        "the page must tell the operator a stop sticks")
+    assert "FLOW.md §9" in src, "the page must point at the operations note"
+    # A stopped run keeps its phase, so "is the graph finished?" is the wrong
+    # question for the button next to it: what the click does depends on whether
+    # anything is running.
+    assert '"续跑到该轮次上限"' in src
+    assert '${isFinished ? "续跑到该轮次上限"' not in src, (
+        "the resume label is gated on the graph again: a stopped run would offer "
+        "「更新上限（下一轮生效）」 while the click actually re-launches it")
+
+
+def test_the_operations_note_is_written_down():
+    """FLOW.md §9 is the operator manual this feature is explained in."""
+    flow = _read("FLOW.md")
+    assert "## 9. 停止与续跑" in flow
+    for needle in ("operator_stop.json", "stopped by operator", "env_retry_attempts",
+                   "Reset", "_run_alive"):
+        assert needle in flow, f"the operations note no longer mentions {needle}"
+    # ... and the README, which is what a maintainer opens first, points at it.
+    readme = _read("README.md")
+    assert "§9" in readme and "FLOW.md" in readme
diff --git a/metainfer/tasks/harness_evolve/tests/test_pipeline_dry.py b/metainfer/tasks/harness_evolve/tests/test_pipeline_dry.py
new file mode 100644
index 00000000..494be4b0
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_pipeline_dry.py
@@ -0,0 +1,132 @@
+"""Offline tests for harness_evolve (dry-run, no GPU/LLM)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from ..orchestrator.attribution import diff_results, evaluate_changes, pass_rate
+from ..orchestrator.cli import run_with_requirements
+from ..orchestrator.config import load_experiment_config
+
+
+def _write_requirements(tmp: Path, answers: dict, task_id: str = "he-test") -> Path:
+    req = tmp / "requirements.json"
+    req.write_text(json.dumps({
+        "task_id": task_id,
+        "answers": answers,
+    }, ensure_ascii=False), encoding="utf-8")
+    return req
+
+
+def test_default_suite_config(tmp_path):
+    req = _write_requirements(tmp_path, {"max_iterations": 1})
+    cfg = load_experiment_config(req, tmp_path / "state", tmp_path / "ws")
+    assert len(cfg.suite) == 2
+    assert cfg.suite[0].id == "hy3_tp8_qkv_m4096_w4"
+
+
+def test_pipeline_dry_run_layout(tmp_path, monkeypatch):
+    monkeypatch.delenv("HARNESS_EVOLVE_FAKE_CHANGE", raising=False)
+    req = _write_requirements(tmp_path, {"max_iterations": 2})
+    rc = run_with_requirements(
+        req, state_dir=tmp_path / "state", workspace_dir=tmp_path / "ws"
+    )
+    assert rc == 0
+    ws = tmp_path / "ws"
+
+    # seeded evolvable workspace from the dcu harness_default seed
+    assert (ws / "workspace" / "manifest.yaml").is_file()
+
+    i1 = ws / "runs" / "iteration_001" / "input"
+    assert (i1 / "workspace" / "gates.yaml").is_file()
+    assert (i1 / "benchmark" / "results.json").is_file()
+    assert (i1 / "diff.json").is_file()
+    assert (i1 / "analysis" / "overview.md").is_file()
+    assert (ws / "best_ever.json").is_file()
+    scores = (ws / "iteration_scores.jsonl").read_text().strip().splitlines()
+    assert len(scores) == 2
+    assert (ws / "report.md").is_file()
+    # iteration 2 attributes iteration 1's (empty) manifest and emits an
+    # explicit NO_SIGNAL decision. Equal performance must NOT be a reject
+    # rollback; the candidate is archived while champion stays active.
+    i2_root = ws / "runs" / "iteration_002"
+    assert (i2_root / "input" / "change_evaluation.json").is_file()
+    decision = json.loads((i2_root / "input" / "decision.json").read_text())
+    assert decision["verdict"] == "NO_SIGNAL"
+    assert not (i2_root / "rollback.json").exists()
+    disposition = json.loads(
+        (i2_root / "candidate_disposition.json").read_text()
+    )
+    assert disposition["action"] == "KEEP_CHAMPION_ARCHIVE_CANDIDATE"
+
+
+def test_attribution_verdicts():
+    results_pass = {"t1": {"passed": True}, "t2": {"passed": False}}
+    results_fail = {"t1": {"passed": False}, "t2": {"passed": False}}
+    diff = diff_results(results_pass, results_fail)
+    assert "t1" in diff["regressed"]
+    stats = pass_rate(results_pass)
+    assert stats["n_pass"] == 1 and stats["n_total"] == 2
+
+    manifest = {
+        "iteration": 2,
+        "changes": [{
+            "id": "chg-1",
+            "description": "fix t1",
+            "files": ["gates.yaml"],
+            "predicted_fixes": ["t1"],
+            "risk_tasks": ["t2"],
+        }],
+    }
+    ev = evaluate_changes(manifest, diff)
+    entry = ev["change_evaluations"][0]
+    assert entry["verdict"] == "INEFFECTIVE"  # predicted t1 regressed, t2 not realized
+
+    # HARMFUL: a declared risk regressed while nothing predicted was fixed
+    diff2 = diff_results(
+        {"t1": {"passed": True}, "t2": {"passed": True}},
+        {"t1": {"passed": True}, "t2": {"passed": False}},
+    )
+    ev2 = evaluate_changes(manifest, diff2)
+    assert ev2["change_evaluations"][0]["verdict"] == "HARMFUL"
+    assert ev2["unattributed_regressions"] == []
+
+
+def test_code_change_defers_to_the_next_round(tmp_path, monkeypatch):
+    """A harness fix must not wait for the run to end — nor apply mid-round.
+
+    The running process already imported the old modules, so the loop stops at
+    the round boundary with ``code_reload_pending:``; the restart records
+    its own revision and resumes there.
+    """
+    from ..orchestrator import pipeline as pl
+    from ..orchestrator.state import loaded_harness_revision
+
+    req = _write_requirements(tmp_path, {"max_iterations": 5})
+    cfg = load_experiment_config(req, tmp_path / "state", tmp_path / "ws")
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=1)
+    loaded = loaded_harness_revision(cfg.state_dir)
+    assert loaded
+
+    # someone edits the loop while the experiment is between rounds
+    monkeypatch.setattr(pl, "harness_revision", lambda: "rev-changed")
+    pl.run_experiment(cfg, start_iteration=2, iterations_to_run=1)
+
+    run = json.loads((cfg.state_dir / "run.json").read_text(encoding="utf-8"))
+    assert run["final_status"] == "code_reload_pending:3"
+    assert pl.completed_iterations(cfg.exp_dir) == 2      # both rounds are real
+    events = [json.loads(l) for l in
+              (cfg.state_dir / "timeline.jsonl").read_text(
+                  encoding="utf-8").splitlines() if l.strip()]
+    deferred = [e for e in events if e["type"] == "code_reload_deferred"]
+    assert deferred and deferred[-1]["payload"]["resumes_at_iteration"] == 3
+    assert deferred[-1]["payload"]["on_disk_revision"] == "rev-changed"
+
+    # the restarted process records its own revision and keeps going
+    monkeypatch.undo()
+    pl.run_experiment(cfg, start_iteration=3, iterations_to_run=1)
+    assert pl.completed_iterations(cfg.exp_dir) == 3
+    run = json.loads((cfg.state_dir / "run.json").read_text(encoding="utf-8"))
+    assert not str(run.get("final_status") or "").startswith("code_reload")
diff --git a/metainfer/tasks/harness_evolve/tests/test_pool_protocol.py b/metainfer/tasks/harness_evolve/tests/test_pool_protocol.py
new file mode 100644
index 00000000..87cd7bef
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_pool_protocol.py
@@ -0,0 +1,333 @@
+"""Tests for suite protocol v2: pool, auto-pass, round_plan guardrails, and
+pool-mode pipeline decisions."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from ..orchestrator.pipeline import run_experiment
+from ..orchestrator.pool import builtin_pool
+from ..orchestrator.round_plan import (
+    RoundPlan,
+    enforce_guardrails,
+    validate_in_pool,
+)
+from ..orchestrator.attribution import pass_rate
+
+
+def _pool():
+    return builtin_pool()
+
+
+def test_builtin_pool_sampling_and_heldout():
+    pool = _pool()
+    ids = pool.stratified_sample(4)
+    assert len(ids) == 4 and len(set(ids)) == 4
+    pool.assert_in_pool(ids)
+    held = pool.heldout_ids(2, exclude=set(ids))
+    assert len(held) == 2 and not (set(held) & set(ids))
+
+
+def test_tau_and_auto_pass():
+    pool = _pool()
+    inst_id = "dsv4_tp8_qkv_proj_m16"  # family decode__qkv_proj
+    inst = pool.get(inst_id)
+    tau = pool.tau_family(inst.family)
+    assert tau > 0.0
+    ok = pool.auto_pass(inst_id, tau * inst.baseline_us * 0.9,
+                        correctness_ok=True, p90_ok=True)
+    assert ok["passed"] is True
+    bad = pool.auto_pass(inst_id, inst.baseline_us * 0.9,
+                         correctness_ok=True, p90_ok=True)
+    assert bad["passed"] is False  # ratio 0.9 > tau
+    noc = pool.auto_pass(inst_id, 1.0, correctness_ok=False, p90_ok=True)
+    assert noc["passed"] is False
+
+
+def test_round_plan_in_pool_validation():
+    pool = _pool()
+    plan = RoundPlan(iteration=1, selected=["dsv4_tp8_qkv_proj_m16",
+                                            "hy3_tp8_qkv_proj_m16"])
+    validate_in_pool(plan, pool)
+    bad = RoundPlan(iteration=1, selected=["not-in-pool"])
+    try:
+        validate_in_pool(bad, pool)
+        raise AssertionError("expected ValueError")
+    except ValueError:
+        pass
+
+
+def test_guardrails_preserve_overlap_and_budget():
+    pool = _pool()
+    prev = pool.stratified_sample(5)
+    # agent proposes a fresh set of 3 different ids
+    fresh = [i for i in pool.stratified_sample(9) if i not in set(prev)][:3]
+    plan = RoundPlan(iteration=2, strategy="explore",
+                     rationale="try fresh", selected=fresh)
+    out = enforce_guardrails(prev, plan, pool, budget=5)
+    overlap = len(set(out) & set(prev))
+    assert overlap >= int(0.6 * len(prev))  # >= 3 of 5
+    assert len(out) == 5 and len(set(out)) == 5
+    pool.assert_in_pool(out)
+
+
+def test_pipeline_pool_mode_dry(tmp_path):
+    """pool-mode: results restricted to scored ids, round_plan written per
+    evolve, overlap decisions recorded."""
+    req = tmp_path / "req.json"
+    req.write_text(json.dumps({
+        "task_id": "pool-dry",
+        "answers": {
+            "suite_yaml": "instances: []",
+            "max_iterations": "2",
+            "execution_mode": "dry-run",
+            "pool_source": "builtin",
+            "per_round_budget": "4",
+        },
+    }), encoding="utf-8")
+    from ..orchestrator.cli import run_with_requirements
+    rc = run_with_requirements(req, state_dir=tmp_path / "st",
+                               workspace_dir=tmp_path / "ws")
+    assert rc == 0
+    ws = tmp_path / "ws"
+    pool = _pool()
+
+    i1 = ws / "runs" / "iteration_001" / "input"
+    i2 = ws / "runs" / "iteration_002" / "input"
+    r1 = json.loads((i1 / "benchmark" / "results.json").read_text())["results"]
+    scored1 = json.loads((i1 / "benchmark" / "scored_ids.json").read_text())
+    assert set(r1) == set(scored1)
+    assert len(scored1) <= 4
+    pool.assert_in_pool(scored1)
+    # evolve writes round_plan each round
+    assert (ws / "runs" / "iteration_001" / "evolve"
+            / "round_plan.json").is_file()
+    assert (ws / "runs" / "iteration_002" / "evolve"
+            / "round_plan.json").is_file()
+    # iteration 2 overlap decision artefacts exist
+    assert (i2 / "diff.json").is_file()
+    stats = pass_rate(r1)
+    assert stats["n_total"] == len(scored1)
+
+
+def _mixed_pool():
+    from ..orchestrator.pool import Pool, PoolInstance
+    insts = {}
+    for iid in ("anchor_a", "anchor_b"):
+        insts[iid] = PoolInstance(
+            id=iid, model="m", tp_size=8, operator="o", M=16, N=1, K=1,
+            baseline_us=100.0, family="decode__o",
+            history=[{"accepted": True, "median_us": 45.0}],
+        )
+    for iid in ("cold_a", "cold_b"):
+        insts[iid] = PoolInstance(
+            id=iid, model="m", tp_size=8, operator="o", M=16, N=1, K=1,
+            baseline_us=100.0, family="decode__o", history=[],
+        )
+    return Pool(insts)
+
+
+def test_anchor_rule_keeps_two_historical_shapes():
+    pool = _mixed_pool()
+    plan = RoundPlan(iteration=1, strategy="explore",
+                     rationale="all cold this round",
+                     selected=["cold_a", "cold_b"])
+    out = enforce_guardrails([], plan, pool, budget=4)
+    assert len(out) == 4
+    anchors = [i for i in out if i in ("anchor_a", "anchor_b")]
+    assert len(anchors) >= 2
+
+
+def test_budget_below_anchor_min_raises():
+    pool = _mixed_pool()
+    try:
+        enforce_guardrails([], None, pool, budget=1)
+        raise AssertionError("expected ValueError")
+    except ValueError:
+        pass
+
+
+def test_historical_ids_helper():
+    pool = _mixed_pool()
+    assert set(pool.historical_ids()) == {"anchor_a", "anchor_b"}
+
+
+def test_pool_dry_run_assigns_question_roles(tmp_path):
+    """Every round carries explicit roles, and the decision records them."""
+    req = tmp_path / "req.json"
+    req.write_text(json.dumps({
+        "task_id": "pool-roles",
+        "answers": {
+            "suite_yaml": "instances: []",
+            "max_iterations": "2",
+            "execution_mode": "dry-run",
+            "pool_source": "builtin",
+            "per_round_budget": "4",
+        },
+    }), encoding="utf-8")
+    from ..orchestrator.cli import run_with_requirements
+    assert run_with_requirements(req, state_dir=tmp_path / "st",
+                                 workspace_dir=tmp_path / "ws") == 0
+    ws = tmp_path / "ws"
+
+    for num in (1, 2):
+        base = ws / "runs" / f"iteration_{num:03d}" / "input"
+        plan = json.loads((base / "benchmark" / "purposes.json")
+                          .read_text(encoding="utf-8"))
+        purposes = plan["purposes"]
+        scored = json.loads((base / "benchmark" / "scored_ids.json")
+                            .read_text(encoding="utf-8"))
+        assert set(purposes) == set(scored)
+        roles = list(purposes.values())
+        assert roles.count("regression") >= 2
+        assert "repair" in roles
+        assert "probe" in roles
+        assert len(set(plan["families"])) >= 2
+
+        decision = json.loads((base / "decision.json").read_text(encoding="utf-8"))
+        if "generalization_gate" in decision:
+            assert decision["probe_ids"]
+            assert "p" not in decision["paired_ids"] or True
+            assert decision["purpose_breakdown"]
+
+    # rotation sidecar records what was examined
+    rotation = json.loads((ws / "pool_rotation.json").read_text(encoding="utf-8"))
+    assert rotation and max(rotation.values()) >= 1
+
+
+def test_wired_scope_violations_detects_changes_to_unwired_components(tmp_path):
+    """Editing a component the manifest marks unwired must be reported."""
+    from ..orchestrator.pipeline import _wired_scope_violations
+
+    before = tmp_path / "before"
+    after = tmp_path / "after"
+    for root in (before, after):
+        root.mkdir()
+    manifest = {
+        "schema_version": 1,
+        "components": {
+            "gates": {"path": "gates.yaml", "wired": True},
+            "prompts": {"path": "prompts.yaml", "wired": False},
+        },
+    }
+    for root in (before, after):
+        (root / "manifest.yaml").write_text(json.dumps(manifest), encoding="utf-8")
+        (root / "gates.yaml").write_text("gates: {}\n", encoding="utf-8")
+        (root / "prompts.yaml").write_text("a: 1\n", encoding="utf-8")
+
+    assert _wired_scope_violations(before, after) == []
+
+    # edit the wired component -> no violation reported
+    (after / "gates.yaml").write_text("gates: {x: 1}\n", encoding="utf-8")
+    assert _wired_scope_violations(before, after) == []
+
+    # edit the unwired component -> violation reported
+    (after / "prompts.yaml").write_text("a: 2\n", encoding="utf-8")
+    violations = _wired_scope_violations(before, after)
+    assert violations == [{"component": "prompts", "path": "prompts.yaml"}]
+
+
+def test_generations_are_recorded_and_champion_is_the_best_measured(tmp_path):
+    """Each round is registered as a generation and best_ever points at it."""
+    req = tmp_path / "req.json"
+    req.write_text(json.dumps({
+        "task_id": "pool-generations",
+        "answers": {
+            "suite_yaml": "instances: []",
+            "max_iterations": "3",
+            "execution_mode": "dry-run",
+            "pool_source": "builtin",
+            "per_round_budget": "4",
+        },
+    }), encoding="utf-8")
+    from ..orchestrator.cli import run_with_requirements
+    from ..orchestrator.champion_selection import (
+        best_generation, load_generations,
+    )
+    assert run_with_requirements(req, state_dir=tmp_path / "st",
+                                 workspace_dir=tmp_path / "ws") == 0
+    ws = tmp_path / "ws"
+    gens = load_generations(ws)
+    assert [g.iteration for g in gens] == [1, 2, 3]
+    best = best_generation(gens)
+    assert best is not None
+    best_ever = json.loads((ws / "best_ever.json").read_text(encoding="utf-8"))
+    assert best_ever["iteration"] == best.iteration
+    assert best_ever["verdict"] in {"BEST_MEASURED", "MANUAL_BASELINE",
+                                    "BASELINE", "PROMOTE",
+                                    "PROMOTE_UNEXPLAINED"}
+
+
+def test_dry_run_writes_best_known_review_and_skips_final_evolve(tmp_path):
+    req = tmp_path / "req.json"
+    req.write_text(json.dumps({
+        "task_id": "pool-review",
+        "answers": {
+            "suite_yaml": "instances: []",
+            "max_iterations": "3",
+            "execution_mode": "dry-run",
+            "pool_source": "builtin",
+            "per_round_budget": "4",
+        },
+    }), encoding="utf-8")
+    from ..orchestrator.cli import run_with_requirements
+    assert run_with_requirements(req, state_dir=tmp_path / "st",
+                                 workspace_dir=tmp_path / "ws") == 0
+    ws = tmp_path / "ws"
+
+    assert (ws / "best_known.json").is_file()
+    review = json.loads((ws / "round_review.json").read_text(encoding="utf-8"))
+    assert review["variant_policy"] == "freeze"
+    assert review["champion_iteration"] in {1, 2, 3}
+    assert (ws / "round_review.md").read_text(encoding="utf-8").startswith(
+        "# Cross-round review")
+
+    manifest = json.loads((ws / "runs" / "iteration_003" / "evolve"
+                           / "change_manifest.json").read_text(encoding="utf-8"))
+    assert manifest["verification"]["status"] == "target_reached_skip_evolve"
+    assert manifest["changes"] == []
+
+
+def test_resume_skips_a_round_that_was_marked_skipped(tmp_path):
+    """A skipped round must not block the loop: the next round still runs.
+
+    Mirrors the live handover (iteration 2 ran while the GPU gate was broken
+    and was marked skipped): its eval artifacts are absent, so the resumed
+    process has to pick the next round up from the previous round's evolve dir.
+    """
+    req = tmp_path / "req.json"
+    req.write_text(json.dumps({
+        "task_id": "pool-skip",
+        "answers": {
+            "suite_yaml": "instances: []",
+            "max_iterations": "4",
+            "execution_mode": "dry-run",
+            "pool_source": "builtin",
+            "per_round_budget": "3",
+        },
+    }), encoding="utf-8")
+    from ..orchestrator.config import load_experiment_config
+    from ..orchestrator.cli import run_experiment
+    from ..orchestrator.pipeline import completed_iterations
+
+    cfg = load_experiment_config(req, tmp_path / "st", tmp_path / "ws")
+    run_experiment(cfg, start_iteration=1, iterations_to_run=1)
+    assert completed_iterations(cfg.exp_dir) == 1
+
+    # mark iteration 2 skipped: results recorded as empty, no evolve artifacts
+    it2 = cfg.exp_dir / "runs" / "iteration_002" / "input"
+    (it2 / "benchmark").mkdir(parents=True, exist_ok=True)
+    (it2 / "benchmark" / "results.json").write_text(
+        json.dumps({"iteration": 2, "results": {}, "skipped": True}),
+        encoding="utf-8")
+    assert completed_iterations(cfg.exp_dir) == 2
+
+    # the resumed process starts at round 3 and runs it end to end
+    run_experiment(cfg, start_iteration=3, iterations_to_run=1)
+    it3 = cfg.exp_dir / "runs" / "iteration_003" / "input"
+    results = json.loads((it3 / "benchmark" / "results.json").read_text())
+    assert results["results"], "round 3 must still select and score questions"
+    assert (it3 / "decision.json").is_file()
+    assert (it3 / "diff.json").is_file()
+    assert completed_iterations(cfg.exp_dir) == 3
diff --git a/metainfer/tasks/harness_evolve/tests/test_promotion_flow.py b/metainfer/tasks/harness_evolve/tests/test_promotion_flow.py
new file mode 100644
index 00000000..01baa79d
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_promotion_flow.py
@@ -0,0 +1,698 @@
+"""End-to-end: a round judged against variants, then a human-approved promotion.
+
+Covers the operator's protocol: round 1 defines the generalization paper, the
+performance gate uses fresh operators, both gates compare against the current
+variant, a passing candidate waits for a human, and approval publishes the
+harness and moves the variants.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+import yaml
+
+from ..orchestrator import pipeline as pl
+from ..orchestrator import rounds as R
+from ..orchestrator import variant as V
+from ..orchestrator import promotion as promo
+from ..orchestrator.config import load_experiment_config
+
+
+class FakeEvaluator:
+    """Deterministic measurements: ``medians`` maps operator id -> us.
+
+    ``medians`` may also be a callable ``(iteration, stage) -> {id: median}``, so
+    a test can make the generalization retake measure differently from the
+    baseline round that fixed the paper — that is the only way to say anything
+    about what an approval publishes.
+    """
+
+    name = "fake"
+
+    def __init__(self, medians):
+        self.medians = medians
+        self.calls = []
+        self.stages = []
+
+    def evaluate(self, cfg, workspace, iteration, instances=None):
+        medians = self.medians
+        stage = str((getattr(cfg, "answers", None) or {}).get("round_stage") or "")
+        if callable(medians):
+            medians = medians(iteration, stage)
+        self.stages.append((iteration, stage))
+        out = {}
+        for inst in (instances or []):
+            iid = inst.id
+            self.calls.append(iid)
+            median = medians.get(iid)
+            if median is None:
+                out[iid] = {"status": "missing_report", "passed": False,
+                            "correctness_ok": False, "median_us": None}
+                continue
+            out[iid] = {"status": "success", "passed": True,
+                        "correctness_ok": True, "p90_ok": True,
+                        "median_us": float(median), "p90_us": float(median),
+                        "repo_path": f"/repos/{iid}"}
+        return out
+
+
+def _write_pool(tmp_path: Path, ops: dict) -> Path:
+    path = tmp_path / "pool.yaml"
+    path.write_text(yaml.safe_dump({
+        "schema_version": 1,
+        "instances": [
+            {"id": iid, "family": f"fam{i}",
+             "contract": {"model": "m", "tp_size": 8, "operator": "o",
+                          "M": 16, "N": 64, "K": 32},
+             "baseline_us": values["baseline"], "best_known_us": values["variant"],
+             "history": []}
+            for i, (iid, values) in enumerate(sorted(ops.items()))
+        ],
+    }, sort_keys=False), encoding="utf-8")
+    return path
+
+
+def _cfg(tmp_path: Path, pool: Path, **answers):
+    req = tmp_path / "requirements.json"
+    req.write_text(json.dumps({
+        "task_id": "variant-e2e",
+        "answers": {
+            "suite_yaml": "instances: []",
+            "execution_mode": "dry-run",
+            "evolve_mode": "dry-run",
+            "pool_source": str(pool),
+            "question_pool": str(pool),
+            "round_questions": "variant",
+            "per_round_budget": "3",
+            "max_iterations": "4",
+            "max_rounds": "3",
+            **answers,
+        },
+    }), encoding="utf-8")
+    return load_experiment_config(req, tmp_path / "state", tmp_path / "ws")
+
+
+def _patch(monkeypatch, tmp_path, evaluator, prod_dir):
+    # A real DKAO child leaves its optimized workspace behind; the kernel
+    # promotion path reads exactly that directory, so the fake evaluator
+    # materializes one per question it measured.
+    original = evaluator.evaluate
+    measured: dict = {}
+
+    def evaluate(cfg, workspace, iteration, instances=None):
+        results = original(cfg, workspace, iteration, instances=instances)
+        from ..orchestrator.rounds import CHILDREN_DIR, stage_dir_name
+        stage = pl.round_stage(
+            iteration=iteration,
+            pending=pl.read_pending_generalization(cfg.exp_dir))
+        override = cfg.answers.pop("_attempt_dir", None)
+        for iid, row in results.items():
+            measured[str(iid)] = row
+            child = (cfg.exp_dir / CHILDREN_DIR
+                     / (override or stage_dir_name(stage, iteration))
+                     / str(iid) / "workspace")
+            child.mkdir(parents=True, exist_ok=True)
+            (child / "manifest.yaml").write_text(
+                f"harness_name: candidate-{iid}\n", encoding="utf-8")
+        return results
+
+    evaluator.evaluate = evaluate
+    # The real child writes an accepted-kernel manifest inside that workspace;
+    # the fake does not, so the DKAO-side promotion is stubbed at its boundary.
+    # Everything the parent does around it (best-known wins on the pool file,
+    # backups, history rows, audit log) stays real.
+    from metainfer.tasks.dcu_kernel_auto_opt.orchestrator import variant_promote
+
+    def promote_variant(*, workspace_dir, answers, shape_id, source_task="",
+                        correctness_ok=None, min_improvement_percent=0.0,
+                        tp=None, m=None, model_label=None, dry_run=False,
+                        **kwargs):
+        # A real child only has an accepted kernel in the workspace that ran.
+        # Reading the manifest keeps this stub honest: pointing it at a
+        # directory that never produced a kernel fails, like DKAO would.
+        try:
+            manifest = (Path(workspace_dir) / "manifest.yaml").read_text(
+                encoding="utf-8")
+        except OSError:
+            manifest = ""
+        if "failed-attempt" in manifest:
+            return {"ok": False, "action": "no_kernel", "shape": shape_id,
+                    "reason": f"no accepted kernel under {workspace_dir}"}
+        row = measured.get(str(shape_id)) or {}
+        action = "would-add" if dry_run else "added"
+        return {"ok": True, "action": action, "shape": shape_id,
+                "new_median_us": row.get("median_us"),
+                "old_median_us": None, "path": f"/variants/{shape_id}"}
+
+    monkeypatch.setattr(variant_promote, "promote_variant", promote_variant)
+    monkeypatch.setattr(pl, "build_evaluator", lambda mode: evaluator)
+    monkeypatch.setattr(promo, "_production_dir", lambda: prod_dir)
+    monkeypatch.setattr(pl, "_head_revision", lambda ws: "rev-candidate")
+
+
+def _steady(ops, share=0.5):
+    """A harness that keeps finding the *same* faster kernel every time.
+
+    The gate compares a measurement with the variant; a tie inside the
+    comparison's noise band is ``NO_SIGNAL``, not a win, so a redrawn operator
+    set needs to be genuinely better than what the last round recorded. Tests
+    that want an unqualified win therefore improve on each round's own result.
+    """
+    def medians(iteration, stage):
+        # Only the baseline round's operators are in the variant pool before
+        # the draw, so a draw's variant is still the original one.
+        return {iid: values["variant"] * share
+                for iid, values in ops.items()}
+    return medians
+
+
+def _ops(n=10, variant=40.0):
+    return {f"op{i:02d}": {"baseline": 100.0 + i, "variant": variant + 4 * i}
+            for i in range(n)}
+
+
+def test_failed_performance_gate_rolls_back_and_keeps_iterating(tmp_path, monkeypatch):
+    """Round 1 judges nothing; the performance gate is what rejects."""
+    pool_path = _write_pool(tmp_path, _ops())
+    cfg = _cfg(tmp_path, pool_path)
+    prod = tmp_path / "production"
+    prod.mkdir()
+    (prod / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+
+    # every operator is far slower than its variant
+    evaluator = FakeEvaluator({iid: 400.0 for iid in _ops()})
+    _patch(monkeypatch, tmp_path, evaluator, prod)
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=1)
+
+    it1 = cfg.exp_dir / "runs" / "iteration_001" / "input"
+    # round 1 is the baseline round: no gate is written at all
+    assert not (it1 / "performance_gate.json").exists()
+    assert not (it1 / "generalization_gate.json").exists()
+    first = json.loads((it1 / "decision.json").read_text())
+    assert first["stage"] == "baseline"
+    assert first["verdict"] == "BASELINE_MEASURED"
+    assert (it1 / "baseline_round.json").is_file()
+
+    # round 2 asks the same (still far too slow) evaluator for the performance
+    # gate, which is the first thing that can actually fail
+    it2 = cfg.exp_dir / "runs" / "iteration_002" / "input"
+    perf = json.loads((it2 / "performance_gate.json").read_text())
+    assert perf["status"] == "FAIL" and perf["wins"] == []
+    assert not (it2 / "generalization_gate.json").exists()
+    assert not (cfg.exp_dir / "pending_promotion.json").exists()
+    assert not (cfg.exp_dir / "pending_generalization.json").exists()
+    # the variant table did not move: a baseline round that beat nothing
+    table = json.loads((cfg.exp_dir / "variant_table.json").read_text())
+    assert table["harness_version"] in {None, "seed"}
+    assert (it2 / "decision.json").is_file()
+
+
+def test_a_retry_leaves_the_failed_attempt_beside_it_and_wins_the_resolution():
+    """A retried question keeps both attempts, and the newest one is "the" one.
+
+    The failed attempt's directory is evidence, so it must survive; the retry
+    gets its own directory because DKAO binds a child workspace's ``main``
+    symlink to the kernel repo it was launched with (reusing the directory makes
+    the retry die instantly). Everything that reads "this question's child"
+    -- the drill-down, the kernel promotion path -- must land on the retry.
+    """
+    import tempfile
+
+    tmp = Path(tempfile.mkdtemp())
+    ops = _ops()
+    pool_path = _write_pool(tmp, ops)
+    exp = tmp / "exp"
+    exp.mkdir()
+    operators = sorted(ops)[:2]
+
+    def plant(attempt_dir: str, harness_name: str) -> None:
+        for iid in operators:
+            workspace = exp / "children" / attempt_dir / iid / "workspace"
+            workspace.mkdir(parents=True, exist_ok=True)
+            (workspace / "manifest.yaml").write_text(
+                f"harness_name: {harness_name}\n", encoding="utf-8")
+
+    plant("iteration_001", "failed-attempt")
+    plant("iteration_001_a2", "retry-attempt")
+
+    # both attempts are on disk, and both are attributable to the iteration
+    names = {d.name for d in R.attempt_dirs(exp / "children", 1)}
+    assert names == {"iteration_001", "iteration_001_a2"}
+    assert R.attempt_dirs(exp / "children", 2) == []
+
+    # the newest attempt is the question's child, for every reader
+    for iid in operators:
+        resolved = R.resolve_child_dir(exp / "children", 1, iid)
+        assert resolved == exp / "children" / "iteration_001_a2" / iid
+    assert R.child_dir_map(exp / "children", 1)[operators[0]] == (
+        exp / "children" / "iteration_001_a2" / operators[0])
+
+    # a question that only ever ran once still resolves to its single attempt
+    only_once = exp / "children" / "iteration_001_a2" / "solo" / "workspace"
+    only_once.mkdir(parents=True)
+    assert R.resolve_child_dir(exp / "children", 1, "solo") == (
+        exp / "children" / "iteration_001_a2" / "solo")
+    assert R.resolve_child_dir(exp / "children", 1, "never-ran") is None
+
+
+def test_round1_judges_no_gate_and_wins_do_not_promote_a_harness(tmp_path, monkeypatch):
+    """A baseline round that beats every variant still promotes nothing."""
+    ops = _ops()
+    pool_path = _write_pool(tmp_path, ops)
+    # exactly one round: the point here is round 1, and a second round is what
+    # would schedule the generalization retake
+    cfg = _cfg(tmp_path, pool_path, rounds="1")
+    prod = tmp_path / "production"
+    prod.mkdir()
+    (prod / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+    evaluator = FakeEvaluator({iid: values["variant"] * 0.5
+                               for iid, values in ops.items()})
+    _patch(monkeypatch, tmp_path, evaluator, prod)
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=1)
+
+    it1 = cfg.exp_dir / "runs" / "iteration_001" / "input"
+    baseline = json.loads((it1 / "baseline_round.json").read_text())
+    assert baseline["verdict"]["status"] == "PASS"
+    assert baseline["pool_promotion"]["promoted"], (
+        "the baseline round's faster kernels must enter the variant pool")
+    # no gate verdict, no pending promotion: only the kernels moved
+    assert not (it1 / "performance_gate.json").exists()
+    assert not (cfg.exp_dir / "pending_promotion.json").exists()
+    assert yaml.safe_load((prod / "manifest.yaml").read_text())["harness_name"] == "seed"
+    # the pool now holds the baseline round's numbers
+    published = yaml.safe_load(pool_path.read_text(encoding="utf-8"))
+    entries = {i["id"]: i for i in published["instances"]}
+    for iid in baseline["pool_promotion"]["promoted"]:
+        assert entries[iid]["best_known_us"] == pytest.approx(
+            ops[iid]["variant"] * 0.5)
+
+
+def test_the_retake_re_hands_the_paper_to_dkao_on_the_same_harness(
+        tmp_path, monkeypatch):
+    """The generalization gate is its own stage: same harness, A1 again."""
+    ops = _ops()
+    pool_path = _write_pool(tmp_path, ops)
+    cfg = _cfg(tmp_path, pool_path)
+    prod = tmp_path / "production"
+    prod.mkdir()
+    (prod / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+    evaluator = FakeEvaluator({iid: values["variant"] * 0.5
+                               for iid, values in ops.items()})
+    _patch(monkeypatch, tmp_path, evaluator, prod)
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=4)
+
+    paper = list(json.loads(
+        (cfg.exp_dir / "variant_table.json").read_text()
+    )["groups"]["generalization_ids"])
+    assert len(paper) == 4
+
+    it1 = cfg.exp_dir / "runs" / "iteration_001" / "input"
+    it2 = cfg.exp_dir / "runs" / "iteration_002" / "input"
+
+    # stage 1: the paper, judged by nothing
+    assert json.loads((it1 / "decision.json").read_text())["stage"] == "baseline"
+    assert sorted(json.loads((it1 / "benchmark" / "scored_ids.json").read_text())) \
+        == sorted(paper)
+
+    # stage 2: a fresh draw that never overlaps the paper, performance gate only
+    a2 = json.loads((it2 / "benchmark" / "scored_ids.json").read_text())
+    assert set(a2).isdisjoint(paper)
+    perf = json.loads((it2 / "performance_gate.json").read_text())
+    assert perf["status"] == "PASS"
+    assert not (it2 / "generalization_gate.json").exists(), (
+        "the paper must not be re-measured in the same pass as the draw")
+
+    # stage 3: the retake — the SAME harness bytes, the paper's operators, sent
+    # to DKAO again, and judged by the generalization gate
+    it3 = cfg.exp_dir / "runs" / "iteration_003" / "input"
+    retake = json.loads((it3 / "stage.json").read_text())
+    assert retake["stage"] == "generalization"
+    assert retake["retake_of_iteration"] == 2, (
+        "the retake belongs to the round that won the performance gate")
+    gen = json.loads((it3 / "generalization_gate.json").read_text())
+    assert gen["status"] == "PASS"
+    assert sorted(gen["operators"]) == sorted(paper)
+    assert json.loads((it3 / "decision.json").read_text())["stage"] == \
+        "generalization"
+    # a fresh DKAO launch, in its own children directory, not the first pass's
+    assert (cfg.exp_dir / "children" / "iteration_003_generalization").is_dir()
+    assert (cfg.exp_dir / "children" / "iteration_002").is_dir(), (
+        "the two stages never share a child directory")
+    assert (cfg.exp_dir / "pending_promotion.json").is_file()
+    run = json.loads((cfg.state_dir / "run.json").read_text())
+    assert run["final_status"] == "awaiting_approval"
+    assert run["rounds_used"] == pytest.approx(1.5), (
+        "the baseline round is free; a performance pass is 1.0 and the "
+        "generalization retake half a round")
+
+
+def test_passing_rounds_wait_for_approval_then_promote(tmp_path, monkeypatch):
+    ops = _ops()
+    pool_path = _write_pool(tmp_path, ops)
+    cfg = _cfg(tmp_path, pool_path)
+    prod = tmp_path / "production"
+    prod.mkdir()
+    (prod / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+
+    # half the variant on every operator: a wide win on both gates
+    evaluator = FakeEvaluator({iid: values["variant"] * 0.5
+                               for iid, values in ops.items()})
+    _patch(monkeypatch, tmp_path, evaluator, prod)
+
+    report = pl.run_experiment(cfg, start_iteration=1, iterations_to_run=4)
+    assert report.is_file()
+
+    it2 = cfg.exp_dir / "runs" / "iteration_002" / "input"
+    perf = json.loads((it2 / "performance_gate.json").read_text())
+    it3 = cfg.exp_dir / "runs" / "iteration_003" / "input"
+    gen = json.loads((it3 / "generalization_gate.json").read_text())
+    assert perf["status"] == "PASS" and gen["status"] == "PASS"
+    assert set(perf["operators"]).isdisjoint(gen["operators"]), (
+        "the performance draw and the paper are different operator sets")
+    pending = json.loads((cfg.exp_dir / "pending_promotion.json").read_text())
+    assert pending["version"] == "h-2"
+    run = json.loads((cfg.state_dir / "run.json").read_text())
+    assert run["final_status"] == "awaiting_approval"
+    # nothing was written to production yet
+    assert yaml.safe_load((prod / "manifest.yaml").read_text())["harness_name"] == "seed"
+
+    # a human approves: the harness lands in production and variants move
+    result = promo.approve_promotion(cfg, cfg.exp_dir, approved_by="tester")
+    assert result["ok"] is True
+    assert result["version"] == "h-2"
+    assert not (cfg.exp_dir / "pending_promotion.json").exists()
+    assert (prod / "manifest.yaml").is_file()
+    assert (prod / "gates.yaml").is_file()
+    promoted = json.loads((cfg.exp_dir / "promoted_harness.json").read_text())
+    assert promoted["approved_by"] == "tester"
+    assert Path(promoted["previous_backup"]).is_dir()
+
+    table = json.loads((cfg.exp_dir / "variant_table.json").read_text())
+    assert table["harness_version"] == "h-2"
+    # the reference never drifts away from production: after approval, every
+    # operator's recorded variant is the pool's own best-known kernel
+    published = {i["id"]: i["best_known_us"]
+                 for i in yaml.safe_load(pool_path.read_text())["instances"]}
+    for iid, row in table["operators"].items():
+        assert row["median_us"] == pytest.approx(published[iid]), (
+            f"{iid}: the experiment's reference must match the pool's best-known")
+    assert set(table["groups"]["generalization_ids"]) == set(gen["operators"])
+
+
+def test_denied_candidate_is_discarded_without_touching_production(
+        tmp_path, monkeypatch):
+    ops = _ops()
+    pool_path = _write_pool(tmp_path, ops)
+    cfg = _cfg(tmp_path, pool_path, max_rounds="9")
+    prod = tmp_path / "production"
+    prod.mkdir()
+    (prod / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+    evaluator = FakeEvaluator({iid: values["variant"] * 0.5
+                               for iid, values in ops.items()})
+    _patch(monkeypatch, tmp_path, evaluator, prod)
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=4)
+    assert (cfg.exp_dir / "pending_promotion.json").is_file()
+
+    denied = promo.deny_promotion(cfg, cfg.exp_dir, reason="not convinced")
+    assert denied["ok"] is True
+    assert not (cfg.exp_dir / "pending_promotion.json").exists()
+    assert not (cfg.exp_dir / "promoted_harness.json").exists()
+    assert yaml.safe_load((prod / "manifest.yaml").read_text())["harness_name"] == "seed"
+    table = json.loads((cfg.exp_dir / "variant_table.json").read_text())
+    assert table["harness_version"] in {None, "seed"}
+
+
+def test_rollback_restores_the_previous_production_tree(tmp_path, monkeypatch):
+    pool_path = _write_pool(tmp_path, _ops())
+    cfg = _cfg(tmp_path, pool_path)
+    prod = tmp_path / "production"
+    prod.mkdir()
+    (prod / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+    _patch(monkeypatch, tmp_path, FakeEvaluator({}), prod)
+
+    # a promoted version replaces production
+    snapshot = tmp_path / "candidate"
+    snapshot.mkdir()
+    (snapshot / "manifest.yaml").write_text("harness_name: evolved\n", encoding="utf-8")
+    promo.promote_harness(cfg.exp_dir, {
+        "iteration": 3, "version": "h-3", "revision": "rev-candidate",
+        "snapshot_dir": str(snapshot),
+        "performance_gate": {"status": "PASS", "wins": ["a"]},
+        "generalization_gate": {"status": "PASS", "wins": ["a"], "operators": ["a"]},
+    }, approved_by="tester")
+    assert yaml.safe_load((prod / "manifest.yaml").read_text()
+                          )["harness_name"] == "evolved"
+
+    back = promo.rollback_harness(cfg.exp_dir, to="last-good")
+    assert back["ok"] is True
+    assert yaml.safe_load((prod / "manifest.yaml").read_text()
+                          )["harness_name"] == "seed"
+    record = json.loads((cfg.exp_dir / "promoted_harness.json").read_text())
+    assert record["version"].startswith("rollback:")
+
+
+def test_generalization_paper_survives_a_failed_first_round(tmp_path, monkeypatch):
+    """Round 1 fixes the paper even when nothing it measures is an improvement."""
+    pool_path = _write_pool(tmp_path, _ops())
+    cfg = _cfg(tmp_path, pool_path)
+    prod = tmp_path / "production"
+    prod.mkdir()
+    evaluator = FakeEvaluator({iid: 900.0 for iid in _ops()})     # all much slower
+    _patch(monkeypatch, tmp_path, evaluator, prod)
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=1)
+
+    table = json.loads((cfg.exp_dir / "variant_table.json").read_text())
+    paper = table["groups"]["generalization_ids"]
+    assert len(paper) == 4
+    # round_questions.json is rewritten by every pass, so the baseline round's
+    # own copy is what says the paper was defined there
+    it1 = cfg.exp_dir / "runs" / "iteration_001" / "input"
+    questions = json.loads((it1 / "stage.json").read_text())
+    assert questions["stage"] == "baseline"
+    assert table["groups"]["generalization_defined_at"] == 1
+    assert len(table["groups"]["generalization_ids"]) == 4
+    # nothing beat its variant, so the pool keeps what it had
+    it1 = cfg.exp_dir / "runs" / "iteration_001" / "input"
+    baseline = json.loads((it1 / "baseline_round.json").read_text())
+    assert (baseline["pool_promotion"] or {}).get("promoted", []) == []
+
+
+def test_the_baseline_round_stops_the_run_when_the_paper_cannot_be_measured(
+        tmp_path, monkeypatch):
+    """Round 1 judges no gate, but it owes a usable number for every operator."""
+    pool_path = _write_pool(tmp_path, _ops())
+    cfg = _cfg(tmp_path, pool_path)
+    _patch(monkeypatch, tmp_path, FakeEvaluator({}), tmp_path / "production")
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=4)
+
+    run = json.loads((cfg.state_dir / "run.json").read_text())
+    assert run["final_status"] == "baseline_round_failed"
+    # stopped at round 1: no later round ran, nothing was parked for approval
+    assert not (cfg.exp_dir / "runs" / "iteration_002" / "input"
+                / "performance_gate.json").exists()
+    assert not (cfg.exp_dir / "pending_promotion.json").exists()
+    events = [json.loads(l)["type"] for l in
+              (cfg.state_dir / "timeline.jsonl").read_text().splitlines()
+              if l.strip()]
+    assert "baseline_round_failed" in events
+
+
+def test_generation_records_carry_their_own_revision(tmp_path, monkeypatch):
+    """A generation's record must name the revision it measured."""
+    from ..orchestrator.champion_selection import load_generations
+
+    pool_path = _write_pool(tmp_path, _ops())
+    cfg = _cfg(tmp_path, pool_path, round_questions="legacy")  # exercise the legacy path
+    evaluator = FakeEvaluator({})
+    _patch(monkeypatch, tmp_path, evaluator, tmp_path / "production")
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=1)
+    generations = load_generations(cfg.exp_dir)
+    assert generations, "a generation must be recorded"
+    assert generations[0].workspace_revision, (
+        "workspace_revision must name this round's commit, not be empty")
+
+
+def test_approval_promotes_the_accepted_rounds_kernels_into_the_pool(
+        tmp_path, monkeypatch):
+    """Only the accepted round's operators reach production — as best-knowns.
+
+    Approving must (a) lower the pool's best_known where the accepted retake beat
+    it, (b) append an auditable history row naming the harness version, and (c)
+    leave the pool recoverable via a backup. The baseline round writes the pool
+    on its own (kernel fact), so the retake has to be *better* than that to show
+    up here.
+    """
+    ops = _ops()
+    pool_path = _write_pool(tmp_path, ops)
+    cfg = _cfg(tmp_path, pool_path)
+    prod = tmp_path / "production"
+    prod.mkdir()
+    (prod / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+
+    def medians(iteration, stage):
+        share = 0.6 if stage == "generalization" else 0.9
+        return {iid: values["variant"] * share for iid, values in ops.items()}
+
+    evaluator = FakeEvaluator(medians)
+    _patch(monkeypatch, tmp_path, evaluator, prod)
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=4)
+    assert (cfg.exp_dir / "pending_promotion.json").is_file()
+    result = promo.approve_promotion(cfg, cfg.exp_dir, approved_by="tester")
+    assert result["ok"] is True
+
+    kernel_update = result["kernel_update"]
+    assert kernel_update["updated"], "the accepted round's kernels must be published"
+    assert kernel_update["pool"] == str(pool_path)
+    assert Path(kernel_update["backup"]).is_file()
+
+    accepted = set(json.loads(
+        (cfg.exp_dir / "variant_table.json").read_text()
+    )["groups"]["generalization_ids"])
+    published = yaml.safe_load(pool_path.read_text(encoding="utf-8"))
+    entries = {i["id"]: i for i in published["instances"]}
+    for change in kernel_update["updated"]:
+        iid = change["operator_id"]
+        assert iid in accepted, "only the accepted round's operators are published"
+        assert entries[iid]["best_known_us"] == pytest.approx(
+            ops[iid]["variant"] * 0.6)
+        history = entries[iid]["history"]
+        assert history[-1]["accepted"] is True
+        assert history[-1]["harness_version"] == "h-2"
+    # the pool is still the original operator-shaped document
+    assert len(published["instances"]) == len(ops)
+
+
+def test_a_slower_round_never_degrades_the_pool(tmp_path, monkeypatch):
+    """Best-known wins: nothing a round measures can make a pool entry worse."""
+    ops = _ops(n=8)
+    pool_path = _write_pool(tmp_path, ops)
+    cfg = _cfg(tmp_path, pool_path)
+    prod = tmp_path / "production"
+    prod.mkdir()
+    (prod / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+
+    def medians(iteration, stage):
+        # A baseline round that does well, then rounds that measure worse
+        share = 0.5 if stage == "baseline" else 0.95
+        return {iid: values["variant"] * share for iid, values in ops.items()}
+
+    _patch(monkeypatch, tmp_path, FakeEvaluator(medians), prod)
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=4)
+
+    published = {i["id"]: i["best_known_us"]
+                 for i in yaml.safe_load(pool_path.read_text())["instances"]}
+    for iid, row in ops.items():
+        assert published[iid] <= row["variant"], (
+            "a pool entry may only ever improve")
+
+
+def test_rejected_round_never_promotes_kernels_in_variant_mode(tmp_path, monkeypatch):
+    """Only the accepted round's optimizations reach production.
+
+    The legacy path promoted each round's kernels as soon as it was evaluated;
+    under the variant protocol a rejected round must leave nothing behind.
+    """
+    ops = _ops()
+    pool_path = _write_pool(tmp_path, ops)
+    cfg = _cfg(tmp_path, pool_path)
+    prod = tmp_path / "production"
+    prod.mkdir()
+    _patch(monkeypatch, tmp_path, FakeEvaluator({}), prod)
+
+    promoted = {"calls": 0}
+
+    def _fake_promote(*a, **k):
+        promoted["calls"] += 1
+        return None
+
+    monkeypatch.setattr(pl, "_promote_iteration", _fake_promote)
+    # the evaluator reports nothing, so every operator fails both gates
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=1)   # round 1 fails
+    assert promoted["calls"] == 0
+    assert not (cfg.exp_dir / "pending_promotion.json").exists()
+    assert not (cfg.exp_dir / "promoted_harness.json").exists()
+
+
+def test_a_run_without_a_usable_pool_stops_loudly(tmp_path, monkeypatch):
+    """No operators to measure must not burn the whole round budget."""
+    pool_path = tmp_path / "empty_pool.yaml"
+    pool_path.write_text(yaml.safe_dump({"schema_version": 1, "instances": []}),
+                         encoding="utf-8")
+    cfg = _cfg(tmp_path, pool_path)
+    _patch(monkeypatch, tmp_path, FakeEvaluator({}), tmp_path / "production")
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=1)
+    run = json.loads((cfg.state_dir / "run.json").read_text())
+    assert run["final_status"] == "no_question_pool"
+    events = [json.loads(l)["type"] for l in
+              (cfg.state_dir / "timeline.jsonl").read_text().splitlines() if l.strip()]
+    assert "variant_pool_empty" in events
+
+
+def test_mechanism_gate_is_recorded_and_does_not_veto_a_win(tmp_path, monkeypatch):
+    """Lenient, but never absent: the mechanism verdict is graded and recorded.
+
+    A contradicted mechanism ("the planner was off, so this planner policy
+    change cannot have caused the win") must show up in the decision, while the
+    counter gates keep deciding. Round 1 is the baseline round and has no gate
+    to grade; the performance round is where a declared change is first judged.
+    """
+    ops = _ops()
+    pool_path = _write_pool(tmp_path, ops)
+    cfg = _cfg(tmp_path, pool_path)
+    prod = tmp_path / "production"
+    prod.mkdir()
+    (prod / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+    evaluator = FakeEvaluator({iid: values["variant"] * 0.5
+                               for iid, values in ops.items()})
+    _patch(monkeypatch, tmp_path, evaluator, prod)
+
+    # a change manifest whose mechanism cannot have fired: the planner was off
+    manifest_dir = cfg.exp_dir / "runs" / "iteration_001" / "evolve"
+    manifest_dir.mkdir(parents=True, exist_ok=True)
+    (manifest_dir / "change_manifest.json").write_text(json.dumps({
+        "schema_version": 1, "iteration": 1,
+        "changes": [{
+            "id": "chg-1",
+            "description": "point the exhausted-budget phase at the bounded plan",
+            "files": ["planner_policy.yaml"],
+            "mechanism_signature": ["planner emits architecture_explore"],
+            "mechanism_check": {"kind": "planner_plan_ids",
+                                "expect_any": ["architecture_explore"]},
+        }],
+    }), encoding="utf-8")
+
+    pl.run_experiment(cfg, start_iteration=1, iterations_to_run=4)
+
+    it1 = cfg.exp_dir / "runs" / "iteration_001" / "input"
+    first = json.loads((it1 / "decision.json").read_text())
+    assert first["stage"] == "baseline"
+    assert first["mechanism_gate"] == {}, "round 1 judges no gate"
+    assert first["performance_gate"] is None
+
+    # round 2 grades that manifest against round 2's children while it judges
+    # the performance gate
+    it2 = cfg.exp_dir / "runs" / "iteration_002" / "input"
+    second = json.loads((it2 / "decision.json").read_text())
+    assert second["stage"] == "performance"
+    mechanism = second["mechanism_gate"]
+    assert mechanism["status"] in {"MISS", "PARTIAL", "UNOBSERVED", "HIT"}
+    assert "detail" in mechanism
+    assert second["performance_gate"]["status"] == "PASS"
+    # a passing performance gate promotes nothing by itself: it hands the same
+    # harness to DKAO again with the frozen paper
+    frozen = json.loads((cfg.exp_dir / "pending_generalization.json").read_text())
+    assert set(frozen["operator_ids"]) == set(
+        json.loads((cfg.exp_dir / "variant_table.json").read_text()
+                   )["groups"]["generalization_ids"])
diff --git a/metainfer/tasks/harness_evolve/tests/test_promotion_pipeline.py b/metainfer/tasks/harness_evolve/tests/test_promotion_pipeline.py
new file mode 100644
index 00000000..820888cd
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_promotion_pipeline.py
@@ -0,0 +1,129 @@
+"""HE pipeline promotion step: better kernels become DKAO variants."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+from metainfer.tasks.dcu_kernel_auto_opt.orchestrator import variant_store as vs
+from metainfer.tasks.harness_evolve.orchestrator.config import (
+    ExperimentConfig, InstanceSpec,
+)
+from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+    _promote_iteration,
+)
+
+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 _cfg(tmp_path: Path, **answers) -> ExperimentConfig:
+    base = {"promote_kernels": "true",
+            "promote_min_improvement_percent": "3"}
+    base.update(answers)
+    return ExperimentConfig(
+        task_id="he-promo", state_dir=tmp_path / "state",
+        workspace_dir=tmp_path / "exp", suite=[], answers=base,
+        execution_mode="dkao-cli", pool_source="pool.yaml",
+        per_round_budget=1, max_iterations=1,
+    )
+
+
+def _inst() -> InstanceSpec:
+    return InstanceSpec(id=SHAPE, shape={
+        "model": "hy3-dsh-tp4-m16-2-368c654c", "tp_size": 4,
+        "operator": "o_proj", "M": 16, "N": 4096, "K": 2048,
+    }, budget_rounds=1, pass_median_us_le=20.0)
+
+
+def _child_workspace(exp: Path, median_us: float) -> Path:
+    leaf = (exp / "children" / "iteration_001" / SHAPE / "workspace"
+            / "workers" / "worker_0" / "accepted" / SHAPE)
+    leaf.mkdir(parents=True, exist_ok=True)
+    (leaf / "kernel.hip").write_text("// promoted kernel\n", encoding="utf-8")
+    (leaf / "manifest.json").write_text(json.dumps({
+        "commit": "abc123",
+        "shape": {"M": 16, "N": 4096, "K": 2048},
+        "metrics": {"median_us": median_us, "p90_us": median_us * 1.01},
+    }), encoding="utf-8")
+    return leaf
+
+
+def _seed_variant(root: Path, median_us: float) -> Path:
+    meta = vs.derive_variant_meta(
+        {"operator": "Quantized GEMM", "dtype": "INT8 W8A8",
+         "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}) + "// old\n// @@end\n", encoding="utf-8")
+    return target
+
+
+def test_promotes_faster_child_and_logs(tmp_path, variant_root):
+    cfg = _cfg(tmp_path)
+    exp = cfg.exp_dir
+    _child_workspace(exp, 8.49)
+    _seed_variant(variant_root, 9.17)          # 7.4% faster -> promote
+    results = {SHAPE: {"correctness_ok": True, "passed": True,
+                       "median_us": 8.49}}
+
+    payload = _promote_iteration(cfg, exp, 1, results, [_inst()])
+
+    assert payload and payload["promoted"] == [SHAPE]
+    out = payload["results"][0]
+    assert out["action"] == "updated"
+    assert out["improvement_percent"] == pytest.approx(7.42, abs=0.05)
+    assert (exp / "runs" / "iteration_001" / "promotion.json").is_file()
+    log = (exp / "promotion_log.jsonl").read_text(encoding="utf-8").strip()
+    assert json.loads(log)["shape"] == SHAPE
+    # the variant leaf now carries the new kernel
+    assert "promoted kernel" in Path(out["path"]).read_text(encoding="utf-8")
+
+
+def test_skips_below_threshold_and_ignores_missing_workspace(tmp_path, variant_root):
+    cfg = _cfg(tmp_path)
+    exp = cfg.exp_dir
+    _child_workspace(exp, 9.10)                # only 0.8% faster than 9.17
+    _seed_variant(variant_root, 9.17)
+    other = InstanceSpec(id="hy3_tp4_qkv_proj_m16", shape={"M": 16}, budget_rounds=1)
+    results = {SHAPE: {"correctness_ok": True, "median_us": 9.10},
+               "hy3_tp4_qkv_proj_m16": {"correctness_ok": True, "median_us": 1.0}}
+
+    payload = _promote_iteration(cfg, exp, 1, results, [_inst(), other])
+    by_shape = {r["shape"]: r for r in payload["results"]}
+    assert payload["promoted"] == []
+    assert by_shape[SHAPE]["action"] == "skipped"
+    assert "required 3.00%" in by_shape[SHAPE]["reason"]
+    # a question with no child directory at all reports "no kernel", the same
+    # as a child directory that produced none
+    assert by_shape["hy3_tp4_qkv_proj_m16"]["action"] == "no_kernel"
+    assert "no accepted kernel" in by_shape["hy3_tp4_qkv_proj_m16"]["reason"]
+
+
+def test_disabled_by_default_switch_returns_none(tmp_path, variant_root):
+    cfg = _cfg(tmp_path, promote_kernels="false")
+    exp = cfg.exp_dir
+    _child_workspace(exp, 1.0)
+    assert _promote_iteration(cfg, exp, 1, {}, [_inst()]) is None
+    assert not (exp / "runs" / "iteration_001" / "promotion.json").exists()
+
+
+def test_failed_correctness_is_never_promoted(tmp_path, variant_root):
+    cfg = _cfg(tmp_path)
+    exp = cfg.exp_dir
+    _child_workspace(exp, 1.0)
+    payload = _promote_iteration(
+        cfg, exp, 1,
+        {SHAPE: {"correctness_ok": False, "median_us": 1.0}}, [_inst()])
+    assert payload["promoted"] == []
+    assert payload["results"][0]["action"] == "skipped"
diff --git a/metainfer/tasks/harness_evolve/tests/test_promotion_routes.py b/metainfer/tasks/harness_evolve/tests/test_promotion_routes.py
new file mode 100644
index 00000000..29b4b9a9
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_promotion_routes.py
@@ -0,0 +1,99 @@
+"""WebUI approval routes: the human gate in front of production DKAO."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from metainfer.tasks.harness_evolve.orchestrator import promotion as promo
+from metainfer.tasks.harness_evolve.server.routes import (
+    _approve_with_config, _deny_with_config,
+)
+
+
+def _task_entry(state_dir: Path, workspace_dir: Path):
+    class Entry:
+        pass
+
+    entry = Entry()
+    entry.id = "promo-task"
+    entry.state_dir = str(state_dir)
+    entry.workspace_dir = str(workspace_dir)
+    return entry
+
+
+def _setup(tmp_path, monkeypatch):
+    """An experiment parked on a pending promotion, plus a fake production dir."""
+    state_dir = tmp_path / "state"
+    workspace = tmp_path / "ws"
+    state_dir.mkdir(parents=True)
+    workspace.mkdir(parents=True)
+    (state_dir / "requirements.json").write_text(json.dumps({
+        "task_id": "promo-task",
+        "answers": {"suite_yaml": "instances: []", "execution_mode": "dry-run"},
+    }), encoding="utf-8")
+    (state_dir / "run.json").write_text(json.dumps({"task_id": "promo-task"}),
+                                        encoding="utf-8")
+
+    snapshot = workspace / "runs" / "iteration_002" / "input" / "workspace"
+    snapshot.mkdir(parents=True)
+    (snapshot / "manifest.yaml").write_text("harness_name: evolved\n", encoding="utf-8")
+    (snapshot / "gates.yaml").write_text("component: gates\n", encoding="utf-8")
+
+    production = tmp_path / "production"
+    production.mkdir()
+    (production / "manifest.yaml").write_text("harness_name: seed\n", encoding="utf-8")
+    monkeypatch.setattr(promo, "_production_dir", lambda: production)
+
+    promo.write_pending_promotion(
+        workspace, iteration=2, revision="rev-2", snapshot_dir=str(snapshot),
+        performance_gate={"status": "PASS", "wins": ["a", "b", "c"],
+                          "operators": ["a", "b", "c", "d"]},
+        generalization_gate={"status": "PASS", "wins": ["a", "b"],
+                             "operators": ["a", "b", "c", "d"]},
+        results={"a": {"median_us": 1.0, "correctness_ok": True}})
+    return _task_entry(state_dir, workspace), state_dir, workspace, production
+
+
+def test_approve_route_publishes_and_clears_the_pending_candidate(tmp_path, monkeypatch):
+    entry, state_dir, workspace, production = _setup(tmp_path, monkeypatch)
+
+    result = _approve_with_config(entry, workspace, "ui-tester")
+    assert result["ok"] is True
+    assert result["version"] == "h-2"
+    # pending cleared, production replaced, record written
+    assert not (workspace / "pending_promotion.json").exists()
+    assert "evolved" in (production / "manifest.yaml").read_text(encoding="utf-8")
+    record = json.loads((workspace / "promoted_harness.json").read_text(encoding="utf-8"))
+    assert record["approved_by"] == "ui-tester"
+    assert Path(record["previous_backup"]).is_dir()
+    # the run is finished from the shell's point of view
+    run = json.loads((state_dir / "run.json").read_text(encoding="utf-8"))
+    assert run["final_status"] == "promoted:h-2"
+
+
+def test_deny_route_leaves_production_alone_and_clears_pending(tmp_path, monkeypatch):
+    entry, state_dir, workspace, production = _setup(tmp_path, monkeypatch)
+
+    result = _deny_with_config(entry, workspace, "not convinced")
+    assert result["ok"] is True
+    assert not (workspace / "pending_promotion.json").exists()
+    assert not (workspace / "promoted_harness.json").exists()
+    assert "seed" in (production / "manifest.yaml").read_text(encoding="utf-8")
+    log = (workspace / "promotion_log.jsonl").read_text(encoding="utf-8")
+    assert "denied" in log and "not convinced" in log
+
+
+def test_approve_route_refuses_when_nothing_is_pending(tmp_path, monkeypatch):
+    state_dir = tmp_path / "state"
+    workspace = tmp_path / "ws"
+    state_dir.mkdir(parents=True)
+    workspace.mkdir(parents=True)
+    (state_dir / "requirements.json").write_text(json.dumps({
+        "task_id": "promo-task",
+        "answers": {"suite_yaml": "instances: []", "execution_mode": "dry-run"},
+    }), encoding="utf-8")
+    monkeypatch.setattr(promo, "_production_dir", lambda: tmp_path / "nope")
+
+    result = _approve_with_config(_task_entry(state_dir, workspace), workspace, "ui")
+    assert result["ok"] is False and "no pending promotion" in result["errors"][0]
diff --git a/metainfer/tasks/harness_evolve/tests/test_publish.py b/metainfer/tasks/harness_evolve/tests/test_publish.py
new file mode 100644
index 00000000..ea26f41f
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_publish.py
@@ -0,0 +1,92 @@
+"""Tests for the explicit publish gate."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import yaml
+
+from ..orchestrator.attribution import copy_tree, save_json
+from ..orchestrator.publish import (
+    check_publishable,
+    publish_harness,
+    publish_kernels,
+)
+
+_DCU_SEED = Path("/root/zth_agent/MetaInfer/metainfer/tasks/"
+                 "dcu_kernel_auto_opt/harness_default")
+
+
+def _fake_experiment(tmp_path, *, verdict="PROMOTE", held="PASS"):
+    exp = tmp_path / "exp"
+    snapshot = exp / "runs" / "iteration_002" / "input" / "workspace"
+    copy_tree(_DCU_SEED, snapshot)
+    save_json(exp / "best_ever.json", {
+        "iteration": 2, "pass_rate": 0.9, "snapshot_dir": str(snapshot),
+        "workspace_revision": "rev-2",
+    })
+    save_json(exp / "runs" / "iteration_002" / "input" / "decision.json", {
+        "iteration": 2, "verdict": verdict, "action": "SET_CHAMPION",
+        "heldout_gate": {"status": held},
+    })
+    return exp
+
+
+def test_publish_gate_rejects_non_promotable(tmp_path):
+    exp = _fake_experiment(tmp_path, verdict="REJECT")
+    gate = check_publishable(exp)
+    assert gate["ok"] is False
+    assert any("not promotable" in e for e in gate["errors"])
+
+
+def test_publish_gate_rejects_overfit(tmp_path):
+    exp = _fake_experiment(tmp_path, verdict="PROMOTE", held="FAIL")
+    gate = check_publishable(exp)
+    assert gate["ok"] is False
+    assert any("held-out" in e for e in gate["errors"])
+
+
+def test_publish_harness_backs_up_and_replaces_target(tmp_path):
+    exp = _fake_experiment(tmp_path)
+    gate = check_publishable(exp)
+    assert gate["ok"] is True and gate["verdict"] == "PROMOTE"
+
+    target = tmp_path / "harness_default"
+    copy_tree(_DCU_SEED, target)
+    marker = target / "old_extra.yaml"
+    marker.write_text("old: true", encoding="utf-8")
+
+    result = publish_harness(exp, target=target)
+    assert result["ok"] and result["published"]
+    assert not marker.exists()  # replace semantics removed stale file
+    assert (target / "planner_policy.yaml").is_file()
+    backups = list(tmp_path.glob("harness_default.pre-publish-*"))
+    assert len(backups) == 1
+    log = (exp / "publish_log.jsonl").read_text().strip().splitlines()
+    assert len(log) == 1 and "harness" in log[0]
+
+
+def test_publish_kernels_updates_pool_with_backup(tmp_path):
+    exp = tmp_path / "exp"
+    child = exp / "children" / "iter001" / "c" / "workspace"
+    child.mkdir(parents=True)
+    report = {
+        "status": "success",
+        "config": {"shapes": [{"id": "s1"}]},
+        "final_validation": {"s1": {"median_us": 30.0, "p90_us": 31.0}},
+    }
+    (child / "final_report.json").write_text(json.dumps(report),
+                                             encoding="utf-8")
+    pool_path = tmp_path / "pool.yaml"
+    pool_path.write_text(yaml.safe_dump({"instances": [{
+        "id": "s1", "contract": {"M": 16}, "family": "f",
+        "baseline_us": 100.0, "best_known_us": 50.0, "history": [],
+    }]}), encoding="utf-8")
+
+    result = publish_kernels(exp, pool_path=pool_path)
+    assert result["ok"] and result["published"]
+    assert result["updated"] == ["s1"]
+    data = yaml.safe_load(pool_path.read_text(encoding="utf-8"))
+    assert data["instances"][0]["best_known_us"] == 30.0
+    assert len(list(tmp_path.glob("pool.yaml.pre-publish-*"))) == 1
diff --git a/metainfer/tasks/harness_evolve/tests/test_question_roles.py b/metainfer/tasks/harness_evolve/tests/test_question_roles.py
new file mode 100644
index 00000000..284d337f
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_question_roles.py
@@ -0,0 +1,187 @@
+"""Role-based question selection: roles, family quota, overlap, rotation."""
+
+from __future__ import annotations
+
+from metainfer.tasks.harness_evolve.orchestrator.question_roles import (
+    PURPOSE_PROBE, PURPOSE_REGRESSION, PURPOSE_REPAIR, main_purposes,
+    overlap_ratio, probe_ids, select_questions,
+)
+
+
+class _Inst:
+    def __init__(self, iid: str, family: str) -> None:
+        self.id = iid
+        self.family = family
+
+
+class _Pool:
+    def __init__(self, rows) -> None:
+        self.instances = {i: _Inst(i, f) for i, f in rows}
+
+    def all_ids(self):
+        return list(self.instances)
+
+    def get(self, iid):
+        return self.instances[iid]
+
+    def historical_ids(self):
+        return list(self.instances)
+
+
+def _pool():
+    return _Pool([
+        ("a1", "famA"), ("a2", "famA"), ("a3", "famA"),
+        ("b1", "famB"), ("b2", "famB"),
+        ("c1", "famC"), ("c2", "famC"),
+        ("d1", "famD"), ("d2", "famD"),
+    ])
+
+
+def _big_pool():
+    """Six families: enough for a genuinely cross-family probe at budget 4."""
+    return _Pool([
+        ("a1", "famA"), ("a2", "famA"), ("a3", "famA"),
+        ("b1", "famB"), ("b2", "famB"), ("b3", "famB"),
+        ("c1", "famC"), ("c2", "famC"), ("c3", "famC"),
+        ("d1", "famD"), ("d2", "famD"), ("d3", "famD"),
+        ("e1", "famE"), ("e2", "famE"), ("e3", "famE"),
+        ("f1", "famF"), ("f2", "famF"), ("f3", "famF"),
+    ])
+
+
+PREV = ["a1", "b1", "c1", "d1"]
+PREV_RESULTS = {
+    "a1": {"passed": True}, "b1": {"passed": True},
+    "c1": {"passed": False}, "d1": {"passed": True},
+}
+
+
+def test_roles_regression_repair_and_probe():
+    pool = _big_pool()
+    sel = select_questions(pool=pool, budget=4, prev_selected=PREV,
+                           prev_results=PREV_RESULTS)
+    assert len(sel.selected) == 4
+    roles = [sel.purposes[i] for i in sel.selected]
+    assert roles.count(PURPOSE_REGRESSION) >= 2      # anchors from passed set
+    assert PURPOSE_REPAIR in roles                   # the failed question
+    assert PURPOSE_PROBE in roles                    # cross-family probe
+    assert "c1" in sel.selected
+    # the probe never shares a family with the rest of the round
+    probe = probe_ids(sel.purposes)[0]
+    others = {pool.get(i).family for i in sel.selected if i != probe}
+    assert pool.get(probe).family not in others
+
+
+def test_family_quota_and_minimum_coverage():
+    pool = _big_pool()
+    sel = select_questions(pool=pool, budget=4, prev_selected=PREV,
+                           prev_results=PREV_RESULTS, min_families=3,
+                           max_per_family=2)
+    families = [pool.get(i).family for i in sel.selected]
+    assert len(set(families)) >= 3
+    assert max(families.count(f) for f in set(families)) <= 2
+
+
+def test_overlap_guardrail_holds():
+    sel = select_questions(pool=_big_pool(), budget=4, prev_selected=PREV,
+                           prev_results=PREV_RESULTS)
+    assert overlap_ratio(sel.selected, PREV) >= 0.6
+
+
+def test_small_pool_degrades_to_the_freshest_question():
+    """With four families and four slots no family is left for a probe."""
+    sel = select_questions(pool=_pool(), budget=4, prev_selected=PREV,
+                           prev_results=PREV_RESULTS)
+    assert len(sel.selected) == 4
+    assert PURPOSE_PROBE in sel.purposes.values()
+    assert any("probe" in n for n in sel.notes)
+
+
+def test_agent_picks_are_honoured_when_they_fit_a_role():
+    class Plan:
+        selected = ["b1", "a1", "c1", "d2"]
+        changes: list = []
+
+    sel = select_questions(pool=_big_pool(), budget=4, prev_selected=PREV,
+                           prev_results=PREV_RESULTS, plan=Plan())
+    # passed agent picks become regression anchors, the failed one repair
+    assert sel.purposes.get("b1") == PURPOSE_REGRESSION
+    assert sel.purposes.get("c1") == PURPOSE_REPAIR
+
+
+def test_agent_picks_are_replaced_when_family_quota_breaks():
+    class Plan:
+        selected = ["a1", "a2", "a3", "b1"]      # three from famA
+        changes: list = []
+
+    sel = select_questions(pool=_big_pool(), budget=4, prev_selected=PREV,
+                           prev_results=PREV_RESULTS, plan=Plan(),
+                           max_per_family=2, min_families=3)
+    pool = _big_pool()
+    families = [pool.get(i).family for i in sel.selected]
+    assert families.count("famA") <= 2
+    assert len(set(families)) >= 3
+
+
+def test_rotation_prefers_least_recently_used_for_the_probe():
+    sel = select_questions(pool=_pool(), budget=4, prev_selected=PREV,
+                           prev_results=PREV_RESULTS,
+                           rotation={"a1": 9, "a2": 9, "a3": 9, "b1": 9,
+                                     "b2": 9, "c1": 9, "c2": 9, "d1": 9,
+                                     "d2": 0})
+    assert "d2" in sel.selected
+
+
+def test_first_round_without_history_still_covers_families():
+    pool = _big_pool()
+    sel = select_questions(pool=pool, budget=4)
+    assert len(sel.selected) == 4
+    families = {pool.get(i).family for i in sel.selected}
+    assert len(families) >= 3
+    assert len(main_purposes(sel.purposes)) >= 3
+
+
+def test_scope_ids_bias_the_regression_anchors():
+    class Plan:
+        selected: list = []
+        changes = [{"scope": {"expected_improve": ["b2"], "at_risk": ["c2"]}}]
+
+    sel = select_questions(pool=_big_pool(), budget=4, prev_selected=PREV,
+                           prev_results=PREV_RESULTS, plan=Plan())
+    assert "b2" in sel.selected or "c2" in sel.selected
+
+
+class _MInst:
+    def __init__(self, iid: str, family: str, m: int) -> None:
+        self.id = iid
+        self.family = family
+        self.M = m
+
+
+class _MPool:
+    def __init__(self, rows) -> None:
+        self.instances = {i: _MInst(i, f, m) for i, f, m in rows}
+
+    def all_ids(self):
+        return list(self.instances)
+
+    def get(self, iid):
+        return self.instances[iid]
+
+    def historical_ids(self):
+        return list(self.instances)
+
+
+def test_fresh_round_covers_both_regimes():
+    """decode (M<=32) and prefill need different kernels: both must appear."""
+    pool = _MPool([
+        ("p1", "prefill__o_proj", 4096), ("p2", "prefill__qkv_proj", 4096),
+        ("p3", "prefill__shared_down_proj", 4096),
+        ("p4", "prefill__fused_qkv_a_proj", 4096),
+        ("d1", "decode__o_proj", 16), ("d2", "decode__qkv_proj", 16),
+        ("d3", "decode__shared_down_proj", 16),
+    ])
+    sel = select_questions(pool=pool, budget=4)
+    regimes = {pool.get(i).M <= 32 and "decode" or "prefill" for i in sel.selected}
+    assert regimes == {"decode", "prefill"}, sel.selected
+    assert len(sel.families) >= 3
diff --git a/metainfer/tasks/harness_evolve/tests/test_real_adapters.py b/metainfer/tasks/harness_evolve/tests/test_real_adapters.py
new file mode 100644
index 00000000..50699036
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_real_adapters.py
@@ -0,0 +1,466 @@
+"""Unit tests for real evaluator/evolver surfaces (no GPU/LLM calls)."""
+
+from __future__ import annotations
+
+import json
+import threading
+import time
+from pathlib import Path
+from typing import List
+
+from ..orchestrator.adapters.eval import (
+    DkaoCliEvaluator, _child_env, _normalize_report,
+    _normalize_workspace_result, _worker_rounds,
+)
+from ..orchestrator.config import ExperimentConfig, InstanceSpec
+from ..orchestrator.evolve import AgentEvolver, _last_result
+
+
+def _inst():
+    return InstanceSpec(
+        id="hy3_tp8_qkv_proj_m16",
+        shape={"model": "hy3", "tp_size": 8, "operator": "qkv_proj",
+               "M": 16, "N": 1280, "K": 4096},
+        budget_rounds=4,
+        pass_median_us_le=60.0,
+    )
+
+
+def _cfg(tmp_path):
+    # HE manages no GPU occupancy: the device is the fixed index % 4 rotation and
+    # there is no gate wait, lease or broker poll to configure.
+    return ExperimentConfig(
+        task_id="ahe-parent", state_dir=tmp_path / "state",
+        workspace_dir=tmp_path / "workspace", suite=[_inst()],
+        execution_mode="dkao-cli", agent_framework="dsh",
+        answers={"ahe_repo_root": str(tmp_path / "ahe-repos")},
+    )
+
+
+def test_dkao_requirements_include_parent_and_shape(tmp_path):
+    ev = DkaoCliEvaluator()
+    req = ev._build_dkao_requirements(
+        _cfg(tmp_path), _inst(), 2, "child-id", "repo-name", 3,
+        tmp_path / "harness",
+    )
+    assert req["parent_ahe_task_id"] == "ahe-parent"
+    assert req["ahe_iteration"] == 2
+    assert req["target_repo_path"] == "repo-name"
+    shape_cfg = __import__("yaml").safe_load(req["shape_config"])
+    # DKAO manual mode requires worker_N <-> GPU N; a single child on GPU 3
+    # must therefore be assigned to worker_3 (not worker_0 with gpu=3).
+    assert shape_cfg["assignments"] == {
+        "worker_3": {"gpu": 3, "shapes": ["hy3_tp8_qkv_proj_m16"]}
+    }
+    assert shape_cfg["shapes"][0]["M"] == 16
+
+
+def test_dkao_child_uses_flash_41_model_label(tmp_path):
+    """The child DKAO task runs on the 4.1 Flash model, and the label it
+    submits is one DKAO's own config resolves (the two task packages must not
+    drift apart on model naming)."""
+    from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.config import (
+        DSH_MODEL_IDS,
+    )
+
+    ev = DkaoCliEvaluator()
+    req = ev._build_dkao_requirements(
+        _cfg(tmp_path), _inst(), 1, "child-id", "repo-name", 0,
+        tmp_path / "harness",
+    )
+    label = req["agent_model"]
+    assert label == "deepseek-flash-4.1"
+    assert DSH_MODEL_IDS[label] == "deepseek/deepseek-flash"
+
+
+def test_normalize_real_final_report():
+    report = {
+        "status": "success",
+        "final_validation": {
+            "hy3_tp8_qkv_proj_m16": {
+                "passed": True, "graph_capture_passed": True,
+                "mismatch_count": 0, "median_us": 42.0, "p90_us": 43.0,
+            },
+        },
+        "workers": {},
+    }
+    out = _normalize_report(_inst(), report)
+    assert out["status"] == "success"
+    assert out["correctness_ok"] is True
+    assert out["median_us"] == 42.0 and out["p90_us"] == 43.0
+
+
+def test_partial_worker_result_is_preserved_but_never_passes(tmp_path):
+    worker = tmp_path / "workers" / "worker_0"
+    worker.mkdir(parents=True)
+    (worker / "result.json").write_text(json.dumps({
+        "shapes": {_inst().id: {"metrics": {
+            "graph_capture_passed": True,
+            "correctness_passed_in_precheck": True,
+            "median_us": 40.0, "p90_us": 41.0,
+        }}},
+    }), encoding="utf-8")
+    out = _normalize_workspace_result(_inst(), tmp_path)
+    assert out["status"] == "partial_worker_result"
+    assert out["median_us"] == 40.0
+    assert out["correctness_ok"] is True
+    assert out["passed"] is False  # no final serial validation
+
+
+def test_real_evaluator_forces_isolated_repo_root(tmp_path, monkeypatch):
+    cfg = _cfg(tmp_path)
+    ev = DkaoCliEvaluator()
+    # The device is a pure layout decision (index % 4): the evaluator takes no
+    # reading of the host, so the test cannot depend on what is running on it.
+    seen = {}
+
+    def fake_run_one(cfg_, workspace, iteration, inst, index, repo_root,
+                     child_root, preflight=None, **kwargs):
+        seen["repo_root"] = repo_root
+        seen["child_root"] = child_root
+        seen["index"] = index
+        return {"status": "completed", "passed": True, "median_us": 1.0}
+
+    monkeypatch.setattr(ev, "_run_one", fake_run_one)
+    out = ev.evaluate(cfg, tmp_path / "harness", 1, [_inst()])
+    assert out[_inst().id]["passed"] is True
+    assert seen["repo_root"] == (tmp_path / "ahe-repos").resolve()
+    assert "children/iteration_001" in str(seen["child_root"])
+    assert seen["index"] == 0
+
+
+def test_agent_evolver_prompt_is_single_role_and_pool_bounded(tmp_path):
+    ev = AgentEvolver()
+    prompt = ev._prompt(tmp_path, 1, {
+        "overview": "overview.md", "selected": ["a"], "pool_ids": ["a", "b"],
+        "results": {}, "diff": {}, "decision": {},
+    })
+    assert "SINGLE AHE Evolve Agent" in prompt
+    assert "two responsibilities in this ONE role" in prompt
+    assert "invent pool ids" in prompt
+    assert "_ahe_change_manifest.json" in prompt
+    assert "_ahe_round_plan.json" in prompt
+
+
+def test_evolve_agent_defaults_to_flash_41(tmp_path):
+    """Both the adapter and build_evolver default to deepseek/deepseek-flash."""
+    from ..orchestrator.evolve import (
+        DEFAULT_EVOLVE_MODEL, LEGACY_EVOLVE_MODEL, build_evolver,
+    )
+
+    assert DEFAULT_EVOLVE_MODEL == "deepseek/deepseek-flash"
+    assert LEGACY_EVOLVE_MODEL == "deepseek/deepseek-v4-flash-0731"
+    assert AgentEvolver().model == DEFAULT_EVOLVE_MODEL
+    assert build_evolver("agent").model == DEFAULT_EVOLVE_MODEL
+    # A form answer still wins over the default.
+    assert build_evolver("agent", model=LEGACY_EVOLVE_MODEL).model == (
+        LEGACY_EVOLVE_MODEL)
+
+
+def test_last_result_parses_stream_json():
+    text = '\n'.join([
+        json.dumps({"type": "assistant", "message": "x"}),
+        json.dumps({"type": "result", "result": "done"}),
+    ])
+    assert _last_result(text) == "done"
+
+
+def test_real_mode_rejects_builtin_pool(tmp_path):
+    from ..orchestrator.config import ExperimentConfig
+    from ..orchestrator.pipeline import _load_pool
+    cfg = ExperimentConfig(
+        task_id="t", state_dir=tmp_path / "s", workspace_dir=tmp_path / "w",
+        suite=[], execution_mode="dkao-cli", agent_framework="dsh",
+        pool_source="builtin",
+    )
+    try:
+        _load_pool(cfg)
+        raise AssertionError("expected ValueError for builtin pool in real mode")
+    except ValueError:
+        pass
+
+
+def test_measured_pool_yaml_loads(tmp_path):
+    from ..orchestrator.pool_builder import build_pool
+    # a minimal fake historical report -> builder -> loader round trip
+    wd = tmp_path / "nodes" / "n" / "workspaces" / "task"
+    (wd / "final_report.json").parent.mkdir(parents=True)
+    (wd / "final_report.json").write_text(json.dumps({
+        "config": {
+            "model": "hy3",
+            "shapes": [{"id": "hy3_tp8_qkv_proj_m16", "M": 16, "N": 1280,
+                        "K": 4096, "tp_size": 8, "operator": "qkv_proj"}],
+        },
+        "initial_metrics": {"hy3_tp8_qkv_proj_m16": {"baseline_us": 100.0}},
+        "final_validation": {"hy3_tp8_qkv_proj_m16": {"median_us": 42.0,
+                                                       "p90_us": 43.0}},
+    }), encoding="utf-8")
+    data = build_pool(tmp_path)
+    assert len(data["instances"]) == 1
+    pool_path = tmp_path / "pool.yaml"
+    import yaml as _y
+    from ..orchestrator.pool import Pool
+    pool_path.write_text(_y.safe_dump(data, sort_keys=False), encoding="utf-8")
+    pool = Pool.from_yaml(pool_path)
+    assert pool.all_ids() == ["hy3_tp8_qkv_proj_m16"]
+    assert pool.get("hy3_tp8_qkv_proj_m16").baseline_us == 100.0
+
+
+def test_worker_rounds_counts_real_attempts(tmp_path):
+    """rounds_used falls back to the workers' own experiment records."""
+    exp = tmp_path / "workers" / "worker_0" / "runs" / "shapeA"
+    exp.mkdir(parents=True)
+    (exp / "experiments.jsonl").write_text(
+        '{"iteration": 1}\n{"iteration": 2}\n{"iteration": 3}\n',
+        encoding="utf-8")
+    assert _worker_rounds(tmp_path) == 3
+    assert _worker_rounds(tmp_path / "missing") == 0
+
+
+def test_child_env_enables_planner_by_default(tmp_path):
+    """The evaluated harness' planner only works with METAINFER_PLANNER=1."""
+    cfg = _cfg(tmp_path)
+    env = _child_env(cfg, tmp_path / "repos", tmp_path / "ws")
+    assert env["METAINFER_PLANNER"] == "1"
+    assert env["METAINFER_KERNEL_REPOS"] == str(tmp_path / "repos")
+    assert env["METAINFER_HARNESS_ROOT"] == str(tmp_path / "ws")
+
+
+def test_child_env_planner_can_be_disabled(tmp_path):
+    cfg = _cfg(tmp_path)
+    cfg.answers["planner_enabled"] = "false"
+    env = _child_env(cfg, tmp_path / "repos", tmp_path / "ws")
+    assert "METAINFER_PLANNER" not in env
+
+
+def test_child_env_uses_task_scope_and_quick_bench_by_default(tmp_path):
+    """One question = one shape: AHE defaults to the cheap validation budget."""
+    cfg = _cfg(tmp_path)
+    env = _child_env(cfg, tmp_path / "repos", tmp_path / "ws")
+    assert env["METAINFER_VALIDATE_SCOPE"] == "task"
+    assert env["METAINFER_BENCH_WARMUPS"] == "30"
+    assert env["METAINFER_BENCH_SAMPLES"] == "30"
+    assert env["METAINFER_BENCH_REPLAYS"] == "50"
+
+
+def test_child_env_full_budget_can_be_requested(tmp_path):
+    cfg = _cfg(tmp_path)
+    cfg.answers["validate_scope"] = "api"
+    cfg.answers["bench_profile"] = "full"
+    env = _child_env(cfg, tmp_path / "repos", tmp_path / "ws")
+    assert env["METAINFER_VALIDATE_SCOPE"] == "api"
+    assert "METAINFER_BENCH_WARMUPS" not in env
+    assert "METAINFER_BENCH_SAMPLES" not in env
+
+
+def test_child_env_explicit_bench_numbers_win(tmp_path):
+    cfg = _cfg(tmp_path)
+    cfg.answers["bench_samples"] = "12"
+    env = _child_env(cfg, tmp_path / "repos", tmp_path / "ws")
+    assert env["METAINFER_BENCH_SAMPLES"] == "12"
+
+
+def test_env_failure_detection_includes_compile_and_timeout():
+    """Compile errors and timeouts are environmental: they earn a retry.
+
+    The rule is "no usable number => the environment owes us a measurement".
+    A reported compile error *with* a number is contradictory, and the number is
+    the stronger evidence: the kernel ran, so the round scores it (as a loss)
+    instead of throwing it away and re-running a question that already measured.
+    """
+    from metainfer.tasks.harness_evolve.orchestrator.adapters.eval import (
+        DkaoCliEvaluator,
+    )
+    ev = DkaoCliEvaluator()
+    assert ev._env_failed({"status": "success", "passed": False,
+                           "median_us": None,
+                           "reason": "ninja: build stopped: subcommand failed"})
+    assert ev._env_failed({"status": "timeout", "passed": False,
+                           "median_us": None, "reason": "timed out after 900s"})
+    measured_with_error_text = {"status": "rejected", "passed": False,
+                                "median_us": 12.0, "correctness_ok": True,
+                                "reason": "hipcc: error: undefined template DUFragment"}
+    assert ev._env_failed(measured_with_error_text) is False
+    # a genuine performance failure is NOT environmental
+    assert not ev._env_failed({"status": "success", "passed": False,
+                               "median_us": 900.0,
+                               "reason": "median above gate"})
+    # a passing result is never an environment failure
+    assert not ev._env_failed({"status": "success", "passed": True,
+                               "median_us": 12.0})
+
+
+def _pool_inst(inst_id: str):
+    from types import SimpleNamespace
+    return SimpleNamespace(id=inst_id, shape={"M": 16, "N": 16, "K": 16},
+                           budget_rounds=3, notes="")
+
+
+def test_a_measured_partial_result_is_not_an_environment_failure():
+    ev = DkaoCliEvaluator()
+    measured = {"status": "partial_worker_result", "passed": False,
+                "correctness_ok": True, "median_us": 8.52,
+                "reason": "final report missing; worker best preserved as evidence"}
+    assert ev._env_failed(measured) is False          # the number is evidence
+    # ... but on a re-run the same partial *is* an environment problem: the
+    # device was shared while it ran.
+    assert ev._env_failed(measured, retry=True) is True
+
+    unmeasured = {"status": "partial_worker_result", "passed": False,
+                  "median_us": None, "reason": "no worker numbers"}
+    assert ev._env_failed(unmeasured) is True
+
+    missing = {"status": "missing_report", "passed": False, "median_us": None}
+    assert ev._env_failed(missing) is True
+    assert ev._env_failed({"status": "success", "passed": True,
+                           "median_us": 5.0}) is False
+
+
+def test_every_attempt_gets_its_own_child_directory(tmp_path, monkeypatch):
+    """A retried question must not reuse its first attempt's child workspace.
+
+    DKAO binds ``children///workspace/main`` to the exact
+    kernel repo it was launched with and refuses to run when that symlink points
+    elsewhere. Re-running into the same directory therefore kills the retry
+    instantly ("workspace main already points to ...") and an environment hiccup
+    turns into a burned retry budget. Each attempt owns a directory instead, so
+    the failed attempt's evidence also survives next to the retry's.
+    """
+    cfg = _cfg(tmp_path)
+    cfg.answers["env_retry_attempts"] = 2
+    ev = DkaoCliEvaluator()
+    dirs: List[str] = []
+
+    def fake_run_one(cfg_, ws, iteration, inst, gpu, rr, cr, layout=None,
+                     **kwargs):
+        dirs.append(cr.name)
+        # A real child leaves a workspace bound to its kernel repo behind.
+        (cr / inst.id / "workspace").mkdir(parents=True, exist_ok=True)
+        (cr / inst.id / "orchestrator-external.log").write_text(
+            f"attempt {len(dirs)}", encoding="utf-8")
+        if len(dirs) == 1:
+            return {"status": "missing_report", "passed": False,
+                    "median_us": None, "returncode": 1}
+        return {"status": "completed", "passed": True, "median_us": 1.0}
+
+    monkeypatch.setattr(ev, "_run_one", fake_run_one)
+    out = ev.evaluate(cfg, tmp_path / "harness", 1, [_inst()])
+
+    assert len(dirs) == 2 and len(set(dirs)) == 2, (
+        f"each attempt needs its own directory, got {dirs}")
+    # the first attempt keeps the plain name; a retry is suffixed
+    assert dirs == ["iteration_001", "iteration_001_a2"]
+    # the retry ran in a fresh directory, so the first attempt's evidence is
+    # still there to explain why it failed
+    first = cfg.exp_dir / "children" / dirs[0] / _inst().id
+    assert (first / "orchestrator-external.log").read_text() == "attempt 1"
+    assert out[_inst().id]["passed"] is True
+
+
+def _events(cfg):
+    path = cfg.state_dir / "timeline.jsonl"
+    if not path.is_file():
+        return []
+    return [json.loads(l)["type"] for l in
+            path.read_text().splitlines() if l.strip()]
+
+
+def test_a_round_hands_every_question_to_dkao_without_checking_the_cards(
+        tmp_path, monkeypatch):
+    """The dispatch never probes a device: it lays questions out and publishes.
+
+    Eight questions are two waves of four, one question per device, and the
+    evaluator never asks the host anything about the cards.
+    """
+    cfg = _cfg(tmp_path)
+    ev = DkaoCliEvaluator()
+    seen = []
+
+    def fake_run_one(cfg_, ws, iteration, inst, gpu, rr, cr, layout=None,
+                     **kwargs):
+        # the device HE hands the question to is the fixed index % 4 rotation
+        seen.append((inst.id, gpu))
+        return {"status": "completed", "passed": True, "median_us": 1.0}
+
+    monkeypatch.setattr(ev, "_run_one", fake_run_one)
+    out = ev.evaluate(cfg, tmp_path / "harness", 1,
+                      [_pool_inst(i) for i in ("a", "b", "c", "d",
+                                               "e", "f", "g", "h")])
+    assert set(out) == set("abcdefgh")
+    # one question per device, in order, wrapping onto device 0 for wave two
+    assert seen == [("a", 0), ("b", 1), ("c", 2), ("d", 3),
+                    ("e", 0), ("f", 1), ("g", 2), ("h", 3)]
+    # nothing in the round reads or leases a device
+    layout = (cfg.exp_dir / "runs" / "iteration_001" / "input" / "benchmark"
+              / "gpu_preflight.json")
+    saved = json.loads(layout.read_text(encoding="utf-8"))
+    assert saved["leases"] == [] and saved["enabled"] is False
+    assert saved["managed_by"] == "dcu_kernel_auto_opt"
+    assert "gpu_layout_fixed" in _events(cfg)
+
+
+def test_a_dkao_gate_verdict_stops_the_round_instead_of_relaunching(
+        tmp_path, monkeypatch):
+    """Exit 75 is the child's gate giving up after its own 24 h budget.
+
+    HE has no cards to reassign, so relaunching the child would mean ignoring a
+    verdict it has no standing to overrule: the round stops and says so.
+    """
+    cfg = _cfg(tmp_path)
+    ev = DkaoCliEvaluator()
+    calls = []
+
+    def fake_run_one(cfg_, ws, iteration, inst, gpu, rr, cr, layout=None,
+                     **kwargs):
+        calls.append(inst.id)
+        return {"status": "environment_failed", "passed": False,
+                "median_us": None, "returncode": 75, "gate_blocked": True}
+
+    monkeypatch.setattr(ev, "_run_one", fake_run_one)
+    out = ev.evaluate(cfg, tmp_path / "harness", 1,
+                      [_pool_inst(i) for i in ("a", "b", "c", "d")])
+    # exactly one attempt per question — no retry storm against a busy card
+    assert calls == ["a", "b", "c", "d"]
+    assert all(r["gate_blocked"] for r in out.values())
+    stop = json.loads((cfg.exp_dir / "stop_requested.json").read_text(
+        encoding="utf-8"))
+    assert stop["reason"] == "round_incomplete"
+    assert stop["cause"] == "gpu_gate_blocked"
+    assert stop["unmeasured"] == ["a", "b", "c", "d"]
+    assert "env_retry_exhausted" not in _events(cfg)
+
+
+def test_an_ordinary_environment_failure_is_still_retried(tmp_path, monkeypatch):
+    """A broken toolchain is not a gate verdict: the environment owes a number."""
+    cfg = _cfg(tmp_path)
+    cfg.answers["env_retry_attempts"] = 2
+    ev = DkaoCliEvaluator()
+    attempts = []
+
+    def fake_run_one(cfg_, ws, iteration, inst, gpu, rr, cr, layout=None,
+                     **kwargs):
+        attempts.append(inst.id)
+        if len(attempts) == 1:
+            return {"status": "missing_report", "passed": False,
+                    "median_us": None, "returncode": 1}
+        return {"status": "completed", "passed": True, "median_us": 1.0}
+
+    monkeypatch.setattr(ev, "_run_one", fake_run_one)
+    out = ev.evaluate(cfg, tmp_path / "harness", 1, [_inst()])
+    assert attempts == [_inst().id, _inst().id]
+    assert out[_inst().id]["passed"] is True
+    assert (cfg.exp_dir / "children" / "iteration_001_a2").is_dir(), (
+        "the retry ran in the second attempt's own directory")
+    assert "env_retry_exhausted" not in _events(cfg)
+
+
+def test_a_gate_verdict_is_never_counted_as_a_measurement(tmp_path, monkeypatch):
+    """A question the gate blocked has no number, so the round is incomplete."""
+    cfg = _cfg(tmp_path)
+    ev = DkaoCliEvaluator()
+    res = {"status": "environment_failed", "passed": False,
+           "median_us": None, "gate_blocked": True}
+    assert ev._is_measured(res) is False
+    assert ev._incomplete_ids([_inst()], {_inst().id: res}) == [_inst().id]
+
+
diff --git a/metainfer/tasks/harness_evolve/tests/test_repo_cleanup.py b/metainfer/tasks/harness_evolve/tests/test_repo_cleanup.py
new file mode 100644
index 00000000..47ccb699
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_repo_cleanup.py
@@ -0,0 +1,152 @@
+"""Tests for orchestrator.repo_cleanup — 任务级候选仓库清理的安全边界。
+
+重点验证"绝不误删"：
+- 只匹配 ``ahe-{task_id}-it`` 精确前缀（task_id 互为前缀时不误伤）
+- 池白名单（registered_pool.yaml 引用）一律保留
+- 符号链接 / 非目录 / 越界路径一律跳过
+- dry_run 不做任何修改
+"""
+
+from __future__ import annotations
+
+import textwrap
+
+import pytest
+
+from ..orchestrator.repo_cleanup import (
+    load_pool_keep_list,
+    owned_repos,
+    purge_task_kernel_repos,
+    repo_name_prefix,
+    resolve_ahe_repo_root,
+)
+
+TASK = "test-9-8-5-e10567a1"
+
+
+def _mk_repo(root, name):
+    d = root / name
+    d.mkdir(parents=True)
+    (d / "kernel.py").write_text("print('x')\n", encoding="utf-8")
+    (d / "sub").mkdir()
+    (d / "sub" / "blob.bin").write_bytes(b"0" * 32)
+    return d
+
+
+def _pool(root, names):
+    lines = ["schema_version: 1", "instances:"]
+    for n in names:
+        lines.append(f"- id: {n}")
+        lines.append(f"  source_report: {root}/{n}/final_report.json")
+    (root / "registered_pool.yaml").write_text(
+        "\n".join(lines) + "\n", encoding="utf-8")
+
+
+def test_repo_name_prefix_uses_it_marker():
+    assert repo_name_prefix(TASK) == f"ahe-{TASK}-it"
+
+
+def test_resolve_ahe_repo_root_prefers_answer(tmp_path):
+    custom = tmp_path / "custom-root"
+    assert resolve_ahe_repo_root({"ahe_repo_root": str(custom)}) == custom.resolve()
+
+
+def test_resolve_ahe_repo_root_defaults_when_absent():
+    assert str(resolve_ahe_repo_root(None)).endswith("ahe-kernel-repos")
+
+
+def test_prefix_protects_repo_when_task_id_is_prefix_of_another(tmp_path):
+    """``ahe--it`` 结尾的 '-it' 保证不误伤同前缀的其他任务。"""
+    _mk_repo(tmp_path, f"ahe-{TASK}-it001-candidate-foo-aaaaaa")
+    other = _mk_repo(tmp_path, f"ahe-{TASK}-extra-it001-candidate-bar-bbbbbb")
+
+    deletable, _ = owned_repos(TASK, tmp_path, keep=set())
+    names = {p.name for p in deletable}
+
+    assert f"ahe-{TASK}-it001-candidate-foo-aaaaaa" in names
+    assert other.name not in names, "task_id 同前缀的其他任务仓库被误列"
+
+
+def test_unrelated_task_repos_untouched(tmp_path):
+    _mk_repo(tmp_path, "ahe-other-task-it001-candidate-foo-aaaaaa")
+    _mk_repo(tmp_path, "dsh-min-1iter-abff697d")  # 人工仓库
+    deletable, _ = owned_repos(TASK, tmp_path, keep=set())
+    assert deletable == []
+
+
+def test_pool_whitelist_is_never_deleted(tmp_path):
+    keep_name = f"ahe-{TASK}-it001-candidate-keepme-aaaaaa"
+    drop_name = f"ahe-{TASK}-it002-candidate-dropme-bbbbbb"
+    _mk_repo(tmp_path, keep_name)
+    _mk_repo(tmp_path, drop_name)
+    _pool(tmp_path, [keep_name])
+
+    keep = load_pool_keep_list(tmp_path)
+    assert keep_name in keep
+
+    deletable, skipped = owned_repos(TASK, tmp_path, keep=keep)
+    names = {p.name for p in deletable}
+    assert names == {drop_name}
+    assert any(keep_name in s for s in skipped)
+
+
+def test_symlink_repos_are_skipped(tmp_path):
+    real = _mk_repo(tmp_path, f"ahe-{TASK}-it003-candidate-real-cccccc")
+    link = tmp_path / f"ahe-{TASK}-it004-candidate-link-dddddd"
+    link.symlink_to(real)
+
+    deletable, skipped = owned_repos(TASK, tmp_path, keep=set())
+    names = {p.name for p in deletable}
+    assert real.name in names
+    assert link.name not in names
+    assert any(link.name in s for s in skipped)
+
+
+def test_dry_run_changes_nothing(tmp_path):
+    name = f"ahe-{TASK}-it005-candidate-x-eeeeee"
+    repo = _mk_repo(tmp_path, name)
+
+    summary = purge_task_kernel_repos(TASK, tmp_path, keep=set(), dry_run=True)
+    assert summary["dry_run"] is True
+    assert summary["matched"] == 1
+    assert summary["deleted"] == []
+    assert summary["bytes"] > 0
+    assert repo.exists(), "dry_run 不应删除任何东西"
+
+
+def test_apply_removes_only_owned_repos(tmp_path):
+    owned = _mk_repo(tmp_path, f"ahe-{TASK}-it006-candidate-x-ffffff")
+    kept = _mk_repo(tmp_path, "ahe-other-it006-candidate-y-111111")
+
+    summary = purge_task_kernel_repos(TASK, tmp_path, keep=set(), dry_run=False)
+    assert len(summary["deleted"]) == 1
+    assert not summary["failed"]
+    assert not owned.exists()
+    assert kept.exists(), "其他任务的仓库被误删"
+
+
+def test_missing_repo_root_is_reported_not_raised(tmp_path):
+    missing = tmp_path / "nope"
+    summary = purge_task_kernel_repos(TASK, missing, dry_run=True)
+    assert summary["matched"] == 0
+    assert summary["skipped"]
+
+
+def test_refuses_when_task_still_registered(tmp_path, monkeypatch):
+    """任务仍在注册表 → 默认拒绝清理（force=False）。"""
+    name = f"ahe-{TASK}-it010-candidate-live-aaaaaa"
+    repo = _mk_repo(tmp_path, name)
+
+    from ..orchestrator import repo_cleanup as rc
+    monkeypatch.setattr(rc, "is_task_registered", lambda tid: True)
+
+    summary = rc.purge_task_kernel_repos(TASK, tmp_path, keep=set(), dry_run=False)
+    assert summary.get("refused") == "task_still_registered"
+    assert summary["deleted"] == []
+    assert repo.exists(), "注册中的任务仓库被误删"
+
+    # force=True 时才允许
+    summary2 = rc.purge_task_kernel_repos(
+        TASK, tmp_path, keep=set(), dry_run=False, force=True)
+    assert len(summary2["deleted"]) == 1
+    assert not repo.exists()
diff --git a/metainfer/tasks/harness_evolve/tests/test_resume.py b/metainfer/tasks/harness_evolve/tests/test_resume.py
new file mode 100644
index 00000000..4e6c4004
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_resume.py
@@ -0,0 +1,356 @@
+"""Resume-the-next-iteration support for harness_evolve."""
+
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+
+from metainfer.tasks.harness_evolve.orchestrator.cli import resume_command
+from metainfer.tasks.harness_evolve.orchestrator.config import (
+    load_experiment_config,
+)
+from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+    completed_iterations, run_experiment,
+)
+from metainfer.tasks.harness_evolve.orchestrator.state import (
+    read_target_iterations, write_target_iterations,
+)
+from metainfer.tasks.harness_evolve.server.routes import (
+    _launch_resume, _resume_running,
+)
+
+
+def _write_req(tmp_path: Path) -> Path:
+    req = tmp_path / "requirements.json"
+    req.write_text(json.dumps({
+        "task_id": "resume-task",
+        "task_type": "harness-evolve",
+        "execution_mode": "dry-run",
+        "max_iterations": 1,
+        "evolve_mode": "dry-run",
+        "per_round_budget": 2,
+        "pool_source": "",
+    }), encoding="utf-8")
+    return req
+
+
+def _cfg(tmp_path: Path):
+    return load_experiment_config(
+        _write_req(tmp_path), tmp_path / "state", tmp_path / "ws")
+
+
+def test_resume_runs_next_iteration_with_champion_from_disk(tmp_path):
+    cfg = _variant_cfg(tmp_path, rounds=2)      # two rounds: baseline + one pass
+    write_target_iterations(cfg.exp_dir, 1)     # this call may run round 1 only
+    run_experiment(cfg)
+    exp = cfg.exp_dir
+    assert completed_iterations(exp) == 1
+    assert (exp / "runs" / "iteration_001" / "input" / "decision.json").is_file()
+
+    write_target_iterations(cfg.exp_dir, 2)     # ... and now it may run round 2
+    run_experiment(cfg, start_iteration=2, iterations_to_run=1)
+
+    assert completed_iterations(exp) == 2
+    it2 = exp / "runs" / "iteration_002" / "input"
+    assert (it2 / "benchmark" / "results.json").is_file()
+    assert (it2 / "decision.json").is_file()
+    # champion/prev state must come from the resumed experiment
+    decision = json.loads((it2 / "decision.json").read_text(encoding="utf-8"))
+    assert decision["iteration"] == 2
+    scores = (exp / "iteration_scores.jsonl").read_text(
+        encoding="utf-8").strip().splitlines()
+    assert len(scores) == 2
+    assert json.loads(scores[-1])["iteration"] == 2
+
+
+def test_cli_resume_starts_round_one_when_nothing_ran(tmp_path):
+    """No completed round and nobody running: resume bootstraps round 1."""
+    state = tmp_path / "state"
+    state.mkdir(parents=True)
+    _write_req(state if False else tmp_path)  # requirements at tmp_path
+    (state / "requirements.json").write_text(
+        (tmp_path / "requirements.json").read_text(encoding="utf-8"),
+        encoding="utf-8")
+    import argparse
+    args = argparse.Namespace(
+        state_dir=state, workspace_dir=tmp_path / "ws",
+        requirements=None, iterations=1)
+    assert resume_command(args) == 0          # not a refusal any more
+    assert (tmp_path / "ws" / "runs" / "iteration_001").is_dir()
+
+
+def test_cli_resume_defers_to_an_already_running_task(tmp_path, monkeypatch):
+    """If the task is already running (started by the server), resume exits."""
+    state = tmp_path / "state"
+    state.mkdir(parents=True)
+    _write_req(tmp_path)
+    (state / "requirements.json").write_text(
+        (tmp_path / "requirements.json").read_text(encoding="utf-8"),
+        encoding="utf-8")
+    import argparse
+    from metainfer.tasks.harness_evolve.orchestrator import supervisor as sup
+    monkeypatch.setattr(sup, "task_process_alive", lambda _state: 4242)
+    args = argparse.Namespace(
+        state_dir=state, workspace_dir=tmp_path / "ws",
+        requirements=None, iterations=1)
+    assert resume_command(args) == 0
+    assert not (tmp_path / "ws" / "runs" / "iteration_001").is_dir()
+
+
+def test_resume_running_detects_live_pid(tmp_path):
+    """Only a live harness_evolve process counts (zombies/pids are rejected)."""
+    import os
+    import subprocess
+    import sys
+
+    state = tmp_path / "state"
+    state.mkdir(parents=True)
+    assert _resume_running(state) is None
+
+    # our own pid is alive but is not a harness run
+    (state / "resume.pid").write_text(str(os.getpid()), encoding="utf-8")
+    assert _resume_running(state) is None
+
+    proc = subprocess.Popen(
+        [sys.executable, "-c", "import time; time.sleep(30)", "harness_evolve"])
+    try:
+        (state / "resume.pid").write_text(str(proc.pid), encoding="utf-8")
+        assert _resume_running(state) == proc.pid
+    finally:
+        proc.kill()
+        proc.wait()
+
+    (state / "resume.pid").write_text("999999", encoding="utf-8")
+    assert _resume_running(state) is None
+
+
+def test_launch_resume_end_to_end_dry(tmp_path):
+    """The resume CLI really runs the next round (dry-run fixture)."""
+    cfg = _cfg(tmp_path)
+    run_experiment(cfg)
+    state_dir = cfg.state_dir
+    (state_dir / "requirements.json").write_text(
+        (tmp_path / "requirements.json").read_text(encoding="utf-8"),
+        encoding="utf-8")
+
+    # The form's Iterations answer is the run's round target; raise it so the
+    # resumed run is allowed to spend another round (a performance pass costs
+    # 1.0, the retake 0.5, the baseline round is free).
+    write_target_iterations(cfg.exp_dir, 3)
+    out = _launch_resume(state_dir, cfg.exp_dir, 1)
+    assert out["ok"] is True and out["pid"] > 0
+
+    deadline = time.time() + 180
+    while time.time() < deadline:
+        if completed_iterations(cfg.exp_dir) >= 2:
+            break
+        time.sleep(1)
+    assert completed_iterations(cfg.exp_dir) >= 2, (state_dir / "resume.log").read_text(
+        encoding="utf-8", errors="replace")[-1500:]
+
+
+def test_target_iterations_override_fallback_rounds(tmp_path):
+    """The target (a round count) decides, not the caller's batch size."""
+    cfg = _variant_cfg(tmp_path)
+    write_target_iterations(cfg.exp_dir, 3)
+    assert read_target_iterations(cfg.exp_dir) == 3
+    run_experiment(cfg, iterations_to_run=1)          # the caller asks for 1
+    assert completed_iterations(cfg.exp_dir) == 3     # the target wins
+
+
+def test_lowered_target_stops_the_loop_immediately(tmp_path):
+    """Lowering the target below the next round stops without a new round."""
+    cfg = _cfg(tmp_path)
+    run_experiment(cfg)                               # iteration 1 done
+    write_target_iterations(cfg.exp_dir, 1)           # target already met
+    run_experiment(cfg, start_iteration=2, iterations_to_run=5)
+    assert completed_iterations(cfg.exp_dir) == 1
+    assert not (cfg.exp_dir / "runs" / "iteration_002").exists()
+
+
+def test_target_can_be_cleared_for_fixed_rounds(tmp_path):
+    cfg = _cfg(tmp_path)
+    write_target_iterations(cfg.exp_dir, 5)
+    from metainfer.tasks.harness_evolve.orchestrator.state import (
+        clear_target_iterations,
+    )
+    clear_target_iterations(cfg.exp_dir)
+    assert read_target_iterations(cfg.exp_dir) is None
+
+
+def test_champion_iteration_pins_the_comparison_baseline(tmp_path):
+    """Resume can compare the new round against a chosen finished iteration."""
+    cfg = _variant_cfg(tmp_path, rounds=3)      # three rounds of headroom
+    # one iteration per call, so the test controls exactly which rounds exist
+    write_target_iterations(cfg.exp_dir, 1)
+    run_experiment(cfg)                                            # iteration 1
+    write_target_iterations(cfg.exp_dir, 2)
+    run_experiment(cfg, start_iteration=2, iterations_to_run=1)     # iteration 2
+    write_target_iterations(cfg.exp_dir, 3)
+    run_experiment(cfg, start_iteration=3, iterations_to_run=1,
+                   champion_iteration=2)                            # iteration 3
+
+    exp = cfg.exp_dir
+    best = json.loads((exp / "best_ever.json").read_text(encoding="utf-8"))
+    assert best["iteration"] == 2
+    assert best["verdict"] == "MANUAL_BASELINE"
+    assert (exp / "champion_override.json").is_file()
+
+    it2 = json.loads((exp / "runs" / "iteration_002" / "input" / "benchmark"
+                      / "results.json").read_text(encoding="utf-8"))["results"]
+    decision = json.loads((exp / "runs" / "iteration_003" / "input"
+                           / "decision.json").read_text(encoding="utf-8"))
+    per = decision["performance_gate"]["comparison"]["per_instance"]
+    assert per, "paired comparison must exist"
+    # The pin is what this test is about: the run's comparison baseline is
+    # iteration 2, so that iteration's results are what a resumed round reads
+    # back (the gate then compares the candidate against the variant table, as
+    # every round does).
+    assert json.loads((exp / "champion_override.json").read_text()
+                      )["champion_iteration"] == 2
+    assert json.loads((exp / "best_ever.json").read_text())["iteration"] == 2
+    assert set(it2) == set(per), (
+        "the pinned round and the resumed round must pair the same operators")
+
+
+def test_resume_clears_a_previous_stop_signal(tmp_path):
+    """A gate-stopped run must be resumable: the stale stop signal is dropped.
+
+    ``stop_requested.json`` is checked after every evaluate, so leaving it in
+    place would make the very next resume stop again without doing any work.
+    """
+    cfg = _cfg(tmp_path)
+    run_experiment(cfg)
+    stop = cfg.exp_dir / "stop_requested.json"
+    stop.write_text(json.dumps({"reason": "gpu_gate_blocked", "checks": 48}),
+                    encoding="utf-8")
+
+    from metainfer.tasks.harness_evolve.orchestrator.cli import clear_stop_request
+    assert clear_stop_request(cfg.exp_dir, state_dir=cfg.state_dir) == [
+        "stop_requested.json"]
+    assert not stop.exists()
+    events = [json.loads(l)["type"] for l in
+              (cfg.state_dir / "timeline.jsonl").read_text(
+                  encoding="utf-8").splitlines() if l.strip()]
+    assert "stop_request_cleared" in events
+    # idempotent: without a signal there is nothing to clear
+    assert clear_stop_request(cfg.exp_dir, state_dir=cfg.state_dir) == []
+
+
+def test_completed_iterations_counts_a_judged_round_without_results_json(tmp_path):
+    """A variant-protocol stage is "done" through its decision artifacts.
+
+    Only a *failing* pass writes ``input/benchmark/results.json``; a round that
+    passed its gate leaves ``decision.json`` + the gate files instead. Reading
+    only ``results.json`` reported 0 completed rounds for a run that had just
+    finished three, and ``resume`` then refused to continue.
+    """
+    exp = tmp_path / "exp"
+    for n, artifact in ((1, "baseline_round.json"),
+                        (2, "performance_gate.json"),
+                        (3, "generalization_gate.json")):
+        d = exp / "runs" / f"iteration_{n:03d}" / "input"
+        d.mkdir(parents=True)
+        (d / "decision.json").write_text("{}", encoding="utf-8")
+        (d / artifact).write_text("{}", encoding="utf-8")
+    # an empty directory is not a completed round
+    (exp / "runs" / "iteration_004" / "input").mkdir(parents=True)
+
+    assert completed_iterations(exp) == 3
+
+    # the classic evidence (a stage that stopped early) still counts
+    (exp / "runs" / "iteration_005" / "input" / "benchmark").mkdir(parents=True)
+    (exp / "runs" / "iteration_005" / "input" / "benchmark"
+     / "results.json").write_text("{}", encoding="utf-8")
+    assert completed_iterations(exp) == 5
+
+
+# ------------------------------------------------- rounds as durable state
+
+def _variant_cfg(tmp_path: Path, **answers):
+    """A config in the variant-round protocol: it needs a real pool file."""
+    import yaml
+
+    ops = {f"op{i:02d}": {"baseline_us": 100.0 + i, "best_known_us": 40.0 + 4 * i}
+           for i in range(8)}
+    pool = tmp_path / "pool.yaml"
+    pool.write_text(yaml.safe_dump({
+        "schema_version": 1,
+        "instances": [
+            {"id": iid, "family": f"decode__fam{i}", "contract": {"M": 16},
+             "baseline_us": v["baseline_us"], "best_known_us": v["best_known_us"],
+             "history": []}
+            for i, (iid, v) in enumerate(sorted(ops.items()))
+        ],
+    }, sort_keys=False), encoding="utf-8")
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    req = tmp_path / "requirements.json"
+    req.write_text(json.dumps({
+        "task_id": "rounds-task",
+        "task_type": "harness-evolve",
+        "execution_mode": "dry-run",
+        "evolve_mode": "dry-run",
+        "pool_source": str(pool),
+        "question_pool": str(pool),
+        "round_questions": "variant",
+        "per_round_budget": "4",
+        "max_iterations": 1,
+        **answers,
+    }), encoding="utf-8")
+    return load_experiment_config(req, tmp_path / "state", tmp_path / "ws")
+
+
+def test_rounds_used_survives_a_resume(tmp_path):
+    """The budget must not reset when the process is restarted.
+
+    ``rounds_used`` used to live only in memory, so a resumed run came back with
+    a full budget and could keep spending rounds the operator never asked for.
+    """
+    from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+        _rounds_used_on_disk,
+    )
+
+    cfg = _variant_cfg(tmp_path, rounds=3)
+    run_experiment(cfg, iterations_to_run=2)          # baseline + one judgment
+    spent = _rounds_used_on_disk(cfg.exp_dir)
+    assert spent >= 1.0, "a judged pass must leave its cost on disk"
+    assert json.loads((cfg.state_dir / "run.json").read_text())["rounds_used"] == spent
+    # a second process (no memory of the first) sees the same number
+    assert _rounds_used_on_disk(cfg.exp_dir) == spent
+
+
+def test_rounds_used_is_derived_when_run_json_predates_it(tmp_path):
+    """Old experiments report their spent rounds instead of a fresh budget."""
+    from metainfer.tasks.harness_evolve.orchestrator.pipeline import (
+        _rounds_used_on_disk,
+    )
+
+    exp = tmp_path / "exp"
+    for n, cost in ((1, 0.0), (2, 1.0), (3, 0.5)):
+        d = exp / "runs" / f"iteration_{n:03d}" / "input"
+        d.mkdir(parents=True)
+        (d / "stage.json").write_text(json.dumps({"round_cost": cost}),
+                                      encoding="utf-8")
+    assert _rounds_used_on_disk(exp) == 1.5
+
+
+def test_a_round_is_an_iteration(tmp_path):
+    """`Rounds: 2` means the run may reach iteration 2 and no further.
+
+    Round 1 is the baseline round; round 2 is a performance pass (plus, when it
+    wins, the generalization retake inside that same iteration — a retake never
+    consumes a round of its own).
+    """
+    cfg = _variant_cfg(tmp_path, rounds=2)
+    write_target_iterations(cfg.exp_dir, 2)
+    run_experiment(cfg, iterations_to_run=5)          # ask for five; rounds say two
+
+    stages = [json.loads(p.read_text())["stage"]
+              for p in sorted((cfg.exp_dir / "runs").glob("iteration_*/input/stage.json"))]
+    assert stages == ["baseline", "performance"]
+    assert completed_iterations(cfg.exp_dir) == 2
+    run = json.loads((cfg.state_dir / "run.json").read_text())
+    assert run["finished"] is True
+    # the reported cost is bookkeeping: one judged pass
+    assert run["rounds_used"] == 1.0
diff --git a/metainfer/tasks/harness_evolve/tests/test_round_incomplete.py b/metainfer/tasks/harness_evolve/tests/test_round_incomplete.py
new file mode 100644
index 00000000..4ed589d0
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_round_incomplete.py
@@ -0,0 +1,144 @@
+"""A round that could not measure every question gets no verdict.
+
+``evaluate_gate`` used to answer ``FAIL`` whenever a drawn question produced no
+number. That reads as "the harness is worse than the variant", which is a claim
+the round did not support — and the loop acted on it: it went through
+analyze→evolve to fix a harness that may have been fine, and on the
+generalization retake it spent the candidate on a missing measurement.
+
+This drives the real loop (dry-run evaluator) with one question's result removed
+from the performance pass and checks what the run does about it.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import yaml
+
+from metainfer.tasks.harness_evolve.orchestrator import pipeline as P
+from metainfer.tasks.harness_evolve.orchestrator.config import (
+    load_experiment_config,
+)
+from metainfer.tasks.harness_evolve.orchestrator.pipeline import run_experiment
+from metainfer.tasks.harness_evolve.orchestrator.state import (
+    write_target_iterations,
+)
+
+
+def _variant_cfg(tmp_path: Path, **answers):
+    """Same shape as the resume tests: a real pool file, four questions."""
+    ops = {f"op{i:02d}": {"baseline_us": 100.0 + i, "best_known_us": 40.0 + 4 * i}
+           for i in range(8)}
+    pool = tmp_path / "pool.yaml"
+    pool.write_text(yaml.safe_dump({
+        "schema_version": 1,
+        "instances": [
+            {"id": iid, "family": f"decode__fam{i}", "contract": {"M": 16},
+             "baseline_us": v["baseline_us"], "best_known_us": v["best_known_us"],
+             "history": []}
+            for i, (iid, v) in enumerate(sorted(ops.items()))
+        ],
+    }, sort_keys=False), encoding="utf-8")
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    req = tmp_path / "requirements.json"
+    req.write_text(json.dumps({
+        "task_id": "incomplete-task",
+        "task_type": "harness-evolve",
+        "execution_mode": "dry-run",
+        "evolve_mode": "dry-run",
+        "pool_source": str(pool),
+        "question_pool": str(pool),
+        "round_questions": "variant",
+        "per_round_budget": "4",
+        "max_iterations": 1,
+        **answers,
+    }), encoding="utf-8")
+    return load_experiment_config(req, tmp_path / "state", tmp_path / "ws")
+
+
+def _timeline_events(cfg) -> list:
+    path = Path(cfg.state_dir) / "timeline.jsonl"
+    if not path.is_file():
+        return []
+    return [json.loads(line)["type"]
+            for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
+
+
+def test_a_round_missing_a_measurement_carries_no_verdict(tmp_path, monkeypatch):
+    cfg = _variant_cfg(tmp_path)
+    write_target_iterations(cfg.exp_dir, 2)
+    real = P.build_evaluator("dry-run")
+    dropped: dict = {}
+
+    def factory(_mode):
+        class Sabotaged:
+            """The dry-run evaluator, minus one question on the pass round."""
+
+            def evaluate(self, cfg_, snapshot, iteration, instances=None):
+                results = real.evaluate(cfg_, snapshot, iteration,
+                                        instances=instances)
+                if (str(cfg_.answers.get("round_stage")) == "performance"
+                        and results):
+                    victim = sorted(results)[0]
+                    dropped["id"] = victim
+                    results = {k: v for k, v in results.items() if k != victim}
+                return results
+
+        return Sabotaged()
+
+    monkeypatch.setattr(P, "build_evaluator", factory)
+
+    run_experiment(cfg)
+
+    exp = cfg.exp_dir
+    # round 1 (baseline) measured everything and was decided as before
+    assert (exp / "runs" / "iteration_001" / "input"
+            / "baseline_round.json").is_file()
+
+    gate = json.loads((exp / "runs" / "iteration_002" / "input"
+                       / "performance_gate.json").read_text(encoding="utf-8"))
+    assert gate["status"] == "INCOMPLETE", gate
+    assert dropped["id"] in gate["unmeasured"]
+
+    events = _timeline_events(cfg)
+    assert "round_incomplete" in events, events
+    # nothing was promoted, and no retake was frozen off a round without a verdict
+    assert not (exp / "pending_promotion.json").is_file()
+    assert not (exp / "pending_generalization.json").is_file()
+
+
+def test_the_incomplete_round_is_recorded_with_its_reason(tmp_path, monkeypatch):
+    cfg = _variant_cfg(tmp_path)
+    write_target_iterations(cfg.exp_dir, 2)
+    real = P.build_evaluator("dry-run")
+    dropped: dict = {}
+
+    def factory(_mode):
+        class Sabotaged:
+            def evaluate(self, cfg_, snapshot, iteration, instances=None):
+                results = real.evaluate(cfg_, snapshot, iteration,
+                                        instances=instances)
+                if (str(cfg_.answers.get("round_stage")) == "performance"
+                        and results):
+                    victim = sorted(results)[0]
+                    dropped["id"] = victim
+                    results = {k: v for k, v in results.items() if k != victim}
+                return results
+
+        return Sabotaged()
+
+    monkeypatch.setattr(P, "build_evaluator", factory)
+    run_experiment(cfg)
+
+    rows = [json.loads(line) for line in
+            (Path(cfg.state_dir) / "timeline.jsonl").read_text(
+                encoding="utf-8").splitlines() if line.strip()]
+    incomplete = [row for row in rows if row["type"] == "round_incomplete"]
+    assert incomplete, rows
+    payload = incomplete[0]["payload"]
+    assert payload["gate"] == "performance"
+    assert payload["unmeasured"] == [dropped["id"]]
+    assert payload["measured"], "the questions that did report are listed"
+    assert "no harness verdict" in payload["policy"]
diff --git a/metainfer/tasks/harness_evolve/tests/test_round_selection_evidence.py b/metainfer/tasks/harness_evolve/tests/test_round_selection_evidence.py
new file mode 100644
index 00000000..bf8a6693
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_round_selection_evidence.py
@@ -0,0 +1,141 @@
+"""The draw has to be explicable afterwards — and a round has to be fillable.
+
+Two small holes around the round's question set:
+
+* the draw used ``random`` without recording anything, so "why these four
+  operators?" could not be answered after the fact (and a pass could not be
+  re-run). The seed is now derived from (task, iteration, stage) unless the form
+  pins ``round_seed``, and it is written into the round's evidence;
+* with fewer judgeable operators than a wave, the count was silently reduced —
+  which changes what the gate means ("at least 75% win" over 3 operators is
+  "all three must win"). The run now stops with a named status instead.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import yaml
+
+from metainfer.tasks.harness_evolve.orchestrator import pipeline as P
+from metainfer.tasks.harness_evolve.orchestrator.config import (
+    load_experiment_config,
+)
+from metainfer.tasks.harness_evolve.orchestrator.pipeline import run_experiment
+from metainfer.tasks.harness_evolve.orchestrator.state import (
+    write_target_iterations,
+)
+
+
+def _cfg(tmp_path: Path, *, operators: int = 8, **answers):
+    pool = tmp_path / "pool.yaml"
+    pool.write_text(yaml.safe_dump({
+        "schema_version": 1,
+        "instances": [
+            {"id": f"op{i:02d}", "family": f"decode__fam{i}",
+             "contract": {"M": 16}, "baseline_us": 100.0 + i,
+             "best_known_us": 40.0 + 4 * i, "history": []}
+            for i in range(operators)
+        ],
+    }, sort_keys=False), encoding="utf-8")
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    req = tmp_path / "requirements.json"
+    req.write_text(json.dumps({
+        "task_id": "selection-task",
+        "task_type": "harness-evolve",
+        "execution_mode": "dry-run",
+        "evolve_mode": "dry-run",
+        "pool_source": str(pool),
+        "question_pool": str(pool),
+        "round_questions": "variant",
+        "per_round_budget": "4",
+        "max_iterations": 1,
+        **answers,
+    }), encoding="utf-8")
+    return load_experiment_config(req, tmp_path / "state", tmp_path / "ws")
+
+
+def _timeline(cfg):
+    path = Path(cfg.state_dir) / "timeline.jsonl"
+    if not path.is_file():
+        return []
+    return [json.loads(line) for line in
+            path.read_text(encoding="utf-8").splitlines() if line.strip()]
+
+
+# ------------------------------------------------------------------- the seed
+
+def test_the_same_round_draws_the_same_questions(tmp_path):
+    cfg = _cfg(tmp_path)
+
+    rng_a, seed_a = P._round_rng(cfg, 2, "performance")
+    rng_b, seed_b = P._round_rng(cfg, 2, "performance")
+
+    assert seed_a == seed_b
+    assert [rng_a.choice(range(1000)) for _ in range(5)] == \
+           [rng_b.choice(range(1000)) for _ in range(5)]
+
+
+def test_different_rounds_draw_differently(tmp_path):
+    cfg = _cfg(tmp_path)
+
+    first = P._round_rng(cfg, 2, "performance")
+    second = P._round_rng(cfg, 3, "performance")
+
+    assert first[1] != second[1]
+    assert [first[0].random() for _ in range(3)] != \
+           [second[0].random() for _ in range(3)]
+
+
+def test_a_pinned_seed_in_the_form_wins(tmp_path):
+    cfg = _cfg(tmp_path, round_seed="4242")
+
+    _rng, label = P._round_rng(cfg, 2, "performance")
+
+    assert label.startswith("round_seed=4242")
+
+
+def test_the_round_records_which_draw_it_was(tmp_path):
+    cfg = _cfg(tmp_path)
+    write_target_iterations(cfg.exp_dir, 1)
+
+    run_experiment(cfg)
+
+    selections = [row for row in _timeline(cfg)
+                  if row["type"] == "variant_round_selection"]
+    assert selections, "the round's question set must be on the record"
+    payload = selections[0]["payload"]
+    assert payload["seed"] == f"selection-task:1:baseline"
+    assert payload["count"] == 4
+    assert payload["judgeable_in_pool"] == 8
+    assert len(payload["performance_ids"]) == 4
+
+
+# --------------------------------------------------- a round needs a full wave
+
+def test_a_pool_that_cannot_fill_a_wave_stops_with_a_named_status(tmp_path):
+    cfg = _cfg(tmp_path, operators=3)          # 3 judgeable, a wave is 4
+
+    run_experiment(cfg)
+
+    run = json.loads((Path(cfg.state_dir) / "run.json").read_text())
+    assert run["final_status"] == "no_question_pool"
+    events = [row["type"] for row in _timeline(cfg)]
+    assert "variant_pool_too_small" in events, events
+    payload = [row for row in _timeline(cfg)
+               if row["type"] == "variant_pool_too_small"][0]["payload"]
+    assert payload["judgeable_count"] == 3
+    assert payload["questions_per_round"] == 4
+
+
+def test_a_pool_with_a_whole_wave_still_runs(tmp_path):
+    """The guard must not fire for a pool that can fill the round."""
+    cfg = _cfg(tmp_path, operators=4)
+
+    run_experiment(cfg)
+
+    run = json.loads((Path(cfg.state_dir) / "run.json").read_text())
+    assert run["final_status"] != "no_question_pool"
+    events = [row["type"] for row in _timeline(cfg)]
+    assert "variant_pool_too_small" not in events
diff --git a/metainfer/tasks/harness_evolve/tests/test_state_routes.py b/metainfer/tasks/harness_evolve/tests/test_state_routes.py
new file mode 100644
index 00000000..9b0a109f
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_state_routes.py
@@ -0,0 +1,97 @@
+"""HE experiment state recording + children-live aggregation."""
+
+from __future__ import annotations
+
+import json
+
+from metainfer.server.state_reader import read_run
+from metainfer.tasks.harness_evolve.orchestrator.state import (
+    append_timeline,
+    read_run as he_read_run,
+    set_run,
+)
+from metainfer.tasks.harness_evolve.server.routes import _children_live
+
+
+def test_set_run_roundtrips_and_matches_shell_reader(tmp_path):
+    state_dir = tmp_path / "state"
+    set_run(state_dir, task_id="t1", task_type="harness-evolve",
+            current_phase="evaluate", current_iteration=2)
+    # shell state_reader must see the same phase (run.json contract).
+    assert read_run(state_dir)["current_phase"] == "evaluate"
+    assert read_run(state_dir)["current_iteration"] == 2
+    assert read_run(state_dir)["task_id"] == "t1"
+    # our own reader agrees
+    assert he_read_run(state_dir)["current_phase"] == "evaluate"
+
+
+def test_set_run_finished_sets_terminal_fields(tmp_path):
+    state_dir = tmp_path / "state"
+    set_run(state_dir, current_phase="finished", finished=True,
+            final_status="success")
+    run = read_run(state_dir)
+    assert run["finished"] is True
+    assert run["final_status"] == "success"
+
+
+def test_append_timeline_rows_are_readable(tmp_path):
+    state_dir = tmp_path / "state"
+    append_timeline(state_dir, "iteration_start", {"iteration": 1})
+    append_timeline(state_dir, "iteration_decision", {"iteration": 1,
+                                                      "verdict": "BASELINE"})
+    from metainfer.tasks.harness_evolve.orchestrator.state import list_timeline
+    rows = list_timeline(state_dir)
+    assert [r["type"] for r in rows] == ["iteration_start",
+                                         "iteration_decision"]
+    assert rows[0]["payload"]["iteration"] == 1
+
+
+def test_children_live_reads_each_dkao_child_run(tmp_path):
+    """Each question is one DKAO child; its state/run.json drives the UI."""
+    exp = tmp_path / "exp"
+    child = exp / "children" / "iteration_001" / "instA" / "state"
+    child2 = exp / "children" / "iteration_001" / "instB" / "state"
+    set_run(child, task_type="dcu-kernel-auto-opt", current_phase="serial_validate")
+    set_run(child2, task_type="dcu-kernel-auto-opt", current_phase="finished",
+            finished=True, final_status="success")
+
+    live = _children_live(exp, 1)
+    assert set(live) == {"instA", "instB"}
+    assert live["instA"]["current_phase"] == "serial_validate"
+    assert live["instA"]["finished"] is False
+    assert live["instB"]["finished"] is True
+    assert live["instB"]["final_status"] == "success"
+
+
+def test_children_live_missing_iteration_returns_empty(tmp_path):
+    assert _children_live(tmp_path / "exp", 3) == {}
+
+
+def test_children_live_shows_a_retried_question_once_at_its_newest_attempt(tmp_path):
+    """One row per question, reading the attempt that ran last.
+
+    A retry lives in its own directory (``iteration_001_a2``), so the detail
+    page must not show the same question twice, and the phase it shows has to
+    come from the retry — the failed attempt is history, not the current state.
+    """
+    exp = tmp_path / "exp"
+    first = exp / "children" / "iteration_001" / "instA" / "state"
+    retry = exp / "children" / "iteration_001_a2" / "instA" / "state"
+    other = exp / "children" / "iteration_001_a2" / "instB" / "state"
+    set_run(first, task_type="dcu-kernel-auto-opt", current_phase="prepare",
+            finished=True, final_status="stopped")
+    set_run(retry, task_type="dcu-kernel-auto-opt",
+            current_phase="serial_validate")
+    set_run(other, task_type="dcu-kernel-auto-opt", current_phase="finished",
+            finished=True, final_status="success")
+
+    live = _children_live(exp, 1)
+    assert set(live) == {"instA", "instB"}, "one row per question, no duplicates"
+    assert live["instA"]["current_phase"] == "serial_validate", (
+        "the row must describe the retry, not the failed attempt")
+    assert live["instA"]["attempt_dir"] == "iteration_001_a2"
+    assert live["instA"]["child_state_dir"] == str(retry)
+    # the failed attempt's evidence is still on disk for the audit
+    assert (first / "run.json").is_file()
+    # a different iteration has no children of its own
+    assert _children_live(exp, 2) == {}
diff --git a/metainfer/tasks/harness_evolve/tests/test_supervisor.py b/metainfer/tasks/harness_evolve/tests/test_supervisor.py
new file mode 100644
index 00000000..a221ec8d
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_supervisor.py
@@ -0,0 +1,541 @@
+"""Incident supervisor: classification, policy and supervision ticks."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from metainfer.tasks.harness_evolve.orchestrator.supervisor import (
+    CATEGORY_AGENT, CATEGORY_CODE, CATEGORY_RESOURCE, CATEGORY_UNKNOWN,
+    Supervisor, classify_incident, decide_action,
+)
+
+
+# ------------------------------------------------------------- classify ----
+
+def test_classify_resource_signals():
+    assert classify_incident(log_tail="hipErrorOutOfMemory: failed")[0] == CATEGORY_RESOURCE
+    assert classify_incident(log_tail="", exit_code=137)[0] == CATEGORY_RESOURCE
+    assert classify_incident(log_tail="no space left on device")[0] == CATEGORY_RESOURCE
+    category, evidence = classify_incident(log_tail="std::bad_alloc")
+    assert category == CATEGORY_RESOURCE and evidence
+
+
+def test_classify_agent_and_code_and_unknown():
+    assert classify_incident(log_tail="TransportClosedError('runtime closed')")[0] == CATEGORY_AGENT
+    assert classify_incident(log_tail="timed out after 1800.0 seconds")[0] == CATEGORY_AGENT
+    assert classify_incident(log_tail="Traceback (most recent call last):")[0] == CATEGORY_CODE
+    assert classify_incident(log_tail="everything is fine")[0] == CATEGORY_UNKNOWN
+
+
+# --------------------------------------------------------------- policy ----
+
+def test_policy_running_and_code_and_resource():
+    assert decide_action(category=CATEGORY_UNKNOWN, same_class_streak=1,
+                         waits_done=0, restarts_done=0, running=True)["action"] == "run"
+    assert decide_action(category=CATEGORY_CODE, same_class_streak=1,
+                         waits_done=0, restarts_done=0, running=False)["action"] == "stop"
+    wait = decide_action(category=CATEGORY_RESOURCE, same_class_streak=1,
+                         waits_done=0, restarts_done=0, running=False,
+                         wait_seconds=1800)
+    assert wait["action"] == "wait" and wait["sleep_s"] == 1800
+    exhausted = decide_action(category=CATEGORY_RESOURCE, same_class_streak=1,
+                              waits_done=6, restarts_done=0, running=False)
+    assert exhausted["action"] == "stop"
+
+
+def test_policy_backoff_restart_and_breaker():
+    agent = decide_action(category=CATEGORY_AGENT, same_class_streak=1,
+                          waits_done=0, restarts_done=0, running=False)
+    assert agent["action"] == "restart" and agent["sleep_s"] > 0
+    unknown = decide_action(category=CATEGORY_UNKNOWN, same_class_streak=1,
+                            waits_done=0, restarts_done=0, running=False)
+    assert unknown["action"] == "restart"
+    breaker = decide_action(category=CATEGORY_AGENT, same_class_streak=3,
+                            waits_done=0, restarts_done=1, running=False)
+    assert breaker["action"] == "stop" and "circuit breaker" in breaker["reason"]
+    gave_up = decide_action(category=CATEGORY_UNKNOWN, same_class_streak=1,
+                            waits_done=0, restarts_done=3, running=False)
+    assert gave_up["action"] == "stop"
+
+
+# ----------------------------------------------------------- supervisor ----
+
+class FakeClock:
+    def __init__(self) -> None:
+        self.now = 1000.0
+
+    def __call__(self) -> float:
+        return self.now
+
+
+def _sup(tmp_path: Path, *, launcher, clock=None, target=4,
+         diagnose=False):
+    """A supervisor whose only inputs are the run's own files.
+
+    There is no free-VRAM getter any more: FLOW.md §0 — HE holds no device
+    verdict, so the supervisor has nothing to ask a GPU about.
+    """
+    slept: list = []
+    sup = Supervisor(
+        tmp_path / "state", tmp_path / "exp", target=target, interval_s=0,
+        clock=clock or FakeClock(), sleeper=lambda s: slept.append(s),
+        resume_launcher=launcher,
+        diagnose_agent=diagnose,
+    )
+    sup.state_dir.mkdir(parents=True, exist_ok=True)
+    sup.exp_dir.mkdir(parents=True, exist_ok=True)
+    sup.slept = slept  # type: ignore[attr-defined]
+    return sup
+
+
+def _write_run(sup: Supervisor, **fields) -> None:
+    (sup.state_dir / "run.json").write_text(
+        json.dumps({"task_id": "t", **fields}), encoding="utf-8")
+
+
+def _pid_alive(sup: Supervisor, alive: bool) -> None:
+    """Register a pid that looks like a live harness_evolve process."""
+    import os
+    import subprocess
+    import sys
+    if not alive:
+        (sup.state_dir / "resume.pid").write_text("99999999", encoding="utf-8")
+        return
+    # argv carries the marker the liveness check looks for
+    proc = subprocess.Popen(
+        [sys.executable, "-c", "import time; time.sleep(60)", "harness_evolve"],
+    )
+    sup._test_proc = proc                       # type: ignore[attr-defined]
+    # a fresh child has no state writes yet: disable the stall guard here
+    sup.stall_seconds = 10 ** 9
+    (sup.state_dir / "resume.pid").write_text(str(proc.pid), encoding="utf-8")
+
+
+def test_step_reports_alive_and_done(tmp_path):
+    sups = _sup(tmp_path, launcher=lambda: 1)
+    _write_run(sups, current_iteration=3, finished=False)
+    _pid_alive(sups, True)
+    try:
+        assert sups.step()["action"] == "run"
+        assert sups.read_state()["status"] == "running"
+
+        _write_run(sups, current_iteration=4, finished=True)
+        done = sups.step()
+        assert done["action"] == "done"
+        assert sups.read_state()["status"] == "done"
+    finally:
+        proc = getattr(sups, "_test_proc", None)
+        if proc is not None:
+            proc.kill()
+            proc.wait()
+
+
+def test_step_waits_then_restarts_on_resource_incident(tmp_path):
+    calls: list = []
+    sups = _sup(tmp_path, launcher=lambda: calls.append(1) or 4242)
+    _write_run(sups, current_iteration=3, finished=False)
+    _pid_alive(sups, False)
+    (sups.state_dir / "resume.log").write_text(
+        "hipErrorOutOfMemory: failed to allocate\n", encoding="utf-8")
+
+    out = sups.step()
+    assert out["action"] == "restart"          # waited out the OOM, then restarted
+    assert sups.slept and sups.slept[0] == 1800
+    assert calls == [1]
+    state = sups.read_state()
+    assert state["status"] == "running" and state["waits_done"] == 1
+    incidents = (sups.exp_dir / "incidents.jsonl").read_text(encoding="utf-8").splitlines()
+    kinds = [json.loads(x)["action"] for x in incidents]
+    assert "wait" in kinds and "restart_after_wait" in kinds
+
+
+def test_the_restart_after_a_wait_does_not_ask_the_cards(tmp_path):
+    """HE holds no device verdict: the wait is the whole condition.
+
+    The supervisor used to re-read free VRAM after the wait and keep waiting
+    while it looked low (``wait_again``) — that is HE managing the environment
+    again (FLOW.md §0). Whether a card is clean now is decided by the *child's*
+    own measurement gate, which waits on the child's behalf.
+    """
+    calls: list = []
+    sups = _sup(tmp_path, launcher=lambda: calls.append(1) or 4242)
+    _write_run(sups, current_iteration=3, finished=False)
+    _pid_alive(sups, False)
+    (sups.state_dir / "resume.log").write_text("bad_alloc\n", encoding="utf-8")
+
+    out = sups.step()
+
+    assert out["action"] == "restart" and calls == [1]
+    assert sups.read_state()["status"] == "running"
+    kinds = [json.loads(x)["action"] for x in
+             (sups.exp_dir / "incidents.jsonl").read_text(
+                 encoding="utf-8").splitlines()]
+    assert "wait_again" not in kinds
+
+
+def test_step_stops_on_deterministic_code_failure(tmp_path):
+    calls: list = []
+    sups = _sup(tmp_path, launcher=lambda: calls.append(1) or 7)
+    _write_run(sups, current_iteration=3, finished=False)
+    _pid_alive(sups, False)
+    (sups.state_dir / "orchestrator.log").write_text(
+        "Traceback (most recent call last):\nValueError: bad config\n",
+        encoding="utf-8")
+    out = sups.step()
+    assert out["action"] == "stop"
+    assert calls == []                         # never restarts a code error
+    assert sups.read_state()["status"] == "stopped"
+
+
+def test_circuit_breaker_after_repeated_agent_failures(tmp_path):
+    calls: list = []
+    sups = _sup(tmp_path, launcher=lambda: calls.append(1) or 5)
+    _write_run(sups, current_iteration=3, finished=False)
+    _pid_alive(sups, False)
+    (sups.state_dir / "resume.log").write_text("TransportClosedError\n",
+                                               encoding="utf-8")
+    assert sups.step()["action"] == "restart"
+    assert sups.step()["action"] == "restart"
+    out = sups.step()
+    assert out["action"] == "stop"
+    assert "circuit breaker" in out["reason"]
+    assert len(calls) == 2                     # third time: stop, no restart
+
+
+def test_stop_flag_wins(tmp_path):
+    sups = _sup(tmp_path, launcher=lambda: 1)
+    (sups.exp_dir / "supervisor.stop").write_text("stop", encoding="utf-8")
+    assert sups.step()["action"] == "stop"
+
+
+def test_project_root_contains_the_package():
+    from metainfer.tasks.harness_evolve.orchestrator.supervisor import Supervisor
+    root = Supervisor.project_root()
+    assert (root / "metainfer" / "tasks").is_dir()
+
+
+def test_pid_is_live_rejects_missing_self_and_zombies():
+    import os
+    import time
+    from metainfer.tasks.harness_evolve.orchestrator import supervisor as sup
+    assert sup.pid_is_live(None) is False
+    assert sup.pid_is_live(99999999) is False
+    assert sup.pid_is_live(os.getpid()) is False          # not a harness process
+
+    pid = os.fork()
+    if pid == 0:                                          # child: exit at once
+        os._exit(0)
+    time.sleep(0.4)                                       # now a zombie
+    try:
+        assert sup.pid_is_live(pid) is False
+    finally:
+        os.waitpid(pid, 0)
+
+
+def test_should_launch_resume_covers_start_crash_and_target():
+    from metainfer.tasks.harness_evolve.server.routes import should_launch_resume
+    assert should_launch_resume({}, None, False) is False
+    assert should_launch_resume({"finished": False}, 5, False) is True
+    assert should_launch_resume({"finished": False}, 5, True) is False
+    assert should_launch_resume({"finished": True, "current_iteration": 1},
+                                5, False) is True
+    assert should_launch_resume({"finished": True, "current_iteration": 5},
+                                5, False) is False
+    assert should_launch_resume({"finished": True, "current_iteration": 5},
+                                None, False) is False
+
+
+def test_supervisor_restarts_when_the_child_is_a_zombie(tmp_path):
+    """A defunct child must not count as 'running' (the stuck-in-analyze bug)."""
+    import os
+    import time
+    from metainfer.tasks.harness_evolve.orchestrator.supervisor import Supervisor
+    state = tmp_path / "st"
+    exp = tmp_path / "exp"
+    state.mkdir()
+    exp.mkdir()
+    pid = os.fork()
+    if pid == 0:
+        os._exit(0)
+    time.sleep(0.4)
+    try:
+        (state / "resume.pid").write_text(str(pid), encoding="utf-8")
+        (state / "run.json").write_text('{"current_iteration": 1, "finished": false}',
+                                        encoding="utf-8")
+        calls = []
+        sup = Supervisor(state, exp, target=5, interval_s=0,
+                         resume_launcher=lambda: calls.append(1) or 4242,
+                         sleeper=lambda s: None)
+        out = sup.step()
+        assert calls == [1], out            # restarted instead of idling
+        assert out["action"] in {"restart", "wait", "wait_again"}
+    finally:
+        os.waitpid(pid, 0)
+
+
+def test_progress_age_counts_child_activity(tmp_path):
+    """A long evaluate round must not look like a stall while children work."""
+    import os
+    import time
+    from metainfer.tasks.harness_evolve.orchestrator.supervisor import Supervisor
+    state = tmp_path / "st"
+    exp = tmp_path / "exp"
+    state.mkdir()
+    exp.mkdir()
+    (state / "run.json").write_text('{"current_iteration": 2, "finished": false}',
+                                    encoding="utf-8")
+    # stale outer state
+    old = time.time() - 7200
+    os.utime(state / "run.json", (old, old))
+    (state / "timeline.jsonl").write_text("{}\n", encoding="utf-8")
+    os.utime(state / "timeline.jsonl", (old, old))
+
+    child_state = exp / "children" / "iteration_002" / "shapeA" / "state"
+    child_state.mkdir(parents=True)
+    (child_state / "timeline.jsonl").write_text("{}\n", encoding="utf-8")
+
+    sup = Supervisor(state, exp, target=5, interval_s=0,
+                     resume_launcher=lambda: 1,
+                     sleeper=lambda s: None)
+    age = sup.progress_age()
+    assert age is not None and age < 60        # child write counts as progress
+
+    # and a genuinely idle run still looks stalled
+    os.utime(child_state / "timeline.jsonl", (old, old))
+    assert sup.progress_age() > 3600
+
+
+def test_janitor_sweep_recovers_unsupervised_tasks(monkeypatch, tmp_path):
+    """A task whose supervisor died must be picked up again by the sweep."""
+    from metainfer.tasks.harness_evolve.server import routes as R
+
+    state = tmp_path / "state"
+    ws = tmp_path / "ws"
+    state.mkdir()
+    ws.mkdir()
+    (state / "run.json").write_text('{"current_iteration": 2, "finished": false}',
+                                    encoding="utf-8")
+
+    class Entry:
+        id = "t-janitor"
+        type = R.PLUGIN_TYPE
+        state_dir = str(state)
+        workspace_dir = str(ws)
+
+    monkeypatch.setattr("metainfer.server.tasks.list_tasks", lambda: [Entry()])
+    monkeypatch.setattr(R, "_resume_running", lambda _s: None)
+    monkeypatch.setattr(R, "_supervisor_running", lambda _w: None)
+    monkeypatch.setattr(R, "read_target_iterations", lambda _w: 5)
+    monkeypatch.setattr(R, "_launch_supervisor",
+                        lambda *a, **k: {"pid": 111, "log": "x"})
+    monkeypatch.setattr(R, "_launch_resume",
+                        lambda *a, **k: {"pid": 222})
+
+    actions = R.janitor_sweep()
+    assert actions and actions[0]["action"] == "supervisor+resume"
+    assert actions[0]["supervisor_pid"] == 111
+    assert actions[0]["resume_pid"] == 222
+
+
+def test_janitor_sweep_leaves_a_healthy_run_alone(monkeypatch, tmp_path):
+    """A live run is nobody's business — including one started by the server.
+
+    The liveness question is ``_run_alive`` (orchestrator.pid first, then a
+    /proc scan). The old version of this test mocked ``_resume_running``, i.e.
+    it only ever described a *resume*-started task — which is why the real
+    production shape (a task the New Task form started, no ``resume.pid``) went
+    unnoticed and got a second driver.
+    """
+    from metainfer.tasks.harness_evolve.server import routes as R
+
+    state = tmp_path / "state"
+    ws = tmp_path / "ws"
+    state.mkdir()
+    ws.mkdir()
+    (state / "run.json").write_text('{"current_iteration": 2, "finished": false}',
+                                    encoding="utf-8")
+
+    class Entry:
+        id = "t-ok"
+        type = R.PLUGIN_TYPE
+        state_dir = str(state)
+        workspace_dir = str(ws)
+
+    monkeypatch.setattr("metainfer.server.tasks.list_tasks", lambda: [Entry()])
+    monkeypatch.setattr(R, "_run_alive", lambda _s: 4242)
+    monkeypatch.setattr(R, "_supervisor_running", lambda _w: 4343)
+    monkeypatch.setattr(R, "read_target_iterations", lambda _w: 5)
+    actions = R.janitor_sweep()
+    assert actions and actions[0]["action"] == "skip"
+    assert actions[0]["reason"] == "run is alive"
+
+
+def test_should_launch_resume_respects_terminal_statuses():
+    """A run stopped on purpose (unmeasured round / abort / breaker) stays stopped.
+
+    ``round_incomplete`` is how a blocked *child* surfaces now: the DKAO gate
+    that stopped it is recorded as the cause, not as HE's own device verdict.
+    """
+    from metainfer.tasks.harness_evolve.server.routes import should_launch_resume
+    for status in ("round_incomplete", "aborted_by_operator", "circuit_breaker"):
+        assert should_launch_resume(
+            {"finished": True, "current_iteration": 6, "final_status": status},
+            10, False) is False
+    # a clean finish below target is resumed as usual
+    assert should_launch_resume(
+        {"finished": True, "current_iteration": 6, "final_status": "success"},
+        10, False) is True
+
+
+def test_task_process_alive_finds_a_server_started_run(tmp_path, monkeypatch):
+    """A task launched by the server (no resume.pid) still counts as running."""
+    import os
+    import subprocess
+    import sys
+    from metainfer.tasks.harness_evolve.orchestrator import supervisor as sup
+
+    state = tmp_path / "state"
+    state.mkdir()
+    assert sup.task_process_alive(state) is None       # nothing running
+
+    # a process whose cmdline mentions this task's state dir and the CLI
+    proc = subprocess.Popen(
+        [sys.executable, "-c", "import time; time.sleep(30)",
+         "harness_evolve.orchestrator.cli", str(state)],
+    )
+    try:
+        assert sup.task_process_alive(state) == proc.pid
+    finally:
+        proc.kill()
+        proc.wait()
+
+    # a stale pid file (dead process) is not "alive"
+    (state / "resume.pid").write_text("99999999", encoding="utf-8")
+    assert sup.task_process_alive(state) is None
+
+
+# ------------------------------------------------- live children / stalls ----
+
+def _fake_proc(tmp_path, pid, cmdline, state="S"):
+    root = tmp_path / "proc"
+    entry = root / str(pid)
+    entry.mkdir(parents=True, exist_ok=True)
+    (entry / "cmdline").write_bytes(cmdline.replace(" ", "\0").encode() + b"\0")
+    (entry / "stat").write_text(f"{pid} (python3) {state} 1 1 1 0 -1 4194560\n")
+    return root
+
+
+def test_live_child_count_reads_processes_not_mtimes(tmp_path):
+    """A silent-but-alive child counts as progress; other rounds do not."""
+    from metainfer.tasks.harness_evolve.orchestrator.supervisor import Supervisor
+
+    state = tmp_path / "st"
+    exp = tmp_path / "exp"
+    state.mkdir()
+    exp.mkdir()
+    (state / "run.json").write_text('{"current_iteration": 2, "finished": false}',
+                                    encoding="utf-8")
+    sup = Supervisor(state, exp, target=5, interval_s=0,
+                     resume_launcher=lambda: 1,
+                     sleeper=lambda s: None)
+
+    root = _fake_proc(tmp_path, 4242,
+                      "python3 -m metainfer.tasks.dcu_kernel_auto_opt."
+                      "orchestrator.cli run /w/children/iteration_002/shapeA")
+    # a child of another iteration must not count for this one
+    _fake_proc(tmp_path, 4243,
+               "python3 -m metainfer.tasks.dcu_kernel_auto_opt.orchestrator.cli"
+               " run /w/children/iteration_001/shapeB")
+    # a zombie is not a live child
+    _fake_proc(tmp_path, 4244,
+               "python3 -m metainfer.tasks.dcu_kernel_auto_opt.orchestrator.cli"
+               " run /w/children/iteration_002/shapeC", state="Z")
+    # something else entirely
+    _fake_proc(tmp_path, 4245, "python3 -m http.server")
+
+    assert sup.live_child_count(root) == 1
+
+
+def test_a_silent_live_child_blocks_the_stall_guard(tmp_path):
+    """The regression that produced two resumes: 40 min of quiet child work."""
+    import os
+    import time
+    from metainfer.tasks.harness_evolve.orchestrator.supervisor import Supervisor
+
+    state = tmp_path / "st"
+    exp = tmp_path / "exp"
+    state.mkdir()
+    exp.mkdir()
+    (state / "run.json").write_text('{"current_iteration": 2, "finished": false}',
+                                    encoding="utf-8")
+    old = time.time() - 7200
+    for name in ("run.json", "timeline.jsonl"):
+        path = state / name
+        if not path.exists():
+            path.write_text("{}", encoding="utf-8")
+        os.utime(path, (old, old))
+    (exp / "children" / "iteration_002").mkdir(parents=True)
+
+    calls = []
+    sup = Supervisor(state, exp, target=5, interval_s=0, stall_seconds=60,
+                     resume_launcher=lambda: calls.append(1) or 4242,
+                     sleeper=lambda s: None)
+    assert sup.process_alive.__self__ is sup        # sanity: bound to this run
+    sup.process_alive = lambda: True                # a live (quiet) orchestrator
+    sup.live_child_count = lambda *a, **k: 1        # ...with a live child
+
+    out = sup.step()
+
+    assert out["action"] == "run", out
+    assert calls == []                              # no duplicate resume
+    actions = [json.loads(l)["action"] for l in
+               (exp / "incidents.jsonl").read_text(encoding="utf-8").splitlines()]
+    assert "stall_ignored" in actions
+
+
+def test_terminating_a_stale_run_kills_its_process_group(tmp_path):
+    import os
+    import time
+    from metainfer.tasks.harness_evolve.orchestrator.supervisor import Supervisor
+
+    state = tmp_path / "st"
+    exp = tmp_path / "exp"
+    state.mkdir()
+    exp.mkdir()
+    sup = Supervisor(state, exp, target=5, interval_s=0,
+                     resume_launcher=lambda: 1,
+                     sleeper=lambda s: time.sleep(min(s, 0.05)))
+
+    pid = os.fork()
+    if pid == 0:                                    # child: own session, then idle
+        try:
+            os.setsid()
+            time.sleep(60)
+        finally:
+            os._exit(0)
+
+    time.sleep(0.3)
+    try:
+        killed = sup.terminate_stale_run([pid])
+        assert killed == [pid]
+        deadline = time.time() + 5
+        while time.time() < deadline:
+            try:
+                wpid, _ = os.waitpid(pid, os.WNOHANG)
+            except ChildProcessError:
+                break
+            if wpid == pid:
+                break
+            time.sleep(0.1)
+        assert not Path(f"/proc/{pid}").exists() or \
+            Path(f"/proc/{pid}/stat").read_text().split(") ", 1)[1].startswith("Z")
+    finally:
+        try:
+            os.kill(pid, 9)
+        except OSError:
+            pass
+        try:
+            os.waitpid(pid, os.WNOHANG)
+        except ChildProcessError:
+            pass
diff --git a/metainfer/tasks/harness_evolve/tests/test_variant_rounds.py b/metainfer/tasks/harness_evolve/tests/test_variant_rounds.py
new file mode 100644
index 00000000..ec2924d6
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tests/test_variant_rounds.py
@@ -0,0 +1,414 @@
+"""Variant-driven rounds: random operators, a fixed generalization paper, and
+the two counter rules (>=75% / >=50% wins with an 80%-of-variant floor)."""
+
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+
+import pytest
+
+from ..orchestrator.decision_engine import gate_verdict
+from ..orchestrator.rounds import (
+    evaluate_gate, record_paper, regime_of, select_round_questions,
+    sample_operators,
+)
+from ..orchestrator.variant import (
+    load_variant_table, record_variants, reference_results, resolve_variant,
+)
+
+
+def _pool(n: int = 12) -> dict:
+    """Alternating decode (M=16) / prefill (M=4096) operators with variants."""
+    out = {}
+    for i in range(n):
+        m = 16 if i % 2 == 0 else 4096
+        regime = "decode" if m == 16 else "prefill"
+        out[f"op{i:02d}"] = {
+            "baseline_us": 100.0 * (i + 1),
+            "best_known_us": 40.0 * (i + 1),      # an existing variant
+            "family": f"{regime}__fam{i}",
+            "contract": {"M": m, "N": 64, "K": 32},
+        }
+    return out
+
+
+def _table() -> dict:
+    return {"operators": {}, "groups": {}, "harness_version": None}
+
+
+# ------------------------------------------------------------------ selection
+
+def test_first_round_defines_the_generalization_paper():
+    pool, table = _pool(), _table()
+    picked = select_round_questions(iteration=1, table=table, instances=pool,
+                                    state={}, rng=random.Random(7))
+    assert len(picked["performance_ids"]) == 4
+    # round 1's operators are the paper, and they are what this round measures
+    assert picked["generalization_ids"] == picked["performance_ids"]
+    assert picked["defines_paper"] is True
+
+
+def test_later_rounds_draw_new_operators_and_reuse_the_paper():
+    pool, table = _pool(), _table()
+    first = select_round_questions(iteration=1, table=table, instances=pool,
+                                   state={}, rng=random.Random(1))
+    groups = record_paper(table, iteration=1,
+                          generalization_ids=first["generalization_ids"],
+                          performance_ids=first["performance_ids"])
+    second = select_round_questions(iteration=2, table=table, instances=pool,
+                                    state={"groups": groups},
+                                    rng=random.Random(2))
+    assert set(second["performance_ids"]).isdisjoint(first["generalization_ids"])
+    assert second["generalization_ids"] == first["generalization_ids"]
+    assert second["defines_paper"] is False
+
+
+def test_operators_without_any_reference_are_not_sampled():
+    pool = {"known": {"baseline_us": 10.0}, "unknown": {}}
+    assert sample_operators(pool, count=2, table=_table()) == ["known"]
+
+
+# ---------------------------------------------------------------- the gates
+
+def _reference(ids, table, pool):
+    return reference_results(table, ids, instances=pool)
+
+
+def test_performance_gate_needs_three_of_four_wins():
+    pool, table = _pool(), _table()
+    ids = ["op00", "op01", "op02", "op03"]
+    ref = _reference(ids, table, pool)          # 40/80/120/160 us variants
+    candidate = {
+        "op00": {"median_us": 30.0, "correctness_ok": True, "status": "success"},
+        "op01": {"median_us": 70.0, "correctness_ok": True, "status": "success"},
+        "op02": {"median_us": 110.0, "correctness_ok": True, "status": "success"},
+        "op03": {"median_us": 170.0, "correctness_ok": True, "status": "success"},
+    }
+    gate = evaluate_gate(kind="performance", operator_ids=ids, results=candidate,
+                         table=table, instances=pool)
+    assert gate["status"] == "PASS" and len(gate["wins"]) == 3
+    assert gate["wins_needed"] == 3 and gate["required_win_ratio"] == 0.75
+
+    # only two wins -> fail
+    candidate["op02"] = {"median_us": 130.0, "correctness_ok": True,
+                         "status": "success"}
+    gate = evaluate_gate(kind="performance", operator_ids=ids, results=candidate,
+                         table=table, instances=pool)
+    assert gate["status"] == "FAIL" and "2/4 operators" in gate["reasons"][-1]
+
+
+def test_a_winning_round_still_fails_if_one_operator_collapses():
+    pool, table = _pool(), _table()
+    ids = ["op00", "op01", "op02", "op03"]
+    candidate = {
+        "op00": {"median_us": 30.0, "correctness_ok": True, "status": "success"},
+        "op01": {"median_us": 70.0, "correctness_ok": True, "status": "success"},
+        "op02": {"median_us": 110.0, "correctness_ok": True, "status": "success"},
+        # variant is 160us; 140us is only 87% ... 220us is 72% (below the floor)
+        "op03": {"median_us": 220.0, "correctness_ok": True, "status": "success"},
+    }
+    gate = evaluate_gate(kind="performance", operator_ids=ids, results=candidate,
+                         table=table, instances=pool)
+    assert gate["status"] == "FAIL"
+    assert gate["below_floor"] == ["op03"]
+
+
+def test_generalization_gate_needs_only_two_of_four():
+    pool, table = _pool(), _table()
+    ids = ["op00", "op01", "op02", "op03"]
+    candidate = {
+        "op00": {"median_us": 30.0, "correctness_ok": True, "status": "success"},
+        "op01": {"median_us": 70.0, "correctness_ok": True, "status": "success"},
+        "op02": {"median_us": 140.0, "correctness_ok": True, "status": "success"},
+        "op03": {"median_us": 180.0, "correctness_ok": True, "status": "success"},
+    }
+    gate = evaluate_gate(kind="generalization", operator_ids=ids,
+                         results=candidate, table=table, instances=pool)
+    assert gate["status"] == "PASS" and gate["wins_needed"] == 2
+    assert gate["required_win_ratio"] == 0.5
+
+    candidate["op01"] = {"median_us": 90.0, "correctness_ok": True,
+                         "status": "success"}
+    gate = evaluate_gate(kind="generalization", operator_ids=ids,
+                         results=candidate, table=table, instances=pool)
+    assert gate["status"] == "FAIL"
+
+
+def test_a_retake_that_reproduces_the_same_kernels_still_holds_the_line():
+    """Ties are not losses in the generalization gate.
+
+    The retake re-measures the paper on the same harness, and DKAO warm-starts
+    from the kernel the baseline round put in the pool — so re-finding that exact
+    kernel (delta 0.0%) is the expected outcome, not a failure. Scoring it as a
+    loss would demand that a *search policy* out-optimize its own best result
+    every single round.
+    """
+    pool, table = _pool(), _table()
+    ids = ["op00", "op01", "op02", "op03"]
+    variants = {"op00": 40.0, "op01": 80.0, "op02": 120.0, "op03": 160.0}
+    candidate = {iid: {"median_us": variants[iid], "correctness_ok": True,
+                       "status": "success"} for iid in ids}
+
+    gate = evaluate_gate(kind="generalization", operator_ids=ids,
+                         results=candidate, table=table, instances=pool)
+    assert gate["tie_policy"] == "neutral"
+    assert gate["ties"] == ids and gate["wins"] == [] and gate["losses"] == []
+    assert gate["status"] == "PASS", gate["reasons"]
+
+    # ... but the floor still bites: one operator genuinely slower than its
+    # variant fails the retake even with ties scored as holding the line
+    candidate["op03"] = {"median_us": 400.0, "correctness_ok": True,
+                         "status": "success"}
+    gate = evaluate_gate(kind="generalization", operator_ids=ids,
+                         results=candidate, table=table, instances=pool)
+    assert gate["status"] == "FAIL"
+    assert gate["losses"] == ["op03"] and gate["below_floor"] == ["op03"]
+
+
+def test_the_performance_gate_still_counts_a_tie_as_a_loss():
+    """Fresh operators must genuinely win: a tie proves nothing there."""
+    pool, table = _pool(), _table()
+    ids = ["op00", "op01", "op02", "op03"]
+    variants = {"op00": 40.0, "op01": 80.0, "op02": 120.0, "op03": 160.0}
+    candidate = {iid: {"median_us": variants[iid], "correctness_ok": True,
+                       "status": "success"} for iid in ids}
+    gate = evaluate_gate(kind="performance", operator_ids=ids,
+                         results=candidate, table=table, instances=pool)
+    assert gate["tie_policy"] == "loss"
+    assert gate["status"] == "FAIL"
+    assert "no operator beat its variant" in " ".join(gate["reasons"])
+
+
+def test_unmeasured_operator_is_never_counted_as_a_loss():
+    """Three wins and one missing number is *not* a verdict about the harness.
+
+    It used to be FAIL, which both blamed the harness for a number the machine
+    never produced and (on the retake) spent the candidate on it. The round now
+    carries ``INCOMPLETE``: no verdict, re-measured later.
+    """
+    pool, table = _pool(), _table()
+    ids = ["op00", "op01", "op02", "op03"]
+    candidate = {
+        "op00": {"median_us": 30.0, "correctness_ok": True, "status": "success"},
+        "op01": {"median_us": 70.0, "correctness_ok": True, "status": "success"},
+        "op02": {"median_us": 110.0, "correctness_ok": True, "status": "success"},
+        # op03 never reported: the round is not evidence
+    }
+    gate = evaluate_gate(kind="performance", operator_ids=ids, results=candidate,
+                         table=table, instances=pool)
+    assert gate["status"] == "INCOMPLETE"
+    assert "op03" in gate["unmeasured"]
+    assert gate["wins_needed"] == 3              # wins already enough to pass...
+    assert gate["wins_needed_over_requested"] == 3   # ... over the drawn round
+    assert any("op03" in reason for reason in gate["reasons"])
+
+
+def test_a_gate_that_cannot_pass_without_the_missing_question_is_a_fail():
+    """Decisive losses are still decided now: no relaunch needed to know."""
+    pool, table = _pool(), _table()               # variants 40/80/120/160 us
+    ids = ["op00", "op01", "op02", "op03"]
+    candidate = {
+        # ~1% slower than each variant: inside the noise band, i.e. a tie → a
+        # loss for the performance gate, but nowhere near the 80% floor.
+        "op00": {"median_us": 40.4, "correctness_ok": True, "status": "success"},
+        "op01": {"median_us": 80.8, "correctness_ok": True, "status": "success"},
+        "op02": {"median_us": 121.2, "correctness_ok": True, "status": "success"},
+        # op03 never reported
+    }
+    gate = evaluate_gate(kind="performance", operator_ids=ids, results=candidate,
+                         table=table, instances=pool)
+    assert gate["status"] == "FAIL", gate
+    assert gate["below_floor"] == []              # the floor is not what decided it
+    assert any("cannot pass" in reason for reason in gate["reasons"])
+
+
+def test_a_question_below_the_floor_is_a_fail_even_with_others_missing():
+    pool, table = _pool(), _table()
+    ids = ["op00", "op01", "op02", "op03"]
+    candidate = {
+        "op00": {"median_us": 30.0, "correctness_ok": True, "status": "success"},
+        "op01": {"median_us": 200.0, "correctness_ok": True, "status": "success"},
+        # way below the 80% floor of its 80us variant
+        "op02": {"median_us": 110.0, "correctness_ok": True, "status": "success"},
+        # op03 never reported
+    }
+    gate = evaluate_gate(kind="performance", operator_ids=ids, results=candidate,
+                         table=table, instances=pool)
+    assert gate["status"] == "FAIL"
+    assert "op01" in gate["below_floor"]
+
+
+# ------------------------------------------------------- variant table rules
+
+def test_variant_table_moves_only_when_a_harness_is_promoted(tmp_path):
+    pool = _pool()
+    table = load_variant_table(tmp_path, pool_path=None)
+    table["operators"] = {}
+    # a rejected round measured something faster, but nothing is promoted
+    before = table["operators"].get("op00")
+    result = record_variants(tmp_path, table, {}, version="h-2", iteration=1)
+    assert result["updated"] == [] and table["operators"] == {}
+
+    # the approved candidate's numbers become the new variants
+    outcome = record_variants(
+        tmp_path, table,
+        {"op00": {"median_us": 11.0, "p90_us": 11.2, "correctness_ok": True,
+                  "repo_path": "/repos/op00"}},
+        version="h-2", iteration=1,
+        groups={"generalization_ids": ["op00"]})
+    assert outcome["updated"] == ["op00"]
+    resolved = resolve_variant(table, "op00", instances=pool)
+    assert resolved["median_us"] == 11.0 and resolved["source"] == "variant"
+    assert resolved["harness_version"] == "h-2"
+    history = [json.loads(l) for l in
+               (tmp_path / "variant_history.jsonl").read_text().splitlines()]
+    assert history[-1]["harness_version"] == "h-2"
+    assert json.loads((tmp_path / "variant_table.json").read_text()
+                      )["harness_version"] == "h-2"
+
+
+def test_a_slower_measurement_never_regresses_a_variant(tmp_path):
+    pool = _pool()
+    table = {"operators": {}, "groups": {}}
+    record_variants(tmp_path, table,
+                    {"op00": {"median_us": 30.0, "correctness_ok": True}},
+                    version="h-2", iteration=1)
+    record_variants(tmp_path, table,
+                    {"op00": {"median_us": 45.0, "correctness_ok": True}},
+                    version="h-3", iteration=2)
+    assert resolve_variant(table, "op00", instances=pool)["median_us"] == 30.0
+
+
+def test_reference_falls_back_to_pool_best_known_then_baseline(tmp_path):
+    pool = {"best": {"baseline_us": 100.0, "best_known_us": 40.0},
+            "baseonly": {"baseline_us": 55.0},
+            "nothing": {}}
+    table = {"operators": {}, "groups": {}}
+    assert resolve_variant(table, "best", instances=pool)["source"] == "best_known"
+    assert resolve_variant(table, "baseonly", instances=pool)["source"] == "baseline"
+    assert resolve_variant(table, "nothing", instances=pool) is None
+
+
+def test_question_count_is_rounded_to_a_multiple_of_the_gpu_count():
+    """A wave must be able to fill all 4 HCUs: counts snap to multiples of 4."""
+    from ..orchestrator.rounds import OPERATORS_PER_GATE, normalized_per_gate
+
+    assert OPERATORS_PER_GATE == 4
+    assert normalized_per_gate(None) == 4          # default
+    assert normalized_per_gate("") == 4
+    assert normalized_per_gate(0) == 4
+    assert normalized_per_gate("bad") == 4
+    assert normalized_per_gate(4) == 4
+    assert normalized_per_gate(5) == 8             # rounds up, never down
+    assert normalized_per_gate(7) == 8
+    assert normalized_per_gate(8) == 8
+    assert normalized_per_gate(12) == 12
+
+
+def test_eight_question_round_uses_eight_operators_in_both_gates():
+    pool, table = _pool(20), _table()
+    picked = select_round_questions(iteration=1, table=table, instances=pool,
+                                    state={}, count=8, rng=random.Random(3))
+    assert len(picked["performance_ids"]) == 8
+    assert picked["generalization_ids"] == picked["performance_ids"]
+
+    groups = record_paper(table, iteration=1,
+                          generalization_ids=picked["generalization_ids"],
+                          performance_ids=picked["performance_ids"])
+    second = select_round_questions(iteration=2, table=table, instances=pool,
+                                    state={"groups": groups}, count=8,
+                                    rng=random.Random(4))
+    assert len(second["performance_ids"]) == 8
+    assert set(second["performance_ids"]).isdisjoint(picked["performance_ids"])
+
+
+def test_eight_question_gates_need_six_and_four_wins():
+    pool, table = _pool(20), _table()
+    ids = [f"op{i:02d}" for i in range(8)]
+    base = {iid: 100.0 for iid in ids}
+    reference = {iid: {"median_us": base[iid], "status": "reference",
+                       "correctness_ok": True} for iid in ids}
+    # 6 wins, 2 small losses -> performance gate passes (75% of 8 = 6)
+    candidate = {iid: {"median_us": 90.0, "correctness_ok": True,
+                       "status": "success"} for iid in ids}
+    for iid in ids[6:]:
+        candidate[iid]["median_us"] = 105.0
+    perf = gate_verdict(reference, candidate, kind="performance")
+    assert perf["wins_needed"] == 6 and perf["status"] == "PASS"
+    # 4 wins -> generalization gate passes (50% of 8 = 4), performance does not
+    for iid in ids[4:]:
+        candidate[iid]["median_us"] = 105.0
+    assert gate_verdict(reference, candidate, kind="performance")["status"] == "FAIL"
+    gen = gate_verdict(reference, candidate, kind="generalization")
+    assert gen["wins_needed"] == 4 and gen["status"] == "PASS"
+
+
+def test_round_size_follows_the_request_and_never_shrinks_to_the_cards():
+    """A round is as big as the operator asked for, whatever the cards are doing.
+
+    ``harness_evolve`` manages no GPU occupancy: it hands every question to a
+    DKAO child and the child's own admission gate waits for a clean device. So
+    there is nothing to shrink the round *to* — no "free devices" number, no
+    ``fit_count_to_devices`` — and a busy machine can no longer silently turn a
+    four-question round into a two-question one.
+    """
+    from ..orchestrator.rounds import (
+        GPU_COUNT, device_for_index, normalized_per_gate,
+    )
+
+    # the requested count is honoured, only rounded up to a whole wave of 4
+    assert normalized_per_gate(4) == 4
+    assert normalized_per_gate(8) == 8
+    assert normalized_per_gate(5) == 8
+
+    # one question per device, repeating cyclically: 8 questions = 2 waves
+    assert [device_for_index(i) for i in range(8)] == [0, 1, 2, 3, 0, 1, 2, 3]
+    assert GPU_COUNT == 4
+
+
+def test_no_gpu_lease_or_occupancy_logic_is_left_in_the_round_path():
+    """The HE round path must not import the broker or probe the cards."""
+    from ..orchestrator import rounds as rounds_mod
+    from ..orchestrator.pipeline import _truthy  # still used for other flags
+
+    assert _truthy(None, default=False) is False      # opt-in, not opt-out
+    assert not hasattr(rounds_mod, "devices_free_now")
+    assert not hasattr(rounds_mod, "fit_count_to_devices")
+
+    import ast
+    import inspect
+    import io
+    import tokenize as _tok
+
+    from ..orchestrator.adapters import eval as eval_mod
+
+    # Judge the code, not the comments/docstrings that document what was removed.
+    source = inspect.getsource(eval_mod)
+    pieces = []
+    for token in _tok.generate_tokens(io.StringIO(source).readline):
+        if token.type in (_tok.COMMENT, _tok.STRING):
+            continue
+        pieces.append(token.string)
+    code = " ".join(pieces)
+    for forbidden in ("gpu_broker", "PRIORITY_HARNESS", "gpu_gate",
+                      "measurement_suspect", "PREFERRED_DEVICES"):
+        assert forbidden not in code, forbidden
+    assert "_evaluate_scheduled" not in code
+    assert "_start_contention_watch" not in code
+    # the gate verdict lives in the child's own measurement_gate.jsonl
+    tree = ast.parse(source)
+    names = {n.id for n in ast.walk(tree) if isinstance(n, ast.Name)}
+    names |= {n.attr for n in ast.walk(tree) if isinstance(n, ast.Attribute)}
+    assert "measurement_suspect" not in names
+
+
+def test_a_three_device_round_draws_three_operators_with_both_regimes():
+    pool, table = _pool(20), _table()
+    picked = select_round_questions(iteration=1, table=table, instances=pool,
+                                    state={}, count=3, rng=random.Random(11))
+    assert len(picked["performance_ids"]) == 3
+    assert picked["generalization_ids"] == picked["performance_ids"]
+    regimes = {regime_of(pool[i]) for i in picked["performance_ids"]}
+    assert regimes == {"decode", "prefill"}
diff --git a/metainfer/tasks/harness_evolve/tools/__init__.py b/metainfer/tasks/harness_evolve/tools/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/metainfer/tasks/harness_evolve/tools/recover_evolve.py b/metainfer/tasks/harness_evolve/tools/recover_evolve.py
new file mode 100644
index 00000000..54e15cd0
--- /dev/null
+++ b/metainfer/tasks/harness_evolve/tools/recover_evolve.py
@@ -0,0 +1,130 @@
+"""Recover a crashed harness_evolve experiment.
+
+When the outer loop dies after an iteration's evaluation (e.g. the evolve
+turn timed out before the robustness fix), the value of that iteration is
+already on disk: benchmark results, decision and best_ever. This tool runs
+the MISSING evolve step for such an iteration, writes the evolve artifacts,
+commits the workspace generation and finishes the report - i.e. it performs
+the tail of ``pipeline.run_experiment`` without re-running any DKAO work.
+
+Usage::
+
+    PYTHONPATH=/root/zth_agent/MetaInfer python3 -m \
+      metainfer.tasks.harness_evolve.tools.recover_evolve \
+       --state-dir DIR --workspace-dir DIR [--iteration N]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from ..orchestrator.attribution import save_json
+from ..orchestrator.config import load_experiment_config
+from ..orchestrator.evolve import DEFAULT_EVOLVE_MODEL, build_evolver
+from ..orchestrator import pipeline as P
+from ..orchestrator.round_plan import RoundPlan
+
+
+def _load_json(path: Path) -> Optional[Dict[str, Any]]:
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    return value if isinstance(value, dict) else None
+
+
+def _candidate_iteration(exp: Path, wanted: Optional[int]) -> Path:
+    runs = exp / "runs"
+    dirs = sorted(d for d in runs.glob("iteration_*") if d.is_dir())
+    if wanted is not None:
+        target = runs / f"iteration_{wanted:03d}"
+        if not target.is_dir():
+            raise SystemExit(f"iteration {wanted} not found under {runs}")
+        return target
+    for d in reversed(dirs):
+        has_decision = (d / "input" / "decision.json").is_file()
+        has_evolve = (d / "evolve" / "change_manifest.json").is_file()
+        if has_decision and not has_evolve:
+            return d
+    raise SystemExit("no iteration with a completed evaluation and missing "
+                     "evolve artifacts; nothing to recover")
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    ap = argparse.ArgumentParser(prog="recover_evolve")
+    ap.add_argument("requirements", type=Path)
+    ap.add_argument("--state-dir", type=Path, required=True)
+    ap.add_argument("--workspace-dir", type=Path, required=True)
+    ap.add_argument("--iteration", type=int, default=None)
+    args = ap.parse_args(argv)
+
+    cfg = load_experiment_config(args.requirements, args.state_dir,
+                                 args.workspace_dir)
+    exp = cfg.exp_dir
+    it_dir = _candidate_iteration(exp, args.iteration)
+    iteration = int(it_dir.name.split("_")[-1])
+    print(f"[recover] iteration {iteration} at {it_dir}")
+
+    input_dir = it_dir / "input"
+    results = ((_load_json(input_dir / "benchmark" / "results.json") or {})
+               .get("results") or {})
+    decision = _load_json(input_dir / "decision.json") or {}
+    diff = _load_json(input_dir / "diff.json") or {}
+    scored = _load_json(input_dir / "benchmark" / "scored_ids.json") or []
+    selected = list(scored) or P._load_ids(exp / P.LAST_SCORED_FILE)
+
+    pool = P._load_pool(cfg)
+    pool_ids = pool.all_ids() if pool is not None else list(selected)
+
+    evolver = build_evolver(
+        str(cfg.answers.get("evolve_mode") or "dry-run"),
+        model=str(cfg.answers.get("evolve_model")
+                  or DEFAULT_EVOLVE_MODEL),
+        timeout_seconds=float(cfg.answers.get("evolve_timeout_seconds")
+                              or 3600),
+    )
+    P._record_state(cfg, P.EVOLVE, iteration, event="evolve_start_recover")
+
+    live_workspace = exp / "workspace"
+    evolve_out = evolver.apply(live_workspace, iteration, {
+        "overview": str(input_dir / "analysis" / "overview.md"),
+        "results": results,
+        "diff": diff,
+        "decision": decision,
+        "selected": selected,
+        "pool_ids": pool_ids,
+    })
+
+    evolve_dir = it_dir / "evolve"
+    evolve_dir.mkdir(parents=True, exist_ok=True)
+    revision = P._commit_workspace(live_workspace, iteration)
+    evolve_out["manifest"]["workspace_revision"] = revision
+    save_json(evolve_dir / "change_manifest.json", evolve_out["manifest"])
+    (evolve_dir / "evolve_summary.md").write_text(
+        evolve_out["summary"] + "\n", encoding="utf-8")
+    if pool is not None:
+        agent_plan = evolve_out.get("round_plan")
+        if isinstance(agent_plan, dict):
+            plan_obj = RoundPlan.from_dict(agent_plan)
+            pool.assert_in_pool(plan_obj.selected)
+        else:
+            plan_obj = RoundPlan(
+                iteration=iteration, strategy="balanced",
+                rationale="recover: keep same scored set",
+                selected=list(selected),
+            )
+        save_json(evolve_dir / "round_plan.json", plan_obj.to_dict())
+    P._record_state(cfg, P.REPORT, event="experiment_report_recover")
+    report = P._final_report(cfg)
+    P._record_state(cfg, P.FINISHED, event="experiment_finished_recover")
+    print(f"[recover] evolve manifest -> {evolve_dir / 'change_manifest.json'}")
+    print(f"[recover] report -> {report}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/metainfer/tasks/opt_kernel/orchestrator/orchestrator.py b/metainfer/tasks/opt_kernel/orchestrator/orchestrator.py
index 9307aeed..4e111b61 100644
--- a/metainfer/tasks/opt_kernel/orchestrator/orchestrator.py
+++ b/metainfer/tasks/opt_kernel/orchestrator/orchestrator.py
@@ -103,6 +103,7 @@ def run_with_requirements(
         effort=effort,
         extra_add_dirs=[notebooks_dir, repo_root, logs_root],
         snapshot_file=paths["agents_file"],
+        state_dir=state_dir,
     )
     orch = Orchestrator(req=req, store=store, cfg=cfg, manager=manager)