From 4ddea3fffee87c2c09d52e7de3c50051e7a5f17a Mon Sep 17 00:00:00 2001 From: Bsgg1 Date: Mon, 20 Jul 2026 14:20:52 +0000 Subject: [PATCH 1/7] feat(orchestrator): support configurable coding agent backends --- README.md | 39 ++++- README_CN.md | 39 ++++- metainfer/orchestrator/_bootstrap.py | 14 ++ metainfer/orchestrator/subagent_manager.py | 74 ++++++++++ .../tests/test_subagent_manager.py | 139 ++++++++++++++++++ .../orchestrator/tests/test_token_budget.py | 18 +++ metainfer/orchestrator/token_budget.py | 7 +- 7 files changed, 324 insertions(+), 6 deletions(-) create mode 100644 metainfer/orchestrator/tests/test_subagent_manager.py diff --git a/README.md b/README.md index 68366e8d..ffd25cf7 100644 --- a/README.md +++ b/README.md @@ -79,13 +79,48 @@ ## 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 +``` + ### 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..ff2451eb 100644 --- a/README_CN.md +++ b/README_CN.md @@ -79,13 +79,48 @@ ## 快速开始 -### 第一步,安装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 +``` + ### 第二步,安装MetaInfer ```bash git clone https://github.com/MetaInfer/MetaInfer.git diff --git a/metainfer/orchestrator/_bootstrap.py b/metainfer/orchestrator/_bootstrap.py index 22b1d68e..3c09b5a9 100644 --- a/metainfer/orchestrator/_bootstrap.py +++ b/metainfer/orchestrator/_bootstrap.py @@ -174,6 +174,8 @@ def restore() -> None: def make_subagent_manager( *, claude_bin: str, + codex_bin: Optional[str] = None, + agent_backend: Optional[str] = None, model: Optional[str], permission_mode: str, effort: str, @@ -190,8 +192,20 @@ def make_subagent_manager( ``budget`` (optional) wires the per-task :class:`TokenBudget` so every agent launch is gated + every result's cost is recorded. """ + 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" + ) return SubAgentManager( claude_bin=claude_bin, + codex_bin=resolved_codex_bin, + agent_backend=resolved_agent_backend, default_model=model, permission_mode=permission_mode, effort=effort, diff --git a/metainfer/orchestrator/subagent_manager.py b/metainfer/orchestrator/subagent_manager.py index 3c79c722..09281237 100644 --- a/metainfer/orchestrator/subagent_manager.py +++ b/metainfer/orchestrator/subagent_manager.py @@ -33,6 +33,16 @@ # 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"] + + +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" + raise ValueError(f"invalid agent backend {value!r}; expected 'claude' or 'codex'") # --------------------------------------------------------------------------- # @@ -149,6 +159,8 @@ class SubAgentManager: def __init__( self, claude_bin: str = "ccb", + codex_bin: str = "codex", + agent_backend: AgentBackend = "claude", default_model: Optional[str] = None, max_concurrent: int = 4, permission_mode: str = "bypassPermissions", @@ -158,7 +170,9 @@ def __init__( budget: Any = None, budget_source: str = "orchestrator", ) -> None: + self.agent_backend = _normalize_agent_backend(agent_backend) self.claude_bin = claude_bin + self.codex_bin = codex_bin self.default_model = default_model self.max_concurrent = max_concurrent # Per-task token / cost budget (TokenBudget instance or None to @@ -616,6 +630,11 @@ def dump_snapshot(self) -> None: # ------------------------------------------------------------------ # def _build_command(self, spec: AgentSpec) -> List[str]: + if self.agent_backend == "codex": + return self._build_codex_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 +677,42 @@ 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_env(self, spec: AgentSpec) -> Dict[str, str]: env = dict(os.environ) env.update(spec.env_overrides) @@ -673,6 +728,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 == "codex": + 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 @@ -704,6 +761,13 @@ def _materialize_result( continue final_text = "" 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): @@ -724,6 +788,8 @@ def _materialize_result( session_id = None 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": @@ -737,6 +803,14 @@ def _materialize_result( 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 diff --git a/metainfer/orchestrator/tests/test_subagent_manager.py b/metainfer/orchestrator/tests/test_subagent_manager.py new file mode 100644 index 00000000..009ac95b --- /dev/null +++ b/metainfer/orchestrator/tests/test_subagent_manager.py @@ -0,0 +1,139 @@ +"""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() 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..f0b208f0 100644 --- a/metainfer/orchestrator/token_budget.py +++ b/metainfer/orchestrator/token_budget.py @@ -567,6 +567,9 @@ 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) return UsageRecord( agent=str(agent), source=str(source), @@ -574,10 +577,10 @@ def usage_from_result_event( 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_read_input_tokens=int(cache_read 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"), + session_id=event.get("session_id") or event.get("thread_id"), ) From 7e1761f9366ef768982099e87faf854a28c3d72a Mon Sep 17 00:00:00 2001 From: flyingdown Date: Wed, 9 Sep 2026 16:40:56 +0800 Subject: [PATCH 2/7] feat(orchestrator): support pi as a coding agent backend Adds `pi` (https://pi.dev) as a third agent backend alongside claude and codex. The orchestrator invokes `pi --mode json -p` and reuses the operator's local ~/.pi provider settings; MetaInfer stores no API keys. orchestrator/subagent_manager.py - AgentBackend Literal gains "pi"; _normalize_agent_backend accepts pi / pi-coding-agent / earendil. - pi_bin constructor param; _build_pi_command maps the orchestrator effort knob onto pi --thinking (low/medium/high/max) and reuses the operator's configured provider/model. - pi session isolation: when state_dir is provided, --session-dir is scoped to /pi-sessions so orchestrator sessions are invisible to a user running `pi --continue` from their shell (pi indexes sessions by cwd; without this a user could attach to a production orchestrator session mid-flight and corrupt the context that later --session --continue turns rely on). The session id is captured from pi's leading events for resume. orchestrator/_bootstrap.py - make_subagent_manager gains pi_bin and state_dir; both resolve from args or env (METAINFER_PI_BIN, METAINFER_AGENT_BACKEND) with claude remaining the default. state_dir is forwarded only to the pi backend. orchestrator/token_budget.py - usage_from_result_event tolerates pi's camelCase usage names (input / output / cacheRead / cacheWrite / cost.total) so a raw pi event is accounted correctly. tasks - Every task orchestrator that builds a SubAgentManager now passes state_dir=state_dir so the pi backend can scope its session store per task (calc_value, find_low_hanging_kernel, gen_cpp_infer_framework, gen_infer_framework, gen_infer_framework_cpp, opt_kernel, dcu_kernel_auto_opt, evolve_kernel). opt_GEMM_kernel and port_model already wired it. Other backends ignore the kwarg. docs - README (en/zh): document the pi backend, install, env vars and the effort->thinking mapping. tests - test_subagent_manager covers pi command building, thinking-level mapping, session-dir scoping and resume-id capture. - test_token_budget covers pi camelCase usage parsing. --- README.md | 29 ++ README_CN.md | 28 ++ metainfer/orchestrator/_bootstrap.py | 14 + metainfer/orchestrator/subagent_manager.py | 268 ++++++++++++++---- .../tests/test_subagent_manager.py | 213 ++++++++++++++ metainfer/orchestrator/token_budget.py | 20 +- .../calc_value/orchestrator/orchestrator.py | 1 + .../orchestrator/orchestrator.py | 1 + .../orchestrator/orchestrator.py | 1 + .../orchestrator/orchestrator.py | 1 + .../orchestrator/orchestrator.py | 1 + .../orchestrator/orchestrator.py | 1 + .../orchestrator/orchestrator.py | 1 + .../opt_kernel/orchestrator/orchestrator.py | 1 + 14 files changed, 525 insertions(+), 55 deletions(-) diff --git a/README.md b/README.md index ffd25cf7..3b94c399 100644 --- a/README.md +++ b/README.md @@ -121,6 +121,35 @@ If the Codex binary is not named `codex` or is not on `PATH`, set: 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 ff2451eb..bd05ef5d 100644 --- a/README_CN.md +++ b/README_CN.md @@ -121,6 +121,34 @@ METAINFER_AGENT_BACKEND=codex ./serve.py 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 3c09b5a9..c8679805 100644 --- a/metainfer/orchestrator/_bootstrap.py +++ b/metainfer/orchestrator/_bootstrap.py @@ -175,6 +175,7 @@ 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, @@ -183,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 @@ -191,6 +193,11 @@ 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 @@ -202,9 +209,15 @@ def make_subagent_manager( 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, @@ -213,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 09281237..5d8c92a1 100644 --- a/metainfer/orchestrator/subagent_manager.py +++ b/metainfer/orchestrator/subagent_manager.py @@ -33,7 +33,19 @@ # 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"] +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: @@ -42,7 +54,11 @@ def _normalize_agent_backend(value: str) -> AgentBackend: return "claude" if v in {"codex", "openai-codex"}: return "codex" - raise ValueError(f"invalid agent backend {value!r}; expected 'claude' or 'codex'") + if v in {"pi", "pi-coding-agent", "earendil"}: + return "pi" + raise ValueError( + f"invalid agent backend {value!r}; expected 'claude', 'codex', or 'pi'" + ) # --------------------------------------------------------------------------- # @@ -160,6 +176,7 @@ 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, @@ -169,10 +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 @@ -222,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() @@ -632,6 +664,8 @@ 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]: @@ -713,6 +747,50 @@ def _build_codex_command(self, spec: AgentSpec) -> List[str]: 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) @@ -728,7 +806,7 @@ 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 == "codex": + 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 @@ -760,57 +838,60 @@ def _materialize_result( except json.JSONDecodeError: continue final_text = "" - 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). session_id = None - 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. 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 - break - if ev.get("type") == "turn.completed" and isinstance(ev, dict): - if "usage" 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 - 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 + # 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 @@ -839,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 index 009ac95b..f87ea94e 100644 --- a/metainfer/orchestrator/tests/test_subagent_manager.py +++ b/metainfer/orchestrator/tests/test_subagent_manager.py @@ -137,3 +137,216 @@ def test_bootstrap_factory_reads_codex_env(monkeypatch): 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/token_budget.py b/metainfer/orchestrator/token_budget.py index f0b208f0..f4c24ca3 100644 --- a/metainfer/orchestrator/token_budget.py +++ b/metainfer/orchestrator/token_budget.py @@ -570,16 +570,28 @@ def usage_from_result_event( 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), + 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(usage.get("cache_creation_input_tokens", 0) or 0), - total_cost_usd=float(event.get("total_cost_usd", 0.0) or 0.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/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/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/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/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) From e3e436308a711c94863ab725cd75bfaadc6e3976 Mon Sep 17 00:00:00 2001 From: 6eanut Date: Fri, 11 Sep 2026 14:31:36 +0800 Subject: [PATCH 3/7] feat(evolve-kernel): kernel source/diff viewer in detail page MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The detail page could open a library kernel's code, but the code was rendered as a plain
 with no line numbers and no way to see what a
kernel actually changed relative to a baseline — the key question when
reviewing an evolved kernel is "what did it change, vs the original or
vs the kernel it was derived from?".

Add a diff endpoint and an inspector panel:

- _state_readers.py: read_kernel_source() resolves one library kernel by
  id; read_kernel_diff() produces a unified diff against the original
  reference kernel or against the kernel's parent_id. Missing kernels, a
  root kernel with no parent, and an evicted parent (library caps at
  MAX_LIBRARY_SIZE) return exists=false with a reason instead of raising.
- routes.py: GET /kernels/{kernel_id}/diff?base=reference|parent.
- ok-runtime-api.js: getKernelDiff() client.
- ok-detail.js: KernelInspector panel with Source / Diff vs reference /
  Diff vs parent modes, a line-numbered source renderer, and a colored
  unified-diff renderer with +/- stats. The parent mode is disabled when
  the kernel has no parent.
- ok.css: .ok-source* / .ok-diff* styles.
- test_kernel_diff.py: 11 tests covering both readers and their
  non-raising failure paths.

Co-Authored-By: deepseek/deepseek-v4-flash-0731 
---
 .../evolve_kernel/server/_state_readers.py    |  92 +++++++++
 .../tasks/evolve_kernel/server/routes.py      |  10 +
 .../tasks/evolve_kernel/static/ok-detail.js   | 107 +++++++++-
 .../evolve_kernel/static/ok-runtime-api.js    |   8 +
 metainfer/tasks/evolve_kernel/static/ok.css   | 104 ++++++++++
 .../evolve_kernel/tests/test_kernel_diff.py   | 189 ++++++++++++++++++
 6 files changed, 503 insertions(+), 7 deletions(-)
 create mode 100644 metainfer/tasks/evolve_kernel/tests/test_kernel_diff.py

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

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

` : null} + `; +} + function ApiDefaultsPreview(_a) { var catalog = _a.catalog; var tpSizes = catalog @@ -502,7 +520,9 @@ function AiShapeSubset(_a) { var catalog = modelCatalog(model); var _b = useState(function () { var parsed = parseShapeRecords(value || ""); - return selectedShapes(parsed, true, catalog, model); + // Start EMPTY (no fallback to the full catalog): only shapes the user + // ticks (or Select-all) enter this task. + return selectedShapes(parsed, false, catalog, model); }), selected = _b[0], setSelected = _b[1]; var emit = useCallback(function (next) { @@ -515,6 +535,10 @@ function AiShapeSubset(_a) { useEffect(function () { emit(selected); }, []); + var count = catalog.filter(function (shape) { + return selected[shape.id]; + }).length; + return html`
@@ -523,6 +547,15 @@ function AiShapeSubset(_a) { shapes stay on the trusted fallback and are checked again before the candidate is published.
+ <${SubsetActions} + count=${count} + total=${catalog.length} + onAll=${function () { + var next = {}; + catalog.forEach(function (shape) { next[shape.id] = true; }); + emit(next); + }} + onNone=${function () { emit({}); }} /> <${ShapeCatalog} shapes=${catalog} selected=${selected} @@ -573,7 +606,8 @@ function ManualGpuAssignment(_a) { var _b = useState(function () { var parsedShapes = parseShapeRecords(value || ""); return subset - ? selectedShapes(parsedShapes, true, catalog, model) + // Start EMPTY for a manual subset: only ticked shapes enter this task. + ? selectedShapes(parsedShapes, false, catalog, model) : selectedMap(catalog.map(function (shape) { return shape.id; })); @@ -620,6 +654,21 @@ function ManualGpuAssignment(_a) { appears exactly once; empty GPU cards are allowed.
${subset && html` + <${SubsetActions} + count=${catalog.filter(function (s) { return selected[s.id]; }).length} + total=${catalog.length} + onAll=${function () { + var nextSelected = {}; + var nextOwners = Object.assign({}, owners); + catalog.forEach(function (shape) { + nextSelected[shape.id] = true; + if (nextOwners[shape.id] == null) { + nextOwners[shape.id] = operatorGroupedOwners(catalog)[shape.id]; + } + }); + emitState(nextOwners, nextSelected); + }} + onNone=${function () { emitState(owners, {}); }} /> <${ShapeCatalog} shapes=${catalog} selected=${selected} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css index 55f7e00f..94b628aa 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css +++ b/metainfer/tasks/dcu_kernel_auto_opt/static/dkao.css @@ -907,3 +907,38 @@ color: var(--muted); font-size: var(--fs-xs); } + +/* Subset (Selected shapes only): quick select-all / clear bar */ +.dkao-subset-actions { + display: flex; + align-items: center; + gap: 10px; + flex-wrap: wrap; + margin: 10px 0 4px; + padding: 8px 10px; + border: 1px dashed var(--border); + border-radius: var(--radius-sm); + background: var(--bg-elev); +} +.dkao-subset-count { + font-weight: 700; + color: var(--fg); + font-size: var(--fs-xs); +} +.dkao-subset-actions .dkao-shape-tab { + padding: 4px 12px; + border: 1px solid var(--border); + border-radius: var(--radius-sm); + background: var(--panel-2); + color: var(--fg); + cursor: pointer; + font-size: var(--fs-xs); +} +.dkao-subset-actions .dkao-shape-tab:disabled { + opacity: 0.5; + cursor: not-allowed; +} +.dkao-subset-hint { + color: var(--muted); + font-size: var(--fs-xxs); +} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gate_policy.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gate_policy.py new file mode 100644 index 00000000..8cb57b72 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gate_policy.py @@ -0,0 +1,158 @@ +"""gates.yaml is wired: the runtime reads it, the manifest decides if it applies.""" + +from __future__ import annotations + +import shutil + +import pytest +import yaml + +from ..orchestrator import gate_policy as gp +from ..orchestrator import planner +from ..orchestrator.harness_io import default_harness_dir + + +@pytest.fixture(autouse=True) +def _clear_cache(): + gp.reset_cache() + yield + gp.reset_cache() + + +def _workspace(tmp_path, *, gates_wired=True, policy_wired=True): + root = tmp_path / "harness" + shutil.copytree(default_harness_dir(), root) + manifest = yaml.safe_load((root / "manifest.yaml").read_text(encoding="utf-8")) + manifest["components"]["gates"]["wired"] = gates_wired + manifest["components"]["planner_policy"]["wired"] = policy_wired + (root / "manifest.yaml").write_text(yaml.safe_dump(manifest, sort_keys=False), + encoding="utf-8") + return root + + +def _edit_gates(root, **values): + data = yaml.safe_load((root / "gates.yaml").read_text(encoding="utf-8")) + for key, value in values.items(): + node = data["gates"] + parts = key.split(".") + for part in parts[:-1]: + node = node.setdefault(part, {}) + node[parts[-1]] = value + (root / "gates.yaml").write_text(yaml.safe_dump(data, sort_keys=False), + encoding="utf-8") + + +def test_seed_matches_builtin_defaults(): + """The seed file mirrors the historical constants (no surprise changes).""" + seed = yaml.safe_load( + (default_harness_dir() / "gates.yaml").read_text(encoding="utf-8") + )["gates"] + defaults = gp.builtin_gates() + assert seed["round_acceptance_improvement_percent"] == ( + defaults["round_acceptance_improvement_percent"]) + assert seed["shadow"]["min_improvement_percent"] == ( + defaults["shadow"]["min_improvement_percent"]) + assert seed["plateau"]["max_regression_percent"] == ( + defaults["plateau"]["max_regression_percent"]) + assert seed["isa_gate"]["required_valid_isa_guided_rounds"] == ( + defaults["isa_gate"]["required_valid_isa_guided_rounds"]) + assert seed["task_budget"]["default_max_iterations"] == ( + defaults["task_budget"]["default_max_iterations"]) + + +def test_editing_gates_changes_runtime_values(tmp_path, monkeypatch): + root = _workspace(tmp_path) + _edit_gates(root, **{ + "round_acceptance_improvement_percent": 2.5, + "p90_guard": "1.05", + "shadow.enabled": False, + "shadow.min_improvement_percent": 0.7, + "plateau.max_regression_percent": 3.5, + "isa_gate.required_valid_isa_guided_rounds": 4, + "task_budget.default_max_iterations": 6, + }) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + + assert gp.round_acceptance_improvement_percent() == 2.5 + assert gp.p90_tolerance() == 1.05 + assert gp.shadow_enabled() is False + assert gp.shadow_min_improvement_percent() == 0.7 + assert gp.plateau_max_regression_percent() == 3.5 + assert gp.isa_required_valid_rounds() == 4 + assert gp.default_max_iterations() == 6 + # the pipeline helpers see the same values + from ..orchestrator import w8a8_pipeline as W + assert W._gates.round_acceptance_improvement_percent() == 2.5 + + +def test_unwired_gates_component_is_ignored(tmp_path, monkeypatch): + """wired:false means the file is inert — the manifest is meaningful.""" + root = _workspace(tmp_path, gates_wired=False) + _edit_gates(root, **{"round_acceptance_improvement_percent": 9.0}) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + assert gp.round_acceptance_improvement_percent() == 1.0 + assert gp.component_wired("gates") is False + + +def test_unwired_planner_policy_falls_back_to_builtin(tmp_path, monkeypatch): + root = _workspace(tmp_path, policy_wired=False) + data = yaml.safe_load((root / "planner_policy.yaml").read_text(encoding="utf-8")) + data["fallback"]["m16"] = ["consolidate"] * 8 + (root / "planner_policy.yaml").write_text(yaml.safe_dump(data, sort_keys=False), + encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + effective = planner.policy() + assert effective["fallback"]["m16"] != ["consolidate"] * 8 + assert effective["fallback"] == planner._builtin_policy()["fallback"] + + +def test_required_hip_rounds_rule_and_literal(tmp_path, monkeypatch): + root = _workspace(tmp_path) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + assert gp.required_hip_rounds(10) == 8 # documented expression + _edit_gates(root, **{"isa_gate.required_hip_rounds_rule": 3}) + gp.reset_cache() + assert gp.required_hip_rounds(10) == 3 # literal override + + +def test_p90_guard_mode_parsing(tmp_path, monkeypatch): + root = _workspace(tmp_path) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + assert gp.p90_guard_mode() == "no_p90_regression" + assert gp.p90_tolerance() == 1.0 + _edit_gates(root, **{"p90_guard": "not-a-number"}) + gp.reset_cache() + assert gp.p90_tolerance() == 1.0 # never loosens by accident + + +def test_snapshot_records_what_a_run_used(tmp_path, monkeypatch): + """Children record their effective gates so the mechanism gate can verify.""" + import json + + root = _workspace(tmp_path) + _edit_gates(root, **{"round_acceptance_improvement_percent": 3.25}) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + state = tmp_path / "state" + out = gp.snapshot(state) + assert out is not None and out.is_file() + payload = json.loads(out.read_text(encoding="utf-8")) + assert payload["gates"]["round_acceptance_improvement_percent"] == 3.25 + assert payload["gates_wired"] is True + assert payload["harness_root"] == str(root) + + +def test_snapshot_marks_unwired_gates(tmp_path, monkeypatch): + import json + + root = _workspace(tmp_path, gates_wired=False) + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + gp.reset_cache() + out = gp.snapshot(tmp_path / "state2") + payload = json.loads(out.read_text(encoding="utf-8")) + assert payload["gates_wired"] is False + assert payload["gates"]["round_acceptance_improvement_percent"] == 1.0 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py new file mode 100644 index 00000000..f3f0b80c --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py @@ -0,0 +1,121 @@ +"""GPU admission pre-check: usage rule + pollution marking.""" + +from __future__ import annotations + +import pytest + +from ..orchestrator import gpu_preflight as gp + +SMI_SAMPLE = """ +================================= System Management Interface ================================== +================================================================================================ +HCU Temp AvgPwr Perf PwrCap VRAM% HCU% Mode +0 51.0C 145.0W manual 300.0W 0% 0.0% Normal +1 49.0C 147.0W manual 300.0W 12% 4.5% Normal +2 57.0C 251.0W manual 300.0W 40% 97.0% Normal +3 58.0C 254.0W manual 300.0W 92% 81.7% Normal +======================================== End of SMI Log ========================================= +""" + + +def test_parse_hy_smi_rows(): + parsed = gp.parse_hy_smi(SMI_SAMPLE) + assert set(parsed) == {0, 1, 2, 3} + assert parsed[0]["vram_percent"] == 0.0 + assert parsed[2]["util_percent"] == 97.0 + assert parsed[3]["power_w"] == 254.0 + assert parsed[1]["temp_c"] == 49.0 + + +def test_usability_is_free_vram_not_percentage(): + """The rule is `free >= 4GB`, so a 91%-used card is still usable.""" + ok = gp.check_gpu(0, {"vram_percent": 91.0, "util_percent": 0.0}, + total_gb=64.0) + too_full = gp.check_gpu(1, {"vram_percent": 94.0, "util_percent": 0.0}, + total_gb=64.0) + assert ok["usable"] is True and ok["free_gb"] == 5.76 + assert ok["measurement_suspect"] is False + assert too_full["usable"] is False and too_full["free_gb"] == 3.84 + assert any("VRAM free" in r for r in too_full["reasons"]) + # a full device is a capacity problem, not a pollution problem + assert too_full["measurement_suspect"] is False + assert gp.DEFAULT_MIN_FREE_GB == 4.0 + + +def test_busy_marks_suspect_without_blocking(): + busy = gp.check_gpu(2, {"vram_percent": 40.0, "util_percent": 97.0, + "power_w": 251.0}, foreign_pids=[12345]) + assert busy["usable"] is True # usage is allowed (VRAM ok) + assert busy["measurement_suspect"] is True + assert any("busy" in r for r in busy["reasons"]) + assert any("power" in r for r in busy["reasons"]) + # foreign pids are informational (idle holders are harmless) + assert any("foreign" in n for n in busy["notes"]) + + +def test_idle_foreign_holders_are_not_suspicion(): + """A device at 0% util / idle power is clean even if other containers hold it.""" + idle = gp.check_gpu(0, {"vram_percent": 0.0, "util_percent": 0.0, + "power_w": 144.0}, foreign_pids=[1, 2, 3]) + assert idle["usable"] is True + assert idle["measurement_suspect"] is False + assert idle["reasons"] == [] + assert any("foreign" in n for n in idle["notes"]) + + +def test_preflight_prefers_clean_then_suspect_then_over_limit(monkeypatch): + monkeypatch.setattr(gp, "sample_gpu_state", + lambda ids=None, **kw: { + 0: {"vram_percent": 1.0, "util_percent": 0.0}, + 1: {"vram_percent": 5.0, "util_percent": 55.0, + "power_w": 250.0}, + 2: {"vram_percent": 30.0, "util_percent": 3.0}, + 3: {"vram_percent": 97.0, "util_percent": 10.0}, + }) + monkeypatch.setattr(gp, "foreign_kfd_pids", lambda *a, **k: []) + monkeypatch.setattr(gp, "total_vram_gb", lambda *a, **k: 64.0) + plan = gp.preflight_gpus([0, 1, 2, 3]) + assert plan["clean_ids"] == [0, 2] + assert plan["suspect_ids"] == [1] + assert plan["over_limit_ids"] == [3] + # usable devices first (clean before suspect); the over-limit device drops + # out while any device is within the limit + assert plan["preferred"] == [0, 2, 1] + assert 3 not in plan["preferred"] + + +def test_over_limit_devices_return_when_nothing_else_is_usable(monkeypatch): + monkeypatch.setattr(gp, "sample_gpu_state", + lambda ids=None, **kw: { + 0: {"vram_percent": 97.0, "util_percent": 5.0}, + 1: {"vram_percent": 99.0, "util_percent": 5.0}, + }) + monkeypatch.setattr(gp, "foreign_kfd_pids", lambda *a, **k: []) + monkeypatch.setattr(gp, "total_vram_gb", lambda *a, **k: 64.0) + plan = gp.preflight_gpus([0, 1]) + assert plan["preferred"] == [0, 1] + + +def test_foreign_kfd_pids_detects_invisible_pids(tmp_path): + (tmp_path / "1").write_text("", encoding="utf-8") # our own pid + (tmp_path / "99999999").write_text("", encoding="utf-8") # foreign + pids = gp.foreign_kfd_pids(tmp_path) + assert 99999999 in pids + assert 1 not in pids + + +def test_preflight_enabled_env_and_answers(monkeypatch): + monkeypatch.delenv("METAINFER_GPU_PREFLIGHT", raising=False) + assert gp.preflight_enabled({}) is True # default on + assert gp.preflight_enabled({"gpu_preflight": "false"}) is False + monkeypatch.setenv("METAINFER_GPU_PREFLIGHT", "1") + assert gp.preflight_enabled({"gpu_preflight": "false"}) is True + monkeypatch.setenv("METAINFER_GPU_PREFLIGHT", "off") + assert gp.preflight_enabled({}) is False + + +def test_disabled_preflight_returns_default_order(): + plan = gp.preflight_gpus([0, 1, 2, 3], enabled=False) + assert plan["enabled"] is False + assert plan["preferred"] == [0, 1, 2, 3] + assert plan["gpus"] == {} diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_harness_io.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_harness_io.py new file mode 100644 index 00000000..013375af --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_harness_io.py @@ -0,0 +1,105 @@ +"""Tests for orchestrator.harness_io (M1 slice 1: harness seed + reader).""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import yaml + +from ..orchestrator import w8a8_pipeline +from ..orchestrator.config import ROUND_ACCEPTANCE_IMPROVEMENT_PERCENT +from ..orchestrator.harness_io import ( + default_harness_dir, + harness_root, + load_gates, + load_manifest, + plugin_dir, + seed_workspace, +) + + +def test_plugin_dir_and_default_seed_exist(): + root = plugin_dir() + assert root.name == "dcu_kernel_auto_opt" + seed = default_harness_dir() + assert seed.is_dir() + assert (seed / "manifest.yaml").is_file() + assert (seed / "gates.yaml").is_file() + + +def test_harness_root_defaults_to_seed(monkeypatch): + monkeypatch.delenv("METAINFER_HARNESS_ROOT", raising=False) + assert harness_root() == default_harness_dir().resolve() + + +def test_harness_root_env_override(monkeypatch, tmp_path): + marker = tmp_path / "marker.txt" + marker.write_text("x", encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(tmp_path)) + assert harness_root() == tmp_path.resolve() + + +def test_manifest_lists_gates_component(): + manifest = load_manifest() + assert manifest.get("schema_version") == 1 + components = manifest.get("components", {}) + assert "gates" in components + # gates is read at runtime now (orchestrator/gate_policy.py), so it is wired + assert components["gates"]["wired"] is True + assert "gate_policy" in (components["gates"].get("role") or "") or ( + "runtime" in (components["gates"].get("role") or "")) + + +def test_gates_seed_matches_builtin_defaults(): + """Drift guard: the seed gates.yaml must equal the built-in defaults. + + The runtime now reads gates.yaml through orchestrator/gate_policy.py (see + tests/test_gate_policy.py); this test keeps the *seed file* honest so a + fresh harness starts exactly where the hard-coded behaviour was. + """ + from ..orchestrator import gate_policy as gp + + gates = load_gates().get("gates", {}) + defaults = gp.builtin_gates() + assert gates["round_acceptance_improvement_percent"] == ( + defaults["round_acceptance_improvement_percent"]) + assert gates["shadow"]["min_improvement_percent"] == ( + defaults["shadow"]["min_improvement_percent"]) + assert gates["shadow"]["max_exclusive_percent"] == ( + defaults["shadow"]["max_exclusive_percent"]) + assert gates["plateau"]["max_regression_percent"] == ( + defaults["plateau"]["max_regression_percent"]) + assert gates["plateau"]["recent_valid_rounds"] == ( + defaults["plateau"]["recent_valid_rounds"]) + assert gates["isa_gate"]["required_valid_isa_guided_rounds"] == ( + defaults["isa_gate"]["required_valid_isa_guided_rounds"]) + assert gates["task_budget"]["default_max_iterations"] == ( + defaults["task_budget"]["default_max_iterations"]) + + +def test_seed_manifest_marks_data_components_wired(): + manifest = load_manifest() + components = manifest.get("components") or {} + assert components["gates"]["wired"] is True + assert components["planner_policy"]["wired"] is True + assert components["planner_catalog"]["wired"] is True + assert components["manifest"]["wired"] is True + + +def test_seed_workspace_copies_tree(tmp_path): + dst = seed_workspace(tmp_path / "ws") + assert dst.is_dir() + for name in ("manifest.yaml", "gates.yaml", "README.md"): + assert (dst / name).is_file() + # content equality with the seed + with (dst / "gates.yaml").open(encoding="utf-8") as fh: + copied = yaml.safe_load(fh) + original = load_gates() + assert copied == original + + +def test_seed_workspace_missing_root_raises(tmp_path): + with pytest.raises(FileNotFoundError): + seed_workspace(tmp_path / "out", root=tmp_path / "does-not-exist") diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plan_render.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plan_render.py new file mode 100644 index 00000000..20475fb3 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_plan_render.py @@ -0,0 +1,102 @@ +"""Tests for planner rendering + controlled-state parity with the legacy menu. + +Controlled-state parity (not corpus heuristic): on states where the planner's +legacy-aligned branches deliberately mirror the current round menus, assert the +mandate texts agree on key direction tokens. +""" + +from __future__ import annotations + +from ..orchestrator.planner import catalog, choose_plan_from_history, render_plan +from ..orchestrator.prompts import w8a8_round_strategy +from ..orchestrator.w8a8_pipeline import isa_round_policy + + +def _valid(improvement: float = 0.5) -> dict: + return { + "build_success": True, + "correctness_passed": True, + "speedup": 1.0, + "metrics": {"graph_capture_passed": True}, + "acceptance": {"improvement_percent": improvement}, + } + + +def test_render_all_catalog_plans_nonempty_and_deterministic(): + cat = catalog() + ctx = {"iteration": 3, "max_iterations": 10, "rounds_left": 8} + seen = {} + for pid in cat: + text = render_plan(pid, ctx=ctx, cat=cat) + assert pid in text and len(text) > 10 + seen[pid] = text + # deterministic + for pid, text in seen.items(): + assert render_plan(pid, ctx=ctx, cat=cat) == text + + +def test_fresh_lane_m16_matches_menu_on_dumma(): + # Fresh m16 lane: planner (establish_arch) and legacy menu (round 1 DUMMA + # bootstrap instruction) both direct at DUMMA. + plan = choose_plan_from_history([], iteration=1, shape={"M": 16}) + menu = w8a8_round_strategy( + {"M": 16, "N": 1536, "K": 4096}, + 1, + [], + {}, + max_iterations=10, + isa_policy=isa_round_policy( + iteration=1, max_iterations=10, history=[] + ), + ) + rendered = render_plan(plan, cat=catalog()) + assert "DUMMA" in rendered and "DUMMA" in menu + + +def test_plateau_open_matches_menu_on_isa(): + history = [_valid(0.5) for _ in range(5)] + [ + _valid(v) for v in (0.5, 0.9, 1.2) + ] + shape = {"M": 16, "N": 1536, "K": 4096} + iteration, max_iterations = 9, 10 + policy = isa_round_policy( + iteration=iteration, max_iterations=max_iterations, history=history + ) + assert policy.get("phase") == "isa_guided_hip" + + plan = choose_plan_from_history( + history, iteration=iteration, max_iterations=max_iterations, + shape=shape, + ) + menu = w8a8_round_strategy( + shape, iteration, history, {}, + max_iterations=max_iterations, isa_policy=policy, + ) + rendered = render_plan(plan, cat=catalog()) + assert plan == "isa_guided_hip" + assert "ISA" in rendered and "ISA" in menu + + +def test_faster_wrong_matches_menu_on_repair(): + history = [_valid(0.5), { + "build_success": True, + "correctness_passed": False, + "speedup": 1.9, + "metrics": {"graph_capture_passed": True}, + }] + shape = {"M": 16, "N": 1536, "K": 4096} + iteration, max_iterations = 4, 10 + plan = choose_plan_from_history( + history, iteration=iteration, max_iterations=max_iterations, + shape=shape, + ) + policy = isa_round_policy( + iteration=iteration, max_iterations=max_iterations, history=history + ) + menu = w8a8_round_strategy( + shape, iteration, history, {}, + max_iterations=max_iterations, isa_policy=policy, + ) + rendered = render_plan(plan, cat=catalog()) + assert plan == "repair_faster_wrong" + assert "repair" in rendered.lower() and "repair" in menu.lower() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner.py new file mode 100644 index 00000000..ffe79f98 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner.py @@ -0,0 +1,204 @@ +"""Tests for orchestrator.planner (M1 slice 2: state-conditioned selector v0).""" + +from __future__ import annotations + +from ..orchestrator import planner + + +def _valid_record(improvement: float = 0.5, *, speedup: float = 1.0) -> dict: + return { + "build_success": True, + "correctness_passed": True, + "speedup": speedup, + "metrics": {"graph_capture_passed": True}, + "acceptance": {"improvement_percent": improvement}, + } + + +def _wrong_fast_record() -> dict: + return { + "build_success": True, + "correctness_passed": False, + "speedup": 1.8, + "metrics": {"graph_capture_passed": True}, + } + + +def _build_fail_record() -> dict: + return {"build_success": False, "failure_reason": "compile error line 42"} + + +def _infra_record() -> dict: + return { + "build_success": False, + "failure_reason": "timed out after 900s", + } + + +def test_catalog_loads_from_yaml_seed(): + cat = planner.catalog() + for pid in ("repair_faster_wrong", "isa_guided_hip", "consolidate"): + assert pid in cat + + +def test_p0_faster_wrong_beats_everything(): + history = [_valid_record(), _wrong_fast_record()] + assert planner.choose_plan_from_history(history, iteration=3) == ( + "repair_faster_wrong" + ) + + +def test_empty_history_is_not_build_failure(): + # A fresh lane (no previous round) must not be routed to fix_build. + assert planner.choose_plan_from_history([]) == "establish_arch" + + +def test_p0_infra_retry_and_build_fix(): + assert planner.choose_plan_from_history([_valid_record(), _infra_record()]) == ( + "retry_same" + ) + assert planner.choose_plan_from_history( + [_valid_record(), _build_fail_record()] + ) == "fix_build" + + +def test_p1_last_round_consolidates(): + history = [_valid_record() for _ in range(3)] + assert planner.choose_plan_from_history( + history, iteration=10, max_iterations=10 + ) == "consolidate" + + +def test_p1_plateau_opens_isa_then_asm(): + # 8 valid HIP rounds + last three improvements inside [-2, 2) prove + # plateau with the ISA gate open (required_hip_rounds = max_iterations - 2). + history = [_valid_record(0.5) for _ in range(5)] + [ + _valid_record(v) for v in (0.5, 0.9, 1.2) + ] + plan = planner.choose_plan_from_history( + history, iteration=9, max_iterations=10 + ) + assert plan == "isa_guided_hip" + plan2 = planner.choose_plan_from_history( + history, + iteration=9, + max_iterations=10, + compiler_limitation_confirmed=True, + ) + assert plan2 == "conditional_inline_asm" + + +def test_p2_occupancy_signature_selects_resource_round(): + history = [_valid_record() for _ in range(2)] + pmc = {"waves_per_cu": 12, "target_waves_per_cu": 16} + assert planner.choose_plan_from_history( + history, pmc=pmc, iteration=5 + ) == "occupancy_resource" + + +def test_p2_bank_conflict_signature(): + history = [_valid_record() for _ in range(2)] + pmc = {"lds_instructions": 1000, "lds_bank_conflicts": 2000} + assert planner.choose_plan_from_history( + history, pmc=pmc, iteration=5 + ) == "memory_layout" + + +def test_p4_fallback_matches_legacy_menu_intent(monkeypatch): + # The no-evidence fallback (uncertainty -> legacy_menu) is on by default; + # disable it here so the P4 portfolio table itself is exercised. + pol = dict(planner._builtin_policy()) + pol["uncertainty"] = {"enabled": False} + monkeypatch.setattr(planner, "policy", lambda root=None: pol) + history = [_valid_record() for _ in range(2)] + plan = planner.choose_plan_from_history( + history, iteration=2, max_iterations=10, shape={"M": 4096} + ) + assert plan == "memory_layout" # large_m legacy menu, step after bootstrap + + +def test_coverage_guard_breaks_same_plan_repetition(): + history = [_valid_record() for _ in range(2)] + pmc = { + "waves_per_cu": 12, + "target_waves_per_cu": 16, + "l2_hit_rate": 40.0, + } + # two prior identical occupancy rounds -> P3 should route to the next + # candidate (memory_layout from l2_low) instead of repeating occupancy. + plan = planner.choose_plan_from_history( + history, + pmc=pmc, + iteration=6, + plan_tags=["occupancy_resource", "occupancy_resource"], + ) + assert plan == "memory_layout" + + +def test_fallback_portfolio_cycles_instead_of_pinning_the_tail(): + """Long runs must not be pinned to the table's last row (consolidate). + + 9-8-8 iteration 3 showed 9 of 11 planner picks were ``consolidate`` + because the fallback indexed past the table end and clamped. The + portfolio now cycles, so a long run keeps revisiting exploration plans. + """ + from ..orchestrator.planner import _legacy_fallback + + picks = [ + _legacy_fallback({"shape": {"M": 16}, "iteration": it}) + for it in range(2, 14) + ] + assert len(set(picks)) > 2, picks + tail = picks[-4:] + assert len(set(tail)) > 1, tail + assert tail.count("consolidate") < len(tail), tail + + +def test_uncertainty_falls_back_to_legacy_menu(monkeypatch): + """No bottleneck evidence -> defer to the hand-tuned menu.""" + from ..orchestrator import planner as P + + pol = dict(P._builtin_policy()) + pol["uncertainty"] = {"enabled": True, "plan": "legacy_menu", + "min_valid_rounds": 1} + monkeypatch.setattr(P, "policy", lambda root=None: pol) + monkeypatch.setattr(P, "bottleneck_tags", lambda pmc: []) + ctx = {"valid_hip_rounds": 2, "pmc": {}, "shape": {"M": 16}, + "iteration": 3, "max_iterations": 11, "rounds_left": 9, + "tried_counts": {}, "consecutive_same": 0, + "last_present": False, "faster_wrong": False} + assert P.choose_plan(ctx) == "legacy_menu" + + +def test_uncertainty_is_skipped_when_a_bottleneck_is_known(monkeypatch): + from ..orchestrator import planner as P + + pol = dict(P._builtin_policy()) + pol["uncertainty"] = {"enabled": True, "plan": "legacy_menu"} + pol["bottleneck_to_plan"] = {"lds_bank_conflicts": "memory_layout"} + monkeypatch.setattr(P, "policy", lambda root=None: pol) + monkeypatch.setattr(P, "bottleneck_tags", + lambda pmc: ["lds_bank_conflicts"]) + ctx = {"valid_hip_rounds": 2, "pmc": {"available": True}, + "shape": {"M": 16}, "iteration": 3, "max_iterations": 11, + "rounds_left": 9, "tried_counts": {}, "consecutive_same": 0, + "last_present": False, "faster_wrong": False} + assert P.choose_plan(ctx) == "memory_layout" + + +def test_legacy_menu_plan_renders_the_legacy_menu(tmp_path, monkeypatch): + import json as _json + from ..orchestrator import w8a8_pipeline as W + + monkeypatch.setenv("METAINFER_PLANNER", "1") + monkeypatch.setattr(W, "choose_plan_from_history", + lambda *a, **k: "legacy_menu") + sink = tmp_path / "planner_plans.jsonl" + text = W._round_strategy_text( + {"M": 16, "N": 4096, "K": 2048}, 2, [], {"available": False}, + {"max_iterations": 5, "skill_allowed": True}, + plan_sink=sink, shape_id="hy3_tp4_o_proj_m16", + ) + assert text and len(text) > 50 + row = _json.loads(sink.read_text(encoding="utf-8").splitlines()[-1]) + assert row["plan_id"] == "legacy_menu" and row["source"] == "planner" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_policy.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_policy.py new file mode 100644 index 00000000..56e4a4cd --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_policy.py @@ -0,0 +1,62 @@ +"""Prove planner_policy.yaml is a live, evolvable state->plan component.""" + +from __future__ import annotations + +import shutil +from pathlib import Path + +import yaml + +from ..orchestrator import planner +from ..orchestrator.harness_io import default_harness_dir + + +def _valid(): + return {"build_success": True, "correctness_passed": True, + "metrics": {"graph_capture_passed": True}, + "acceptance": {"improvement_percent": 5.0}} + + +def test_policy_seed_loads(): + pol = planner.policy() + assert pol["repair_priority"]["faster_wrong"] == "repair_faster_wrong" + assert pol["bottleneck_to_plan"]["occupancy_limited"] == "occupancy_resource" + + +def test_workspace_policy_change_changes_selected_plan(monkeypatch, tmp_path): + root = tmp_path / "harness" + shutil.copytree(default_harness_dir(), root) + policy_path = root / "planner_policy.yaml" + data = yaml.safe_load(policy_path.read_text(encoding="utf-8")) + # AHE-like change: route occupancy pressure to pipeline_tune instead of + # occupancy_resource. The selector must obey the evolved workspace file. + data["bottleneck_to_plan"]["occupancy_limited"] = "pipeline_tune" + policy_path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + + plan = planner.choose_plan_from_history( + [_valid(), _valid()], + pmc={"waves_per_cu": 8, "target_waves_per_cu": 16}, + iteration=4, + max_iterations=10, + shape={"M": 4096}, + ) + assert plan == "pipeline_tune" + + +def test_policy_fallback_order_is_evolvable(monkeypatch, tmp_path): + root = tmp_path / "harness" + shutil.copytree(default_harness_dir(), root) + path = root / "planner_policy.yaml" + data = yaml.safe_load(path.read_text(encoding="utf-8")) + data["fallback"]["prefill"][0] = "occupancy_resource" + # exercise the evolvable portfolio itself, not the no-evidence fallback + data.setdefault("uncertainty", {})["enabled"] = False + path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") + monkeypatch.setenv("METAINFER_HARNESS_ROOT", str(root)) + + plan = planner.choose_plan_from_history( + [_valid(), _valid()], pmc={}, iteration=2, max_iterations=10, + shape={"M": 4096}, + ) + assert plan == "occupancy_resource" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_wiring.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_wiring.py new file mode 100644 index 00000000..9261414f --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_planner_wiring.py @@ -0,0 +1,84 @@ +"""Tests for planner wiring env toggle (METAINFER_PLANNER).""" + +from __future__ import annotations + +from ..orchestrator.w8a8_pipeline import ( + _planner_enabled, + _round_strategy_text, + isa_round_policy, +) +from ..orchestrator.prompts import w8a8_round_strategy + + +def _isa_policy(iteration: int, max_iterations: int, history) -> dict: + return isa_round_policy( + iteration=iteration, max_iterations=max_iterations, history=history + ) + + +def test_planner_disabled_by_default(monkeypatch): + monkeypatch.delenv("METAINFER_PLANNER", raising=False) + assert _planner_enabled() is False + + +def test_env_toggle_enables(monkeypatch): + monkeypatch.setenv("METAINFER_PLANNER", "1") + assert _planner_enabled() is True + + +def test_disabled_returns_exact_legacy_text(monkeypatch): + monkeypatch.delenv("METAINFER_PLANNER", raising=False) + shape = {"M": 16, "N": 1536, "K": 4096} + history = [] + policy = _isa_policy(1, 10, history) + got = _round_strategy_text(shape, 1, history, {}, policy) + expected = w8a8_round_strategy( + shape, 1, history, {}, + max_iterations=10, isa_policy=policy, + ) + assert got == expected + + +def test_enabled_renders_planner_mandate(monkeypatch): + monkeypatch.setenv("METAINFER_PLANNER", "1") + shape = {"M": 16, "N": 1536, "K": 4096} + history = [] + policy = _isa_policy(1, 10, history) + got = _round_strategy_text(shape, 1, history, {}, policy) + assert got.startswith("Mandatory decision for this round:") + assert "establish_arch" in got + + +def test_round_strategy_text_records_planner_decision(tmp_path, monkeypatch): + """The planner's own choice is written to planner_plans.jsonl (hard evidence).""" + import json as _json + from ..orchestrator.w8a8_pipeline import _round_strategy_text + + monkeypatch.setenv("METAINFER_PLANNER", "1") + sink = tmp_path / "planner_plans.jsonl" + text = _round_strategy_text( + {"M": 16, "N": 4096, "K": 2048}, 1, [], {"available": False}, + {"max_iterations": 3, "skill_allowed": True}, + plan_sink=sink, shape_id="hy3_tp4_o_proj_m16", + ) + assert text + assert sink.is_file() + row = _json.loads(sink.read_text(encoding="utf-8").splitlines()[0]) + assert row["source"] == "planner" + assert row["plan_id"] + assert row["iteration"] == 1 + assert row["shape_id"] == "hy3_tp4_o_proj_m16" + + +def test_round_strategy_text_writes_nothing_without_planner(tmp_path, monkeypatch): + from ..orchestrator.w8a8_pipeline import _round_strategy_text + + monkeypatch.delenv("METAINFER_PLANNER", raising=False) + sink = tmp_path / "planner_plans.jsonl" + text = _round_strategy_text( + {"M": 16, "N": 4096, "K": 2048}, 1, [], {"available": False}, + {"max_iterations": 3, "skill_allowed": True}, + plan_sink=sink, shape_id="hy3_tp4_o_proj_m16", + ) + assert text + assert not sink.exists() diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_predictions.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_predictions.py new file mode 100644 index 00000000..6b047a5b --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_predictions.py @@ -0,0 +1,94 @@ +"""Tests for orchestrator.predictions (M1 slice 3: inner decision hook).""" + +from __future__ import annotations + +from ..orchestrator.predictions import check_prediction, parse_prediction, plan_tag + + +def test_no_structured_prediction_returns_none(): + assert check_prediction( + {"hypothesis": "plain prose"}, + candidate_us=100.0, + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) is None + assert parse_prediction({"hypothesis": "x"}) is None + + +def test_infra_failure_is_na(): + out = check_prediction( + {"prediction": {"direction": "improve"}}, + candidate_us=float("inf"), + best_us=100.0, + passed=False, + p90_guard_passed=False, + failure_reason="timed out after 900s", + ) + assert out is not None and out["checked"] == "na" + + +def test_range_hit_and_miss(): + proposal = {"prediction": {"expected_us_range": [90.0, 105.0]}} + hit = check_prediction( + proposal, + candidate_us=98.0, + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) + assert hit["checked"] == "hit" + miss = check_prediction( + proposal, + candidate_us=120.0, + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) + assert miss["checked"] == "miss" + assert "outside declared range" in miss["reason"] + + +def test_direction_improve_requires_band(): + ok = check_prediction( + {"prediction": {"direction": "improve"}}, + candidate_us=100.0, + best_us=110.0, # ~9% better + passed=True, + p90_guard_passed=True, + ) + assert ok["checked"] == "hit" + flat = check_prediction( + {"prediction": {"direction": "flat"}}, + candidate_us=109.0, # ~0.9% delta -> inside flat band + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) + assert flat["checked"] == "hit" + wrong = check_prediction( + {"prediction": {"direction": "improve"}}, + candidate_us=109.0, # only ~0.9% better -> below the 2% band + best_us=110.0, + passed=True, + p90_guard_passed=True, + ) + assert wrong["checked"] == "miss" + + +def test_correctness_failure_means_miss(): + out = check_prediction( + {"prediction": {"expected_us_range": [90.0, 105.0]}}, + candidate_us=98.0, + best_us=110.0, + passed=False, + p90_guard_passed=False, + ) + assert out["checked"] == "miss" + assert "correctness failed" in out["reason"] + + +def test_plan_tag_helpers(): + assert plan_tag({"hypothesis": "x"}) is None + assert plan_tag({"plan_id": "occupancy_resource"}) == "occupancy_resource" + assert plan_tag({"plan_id": " "}) is None diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_prompt_schema.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_prompt_schema.py new file mode 100644 index 00000000..a3dd5acd --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_prompt_schema.py @@ -0,0 +1,47 @@ +"""Prompt-schema tests: worker prompt now carries optional decision fields.""" + +from __future__ import annotations + +from pathlib import Path + +from ..orchestrator.config import WorkerAssignment +from ..orchestrator.w8a8_pipeline import RealW8A8OptimizationPipeline + + +def _round_prompt(**overrides): + args = dict( + assignment=WorkerAssignment("worker_0", 0, ["m2"]), + shape_id="m2", + shape={"M": 2, "N": 16, "K": 32}, + best={"median_us": 10.0}, + root=Path("/tmp/worker"), + iteration=1, + guidance=None, + history=[], + pmc_evidence={}, + ) + args.update(overrides) + return RealW8A8OptimizationPipeline._worker_prompt(**args) + + +def test_first_turn_prompt_mentions_decision_fields(): + prompt = _round_prompt() + assert '"plan_id"' in prompt + assert '"prediction"' in prompt + assert '"expected_us_range"' in prompt + assert '"direction": "improve"' in prompt + assert "checks `prediction` against the measured round" in prompt + + +def test_continuation_prompt_mentions_optional_fields(): + prompt = _round_prompt( + continuation=True, + history=[{ + "iteration": 1, + "accepted": True, + "build_success": True, + "correctness_passed": True, + "metrics": {"median_us": 9.0}, + }], + ) + assert "optional plan_id and prediction" in prompt diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_skill_store.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_skill_store.py index 54333e68..9eda815c 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_skill_store.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_skill_store.py @@ -3,19 +3,31 @@ import pytest from ..orchestrator.config import load_config +from ..orchestrator import skill_store as skill_store_mod from ..orchestrator.skill_store import ( _apply_fuse_decision, _parse_fuse_decision, + bundled_skills_root, dsh_skills_root, generate_merged_skill, generate_worker_skill, list_skill_library, publish_skill, rollback_skill, + seed_bundled_skills, sync_skill_libraries, ) +@pytest.fixture(autouse=True) +def _isolated_bundled_skills(tmp_path, monkeypatch): + """Point the plugin's bundled-skills root at an empty tmp dir by default, + so seeding never pollutes library assertions in unrelated tests.""" + bundled = tmp_path / "bundled-skills" + bundled.mkdir(exist_ok=True) + monkeypatch.setattr(skill_store_mod, "bundled_skills_root", lambda: bundled) + + def test_skill_library_lists_and_publishes_without_overwrite( tmp_path, monkeypatch ): @@ -84,6 +96,28 @@ def test_sync_skill_libraries_mirrors_dsh_to_ccb(tmp_path, monkeypatch): assert summary2["skipped"] == ["alpha", "beta"] +def test_seed_bundled_skills_copies_missing_only(tmp_path, monkeypatch): + dsh_root = tmp_path / "dsh" + bundled = tmp_path / "bundled" + (bundled / "zeta").mkdir(parents=True) + (bundled / "zeta" / "SKILL.md").write_text("# zeta seed\n", encoding="utf-8") + (bundled / "eta").mkdir(parents=True) + (bundled / "eta" / "SKILL.md").write_text("# eta seed\n", encoding="utf-8") + # dsh already has eta (different content) -> must not be overwritten + (dsh_root / "eta").mkdir(parents=True) + (dsh_root / "eta" / "SKILL.md").write_text("# eta live\n", encoding="utf-8") + monkeypatch.setenv("DSH_SKILLS_DIR", str(dsh_root)) + monkeypatch.setattr(skill_store_mod, "bundled_skills_root", lambda: bundled) + + seeded = seed_bundled_skills() + assert seeded["added"] == ["zeta"] + assert (dsh_root / "zeta" / "SKILL.md").read_text() == "# zeta seed\n" + # an existing live skill is never clobbered by the bundled seed + assert (dsh_root / "eta" / "SKILL.md").read_text() == "# eta live\n" + # idempotent + assert seed_bundled_skills()["added"] == [] + + def test_rollback_skill_restores_backup(tmp_path, monkeypatch): dsh_root = tmp_path / "dsh" ccb_root = tmp_path / "ccb" diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_validation_budget.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_validation_budget.py new file mode 100644 index 00000000..fbd26361 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_validation_budget.py @@ -0,0 +1,54 @@ +"""Validation-budget knobs (scope + bench sampling) shared by DKAO and AHE.""" + +from __future__ import annotations + +import pytest + +from ..orchestrator.gen_and_opt_pipeline import _validation_shape_list +from ..orchestrator.validation_budget import ( + ENV_BENCH_SAMPLES, ENV_BENCH_WARMUPS, ENV_VALIDATE_SCOPE, QUICK_BENCH, + resolve_bench_kwargs, resolve_validation_scope, +) + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + for name in (ENV_VALIDATE_SCOPE, ENV_BENCH_WARMUPS, ENV_BENCH_SAMPLES, + "METAINFER_BENCH_REPLAYS"): + monkeypatch.delenv(name, raising=False) + + +def test_scope_defaults_to_full_api(monkeypatch): + assert resolve_validation_scope({}) == "api" + assert resolve_validation_scope(None) == "api" + + +def test_scope_reads_form_labels(monkeypatch): + assert resolve_validation_scope({"validation_scope": "Task shapes only"}) == "task" + assert resolve_validation_scope({"validation_scope": "All API shapes"}) == "api" + assert resolve_validation_scope({"validation_scope": "task"}) == "task" + + +def test_scope_env_wins(monkeypatch): + monkeypatch.setenv(ENV_VALIDATE_SCOPE, "task") + assert resolve_validation_scope({"validation_scope": "All API shapes"}) == "task" + + +def test_bench_kwargs_default_and_quick(monkeypatch): + assert resolve_bench_kwargs({}) == {} + assert resolve_bench_kwargs({"bench_profile": "quick"}) == QUICK_BENCH + assert resolve_bench_kwargs({"bench_profile": "full"}) == {} + + +def test_bench_kwargs_explicit_and_env_priority(monkeypatch): + out = resolve_bench_kwargs({"bench_profile": "quick", "bench_samples": 7}) + assert out["samples"] == 7 and out["warmups"] == QUICK_BENCH["warmups"] + monkeypatch.setenv(ENV_BENCH_SAMPLES, "5") + assert resolve_bench_kwargs({"bench_samples": 7})["samples"] == 5 + + +def test_validation_shape_list_scopes(): + optimized = [{"id": "a"}] + fallback = [{"id": "b"}, {"id": "c"}] + assert [s["id"] for s in _validation_shape_list(optimized, fallback, "task")] == ["a"] + assert [s["id"] for s in _validation_shape_list(optimized, fallback, "api")] == ["a", "b", "c"] diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_promote.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_promote.py new file mode 100644 index 00000000..cd23a824 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_promote.py @@ -0,0 +1,149 @@ +"""Kernel promotion gates (variant tree write) — no GPU calls.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from ..orchestrator import variant_store as vs +from ..orchestrator.variant_promote import ( + accepted_kernel_for, model_label_for, promote_variant, +) + +ANSWERS = {"operator": "Quantized GEMM", "dtype": "INT8 W8A8"} +SHAPE = "hy3_tp4_o_proj_m16" + + +@pytest.fixture() +def variant_root(tmp_path, monkeypatch): + root = tmp_path / "variant" + root.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(vs, "variant_root", lambda: root) + return root + + +def _workspace(tmp_path: Path, median_us: float, *, shape: str = SHAPE, + commit: str = "deadbeef") -> Path: + ws = tmp_path / "ws" + leaf = ws / "workers" / "worker_0" / "accepted" / shape + leaf.mkdir(parents=True, exist_ok=True) + (leaf / "kernel.hip").write_text("// kernel source\n", encoding="utf-8") + (leaf / "manifest.json").write_text(json.dumps({ + "commit": commit, + "shape": {"M": 16, "N": 4096, "K": 2048}, + "metrics": {"median_us": median_us, "p90_us": median_us * 1.02}, + }), encoding="utf-8") + return ws + + +def _seed_variant(variant_root: Path, median_us: float) -> Path: + meta = vs.derive_variant_meta(ANSWERS | {"model": "Hy3 (Hunyuan 3)"}, SHAPE) + target = vs.variant_path(meta) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(vs.section_header(meta, commit="old", metrics={ + "median_us": median_us, "p90_us": median_us}) + "// old kernel\n// @@end\n", + encoding="utf-8") + return target + + +def test_adds_when_no_existing_variant(tmp_path, variant_root): + ws = _workspace(tmp_path, 20.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=3.0, + source_task="he-test") + assert out["action"] == "added" + assert Path(out["path"]).is_file() + assert "he-test" in Path(out["path"]).read_text(encoding="utf-8") + + +def test_updates_when_faster_than_threshold(tmp_path, variant_root): + _seed_variant(variant_root, 100.0) + ws = _workspace(tmp_path, 90.0) # 10% faster + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=3.0) + assert out["action"] == "updated" + assert out["improvement_percent"] == pytest.approx(10.0, abs=1e-6) + assert out["backup"] and Path(out["backup"]).is_file() + assert "median_us=90" in Path(out["path"]).read_text(encoding="utf-8") + + +def test_skips_improvement_below_threshold(tmp_path, variant_root): + _seed_variant(variant_root, 100.0) + ws = _workspace(tmp_path, 98.5) # only 1.5% faster + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=3.0) + assert out["action"] == "skipped" + assert "required 3.00%" in out["reason"] + # existing variant untouched (still the old median) + assert "median_us=100" in Path(out["path"]).read_text(encoding="utf-8") + + +def test_slower_candidate_is_skipped_by_default_gate(tmp_path, variant_root): + """With the default gate (min 0%) any regression is skipped up front.""" + _seed_variant(variant_root, 100.0) + ws = _workspace(tmp_path, 120.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=0.0) + assert out["action"] == "skipped" + assert "required 0.00%" in out["reason"] + + +def test_add_variant_guard_rejects_slower_when_gate_is_loosened(tmp_path, variant_root): + """A loosened threshold still cannot overwrite with a slower kernel.""" + _seed_variant(variant_root, 100.0) + ws = _workspace(tmp_path, 101.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=-10.0) + assert out["action"] == "rejected" + assert "slower" in out["reason"] + + +def test_skips_when_correctness_failed(tmp_path, variant_root): + ws = _workspace(tmp_path, 10.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=False, min_improvement_percent=3.0) + assert out["action"] == "skipped" + assert "correctness" in out["reason"] + + +def test_no_kernel_found(tmp_path, variant_root): + out = promote_variant(workspace_dir=tmp_path / "empty", answers=ANSWERS, + shape_id=SHAPE) + assert out["action"] == "no_kernel" + + +def test_dry_run_does_not_write(tmp_path, variant_root): + ws = _workspace(tmp_path, 10.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True, min_improvement_percent=3.0, + dry_run=True) + assert out["action"] == "would-add" + assert not Path(out["path"]).exists() + + +def test_meta_taxonomy_and_overrides(tmp_path, variant_root): + ws = _workspace(tmp_path, 10.0) + out = promote_variant(workspace_dir=ws, answers=ANSWERS, shape_id=SHAPE, + correctness_ok=True) + meta = out["meta"] + assert (meta["family"], meta["model"], meta["tp"], meta["m"], + meta["operator_name"]) == ("int8w8a8-gemm", "hy3", 4, 16, "o_proj") + + # shapes without tp/M in the id take them from the caller (pool contract) + ws2 = _workspace(tmp_path / "two", 10.0, shape="m16_wqkv_a") + out2 = promote_variant(workspace_dir=ws2, answers=ANSWERS, + shape_id="m16_wqkv_a", correctness_ok=True, + tp=8, m=16, model_label="DeepSeek V4 Flash") + assert out2["meta"]["tp"] == 8 and out2["meta"]["m"] == 16 + assert "/TP8/M16/" in out2["path"] + + +def test_model_label_and_accepted_lookup(tmp_path): + assert model_label_for("minimax_tp8_o_proj_m16") == "MiniMax M3" + assert model_label_for("glm_tp8_o_proj_m4096") == "GLM5.2" + ws = _workspace(tmp_path, 5.0) + found = accepted_kernel_for(ws, SHAPE) + assert found and found["metrics"]["median_us"] == 5.0 + assert accepted_kernel_for(ws, "nope") is None diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_store.py b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_store.py index 6310e5e2..0f5a29c4 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_store.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_variant_store.py @@ -196,3 +196,31 @@ def test_reject_slower_guard_skips_when_median_unknown(tmp_path, monkeypatch): ) assert result["action"] == "updated" assert "NEW" in vs.variant_path(meta).read_text() + + +def test_backfill_variant_baselines(tmp_path, monkeypatch): + monkeypatch.setattr(vs, "variant_root", lambda: tmp_path) + target = ( + tmp_path / "int8w8a8-gemm" / "deepseek-v4" / "TP4" / "M16" + / "wqkv_a.hip" + ) + target.parent.mkdir(parents=True) + target.write_text( + "// @@variant shape=tp4_wqkv_a_m16 commit=abc added=2026-08-26\n" + "// median_us=20.0 p90_us=21.0\n" + "// source=task-1\n" + "// @@end\nkernel\n", + encoding="utf-8", + ) + r = vs.backfill_variant_baselines() + # tp4 wqkv_a m16 fixed Triton baseline is 66.597 us -> speedup ~3.33 + assert r["updated"] == ["int8w8a8-gemm/deepseek-v4/TP4/M16/wqkv_a.hip"] + text = target.read_text() + assert "baseline_us=66.597" in text + assert "speedup=3.329" in text + assert text.index("baseline_us") < text.index("median_us") + # idempotent and surfaced by the index + assert vs.backfill_variant_baselines()["updated"] == [] + v = vs.list_variant_index()[0] + assert v["speedup"] is not None + assert abs(v["baseline_us"] - 66.597) < 1e-3 diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp4_m4096.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp4_m4096.py new file mode 100644 index 00000000..6da463c1 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp4_m4096.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +"""Fair Triton W8A8 GEMM baseline for the six TP4 logical shapes at M=4096. + +Timed region launches only lmslim-style int8_utils.matmul_kernel into a +preallocated bf16 output (output allocation/clear excluded), matching the +MetaInfer fair-baseline scope. Uses GPU events, hot cache, and reports +median/P90/min in microseconds plus effective INT8 TOPS. + +The Triton config replicates matmul_int8's built-in default for M > 1024: + BLOCK_SIZE_M=256, BLOCK_SIZE_N=256, BLOCK_SIZE_K=64, + GROUP_SIZE_M=8, SPLIT_K=1, num_stages=0, num_warps=8 +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import math +import statistics +import sys +from pathlib import Path + +import torch +import triton + + +UTILS_PATH = str(Path(__file__).resolve().parent / "int8_utils.py") +M = 4096 + +# (shape_id, operator, K, N) -- the six TP4 logical operators. +TP4_CASES = [ + ("tp4_wqkv_a_m4096", "wqkv_a", 4096, 1536), + ("tp4_wq_b_m4096", "wq_b", 1024, 8192), + ("tp4_indexer_wq_b_m4096", "indexer.wq_b", 1024, 8192), + ("tp4_wo_b_m4096", "wo_b", 2048, 4096), + ("tp4_shared_gate_up_proj_m4096", "shared_gate_up_proj", 4096, 1024), + ("tp4_shared_down_proj_m4096", "shared_down_proj", 512, 4096), +] + +# Exact default config used by matmul_int8 when M > 1024. +CONFIG = { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8, +} + + +def load_int8_utils(path: str): + spec = importlib.util.spec_from_file_location( + "baseline_int8_utils", path + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def launch(out, a, a_scale, b, b_scale, m, n, k, utils) -> None: + grid = ( + triton.cdiv(m, CONFIG["BLOCK_SIZE_M"]) + * triton.cdiv(n, CONFIG["BLOCK_SIZE_N"]), + CONFIG["SPLIT_K"], + ) + utils.matmul_kernel[grid]( + a, + a_scale, + b, + b_scale, + out, + m, + n, + k, + a.stride(0), + a.stride(1), + a_scale.stride(0), + b.stride(0), + b.stride(1), + b_scale.stride(0), + out.stride(0), + out.stride(1), + **CONFIG, + ) + + +def percentile(values, fraction): + ordered = sorted(values) + return ordered[max(0, math.ceil(fraction * len(ordered)) - 1)] + + +def measure(fn, warmups, samples, launches_per_sample): + for _ in range(warmups): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + times_us = [] + for _ in range(samples): + start.record() + for _ in range(launches_per_sample): + fn() + end.record() + end.synchronize() + times_us.append( + start.elapsed_time(end) * 1000.0 / launches_per_sample + ) + return { + "median_us": statistics.median(times_us), + "p90_us": percentile(times_us, 0.90), + "min_us": min(times_us), + "max_us": max(times_us), + "samples_us": [round(v, 3) for v in times_us], + } + + +@torch.no_grad() +def reference_us(a, b, a_scale, b_scale, device): + dot = torch.mm( + a.to("cpu", dtype=torch.int64), + b.to("cpu", dtype=torch.int64), + ) + scaled = ( + dot.to(torch.float32) + * a_scale.to("cpu", dtype=torch.float32) + * b_scale.to("cpu", dtype=torch.float32).T + ) + return scaled.to(torch.bfloat16).to(device) + + +def run_case(utils, case_id, operator, k, n, warmups, samples, + launches_per_sample, skip_check, mode): + device = torch.device("cuda:0") + torch.manual_seed(20260724 + M + n + k) + a = torch.randint(-127, 128, (M, k), device=device, dtype=torch.int8) + b = torch.randint(-127, 128, (k, n), device=device, dtype=torch.int8) + a_scale = torch.rand((M, 1), device=device, dtype=torch.float32) + 0.01 + b_scale = torch.rand((n, 1), device=device, dtype=torch.float32) + 0.01 + out = torch.empty((M, n), device=device, dtype=torch.bfloat16) + + fn = lambda: launch(out, a, a_scale, b, b_scale, M, n, k, utils) + + # One untimed JIT compile + correctness sanity pass. + fn() + torch.cuda.synchronize() + errors = None + if not skip_check: + ref = reference_us(a, b, a_scale, b_scale, device) + diff = (out.float() - ref.float()).abs() + denom = ref.float().abs().clamp_min(1e-6) + errors = { + "max_abs": float(diff.max().item()), + "max_rel": float((diff / denom).max().item()), + "mismatch_count": int((out != ref).sum().item()), + } + del ref + + result = { + "shape_id": case_id, + "operator": operator, + "M": M, + "N": n, + "K": k, + "config": CONFIG, + "errors": errors, + } + + if mode in ("eager", "both"): + result["eager"] = measure( + fn, warmups, samples, launches_per_sample + ) + if mode in ("graph", "both"): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + torch.cuda.synchronize() + result["graph"] = measure( + graph.replay, warmups, samples, launches_per_sample + ) + del graph + + del a, b, a_scale, b_scale, out + return result + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--utils", default=UTILS_PATH) + parser.add_argument("--warmups", type=int, default=10) + parser.add_argument("--samples", type=int, default=20) + parser.add_argument("--launches-per-sample", type=int, default=5) + parser.add_argument("--mode", choices=["eager", "graph", "both"], + default="both") + parser.add_argument("--skip-check", action="store_true") + parser.add_argument("--json", type=str, default="") + args = parser.parse_args() + + utils = load_int8_utils(args.utils) + print( + "BENCHMARK triton=int8_utils.matmul_kernel output=preallocated " + f"quantization=excluded allocation=excluded epilogue=included " + f"M={M} config={json.dumps(CONFIG)} mode={args.mode}", + flush=True, + ) + print( + f"{'shape_id':28s} {'N':5s} {'K':5s} {'eager_us':>10s} " + f"{'graph_us':>10s} {'tops':>9s}", + flush=True, + ) + results = [] + for case_id, operator, k, n in TP4_CASES: + res = run_case( + utils, + case_id, + operator, + k, + n, + args.warmups, + args.samples, + args.launches_per_sample, + args.skip_check, + args.mode, + ) + results.append(res) + eager_us = res.get("eager", {}).get("median_us") + graph_us = res.get("graph", {}).get("median_us") + tops = ( + 2.0 * M * n * k / ((graph_us or eager_us) * 1.0e-6) / 1.0e12 + if (graph_us or eager_us) else float("nan") + ) + print( + f"{case_id:28s} {n:5d} {k:5d} " + f"{(eager_us if eager_us is not None else float('nan')):10.3f} " + f"{(graph_us if graph_us is not None else float('nan')):10.3f} " + f"{tops:9.3f}", + flush=True, + ) + for key in ("eager", "graph"): + if key in res: + m_ = res[key] + print( + f" {key}: median={m_['median_us']:.3f}us " + f"p90={m_['p90_us']:.3f}us min={m_['min_us']:.3f}us " + f"max={m_['max_us']:.3f}us errors={res['errors']}", + flush=True, + ) + if args.json: + with open(args.json, "w", encoding="utf-8") as handle: + json.dump( + { + "scope": "gemm_out", + "triton_source": args.utils, + "M": M, + "warmups": args.warmups, + "samples": args.samples, + "launches_per_sample": args.launches_per_sample, + "mode": args.mode, + "results": results, + }, + handle, + indent=2, + ensure_ascii=False, + ) + print(f"wrote {args.json}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m4096.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m4096.py new file mode 100644 index 00000000..3dc94b39 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m4096.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +"""Fair Triton W8A8 GEMM baselines for the model-catalog TP8 shapes at M=4096. + +Same protocol as bench_triton_tp4_m4096.py / the 2026-08-12 catalog entries: +lmslim-style int8_utils.matmul_kernel into a preallocated bf16 output (output +allocation/clear excluded), GPU events, hot cache, CUDA-Graph replay, +warmups=10 / samples=20 / launches_per_sample=5. Reports median/P90/min in +microseconds. + +Triton config replicates matmul_int8's built-in default for M > 1024: + BLOCK_SIZE_M=256, BLOCK_SIZE_N=256, BLOCK_SIZE_K=64, + GROUP_SIZE_M=8, SPLIT_K=1, num_stages=0, num_warps=8 +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import math +import statistics +import sys +from pathlib import Path + +import torch +import triton + + +UTILS_PATH = str(Path(__file__).resolve().parent / "int8_utils.py") +M = 4096 + +# (shape_id, operator, K, N) -- Hy3 / MiniMax M3 / GLM5.2 TP8 at M=4096. +TP8_CASES = [ + ("hy3_tp8_qkv_proj_m4096", "qkv_proj", 4096, 1280), + ("hy3_tp8_o_proj_m4096", "o_proj", 1024, 4096), + ("hy3_tp8_shared_gate_up_proj_m4096", "shared_gate_up_proj", 4096, 384), + ("hy3_tp8_shared_down_proj_m4096", "shared_down_proj", 192, 4096), + ("minimax_tp8_qkv_proj_m4096", "qkv_proj", 6144, 1280), + ("minimax_tp8_qkv_proj_and_indexer_qk_m4096", "qkv_proj_and_indexer_qk", 6144, 1536), + ("minimax_tp8_o_proj_m4096", "o_proj", 1024, 6144), + ("minimax_tp8_shared_gate_up_proj_m4096", "shared_gate_up_proj", 6144, 768), + ("minimax_tp8_shared_down_proj_m4096", "shared_down_proj", 384, 6144), + ("glm52_tp8_fused_qkv_a_proj_m4096", "fused_qkv_a_proj", 6144, 2624), + ("glm52_tp8_q_b_proj_m4096", "q_b_proj", 2048, 2048), + ("glm52_tp8_kv_b_proj_m4096", "kv_b_proj", 512, 3584), + ("glm52_tp8_o_proj_m4096", "o_proj", 2048, 6144), + ("glm52_tp8_shared_gate_up_proj_m4096", "shared_gate_up_proj", 6144, 512), + ("glm52_tp8_shared_down_proj_m4096", "shared_down_proj", 256, 6144), +] + +# Exact default config used by matmul_int8 when M > 1024. +CONFIG = { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8, +} + + +def load_int8_utils(path: str): + spec = importlib.util.spec_from_file_location( + "baseline_int8_utils", path + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def launch(out, a, a_scale, b, b_scale, m, n, k, utils) -> None: + grid = ( + triton.cdiv(m, CONFIG["BLOCK_SIZE_M"]) + * triton.cdiv(n, CONFIG["BLOCK_SIZE_N"]), + CONFIG["SPLIT_K"], + ) + utils.matmul_kernel[grid]( + a, + a_scale, + b, + b_scale, + out, + m, + n, + k, + a.stride(0), + a.stride(1), + a_scale.stride(0), + b.stride(0), + b.stride(1), + b_scale.stride(0), + out.stride(0), + out.stride(1), + **CONFIG, + ) + + +def percentile(values, fraction): + ordered = sorted(values) + return ordered[max(0, math.ceil(fraction * len(ordered)) - 1)] + + +def measure(fn, warmups, samples, launches_per_sample): + for _ in range(warmups): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + times_us = [] + for _ in range(samples): + start.record() + for _ in range(launches_per_sample): + fn() + end.record() + end.synchronize() + times_us.append( + start.elapsed_time(end) * 1000.0 / launches_per_sample + ) + return { + "median_us": statistics.median(times_us), + "p90_us": percentile(times_us, 0.90), + "min_us": min(times_us), + "max_us": max(times_us), + "samples_us": [round(v, 3) for v in times_us], + } + + +def run_case(utils, case_id, operator, k, n, warmups, samples, + launches_per_sample, mode): + device = torch.device("cuda:0") + torch.manual_seed(20260724 + M + n + k) + a = torch.randint(-127, 128, (M, k), device=device, dtype=torch.int8) + b = torch.randint(-127, 128, (k, n), device=device, dtype=torch.int8) + a_scale = torch.rand((M, 1), device=device, dtype=torch.float32) + 0.01 + b_scale = torch.rand((n, 1), device=device, dtype=torch.float32) + 0.01 + out = torch.empty((M, n), device=device, dtype=torch.bfloat16) + + fn = lambda: launch(out, a, a_scale, b, b_scale, M, n, k, utils) + + # One untimed JIT compile pass (correctness is the production Triton path). + fn() + torch.cuda.synchronize() + + result = { + "shape_id": case_id, + "operator": operator, + "M": M, + "N": n, + "K": k, + "config": CONFIG, + } + if mode in ("eager", "both"): + result["eager"] = measure( + fn, warmups, samples, launches_per_sample + ) + if mode in ("graph", "both"): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + torch.cuda.synchronize() + result["graph"] = measure( + graph.replay, warmups, samples, launches_per_sample + ) + del graph + + del a, b, a_scale, b_scale, out + return result + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--utils", default=UTILS_PATH) + parser.add_argument("--warmups", type=int, default=10) + parser.add_argument("--samples", type=int, default=20) + parser.add_argument("--launches-per-sample", type=int, default=5) + parser.add_argument("--mode", choices=["eager", "graph", "both"], + default="graph") + parser.add_argument("--json", type=str, default="") + args = parser.parse_args() + + utils = load_int8_utils(args.utils) + print( + "BENCHMARK triton=int8_utils.matmul_kernel output=preallocated " + f"quantization=excluded allocation=excluded epilogue=included " + f"M={M} config={json.dumps(CONFIG)} mode={args.mode}", + flush=True, + ) + print( + f"{'shape_id':42s} {'N':5s} {'K':5s} {'graph_us':>10s} " + f"{'p90_us':>10s} {'tops':>9s}", + flush=True, + ) + results = [] + for case_id, operator, k, n in TP8_CASES: + res = run_case( + utils, + case_id, + operator, + k, + n, + args.warmups, + args.samples, + args.launches_per_sample, + args.mode, + ) + results.append(res) + graph_us = (res.get("graph") or res.get("eager") or {}).get("median_us") + p90_us = (res.get("graph") or res.get("eager") or {}).get("p90_us") + tops = ( + 2.0 * M * n * k / (graph_us * 1e-6) / 1e12 + if graph_us else None + ) + print( + f"{case_id:42s} {n:5d} {k:5d} " + f"{graph_us:>10.3f} {p90_us:>10.3f} " + f"{(f'{tops:.3f}' if tops else '—'):>9s}", + flush=True, + ) + if args.json: + with open(args.json, "w", encoding="utf-8") as fh: + json.dump(results, fh, indent=2) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/int8_utils.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/int8_utils.py new file mode 100644 index 00000000..11886719 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/int8_utils.py @@ -0,0 +1,602 @@ +import functools +import json +import logging +import os +from typing import Any, Dict, List, Optional, Tuple +import json + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice +import lmslim.envs as lsenvs + +use_lightop = lsenvs.LMSLIM_USE_LIGHTOP +device_name = lsenvs.LMSLIM_GPU_NAME + +@triton.jit +def _per_token_quant_int8( + x_ptr, + xq_ptr, + scale_ptr, + stride_x, + stride_xq, + N, + BLOCK: tl.constexpr, +): + # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282 + row_id = tl.program_id(0) + + cols = tl.arange(0, BLOCK) + mask = cols < N + + x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask, + other=0.0).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10) + scale_x = absmax / 127 + x_q = x * (127 / absmax) + x_q = libdevice.nearbyint(x_q).to(tl.int8) + + tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask) + tl.store(scale_ptr + row_id, scale_x) + + +def per_token_quant_int8(x): + M = x.numel() // x.shape[-1] + N = x.shape[-1] + x_q = torch.empty_like(x, device=x.device, dtype=torch.int8) + scales = torch.empty(x.shape[:-1] + (1, ), + device=x.device, + dtype=torch.float32) + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + + #assert x.is_contiguous() + _per_token_quant_int8[(M, )]( + x, + x_q, + scales, + stride_x=x.stride(-2), + stride_xq=x_q.stride(-2), + N=N, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=1, + ) + + return x_q, scales + + +@triton.jit +def _per_token_group_quant_int8( + # Pointers to inputs and output + y_ptr, + y_q_ptr, + y_s_ptr, + # Stride of input + group_size, + # M, + # K, + # # Collums of input + # N, + SIZE, + # Avoid to divide zero + eps, + # Information for int8 + int8_min, + int8_max, + # Meta-parameters + BLOCK: tl.constexpr, + s_num : tl.constexpr, +): + """A Triton-accelerated function to perform per-token-group + quantization on a tensor. + + This function converts the tensor values into int8 values. + """ + # Map the program id to the row of X and Y it should compute. + g_id = tl.program_id(0) + y_ptr += g_id * BLOCK + y_q_ptr += g_id * BLOCK + y_s_ptr += g_id * s_num + + cols = tl.arange(0, BLOCK) # N <= BLOCK + s_cols = tl.arange(0, s_num) + mask = g_id * BLOCK + cols < SIZE + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + y = tl.reshape(y, (s_num, 128)) + # Quant + _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps) + y_s = (_absmax / int8_max).reshape(s_num, 1) + y_q = tl.clamp(y / y_s, int8_min, int8_max).to(y_q_ptr.dtype.element_ty) + + y_q = tl.reshape(y_q, (s_num*128)) + y_s = tl.reshape(y_s, (s_num)) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr + s_cols, y_s.to(y_s_ptr.dtype.element_ty)) + + +def per_token_group_quant_int8( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: torch.dtype = torch.int8, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Function to perform per-token-group quantization on an input tensor `x`. + + It converts the tensor values into signed int8 values and returns the + quantized tensor along with the scaling factor used for quantization. + + Args: + x: The input tenosr with ndim >= 2. + group_size: The group size used for quantization. + eps: The minimum to avoid dividing zero. + dtype: The dype of output tensor. Note that only `torch.int8` + is supported for now. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the + scaling factor for quantization. + """ + assert (x.shape[-1] % group_size == 0 + ), "the last dimension of `x` cannot be divisible by `group_size`" + #assert x.is_contiguous(), "`x` is not contiguous" + + iinfo = torch.iinfo(dtype) + int8_max = iinfo.max + int8_min = iinfo.min + + x_q = torch.empty_like(x, device=x.device, dtype=dtype) + N = group_size + + m = x.shape[0] + if m<=16: + config={"BLOCK":128,"s_num":1,"num_warps":1,"num_stages":1} + elif m<=256: + config={"BLOCK":1024,"s_num":8,"num_warps":4,"num_stages":1} + else: + config={"BLOCK":2048,"s_num":16,"num_warps":4,"num_stages":2} + + grid = lambda META: ( + triton.cdiv(x.numel(), META['BLOCK']), + ) + + x_s = torch.empty( + x.shape[:-1] + (x.shape[-1] // group_size, ), + device=x.device, + dtype=torch.float32, + ) + + BLOCK = triton.next_power_of_2(N) + # heuristics for number of warps + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + _per_token_group_quant_int8[grid]( + x, + x_q, + x_s, + group_size, + # M, + # K, + # N, + x.numel(), + eps, + int8_min=int8_min, + int8_max=int8_max, + **config + ) + + return x_q, x_s + + +@triton.jit +def _w8a8_block_int8_matmul( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, + and store the result in output tensor `C`. + """ + + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + + # offs_bsn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_bsn = pid_n * BLOCK_SIZE_N // group_n + + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # a_ptrs = A + (offs_am[:, None] * stride_am) + # b_ptrs = B + (offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + # offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + a = tl.load(a_ptrs, + mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, + other=0.0) + b = tl.load(b_ptrs, + mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, + other=0.0) + + + accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :] + + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +def matmul_block_int8( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype = torch.float16, + bias:Optional[torch.Tensor]=None, + config:Optional[dict] = None +) -> torch.Tensor: + """This function performs matrix multiplication with block-wise + quantization. + + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should be + 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + + Returns: + torch.Tensor: The result of matmul. + """ + + block_n, block_k = block_size[0], block_size[1] + M = A.numel() // A.shape[-1] + N, K = B.shape + + C_shape = A.shape[:-1] + (N, ) + C = A.new_empty(C_shape, dtype=output_dtype) + + if config is None: + if M<=64: + config = { + "BLOCK_SIZE_M": 16, #64 + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_stages": 0, + } + elif M<128: + config = { + "BLOCK_SIZE_M": 32, #64 + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_stages": 0, + } + elif M<=256: + config = { + "BLOCK_SIZE_M": 64, #64 + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 2, + "num_warps": 4, + "num_stages": 0, + } + else : + config = { + "BLOCK_SIZE_M": 64, #64 + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 0, + } + + def grid(META): + return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * + triton.cdiv(N, META["BLOCK_SIZE_N"]), ) + + + _w8a8_block_int8_matmul[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + **config, + ) + + if bias is not None: + C = C + bias + + return C + +def apply_w8a8_block_int8_linear( + input: torch.Tensor, + weight: torch.Tensor, + block_size: List[int], + weight_scale: torch.Tensor, + input_scale: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + config: Optional[dict] = None, +) -> torch.Tensor: + assert input_scale is None + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[0]] + + q_input, x_scale = per_token_group_quant_int8(input_2d, block_size[1]) + + output = matmul_block_int8( + q_input, weight, x_scale, weight_scale, block_size, + output_dtype=input.dtype,config=config, + ) + + if bias is not None: + output = output + bias + return output.to(dtype=input.dtype).view(*output_shape) + + +def input_to_int8( + x: torch.Tensor, dtype: torch.dtype = torch.int8 +) -> Tuple[torch.Tensor, torch.Tensor]: + """This function quantizes input values to + int8 values with tensor-wise quantization. + """ + iinfo = torch.iinfo(dtype) + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) + int8_min, int8_max = iinfo.min, iinfo.max + scale = int8_max / amax + x_scl_sat = (x * scale).clamp(min=int8_min, max=int8_max) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + + +def block_dequant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: List[int], +) -> torch.Tensor: + """This function conducts block-wise dequantization. + The inputs are block-wise quantization tensor `x_q_block`, + block-wise quantization scale and the block size. + The outputs are dequantized tensor. + """ + block_n, block_k = block_size[0], block_size[1] + n, k = x_q_block.shape + n_tiles = (n + block_n - 1) // block_n + k_tiles = (k + block_k - 1) // block_k + assert n_tiles == x_s.shape[0] + assert k_tiles == x_s.shape[1] + + x_dq_block = x_q_block.to(torch.float32) + + for i in range(k_tiles): + for j in range(n_tiles): + x_dq_block[ + j * block_n : min((j + 1) * block_n, n), + i * block_k : min((i + 1) * block_k, k), + ] *= x_s[j][i] + + return x_dq_block + +@triton.jit +def matmul_kernel( + # Pointers to matrices + a_ptr, as_ptr, b_ptr, bs_ptr, c_ptr, + # Matrix dimensions + M, N, K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` + # by to get the element one row down (A has M rows). + stride_am, stride_ak, + stride_asm, + stride_bk, stride_bn, + stride_bsn, + stride_cm, stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, SPLIT_K: tl.constexpr, +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + # See above `L2 Cache Optimizations` section for details. + pid = tl.program_id(axis=0) + pid_sp_k = tl.program_id(axis=1) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + # See above `Pointer Arithmetics` section for details + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = pid_sp_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + as_ptrs = as_ptr + offs_am * stride_asm + bs_ptrs = bs_ptr + offs_bn * stride_bsn + a_scale = tl.load(as_ptrs, mask=offs_am < M, other=0.0) + b_scale = tl.load(bs_ptrs, mask=offs_bn < N, other=0.0) + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.int32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K * SPLIT_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K * SPLIT_K, other=0.0) + # We accumulate along the K dimension. + accumulator += tl.dot(a, b) + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak + b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk + # You can fuse arbitrary activation functions here + # while the accumulator is still in FP32! + if SPLIT_K == 1: + c = (accumulator.to(tl.float32) * a_scale[:, None] * b_scale[None, :]).to(c_ptr.dtype.element_ty) + else: + c = (accumulator.to(tl.float32) * a_scale[:, None] * b_scale[None, :])#.to(c_ptr.dtype.element_ty) + # ----------------------------------------------------------- + # Write back the block of the output matrix C with masks. + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + if SPLIT_K == 1: + tl.store(c_ptrs, c, mask=c_mask) + else: + tl.atomic_add(c_ptrs, c, mask=c_mask) + +def matmul_int8(a, a_scale, b, b_scale,out_dtype,config): + # Check constraints. + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + M, K = a.shape + K, N = b.shape + + if config is None: + #print("no config") + if M<=32: + config= {'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4,'SPLIT_K': 1,'num_stages':0, 'num_warps':4} + elif M<=64: + config= {'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4,'SPLIT_K': 1,'num_stages':0, 'num_warps':4} + elif M<=128: + config= {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4,'SPLIT_K': 1,'num_stages':0, 'num_warps':4} + elif M<=1024: + config= {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4,'SPLIT_K': 1,'num_stages':0, 'num_warps':8} + else: + config= {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8,'SPLIT_K': 1,'num_stages':0, 'num_warps':8} + if config['SPLIT_K']==1: + c = torch.zeros((M, N), device=a.device, dtype=out_dtype) + else: + c = torch.zeros((M, N), device=a.device, dtype=torch.float32) + + grid = lambda META: ( + triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), + META['SPLIT_K'], + ) + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + + + matmul_kernel[grid]( + a, a_scale, b, b_scale, c, + M, N, K, + a.stride(0), a.stride(1), + a_scale.stride(0), + b.stride(0), b.stride(1), + b_scale.stride(0), + c.stride(0), c.stride(1), + **config + ) + + if config['SPLIT_K']!=1: + c=c.to(out_dtype) + return c + + diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m4096_graph.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m4096_graph.json new file mode 100644 index 00000000..ccbf2f8d --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m4096_graph.json @@ -0,0 +1,662 @@ +[ + { + "shape_id": "hy3_tp8_qkv_proj_m4096", + "operator": "qkv_proj", + "M": 4096, + "N": 1280, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 19381.382751464844, + "p90_us": 20706.951904296875, + "min_us": 15527.166748046875, + "max_us": 33025.71716308594, + "samples_us": [ + 33025.717, + 19930.211, + 15664.891, + 17588.31, + 18619.633, + 18823.849, + 21526.091, + 15859.666, + 19238.075, + 20260.184, + 15527.167, + 19733.803, + 19524.69, + 19786.154, + 20706.952, + 17930.505, + 20039.809, + 19157.535, + 16103.629, + 19555.281 + ] + } + }, + { + "shape_id": "hy3_tp8_o_proj_m4096", + "operator": "o_proj", + "M": 4096, + "N": 4096, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 14151.79214477539, + "p90_us": 15978.352355957031, + "min_us": 10805.253601074219, + "max_us": 17007.850646972656, + "samples_us": [ + 14512.738, + 17007.851, + 11326.995, + 10805.254, + 13926.711, + 14670.845, + 14319.626, + 14161.777, + 13924.182, + 12676.867, + 13472.328, + 11711.237, + 14111.505, + 14375.624, + 14141.808, + 16150.25, + 14195.662, + 14111.92, + 15978.352, + 14732.059 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 384, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 16712.230682373047, + "p90_us": 17514.361572265625, + "min_us": 15547.166442871094, + "max_us": 18294.686889648438, + "samples_us": [ + 16706.294, + 18294.687, + 16718.167, + 17544.826, + 17514.362, + 16727.061, + 15547.166, + 16738.069, + 16576.123, + 16353.539, + 16856.978, + 16953.551, + 16529.98, + 16421.12, + 16027.951, + 16203.271, + 16812.18, + 16739.862, + 16485.535, + 16286.533 + ] + } + }, + { + "shape_id": "hy3_tp8_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 4096, + "K": 192, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 10762.615585327148, + "p90_us": 12987.09716796875, + "min_us": 9852.87094116211, + "max_us": 13697.439575195312, + "samples_us": [ + 13056.47, + 10556.271, + 11678.119, + 9955.428, + 10320.246, + 10068.512, + 13697.44, + 11907.071, + 10484.849, + 10555.374, + 10817.062, + 12246.387, + 10072.864, + 10708.17, + 11191.416, + 12355.088, + 12987.097, + 10983.04, + 9852.871, + 9935.653 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_m4096", + "operator": "qkv_proj", + "M": 4096, + "N": 1280, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 41000.61798095703, + "p90_us": 41881.70166015625, + "min_us": 40260.833740234375, + "max_us": 42669.35729980469, + "samples_us": [ + 40708.112, + 41134.851, + 41381.177, + 41881.702, + 41801.389, + 42043.939, + 42669.357, + 40753.168, + 40260.834, + 41577.524, + 40839.377, + 40447.772, + 40937.708, + 41451.099, + 40329.794, + 40462.268, + 41187.875, + 40845.648, + 40568.027, + 41063.528 + ] + } + }, + { + "shape_id": "minimax_tp8_qkv_proj_and_indexer_qk_m4096", + "operator": "qkv_proj_and_indexer_qk", + "M": 4096, + "N": 1536, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 35604.698181152344, + "p90_us": 45884.57946777344, + "min_us": 27749.08447265625, + "max_us": 46322.772216796875, + "samples_us": [ + 46322.772, + 45884.579, + 45077.631, + 45431.604, + 45714.41, + 44870.535, + 45646.893, + 46203.641, + 36398.19, + 27957.782, + 32282.748, + 31533.688, + 35115.964, + 29992.844, + 32868.619, + 27749.084, + 32278.049, + 31587.637, + 36093.433, + 28834.07 + ] + } + }, + { + "shape_id": "minimax_tp8_o_proj_m4096", + "operator": "o_proj", + "M": 4096, + "N": 6144, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 19897.51739501953, + "p90_us": 20350.270080566406, + "min_us": 16113.107299804688, + "max_us": 22143.16864013672, + "samples_us": [ + 18625.305, + 16587.906, + 19008.813, + 19965.579, + 19805.585, + 20406.717, + 20196.259, + 20116.774, + 22143.169, + 19389.856, + 17160.941, + 20140.356, + 19829.456, + 20350.27, + 20321.152, + 19153.064, + 16113.107, + 18237.416, + 20135.303, + 20016.586 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 768, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 31882.257080078125, + "p90_us": 32810.479736328125, + "min_us": 28815.545654296875, + "max_us": 33757.51953125, + "samples_us": [ + 28815.546, + 32810.48, + 33757.52, + 33088.87, + 31372.415, + 32189.896, + 31768.692, + 31452.863, + 31608.185, + 32180.679, + 31995.822, + 31029.648, + 32477.213, + 32050.54, + 31519.229, + 31712.759, + 32021.133, + 31480.319, + 31442.175, + 32122.824 + ] + } + }, + { + "shape_id": "minimax_tp8_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 6144, + "K": 384, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 22389.16244506836, + "p90_us": 23217.950439453125, + "min_us": 19888.59405517578, + "max_us": 26559.368896484375, + "samples_us": [ + 26559.369, + 23217.95, + 23532.243, + 22366.364, + 22783.244, + 22659.154, + 22893.546, + 20970.54, + 22374.107, + 21844.141, + 21321.919, + 21402.269, + 22404.218, + 19888.594, + 23014.03, + 22565.503, + 20996.107, + 22041.0, + 22622.835, + 20145.161 + ] + } + }, + { + "shape_id": "glm52_tp8_fused_qkv_a_proj_m4096", + "operator": "fused_qkv_a_proj", + "M": 4096, + "N": 2624, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 68661.41662597656, + "p90_us": 70210.54077148438, + "min_us": 47223.052978515625, + "max_us": 70451.49536132812, + "samples_us": [ + 69346.32, + 68621.863, + 68886.328, + 70219.086, + 68263.092, + 70210.541, + 70451.495, + 68671.973, + 69810.4, + 69454.895, + 69113.562, + 68650.861, + 69821.857, + 58982.715, + 47223.053, + 49626.874, + 51567.096, + 52752.594, + 53318.494, + 49815.668 + ] + } + }, + { + "shape_id": "glm52_tp8_q_b_proj_m4096", + "operator": "q_b_proj", + "M": 4096, + "N": 2048, + "K": 2048, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 14385.5712890625, + "p90_us": 15009.501647949219, + "min_us": 11206.52847290039, + "max_us": 16653.09295654297, + "samples_us": [ + 14075.645, + 12226.59, + 11618.449, + 14454.031, + 14636.586, + 14437.842, + 14322.932, + 14592.075, + 16653.093, + 15187.801, + 11206.528, + 14112.669, + 14517.839, + 15009.502, + 14548.204, + 12725.101, + 11365.851, + 13823.718, + 14589.165, + 14333.301 + ] + } + }, + { + "shape_id": "glm52_tp8_kv_b_proj_m4096", + "operator": "kv_b_proj", + "M": 4096, + "N": 3584, + "K": 512, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 6939.938735961914, + "p90_us": 7771.638488769531, + "min_us": 5538.338470458984, + "max_us": 8975.7568359375, + "samples_us": [ + 6990.929, + 8975.757, + 8890.351, + 5578.849, + 5567.842, + 5553.794, + 7100.013, + 7378.083, + 7270.152, + 7218.505, + 6882.964, + 7771.638, + 6567.775, + 6888.949, + 5586.017, + 5567.042, + 5538.338, + 6640.636, + 7402.626, + 7460.033 + ] + } + }, + { + "shape_id": "glm52_tp8_o_proj_m4096", + "operator": "o_proj", + "M": 4096, + "N": 6144, + "K": 2048, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 54775.79650878906, + "p90_us": 56661.138916015625, + "min_us": 32334.09423828125, + "max_us": 57705.450439453125, + "samples_us": [ + 36803.061, + 38816.177, + 32334.094, + 36348.23, + 36142.572, + 57260.956, + 56661.139, + 57705.45, + 54939.215, + 56308.258, + 55439.417, + 54385.504, + 55639.667, + 55714.325, + 55518.042, + 53696.216, + 54612.378, + 54527.069, + 54463.104, + 55575.513 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 512, + "K": 6144, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 21254.589080810547, + "p90_us": 21705.372619628906, + "min_us": 20394.827270507812, + "max_us": 23713.912963867188, + "samples_us": [ + 20852.121, + 20782.909, + 20957.654, + 21016.469, + 21695.837, + 20957.654, + 20914.935, + 21597.601, + 22041.522, + 23713.913, + 21585.057, + 21181.616, + 21223.471, + 21285.707, + 21705.373, + 21578.625, + 21093.587, + 20394.827, + 21346.121, + 21464.134 + ] + } + }, + { + "shape_id": "glm52_tp8_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 6144, + "K": 256, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "graph": { + "median_us": 17700.262451171875, + "p90_us": 19782.623291015625, + "min_us": 15626.028442382812, + "max_us": 22184.974670410156, + "samples_us": [ + 17591.882, + 19158.612, + 19004.921, + 16231.48, + 17128.282, + 19256.914, + 19557.191, + 19782.623, + 22184.975, + 19914.011, + 17808.643, + 18663.495, + 15845.126, + 17278.868, + 17352.115, + 17899.68, + 16849.059, + 15626.028, + 17149.53, + 17040.669 + ] + } + } +] \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_eager.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_eager.json new file mode 100644 index 00000000..48a063d7 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_eager.json @@ -0,0 +1,281 @@ +{ + "scope": "gemm_out", + "triton_source": "/tmp/baseline_int8_utils.py", + "M": 4096, + "warmups": 10, + "samples": 20, + "launches_per_sample": 5, + "mode": "eager", + "results": [ + { + "shape_id": "tp4_wqkv_a_m4096", + "operator": "wqkv_a", + "M": 4096, + "N": 1536, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 13287.435150146484, + "p90_us": 13315.707397460938, + "min_us": 13170.518493652344, + "max_us": 13707.110595703125, + "samples_us": [ + 13226.106, + 13308.22, + 13230.266, + 13292.188, + 13315.707, + 13190.199, + 13218.904, + 13296.187, + 13313.756, + 13294.81, + 13282.683, + 13296.956, + 13353.532, + 13170.518, + 13203.448, + 13707.111, + 13215.833, + 13229.721, + 13219.096, + 13292.923 + ] + } + }, + { + "shape_id": "tp4_wq_b_m4096", + "operator": "wq_b", + "M": 4096, + "N": 8192, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 20623.723602294922, + "p90_us": 20699.789428710938, + "min_us": 20575.30517578125, + "max_us": 21157.14569091797, + "samples_us": [ + 20627.628, + 20606.154, + 20615.851, + 20721.645, + 20621.803, + 20619.339, + 20625.644, + 21157.146, + 20654.027, + 20667.628, + 20578.441, + 20637.163, + 20699.789, + 20575.305, + 20658.06, + 20611.241, + 20650.603, + 20599.562, + 20597.418, + 20594.986 + ] + } + }, + { + "shape_id": "tp4_indexer_wq_b_m4096", + "operator": "indexer.wq_b", + "M": 4096, + "N": 8192, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 20654.634857177734, + "p90_us": 20685.16387939453, + "min_us": 20580.072021484375, + "max_us": 21134.840393066406, + "samples_us": [ + 20639.98, + 20622.218, + 20637.993, + 20662.732, + 20598.09, + 20663.66, + 20658.06, + 20632.042, + 20633.29, + 20651.21, + 20665.291, + 21134.84, + 20685.164, + 20580.072, + 20664.265, + 20590.665, + 20696.204, + 20669.418, + 20670.186, + 20642.058 + ] + } + }, + { + "shape_id": "tp4_wo_b_m4096", + "operator": "wo_b", + "M": 4096, + "N": 4096, + "K": 2048, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 19915.618896484375, + "p90_us": 19963.380432128906, + "min_us": 19849.97100830078, + "max_us": 20379.873657226562, + "samples_us": [ + 19945.428, + 19849.971, + 20379.874, + 19900.435, + 19941.397, + 19923.093, + 19873.17, + 19899.539, + 19913.683, + 19905.396, + 19963.38, + 19887.508, + 19876.372, + 19917.555, + 19924.883, + 20352.672, + 19895.027, + 19882.227, + 19918.611, + 19960.02 + ] + } + }, + { + "shape_id": "tp4_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 1024, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 8774.919891357422, + "p90_us": 8838.265228271484, + "min_us": 8717.046356201172, + "max_us": 9269.158172607422, + "samples_us": [ + 8721.815, + 8792.313, + 8838.265, + 8745.751, + 8775.352, + 8793.529, + 8774.488, + 8747.095, + 8717.046, + 8845.21, + 8749.047, + 8780.152, + 8805.529, + 8770.008, + 8790.744, + 8806.841, + 8736.246, + 8773.912, + 9269.158, + 8771.128 + ] + } + }, + { + "shape_id": "tp4_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 4096, + "K": 512, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "eager": { + "median_us": 5582.3015213012695, + "p90_us": 5592.509841918945, + "min_us": 5555.356979370117, + "max_us": 5619.230270385742, + "samples_us": [ + 5585.566, + 5592.51, + 5555.357, + 5586.398, + 5571.326, + 5581.853, + 5582.75, + 5586.045, + 5586.43, + 5582.973, + 5576.606, + 5571.358, + 5567.805, + 5563.549, + 5576.189, + 5619.23, + 5587.838, + 5566.749, + 5568.093, + 5595.134 + ] + } + } + ] +} \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_graph.json b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_graph.json new file mode 100644 index 00000000..a6464631 --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/triton_tp4_m4096_graph.json @@ -0,0 +1,281 @@ +{ + "scope": "gemm_out", + "triton_source": "/tmp/baseline_int8_utils.py", + "M": 4096, + "warmups": 10, + "samples": 20, + "launches_per_sample": 5, + "mode": "graph", + "results": [ + { + "shape_id": "tp4_wqkv_a_m4096", + "operator": "wqkv_a", + "M": 4096, + "N": 1536, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 13247.252655029297, + "p90_us": 13320.054626464844, + "min_us": 13144.914245605469, + "max_us": 13643.904113769531, + "samples_us": [ + 13336.472, + 13262.518, + 13314.262, + 13197.523, + 13320.055, + 13643.904, + 13300.246, + 13269.365, + 13284.662, + 13253.653, + 13234.421, + 13224.02, + 13218.292, + 13228.885, + 13147.346, + 13240.852, + 13214.676, + 13229.396, + 13144.914, + 13285.077 + ] + } + }, + { + "shape_id": "tp4_wq_b_m4096", + "operator": "wq_b", + "M": 4096, + "N": 8192, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 20590.225219726562, + "p90_us": 20648.818969726562, + "min_us": 20520.016479492188, + "max_us": 21081.182861328125, + "samples_us": [ + 20596.434, + 21081.183, + 20520.016, + 20586.673, + 20577.713, + 20597.105, + 20556.367, + 20648.819, + 20589.33, + 20599.345, + 20590.61, + 20578.609, + 20576.752, + 20544.048, + 21078.076, + 20597.136, + 20589.841, + 20524.046, + 20621.681, + 20624.913 + ] + } + }, + { + "shape_id": "tp4_indexer_wq_b_m4096", + "operator": "indexer.wq_b", + "M": 4096, + "N": 8192, + "K": 1024, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 20602.671813964844, + "p90_us": 20697.555541992188, + "min_us": 20566.477966308594, + "max_us": 21020.53985595703, + "samples_us": [ + 20609.616, + 20639.314, + 20574.065, + 20697.556, + 20587.505, + 21020.54, + 20601.36, + 20644.786, + 20570.096, + 20603.824, + 20603.375, + 20586.96, + 20610.159, + 20601.968, + 20566.478, + 20597.328, + 20613.904, + 20600.912, + 20981.018, + 20566.574 + ] + } + }, + { + "shape_id": "tp4_wo_b_m4096", + "operator": "wo_b", + "M": 4096, + "N": 4096, + "K": 2048, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 19881.948852539062, + "p90_us": 19916.25213623047, + "min_us": 19814.553833007812, + "max_us": 20356.455993652344, + "samples_us": [ + 19870.557, + 19886.205, + 19901.534, + 19921.021, + 19859.677, + 19874.237, + 19911.804, + 19833.531, + 19829.051, + 20356.456, + 19863.837, + 19898.973, + 19877.693, + 19857.532, + 19838.747, + 19896.861, + 19889.693, + 19896.861, + 19814.554, + 19916.252 + ] + } + }, + { + "shape_id": "tp4_shared_gate_up_proj_m4096", + "operator": "shared_gate_up_proj", + "M": 4096, + "N": 1024, + "K": 4096, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 8790.192031860352, + "p90_us": 8815.552520751953, + "min_us": 8717.342376708984, + "max_us": 9243.083953857422, + "samples_us": [ + 8815.553, + 8757.567, + 8770.112, + 9243.084, + 8726.078, + 8750.815, + 8791.904, + 8795.233, + 8717.342, + 8759.872, + 8832.194, + 8806.433, + 8799.168, + 8779.68, + 8723.199, + 8802.465, + 8798.4, + 8788.48, + 8806.753, + 8740.222 + ] + } + }, + { + "shape_id": "tp4_shared_down_proj_m4096", + "operator": "shared_down_proj", + "M": 4096, + "N": 4096, + "K": 512, + "config": { + "BLOCK_SIZE_M": 256, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "SPLIT_K": 1, + "num_stages": 0, + "num_warps": 8 + }, + "errors": null, + "graph": { + "median_us": 5545.821189880371, + "p90_us": 5561.069488525391, + "min_us": 5528.492736816406, + "max_us": 6035.929489135742, + "samples_us": [ + 5570.702, + 5552.142, + 5551.213, + 5546.317, + 5534.701, + 5537.453, + 5546.029, + 5541.453, + 5545.197, + 5557.133, + 5537.933, + 5528.493, + 5561.069, + 5544.845, + 5552.333, + 5551.566, + 5545.613, + 5533.453, + 5544.685, + 6035.929 + ] + } + } + ] +} \ No newline at end of file diff --git a/metainfer/tasks/dcu_kernel_auto_opt/tools/planner_parity.py b/metainfer/tasks/dcu_kernel_auto_opt/tools/planner_parity.py new file mode 100644 index 00000000..046255fd --- /dev/null +++ b/metainfer/tasks/dcu_kernel_auto_opt/tools/planner_parity.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python3 +"""Offline parity analysis: planner (v0) vs the legacy round menu. + +Read-only tool over historical task repos under ``kernel-repos`` (and optionally +workspace ``iterations``). For each recorded round it reconstructs the state +(prior same-worker/same-shape records + pmc evidence), reproduces what the +legacy menu (``w8a8_round_strategy``) would mandate, asks what the +state-conditioned ``planner`` would pick, maps both to coarse plan families, and +reports agreement/divergence by M regime and iteration bucket, plus samples. + +Caveats (read the output with these in mind): +- State reconstruction is approximate: pmc evidence is frequently absent in + records, so the planner often falls back to its P4 legacy approximation. +- The menu text -> family mapping is a keyword heuristic. +- This tool never writes into the repos; it only prints / dumps a report. + +Usage: + python3 tools/planner_parity.py [--kernel-repos ROOT] [--limit-repos N] + [--out /tmp/planner_parity_report.json] +""" + +from __future__ import annotations + +import argparse +import collections +import json +import re +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +# Make the plugin importable when run as a script from the plugin tree. +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) # .../dcu_kernel_auto_opt +sys.path.insert(0, str(Path(__file__).resolve().parents[4])) # MetaInfer/ + +from metainfer.tasks.dcu_kernel_auto_opt.orchestrator import planner as planner_mod # noqa: E402 +from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.planner import choose_plan_from_history # noqa: E402 +from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.prompts import w8a8_round_strategy # noqa: E402 +from metainfer.tasks.dcu_kernel_auto_opt.orchestrator.w8a8_pipeline import isa_round_policy # noqa: E402 + +_MENU_FAMILY_RULES: List[Tuple[str, str]] = [ + # Explicit markers first so generic words inside round prose do not win. + ("faster but incorrect", "repair"), + ("faster_wrong", "repair"), + ("repair", "repair"), + ("conditional inline-asm", "inline_asm"), + ("inline asm", "inline_asm"), + ("ISA-guided", "isa_guided"), + ("ISA", "isa"), + ("split-K", "grid_splitk"), + ("split_k", "grid_splitk"), + ("consolidat", "consolidate"), + ("epilogue", "epilogue"), + ("Architecture round", "architecture"), + ("Architecture/pipeline round", "architecture"), + ("macro-tile", "architecture"), + ("tile-shape", "architecture"), + ("geometry", "architecture"), + ("tile", "architecture"), + ("architecture", "architecture"), + ("packing", "memory_layout"), + ("packed", "memory_layout"), + ("swizzle", "memory_layout"), + ("bank", "memory_layout"), + ("layout", "memory_layout"), + ("double buff", "pipeline"), + ("prefetch", "pipeline"), + ("staging", "pipeline"), + ("barrier", "pipeline"), + ("pipeline", "pipeline"), + ("occupancy", "occupancy"), + ("waves per block", "occupancy"), + ("VGPR", "occupancy"), + ("register", "occupancy"), + ("bootstrap", "bootstrap"), + ("scalar", "bootstrap"), +] + +_PLAN_TO_FAMILY = { + "repair_faster_wrong": "repair", + "fix_build": "repair", + "retry_same": "repair", + "bootstrap_correctness": "bootstrap", + "establish_arch": "architecture", + "architecture_explore": "architecture", + "grid_splitk": "grid_splitk", + "pipeline_tune": "pipeline", + "memory_layout": "memory_layout", + "occupancy_resource": "occupancy", + "epilogue_fusion": "epilogue", + "isa_guided_hip": "isa_guided", + "conditional_inline_asm": "inline_asm", + "consolidate": "consolidate", +} + + +def menu_family(text: str) -> str: + lowered = text.lower() + for needle, family in _MENU_FAMILY_RULES: + if needle in lowered: + return family + return "unknown" + + +def iter_number_from_path(path: Path) -> int: + m = re.search(r"iteration(\d+)", path.name) + if m: + return int(m.group(1)) + # parent dirs like .../iteration13/iteration.json + m = re.search(r"iteration(\d+)", path.parent.name) + return int(m.group(1)) if m else -1 + + +def load_records(root: Path, limit_repos: Optional[int]) -> List[Tuple[str, Dict[str, Any]]]: + """Return (repo_dir_name, record) for every candidate iteration.json.""" + records: List[Tuple[str, Dict[str, Any]]] = [] + repos = sorted(p for p in root.iterdir() if p.is_dir()) + if limit_repos: + repos = repos[:limit_repos] + for repo in repos: + for path in sorted( + p for p in repo.rglob("iteration.json") + if "candidates" in p.parts + ): + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + continue + if isinstance(data, dict): + records.append((repo.name, data)) + return records + + +def analyze(root: Path, limit_repos: Optional[int]) -> Dict[str, Any]: + records = load_records(root, limit_repos) + chains: Dict[Tuple[str, str, str], List[Dict[str, Any]]] = collections.defaultdict(list) + for repo, rec in records: + worker = str(rec.get("worker_id", "?")) + shape = rec.get("shape_id", "?") + iteration = iter_number_from_path(Path(str(rec.get("record_path", "iteration0")))) if False else int(rec.get("iteration", 0)) + chains[(repo, worker, shape)].append( + {"record": rec, "iteration": iteration} + ) + + stats = collections.Counter() + buckets: Dict[str, collections.Counter] = collections.defaultdict(collections.Counter) + samples: List[Dict[str, Any]] = [] + + for (repo, worker, shape_id), entries in chains.items(): + entries.sort(key=lambda e: e["iteration"]) + history: List[Dict[str, Any]] = [] + for entry in entries: + rec = entry["record"] + iteration = entry["iteration"] + shape = rec.get("shape", {}) or {} + m = int(shape.get("M", 0) or 0) + regime = "prefill(M>=128)" if m >= 128 else ("m16" if m >= 16 else "small(M<16)") + pmc = rec.get("pmc_evidence") or {} + + try: + policy = isa_round_policy( + iteration=iteration, + max_iterations=10, + history=history, + ) + menu_text = w8a8_round_strategy( + shape, iteration, history, pmc, + max_iterations=10, isa_policy=policy, + ) + except Exception as exc: # noqa: BLE001 + stats["menu_error"] += 1 + history.append(rec) + continue + + menu_fam = menu_family(menu_text) + try: + plan_id = choose_plan_from_history( + history, pmc, iteration=iteration, max_iterations=10, + shape=shape, + compiler_limitation_confirmed=( + policy.get("phase") == "conditional_inline_asm" + ), + ) + except Exception as exc: # noqa: BLE001 + stats["planner_error"] += 1 + history.append(rec) + continue + + plan_fam = _PLAN_TO_FAMILY.get(plan_id, "unknown") + agree = (menu_fam == plan_fam) + stats["total"] += 1 + stats["agree" if agree else "disagree"] += 1 + bucket_key = f"{regime}|iter{min(iteration, 12)}" + buckets[bucket_key]["agree" if agree else "disagree"] += 1 + + if not agree and len(samples) < 40: + samples.append({ + "repo": repo, + "shape_id": shape_id, + "iteration": iteration, + "regime": regime, + "menu_family": menu_fam, + "plan_family": plan_fam, + "plan_id": plan_id, + "hypothesis": str(rec.get("hypothesis") or "")[:200], + }) + history.append(rec) + + total = stats.get("total", 0) + return { + "corpus": {"repos_scanned": len(set(r for r, _ in records)) if records else 0, + "records": len(records)}, + "stats": dict(stats), + "agree_rate": (stats.get("agree", 0) / total) if total else None, + "buckets": { + k: dict(v) for k, v in sorted(buckets.items()) + }, + "samples": samples, + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--kernel-repos", type=Path, + default=Path(__file__).resolve().parents[2].parent.parent.parent / "kernel-repos") + parser.add_argument("--limit-repos", type=int, default=None) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + if not args.kernel_repos.is_dir(): + print(f"kernel-repos not found at {args.kernel_repos}", file=sys.stderr) + return 2 + + report = analyze(args.kernel_repos, args.limit_repos) + total = report["stats"].get("total", 0) + print(f"corpus: {report['corpus']['records']} records / {report['corpus']['repos_scanned']} repos") + print(f"agree={report['stats'].get('agree',0)} disagree={report['stats'].get('disagree',0)} " + f"agree_rate={report['agree_rate']:.1%}" if report["agree_rate"] is not None + else f"agree={report['stats'].get('agree',0)} disagree={report['stats'].get('disagree',0)}") + print("--- buckets (regime|iteration): agree/disagree ---") + for k, v in list(report["buckets"].items())[:40]: + print(f" {k}: {v}") + print("--- divergence samples ---") + for s in report["samples"][:10]: + print(f" [{s['repo']} {s['shape_id']} it{s['iteration']} {s['regime']}] " + f"menu={s['menu_family']} plan={s['plan_family']} ({s['plan_id']})") + print(f" hypo: {s['hypothesis'][:150]}") + + if args.out: + args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2), + encoding="utf-8") + print(f"report -> {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/indexer_wq_b.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/indexer_wq_b.hip index 143caa0b..c6a575d2 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/indexer_wq_b.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/indexer_wq_b.hip @@ -1,4 +1,5 @@ // @@variant shape=tp8_indexer_wq_b_m16 commit=9b5b710d0a2c4cd7ac22e2037dfe8735059eec0b added=2026-08-26 +// baseline_us=48.274 speedup=2.70443 // median_us=17.85 p90_us=17.88 // source=hy3-dsh-tp8-m16-1-adf021ab // W8A8 INT8 GEMM — gfx928 implementation (worker_1, physical GPU 1) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_down_proj.hip index 487c6b3a..2cce19a8 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_down_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_down_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=tp8_shared_down_proj_m16 commit=119c281feefbb681f88dba89897484424bb07937 added=2026-08-26 +// baseline_us=14.849 speedup=1.75665 // median_us=8.453 p90_us=9.13 // source=hy3-dsh-tp8-m16-1-adf021ab // W8A8 INT8 GEMM bootstrap kernel for Hygon K500SM_AI / gfx928. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_gate_up_proj.hip index 15e74f81..4c63671f 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_gate_up_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/shared_gate_up_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=tp8_shared_gate_up_proj_m16 commit=ea0f87b2e586e93141907abbdfc541b1585f4c50 added=2026-08-26 +// baseline_us=59.21 speedup=5.75973 // median_us=10.28 p90_us=10.33 // source=hy3-dsh-tp8-m16-1-adf021ab // W8A8 INT8 GEMM bootstrap kernel for Hygon K500SM_AI / gfx928. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wo_b.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wo_b.hip index f420f456..c67c8445 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wo_b.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wo_b.hip @@ -1,4 +1,5 @@ // @@variant shape=tp8_wo_b_m16 commit=eb0fbbad60c58b1cf09e6ceff6cb9aaa68a2c1cb added=2026-08-26 +// baseline_us=32.137 speedup=3.27594 // median_us=9.81 p90_us=9.84 // source=hy3-dsh-tp8-m16-1-adf021ab // MetaInfer INT8 W8A8 GEMM for gfx928 (worker_2 / tp8_wo_b_m16). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wq_b.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wq_b.hip index ff0598e7..0b19eb88 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wq_b.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wq_b.hip @@ -1,4 +1,5 @@ // @@variant shape=tp8_wq_b_m16 commit=bfaef2fc76d0e4a095b2ffae2b466f2863511d43 added=2026-08-26 +// baseline_us=32.137 speedup=2.19066 // median_us=14.67 p90_us=14.7 // source=hy3-dsh-tp8-m16-1-adf021ab // W8A8 INT8 GEMM — gfx928 implementation (worker_1, physical GPU 1) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wqkv_a.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wqkv_a.hip index bfa665ac..a383ad61 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wqkv_a.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/deepseek-v4/TP8/M16/wqkv_a.hip @@ -1,4 +1,5 @@ // @@variant shape=tp8_wqkv_a_m16 commit=b67e227c20b3349d46020f9f52cd2484b2142ff3 added=2026-08-26 +// baseline_us=66.591 speedup=3.8052 // median_us=17.5 p90_us=17.52 // source=hy3-dsh-tp8-m16-1-adf021ab // MetaInfer W8A8 GEMM - worker_0 HIP implementation (gfx928). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/fused_qkv_a_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/fused_qkv_a_proj.hip index a7e83daf..7748f78d 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/fused_qkv_a_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/fused_qkv_a_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_fused_qkv_a_proj_m16 commit=bf306c40c48713a81efbd7ace430c9cf21cbb429 added=2026-08-31 +// baseline_us=97.99 speedup=2.33476 // median_us=41.97 p90_us=42.09 // source=glm5-2-dsh-tp8-m16-test1-cb4d262c // INT8 W8A8 GEMM kernel for Hygon K500SM_AI / gfx928 (worker_0). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/kv_b_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/kv_b_proj.hip index 33985f73..93be3411 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/kv_b_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/kv_b_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_kv_b_proj_m16 commit=a150485d1a63d39ca9251b13d3da304013431efd added=2026-08-31 +// baseline_us=18.833 speedup=2.1881 // median_us=8.607 p90_us=9.766 // source=glm5-2-dsh-tp8-m16-test1-cb4d262c // csrc/w8a8_gemm_hip.hip diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/o_proj.hip index 3ec28b39..caab0383 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/o_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/o_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_o_proj_m16 commit=c1cfe1975c2fcb8d4bc23c86e8026afae225a4c4 added=2026-08-31 +// baseline_us=66.98 speedup=2.79083 // median_us=24 p90_us=24.02 // source=glm5-2-dsh-tp8-m16-test1-cb4d262c // MetaInfer INT8 W8A8 GEMM for Hygon K500SM_AI / gfx928. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/q_b_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/q_b_proj.hip index 8d888111..ad036b39 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/q_b_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/q_b_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_q_b_proj_m16 commit=8a9a28e409277a6f266a59e0c053a3d2bb516f12 added=2026-08-31 +// baseline_us=42.035 speedup=2.97698 // median_us=14.12 p90_us=14.14 // source=glm5-2-dsh-tp8-m16-test1-cb4d262c // csrc/w8a8_gemm_hip.hip diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_down_proj.hip index fdf64fd2..b841e219 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_down_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_down_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_shared_down_proj_m16 commit=42d49afdb24af21a5480786c88ca11a84d57f44a added=2026-08-31 +// baseline_us=18.177 speedup=2.20621 // median_us=8.239 p90_us=9.122 // source=glm5-2-dsh-tp8-m16-test1-cb4d262c // W8A8 INT8 GEMM for Hygon K500SM_AI / gfx928 (worker_3). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_gate_up_proj.hip index 3455553e..b33439e2 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_gate_up_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M16/shared_gate_up_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=glm_tp8_shared_gate_up_proj_m16 commit=fa6ea50630e35036b22141edc48f72975fca041d added=2026-08-31 +// baseline_us=99.318 speedup=9.17911 // median_us=10.82 p90_us=10.87 // source=glm5-2-dsh-tp8-m16-test1-cb4d262c // W8A8 INT8 GEMM for Hygon K500SM_AI / gfx928 (worker_3). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/fused_qkv_a_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/fused_qkv_a_proj.hip index aaa4a2a8..f5387ae1 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/fused_qkv_a_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/fused_qkv_a_proj.hip @@ -1,8 +1,8 @@ -// @@variant shape=glm_tp8_fused_qkv_a_proj_m4096 commit=935c285f78b3c36ce8f49ef13a68b0cf4e8ff4d5 added=2026-08-31 -// median_us=937.6 p90_us=939.5 speedup=73.23 baseline_us=6.866e+04 -// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 -// @@variant shape=glm_tp8_fused_qkv_a_proj_m4096 -// MetaInfer W8A8 INT8 GEMM for Hygon K500SM_AI / gfx928 (worker_0, GPU 0). +// @@variant shape=glm_tp8_fused_qkv_a_proj_m4096 commit=99b9a733eddfc7196f3d186b80affe4324351ac4 added=2026-09-11 +// median_us=883.7 p90_us=885.5 speedup=77.7 baseline_us=6.866e+04 +// source=test-9-11-1-b01955b3/iteration_001/glm_tp8_fused_qkv_a_proj_m4096 +// MetaInfer W8A8 INT8 GEMM -- round-1 architecture establishment for +// Hygon K500SM_AI / gfx928 (worker_0, physical GPU 0). // // Assigned shape: // glm_tp8_fused_qkv_a_proj_m4096 : M=4096, N=2624, K=6144 @@ -11,255 +11,446 @@ // out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) // * x_scale[m] * weight_scale[n] ) // -// Iteration 1 (correctness-first DUMMA bootstrap): -// * Packing bootstrap: launch_pack_w8a8_weight is the mandated identity -// device-to-device copy for every (K, N) (raw [K, N] row-major weight is -// copied unchanged and the fp32 scales are copied unchanged). The GEMM -// therefore interprets packed_weight as the logical [K, N] row-major -// layout everywhere. Later Parallel explore rounds may change the pack -// kernel and the matching GEMM interpretation together. -// * Iteration 1 dispatched tile <64, 64, 128> (41 x 64 = 2624 blocks, -// 16,384 B LDS/block, 4 blocks/CU, official median 6391.92 us = 20.66 -// logical TOPS). Its exact-source ISA shows the stage loop is -// LDS-issue-bound, not MMA-bound: per stage per wavefront 136 ds_read -// (mostly du_load_matrix_sync byte-reassembly ds_read_u8) + 107 waitcnt -// vs only 32 v_mmac and 2 barriers; 2624 blocks x 96 barriers = -// 251,904 barriers and 2.06 GB of staged global traffic per replay. -// * The epilogue is the coalesced fragment store: the verified gfx928 -// accumulator mapping (row = lane & 15, col_mod4 = lane >> 4, -// frag.x[i] -> column col_mod4 + 4*i) is transposed in registers (8 -// shfl_xor + 8 v_cndmask) so each lane owns four contiguous bf16 columns -// and writes ONE 8-byte store (100% store sector efficiency). The -// per-element float multiply order (float(dot) * x_scale[m] * -// weight_scale[n]) and the bf16 rounding are the reference order, so the -// stored bits are identical to the scalar path. -// * Generic arms: the <64,64,128> tiled kernel also serves every -// large-M shape with N % 64 == 0 and K % 128 == 0 (identity layout); a -// scalar int8/int32 grid-stride kernel is the fallback for every other -// (m, n, k), including the paired M=2 and M=16 API shapes with the same -// (N, K), which must never reach the tiled path (m >= 128 guard). -// * The timed operator (launch_w8a8_gemm) performs no allocation, -// compilation, autotuning, packing, host/device synchronization, or -// default-stream launch: it only dispatches kernels on the -// caller-provided HIP stream and ignores the workspace (split-K capacity -// is zero for this shape: m*n*4 = 42,990,080 B > 16 MiB budget). +// Round 1 (mandatory establish_arch) -- packed_weight architecture family: // -// Iteration 1 (mandated 2-D macro-tile baseline: benchmark 64x64 / 64x128 / -// 128x64, pick with code-object resources + occupancy): -// * Tile comparison (all three instantiations already compiled by the -// iteration-1 build; VGPR/SGPR/LDS/spills are exact code-object facts -// from profiles/.../iteration1/current-best-isa/metadata.txt; grid and -// staged traffic are arithmetic for M=4096, N=2624, K=6144): -// tile VGPR SGPR LDS B blocks/CU waves/CU grid staged GB -// <64, 64,128> 57 36 16384 4 16 2624 2.06 -// <64,128,128> 92 36 24576 2 8 1344 1.56* -// <128,64,128> 88 34 24576 2 8 1312 1.55 -// (*) 64x128 needs an N tail: N=2624 = 41*64 is not a multiple of -// 128, so the n0=2560 block would zero-fill B columns >= 2624 and an -// epilogue column guard, wasting ~2.4% of all MMAs. No split-K in any -// arm: 1312 blocks dwarf the 120 CUs. -// * Iteration 1 dispatched the exact-fit 2-D macro-tile <128, 64, 128>: -// M=4096 = 32*128 and N=2624 = 41*64 are both exact, so no tail -// predication is introduced anywhere (the generic A row guard and the -// epilogue row guard are retained for the generic path). Grid = (N/64) x -// (M/128) = 41 x 32 = 1312 blocks; 256 threads = 4 wavefronts, each wave -// owns a 64x32 quadrant = eight m16n16k32 int32 accumulator fragments; -// single-buffered 128-K LDS stage (A[128,128] + B[128,64] = 24,576 -// B/block) with TWO __syncthreads per stage, K=6144 = 48 stages. Versus -// the iteration-1 64x64 tile this halves B re-reads (B staged by 32 -// M-blocks instead of 64; staged global traffic 2.06 -> 1.55 GB per -// replay) and cuts total barriers 251,904 -> 125,952, at the cost of -// occupancy 4 -> 2 blocks/CU (16 -> 8 resident waves/CU). The int32 -// accumulation order (k0-outer over 128-K stages, kk-inner ascending, -// same m16n16k32 element-to-slot mapping) is bit-identical to the -// iteration-1 kernel and the scalar reference. -// * The <64,128,128> instantiation stays compiled as the next benchmark -// candidate (needs the N-tail handling described above before dispatch); -// the generic <64,64,128> arm and the scalar fallback are untouched. +// Bottleneck hypothesis, from the trusted PMC of the current best source +// (the correctness-first bootstrap: w8a8_dumma_prefill_tiled_kernel +// <128, 64, 128>, identity [K, N] weight, library du_load_matrix_sync +// loaders, LDS row strides 128 (A) / 64 (B), grid 41 x 32 = 1312 blocks, +// 256 threads, 2 blocks/CU): +// median 2920.88 us / p90 2962.82 us, lds_instructions 21,999,616, +// lds_bank_conflicts 274,071,552 (~12.5 per LDS instruction), +// lds_wait_instructions 50,181,819, valu_instructions 86,560,512, +// vmem_read_instructions 1,574,400, and a counter-derived operator HBM +// bandwidth of only 50.6 GB/s (147.8 MB per replay at a 90.7% L2 hit +// rate). The exact-shape stage loop is LDS-issue/conflict bound, not +// MMA- or HBM-bound: both LDS row strides alias every row onto one bank +// phase (128 mod 32 == 0, 64 mod 32 == 0), and the row_major matrix_b +// loader of du_load_matrix_sync lowers each m16n16k32 int8 B fragment to +// 8 ds_read_u8 at 16-byte spacing plus a byte-reassembly chain. // -// Iteration 2 (mandated operand-reuse round: cooperative A+B LDS staging -// vs direct loads, quantified A/B reuse, vectorized coalesced global loads, -// bank-safe LDS layout): -// * Operand-reuse accounting for the <128,64,128> macro-tile (M=4096, -// N=2624, K=6144): per block A = 128 x 6144 = 786,432 B and B = 6144 x -// 64 = 393,216 B of unique bytes; the 41 N-blocks re-read every A -// row-slab (A global traffic M*K*41 = 1.03 GB/replay, 41x amplification -// of the 25.2 MB unique A) and the 32 M-blocks re-read every B -// column-slab (K*N*32 = 0.52 GB/replay, 32x of the 16.1 MB unique B); -// total staged traffic 1.55 GB/replay. Cooperative staging additionally -// buys 4x wave-level reuse inside each block (one 16-B vectorized global -// load per byte, four waves re-read it from LDS), so it is retained: a -// direct per-wave global fragment path would issue 4x more vmem (4 waves -// x 48 x 8-B fragment loads per stage = ~9,216 vs ~288 staged -// 16-B loads per wave) and every 8-B fragment load would be a -// row-strided 16-sector access (vs 4 sectors for the coalesced 16-B -// staging) -- both directions strictly worse. -// * The exact-source ISA of the iteration-1 <128,64,128> kernel shows the -// per-stage per-wavefront consume issues 80 ds_read (64 ds_read_u8 + -// 16 ds_read2_b32) + 99 waitcnt + ~96 byte-reassembly VALU -// (v_and/v_or3/v_lshlrev) per stage against only 32 v_mmac: the -// row_major matrix_b loader of du_load_matrix_sync lowers each B -// fragment to 8 ds_read_u8 at 16-B spacing plus a byte-reassembly chain, -// and the LDS row strides 128 (A) / 64 (B) alias every row onto one -// bank phase (32r and 16r mod 32 == 0), matching the PMC -// lds_bank_conflicts = 274,071,552 (~12.5 per LDS instruction). The -// kernel is LDS-issue/conflict-bound, not MMA- or HBM-bound. -// * Focused change (one mechanism, all inside csrc/w8a8_gemm_hip.hip): -// (1) launch_pack_w8a8_weight now packs the exact (k,n) == (6144,2624) -// weight once, outside the timed region and out of Graph capture, into -// the n-major layout packed[n*K + kk] = raw[kk*N + n] (same byte count -// and buffer, captured addresses unchanged; every other (K, N) keeps -// the identity copy); -// (2) the exact-shape guard dispatches a new -// w8a8_dumma_prefill_packedb_kernel<128,64,128> (grid 41x32 = 1312 -// blocks, 256 threads = 4 waves, same single-buffered 128-K stage with -// two __syncthreads, same k0-outer/kk-inner int32 order, same epilogue) -// that stages A row-major and B n-major into bank-safe LDS rows of 136 B -// (StageK + 8: 8-byte-aligned so every fragment read is one ds_read_b64, -// 34 dwords so the 16 rows of each 16-lane ds_read_b64 phase land on 16 -// distinct bank pairs -> 4-phase conflict-free floor) and replaces both -// library loaders with the lineage-validated explicit load_frag8 (one -// aligned 8-B load per lane per fragment: x[0..7] = 8 consecutive bytes -// of row (lane&15) at k-offset ((lane>>4)*8) for row_major A and the -// same 8 k-bytes of column n0+(lane&15) for col_major B over the -// packed [N, K] tile -- identical fragment values, so the m16n16k32 -// v_mmac inputs and the exact int32 accumulation order are bit-identical -// to the iteration-1 kernel and the scalar reference); -// (3) the scalar fallback decodes the n-major layout when -// (k,n) == (6144,2624) so the paired M=2/M=16 API shapes with the same -// (N, K) stay byte-exact; the generic <64,64,128> identity arm and the -// <64,128,128>/<128,64,128> sibling instantiations are untouched. -// * Expected ISA/PMC deltas: ds_read per stage per wavefront 80 -> 24 -// (16 A + 8 B ds_read_b64), waitcnt 99 -> ~24-32, reassembly VALU ~0, -// lds_bank_conflicts 274M -> conflict-free floor (~1-3M), LDS/block -// 24,576 -> 26,112 B (2 blocks/CU unchanged; no spills). +// Focused change (ONE architecture family -- packed_weight -- entirely +// inside csrc/w8a8_gemm_hip.hip): +// (1) launch_pack_w8a8_weight packs the exact (k, n) == (6144, 2624) +// weight ONCE, outside the timed region and out of Graph capture, +// into the n-major layout packed[n * k + kk] = raw[kk * n + n] +// (same byte count, same buffer allocation, so every captured +// address is unchanged). Every other (K, N) keeps the identity +// device-to-device copy. +// (2) the exact-shape guard dispatches the new +// w8a8_dumma_prefill_packedb_kernel<128, 64, 128>. It keeps the +// EXACT bootstrap geometry -- grid (N/64) x (M/128) = 41 x 32 = 1312 +// blocks, 256 threads = 4 gfx928 wavefronts, one 64 x 32 quadrant of +// eight m16n16k32 int32 accumulator fragments per wave, the same +// single-buffered 128-K LDS stage with two __syncthreads per stage, +// the same k0-outer / kk-inner ascending int32 accumulation order and +// the same coalesced bf16 epilogue -- and changes only the operand +// path to the packed-weight architecture: +// * B is staged n-major (b_tile[n_row][k]) from the packed [N, K] +// weight, so a col_major matrix_b fragment is 8 contiguous K +// bytes per lane and every staging vector is one coalesced +// 16-byte global read; +// * both LDS tiles use the bank-safe row stride StageK + 8 = 136 +// bytes (34 dwords): 8-byte aligned with an odd 8-byte count +// (136/8 = 17), so the 16 rows of each 16-lane ds_read_b64 phase +// land on 16 distinct bank pairs (the conflict-free 4-phase +// floor), and 136 % 16 == 8 keeps the two-half int64 staging +// form; +// * both fragment loaders are the explicit load_frag8: ONE aligned +// 8-byte load per lane per fragment (row = lane & 15, k-quarter = +// lane >> 4), replacing the byte-reassembly library path with a +// single ds_read_b64 while producing identical fragment values. +// (3) LDS/block grows 24,576 -> 26,112 B (A[128, 136] + B[64, 136]); +// 2 x 26,112 = 52,224 <= 65,536 B keeps TWO resident blocks/CU = +// 8 waves/CU, so grid_blocks (1312), waves_per_block (4) and the +// active CU count (120) are unchanged. +// (4) the scalar fallback decodes the n-major layout when +// (k, n) == (6144, 2624) with ascending kk, so the paired M=2 / +// M=16 API shapes with the same (N, K) stay byte-exact; the generic +// <64, 64, 128> identity arm is untouched. // -// Iteration 3 (mandated pipeline round: single vs double buffering across K -// tiles; retain double buffering only on ISA/PMC evidence of reduced VMEM +// Correctness: the same weight elements reach the same fragment slots in +// the same ascending k-chunk sequence (k0 outer over 48 stages of 128 K, +// kk inner ascending over four m16n16k32 steps), so every int32 +// accumulator -- and therefore every bf16 output byte -- is identical to +// the bootstrap kernel and to the scalar reference. +// +// Round 2 (mandated operand-reuse round: cooperative A+B LDS staging vs direct +// loads, quantified A/B reuse per macro-tile, vectorized coalesced global loads +// with a bank-safe LDS layout): +// +// Operand-reuse accounting for the accepted <128, 64, 128> macro-tile +// (M=4096, N=2624, K=6144): +// * per block the unique operand bytes are A = 128 x 6144 = 786,432 and +// B = 6144 x 64 = 393,216; the 41 N-blocks re-read every A row-slab +// (M*K*41 = 1.03 GB/replay, 41x amplification of the 25.2 MB unique A) +// and the 32 M-blocks re-read every B column-slab (K*N*32 = 0.52 +// GB/replay, 32x of the 16.1 MB unique B): 1.55 GB/replay of staged +// global traffic; +// * inside one block each staged 16-B vector is re-read from LDS by the two +// wave columns (A) / the two wave rows (B), i.e. exactly 2x wave-level +// reuse per staged byte, so the LDS read traffic is 2x the staged bytes +// (4.72 GB/replay, ~23 B/cycle/CU at the measured 1.271 GHz -- far below +// the LDS bandwidth), and the counter-derived HBM rate is only +// 108.4 GB/s (147.8 MB/replay at a 90.7% L2 hit rate). The kernel is +// therefore NOT HBM- or LDS-bandwidth bound, so a direct per-wave global +// fragment path is rejected up front: it would issue 4x more vmem +// instructions (4 waves x 48 stages of 8-B fragment loads) and every 8-B +// fragment load would be a row-strided multi-line access instead of the +// coalesced 16-B staging vectors of the cooperative path. Cooperative +// A+B LDS staging is RETAINED -- the round improves it instead. +// +// Remaining exposed cost of the retained path: the VMEM round trip in front +// of every stage consume. The accepted source (PMC digest daf403b3..., +// profiles/.../iteration2/pmc.json: profiled 1368.793 us, grid 1312 blocks, +// 2 blocks/CU, 26,112 B LDS, lds_instructions 7,892,992 with +// lds_wait 4,535,391, lds_bank_conflicts 0, vmem_read 1,574,400, +// grbm_gui_active == grbm_count) issues per 128-K stage 24 wave-level +// global_load_dwordx4 + 24 ds_write2_b64 and then, per the exact-source ISA, +// strictly serializes global_load_dwordx4 -> s_waitcnt vmcnt(0) -> +// ds_write2_b64 -> s_barrier -> 24x ds_read_b64 -> 32x v_mmac -> s_barrier: +// with 8 resident waves/CU the GPU is 100% GRBM GUI-active yet issues only +// ~4.6% of its slots, i.e. the stage boundary -- not HBM, LDS bandwidth or +// MMA throughput -- sets the pace. +// +// Focused change (ONE mechanism -- the K stage pipeline -- entirely inside +// csrc/w8a8_gemm_hip.hip): the exact-shape guard now dispatches +// w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64>. It keeps the SAME +// macro-tile (grid 41 x 32 = 1312 blocks, 256 threads = 4 wavefronts, 64 x 32 +// quadrant of eight m16n16k32 int32 accumulators per wave), the SAME packed +// n-major B (16-B coalesced staging vectors), the SAME explicit load_frag8 +// fragment reads, the SAME bank-safe LDS stride form and the SAME coalesced +// epilogue -- and changes only how the K stages are staged: +// * the single-buffered 128-K stage (TWO __syncthreads per stage, full +// VMEM latency exposed in front of every consume) becomes a +// DOUBLE-buffered 64-K stage with ONE __syncthreads per stage: the +// stage-(s+1) global loads (2 int4 A + 1 int4 B per thread) are issued +// into staging VGPRs BEFORE the stage-s consume and flushed into the idle +// LDS buffer after it, so the stage-(s+1) VMEM latency overlaps the +// stage-s LDS/v_mmac work. +// +// Resources: LDS = 2 x (A[128, 72] + B[64, 72]) = 2 x 13,824 = 27,648 B/block +// (26,112 -> 27,648 B, +5.9%); 2 x 27,648 = 55,296 <= 65,536 B keeps TWO +// resident blocks/CU = 8 waves/CU, occupancy UNCHANGED. kLdsStride = 72 = +// StageK + 8: 8-aligned with an odd 8-byte count (72/8 = 9) so the 16 rows of +// each 16-lane ds_read_b64 phase land on 16 distinct bank pairs (18r mod 32, +// r = 0..15) -- the same conflict-free 4-phase floor as the accepted kernel -- +// and 72 % 16 == 8 keeps the two-half ds_write2_b64 staging form. 12 extra +// VGPR hold the prefetch payload (~72 -> ~84, still <= the 128/thread budget +// of 2 blocks/CU; 0 spills expected). +// +// Correctness: k0-outer over 96 ASCENDING 64-K stages with kk-inner ascending +// over two m16n16k32 steps gives the identical (k0 + kk) k-chunk sequence +// 0, 32, 64, 96, ... as the accepted 48 x 128-K stages, with the identical +// element-to-slot fragment mapping, so every int32 accumulator -- and every +// bf16 output byte -- is bit-identical to the accepted kernel and to the +// scalar reference. The n-major pack, the scalar-fallback n-major decode, the +// generic <64, 64, 128> identity arm and every shape guard are untouched. +// +// Round 3 (mandated pipeline round: compare single and double buffering across +// K tiles; retain double buffering only on ISA/PMC evidence of reduced VMEM // stalls without harmful LDS or occupancy growth): -// * Bottleneck evidence from the accepted iteration-2 source (fresh PMC -// digest 20f3bc9a..., profiles/.../iteration3/pmc.json: profiled -// 2041.095 us, vmem_read 1,574,400, lds_instructions 7,892,992 with -// lds_wait 4,571,574, lds_bank_conflicts 0, 26,112 B LDS, 2 blocks/CU): -// the exact-source ISA of w8a8_dumma_prefill_packedb_kernel<128,64,128> -// shows the stage loop is [global_load_dwordx4 ... -> s_waitcnt -// vmcnt(0) -> ds_write2_b64 -> s_barrier -> 24x ds_read_b64 -> 32x -// v_mmac -> s_barrier] -- the FULL global-load round trip (issue, vmcnt -// wait, LDS flush) is serialized in front of every consume with ZERO -// compute overlap, and the 14 vmcnt + 57 lgkmcnt waits in the symbol sit -// on that critical path. The K=6144 grid (1312 blocks) and 2-block -// co-residency do not hide this per-stage latency. -// * Focused change (one mechanism, all inside csrc/w8a8_gemm_hip.hip): -// the exact-shape guard now dispatches a new -// w8a8_dumma_prefill_packedb_db_kernel<128,64,64> -- SAME macro-tile -// (grid 41x32 = 1312 blocks, 256 threads = 4 waves, 64x32 quadrant per -// wave, eight m16n16k32 int32 accumulators), SAME load_frag8 fragment -// reads, SAME coalesced epilogue -- with the K stage changed from a -// single-buffered 128-K stage (2 __syncthreads, full VMEM latency -// exposed) to a DOUBLE-buffered 64-K stage with ONE __syncthreads per -// stage: the stage-(s+1) global loads are issued into 3 int4 staging -// VGPR per thread BEFORE the stage-s consume, and the prefetched vectors -// are flushed into the idle LDS buffer after the consume (the validated -// M=3072 prefetch order global_load_dwordx4 -> v_mmac -> vmcnt(0) -> -// ds_write2_b64, here in pure HIP: the vmcnt wait lands on the data -// dependency before the ds_write, i.e. AFTER the MMAC burst). -// * Resources: LDS = 2 x (A[128,72] + B[64,72]) = 2 x 13,824 = 27,648 -// B/block (26,112 -> 27,648 B, +5.9%); 2 x 27,648 = 55,296 <= 65,536 B -// keeps TWO resident blocks/CU = 8 waves/CU, occupancy UNCHANGED (the -// StageK=128 double-buffer alternative needs 52,224 B/block -> 1 -// block/CU = 4 waves/CU and is rejected up front as harmful occupancy -// loss). kLdsStride = 72 = 64 + 8: 8-aligned with an odd 8-byte count -// (72/8 = 9) so the 16 rows of each 16-lane ds_read_b64 phase land on 16 -// distinct bank pairs (18r mod 32, r = 0..15) -- the same conflict-free -// 4-phase floor as iteration 2 -- and 72 % 16 == 8 keeps the two-half -// ds_write2_b64 staging form. VGPR grows by the 12 staging registers -// (~72 -> ~84), far below the 128/thread budget for 2 blocks/CU; 0 -// spills expected. Barriers per block: 1 (prologue) + 95 (loop) = 96, -// the same as iteration 2's 2 x 48 (the final stage is consumed without -// a trailing barrier). -// * Correctness: k0-outer over 96 ASCENDING 64-K stages with kk-inner -// ascending over two m16n16k32 steps gives the identical (k0+kk) -// k-chunk sequence 0, 32, 64, 96, ... as iteration 2's 48 x 128-K -// stages, and the element-to-slot fragment mapping is unchanged, so -// every int32 accumulator is bit-identical (mismatch 0 / max_abs_error -// 0.0 expected). -// * Falsifiable retention: double buffering is retained only if the exact -// shape shows median_us < 2276.930923461914 AND p90_us < -// 2540.756378173828 with mismatch_count == 0, max_abs_error == 0.0, +// +// * Mandatory comparison (exact trusted PMC/ISA of the two accepted arms): +// arm median us LDS B/blk blocks/CU VGPR +// single 128-K <128,64,128> 1363.078 26,112 2/8 waves 72 +// double 64-K <128,64,64> 1153.970 27,648 2/8 waves 96 +// A single 64-K arm is not built: with NO prefetch it can only add +// barriers (2 per 64-K stage = 192 per block vs the 128-K arm's 96) while +// exposing the same VMEM round trip, so it is strictly dominated and +// rejected by construction. +// RETAINED -- double buffering. Evidence: median 1363.078 -> 1153.970 us +// (-15.3%) and p90 1364.825 -> 1164.522 us (-14.7%); the exact-source PMC +// of the double-buffered arm (digest ee609c3f..., iteration3/pmc.json) +// shows lds_wait_instructions 4,562,210 -> 3,590,796 (-21.3%) with +// grbm_gui_active == grbm_count (100% GUI-active), LDS 27,648 B/block +// (2 x 27,648 = 55,296 <= 65,536 so the trusted occupancy-probe split 2 +// is unchanged = 8 waves/CU), arch_vgpr 96 with 0 spills (<= the +// 128/thread budget of 2 blocks/CU) and sgpr 32. The ISA confirms the +// mechanism: the stage-(s+1) global_load_dwordx4s are issued BEFORE the +// stage-s ds_read_b64 / v_mmac burst and the vmcnt(0) wait lands on the +// ds_write2_b64 flush AFTER the MMACs, while the 128-K arm serializes +// global_load_dwordx4 -> s_waitcnt vmcnt(0) -> ds_write2_b64 -> s_barrier +// with zero compute overlap. So the prefetch converted exposed VMEM +// latency into overlapped work without trading away occupancy. +// LDS-growth verdict: the +1,536 B/block stays inside the 2-block budget +// and lds_wait FELL (the LDS path is not harmed in the net); the new +// lds_bank_conflicts 12,091,392 is exactly 6 x 4 x 4 x 96 x 1312 = one +// conflict per 16-lane phase of every merged ds_read2_b64, i.e. the +// bank-preserving alias of the +16-row second fragment offset +// (16 x 72 = 1152 B = 9 x 128 B; for ANY 8-aligned uniform row stride S, +// 16 rows = 4S dwords with 4S % 32 == 0). The unmerged ds_read_b64 +// control at the same 72/136-byte strides measured 0 conflicts, so this +// is an alias of the merged read form, not a layout defect. +// Every other K-tile pipeline is rejected inside the 65,536 B budget: +// StageK=96 double buffer = 39,936 B/block, StageK=128 double buffer = +// 52,224 B/block and 3 x 64-K = 41,472 B/block all exceed the 32,768 +// B/block that 2 resident blocks/CU permit (each would drop to 1 +// block/CU = 4 waves/CU, a 2x occupancy loss, explicitly excluded), and a +// 32-K stage leaves only 128 B staging vectors per block (half the 256 +// threads idle) while carrying a single m16n16k32 step per stage (too +// little consume work to amortize the barrier). StageK=64 x 2 buffers is +// therefore the unique 2-resident-block K pipeline and is retained. +// +// * Focused change (ONE mechanism -- the B operand's global-memory layout -- +// entirely inside csrc/w8a8_gemm_hip.hip). The pipeline round's own PMC +// identifies the next exposed cost: the accepted source reads +// 170,221,184 B/replay at only an 86.938% L2 hit rate (3,250,848 misses) +// and 166.147 GB/s counter-derived HBM, while the unique operands are only +// 25.2 MB (A) + 16.1 MB (B). The B side causes it: the n-major pack +// packed[n*K + kk] lays a [64, 64] B tile out as 64 separate 64-byte +// column chunks at a 6,144-byte stride, so every one of the 8 warp-level +// 16-byte staging loads touches 32 half-used 128-byte lines. +// launch_pack_w8a8_weight now packs the exact (k, n) == (6144, 2624) +// weight once, outside the timed region and out of Graph capture, into 41 +// B panels of [K, 64] int8 with 16-byte k-runs per column: +// packed[(n>>6)*(k*64) + (kk>>4)*1024 + (n&63)*16 + (kk&15)] +// = raw[kk*n + n] +// (byte-exact bijection 41 x 384 x 64 x 16 = k*n = 16,121,856; same byte +// count, same buffer and same captured addresses; every other (K, N) +// keeps the identity copy). Each [64, 64] B tile of a 64-K stage is then +// FOUR contiguous 1-KiB slabs (one per 16-k-run) and every staging vector +// is a 16-byte-aligned contiguous k-run of one column, so a warp's 32 +// x 16-byte B staging loads touch 4 FULLY-USED 128-byte lines instead of +// 32 half-used ones (8x fewer line requests per warp-load; the tile's line +// footprint is 32 fully-used lines instead of 64 half-used ones). +// The ONLY kernel change is the two B global-load sites (prologue + +// steady-state prefetch) of the dispatched +// w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64>: byte +// (n0>>6)*(k*64) + ((k0>>4)+k16)*1024 + n_row*16 replaces +// (n0+n_row)*k + k0 + k16*16. The LDS staging layout b_tile[2][64][72] +// (two-half ds_write2_b64), the col_major load_frag8 consumes, the +// k0-outer/kk-inner ascending int32 accumulation order, the coalesced +// epilogue, the grid (41 x 32 = 1312 blocks), the block size (256 = 4 +// wavefronts), the retained double-buffered 64-K pipeline (one +// __syncthreads per stage), LDS/block 27,648 B (2 blocks/CU = 8 waves/CU) +// and the VGPR budget are ALL unchanged, so this is a pure global-layout +// A/B. +// * Correctness: the same weight elements reach the same fragment slots in +// the same ascending k-chunk sequence, so every int32 accumulator -- and +// every bf16 output byte -- is bit-identical to the accepted kernel and to +// the scalar reference. The scalar fallback for (k, n) == (6144, 2624) +// decodes the panel layout in ascending 16-byte k-runs (the paired M=2 / +// M=16 API shapes with the same (N, K) stay byte-exact); the generic +// identity arm, the round-2 n-major pack kernel (kept compiled as the +// layout control) and every shape guard are untouched. Packing stays +// outside the timed region and out of Graph capture. +// * Predicted deltas: vmem_read_instructions unchanged (1,574,400), LDS +// counters unchanged (lds_instructions 4,901,632, lds_wait 3,590,796, +// conflicts at the merged-read alias floor 12,091,392), L2 misses strictly +// below 3,250,848 with the hit rate up, operator read bytes strictly below +// 170,221,184 (B line fetches halve), median < 1153.970 us. +// * Superseded by Round 5: the B-panel pack and the b_tile[2][64][72] LDS +// form below are replaced by the fragment-interleaved swizzle. +// +// Round 5 (mandated packing round: one weight packing/swizzle that makes each +// DUMMA B tile vector-loadable and LDS-bank-safe): +// +// * Baseline: the accepted double-buffered <128, 64, 64> kernel (official +// median 937.431 us / p90 939.452 us; exact-source PMC +// profiles/.../iteration4/pmc.json: profiled 889.453 us, grid 41x32 = +// 1312 blocks, 2 blocks/CU = 8 waves/CU, 27,648 B LDS, arch_vgpr 96 with +// 0 spills, lds_instructions 4,901,632 = 9.7 per wave-stage (6 merged +// ds_read2_b64 + 3 ds_write2_b64), lds_wait 3,828,676, +// lds_bank_conflicts 12,091,392 = 6 x 4 x 4 x 96 x 1312, vmem_read +// 1,574,400, L2 hit 90.39%). The surviving exposed LDS cost is the +// b_tile row-major layout: each col_major B fragment is one ds_read_b64 +// at a 72-byte row stride, and the compiler merges the fragment pairs +// into ds_read2_b64 whose +16-row second offset (16 x 72 = 1152 B = +// 9 x 128 B) aliases the same bank pairs -- one conflict per 16-lane +// phase, exactly the measured 12,091,392 conflicts. +// * Mechanism (ONE packing/swizzle, entirely in this file): replace the +// Round-3 B-panel pack with a fragment-interleaved swizzle. For each +// 32-k x 16-n DUMMA B fragment block the 4 k8 groups x 16 n-rows x 8 B +// are stored n-interleaved as [k8][n][8] (512 B per fragment block), +// 32-k major inside the 64-K stage: +// packed[((kk>>5)*n16g + (col>>4))*512 + +// (((kk>>3)&3)*16 + (col&15))*8 + (kk&7)] = raw[kk*n + col], +// with n16g = n>>4 (164 for the exact shape). Byte-exact bijection +// 192 x 164 x 512 = k*n = 16,121,856: same byte count, same buffer and +// same captured addresses (graph-stable), and packing still runs once, +// outside the timed region and out of Graph capture. +// * Kernel-side effect: every 16-byte B staging vector is now 2 n-rows x +// 8 B of one k8 of one fragment block, so the staging store is ONE +// 16-byte-aligned int4 (ds_write_b128) flat copy per thread (256 chunks +// = one per thread per 64-K stage; a warp reads 8 fully-used 128-byte +// lines) and every col_major B fragment is load_frag8_swz: lane +// (n = lane&15, k8 = lane>>4) reads p + k8*128 + n*8, i.e. one +// ds_read_b64 whose 16-lane phase spans exactly one 128-byte k8 block -- +// all 32 banks exactly once per phase, conflict-free by construction. +// The A tile keeps its bank-safe 72-byte stride, and the double-buffered +// 64-K pipeline, the k0-outer/kk-inner ascending int32 order, the +// coalesced epilogue, the grid (41x32 = 1312 blocks), the block +// (256 threads = 4 wavefronts) and the occupancy (2 blocks/CU = 8 +// waves/CU) are all unchanged, so this is a pure packing/swizzle A/B. +// * LDS: b_tile shrinks from 2 x [64, 72] = 9,216 B to 2 x 4,096 B; the +// block total is 18,432 (A) + 8,192 (B) = 26,624 B, so 2 x 26,624 = +// 53,248 <= 65,536 B keeps 2 resident blocks/CU = 8 waves/CU. +// * Correctness: the same weight elements reach the same fragment slots in +// the same k0-outer/kk-inner ascending k-chunk sequence, so every int32 +// accumulator -- and every bf16 output byte -- is bit-identical to the +// accepted kernel and to the scalar reference. The scalar fallback for +// (k, n) == (6144, 2624) decodes the swizzle element-wise in ascending +// kk (the paired M=2 / M=16 API shapes stay byte-exact); the generic +// identity arm, the Round-2 n-major pack control and every shape guard +// are untouched. +// * Predicted deltas: vmem_read_instructions unchanged (~1,574,400), +// LDS/block 27,648 -> 26,624 B, lds_bank_conflicts strictly below +// 12,091,392, lds_wait below 3,828,676, median < 937.431 us. +// Falsifiable: retained only if median_us < 937.4308013916016, +// mismatch_count == 0, max_abs_error == 0.0, // graph_capture_passed == true, the exact code object shows 0 spills / -// LDS == 27,648 B (2 blocks/CU still fits: 55,296 <= 65,536) and the -// stage ISA shows the prefetch loads issued BEFORE the ds_read_b64 / -// v_mmac burst with the vmcnt wait after the MMACs (load -> mmac -> -// vmcnt(0) -> ds_write2_b64), i.e. reduced exposed VMEM stalls; any -// median/p90 >= current best, any mismatch, any spill, any occupancy -// loss, or any ISA that still serializes the vmcnt wait in front of the -// consume falsifies the round and the control plane restores the -// iteration-2 source (single buffering retained). +// LDS == 26,624 B / VGPR <= 128 (2 blocks/CU), and the PMC shows +// lds_bank_conflicts and lds_wait strictly below the baseline with +// vmem_read not growing; any regression restores the iteration-3 source. +// +// Round 6 (mandated epilogue round: fuse per-row/per-column scales, bf16 +// conversion and the final coalesced store into the compute kernel; remove any +// unnecessary workspace/combine pass): +// +// * Fusion audit (nothing to remove, so the fused epilogue is made cheaper). +// The accepted source ALREADY fuses everything: store_prefill_fragment_ +// coalesced applies x_scale[row](per-row) * weight_scale[col](per-column), +// converts to bf16 and writes one coalesced 8-byte store per lane, and the +// exact-shape operator replay is a SINGLE dispatch -- the trusted PMC +// (profiles/glm_tp8_fused_qkv_a_proj_m4096/iteration6/pmc.json, digest +// 7e5731dc..., compile key 120c807e...) reports operator_aggregate +// kernel_count == 1, scratch_bytes == 0, vmem_write_instructions 41,984 = +// 8 x 4 x 1312 and write_bytes_per_operator_replay 21,500,384 = m*n*2, i.e. +// there is no workspace, no split-K combine pass and no second Graph node +// to remove; workspace_bytes is unused and workspace_split_k_capacity is +// zero for this shape. The remaining epilogue cost is therefore the +// conversion itself. +// * Measured cost of the fused epilogue (exact gfx928 code object of the +// current best, iteration6/current-best-isa/isa.txt). The dispatched +// kernel is ~1,920 instructions; the steady-state K-stage loop body is 107 +// instructions (one `s_branch 65428` back-edge at 0x6BB8 -> 0x6A0C) and +// the prologue is ~259, so the fused epilogue is ~1,490 static +// instructions -- roughly 13% of the dynamic instruction stream over +// 96 stages. The dominant term is that this DTK lowers EACH of the 32 +// per-lane __float2bfloat16 calls to a ~13-instruction EXEC-masked +// Inf/NaN sequence: +// v_and_b32 (b & 0x7f800000) / v_cmp_ne_u32 / s_and_saveexec_b64 / +// s_xor_b64 / v_bfe_u32 (bit 16) / v_add3_u32 (b + lsb + 0x7fff) / +// s_andn2_saveexec_b64 / v_mov_b32 0 / v_or_b32 (b | 0x10000) / +// v_cmp_eq_u32_sdwa / v_cndmask_b64 / s_or_b64 (restore EXEC) +// so ~420 instructions per lane plus 4 x (3 shifts + 3 ors) of packing, +// with 32 dependent SALU EXEC save/restore pairs that cannot be software- +// pipelined and serialise the wavefront (the epilogue is the only place +// VALU issues drop while the exact-shape kernel is 100% GRBM GUI-active). +// * Focused change (ONE mechanism -- the bf16 conversion of the fused +// epilogue -- entirely inside csrc/w8a8_gemm_hip.hip). Replace the library +// conversion + shift/or pack in store_prefill_fragment_coalesced with the +// branch-free round-to-nearest-even integer identity +// bf16_rne_high16(v) = (b + 0x7fff + ((b >> 16) & 1)) >> 16, +// b = __float_as_uint(v), +// which the disassembly proves is exactly what the normal (finite) path of +// the library conversion computes (v_bfe bit 16, v_add3 b + lsb + 0x7fff). +// Adding 0x7fff + the retained bit-16 LSB cannot carry out of the sign, the +// +((b>>16)&1) term rounds exact ties to an even LSB (RNE), max-float still +// rounds to +/-Inf and 0.0/-0.0 and subnormal encodings are preserved, so +// every 16-bit payload is identical. Every value converted here is finite +// by construction (|int32 dot| < 2^31 times finite dequant scales), so the +// removed Inf/NaN special case is unreachable. Four values now pack into +// one 8-byte word with three VALU ops (v_bfe/v_and + v_add3 + v_lshrrev) +// instead of ~13 EXEC-bracketed instructions each, and the 32 EXEC +// save/restore chains disappear. The epilogue structure, the per-row +// x_scale gather, the per-column float4 weight_scale gather, the register +// 4x4 shuffle transpose, the multiply order float(dot) * x_scale * +// weight_scale, the 8-byte-aligned coalesced store address, the grid +// (41 x 32 = 1312), the block (256 = 4 wavefronts), the double-buffered +// 64-K pipeline, the swizzled packed B, LDS/block 26,624 B (2 blocks/CU = +// 8 waves/CU) and the 65,536 B workspace budget are ALL unchanged, so this +// is a pure conversion A/B. The scalar fallback keeps __float2bfloat16 (it +// is not on the timed path); the generic identity arm and every shape +// guard are untouched. +// * Predicted deltas: grid_blocks 1312, LDS/block 26,624 B, 2 blocks/CU = +// 8 waves/CU, vmem_read 1,574,400 and vmem_write 41,984 unchanged; +// valu_instructions strictly below 23,920,384 and the epilogue's static +// EXEC/SALU count down by >1,000; median < 918.191 us with no p90 +// regression. Falsifiable: retained only if median_us <= 909.0 us +// (>= 1% round acceptance) and p90_us <= 919.530 us, mismatch_count == 0, +// max_abs_error == 0.0, graph_capture_passed == true, the exact code +// object shows 0 spills / LDS == 26,624 B / VGPR <= 128 (2 blocks/CU) and +// no __float2bfloat16 EXEC sequence remains in the epilogue; otherwise the +// iteration-5 source is restored. +// +// Round 9 (late ISA-diagnosis round; the control-plane plateau gate is CLOSED +// for this round -- phase hip_only, plateau false, raw inline asm and ISA +// Skills forbidden -- so this stays a HIP-only round): // -// Iteration 5 (mandated packing round: test one weight packing/swizzle that -// makes each DUMMA B tile vector-loadable and LDS-bank-safe): -// * Baseline: the accepted iteration-3 kernel -// w8a8_dumma_prefill_packedb_db_kernel<128,64,64> (official median -// 1070.23 us = 123.4 logical TOPS; fresh PMC digest 47542b38..., -// profiles/.../iteration4/pmc.json of the current source: profiled -// 1101.96 us, grid 1312, 96 arch_vgpr / 27,648 B LDS, lds_instructions -// 4,901,632 = 9.7 per wave-stage (6 merged ds_read2_b64 + 3 -// ds_write2_b64), lds_wait 3,701,906, lds_bank_conflicts 12,091,392 = -// 6 x 4 x 4 x 96 x 1312 EXACTLY -- one conflict per 16-lane phase of -// every merged ds_read2_b64, the inherent 2-cycle alias of the +16-row -// (16 x 72 = 1152 B = 9 x 128 B) second-fragment offset, NOT a -// layout defect (the iteration-2 unmerged ds_read_b64 control had 0 -// conflicts at the same 72/136-B strides), vmem_read 1,574,400, L2 hit -// 79.3% with 5,155,404 misses, TCC read 269.3 MB/replay). The kernel is -// LDS-latency-bound, not HBM-bound; the global B operand is the -// remaining layout lever. -// * Mechanism: replace the iteration-2 n-major pack -// packed[n*K + kk] (each 64-K B tile = 64 column chunks of 64 B at -// 6144-B stride -> 64 x 128-B lines touched per tile, 50% line -// utilization, scattered L2 footprint) with a B-PANEL pack -// packed[(n>>6)*(k*64) + (kk>>4)*1024 + (n&63)*16 + (kk&15)] = -// raw[kk*n + n] (41 panels of [K, 64] int8, 16-byte k-runs per column, -// byte-exact bijection verified by simulation): every [64,64] B tile of -// a 64-K stage is then FOUR CONTIGUOUS 1-KiB slabs (one per 16-k-run), -// each staging vector is a 16-byte-aligned contiguous k-run, and a wave -// reads 4 x 256 B of one panel region per stage -- the tile is -// vector-loadable as one coalesced stream with 8x fewer 128-B lines -// touched per tile (32 vs 256) and full line utilization, so L2 tag -// pressure and B HBM fetch traffic drop without any change in LDS -// layout or instruction count. -// * Kernel-side changes are ONLY the two B global-load sites (prologue + -// steady-state prefetch) of w8a8_dumma_prefill_packedb_db_kernel -// (dispatched solely by the exact-shape guard; static_assert BN == 64): -// vector (n_row, k16) is loaded from byte -// (n0>>6)*(k*64) + ((k0>>4)+k16)*1024 + n_row*16 instead of -// (n0+n_row)*k + k0 + k16*16. LDS staging layout (b_tile[64][72], -// two-half ds_write2_b64), the col_major load_frag8 consumes, the -// k0-outer/kk-inner int32 accumulation order, the coalesced epilogue, -// the grid (41x32 = 1312 blocks), threads (256 = 4 waves), occupancy -// (2 blocks/CU, LDS 27,648 B) and VGPR budget are ALL unchanged, so the -// round is a pure packing/swizzle A/B. -// * Correctness: the same weight elements reach the same fragment slots -// in the same ascending k-chunk sequence, so every int32 accumulator is -// bit-identical (mismatch 0 / max_abs_error 0.0 expected); the scalar -// fallback decodes the panel layout when (k,n) == (6144,2624) (ascending -// 16-byte k-runs keep the exact int32 order for the paired M=2/M=16 API -// shapes); the generic identity arms and sibling kernels are untouched. -// * Packing stays outside the timed region and out of Graph capture -// (launch_pack_w8a8_weight, same byte count k*n = 16,121,856 and same -// buffer -> graph-stable captured addresses unchanged; pack relation -// verified by a byte-exact index simulation before dispatch). -// * Predicted deltas: vmem_read_instructions unchanged (~1,574,400), LDS -// counters unchanged (conflicts stay at the 12,091,392 merged-read alias -// floor), L2 misses < 5,155,404 with hit rate UP, TCC read bytes < -// 269,329,728 (B line fetches halve), median ~1000-1070 us. Falsifiable: -// retained only if median_us < 1070.2328491210938 AND p90_us < -// 1081.9308471679688 with mismatch_count == 0, max_abs_error == 0.0, +// * Bottleneck evidence (exact trusted gfx928 ISA of the current best, +// profiles/glm_tp8_fused_qkv_a_proj_m4096/iteration6/current-best-isa/ +// isa.txt, digest 7e5731dc...; the round-6 source differs from that +// profiled official best only in the bf16 conversion, so the steady-state +// loop analysed here is the dispatched loop). The steady-state stage body +// is 107 instructions (0x6A0C..0x6BB8, one back-edge `s_branch 65428` at +// 0x6BB8 -> 0x6A0C) and roughly 14 of them are pure A-row tail +// predication: +// 0x6A0C..0x6A20 six unconditional `v_mov_b32` zero-fills of the two +// int4 staging vectors (v33/v34 and v37..v40), +// 0x6A24..0x6A34 `s_and_saveexec_b64` / `s_cbranch_execz` / +// `global_load_dwordx4` / `s_or_b64` for A vector 0, +// 0x6A38..0x6A50 two more `v_mov_b32` zero-fills (v35/v36) plus the +// same EXEC save/restore + uniform branch for A +// vector 1. +// The guard source `a_pre[v] = (g_row < m) ? load : int4{0,0,0,0}` is +// never false for the dispatched M=4096 (M / BM = 4096 / 128 = 32 exact), +// yet its zero-fills execute unconditionally in every one of the 96 +// stages and its two EXEC read-modify-writes serialize the wavefront in +// front of the very global loads that prefetch the next tile. That is +// ~14/107 = 13% of the hot loop's instruction stream, on top of the +// epilogue's per-fragment `row >= m` EXEC bracket (s_and_saveexec / +// s_xor_b64 / s_andn2_saveexec_b64 / s_cbranch_execz / s_or_b64 around +// each of the eight unrolled fragment stores), for a kernel the PMC shows +// 100% GRBM GUI-active yet issuing only ~0.42 instructions per cycle per +// CU (lds_wait_instructions 5,302,705 against lds_instructions +// 4,901,632). +// * Focused change (ONE mechanism -- exact-shape full-tile specialization, +// entirely inside csrc/w8a8_gemm_hip.hip): the dispatched kernel gains a +// compile-time `bool kFullTile` template parameter. The exact-shape guard +// keeps `(m >= 128, n == 2624, k == 6144)` and the same grid (N/64 x +// ceil(M/128)) and block (256 = 4 wavefronts), and selects +// `kFullTile == true` exactly when the caller's m is a multiple of +// BM = 128 -- i.e. when every block tile lies fully inside the output (the +// assigned M=4096 case). In that instantiation the A staging loads +// (prologue AND steady state) drop the `(g_row < m)` select / zero-fill +// and the epilogue drops the per-fragment `row >= m` early return, so the +// loop body keeps only the address math and the three coalesced +// `global_load_dwordx4`. Every other m keeps the guarded instantiation +// byte-for-byte (zero-fill, epilogue row mask, tail-block coverage), and +// the paired M=2/M=16 API shapes with the same (N, K) still take the +// scalar fallback because the `m >= 128` guard is untouched. +// * Why this is the remaining lever and not a new pipeline: rounds 7 +// (deeper register prefetch / stage pairs) and 8 (smaller macro-tile, more +// resident waves) both regressed, so the double-buffered 64-K stage (ONE +// __syncthreads per stage, LDS 26,624 B/block, 2 blocks/CU = 8 waves/CU), +// the 128 x 64 macro-tile, the fragment-interleaved packed B swizzle, the +// load_frag8 / load_frag8_swz consumes, the k0-outer / kk-inner ascending +// m16n16k32 int32 accumulation order and the fused branch-free bf16 +// epilogue are ALL retained unchanged. The round removes instructions that +// are provably dead for this shape without touching any resource, any +// memory access order or any MMA input. +// * Correctness: for m % 128 == 0 no block tile crosses m, so the removed +// guard predicate is never false and the removed zero-fill never reaches an +// accumulator; the surviving assignments, the ascending k-chunk sequence +// and the exact `float(dot) * x_scale * weight_scale` multiply order are +// bit-identical to the guarded path (mismatch 0 / max_abs_error 0.0 +// expected). The guarded instantiation preserves the exact current +// semantics for every m % 128 != 0. +// * Predicted deltas: steady-state loop 107 -> ~93 instructions (the eight +// zero-fills and the two EXEC save/restore + uniform-branch pairs removed), +// valu_instructions strictly below 23,920,384, grid_blocks 1312, +// block_threads 256, LDS/block 26,624 B, arch_vgpr <= 128 with 0 spills, +// vmem_read_instructions 1,574,400, vmem_write_instructions 41,984 and the +// epilogue store count all unchanged. Falsifiable: retained only if +// median_us < 912.7147674560547 (the current shadow base) with no p90 +// regression, mismatch_count == 0, max_abs_error == 0.0, // graph_capture_passed == true, the exact code object shows 0 spills / -// LDS == 27,648 B / VGPR <= 128 (2 blocks/CU), and PMC confirms L2 -// misses and TCC read bytes strictly below the iteration-4 values with -// vmem_read/lds/conflicts not growing; any median/p90 >= current best, -// any mismatch, any spill, any occupancy loss, or any L2/HBM regression -// falsifies the pack and the control plane restores the iteration-3 +// LDS == 26,624 B / VGPR <= 128 (2 blocks/CU) and the kFullTile symbol's +// loop body no longer contains the A-guard +// `s_and_saveexec`/`s_cbranch_execz`/zero-fill sequence; any median/p90 +// regression, mismatch, spill or occupancy loss restores the round-6 // source. // // Include order is fixed by the control plane: hip_runtime, hip_bfloat16, @@ -273,7 +464,7 @@ namespace { -constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kWaveSize = 64; // gfx928 native wavefront size constexpr int kTileM = 16; constexpr int kTileN = 16; constexpr int kTileK = 32; // INT8 DUMMA unit: m16n16k32, int32 accumulation @@ -286,8 +477,7 @@ constexpr int kExactK = 6144; using namespace du::dumma; -// Coalesced accumulator epilogue for gfx928 INT8 m16n16k32 (verified against -// du_store_matrix_sync in the accepted worker-29 lineage): the accumulator +// Coalesced accumulator epilogue for gfx928 INT8 m16n16k32: the accumulator // lane mapping is row = lane & 15, col_mod4 = lane >> 4, frag.x[i] maps to // columns col_mod4 + 4*i. A register 4x4 transpose (two 2x2 steps with // shfl_xor 16 then 32, one v_cndmask per element per step) re-routes the @@ -298,9 +488,43 @@ using namespace du::dumma; // wavefront-uniform (lane & 15 cycles the same 16 rows in every 16-lane // group), so the shuffles never mix active and inactive lanes. base_col is a // multiple of 64 for the dispatched tiles, so the float4 weight_scale load -// (col0 % 4 == 0) and the 8-byte store (n even -> out offset a multiple of -// 8) are aligned. -template +// (col0 % 4 == 0) and the 8-byte store (n even -> offset a multiple of 8) are +// aligned. +// +// Round 6: the fused bf16 conversion itself is made branch-free. The exact +// gfx928 code object of the previous best lowers each of the 32 per-lane +// __float2bfloat16 calls to a ~13-instruction sequence that brackets an +// Inf/NaN special case with EXEC save/restore: +// v_and_b32 (b & 0x7f800000) / v_cmp_ne_u32 / s_and_saveexec_b64 / s_xor_b64 +// / v_bfe_u32 (bit 16) / v_add3_u32 (b + lsb + 0x7fff) / s_andn2_saveexec_b64 +// / v_mov_b32 0 / v_or_b32 (b | 0x10000) / v_cmp_eq_u32_sdwa / v_cndmask_b64 +// / s_or_b64 (restore EXEC) +// The EXEC manipulation cannot be software-pipelined and serialises the +// wavefront. For every FINITE fp32 the library conversion is exactly +// round-to-nearest-even at bit 16, which is the branch-free integer identity +// bf16 = (b + 0x7fff + ((b >> 16) & 1)) >> 16, b = __float_as_uint(v) +// (adding 0x7fff + the retained bit-16 LSB never carries out of the sign, and +// the +((b>>16)&1) term breaks exact ties toward an even LSB, so max-float +// still rounds to +/-Inf and -0.0/0.0 are preserved). Every value converted +// here is finite by construction: |int32 dot| < 2^31 and the dequant scales +// are finite, so no Inf/NaN can reach the conversion. This yields the exact +// same 16-bit payloads as __float2bfloat16 while dropping the EXEC sequence +// and the separate shift/or pack (four bf16 pairs become three VALU ops per +// 8-byte word instead of 4 x 13 + 3). +__device__ __forceinline__ unsigned int bf16_rne_high16(float v) { + const unsigned int b = __float_as_uint(v); + return (b + 0x7fffu + ((b >> 16) & 1u)) >> 16; +} + +// Round 9 adds the compile-time `kGuarded` flag: kGuarded == false is the +// exact-shape full-tile specialization (m % BM == 0), for which every row of +// every block tile is strictly inside the output, so the `row >= m` early +// return -- which the exact code object lowers to an EXEC save/restore pair +// plus a uniform branch around each of the eight unrolled fragment stores -- +// is removed at compile time. The default kGuarded == true keeps the exact +// current semantics for every other m (including the generic arm and the +// single-buffered siblings, which never pass the flag). +template __device__ __forceinline__ void store_prefill_fragment_coalesced( const AccFragment& frag, const float* __restrict__ x_scale, @@ -312,8 +536,10 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced( int base_col, int lane) { const int row = base_row + (lane & 15); - if (row >= m) { - return; + if (kGuarded) { + if (row >= m) { + return; + } } const int c4 = lane >> 4; // 0..3 const int x0 = frag.x[0]; @@ -352,53 +578,105 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced( const float v3 = static_cast(f3) * xs * ws.w; const uint64_t packed = static_cast( - static_cast(__float2bfloat16(v0))) | - (static_cast( - static_cast(__float2bfloat16(v1))) - << 16) | - (static_cast( - static_cast(__float2bfloat16(v2))) - << 32) | + bf16_rne_high16(v0) | + (bf16_rne_high16(v1) << 16)) | (static_cast( - static_cast(__float2bfloat16(v3))) - << 48); + bf16_rne_high16(v2) | + (bf16_rne_high16(v3) << 16)) + << 32); *reinterpret_cast(out + static_cast(row) * n + col0) = packed; } -// Simple native INT8 DUMMA m16n16k32 tiled prefill kernel (identity [K, N] -// weight layout; launch_pack_w8a8_weight is the identity bootstrap, so -// packed_weight == raw logical weight). Template parameters: BM x BN output +// Explicit 8-byte int8 fragment loader (lineage-validated load_frag8): for +// BOTH the row_major matrix_a and the col_major matrix_b m16n16k32 fragments, +// lane l -> row (l & 15), k-quarter (l >> 4) holds the EIGHT CONTIGUOUS bytes +// at p[row * ldm + (l >> 4) * 8 .. +7]. Loading the 8 bytes as one int64 +// forces the ds_read_b64 form (16-lane phases, distinct bank pairs per group +// -> the 4-phase conflict-free floor with an odd 8-byte LDS row stride), +// identical fragment values to the library loader, so the v_mmac inputs and +// the exact int32 accumulation order are unchanged. `lane` is +// threadIdx.x % kWaveSize, so no wave-id builtin is required. +template +__device__ __forceinline__ void load_frag8( + Frag& f, + const int8_t* __restrict__ p, + unsigned ldm, + int lane) { + const unsigned row = static_cast(lane) & 0xfu; + const unsigned kq = static_cast(lane) >> 4; + const int64_t v = *reinterpret_cast(p + row * ldm + (kq << 3)); + reinterpret_cast(f.x)[0] = v; +} + +// Round-5 fragment-interleaved B swizzle helpers. +// +// Each 32-k x 16-n DUMMA B fragment block is 512 B = [k8 0..3][n 0..15][8 B] +// (n-interleaved k8 groups). The 16-byte chunk (k32, n16, k8, n_pair) holds +// the two 8-byte k-runs of n-rows 2*n_pair and 2*n_pair+1 of one k8 group, +// and sits at byte (k32 * n16g + n16) * 512 + k8 * 128 + n_pair * 16 of the +// packed [n>>4][k>>5][512 B] array (n16g = n>>4). Global elements satisfy +// packed element (col, kk) at ((kk>>5)*n16g + (col>>4))*512 + +// (((kk>>3)&3)*16 + (col&15))*8 + (kk&7) -- the exact inverse of this chunk +// map, so staging is a flat 16-byte copy and the fragment load reads the +// same bytes the Round-3 layout delivered. +__device__ __forceinline__ int64_t packed_b_swz_offset( + int k32, + int n16, + int n16g, + int k8, + int n_pair) { + return (static_cast(k32) * n16g + n16) * 512 + + static_cast(k8) * 128 + (n_pair << 4); +} + +// Round-5 col_major B fragment loader on the swizzled tile: lane (n = lane&15, +// k8 = lane>>4) reads p + k8*128 + n*8. One ds_read_b64 per fragment and each +// 16-lane phase spans exactly one 128-byte k8 block (each bank exactly once), +// so the read is conflict-free by construction. Byte placement is identical to +// load_frag8, so the v_mmac operands and the int32 accumulation stay +// bit-exact. +template +__device__ __forceinline__ void load_frag8_swz( + Frag& f, + const int8_t* __restrict__ p, + int lane) { + const unsigned n = static_cast(lane) & 0xfu; + const unsigned k8 = static_cast(lane) >> 4; + const int64_t v = + *reinterpret_cast(p + (k8 << 7) + (n << 3)); + reinterpret_cast(f.x)[0] = v; +} + +// Generic native INT8 DUMMA m16n16k32 tiled prefill kernel with the identity +// [K, N] weight layout (packed_weight == raw logical weight for every (K, N) +// other than the exact assigned pair). Template parameters: BM x BN output // tile per block, StageK K rows per single-buffered LDS stage. 256 threads = -// 4 wavefronts; each wave owns a (BM/2) x (BN/2) quadrant of m16n16k32 -// int32 accumulator fragments. Grid = (N/BN) x (ceil(M/BM)). +// 4 wavefronts; each wave owns a (BM/2) x (BN/2) quadrant of m16n16k32 int32 +// accumulator fragments. Grid = (N/BN) x (ceil(M/BM)). // -// K is staged cooperatively in a SINGLE LDS buffer at StageK granularity -// with TWO __syncthreads per stage (one after the staging stores, one after -// all waves consumed the buffer): -// * A[BM, StageK] is staged row-major with flat 16-byte vectors (each -// thread strides over kAVectors = BM*StageK/16 vectors); rows past M are +// K is staged cooperatively in a SINGLE LDS buffer at StageK granularity with +// TWO __syncthreads per stage (one after the staging stores, one after all +// waves consumed the buffer): +// * A[BM, StageK] is staged row-major with flat 16-byte vectors (each thread +// strides over kAVectors = BM*StageK/16 vectors); rows past M are // zero-filled and masked in the epilogue. // * B[StageK, BN] is staged row-major from the logical [K, N] weight with -// flat 16-byte vectors (each thread strides over kBVectors = -// StageK*BN/16 vectors). The dispatch guarantees n % BN == 0 and +// flat 16-byte vectors. The dispatch guarantees n % BN == 0 and // k % StageK == 0, so every global vector is aligned and in-bounds. // * Fragment loads use the library du_load_matrix_sync (row_major for both // operands) and accumulate with du_mma_sync in ascending k order -- -// k0-outer over stages, kk-inner over kTileK steps -- preserving the -// exact int32 accumulation order of the scalar reference. -// * The epilogue is the coalesced fragment store. -// -// Bootstrap keeps this direct/single-buffered by design (no double buffering, -// no split-K, no raw asm); later rounds tune tile, stage, and packing. +// k0-outer over stages, kk-inner over kTileK steps -- preserving the exact +// int32 accumulation order of the scalar reference. +// * The epilogue is the coalesced fragment store above. template __global__ __launch_bounds__(kThreadsPerBlock) void w8a8_dumma_prefill_tiled_kernel( - const int8_t* __restrict__ x_q, // [M, K] row-major - const int8_t* __restrict__ weight, // [K, N] row-major (identity pack) - const float* __restrict__ x_scale, // [M] + const int8_t* __restrict__ x_q, // [M, K] row-major + const int8_t* __restrict__ weight, // [K, N] row-major (identity pack) + const float* __restrict__ x_scale, // [M] const float* __restrict__ weight_scale, // [N] - hip_bfloat16* __restrict__ out, // [M, N] + hip_bfloat16* __restrict__ out, // [M, N] int m, int n, int k) { @@ -509,56 +787,42 @@ w8a8_dumma_prefill_tiled_kernel( } } -// Explicit 8-byte int8 fragment loader (lineage-validated load_frag8): for -// both the row_major matrix_a and the col_major matrix_b m16n16k32 -// fragments, lane l -> row (l & 15), k-quarter (l >> 4) holds the EIGHT -// CONTIGUOUS bytes at p[row*ldm + (l>>4)*8 .. +7]. Loading the 8 bytes as -// one int64 forces the ds_read_b64 form (16-lane phases, distinct bank -// pairs per group -> 4-phase conflict-free floor with an odd 8-byte LDS row -// stride), identical fragment values to the library loader, so the v_mmac -// inputs and the exact int32 accumulation order are unchanged. -template -__device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, - unsigned ldm) { - const unsigned row = __lane_id() & 0xf; - const unsigned kq = __lane_id() >> 4; - const int64_t v = - *reinterpret_cast(p + row * ldm + (kq << 3)); - reinterpret_cast(f.x)[0] = v; -} - -// Iteration-2 exact-shape packed-B prefill kernel (dispatched only for -// (m >= 128, n == 2624, k == 6144), where launch_pack_w8a8_weight produced -// the n-major packed_weight[n*K + kk] = raw[kk*N + n] layout). Same -// geometry as w8a8_dumma_prefill_tiled_kernel<128,64,128> (grid (N/64) x -// (M/128), 256 threads = 4 wavefronts, 64x32 quadrant per wave, eight -// m16n16k32 int32 accumulators, single-buffered 128-K LDS stage with TWO -// __syncthreads per stage, same k0-outer/kk-inner int32 accumulation order, -// same coalesced fragment epilogue). The two differences, both inside the -// operand path: -// * B is staged n-major (b_tile[n][kk] from packed [N, K]) so col_major -// B fragments read 8 contiguous K bytes per lane; -// * A and B LDS row strides are 136 B (StageK + 8): 8-byte aligned for -// the ds_read_b64 fragment loads (odd 8-byte stride count -> the 16 -// rows of each 16-lane phase land on 16 distinct bank pairs, 34r mod 32 -// = 2r mod 32, r = 0..15 -> all 16 even dword starts, conflict-free -// 4-phase floor) and 136 % 16 == 8 so staging stores are the two-half -// int64 ds_write2_b64 form (single ds_write_b128 would be misaligned on -// odd rows). +// Round-1 packed-weight architecture kernel, kept as the single-buffered +// 128-K sibling of the dispatched <128, 64, 64> double-buffered kernel below +// (Round 2 moved the exact-shape dispatch to that sibling; Round 3 keeps this +// kernel compiled but undispatched as the n-major layout / single-buffer +// control, and only the double-buffered kernel reads the packed B; Round 5 +// replaces that packed B layout with the fragment-interleaved swizzle, so this +// n-major control no longer matches the pack either). It is defined for +// the n-major packed_weight[n * K + kk] = raw[kk * N + n] layout produced by +// launch_pack_w8a8_weight. SAME geometry as the bootstrap +// w8a8_dumma_prefill_tiled_kernel<128, 64, 128>: +// grid (N/64) x (M/128) = 41 x 32 = 1312 blocks, 256 threads = 4 wavefronts, +// one 64 x 32 quadrant of eight m16n16k32 int32 accumulators per wave, +// single-buffered 128-K LDS stage with TWO __syncthreads per stage, same +// k0-outer/kk-inner ascending int32 accumulation order, same coalesced +// fragment epilogue. The two differences, both inside the operand path: +// * B is staged n-major (b_tile[n_row][kk] from the packed [N, K] weight) +// so col_major B fragments are 8 contiguous K bytes per lane and staging +// is one coalesced 16-byte vector per (n_row, k16) pair; +// * A and B LDS row strides are 136 B (StageK + 8): 8-byte aligned with an +// odd 8-byte count (136/8 = 17) so the 16 rows of each 16-lane ds_read_b64 +// phase land on 16 distinct bank pairs (the conflict-free 4-phase floor), +// and 136 % 16 == 8 keeps the two-half int64 staging form. // Both fragment loaders are the explicit load_frag8 above; the staged bytes // are the same matrix elements in the same fragment slots as the library // loaders produced from the identity layout, so every v_mmac input and the -// exact int32 accumulation (k0-outer over 48 stages, kk-inner ascending -// over four m16n16k32 steps) are bit-identical to the accepted iteration-1 -// kernel and to the scalar reference. +// exact int32 accumulation (k0-outer over 48 stages, kk-inner ascending over +// four m16n16k32 steps) are bit-identical to the bootstrap kernel and to the +// scalar reference. template __global__ __launch_bounds__(kThreadsPerBlock) void w8a8_dumma_prefill_packedb_kernel( - const int8_t* __restrict__ x_q, // [M, K] row-major - const int8_t* __restrict__ packed_w, // [N, K] n-major (exact shape) - const float* __restrict__ x_scale, // [M] + const int8_t* __restrict__ x_q, // [M, K] row-major + const int8_t* __restrict__ packed_w, // [N, K] n-major (exact shape) + const float* __restrict__ x_scale, // [M] const float* __restrict__ weight_scale, // [N] - hip_bfloat16* __restrict__ out, // [M, N] + hip_bfloat16* __restrict__ out, // [M, N] int m, int n, int k) { @@ -568,9 +832,10 @@ w8a8_dumma_prefill_packedb_kernel( "block tile must be a multiple of the wave quadrant"); static_assert(kThreadsPerBlock % kWaveSize == 0, "blockDim must be a multiple of the gfx928 wavefront size"); - // Bank-safe LDS row stride: StageK + 8 bytes = 34 dwords. 8-aligned with - // an odd 8-byte count -> every lane's 8-byte fragment is ds_read_b64 at - // the 4-phase conflict-free floor (see the kernel comment above). + // Bank-safe LDS row stride: StageK + 8 bytes = 34 dwords for StageK = 128. + // 8-aligned with an odd 8-byte count -> every lane's 8-byte fragment is a + // ds_read_b64 at the 4-phase conflict-free floor (see the kernel comment + // above); 8 mod 16 keeps the two-half staging store form. constexpr int kLdsStride = StageK + 8; static_assert(kLdsStride % 8 == 0 && (kLdsStride / 8) % 2 == 1, "LDS row stride must be 8-aligned with an odd 8-byte count"); @@ -602,14 +867,14 @@ w8a8_dumma_prefill_packedb_kernel( } } - // Flat 16-byte staging vectors per stage (dispatch guarantees - // k % StageK == 0 and n % BN == 0). 8 vectors per 128-K row. + // Flat 16-byte staging vectors per stage (the dispatch guarantees + // k % StageK == 0 and n % BN == 0). 8 vectors per StageK row for 128. + constexpr int kVecsPerRow = StageK / static_cast(sizeof(int4)); constexpr int kAVectors = BM * StageK / static_cast(sizeof(int4)); constexpr int kBVectors = StageK * BN / static_cast(sizeof(int4)); - constexpr int kVecsPerRow = StageK / static_cast(sizeof(int4)); for (int k0 = 0; k0 < k; k0 += StageK) { - // Stage A[BM, StageK] row-major with row stride kLdsStride: thread + // Stage A[BM, StageK] row-major with the bank-safe row stride: thread // `vec` owns the 16-byte chunk at (local_row, k16); rows past M are // zero-filled and masked in the epilogue. for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { @@ -621,18 +886,17 @@ w8a8_dumma_prefill_packedb_kernel( v = *reinterpret_cast( x_q + static_cast(g_row) * k + k0 + k16 * 16); } - // Two 8-byte halves (stride % 16 == 8 -> odd rows are not 16-byte - // aligned, so a single int4 store is illegal; the int64 pair lowers - // to one ds_write2_b64 per thread). + // Two 8-byte halves: the row stride is 8 mod 16, so odd rows are not + // 16-byte aligned and a single int4 LDS store is illegal. The int64 + // pair lowers to one ds_write2_b64 per thread. const int64_t* src64 = reinterpret_cast(&v); int64_t* dst64 = reinterpret_cast( a_tile + local_row * kLdsStride + k16 * 16); dst64[0] = src64[0]; dst64[1] = src64[1]; } - // Stage B[BN, StageK] n-major from the packed [N, K] weight: thread - // `vec` owns the 16-byte chunk at (n_row, k16); the K axis is always - // in range. + // Stage B[BN, StageK] n-major from the packed [N, K] weight: thread `vec` + // owns the 16-byte chunk at (n_row, k16); the K axis is always in range. for (int vec = tid; vec < kBVectors; vec += kThreadsPerBlock) { const int n_row = vec / kVecsPerRow; const int k16 = vec - n_row * kVecsPerRow; @@ -658,14 +922,16 @@ w8a8_dumma_prefill_packedb_kernel( load_frag8( a_frag[j], a_tile + (local_row + j * kTileM) * kLdsStride + kk, - kLdsStride); + static_cast(kLdsStride), + lane); } #pragma unroll for (int i = 0; i < kWaveN16; ++i) { load_frag8( b_frag[i], b_tile + (local_col + i * kTileN) * kLdsStride + kk, - kLdsStride); + static_cast(kLdsStride), + lane); } #pragma unroll for (int j = 0; j < kWaveM16; ++j) { @@ -675,8 +941,8 @@ w8a8_dumma_prefill_packedb_kernel( } } } - // Single buffer: every wavefront must finish reading LDS before the - // next stage's stores overwrite it. + // Single buffer: every wavefront must finish reading LDS before the next + // stage's stores overwrite it. __syncthreads(); } @@ -693,51 +959,74 @@ w8a8_dumma_prefill_packedb_kernel( } } -// Iteration-3 double-buffered packed-B prefill kernel (dispatched only for -// (m >= 128, n == 2624, k == 6144)): SAME macro-tile as the accepted -// iteration-2 w8a8_dumma_prefill_packedb_kernel<128,64,128> (grid (N/64) x -// (M/128) = 41x32 = 1312 blocks, 256 threads = 4 wavefronts, 64x32 quadrant -// per wave, eight m16n16k32 int32 accumulators, same load_frag8 fragment -// reads, same bank-safe LDS row strides, same coalesced epilogue). The ONLY -// change is the K pipeline (see the file-header iteration-3 section); -// iteration 5 additionally reads the packed B operand from the B-panel -// layout (see the file-header iteration-5 section and -// w8a8_pack_panel_i8_kernel) -- the LDS staging layout, fragment consumes, -// and accumulation order are untouched: -// * The single-buffered 128-K stage (two __syncthreads; the exact-source -// ISA serializes global_load_dwordx4 -> s_waitcnt vmcnt(0) -> -// ds_write2_b64 -> s_barrier -> consume with no compute overlapping the -// VMEM round trip) becomes a 64-K DOUBLE-buffered stage with ONE -// __syncthreads per stage: the stage-(s+1) global loads are issued into -// 3 int4 staging VGPR per thread BEFORE the stage-s consume, and the -// prefetched vectors are flushed into the idle buffer after the consume, -// so the VMEM latency of stage s+1 overlaps the LDS/MMAC work of stage s -// (validated M=3072 prefetch order, in pure HIP: the vmcnt wait lands on -// the register data dependency right before the ds_write2_b64, i.e. -// after the v_mmac burst). -// * LDS = 2 x (A[128,72] + B[64,72]) = 2 x 13,824 = 27,648 B/block -// (26,112 -> 27,648 B, +5.9%); 2 x 27,648 = 55,296 <= 65,536 B keeps -// TWO resident blocks/CU = 8 waves/CU, occupancy unchanged. kLdsStride = -// 72 = StageK + 8 (8-aligned, odd 8-byte count 9 -> the 16 rows of each -// 16-lane ds_read_b64 phase land on 16 distinct bank pairs, 18r mod 32, -// r = 0..15 -> the same 4-phase conflict-free floor; 72 % 16 == 8 keeps -// the two-half ds_write2_b64 staging form). 12 extra VGPR for the -// staging registers (~72 -> ~84, still <= 128 for 2 blocks/CU). -// * k0-outer over 96 ASCENDING 64-K stages with kk-inner ascending over -// two m16n16k32 steps: the (k0+kk) k-chunk sequence 0, 32, 64, 96, ... -// is identical to iteration 2 (48 x 128-K stages), so every int32 -// accumulator is bit-identical (mismatch 0 / max_abs_error 0.0). -// * Barriers per block: 1 (prologue) + 95 (loop) = 96, the same as -// iteration 2's 2 x 48; the final stage is consumed without a trailing -// barrier. -template +// Round 9 exact-shape A-staging load. For the dispatched M=4096 every 128-row +// block tile lies fully inside the output, so the `(g_row < m)` row guard is +// dead: the exact code object lowers it to eight unconditional zero-fill +// `v_mov_b32` plus two EXEC save/restore (`s_and_saveexec_b64` / +// `s_cbranch_execz` / `s_or_b64`) uniform-branch pairs per K stage, i.e. ~14 +// of the 107 steady-state loop instructions and an EXEC serialization in +// front of the very `global_load_dwordx4`s that prefetch the next tile. The +// kFullTile instantiation removes the select (one coalesced load per staging +// vector, no zero-fill), while kFullTile == false keeps the zero-fill + row +// mask semantics identical to the current kernel for every m % 128 != 0. +template +__device__ __forceinline__ int4 load_a_stage_vec( + const int8_t* __restrict__ x_q, + int64_t offset, + int m, + int g_row) { + if (kFullTile) { + return *reinterpret_cast(x_q + offset); + } + return (g_row < m) ? *reinterpret_cast(x_q + offset) + : int4{0, 0, 0, 0}; +} + +// Round-2 double-buffered packed-B prefill kernel, dispatched ONLY for +// (m >= 128, n == 2624, k == 6144). SAME macro-tile and SAME operand path as +// the accepted w8a8_dumma_prefill_packedb_kernel<128, 64, 128>: grid +// (N/64) x (M/128) = 41 x 32 = 1312 blocks, 256 threads = 4 wavefronts, one +// 64 x 32 quadrant of eight m16n16k32 int32 accumulators per wave, packed +// n-major B, explicit load_frag8 reads, coalesced epilogue. The ONLY change is +// the K pipeline (see the Round-2 section of the file header): +// * 64-K DOUBLE-buffered stages (2 x 13,824 = 27,648 B LDS/block, still +// 2 blocks/CU = 8 waves/CU): the stage-(s+1) global loads are issued into +// 3 int4 staging registers per thread BEFORE the stage-s consume and +// flushed into the idle LDS buffer after it, so the VMEM latency of stage +// s+1 overlaps the LDS/v_mmac work of stage s (one __syncthreads per +// stage instead of two); +// * k0-outer over 96 ASCENDING 64-K stages with kk-inner ascending over two +// m16n16k32 steps, i.e. the same (k0 + kk) k-chunk sequence +// 0, 32, 64, 96, ... as the accepted 48 x 128-K stages, so every int32 +// accumulator is bit-identical. +// kLdsStride = StageK + 8 = 72 B: 8-aligned with an odd 8-byte count (9) so +// the 16 rows of each 16-lane ds_read_b64 phase land on 16 distinct bank pairs +// (18r mod 32, r = 0..15) -- the conflict-free 4-phase floor -- and +// 72 % 16 == 8 keeps the two-half ds_write2_b64 staging form. +// +// Round 5 additionally reads the packed B operand from the fragment- +// interleaved swizzle (see w8a8_pack_bswz_i8_kernel and the Round-5 header +// section): the LDS b_tile is now the SAME [k32][n16][k8][n][8] layout as the +// pack (2 x 256 16-byte chunks per 64-K stage), so staging is one flat +// 16-byte copy per thread and each col_major B fragment is one conflict-free +// ds_read_b64 (load_frag8_swz). The int32 accumulation order is untouched. +// +// Round 9 adds the compile-time `kFullTile` flag: it is true exactly when the +// caller's m is a multiple of BM, so every block tile is fully inside the +// output. In that instantiation the A staging loads drop the row guard (see +// load_a_stage_vec) and the epilogue drops the per-fragment row mask (see +// store_prefill_fragment_coalesced); false keeps the guarded form +// byte-for-byte. Geometry, resources, operand path and accumulation order are +// identical in both instantiations. +template __global__ __launch_bounds__(kThreadsPerBlock) void w8a8_dumma_prefill_packedb_db_kernel( - const int8_t* __restrict__ x_q, // [M, K] row-major - const int8_t* __restrict__ packed_w, // [N/64][K][64] B panels (exact shape, iter 5) - const float* __restrict__ x_scale, // [M] + const int8_t* __restrict__ x_q, // [M, K] row-major + const int8_t* __restrict__ packed_w, // [n>>4][k>>5][512 B] swizzled B + // fragment blocks (round-5 layout) + const float* __restrict__ x_scale, // [M] const float* __restrict__ weight_scale, // [N] - hip_bfloat16* __restrict__ out, // [M, N] + hip_bfloat16* __restrict__ out, // [M, N] int m, int n, int k) { @@ -747,19 +1036,19 @@ w8a8_dumma_prefill_packedb_db_kernel( "block tile must be a multiple of the wave quadrant"); static_assert(kThreadsPerBlock % kWaveSize == 0, "blockDim must be a multiple of the gfx928 wavefront size"); - // Bank-safe LDS row stride: StageK + 8 bytes = 18 dwords for StageK = 64. constexpr int kLdsStride = StageK + 8; static_assert(kLdsStride % 8 == 0 && (kLdsStride / 8) % 2 == 1, "LDS row stride must be 8-aligned with an odd 8-byte count"); static_assert(kLdsStride % 16 == 8, "LDS row stride must be 8 mod 16 for the two-half staging"); - // Iteration 5: the packed B weight is the panel layout (see - // w8a8_pack_panel_i8_kernel), which is defined for BN == 64 columns per - // panel and 16-byte k-runs (StageK % 16 == 0). Only <128,64,64> is ever - // dispatched by the exact-shape guard, so the panel index math below is - // exact for every dispatch. - static_assert(BN == 64 && StageK % 16 == 0, - "panel B pack requires BN == 64 and StageK % 16 == 0"); + // Round 5: the packed B weight is the fragment-interleaved swizzle (see + // w8a8_pack_bswz_i8_kernel), defined for BN == 64 and StageK == 64 (one + // 32-k x 16-n fragment block per 512 B; 2 k32 blocks x 4 n16 groups x 4 k8 + // groups x 8 n-pairs = 256 16-byte chunks per stage, one per thread). Only + // <128, 64, 64> is ever dispatched by the exact-shape guard, so the chunk + // index math below is exact. + static_assert(BN == 64 && StageK == 64, + "swizzled B pack requires BN == 64 and StageK == 64"); const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; @@ -769,9 +1058,12 @@ w8a8_dumma_prefill_packedb_db_kernel( const int m0 = static_cast(blockIdx.y) * BM; const int n0 = static_cast(blockIdx.x) * BN; const int kStages = k / StageK; // dispatch guarantees k % StageK == 0 + const int n16g = n >> 4; // 16-n row groups per 32-k block of the pack __shared__ __align__(16) int8_t a_tile[2][BM * kLdsStride]; - __shared__ __align__(16) int8_t b_tile[2][BN * kLdsStride]; + // Round-5 swizzled B tile: flat [k32][n16][k8][n_pair][16 B] = 4,096 B per + // 64-K stage, byte-identical to the packed global layout of that stage. + __shared__ __align__(16) int8_t b_tile[2][BN * StageK]; DUFragment a_frag[kWaveM16]; @@ -787,10 +1079,10 @@ w8a8_dumma_prefill_packedb_db_kernel( } } - // Flat 16-byte staging vectors per stage. For the dispatched <128,64,64> + // Flat 16-byte staging vectors per stage. For the dispatched <128, 64, 64> // instantiation the 512 A vectors and 256 B vectors divide evenly over the - // 256 threads (2 + 1 per thread), so the prefetch payload is 3 int4 - // staging registers held live across the consume. + // 256 threads (2 + 1 per thread), so the prefetch payload is 3 int4 staging + // registers held live across the consume. constexpr int kVecsPerRow = StageK / static_cast(sizeof(int4)); constexpr int kAVectors = BM * StageK / static_cast(sizeof(int4)); constexpr int kBVectors = StageK * BN / static_cast(sizeof(int4)); @@ -799,12 +1091,16 @@ w8a8_dumma_prefill_packedb_db_kernel( static_assert(kAVectors % kThreadsPerBlock == 0 && kBVectors % kThreadsPerBlock == 0, "staging vectors must divide evenly across the block"); + // Round 5: the swizzled B stage is exactly one 16-byte chunk per thread, so + // thread `vec` = tid is also the flat [k32][n16][k8][n_pair] chunk index. + static_assert(kBVecPerThread == 1 && kBVectors == kThreadsPerBlock, + "swizzled B staging must be one 16-B chunk per thread"); int4 a_pre[kAVecPerThread]; int4 b_pre[kBVecPerThread]; - // ---- Prologue: issue the stage-0 global loads, flush them into buffer 0 - // and make the tile visible to every wavefront. + // ---- Prologue: issue the stage-0 global loads, flush them into buffer 0 and + // make the tile visible to every wavefront. { const int k0 = 0; #pragma unroll @@ -813,25 +1109,22 @@ w8a8_dumma_prefill_packedb_db_kernel( const int local_row = vec / kVecsPerRow; const int k16 = vec - local_row * kVecsPerRow; const int g_row = m0 + local_row; - a_pre[v] = (g_row < m) - ? *reinterpret_cast( - x_q + static_cast(g_row) * k + k0 + - k16 * 16) - : int4{0, 0, 0, 0}; + a_pre[v] = load_a_stage_vec( + x_q, static_cast(g_row) * k + k0 + k16 * 16, m, g_row); } #pragma unroll for (int v = 0; v < kBVecPerThread; ++v) { const int vec = tid + v * kThreadsPerBlock; - const int n_row = vec / kVecsPerRow; - const int k16 = vec - n_row * kVecsPerRow; - // Iteration 5: B is read from the panel-packed layout; the [64, 64] - // tile is 4 contiguous 1-KiB slabs (one per 16-k-run) and vector - // (n_row, k16) sits at byte (n0>>6)*(k*64) + ((k0>>4)+k16)*1024 + - // n_row*16 (16-byte aligned, contiguous k-run of column n0+n_row). - const int64_t panel_off = - static_cast(n0 >> 6) * (static_cast(k) * 64) + - (static_cast(k0 >> 4) + k16) * 1024 + n_row * 16; - b_pre[v] = *reinterpret_cast(packed_w + panel_off); + // Round 5: the 16-byte chunk (k32 = vec>>7, n16 = (vec>>5)&3, + // k8 = (vec>>3)&3, n_pair = vec&7) of the swizzled pack sits at byte + // ((k0>>5)+k32)*n16g*512 + ((n0>>4)+n16)*512 + k8*128 + n_pair*16; + // consecutive lanes read consecutive 16-byte chunks, so a warp's + // staging loads touch 8 fully-used 128-byte lines. + b_pre[v] = *reinterpret_cast( + packed_w + + packed_b_swz_offset((k0 >> 5) + (vec >> 7), + (n0 >> 4) + ((vec >> 5) & 3), n16g, + (vec >> 3) & 3, vec & 7)); } #pragma unroll for (int v = 0; v < kAVecPerThread; ++v) { @@ -847,26 +1140,22 @@ w8a8_dumma_prefill_packedb_db_kernel( #pragma unroll for (int v = 0; v < kBVecPerThread; ++v) { const int vec = tid + v * kThreadsPerBlock; - const int n_row = vec / kVecsPerRow; - const int k16 = vec - n_row * kVecsPerRow; - const int64_t* src64 = reinterpret_cast(&b_pre[v]); - int64_t* dst64 = reinterpret_cast( - b_tile[0] + n_row * kLdsStride + k16 * 16); - dst64[0] = src64[0]; - dst64[1] = src64[1]; + // Round 5: the swizzled tile is a flat copy of the packed stage, so the + // commit is ONE 16-byte-aligned int4 (ds_write_b128) per thread. + *reinterpret_cast(b_tile[0] + vec * 16) = b_pre[v]; } } __syncthreads(); - // ---- Steady state: for s = 0..kStages-2 issue the stage-(s+1) global - // loads into the staging registers BEFORE consuming stage s from buffer - // s&1, then flush the prefetched vectors into the idle buffer (s+1)&1 and - // close the stage with ONE __syncthreads (the barrier's vmcnt wait lands - // after the consume, so the VMEM latency of stage s+1 overlaps the - // LDS/MMAC work of stage s). The last stage is consumed after the loop - // without a trailing barrier. + // ---- Steady state: for stage s, issue the stage-(s+1) global loads into the + // staging registers BEFORE consuming stage s from buffer s&1, then flush the + // prefetched vectors into the idle buffer (s+1)&1 and close the stage with + // ONE __syncthreads (so the vmcnt wait for stage s+1 lands after the MMAC + // burst of stage s, overlapping the VMEM latency of s+1 with the stage-s + // LDS/v_mmac work). The last stage is consumed after the loop without a + // trailing barrier. for (int s = 0; s < kStages - 1; ++s) { - // Prefetch stage s+1: issue the global loads (kept in VGPR while the + // Prefetch stage s+1: issue the global loads (held in VGPRs while the // stage-s consume runs). const int k0 = (s + 1) * StageK; #pragma unroll @@ -875,29 +1164,28 @@ w8a8_dumma_prefill_packedb_db_kernel( const int local_row = vec / kVecsPerRow; const int k16 = vec - local_row * kVecsPerRow; const int g_row = m0 + local_row; - a_pre[v] = (g_row < m) - ? *reinterpret_cast( - x_q + static_cast(g_row) * k + k0 + - k16 * 16) - : int4{0, 0, 0, 0}; + a_pre[v] = load_a_stage_vec( + x_q, static_cast(g_row) * k + k0 + k16 * 16, m, g_row); } #pragma unroll for (int v = 0; v < kBVecPerThread; ++v) { const int vec = tid + v * kThreadsPerBlock; - const int n_row = vec / kVecsPerRow; - const int k16 = vec - n_row * kVecsPerRow; - // Iteration 5: panel-packed B tile read (see the prologue comment). - const int64_t panel_off = - static_cast(n0 >> 6) * (static_cast(k) * 64) + - (static_cast(k0 >> 4) + k16) * 1024 + n_row * 16; - b_pre[v] = *reinterpret_cast(packed_w + panel_off); + // Round 5: the 16-byte chunk (k32 = vec>>7, n16 = (vec>>5)&3, + // k8 = (vec>>3)&3, n_pair = vec&7) of the swizzled pack sits at byte + // ((k0>>5)+k32)*n16g*512 + ((n0>>4)+n16)*512 + k8*128 + n_pair*16; + // consecutive lanes read consecutive 16-byte chunks, so a warp's + // staging loads touch 8 fully-used 128-byte lines. + b_pre[v] = *reinterpret_cast( + packed_w + + packed_b_swz_offset((k0 >> 5) + (vec >> 7), + (n0 >> 4) + ((vec >> 5) & 3), n16g, + (vec >> 3) & 3, vec & 7)); } // Consume stage s from buffer s&1: each wave accumulates its (BM/2) x // (BN/2) quadrant with kk-ascending m16n16k32 MMAs; one ds_read_b64 per // fragment. const int local_row = wave_row * (BM / 2); - const int local_col = wave_col * (BN / 2); #pragma unroll for (int kk = 0; kk < StageK; kk += kTileK) { #pragma unroll @@ -905,14 +1193,19 @@ w8a8_dumma_prefill_packedb_db_kernel( load_frag8( a_frag[j], a_tile[s & 1] + (local_row + j * kTileM) * kLdsStride + kk, - kLdsStride); + static_cast(kLdsStride), + lane); } #pragma unroll for (int i = 0; i < kWaveN16; ++i) { - load_frag8( + // Round 5: fragment block (k32 = kk>>5, n16 = wave_col*kWaveN16 + i) + // of the swizzled tile -- 512 B per 16-n x 32-k DUMMA B fragment, one + // conflict-free ds_read_b64 per lane. + load_frag8_swz( b_frag[i], - b_tile[s & 1] + (local_col + i * kTileN) * kLdsStride + kk, - kLdsStride); + b_tile[s & 1] + + ((kk >> 5) * (BN / 16) + wave_col * kWaveN16 + i) * 512, + lane); } #pragma unroll for (int j = 0; j < kWaveM16; ++j) { @@ -940,13 +1233,8 @@ w8a8_dumma_prefill_packedb_db_kernel( #pragma unroll for (int v = 0; v < kBVecPerThread; ++v) { const int vec = tid + v * kThreadsPerBlock; - const int n_row = vec / kVecsPerRow; - const int k16 = vec - n_row * kVecsPerRow; - const int64_t* src64 = reinterpret_cast(&b_pre[v]); - int64_t* dst64 = reinterpret_cast( - b_tile[buf] + n_row * kLdsStride + k16 * 16); - dst64[0] = src64[0]; - dst64[1] = src64[1]; + // Round 5: flat 16-byte swizzled-tile commit (ds_write_b128). + *reinterpret_cast(b_tile[buf] + vec * 16) = b_pre[v]; } } __syncthreads(); @@ -957,7 +1245,6 @@ w8a8_dumma_prefill_packedb_db_kernel( { const int s = kStages - 1; const int local_row = wave_row * (BM / 2); - const int local_col = wave_col * (BN / 2); #pragma unroll for (int kk = 0; kk < StageK; kk += kTileK) { #pragma unroll @@ -965,14 +1252,19 @@ w8a8_dumma_prefill_packedb_db_kernel( load_frag8( a_frag[j], a_tile[s & 1] + (local_row + j * kTileM) * kLdsStride + kk, - kLdsStride); + static_cast(kLdsStride), + lane); } #pragma unroll for (int i = 0; i < kWaveN16; ++i) { - load_frag8( + // Round 5: fragment block (k32 = kk>>5, n16 = wave_col*kWaveN16 + i) + // of the swizzled tile -- 512 B per 16-n x 32-k DUMMA B fragment, one + // conflict-free ds_read_b64 per lane. + load_frag8_swz( b_frag[i], - b_tile[s & 1] + (local_col + i * kTileN) * kLdsStride + kk, - kLdsStride); + b_tile[s & 1] + + ((kk >> 5) * (BN / 16) + wave_col * kWaveN16 + i) * 512, + lane); } #pragma unroll for (int j = 0; j < kWaveM16; ++j) { @@ -990,7 +1282,9 @@ w8a8_dumma_prefill_packedb_db_kernel( for (int j = 0; j < kWaveM16; ++j) { #pragma unroll for (int i = 0; i < kWaveN16; ++i) { - store_prefill_fragment_coalesced( + // Round 9: kGuarded == !kFullTile -- the full-tile instantiation skips + // the dead per-fragment `row >= m` EXEC bracket (see the helper above). + store_prefill_fragment_coalesced( acc[j][i], x_scale, weight_scale, out, m, n, base_row + j * kTileM, base_col + i * kTileN, lane); } @@ -999,23 +1293,37 @@ w8a8_dumma_prefill_packedb_db_kernel( // Generic scalar int8/int32 fallback for unmatched shapes and small-M API // cases (including the paired M=2 and M=16 shapes with the same (N, K) -- -// they never reach the m >= 128 tiled path). One output element per -// grid-stride step; exact int32 accumulation (max assigned K = 6144 keeps -// the int8 dot well within int32 range: |dot| <= 6144*127*127 < 2^31), then -// the fused float scale (dot * x_scale[m] * weight_scale[n]) and bf16 store. -// The bootstrap pack is the identity copy, so packed_weight is always the -// logical [K, N] row-major layout here. +// they never reach the m >= 128 tiled/packed paths). One output element per +// grid-stride step; exact int32 accumulation (max assigned K = 6144 keeps the +// int8 dot well within int32 range: |dot| <= 6144*127*127 < 2^31), then the +// fused float scale (dot * x_scale[m] * weight_scale[n]) and a manual uint16 +// bf16 store. +// +// Round 1: for the exact (k, n) == (6144, 2624) the packed_weight buffer +// holds a packed layout (produced once, outside the timed region), so the +// fallback decodes that layout there with ascending kk (every other (K, N) +// keeps the identity [K, N] row-major copy). The int32 accumulation stays +// kk-ascending 0..k-1 in both arms; integer addition is exact and commutative, +// so the int32 dot -- and the paired M=2 / M=16 API shapes -- remain +// byte-exact against the identity-layout reference. Round 5 replaces the +// exact-shape pack with the fragment-interleaved swizzle (see +// w8a8_pack_bswz_i8_kernel), so the exact-shape decode below walks +// packed[((kk>>5)*n16g + (col>>4))*512 + +// (((kk>>3)&3)*16 + (col&15))*8 + (kk&7)], n16g = n>>4, +// element-by-element in ascending kk -- identical elements, identical int32 +// accumulators. __global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( - const int8_t* __restrict__ a, // [M, K] row-major - const int8_t* __restrict__ b, // [K, N] row-major (identity pack) - const float* __restrict__ x_scale, // [M] + const int8_t* __restrict__ a, // [M, K] row-major + const int8_t* __restrict__ b, // packed weight (swizzled/identity) + const float* __restrict__ x_scale, // [M] const float* __restrict__ weight_scale, // [N] - hip_bfloat16* __restrict__ out, // [M, N] + hip_bfloat16* __restrict__ out, // [M, N] int m, int n, int k) { const int64_t total = static_cast(m) * n; const int64_t stride = static_cast(gridDim.x) * blockDim.x; + const bool packed_swz = (k == kExactK && n == kExactN); for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < total; @@ -1023,38 +1331,29 @@ __global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( const int row = static_cast(idx / n); const int col = static_cast(idx - static_cast(row) * n); const int8_t* a_row = a + static_cast(row) * k; - // Iteration 5: for the exact (k, n) == (6144, 2624) the packed_weight - // buffer holds the B-panel layout packed[(col>>6)*(k*64) + (kk>>4)*1024 - // + (col&63)*16 + (kk&15)] = raw[kk*n + col] (41 panels of [K, 64] with - // 16-byte k-runs per column, produced once out of the timed region), so - // the fallback decodes that layout there in ascending-kk 16-byte runs - // (keeps the paired M=2/M=16 API shapes with the same (N, K) - // byte-exact); every other (k, n) keeps the identity [K, N] row-major - // copy. The int32 accumulation stays kk-ascending 0..k-1 in both arms. - const bool packed_panel = (k == kExactK && n == kExactN); int32_t acc = 0; - if (packed_panel) { - const int8_t* panel = - b + static_cast(col >> 6) * (static_cast(k) * 64) + - (col & 63) * 16; - for (int kkb = 0; kkb < k; kkb += 16) { - const int8_t* p = panel + (kkb >> 4) * 1024; - for (int bb = 0; bb < 16; ++bb) { - acc += static_cast(a_row[kkb + bb]) * - static_cast(p[bb]); - } + if (packed_swz) { + // Round-5 swizzle: one 8-byte k-run per (n, k8) inside a 512-byte + // 32-k x 16-n fragment block; walk it in ascending kk. + const int64_t n16g = static_cast(n) >> 4; + for (int kk = 0; kk < k; ++kk) { + const int64_t off = + ((static_cast(kk >> 5) * n16g + (col >> 4)) << 9) + + ((((kk >> 3) & 3) * 16 + (col & 15)) << 3) + (kk & 7); + acc += static_cast(a_row[kk]) * + static_cast(b[off]); } } else { const int8_t* b_ptr = b + col; for (int kk = 0; kk < k; ++kk) { - acc += static_cast(a_row[kk]) * - static_cast(b_ptr[0]); + acc += static_cast(a_row[kk]) * static_cast(b_ptr[0]); b_ptr += n; } } - const float scaled = static_cast(acc) * - x_scale[row] * weight_scale[col]; - out[idx] = __float2bfloat16(scaled); + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + reinterpret_cast(out)[idx] = + static_cast(__float2bfloat16(scaled)); } } @@ -1073,14 +1372,13 @@ __global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( } } -// Exact-shape n-major pack (iteration 2): for (k, n) == (6144, 2624) the -// weight is transposed once, outside the timed region and out of Graph -// capture, into packed[n * k + kk] = raw[kk * n + n_row] so every col_major -// B fragment is 8 contiguous K bytes in the packed buffer. Same byte count -// and buffer as the identity pack, so captured addresses are unchanged. One -// thread per output byte; runs once during weight prep. (Kept compiled as -// the iteration-2/3 layout control; iteration 5 dispatches the panel pack -// below for the exact shape.) +// Round-1 exact-shape n-major pack: for (k, n) == (6144, 2624) the weight is +// transposed once, outside the timed region and out of Graph capture, into +// packed[n * k + kk] = raw[kk * n + n] so every col_major B fragment is 8 +// contiguous K bytes in the packed buffer and every staging vector is one +// coalesced 16-byte read. Same byte count and buffer as the identity pack, so +// the captured addresses are unchanged. One thread per output byte; runs once +// during weight preparation. __global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, @@ -1098,36 +1396,52 @@ __global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( } } -// Exact-shape B-panel pack (iteration 5): for (k, n) == (6144, 2624) the -// weight is packed once, outside the timed region and out of Graph capture, -// into 41 panels of [K, 64] int8 with 16-byte k-runs per column: -// packed[(n >> 6) * (k * 64) + (kk >> 4) * 1024 + (n & 63) * 16 + -// (kk & 15)] = raw[kk * n + n] -// (the byte-exact bijection 41 x 384 x 64 x 16 = k*n is verified by -// simulation; same byte count and buffer, so captured addresses are -// unchanged). Every [64, 64] B tile of a 64-K stage is then 4 CONTIGUOUS -// 1-KiB slabs (one per 16-k-run) and every 16-byte staging vector is a -// 16-byte-aligned contiguous k-run of one column, so each DUMMA B tile is -// vector-loadable as one coalesced stream with full 128-B line utilization -// (vs 64 scattered 64-B chunks at 6144-B stride in the n-major layout). -// One thread per output byte; runs once during weight prep. -__global__ __launch_bounds__(256) void w8a8_pack_panel_i8_kernel( +// Round-5 exact-shape fragment-interleaved B swizzle: for (k, n) == +// (6144, 2624) the weight is packed once, outside the timed region and out of +// Graph capture, into 512-byte 32-k x 16-n DUMMA B fragment blocks laid out +// [k32][n16][k8][n][8]: +// packed[((kk >> 5) * n16g + (n >> 4)) * 512 + +// (((kk >> 3) & 3) * 16 + (n & 15)) * 8 + (kk & 7)] = raw[kk * n + n], +// where n16g = n >> 4 (164 for the exact shape). This is a byte-exact +// bijection 192 x 164 x 512 = k*n = 16,121,856: same byte count and buffer as +// the identity/n-major/panel packs, so captured addresses are unchanged. Each +// thread packs one 16-byte chunk = the 8-byte k-runs of two n-rows of one k8 +// group (a strided 2x8 gather on the read side; packing is outside timing, so +// the gather is off the critical path). Consecutive chunks are consecutive +// packed bytes, so the kernel's B staging is a flat 16-byte copy and every +// col_major B fragment read is one conflict-free ds_read_b64 (see +// load_frag8_swz). The round-2 n-major pack kernel above stays compiled as the +// layout control. +__global__ __launch_bounds__(256) void w8a8_pack_bswz_i8_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, int k, int n) { - const int64_t total = static_cast(k) * n; + const int64_t chunks = (static_cast(k) * n) >> 4; const int64_t stride = static_cast(gridDim.x) * blockDim.x; - for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + - threadIdx.x; - idx < total; - idx += stride) { - const int col = static_cast(idx / k); - const int kk = static_cast(idx - static_cast(col) * k); - // (col >> 6) < n/64 <= 41, (kk >> 4) < k/16 <= 384, (col & 63) < 64, - // (kk & 15) < 16 -> the mixed-radix index stays < k*n in int. - packed[(col >> 6) * (k * 64) + (kk >> 4) * 1024 + (col & 63) * 16 + - (kk & 15)] = raw[static_cast(kk) * n + col]; + const int n16g = n >> 4; + for (int64_t c = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + c < chunks; + c += stride) { + const int k32 = + static_cast(c / (static_cast(n16g) * 32)); + const int rem = + static_cast(c - static_cast(k32) * n16g * 32); + const int n16 = rem >> 5; + const int w = rem & 31; + const int k8 = w >> 3; + const int np = w & 7; + const int n0 = (n16 << 4) + (np << 1); + const int k0 = (k32 << 5) + (k8 << 3); + int8_t* dst16 = packed + (c << 4); +#pragma unroll + for (int i = 0; i < 8; ++i) { + // 16-byte chunk = [n0 row: k0..k0+7][n0+1 row: k0..k0+7] (n-major + // within the k8 group, matching the [k8][n][8] tile layout). + dst16[i] = raw[static_cast(k0 + i) * n + n0]; + dst16[8 + i] = raw[static_cast(k0 + i) * n + n0 + 1]; + } } } @@ -1149,7 +1463,9 @@ __global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( // Stable host launch symbol consumed by csrc/bindings.cpp. Dispatches on the // caller-provided stream only; never allocates, synchronizes, packs, or -// touches the default stream. +// touches the default stream. The workspace is unused for this shape +// (single-buffered, no split-K): m*n*4 = 42,990,080 B exceeds the 16 MiB +// workspace budget, so workspace_split_k_capacity is zero. extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -1170,96 +1486,77 @@ extern "C" void launch_w8a8_gemm( auto* out_ptr = static_cast(out); // Exact assigned shape guard: glm_tp8_fused_qkv_a_proj_m4096 - // (M=4096, N=2624, K=6144), m >= 128 so M=2/M=16 paired API shapes with - // the same (N, K) keep reaching the scalar fallback below. Iteration 5: - // dispatch the exact-fit 2-D macro-tile <128, 64, 128> geometry with the - // B-PANEL packed weight (launch_pack_w8a8_weight packs (6144, 2624) once, - // out of the timed region, into packed[(n>>6)*(k*64) + (kk>>4)*1024 + - // (n&63)*16 + (kk&15)] = raw[kk*N + n] so every 64-K B tile is 4 - // contiguous 1-KiB slabs), the bank-safe LDS row strides, the explicit - // 8-byte load_frag8 fragment loads, and the double-buffered 64-K K - // pipeline (one __syncthreads per stage, stage-(s+1) global loads - // prefetched into staging VGPR before the stage-s consume) -- the - // mandated packing/swizzle round at unchanged tile/occupancy (2 - // blocks/CU, 27,648 B LDS/block). The single-buffered iteration-2 - // control kernel stays compiled (see the m < 0 forcing block below) for - // the A/B and for rollback. + // (M=4096, N=2624, K=6144). The m >= 128 guard keeps the paired M=2/M=16 + // API shapes with the same (N, K) on the scalar fallback below. Round 2 + // keeps the accepted packed-weight architecture (the exact-fit 2-D macro-tile + // <128, 64, ...> with the packed weight produced once, out of the + // timed region, by launch_pack_w8a8_weight, the bank-safe LDS row strides and + // the explicit 8-byte load_frag8 fragment loads) and changes ONLY the K-stage + // pipeline: the dispatch goes to the 64-K double-buffered prefetch kernel + // <128, 64, 64> instead of the single-buffered <128, 64, 128> one, so the + // stage-(s+1) VMEM latency overlaps the stage-s LDS/v_mmac work. Round 3 + // retains that double-buffered 64-K pipeline unchanged (it is the unique + // 2-resident-block K pipeline inside the 65,536 B LDS budget: 2 x 27,648 = + // 55,296 <= 65,536) and changes only the packed B layout consumed by the two + // B global-load sites -- the weight is now the B-panel layout, so a warp's B + // staging loads touch 4 fully-used 128-byte lines instead of 32 half-used + // ones. Round 5 keeps the kernel, geometry and occupancy identical and + // changes only the packed B layout again: the B-panel pack is replaced by the + // fragment-interleaved swizzle, so the B stage is a flat 16-byte copy and + // every col_major B fragment is one conflict-free ds_read_b64 + // (load_frag8_swz). M = 32*128 and N = 41*64 are both exact, so no tail + // predication is introduced. Round 9 keeps that kernel, geometry and + // occupancy and only selects a compile-time tail-predication specialization: + // when m is a multiple of BM = 128 every block tile is fully inside the + // output, so the instantiation with kFullTile == true drops the dead A-row + // zero-fill/select in both staging sites and the dead epilogue row mask + // (see the Round-9 header section); every other m keeps the guarded + // instantiation byte-for-byte. if (m >= 128 && n == kExactN && k == kExactK) { const dim3 grid( static_cast(n / 64), static_cast((m + 127) / 128)); const dim3 block(static_cast(kThreadsPerBlock)); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64>), - grid, - block, - 0, - stream, - a, - b, - x_scale, - weight_scale, - out_ptr, - m, - n, - k); - return; - } - - // Force the sibling tile instantiations to compile (<64,128,128> is the - // next benchmark candidate; this branch can never run because m <= 0 - // already returned above). The single-buffered iteration-2 packed-B kernel - // is kept here so the control arm of the iteration-3 A/B stays compiled - // (and rollback only needs the guard repointed). - if (m < 0) { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<64, 128, 128>), - dim3(1), - dim3(static_cast(kThreadsPerBlock)), - 0, - stream, - a, - b, - x_scale, - weight_scale, - out_ptr, - m, - n, - k); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tiled_kernel<128, 64, 128>), - dim3(1), - dim3(static_cast(kThreadsPerBlock)), - 0, - stream, - a, - b, - x_scale, - weight_scale, - out_ptr, - m, - n, - k); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_packedb_kernel<128, 64, 128>), - dim3(1), - dim3(static_cast(kThreadsPerBlock)), - 0, - stream, - a, - b, - x_scale, - weight_scale, - out_ptr, - m, - n, - k); + if ((m & 127) == 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64, true>), + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + } else { + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64, false>), + grid, + block, + 0, + stream, + a, + b, + x_scale, + weight_scale, + out_ptr, + m, + n, + k); + } return; } - // Generic native INT8 DUMMA prefill path for every large-M shape with - // N % 64 == 0 and K % 128 == 0 (identity [K, N] layout). The exact shape - // above already returned; this covers compatible sibling shapes. + // Generic native INT8 DUMMA prefill path for every other large-M shape with + // N % 64 == 0 and K % 128 == 0 (identity [K, N] layout). The A row guard + // zero-fills rows past M and the epilogue masks them, so only the exact + // shape above avoids predication. if (m >= 128 && (n % 64) == 0 && (k % 128) == 0) { const dim3 grid( static_cast(n / 64), @@ -1282,8 +1579,8 @@ extern "C" void launch_w8a8_gemm( return; } - // Scalar int8/int32 fallback for every unmatched (m, n, k) and all - // small-M API cases. + // Scalar int8/int32 fallback for every unmatched (m, n, k) and all small-M + // API cases. const int64_t total = static_cast(m) * n; constexpr int kFallbackThreads = 256; int64_t blocks = (total + kFallbackThreads - 1) / kFallbackThreads; @@ -1311,15 +1608,15 @@ extern "C" void launch_w8a8_gemm( k); } -// Stable host launch symbol consumed by csrc/bindings.cpp. Iteration 5: the -// exact (k, n) == (6144, 2624) int8 weight is packed once (outside the -// timed region and out of Graph capture) into the B-panel layout -// packed[(n>>6)*(k*64) + (kk>>4)*1024 + (n&63)*16 + (kk&15)] = raw[kk*N + n] -// consumed by the packed-B prefill kernel's panel tile reads and by the -// scalar fallback's panel decode; every other (K, N) keeps the identity copy -// (the logical [K, N] row-major weight). The [N] fp32 scales are always -// copied unchanged. Same byte count (k*n and n), same allocation, -// graph-stable addresses. +// Stable host launch symbol consumed by csrc/bindings.cpp. Round 5: the exact +// (k, n) == (6144, 2624) int8 weight is packed once (outside the timed region +// and out of Graph capture) into the fragment-interleaved swizzle +// packed[((kk>>5)*n16g + (n>>4))*512 + (((kk>>3)&3)*16 + (n&15))*8 + (kk&7)] +// = raw[kk*n + n] with n16g = n>>4, consumed by the packed-weight prefill +// kernel's flat B stage copy and by the scalar fallback's element decode; every +// other (K, N) keeps the identity copy of the logical [K, N] row-major weight. +// The [N] fp32 scales are always copied unchanged. Same byte count (k*n and n), +// same allocations, so the captured addresses stay graph-stable. extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -1330,14 +1627,10 @@ extern "C" void launch_pack_w8a8_weight( hipStream_t stream) { constexpr int kPackThreads = 256; const int64_t weight_count = static_cast(k) * n; - // Iteration 5: the exact (k, n) == (6144, 2624) weight is packed once, - // outside the timed region and out of Graph capture, into the B-panel - // layout (same byte count and buffer addresses, so captured pointers stay - // valid; the packed-B prefill kernel and the scalar fallback for this - // (k, n) both decode it). Every other (K, N) keeps the identity copy - // (raw row-major [K, N]). if (k == kExactK && n == kExactN) { - int64_t blocks = (weight_count + kPackThreads - 1) / kPackThreads; + // Round 5: one thread packs one 16-byte chunk (16 bytes = 2 n-rows x 8 B). + const int64_t chunks = weight_count >> 4; + int64_t blocks = (chunks + kPackThreads - 1) / kPackThreads; if (blocks > 4096) { blocks = 4096; } @@ -1345,7 +1638,7 @@ extern "C" void launch_pack_w8a8_weight( blocks = 1; } hipLaunchKernelGGL( - w8a8_pack_panel_i8_kernel, + w8a8_pack_bswz_i8_kernel, dim3(static_cast(blocks)), dim3(kPackThreads), 0, @@ -1393,4 +1686,3 @@ extern "C" void launch_pack_w8a8_weight( } } // @@end -// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip index 7de05137..af1b458c 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip @@ -1,242 +1,320 @@ -// @@variant shape=glm_tp8_kv_b_proj_m4096 commit=bb4aec7f2d04a07b6c3c5cf4c19e1b9fc97fc8a3 added=2026-08-31 -// median_us=137.4 p90_us=137.9 speedup=50.51 baseline_us=6940 -// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// @@variant shape=glm_tp8_kv_b_proj_m4096 commit=547eec13f53e3cee7da02bbdd17763f46e71761e added=2026-09-11 +// median_us=120.1 p90_us=120.9 speedup=57.78 baseline_us=6940 +// source=test-9-11-1-b01955b3/iteration_001/glm_tp8_kv_b_proj_m4096 // INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. // -// Worker: worker_1 (physical GPU 1), assigned shapes: -// glm_tp8_q_b_proj_m4096 : M=4096, N=2048, K=2048 +// Worker: worker_2 (physical GPU 2). Assigned shape: // glm_tp8_kv_b_proj_m4096 : M=4096, N=3584, K=512 // // Logical operation (exact contract): // out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) // * x_scale[m] * weight_scale[n] ) // -// Iteration 3 (first valid HIP experiment; iterations 1-2 were killed in the -// agent infrastructure before any proposal). Accepted-best source digest -// a64b6250... with one bounded mechanism: n-major packed B + col_major B -// fragments for the exact glm_tp8_q_b_proj_m4096 (k,n) == (2048, 2048). -// * Large-prefill path (m >= 128 with exact 128x64x64 geometry): native -// INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 output tile -// per block; four wavefronts (256 threads); each wave owns a 64x32 -// quadrant built from eight m16n16k32 int32 accumulator fragments; the -// block cooperatively vector-loads A[128,64] and B[64,64] into one -// single-buffered LDS stage (15,360 B total: 128*80 + 64*80 with 16 B -// padding per row for bank skew); two __syncthreads per stage; fused -// dot * x_scale[m] * weight_scale[n] epilogue stored directly as bf16 -// from the accumulator fragments using the verified gfx928 int8 -// m16n16k32 lane mapping (row = lane & 15, col = (lane >> 4) + 4*i). -// A is read with the library du_load_matrix_sync row-major loader. -// B is layout-templated: -// - kNMajorB == true (exact (k,n) == (2048, 2048), q_b_proj): weight -// is packed once to n-major packed[n*K + k] == raw[k*N + n] outside -// the timed region; the stage stores B n-major (n row, k contiguous, -// 80 B row stride); each B fragment is one contiguous 8-byte LDS run -// per lane, loaded with the lineage-validated load_fragment8 (one -// ds_read2_b64) - kills the 32 ds_read_u8 byte-gathers + mask/OR -// reassembly VALU seen in the accepted-best ISA. -// - kNMajorB == false (any other shape, e.g. kv_b_proj (512, 3584)): -// raw [K, N] row-major identity layout, row-major B loader (the -// accepted-best path, byte-identical). -// Grid dim3(N/64, M/128) = 1024 blocks (q_b_proj) / 1792 blocks -// (kv_b_proj) dwarfs the 120 CUs, so no split-K is needed. -// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), -// including all small-M API cases (M=2, M=16), M tails, and M in -// (0, 128) with the same (K, N); it decodes the n-major packed layout -// for the exact (2048, 2048) and the identity [K, N] layout otherwise. -// * launch_pack_w8a8_weight: n-major device-to-device permutation for -// (k, n) == (2048, 2048), identity copy for every other (k, n) (both the -// int8 weight and the fp32 scale). Packing never happens inside the -// timed GEMM and keeps the same byte count (graph-stable addresses). +// Iteration 9 (consolidate round): the exact-shape arm keeps its winning +// 128x128 macro-tile, its 512-thread / eight-wavefront geometry, its +// 16,384-B single-buffered swizzled-A + plane-B staging, its direct 8-byte +// fragment loads, its k0-outer/kk-inner int32 order and its coalesced +// 8-byte epilogue. The one consolidated defect is the bf16 CONVERSION +// lowering: the exact gfx928 code object of the current best +// (compile_cache_key 0ece2e0cdbddfe2f512228e2bde5e23189f88fc22d01ae5ae49daf4a97a80f77, +// source_hip_digest 95dab478d6509cacbb6ca0888d90635197ed82d9202aead00ed49f50c4c7c9ac) +// shows hip's __float2bfloat16 emitting, for EVERY element each lane +// converts, an exec-masked inf/NaN fixup - v_and 0x7f800000, v_cmp_ne, +// s_and_saveexec, v_bfe, v_add3, v_or 0x10000, v_cndmask, s_andn2_saveexec +// (code-object addresses 0x5EEC..0x77E0 of +// w8a8_dumma_prefill_128x128_nmajor_kernel) - i.e. ~9 vector instructions +// plus two scalar EXEC manipulations per element against the 3-instruction +// rounding the value actually needs. With M*N = 14.68M converted elements +// and 6,931,456 measured valu_instructions, the block-tail epilogue is the +// single largest VALU consumer in the kernel, and the +// s_and_saveexec/s_andn2_saveexec pairs serialize the wavefront there where +// the only other work left is the second co-resident block. This round +// replaces ONLY that conversion, ONLY in the exact-shape epilogue, with a +// plain RNE form (bf16_rne_u16: u += 0x7fff + ((u >> 16) & 1); u >> 16) - +// the same lineage-validated fix the same-shape sibling reference +// references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip applies on this +// exact DTK 26.04 / gfx928. The formula is bit-identical to hip's own +// non-masked branch for every finite float, and int32 accumulators times +// finite scales can only produce finite floats, so mismatch_count 0 / +// max_abs_error 0 is expected on the unchanged pre-seeded reference data. +// Everything else - the dense swizzle, the tile-contiguous pack and its +// scalar decode, the fragment placement, the k0/kk accumulation order, the +// exact (k, n) == (512, 3584) guard, the identity-layout 128x64 arm, the +// pack kernels and the generic scalar fallback - is byte-identical. // -// Iteration 4 (tile-aspect round): the exact (k, n) == (2048, 2048) q_b_proj -// arm switches from the 128x64 tile to a 128x128 tile -// (w8a8_dumma_prefill_128x128_kernel). With the n-major packed B + 8-byte -// load_fragment8 B reads in place (iteration 3), the remaining LDS-read cost -// is dominated by the A-side row-major fragment loads (4 x ds_read2_b32 per -// kk per wave). A 128x128 tile with four 64x64 quadrants keeps the same -// 4-wave x 8-accumulator-fragment count per wave (8 -> 16 m16n16k32 MMACs -// per kk from 8 LDS fragment reads, i.e. per-MMAC LDS reads drop 6/8 -> 8/16 -// = -33%), and the A-side global re-read halves (A reuse 32 -> 16, A traffic -// 268 -> 134 MiB) while B reuse stays 32 (134 MiB), balancing the per-byte -// A:B traffic at 1:1 and cutting total tile traffic 402 -> 268 MiB (-33%). -// LDS grows 15,360 -> 20,480 B/block (still 2 blocks/CU = 40,960 B <= 64 -// KiB); the grid becomes (N/128, M/128) = 512 blocks (~4.3/CU). The 16 -// accumulator fragments cost ~+32 VGPR; arch VGPR is expected <= 128 so the -// accepted kernel's 2-blocks/CU co-residency is preserved (falsified if the -// exact code object shows > 128 VGPR -> 1 block/CU). kv_b_proj keeps the -// byte-identical 128x64 kNMajorB=false arm and the scalar fallback is -// untouched. +// Iteration 10 (consolidate round): final HIP-only polish on the same winning +// architecture. The one consolidated defect is the dead per-fragment tail-M +// masking in the exact-shape epilogue. The exact gfx928 code object of the +// current best (compile_cache_key 0f7ddb1a6c31e6f3cdd7772a8760e5b20dcea6da9ffca5d556e35790d7e25113, +// source_hip_digest 8c2468b4dfc603de4dd4a957ebf58d891830cbac940a07c58aed7c52030a8def, +// arch_vgpr 64, 0 scratch, 16,384 B LDS, 896 blocks) shows the block tail +// split into eight exec-masked basic blocks +// (s_and_saveexec / s_cbranch_execz / s_or_b64 around each of the eight +// fragment epilogues) even though the dispatch guard for this arm is +// m >= 128 with m % 128 == 0, so every row of every 128-row tile is in bounds +// and neither the (row >= m) early return in the store helper nor the two +// (rs < m) x_scale ternaries in the kernel can ever fire. Because each +// fragment epilogue sits in its own basic block, the compiler also repeats the +// fragment-independent shuffle scaffolding per fragment (v_mbcnt_lo / +// v_mbcnt_hi lane-id, the two __shfl_xor partner indices, the two comparison +// and cndmask pairs, ~13 VALU per fragment x 8). This round removes ONLY that +// dead masking from the exact-shape arm: the store helper drops its +// (row >= m) early return and its now-unused m argument, and the caller loads +// x_scale[rs0] / x_scale[rs1] directly. Nothing else changes: the 128x128 +// macro-tile, the 512-thread / eight-wavefront block, the 16,384-B +// single-buffered swizzled-A + plane-B staging, the direct 8-byte fragment +// loads, the k0-outer/kk-inner int32 order, the 4x4 shfl_xor transpose and +// the single 8-byte bf16 store per lane are all untouched, so the written +// bytes are bit-identical (removing a mask that is never asserted cannot +// change any output). The 128x64 identity-layout arm keeps its own +// store_fragment_scaled bounds check and the generic scalar fallback keeps its +// linear bounds check, both byte-identical. Falsifiable signature: the eight +// exec-masked guard blocks disappear from the 128x128 code object, the +// per-fragment duplicate lane-id / shuffle-index VALU is hoisted, and the +// dynamic valu_instructions fall from 5,512,192 without changing +// lds_instructions 1,089,536, lds_bank_conflicts 458,752, +// vmem_write_instructions 57,344 or the 16 v_mmac per stage. // -// Iteration 5 (packing round): the exact (k, n) == (2048, 2048) q_b_proj -// weight is re-packed once (outside timing/Graph, same byte count, same -// graph-stable buffer) into a swizzled 64-k-stage-major layout -// packed[((k0*8 + kc)*n + col)*8 + b] == raw[kk*n + col], -// kk = k0*64 + kc*8 + b (k0 = 64-k stage, kc = 8-byte sub-chunk, b = byte), -// and the 128x128 kernel stages each B tile as 16-byte plane chunks into an -// LDS plane layout b_tile[kc][n][8] (plane stride 1024 B, n stride 8 B). -// The exact code object of the accepted kernel shows every B fragment read -// is an 8-byte LDS access at (lane&15)*80 + (lane>>4)*8 (n-major stride-80 -// rows), and PMC counts 6,291,456 bank conflicts over 1,048,576 LDS slots -// (6.0/slot) with 4,078,494 LDS waits (3.9/slot): the loop is LDS-latency -// bound and the stride-80 rows alias bank phases every 8 rows (80*8 == 640 -// == 0 mod 128), so every fragment read conflicts ~4-6-way. In the plane -// layout a lane's 8 bytes sit at (lane>>4)*1024 + (lane&15)*8 relative to -// the fragment origin: 16 lanes per 128-B phase cover all 32 banks exactly -// once, i.e. B fragment reads become lane-linear and zero-conflict (4-cycle -// minimum for 512 B). Staging reads stay coalesced (each 16-B chunk is two -// n rows of one sub-chunk -> one aligned int4) and LDS writes are 2-way -// (the 16-B optimum). Same fragment operand bytes, same k0-outer/kk-inner -// int32 accumulation -> bit-identical results. The scalar fallback decodes -// the new pack for (2048,2048) only; kv_b_proj keeps the byte-identical -// 128x64 kNMajorB=false identity path. +// Iteration 13 (consolidate round, 13/13): final HIP-only polish on the +// unchanged winning exact-shape architecture, one bounded mechanism. The +// exact current-best code object (source_hip_digest +// 8010314fdb6ba5ec0da98d662fc3911763b088245fa7c316bf0986572499207d, +// compile_cache_key 3fbcc466cfb5087e6f7057b3f26684e598046f28d0f71e0d794962688901c1c7, +// arch_vgpr 64, scratch 0, 16,384 B static LDS, 896 blocks x 512 threads) puts +// the whole 64-K stage wall on the EXPOSED cooperative-staging round trip: the +// two global_load_dwordx4 A/B chunks are issued only after the previous stage's +// trailing s_barrier and are then drained within ~8 instructions by +// s_waitcnt vmcnt(1) / vmcnt(0) in front of the three ds_write_b64/b128 +// commits (0x5D1C..0x5D5C), so no part of the A+B L2 latency overlaps any +// arithmetic or the barrier wait. Iteration 7 raised the staging depth to 2 to +// overlap exactly that round trip, but it doubled LDS to 32,768 B/block and +// gave up the 2-blocks/CU residency point, so the prefetch mechanism itself was +// never measured. This round re-introduces ONLY the single-buffered register +// prefetch: the A row base and the B per-stage stride are hoisted out of the K +// loop, the next stage's A and B int4 chunks are loaded BEFORE the trailing +// barrier and consumed by the next stage's ds_write commits, while LDS stays +// 16,384 B, __launch_bounds__(kNMBlockThreads, 2), the 64-VGPR residency point +// and the two barriers per stage are all untouched, and the last stage re-reads +// its own chunk (k1 == k0) instead of reading past the last K block. The +// swizzled A tile, the tile-contiguous plane-B layout, the direct 8-byte +// fragment loads, the k0-outer/kk-inner ascending int32 order, the coalesced +// RNE epilogue, the exact (k, n) == (512, 3584) dispatch guard, the 128x64 +// identity-layout arm, the pack kernels and the generic scalar fallback are +// byte-identical, so every stored bit is unchanged. Falsifiable signature: the +// global_load_dwordx4 pair for stage s+1 is scheduled before the trailing +// s_barrier and its vmcnt drain moves to the next stage's top, while arch_vgpr +// stays 64 with vgpr_spill_count 0 and scratch_bytes 0; if the machine +// scheduler keeps the loads where they already are the object is unchanged and +// the round is flat. // -// Iteration 6 (epilogue round): the fused dot * x_scale[m] * weight_scale[n] -// -> bf16 epilogue has been in-kernel since iteration 1 and the workspace is -// unused ((void)workspace; no split-K, no combine pass anywhere in the call -// chain), so the remaining epilogue inefficiency of the exact (2048, 2048) -// 128x128 arm is the STORE pattern: store_prefill_fragment issues one -// 2-byte bf16 store per lane per element - 64 scattered stores per wave per -// fragment set, each wavefront store touching 16 rows x 4 columns so every -// 32-B sector is only 25% utilized (PMC: 131,072 vmem_write_instructions = -// 512 blocks x 4 waves x 64). The new epilogue transposes each 16x16 -// fragment's 4-element groups inside the 4-lane column group (lanes r, -// r+16, r+32, r+48; two 2x2 steps with __shfl_xor 16 then 32 - the -// lineage-validated 4x4 register transpose accepted on the sibling TP8 -// workers 0/2, same DTK 26.04/gfx928), so lane (r, c4) owns the four -// CONTIGUOUS columns 4*c4 .. +3 and writes ONE 8-byte store per lane (100% -// store sector efficiency; vmem_write 131,072 -> 32,768). Only int32 -// values move between lanes: the per-element multiply order -// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are -// unchanged, so stored bits are identical. The per-row x_scale (4 rows per -// lane) and per-column weight_scale float4s (4 per lane) are additionally -// register-batched on the last K stage (before the final protective -// __syncthreads), so the epilogue is pure compute + 16 coalesced stores with -// no interleaved vmem loads (vmem_read ~413,696 -> ~278k). kv_b_proj keeps -// the byte-identical 128x64 kNMajorB=false arm (old per-element epilogue) -// and the generic scalar fallback is untouched. +// Iteration 5 (packing round): the exact (k, n) == (512, 3584) weight is +// re-packed once, outside the timed/Graph region, from the n-major [N, K] +// transpose to a TILE-CONTIGUOUS 64-K-stage layout +// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] = raw[kk*n + col] +// kk = k0*64 + kc*8 + b, nt = col/128, row = col%128, +// i.e. every (64-K stage, 128-N tile) DUMMA B tile is ONE contiguous 8192-B +// global region laid out as [kc = 8-byte k sub-chunk][128 N rows][8 bytes], +// which is exactly the plane order the B fragments are read from in LDS. The +// GEMM's cooperative B staging therefore becomes a plain 2 x 16-B contiguous +// vector copy per thread (global_load_dwordx4 -> ds_write_b128) with +// consecutive lanes on consecutive 16-B slots: the global reads are fully +// coalesced (the n-major [N, K] layout made a wave touch 16 rows 512 B apart) +// and the 16-B LDS stores have no bank conflicts. The B fragment load becomes +// lane-linear (kc*(128*8) + row*8, no swizzle arithmetic) and each 16-lane +// phase covers all 32 banks exactly once, so B is now both vector-loadable and +// LDS-bank-safe; the swizzled 8-byte-group map is retained for A only (x_q is +// a runtime activation and cannot be pre-packed). The int32 accumulation is +// untouched: identical per-lane fragment placement (row = row0 + (lane & 15), +// k group kc = (lane >> 4) + (kk >> 3), 8 consecutive natural-order k bytes) +// and identical k0-outer/kk-inner ascending order, so every output byte is +// bit-identical. Falsifiable signature: the B half of the 917,504 measured +// lds_bank_conflicts disappears (32 -> 16 per stage per wave), the B staging's +// 4 x ds_write_b64 become 2 x ds_write_b128, lds_instructions fall +// 688,128 -> 630,784, and the median improves only if the global/LDS staging +// latency chain, not the 16-MMAC burst, is the limiter. // -// Iteration 7 (compute-pipeline round): the exact (k, n) == (2048, 2048) -// 128x128 arm loads its four A fragments per kk with the lineage-validated -// direct load_fragment8 (one ds_read2_b64 straight into the v_mmac operand) -// instead of the library du_load_matrix_sync row_major loader. The exact -// code object shows the library loader lowers to 8 x ds_read2_b32 plus a -// redundant byte-reassembly chain per fragment (~7 VALU: v_and 0xff00 / -// 0xff0000 / 0xff000000 + v_or_b32_sdwa + v_or3, ~50-56 VALU per stage per -// wave) between the LDS read and the first MMAC, and PMC counts 3,009,148 -// LDS waits (2.3/slot) against a latency-bound loop running at 1 block/CU -// (160 VGPR). du_mma.hpp defines matrix_a row_major int8 as x[i] = -// p[(lane&15)*ldm + (lane>>4)*8 + i] (8 consecutive bytes, memory order) -// and du_mma_sync passes reinterpret(x) unchanged to v_mmac, so -// load_fragment8 produces byte-identical operand values and the int32 -// accumulation is bit-identical; only the redundant VALU reassembly (and its -// lgkmcnt wait states) is removed from the LDS->MMAC critical path, letting -// the 32-MMAC burst issue back-to-back. Expected PMC: valu_instructions -// 8.66M -> ~4.8-5.6M (-35..-45%), lds_instructions 1.31M -> ~1.05M (-20%, -// 8 A ds_read2_b32 + 4 B ds_read2_b64 -> 8 ds_read2_b64 per stage), -// lds_wait_instructions 3.01M -> ~2.2-2.5M, lds_bank_conflicts -// approximately unchanged (A rows keep the 2-way stride-80 aliasing), -// vmem_read/vmem_write and the 32 v_mmac/stage unchanged. kv_b_proj keeps -// the byte-identical 128x64 kNMajorB=false arm (unchanged lowering) and the -// generic scalar fallback is untouched. +// Iteration 6 (mandated epilogue round): the fused scale/bf16 epilogue has +// been inside the compute kernel since iteration 1 and the workspace is +// already ignored ((void)workspace in launch_w8a8_gemm), so the remaining +// epilogue inefficiency of the exact (512, 3584) arm is the STORE pattern: +// store_fragment_scaled issues one 2-byte bf16 store per lane per element - +// 64 scattered stores per wave, each wavefront store touching 16 rows x 4 +// columns so every 32-B store sector is only 25% utilized (the bootstrapped +// PMC counts vmem_write_instructions == 229,376 == 896 x 4 x 16 x 4 exactly). +// The new epilogue is register-batched: the per-row x_scale is loaded once per +// 16-row row-block and the four per-column scales of each fragment's owned +// columns are one aligned float4 load, so the epilogue issues 4 scalar + 16 +// vector vmem loads per lane instead of 80 scalar loads. It then transposes +// the four-element groups within each 4-lane column group in registers (two +// 2x2 shfl_xor steps, 16 then 32 - the lineage-validated pattern ported from +// references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip, which lowers to +// ds_bpermute at the block tail where the LDS pipe is idle) so lane (r, c4) +// ends up holding the four CONTIGUOUS columns 4*c4 .. +3, converts them to +// bf16 and writes ONE 8-byte store per lane per fragment: 16 coalesced stores +// per wave instead of 64 scattered ones at 100% sector efficiency, i.e. +// vmem_write 229,376 -> 57,344. The scale multiply order +// (float(dot) * x_scale[row] * weight_scale[col]) and the __float2bfloat16 +// rounding are unchanged, so the stored bits are identical; no LDS, no +// barrier, and no workspace/combine pass is added. Everything else - the A +// swizzle, the tile-contiguous plane-B staging, the k0-outer/kk-inner +// ascending int32 order, the exact (k, n) == (512, 3584) guard, the +// identity-layout 128x64 arm, the pack kernels and the generic scalar +// fallback - is byte-identical. // -// Iteration 1 (kv_b_proj tile-baseline round): the exact (k, n) == (512, -// 3584) kv_b_proj arm is still the iteration-2 bootstrap pipeline (library -// fragment loaders, identity [K, N] B, per-element scattered epilogue) and -// runs at 65 TOPS vs 156 TOPS on the fully-modernized q_b_proj arm (same -// GPU, same DTK, measured 220.14 us on K=2048/N=2048). This round ports the -// proven pipeline (swizzled plane-pack B, direct 8-byte A/B fragment loads, -// register-batched coalesced epilogue, fully unrolled kk) to kv_b_proj as -// one templated 2-D macro-tile family (w8a8_dumma_prefill_tile_kernel) over -// the three mandated block tiles plus the iteration-4 aspect flip, all -// compiled, active selected by kActiveKvTile (0 = 64x64, 1 = 64x128, -// 2 = 128x64, 3 = 64x128 with 32x32 quadrants): -// config tile waves threads LDS B grid (blocks) est. VGPR -// 0 64x64 1 64 9,216 (56, 64) = 3584 ~128-140 -// 1 64x128 2 128 13,312 (28, 64) = 1792 ~128-140 -// 2 128x64 4 256 14,336 (56, 32) = 1792 ~88-100 -// 3 64x128 8 512 13,312 (28, 64) = 1792 ~55-60 -// Per-wave LDS fragment-read bytes per MMAC: 256 B for configs 0/1 (4 A + 4 -// B fragments feed 16 MMACs per kk) vs 384 B for config 2 (4 A + 2 B feed 8 -// MMACs) and vs 384 B for the old 128x64 arm; B reads are zero-conflict in -// the plane layout for every plane stride, and the A stride-80 2-way -// aliasing matches the accepted q_b_proj arm. Occupancy: 64x64 -> 7 -// blocks/CU = 7 waves/CU (LDS- and VGPR-limited), 64x128 -> 4 blocks/CU = 8 -// waves/CU (LDS-limited; VGPR <= 128 keeps 4 blocks), 128x64 -> 2 blocks/CU -// = 8 waves/CU (VGPR-limited) - vs 2 blocks/CU = 8 waves/CU today but with -// ~2.9x fewer LDS instructions (direct 8-byte reads vs library byte-gathers) -// and 4x fewer store instructions (coalesced epilogue, 229,376 -> 57,344 -// vmem_write). Active config for the measurement: 64x128 (config 1): -// deepest per-kk MMAC burst of the three alongside the lowest per-MMAC LDS -// read volume, half the A-side global re-read of 128x64, and 8 waves/CU at -// 4 blocks/CU. -// * Iteration-1 measurement: config 1 landed at arch_vgpr 136 (code object -// 135) -> 3 blocks/CU = 6 waves/CU (1.5/SIMD), median 209.4 us = 71.8 -// TOPS. The kernel is latency-bound (per-stage wall ~2,454 cycles vs -// ~250 cycles of pure issue), so 6 waves/CU leaves the barrier + LDS + -// global-latency chain exposed on the SIMDs that hold only one wave. -// * Iteration-2 (operand-reuse) round selects config 2 (128x64): same -// cooperative A+B staging machinery and swizzled-pack B, same -// bit-identical int32 order; its +50% A-side LDS fragment-read bytes per -// MMAC (384 vs 256 B; A rows are re-read by the two n-half waves in both -// configs, but config 2's narrower 64x32 quadrant halves the MMACs per -// A-frag load) is absorbed by the LDS pipe (~10% of per-stage cycles) -// while the compiled 78 VGPR / 14,336 B gives 3 blocks x 4 waves = 12 -// waves/CU (3/SIMD) - a 2x occupancy increase over the measured config-1 -// residency. -// The swizzled pack is extended to (512, 3584) in -// launch_pack_w8a8_weight (same byte count, outside the timed region) and -// the scalar fallback decodes it for both swizzled pairs, so the paired -// M=2/(3584,512) API case stays exact. All other arms (q_b_proj 128x128, -// generic 128x64, identity packs) are byte-identical. +// Iteration 8 (resource round): the exact-shape 128x128 n-major arm keeps its +// winning tile family, its tile-contiguous plane-B pack, its A swizzle, its +// direct 8-byte fragment loads, its k0-outer/kk-inner int32 order and its +// coalesced epilogue, and changes only the occupancy resources: waves per +// block, the VGPR live range and the per-thread LDS staging payload. +// * Measured occupancy of the current best (iteration-6 source): arch_vgpr +// 120, static LDS 16,384 B/block, __launch_bounds__(256, 2) -> 2 blocks/CU +// x 4 waves = 8 waves/CU = 2 waves/SIMD with 16 int32 accumulator VGPR +// fragments + 8 operand fragments + 4 live int4 staging registers per +// thread. The PMC says the kernel is issue-stalled, not LDS-bandwidth- or +// MMAC-bound: 195,964 cycles for 7,168 block-stages, 974,848 +// lds_instructions (4.97 per GPU cycle over 120 CUs) and 1,404,478 +// lds_wait_instructions against a theoretical LDS issue rate of 1/cycle/CU +// - i.e. the SIMDs are starved of co-resident waves, not of LDS slots. +// * Focused change: the same 128x128 macro-tile is covered by EIGHT +// wavefronts of 64 lanes (512 threads, kNMBlockThreads = 8*64) instead of +// four; each wave now owns a 32x64 quadrant = kAccM 2 A fragments x +// kAccN 4 B fragments = 8 m16n16k32 int32 accumulators, i.e. 32 +// accumulator VGPR per thread instead of 64. The staging payload per +// thread halves to exactly one int4 of A (row = tid >> 2, 16-byte K group +// (tid & 3) * 16) and one int4 of B (offset tid * 16), and the arm is +// pinned by __launch_bounds__(kNMBlockThreads, 2) to the 64-VGPR / +// 2-blocks/CU point (65536 VGPR per CU / (2 x 512 lanes) = 64), i.e. +// 2 blocks x 8 waves = 16 waves/CU = 4 waves/SIMD, twice the currently +// measured wave count. LDS stays 16,384 B/block, so 2 blocks use 32,768 B +// of the 65,536-B per-CU LDS (no exact-fit cliff like iteration 7's +// 32,768 B/block double buffer). +// * Why exactly this point: the trusted occupancy-probe candidate list for +// the measured 120-CU device is [2, 3, 4, 5, 6, 7, 8] waves per block; 8 +// is the largest value that still tiles the 128x128 macro-tile with +// rectangular m16n16k32 quadrants (4 M row groups x 2 N column groups) and +// it gives the top of the trusted residency range (8 waves/CU -> 16 +// waves/CU = 4/SIMD). The direct sibling evidence is the lineage reference +// references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip (same exact +// M=4096/N=3584/K=512 shape, same DTK 26.04/gfx928): its 8-wave/512-thread +// 16-waves/CU recipe (config 3, arch_vgpr <= 64, 2 blocks/CU) is the +// recorded 137.4 us winner of that family versus 169.3 us for its +// 12-waves/CU 128x64 form, which is the same occupancy axis this round +// probes while keeping our better 1:1 macro-tile and 0.75-fragment/MMAC +// quadrant instead of that reference's 1.0-fragment/MMAC 32x32 quadrants. +// * Falsifiers: the exact code object must show vgpr_spill_count == 0, +// scratch_bytes == 0 and arch_vgpr <= 64 (else residency falls to 1 +// block/CU = 8 waves/CU, cancelling the whole point); the LDS-read volume +// per MMAC rises 0.5 -> 0.75 fragments, so the median improves only if the +// doubled wave count hides more LDS/barrier latency than the +50% fragment +// reads cost. The int32 accumulation is untouched (exact integer +// additions of the same products), every fragment placement is the same +// 8-byte per-lane operand, and the exact (k, n) == (512, 3584) guard, the +// identity-layout 128x64 arm, the pack kernels and the generic scalar +// fallback are byte-identical. // -// Iteration 5 (kv_b_proj packing round): the exact (k, n) == (512, 3584) -// weight is re-packed once (outside timing/Graph, same byte count, same -// graph-stable buffer) from the k-stage-major swizzle to a TILE-CONTIGUOUS -// stage-major layout -// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] == raw[kk*n + col] -// kk = k0*64 + kc*8 + b, nt = col/128, row = col%128, -// so every (64-k stage, 128-n tile) B tile is ONE contiguous 8192-B global -// region. The tile-family kernel (w8a8_dumma_prefill_tile_kernel, all four -// configs) stages it as 8 consecutive 1024-B wavefront streams instead of 8 -// streams 28,672 B apart in the k-stage-major layout; the LDS plane tile -// [kc][n][8], the zero-conflict lane-linear 8-byte B fragment reads -// (load_fragment8_plane_stride), the 16-B staging writes and the -// k0-outer/kk-inner int32 accumulation are byte-identical, so mismatch 0 / -// max_abs_error 0 is expected without tolerance debate. q_b_proj (2048, -// 2048) keeps the old k-stage-major swizzle byte-identical (its 128x128 -// kernel, pack kernel and fallback decode are untouched); the scalar -// fallback decodes the new tile-contiguous layout for (512,3584) and the -// old swizzle for (2048,2048). +// Iteration 4 (tile-shape round): the exact-shape n-major arm moves from the +// 128x64 macro-tile (M:N = 2:1) to a 128x128 macro-tile (M:N = 1:1) for the +// exact (M, N, K) = (4096, 3584, 512) shape; the four-wavefront/256-thread +// block and the direct 8-byte swizzled fragment path are unchanged. Each +// wave now owns a 64x64 quadrant built from sixteen m16n16k32 int32 +// accumulators (kAccM = 4 A fragments x kAccN = 4 B fragments = 16 MMACs per +// kk) instead of a 64x32 quadrant with eight. Why the aspect/area change is +// the lever for this exact M/N/K: +// * Per-block-stage cooperative-load bytes are fixed by the tile shape: +// 128*64 (A) + 128*64 (B) = 16 KiB for 128x128 versus 8 KiB + 4 KiB = +// 12 KiB for 128x64, but the 128x128 tile holds 2x the outputs, so the +// global/L2 read traffic for the whole operator drops from +// M*N*K*(1/128 + 1/64) = 176.2 MB to M*N*K*(1/128 + 1/128) = 117.4 MB +// (-33%). Concretely: A is re-read once per N-tile column (56 -> 28) so +// A traffic halves 117.4 -> 58.7 MB, while B keeps its 32 M-tile re-reads +// at 58.7 MB. The 128x64 arm's 2:1 A:B read imbalance becomes 1:1, which +// is the aspect balance this round is for; the smaller A (2 MiB) and B +// (1.75 MiB) both stay L2-resident either way. +// * LDS fragment reads per MMAC drop 6/8 -> 8/16 (-33%): the 64x32 quadrant +// needed 4 A + 2 B = 6 direct 8-byte reads for 8 MMACs, the 64x64 +// quadrant needs 4 A + 4 B = 8 reads for 16 MMACs. +// * The measured per-stage wall is dominated by the exposed A+B L2 round +// trip in front of the MMAC burst (exact code object: 3 x +// global_load_dwordx4 -> s_waitcnt vmcnt(2)/(1)/(0) -> 6 x ds_write_b64, +// per 64-K stage). Halving the macro-tile count (grid 1792 -> 896) halves +// the number of exposed stage walls for the same total work. +// * 896 blocks is still ~3.7 waves of the 240 co-resident blocks +// (2 blocks/CU x 120 CU), i.e. enough blocks to occupy every CU; no +// split-K and the workspace stays unused. +// Residency: LDS 16,384 B/block (2 x 16,384 = 32,768 <= 64 KiB) and the +// sixteen accumulators cost 64 VGPR, so the arm is pinned to 2 blocks/CU = +// 8 waves/CU by an explicit __launch_bounds__(256, 2) register cap at +// 128 VGPR (floor(16384 VGPR/SIMD / (128 x 64 lanes)) = 2 waves/SIMD); +// falsified if the exact code object shows > 128 VGPR or any spill -> 1 +// block/CU = 4 waves/CU. +// Everything else is byte-identical to iteration 2/3: the dense 8-byte-group +// swizzle, load_frag8_swizzled placement, the one-time n-major packed B, the +// k0-outer/kk-inner ascending int32 order, the exact (k, n) == (512, 3584) +// guard, the identity-layout 128x64 arm, the pack kernels and the generic +// scalar fallback. // -// Iteration 6 (kv_b_proj epilogue round): the mandate (per-row x_scale, -// per-column weight_scale, bf16 conversion and the final coalesced store -// fused into the compute kernel; no workspace/combine pass) has been in -// place since iteration 1 - the launcher is a single kernel with -// (void)workspace and the epilogue is register-batched (xs scalar + ws -// float4 prefetched on the last K stage) with one 8-byte bf16 store per lane -// per fragment (57,344 vmem_write). The remaining epilogue inefficiency of -// the exact (512, 3584) tile family is the bf16 CONVERSION lowering: the -// exact code object shows hip's __float2bfloat16 emits, per element, an -// exec-masked inf/NaN fixup (v_and exp + v_cmp_ne + s_and_saveexec + s_xor + -// v_mov 0 + v_or 0x10000 + v_cmp_eq_u32_sdwa + v_cndmask + s_or_b64) around -// the 3-instruction RNE rounding (v_bfe + v_add3 + shift-in-pack) - ~11 -// instructions and three scalar exec manipulations that serialize the -// block-tail epilogue. This round replaces that path, ONLY in the -// w8a8_dumma_prefill_tile_kernel epilogue, with a plain RNE conversion -// (bf16_rne_u16: u += 0x7fff + ((u >> 16) & 1); u >> 16) behind a new -// kRneOnly template flag on store_prefill_fragment_coalesced_scaled -// (default false, so the q_b_proj 128x128 / 128x64 arms and the scalar -// fallback stay byte-identical). The RNE formula is bit-identical to hip's -// own non-masked branch for every finite float - the only values this GEMM -// can produce (int32 accumulators x finite scales) - so mismatch 0 / -// max_abs_error 0 is expected on the unchanged pre-seeded reference data. -// PMC prediction: valu_instructions 7.31M -> ~5.9-6.1M (-15..-20%), -// lds_instructions 1,089,536 unchanged, lds_bank_conflicts 2,293,760 -// unchanged, vmem_read 229,376 / vmem_write 57,344 unchanged, arch_vgpr <= -// 64 / 0 scratch / 13,312 B LDS / 2 blocks/CU = 16 waves/CU residency -// unchanged; the exact code object must show no s_and_saveexec/s_xor/ -// s_andn2_saveexec around the conversion (v_bfe + v_add3 straight into the -// 8-byte pack). +// Iteration 2 (operand-reuse round): keep the cooperative A+B LDS staging +// (the reuse argument below favours it over per-fragment global loads) but +// replace the two structural costs the bootstrap PMC exposes on the assigned +// (M, N, K) = (4096, 3584, 512) shape: +// * lds_instructions 2,465,792 and lds_wait_instructions 7,165,860: +// the installed du_load_matrix_sync int8 loaders lower to a per-byte +// ds_read_u8 gather plus mask/OR reassembly VALU for EVERY fragment. +// * lds_bank_conflicts 6,881,280: the 80-byte padded row stride has bank +// period 8 over a fragment's 16 rows, so lanes r and r+8 collide. +// Reuse accounting for one 128x64 macro-tile / 64-K stage (the reason direct +// global operand loads are rejected): +// A[128,64] = 8 KiB staged; each A fragment row-block feeds two B fragments +// (kAccN = 2) and two kk sub-steps, i.e. every staged A byte is consumed +// by 4 or more MMAC operand reads; a direct global fragment load would +// re-fetch it per fragment (4x redundant global traffic). +// B[64,64] = 4 KiB staged; each B fragment feeds all four A row-blocks +// (kAccM = 4), i.e. 4x reuse. So cooperative staging is retained and the +// real lever is the *read* path and the LDS conflict structure. +// Change (one mechanism, both tiles): +// * Both LDS tiles become dense 64-byte-row swizzled 8-byte-group layouts +// (no pad). group(row, c) = (((row & 15) + 4*(c & 3)) & 15) + 16*c + +// 128*(row >> 4); the writers emit two 8-byte LDS stores per 16-byte +// global chunk at group(row, 2j) / group(row, 2j+1) and the fragment +// loader reads the swizzle back. Every 16-lane fragment load then touches +// 16 distinct bank pairs (zero conflicts) instead of the period-8 +// collision. +// * Every A and B fragment read becomes ONE 8-byte LDS load per lane +// (load_frag8_swizzled), reproducing du_load_matrix_sync's per-lane +// (row = row0 + (lane & 15), c = (lane >> 4) + (kk >> 3), 8 consecutive +// k bytes in natural order) placement, so the int32 accumulation is +// bit-identical. +// * B is read col_major from a one-time n-major transpose +// packed[n*k + kk] = raw[kk*n + n] (produced outside the timed/Graph +// region by launch_pack_w8a8_weight for exactly (k, n) == (512, 3584)); +// this makes each lane's 8 k bytes contiguous. The identity [K, N] copy +// is kept for every other (k, n). +// Iteration-2 geometry (superseded for the n-major arm by the iteration-4 +// 128x128 tile above; the identity-layout arm still uses it): one 128x64 +// tile/block, four gfx928 wavefronts of 64 lanes (256 threads), each wave +// owning a 64x32 quadrant (kAccM = 4, kAccN = 2, eight m16n16k32 int32 +// accumulators), 2 barriers per 64-K stage, no split-K, grid +// dim3(N/64, M/128) = (56, 32) = 1792 blocks, LDS = 8 KiB + 4 KiB = 12,288 +// B/block. +// * The fused int32 * x_scale * weight_scale -> bf16 epilogue keeps the +// verified gfx928 int8 m16n16k32 accumulator ownership (row = lane & 15, +// col = (lane >> 4) + 4*i), multiply order (float(dot) * x_scale[row] * +// weight_scale[col]) and bf16 rounding of the harness reference. +// * The identity-layout 128x64 DUMMA arm is retained for every non-packed +// large shape, and the generic scalar int8/int32 fallback covers every +// unmatched (m, n, k), including all small-M API cases (M=2, M=16) and +// M < 128 with the same (K, N): one thread per output element, exact +// int32 dot over K, then float(acc) * x_scale[row] * weight_scale[col] -> +// bf16. For exactly (512, 3584) the fallback decodes the n-major pack. +// * launch_pack_w8a8_weight runs outside the timed/Graph region, allocates +// nothing, keeps the same byte count and graph-stable addresses: the +// n-major transpose for exactly (k, n) == (512, 3584), identity for every +// other (k, n), fp32 scale always copied identity. +// +// The timed operator performs no allocation, compilation, autotuning, weight +// packing, host/device synchronization, or default-stream launch; it uses only +// the caller-provided out and workspace and launches on PyTorch's current HIP +// stream via hipLaunchKernelGGL. #include #include @@ -251,71 +329,55 @@ using namespace du::dumma; constexpr int kTileM = 16; constexpr int kTileN = 16; constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 -constexpr int kWaveSize = 64; // gfx928 native wavefront +constexpr int kWaveSize = 64; // gfx928 native wavefront, not 32 constexpr int kBlockM = 128; constexpr int kBlockN = 64; constexpr int kStageK = 64; -constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) -constexpr int kBStride = kBlockN + kBPad; -constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride -constexpr int kBlockThreads = 4 * kWaveSize; +constexpr int kPad = 16; // 16 B LDS row padding (64 -> 80 B rows) +constexpr int kAStride = kStageK + kPad; +constexpr int kBStride = kBlockN + kPad; +constexpr int kBlockThreads = 4 * kWaveSize; // 256 = four gfx928 wavefronts -// Exact (k, n) pairs whose weight buffer is packed with the iteration-5 -// swizzled 64-k-stage layout (packed[((k0*8+kc)*n+col)*8+b] == raw[kk*n+col]): -// the q_b_proj pair (2048, 2048) and - since this kv_b_proj tile round - the -// kv_b_proj pair (512, 3584). Every other (k, n) keeps the identity layout. -constexpr int kPackSwizzleK1 = 2048; -constexpr int kPackSwizzleN1 = 2048; -constexpr int kPackSwizzleK2 = 512; -constexpr int kPackSwizzleN2 = 3584; +// Iteration-4 tile-shape round: the exact-shape n-major arm uses a 128x128 +// macro-tile (M:N = 1:1) instead of the 128x64 identity-arm geometry. +// Iteration-8 resource round: the same 128x128 macro-tile is covered by EIGHT +// wavefronts of 64 lanes (512 threads), each owning a 32x64 quadrant built +// from kAccM = 2 A fragments x kAccN = 4 B fragments = eight m16n16k32 int32 +// accumulators (32 VGPR of accumulator per thread instead of 64). The block +// is pinned to 2 blocks/CU by __launch_bounds__(512, 2), i.e. a 64-VGPR cap: +// 2 blocks x 8 waves = 16 waves/CU = 4 waves/SIMD instead of the +// measured 2 blocks x 4 waves = 8 waves/CU = 2 waves/SIMD. +constexpr int kNMBlockM = 128; +constexpr int kNMBlockN = 128; +constexpr int kNMQuadM = 32; // per-wave M extent (4 row groups) +constexpr int kNMQuadN = 64; // per-wave N extent (2 col groups) +constexpr int kNMBlockThreads = 8 * kWaveSize; // 512 = eight gfx928 wavefronts -inline bool is_swizzled_pack(int k, int n) { - return (k == kPackSwizzleK1 && n == kPackSwizzleN1) || - (k == kPackSwizzleK2 && n == kPackSwizzleN2); -} +// Exact (k, n) pair whose weight buffer uses the iteration-5 tile-contiguous +// 64-K-stage pack packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] = +// raw[kk*n + col] (identity [K, N] for every other pair). +constexpr int kPackK = 512; +constexpr int kPackN = 3584; +constexpr int kPackTileN = 128; // N rows per tile-contiguous pack region +constexpr int kPackStageK = 64; // k bytes per pack stage + +using bf16_t = hip_bfloat16; -// Iteration 5 (kv_b_proj packing round): the (512, 3584) pair switched from -// the k-stage-major swizzle to the tile-contiguous layout -// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b]; (2048, 2048) keeps -// the k-stage-major swizzle byte-identical. inline bool is_tile_major_pack(int k, int n) { - return k == kPackSwizzleK2 && n == kPackSwizzleN2; + return k == kPackK && n == kPackN; } -// Active 2-D macro-tile for the exact glm_tp8_kv_b_proj_m4096 shape among the -// three mandated candidates plus the iteration-4 aspect flip (0 = 64x64, -// 1 = 64x128, 2 = 128x64, 3 = 64x128 with 32x32 quadrants). All four -// instantiations are compiled and dispatched behind the exact (k, n) guard; -// flipping this constant re-routes the kv_b_proj arm without source surgery. -constexpr int kActiveKvTile = 3; // 3 = 64x128 with 32x32 quadrants - // (iteration-4 tile-shape round): the - // down_proj-validated 16-waves/CU recipe - // (512 threads, 8 wavefronts/block, 4 - // accs/wave, arch_vgpr <= 64 -> 2 blocks/CU - // = 16 waves/CU = 4/SIMD) applied as the - // M-tile 128 -> 64 / N-tile 64 -> 128 aspect - // flip; config 2 (128x64, 3 blocks x 4 waves - // = 12 waves/CU, accepted best 169.257 us) - // stays compiled and is the retention - // fallback if config 3 regresses (arch_vgpr - // > 64 -> 1 block/CU = 8 waves -> reject). -constexpr int kKvTileM = 64; -constexpr int kKvTileN = 128; -constexpr int kKvQuadN = 64; - -using bf16_t = hip_bfloat16; - // --------------------------------------------------------------------------- // Direct fragment epilogue for one m16n16k32 accumulator fragment. // Verified gfx928 int8 m16n16k32 accumulator ownership (matches // du_store_matrix_sync): lane & 15 selects the row, lane >> 4 selects // col % 4, and x[i] maps to columns col%4 + 4*i. The scale multiply order -// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are -// identical to the harness reference ((dot.float() * x_scale) * ws.T then -// .to(bfloat16)), so stored bf16 bits match exactly. +// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding match +// the harness reference ((dot.float() * x_scale) * weight_scale.T -> bf16), so +// the stored bf16 bits are identical. // --------------------------------------------------------------------------- template -__device__ __forceinline__ void store_prefill_fragment( +__device__ __forceinline__ void store_fragment_scaled( const AccFragment& frag, const float* __restrict__ x_scale, const float* __restrict__ weight_scale, @@ -340,45 +402,75 @@ __device__ __forceinline__ void store_prefill_fragment( } // --------------------------------------------------------------------------- -// Iteration 6 (epilogue round): register-batched COALESCED direct-fragment -// epilogue for the exact (2048, 2048) 128x128 arm. The per-element variant -// (store_prefill_fragment) above issues one 2-byte bf16 store per lane per -// element: 64 scattered stores per wave per fragment set, each touching 16 -// rows x 4 columns, i.e. every 32-B store sector only 25% utilized. This -// epilogue transposes the 4-element groups within each 4-lane column group -// (lanes r, r+16, r+32, r+48 - a 4x4 transpose, two 2x2 steps with -// shfl_xor 16 then 32; the lineage-validated pattern accepted on the sibling -// TP8 workers 0/2 on this DTK, which lowers __shfl_xor to ds_bpermute at the -// block tail where the LDS pipe is idle), so lane (r, c4) ends up holding -// the four CONTIGUOUS columns 4*c4 .. 4*c4+3, converts them to bf16, packs -// 4 bf16 (8 B) and writes ONE 8-byte store per lane (100% store sector -// efficiency). Only the int32 values are re-routed between lanes; the -// per-element scale multiply order (float(dot) * x_scale[row] * -// weight_scale[col]) and the __float2bfloat16 rounding are unchanged, so -// the stored bits are identical to the per-element store. xs/ws come from -// caller registers (batched per wave on the last K stage), so the epilogue -// issues no vmem loads. The row >= m guard is wavefront-uniform (all 64 -// lanes of a wave share the same 16-row window), so the shuffles never mix -// active and inactive lanes; base_col is a multiple of 64 and n*2 a -// multiple of 8, so the 8-byte store is aligned. +// Bank-safe dense 8-byte-group LDS map for the exact-shape swizzled arm. +// group(row, c) is the 8-byte group index of tile position (row, k 8-byte +// group c). For a fragment load the 16 lanes of a group share row0 and kk, so +// row & 15 ranges over 0..15 and c is fixed; (row & 15) -> ((row & 15) + +// 4*(c & 3)) & 15 is a permutation, giving 16 distinct bank pairs and zero +// bank conflicts over the whole fragment read. +// --------------------------------------------------------------------------- +__device__ __forceinline__ int swizzle_group(int row, int c) { + return (((row & 15) + 4 * (c & 3)) & 15) + 16 * c + 128 * (row >> 4); +} + +// One 8-byte LDS store of each half of a 16-byte global chunk. +__device__ __forceinline__ void store_chunk16_swizzled( + int8_t* __restrict__ p, int row, int j, const int4& v) { + *reinterpret_cast(p + swizzle_group(row, 2 * j) * 8) = + *reinterpret_cast(&v); + *reinterpret_cast(p + swizzle_group(row, 2 * j + 1) * 8) = + *reinterpret_cast(reinterpret_cast(&v) + 8); +} + +// Direct 8-byte fragment load that reproduces du_load_matrix_sync's per-lane +// placement exactly: row = row0 + (lane & 15), k group c = (lane >> 4) + +// (kk >> 3), eight consecutive k bytes in natural order. Writing them +// straight into the fragment storage lets the compiler emit one 8-byte LDS +// read per fragment with none of the library's per-byte gather/reassembly +// VALU, so the int32 accumulation is bit-identical. +__device__ __forceinline__ void load_frag8_swizzled( + DUFragmentBase& frag, + const signed char* __restrict__ p, + int lane, + int row0, + int kk) { + const int row = row0 + (lane & 15); + const int c = (lane >> 4) + (kk >> 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast(p + swizzle_group(row, c) * 8); +} + +// Iteration-5 (packing round) B fragment loader for the tile-contiguous plane +// LDS layout b_tile[kc][n][8] (kc = 8-byte k sub-chunk; plane stride +// kNMBlockN*8 = 1024 B). Lane l reads the same 8 consecutive k bytes as +// load_frag8_swizzled's placement (n row = row0 + (lane & 15), k group +// kc = (lane >> 4) + (kk >> 3)) but from the lane-linear address +// kc * (kNMBlockN * 8) + row * 8, +// i.e. ((lane >> 4) * 1024) + ((lane & 15) * 8) relative to row0's plane +// origin, so each 16-lane group of a 128-B LDS phase spans all 32 banks +// exactly once (zero bank conflicts) and the compiler emits one ds_read_b64. +__device__ __forceinline__ void load_frag8_plane( + DUFragmentBase& frag, + const signed char* __restrict__ p, + int lane, + int row0, + int kk) { + const int row = row0 + (lane & 15); + const int kc = (lane >> 4) + (kk >> 3); + *reinterpret_cast(&frag.x[0]) = + *reinterpret_cast( + p + kc * (kNMBlockN * 8) + row * 8); +} + // --------------------------------------------------------------------------- -// Iteration 6 (kv_b_proj epilogue round): plain round-to-nearest-even bf16 -// conversion without the inf/NaN fixup path. hip's __float2bfloat16 lowers -// (verified on this exact DTK 26.04 / gfx928 toolchain with a micro compile) -// to an exec-masked sequence per element: v_and (exp mask) + v_cmp_ne + -// s_and_saveexec + s_xor + v_bfe + v_add3 + v_mov 0 + v_or 0x10000 + -// v_cmp_eq_u32_sdwa + v_cndmask + s_or_b64 (~11 instructions including three -// scalar exec manipulations that serialize the wavefront). The masking only -// changes the result for inf/NaN inputs; for every finite float - the only -// values this GEMM can produce from int32 accumulators and finite scales - -// the RNE rounding is exactly u += 0x7fff + ((u >> 16) & 1); return u >> 16, -// bit-identical to hip's own non-masked branch (same v_bfe + v_add3 + -// shift). The kernel epilogue is the only place that converts all M*N = -// 14.68M elements; dropping the mask removes ~6 VALU + 5 SALU per element -// (~1.38M VALU + 1.15M SALU per launch, ~19% of the measured 7.31M VALU -// stream) and - the falsifiable axis - removes the -// s_and_saveexec / s_xor / s_andn2_saveexec / s_or_b64 pairs that serialize -// the block-tail epilogue behind scalar exec manipulation. +// Iteration 9 (consolidate round): plain round-to-nearest-even fp32 -> bf16 +// conversion, bit-identical to hip's __float2bfloat16 for every finite input +// (the only values this GEMM can produce: int32 accumulators times finite +// scales), but without the exec-masked inf/NaN fixup sequence the DTK 26.04 +// gfx928 lowering emits around it (see the header note). The canonical RNE +// bit trick: add 0x7fff plus the retained LSB so ties round to even, then +// truncate. Same form as the sibling reference +// references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip on this DTK. // --------------------------------------------------------------------------- __device__ __forceinline__ unsigned short bf16_rne_u16(float x) { union { @@ -390,21 +482,52 @@ __device__ __forceinline__ unsigned short bf16_rne_u16(float x) { return static_cast(u >> 16); } -template -__device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( +// --------------------------------------------------------------------------- +// Iteration-6 (epilogue round): register-batched COALESCED direct-fragment +// epilogue for the exact-shape arm, ported from the lineage-validated sibling +// reference references/int8w8a8-gemm/glm52/TP8/M4096/kv_b_proj.hip. The +// per-element variant (store_fragment_scaled) issues one 2-byte bf16 store per +// lane per element: four scattered stores per fragment, 64 per wave, each +// wavefront store touching 16 rows x 4 columns so every 32-B store sector is +// only 25% utilized (the bootstrap PMC's 229,376 vmem_write instructions = +// 896 blocks x 4 waves x 16 fragments x 4). This epilogue transposes the +// four-element groups within each 4-lane column group (lanes r, r+16, r+32, +// r+48 - a 4x4 transpose done as two 2x2 steps with shfl_xor 16 then 32, the +// lineage-validated pattern that lowers to ds_bpermute at the block tail +// where the LDS pipe is idle), so lane (r, c4) ends up holding the four +// CONTIGUOUS columns base_col + 4*c4 .. +3, converts them to bf16 and writes +// ONE 8-byte store (100% store sector efficiency; 16 stores per wave instead +// of 64). Only the int32 values are re-routed between lanes; the per-element +// scale multiply order (float(dot) * x_scale[row] * weight_scale[col]) and the +// __float2bfloat16 rounding are unchanged, so the stored bits are identical. +// xs is batched by the caller (one load per 16-row row-block) and ws is the +// float4 of the four per-column scales at col0, 16-B aligned because base_col +// is a multiple of 16, so the epilogue issues 2 scalar + 8 vector vmem loads +// per lane instead of 80 scalar loads. Only the exact-shape n-major arm calls +// this helper, and its dispatch guard is m >= 128 with m % 128 == 0, so every +// row of every block is in bounds and no tail-M mask is needed: the shuffles +// never mix active and inactive lanes, and n*2 is a multiple of 8, so the +// 8-byte store is aligned. +// +// Iteration 9 (consolidate round): the bf16 conversion below is a plain RNE +// (bf16_rne_u16) instead of __float2bfloat16. The exact code object of the +// current best showed the DTK 26.04 lowering wrapping every element in an +// exec-masked inf/NaN fixup (~9 vector instructions + two s_and_saveexec / +// s_andn2_saveexec scalar manipulations) whose only effect vs RNE is on +// inf/NaN inputs; int32 accumulators times finite scales are always finite, +// and for finite floats the two forms produce the identical 16-bit pattern. +// --------------------------------------------------------------------------- +template +__device__ __forceinline__ void store_fragment_coalesced_scaled( const AccFragment& frag, float xs, - const float4& ws, + const float* __restrict__ weight_scale, bf16_t* __restrict__ out, int base_row, int base_col, - int m, int n, int lane) { const int row = base_row + (lane & 15); - if (row >= m) { - return; // tail-M masking: padded rows never write (wavefront-uniform) - } const int c4 = lane >> 4; // 0..3 const int x0 = frag.x[0]; const int x1 = frag.x[1]; @@ -431,165 +554,62 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( const int f2 = hi ? t0 : a2; const int f3 = hi ? t1 : a3; - // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte - // aligned: base_col is a multiple of 16, n*2 is a multiple of 8); ws is - // the float4 at weight_scale + base_col + 4*c4 preloaded by the caller. + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte aligned: + // base_col is a multiple of 16 and n is a multiple of 4); ws is the float4 + // of per-column scales at col0. const int col0 = base_col + 4 * c4; + const float4 ws = *reinterpret_cast(weight_scale + col0); const float v0 = static_cast(f0) * xs * ws.x; const float v1 = static_cast(f1) * xs * ws.y; const float v2 = static_cast(f2) * xs * ws.z; const float v3 = static_cast(f3) * xs * ws.w; - uint64_t packed; - if constexpr (kRneOnly) { - // Iteration 6 (kv_b_proj epilogue round): plain RNE (see bf16_rne_u16) - - // bit-identical to __float2bfloat16 for every finite input, without the - // exec-masked inf/NaN fixup path. - packed = static_cast(bf16_rne_u16(v0)) | - (static_cast(bf16_rne_u16(v1)) << 16) | - (static_cast(bf16_rne_u16(v2)) << 32) | - (static_cast(bf16_rne_u16(v3)) << 48); - } else { - packed = - static_cast( - static_cast(__float2bfloat16(v0))) | - (static_cast( - static_cast(__float2bfloat16(v1))) - << 16) | - (static_cast( - static_cast(__float2bfloat16(v2))) - << 32) | - (static_cast( - static_cast(__float2bfloat16(v3))) - << 48); - } - *reinterpret_cast(out + static_cast(row) * n + col0) = - packed; -} - -// --------------------------------------------------------------------------- -// Direct 8-byte LDS fragment load (lineage-validated on gfx928/DTK 26.04 for -// m16n16k32 int8 fragments). du_load_matrix_sync's int8 loaders assign -// x[0..7] = 8 consecutive bytes at (lane & 15) * ldm + ((lane >> 4) << 3) for -// matrix_b col_major, but the compiler lowers that to per-byte ds_read_u8 + -// mask/OR reassembly VALU. Writing the same 8 bytes directly into the -// fragment storage keeps the operand bit pattern identical (exact int32 -// accumulation unchanged) and lets the compiler emit one ds_read2_b64 per -// fragment straight into the v_mmac operand. -// --------------------------------------------------------------------------- -__device__ __forceinline__ void load_fragment8( - DUFragmentBase& frag, - const signed char* __restrict__ base, - int stride, - int lane) { - const int off = (lane & 15) * stride + ((lane >> 4) << 3); - *reinterpret_cast(&frag.x[0]) = - *reinterpret_cast(base + off); -} - -// --------------------------------------------------------------------------- -// Iteration 5 (packing round): 8-byte fragment load for the swizzled plane -// LDS layout b_tile[kc][n][8] (kc = 8-byte k sub-chunk, n row stride 8 B, -// plane stride 1024 B). Lane l reads the same 8 consecutive k bytes as -// load_fragment8 (n row = lane&15, k chunk = lane>>4) but from -// (lane >> 4) * 1024 + (lane & 15) * 8 -// relative to the fragment origin, so the 64 lanes hit 16 distinct bank -// pairs per 128-B phase (all 32 banks once) -> zero bank conflicts, the -// 4-cycle minimum for a 512-B fragment. -// --------------------------------------------------------------------------- -__device__ __forceinline__ void load_fragment8_plane( - DUFragmentBase& frag, - const signed char* __restrict__ base, - int lane) { - const int off = ((lane >> 4) << 10) + ((lane & 15) << 3); - *reinterpret_cast(&frag.x[0]) = - *reinterpret_cast(base + off); -} - -// --------------------------------------------------------------------------- -// Plane-layout 8-byte fragment load with a parameterized plane stride -// (plane_stride = kBlockN * 8 bytes). Identical to load_fragment8_plane for -// kBlockN == 128 (plane stride 1024 B); for narrower tiles the plane is -// smaller but the bank pattern is unchanged: lane l reads the 8 bytes at -// (lane >> 4) * plane_stride + (lane & 15) * 8, so each 16-lane group of a -// 128-B phase covers all 32 banks exactly once (zero bank conflicts) for any -// plane stride. -// --------------------------------------------------------------------------- -__device__ __forceinline__ void load_fragment8_plane_stride( - DUFragmentBase& frag, - const signed char* __restrict__ base, - int plane_stride, - int lane) { - const int off = ((lane >> 4) * plane_stride) + ((lane & 15) << 3); - *reinterpret_cast(&frag.x[0]) = - *reinterpret_cast(base + off); + const uint64_t packed = + static_cast(bf16_rne_u16(v0)) | + (static_cast(bf16_rne_u16(v1)) << 16) | + (static_cast(bf16_rne_u16(v2)) << 32) | + (static_cast(bf16_rne_u16(v3)) << 48); + *reinterpret_cast( + out + static_cast(row) * n + col0) = packed; } -// --------------------------------------------------------------------------- -// Compile-time B-fragment layout selector. row_major / col_major are tag -// types in du::dumma, so the layout template argument must be selected as a -// type (a conditional expression over type names is not a valid template -// argument). kNMajorB == true (n-major packed B, exact (2048, 2048)) -> -// col_major fragments (8 contiguous k bytes per lane, load_fragment8); -// kNMajorB == false (raw [K, N] row-major B) -> row_major fragments (the -// accepted-best loader). -// --------------------------------------------------------------------------- -template -struct b_frag_layout { - using type = row_major; -}; -template <> -struct b_frag_layout { - using type = col_major; -}; - // --------------------------------------------------------------------------- // Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. // Each wave owns a 64x32 quadrant (eight m16n16k32 int32 accumulators); the // block cooperatively stages A[128,64] from x_q (row-major, stride k) and -// B[64,64] into a single-buffered LDS stage. Two barriers per stage: one -// after the cooperative load, one before the next stage overwrites LDS. -// A fragments use the library du_load_matrix_sync row-major loader. B is -// templated on its staged layout: -// * kNMajorB == true: B is packed n-major (packed[n*K + k]) and staged -// n-major (n row, 80 B stride, k contiguous); each B fragment is 8 -// contiguous k bytes per lane, loaded by load_fragment8 (one 8-byte LDS -// read per fragment), eliminating the ds_read_u8 byte-gather + VALU -// reassembly of the accepted-best kernel. -// * kNMajorB == false: B stays raw [K, N] row-major, staged k-major with -// row-major library fragments (accepted-best path, unchanged). -// Dispatch guarantees m % 128 == 0, n % 64 == 0, k % 64 == 0, so every -// global load/store is in-bounds and 16-byte aligned. +// B[64,64] from the raw [K, N] row-major weight into a single-buffered LDS +// stage. Two barriers per 64-K stage: one after the cooperative load, one +// before the next stage overwrites the buffer. Dispatch guarantees +// m % 128 == 0, n % 64 == 0, k % 64 == 0, so every global load/store is +// in-bounds and 16-byte aligned. // --------------------------------------------------------------------------- -template -__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - bf16_t* __restrict__ out, - int m, - int n, - int k) { +__global__ __launch_bounds__(kBlockThreads) void +w8a8_dumma_prefill_128x64_kernel(const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; const int lane = tid % kWaveSize; - const int wave_row = wave >> 1; - const int wave_col = wave & 1; + const int wave_row = wave >> 1; // 0 or 1 + const int wave_col = wave & 1; // 0 or 1 const int m0 = static_cast(blockIdx.y) * kBlockM; const int n0 = static_cast(blockIdx.x) * kBlockN; const int local_row = wave_row * 64; const int local_col = wave_col * 32; - // Single-buffered stage: A[128, 80] + B[64, 80] = 15,360 B/block - // (4 blocks/CU fit the 64 KiB LDS budget; the padded 80-byte strides are - // 16-byte-aligned and break the 64-byte LDS bank periodicity). + // Single-buffered stage: A[128, 80] + B[64, 80] = 15,360 B/block (four + // blocks fit the 64 KiB per-CU LDS budget; the padded 80-byte strides are + // 16-byte aligned and break the 64-byte LDS bank periodicity). __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; DUFragment a_frag0, a_frag1; - DUFragment::type> + DUFragment b_frag0, b_frag1; DUFragment acc00, acc01, acc10, acc11, acc20, acc21, acc30, acc31; @@ -602,21 +622,18 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne du_fill_fragment(acc30, 0); du_fill_fragment(acc31, 0); - // Cooperative staging: A[128,64] is 512 int4s (two per thread), B[64,64] - // is 256 int4s (one per thread). - // A: thread tid owns row (tid*16)/64 and 16-byte column group + // Cooperative staging: A[128,64] is 512 int4s (two per thread), B[64,64] is + // 256 int4s (one per thread). + // A: thread tid owns row (tid*16)/64 and the 16-byte K-column group // (tid*16)%64 for rows [0,64) and [64,128). - // B (kNMajorB == false): raw [K, N] row-major; thread tid owns K row - // tid>>2 and the 16-byte N-column group (tid&3)*16 (16 consecutive N - // values at a fixed K row -> one aligned int4 in global and in LDS). - // B (kNMajorB == true): packed n-major; thread tid owns N row tid>>2 and - // the 16-byte K-column group (tid&3)*16 (16 consecutive K values at a - // fixed N row -> one aligned int4 in global and in the n-major LDS). + // B: raw [K, N] row-major; thread tid owns K row tid>>2 and the 16-byte + // N-column group (tid&3)*16 (16 consecutive N values at a fixed K row, + // one aligned int4 in global and in LDS). const int vector_byte_offset = tid * static_cast(sizeof(int4)); const int stage_row = vector_byte_offset / kStageK; // 0..63 const int stage_col = vector_byte_offset - stage_row * kStageK; - const int b_k = tid >> 2; // K row (false) / N row (true) - const int b_nc = (tid & 3) * 16; // N group (false) / K group (true) + const int b_k = tid >> 2; // K row 0..63 + const int b_nc = (tid & 3) * 16; // N-column group 0/16/32/48 for (int k0 = 0; k0 < k; k0 += kStageK) { *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = @@ -625,43 +642,23 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne *reinterpret_cast(a_tile + (stage_row + kBlockM / 2) * kAStride + stage_col) = *reinterpret_cast( - x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + - k0 + stage_col); - if constexpr (kNMajorB) { - // n-major packed B: 16 consecutive k bytes of the (n0 + b_k) row. - *reinterpret_cast(b_tile + b_k * kBStride + b_nc) = - *reinterpret_cast( - weight + static_cast(n0 + b_k) * k + k0 + b_nc); - } else { - // raw [K, N] row-major B: 16 consecutive n bytes of the (k0 + b_k) row. - *reinterpret_cast(b_tile + b_k * kBStride + b_nc) = - *reinterpret_cast( - weight + static_cast(k0 + b_k) * n + n0 + b_nc); - } + x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + k0 + + stage_col); + *reinterpret_cast(b_tile + b_k * kBStride + b_nc) = + *reinterpret_cast( + weight + static_cast(k0 + b_k) * n + n0 + b_nc); __syncthreads(); // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. - // Accumulation order: k0-outer over 64-K stages, kk-inner (kk=0 then - // kk=32), matching the reference int32 accumulation. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk = 0 then + // 32), matching the reference int32 accumulation exactly. #pragma unroll for (int kk = 0; kk < kStageK; kk += kTileK) { - if constexpr (kNMajorB) { - // n-major LDS tile: n rows stride kBStride, k contiguous within a - // row; col_major fragment slots are 8 consecutive k bytes at - // (lane & 15) * kBStride + ((lane >> 4) << 3) relative to the - // fragment origin -> one ds_read2_b64 per fragment per lane. - load_fragment8(b_frag0, b_tile + local_col * kBStride + kk, kBStride, - lane); - load_fragment8(b_frag1, - b_tile + (local_col + kTileN) * kBStride + kk, - kBStride, lane); - } else { - du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, - kBStride); - du_load_matrix_sync(b_frag1, - b_tile + kk * kBStride + local_col + kTileN, - kBStride); - } + du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, + kBStride); + du_load_matrix_sync(b_frag1, + b_tile + kk * kBStride + local_col + kTileN, + kBStride); du_load_matrix_sync(a_frag0, a_tile + local_row * kAStride + kk, kAStride); du_load_matrix_sync(a_frag1, @@ -689,101 +686,84 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne const int base_row = m0 + local_row; const int base_col = n0 + local_col; - store_prefill_fragment(acc00, x_scale, weight_scale, out, base_row, - base_col, m, n, lane); - store_prefill_fragment(acc01, x_scale, weight_scale, out, base_row, - base_col + kTileN, m, n, lane); - store_prefill_fragment(acc10, x_scale, weight_scale, out, - base_row + kTileM, base_col, m, n, lane); - store_prefill_fragment(acc11, x_scale, weight_scale, out, - base_row + kTileM, base_col + kTileN, m, n, lane); - store_prefill_fragment(acc20, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col, m, n, lane); - store_prefill_fragment(acc21, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col + kTileN, m, n, - lane); - store_prefill_fragment(acc30, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col, m, n, lane); - store_prefill_fragment(acc31, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col + kTileN, m, n, - lane); + store_fragment_scaled(acc00, x_scale, weight_scale, out, base_row, base_col, + m, n, lane); + store_fragment_scaled(acc01, x_scale, weight_scale, out, base_row, + base_col + kTileN, m, n, lane); + store_fragment_scaled(acc10, x_scale, weight_scale, out, + base_row + kTileM, base_col, m, n, lane); + store_fragment_scaled(acc11, x_scale, weight_scale, out, + base_row + kTileM, base_col + kTileN, m, n, lane); + store_fragment_scaled(acc20, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col, m, n, lane); + store_fragment_scaled(acc21, x_scale, weight_scale, out, + base_row + 2 * kTileM, base_col + kTileN, m, n, lane); + store_fragment_scaled(acc30, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col, m, n, lane); + store_fragment_scaled(acc31, x_scale, weight_scale, out, + base_row + 3 * kTileM, base_col + kTileN, m, n, lane); } // --------------------------------------------------------------------------- -// Iteration 4 (tile-aspect round) + iteration 5 (packing round): 128x128 -// tile for the exact q_b_proj (k, n) == (2048, 2048) swizzled-pack B arm. -// Four wavefronts of 64 lanes; each wave owns a 64x64 quadrant (four A x -// four B m16n16k32 fragments = 16 MMACs per kk, 16 int32 accumulator -// fragments). The block cooperatively stages A[128,64] (2 int4/thread, -// 80-byte-strided rows) and B[128,64] from the iteration-5 swizzled pack (2 -// int4/thread, 16-byte plane chunks) into a single-buffered LDS stage: -// A[128,80] (10,240 B) + B plane tile [8][128][8] (8,192 B) = 18,432 -// B/block (2 blocks/CU = 36,864 B <= 64 KiB); two __syncthreads per stage. -// A fragments use the library du_load_matrix_sync row-major loader -// (4 x ds_read2_b32 per kk per wave), B fragments use load_fragment8_plane -// (lane-linear 8-byte reads, zero LDS bank conflicts). A-side global -// re-read halves (A reuse 32 -> 16, 268 -> 134 MiB) while B reuse stays 32 -// (134 MiB): per-byte A:B tile traffic balances 1:1, total 402 -> 268 MiB -// (-33%). Grid dim3(N/128, M/128) = 512 blocks (~4.3/CU). Dispatch -// guarantees m % 128 == 0, n == 2048, k % 64 == 0, so every global -// load/store is in-bounds and 16-byte aligned. Accumulation order stays -// k0-outer / kk-inner with the same element-to-slot fragment mapping, so the -// int32 accumulation is bit-identical to the accepted kernel. +// Exact-shape (k, n) == (512, 3584) large-M prefill: 128x128 macro-tile per +// block, EIGHT gfx928 wavefronts of 64 lanes (512 threads). Each wave owns a +// 32x64 quadrant = kAccM 2 A fragments x kAccN 4 B fragments = eight +// m16n16k32 int32 accumulators (8 MMACs per kk). Cooperative A+B LDS staging +// is retained (each A row-block feeds 4 B fragments x 2 kk, each B fragment +// feeds 2 A row-blocks x 2 kk) and every fragment read is one direct 8-byte +// LDS load. A comes from the identity [M, K] x_q (row-major, stride k) and +// keeps the dense swizzled 8-byte-group LDS layout; B comes from the +// iteration-5 one-time tile-contiguous 64-K-stage pack +// packed[(((k0*(n/128)+nt)*8+kc)*128+row)*8+b] = raw[kk*n+col] and uses the +// plane LDS layout [kc][128][8], so its staging is a plain contiguous 16-B +// vector copy and its fragment loads are lane-linear and bank-conflict-free. +// Two barriers per 64-K stage, single-buffered, int32 accumulation resident +// through the whole K loop, no split-K (grid dim3(28, 32) = 896 blocks over +// 120 CUs, ~3.7 waves of the 240 co-resident blocks). Dispatch guarantees +// k == 512, n == 3584, m >= 128 and m % 128 == 0, so every global load is +// in-bounds and 16-byte aligned. The two-argument __launch_bounds__ caps the +// register file at the 64-VGPR 2-blocks/CU residency point for 512-thread +// blocks (65536 VGPR per CU / (2 x 512 lanes) = 64): 2 blocks x 8 waves = +// 16 waves/CU = 4 waves/SIMD instead of the measured 120-VGPR 128x128 +// four-wave form at 2 waves/SIMD. The iteration-6 epilogue is kept +// register-only: 2 batched per-row scale loads, one aligned float4 per-column +// scale load per fragment and one 8-byte coalesced store per lane per +// fragment, so it adds no LDS, no barrier and no occupancy pressure. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - bf16_t* __restrict__ out, - int m, - int n, - int k) { - constexpr int kBlockM128 = 128; - constexpr int kBlockN128 = 128; - constexpr int kTileStride = kStageK + kBPad; // 64 -> 80-byte row stride - constexpr int kTileBytes = kBlockM128 * kTileStride; // 10,240 B per tile - // Iteration 5 (packing round): B tile is the swizzled plane layout - // [kc][n][8] (8 planes of 128 n rows x 8 B, plane stride 1024 B) so each - // lane's 8-byte fragment chunk is lane-linear and bank-conflict-free. - constexpr int kBTileBytes = 8 * kBlockN128 * 8; // 8,192 B/block - constexpr int kBTilePlane = 1024; // 128 n rows x 8 B - +__global__ __launch_bounds__(kNMBlockThreads, 2) void +w8a8_dumma_prefill_128x128_nmajor_kernel(const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + // Iteration 10 (consolidate round): the dispatch guard for this arm is + // is_tile_major_pack(k, n) && m >= kNMBlockM && m % kNMBlockM == 0, so the + // 128-row tile is always fully in bounds and no tail-M masking is needed. + (void)m; const int tid = static_cast(threadIdx.x); - const int wave = tid / kWaveSize; + const int wave = tid / kWaveSize; // 0..7 const int lane = tid % kWaveSize; - const int wave_row = wave >> 1; - const int wave_col = wave & 1; - const int m0 = static_cast(blockIdx.y) * kBlockM128; - const int n0 = static_cast(blockIdx.x) * kBlockN128; - const int local_row = wave_row * (kBlockM128 / 2); // 0 or 64 - const int local_col = wave_col * (kBlockN128 / 2); // 0 or 64 - const int base_row = m0 + local_row; - const int base_col = n0 + local_col; + const int wave_row = wave >> 1; // 0..3 + const int wave_col = wave & 1; // 0..1 + const int m0 = static_cast(blockIdx.y) * kNMBlockM; + const int n0 = static_cast(blockIdx.x) * kNMBlockN; + const int local_row = wave_row * kNMQuadM; // 0, 32, 64, 96 + const int local_col = wave_col * kNMQuadN; // 0 or 64 - // Iteration 6 (epilogue round): per-lane register-batched scales. Each - // lane owns 4 rows (base_row + 16*i + (lane & 15), i = 0..3) and 4 - // weight_scale float4s (base_col + 16*j + 4*(lane >> 4), j = 0..3) across - // its 16 fragments; they are loaded exactly once on the last K stage - // before the final protective barrier, so the coalesced epilogue below is - // pure compute + 16 eight-byte stores with no interleaved vmem loads. - float xs_m[4]; - float4 ws_m[4]; - - // Single-buffered stage: A[128,80] (10,240 B) + B plane tile [8][128][8] - // (8,192 B) = 18,432 B/block (2 blocks/CU = 36,864 B <= 64 KiB; the A rows - // keep the padded 80-byte stride, the B tile is plane-swizzled so fragment - // reads are lane-linear with zero bank conflicts). - __shared__ __align__(16) int8_t a_tile[kTileBytes]; - __shared__ __align__(16) int8_t b_tile[kBTileBytes]; + // A 8 KiB (dense swizzled 8-byte groups) + B 8 KiB (iteration-5 + // tile-contiguous plane [kc][128][8]) = 16,384 B/block. + __shared__ __align__(16) int8_t a_tile[kNMBlockM * kStageK]; + __shared__ __align__(16) int8_t b_tile[kNMBlockN * kStageK]; DUFragment - a_frag0, a_frag1, a_frag2, a_frag3; + a_frag0, a_frag1; DUFragment b_frag0, b_frag1, b_frag2, b_frag3; DUFragment - acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13, - acc20, acc21, acc22, acc23, acc30, acc31, acc32, acc33; + acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13; du_fill_fragment(acc00, 0); du_fill_fragment(acc01, 0); du_fill_fragment(acc02, 0); @@ -792,102 +772,59 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern du_fill_fragment(acc11, 0); du_fill_fragment(acc12, 0); du_fill_fragment(acc13, 0); - du_fill_fragment(acc20, 0); - du_fill_fragment(acc21, 0); - du_fill_fragment(acc22, 0); - du_fill_fragment(acc23, 0); - du_fill_fragment(acc30, 0); - du_fill_fragment(acc31, 0); - du_fill_fragment(acc32, 0); - du_fill_fragment(acc33, 0); - // Cooperative staging: A[128,64] is 512 int4s (two per thread) and - // B[128,64] from the iteration-5 swizzled pack is 512 int4s (two per - // thread). Thread tid owns linear int4 slots tid and tid + 256; for A a - // slot maps to row linear >> 2 and the 16-byte column group - // (linear & 3) * 16 (an M row of x_q, one aligned int4 in global and in - // LDS); for B a slot maps to plane linear >> 6 and n-pair linear & 63 - // (two n rows of one 8-byte k sub-chunk, one aligned int4 in the pack and - // in the plane LDS tile). + // Cooperative staging (iteration-8 resource round): with 512 threads each + // thread owns exactly ONE 16-byte chunk of A and ONE 16-byte chunk of B, so + // the staging payload per thread halves (1 int4 + 1 int4 instead of + // 2 + 2) - this is what keeps the live-range budget compatible with the + // 64-VGPR / 2-blocks/CU occupancy point. + // A: thread tid owns row tid >> 2 and the 16-byte K column group + // (tid & 3) * 16; the swizzled 8-byte-group LDS store is unchanged. + // B (iteration-5 packing round): the tile-contiguous pack lays each + // (64-K stage, 128-N tile) B tile out as one contiguous 8192-B region + // already in the LDS plane order [kc][128 N rows][8 bytes], so staging + // is a plain 1 x 16-B contiguous vector copy per thread. + const int vector_byte_offset = tid * static_cast(sizeof(int4)); + const int stage_row = vector_byte_offset / kStageK; // 0..127 + const int stage_col = vector_byte_offset - stage_row * kStageK; // 0,16,32,48 + const int b_tile_off = vector_byte_offset; // 0..8176 + + // Iteration 13 (consolidate round): one-stage register prefetch of the + // cooperative staging payload. The A row base and the B per-stage stride are + // loop-invariant, so the next 64-K A/B chunks are loaded while the current + // stage's fragment reads and MMACs are still draining; at the consumer side + // the load is already in flight across the trailing __syncthreads() instead + // of being issued and drained within a handful of instructions at the stage + // top. The loaded values, their order, the LDS bytes and the int32 + // accumulation are unchanged, so every stored bit is unchanged. + const int8_t* const a_stage_base = + x_q + static_cast(m0 + stage_row) * k + stage_col; + const int64_t b_stage_step = + static_cast(n / kNMBlockN) * (kNMBlockN * kStageK); + const int64_t b_stage_base = + static_cast(n0 / kNMBlockN) * (kNMBlockN * kStageK); + int4 a_chunk = *reinterpret_cast(a_stage_base); + int4 b_chunk = + *reinterpret_cast(weight + b_stage_base + b_tile_off); + for (int k0 = 0; k0 < k; k0 += kStageK) { - const int kstage = k0 >> 6; // 64-k group index inside the swizzled pack -#pragma unroll - for (int i = 0; i < 2; ++i) { - const int linear = tid + i * kBlockThreads; - const int row = linear >> 2; // 0..127 - const int col = (linear & 3) * 16; // 0/16/32/48 - *reinterpret_cast(a_tile + row * kTileStride + col) = - *reinterpret_cast( - x_q + static_cast(m0 + row) * k + k0 + col); - // Swizzled pack: slot linear -> (kc = linear >> 6, j = linear & 63); - // the 16-byte chunk is two n rows (n0 + 2j, n0 + 2j + 1) of the 8-byte - // k sub-chunk kc of stage kstage: one aligned int4 in global (two - // consecutive n rows of one plane) and one aligned int4 in the plane - // LDS tile (plane kc, rows 2j..2j+1 -> contiguous 16 B). - const int b_kc = linear >> 6; // 0..7 - const int b_j = linear & 63; // 0..63 (n-pair within the plane) - *reinterpret_cast(b_tile + b_kc * kBTilePlane + b_j * 16) = - *reinterpret_cast( - weight + (static_cast(kstage * 8 + b_kc) * n + - (n0 + 2 * b_j)) * 8); - } + store_chunk16_swizzled(a_tile, stage_row, stage_col >> 4, a_chunk); + *reinterpret_cast(b_tile + b_tile_off) = b_chunk; __syncthreads(); - // Each wave consumes its 64x64 quadrant: sixteen m16n16k32 MMACs per kk. - // Accumulation order: k0-outer over 64-K stages, kk-inner (kk=0 then - // kk=32), matching the reference int32 accumulation. All eight fragment - // loads are issued before the sixteen MMACs (load-all-then-MMAC-all). - // - // Iteration 7 (compute-pipeline round): the A fragments are loaded with - // the same direct 8-byte LDS read as B (load_fragment8) instead of the - // library du_load_matrix_sync row_major loader. The library's int8 - // matrix_a row_major loader assigns x[i] = p[(lane&15)*ldm + - // (lane>>4)*8 + i] (8 consecutive bytes, memory order) and du_mma_sync - // feeds reinterpret(x) straight into v_mmac, but on this DTK the - // loader lowers to 8 x ds_read2_b32 + a redundant per-dword byte - // reassembly chain (~7 VALU: v_and 0xff00/0xff0000/0xff000000 + - // v_or_b32_sdwa + v_or3 per second dword) sitting between the LDS read - // and the MMAC issue (exact code object, both the 128x128 and the 128x64 - // symbols). load_fragment8 fills the same x[0..7] with one 64-bit - // little-endian write, so the operand bytes are bit-identical and the - // compiler emits one ds_read2_b64 straight into the v_mmac operand (the - // same lineage as the B side since iteration 3/5): the LDS->MMAC - // critical path shortens by ~50 VALU + their lgkmcnt wait states per - // stage, and the 32-MMAC burst can issue back-to-back after the barrier. + // Each wave consumes its 32x64 quadrant: eight m16n16k32 MMACs per kk. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk = 0 then + // 32), matching the reference int32 accumulation exactly (the DUMMA + // accumulator is exact int32 addition, so the 2 x 4 fragment partition + // over the same products is bit-identical). #pragma unroll for (int kk = 0; kk < kStageK; kk += kTileK) { - load_fragment8(a_frag0, a_tile + local_row * kTileStride + kk, - kTileStride, lane); - load_fragment8(a_frag1, - a_tile + (local_row + kTileM) * kTileStride + kk, - kTileStride, lane); - load_fragment8(a_frag2, - a_tile + (local_row + 2 * kTileM) * kTileStride + kk, - kTileStride, lane); - load_fragment8(a_frag3, - a_tile + (local_row + 3 * kTileM) * kTileStride + kk, - kTileStride, lane); - // B fragments in the plane layout: fragment (f, kk) origin is - // b_tile + (kk >> 3) * 1024 + (local_col + 16f) * 8 and lane l reads - // the 8 bytes at (lane >> 4) * 1024 + (lane & 15) * 8 (lane-linear, - // zero bank conflicts); operand bytes are identical to the accepted - // n-major load_fragment8, so the int32 accumulation is bit-identical. - const int b_kc0 = kk >> 3; // 0 (kk == 0) or 4 (kk == 32) - load_fragment8_plane(b_frag0, - b_tile + b_kc0 * kBTilePlane + local_col * 8, - lane); - load_fragment8_plane(b_frag1, - b_tile + b_kc0 * kBTilePlane + - (local_col + kTileN) * 8, - lane); - load_fragment8_plane(b_frag2, - b_tile + b_kc0 * kBTilePlane + - (local_col + 2 * kTileN) * 8, - lane); - load_fragment8_plane(b_frag3, - b_tile + b_kc0 * kBTilePlane + - (local_col + 3 * kTileN) * 8, - lane); + load_frag8_plane(b_frag0, b_tile, lane, local_col, kk); + load_frag8_plane(b_frag1, b_tile, lane, local_col + kTileN, kk); + load_frag8_plane(b_frag2, b_tile, lane, local_col + 2 * kTileN, kk); + load_frag8_plane(b_frag3, b_tile, lane, local_col + 3 * kTileN, kk); + load_frag8_swizzled(a_frag0, a_tile, lane, local_row, kk); + load_frag8_swizzled(a_frag1, a_tile, lane, local_row + kTileM, kk); du_mma_sync(acc00, a_frag0, b_frag0, acc00); du_mma_sync(acc01, a_frag0, b_frag1, acc01); du_mma_sync(acc02, a_frag0, b_frag2, acc02); @@ -896,626 +833,83 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern du_mma_sync(acc11, a_frag1, b_frag1, acc11); du_mma_sync(acc12, a_frag1, b_frag2, acc12); du_mma_sync(acc13, a_frag1, b_frag3, acc13); - du_mma_sync(acc20, a_frag2, b_frag0, acc20); - du_mma_sync(acc21, a_frag2, b_frag1, acc21); - du_mma_sync(acc22, a_frag2, b_frag2, acc22); - du_mma_sync(acc23, a_frag2, b_frag3, acc23); - du_mma_sync(acc30, a_frag3, b_frag0, acc30); - du_mma_sync(acc31, a_frag3, b_frag1, acc31); - du_mma_sync(acc32, a_frag3, b_frag2, acc32); - du_mma_sync(acc33, a_frag3, b_frag3, acc33); - } - - // Last stage only: prefetch the epilogue's per-row x_scale and - // per-column weight_scale values into registers (uniform branch: k0 is - // block-uniform, so the barrier below is reached by every thread; the - // row < m guard keeps the x_scale reads in bounds for any m tail). The - // vmem latency overlaps the barrier below (s_barrier waits on - // lgkmcnt/arrival, not vmcnt) and the dead a_frag/b_frag VGPR slots are - // reused, so the epilogue is pure compute + coalesced stores. - if (k0 + kStageK >= k) { - const int r = lane & 15; - const int c4 = lane >> 4; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int row = base_row + i * kTileM + r; - xs_m[i] = (row < m) ? x_scale[row] : 0.0f; - } -#pragma unroll - for (int j = 0; j < 4; ++j) { - ws_m[j] = *reinterpret_cast( - weight_scale + base_col + j * kTileN + 4 * c4); - } - } - // Protect the LDS buffers from the next stage's cooperative overwrite. - __syncthreads(); - } - - // Iteration 6 (epilogue round): coalesced direct-fragment epilogue. Each - // lane owns four contiguous bf16 columns of its fragment row and writes - // ONE 8-byte store per fragment (16 stores per wave vs 64 scattered 2-byte - // stores before); the scales come from the registers batched on the last - // K stage, so no vmem loads are interleaved with the stores. The int32 - // accumulation is untouched, and the multiply order / bf16 rounding are - // identical to the old per-element store, so output bits are unchanged. - store_prefill_fragment_coalesced_scaled(acc00, xs_m[0], ws_m[0], out, - base_row, base_col, m, n, lane); - store_prefill_fragment_coalesced_scaled(acc01, xs_m[0], ws_m[1], out, - base_row, base_col + kTileN, m, n, - lane); - store_prefill_fragment_coalesced_scaled(acc02, xs_m[0], ws_m[2], out, - base_row, base_col + 2 * kTileN, m, - n, lane); - store_prefill_fragment_coalesced_scaled(acc03, xs_m[0], ws_m[3], out, - base_row, base_col + 3 * kTileN, m, - n, lane); - store_prefill_fragment_coalesced_scaled(acc10, xs_m[1], ws_m[0], out, - base_row + kTileM, base_col, m, n, - lane); - store_prefill_fragment_coalesced_scaled(acc11, xs_m[1], ws_m[1], out, - base_row + kTileM, - base_col + kTileN, m, n, lane); - store_prefill_fragment_coalesced_scaled(acc12, xs_m[1], ws_m[2], out, - base_row + kTileM, - base_col + 2 * kTileN, m, n, lane); - store_prefill_fragment_coalesced_scaled(acc13, xs_m[1], ws_m[3], out, - base_row + kTileM, - base_col + 3 * kTileN, m, n, lane); - store_prefill_fragment_coalesced_scaled(acc20, xs_m[2], ws_m[0], out, - base_row + 2 * kTileM, base_col, m, - n, lane); - store_prefill_fragment_coalesced_scaled(acc21, xs_m[2], ws_m[1], out, - base_row + 2 * kTileM, - base_col + kTileN, m, n, lane); - store_prefill_fragment_coalesced_scaled(acc22, xs_m[2], ws_m[2], out, - base_row + 2 * kTileM, - base_col + 2 * kTileN, m, n, lane); - store_prefill_fragment_coalesced_scaled(acc23, xs_m[2], ws_m[3], out, - base_row + 2 * kTileM, - base_col + 3 * kTileN, m, n, lane); - store_prefill_fragment_coalesced_scaled(acc30, xs_m[3], ws_m[0], out, - base_row + 3 * kTileM, base_col, m, - n, lane); - store_prefill_fragment_coalesced_scaled(acc31, xs_m[3], ws_m[1], out, - base_row + 3 * kTileM, - base_col + kTileN, m, n, lane); - store_prefill_fragment_coalesced_scaled(acc32, xs_m[3], ws_m[2], out, - base_row + 3 * kTileM, - base_col + 2 * kTileN, m, n, lane); - store_prefill_fragment_coalesced_scaled(acc33, xs_m[3], ws_m[3], out, - base_row + 3 * kTileM, - base_col + 3 * kTileN, m, n, lane); -} - -// --------------------------------------------------------------------------- -// Iteration 1 (kv_b_proj tile-baseline round): 2-D macro-tile DUMMA family -// for the exact glm_tp8_kv_b_proj_m4096 (k, n) == (512, 3584). The prior -// kv_b_proj arm (w8a8_dumma_prefill_128x64_kernel) is the iteration-2 -// bootstrap pipeline: library du_load_matrix_sync fragment loaders (per-byte -// ds_read_u8 byte-gather + mask/OR reassembly VALU, ~2.47M LDS instructions -// and 6.88M bank conflicts measured), raw [K, N] B identity layout, and a -// per-element 2-byte scattered epilogue (229,376 vmem_write instructions = -// 25% store sector efficiency) - i.e. 65 TOPS vs 156 TOPS on the sibling -// q_b_proj arm that already carries the full modern pipeline (plane-swizzled -// B, direct 8-byte fragment loads, coalesced epilogue). This round ports -// that proven pipeline to kv_b_proj as one templated family over the three -// mandated macro-tiles: -// config 0: 64x64 (1 wavefront, 64 threads; wave owns the whole tile) -// config 1: 64x128 (2 wavefronts, 128 threads; each wave owns 64x64) -// config 2: 128x64 (4 wavefronts, 256 threads; each wave owns 64x32) -// config 3: 64x128 (8 wavefronts, 512 threads; each wave owns 32x32; -// iteration-4 tile-shape round) -// (kActiveKvTile selects the active config; all four are compiled.) -// -// Pipeline (byte-identical to the validated q_b_proj arm): -// * B is staged from the iteration-5 swizzled 64-k-stage pack, now also -// produced for (512, 3584) by launch_pack_w8a8_weight, into the plane -// LDS tile b_tile[kc][n][8] (plane stride kBlockN*8 B, n stride 8 B); -// every B fragment read is one 8-byte lane-linear ds_read (zero bank -// conflicts) via load_fragment8_plane_stride. -// * A is staged [kBlockM, 80] (16 B padded row stride) and each A fragment -// is one direct 8-byte LDS read (load_fragment8), no VALU reassembly. -// * int32 accumulation stays resident across the full K loop (k0-outer over -// 64-K stages, kk-inner over two m16n16k32 DUMMA steps), fully unrolled -// per stage like the q_b_proj winner (deep 16-MMAC bursts per kk). -// * The fused dot * x_scale * weight_scale -> bf16 epilogue is the -// register-batched coalesced variant (one 8-byte store per lane per -// fragment; scales prefetched on the last K stage). -// * Single-buffered LDS, two __syncthreads per stage, grid dim3(N/BN, -// M/BM) = 1792 blocks (64x128 / 128x64) or 3584 blocks (64x64) - no -// split-K (the MxN grid already over-subscribes the 120 CUs). -// -// Dispatch guarantees m % kBlockM == 0, n % kBlockN == 0, k % 64 == 0 for the -// exact (512, 3584) pair (m == 4096 in the assigned shape), so every global -// load/store is in bounds and 16-byte aligned; the row < m guards stay for -// generic safety (wavefront-uniform, so the coalesced shuffles never mix -// active and inactive lanes). -// -// Iteration 16 (kv_b_proj load-placement round): the exact (k, n) == (512, -// 3584) active config-3 kernel (w8a8_dumma_prefill_tile_kernel<64,128,32,32, -// true>, median 139.567 us = 107.7 TOPS, 49.72x vs fixed Triton 6939.939 us; -// code object 61 VGPR / 28 SGPR / 13,312 B LDS / 0 scratch / 2 blocks/CU = -// 16 waves/CU; grid 1792, workgroup 512) is LDS-latency/issue-bound (lineage -// rule 5) and its stage critical path still exposes, per 64-K stage, one -// grouped A+B staging L2 round trip at the stage top (iteration-7 grouping -// gained +2.24% by merging the two serialized round trips of iteration 6, -// so the round trip sits on the barrier-to-barrier wall). The double- -// buffered one-barrier pipeline is arithmetic-infeasible at 2 blocks/CU -// (iters 9/10/12 compiled to 72/72/68 VGPR -> 1 block/CU = 8 waves/CU, the -// measured regression regime; every full double-buffer costs >= 7 VGPR on -// this 64-VGPR-constrained 512-thread config) and the 128-K stage-depth -// rework was attempted in the killed iteration-15 candidate (NOT replayed). -// This round keeps the SINGLE-buffered two-barrier stage byte-identical and -// moves ONLY the global-load ISSUE point: the grouped A+B staging loads for -// stage s+1 are issued at the END of stage s (after the MMAC burst, before -// the protective trailing __syncthreads) instead of at the top of stage s+1. -// Global loads never touch LDS, so the early issue is hazard-free; the -// vmcnt wait is inserted by the compiler before the first LDS write of stage -// s+1 that consumes the chunk, and the in-flight round trip overlaps the -// trailing barrier + loop back-edge of stage s (partial hiding of each of -// the 8 exposed staging round trips). Crucially NO register payload is -// carried across the burst: a_chunk/b_chunk are dead during the burst and -// are re-issued after it, so peak VGPR stays <= 64 -> 2 blocks/CU = 16 -// waves/CU (this is NOT the iteration-9/10/12 register-carrying placement, -// and NOT the iteration-12 double-buffer: LDS stays 13,312 B/block, 16 -// s_barriers/block). Bit-identity: identical global bytes -> identical LDS -// writes at identical (row, kk) slots -> identical operand bytes in the -// same k0-outer/kk-inner int32 accumulation -> identical register-batched -// RNE epilogue (only the load issue point moves), so mismatch 0 / -// max_abs_error 0 is expected without tolerance debate. Scoping: a new -// defaulted template parameter kPrefetchNext = false keeps configs 0-2 on -// the exact iteration-7 issue-at-stage-top placement (compiled identically -// via the if constexpr retention branch); the q_b_proj 128x128 / 128x64 -// arms, the generic scalar fallback and the pack kernels are untouched. -// Iteration 17 (kv_b_proj consolidation round): the two independently -// measured single-buffered wins are consolidated into ONE active config-3 -// instantiation (w8a8_dumma_prefill_tile_kernel<64,128,32,32,true,24>): -// (a) iteration 16's next-stage load-issue prefetch (kPrefetchNext = true, -// the current shadow candidate, median 138.229 us / p90 138.781, -// +0.968% vs the accepted best): the grouped A+B staging loads for -// stage s+1 are issued after the MMAC burst of stage s, so each of the -// 8 staging L2 round trips overlaps the trailing barrier + loop -// back-edge instead of stalling the stage-top LDS writes; -// (b) iteration 13's A-side LDS bank skew (kALdsPad = 24 -> 88-byte A row -// stride, standalone median 138.783 us / p90 139.185, +0.565% vs the -// accepted best): stride-80 rows are 8-periodic in the 128-B bank -// phase, so the 16 fragment rows (lane & 15) collapse onto 8 phases -// (r and r+8 collide, 2-way); 88 = 64+24 is 8-byte aligned and 88r mod -// 128 is 16-periodic, so the 16 fragment rows hit 16 DISTINCT phases -// -> zero-conflict A fragment reads (lineage rule 3, the qkv stride-68 -// / down_proj stride-72 pattern). Because stride % 16 == 8, the -// 16-byte A staging chunk is written as two int64 halves -// (ds_write2_b64, lineage rule 3), with the SAME 16 bytes at the SAME -// logical (row, kk) slots. -// The two mechanisms are orthogonal: (a) removes the staging L2 round trip -// from the stage-top critical path, (b) removes the 2,293,760-cycle LDS -// bank-conflict term from the LDS pipe (the conflict model: A fragment reads -// 32 x 2-way 8-B reads/stage -> 8 vs 4 cycles, + A staging writes 4 x 2-way -// 16-B writes/stage -> 16 vs 8 cycles, x 8 stages = 1,280 extra cycles/block -// x 1792 blocks). Both keep the single-buffered two-barrier stage (16 -// s_barriers/block), 2 blocks/CU = 16 waves/CU (vgpr <= 64: the write split -// reuses the same 4-VGPR a_chunk and NO payload crosses the burst), and LDS -// 13,824 B/block (2 blocks = 27,648 B <= 64 KiB). Bit-identity: identical -// global bytes -> identical logical A tile content at every (row, kk) (only -// inter-row padding changes) -> identical 8-byte m16n16k32 operand bytes per -// lane -> identical k0-outer/kk-inner int32 accumulation and register-batched -// RNE epilogue -> mismatch 0 / max_abs_error 0 expected without tolerance -// debate. Scoping: kALdsPad defaults to kBPad = 16, so configs 0-2 keep the -// exact iteration-7 stride-80 int4 write path (token-identical); the -// q_b_proj 128x128 / 128x64 arms, the generic scalar fallback and the pack -// kernels are untouched; the exact-shape guard (k == 512 && n == 3584 && -// m >= 64 && m % 64 == 0) is preserved. -// --------------------------------------------------------------------------- -template -__global__ __launch_bounds__( - ((kBlockM / kQuadM) * (kBlockN / kQuadN)) * kWaveSize) -void w8a8_dumma_prefill_tile_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - bf16_t* __restrict__ out, - int m, - int n, - int k) { - // kQuadM = rows owned by each wavefront (64 in configs 0-2; 32 in the - // iteration-4 config 3 so a 64x128 block hosts 8 wavefronts). - constexpr int kThreads = - (kBlockM / kQuadM) * (kBlockN / kQuadN) * kWaveSize; - // Iteration 17 (kv_b_proj consolidation round): A row stride is - // kStageK + kALdsPad. The default pad 16 keeps stride 80 (16-byte - // aligned, 2-way bank aliasing) for the retention configs 0-2; the active - // config-3 instantiation passes kALdsPad = 24 -> stride 88 (8-byte - // aligned, 16 distinct 128-B phases for the 16 fragment rows -> zero - // conflict A fragment reads; lineage rule: stride % 16 == 8 stages with - // two int64 halves, see the A staging write below). - constexpr int kTileStride = kStageK + kALdsPad; // 64 -> 80/88 B A stride - constexpr int kBPlaneStride = kBlockN * 8; // n rows x 8 B per plane - constexpr int kATileBytes = kBlockM * kTileStride; - constexpr int kBTileBytes = 8 * kBlockN * 8; // 8 k-planes of n rows x 8 B - constexpr int kASlots = (kBlockM * kStageK) / 16; // int4 A chunks - constexpr int kBSlots = kBTileBytes / 16; // int4 B chunks - constexpr int kALoads = (kASlots + kThreads - 1) / kThreads; - constexpr int kBLoads = (kBSlots + kThreads - 1) / kThreads; - constexpr int kAccM = kQuadM / kTileM; // A fragments per wave (4 / 2) - constexpr int kAccN = kQuadN / kTileN; // B fragments per wave - constexpr int kWavesN = kBlockN / kQuadN; - - const int tid = static_cast(threadIdx.x); - const int wave = tid / kWaveSize; - const int lane = tid % kWaveSize; - const int wave_row = wave / kWavesN; - const int wave_col = wave % kWavesN; - const int m0 = static_cast(blockIdx.y) * kBlockM; - const int n0 = static_cast(blockIdx.x) * kBlockN; - const int local_row = wave_row * kQuadM; - const int local_col = wave_col * kQuadN; - const int base_row = m0 + local_row; - const int base_col = n0 + local_col; - - // Single-buffered stage: A[kBlockM, 80] + B plane tile [8][kBlockN][8]. - __shared__ __align__(16) int8_t a_tile[kATileBytes]; - __shared__ __align__(16) int8_t b_tile[kBTileBytes]; - - DUFragment - a_frag[kAccM]; - DUFragment - b_frag[kAccN]; - DUFragment acc[kAccM][kAccN]; -#pragma unroll - for (int i = 0; i < kAccM; ++i) { -#pragma unroll - for (int j = 0; j < kAccN; ++j) { - du_fill_fragment(acc[i][j], 0); } - } - - // Register-batched epilogue scales (prefetched on the last K stage). - float xs_m[kAccM]; - float4 ws_m[kAccN]; - if constexpr (kPrefetchNext) { - // ----------------------------------------------------------------------- - // Iteration 16 (kv_b_proj load-placement round) - ACTIVE config-3 path: - // single-buffered two-barrier stage with the grouped A+B staging loads - // for stage s+1 issued at the END of stage s (after the MMAC burst, - // before the protective trailing barrier). a_chunk/b_chunk live across - // the trailing barrier + the loop back-edge only (never across a burst), - // so peak VGPR stays <= 64 -> 2 blocks/CU = 16 waves/CU unchanged; LDS - // 13,312 B/block and 16 s_barriers/block unchanged. The vmcnt wait for - // the in-flight loads is inserted by the compiler before the first LDS - // write of the next stage, i.e. each staging L2 round trip overlaps the - // trailing barrier + back-edge instead of stalling the LDS writes. - // ----------------------------------------------------------------------- - int4 a_chunk[kALoads]; - bool a_ok[kALoads]; - int4 b_chunk[kBLoads]; + // Iteration 13 (consolidate round): issue the next 64-K stage's A and B + // chunks before the trailing barrier so their L2 round trip overlaps the + // barrier wait and the next stage's commit instead of being drained + // immediately at the stage top. The last stage re-reads its own chunk + // (k1 == k0) rather than running past the final K block, so every + // prefetch address stays in bounds. { - // Prologue: issue the stage-0 A+B loads (the same grouped pair and the - // same addresses as the iteration-7 stage-top placement) so the first - // iteration starts directly with the LDS writes. - const int k0 = 0; - const int kstage = 0; // 64-k stage index inside the swizzled pack -#pragma unroll - for (int i = 0; i < kALoads; ++i) { - const int slot = tid + i * kThreads; - a_ok[i] = slot < kASlots; - if (a_ok[i]) { - const int row = slot >> 2; // 0..kBlockM-1 - const int col = (slot & 3) * 16; // 0/16/32/48 - a_chunk[i] = *reinterpret_cast( - x_q + static_cast(m0 + row) * k + k0 + col); - } - } -#pragma unroll - for (int i = 0; i < kBLoads; ++i) { - const int slot = tid + i * kThreads; - if (slot < kBSlots) { - const int b_kc = slot / (kBlockN / 2); // plane 0..7 - const int b_j = slot - b_kc * (kBlockN / 2); // n-pair 0..BN/2-1 - const int b_row0 = (n0 + 2 * b_j) & 127; - const int b_nt = (n0 + 2 * b_j) >> 7; - b_chunk[i] = *reinterpret_cast( - weight + (((static_cast(kstage) * (n >> 7) + b_nt) * 8 + - b_kc) * 128 + b_row0) * 8); - } - } - } - for (int k0 = 0; k0 < k; k0 += kStageK) { - const int kstage = k0 >> 6; // 64-k stage index inside the swizzled pack -#pragma unroll - for (int i = 0; i < kALoads; ++i) { - if (a_ok[i]) { - const int slot = tid + i * kThreads; - const int row = slot >> 2; - const int col = (slot & 3) * 16; - if constexpr (kTileStride % 16 == 8) { - // Iteration 17 (kv_b_proj consolidation round): 88-byte A rows - // are only 8-byte aligned, so stage the 16-byte chunk as two - // int64 halves (ds_write2_b64, not ds_write_b128 - lineage rule - // 3). Same 16 staging bytes at the same logical (row, kk) - // slots, so bit-identity holds. - int8_t* dst = a_tile + row * kTileStride + col; - const int64_t* src = - reinterpret_cast(&a_chunk[i]); - *reinterpret_cast(dst) = src[0]; - *reinterpret_cast(dst + 8) = src[1]; - } else { - *reinterpret_cast(a_tile + row * kTileStride + col) = - a_chunk[i]; - } - } - } -#pragma unroll - for (int i = 0; i < kBLoads; ++i) { - const int slot = tid + i * kThreads; - if (slot < kBSlots) { - const int b_kc = slot / (kBlockN / 2); - const int b_j = slot - b_kc * (kBlockN / 2); - *reinterpret_cast(b_tile + b_kc * kBPlaneStride + b_j * 16) = - b_chunk[i]; - } - } - __syncthreads(); - - // Each wave consumes its 64 x kQuadN quadrant: kAccM * kAccN - // m16n16k32 MMACs per kk. Accumulation order: k0-outer over 64-K - // stages, kk-inner (kk = 0 then 32), matching the reference int32 - // accumulation. All fragment loads are issued before the MMAC burst - // (load-all-then-MMAC-all, fully unrolled like the q_b_proj winner). -#pragma unroll - for (int kk = 0; kk < kStageK; kk += kTileK) { -#pragma unroll - for (int i = 0; i < kAccM; ++i) { - load_fragment8(a_frag[i], - a_tile + (local_row + i * kTileM) * kTileStride + kk, - kTileStride, lane); - } -#pragma unroll - for (int j = 0; j < kAccN; ++j) { - load_fragment8_plane_stride( - b_frag[j], - b_tile + (kk >> 3) * kBPlaneStride + - (local_col + j * kTileN) * 8, - kBPlaneStride, lane); - } -#pragma unroll - for (int j = 0; j < kAccN; ++j) { -#pragma unroll - for (int i = 0; i < kAccM; ++i) { - du_mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]); - } - } - } - - // Last stage only: prefetch the epilogue's per-row x_scale and - // per-column weight_scale values into registers (uniform branch: k0 is - // block-uniform, so the barrier below is reached by every thread; the - // row < m guard keeps the x_scale reads in bounds for any m tail). - if (k0 + kStageK >= k) { - const int r = lane & 15; - const int c4 = lane >> 4; -#pragma unroll - for (int i = 0; i < kAccM; ++i) { - const int row = base_row + i * kTileM + r; - xs_m[i] = (row < m) ? x_scale[row] : 0.0f; - } -#pragma unroll - for (int j = 0; j < kAccN; ++j) { - ws_m[j] = *reinterpret_cast( - weight_scale + base_col + j * kTileN + 4 * c4); - } - } - // Iteration 16 (kv_b_proj load-placement round): issue the grouped A+B - // staging loads for the NEXT stage here - after the burst and the - // epilogue scale prefetch, before the protective trailing barrier - so - // their L2 round trip overlaps the barrier + loop back-edge instead of - // stalling the next stage's LDS writes. Uniform branch (k0 is - // block-uniform); only GLOBAL memory is touched, so the barrier - // semantics are unchanged. On the last stage no loads are issued and - // the placement degenerates to the accepted schedule. - if (k0 + kStageK < k) { - const int k0n = k0 + kStageK; - const int kstagen = k0n >> 6; // 64-k stage index of the next stage -#pragma unroll - for (int i = 0; i < kALoads; ++i) { - const int slot = tid + i * kThreads; - a_ok[i] = slot < kASlots; - if (a_ok[i]) { - const int row = slot >> 2; // 0..kBlockM-1 - const int col = (slot & 3) * 16; // 0/16/32/48 - a_chunk[i] = *reinterpret_cast( - x_q + static_cast(m0 + row) * k + k0n + col); - } - } -#pragma unroll - for (int i = 0; i < kBLoads; ++i) { - const int slot = tid + i * kThreads; - if (slot < kBSlots) { - const int b_kc = slot / (kBlockN / 2); // plane 0..7 - const int b_j = slot - b_kc * (kBlockN / 2); // n-pair 0..BN/2-1 - const int b_row0 = (n0 + 2 * b_j) & 127; - const int b_nt = (n0 + 2 * b_j) >> 7; - b_chunk[i] = *reinterpret_cast( - weight + - (((static_cast(kstagen) * (n >> 7) + b_nt) * 8 + - b_kc) * 128 + b_row0) * 8); - } - } - } - // Protect the LDS buffers from the next stage's cooperative overwrite. - __syncthreads(); + const int k1 = (k0 + kStageK < k) ? (k0 + kStageK) : k0; + a_chunk = *reinterpret_cast(a_stage_base + k1); + b_chunk = *reinterpret_cast( + weight + static_cast(k1 / kStageK) * b_stage_step + + b_stage_base + b_tile_off); } - } else { - // ----------------------------------------------------------------------- - // Retention path (configs 0-2): EXACT iteration-7 loop, byte-identical. - // ----------------------------------------------------------------------- - for (int k0 = 0; k0 < k; k0 += kStageK) { - const int kstage = k0 >> 6; // 64-k stage index inside the swizzled pack - // Iteration 7 (compute-pipeline round): issue the A and B staging global - // loads for the WHOLE stage before either vmcnt wait (issue grouping / - // prefetch distance = one VMEM round trip), then one wait and both LDS - // writes. The iteration-6 exact code object serializes the staging as - // global_load_dwordx4 A -> s_waitcnt vmcnt(0) -> ds_write_b128 A - // global_load_dwordx4 B -> s_waitcnt vmcnt(0) -> ds_write_b128 B - // (zero instructions between each load and its wait, and the compiler - // reuses one 4-VGPR slot for both loads), so each 64-K stage exposes ~2 - // L2 round trips back-to-back on the barrier-to-barrier critical path. - // The A and B loads are independent (A: x_q rows; B: swizzled-pack tile), - // so grouping them as one issued pair overlaps the two latencies into ~1 - // round trip. Staging content, LDS addresses, the int32 accumulation - // order (k0-outer / kk-inner) and the epilogue are byte-identical; the - // loads simply complete concurrently instead of serially. - int4 a_chunk[kALoads]; - bool a_ok[kALoads]; -#pragma unroll - for (int i = 0; i < kALoads; ++i) { - const int slot = tid + i * kThreads; - a_ok[i] = slot < kASlots; - if (a_ok[i]) { - const int row = slot >> 2; // 0..kBlockM-1 - const int col = (slot & 3) * 16; // 0/16/32/48 - a_chunk[i] = *reinterpret_cast( - x_q + static_cast(m0 + row) * k + k0 + col); - } - } - int4 b_chunk[kBLoads]; -#pragma unroll - for (int i = 0; i < kBLoads; ++i) { - const int slot = tid + i * kThreads; - if (slot < kBSlots) { - const int b_kc = slot / (kBlockN / 2); // plane 0..7 - const int b_j = slot - b_kc * (kBlockN / 2); // n-pair 0..BN/2-1 - // Iteration 5 (packing round): tile-contiguous (512, 3584) pack - - // the (kstage, n-tile) B tile is ONE contiguous 8192-B region - // packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b]; the 16-B - // chunk is two consecutive n rows (n0 + 2j, +1) of one 8-byte k - // sub-chunk. It never crosses a 128-row tile boundary (n0 % 64 == 0 - // and a 64/128-row window starting at a multiple of 64 stays inside - // one 128-row pack tile), so the int4 is aligned in global and LDS. - const int b_row0 = (n0 + 2 * b_j) & 127; - const int b_nt = (n0 + 2 * b_j) >> 7; - b_chunk[i] = *reinterpret_cast( - weight + (((static_cast(kstage) * (n >> 7) + b_nt) * 8 + - b_kc) * 128 + b_row0) * 8); - } - } -#pragma unroll - for (int i = 0; i < kALoads; ++i) { - if (a_ok[i]) { - const int slot = tid + i * kThreads; - const int row = slot >> 2; - const int col = (slot & 3) * 16; - if constexpr (kTileStride % 16 == 8) { - // Iteration 17 (kv_b_proj consolidation round): 88-byte A rows - // are only 8-byte aligned, so stage the 16-byte chunk as two - // int64 halves (ds_write2_b64, not ds_write_b128 - lineage rule - // 3). Same 16 staging bytes at the same logical (row, kk) - // slots, so bit-identity holds. - int8_t* dst = a_tile + row * kTileStride + col; - const int64_t* src = - reinterpret_cast(&a_chunk[i]); - *reinterpret_cast(dst) = src[0]; - *reinterpret_cast(dst + 8) = src[1]; - } else { - *reinterpret_cast(a_tile + row * kTileStride + col) = - a_chunk[i]; - } - } - } -#pragma unroll - for (int i = 0; i < kBLoads; ++i) { - const int slot = tid + i * kThreads; - if (slot < kBSlots) { - const int b_kc = slot / (kBlockN / 2); - const int b_j = slot - b_kc * (kBlockN / 2); - *reinterpret_cast(b_tile + b_kc * kBPlaneStride + b_j * 16) = - b_chunk[i]; - } - } - __syncthreads(); - - // Each wave consumes its 64 x kQuadN quadrant: kAccM * kAccN - // m16n16k32 MMACs per kk. Accumulation order: k0-outer over 64-K - // stages, kk-inner (kk = 0 then 32), matching the reference int32 - // accumulation. All fragment loads are issued before the MMAC burst - // (load-all-then-MMAC-all, fully unrolled like the q_b_proj winner). -#pragma unroll - for (int kk = 0; kk < kStageK; kk += kTileK) { -#pragma unroll - for (int i = 0; i < kAccM; ++i) { - load_fragment8(a_frag[i], - a_tile + (local_row + i * kTileM) * kTileStride + kk, - kTileStride, lane); - } -#pragma unroll - for (int j = 0; j < kAccN; ++j) { - load_fragment8_plane_stride( - b_frag[j], - b_tile + (kk >> 3) * kBPlaneStride + - (local_col + j * kTileN) * 8, - kBPlaneStride, lane); - } -#pragma unroll - for (int j = 0; j < kAccN; ++j) { -#pragma unroll - for (int i = 0; i < kAccM; ++i) { - du_mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]); - } - } - } - // Last stage only: prefetch the epilogue's per-row x_scale and - // per-column weight_scale values into registers (uniform branch: k0 is - // block-uniform, so the barrier below is reached by every thread; the - // row < m guard keeps the x_scale reads in bounds for any m tail). - if (k0 + kStageK >= k) { - const int r = lane & 15; - const int c4 = lane >> 4; -#pragma unroll - for (int i = 0; i < kAccM; ++i) { - const int row = base_row + i * kTileM + r; - xs_m[i] = (row < m) ? x_scale[row] : 0.0f; - } -#pragma unroll - for (int j = 0; j < kAccN; ++j) { - ws_m[j] = *reinterpret_cast( - weight_scale + base_col + j * kTileN + 4 * c4); - } - } - // Protect the LDS buffers from the next stage's cooperative overwrite. - __syncthreads(); - } + // Protect the LDS buffers from the next stage's cooperative overwrite. + __syncthreads(); } - // Coalesced direct-fragment epilogue: one 8-byte bf16 store per lane per - // fragment (100% store sector efficiency), scales from caller registers. - // Iteration 6 (kv_b_proj epilogue round): kRneOnly = true - plain RNE bf16 - // conversion without the exec-masked inf/NaN fixup (bit-identical for all - // finite inputs; see bf16_rne_u16). -#pragma unroll - for (int j = 0; j < kAccN; ++j) { -#pragma unroll - for (int i = 0; i < kAccM; ++i) { - store_prefill_fragment_coalesced_scaled( - acc[i][j], xs_m[i], ws_m[j], out, base_row + i * kTileM, - base_col + j * kTileN, m, n, lane); - } - } + // --------------------------------------------------------------------- + // Iteration-6 (epilogue round): register-batched coalesced epilogue, kept + // unchanged; the iteration-8 eight-wave quadrant only changes how many + // fragments each wave owns (2 row-blocks x 4 col-blocks = 8 fragments). + // Iteration 10 (consolidate round): the dead per-fragment tail-M guard is + // removed from this arm. The dispatch guard is m >= 128 with m % 128 == 0, + // so base_row + (lane & 15) is always < m and the two (rs < m) x_scale + // ternaries and the (row >= m) early return in the store helper never fire; + // as written they forced each of the eight fragment epilogues into its own + // exec-masked basic block (eight s_and_saveexec / s_or_b64 / s_cbranch + // groups and eight redundant lane-id / shuffle-index computations in the + // iteration-10 code object). Removing them keeps the exact 128x128 tile, + // the eight-wave geometry, the 16,384-B single-buffered staging, the LDS + // fragment loads, the int32 accumulation and the coalesced 8-byte store, + // and lets the eight independent fragment epilogues issue as one + // straight-line block. The 128x64 generic arm and the scalar fallback keep + // their own bounds checks unchanged. + // --------------------------------------------------------------------- + const int gcol_base = n0 + local_col; + const int rs0 = m0 + local_row + (lane & 15); + const float xs0 = x_scale[rs0]; + const int rs1 = rs0 + kTileM; + const float xs1 = x_scale[rs1]; + store_fragment_coalesced_scaled(acc00, xs0, weight_scale, out, + m0 + local_row, gcol_base, n, lane); + store_fragment_coalesced_scaled(acc01, xs0, weight_scale, out, + m0 + local_row, gcol_base + kTileN, n, lane); + store_fragment_coalesced_scaled(acc02, xs0, weight_scale, out, + m0 + local_row, gcol_base + 2 * kTileN, n, + lane); + store_fragment_coalesced_scaled(acc03, xs0, weight_scale, out, + m0 + local_row, gcol_base + 3 * kTileN, n, + lane); + store_fragment_coalesced_scaled(acc10, xs1, weight_scale, out, + m0 + local_row + kTileM, gcol_base, n, lane); + store_fragment_coalesced_scaled(acc11, xs1, weight_scale, out, + m0 + local_row + kTileM, gcol_base + kTileN, + n, lane); + store_fragment_coalesced_scaled(acc12, xs1, weight_scale, out, + m0 + local_row + kTileM, + gcol_base + 2 * kTileN, n, lane); + store_fragment_coalesced_scaled(acc13, xs1, weight_scale, out, + m0 + local_row + kTileM, + gcol_base + 3 * kTileN, n, lane); } // --------------------------------------------------------------------------- // Generic scalar fallback: one thread per output element. Exact int32 dot -// over K, then fused x_scale * weight_scale, then bf16 store. Handles every -// unmatched (m, n, k), including all small-M API cases (M=2, M=16), M tails -// and M < 128 with the same (K, N). kNMajorPack == true decodes the -// iteration-5 swizzled 64-k-stage layout -// packed[((k0*8+kc)*n+col)*8+b] == raw[kk*n+col] (kk = k0*64+kc*8+b) for the -// exact (k, n) == (2048, 2048); otherwise the weight is the raw [K, N] -// row-major identity layout. +// over K, then fused x_scale * weight_scale, then bf16 store. kPackedNMajor +// == true decodes the iteration-5 tile-contiguous 64-K-stage pack +// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] = raw[kk*n + col] for +// exactly (k, n) == (512, 3584); false reads the identity [K, N] row-major +// layout. Handles every unmatched (m, n, k), including all small-M API cases +// (M=2, M=16), the paired same-(K, N) shapes, and M < 128 with the exact +// packed pair. // --------------------------------------------------------------------------- -template +template __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, @@ -1535,39 +929,21 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( const int col = static_cast(linear - static_cast(row) * n); int32_t acc = 0; const int8_t* a_row = x_q + static_cast(row) * k; - if constexpr (kNMajorPack) { - if (k == kPackSwizzleK2 && n == kPackSwizzleN2) { - // Tile-contiguous pack (iteration-5 kv_b_proj packing round, exact - // (512, 3584)): packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] == - // raw[kk*n + col] with kk = k0*64 + kc*8 + b, nt = col/128, - // row = col%128, so raw[kk][col] = packed[...]. - for (int kk = 0; kk < k; ++kk) { - const int k0 = kk >> 6; - const int kc = (kk >> 3) & 7; - const int b = kk & 7; - const int nt = col >> 7; - const int row = col & 127; - const int8_t* p = - weight + - ((((static_cast(k0) * (n >> 7) + nt) * 8 + kc) * 128 + - row) * 8) + - b; - acc += static_cast(a_row[kk]) * static_cast(*p); - } - } else { - // Swizzled 64-k-stage pack (iteration 5; exact (2048,2048) q_b_proj - // pair): packed[((k0*8 + kc)*n + col)*8 + b] == raw[kk*n + col] with - // kk = k0*64 + kc*8 + b, so raw[kk][col] = - // packed[((k0*8+kc)*n+col)*8+b]. - for (int kk = 0; kk < k; ++kk) { - const int k0 = kk >> 6; - const int kc = (kk >> 3) & 7; - const int b = kk & 7; - const int8_t* p = weight + - (((static_cast(k0) * 8 + kc) * n + col) * 8) + - b; - acc += static_cast(a_row[kk]) * static_cast(*p); - } + if constexpr (kPackedNMajor) { + // Tile-contiguous pack: packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8+b] + // == raw[kk*n + col] with kk = k0*64 + kc*8 + b, nt = col/128, + // row = col%128. + const int n_row = col & (kPackTileN - 1); + const int nt = col / kPackTileN; + for (int kk = 0; kk < k; ++kk) { + const int k0 = kk / kPackStageK; + const int kc = (kk >> 3) & 7; + const int b = kk & 7; + const int64_t tile = + static_cast(k0) * (n / kPackTileN) + nt; + const int8_t* p = + weight + ((tile * 8 + kc) * kPackTileN + n_row) * 8 + b; + acc += static_cast(a_row[kk]) * static_cast(*p); } } else { // identity [K, N] row-major: column col is strided by n. @@ -1583,37 +959,20 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( } // --------------------------------------------------------------------------- -// Identity device-to-device weight packing (outside the timed region and -// outside Graph capture) for every (k, n) except the exact swizzled pairs -// (2048, 2048) q_b_proj and (512, 3584) kv_b_proj (which use the -// iteration-5 swizzled pack). The packed buffer keeps the same byte count -// K*N and the same allocated address, so the layout is graph-stable. +// Iteration-5 packing round: one-time tile-contiguous 64-K-stage permutation +// for exactly (k, n) == (512, 3584): +// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] = raw[kk*n + col], +// kk = k0*64 + kc*8 + b, nt = col/128, row = col%128. +// Every (64-K stage, 128-N tile) DUMMA B tile is ONE contiguous 8192-B global +// region laid out as [kc = 8-byte k sub-chunk][128 N rows][8 bytes], which is +// exactly the plane order the GEMM stages into LDS, so the staging is a plain +// contiguous 16-B vector copy and the B fragment reads are lane-linear and +// bank-conflict-free. Byte-wise (one thread per byte) so the permutation is +// trivially correct; runs once outside the timed region and outside Graph +// capture, keeping the same byte count and the same graph-stable buffer +// address. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( - const int8_t* __restrict__ raw, - int8_t* __restrict__ packed, - int64_t count) { - const int64_t linear = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (linear < count) { - packed[linear] = raw[linear]; - } -} - -// Iteration 5 (packing round): swizzled 64-k-stage-major permutation for the -// exact swizzled (k, n) pair q_b_proj (2048, 2048) only (kv_b_proj -// (512, 3584) switched to the tile-contiguous w8a8_pack_tile_major_kernel -// in the kv_b_proj iteration-5 packing round): -// packed[((k0*8 + kc)*n + col)*8 + b] = raw[kk*n + col], -// kk = k0*64 + kc*8 + b (k0 = 64-k stage, kc = 8-byte sub-chunk, b = byte). -// Byte-wise (one thread per byte) so the permutation is trivially correct; -// runs once outside the timed region and outside Graph capture, keeping the -// same byte count K*N and the same graph-stable buffer address. The layout -// lets the 128x128 kernel stage each B tile as 16-byte plane chunks (one -// aligned int4 global read per chunk, perfectly coalesced) into the plane -// LDS tile [kc][n][8] whose fragment loads are lane-linear with zero bank -// conflicts. -__global__ __launch_bounds__(256) void w8a8_pack_swizzle_kernel( +__global__ __launch_bounds__(256) void w8a8_pack_tile_major_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, int k, @@ -1624,46 +983,31 @@ __global__ __launch_bounds__(256) void w8a8_pack_swizzle_kernel( if (linear < total) { const int kk = static_cast(linear / n); const int col = static_cast(linear - static_cast(kk) * n); - const int k0 = kk >> 6; + const int k0 = kk / kPackStageK; const int kc = (kk >> 3) & 7; const int b = kk & 7; - packed[(((static_cast(k0) * 8 + kc) * n + col) * 8) + b] = - raw[linear]; + const int nt = col / kPackTileN; + const int row = col & (kPackTileN - 1); + const int64_t tile = static_cast(k0) * (n / kPackTileN) + nt; + packed[((tile * 8 + kc) * kPackTileN + row) * 8 + b] = raw[linear]; } } -// Iteration 5 (kv_b_proj packing round): tile-contiguous 64-k-stage-major -// permutation for the exact (k, n) == (512, 3584) pair only: -// packed[(((k0*(n/128) + nt)*8 + kc)*128 + row)*8 + b] = raw[kk*n + col], -// kk = k0*64 + kc*8 + b, nt = col/128, row = col%128. -// Every (64-k stage, 128-n tile) B tile is ONE contiguous 8192-B global -// region; the tile-family kernel stages it as 8 consecutive 1024-B -// wavefront streams (the k-stage-major swizzle read 8 streams 28,672 B -// apart). Byte-wise (one thread per byte) so the permutation is trivially -// correct; runs once outside the timed region and outside Graph capture, -// keeping the same byte count K*N and the same graph-stable buffer address. -// The (2048, 2048) q_b_proj pair keeps the k-stage-major swizzle -// (w8a8_pack_swizzle_kernel) byte-identical. -__global__ __launch_bounds__(256) void w8a8_pack_tile_major_kernel( +// --------------------------------------------------------------------------- +// Bootstrap identity device-to-device packing for every (k, n) except the +// exact tile-major-packed pair. Runs once outside the timed +// region and outside Graph capture, keeps the same byte count and +// graph-stable buffer addresses, and is the generic fallback for every +// (k, n). +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, - int k, - int n) { + int64_t count) { const int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t total = static_cast(k) * n; - if (linear < total) { - const int kk = static_cast(linear / n); - const int col = static_cast(linear - static_cast(kk) * n); - const int k0 = kk >> 6; - const int kc = (kk >> 3) & 7; - const int b = kk & 7; - const int nt = col >> 7; - const int row = col & 127; - packed[((((static_cast(k0) * (n >> 7) + nt) * 8 + kc) * 128 + - row) * - 8) + - b] = raw[linear]; + if (linear < count) { + packed[linear] = raw[linear]; } } @@ -1685,187 +1029,89 @@ __global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( // synchronization, no default-stream launch; they run on the caller-provided // PyTorch stream and are CUDA-Graph safe. // --------------------------------------------------------------------------- -extern "C" void launch_w8a8_gemm( - const int8_t* a, - const int8_t* b, - const float* x_scale, - const float* weight_scale, - void* out, - void* workspace, - int64_t workspace_bytes, - int m, - int n, - int k, - hipStream_t stream) { +extern "C" void launch_w8a8_gemm(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { (void)workspace; // no split-K: the GEMM does not use the workspace (void)workspace_bytes; auto* out_bf16 = reinterpret_cast(out); - // Explicit dispatch. Large-M shapes with exact tiled geometry (both - // assigned shapes: M=4096, N in {2048, 3584}, K in {2048, 512}) take the - // native INT8 DUMMA tiled path; every other (m, n, k) - including small-M - // API cases (M=2, M=16) and M < 64 - takes the scalar fallback. The - // exact (k, n) == (2048, 2048) q_b_proj pair takes the iteration-4 128x128 - // tile with the iteration-5 swizzled-pack plane B layout; the exact - // (k, n) == (512, 3584) kv_b_proj pair takes the iteration-1 tile family - // (w8a8_dumma_prefill_tile_kernel, active config kActiveKvTile, also - // swizzled-pack B); every other (k, n) uses the byte-identical 128x64 - // kNMajorB=false arm. The scalar fallback decodes the swizzled pack for - // both swizzled (k, n) pairs. - if (k == kPackSwizzleK2 && n == kPackSwizzleN2 && m >= 64 && m % 64 == 0) { - // Exact kv_b_proj (512, 3584): 2-D macro-tile family. All four - // instantiations are referenced (so all are compiled); only - // kActiveKvTile is launched. grid (N/BN, M/BM), block (waves*64): - // 0 -> 64x64, 1 wavefront/block, (56, 64) = 3584 blocks - // 1 -> 64x128, 2 wavefronts/block, (28, 64) = 1792 blocks - // 2 -> 128x64, 4 wavefronts/block, (56, 32) = 1792 blocks - // 3 -> 64x128, 8 wavefronts/block, (28, 64) = 1792 blocks - const dim3 grid64x64(static_cast(n / 64), - static_cast(m / 64)); - const dim3 block64x64(1 * kWaveSize); - const dim3 grid64x128(static_cast(n / 128), - static_cast(m / 64)); - const dim3 block64x128(2 * kWaveSize); - const dim3 block64x128_8w(8 * kWaveSize); - const dim3 grid128x64(static_cast(n / 64), - static_cast(m / 128)); - const dim3 block128x64(4 * kWaveSize); - if (kActiveKvTile == 0) { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tile_kernel<64, 64, 64>), - grid64x64, block64x64, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - } else if (kActiveKvTile == 1) { - hipLaunchKernelGGL( - HIP_KERNEL_NAME( - w8a8_dumma_prefill_tile_kernel), - grid64x128, block64x128, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - } else if (kActiveKvTile == 2) { - // config 2 (128x64) additionally needs m % 128 == 0; any (512, 3584) - // m not covered by the active tile falls back to the swizzle decode. - if (m % 128 == 0) { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_tile_kernel<128, 64, 32>), - grid128x64, block128x64, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - } else { - constexpr int kBlock = 256; - const int64_t total = static_cast(m) * n; - const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); - const dim3 block(kBlock); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - } - } else { - // config 3 (iteration-4 tile-shape round): 64x128 with 32x32 - // quadrants, 512 threads. m % 64 == 0 is guaranteed by the exact-shape - // guard above, so no tail fallback is needed. - // Iteration 16 (kv_b_proj load-placement round): kPrefetchNext = true - // -> the grouped A+B staging loads for stage s+1 are issued after the - // burst of stage s (before the protective trailing barrier), partially - // hiding each of the 8 staging L2 round trips behind the barrier + - // loop back-edge at unchanged 2 blocks/CU = 16 waves/CU (no register - // payload crosses the burst). Iteration 17 (kv_b_proj consolidation - // round): kALdsPad = 24 additionally applies the iteration-13 A-side - // LDS bank skew (88-byte A row stride -> zero-conflict A fragment - // reads; ds_write2_b64 halves at the staging writes) on top of the - // prefetch. LDS 13,312 -> 13,824 B/block (2 blocks = 27,648 B <= 64 - // KiB); 16 s_barriers/block unchanged. Configs 0-2 keep the defaults - // (kPrefetchNext = false, kALdsPad = 16) -> the exact iteration-7 - // issue-at-stage-top / stride-80 placement (byte-identical). - hipLaunchKernelGGL( - HIP_KERNEL_NAME( - w8a8_dumma_prefill_tile_kernel<64, 128, 32, 32, true, 24>), - grid64x128, block64x128_8w, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - } + // Explicit dispatch. For exactly (k, n) == (512, 3584) the weight is the + // one-time tile-contiguous 64-K-stage pack: large-M uses the 128x128 + // plane-B direct-load DUMMA arm (iteration-8 resource round: eight + // wavefronts / 512 threads per block, the wavefront=64 geometry of the + // exact-shape arm), every other m (small-M API cases and M < 128) decodes + // the packed layout in the scalar fallback. Every other (m, n, k) keeps + // the bootstrap behaviour: the identity-layout 128x64 DUMMA arm (four + // wavefronts / 256 threads) for large shapes, the scalar fallback + // otherwise. + if (is_tile_major_pack(k, n) && m >= kNMBlockM && m % kNMBlockM == 0) { + const dim3 grid(static_cast(n / kNMBlockN), + static_cast(m / kNMBlockM)); + const dim3 block(kNMBlockThreads); + hipLaunchKernelGGL(w8a8_dumma_prefill_128x128_nmajor_kernel, grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else if (is_tile_major_pack(k, n)) { + constexpr int kBlock = 256; + const int64_t total = static_cast(m) * n; + const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); + const dim3 block(kBlock); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); } else if (m >= kBlockM && m % kBlockM == 0 && n % kBlockN == 0 && k % kStageK == 0) { const dim3 grid(static_cast(n / kBlockN), static_cast(m / kBlockM)); const dim3 block(kBlockThreads); - if (is_swizzled_pack(k, n)) { - // Exact q_b_proj (2048, 2048): iteration-4 tile-aspect arm, 128x128 - // tile with the iteration-5 swizzled-pack plane B layout; grid - // (N/128, M/128) = 512 blocks for M=4096. n == 2048 is divisible by - // 128 by the exact guard; the outer arm already guarantees - // m % 128 == 0, k % 64 == 0. - const dim3 grid128(static_cast(n / 128), - static_cast(m / 128)); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_128x128_kernel), - grid128, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - } else { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_128x64_kernel), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - } + hipLaunchKernelGGL(w8a8_dumma_prefill_128x64_kernel, grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); } else { constexpr int kBlock = 256; const int64_t total = static_cast(m) * n; const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); const dim3 block(kBlock); - if (is_swizzled_pack(k, n)) { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - } else { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); - } + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), grid, block, + 0, stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); } - (void)hipGetLastError(); } -extern "C" void launch_pack_w8a8_weight( - const int8_t* raw_weight, - const float* weight_scale, - int8_t* packed_weight, - float* packed_weight_scale, - int k, - int n, - hipStream_t stream) { +extern "C" void launch_pack_w8a8_weight(const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { constexpr int kBlock = 256; const dim3 block(kBlock); - // Pack: iteration-5 kv_b_proj packing round - tile-contiguous permutation - // for (512, 3584) kv_b_proj, k-stage-major swizzle for (2048, 2048) - // q_b_proj, identity device-to-device copy for every other (k, n). - // Packing runs once outside the timed region and keeps the same byte count - // (graph-stable addresses). const int64_t weight_count = static_cast(k) * n; - const dim3 weight_grid( - static_cast((weight_count + kBlock - 1) / kBlock)); if (is_tile_major_pack(k, n)) { - // Iteration 5 (kv_b_proj packing round): tile-contiguous pack for the - // exact (512, 3584) pair (one contiguous 8192-B region per (stage, - // n-tile) B tile); (2048, 2048) keeps the k-stage-major swizzle below. - hipLaunchKernelGGL(w8a8_pack_tile_major_kernel, - weight_grid, block, 0, stream, - raw_weight, packed_weight, k, n); - } else if (is_swizzled_pack(k, n)) { - hipLaunchKernelGGL(w8a8_pack_swizzle_kernel, - weight_grid, block, 0, stream, + // Iteration-5 packing round: one-time tile-contiguous 64-K-stage pack so + // the DUMMA arm stages each B tile with contiguous 16-B vector loads and + // LDS stores and loads B fragments lane-linearly with zero bank conflicts. + const dim3 tile_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_tile_major_kernel, tile_grid, block, 0, stream, raw_weight, packed_weight, k, n); } else { - hipLaunchKernelGGL(w8a8_pack_identity_kernel, - weight_grid, block, 0, stream, + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_identity_kernel, weight_grid, block, 0, stream, raw_weight, packed_weight, weight_count); } - const int64_t scale_count = n; - const dim3 scale_grid( - static_cast((scale_count + kBlock - 1) / kBlock)); - hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, - scale_grid, block, 0, stream, - weight_scale, packed_weight_scale, n); - (void)hipGetLastError(); + const dim3 scale_grid(static_cast((n + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, scale_grid, block, 0, + stream, weight_scale, packed_weight_scale, n); } // @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/q_b_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/q_b_proj.hip index 72932e37..b327ab35 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/q_b_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/glm52/TP8/M4096/q_b_proj.hip @@ -1,143 +1,267 @@ -// @@variant shape=glm_tp8_q_b_proj_m4096 commit=a9fcfc5e2ac1234af51849905834a850b5ef3448 added=2026-08-31 -// median_us=220.1 p90_us=220.5 speedup=65.35 baseline_us=1.439e+04 -// source=glm5-2-dsh-tp8-m4096-1-e6a280a2 +// @@variant shape=glm_tp8_q_b_proj_m4096 commit=9510ab4d00ea72a47db8821d2b56b35395c23ba5 added=2026-09-11 +// median_us=209.5 p90_us=210.2 speedup=68.68 baseline_us=1.439e+04 +// source=test-9-11-1-b01955b3/iteration_001/glm_tp8_q_b_proj_m4096 // INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. // -// Worker: worker_1 (physical GPU 1), assigned shapes: -// glm_tp8_q_b_proj_m4096 : M=4096, N=2048, K=2048 -// glm_tp8_kv_b_proj_m4096 : M=4096, N=3584, K=512 +// Worker worker_3 (physical GPU 3), assigned shapes: +// glm_tp8_q_b_proj_m4096 : M=4096, N=2048, K=2048 // // Logical operation (exact contract): // out[m, n] = bf16( int32_dot(x_q[m, :], raw_weight[:, n]) // * x_scale[m] * weight_scale[n] ) // -// Iteration 3 (first valid HIP experiment; iterations 1-2 were killed in the -// agent infrastructure before any proposal). Accepted-best source digest -// a64b6250... with one bounded mechanism: n-major packed B + col_major B -// fragments for the exact glm_tp8_q_b_proj_m4096 (k,n) == (2048, 2048). -// * Large-prefill path (m >= 128 with exact 128x64x64 geometry): native -// INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 output tile -// per block; four wavefronts (256 threads); each wave owns a 64x32 -// quadrant built from eight m16n16k32 int32 accumulator fragments; the -// block cooperatively vector-loads A[128,64] and B[64,64] into one -// single-buffered LDS stage (15,360 B total: 128*80 + 64*80 with 16 B -// padding per row for bank skew); two __syncthreads per stage; fused -// dot * x_scale[m] * weight_scale[n] epilogue stored directly as bf16 +// Iteration 1 bootstrap (correctness-first): +// * Large-M DUMMA path (m >= 128 && m % 128 == 0 && n % 64 == 0 && +// k % 64 == 0): native INT8 DUMMA m16n16k32 tiled kernel, one 128x64 +// output tile per block, four gfx928 wavefronts of 64 lanes (256 +// threads). The block cooperatively stages A[128,64] (row-major, +// 16-byte int4 loads) and B[64,64] (raw [K, N] row-major, 16-byte int4 +// loads) into a single-buffered LDS stage (128*80 + 64*80 = 15,360 B +// with 16-byte per-row padding for bank skew) and performs two +// __syncthreads per K stage. Each wave owns a 64x32 quadrant built from +// eight m16n16k32 int32 accumulator fragments (four M x two N). Int32 +// accumulation order is k0-outer over 64-K stages, kk-inner (kk=0 then +// kk=32). The fused epilogue computes +// float(dot) * x_scale[row] * weight_scale[col] and stores bf16 directly // from the accumulator fragments using the verified gfx928 int8 // m16n16k32 lane mapping (row = lane & 15, col = (lane >> 4) + 4*i). -// A is read with the library du_load_matrix_sync row-major loader. -// B is layout-templated: -// - kNMajorB == true (exact (k,n) == (2048, 2048), q_b_proj): weight -// is packed once to n-major packed[n*K + k] == raw[k*N + n] outside -// the timed region; the stage stores B n-major (n row, k contiguous, -// 80 B row stride); each B fragment is one contiguous 8-byte LDS run -// per lane, loaded with the lineage-validated load_fragment8 (one -// ds_read2_b64) - kills the 32 ds_read_u8 byte-gathers + mask/OR -// reassembly VALU seen in the accepted-best ISA. -// - kNMajorB == false (any other shape, e.g. kv_b_proj (512, 3584)): -// raw [K, N] row-major identity layout, row-major B loader (the -// accepted-best path, byte-identical). -// Grid dim3(N/64, M/128) = 1024 blocks (q_b_proj) / 1792 blocks -// (kv_b_proj) dwarfs the 120 CUs, so no split-K is needed. -// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), -// including all small-M API cases (M=2, M=16), M tails, and M in -// (0, 128) with the same (K, N); it decodes the n-major packed layout -// for the exact (2048, 2048) and the identity [K, N] layout otherwise. -// * launch_pack_w8a8_weight: n-major device-to-device permutation for -// (k, n) == (2048, 2048), identity copy for every other (k, n) (both the -// int8 weight and the fp32 scale). Packing never happens inside the -// timed GEMM and keeps the same byte count (graph-stable addresses). +// Grid = (N/64, M/128); for the assigned shape that is 32 x 32 = 1024 +// blocks, far more than the 120 CUs, so no split-K is needed. +// * launch_pack_w8a8_weight is the bootstrap identity device-to-device copy +// for both the int8 weight and the fp32 scale (for every (k, n) except the +// exact pair that iteration 2 packs n-major below), so the GEMM consumes +// the raw [K, N] row-major weight directly. +// * Generic scalar int8/int32 fallback (one thread per output element) for +// every unmatched (m, n, k), including all small-M API cases (M=2, M=16) +// and any M in (0, 128) sharing the assigned (N, K). // -// Iteration 4 (tile-aspect round): the exact (k, n) == (2048, 2048) q_b_proj -// arm switches from the 128x64 tile to a 128x128 tile -// (w8a8_dumma_prefill_128x128_kernel). With the n-major packed B + 8-byte -// load_fragment8 B reads in place (iteration 3), the remaining LDS-read cost -// is dominated by the A-side row-major fragment loads (4 x ds_read2_b32 per -// kk per wave). A 128x128 tile with four 64x64 quadrants keeps the same -// 4-wave x 8-accumulator-fragment count per wave (8 -> 16 m16n16k32 MMACs -// per kk from 8 LDS fragment reads, i.e. per-MMAC LDS reads drop 6/8 -> 8/16 -// = -33%), and the A-side global re-read halves (A reuse 32 -> 16, A traffic -// 268 -> 134 MiB) while B reuse stays 32 (134 MiB), balancing the per-byte -// A:B traffic at 1:1 and cutting total tile traffic 402 -> 268 MiB (-33%). -// LDS grows 15,360 -> 20,480 B/block (still 2 blocks/CU = 40,960 B <= 64 -// KiB); the grid becomes (N/128, M/128) = 512 blocks (~4.3/CU). The 16 -// accumulator fragments cost ~+32 VGPR; arch VGPR is expected <= 128 so the -// accepted kernel's 2-blocks/CU co-residency is preserved (falsified if the -// exact code object shows > 128 VGPR -> 1 block/CU). kv_b_proj keeps the -// byte-identical 128x64 kNMajorB=false arm and the scalar fallback is -// untouched. +// Iteration 2 (operand-reuse round): the exact (k, n) == (2048, 2048) +// q_b_proj pair keeps the 128x64 macro-tile and the cooperative A+B LDS +// staging (each staged byte is consumed by two of the four waves and each +// staged fragment by two MMACs, so cooperative staging halves global operand +// traffic versus direct-to-register global fragment loads: ~402 MiB instead of +// ~805 MiB for the assigned shape). The staged B path is made direct, +// vectorized and bank-safe instead: +// * The weight is packed once (outside the timed region, outside Graph +// capture, same byte count and same graph-stable buffer) to n-major +// packed[n * K + k] == raw[k * N + n], so 16 consecutive k bytes of one n +// row are one aligned, perfectly coalesced int4 global read. +// * The B LDS tile is stored n-major with a 72-byte (18-word) bank-skewed +// row stride; each lane's 8-byte col_major fragment slot then starts on +// bank (18 * (lane & 15) + 2 * (lane >> 4)) mod 32, and 18 * r mod 32 is a +// permutation of the 16 even bank residues, so a 16-lane LDS phase covers +// all 32 banks exactly once (zero conflicts) and every staged access is +// 8-byte aligned. +// * Each B fragment is loaded with one 8-byte LDS read (load_fragment8) +// instead of the row_major loader's eight strided ds_read_u8 byte gathers +// plus mask/OR reassembly VALU. The 8 operand bytes are identical in the +// same fragment slots, so the int32 accumulation is still bit-identical. +// The exact (k, n) == (2048, 2048) pair takes this arm; every other (k, n) +// keeps the identity pack and the byte-identical raw [K, N] k-major arm. +// +// Iteration 4 (tile-aspect round): for the exact (k, n) == (2048, 2048) +// q_b_proj arm the 128x64 macro-tile is widened along N to a 128x128 tile +// (w8a8_dumma_prefill_128x128_nmajor_kernel), giving an M:N tile aspect of +// 1:1. The 128x64 best kernel re-reads A once per N/64 = 32 column tiles +// (~268 MiB A traffic) against B's ~134 MiB (A:B = 2:1) and issues 6 LDS +// fragment reads per 8 MMACs in each 64x32 wave quadrant; the 128x128 tile +// halves the A-side re-reads (A reuse 32 -> 16, ~134 MiB) while B reuse stays +// 32 (~134 MiB), balancing per-byte A:B tile traffic at 1:1 and cutting total +// staged tile traffic ~402 -> ~268 MiB (-33%). Each wave now owns a 64x64 +// quadrant (four A x four B fragments, sixteen int32 accumulators) and issues +// 8 LDS fragment reads per 16 MMACs (-33% per MMAC). The n-major stride-72 B +// layout and the direct load_fragment8 reads from iteration 2 are kept +// unchanged; the block stages A[128,64] (two int4/thread, 80-byte-strided +// rows) and the packed B[128,64] (two 16-byte n-row chunks/thread into the +// 72-byte-strided n-major tile) into a 19,456 B single buffer, so two +// blocks/CU still co-reside. Accumulation order (k0-outer, kk-inner) and the +// fragment element-to-slot mapping are unchanged, so the int32 accumulation is +// bit-identical; the scalar fallback and the 128x64 raw [K, N] arm for every +// other (k, n) are untouched. // // Iteration 5 (packing round): the exact (k, n) == (2048, 2048) q_b_proj -// weight is re-packed once (outside timing/Graph, same byte count, same -// graph-stable buffer) into a swizzled 64-k-stage-major layout -// packed[((k0*8 + kc)*n + col)*8 + b] == raw[kk*n + col], -// kk = k0*64 + kc*8 + b (k0 = 64-k stage, kc = 8-byte sub-chunk, b = byte), -// and the 128x128 kernel stages each B tile as 16-byte plane chunks into an -// LDS plane layout b_tile[kc][n][8] (plane stride 1024 B, n stride 8 B). -// The exact code object of the accepted kernel shows every B fragment read -// is an 8-byte LDS access at (lane&15)*80 + (lane>>4)*8 (n-major stride-80 -// rows), and PMC counts 6,291,456 bank conflicts over 1,048,576 LDS slots -// (6.0/slot) with 4,078,494 LDS waits (3.9/slot): the loop is LDS-latency -// bound and the stride-80 rows alias bank phases every 8 rows (80*8 == 640 -// == 0 mod 128), so every fragment read conflicts ~4-6-way. In the plane -// layout a lane's 8 bytes sit at (lane>>4)*1024 + (lane&15)*8 relative to -// the fragment origin: 16 lanes per 128-B phase cover all 32 banks exactly -// once, i.e. B fragment reads become lane-linear and zero-conflict (4-cycle -// minimum for 512 B). Staging reads stay coalesced (each 16-B chunk is two -// n rows of one sub-chunk -> one aligned int4) and LDS writes are 2-way -// (the 16-B optimum). Same fragment operand bytes, same k0-outer/kk-inner -// int32 accumulation -> bit-identical results. The scalar fallback decodes -// the new pack for (2048,2048) only; kv_b_proj keeps the byte-identical -// 128x64 kNMajorB=false identity path. +// weight is packed FRAGMENT-MAJOR instead of n-major. One DUMMA matrix_b +// tile (m16n16k32 -> 16 n x 32 k int8 = 512 contiguous bytes) is one +// graph-stable, 16-byte-aligned block, so: +// * the packer writes a pure permutation of the same K * N bytes into the +// same buffer: byte (`n`, `k`) of the tile (`n/16`, `k/32`) sits at +// block * 512 + (k % 32) / 8 * 128 + (n % 16) * 8 + (k % 8); +// * the cooperative stage copies each packed 16-byte chunk with ONE aligned +// int4 global read and ONE aligned int4 LDS store (no 8-byte halves and +// no shift/or recombine VALU), and every staged 16-byte chunk is +// 16-byte aligned by construction; +// * each wave-level B fragment is one 8-byte LDS read per lane at +// (lane & 15) * 8 + (lane >> 4) * 128, i.e. word 2 * (lane & 15) plus a +// 32-word k-group stride inside the 512-byte tile, so a 16-lane LDS +// phase (lane & 15 = 0..15) covers all 32 banks exactly once with zero +// padding (bank-conflict-free, fully vector-loadable). +// The fragment bytes, their (k, n) slots and the k0-outer/kk-inner int32 +// accumulation order are unchanged, so correctness is bit-identical; the +// scalar fallback decodes the same permutation for the unmatched-M cases of +// the exact (k, n) pair. // // Iteration 6 (epilogue round): the fused dot * x_scale[m] * weight_scale[n] // -> bf16 epilogue has been in-kernel since iteration 1 and the workspace is -// unused ((void)workspace; no split-K, no combine pass anywhere in the call -// chain), so the remaining epilogue inefficiency of the exact (2048, 2048) -// 128x128 arm is the STORE pattern: store_prefill_fragment issues one -// 2-byte bf16 store per lane per element - 64 scattered stores per wave per -// fragment set, each wavefront store touching 16 rows x 4 columns so every -// 32-B sector is only 25% utilized (PMC: 131,072 vmem_write_instructions = -// 512 blocks x 4 waves x 64). The new epilogue transposes each 16x16 -// fragment's 4-element groups inside the 4-lane column group (lanes r, -// r+16, r+32, r+48; two 2x2 steps with __shfl_xor 16 then 32 - the -// lineage-validated 4x4 register transpose accepted on the sibling TP8 -// workers 0/2, same DTK 26.04/gfx928), so lane (r, c4) owns the four -// CONTIGUOUS columns 4*c4 .. +3 and writes ONE 8-byte store per lane (100% -// store sector efficiency; vmem_write 131,072 -> 32,768). Only int32 -// values move between lanes: the per-element multiply order -// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are -// unchanged, so stored bits are identical. The per-row x_scale (4 rows per -// lane) and per-column weight_scale float4s (4 per lane) are additionally -// register-batched on the last K stage (before the final protective -// __syncthreads), so the epilogue is pure compute + 16 coalesced stores with -// no interleaved vmem loads (vmem_read ~413,696 -> ~278k). kv_b_proj keeps -// the byte-identical 128x64 kNMajorB=false arm (old per-element epilogue) -// and the generic scalar fallback is untouched. +// never used ((void)workspace; no split-K, no combine pass anywhere in the +// call chain), so the remaining epilogue inefficiency of the exact +// (k, n) == (2048, 2048) 128x128 arm is the STORE pattern: +// store_prefill_fragment issues one 2-byte bf16 store per lane per element - +// 64 scattered stores per wave per fragment set, each touching 16 rows x 4 +// columns, so every 32-byte store sector is only 25% utilized (PMC: 131,072 +// vmem_write_instructions = 512 blocks x 4 waves x 64). The new epilogue +// transposes each 16x16 fragment's 4-element groups inside the 4-lane column +// group (lanes r, r+16, r+32, r+48; a 4x4 register transpose in two 2x2 steps +// with __shfl_xor 16 then 32), so lane (r, c4) owns the four CONTIGUOUS +// columns 4*c4 .. +3 and writes ONE aligned 8-byte store per lane (100% store +// sector efficiency; vmem_write 131,072 -> 32,768). Only int32 values move +// between lanes: the per-element multiply order (float(dot) * x_scale[row] * +// weight_scale[col]) and the bf16 rounding are unchanged, so stored bits are +// identical. The per-row x_scale (4 rows per lane) and per-column +// weight_scale float4s (4 per lane) are additionally register-batched on the +// last K stage before the final protective __syncthreads, so the epilogue is +// pure compute + 16 coalesced stores with no interleaved vmem loads. The +// 128x64 arm keeps the byte-identical per-element epilogue and the generic +// scalar fallback is untouched. +// +// Iteration 7 (compute-pipeline round): the 128x128 packed arm's per-K-stage +// critical path still begins with four cooperative 16-byte global staging loads +// whose latency is drained immediately by the s_waitcnt vmcnt(3)/(2)/(1)/(0) +// sequence that precedes the four ds_write_b128 (the code object shows the load +// and its store in the same step, so the L2/HBM round trip sits inside the +// barrier-to-barrier wall of every one of the 32 K stages; the PMC LDS wait +// count and the ~0.05 MMAC/cycle/SIMD issue rate show the warp stalling far +// more than issuing). The staging loop now carries a ONE-K-STAGE register +// lookahead: this thread's four staging chunks (A rows stage_row and +// stage_row + 64, packed B tiles b_blk and b_blk + 8) are loaded for K stage +// s + 1 right AFTER the publish barrier of stage s and BEFORE stage s's twelve +// fragment ds_reads and thirty-two v_mmac, so their latency is covered by the +// whole stage-s compute burst and the trailing WAR barrier, and the stage-s + 1 +// ds_write_b128 no longer waits on a just-issued load. The payload is four +// int4 (16 VGPRs) live across the compute; LDS size, barrier count (still two +// per K stage), staged bytes, the k0-outer/kk-inner int32 accumulation order +// and every fragment address are unchanged, so the int32 accumulation and the +// stored bits are bit-identical. The 128x64 packed/k-major arms and the +// generic scalar fallback are untouched. +// +// Iteration 9 (A-fragment direct-load round): the 128x128 packed arm reads its +// four A fragments per kk with the library du_load_matrix_sync row_major +// loader. du_mma.h defines that loader with the same +// x[i] = p[(lane & 15) * ldm + ((lane >> 4) << 3) + i] +// eight-contiguous-bytes mapping as the matrix_b col_major loader (the comment +// on load_fragment8_block below records exactly that mapping for both), but on +// this gfx928 code object the A fragments come out as pairs of ds_read2_b32 - +// two 4-byte LDS reads plus a byte reassembly chain per 8-byte fragment - and +// they sit between the LDS read and the dependent v_mmac, while the B side, +// direct since iterations 2 and 5, reads each fragment with one wide LDS +// access. Iteration 9 replaces the four library A loads with the direct +// 8-byte load_fragment8 on the A tile (row stride kAStride = 80 B), so each +// lane's x[0..7] is filled by one aligned 64-bit LDS read at +// (lane & 15) * kAStride + ((lane >> 4) << 3). The operand bytes, the +// k0-outer/kk-inner int32 accumulation order and every fragment address are +// unchanged, so the int32 accumulation stays bit-identical; only the +// redundant reassembly VALU and its lgkmcnt wait states leave the LDS->MMAC +// critical path. The A-tile layout, both barriers per K step, the staged +// bytes, the one-K-stage staging lookahead, the coalesced epilogue, the +// 128x64 packed/k-major arms and the generic scalar fallback are untouched. +// +// Iteration 10 (consolidation round, final HIP-only polish on the fixed +// 128x128 packed architecture): the arm's last measured LDS inefficiency is +// the A tile's 80-byte (20-word) row stride. The exact-source PMC of the +// accepted iteration-9 source (digest 4f7eee6f..., arch_vgpr 152, workgroup +// 256, grid 512, lds_bytes 18,432) reports lds_instructions 1,048,576 = the +// eight 8-byte fragment reads per kk (4 A + 4 B) x 32 stages x 4 waves x 512 +// blocks, and lds_bank_conflicts 3,145,728 = 3 x lds_instructions exactly. +// The B side is conflict-free by construction (fragment-major tiles: a lane's +// 8 bytes start at word 2*(lane&15), so a 16-lane phase touches all 32 banks +// once), so all 3,145,728 conflicts are A-side, i.e. 6 per A wave-instruction: +// with a 20-word stride the start word of row r is 20*r mod 32, which has +// period 8 (20*8 == 160 == 0 mod 32 == 0 mod 128 B), so lanes r and r+8 always +// land on the same bank pair and each 16-lane LDS phase runs at 2x the minimum +// (the exact aliasing the iteration-7 and iteration-9 hypotheses recorded and +// the sibling reference also left in place). A 16-byte-aligned row start can +// only ever produce 8 distinct bank residues mod 32 (starts are multiples of +// 4 words), so no padding of the existing 80-byte layout can remove it; the +// fix has to be an 8-byte-granular row stride. Focused change, 128x128 packed +// arm only: the A tile row stride becomes kAStride128 = 72 (18 words). 18 has +// gcd(18, 32) = 2, so 18*r mod 32 has period 16 and r = 0..15 is a permutation +// of the 16 even bank residues: every lane's 8-byte read occupies one distinct +// even bank and its following odd bank, i.e. all 32 banks exactly once per +// 16-lane phase (zero conflicts, the 4-cycle minimum for a 512-B fragment). +// Because 72 is only 8-byte aligned, the cooperative A staging keeps its two +// aligned int4 global loads per thread but publishes each 16-byte chunk with +// two aligned 8-byte (`int2`) LDS stores at +0/+8, so all LDS addresses stay +// naturally aligned; the whole-kernel LDS stage shrinks 18,432 -> 17,408 B +// (A 10,240 -> 9,216 B) and every staged byte keeps its exact (row, k) slot. +// The A fragment operand bytes, the 128-byte global A reads, the packed-B +// path, both barriers per K step, the one-K-stage staging lookahead, the +// coalesced epilogue, the k0-outer/kk-inner int32 accumulation order and the +// 128x64 packed/k-major arms plus the generic scalar fallback are untouched, +// so the int32 accumulation and the stored bits are bit-identical. // -// Iteration 7 (compute-pipeline round): the exact (k, n) == (2048, 2048) -// 128x128 arm loads its four A fragments per kk with the lineage-validated -// direct load_fragment8 (one ds_read2_b64 straight into the v_mmac operand) -// instead of the library du_load_matrix_sync row_major loader. The exact -// code object shows the library loader lowers to 8 x ds_read2_b32 plus a -// redundant byte-reassembly chain per fragment (~7 VALU: v_and 0xff00 / -// 0xff0000 / 0xff000000 + v_or_b32_sdwa + v_or3, ~50-56 VALU per stage per -// wave) between the LDS read and the first MMAC, and PMC counts 3,009,148 -// LDS waits (2.3/slot) against a latency-bound loop running at 1 block/CU -// (160 VGPR). du_mma.hpp defines matrix_a row_major int8 as x[i] = -// p[(lane&15)*ldm + (lane>>4)*8 + i] (8 consecutive bytes, memory order) -// and du_mma_sync passes reinterpret(x) unchanged to v_mmac, so -// load_fragment8 produces byte-identical operand values and the int32 -// accumulation is bit-identical; only the redundant VALU reassembly (and its -// lgkmcnt wait states) is removed from the LDS->MMAC critical path, letting -// the 32-MMAC burst issue back-to-back. Expected PMC: valu_instructions -// 8.66M -> ~4.8-5.6M (-35..-45%), lds_instructions 1.31M -> ~1.05M (-20%, -// 8 A ds_read2_b32 + 4 B ds_read2_b64 -> 8 ds_read2_b64 per stage), -// lds_wait_instructions 3.01M -> ~2.2-2.5M, lds_bank_conflicts -// approximately unchanged (A rows keep the 2-way stride-80 aliasing), -// vmem_read/vmem_write and the 32 v_mmac/stage unchanged. kv_b_proj keeps -// the byte-identical 128x64 kNMajorB=false arm (unchanged lowering) and the -// generic scalar fallback is untouched. +// Iteration 12 (consolidation round, final HIP-only polish on the fixed +// 128x128 packed architecture): the last measured LDS inefficiency is the +// A-staging publish mapping. The exact-source PMC of the accepted +// iteration-10 source (source_hip_digest 9386d564..., arch_vgpr 152, workgroup +// 256, grid 512, lds_bytes 17,408) reports lds_bank_conflicts 1,048,576 with +// lds_instructions 1,048,576, and that conflict total is exactly the A publish: +// four 8-byte staging halves per thread per K stage (two int2 halves for each +// of the two payloads) x 4 LDS phases x 1 extra cycle x 32 stages x 4 waves x +// 512 blocks = 1,048,576, while the B publish (linear tid * 16), both fragment +// ds_reads and the coalesced epilogue are conflict-free. 18 has gcd(18, 32) +// = 2, so the 72-byte A row stride makes the fragment reads conflict-free, but +// the publish mapping (thread tid owns A row tid / 4 with 16-byte group +// (tid % 4) * 16) puts the four rows of one 16-lane phase at 18 r mod 32 for +// r = 0..3 - only four of the sixteen even bank residues - so every 8-byte +// staging half needs two LDS cycles instead of one; no 16-byte-aligned stride +// can remove it (such strides are multiples of four words, so a phase reaches +// at most eight residues). Focused change, 128x128 packed arm only: the A +// publish mapping becomes +// row = 16 * (tid / 64) + 2 * ((tid % 64) / 16) +// + 8 * (((tid / 4) % 4) & 1) + (((tid / 4) % 4) >> 1) +// col = 16 * (tid % 4) +// so the four rows of each 16-lane phase are (2p, 2p + 8, 2p + 1, 2p + 9) and +// their publish word starts 18 r mod 32 are exactly the sixteen even residues +// for both the +0 and the +8 half: each phase touches all 32 LDS banks once +// (one cycle, zero conflicts). Every 4-lane group still publishes one A row's +// 64 contiguous bytes, so the global A read set, its L1 line/sector pattern +// and the 128-byte row reads are identical to the previous tid-linear mapping; +// the A tile layout (72-byte rows), both barrier-separated fragment reads, the +// B staging, the one-K-stage staging lookahead, the coalesced epilogue, the +// k0-outer/kk-inner int32 accumulation order and the 128x64 arms plus the +// generic scalar fallback are untouched, so the int32 accumulation and the +// stored bits are bit-identical. +// +// Iteration 18 (consolidation round, final HIP-only polish on the fixed +// 128x128 packed architecture): the last measured fixed-region waste of the +// arm is the bf16 conversion in the fused epilogue. The exact gfx928 code +// object of the accepted iteration-12 source shows every __float2bfloat16(v) +// inside store_prefill_fragment_coalesced_scaled lowered to the Inf/NaN-guarded +// software sequence of amd_hip_bfloat16.h float_to_bfloat16: v_and +// 0x7f800000, v_cmp_ne, s_and_saveexec_b64, s_xor_b64, v_bfe_u32 (round bit), +// v_add3_u32 (0x7fff), s_andn2_saveexec_b64, v_mov_b32 0, v_or_b32 0x10000, +// v_cmp_eq_u32_sdwa, v_cndmask_b32, s_or_b64 exec - roughly a dozen +// VALU/scalar instructions and two exec-mask branches for each of the 64 +// values a wave converts (16 fragments x 4 elements), and those branches split +// the epilogue into many small basic blocks. That Inf/NaN arm is unreachable +// on this arm: its operand is an int32 dot product (|dot| <= 2048 * 127 * 127 +// < 2^25) times the finite x_scale and weight_scale values of this operator +// family, so every converted value is finite and the library's finite path is +// just the round-to-nearest-even increment u += 0x7fff + ((u >> 16) & 1) +// followed by u >> 16. Focused change, 128x128 packed arm only: the coalesced +// epilogue converts its four accumulators per fragment with a straight-line +// bf16_rne_finite() helper (four integer instructions, no s_mov_b32 +// 0x7f800000 constant, no exec mask, no branch), which stores bit-identical +// bf16 for every finite input (the 8-byte coalesced stores are untouched). The +// K loop, the LDS layouts, both barriers per K step, the one-K-stage staging +// lookahead, the in-loop scale batch, the fragment mapping, the store pattern +// and the 128x64 packed/k-major arms plus the generic scalar fallback (which +// keep the library __float2bfloat16 for their per-element stores) are +// byte-identical. +// +// The timed operator performs no allocation, compilation, autotuning, weight +// packing, host synchronization, device synchronization, or default-stream +// launch; both launchers run on the caller-provided PyTorch HIP stream and are +// CUDA/HIP-Graph safe. #include #include @@ -155,27 +279,50 @@ constexpr int kTileK = 32; // gfx928 INT8 DUMMA primitive: m16n16k32 constexpr int kWaveSize = 64; // gfx928 native wavefront constexpr int kBlockM = 128; constexpr int kBlockN = 64; +// Iteration 4 tile-aspect arm: the exact (k, n) == (2048, 2048) packed pair +// takes a 128-wide N tile (M:N tile aspect 1:1) instead of kBlockN == 64. +constexpr int kBlockN128 = 128; constexpr int kStageK = 64; -constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) -constexpr int kBStride = kBlockN + kBPad; -constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride +// Fragment-major packed B (iteration 5). One DUMMA matrix_b tile +// (m16n16k32 -> 16 n x 32 k int8) is kTileN * kTileK = 512 contiguous bytes: +// byte (n, k) of tile (n/16, k/32) sits at +// (k % 32) / 8 * 128 + (n % 16) * 8 + (k % 8). +// Every 16-byte chunk of a tile is therefore one aligned int4 LDS store / +// global read, and a lane's 8-byte fragment slot inside the tile is at +// (lane & 15) * 8 + (lane >> 4) * 128 +// i.e. word 2 * (lane & 15) plus a 32-word k-group stride, so the 16-lane LDS +// phase (lane & 15 = 0..15) covers all 32 banks exactly once: bank-conflict- +// free with zero padding. +constexpr int kBFragBytes = kTileN * kTileK; // 512 +// Legacy raw [K, N] k-major arm: 64 -> 80-byte row stride (16-byte aligned so +// the cooperative staging stays one 128-bit LDS store). +constexpr int kBStride = kBlockN + 16; +constexpr int kAStride = kStageK + 16; // 64 -> 80-byte A row stride +// Iteration 10 (consolidation round): the 128x128 packed arm uses a 72-byte +// (18-word) A row stride. 72 is 8-byte aligned (so every 8-byte fragment read +// and every split int2 staging store is naturally aligned) but NOT 16-byte +// aligned, and 18*r mod 32 for r = 0..15 is exactly the 16 even bank residues, +// so a 16-lane 8-byte fragment phase touches all 32 LDS banks once instead of +// the 2-way aliasing of the 20-word (80-byte) stride. The 128x64 arm keeps +// the accepted 80-byte stride via kAStride above. +constexpr int kAStride128 = kStageK + 8; // 64 -> 72-byte A row stride constexpr int kBlockThreads = 4 * kWaveSize; -// Exact (k, n) pair whose weight buffer is packed n-major -// (packed[n*K + k] == raw[k*N + n]): glm_tp8_q_b_proj_m4096 only. -constexpr int kPackNMajorK = 2048; -constexpr int kPackNMajorN = 2048; +// Exact (k, n) pair whose packed weight buffer is fragment-major for this +// worker's assigned shape (glm_tp8_q_b_proj_m4096). Every other (k, n) keeps +// the identity device-to-device pack and the raw [K, N] k-major arm. +constexpr int kPackFragMajorK = 2048; +constexpr int kPackFragMajorN = 2048; using bf16_t = hip_bfloat16; // --------------------------------------------------------------------------- // Direct fragment epilogue for one m16n16k32 accumulator fragment. -// Verified gfx928 int8 m16n16k32 accumulator ownership (matches -// du_store_matrix_sync): lane & 15 selects the row, lane >> 4 selects -// col % 4, and x[i] maps to columns col%4 + 4*i. The scale multiply order -// (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are -// identical to the harness reference ((dot.float() * x_scale) * ws.T then -// .to(bfloat16)), so stored bf16 bits match exactly. +// Verified gfx928 int8 m16n16k32 accumulator ownership: lane & 15 selects the +// row, lane >> 4 selects col % 4, and frag.x[i] maps to columns col%4 + 4*i. +// The scale multiply order (float(dot) * x_scale[row] * weight_scale[col]) and +// the bf16 rounding match the harness reference +// ((dot.float() * x_scale) * ws.T).to(bfloat16), so the stored bits match. // --------------------------------------------------------------------------- template __device__ __forceinline__ void store_prefill_fragment( @@ -202,28 +349,48 @@ __device__ __forceinline__ void store_prefill_fragment( } } +// --------------------------------------------------------------------------- +// Iteration 18 (consolidation round): straight-line round-to-nearest-even +// float -> bfloat16 for the coalesced epilogue. This is exactly the finite +// path of amd_hip_bfloat16.h float_to_bfloat16 +// u += 0x7fff + ((u >> 16) & 1); return u >> 16; +// with the Inf/NaN special case omitted; the 128x128 packed q_b_proj arm only +// ever converts finite values (int32 dot x finite scales), so the bf16 bits +// are identical to __float2bfloat16 while the exec-masked compare/branch pair +// and the s_mov_b32 0x7f800000 / v_or_b32 0x10000 fixups leave the epilogue. +// The per-element 128x64 arm and the generic scalar fallback keep the library +// conversion byte-identically. +// --------------------------------------------------------------------------- +__device__ __forceinline__ unsigned short bf16_rne_finite(float v) { + unsigned int u = __float_as_uint(v); + u += 0x7fffu + ((u >> 16) & 1u); // round to nearest, ties to even + return static_cast(u >> 16); +} + // --------------------------------------------------------------------------- // Iteration 6 (epilogue round): register-batched COALESCED direct-fragment -// epilogue for the exact (2048, 2048) 128x128 arm. The per-element variant -// (store_prefill_fragment) above issues one 2-byte bf16 store per lane per -// element: 64 scattered stores per wave per fragment set, each touching 16 -// rows x 4 columns, i.e. every 32-B store sector only 25% utilized. This -// epilogue transposes the 4-element groups within each 4-lane column group -// (lanes r, r+16, r+32, r+48 - a 4x4 transpose, two 2x2 steps with -// shfl_xor 16 then 32; the lineage-validated pattern accepted on the sibling -// TP8 workers 0/2 on this DTK, which lowers __shfl_xor to ds_bpermute at the -// block tail where the LDS pipe is idle), so lane (r, c4) ends up holding -// the four CONTIGUOUS columns 4*c4 .. 4*c4+3, converts them to bf16, packs -// 4 bf16 (8 B) and writes ONE 8-byte store per lane (100% store sector -// efficiency). Only the int32 values are re-routed between lanes; the -// per-element scale multiply order (float(dot) * x_scale[row] * -// weight_scale[col]) and the __float2bfloat16 rounding are unchanged, so -// the stored bits are identical to the per-element store. xs/ws come from -// caller registers (batched per wave on the last K stage), so the epilogue -// issues no vmem loads. The row >= m guard is wavefront-uniform (all 64 -// lanes of a wave share the same 16-row window), so the shuffles never mix -// active and inactive lanes; base_col is a multiple of 64 and n*2 a -// multiple of 8, so the 8-byte store is aligned. +// epilogue for the exact (k, n) == (2048, 2048) q_b_proj 128x128 arm. The +// per-element variant above issues one 2-byte bf16 store per lane per +// element: 64 scattered stores per wave per fragment set, each store touching +// 16 rows x 4 columns, so every 32-byte store sector is only 25% utilized +// (PMC: 131,072 vmem_write_instructions = 512 blocks x 4 waves x 64). +// +// This epilogue transposes the four int32 values inside each 4-lane column +// group (lanes r, r+16, r+32, r+48 - a 4x4 register transpose, two 2x2 steps +// with __shfl_xor 16 then 32, which lowers to ds_bpermute at the block tail +// where the LDS pipe is idle), so lane (r, c4) ends up owning the four +// CONTIGUOUS columns base_col + 4*c4 .. +3. It converts those to bf16, packs +// four bf16 into one 8-byte value and writes ONE 8-byte store per lane (100% +// store sector efficiency; the previous 64 stores become 16 per lane). +// +// Only the int32 accumulator values are re-routed between lanes; the +// per-element multiply order (float(dot) * x_scale[row] * weight_scale[col]) +// and the __float2bfloat16 rounding are unchanged, so the stored bits are +// identical. xs and ws come from caller registers batched on the last K +// stage, so the epilogue issues no interleaved vmem loads. base_col is a +// multiple of 16 and n*2 is a multiple of 8, so the 8-byte store is aligned; +// the 128x128 arm is only dispatched for m % 128 == 0, so the row >= m guard +// is always false and the shuffles never mix active and inactive lanes. // --------------------------------------------------------------------------- template __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( @@ -238,9 +405,9 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( int lane) { const int row = base_row + (lane & 15); if (row >= m) { - return; // tail-M masking: padded rows never write (wavefront-uniform) + return; // tail-M masking: never fires for m % 128 == 0 } - const int c4 = lane >> 4; // 0..3 + const int c4 = lane >> 4; // 0..3 (column group inside the m16n16 tile) const int x0 = frag.x[0]; const int x1 = frag.x[1]; const int x2 = frag.x[2]; @@ -266,77 +433,98 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced_scaled( const int f2 = hi ? t0 : a2; const int f3 = hi ? t1 : a3; - // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 (8 B, 8-byte - // aligned: base_col is a multiple of 16, n*2 is a multiple of 8); ws is + // Lane (r, c4) now owns columns base_col + 4*c4 .. +3 of row `row`; ws is // the float4 at weight_scale + base_col + 4*c4 preloaded by the caller. const int col0 = base_col + 4 * c4; const float v0 = static_cast(f0) * xs * ws.x; const float v1 = static_cast(f1) * xs * ws.y; const float v2 = static_cast(f2) * xs * ws.z; const float v3 = static_cast(f3) * xs * ws.w; + // Iteration 18: straight-line RNE conversion (finite path), bit-identical + // to the library __float2bfloat16 for this arm's finite values. const uint64_t packed = - static_cast( - static_cast(__float2bfloat16(v0))) | - (static_cast( - static_cast(__float2bfloat16(v1))) - << 16) | - (static_cast( - static_cast(__float2bfloat16(v2))) - << 32) | - (static_cast( - static_cast(__float2bfloat16(v3))) - << 48); + static_cast(bf16_rne_finite(v0)) | + (static_cast(bf16_rne_finite(v1)) << 16) | + (static_cast(bf16_rne_finite(v2)) << 32) | + (static_cast(bf16_rne_finite(v3)) << 48); *reinterpret_cast(out + static_cast(row) * n + col0) = packed; } // --------------------------------------------------------------------------- -// Direct 8-byte LDS fragment load (lineage-validated on gfx928/DTK 26.04 for -// m16n16k32 int8 fragments). du_load_matrix_sync's int8 loaders assign -// x[0..7] = 8 consecutive bytes at (lane & 15) * ldm + ((lane >> 4) << 3) for -// matrix_b col_major, but the compiler lowers that to per-byte ds_read_u8 + -// mask/OR reassembly VALU. Writing the same 8 bytes directly into the -// fragment storage keeps the operand bit pattern identical (exact int32 -// accumulation unchanged) and lets the compiler emit one ds_read2_b64 per -// fragment straight into the v_mmac operand. +// Direct 8-byte staged-fragment load. du_mma.h defines both the matrix_a +// row_major and the matrix_b col_major int8 loaders as +// x[i] = p[(lane & 15) * ldm + ((lane >> 4) << 3) + i], +// i.e. eight CONTIGUOUS bytes per lane, but the generic slot-by-slot loader +// lowers to per-byte LDS byte gathers plus a mask/shift/OR reassembly chain. +// Writing those same 8 bytes with one 64-bit LDS read keeps the fragment +// storage bit-identical (the int32 accumulation is unchanged) and shortens the +// LDS -> MMAC critical path. +// +// Iteration 5 form: `block` is the base of one 512-byte fragment-major DUMMA +// matrix_b tile, whose layout puts the lane's 8 bytes at +// (lane & 15) * 8 + (lane >> 4) * 128 +// (word 2 * (lane & 15) plus a 32-word k-group stride). A 16-lane LDS phase +// therefore touches each of the 32 banks exactly once, and the access is a +// single aligned 64-bit read. Operand bytes are byte-for-byte identical to +// the n-major layout, so correctness is unaffected. // --------------------------------------------------------------------------- -__device__ __forceinline__ void load_fragment8( +__device__ __forceinline__ void load_fragment8_block( DUFragmentBase& frag, - const signed char* __restrict__ base, - int stride, + const signed char* __restrict__ block, int lane) { - const int off = (lane & 15) * stride + ((lane >> 4) << 3); + const int off = (lane & 15) * 8 + ((lane >> 4) << 7); *reinterpret_cast(&frag.x[0]) = - *reinterpret_cast(base + off); + *reinterpret_cast(block + off); } // --------------------------------------------------------------------------- -// Iteration 5 (packing round): 8-byte fragment load for the swizzled plane -// LDS layout b_tile[kc][n][8] (kc = 8-byte k sub-chunk, n row stride 8 B, -// plane stride 1024 B). Lane l reads the same 8 consecutive k bytes as -// load_fragment8 (n row = lane&15, k chunk = lane>>4) but from -// (lane >> 4) * 1024 + (lane & 15) * 8 -// relative to the fragment origin, so the 64 lanes hit 16 distinct bank -// pairs per 128-B phase (all 32 banks once) -> zero bank conflicts, the -// 4-cycle minimum for a 512-B fragment. +// Iteration 9 (A-fragment direct-load round): direct 8-byte staged A-fragment +// load, the A-side twin of load_fragment8_block above. du_mma.h defines the +// matrix_a row_major int8 loader with the same +// x[i] = p[(lane & 15) * ldm + ((lane >> 4) << 3) + i] +// mapping as the matrix_b col_major loader, i.e. eight contiguous memory-order +// bytes per lane: the identical mapping the load_fragment8_block comment above +// records for both loaders. The library wrapper passes base = the fragment +// origin and ldm = the LDS row stride, so the direct read below (base + off, +// off = (lane & 15) * stride + ((lane >> 4) << 3)) addresses exactly the same +// bytes. It fills the same x[0..7] storage with one 64-bit LDS read, letting +// the compiler feed v_mmac directly instead of materializing the value through +// a pair of ds_read2_b32 plus the byte reassembly chain; the operand bytes and +// therefore the int32 accumulation are bit-identical. `base + off` is 8-byte +// aligned for the 128x128 arm: a_tile is 16-byte aligned, the row stride is +// 72 (kAStride128, iteration 10) or 80 (kAStride, the 128x64 arm), local_row +// is a multiple of 64, kk is 0 or 32 and every term is a multiple of 8. // --------------------------------------------------------------------------- -__device__ __forceinline__ void load_fragment8_plane( +__device__ __forceinline__ void load_fragment8( DUFragmentBase& frag, const signed char* __restrict__ base, + int stride, int lane) { - const int off = ((lane >> 4) << 10) + ((lane & 15) << 3); + const int off = (lane & 15) * stride + ((lane >> 4) << 3); *reinterpret_cast(&frag.x[0]) = *reinterpret_cast(base + off); } +// Byte offset of the 512-byte fragment-major DUMMA B tile for output-column +// fragment n_frag (= n / 16) and k fragment k_frag (= k / 32). Used by both +// the cooperative staging (global read) and by the scalar fallback decode. +__device__ __forceinline__ int64_t b_frag_block_offset(int n_frag, + int k_frag, + int k) { + return (static_cast(n_frag) * (k / kTileK) + k_frag) * + kBFragBytes; +} + // --------------------------------------------------------------------------- -// Compile-time B-fragment layout selector. row_major / col_major are tag -// types in du::dumma, so the layout template argument must be selected as a -// type (a conditional expression over type names is not a valid template -// argument). kNMajorB == true (n-major packed B, exact (2048, 2048)) -> -// col_major fragments (8 contiguous k bytes per lane, load_fragment8); -// kNMajorB == false (raw [K, N] row-major B) -> row_major fragments (the -// accepted-best loader). +// Compile-time B staged-layout selector. row_major / col_major are tag types +// in du::dumma, so the layout template argument has to be selected as a type. +// kNMajorB == true (fragment-major packed B, exact (k, n) == (2048, 2048)): +// the LDS tile is a contiguous array of 512-byte DUMMA B tiles, each +// staged with one aligned int4 read/store and each read with the +// bank-conflict-free load_fragment8_block. +// kNMajorB == false (raw [K, N] row-major B): k-major tile with the accepted +// library loader (unchanged). // --------------------------------------------------------------------------- template struct b_frag_layout { @@ -348,22 +536,27 @@ struct b_frag_layout { }; // --------------------------------------------------------------------------- -// Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. -// Each wave owns a 64x32 quadrant (eight m16n16k32 int32 accumulators); the -// block cooperatively stages A[128,64] from x_q (row-major, stride k) and -// B[64,64] into a single-buffered LDS stage. Two barriers per stage: one -// after the cooperative load, one before the next stage overwrites LDS. -// A fragments use the library du_load_matrix_sync row-major loader. B is -// templated on its staged layout: -// * kNMajorB == true: B is packed n-major (packed[n*K + k]) and staged -// n-major (n row, 80 B stride, k contiguous); each B fragment is 8 -// contiguous k bytes per lane, loaded by load_fragment8 (one 8-byte LDS -// read per fragment), eliminating the ds_read_u8 byte-gather + VALU -// reassembly of the accepted-best kernel. -// * kNMajorB == false: B stays raw [K, N] row-major, staged k-major with -// row-major library fragments (accepted-best path, unchanged). -// Dispatch guarantees m % 128 == 0, n % 64 == 0, k % 64 == 0, so every -// global load/store is in-bounds and 16-byte aligned. +// Large-M prefill: one 128x64 output tile per block, four wavefronts of 64 +// lanes. Each wave owns a 64x32 quadrant (eight m16n16k32 int32 +// accumulators); the block cooperatively stages A[128,64] from x_q (row-major, +// stride k) and B[64,64] into a single-buffered LDS stage. Two barriers per +// stage: one after the cooperative load, one before the next stage overwrites +// LDS. Dispatch guarantees m % 128 == 0, n % 64 == 0, k % 64 == 0, so every +// global load / store is in-bounds and 16-byte aligned. +// +// B is templated on its staged layout: +// * kNMajorB == true (exact (k, n) == (2048, 2048)): B is the +// fragment-major packed weight; each thread stages one aligned 16-byte +// chunk of the 8 contiguous 512-byte DUMMA B tiles (4 n fragments x 2 k +// fragments) with one coalesced int4 global read and one aligned int4 LDS +// store, and each col_major fragment is one bank-conflict-free 8-byte LDS +// read per lane (load_fragment8_block). +// * kNMajorB == false: B is the raw [K, N] identity layout staged k-major +// with the library row_major loader (unchanged arm). +// +// LDS budget: A[128, 80] = 10,240 B + B (8 * 512 = 4,096 B on the packed arm, +// B[64,80] = 5,120 B on the k-major arm) <= 15,360 B/block, well inside the +// 64 KiB per-CU LDS capacity (4 blocks/CU by LDS alone). // --------------------------------------------------------------------------- template __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kernel( @@ -375,6 +568,9 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne int m, int n, int k) { + constexpr int kBBlocksB = + kNMajorB ? (kBlockN / kTileN) * (kStageK / kTileK) : kStageK; + constexpr int kBStrideB = kNMajorB ? kBFragBytes : kBStride; const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; const int lane = tid % kWaveSize; @@ -385,11 +581,8 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne const int local_row = wave_row * 64; const int local_col = wave_col * 32; - // Single-buffered stage: A[128, 80] + B[64, 80] = 15,360 B/block - // (4 blocks/CU fit the 64 KiB LDS budget; the padded 80-byte strides are - // 16-byte-aligned and break the 64-byte LDS bank periodicity). __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; - __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + __shared__ __align__(16) int8_t b_tile[kBBlocksB * kBStrideB]; DUFragment a_frag0, a_frag1; @@ -407,21 +600,14 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne du_fill_fragment(acc30, 0); du_fill_fragment(acc31, 0); - // Cooperative staging: A[128,64] is 512 int4s (two per thread), B[64,64] - // is 256 int4s (one per thread). - // A: thread tid owns row (tid*16)/64 and 16-byte column group - // (tid*16)%64 for rows [0,64) and [64,128). - // B (kNMajorB == false): raw [K, N] row-major; thread tid owns K row - // tid>>2 and the 16-byte N-column group (tid&3)*16 (16 consecutive N - // values at a fixed K row -> one aligned int4 in global and in LDS). - // B (kNMajorB == true): packed n-major; thread tid owns N row tid>>2 and - // the 16-byte K-column group (tid&3)*16 (16 consecutive K values at a - // fixed N row -> one aligned int4 in global and in the n-major LDS). + // Cooperative staging: A[128,64] is 512 int4s (two per thread, rows + // [0,64) and [64,128)); B[64,64] is 256 int4s (one per thread). Thread tid + // owns A row tid>>2 and 16-byte column group (tid & 3) * 16, and B row + // tid>>2 with 16-byte group (tid & 3) * 16 (k-major K row for the identity + // arm, packed fragment-major B tile for the packed arm). const int vector_byte_offset = tid * static_cast(sizeof(int4)); const int stage_row = vector_byte_offset / kStageK; // 0..63 const int stage_col = vector_byte_offset - stage_row * kStageK; - const int b_k = tid >> 2; // K row (false) / N row (true) - const int b_nc = (tid & 3) * 16; // N group (false) / K group (true) for (int k0 = 0; k0 < k; k0 += kStageK) { *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = @@ -433,39 +619,51 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + k0 + stage_col); if constexpr (kNMajorB) { - // n-major packed B: 16 consecutive k bytes of the (n0 + b_k) row. - *reinterpret_cast(b_tile + b_k * kBStride + b_nc) = + // Fragment-major packed B: thread tid stages 16-byte chunk tid of the + // 8 contiguous 512-byte DUMMA B tiles (4 n fragments x 2 k fragments): + // one aligned int4 global read and one aligned int4 LDS store at LDS + // byte offset tid * 16 (contiguous across the wave). + const int b_blk = tid >> 5; // 0..7 packed B tile index + const int b_slot = tid & 31; // 0..31 16-byte chunk inside the tile + *reinterpret_cast(b_tile + b_blk * kBFragBytes + b_slot * 16) = *reinterpret_cast( - weight + static_cast(n0 + b_k) * k + k0 + b_nc); + weight + b_frag_block_offset(n0 / kTileN + (b_blk >> 1), + k0 / kTileK + (b_blk & 1), k) + + b_slot * 16); } else { - // raw [K, N] row-major B: 16 consecutive n bytes of the (k0 + b_k) row. - *reinterpret_cast(b_tile + b_k * kBStride + b_nc) = + // Raw [K, N] row-major B: 16 consecutive n bytes of the (k0 + b_n) row. + const int b_n = tid >> 2; // B row (K) + const int b_nc = (tid & 3) * 16; // 16-byte group inside that row + *reinterpret_cast(b_tile + b_n * kBStrideB + b_nc) = *reinterpret_cast( - weight + static_cast(k0 + b_k) * n + n0 + b_nc); + weight + static_cast(k0 + b_n) * n + n0 + b_nc); } __syncthreads(); // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. - // Accumulation order: k0-outer over 64-K stages, kk-inner (kk=0 then - // kk=32), matching the reference int32 accumulation. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk = 0 then + // kk = 32), matching the reference int32 accumulation. #pragma unroll for (int kk = 0; kk < kStageK; kk += kTileK) { if constexpr (kNMajorB) { - // n-major LDS tile: n rows stride kBStride, k contiguous within a - // row; col_major fragment slots are 8 consecutive k bytes at - // (lane & 15) * kBStride + ((lane >> 4) << 3) relative to the - // fragment origin -> one ds_read2_b64 per fragment per lane. - load_fragment8(b_frag0, b_tile + local_col * kBStride + kk, kBStride, - lane); - load_fragment8(b_frag1, - b_tile + (local_col + kTileN) * kBStride + kk, - kBStride, lane); + // Fragment-major LDS tile: the wave's two B fragments are the two + // contiguous 512-byte DUMMA B tiles (local_col/16 + i, kk/32), each + // read with one aligned bank-conflict-free 8-byte LDS read per lane. + load_fragment8_block( + b_frag0, + b_tile + ((local_col / kTileN) * 2 + (kk / kTileK)) * kBFragBytes, + lane); + load_fragment8_block( + b_frag1, + b_tile + (((local_col / kTileN) + 1) * 2 + (kk / kTileK)) * + kBFragBytes, + lane); } else { - du_load_matrix_sync(b_frag0, b_tile + kk * kBStride + local_col, - kBStride); + du_load_matrix_sync(b_frag0, b_tile + kk * kBStrideB + local_col, + kBStrideB); du_load_matrix_sync(b_frag1, - b_tile + kk * kBStride + local_col + kTileN, - kBStride); + b_tile + kk * kBStrideB + local_col + kTileN, + kBStrideB); } du_load_matrix_sync(a_frag0, a_tile + local_row * kAStride + kk, kAStride); @@ -505,37 +703,42 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x64_kerne store_prefill_fragment(acc20, x_scale, weight_scale, out, base_row + 2 * kTileM, base_col, m, n, lane); store_prefill_fragment(acc21, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col + kTileN, m, n, - lane); + base_row + 2 * kTileM, base_col + kTileN, m, n, lane); store_prefill_fragment(acc30, x_scale, weight_scale, out, base_row + 3 * kTileM, base_col, m, n, lane); store_prefill_fragment(acc31, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col + kTileN, m, n, - lane); + base_row + 3 * kTileM, base_col + kTileN, m, n, lane); } // --------------------------------------------------------------------------- -// Iteration 4 (tile-aspect round) + iteration 5 (packing round): 128x128 -// tile for the exact q_b_proj (k, n) == (2048, 2048) swizzled-pack B arm. -// Four wavefronts of 64 lanes; each wave owns a 64x64 quadrant (four A x -// four B m16n16k32 fragments = 16 MMACs per kk, 16 int32 accumulator -// fragments). The block cooperatively stages A[128,64] (2 int4/thread, -// 80-byte-strided rows) and B[128,64] from the iteration-5 swizzled pack (2 -// int4/thread, 16-byte plane chunks) into a single-buffered LDS stage: -// A[128,80] (10,240 B) + B plane tile [8][128][8] (8,192 B) = 18,432 -// B/block (2 blocks/CU = 36,864 B <= 64 KiB); two __syncthreads per stage. -// A fragments use the library du_load_matrix_sync row-major loader -// (4 x ds_read2_b32 per kk per wave), B fragments use load_fragment8_plane -// (lane-linear 8-byte reads, zero LDS bank conflicts). A-side global -// re-read halves (A reuse 32 -> 16, 268 -> 134 MiB) while B reuse stays 32 -// (134 MiB): per-byte A:B tile traffic balances 1:1, total 402 -> 268 MiB -// (-33%). Grid dim3(N/128, M/128) = 512 blocks (~4.3/CU). Dispatch -// guarantees m % 128 == 0, n == 2048, k % 64 == 0, so every global -// load/store is in-bounds and 16-byte aligned. Accumulation order stays -// k0-outer / kk-inner with the same element-to-slot fragment mapping, so the -// int32 accumulation is bit-identical to the accepted kernel. +// Iteration 4 (tile-aspect round): 128x128 macro-tile for the exact +// (k, n) == (2048, 2048) q_b_proj arm (fragment-major packed weight), four +// wavefronts of 64 lanes, each wave owning a 64x64 quadrant built from four A +// x four B m16n16k32 fragments (sixteen int32 accumulators). Widening the +// N-tile 64 -> 128 makes the M:N tile aspect 1:1: the A-side global re-read +// halves (A reuse 32 -> 16, ~268 -> ~134 MiB) while B reuse stays 32 +// (~134 MiB), so total staged tile traffic drops ~402 -> ~268 MiB (-33%), and +// each wave issues 8 LDS fragment reads per 16 MMACs instead of 6 per 8. +// +// The block stages A[128,64] (two int4/thread, rows tid>>2 and tid>>2 + 64; +// iteration 10 makes the A row stride 72 B = kAStride128, so each 16-byte +// chunk is published with two aligned int2 LDS stores) and the fragment-major +// packed B[128,64] (16 contiguous 512-byte DUMMA B tiles: two 16-byte +// chunks/thread, one aligned int4 global read and one aligned int4 LDS store +// each) into one single-buffered LDS stage: A[128,72] (9,216 B) + B +// (16 * 512 = 8,192 B) = 17,408 B/block, well inside the 64 KiB LDS budget. +// A fragments use the iteration-9 direct 8-byte load_fragment8 (the 72-byte +// stride makes each 16-lane 8-byte phase touch all 32 banks exactly once, i.e. +// zero A-side bank conflicts); B fragments use the bank-conflict-free +// fragment-major load_fragment8_block. Two __syncthreads per K stage. +// Dispatch guarantees m % 128 == 0, n % 128 == 0 and k % 64 == 0 for this +// arm, so every global load/store is in-bounds and 16-byte aligned. +// Accumulation order (k0-outer over 64-K stages, kk-inner 0 then 32) and the +// fragment element-to-slot mapping are identical to the 128x64 arm, so the +// int32 accumulation is bit-identical. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kernel( +__global__ __launch_bounds__(kBlockThreads) void +w8a8_dumma_prefill_128x128_nmajor_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, const float* __restrict__ x_scale, @@ -544,43 +747,32 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern int m, int n, int k) { - constexpr int kBlockM128 = 128; - constexpr int kBlockN128 = 128; - constexpr int kTileStride = kStageK + kBPad; // 64 -> 80-byte row stride - constexpr int kTileBytes = kBlockM128 * kTileStride; // 10,240 B per tile - // Iteration 5 (packing round): B tile is the swizzled plane layout - // [kc][n][8] (8 planes of 128 n rows x 8 B, plane stride 1024 B) so each - // lane's 8-byte fragment chunk is lane-linear and bank-conflict-free. - constexpr int kBTileBytes = 8 * kBlockN128 * 8; // 8,192 B/block - constexpr int kBTilePlane = 1024; // 128 n rows x 8 B - + constexpr int kBBlocksB = + (kBlockN128 / kTileN) * (kStageK / kTileK); // 16 packed B tiles const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; const int lane = tid % kWaveSize; - const int wave_row = wave >> 1; - const int wave_col = wave & 1; - const int m0 = static_cast(blockIdx.y) * kBlockM128; + const int wave_row = wave >> 1; // 0..1 + const int wave_col = wave & 1; // 0..1 + const int m0 = static_cast(blockIdx.y) * kBlockM; const int n0 = static_cast(blockIdx.x) * kBlockN128; - const int local_row = wave_row * (kBlockM128 / 2); // 0 or 64 + const int local_row = wave_row * (kBlockM / 2); // 0 or 64 const int local_col = wave_col * (kBlockN128 / 2); // 0 or 64 const int base_row = m0 + local_row; const int base_col = n0 + local_col; - // Iteration 6 (epilogue round): per-lane register-batched scales. Each - // lane owns 4 rows (base_row + 16*i + (lane & 15), i = 0..3) and 4 + // Iteration 6 (epilogue round): per-lane register-batched scales. Each lane + // owns 4 rows (base_row + 16*i + (lane & 15), i = 0..3) and 4 // weight_scale float4s (base_col + 16*j + 4*(lane >> 4), j = 0..3) across - // its 16 fragments; they are loaded exactly once on the last K stage - // before the final protective barrier, so the coalesced epilogue below is - // pure compute + 16 eight-byte stores with no interleaved vmem loads. + // its 16 accumulator fragments; they are loaded exactly once on the last K + // stage before the final protective __syncthreads, so the coalesced + // epilogue below is pure compute + 16 eight-byte stores with no + // interleaved vmem loads. float xs_m[4]; float4 ws_m[4]; - // Single-buffered stage: A[128,80] (10,240 B) + B plane tile [8][128][8] - // (8,192 B) = 18,432 B/block (2 blocks/CU = 36,864 B <= 64 KiB; the A rows - // keep the padded 80-byte stride, the B tile is plane-swizzled so fragment - // reads are lane-linear with zero bank conflicts). - __shared__ __align__(16) int8_t a_tile[kTileBytes]; - __shared__ __align__(16) int8_t b_tile[kBTileBytes]; + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride128]; + __shared__ __align__(16) int8_t b_tile[kBBlocksB * kBFragBytes]; DUFragment a_frag0, a_frag1, a_frag2, a_frag3; @@ -606,93 +798,157 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern du_fill_fragment(acc32, 0); du_fill_fragment(acc33, 0); - // Cooperative staging: A[128,64] is 512 int4s (two per thread) and - // B[128,64] from the iteration-5 swizzled pack is 512 int4s (two per - // thread). Thread tid owns linear int4 slots tid and tid + 256; for A a - // slot maps to row linear >> 2 and the 16-byte column group - // (linear & 3) * 16 (an M row of x_q, one aligned int4 in global and in - // LDS); for B a slot maps to plane linear >> 6 and n-pair linear & 63 - // (two n rows of one 8-byte k sub-chunk, one aligned int4 in the pack and - // in the plane LDS tile). + // Cooperative staging: A[128,64] is 512 int4s (two per thread: M rows + // stage_row and stage_row + 64, 16-byte k group stage_col) and the packed + // B[128,64] is 16 contiguous 512-byte DUMMA B tiles = 512 16-byte chunks + // (two per thread: tile b_blk = tid>>5 and b_blk + 8, 16-byte slot + // tid & 31). Each chunk is one aligned int4 global read; the B chunks are + // also one aligned int4 LDS store, while the A chunks (72-byte strided rows, + // iteration 10) are two aligned int2 LDS stores at +0/+8. Iteration 12 + // remaps which lane publishes which A chunk (stage_row / stage_col below) so + // that every 16-lane publish phase covers all 32 LDS banks exactly once while + // each 4-lane group still keeps one A row's 64 contiguous bytes. + // + // Iteration 7 (compute-pipeline round): the four chunks are loaded ONE K + // STAGE AHEAD. The prologue loads K stage 0 into the payload registers; in + // the loop the payload of stage s is stored to LDS first, then (for + // k0 + kStageK < k) stage s + 1 is loaded into the same registers BEFORE the + // fragment ds_reads and the 32 v_mmac of stage s, so the global-load latency + // is covered by the stage-s compute burst plus the trailing WAR barrier + // rather than being drained by a just-issued `s_waitcnt vmcnt(N)` at the + // stage top. The payload stays four int4 (16 VGPRs); the LDS stage, both + // barriers per K step, every staged byte and every fragment address are + // unchanged. + // Iteration 12 (consolidation round): A-staging lane mapping. The four rows + // of a 16-lane LDS phase are (2p, 2p + 8, 2p + 1, 2p + 9), so the publish + // word starts 18 * row mod 32 of the phase are exactly the sixteen even bank + // residues (4p + {0, 2, 4, ..., 30}): each 8-byte staging half touches all + // 32 LDS banks once instead of twice. + const int a_phase = lane >> 4; // 0..3 16-lane LDS phase + const int a_slot = (lane >> 2) & 3; // 0..3 4-lane row group + const int stage_row = wave * 16 + 2 * a_phase + 8 * (a_slot & 1) + + (a_slot >> 1); // 0..63 + const int stage_col = (lane & 3) << 4; // 0,16,32,48 + const int b_blk = tid >> 5; // 0..7 packed B tile index + const int b_slot = tid & 31; // 0..31 16-byte chunk inside the tile + + const int64_t a_row0 = static_cast(m0 + stage_row) * k; + const int64_t a_row1 = + static_cast(m0 + stage_row + kBlockM / 2) * k; + const int64_t b_step = + 4 * (static_cast(k / kTileK) * kBFragBytes); // +4 n fragments + int4 a_payload0 = make_int4(0, 0, 0, 0); + int4 a_payload1 = make_int4(0, 0, 0, 0); + int4 b_payload0 = make_int4(0, 0, 0, 0); + int4 b_payload1 = make_int4(0, 0, 0, 0); + if (k > 0) { + a_payload0 = *reinterpret_cast(x_q + a_row0 + stage_col); + a_payload1 = *reinterpret_cast(x_q + a_row1 + stage_col); + const int64_t b_off0 = + b_frag_block_offset(n0 / kTileN + (b_blk >> 1), (b_blk & 1), k); + b_payload0 = + *reinterpret_cast(weight + b_off0 + b_slot * 16); + b_payload1 = + *reinterpret_cast(weight + b_off0 + b_step + b_slot * 16); + } + for (int k0 = 0; k0 < k; k0 += kStageK) { - const int kstage = k0 >> 6; // 64-k group index inside the swizzled pack -#pragma unroll - for (int i = 0; i < 2; ++i) { - const int linear = tid + i * kBlockThreads; - const int row = linear >> 2; // 0..127 - const int col = (linear & 3) * 16; // 0/16/32/48 - *reinterpret_cast(a_tile + row * kTileStride + col) = - *reinterpret_cast( - x_q + static_cast(m0 + row) * k + k0 + col); - // Swizzled pack: slot linear -> (kc = linear >> 6, j = linear & 63); - // the 16-byte chunk is two n rows (n0 + 2j, n0 + 2j + 1) of the 8-byte - // k sub-chunk kc of stage kstage: one aligned int4 in global (two - // consecutive n rows of one plane) and one aligned int4 in the plane - // LDS tile (plane kc, rows 2j..2j+1 -> contiguous 16 B). - const int b_kc = linear >> 6; // 0..7 - const int b_j = linear & 63; // 0..63 (n-pair within the plane) - *reinterpret_cast(b_tile + b_kc * kBTilePlane + b_j * 16) = - *reinterpret_cast( - weight + (static_cast(kstage * 8 + b_kc) * n + - (n0 + 2 * b_j)) * 8); - } + // Publish the payload prefetched one K stage ago (prologue for k0 == 0). + // Iteration 10: the 128x128 arm's A rows are 72-byte strided, so each + // 16-byte chunk is published with two 8-byte (`int2`) LDS stores at +0 and + // +8 (row*72 + stage_col is 8-byte aligned for every row); the staged + // bytes keep their exact (row, k) slots. + int8_t* const a_dst0 = a_tile + stage_row * kAStride128 + stage_col; + int8_t* const a_dst1 = + a_tile + (stage_row + kBlockM / 2) * kAStride128 + stage_col; + *reinterpret_cast(a_dst0) = make_int2(a_payload0.x, a_payload0.y); + *reinterpret_cast(a_dst0 + 8) = + make_int2(a_payload0.z, a_payload0.w); + *reinterpret_cast(a_dst1) = make_int2(a_payload1.x, a_payload1.y); + *reinterpret_cast(a_dst1 + 8) = + make_int2(a_payload1.z, a_payload1.w); + *reinterpret_cast(b_tile + b_blk * kBFragBytes + b_slot * 16) = + b_payload0; + *reinterpret_cast(b_tile + (b_blk + 8) * kBFragBytes + + b_slot * 16) = b_payload1; __syncthreads(); + if (k0 + kStageK < k) { + // Prefetch K stage s + 1 before consuming stage s: four aligned int4 + // global reads whose latency overlaps the fragment reads and MMACs. + const int k1 = k0 + kStageK; + a_payload0 = + *reinterpret_cast(x_q + a_row0 + k1 + stage_col); + a_payload1 = + *reinterpret_cast(x_q + a_row1 + k1 + stage_col); + // Tiles 0..7 cover n fragments 0..3 (k fragment = b_blk & 1); tiles + // 8..15 cover n fragments 4..7 with the same two k fragments. + const int64_t b_off0 = + b_frag_block_offset(n0 / kTileN + (b_blk >> 1), + k1 / kTileK + (b_blk & 1), k); + b_payload0 = + *reinterpret_cast(weight + b_off0 + b_slot * 16); + b_payload1 = + *reinterpret_cast(weight + b_off0 + b_step + + b_slot * 16); + } + // Each wave consumes its 64x64 quadrant: sixteen m16n16k32 MMACs per kk. - // Accumulation order: k0-outer over 64-K stages, kk-inner (kk=0 then - // kk=32), matching the reference int32 accumulation. All eight fragment - // loads are issued before the sixteen MMACs (load-all-then-MMAC-all). - // - // Iteration 7 (compute-pipeline round): the A fragments are loaded with - // the same direct 8-byte LDS read as B (load_fragment8) instead of the - // library du_load_matrix_sync row_major loader. The library's int8 - // matrix_a row_major loader assigns x[i] = p[(lane&15)*ldm + - // (lane>>4)*8 + i] (8 consecutive bytes, memory order) and du_mma_sync - // feeds reinterpret(x) straight into v_mmac, but on this DTK the - // loader lowers to 8 x ds_read2_b32 + a redundant per-dword byte - // reassembly chain (~7 VALU: v_and 0xff00/0xff0000/0xff000000 + - // v_or_b32_sdwa + v_or3 per second dword) sitting between the LDS read - // and the MMAC issue (exact code object, both the 128x128 and the 128x64 - // symbols). load_fragment8 fills the same x[0..7] with one 64-bit - // little-endian write, so the operand bytes are bit-identical and the - // compiler emits one ds_read2_b64 straight into the v_mmac operand (the - // same lineage as the B side since iteration 3/5): the LDS->MMAC - // critical path shortens by ~50 VALU + their lgkmcnt wait states per - // stage, and the 32-MMAC burst can issue back-to-back after the barrier. + // Accumulation order: k0-outer over 64-K stages, kk-inner (kk = 0 then + // kk = 32), matching the reference int32 accumulation. #pragma unroll for (int kk = 0; kk < kStageK; kk += kTileK) { - load_fragment8(a_frag0, a_tile + local_row * kTileStride + kk, - kTileStride, lane); + // Iteration 9 (A-fragment direct-load round): the four A fragments are + // read with the direct 8-byte load_fragment8 instead of the library + // du_load_matrix_sync row_major loader. The loader has the identical + // x[i] = p[(lane & 15) * stride + ((lane >> 4) << 3) + i] mapping + // (eight contiguous memory-order bytes per lane), but on this gfx928 + // code object it lowers to a pair of ds_read2_b32 plus a byte + // reassembly chain between the LDS read and the dependent v_mmac, + // whereas the B side (already direct since iterations 2/5) reads its + // fragment with one wide LDS access. Filling the same x[0..7] with one + // aligned 64-bit LDS read keeps the operand bytes bit-identical (the + // int32 accumulation is unchanged) and removes the reassembly VALU and + // its lgkmcnt wait states from the LDS->MMAC critical path. + // + // Iteration 10 (consolidation round): the A tile row stride is now + // kAStride128 = 72 (18 words) instead of 80, so the start word of row r + // is 18*r mod 32, a permutation of the 16 even bank residues: a 16-lane + // 8-byte fragment phase occupies each of the 32 banks exactly once and + // the A-side 2-way bank aliasing (80*8 == 640 == 0 mod 128) is gone. + load_fragment8(a_frag0, a_tile + local_row * kAStride128 + kk, + kAStride128, lane); load_fragment8(a_frag1, - a_tile + (local_row + kTileM) * kTileStride + kk, - kTileStride, lane); + a_tile + (local_row + kTileM) * kAStride128 + kk, + kAStride128, lane); load_fragment8(a_frag2, - a_tile + (local_row + 2 * kTileM) * kTileStride + kk, - kTileStride, lane); + a_tile + (local_row + 2 * kTileM) * kAStride128 + kk, + kAStride128, lane); load_fragment8(a_frag3, - a_tile + (local_row + 3 * kTileM) * kTileStride + kk, - kTileStride, lane); - // B fragments in the plane layout: fragment (f, kk) origin is - // b_tile + (kk >> 3) * 1024 + (local_col + 16f) * 8 and lane l reads - // the 8 bytes at (lane >> 4) * 1024 + (lane & 15) * 8 (lane-linear, - // zero bank conflicts); operand bytes are identical to the accepted - // n-major load_fragment8, so the int32 accumulation is bit-identical. - const int b_kc0 = kk >> 3; // 0 (kk == 0) or 4 (kk == 32) - load_fragment8_plane(b_frag0, - b_tile + b_kc0 * kBTilePlane + local_col * 8, - lane); - load_fragment8_plane(b_frag1, - b_tile + b_kc0 * kBTilePlane + - (local_col + kTileN) * 8, - lane); - load_fragment8_plane(b_frag2, - b_tile + b_kc0 * kBTilePlane + - (local_col + 2 * kTileN) * 8, - lane); - load_fragment8_plane(b_frag3, - b_tile + b_kc0 * kBTilePlane + - (local_col + 3 * kTileN) * 8, - lane); + a_tile + (local_row + 3 * kTileM) * kAStride128 + kk, + kAStride128, lane); + // Fragment-major LDS tile: the wave's four B fragments are four + // contiguous 512-byte DUMMA B tiles (local_col/16 + i, kk/32), each read + // with one aligned bank-conflict-free 8-byte LDS read per lane. + load_fragment8_block( + b_frag0, + b_tile + ((local_col / kTileN) * 2 + (kk / kTileK)) * kBFragBytes, + lane); + load_fragment8_block( + b_frag1, + b_tile + (((local_col / kTileN) + 1) * 2 + (kk / kTileK)) * + kBFragBytes, + lane); + load_fragment8_block( + b_frag2, + b_tile + (((local_col / kTileN) + 2) * 2 + (kk / kTileK)) * + kBFragBytes, + lane); + load_fragment8_block( + b_frag3, + b_tile + (((local_col / kTileN) + 3) * 2 + (kk / kTileK)) * + kBFragBytes, + lane); du_mma_sync(acc00, a_frag0, b_frag0, acc00); du_mma_sync(acc01, a_frag0, b_frag1, acc01); du_mma_sync(acc02, a_frag0, b_frag2, acc02); @@ -711,13 +967,12 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern du_mma_sync(acc33, a_frag3, b_frag3, acc33); } - // Last stage only: prefetch the epilogue's per-row x_scale and - // per-column weight_scale values into registers (uniform branch: k0 is - // block-uniform, so the barrier below is reached by every thread; the - // row < m guard keeps the x_scale reads in bounds for any m tail). The - // vmem latency overlaps the barrier below (s_barrier waits on - // lgkmcnt/arrival, not vmcnt) and the dead a_frag/b_frag VGPR slots are - // reused, so the epilogue is pure compute + coalesced stores. + // Last stage only: prefetch the epilogue's per-row x_scale and per-column + // weight_scale values into registers. k0 is block-uniform, so the + // barrier below is still reached by every thread; the row < m guard keeps + // the x_scale reads in bounds (never fires for m % 128 == 0). The vmem + // latency overlaps the barrier below and the dead a_frag/b_frag VGPR slots + // are reused, so the epilogue is pure compute + coalesced stores. if (k0 + kStageK >= k) { const int r = lane & 15; const int c4 = lane >> 4; @@ -732,16 +987,17 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern weight_scale + base_col + j * kTileN + 4 * c4); } } - // Protect the LDS buffers from the next stage's cooperative overwrite. + + // Protect the LDS buffer from the next stage's cooperative overwrite. __syncthreads(); } // Iteration 6 (epilogue round): coalesced direct-fragment epilogue. Each - // lane owns four contiguous bf16 columns of its fragment row and writes - // ONE 8-byte store per fragment (16 stores per wave vs 64 scattered 2-byte - // stores before); the scales come from the registers batched on the last - // K stage, so no vmem loads are interleaved with the stores. The int32 - // accumulation is untouched, and the multiply order / bf16 rounding are + // lane owns four contiguous bf16 columns of its fragment row and writes ONE + // 8-byte store per fragment (16 stores per wave vs 64 scattered 2-byte + // stores before); the scales come from the registers batched on the last K + // stage, so no vmem loads are interleaved with the stores. The int32 + // accumulation is untouched and the multiply order / bf16 rounding are // identical to the old per-element store, so output bits are unchanged. store_prefill_fragment_coalesced_scaled(acc00, xs_m[0], ws_m[0], out, base_row, base_col, m, n, lane); @@ -793,14 +1049,13 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_128x128_kern } // --------------------------------------------------------------------------- -// Generic scalar fallback: one thread per output element. Exact int32 dot -// over K, then fused x_scale * weight_scale, then bf16 store. Handles every -// unmatched (m, n, k), including all small-M API cases (M=2, M=16), M tails -// and M < 128 with the same (K, N). kNMajorPack == true decodes the -// iteration-5 swizzled 64-k-stage layout -// packed[((k0*8+kc)*n+col)*8+b] == raw[kk*n+col] (kk = k0*64+kc*8+b) for the -// exact (k, n) == (2048, 2048); otherwise the weight is the raw [K, N] -// row-major identity layout. +// Generic scalar fallback: one thread per output element, exact int32 dot over +// K, then the fused x_scale[row] * weight_scale[col] epilogue and one bf16 +// store. Handles every unmatched (m, n, k), including all small-M API cases +// (M=2, M=16) and any M < 128 with the same (N, K) as the tiled path. +// kNMajorPack == true decodes the fragment-major packed weight of the exact +// (k, n) == (2048, 2048) pair; otherwise the weight is the raw [K, N] +// identity layout. // --------------------------------------------------------------------------- template __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( @@ -820,23 +1075,26 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( } const int row = static_cast(linear / n); const int col = static_cast(linear - static_cast(row) * n); - int32_t acc = 0; const int8_t* a_row = x_q + static_cast(row) * k; + int32_t acc = 0; if constexpr (kNMajorPack) { - // Swizzled 64-k-stage pack (iteration 5, exact (2048,2048) only): - // packed[((k0*8 + kc)*n + col)*8 + b] == raw[kk*n + col] with - // kk = k0*64 + kc*8 + b, so raw[kk][col] = packed[((k0*8+kc)*n+col)*8+b]. + // Fragment-major pack: decode the 512-byte DUMMA B tile offset of + // (n = col, k = kk) exactly as w8a8_pack_fragment_major_i8_kernel wrote it. + const int n_frag = col / kTileN; + const int n_rem = col % kTileN; + const int k_frags = k / kTileK; for (int kk = 0; kk < k; ++kk) { - const int k0 = kk >> 6; - const int kc = (kk >> 3) & 7; - const int b = kk & 7; - const int8_t* p = weight + - (((static_cast(k0) * 8 + kc) * n + col) * 8) + - b; - acc += static_cast(a_row[kk]) * static_cast(*p); + const int k_frag = kk / kTileK; + const int kk8 = (kk % kTileK) / 8; + const int kk8r = kk % 8; + const int64_t idx = + ((static_cast(n_frag) * k_frags + k_frag) * kBFragBytes) + + (kk8 * 128) + (n_rem * 8) + kk8r; + acc += static_cast(a_row[kk]) * + static_cast(weight[idx]); } } else { - // identity [K, N] row-major: column col is strided by n. + // Identity [K, N]: column col is strided by n. const int8_t* b_col = weight + col; for (int kk = 0; kk < k; ++kk) { acc += static_cast(a_row[kk]) * @@ -849,54 +1107,63 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( } // --------------------------------------------------------------------------- -// Identity device-to-device weight packing (outside the timed region and -// outside Graph capture) for every (k, n) except the exact (2048, 2048) -// q_b_proj pair (which uses the iteration-5 swizzled pack). The packed -// buffer keeps the same byte count K*N and the same allocated address, so -// the layout is graph-stable. +// Iteration 5 fragment-major permutation for the exact (k, n) == (2048, 2048) +// q_b_proj weight: one DUMMA matrix_b tile (16 n x 32 k int8 = kBFragBytes = +// 512 B) is one contiguous, 16-byte-aligned block, +// packed[((n / 16) * (K / 32) + k / 32) * 512 +// + (k % 32) / 8 * 128 + (n % 16) * 8 + (k % 8)] +// == raw[k * N + n], +// i.e. a pure permutation of the same K * N bytes into the same graph-stable +// buffer. Runs once outside the timed region and outside Graph capture. The +// layout lets every 16-byte staged chunk become one aligned int4 global read +// plus one aligned int4 LDS store, and every per-lane 8-byte fragment read +// covers all 32 LDS banks exactly once per 16-lane phase. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( +__global__ __launch_bounds__(256) void w8a8_pack_fragment_major_i8_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, - int64_t count) { + int k, + int n) { const int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (linear < count) { - packed[linear] = raw[linear]; + const int64_t total = static_cast(n) * k; + if (linear >= total) { + return; } + const int k_frags = k / kTileK; // K / 32 + const int blk = static_cast(linear / kBFragBytes); + const int off = + static_cast(linear - static_cast(blk) * kBFragBytes); + const int kk8 = off >> 7; // (k % 32) / 8, 0..3 + const int rem = off & 127; // 0..127 + const int n_rem = rem >> 3; // n % 16 + const int kk8r = rem & 7; // k % 8 + const int n_frag = blk / k_frags; // n / 16 + const int k_frag = blk - n_frag * k_frags; // k / 32 + const int n_idx = (n_frag * kTileN) + n_rem; + const int k_idx = (k_frag * kTileK) + (kk8 * 8) + kk8r; + packed[linear] = raw[static_cast(k_idx) * n + n_idx]; } -// Iteration 5 (packing round): swizzled 64-k-stage-major permutation for the -// exact (k, n) == (2048, 2048) q_b_proj weight: -// packed[((k0*8 + kc)*n + col)*8 + b] = raw[kk*n + col], -// kk = k0*64 + kc*8 + b (k0 = 64-k stage, kc = 8-byte sub-chunk, b = byte). -// Byte-wise (one thread per byte) so the permutation is trivially correct; -// runs once outside the timed region and outside Graph capture, keeping the -// same byte count K*N and the same graph-stable buffer address. The layout -// lets the 128x128 kernel stage each B tile as 16-byte plane chunks (one -// aligned int4 global read per chunk, perfectly coalesced) into the plane -// LDS tile [kc][n][8] whose fragment loads are lane-linear with zero bank -// conflicts. -__global__ __launch_bounds__(256) void w8a8_pack_swizzle_kernel( +// --------------------------------------------------------------------------- +// Bootstrap identity device-to-device packing (outside the timed region and +// outside Graph capture) for every (k, n) except the exact fragment-major pair +// above. +// The packed buffers keep the same byte count and the same graph-stable +// addresses, so the GEMM consumes the raw [K, N] row-major layout directly. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, - int k, - int n) { + int64_t count) { const int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t total = static_cast(k) * n; - if (linear < total) { - const int kk = static_cast(linear / n); - const int col = static_cast(linear - static_cast(kk) * n); - const int k0 = kk >> 6; - const int kc = (kk >> 3) & 7; - const int b = kk & 7; - packed[(((static_cast(k0) * 8 + kc) * n + col) * 8) + b] = - raw[linear]; + if (linear < count) { + packed[linear] = raw[linear]; } } -__global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( +__global__ __launch_bounds__(256) void w8a8_pack_identity_f32_kernel( const float* __restrict__ raw, float* __restrict__ packed, int count) { @@ -910,9 +1177,9 @@ __global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( // --------------------------------------------------------------------------- // Stable host launch symbols consumed by csrc/bindings.cpp. -// Both launchers are pure dispatch: no allocation, no packing, no +// Both launchers are pure dispatch: no allocation, no packing decision, no // synchronization, no default-stream launch; they run on the caller-provided -// PyTorch stream and are CUDA-Graph safe. +// PyTorch stream and are CUDA/HIP-Graph safe. // --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, @@ -930,32 +1197,39 @@ extern "C" void launch_w8a8_gemm( (void)workspace_bytes; auto* out_bf16 = reinterpret_cast(out); - // Explicit dispatch. Large-M shapes with exact tiled geometry (both - // assigned shapes: M=4096, N in {2048, 3584}, K in {2048, 512}) take the - // native INT8 DUMMA tiled path; every other (m, n, k) - including small-M - // API cases (M=2, M=16) and M < 128 - takes the scalar fallback. The - // exact (k, n) == (2048, 2048) q_b_proj pair takes the iteration-4 128x128 - // tile with the iteration-5 swizzled-pack B layout; every other (k, n) - - // including kv_b_proj (512, 3584) - uses the byte-identical 128x64 - // kNMajorB=false arm (kNMajorPack == true for the scalar fallback's - // (2048,2048) decode). + // Explicit dispatch. The assigned large-prefill shape (M=4096, N=2048, + // K=2048) takes the native INT8 DUMMA tiled path; every other (m, n, k) - + // including all small-M API cases (M=2, M=16) and M < 128 - takes the + // scalar fallback. The geometry guard is exact so the paired small-M API + // shape with the same (N, K) can never enter the tiled kernel. Within the + // tiled path, only the exact (k, n) == (2048, 2048) pair - whose weight the + // packer permuted fragment-major - takes the packed arm: iteration 4 gives + // it the 128x128 tile-aspect arm when n is a multiple of 128, and the + // 128x64 arm otherwise. Every other pair keeps the byte-identical raw + // [K, N] identity 128x64 arm. + const bool packed_pair = (k == kPackFragMajorK && n == kPackFragMajorN); if (m >= kBlockM && m % kBlockM == 0 && n % kBlockN == 0 && k % kStageK == 0) { const dim3 grid(static_cast(n / kBlockN), static_cast(m / kBlockM)); const dim3 block(kBlockThreads); - if (k == kPackNMajorK && n == kPackNMajorN) { + if (packed_pair && n % kBlockN128 == 0) { // Exact q_b_proj (2048, 2048): iteration-4 tile-aspect arm, 128x128 - // tile with the iteration-5 swizzled-pack plane B layout; grid - // (N/128, M/128) = 512 blocks for M=4096. n == 2048 is divisible by - // 128 by the exact guard; the outer arm already guarantees - // m % 128 == 0, k % 64 == 0. - const dim3 grid128(static_cast(n / 128), - static_cast(m / 128)); + // tile (M:N aspect 1:1) with the iteration-5 fragment-major packed B + // layout; grid (N/128, M/128) = 512 blocks for M=4096. n == 2048 is + // divisible by 128 by the exact guard, and the outer tiled guard already + // implies m % 128 == 0 and k % 64 == 0. + const dim3 grid128(static_cast(n / kBlockN128), + static_cast(m / kBlockM)); hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_128x128_kernel), + HIP_KERNEL_NAME(w8a8_dumma_prefill_128x128_nmajor_kernel), grid128, block, 0, stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else if (packed_pair) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_prefill_128x64_kernel), + grid, block, 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); } else { hipLaunchKernelGGL( HIP_KERNEL_NAME(w8a8_dumma_prefill_128x64_kernel), @@ -967,7 +1241,7 @@ extern "C" void launch_w8a8_gemm( const int64_t total = static_cast(m) * n; const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); const dim3 block(kBlock); - if (k == kPackNMajorK && n == kPackNMajorN) { + if (packed_pair) { hipLaunchKernelGGL( HIP_KERNEL_NAME(w8a8_gemm_scalar_fallback_kernel), grid, block, 0, stream, @@ -992,18 +1266,21 @@ extern "C" void launch_pack_w8a8_weight( hipStream_t stream) { constexpr int kBlock = 256; const dim3 block(kBlock); - // Bootstrap pack: iteration-5 swizzled 64-k-stage permutation for the exact - // (2048, 2048) q_b_proj weight; identity device-to-device copy for every - // other (k, n). + + // Iteration 5 pack: fragment-major permutation of the exact + // (k, n) == (2048, 2048) q_b_proj weight (one 512-byte DUMMA B tile per + // (n/16, k/32) pair); identity device-to-device copy for every other + // (k, n). Same byte count as the raw weight either way, so the buffer size + // / address is stable across Graph capture and replay. const int64_t weight_count = static_cast(k) * n; const dim3 weight_grid( static_cast((weight_count + kBlock - 1) / kBlock)); - if (k == kPackNMajorK && n == kPackNMajorN) { - hipLaunchKernelGGL(w8a8_pack_swizzle_kernel, + if (k == kPackFragMajorK && n == kPackFragMajorN) { + hipLaunchKernelGGL(w8a8_pack_fragment_major_i8_kernel, weight_grid, block, 0, stream, raw_weight, packed_weight, k, n); } else { - hipLaunchKernelGGL(w8a8_pack_identity_kernel, + hipLaunchKernelGGL(w8a8_pack_identity_i8_kernel, weight_grid, block, 0, stream, raw_weight, packed_weight, weight_count); } @@ -1011,9 +1288,9 @@ extern "C" void launch_pack_w8a8_weight( const int64_t scale_count = n; const dim3 scale_grid( static_cast((scale_count + kBlock - 1) / kBlock)); - hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, + hipLaunchKernelGGL(w8a8_pack_identity_f32_kernel, scale_grid, block, 0, stream, - weight_scale, packed_weight_scale, n); + weight_scale, packed_weight_scale, static_cast(n)); (void)hipGetLastError(); } // @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/o_proj.hip index fc602161..30ff5116 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/o_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/o_proj.hip @@ -1,96 +1,58 @@ -// @@variant shape=hy3_tp4_o_proj_m16 commit=cb7af9e4e0a3adf58bc62ec3cf8be1e1a3f4c010 added=2026-08-26 -// median_us=15.51 p90_us=15.55 -// source=hy3-dsh-tp4-m16-2-368c654c +// @@variant shape=hy3_tp4_o_proj_m16 commit=71d1281fd4cc5068bc425a5f9a49d5b8d40e90ec added=2026-09-11 +// median_us=14.73 p90_us=14.79 speedup=3.707 baseline_us=54.62 +// source=test-9-11-1-b01955b3/iteration_001/hy3_tp4_o_proj_m16 // MetaInfer W8A8 INT8 GEMM - HIP implementation for K500SM_AI / gfx928. // -// Worker: worker_1 (physical GPU 1) -// Assigned shape: hy3_tp4_o_proj_m16 (M=16, N=4096, K=2048) +// Assigned shape: hy3_tp4_o_proj_m16 (M=16, N=4096, K=2048). // -// Iteration 1 (accepted) = minimal 16x16x32 DUMMA bootstrap: 256 blocks x 1 -// 64-thread wavefront, one N tile per block, library byte-loads of A/B -// straight from the logical row-major globals, 64 ascending-K int32 steps, -// LDS epilogue plane, zero in-loop barriers. -// -> median 89.90 us / p90 90.10 us, 95.3 GB/s algorithmic, 263,424 -// vmem_read/replay = 16.08 loads per K-step per wave, all fenced -// vmcnt(15)..vmcnt(0) before the single v_mmac. -// Iteration 2 (rejected) = grid-parallelism probe: 128 blocks x 2 waves x 2 -// adjacent N tiles (SAME 256 waves, no split-K, no packing). -// -> median 117.81 us (REGRESSION): per-CU wave count is NOT the lever on -// the byte-load design at 2.13 waves/CU; the named next lever was load -// vectorization (packed-B dwordx2 fragment loads) and/or in-block -// split-K. -// Iteration 3 (rejected, measurement-polluted) = architecture round: 256 -// blocks x 2 waves, in-block split-K=2 (uniform K=1024 slices), packed-B -// [n_tile][k_step][lane][8] fragment-slot loads, register-only K-loop -// transport with depth-1 prefetch, K=128 stages. This is byte-for-byte the -// authoritative lineage's ACCEPTED final kernel for this exact shape -// (w8a8_gemm_m16_dumma_packedb_kernel, 18.221/18.242 us median/p90 there). -// -> min 21.37 us (9 uncontended samples tight at 21.4-22.5 us) but median -// 166.13 / p90 216.62: the 30-sample window was bimodal, with ~19 -// samples at ~190-240 us matching an ~8-10x-slower interleaved workload -// on the GPU (external interference, not kernel behavior: the same -// captured graph cannot run 10x slower under its own resources at 2.13 -// blocks/CU all-resident). The fast tail proves the design; the slow -// head is the measurement window. -// Iteration 4 (accepted, 17.155 us median / 17.185 us p90, 499.5 GB/s -// algorithmic - the fastest accepted W8A8 kernel in every worker ledger on -// this node) = architecture/pipeline round, multi-N-tile reuse: each block -// computes TWO adjacent 16x16 N tiles (T=2) so the A fragments loaded per -// K=128 stage are shared across both tiles in registers; everything else -// frozen from iteration 3 (packed-B layout, in-block split-K=2, K=128 -// stages, depth-1 register prefetch). HBM floor analysis: 8,484,352 B of -// unavoidable traffic (8 MiB B read-once + 32 KiB A + 128 KiB bf16 C) at -// 17.155 us = 499.5 GB/s, and no accepted kernel on this device has ever -// exceeded that figure - the kernel sits at the practical streaming -// ceiling. -// Iteration 5 (THIS round) = HIP-only packed-layout comparison (mandate: -// packed layout / A-only staging / B-only staging): repack B from -// [n_tile][k_step][lane][8] to a TILE-PAIR-INTERLEAVED dwordx4 layout -// [n_tile_pair][k_step][lane][16] in which lane's 16 bytes are tile0's 8 -// fragment bytes followed by tile1's 8 fragment bytes for the same -// (k_step, lane). Each lane then covers BOTH tiles with ONE aligned 16-byte -// load per k_step (a single 1 KiB contiguous region per load, a single -// 4 KiB contiguous region per K=128 stage, and one sequential 64 KiB -// stream per block lifetime - was two 2 KiB regions 32 KiB apart per -// stage); B vmem instructions drop from 8 to 4 per stage per wave at an -// identical 24 VGPR in-flight footprint and identical 8 MiB of HBM bytes. -// Geometry frozen from iteration 4 (128 blocks x 2 waves x T=2, in-block -// split-K=2, K=128 stages, depth-1 register prefetch, partial-planes-only -// LDS, one end-of-K barrier). See the kernel comment below. +// Round 1 (mandate: establish_arch - first correct DUMMA geometry). The +// scalar bootstrap is replaced by a native gfx928 DUMMA INT8 m16n16k32 +// architecture for the exact assigned shape: +// * grid = N / (16 * 2) = 128 blocks, one block per pair of adjacent 16x16 +// N tiles (128 >= 120 device CUs), 128 threads = 2 x 64-lane wavefronts; +// * in-block split-K=2: wave w owns 8 interleaved K=128 stages (K=1024, +// stages w, w + 2, ...) of BOTH tiles as two independent int32 accumulator +// fragments, so the A fragments loaded per K stage are shared across the +// two tiles from registers; combined with the group-major packed-B layout +// the block's aggregate walk is one 8 KiB window per iteration and the whole +// 128-block grid tiles those windows into a single dense advancing stream; +// * B is pre-packed once (untimed, outside Graph capture) into a +// tile-pair-interleaved [k_step_group][n_tile_pair][k_step_within][lane][16] +// layout: one lane's 16 bytes hold tile0's 8 B fragment bytes followed by +// tile1's 8 B fragment bytes for the same (k_step, lane), so a single +// aligned 16-byte dwordx4 load covers both tiles and every B byte is read +// exactly once. Grouping the kStepsPerGroup = 8 k_steps one K-loop +// iteration consumes outermost makes the 128 blocks' concurrent 8 KiB +// packed-B windows tile ONE contiguous monotonic group region (1 MiB at +// n = 4096) instead of 128 windows 64 KiB apart, so the whole grid presents +// the DRAM/L2 queues a single dense advancing stream; +// * register-only K=128-stage transport with depth-1 prefetch (4 A dwordx2 +// shared by both tiles + 4 B dwordx4 = 8 loads, 6 KiB in flight per wait); +// * the only barrier is one END-of-K __syncthreads() before the fused +// partial combine + scale -> bf16 epilogue; barriers per K step = 0. // -// Scalar fallback strategy (one thread per output element; still the path for -// every shape OTHER than the exact assigned M=16 shape): -// * blockDim = 256 (multiple of the gfx928 64-lane wavefront); -// * grid = (ceil(N/256), M): each block covers 256 adjacent N columns of a -// single M row, so adjacent lanes touch adjacent addresses in the -// fastest-changing N dimension (coalesced B reads / bf16 stores); -// * the complete K loop accumulates exactly in int32. For the assigned -// K=2048 the max |dot| is 2048*127*127 = 33,032,192 << 2^31, so the -// integer accumulation never overflows; -// * after the K loop the two float scales are applied as -// float(dot) * x_scale[m] * weight_scale[n] and the result is stored as -// bf16 (hip_bfloat16 via __float2bfloat16); -// * for the exact (n, k) == (4096, 2048) pair the weight tensor is the -// packed [n_tile_pair][k_step][lane][16] tile-pair-interleaved layout -// (16 bytes = tile0's 8 fragment bytes ++ tile1's 8 fragment bytes for -// the same k_step/lane), which the fallback decodes elementwise (keeps -// the paired M=2 validation exact with the new pack); -// * workspace is untouched by the timed operator (in-block split-K needs no -// partial planes). +// The timed operator (launch_w8a8_gemm) only launches kernels on the +// caller-provided stream and touches only caller-provided buffers: no +// allocation, no compilation, no packing, no host or device synchronization, +// no default-stream fallback. In-block split-K needs no cross-block partials, +// so the caller workspace is deliberately untouched. // -// pack_weight: for the exact (k, n) == (2048, 4096) o_proj weight it is a -// one-time device permutation into packed[n_tile_pair][k_step][lane][16] B -// fragment slots (byte count unchanged: 8 MiB; runs outside the timed GEMM -// and outside Graph capture). Every other (K, N) keeps the generic identity -// device-to-device copy; packed_weight_scale[n] = weight_scale[n] always. +// Scalar fallback: one thread per output element for every shape other than +// the exact M=16 assignment, including the paired M=2 shape that shares the +// same (N, K); for (n, k) == (4096, 2048) it decodes the packed tile-pair +// layout elementwise, so the paired validation stays exact. +// +// pack_weight: for the exact (k, n) == (2048, 4096) o_proj weight it is the +// one-time permutation into the packed fragment-slot layout (byte count +// unchanged, 8 MiB); every other (K, N) keeps the generic identity +// device-to-device copy. packed_weight_scale[n] = weight_scale[n] always. // // Host launch symbols (stable, C linkage, consumed by csrc/bindings.cpp): // launch_w8a8_gemm(...) // launch_pack_w8a8_weight(...) // // Include order is known-good for this DTK: hip_runtime.h first, then -// hip_bfloat16.h, then du_mma.h (du_mma.h is not self-contained when -// included before the HIP runtime headers). +// hip_bfloat16.h, then du_mma.h. #include #include @@ -109,11 +71,10 @@ constexpr int kTargetK = 2048; // --------------------------------------------------------------------------- // DUMMA geometry for the exact assigned shape: gfx928 INT8 m16n16k32 tile. -// Iteration 4 multi-N-tile reuse: each block owns TWO adjacent 16x16 N tiles -// (T=2, grid = N/(16*2) = 128 blocks >= 120 device CUs) and 2 x 64-thread -// wavefronts; wave w accumulates the int32 dot over its UNIFORM K=1024 slice -// for BOTH tiles, so the A fragments of each K=128 stage are loaded once and -// reused across the two tiles from registers. +// Each block owns TWO adjacent 16x16 N tiles (T=2) and 2 x 64-thread +// wavefronts; wave w accumulates the int32 dot over its 8-of-16 interleaved +// K=128 stages (K=1024) for BOTH tiles, so the A fragments of each K=128 stage +// are loaded once and reused across the two tiles from registers. // --------------------------------------------------------------------------- constexpr int kDummaBlockThreads = 64; // one gfx928 wavefront constexpr int kDummaWavesPerBlock = 2; // in-block split-K=2 @@ -122,30 +83,49 @@ constexpr int kDummaTileM = 16; constexpr int kDummaTileN = 16; constexpr int kDummaTileK = 32; -// Bytes of one (n_tile, k_step) B fragment in the packed layout: -// 16 columns * 32 k rows = 512 B, stored as 64 lanes * 8 B. -constexpr int kPackedFragBytes = kDummaTileN * kDummaTileK; // 512 - -// Iteration-5 tile-pair-interleaved layout: one (k_step) slot holds BOTH -// tiles of a block in 16 B per lane - tile0's 8 fragment bytes (bytes 0..7) -// followed by tile1's 8 fragment bytes (bytes 8..15) - so a single aligned -// 16-byte (dwordx4) load per lane covers one k_step of both tiles: 64 lanes -// * 16 B = 1 KiB contiguous per load. Per n_tile_pair the 64 k_steps are -// stored ascending (64 KiB sequential per block lifetime). +// Tile-pair-interleaved packed-B slot: one (k_step) slot holds BOTH tiles of a +// block in 16 B per lane - tile0's 8 fragment bytes (bytes 0..7) followed by +// tile1's 8 fragment bytes (bytes 8..15) - so one aligned 16-byte (dwordx4) +// load per lane covers one k_step of both tiles: 64 lanes * 16 B = 1 KiB +// contiguous per load. k_steps are ordered GROUP-major: the kStepsPerGroup = 8 +// k_steps one K-loop iteration of a block consumes are adjacent across pairs, +// and the per-pair slot stride inside a group is kStepsPerGroup * +// kPackedPairBytes (see kStepsPerGroup below). constexpr int kPackedPairBytes = kDummaBlockThreads * 16; // 1024 +// Padded row stride (in int32 lanes) of one 16x16 int32 combine plane. +// du_store_matrix_sync stores the int8 m16n16k32 accumulator fragment with +// row = lane & 15 and col = (lane >> 4) + 4 * i (x[i] at +4i dwords), so with +// the natural stride 16 every lane lands on bank 16*(row & 1) + (col & 7): +// only 8 of the 32 LDS banks are touched and the store serializes (~28 extra +// bank-conflict cycles per store, 28,672 per operator replay - the only +// non-compulsory counter left in the current best). +// Stride 24 = 16 + 8 spreads the same 128 dwords exactly 4-per-bank (the +// hardware minimum) for the store, and keeps the epilogue's (row, row + 8) +// read pair 8 * 24 = 192 = 3 * 64 dwords apart, so the compiler can still emit +// one ds_read2st64 per pair and those reads also stay at the 4-phase minimum. +// Same int32 values, same row/col meaning, order-independent combine. +constexpr int kCombineStride = kDummaTileN + 8; // 24 + // K=128 stage: 4 DUMMA steps per stage, 8 loads per stage per wave (4 A // dwordx2 shared by both tiles + 4 B dwordx4, one per k_step covering both // tiles = 6 KB in flight), 8 v_mmac per stage per wave (4 steps x 2 tiles). -constexpr int kStageK = 128; // K bytes per stage +constexpr int kStageK = 128; // K elements per stage constexpr int kStageSteps = kStageK / kDummaTileK; // 4 DUMMA steps per stage +// Packed-B group: the kStepsPerGroup = 8 k_steps (2 waves x kStageSteps) that +// one K-loop iteration of a block consumes. This is the outermost packed +// dimension, so at loop iteration s the whole grid reads one contiguous +// group_stride region: block b's 8 KiB window sits at pair-slot b inside it. +// Dense monotonic footprint instead of 128 windows 64 KiB apart. +constexpr int kStepsPerGroup = + kDummaWavesPerBlock * (kStageK / kDummaTileK); // 8 + // --------------------------------------------------------------------------- // Packed-B decode for the generic scalar fallback: returns logical -// weight[kk, n] from the packed [n_tile_pair][k_step][lane][16] -// tile-pair-interleaved layout (see the pack kernel for the exact -// permutation). The 16 bytes of a lane slot are tile0's 8 fragment bytes -// (bytes 0..7) followed by tile1's 8 fragment bytes (bytes 8..15). +// weight[kk, n] from the group-major +// [k_step_group][n_tile_pair][k_step_within][lane][16] tile-pair-interleaved +// layout (see the pack kernel for the exact permutation). // --------------------------------------------------------------------------- __device__ __forceinline__ int8_t packed_b_element(const int8_t* __restrict__ packed, int n, int kk) { @@ -157,18 +137,23 @@ packed_b_element(const int8_t* __restrict__ packed, int n, int kk) { const int kk8 = kk & 31; const int lane = (kk8 >> 3) * 16 + nn; const int i = kk8 & 7; - return packed[(static_cast(pair) * (kTargetK / kDummaTileK) + - k_step) * + // Group-major packed layout [k_step_group][n_tile_pair][k_step_within][lane]. + const int n_pairs = kTargetN / (2 * kDummaTileN); + const int group = k_step / kStepsPerGroup; + const int within = k_step % kStepsPerGroup; + return packed[static_cast(group) * n_pairs * kStepsPerGroup * kPackedPairBytes + - lane * 16 + sub * 8 + i]; + static_cast(pair) * kStepsPerGroup * kPackedPairBytes + + static_cast(within) * kPackedPairBytes + lane * 16 + + sub * 8 + i]; } // One thread computes one output element out[row, col]: // acc = sum_k a[row,k] * b[k,col] (exact int32, k ascending) // out = bf16(acc * x_scale[row] * weight_scale[col]) -// Bounds-checked, so it is a valid generic fallback for every (m, n, k). -// For the exact (n, k) == (4096, 2048) pair the weight tensor is the packed -// [n_tile_pair][k_step][lane][16] layout (o_proj M=16 pack), decoded +// Bounds-checked, so it is a valid generic fallback for every (m, n, k). For +// the exact (n, k) == (4096, 2048) pair the weight tensor is the group-major +// [k_step_group][n_tile_pair][k_step_within][lane][16] layout, decoded // elementwise. __global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( const int8_t* __restrict__ a, // [m, k] row-major @@ -203,55 +188,59 @@ __global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( } // --------------------------------------------------------------------------- -// Iteration-5 stage prefetch helper (multi-N-tile reuse): issue the eight -// loads for one K=128 stage - the 4 A dwordx2 fragments (shared by both -// tiles) plus 4 B dwordx4 fragment-pair loads (one per k_step, each covering -// BOTH tiles) - into caller-provided u64 registers. +// Stage prefetch helper (multi-N-tile reuse): issue the eight loads for one +// K=128 stage - the 4 A dwordx2 fragments (shared by both tiles) plus 4 B +// dwordx4 fragment-pair loads (one per k_step, each covering BOTH tiles) - +// into caller-provided u64 registers. // * A: four 8-byte loads of one activation row at ldm = k, 32 B apart (the -// four DUMMA steps of the stage) - lane ownership is loop-invariant +// four DUMMA steps of the stage). Lane ownership is loop-invariant // (du_mma.hpp matrix_a row_major loader: row = lane & 0xf, // col = (lane >> 4) << 3, x[i] = p[row*ldm + col + i], 8 consecutive -// bytes; du_mma_sync consumes a.x as one 64-bit value, du_mma.hpp line -// 1198). A stays the logical row-major 32 KiB activation (L2-hot). -// These 4 values feed BOTH accumulator tiles - A is read once per block -// instead of once per N tile. -// * B (iteration-5 tile-pair-interleaved layout): per k_step one aligned -// 16-byte load from the [n_tile_pair][k_step][lane][16] slot whose bytes -// 0..7 are tile0's fragment bytes and bytes 8..15 tile1's fragment bytes -// for the same (k_step, lane) - byte-order identical to the library -// matrix_b row_major loader (x[i] = p[(col + i) * ldm + row]) per tile, -// so the two 8-byte halves feed b_frag0.x / b_frag1.x directly. A stage's -// four k_step loads sit at consecutive 1 KiB slot offsets: one 4 KiB -// contiguous region per stage per wave; each wave's 8 stages walk one -// sequential 32 KiB half of the block's 64 KiB pair region, so every B -// byte is read exactly once per replay. -// * The caller's k0_base K-slice offset makes each split wave walk its own -// packed k_step range. The 16-byte loads are 16 B aligned (slot offset is -// k_step * 1024 + lane * 16 on a 16 B-aligned packed buffer). +// bytes). A stays the logical row-major 32 KiB activation (L2-hot) and +// feeds BOTH accumulator tiles. +// * B (tile-pair-interleaved, group-major layout): per k_step one aligned +// 16-byte load from the [k_step_group][n_tile_pair][k_step_within][lane][16] +// slot whose bytes 0..7 are tile0's fragment bytes and bytes 8..15 tile1's +// fragment bytes for the same (k_step, lane) - byte-order identical to the +// library matrix_b row_major loader (x[i] = p[(col + i) * ldm + row]) per +// tile, so the two 8-byte halves feed b_frag0.x / b_frag1.x directly. +// * The caller's k0_base K-slice offset plus the k0_stride stage stride make +// each split wave walk its own packed k_step sequence; group_stride is the +// byte size of one k_step_group across all n_tile_pairs. The four k_steps +// of one stage are consecutive inside their group, so all four 16-byte +// loads are 16 B aligned (kPackedPairBytes apart) on a 16 B-aligned +// packed buffer. // --------------------------------------------------------------------------- __device__ __forceinline__ void load_stage_fragments_t2( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, int lane, int a_row, int a_col, int64_t b_base, int k, - int64_t k0_base, int s, + int64_t k0_base, int64_t k0_stride, int64_t group_stride, int s, uint64_t& a0, uint64_t& a1, uint64_t& a2, uint64_t& a3, uint64_t& b00, uint64_t& b01, uint64_t& b02, uint64_t& b03, uint64_t& b10, uint64_t& b11, uint64_t& b12, uint64_t& b13) { - const int64_t k0 = k0_base + static_cast(s) * kStageK; + const int64_t k0 = k0_base + static_cast(s) * k0_stride; const int64_t a_off = static_cast(a_row) * k + k0 + a_col; a0 = *reinterpret_cast(x_q + a_off); a1 = *reinterpret_cast(x_q + a_off + kDummaTileK); a2 = *reinterpret_cast(x_q + a_off + 2 * kDummaTileK); a3 = *reinterpret_cast(x_q + a_off + 3 * kDummaTileK); - const int64_t pair_off = (k0 / kDummaTileK) * kPackedPairBytes + + // Group-major packed-B slot: k_step decides the group (outermost) and the + // position inside it; b_base is this block's n_tile_pair offset in a group. + const int64_t k_step = k0 / kDummaTileK; + const int64_t pair_off = (k_step / kStepsPerGroup) * group_stride + + (k_step % kStepsPerGroup) * kPackedPairBytes + static_cast(lane) * 16; const int64_t b_off = b_base + pair_off; // One 16-byte load per k_step covers both tiles: bytes 0..7 = tile0's // fragment, bytes 8..15 = tile1's fragment (little-endian word order). const uint4 bv0 = *reinterpret_cast(weight + b_off); - const uint4 bv1 = *reinterpret_cast(weight + b_off + kPackedPairBytes); - const uint4 bv2 = *reinterpret_cast(weight + b_off + 2 * kPackedPairBytes); - const uint4 bv3 = *reinterpret_cast(weight + b_off + 3 * kPackedPairBytes); + const uint4 bv1 = + *reinterpret_cast(weight + b_off + kPackedPairBytes); + const uint4 bv2 = + *reinterpret_cast(weight + b_off + 2 * kPackedPairBytes); + const uint4 bv3 = + *reinterpret_cast(weight + b_off + 3 * kPackedPairBytes); b00 = (static_cast(bv0.y) << 32) | bv0.x; // tile0, k_step 0 b01 = (static_cast(bv1.y) << 32) | bv1.x; // tile0, k_step 1 b02 = (static_cast(bv2.y) << 32) | bv2.x; // tile0, k_step 2 @@ -262,22 +251,14 @@ __device__ __forceinline__ void load_stage_fragments_t2( b13 = (static_cast(bv3.w) << 32) | bv3.z; // tile1, k_step 3 } -// --------------------------------------------------------------------------- -// Iteration 17 (HIP-only tail consolidation): float32 -> bf16 round-to-nearest- -// even WITHOUT the __float2bfloat16 NaN/Inf branch. The exact iteration-13 -// code object shows each of the epilogue's four conversions as a ~10- -// instruction guarded sequence ON the block's post-barrier finishing tail: -// the finite path is v_bfe_u32 (bit 16) + v_add3_u32 (u + bit + 0x7fff) with -// the store's d16_hi keeping the high 16 bits, plus the dead-for-this-operator -// s_and_saveexec / v_or_b32 / v_cmp_eq_sdwa / v_cndmask NaN/Inf handling. -// Every converted value here is finite (the int32 dot is bounded by +// Float32 -> bf16 round-to-nearest-even WITHOUT the __float2bfloat16 NaN/Inf +// branch. Every converted value here is finite (the int32 dot is bounded by // 2048*127*127 << 2^31 and both scales are finite), so the guarded branch is -// pure tail cost sitting between the LDS reads and the stores. This helper +// pure tail cost sitting between the LDS reads and the stores. This helper // emits only the finite RNE path: r = u + 0x7fff + ((u >> 16) & 1), high 16 // bits kept - bit-identical to __float2bfloat16 for every finite input (same // rounding, same carry into the exponent at the bf16 range edge); the only // divergence is NaN payload bits, which cannot occur here. -// --------------------------------------------------------------------------- __device__ __forceinline__ hip_bfloat16 bf16_rne_finite(float f) { uint32_t u; __builtin_memcpy(&u, &f, sizeof(u)); @@ -293,48 +274,44 @@ __device__ __forceinline__ hip_bfloat16 bf16_rne_finite(float f) { // tile, weight pre-packed into tile-pair-interleaved B fragment slots, // MULTI-N-TILE REUSE. // * Geometry: 128 blocks (grid = N/(16*2), one block per pair of adjacent -// 16x16 N tiles; 128 >= 120 device CUs, so every CU gets at least one -// block) x 128 threads = 2 x 64-lane wavefronts. -// * In-block split-K=2: wave w owns the UNIFORM K slice of 8 x K=128 stages -// (stages [floor(w*16/2), floor((w+1)*16/2)) = exactly 8 stages, K=1024) -// and accumulates BOTH tiles over that slice - 2 independent accumulator -// fragments per wave. Per K=128 stage per wave: 4 A dwordx2 + 4 B dwordx4 -// (one 16-byte load per k_step covering both tiles) = 8 loads (6 KB in -// flight, identical byte footprint to iteration 4) covering an 8-v_mmac -// burst (4 steps x 2 tiles): load:compute ratio 1.5 (iteration 4) -> 1.0, -// B vmem instructions per replay drop from 16,384 dwordx2 to 8,192 -// dwordx4, and each block's B lifetime stream is ONE sequential 64 KiB -// region (4 KiB contiguous per stage) instead of two 32 KiB regions 32 -// KiB apart; A global traffic unchanged (L2-hot 32 KiB footprint, read -// once per block), B HBM traffic unchanged (8 MiB, every byte read -// exactly once). -// * REGISTER-ONLY K-loop transport with depth-1 prefetch, identical to -// iteration 4: current stage's 8 loads' fragments live in registers, the -// NEXT stage's 8 loads are issued BEFORE the current 8-MMAC burst, -// and the compiler-inserted vmcnt wait lands at the next iteration's -// first fragment fill (global_load -> v_mmac burst -> wait, no LDS hop, -// no lgkmcnt wait in the K loop). The 4 A dwordx2 + 4 B dwordx4 hold the -// same 24 VGPR as iteration 4's 12 dwordx2, so occupancy is unchanged. -// * LDS holds only the four 1 KiB int32 partial planes (2 waves x 2 tiles, -// 4 KiB/block); the ONLY barrier is the single END-of-K __syncthreads() -// before the fused partial combine + scale -> bf16 epilogue (all 128 -// threads, 4 elements each). Barriers per K step = 0. -// * Correctness: per element the int32 accumulation is ascending-K within -// the wave's slice (4 ascending steps per stage, stages ascending), then -// the two waves' partials are summed in LDS (order-independent int32, no -// overflow: max |dot| = 2048*127*127 << 2^31) - bit-identical to -// iteration 4. The iteration-5 repack changed the packed layout -// COHERENTLY across the pack kernel, the DUMMA kernel's loads and the -// scalar fallback decode (same (tile, k_step, lane, byte) -> logical -// (k, n) mapping, new memory order), so the paired M=2 validation stays -// exact. +// 16x16 N tiles; 128 >= 120 device CUs) x 128 threads = 2 x 64-lane +// wavefronts. +// * In-block split-K=2: wave w owns 8 x K=128 stages (K=1024) and +// accumulates BOTH tiles over that slice - two independent accumulator +// fragments per wave. The two waves' stage ownership is INTERLEAVED +// (wave w owns stages w, w + 2, ...) and the packed-B layout is +// group-major, so at each K-loop iteration the two waves cover one +// contiguous 8 KiB window and the 128 blocks' windows tile one dense +// group region that advances monotonically (a single dense DRAM stream). +// Per K=128 stage per wave: 4 A dwordx2 + 4 B dwordx4 (one 16-byte load per +// k_step covering both tiles) = 8 loads (6 KB in flight) covering an +// 8-v_mmac burst (4 steps x 2 tiles). Every B byte is read exactly once; A +// is read once per block. +// * REGISTER-ONLY K-loop transport with depth-1 prefetch: the current +// stage's fragments live in registers, the NEXT stage's 8 loads are issued +// BEFORE the current 8-MMAC burst, and the compiler-inserted vmcnt wait +// lands at the next iteration's first fragment fill (global_load -> +// v_mmac burst -> wait, no LDS hop, no lgkmcnt wait in the K loop). +// * LDS holds only the four int32 partial planes (2 waves x 2 tiles, each +// 16 x kCombineStride = 16 x 24 int32, 6 KiB/block, padded so the fragment +// store and the epilogue read are bank-conflict-free); the ONLY barrier is +// the single END-of-K __syncthreads() before the fused partial combine + +// scale -> bf16 epilogue (all 128 threads, 4 elements each). Barriers per +// K step = 0. +// * Correctness: per element the int32 accumulation is ascending over each +// wave's own stage sequence (the two waves own interleaved K=128 stages), +// then the two waves' partials are summed in LDS +// (order-independent int32, no overflow: max |dot| = 2048*127*127 << +// 2^31). The packed layout is coherent across the pack kernel, the DUMMA +// loads and the scalar fallback decode (same (tile, k_step, lane, byte) +// -> logical (k, n) mapping), so the paired M=2 validation stays exact. // * Graph-safe: static smem only, no allocation/sync/default-stream launch, -// pure dispatch on the caller-provided stream. Workspace untouched. +// pure dispatch on the caller-provided stream. Workspace untouched. // --------------------------------------------------------------------------- __global__ __launch_bounds__(kDummaWavesPerBlock * kDummaBlockThreads) void w8a8_gemm_m16_dumma_packedb_multint_kernel( const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, // packed [n_tile_pair][k_step][lane][16] + const int8_t* __restrict__ weight, // packed group-major B fragment slots const float* __restrict__ x_scale, const float* __restrict__ weight_scale, hip_bfloat16* __restrict__ out, @@ -363,28 +340,29 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( du::dumma::du_fill_fragment(acc0, 0); du::dumma::du_fill_fragment(acc1, 0); - // Each n_tile_pair owns one contiguous packed region of k_steps * 1024 B - // (64 KiB: both tiles' fragments interleaved at 16 B per lane per k_step). - const int64_t k_steps = k / kDummaTileK; + // Group-major packed-B layout: one k_step_group (the kStepsPerGroup = 8 + // k_steps a K-loop iteration consumes) holds ALL n_tile_pairs, one + // kStepsPerGroup * kPackedPairBytes slot each, so the 128 blocks' concurrent + // 8 KiB windows tile one contiguous group_stride region (1 MiB at n = 4096) + // and consecutive loop iterations advance that dense region monotonically. + // b_base is this block's pair slot inside every group. + const int64_t n_pairs = static_cast(n) / (2 * kDummaTileN); + const int64_t group_stride = + n_pairs * kStepsPerGroup * kPackedPairBytes; const int64_t b_base = - static_cast(blockIdx.x) * k_steps * kPackedPairBytes; + static_cast(blockIdx.x) * kStepsPerGroup * kPackedPairBytes; - // Per-lane A fragment ownership is loop-invariant (du_mma.hpp lines - // 447-460): row = lane & 0xf, col = (lane >> 4) << 3. + // Per-lane A fragment ownership is loop-invariant: row = lane & 0xf, + // col = (lane >> 4) << 3. const int a_row = lane & 15; const int a_col = (lane >> 4) << 3; - // Iteration 13 (HIP-only consolidation): prefetch the four per-thread scale - // values at the kernel top, BEFORE the stage-0 fragment loads. The fused - // epilogue then has ZERO exposed global latency on the block's finishing - // tail: the same 4 vmem reads (x_scale[row], x_scale[row2], - // weight_scale[n0+col], weight_scale[n0+16+col] - each used twice) move - // from after the END-of-K __syncthreads() into the cold-start window, where - // their L2 round trip overlaps the first B-load DRAM latency. Same 4 loads, - // same values, bit-identical math - only the issue time moves (validated - // pattern: the accepted down_proj kernel prefetches its scales at kernel - // top too). Cost: +4 VGPR held live across the K loop (64 -> ~68, register - // file non-binding); scratch must stay 0. + // Prefetch the four per-thread scale values BEFORE the stage-0 fragment + // loads so the fused epilogue has zero exposed global latency on the block's + // finishing tail: the four vmem reads (x_scale[row], x_scale[row2], + // weight_scale[n0+col], weight_scale[n0+16+col] - each used twice) move into + // the cold-start window, where their L2 round trip overlaps the first + // B-load DRAM latency. Same 4 loads, same values, bit-identical math. const int row = tid >> 4; // epilogue rows 0..7 (and row2 = +8) const int col = tid & 15; // epilogue column within the tile const int row2 = row + (kDummaTileM >> 1); @@ -393,31 +371,45 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( const float ws_lo = weight_scale[n0 + col]; const float ws_hi = weight_scale[n0 + kDummaTileN + col]; - // Four 1 KiB row-major int32 partial planes: [wave][tile][256]. - __shared__ __align__(16) int32_t acc_planes[kDummaWavesPerBlock * - kDummaTilesPerBlock * - kDummaTileM * kDummaTileN]; + // Four bank-conflict-free int32 partial planes: [wave][tile][16 rows x + // kCombineStride] (16 x 24 int32 = 1.5 KiB per plane, 6 KiB static LDS per + // block). The padded row stride is the only change versus the natural + // 16-column layout: the same fragment values are stored and the same 2-wave + // order-independent int32 sums are read back, so the result is bit-identical. + __shared__ __align__(16) int32_t + acc_planes[kDummaWavesPerBlock * kDummaTilesPerBlock * kDummaTileM * + kCombineStride]; // 16 stages of K=128 partitioned UNIFORMLY across the 2 waves (8 stages of // K=1024 per wave); ascending-K order per wave and the byte-exact fragment - // loads (A dwordx2, B dwordx4 tile-pair) are unchanged, the end-of-K int32 - // combine is order-independent, so the result is bit-identical to - // iteration 4. + // loads are unchanged, and the end-of-K int32 combine is order-independent. + // + // Stage ownership is INTERLEAVED: wave w owns stages w, w + 2, ..., so the + // two waves of a block walk ADJACENT K=128 stages (4 KiB apart) in lockstep + // instead of two 32 KiB-apart streams. Combined with the group-major packed + // layout, at loop iteration s the block reads its 8 KiB window + // [s * group_stride, (s+1) * group_stride) at pair-slot offset, and because + // blockIdx.x tiles those slots the WHOLE GRID reads one contiguous dense + // group_stride region per iteration that advances monotonically - a single + // dense DRAM/L2 stream instead of 128 windows 64 KiB apart. Every wave still + // covers exactly 8 stages = K=1024 and the per-wave k order is still + // ascending; the bytes, the addresses (as a permutation) and the LDS combine + // are unchanged, and the int32 sums are order-independent. const int total_stages = k / kStageK; // 16 - const int stage_start = (wave * total_stages) / kDummaWavesPerBlock; - const int stage_end = ((wave + 1) * total_stages) / kDummaWavesPerBlock; - const int n_stages = stage_end - stage_start; // 8 stages of K=128 - const int64_t k0_base = static_cast(stage_start) * kStageK; - - // Prologue: stage-0 fragments land directly in registers (one cold-start - // DRAM latency per wave, amortized over 8 stages and staggered across the - // 2 waves). + const int n_stages = total_stages / kDummaWavesPerBlock; // 8 stages of K=128 + const int64_t k0_base = static_cast(wave) * kStageK; + const int64_t k0_stage_stride = + static_cast(kDummaWavesPerBlock) * kStageK; + + // Prologue: the wave's first stage fragments land directly in registers (one + // cold-start DRAM latency per wave, amortized over 8 stages and staggered + // across the 2 waves). uint64_t ca0, ca1, ca2, ca3; // current stage's A uint64_t cb00, cb01, cb02, cb03; // current stage's B, tile 0 uint64_t cb10, cb11, cb12, cb13; // current stage's B, tile 1 load_stage_fragments_t2(x_q, weight, lane, a_row, a_col, b_base, k, k0_base, - 0, ca0, ca1, ca2, ca3, cb00, cb01, cb02, cb03, cb10, - cb11, cb12, cb13); + k0_stage_stride, group_stride, 0, ca0, ca1, ca2, ca3, + cb00, cb01, cb02, cb03, cb10, cb11, cb12, cb13); for (int s = 0; s < n_stages; ++s) { // Issue the NEXT stage's eight global loads (4 A dwordx2 + 4 B dwordx4) @@ -429,8 +421,9 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( uint64_t nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13; if (s + 1 < n_stages) { load_stage_fragments_t2(x_q, weight, lane, a_row, a_col, b_base, k, - k0_base, s + 1, na0, na1, na2, na3, nb00, nb01, - nb02, nb03, nb10, nb11, nb12, nb13); + k0_base, k0_stage_stride, group_stride, s + 1, + na0, na1, na2, na3, nb00, nb01, nb02, nb03, nb10, + nb11, nb12, nb13); } // Stage-s MMAC burst straight from registers, in the exact a_frag.x / // b_frag.x byte order du_mma_sync consumes: the same A fragment feeds @@ -449,9 +442,9 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( du::dumma::du_mma_sync(acc0, a_frag, b_frag0, acc0); du::dumma::du_mma_sync(acc1, a_frag, b_frag1, acc1); } - // Rotate: the prefetched stage s+1 becomes the current stage. No LDS - // WAR hazard; the vmcnt wait for the prefetch is satisfied before the - // next iteration's first use. + // Rotate: the prefetched stage s+1 becomes the current stage. No LDS WOR + // hazard; the vmcnt wait for the prefetch is satisfied before the next + // iteration's first use. if (s + 1 < n_stages) { ca0 = na0; ca1 = na1; @@ -469,36 +462,41 @@ w8a8_gemm_m16_dumma_packedb_multint_kernel( } // Combine the two wavefronts' int32 partials in LDS (exact int32 sums, - // order independent) after a single barrier at the END of K - barrier - // count per K step stays 0. Each wave stores its two tiles' planes. + // order independent) after a single barrier at the END of K - barrier count + // per K step stays 0. Each wave stores its two tiles' planes with the + // padded kCombineStride leading dimension, which spreads the fragment lanes + // over all 32 banks instead of 8. du::dumma::du_store_matrix_sync( - acc_planes + (wave * kDummaTilesPerBlock) * kDummaTileM * kDummaTileN, - acc0, kDummaTileN, du::dumma::mem_row_major); + acc_planes + (wave * kDummaTilesPerBlock) * kDummaTileM * kCombineStride, + acc0, kCombineStride, du::dumma::mem_row_major); du::dumma::du_store_matrix_sync( acc_planes + (wave * kDummaTilesPerBlock + 1) * kDummaTileM * - kDummaTileN, - acc1, kDummaTileN, du::dumma::mem_row_major); + kCombineStride, + acc1, kCombineStride, du::dumma::mem_row_major); __syncthreads(); - // Cooperative epilogue over the two 16x16 tiles: all 128 threads, 4 - // elements each - for tile t: (row = tid >> 4, col = tid & 15) for rows - // 0..7 and (row + 8, col) for rows 8..15; sum the 2 planes per element - // (int32 add is order-independent), then fused dot * x_scale[row] * - // weight_scale[col] -> bf16 store. The four scale values were prefetched - // at the kernel top (iteration 13), so no global load sits on this tail; - // iteration 17 replaces the four __float2bfloat16 calls with the finite - // RNE-only bf16_rne_finite (same bits for every finite input, no NaN/Inf - // branch), shortening each element's dependent tail chain. + // Cooperative epilogue over the two 16x16 tiles: all 128 threads, 4 elements + // each - for tile t: (row = tid >> 4, col = tid & 15) for rows 0..7 and + // (row + 8, col) for rows 8..15; sum the 2 planes per element (int32 add is + // order-independent), then fused dot * x_scale[row] * weight_scale[col] -> + // bf16 store. The four scale values were prefetched at the kernel top, and + // bf16_rne_finite replaces the branchy __float2bfloat16 conversion. int32_t t0_lo = 0, t0_hi = 0, t1_lo = 0, t1_hi = 0; #pragma unroll for (int w = 0; w < kDummaWavesPerBlock; ++w) { const int plane_base = - w * kDummaTilesPerBlock * kDummaTileM * kDummaTileN; - t0_lo += acc_planes[plane_base + tid]; - t0_hi += acc_planes[plane_base + tid + kDummaTileM * kDummaTileN / 2]; - t1_lo += acc_planes[plane_base + kDummaTileM * kDummaTileN + tid]; - t1_hi += acc_planes[plane_base + kDummaTileM * kDummaTileN + tid + - kDummaTileM * kDummaTileN / 2]; + w * kDummaTilesPerBlock * kDummaTileM * kCombineStride; + // Element (row, col) lives at row * kCombineStride + col; the second + // element of each thread is (row + 8, col), exactly 8 * kCombineStride = + // 192 dwords away (an integral ds_read2st64 offset), and tile 1's plane is + // one full 16 x kCombineStride plane further on. + const int lo = plane_base + row * kCombineStride + col; + const int hi = lo + (kDummaTileM / 2) * kCombineStride; + const int tile1 = kDummaTileM * kCombineStride; + t0_lo += acc_planes[lo]; + t0_hi += acc_planes[hi]; + t1_lo += acc_planes[lo + tile1]; + t1_hi += acc_planes[hi + tile1]; } const float s0_lo = static_cast(t0_lo) * x_scale_lo * ws_lo; out[static_cast(row) * n + n0 + col] = bf16_rne_finite(s0_lo); @@ -541,16 +539,19 @@ constexpr int kMaxCopyBlocks = 65535; // --------------------------------------------------------------------------- // One-time pack for the exact (k, n) == (2048, 4096) o_proj weight: permute -// the logical [K, N] int8 layout into packed[n_tile_pair][k_step][lane][16] -// B fragment slots - tile-pair-interleaved: for the same (k_step, lane), -// bytes 0..7 are the 8 fragment bytes of tile (2*pair) and bytes 8..15 the 8 -// fragment bytes of tile (2*pair+1) (see the file header and the kernel -// comment for the exact du_mma.hpp matrix_b row_major lane mapping: -// row = lane & 0xf, col = (lane >> 4) << 3, x[i] = p[(col + i) * ldm + row]). -// One thread per (pair, k_step, lane): 64 k_steps x 128 pairs x 64 lanes = -// 524,288 threads, each writing one aligned 16-byte slot. Runs once per -// weight, outside the timed GEMM and outside Graph capture; byte count is -// unchanged (8 MiB), so the packed tensor is a same-size permutation. +// the logical [K, N] int8 layout into the group-major +// packed[k_step_group][n_tile_pair][k_step_within][lane][16] B fragment slots - +// tile-pair-interleaved: for the same (k_step, lane), bytes 0..7 are the 8 +// fragment bytes of tile (2*pair) and bytes 8..15 the 8 fragment bytes of tile +// (2*pair+1) (du_mma.hpp matrix_b row_major lane mapping: row = lane & 0xf, +// col = (lane >> 4) << 3, x[i] = p[(col + i) * ldm + row]). k_step_group is +// k_step / kStepsPerGroup (the 8 k_steps one K-loop iteration consumes), so all +// n_tile_pairs of a group are adjacent: the grid's concurrent 8 KiB windows +// tile one dense monotonic group region. One thread per (pair, k_step, lane): +// 64 k_steps x 128 pairs x 64 lanes = 524,288 threads, each writing one aligned +// 16-byte slot. Runs once per weight, outside the timed GEMM and outside Graph +// capture; byte count is unchanged (8 MiB), so the packed tensor is a same-size +// permutation. // --------------------------------------------------------------------------- __global__ __launch_bounds__(kScalarThreads) void w8a8_pack_o_proj_m16_kernel( const int8_t* __restrict__ raw, @@ -590,18 +591,23 @@ __global__ __launch_bounds__(kScalarThreads) void w8a8_pack_o_proj_m16_kernel( slot.y = static_cast(v0 >> 32); slot.z = static_cast(v1 & 0xffffffffu); slot.w = static_cast(v1 >> 32); + // Group-major destination: [k_step_group][n_tile_pair][k_step_within][lane]. + const int64_t n_pairs = static_cast(n) / (2 * kDummaTileN); + const int group = k_step / kStepsPerGroup; + const int within = k_step % kStepsPerGroup; *reinterpret_cast( - packed + (static_cast(pair) * (k / kDummaTileK) + k_step) * + packed + static_cast(group) * n_pairs * kStepsPerGroup * kPackedPairBytes + - lane * 16) = slot; + static_cast(pair) * kStepsPerGroup * kPackedPairBytes + + static_cast(within) * kPackedPairBytes + lane * 16) = slot; } } // namespace // Stable host launch symbol used by the trusted binding for -// torch.ops.zth_w8a8.gemm_out. Runs entirely on the caller-provided stream: +// torch.ops.zth_w8a8.gemm_out. Runs entirely on the caller-provided stream: // no allocation, no synchronization, no default-stream launch, no workspace -// writes. Returns nothing; the binding returns the caller-provided out. +// writes. Returns nothing; the binding returns the caller-provided out. extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -621,12 +627,11 @@ extern "C" void launch_w8a8_gemm( if (m == kTargetM && n == kTargetN && k == kTargetK) { // Exact assigned shape hy3_tp4_o_proj_m16: packed-B two-wave in-block - // split-K=2 DUMMA m16n16k32 kernel with MULTI-N-TILE REUSE (128 blocks, + // split-K=2 DUMMA m16n16k32 kernel with multi-N-tile reuse (128 blocks, // two adjacent 16x16 N tiles per block, 2 x 64-thread wavefronts, K=128 - // stages, depth-1 register prefetch). The scalar fallback below stays + // stages, depth-1 register prefetch). The scalar fallback below stays // untouched and still covers every unmatched shape, including the paired - // M=2 shape with the same (N, K) (the guard is shape-exact on m too; the - // fallback decodes the packed weight for (n, k) == (4096, 2048)). + // M=2 shape with the same (N, K). const dim3 grid(static_cast( kTargetN / (kDummaTileN * kDummaTilesPerBlock))); const dim3 block(kDummaBlockThreads * kDummaWavesPerBlock); @@ -636,7 +641,7 @@ extern "C" void launch_w8a8_gemm( return; } - // Generic scalar fallback for every unmatched (m, n, k). This also covers + // Generic scalar fallback for every unmatched (m, n, k). This also covers // the paired M=2 API shape that shares the same (N, K): it must never be // routed through an M=16-only specialization. const dim3 grid( @@ -648,10 +653,11 @@ extern "C" void launch_w8a8_gemm( } // Stable host launch symbol used by the optional out-of-timed-region -// torch.ops.zth_w8a8.pack_weight. For the exact (k, n) == (2048, 4096) o_proj -// weight this permutes into the packed-B fragment-slot layout (one-time, -// untimed, outside Graph capture); every other (K, N) keeps the generic -// identity device-to-device copy. Weight scales are always copied through. +// torch.ops.zth_w8a8.pack_weight. For the exact (k, n) == (2048, 4096) +// o_proj weight this permutes into the packed-B fragment-slot layout +// (one-time, untimed, outside Graph capture); every other (K, N) keeps the +// generic identity device-to-device copy. Weight scales are always copied +// through. extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -661,8 +667,9 @@ extern "C" void launch_pack_w8a8_weight( int n, hipStream_t stream) { if (k == kTargetK && n == kTargetN) { - // Exact o_proj M=16 weight: permute into [n_tile_pair][k_step][lane][16] - // B fragment slots (one-time, untimed, outside Graph capture). + // Exact o_proj M=16 weight: permute into the group-major + // [k_step_group][n_tile_pair][k_step_within][lane][16] B fragment slots + // (one-time, untimed, outside Graph capture). const int64_t total = (static_cast(k) / kDummaTileK) * (n / (2 * kDummaTileN)) * kDummaBlockThreads; const dim3 grid( diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/qkv_proj.hip index 6d9acfa9..c42c939c 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/qkv_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/qkv_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=hy3_tp4_qkv_proj_m16 commit=7d6959345400cb268cfd9289338d0891db2297f9 added=2026-08-26 +// baseline_us=80.19 speedup=3.07831 // median_us=26.05 p90_us=26.09 // source=hy3-dsh-tp4-m16-1-7f1fb1d1 // MetaInfer W8A8 INT8 GEMM - worker_0 (iteration 17, fused combine tail). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_down_proj.hip index 4a5f7a29..456eadd0 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_down_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_down_proj.hip @@ -1,53 +1,100 @@ -// @@variant shape=hy3_tp4_shared_down_proj_m16 commit=8f7319d888b0772736fa86ed3dc42afb6b56aa7b added=2026-08-26 -// median_us=9.174 p90_us=9.754 -// source=hy3-dsh-tp4-m16-2-368c654c -// INT8 W8A8 GEMM - worker_3, round 1 (DUMMA bootstrap). +// @@variant shape=hy3_tp4_shared_down_proj_m16 commit=49c8cf883f92c0d9a791c78c94e84604b8eafc17 added=2026-09-10 +// median_us=8.489 p90_us=8.888 speedup=2.735 baseline_us=23.22 +// source=test-9-8-8-803e5cee/iteration_001/hy3_tp4_shared_down_proj_m16 +// INT8 W8A8 GEMM implementation for Hygon K500SM_AI / gfx928. // -// Assigned shape (worker_3, physical GPU 3): -// hy3_tp4_shared_down_proj_m16 : M=16, N=4096, K=384 +// Parallel-explore child worker_3, physical GPU 3. +// Assigned shape: hy3_tp4_shared_down_proj_m16 (M=16, N=4096, K=384). // -// Mandated architecture: minimal native gfx928 DUMMA INT8 m16n16k32 tile, -// ONE 64-lane wavefront per block, one 16x16 output tile per block -// (grid = N/16 = 256 blocks), explicit int32 accumulation, NO cross-wave -// barrier anywhere (single wave per block: no LDS, no __syncthreads). +// Iteration 11 = HIP-only occupancy-consolidation round. The accepted +// iteration-3 kernel (median 9.092 us, one 64-lane wave per 16x16 tile, +// grid 256 = 256 wavefronts over the measured 120 CUs -> only ~2.13 resident +// waves/CU, i.e. ~0.53 waves per SIMD) is per-wave memory-latency bound: +// iteration 2 doubled resident wavefronts with in-block split-K=2 and moved +// nothing (0.148%, 512 waves = only ~1.07 waves/SIMD - still no inter-wave +// latency hiding), iteration 4 halved the grid to 128 (two N-tiles per wave +// sharing A) and regressed -25.1% (12.140 us), iteration 5 replaced the +// per-step B global waits with a block-head B-only LDS stage and regressed +// -18.97% (11.220 us), iteration 7 split the accumulator into two independent +// 6-deep v_mmac chains (regressed -2.67%, 9.341 us - the serial mma chain is +// NOT the binding constraint), and iteration 10 reordered all 12 cold-B loads +// into a one-burst register preload and regressed -38% (14.719 us - a +// front-loaded load burst starves the interleaved schedule). The fresh +// current-best ISA (arch_vgpr 24, sgpr 32, 0 scratch, 0 LDS, grid 256) shows +// each wave issues ~220 VALU + 24 global_load_dwordx2 + 12 v_mmac around 27 +// s_waitcnt and that the packed [N][K] B stream is cold per replay; with only +// ~0.5 resident wavefronts per SIMD each wave's per-step B-load vmcnt wait is +// exposed with no co-resident wave to fill it. // -// * A (x_q, logical [16][384] row-major) is loaded with -// du::dumma::du_load_matrix_sync(matrix_a, row_major): each lane's 8 -// fragment bytes A[row][k0 + i] (row = lane&15, k0 = (lane>>4)*8) are -// contiguous in global memory, so one 8-byte load per lane per step. -// * B (weight, logical [384][4096] row-major) is packed ONCE outside the -// timed region by launch_pack_w8a8_weight for the exact (k,n) == -// (384,4096) pair into the [N][K] n-major transpose -// P[n*384 + k] = W[k*4096 + n]. -// The B fragment is loaded with du::dumma::du_load_matrix_sync( -// matrix_b, col_major) over P: lane (row = n column, col = 8-k group) -// reads 8 contiguous bytes P[(n0+row)*384 + k0 + (lane>>4)*8 + i], i.e. -// one 8-byte load per lane per step instead of the strided byte-load -// expansion of a row-major [K][N] B fragment. -// * K=384 = 12 fully-unrolled m16n16k32 steps accumulate into one int32 -// accumulator fragment (du_fill_fragment + du_mma_sync). -// * Fused direct epilogue using the verified gfx928 int8 m16n16k32 -// accumulator ownership (lane%16 = row, lane/16 = col%4, x[i] = columns -// col%4 + 4*i): each lane scales its four int32 accumulators by -// (float(dot) * x_scale[row]) * weight_scale[col] in the same -// left-associative float32 order as the harness reference, then stores -// round-to-nearest-even bf16 via __float2bfloat16. +// This round consolidates the two best-measured mechanisms into one kernel: +// the iteration-3 packed-B fast path (direct col_major 8-byte fragment loads, +// exact int32 accumulation) is launched with the iteration-2 in-block split-K +// combine (validated, bit-exact LDS int32 plane merge) raised to split-K=4 +// from the trusted occupancy-probe set {2,3,4,6} that divides the 12-step +// K=384 loop: blockDim 256 = four 64-lane wavefronts, grid 256, each wave +// accumulates one K quarter (96 = 3 m16n16k32 steps, K ascending) and one +// __syncthreads() orders the four int32 partial planes before all 256 lanes +// sum them in ascending-K order and store coalesced bf16. 256 blocks x 4 +// waves = 1024 resident wavefronts (~8.5 waves/CU = ~2.1 waves/SIMD) so each +// SIMD can overlap one wave's B-load/mma stalls with a co-resident wave's +// work, while per-wave serial K depth drops 12 -> 3 steps. HBM traffic is +// unchanged (each block's [16 n x 384 k] B pack slice is still read exactly +// once, A is read once per tile from L2, the 4 KiB partial planes stay +// on-chip), the epilogue store becomes fully coalesced row-major, and the +// int32 result is bit-identical to iteration 3 (no overflow: each partial is +// bounded by 96*127*127 << 2^31, and the quarter sums are combined in +// ascending-K wave order), so the scaled bf16 output stays bit-identical. // -// Exact int32 accumulation order (k-ascending, no split-K) is preserved, so -// the bf16 output is expected bit-identical to the scalar bootstrap. +// Iteration 3 = data-layout + epilogue round. The accepted iteration-1/2 +// one-wave and two-wave split-K=2 DUMMA kernels (median 21.656 / 21.624 us) +// both load the B fragment directly from the *logical row-major [K][N]* +// weight with du_load_matrix_sync(matrix_b, row_major). In that layout each +// lane's 8 fragment bytes are strided by N, so the compiler expands every B +// fragment load into ~8 scalar `global_load_ubyte` instructions plus +// reassembly, with a vmcnt wait burst before each v_mmac (PMC: 50,432 vmem +// reads / 27.36 us profiled; static ISA shows 16+ ubyte loads + 13-16 +// waitcnts per K step). Iteration-2 measured that doubling resident +// wavefronts via in-block split-K=2 changes nothing (0.148%) -> the binding +// constraint is per-wave load vectorization / issue cost, NOT grid +// parallelism. This round removes both remaining structural costs: // -// Dispatch: -// * exact-shape guard (m == 16 && n == 4096 && k == 384) -> DUMMA arm; -// * every other (m, n, k) - including the paired M=2 API shape with the -// same (N, K) - reaches the generic scalar fallback, which decodes the -// n-major pack for (n, k) == (4096, 384) via a b_transposed flag and -// keeps the identity row-major read for every other pair. +// * B is packed ONCE, outside the timed region, by launch_pack_w8a8_weight +// for the exact (k, n) == (384, 4096) pair into the [N][K] n-major +// transpose P[n*384 + k] = W[k*4096 + n]. The M=16 DUMMA arm then loads +// B fragments with du_load_matrix_sync(matrix_b, col_major) over P: each +// lane's 8 fragment bytes are contiguous in P -> one 8-byte load per lane +// per m16n16k32 step instead of the strided byte-load expansion. +// * The epilogue drops the LDS int32 materialization + __syncthreads + +// LDS read-back entirely and instead uses the verified gfx928 int8 +// m16n16k32 accumulator ownership (lane%16 = row, lane/16 = col%4, +// frag.x[i] = column col%4 + 4*i): every lane fuses the fp32 scale and +// bf16 conversion directly on its four register accumulators. Kernel +// static LDS becomes 0 B and there are no barriers anywhere. // -// Header order is the known-good DTK order: HIP runtime, then bfloat16, -// then du_mma.h (du_mma.h is not self-contained before the HIP runtime -// headers). gfx928 wavefront is 64 lanes; every block size is a multiple -// of 64. +// Geometry stays the proven iteration-1 shape: one 64-lane wavefront owns one +// 16x16 output tile, grid = N/16 = 256 blocks, explicit int32 accumulation +// over K ascending in 12 fully-unrolled m16n16k32 steps (bit-exact int8 dot +// -> bit-identical bf16 output, 0 mismatches expected), no cross-wave +// barrier. The scalar kernel below is retained as the generic fallback for +// every unmatched (m, n, k), especially the paired M=2 API shape with the +// same (N, K); it decodes the n-major pack for (k, n) == (384, 4096) via a +// b_transposed flag and reads the identity row-major [K][N] layout for every +// other pair. Shape-specialized launches stay behind exact (m, n, k) +// guards. +// +// Two stable host symbols are provided (declared in csrc/bindings.cpp): +// launch_w8a8_gemm(...) -- timed / CUDA-Graph-captured GEMM +// launch_pack_w8a8_weight(...) -- out-of-timed-region weight packing +// The pack is the [N][K] n-major transpose for the exact (k, n) == (384,4096) +// pair and an identity device copy otherwise; packing never occurs inside +// gemm_out. +// +// Graph-safety: gemm_out launches only on the caller-provided stream, touches +// only caller-provided tensors/workspace, and performs no allocation, +// compilation, autotuning, packing, or host/device synchronization. +// Known-good include order for this DTK: hip_runtime.h must precede +// hip_bfloat16.h and du_mma.h (du_mma.h is not self-contained otherwise). #include #include #include @@ -56,39 +103,164 @@ namespace { -// One block per 256 threads; 256 is a multiple of the gfx928 wavefront (64). +// gfx928 wavefront is 64 lanes; block size must be a multiple of 64. +constexpr int kWaveSize = 64; constexpr int kScalarBlockThreads = 256; -// DUMMA INT8 m16n16k32 tile constants for the M=16 exact-shape arm. -constexpr int kDummaTileM = 16; -constexpr int kDummaTileN = 16; -constexpr int kDummaTileK = 32; -constexpr int kDummaBlockThreads = 64; // one wavefront per block -constexpr int kDummaK = 384; // exact-guard K; 12 m16n16k32 steps +// DUMMA INT8 hardware tile consumed by one wavefront per mma step. +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; -// Exact pack guard for the assigned (K, N) pair: (384, 4096). +// Exact guard constants for the assigned (K, N) pair: (384, 4096). constexpr int kPackK = 384; constexpr int kPackN = 4096; // --------------------------------------------------------------------------- -// Generic scalar fallback: one thread per output element. -// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) -// * x_scale[m] * weight_scale[n] ) -// x_q is [M, K] int8 row-major (stride K). -// weight is the packed [K, N] int8 buffer (stride N in the identity layout). -// b_transposed selects the exact-shape packed [N][K] layout -// (P[n*K + k] = W[k*N + n], produced by launch_pack_w8a8_weight for -// (k,n) == (384,4096)): the logical column col starts at P[col*K] with unit -// stride; the identity pack keeps the legacy row-major [K][N] layout with -// stride n. The branch is hoisted out of the K loop by the compiler. -// Consecutive threads own consecutive N columns, so adjacent lanes touch -// adjacent addresses in the fastest-changing N dimension. -__global__ __launch_bounds__(kScalarBlockThreads) void w8a8_scalar_gemm_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - __hip_bfloat16* __restrict__ out, +// gfx928 DUMMA INT8 m16n16k32 in-block split-K=4 kernel (assigned-shape fast +// path, iteration 11). +// +// acc[m, n] = sum_{kk=0}^{k-1} int32(a[m* k + kk]) * int32(b[kk*n + n0+n]) +// out[m, n0+n] = bf16( float(acc[m, n]) * x_scale[m] * weight_scale[n0+n] ) +// +// Block geometry: blockDim = 256 (four 64-lane wavefronts), grid = n/16 +// output tiles, split-K = 4 along K. Wave w accumulates the K quarter +// [w*96, (w+1)*96) in ascending kTileK steps (K=384 -> 96 per wave = 3 +// du_mma_sync steps), reading A row_major from the logical x_q and B +// col_major from the exact-shape [N][K] n-major pack exactly as the accepted +// iteration-3 one-wave kernel does at the same k offsets (one contiguous +// 8-byte global load per lane per step). All four waves then publish their +// int32 partial to LDS with du_store_matrix_sync(mem_row_major), a single +// __syncthreads() orders the planes, and every one of the 256 lanes sums the +// four partials in ascending-K wave order (partial 0 = K [0,96), ..., partial +// 3 = K [288,384) -> bit-exact full-K int8 dot, identical to the iteration-3 +// 12-step chain because no intermediate sum overflows int32) and writes its +// one scaled bf16 element with a coalesced row-major store. No barrier +// exists inside the K loop; the sole __syncthreads() is reached by every +// block thread between the partial store and the combine. +// +// Occupancy consolidation: 256 blocks x 4 waves = 1024 resident wavefronts +// (~8.5 waves/CU ~= 2.1 waves/SIMD on the measured 120 CUs) instead of the +// iteration-3 ~0.53 waves/SIMD, so a wave's per-step cold-B vmcnt waits can +// be overlapped by the co-resident wave on the same SIMD, and each wave's +// serial mma depth is 3 steps instead of 12. The 4 KiB LDS partial planes +// are on-chip (no HBM traffic change) and the kernel stays one launch, one +// barrier, graph-safe, behind the exact (m, n, k) guard. +// --------------------------------------------------------------------------- +constexpr int kSplitK = 4; // trusted occupancy-probe split that divides the + // 12-step K loop and reaches ~2 waves/SIMD +constexpr int kStepsPerWave = (kPackK / kTileK) / kSplitK; // 12 / 4 = 3 + +__global__ __launch_bounds__(kSplitK * kWaveSize) void +w8a8_dumma_m16n16k32_sk4_kernel( + const int8_t* __restrict__ a, // [16, 384] row-major int8 activations + const int8_t* __restrict__ b, // packed [4096, 384] n-major int8 weight + const float* __restrict__ x_scale, // [16] fp32 per-row scale + const float* __restrict__ weight_scale, // [4096] fp32 per-col scale + __hip_bfloat16* __restrict__ out, // [16, 4096] bf16 result + int n) { + const int tid = static_cast(threadIdx.x); // 0..255 (four 64-lane waves) + const int wave = tid / kWaveSize; // hardware wave id: 0..3 + const int n0 = static_cast(blockIdx.x) * kTileN; // this tile's columns + + __shared__ __align__(16) int32_t part[kSplitK][kTileM * kTileN]; // 4 KiB + + using namespace du::dumma; + DUFragment a_frag; + DUFragment b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + // The col_major B loader adds (lane%16)*ldm internally, so the pointer is + // the tile base only: P[n0*384 + k + (lane%16)*384 + (lane>>4)*8 + i], + // i.e. each lane's 8 fragment bytes are contiguous in P (8-byte aligned) -> + // one 8-byte global load per lane per step. This launch is guarded to the + // exact (m, n, k) == (16, 4096, 384) shape, so each wave's K quarter + // [wave*96, wave*96 + 96) is a whole number of kTileK steps. + const int64_t b_tile_base = static_cast(n0) * kPackK; + const int kStart = wave * kStepsPerWave * kTileK; + + // K quarter, ascending, exactly as the iteration-3 one-wave kernel at the + // same k offsets (per-wave partial is bit-identical to that chain's prefix). +#pragma unroll + for (int s = 0; s < kStepsPerWave; ++s) { + const int koff = kStart + s * kTileK; + // a + koff: 16 A rows (stride 384), columns = K slice [koff, koff + 32). + du_load_matrix_sync(a_frag, a + koff, kPackK); + du_load_matrix_sync(b_frag, b + b_tile_base + koff, kPackK); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Publish this wave's int32 partial for K quarter [kStart, kStart + 96) + // (ascending) into its own LDS plane. Row-major [16][16] plane layout with + // ldm = 16 matches du_store_matrix_sync(mem_row_major) exactly as in the + // accepted split-K=2 lineage. + du_store_matrix_sync(&part[wave][0], acc_frag, kTileN, mem_row_major); + __syncthreads(); + + // All 256 lanes combine the four partials in ascending split order (wave 0 + // = K [0, 96) ... wave 3 = K [288, 384) -> bit-exact full-K int32 dot) and + // emit one scaled bf16 element each with a coalesced row-major store: + // adjacent 16-lane groups write consecutive columns of one output row. + const int row = tid / kTileN; // 0..15 + const int col = tid - row * kTileN; // 0..15 + const int32_t acc = + part[0][tid] + part[1][tid] + part[2][tid] + part[3][tid]; + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[n0 + col]; + out[static_cast(row) * n + n0 + col] = __float2bfloat16(scaled); +} + +void launch_dumma_m16n16_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + int n, + hipStream_t stream) { + const unsigned int blocks = static_cast(n / kTileN); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_dumma_m16n16k32_sk4_kernel), + dim3(blocks), + dim3(kSplitK * kWaveSize), + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast<__hip_bfloat16*>(out), + n); +} + +// --------------------------------------------------------------------------- +// Scalar int8 dot-product kernel (generic fallback for every unmatched shape). +// +// acc(row,col) = sum_{kk=0}^{k-1} int32(a[row*k + kk]) * int32(b[kk*n + col]) +// out[row*n + col] = bf16( (float)acc * x_scale[row] * weight_scale[col] ) +// +// x_scale is the [M,1] fp32 per-token scale (contiguous -> M floats); +// weight_scale is the [N,1] fp32 per-output-channel scale (-> N floats). +// b is the packed buffer produced by launch_pack_w8a8_weight: for the exact +// (k, n) == (384, 4096) pair it is the [N][K] n-major transpose +// P[n*384 + k] = W[k*4096 + n] (logical column col starts at P[col*384] with +// unit stride); every other pair keeps the identity row-major [K][N] layout +// (logical column col has stride n). The b_transposed branch is hoisted out +// of the K loop by the compiler. +// +// The maximum assigned K keeps the exact int8 dot well inside int32 range; +// int32 -> float is exact here, so converting only at the very end after the +// two fp32 scale multiplies matches the reference +// (A.float() @ B.float()) * x_scale * weight_scale.T -> bfloat16. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kScalarBlockThreads) void +w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 activations + const int8_t* __restrict__ b, // packed weight buffer (see above) + const float* __restrict__ x_scale, // [M] fp32 per-row scale + const float* __restrict__ weight_scale, // [N] fp32 per-col scale + __hip_bfloat16* __restrict__ out, // [M, N] bf16 result int m, int n, int k, @@ -97,92 +269,67 @@ __global__ __launch_bounds__(kScalarBlockThreads) void w8a8_scalar_gemm_kernel( static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const int64_t total = static_cast(m) * n; if (linear >= total) { - return; + return; // tail guard for shapes where M*N is not a block multiple } - const int row = static_cast(linear / n); const int col = static_cast(linear - static_cast(row) * n); - const int8_t* a_row = x_q + static_cast(row) * k; - const int8_t* b_col = - weight + (b_transposed ? static_cast(col) * k : col); + const int8_t* __restrict__ a_row = a + static_cast(row) * k; + const int8_t* __restrict__ b_col = + b + (b_transposed ? static_cast(col) * k : col); const int64_t b_stride = b_transposed ? 1 : static_cast(n); - // Exact int32 accumulation over the full K loop. The maximum assigned K - // keeps the exact int8 dot within int32 range (K * 127 * 127 << 2^31). int32_t acc = 0; for (int kk = 0; kk < k; ++kk) { acc += static_cast(a_row[kk]) * static_cast(b_col[static_cast(kk) * b_stride]); } - // Convert to float only for the two scales, then store bf16. - const float scaled = static_cast(acc) * x_scale[row] * - weight_scale[col]; + // Scale in fp32 exactly as the reference does (dot first, then x_scale, + // then weight_scale), then store bf16. + const float scaled = static_cast(acc) * + x_scale[row] * weight_scale[col]; out[linear] = __float2bfloat16(scaled); } -// --------------------------------------------------------------------------- -// Exact-shape M=16 DUMMA arm (mandated minimal geometry): blockDim = 64 = -// one 64-lane wavefront, one 16x16 output tile per block, grid = N/16 = 256 -// blocks. Direct global fragment loads (A row_major from logical x_q, B -// col_major from the exact-shape [N][K] n-major pack), explicit int32 -// accumulation over the full K=384 (12 unrolled m16n16k32 steps), fused -// direct scale+bf16 epilogue. No LDS and no barriers anywhere. Only launched -// for (m, n, k) == (16, 4096, 384). -// --------------------------------------------------------------------------- -__global__ __launch_bounds__(kDummaBlockThreads) void -w8a8_dumma_m16n16k32_1wave_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ packed_b, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - __hip_bfloat16* __restrict__ out, - int n) { - constexpr int kSteps = kDummaK / kDummaTileK; // 12 - - const int lane = static_cast(threadIdx.x); // 0 .. 63 - const int n0 = static_cast(blockIdx.x) * kDummaTileN; - - du::dumma::DUFragment - a_frag; - du::dumma::DUFragment - b_frag; - du::dumma::DUFragment - acc_frag; - du::dumma::du_fill_fragment(acc_frag, 0); - - // The col_major B loader adds (lane%16)*ldm internally, so the pointer is - // the tile base only: P[n0*384 + s*32 + (lane%16)*384 + (lane>>4)*8 + i], - // i.e. each lane's 8 fragment bytes are contiguous in P (8-byte aligned). - const int64_t b_tile_base = static_cast(n0) * kDummaK; - -#pragma unroll - for (int s = 0; s < kSteps; ++s) { - du::dumma::du_load_matrix_sync(a_frag, x_q + s * kDummaTileK, kDummaK); - du::dumma::du_load_matrix_sync( - b_frag, packed_b + b_tile_base + s * kDummaTileK, kDummaK); - du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - - // gfx928 int8 m16n16k32 accumulator ownership, established against - // du_store_matrix_sync: lane%16 selects the row, lane/16 selects col%4, - // and x[i] holds the columns col%4 + 4*i. - const int row = lane & 15; - const int col_mod4 = lane >> 4; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int out_col = n0 + col_mod4 + 4 * i; - const float scaled = static_cast(acc_frag.x[i]) * - x_scale[row] * weight_scale[out_col]; - out[static_cast(row) * n + out_col] = - __float2bfloat16(scaled); +void launch_scalar_gemm( + const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + int m, + int n, + int k, + int b_transposed, + hipStream_t stream) { + const int64_t total = static_cast(m) * static_cast(n); + if (total == 0) { + return; } + const unsigned int blocks = static_cast( + (total + kScalarBlockThreads - 1) / kScalarBlockThreads); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_scalar_kernel), + dim3(blocks), + dim3(kScalarBlockThreads), + 0, + stream, + a, + b, + x_scale, + weight_scale, + static_cast<__hip_bfloat16*>(out), + m, + n, + k, + b_transposed); } +// --------------------------------------------------------------------------- +// Packing helpers (run out of the timed region / Graph). +// --------------------------------------------------------------------------- + // Simple device-to-device copy used by the identity packing (every (K, N) // pair that is not the exact-shape transpose pair, plus all scales). template @@ -199,9 +346,9 @@ w8a8_identity_copy_kernel(const T* __restrict__ src, // Exact-shape pack: writes the [N][K] transpose // dst[n*K + k] = src[k*N + n] -// for (k, n) == (384, 4096). One thread per 16-byte k-chunk of one column; +// for (k, n) == (384, 4096). One thread per 16-byte k-chunk of one column; // runs once per weight outside the timed region/Graph, so the strided byte -// loads are acceptable. All other shapes keep the identity copy kernel. +// loads are acceptable. All other shapes keep the identity copy kernel. __global__ __launch_bounds__(kScalarBlockThreads) void w8a8_pack_b_transpose_kernel(const int8_t* __restrict__ src, int8_t* __restrict__ dst, @@ -209,7 +356,8 @@ w8a8_pack_b_transpose_kernel(const int8_t* __restrict__ src, int n) { const int kchunks = (k + 15) >> 4; const int64_t total = static_cast(n) * kchunks; - const int64_t stride = static_cast(gridDim.x) * blockDim.x; + const int64_t stride = + static_cast(gridDim.x) * static_cast(blockDim.x); for (int64_t t = static_cast(blockIdx.x) * blockDim.x + static_cast(threadIdx.x); t < total; t += stride) { @@ -227,12 +375,16 @@ w8a8_pack_b_transpose_kernel(const int8_t* __restrict__ src, } // namespace -// Optional out-of-timed-region weight packing. For the exact (k, n) == -// (384, 4096) pair it produces the [N][K] n-major transpose -// P[n*384 + k] = W[k*4096 + n] that the M=16 DUMMA arm (col_major B -// fragments) and the scalar fallback (b_transposed flag) both decode; every +// --------------------------------------------------------------------------- +// Optional out-of-timed-region weight packing entry point (registered as +// torch.ops.zth_w8a8.pack_weight through the trusted csrc/bindings.cpp). +// For the exact (k, n) == (384, 4096) pair it produces the [N][K] n-major +// transpose P[n*384 + k] = W[k*4096 + n] that the M=16 DUMMA arm (col_major +// B fragments) and the scalar fallback (b_transposed flag) both decode; every // other (K, N) keeps the identity layout, and all scales keep the identity -// copy. +// copy. Asynchronous on the caller's stream: no host/device +// synchronization. +// --------------------------------------------------------------------------- extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -267,9 +419,12 @@ extern "C" void launch_pack_w8a8_weight( weight_scale, packed_weight_scale, n); } -// Timed GEMM operator: no allocation, no compilation, no autotuning, no -// weight packing, no host/device synchronization, no default-stream launch. -// Uses only the caller-provided out and workspace, on the caller's stream. +// --------------------------------------------------------------------------- +// Timed GEMM entry point (registered as torch.ops.zth_w8a8.gemm_out through +// the trusted csrc/bindings.cpp). Must stay Graph-safe: launches only on the +// current PyTorch stream, no allocation/sync/packing/autotuning, and uses +// only caller-provided `out` and `workspace`. +// --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -282,40 +437,33 @@ extern "C" void launch_w8a8_gemm( int n, int k, hipStream_t stream) { - // The one-wave DUMMA arm performs no split-K, so the workspace is unused. + // No inter-block split-K in this architecture (the in-block split-K=4 + // partials are combined on-chip in LDS), so the workspace is unused; it is + // still accepted (and required to be preallocated) for API compatibility. (void)workspace; (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0) { + return; + } + // The exact (k, n) == (384, 4096) pack is the [N][K] n-major transpose; - // any other (k, n) keeps the identity pack (row-major [K][N]). + // any other (k, n) keeps the identity pack (row-major [K][N]). Decoding + // must follow the same rule in the scalar fallback below. const int b_transposed = (k == kPackK && n == kPackN) ? 1 : 0; - // Exact-shape M=16 DUMMA specialization: one 64-lane wavefront per 16x16 - // output tile -> grid = N/16 = 256 blocks, full K=384 reduction, no - // barriers. Guarded by the full (m, n, k) triple so the paired M=2 API - // shape with the same (N, K) still reaches the generic scalar fallback. + // Explicit assigned-shape dispatch point. The M=16 / N=4096 / K=384 + // assigned shape runs on the DUMMA m16n16k32 split-K=4 path. The generic + // scalar fallback below serves every unmatched (m, n, k), including the + // paired M=2 API shape with the same (N, K). if (m == 16 && n == kPackN && k == kPackK) { - const dim3 grid(static_cast(n / kDummaTileN)); - const dim3 block(kDummaBlockThreads); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_m16n16k32_1wave_kernel), - grid, block, 0, stream, - a, b, x_scale, weight_scale, - reinterpret_cast<__hip_bfloat16*>(out), n); + launch_dumma_m16n16_gemm(a, b, x_scale, weight_scale, out, n, stream); return; } - // Generic scalar fallback for every other (m, n, k), including the paired - // M=2 shape with the same (N, K): decodes the n-major pack when - // b_transposed, row-major otherwise. - const int64_t total = static_cast(m) * n; - const dim3 grid(static_cast( - (total + kScalarBlockThreads - 1) / kScalarBlockThreads)); - const dim3 block(kScalarBlockThreads); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_scalar_gemm_kernel), - grid, block, 0, stream, - a, b, x_scale, weight_scale, - reinterpret_cast<__hip_bfloat16*>(out), m, n, k, b_transposed); + // Generic scalar fallback for every unmatched (m, n, k), including M=2: + // decodes the n-major pack when b_transposed, row-major otherwise. + launch_scalar_gemm(a, b, x_scale, weight_scale, out, m, n, k, b_transposed, + stream); } // @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_gate_up_proj.hip index 7d26ecac..b98d2cc7 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_gate_up_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP4/M16/shared_gate_up_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=hy3_tp4_shared_gate_up_proj_m16 commit=f52225ed98ffd619fa1318d3eca7e56addd32f82 added=2026-08-26 +// baseline_us=73.262 speedup=6.47763 // median_us=11.31 p90_us=11.32 // source=hy3-dsh-tp4-m16-1-7f1fb1d1 // MetaInfer W8A8 INT8 GEMM for gfx928 (K500SM_AI). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/o_proj.hip index 40e1cdef..b5950e08 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/o_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/o_proj.hip @@ -1,208 +1,170 @@ -// @@variant shape=hy3_tp8_o_proj_m16 commit=4ddcb92e2032eb19fef17c7d4e19ac6739780ba9 added=2026-08-27 -// median_us=12.43 p90_us=12.59 -// source=hy3-dsh-tp8-m16-2-dc53295a -// MetaInfer W8A8 INT8 GEMM HIP implementation for gfx928 (worker_1). +// @@variant shape=hy3_tp8_o_proj_m16 commit=0b2614f457656239125b5d510d60e65275d45e24 added=2026-09-09 +// median_us=11.04 p90_us=11.16 speedup=2.91 baseline_us=32.14 +// source=hy3-dsh-tp8-m16-9-8-0161e718 +// MetaInfer W8A8 INT8 GEMM HIP implementation. // -// Iteration 4: bounded LDS prefetch pipeline for the assigned shape -// hy3_tp8_o_proj_m16 (M=16, N=4096, K=1024). -// * Iteration 1 (accepted): 256 blocks x 1 wavefront x 1 16x16 N tile, -// direct global fragment loads, 32 K-steps; official median -// 97.3785 us / p90 99.7009 us. -// * Iteration 2 (rejected): 128 blocks x 2 wavefronts x 2 N tiles, -// median 168.109 us / p90 177.434 us -- intra-block wave pairing with -// the direct byte-load B path regressed, so launch-geometry co- -// residency alone is not the lever. -// * Iteration 3 (accepted official best): split-K=2 pipeline with int32 -// partials in the caller workspace -- (1) -// w8a8_dumma_m16_n16_splitk_partial_kernel<2>: grid = 2*256 = 512 -// one-wave blocks (one 16x16 N tile per block; split s owns the K-half -// [s*512, s*512+512), 16 ascending m16n16k32 steps with the -// iteration-1 direct-load double-buffered K loop), each block storing -// its int32 tile to workspace plane s via the API LDS plane -// (mem_row_major); (2) w8a8_dumma_m16_combine_kernel<2>: grid = 256 -// one-wave blocks summing the planes in ascending split order and -// applying the exact reference fp32 scale/bf16 chain. Official median -// 31.446 us / p90 31.509 us (beats the fixed Triton Graph baseline -// 32.137/35.969 us). -// * This round (mandated architecture/pipeline axis: bounded -// register/LDS prefetch): the exact code object of the accepted -// iteration-3 partial kernel shows per K-step ONE global_load_dwordx2 -// per lane for A (8 B, immediate K-offsets -- already vectorized and -// L2-hot) but EIGHT per-lane global_load_ubyte for B (one byte per -// k-row, k-rows strided by N=4096) with s_waitcnt vmcnt(N) interleaved -// with the byte reassembly -- 16 dwordx2 + 128 ubyte + 150 waits + 16 -// mmac per block, i.e. ~9 serialized global-latency rounds per K-step. -// A-only staging would remove only ~1 of those 9 rounds, and -// multi-N-tile reuse would add 8 B-ubyte rounds per extra tile per -// step (per-step chain grows while blocks/CU halves), so the load path -// itself is replaced with a BOUNDED depth-1 LDS prefetch pipeline: the -// partial kernel keeps the exact split-K=2 geometry (512 one-wave -// blocks = 4.27/CU, zero barriers, same combine kernel, same workspace -// planes, same ascending-K int32 order) and only the B load path -// changes. Each block stages its 8-KiB B tile slice into a -// double-buffered LDS window (2 x 64 rows x 48-B padded stride; one -// window = 2 m16n16k32 steps; 8 windows for kLen=512) with ONE 16-B -// global_load_dwordx4 per lane per window, issued one full window -// ahead of consumption; the K loop then reads B fragments from LDS -// (ds_read_ubyte x8 + reassembly, ~30-cycle LDS latency, no vmcnt) -// while A stays on its proven single-dwordx2 direct path. Byte reuse: -// A (8 KiB/block slice; 16 KiB total, L2-hot) is re-read per K-step -// exactly as before; B (8 KiB per block) is read from global exactly -// once per block through the vectorized prefetch and every staged byte -// is consumed by exactly one m16n16k32 step (B is inherently -// single-use per block -- the change converts 128 scattered -// ubyte-load+wait rounds into 8 vectorized prefetch rounds plus LDS -// reads). -// * Falsifiable gate: median_us < 31.446 AND p90_us < 31.509 (strict -// improvement over the iteration-3 official best with the p90 noise -// guard); expected effect ~2x (partial kernel ~3-4x faster, combine -// ~2-3 us unchanged). -// * The CU-aligned non-power-of-two candidate (SPLIT_K=3, grid = 768 -// blocks = 6.4 blocks/CU, non-uniform 32-aligned K slices 384/384/256) -// remains implemented in the same templated partial+combine pair and -// explicitly instantiated; this round launches SPLIT_K=2 as the -// primary candidate. -// * Workspace contract unchanged: partials occupy SPLIT_K * 65536 int32 -// planes (512 KiB for S=2) and are overwritten on every launch; stream -// order makes combine read the planes written by the same replay. -// Graph capture/replay safe: two kernel launches, no allocation, no -// sync, no host reads, no default stream. -// * Iteration 6 (this round): the exact code object of the accepted -// iteration-4 partial kernel shows the B fragment is still reassembled -// per m16n16k32 step from EIGHT per-lane ds_read_u8 (k-rows strided by -// the 48-B padded window stride, one byte per k-row) with ~8 -// s_waitcnt lgkmcnt + ~15 VALU byte-reassembly ops per step, while A is -// already one dwordx2 per lane per step and software-pipelined 4-6 -// steps ahead by the compiler (offsets 128..288 issued early). The -// remaining per-step serial chain is therefore the B-from-LDS byte -// path, so this round replaces the B data path with the validated -// gate_up/down_proj recipe (int8-w8a8-gemm-decode skill): pack the -// exact-shape weight OUT of the timed region into n-major fragment -// layout packed[n*K + k] = raw[k*N + n] (same 4 MiB byte count, same -// buffer, graph-stable addresses; identity pack retained for every -// other (K, N)), keep the depth-1 double-buffered LDS window pipeline, -// but store each window n-major (16 n-cols x 64 k-rows, 80-B padded -// k-stride = 64+16, 16-B aligned non-power-of-two bank skew) so the -// col_major B fragment load (du_mma.hpp: lane l reads -// p[(l&15)*ldm + (l>>4)*8 + i], i=0..7 -- 8 CONTIGUOUS bytes) becomes -// one 8-B vectorized read per lane per step (emitted as two -// ds_read_b32) with no reassembly. Stage-in stays -// one 16-B global_load_dwordx4 per lane per window (now from the packed -// n-major buffer), issued one window ahead of consumption. A path, -// split-K=2 geometry (512 + 256 one-wave blocks), combine kernel, -// workspace planes, ascending-K exact int32 order, and zero-barrier -// structure are held fixed; the scalar fallback decodes the n-major -// pack for (k,n)==(1024,4096) (keeps the paired M=2 API shape exact). -// * Iteration 7 (rejected, reverted): single-buffering leg of the -// mandated double-vs-single buffer comparison -- all 8 B windows staged -// in a prologue burst (LDS 11264 B, zero B-related vmcnt in the K loop, -// 525 -> 371 instructions). Median 14.7632 / p90 14.7728 us vs the -// iteration-6 gate 14.7233/14.7345 -> NOT accepted. Conclusion: the B -// stage-in latency was not the binding constraint; the residual -// per-wave serial latency is the direct A-path global loads waited -// inside the K loop (vmcnt(0..6) pipelined waits) with only 1.07 -// waves/SIMD (4.27 one-wave blocks/CU) to overlap them. -// * Iteration 8 (this round, mandated HIP-only occupancy probe): PMC for -// the accepted iteration-6 build (arch_vgpr 40, sgpr 16, LDS 3584 B, -// scratch 0, 512 blocks, partial kernel 12.8 us) shows LDS allows 18 -// blocks/CU and VGPR 24 blocks/CU with zero spills -- neither resource -// binds at 4.27 blocks/CU, so the only occupancy limiter is the grid -// (total one-wave blocks). This round flips ONLY the split factor to -// SPLIT_K=4: 1024 one-wave partial blocks = 8.53 blocks/CU = 2.13 -// wavefronts/SIMD (first candidate in the trusted [2..8] probe set that -// crosses 2 waves/SIMD), uniform 64-window-aligned K slices 256/256/ -// 256/256 (numWindows = 4), all blocks resident (LDS 18.3 >= 8.53, -// VGPR 24 >= 8.53). The kernel body is bit-identical to the accepted -// iteration-6 code: depth-1 double-buffered LDS window pipeline -// (b_lds[2][1280], one 16-B global_load_dwordx4 per lane per window -// issued one window ahead), col_major B fragment (one 8-B vectorized -// LDS read per lane per step), direct A dwordx2 path, zero barriers, -// same combine/scale chain, same n-major pack, same fallbacks. A and B -// per-replay HBM totals stay constant (A: 1024 x 4 KiB, B: 1024 x -// 4 KiB); only the combine kernel now sums 4 freshly-written -// L2-resident planes (+512 KiB reads, no repeated A/B HBM reads). -// Workspace: 4 planes x 256 KiB = 1 MiB <= the 16-plane (4 MiB) caller -// budget; the launcher workspace guard keeps the scalar fallback for an -// undersized workspace. Expected effect: the ~2x resident waves overlap -// the A-load global stalls (and the 4 instead of 8 per-block stage-in -// rounds halve the per-wave exposed-latency chain), partial kernel -// ~12.8 us -> ~7-10 us, combine ~2 us -> ~2.5-3.5 us, total ~10-13 us. -// Falsifiable gate: median_us < 14.72327470779419 AND p90_us < -// 14.734469652175903 (strict improvement over the iteration-6 official -// best with the p90 noise guard). If the median stays ~14.7 us, the -// kernel is not occupancy-limited on this axis and the next round must -// attack the A-path itself (LDS-stage A, or A layout) or probe SPLIT_K=8 -// (4.27 waves/SIMD, combine reads 2 MiB). -// * Iteration 11 (rejected, reverted): SPLIT_K=8 occupancy probe (2048 -// one-wave partial blocks = 17.07 blocks/CU = 4.27 waves/SIMD, kernel -// body bit-identical to iteration 8) regressed to 17.3119/17.3270 vs the -// 14.4528/14.4672 gate -- the occupancy axis is falsified (1.07 through -// 4.27 waves/SIMD all flat or worse; 2.13 is the peak). Per-kernel -// profile of the accepted build: partial kernel 12.16 us + combine -// kernel 2.88 us (profiled; operator 15.04 us), i.e. the combine is a -// ~20% SERIAL second phase plus its launch gap. -// * Iteration 12 (this round, HIP-only consolidation): fuse the combine -// into the partial kernel with the qkv_proj-lineage per-tile -// last-arrival protocol -- monotonic counters in the workspace tail -// (256 x int32 after plane 3; one async zero per workspace pointer on -// its first eager use, before capture), each block does one -// __threadfence() + atomicAdd(&counters[tile], 1), and the arrival with -// (arrived % SPLIT_K) == SPLIT_K - 1 sums the tile's 4 planes in -// ascending split order and writes the scaled bf16 output with the exact -// combine-kernel math (bit-identical bytes). The operator becomes ONE -// launch per replay: the ~2.9-us serial combine phase and its launch gap -// disappear, each tile's combine runs as soon as its 4 splits land -// (overlapped with sibling blocks), and the plane reads happen while the -// planes are L2-fresh. SPLIT_K=4 geometry, kernel body, n-major pack, -// workspace planes, exact int32 order, and fallbacks are unchanged; the -// two-kernel path remains for workspaces that fit the planes but not the -// counters. Expected: partial kernel ~12 us (combine tail on the -// last-arrivers only) + no second kernel -> total ~12-13 us. Falsifiable -// gate: median_us < 14.452790021896362 AND p90_us < 14.467190504074097 -// (strict improvement over the iteration-8 official best with the p90 -// noise guard). If the median stays ~14.4 us, the serial combine phase -// was not the binding cost and the next round must attack the partial -// kernel's own issue chain (dual accumulator fragments) or geometry. -// * Iteration 21 (this round, HIP-only): A-LDS double-buffered window -// staging -- A is staged through the same depth-1 double-buffered LDS -// window pipeline as B (a_lds[2][1280], one 16-B dwordx4 per lane per -// window issued one window ahead) and the per-step a_frag loads read LDS -// (ldm = 80) instead of global (ldm = 1024). Iteration 7 showed removing -// every B-related vmcnt wait from the K loop does not move the time, and -// identified the DIRECT A-PATH global loads waited inside the loop as the -// residual per-wave serial latency; iteration 10 hoisted all A dwordx2 -// loads into a prologue burst and regressed (14.8639) because the 8 -// steps of A sat in VGPRs. Staging A in LDS (not VGPRs) keeps the -// prefetch in registers only (the compiler hoists all 8 A/B window -// dwordx4 prefetches into the prologue with progressive vmcnt waits; -// arch_vgpr 40 -> 50, 2-3 waves/SIMD -- lockstep evidence says extra -// co-residency does not change the serial-chain makespan), LDS 3588 -> -// 6148 B/block (64 KiB / 6.15 KiB = 10.4 blocks/CU >= 8.53 needed, grid -// stays fully resident), and moves the A vmcnt wait out of the per-step -// mmac chain to the top-of-window ds_write where it overlaps a full -// compute window -- the exact recipe that took the B path from 31.45 to -// 19.01 us in iteration 4. Fragment values are byte-identical (same -// global bytes staged, same col_major-style 8-B fragment mapping), so -// the k-ascending int32 accumulation, the ascending-split sum, the fp32 -// scale/bf16 chain and every output byte are unchanged. Geometry -// (SPLIT_K=4, 1024 one-wave blocks, one launch per replay), B path, -// fused last-arrival counter protocol, workspace/planes/counters budget, -// Graph safety, two-kernel fallback (counters = nullptr) and generic -// scalar fallback are UNCHANGED. Expected: per-block serial latency -// drops by the exposed A-load round trips (~8 waited dwordx2 -> 4 -// prologue-hidden dwordx4), operator median toward ~8-11 us. -// Falsifiable gate: median_us < 13.084909915924072 AND p90_us < -// 13.115299940109253 (strict improvement over the iteration-12 official -// best with the p90 noise guard). If the median stays ~13 us, the A path -// is falsified as the binding cost and the residual is the fused -// atomic/combine tail, irreducible in this protocol. +// Worker: worker_1 (physical GPU 1) +// Task: hy3-dsh-tp8-m16-9-8-0161e718, parallel-explore child +// Shape: hy3_tp8_o_proj_m16 M=16, N=4096, K=1024 // -// The guarded exact-shape branch keeps every other (m, n, k) -- including the -// paired M=2 API shape with the same (N, K) -- on the generic scalar fallback -// (which decodes the n-major pack only for (k,n)==(1024,4096) and otherwise -// reads the row-major weight, matching the identity pack below). -// Later optimization rounds may replace only the exact-shape branch -// implementation (and the pack op) while this dispatch structure stays. +// Iteration 0 = correctness bootstrap: one scalar int8 dot-product kernel, +// one thread -> one output element (no DUMMA). Official median 78.99 us. +// +// Iteration 1 = minimal native gfx928 DUMMA m16n16k32 INT8 tile with the +// installed API: one 64-thread wavefront per block, one 16x16 N +// tile per block (grid = N/16 = 256 blocks; M = 16 fits exactly one M-tile), +// 32 strictly ascending du_mma_sync steps over K = 1024 in a single int32 +// accumulator fragment, zero barriers anywhere (single wavefront per block; +// the epilogue LDS store->read round-trip is intra-wavefront program order), +// exact fp32 scale chain (float(acc) * x_scale[m] * weight_scale[n]) and +// round-to-nearest-even bf16 store, one launch per replay on the caller +// stream, no workspace use. Official median 18.2465 us (1.76x vs the fixed +// Triton 32.137 us Graph baseline), but the grid is only 256 resident +// one-wave blocks = 0.53 wavefronts/SIMD (256 waves over 120 CUs x 4 +// SIMDs): about half of every SIMD's issue slots are empty, so the kernel +// is co-residency / load-latency bound, not compute bound. +// +// Iteration 2 (accepted official best 15.118930/15.126930 us; architecture +// mandate: one complete launch geometry among 1/2/4 waves per block x 1/2/4 adjacent N tiles): choose +// 2 wavefronts per block x 1 adjacent N tile per block = in-block +// split-K=2. Grid stays kNTiles = 256 blocks (>= 120 device CUs, ~2.13 +// blocks/CU); each 128-thread block holds two 64-thread wavefronts. Wave 0 +// accumulates the contiguous ascending K slice [0, 512), wave 1 accumulates +// [512, 1024) (16 m16n16k32 steps each, direct 8-B vectorized global +// fragment loads as in iteration 1). Each wave publishes its int32 partial +// tile to a per-block LDS plane pair (2 x 1 KiB), one END-of-K +// __syncthreads(), then wave 1 adds the planes (s=0 then s=1; int32 +// addition is exact, so the combined total is bit-identical to iteration +// 1's single full-K accumulator) and runs the unchanged scale/RNE-bf16 +// epilogue. This doubles the wavefront count to 512 (~1.07/SIMD, 4-5 +// resident waves/CU) without any extra global traffic, needs no workspace, +// and mirrors the validated N=4096 M=16 o_proj winner of the TP4 decode +// lineage (256 blocks x 2 waves, one tile/block, one END-of-K barrier). +// The exact-shape weight pack (untimed, launch_pack_w8a8_weight) switches +// from the identity copy to an n-major transpose packed[n*K + k] = +// raw[k*N + n] for exactly (k, n) == (1024, 4096) so each lane's 8-byte +// matrix_b col_major fragment is one contiguous 8-B vectorized global load +// (identity [K, N] would scatter it one byte per k row); the generic scalar +// fallback decodes that n-major layout for the exact (k, n) pair so every +// other (m, n, k) -- including the paired M=2 API shape with the same +// (N, K) -- stays byte-exact. Identity packing is retained for every other +// (K, N). +// +// Iteration 9 (this round): the accepted iteration-2 kernel (official +// 15.118930/15.126930 us) is replaced by the exact-shape reference lineage's +// validated final architecture for this geometry (hy3-dsh-tp8-m16-2-dc53295a, +// official 12.43/12.59 us on the same fixed 32.137 us Triton baseline): +// grid-level split-K=4 (1024 one-wave blocks, ~2.13 resident wavefronts/ +// SIMD), each block accumulating one contiguous ascending K quarter with A +// AND B both staged through the depth-1 double-buffered LDS window pipeline +// (one 16-B uint4 per lane per window issued one window ahead; per-step +// fragment loads read LDS at the 80-B padded stride), publishing its int32 +// tile to a workspace plane, with the split-K combine FUSED into the same +// kernel as a per-tile last-arrival tail (monotonic counters in the +// workspace tail) so the operator stays one launch per replay. In-lineage +// experiments that failed to beat iteration 2 are reverted and NOT replayed: +// iteration 5 (wave-1 tail scale prefetch, 15.2968/15.3360 us), iteration 7 +// (rolled depth-1 register stage transport, 19.2353/19.2721 us), iteration 8 +// (in-block split-K=4 occupancy probe, 15.4697/15.5025 us). +// +// Iteration 17 (this round): B-stream-locality round on the accepted +// iteration-9 kernel (official 11.320700/11.428699 us, digest 71ac0c9a, +// verified == live source before this edit; git HEAD 2dc3fd7 clean). Rounds +// 11 (publish tail surgery, 12.7361/12.7681), 14 (A register-resident, +// 12.3341/12.4781) and 15 (grid split-K 8, 14.8350/14.8526) were rejected +// and rolled back byte-for-byte, so this is a NEW bounded experiment from +// the pristine accepted source. Round-14/15 evidence leaves ONE untested +// axis of the accepted architecture: the physical layout of the once-read +// 4-MiB packed-B stream. The accepted kernel's exact-shape pack is the +// n-major transpose packed[n*K+k] = raw[k*N+n] (column stride K = 1024), so +// every one of a block's four 1-KiB window prefetches issues 16 FULL 64-B +// L2 lines strided 1 KiB apart (each block's B quarter is 16 scattered +// column chunks of 256 B over 16 KiB of address space), and the replay-wide +// 4-MiB B stream is served as ~65k single lines with no run structure at +// the DRAM row level (~377 GB/s effective read for a per-replay 4.27-MB +// footprint in 11.32 us). This round REPACKS the exact-shape buffer (same +// 4-MiB buffer, same byte count, one-time untimed pack, byte-permutation +// only) so each (split, tile) block owns one contiguous 4-KiB slot holding +// its whole B slice: slot = split*256 + tile, then window (1 KiB), then +// column (64 B), then k. Each window prefetch then reads 16 SEQUENTIAL +// 64-B lines (one contiguous 1-KiB run per block per matrix per window), +// block slots are issued in launch order (blockIdx = split*256 + tile), and +// the replay B stream decomposes into 1024 advancing 4-KiB sequential runs +// instead of 4096 scattered 1-KiB column runs -- strictly better DRAM +// row-buffer / prefetch structure for byte-identical traffic. LDS staging, +// fragment order, plane/workspace protocol, epilogue chain, launch geometry +// and every byte VALUE the kernel consumes are unchanged (the slot is +// exactly the old 16 columns x 256 k of the block in the same ascending-k +// order), so int32 accumulation is bit-identical to iteration 9; the +// generic scalar fallback and the paired M=2 path decode the same new pack +// (w8a8_pack_b_off), keeping every output byte exact. Falsifiable +// prediction: official unprofiled Graph-replay median below 11.207493 us +// (>= 1% under 11.320700 us) with p90 <= 1.05 x 11.428699 = 11.999134 us, +// 0 mismatches, Graph capture passed and the paired M=2 fallback validation +// unchanged; flat 11.2-11.6 us falsifies DRAM run-locality as the limiter +// (kernel is latency-chain bound instead) and directs the next round to +// per-wave latency reduction at unchanged geometry/traffic. +// +// Iteration 19 (this round): two-tile-per-wave A-sharing round on the +// shadow-active iteration-17 source (digest 19ec89c8, official Graph-replay +// median 11.272130/11.289730 us on the same fixed 32.137 us Triton baseline; +// git HEAD 2dc3fd7 = accepted iteration 9). Rounds 11 (publish tile_lds +// removal, 12.7361/12.7681), 14 (A register-resident with 8-B global +// fragment preloads, 12.3341/12.4781), 15 (grid split-K 8, 2048 blocks, +// 14.8350/14.8526), 17 (B block-contiguous repack, 11.2721/11.2897 -- only +// +0.43%, not accepted) and 18 (split-K 8 + A register-resident at 3588 B +// LDS, 15.0071/16.1807) bound this round. The surviving limiter reading: +// the replay wall is dominated by per-replay L2 line-REQUEST service, not by +// DRAM bytes, LDS path instructions or occupancy. Every falsified round kept +// the per-replay line-request count at 131,072 read requests (A 65,536 = 4 +// splits x 256 tile-blocks x one 4-KiB/64-line A slice each, B 65,536 +// compulsory once-read lines), and round 14 is the one negative control that +// MOVED the request count: its 8 x 8-B dwordx2 A preloads re-requested each +// A 64-B line twice (request count up ~50%) and the wall went up ~9%. +// This round halves the A-side requests by pairing the two adjacent 16-col +// tiles of each (split, pair) in ONE 64-lane wave that stages and loads its +// 4-KiB A slice ONCE and feeds two 16x16 accumulator tiles from the same +// a_lds window stream (grid 512 = 4 splits x 128 tile pairs; b_lds gains a +// second per-stage 1280-B plane for the paired tile; per-window MMA work +// doubles 2 -> 4 so the depth-1 window-prefetch latency cover per wave +// grows; per-block plane publish and arrival machinery now serve two tiles). +// Per-replay read requests drop to ~98,304 (-25%: A 32,768, B 65,536) at +// byte-identical per-(tile, split) traffic, DRAM bytes, workspace protocol, +// int32 accumulation order and every output byte. Falsifiable prediction: +// official unprofiled Graph-replay median below 11.207493 us (>= 1% under +// 11.320700 us) with p90 <= 11.999134 us, 0 mismatches, Graph capture passed +// and the paired M=2 fallback validation unchanged; point expectation +// ~9.6-10.9 us if request service is the wall (the tp4 o_proj M16 lineage's +// 512-wavefront final kernel streamed its once-read B at ~460-520 GB/s vs +// ~372-450 GB/s here at 1024 waves -- fewer waves with more per-wave work is +// the proven-faster regime on this exact N=4096 geometry); flat ~11.2-11.6 +// us or worse falsifies request service as the limiter and shows that >= ~2 +// resident waves/SIMD are required to hide the per-wave latency chain, +// directing the next round back to grid-1024 per-wave latency/tail work. +// +// Logical operator (fixed contract, see int8_w8a8_gemm_api.py): +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// * x_scale[m, 0] * weight_scale[n, 0] ) +// with x_q [M, K] int8, packed_weight the exact-shape n-major [N, K] int8 +// transpose (identity [K, N] for every other shape), x_scale [M, 1] fp32, +// weight_scale [N, 1] fp32 and out [M, N] bf16. +// +// Host symbols (called from csrc/bindings.cpp, which owns TORCH_LIBRARY +// registration and the current-stream lookup): +// launch_w8a8_gemm(...) +// launch_pack_w8a8_weight(...) +// +// Graph-safety: the gemm launch uses only caller-provided tensors and the +// passed hipStream_t (PyTorch's current stream); it never allocates, packs, +// synchronizes, or touches the default stream. The exact-shape arm uses the +// caller workspace as int32 partial planes plus a 1-KiB monotonic +// arrival-counter tail that the launcher zeros once per workspace pointer on +// its first eager use (before Graph capture; never during capture/replay). +// pack_weight runs outside the timed region and outside Graph capture. +// +// Include order is the known-good DTK order: hip_runtime.h first, then +// hip_bfloat16.h, then du_mma.h (du_mma.h is not self-contained when included +// before the HIP runtime headers; du_mma_common.h is not a public header). #include #include @@ -212,619 +174,641 @@ namespace { -// gfx928 wavefront is 64; block dimensions must be multiples of 64. -constexpr int kScalarBlockThreads = 256; +// --------------------------------------------------------------------------- +// Geometry / tuning constants +// --------------------------------------------------------------------------- + +// Generic scalar fallback block size (multiple of the 64-lane wavefront). +constexpr int kScalarThreads = 256; + +// gfx928 wavefront / DUMMA tile geometry. constexpr int kWaveSize = 64; +constexpr int kTileM = 16; +constexpr int kTileN = 16; +constexpr int kTileK = 32; + +// Exact assigned decode shape: hy3_tp8_o_proj_m16 (worker_1 / GPU 1). +// (K, N) = (1024, 4096) is also a valid paired M=2 API shape; the exact-shape +// guard below deliberately includes m == kExactM so those paired shapes keep +// reaching the generic scalar fallback (which decodes the exact-shape n-major +// pack there). +constexpr int kExactM = 16; +constexpr int kExactN = 4096; +constexpr int kExactK = 1024; + +// Exact-shape DUMMA launch plan (iteration 9): one 16x16 N tile per block, +// grid-level split-K = 4 along K (uniform contiguous ascending K quarters), +// one 64-thread wavefront per block, fused per-tile last-arrival combine +// (the exact-shape reference lineage's validated final architecture for this +// geometry: 12.43/12.59 us in hy3-dsh-tp8-m16-2-dc53295a). Each block stages +// its A+B slice through a depth-1 double-buffered LDS window pipeline (see +// the kernel below); the int32 partial planes and the per-tile arrival +// counters live in the caller workspace. +constexpr int kNTiles = kExactN / kTileN; // 256 tiles (>= 120 device CUs) +constexpr int kSkSplitK = 4; // grid split factor along K +constexpr int kSkKPerBlock = kExactK / kSkSplitK; // 256 k per block +constexpr int kSkTilesPerBlock = 2; // iteration 19: adjacent N tiles per wave +constexpr int kSkTPairs = kNTiles / kSkTilesPerBlock; // 128 tile pairs per split +constexpr int kSkGrid = kSkSplitK * kSkTPairs; // 512 one-wave blocks (1.07/SIMD) +constexpr int kWindowKRows = 64; // k-rows per LDS window (2 m16n16k32 steps) +constexpr int kWindowKStride = 80; // 64 k-rows + 16 pad; 16-B aligned, + // non-power-of-two bank skew (<=2-way) +constexpr int kWindowBytes = kTileN * kWindowKStride; // 1280 B per stage +constexpr int kNumLdsStages = 2; // depth-1 double buffering +constexpr int kSkWindows = kSkKPerBlock / kWindowKRows; // 4 windows per block + +// Iteration-17 B-pack slot geometry (see w8a8_pack_b_off / the pack kernel): +// the exact-shape packed buffer stores every (split, tile) block's whole B +// slice as one contiguous 4-KiB slot (4 windows x 16 columns x 64 k), so +// each window prefetch is one contiguous 1-KiB run (16 sequential 64-B +// lines) per matrix instead of 16 lines strided by the old column stride +// K = 1024. +constexpr int kBPackSliceBytes = 1 << 12; // 4096 B per slot +constexpr int kBPackWindowBytes = kTileN * kWindowKRows; // 1024 B per window +constexpr int kBPackColBytes = kWindowKRows; // 64 B per column per window +constexpr int64_t kPlaneInts = + static_cast(kExactM) * kExactN; // 65536 int32 per split plane +constexpr int64_t kPlaneBytes = kPlaneInts * sizeof(int32_t); // 256 KiB +constexpr int64_t kCounterInts = kNTiles; // one arrival counter per tile -// Exact-shape DUMMA geometry for hy3_tp8_o_proj_m16 (M=16, N=4096, K=1024). -constexpr int kTargetM = 16; -constexpr int kTargetN = 4096; -constexpr int kTargetK = 1024; -constexpr int kDummaTileM = 16; -constexpr int kDummaTileN = 16; -constexpr int kDummaTileK = 32; - -// Split-K planning for the exact shape: one 16x16 N tile per block means -// kNTiles = N/16 = 256 tiles; a split-K=S partial grid has S*kNTiles -// one-wave blocks. Every K boundary is aligned to the 32-element DUMMA step -// and to the 64-k-row LDS window (uniform 512/512 for S=2; uniform 256/256/ -// 256/256 for S=4; non-uniform 384/384/256 for S=3). -constexpr int kNTiles = kTargetN / kDummaTileN; // 256 -constexpr int kPlaneInts = kTargetM * kTargetN; // 65536 - -__device__ __forceinline__ constexpr int kSplitStart(int split_k, int s) { - if (split_k == 2) return s * 512; - if (split_k == 4) return s * 256; - return s * 384; // S=3 +// --------------------------------------------------------------------------- +// Device helpers +// --------------------------------------------------------------------------- + +// float32 -> bfloat16 with round-to-nearest-even, stored as its raw uint16 +// bit pattern. This is bit-identical to __float2bfloat16() / torch's bf16 +// cast for finite inputs and keeps the epilogue independent of any one +// bfloat16 class/family: out tensors are plain 16-bit bf16 storage. +__device__ __forceinline__ uint16_t float_to_bf16_rn(const float v) { + const uint32_t bits = __float_as_uint(v); + // Add 0x7FFF plus the LSB of the kept half (ties round to even), then keep + // the high 16 bits. + const uint32_t rounded = (bits + 0x7FFFu + ((bits >> 16) & 1u)) >> 16; + return static_cast(rounded); } -__device__ __forceinline__ constexpr int kSplitLen(int split_k, int s) { - if (split_k == 2) return 512; - if (split_k == 4) return 256; - return (s < 2) ? 384 : 256; // S=3 + +// --------------------------------------------------------------------------- +// Exact-shape block-contiguous B pack decode (iteration 17) +// --------------------------------------------------------------------------- +// Byte offset of column `col` / k-row `kk` inside the exact-shape packed +// buffer. The iteration-17 pack layout is a pure byte permutation of the +// iteration-1/9 n-major transpose packed[n*K+k] = raw[k*N+n] (same buffer, +// same byte count, same logical (col, k) byte values): +// packed[ split*256 + tile ][ window ][ col&15 ][ kk&63 ] +// with split = kk >> 8 (K quarter 0..3), window = (kk >> 6) & 3 (1-KiB +// window inside the block's 4-KiB slice), tile = col >> 4 (0..255) and +// column c = col & 15. Slot = split*256 + tile holds one block's whole B +// slice: 4 windows x 1024 B; a window holds 16 columns x 64 contiguous k. +__device__ __forceinline__ int64_t w8a8_pack_b_off(const int col, + const int kk) { + const int64_t split = static_cast(kk) >> 8; // 0..3 + const int64_t w = (static_cast(kk) >> 6) & 3; // 0..3 + const int64_t tile = static_cast(col) >> 4; // 0..255 + const int64_t c = static_cast(col) & 15; // 0..15 + return (split * kNTiles + tile) * kBPackSliceBytes + + w * kBPackWindowBytes + c * kBPackColBytes + + (static_cast(kk) & (kBPackColBytes - 1)); } // --------------------------------------------------------------------------- -// Generic scalar int8 dot-product kernel: one thread per output element. -// Fallback for every (m, n, k) not handled by an exact-shape specialization. -// For the exact (k, n) == (1024, 4096) the packed weight buffer holds the -// n-major transpose packed[n*K + k] = raw[k*N + n] (iteration 6), so the -// fallback decodes that layout there (this keeps the paired M=2 API shape -// with the same (N, K) byte-exact); every other (k, n) keeps the identity -// pack (raw row-major [K, N]). +// Generic scalar correctness kernel (fallback for every unmatched shape) // --------------------------------------------------------------------------- -__global__ __launch_bounds__(kScalarBlockThreads) void -w8a8_scalar_gemm_kernel( - const int8_t* __restrict__ a, - const int8_t* __restrict__ b, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - hip_bfloat16* __restrict__ out, - int m, - int n, - int k) { - const int64_t linear = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t total = static_cast(m) * n; - if (linear >= total) { + +// Scalar int8 dot-product GEMM, one thread per output element. Grid: +// blockIdx.x tiles the N dimension with kScalarThreads columns per block, +// blockIdx.y selects the M row, so adjacent lanes (threadIdx.x) read +// adjacent addresses along the fastest-varying N dimension (coalesced B +// reads) and every lane of a block broadcasts the same A row. +// For the exact (k, n) == (1024, 4096) the packed weight buffer holds the +// iteration-17 block-contiguous pack (one contiguous 4-KiB slot per +// (split, tile), see w8a8_pack_b_off), so the fallback decodes that layout +// there (this keeps the paired M=2 API shape with the same (N, K) +// byte-exact; the decode is a pure permutation of the iteration-1 n-major +// transpose it supersedes); every other (k, n) keeps the identity pack +// (raw row-major [K, N]). +__global__ __launch_bounds__(kScalarThreads) void w8a8_gemm_scalar_kernel( + const int8_t* __restrict__ a, // [M, K] row-major int8 + const int8_t* __restrict__ b, // [K, N] (identity) or [N, K] (exact pair) + const float* __restrict__ x_scale, // [M] fp32 + const float* __restrict__ weight_scale, // [N] fp32 + uint16_t* __restrict__ out, // [M, N] bf16 raw storage + const int m, + const int n, + const int k) { + const int col = static_cast(blockIdx.x) * kScalarThreads + + static_cast(threadIdx.x); + const int row = static_cast(blockIdx.y); + if (row >= m || col >= n) { return; } - const int row = static_cast(linear / n); - const int col = static_cast(linear - static_cast(row) * n); + // Exact k-ascending int32 accumulation. Assigned K keeps the exact int8 + // dot well inside int32 range. const int8_t* a_row = a + static_cast(row) * k; - const bool packed_nmajor = (k == kTargetK && n == kTargetN); - const int8_t* b_col = b + col; + const bool packed_nmajor = (k == kExactK && n == kExactN); int32_t acc = 0; if (packed_nmajor) { - // packed[col * 1024 + kk] == raw[kk * 4096 + col] (iteration-6 pack). - const int8_t* b_pack_col = b + static_cast(col) * kTargetK; + // Iteration-17 block-contiguous pack decode (w8a8_pack_b_off): the + // exact-shape buffer was repacked (same bytes, new permutation) so each + // block's B slice is one contiguous 4-KiB slot; the fallback reads the + // same logical (col, kk) byte values as the iteration-1/9 n-major + // transpose it supersedes. for (int kk = 0; kk < k; ++kk) { acc += static_cast(a_row[kk]) * - static_cast(b_pack_col[kk]); + static_cast(b[w8a8_pack_b_off(col, kk)]); } } else { + const int8_t* b_col = b + col; // advances by n every k step for (int kk = 0; kk < k; ++kk) { acc += static_cast(a_row[kk]) * - static_cast(b_col[static_cast(kk) * n]); + static_cast(*b_col); + b_col += n; } } + + // Convert to float only for the scaled product, then store bf16: + // (int32 dot) * x_scale[m] * weight_scale[n] == ((dot * xs) * ws). const float scaled = static_cast(acc) * x_scale[row] * weight_scale[col]; - out[linear] = __float2bfloat16(scaled); + out[static_cast(row) * n + col] = float_to_bf16_rn(scaled); } // --------------------------------------------------------------------------- -// Split-K partial kernel: one 64-thread wavefront per block, one 16x16 N -// tile per block, one K-slice per split. Iteration 6 replaces the B data -// path with the validated n-major packed + col_major fragment recipe: the -// exact-shape weight is packed (out of the timed region) as -// packed[n*K + k] = raw[k*N + n], and each block stages its B slice -// (kLen k-rows x 16 n-cols) into a double-buffered LDS window stored -// n-major (16 n-cols x 64 k-rows, 80-B padded k-stride) with ONE 16-B -// global_load_dwordx4 per lane per window, issued one full window (2 -// m16n16k32 steps) ahead of consumption; the per-step K loop then reads B -// fragments from LDS with the col_major load (du_mma.hpp lane mapping: -// p[(l&15)*ldm + (l>>4)*8 + i], i=0..7), i.e. one 8-B vectorized read per -// lane per step -- emitted by the compiler as two ds_read_b32 (8 contiguous -// bytes) -- so the 8 ds_read_u8 + ~8 lgkmcnt + ~15 VALU byte-reassembly -// chain of iteration 4 is gone (validated on gate_up/down_proj in this -// lineage). -// Iteration 21: A takes the same depth-1 double-buffered LDS window path as -// B (a_lds[2][1280], one 16-B dwordx4 per lane per window issued one window -// ahead; per-step a_frag reads from LDS at ldm = kWindowKStride), so the -// direct A-path vmcnt wait leaves the per-step mmac chain (iteration 7 -// identified those waited A loads as the residual per-wave serial latency; -// iteration 10's register hoist regressed on VGPR pressure -- LDS staging -// avoids it). Explicit int32 accumulation (max |dot| per slice = -// 16,516,096 << 2^31, k-ascending). -// Each block publishes its int32 tile through the API LDS plane -// (mem_row_major) and stores it to workspace plane `split` with one 16-B -// vector store per lane. Single wavefront -> no barrier anywhere (LDS -// store->read ordering is per-wave program order; buffer w&1 is stored at -// the top of iteration w and read only after that store). Guarded in the -// launcher: only this exact (m, n, k) reaches it. -// Iteration 12 (fused last-arrival combine, qkv_proj lineage recipe): with -// the counters pointer the tile's LAST arrival (monotonic per-tile counters -// in the workspace tail, (arrived % SPLIT_K) == SPLIT_K - 1) sums the -// tile's SPLIT_K planes in ascending split order and emits the scaled bf16 -// output in-kernel, byte-identical to the separate combine kernel, so the -// operator becomes ONE launch per replay; counters = nullptr keeps the -// iteration-8 two-kernel behavior for undersized workspaces. // --------------------------------------------------------------------------- -constexpr int kWindowKRows = 64; // one window = 2 m16n16k32 steps -constexpr int kWindowNCols = 16; // block tile width (n-cols per window) -constexpr int kWindowKStride = 80; // 64 k-rows + 16 pad; 16-B aligned, - // non-power-of-two bank skew (<=2-way) -constexpr int kNumBStages = 2; // double-buffered (depth-1 prefetch) - -template -__global__ __launch_bounds__(kWaveSize) void -w8a8_dumma_m16_n16_splitk_partial_kernel( - const int8_t* __restrict__ a, - const int8_t* __restrict__ b, - int32_t* __restrict__ partials, - int32_t* __restrict__ counters, +// Exact-shape DUMMA decode kernel (iteration 19: iteration-17 source plus +// two-tile-per-wave A-sharing; grid-level split-K = 4, A+B LDS-window +// staged, fused per-tile last-arrival combine) +// --------------------------------------------------------------------------- +// Architecture lineage: iteration 9 ported this geometry's validated +// reference kernel (hy3-dsh-tp8-m16-2-dc53295a, 12.43/12.59 us) -- one 16x16 +// N tile per ONE-wave block, SPLIT_K = 4 grid-level K quarters (1024 blocks, +// ~2.13 resident waves/SIMD, zero intra-pipeline barriers), A and B both +// staged through a depth-1 double-buffered LDS window pipeline, split-K +// combine FUSED into the kernel (per-tile last-arrival tail) -- and reached +// the official best 11.320700/11.428699 us; iteration 17 repacked B into +// block-contiguous 4-KiB slots (shadow 11.272130/11.289730 us, +0.43% only). +// +// Iteration 19 geometry: each of the 128 tile PAIRS per split (two adjacent +// 16-col tiles) is computed by ONE 64-lane wave, so the grid is 512 one-wave +// blocks = 4 splits x 128 pairs (1.07 waves/SIMD). The wave stages and +// loads its 4-KiB A slice ONCE per window and feeds two 16x16 accumulator +// tiles from the same a_lds window stream (a_frag shared, b_lds holds a +// second per-stage 1280-B plane for the paired tile); per-window MMA work +// doubles (2 -> 4 m16n16k32 steps) so each window's compute covers its own +// depth-1 prefetch latency longer; per-block prologue, plane-publish and +// arrival machinery now serve two tiles. This is the first lineage round +// that reduces the per-replay L2 line-request count: A-side requests halve +// (65,536 -> 32,768: one 64-line A-slice request set per 2 tiles instead of +// per tile) and total read requests drop ~25% to ~98,304 (B stays 65,536 +// compulsory once-read lines; DRAM bytes, workspace protocol, plane layout, +// counter semantics and every output byte are unchanged). +// +// Window pipeline per tile (byte-for-byte the accepted data path): tile t of +// a pair owns the contiguous ascending K slice [256*split, 256*split + 256) +// and k-rows are covered window by window ([64w, 64w + 64), w = 0..3 = two +// m16n16k32 steps per tile). Lane l prefetches ONE 16-B uint4 per matrix per +// window -- 16 contiguous k-bytes at k-offset 16*(lane>>4) of the window: +// A from the logical row-major ldm-K=1024 activation at row (lane&15), B +// from the iteration-17 block-contiguous pack slots (split*256 + 2*pair)<<12 +// and +4096 (each (split, tile) slot = one contiguous 4-KiB run; the pair's +// two slots are adjacent) at column (lane&15) -- issued one full window +// ahead of consumption, stored into +// a_lds[w&1][(lane&15)*80 + 16*(lane>>4)] and +// b_lds[w&1][t&1][(lane&15)*80 + 16*(lane>>4)] (16 x 64-k window planes at +// the 80-B padded k-stride = 64+16 bank skew). The per-step fragment loads +// (du_mma.hpp: lane l reads 8 contiguous bytes at row (l&15)*ldm + +// (l>>4)*8 + i) read LDS at ldm = 80 -- no global vmcnt in the mmac chain. +// +// Exact int32 accumulation: per-tile dots stay strictly k-ascending over the +// windowed m16n16k32 steps (max |partial| = 256*127*127 << 2^31), so each +// published plane tile is bit-identical to the iteration-9/17 per-(split, +// tile) blocks; the last-arrival tail sums the SPLIT_K = 4 planes in +// ascending split order (s = 0..3); int32 addition is exact, so every output +// byte is bit-identical to the accepted kernels (0 mismatches expected). +// +// Combine + epilogue (fused tail, counters != nullptr): the wave publishes +// BOTH tiles' int32 results to workspace planes `split` (per-tile contiguous +// layout workspace[split][tile][256], one 16-B vector store per lane per +// tile, through two API LDS planes tile_lds0/tile_lds1 with +// du_store_matrix_sync mem_row_major ldm 16), then lane 0 __threadfence()s +// (release), increments BOTH tiles' monotonic arrival counters (atomicAdd), +// __threadfence()s (acquire) and records for each tile whether the +// pre-increment value satisfies (arrived % SPLIT_K) == SPLIT_K - 1 in +// shared s_last0/s_last1. Each tile's LAST arrival of the replay then reads +// its 4 row-major int32 chunks of each plane in ascending split order and +// runs the UNCHANGED accepted scale/RNE-bf16 epilogue +// float(total) * x_scale[row] * weight_scale[n0 + col] for its 4 columns +// (row = lane>>2, cbase = 4*(lane&3)); the other three splits of each tile +// exit. Counters are monotonic (each replay adds exactly SPLIT_K per tile), +// so no per-replay reset is needed; the launcher zeroes them once per +// workspace pointer on first eager use, before Graph capture. +// --------------------------------------------------------------------------- +__device__ __forceinline__ void w8a8_tail_emit_tile( + const int32_t* __restrict__ partials, const int tile, const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - hip_bfloat16* __restrict__ out) { - static_assert(SPLIT_K == 2 || SPLIT_K == 3 || SPLIT_K == 4, - "exact-shape split-K plan supports SPLIT_K=2, 3 or 4"); + const float* __restrict__ weight_scale, uint16_t* __restrict__ out) { + // Fused per-tile last-arrival epilogue (byte-for-byte the iteration-9 tail): + // 16-B aligned reads of the per-tile-contiguous plane layout + // workspace[split][tile][256]; plane stride kPlaneInts/4 uint4s. Only the + // tile's last arrival calls this (guarded by the caller), so the planes it + // reads are all released (acquire fence in the caller) and L2-fresh. const int lane = static_cast(threadIdx.x); - const int split = static_cast(blockIdx.x) / kNTiles; - const int tile = static_cast(blockIdx.x) % kNTiles; - const int n0 = tile * kDummaTileN; - const int kStart = kSplitStart(SPLIT_K, split); - const int kLen = kSplitLen(SPLIT_K, split); - const int numWindows = kLen / kWindowKRows; // 8 (S=2), 6/4 (S=3) - - __shared__ __align__(16) int8_t b_lds[kNumBStages][kWindowNCols * - kWindowKStride]; - // Iteration 21: A is staged through the same double-buffered LDS window - // pipeline as B (16 m-rows x 64 k-rows per window, 80-B padded stride, - // one 16-B dwordx4 per lane per window issued one window ahead), so the - // per-step a_frag loads read LDS (ldm = kWindowKStride) instead of global - // (ldm = kTargetK). The A-path vmcnt wait leaves the per-step mmac chain - // and lands at the top-of-window ds_write alongside B's (iteration 7 - // identified the direct A-path global loads as the residual per-wave - // serial latency; iteration 10's register hoist was rejected -- staging in - // LDS, not VGPRs, keeps occupancy and the windowed depth-1 prefetch). - __shared__ __align__(16) int8_t a_lds[kNumBStages][kWindowNCols * - kWindowKStride]; - - du::dumma::DUFragment + const int row = lane >> 2; // 0..15 + const int cbase = (lane & 3) * 4; // 0, 4, 8, 12 within the tile + const uint4* plane = reinterpret_cast(partials) + + static_cast(tile) * (kTileM * kTileN / 4) + + row * 4 + (lane & 3); + int32_t acc[4] = {0, 0, 0, 0}; +#pragma unroll + for (int s = 0; s < kSkSplitK; ++s) { + const uint4 v = plane[s * (kPlaneInts / 4)]; + acc[0] += v.x; + acc[1] += v.y; + acc[2] += v.z; + acc[3] += v.w; + } + const float xs = x_scale[row]; + const int n0 = tile * kTileN; + uint16_t* out_p = + out + static_cast(row) * kExactN + n0 + cbase; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float scaled = + static_cast(acc[i]) * xs * weight_scale[n0 + cbase + i]; + out_p[i] = float_to_bf16_rn(scaled); + } +} +__global__ __launch_bounds__(kWaveSize) void w8a8_dumma_m16_n16_sk4_lds_fused_kernel( + const int8_t* __restrict__ a, // [M=16, K=1024] row-major int8 + const int8_t* __restrict__ b, // packed [N, K] n-major int8 + int32_t* __restrict__ partials, // workspace [kSkSplitK][kPlaneInts] + int32_t* __restrict__ counters, // workspace tail, 256 int32 + const float* __restrict__ x_scale, // [16] fp32 + const float* __restrict__ weight_scale, // [N] fp32 + uint16_t* __restrict__ out) { // [16, N] bf16 raw storage + const int lane = static_cast(threadIdx.x); // 0..63 + const int pair = static_cast(blockIdx.x) % kSkTPairs; // 0..127 + const int split = static_cast(blockIdx.x) / kSkTPairs; // 0..3 (K quarter) + const int tile0 = pair * kSkTilesPerBlock; // even N tile of the pair + const int tile1 = tile0 + 1; // adjacent odd N tile + const int kStart = split * kSkKPerBlock; // contiguous ascending K quarter + + // Double-buffered LDS windows (private to this single-wave block; no + // barrier inside the pipeline: store->read ordering is per-wave program + // order, enforced by the compiler's lgkmcnt waits). b_lds holds one + // 1280-B window plane per stage per tile of the pair; a_lds is SHARED by + // the two tiles (one A window load/stage feeds two 16x16 accumulators). + __shared__ __align__(16) int8_t a_lds[kNumLdsStages][kWindowBytes]; + __shared__ __align__(16) int8_t b_lds[kNumLdsStages][kSkTilesPerBlock] + [kWindowBytes]; + + du::dumma::DUFragment a_frag[2]; - du::dumma::DUFragment - b_frag[2]; - du::dumma::DUFragment - acc_frag; - du::dumma::du_fill_fragment(acc_frag, 0); - - // Depth-1 double-buffered LDS prefetch. Window w covers B k-rows - // [kStart + w*64, kStart + w*64 + 64) at the block's 16 n-cols. The - // packed buffer is n-major, so lane l fetches 16 CONTIGUOUS k-rows - // (16 B, one dwordx4) at n-col (n0 + (l&15)) and k-offset - // 16*(l>>4) within the window; the data for window w is issued one - // iteration ahead (prologue for w=0) so the vmcnt wait lands at the - // top-of-iteration ds_write_b128, a full compute window (~2 mmac steps) - // after issue. - const int8_t* b_packed_col = - b + static_cast(n0 + (lane & 15)) * kTargetK; - uint4 prefetch = *reinterpret_cast( - b_packed_col + kStart + 16 * (lane >> 4)); - // A prefetch mirrors B: lane l fetches 16 contiguous k-rows (16 B, one - // dwordx4) of its A row (lane & 15) at k-offset 16*(lane >> 4) within the - // window; every staged byte is consumed by exactly one m16n16k32 step via - // the row_major col_major-style fragment (the 16-B chunk covers two - // 8-B k-groups, exactly like the B window). - const int8_t* a_row = a + static_cast(lane & 15) * kTargetK; + du::dumma::DUFragment + b0_frag[2]; + du::dumma::DUFragment + b1_frag[2]; + du::dumma::DUFragment + acc0_frag; + du::dumma::DUFragment + acc1_frag; + du::dumma::du_fill_fragment(acc0_frag, 0); + du::dumma::du_fill_fragment(acc1_frag, 0); + + // Per-lane base pointers: A slice = activation row (lane&15). B = the + // pair's two contiguous 4-KiB slots of the iteration-17 block-contiguous + // pack at (split*256 + tile0) << 12 and +4096; inside a window (1 KiB), + // lane (lane&15) owns column c = lane&15 whose 64 k-bytes sit at c << 6, + // so one window prefetch is one contiguous 1-KiB run per tile (a pair + // covers 8 KiB of sequential pack slots per block). Byte values are + // identical to the iteration-9/17 layouts. + const int8_t* b_slot = + b + (static_cast(split) * kNTiles + tile0) * kBPackSliceBytes + + (lane & 15) * kBPackColBytes; + const int8_t* a_row = a + static_cast(lane & 15) * kExactK; + + // Prologue: prefetch window 0 for A and for each tile's B plane + // (16 contiguous k-bytes per lane per matrix per tile). uint4 a_prefetch = *reinterpret_cast( a_row + kStart + 16 * (lane >> 4)); + uint4 b0_prefetch = *reinterpret_cast( + b_slot + 16 * (lane >> 4)); + uint4 b1_prefetch = *reinterpret_cast( + b_slot + kBPackSliceBytes + 16 * (lane >> 4)); + #pragma unroll - for (int w = 0; w < numWindows; ++w) { + for (int w = 0; w < kSkWindows; ++w) { // Store the previously-prefetched window w into buffer w&1 (its vmcnt // wait has had the prior iteration's compute to complete), then issue // the next window's prefetches before computing this window from LDS. - *reinterpret_cast( - &b_lds[w & 1][(lane & 15) * kWindowKStride + 16 * (lane >> 4)]) = - prefetch; *reinterpret_cast( &a_lds[w & 1][(lane & 15) * kWindowKStride + 16 * (lane >> 4)]) = a_prefetch; - if (w + 1 < numWindows) { - prefetch = *reinterpret_cast( - b_packed_col + kStart + (w + 1) * kWindowKRows + 16 * (lane >> 4)); + *reinterpret_cast( + &b_lds[w & 1][0][(lane & 15) * kWindowKStride + 16 * (lane >> 4)]) = + b0_prefetch; + *reinterpret_cast( + &b_lds[w & 1][1][(lane & 15) * kWindowKStride + 16 * (lane >> 4)]) = + b1_prefetch; + if (w + 1 < kSkWindows) { a_prefetch = *reinterpret_cast( a_row + kStart + (w + 1) * kWindowKRows + 16 * (lane >> 4)); + b0_prefetch = *reinterpret_cast( + b_slot + (w + 1) * kBPackWindowBytes + 16 * (lane >> 4)); + b1_prefetch = *reinterpret_cast( + b_slot + kBPackSliceBytes + (w + 1) * kBPackWindowBytes + + 16 * (lane >> 4)); } - const int8_t* bwin = &b_lds[w & 1][0]; + // Two ascending m16n16k32 steps of this window PER TILE, fragments from + // LDS. The A fragments are loaded once and shared by both tiles' MMAs, + // so one A window global load + LDS stage serves two 16x16 tiles: the + // per-replay A-side line requests halve vs one tile per wave. Each + // tile's step order stays k-ascending (step 0 = lower k32 of the window, + // step 1 = upper k32), bit-identical to the iteration-9 per-(split, + // tile) blocks. const int8_t* awin = &a_lds[w & 1][0]; + const int8_t* bwin0 = &b_lds[w & 1][0][0]; + const int8_t* bwin1 = &b_lds[w & 1][1][0]; du::dumma::du_load_matrix_sync(a_frag[0], awin, kWindowKStride); - du::dumma::du_load_matrix_sync(b_frag[0], bwin, kWindowKStride); - du::dumma::du_load_matrix_sync(a_frag[1], awin + kDummaTileK, + du::dumma::du_load_matrix_sync(a_frag[1], awin + kTileK, + kWindowKStride); + du::dumma::du_load_matrix_sync(b0_frag[0], bwin0, kWindowKStride); + du::dumma::du_load_matrix_sync(b0_frag[1], bwin0 + kTileK, kWindowKStride); - du::dumma::du_load_matrix_sync(b_frag[1], bwin + kDummaTileK, + du::dumma::du_load_matrix_sync(b1_frag[0], bwin1, kWindowKStride); + du::dumma::du_load_matrix_sync(b1_frag[1], bwin1 + kTileK, kWindowKStride); - du::dumma::du_mma_sync(acc_frag, a_frag[0], b_frag[0], acc_frag); - du::dumma::du_mma_sync(acc_frag, a_frag[1], b_frag[1], acc_frag); + du::dumma::du_mma_sync(acc0_frag, a_frag[0], b0_frag[0], acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag[0], b1_frag[0], acc1_frag); + du::dumma::du_mma_sync(acc0_frag, a_frag[1], b0_frag[1], acc0_frag); + du::dumma::du_mma_sync(acc1_frag, a_frag[1], b1_frag[1], acc1_frag); } - // Publish the int32 tile through the API (mem_row_major), then one 16-B - // vector store per lane. Plane layout is per-tile contiguous - // workspace[split][tile][256] (e = row*16 + col), so each block's 1-KiB - // int32 tile is one coalesced 64-lane x 16-B store. Single wavefront per - // block: the LDS write->read dependency is ordered by the compiler's - // lgkmcnt wait, no barrier needed. - __shared__ __align__(16) int32_t tile_lds[kDummaTileM * kDummaTileN]; - du::dumma::du_store_matrix_sync(tile_lds, acc_frag, kDummaTileN, + // Publish this wave's two int32 partial tiles to workspace plane `split` + // (per-tile contiguous plane layout workspace[split][tile][256], row-major + // via the API LDS planes), one 16-B vector store per lane per tile. + __shared__ __align__(16) int32_t tile_lds0[kTileM * kTileN]; + __shared__ __align__(16) int32_t tile_lds1[kTileM * kTileN]; + const int e = lane * 4; // this lane's 4 consecutive row-major int32 + du::dumma::du_store_matrix_sync(&tile_lds0[0], acc0_frag, kTileN, du::dumma::mem_row_major); - - const int e = lane * 4; - const uint4 packed = - make_uint4(tile_lds[e], tile_lds[e + 1], tile_lds[e + 2], - tile_lds[e + 3]); - *reinterpret_cast(partials + split * kPlaneInts + - tile * (kDummaTileM * kDummaTileN) + e) = packed; - - // Fused last-arrival combine tail (iteration 12): when counters != nullptr - // the separate combine kernel is removed and the operator is ONE launch per - // replay. Every block signals its arrival for its tile with one lane-0 - // __threadfence() + atomicAdd(&counters[tile], 1); the arrival whose - // pre-increment value satisfies (arrived % SPLIT_K) == SPLIT_K - 1 is the - // tile's LAST arrival of this replay (counters are monotonic: each replay - // adds exactly SPLIT_K to every tile counter) and it sums the tile's - // SPLIT_K planes in ascending split order and writes the scaled bf16 - // output -- byte-identical to the separate combine kernel (same plane - // chunk reads, same ascending int32 sum, same - // float(acc) * x_scale[row] * weight_scale[col] -> __float2bfloat16 chain, - // same store addresses). Each tile's combine now runs as soon as its 4 - // splits land, while sibling blocks are still computing, so the ~2.9-us - // serial combine phase and its launch gap disappear and the plane reads - // happen while the planes are L2-fresh (per-tile, immediately after the - // tile's last write). Exact int32 accumulation unchanged: per-slice max - // |dot| << 2^31, ascending-split sum order (same as the reference's - // ascending-K order). The counters live in the workspace tail (last 256 - // int32 = 1 KiB, past plane SPLIT_K-1) and are zeroed once per workspace - // pointer by the launcher before the first eager use (never inside Graph - // capture); the two-kernel fallback passes counters = nullptr and the tail - // branch is dead there. + *reinterpret_cast( + partials + split * kPlaneInts + + static_cast(tile0) * (kTileM * kTileN) + e) = + make_uint4(tile_lds0[e], tile_lds0[e + 1], tile_lds0[e + 2], + tile_lds0[e + 3]); + du::dumma::du_store_matrix_sync(&tile_lds1[0], acc1_frag, kTileN, + du::dumma::mem_row_major); + *reinterpret_cast( + partials + split * kPlaneInts + + static_cast(tile1) * (kTileM * kTileN) + e) = + make_uint4(tile_lds1[e], tile_lds1[e + 1], tile_lds1[e + 2], + tile_lds1[e + 3]); + + // Fused last-arrival combine tail (unchanged per-tile protocol, now for + // two tiles). Only a tile's LAST arrival of this replay (pre-increment + // counter value % SPLIT_K == SPLIT_K - 1) sums its 4 planes in ascending + // split order and emits the scaled bf16 output, so each tile's combine + // runs as soon as its splits land while sibling blocks are still + // computing and the plane reads hit L2-fresh data. if (counters != nullptr) { __syncthreads(); - __shared__ int s_is_last; + __shared__ int s_last0; + __shared__ int s_last1; if (lane == 0) { - __threadfence(); // release: this block's plane store is visible to the - // observer of this block's arrival atomic - const int arrived = atomicAdd(&counters[tile], 1); - __threadfence(); // acquire: the reads below (after the barrier) see - // every sibling plane store released before its - // arrival atomic - s_is_last = ((arrived % SPLIT_K) == SPLIT_K - 1); + __threadfence(); // release: this wave's plane stores (both tiles) are + // visible to the observers of the arrival atomics + const int arrived0 = atomicAdd(&counters[tile0], 1); + const int arrived1 = atomicAdd(&counters[tile1], 1); + __threadfence(); // acquire: the plane reads below see every sibling + // block's store released before its arrival atomic + s_last0 = ((arrived0 % kSkSplitK) == kSkSplitK - 1); + s_last1 = ((arrived1 % kSkSplitK) == kSkSplitK - 1); } __syncthreads(); - if (s_is_last) { - const int row = lane >> 2; // 0..15 - const int cbase = (lane & 3) * 4; // 0, 4, 8, 12 within the tile - // 16-B aligned reads of the per-tile-contiguous plane layout - // workspace[split][tile][256]; plane stride is kPlaneInts/4 uint4s. - const uint4* plane = - reinterpret_cast(partials) + - tile * (kDummaTileM * kDummaTileN / 4) + row * (kDummaTileN / 4) + - (lane & 3); - int32_t acc[4] = {0, 0, 0, 0}; -#pragma unroll - for (int s = 0; s < SPLIT_K; ++s) { - const uint4 v = plane[s * (kPlaneInts / 4)]; - acc[0] += v.x; - acc[1] += v.y; - acc[2] += v.z; - acc[3] += v.w; - } - const float xs = x_scale[row]; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int col = n0 + cbase + i; - const float scaled = - static_cast(acc[i]) * xs * weight_scale[col]; - out[static_cast(row) * kTargetN + col] = - __float2bfloat16(scaled); - } + if (s_last0) { + w8a8_tail_emit_tile(partials, tile0, x_scale, weight_scale, out); + } + if (s_last1) { + w8a8_tail_emit_tile(partials, tile1, x_scale, weight_scale, out); } } } // --------------------------------------------------------------------------- -// Combine+scale kernel: one 64-thread wavefront per block, one 16x16 N tile -// per block. Each lane loads its 4-element chunk from every split plane with -// one 16-B vector load per plane, sums the SPLIT_K planes in ascending split -// order (exact int32; ascending-split sum == the reference's ascending-K -// order), then applies the exact reference fp32 chain -// float(acc) * x_scale[row] * weight_scale[n0 + col] -> __float2bfloat16 -// (RN). Same expression order as the reference -> bit-identical output for -// the same int32 dot. Zero barriers. Reads only planes written by the -// partial kernel of the same stream-ordered launch. +// Pack kernels (both run outside the timed region, in launch_pack_w8a8_weight) // --------------------------------------------------------------------------- -template -__global__ __launch_bounds__(kWaveSize) void w8a8_dumma_m16_combine_kernel( - const int32_t* __restrict__ partials, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - hip_bfloat16* __restrict__ out) { - const int lane = static_cast(threadIdx.x); - const int n0 = static_cast(blockIdx.x) * kDummaTileN; - const int row = lane >> 2; // 0..15 - const int cbase = (lane & 3) * 4; // 0, 4, 8, 12 within the tile - // 16-B aligned reads of the per-tile-contiguous plane layout - // workspace[split][tile][256]; plane stride is kPlaneInts/4 uint4s. - const uint4* plane = - reinterpret_cast(partials) + - blockIdx.x * (kDummaTileM * kDummaTileN / 4) + row * (kDummaTileN / 4) + - (lane & 3); - int32_t acc[4] = {0, 0, 0, 0}; -#pragma unroll - for (int s = 0; s < SPLIT_K; ++s) { - const uint4 v = plane[s * (kPlaneInts / 4)]; - acc[0] += v.x; - acc[1] += v.y; - acc[2] += v.z; - acc[3] += v.w; +// Identity device-to-device copy (generic fallback for every (K, N) and the +// scale copy for all shapes). +template +__global__ __launch_bounds__(kScalarThreads) void w8a8_identity_copy_kernel( + const T* __restrict__ src, T* __restrict__ dst, const int64_t numel) { + const int64_t stride = + static_cast(gridDim.x) * static_cast(blockDim.x); + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < numel; i += stride) { + dst[i] = src[i]; } +} - const float xs = x_scale[row]; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int col = n0 + cbase + i; - const float scaled = static_cast(acc[i]) * xs * weight_scale[col]; - out[static_cast(row) * kTargetN + col] = - __float2bfloat16(scaled); +void launch_identity_copy_int8(const int8_t* src, + int8_t* dst, + const int64_t numel, + hipStream_t stream) { + constexpr int64_t kMaxBlocks = 65535; + int64_t blocks = (numel + kScalarThreads - 1) / kScalarThreads; + if (blocks > kMaxBlocks) { + blocks = kMaxBlocks; + } + if (blocks < 1) { + blocks = 1; } + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(static_cast(blocks)), dim3(kScalarThreads), 0, stream, + src, dst, numel); } -// --------------------------------------------------------------------------- -// Identity device-to-device packing: byte-for-byte copy. This is the generic -// fallback for any (K, N) and is valid for every shape. -// pack_weight runs outside the timed region and outside Graph capture. -// --------------------------------------------------------------------------- -__global__ __launch_bounds__(kScalarBlockThreads) void -w8a8_pack_identity_bytes_kernel( - const uint8_t* __restrict__ src, - uint8_t* __restrict__ dst, - int64_t num_bytes) { - const int64_t stride = - static_cast(gridDim.x) * blockDim.x; - for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - i < num_bytes; i += stride) { - dst[i] = src[i]; +void launch_identity_copy_f32(const float* src, + float* dst, + const int64_t numel, + hipStream_t stream) { + constexpr int64_t kMaxBlocks = 65535; + int64_t blocks = (numel + kScalarThreads - 1) / kScalarThreads; + if (blocks > kMaxBlocks) { + blocks = kMaxBlocks; } + if (blocks < 1) { + blocks = 1; + } + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_identity_copy_kernel), + dim3(static_cast(blocks)), dim3(kScalarThreads), 0, stream, + src, dst, numel); } -// --------------------------------------------------------------------------- -// Exact-shape n-major pack (iteration 6): for (K, N) == (1024, 4096) the -// weight is transposed one-time (outside the timed region / Graph) into -// packed[n * K + k] = raw[k * N + n], so each lane's 8-byte col_major B -// fragment (du_mma.hpp lane mapping) is contiguous and every m16n16k32 -// fragment is a single vectorized load. Same byte count and buffer as the -// identity pack, so captured addresses are unchanged. One thread per output -// byte; runs once during weight prep. -// --------------------------------------------------------------------------- -__global__ __launch_bounds__(kScalarBlockThreads) void -w8a8_pack_nmajor_bytes_kernel( +// Exact-shape block-contiguous pack (iteration 17): for (K, N) == +// (1024, 4096) the weight is transposed one-time (outside the timed region +// / Graph) into the iteration-17 slot layout consumed by the DUMMA kernel +// and decoded by the scalar fallback (w8a8_pack_b_off): every (split, tile) +// block owns one contiguous 4-KiB slot = 4 windows x (16 columns x 64 k), +// with the k-byte run of each column inside a window contiguous so the +// kernel's 16-B per-lane window prefetches are one contiguous 1-KiB run per +// matrix per window (16 sequential 64-B lines). Same byte count and buffer +// as the identity/n-major pack, so captured addresses are unchanged; the +// mapping is a pure byte permutation of the iteration-1 n-major transpose. +// One thread per output byte; runs once during weight prep. Only correct / +// invoked for (k, n) == (1024, 4096) (the launcher guard below). +__global__ __launch_bounds__(kScalarThreads) void w8a8_pack_nmajor_bytes_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, - int k, - int n) { + const int k, + const int n) { const int64_t total = static_cast(k) * n; const int64_t stride = - static_cast(gridDim.x) * blockDim.x; + static_cast(gridDim.x) * static_cast(blockDim.x); for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < total; idx += stride) { - const int col = static_cast(idx / k); - const int kk = static_cast(idx - static_cast(col) * k); - packed[idx] = raw[static_cast(kk) * n + col]; + // idx = slot * 4096 + w * 1024 + c * 64 + kk + const int64_t slot = idx / kBPackSliceBytes; // split*256+tile (0..1023) + const int64_t rem = idx % kBPackSliceBytes; + const int64_t w = rem / kBPackWindowBytes; // 0..3 window in the slot + const int64_t c = (rem / kBPackColBytes) & 15; // 0..15 col in the tile + const int64_t kk = rem % kBPackColBytes; // k inside the 64-k window + const int k_src = + static_cast(((slot / kNTiles) << 8) | (w << 6) | kk); // 0..1023 + const int col = static_cast(((slot % kNTiles) << 4) | c); // 0..4095 + packed[idx] = raw[static_cast(k_src) * n + col]; } } -// Explicit instantiations: the CU-aligned non-power-of-two candidate -// (SPLIT_K=3, grid = 3*kNTiles = 768 blocks = 6.4 blocks/CU, non-uniform -// 32-aligned K slices 384/384/256) and the SPLIT_K=4 occupancy candidate -// (grid = 4*kNTiles = 1024 blocks = 8.53 blocks/CU, uniform 64-window-aligned -// K slices 256/256/256/256) are compiled so a future round flips one -// constant; this round launches SPLIT_K=4 as the primary mandated candidate. -template __global__ void w8a8_dumma_m16_n16_splitk_partial_kernel<3>( - const int8_t*, const int8_t*, int32_t*, int32_t*, const float*, - const float*, hip_bfloat16*); -template __global__ void w8a8_dumma_m16_combine_kernel<3>( - const int32_t*, const float*, const float*, hip_bfloat16*); -template __global__ void w8a8_dumma_m16_n16_splitk_partial_kernel<4>( - const int8_t*, const int8_t*, int32_t*, int32_t*, const float*, - const float*, hip_bfloat16*); -template __global__ void w8a8_dumma_m16_combine_kernel<4>( - const int32_t*, const float*, const float*, hip_bfloat16*); +// Generic scalar launch shared by the exact-shape-adjacent fallback path. +void launch_w8a8_gemm_scalar(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + const int m, + const int n, + const int k, + hipStream_t stream) { + const dim3 block(kScalarThreads); + const dim3 grid((static_cast(n) + kScalarThreads - 1) / + kScalarThreads, + static_cast(m)); + hipLaunchKernelGGL(w8a8_gemm_scalar_kernel, grid, block, 0, stream, a, b, + x_scale, weight_scale, + reinterpret_cast(out), m, n, k); +} } // namespace // --------------------------------------------------------------------------- -// Host launchers (stable symbols consumed by csrc/bindings.cpp). +// Host launch symbols (stable contract with csrc/bindings.cpp) // --------------------------------------------------------------------------- -extern "C" void launch_w8a8_gemm( - const int8_t* a, - const int8_t* b, - const float* x_scale, - const float* weight_scale, - void* out, - void* workspace, - int64_t workspace_bytes, - int m, - int n, - int k, - hipStream_t stream) { - hip_bfloat16* out_bf16 = static_cast(out); - - // Guarded exact-shape dispatch: only hy3_tp8_o_proj_m16 (16, 4096, 1024) - // reaches the split-K DUMMA pipeline. Every other (m, n, k) -- including - // the paired M=2 API shape with the same (N, K) -- takes the generic - // scalar fallback (which reads the row-major weight, matching the - // identity pack). - if (m == kTargetM && n == kTargetN && k == kTargetK) { - // Split-K pipeline (iteration 8: occupancy probe -- SPLIT_K=4, grid = - // 4*kNTiles = 1024 one-wave partial blocks = 8.53 blocks/CU = 2.13 - // wavefronts/SIMD. PMC evidence for the accepted iteration-6 SPLIT_K=2 - // build (arch_vgpr 40, sgpr 16, LDS 3584 B, scratch 0, 512 blocks) shows - // LDS allows 18 blocks/CU and VGPR 24 blocks/CU -- neither binds at the - // 4.27 blocks/CU needed, so the only occupancy limiter is the grid - // itself (1.07 waves/SIMD). Iteration 7 proved removing every B-related - // vmcnt wait from the K loop does not move the time (14.763 vs 14.723 - // us), i.e. the per-wave serial latency is the direct A-path global - // loads waited inside the loop with no sibling wave to overlap them. - // SPLIT_K=4 doubles resident waves to 2.13/SIMD (all 1024 blocks stay - // resident: LDS 18.3 >= 8.53, VGPR 24 >= 8.53) while keeping the exact - // iteration-6 depth-1 double-buffered window pipeline, A/B totals and - // per-replay HBM traffic constant (only the combine reads 4 instead of - // 2 freshly-written L2-resident planes, +512 KiB). The partial kernel - // reads B from the n-major packed buffer through the col_major fragment - // path (one 8-B vectorized LDS read per lane per step, staged via the - // double-buffered LDS window prefetched one window ahead): int32 - // partials in the caller workspace, combine+scale kernel in the timed - // Graph. Both launches are stream ordered inside the same captured - // region; the partial kernel overwrites every partial element on each - // launch (no workspace clear needed). - constexpr int kSplitK = 4; - constexpr int kPartialBlocks = kSplitK * kNTiles; // 1024 - constexpr int kCombineBlocks = kNTiles; // 256 - const int64_t plane_bytes = - static_cast(kPlaneInts) * sizeof(int32_t); // 256 KiB - const int64_t counters_bytes = - static_cast(kNTiles) * sizeof(int32_t); // 1 KiB +extern "C" void launch_w8a8_gemm(const int8_t* a, + const int8_t* b, + const float* x_scale, + const float* weight_scale, + void* out, + void* workspace, + int64_t workspace_bytes, + int m, + int n, + int k, + hipStream_t stream) { + // Timed operator constraints: no allocation, compilation, autotuning, + // packing, host/device synchronization, or default-stream launch. Only the + // caller-provided out and the caller-provided workspace (int32 partial + // planes + arrival-counter tail, iteration 9) may be written; the one-time + // counter hipMemsetAsync below is a stream op issued before Graph capture + // (first eager use), never during a capture/replay call. + // + // Exact assigned decode shape arm (hy3_tp8_o_proj_m16). Iteration 19 runs + // the grid-level split-K=4 kernel w8a8_dumma_m16_n16_sk4_lds_fused_kernel + // (512 one-wave blocks = 4 splits x 128 two-tile pairs, A+B LDS-window + // staged, two adjacent 16-col tiles per wave sharing one A window stream, + // fused last-arrival combine; iteration-9 reference architecture plus the + // iteration-19 A-sharing request-halving round). The guard keeps + // `m == kExactM`, so paired shapes with the same (N, K) (e.g. the M=2 API + // shape) fall through to the generic scalar fallback below (which decodes + // the exact-shape block-contiguous pack). + if (m == kExactM && n == kExactN && k == kExactK) { + // Split-K=4 fused path: 4 int32 planes (4 x 256 KiB) + 256 per-tile + // arrival counters (1 KiB) in the caller workspace. The fixed API + // validates workspaces for this (M=16, N=4096) shape at the 16-plane / + // 4-MiB capacity, so the fused single-launch path is always taken; an + // undersized workspace falls through to the generic scalar fallback and + // never runs the tuned pipeline. int32_t* partials = static_cast(workspace); - if (workspace_bytes >= kSplitK * plane_bytes + counters_bytes) { - // Fused single-launch path (iteration 12): the combine runs inside the - // partial kernel (per-tile last arrival), so the operator is ONE kernel - // launch per replay and the separate combine kernel + its launch gap - // disappear. The per-tile arrival counters (256 x int32 = 1 KiB) live - // immediately after plane kSplitK-1 (the partial kernel writes planes - // 0..kSplitK-1 only; the two-kernel fallback below never touches the - // counters). Counters are monotonic -- each replay adds exactly SPLIT_K - // to every tile counter -- so the last-arrival test is - // (arrived % SPLIT_K) == SPLIT_K - 1 and NO per-replay reset is needed. - // One async zero per workspace pointer is issued on the workspace's - // first eager use (the Graph capture flow always warms up eagerly - // first): the guarded hipMemsetAsync is a stream op on the caller's - // stream, no host sync, and is NOT part of the captured graph -- the - // qkv_proj lineage validated this exact protocol. - int32_t* counters = partials + kSplitK * kPlaneInts; + constexpr int64_t kFusedBytes = + static_cast(kSkSplitK) * kPlaneBytes + + kCounterInts * static_cast(sizeof(int32_t)); + if (workspace_bytes >= kFusedBytes) { + int32_t* counters = partials + kSkSplitK * kPlaneInts; + // Counters are monotonic (each replay adds exactly kSkSplitK per tile + // counter), so zeroing once per workspace pointer -- on its first + // eager use, before Graph capture -- is sufficient; the zero is never + // issued inside capture or on replay. static const void* s_fused_counters_ws = nullptr; if (s_fused_counters_ws != workspace) { - hipMemsetAsync(counters, 0, static_cast(counters_bytes), - stream); + (void)hipMemsetAsync( + counters, 0, static_cast(kCounterInts) * sizeof(int32_t), + stream); s_fused_counters_ws = workspace; } hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_m16_n16_splitk_partial_kernel), - dim3(kPartialBlocks), - dim3(kWaveSize), - 0, - stream, - a, - b, - partials, - counters, - x_scale, - weight_scale, - out_bf16); - return; - } - if (workspace_bytes >= kSplitK * plane_bytes) { - // Two-kernel fallback for workspaces that fit the partial planes but - // not the counters: partial kernel without the fused tail - // (counters = nullptr -> the tail branch is dead) then the combine - // kernel, byte-identical to the accepted iteration-8 behavior. - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_m16_n16_splitk_partial_kernel), - dim3(kPartialBlocks), - dim3(kWaveSize), - 0, - stream, - a, - b, - partials, - static_cast(nullptr), - x_scale, - weight_scale, - out_bf16); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_m16_combine_kernel), - dim3(kCombineBlocks), - dim3(kWaveSize), - 0, - stream, - partials, - x_scale, - weight_scale, - out_bf16); + w8a8_dumma_m16_n16_sk4_lds_fused_kernel, + dim3(static_cast(kSkGrid)), dim3(kWaveSize), 0, stream, + a, b, partials, counters, x_scale, weight_scale, + reinterpret_cast(out)); return; } - // Workspace smaller than the split-K=4 partial planes: fall through to - // the generic scalar fallback (correct for every (m, n, k)); never run - // the tuned pipeline with an undersized workspace. } - const int64_t total = static_cast(m) * n; - const unsigned blocks = static_cast( - (total + kScalarBlockThreads - 1) / kScalarBlockThreads); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_scalar_gemm_kernel), - dim3(blocks), - dim3(kScalarBlockThreads), - 0, - stream, - a, - b, - x_scale, - weight_scale, - out_bf16, - m, - n, - k); + // Generic scalar fallback for every unmatched (m, n, k). + launch_w8a8_gemm_scalar(a, b, x_scale, weight_scale, out, m, n, k, stream); } -extern "C" void launch_pack_w8a8_weight( - const int8_t* raw_weight, - const float* weight_scale, - int8_t* packed_weight, - float* packed_weight_scale, - int k, - int n, - hipStream_t stream) { +extern "C" void launch_pack_w8a8_weight(const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, + int k, + int n, + hipStream_t stream) { + // Exact-shape n-major pack (iteration 1): only (k, n) == (1024, 4096) gets + // the transposed fragment layout consumed by the exact-shape DUMMA kernel + // (and decoded by the scalar fallback for the same (k, n)); every other + // (K, N) keeps the byte-identical identity copy so the generic scalar + // fallback stays correct for all other shapes. const int64_t weight_bytes = static_cast(k) * n; - const int64_t scale_bytes = static_cast(n) * sizeof(float); - constexpr int kPackThreads = kScalarBlockThreads; - - // Exact-shape n-major pack (iteration 6): only (k, n) == (1024, 4096) - // gets the transposed fragment layout consumed by the tuned partial - // kernel; every other (K, N) keeps the byte-identical identity pack so - // the generic scalar fallback stays correct for all other shapes. - if (k == kTargetK && n == kTargetN) { - const unsigned pack_blocks = static_cast( - (weight_bytes + kPackThreads - 1) / kPackThreads); + constexpr int64_t kMaxPackBlocks = 65535; + int64_t weight_blocks = (weight_bytes + kScalarThreads - 1) / kScalarThreads; + if (weight_blocks > kMaxPackBlocks) { + weight_blocks = kMaxPackBlocks; + } + if (weight_blocks < 1) { + weight_blocks = 1; + } + if (k == kExactK && n == kExactN) { hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_pack_nmajor_bytes_kernel), - dim3(pack_blocks), - dim3(kPackThreads), - 0, - stream, - raw_weight, - packed_weight, - k, - n); + w8a8_pack_nmajor_bytes_kernel, + dim3(static_cast(weight_blocks)), dim3(kScalarThreads), 0, + stream, raw_weight, packed_weight, k, n); } else { - const unsigned weight_blocks = static_cast( - (weight_bytes + kPackThreads - 1) / kPackThreads); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_pack_identity_bytes_kernel), - dim3(weight_blocks), - dim3(kPackThreads), - 0, - stream, - reinterpret_cast(raw_weight), - reinterpret_cast(packed_weight), - weight_bytes); + launch_identity_copy_int8(raw_weight, packed_weight, weight_bytes, + stream); } - - const unsigned scale_blocks = static_cast( - (scale_bytes + kPackThreads - 1) / kPackThreads); - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_pack_identity_bytes_kernel), - dim3(scale_blocks), - dim3(kPackThreads), - 0, - stream, - reinterpret_cast(weight_scale), - reinterpret_cast(packed_weight_scale), - scale_bytes); + launch_identity_copy_f32(weight_scale, packed_weight_scale, + static_cast(n), stream); } // @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip index 7ff2b06e..e0c06dd2 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip @@ -1,757 +1,1123 @@ -// @@variant shape=hy3_tp8_qkv_proj_m16 commit=467827c4a33a7e06fe4b3d9871cad428d75648fe added=2026-08-27 -// median_us=23.17 p90_us=23.21 -// source=hy3-dsh-tp8-m16-2-dc53295a -// INT8 W8A8 GEMM for Hygon K500SM_AI / gfx928. +// @@variant shape=hy3_tp8_qkv_proj_m16 commit=f2b771e350c9fea3b13d59b5fec4219786976bcc added=2026-09-09 +// median_us=19.3 p90_us=19.32 speedup=3.875 baseline_us=74.78 +// source=hy3-dsh-tp8-m16-9-8-0161e718 +// W8A8 INT8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. // -// Iteration 1 (DUMMA bootstrap): replace the correctness-first scalar kernel -// for the assigned shape hy3_tp8_qkv_proj_m16 (M=16, N=1280, K=4096) with the -// minimal native gfx928 DUMMA INT8 m16n16k32 tile: one 64-lane wavefront per -// block, one 16x16 output tile per block (grid = N/16 = 80 blocks), direct -// global-to-fragment loads (no LDS staging), explicit int32 accumulation, and -// a single-wave LDS epilogue with no cross-wave barrier. Measured 176.98 us. +// Worker: worker_0 (physical GPU 0), shape hy3_tp8_qkv_proj_m16 +// M = 16, N = 1280, K = 4096 (logical weight layout [K, N]). // -// Iteration 2 (architecture round: launch geometry): keep the direct-load -// data path unchanged and change the launch geometry to 2 wavefronts per -// block (128 threads) with in-block split-K=2. Each wavefront accumulates -// exactly one K-half (2048 = 64 m16n16k32 steps) in int32 k-ascending order, -// publishes its int32 partial to its own LDS plane, and after one END-of-K -// barrier the two wavefronts add partial 0 + partial 1 (ascending split -// order, bit-identical to the unsplit k-ascending int32 accumulation) and -// emit the scaled bf16 tile. grid stays N/16 = 80 blocks. Measured -// 102.12 us median / 105.59 us p90 (vs the 74.782 us Triton Graph baseline); -// the ~1.73x geometry gain shows each active CU still runs one serial -// load->wait->MMA chain per wave (64 steps after split-K=2), with the two -// chains of one block co-scheduled on the same CU. +// Iteration 0 (bootstrap): correctness-first scalar kernel (one thread per +// output element) served every shape and measured 336 us on this shape. // -// Iteration 3 (architecture round: workspace split-K pair). Mandate: the -// unsplit grid (80 blocks) has fewer than two blocks per device CU (120 CUs) -// and K=4096 >= 1024, so implement split-K=2 plus at least one CU-aligned -// candidate (non-power-of-two allowed), writing int32 partials into the -// caller workspace and including the combine+scale kernel in the timed Graph. -// Replace the iteration-2 single-launch in-block-combine kernel with a -// workspace split-K pair (the pattern validated on the hy3 TP4 qkv lineage, -// which took the same architecture round at its iteration 3): +// Iteration 1 (DUMMA bootstrap, this round): replace the scalar kernel for +// the exact guard (m, n, k) == (16, 1280, 4096) with the minimal native +// gfx928 DUMMA INT8 m16n16k32 tile per the round mandate: +// - one 64-lane wavefront per block, one 16x16 output tile per block +// (grid = N/16 = 80 blocks), no cross-wave barrier; +// - A and B fragments are loaded DIRECTLY from their row-major global +// tensors every 32-K step (no LDS staging, no split-K, no workspace, +// no combine kernel in the timed region); +// - explicit int32 accumulation: du_fill_fragment(acc, 0) then one +// du_mma_sync per ascending 32-K step (exact int32 dot, same sum as +// the scalar kernel since int32 addition is exact and order-free); +// - single-wave LDS epilogue: du_store_matrix_sync materializes the +// 16x16 int32 tile in LDS (library-defined ownership), one +// __syncthreads() reached by all 64 threads, then the 64 lanes scale +// with x_scale[row] * weight_scale[n0+col] (scalar-path float order) +// and store bf16 through the manual round-to-nearest-even conversion. +// This establishes the exact DUMMA API family below so later rounds can +// tune geometry (wavefronts/block, LDS staging, split-K planes). // -// 1. w8a8_dumma_m16_splitk_partial_kernel: one 64-thread -// wavefront per block, grid = (N/16 = 80 tiles, SPLIT_K K-slices). -// Each block computes the int32 partial dot for one 16x16 output tile -// over one 32-aligned K slice and du_store_matrix_sync's the partial -// straight into the caller workspace. ZERO barriers, ZERO LDS, ZERO -// atomics in the main kernel. -// 2. w8a8_dumma_m16_combine_scale_kernel: 80 blocks x 256 threads, one -// thread per output element; exact int32 sum of the SPLIT_K partials -// (ascending split order), then x_scale/weight_scale multiply and the -// bf16 (RN-even) epilogue. Barrier-free (each thread reads only its -// own 16x16 element's S partial values, written by the previous kernel -// on the same stream). +// Iteration 2 (architecture round, this round): measure grid parallelism +// before polishing. The iteration-1 geometry (80 one-wave blocks over 120 +// CUs, 0.67 blocks/CU) is far below the two-blocks-per-CU latency-hiding +// target, and the measured 130.43 us median (vs 74.782 us Triton Graph +// baseline) is dominated by underfilled CUs and per-block serial load->mma +// chains (128 direct-load 32-K steps per block). This round replaces the +// exact-guard dispatch with a grid-level split-K=3 geometry, still one +// 64-lane wavefront per block and one 16x16 N tile per wave (zero in-block +// barriers): +// - grid = (N/16) * SPLIT_K = 80 * 3 = 240 blocks = exactly 2 blocks per +// device CU (120 CUs), split id on gridDim.y, tile id on gridDim.x; +// - the 4,096-long K axis is partitioned into three legal 32-aligned, +// non-uniform, ascending slices of {1376, 1376, 1344} rows (43/43/42 +// m16n16k32 steps), each block du_fill_fragment(acc,0) then runs its +// own ascending 32-K DUMMA chain over its slice only (bit-exact int32 +// partial dot: |partial| <= 6.6e7 << 2^31, order-free); +// - each block materializes its 16x16 int32 partial (LDS epilogue, same +// single-wave store/barrier pattern as iteration 1) into its own plane +// of the caller's int32 workspace: partials[split][16][N]; every launch +// overwrites every plane, so no workspace clear is needed; +// - a small combine kernel (80 blocks x 64 lanes, second launch on the +// same stream inside the timed API) sums the SPLIT_K planes ascending +// per output element (exact int32, identical to the scalar k-ascending +// sum) and applies x_scale[row] * weight_scale[col] then bf16 RNE, +// keeping every output bit-identical to the CPU reference. +// Split-K combine is part of the operator wall and stays inside the timed +// region/Graph. Workspace need = SPLIT_K*M*N*4 = 245,760 B, well under the +// contract's 16-plane 1,310,720 B allocation for this shape. // -// Both kernels launch on the caller's stream inside the timed Graph, so the -// combine cost is included in the measured operator latency. +// Iteration 3 (architecture round, this round): continue the mandated +// split-K occupancy sweep. Iteration 2's falsifiable claim held (240 +// one-wave blocks = 2 blocks/CU cut the wall from 130.43 to 68.65 us +// median, 1.089x over the 74.782 us Triton Graph baseline), but the +// measured logical rate (86 GB/s, 2.44 TOPS) is still far below the +// headroom of this decode family, so the sweep goes finer inside the +// trusted CU-aligned probe set {2,3,4,5,6,8,9,12}: SPLIT_K=6 -> grid = +// 80 N tiles x 6 K-slices = 480 one-wave zero-in-block-barrier blocks = +// exactly 4 blocks per device CU, and each block's serial 32-K chain +// halves to {22,22,21,21,21,21} steps = rows {704,704,672,672,672,672} +// (the existing generic slice formula; disjoint, ascending, 32-aligned, +// exact cover of [0,4096)). The partial and combine kernels are +// UNCHANGED code -- both were already generic over gridDim.y / +// split_count -- so this round only moves the launcher geometry constant +// kSplitK from 3 to 6. Workspace need = 6*M*N*4 = 491,520 B, still +// inside the caller's 16-plane 1,310,720 B allocation; the combine kernel +// (80 blocks x 64 lanes) sums the six planes ascending per output element +// (exact int32 -> bit-identical to the scalar k-ascending sum) and stays +// in the timed Graph. // -// Iteration 4 (register prefetch, rejected): double-buffered a_frag/b_frag in -// registers one K step ahead. Regressed to 120.5 us median: the exact code -// object kept the full vmcnt drain before each v_mmac, so prefetching could -// not remove the per-step global round trip from the MMA critical path. -// Iterations 5-6 were killed by the agent infrastructure (no evidence). +// Iteration 4 (architecture/pipeline round, this round): multi-N-tile +// reuse. The iteration-3 code object (PMC digest 26d6df45...) shows the +// poison direct-load pattern per 32-K step: 16 global_load_ubyte (8 per +// A/B fragment) with ~10 scattered s_waitcnt vmcnt before byte reassembly +// and one v_mmac_i32_16x16x32_i8, and no cross-iteration load overlap, so +// at 4 resident one-wave blocks/CU every block's serial chain is one +// ~load-latency step per m16n16k32 mma (52.19 us median; the wall is the +// per-step latency of a one-mma step, not raw issue or HBM traffic). +// This round keeps the same split-K partial + combine architecture and +// the same 4 blocks/CU co-residency but widens each one-wave block to TWO +// adjacent 16-wide N tiles (32 columns): per 32-K step the wave loads the +// A fragment ONCE and issues two du_mma_sync (independent accumulators), +// so the x_q A bytes [16 x 32] of the shared step are reused across the +// two output tiles (per-output A global loads/reassembly halve) while +// each B byte is still read exactly once (no B reuse). To keep 480 +// one-wave blocks = exactly 4/CU over 120 CUs the grid becomes 40 N +// groups x SPLIT_K=12 K-slices ({11x8, 10x4} 32-K steps = rows +// {352x8, 320x4}, disjoint ascending exact cover of [0,4096), every +// boundary a DUMMA-k multiple; existing generic slice formula). Per-block +// serial chain drops to <= 11 steps of TWO mmas, so if the iteration-3 +// wall is per-step load->mma latency the per-output step latency halves +// and the unprofiled Graph median/P90 must drop materially below +// 52.194/52.269 us. Workspace need = 12*M*N*4 = 983,040 B, still inside +// the caller's 16-plane 1,310,720 B allocation; the combine kernel (80 +// blocks x 64 lanes) is UNCHANGED code and sums the twelve planes +// ascending per output element (exact int32 -> bit-identical to the +// scalar k-ascending sum), staying in the timed Graph. The only edits +// are the widened per-wave N tile in the partial kernel plus the launcher +// geometry; the workspace-less fallback keeps the iteration-1 one-wave +// direct kernel byte-for-byte and the scalar fallback still serves every +// unmatched shape including the paired M=2 API shape. // -// Iteration 7 (architecture round: LDS slice staging). The exact gfx928 code -// object of the iteration-3 kernel shows every one of the 64 K steps compiles -// to 16 narrow global_load_ubyte + byte-reassembly VALU + a full progressive -// vmcnt drain before the single v_mmac: the whole global round trip is paid -// serially per step with only ~1.3 waves/CU to hide it. This round applies -// the TP4 qkv lineage remedy (validated 26.05 us on N=2560) scaled to this -// shape: stage the block's entire 32-aligned K slice into LDS once with -// coalesced dword global loads, one __syncthreads, then a zero-barrier -// LDS-only K loop (du_load_matrix_sync reads both fragments from LDS, so the -// MMA path sees only ~30-cycle LDS latency instead of a global round trip -// per step): +// Iteration 5 (packed-weight/staging round, this round): B-only LDS slice +// staging. The iteration-4 code object (PMC digest 1a4e30d0..., vmem_read +// 122,880 = 24 byte-granular global_load_ubyte per 32-K step x 5120 +// block-steps) shows the remaining wall is the poison direct-global B half +// of the per-step load->wait->reassembly->mma chain: per step the wave +// still pays 16 global byte loads (B0/B1) plus ~10 scattered s_waitcnt +// vmcnt on the MMA critical path even though every byte is L2-hot in the +// timed Graph replay. Per the lineage rule, direct-global fragment loads +// are poison: this round removes the B half by staging each block's whole +// 32-column B slice (<= 352 rows x 32 B = 11,264 B) into LDS ONCE at block +// start with 16-B int4 cooperative loads (64 lanes x 16 B per wavefront +// instruction, 4-deep source batching so one vmcnt drain covers 4 +// in-flight round trips), one same-wave __syncthreads(), then every +// 32-K step reads the B0/B1 fragments from LDS with du_load_matrix_sync +// (~30-cycle LDS latency, no global round trip) at row stride 32 -- the +// staged tile is an exact byte mirror of the block's global B strip, so +// fragment contents, the k-ascending int32 order and the partial planes +// are bit-identical to iteration 4. A stays direct-global on purpose: +// this round isolates the B-side hypothesis (A-only staging and the +// packed [N,K] layout remain untested axes). Geometry is UNCHANGED (40 +// N groups x SPLIT_K=12 = 480 one-wave blocks = exactly 4 blocks/CU over +// 120 CUs; per-block LDS 11,264 B + 2 x 1,024 B acc tiles = 13,312 B, +// 65,536/13,312 = 4.9, so the 4/CU co-residency is preserved and LDS does +// not become the occupancy limiter). If the per-step wall is the B-load +// half of the global wait ladder, the official unprofiled Graph median/P90 +// must drop materially below 41.154/41.233 us (partial-kernel profiled +// 36.32 us); if instead the median stays flat or regresses, the wall at 4 +// waves/CU is A-side latency or per-CU issue (16 vs 24 loads immaterial), +// and the next round must stage A too (LDS-only K loop) or raise +// co-residency -- never a scalar-fallback change. Exactness is preserved +// for every shape: the (m,n,k)==(16,1280,4096) guard is unchanged; the +// staged bytes are a copy of the same global bytes the direct loader +// read, the workspace-less fallback keeps the iteration-1 direct kernel +// byte-for-byte (still correct for this shape), the scalar fallback still +// serves every unmatched shape including the paired M=2 API shape, and no +// pack/layout/env/stream/config change is made; all launches stay on the +// caller stream inside the timed API/Graph. // -// w8a8_dumma_m16_slicestage_partial_kernel: one 64-thread -// wavefront per block, grid = (N/16 = 80 tiles, SPLIT_K K-slices), same -// per-tile 16x16 output and the same workspace int32 plane layout, so the -// combine+scale kernel below is unchanged. A (16 x slice) and B -// (slice x 16) are staged into LDS with the same 32-aligned non-uniform -// slice bounds and the same k-ascending int32 MMA order as iteration 3, so -// outputs stay bit-identical (0 mismatches) and Graph capture/replay with -// changed contents must pass. +// Iteration 6 (infrastructure failure, no valid artifact): the round-6 +// attempt was killed by the agent infrastructure before it produced a +// candidate result; no round-6 source change was accepted and the session +// was reset to the iteration-5 accepted best source below. The failure does +// not count as a completed optimization round. // -// Round-8 geometry (default): split-K=16 -> grid = (80 tiles, 16 slices) = -// 1,280 one-wave blocks = 10.67 blocks/CU on the 120-CU device. Each block -// stages exactly 256 K (A 16x256 -> 4,160 B + B 256x16 -> 5,120 B = 9,280 B -// LDS = 7 resident blocks/CU by the 64 KiB LDS) and runs exactly 8 uniform -// zero-barrier LDS MMAs (128/16 = 8, so no 11-step stragglers). This is an -// occupancy (LDS-footprint) tune against the round-7 PMC evidence: at -// split-K=12 the profile shows grid 960 = 8 waves/CU available but -// 12,736 B LDS/block caps residency at 5 blocks/CU (62.5%), with VGPR 24 / -// scratch 0 / SGPR 32 proving registers and spills are NOT limiters -- the -// per-block LDS footprint is the binding occupancy limiter, and S=16 is the -// largest split that fits the 16-plane contract workspace (1,310,720 B -// exactly), so residency rises to 7 of 10.67 waves/CU. A bytes are reused -// across the 80 tile blocks of each slice (1,280 x 16 x 256 = 5.24 MB total -// A reads vs 64 KB unique, L2-hot); B 16-column strips are read exactly once -// (5.24 MB unique), with the 64-B L2 sectors shared with the adjacent -// 16-column tiles, so HBM A/B traffic is byte-identical to split-K=12 -// (~5.30 MB unique; no repeated HBM reads traded for occupancy -- only the -// split-K combine workspace round trip grows from 12 to 16 planes, +0.33 MB -// write + read). vmem_read_instructions stay ~40,960/replay (slightly fewer -// than S=12's 42,240: each 256-K slice needs 32 wavefront loads vs 44). +// Iteration 7 (staging round, this round): A+B whole-slice LDS staging +// (LDS-only K loop). The iteration-5 result (39.404/39.485 us official +// Graph median/P90 vs iteration-4 41.154/41.233) falsified its own claim: +// the predicted 25-33 us for the B-half hypothesis did not appear, so per +// the iteration-5 contingency text this round stages the shared A operand +// too. The iteration-5 PMC (source digest 5fb42aff..., grid 480 = 40 N +// groups x SPLIT_K=12, vmem_read 46,080, l2_misses 98,879, counter-derived +// HBM ~187.9 GB/s << peak) still shows 8 byte-granular global_load_ubyte of +// the shared A fragment plus ~8 scattered s_waitcnt vmcnt before shift/or +// reassembly on the MMA path every 32-K step (~40,960 of the 46,080 +// vmem_read over 5,120 block-steps): pure per-step A-side global latency +// serialization, not traffic. This round stages BOTH operands once per +// block with 16-B int4 cooperative loads (A slice: 16 m-rows x 256 B + 16 B +// pad/row = 4,352 B at row stride 272 for bank skew; B strip: 256 x 32 = +// 8,192 B), one same-wave __syncthreads(), then an LDS-only K loop (A +// fragment at row stride 272, B0/B1 at row stride 32 via +// du_load_matrix_sync) with zero per-step global loads/waits. Staged tiles +// are exact byte mirrors of the global operands, so fragment bytes, the +// k-ascending int32 order and the partial planes are bit-identical to +// iteration 5. Geometry forced by LDS capacity: a whole-slice A stage does +// not fit at SPLIT_K=12 with 4 resident blocks/CU (A 16 x 352 = 5,632 B + +// B 11,264 B + acc 2,048 B = 18,944 B -> floor(65,536/18,944) = 3), so the +// guard launch moves to SPLIT_K=16: sixteen UNIFORM 8-step slices (128 == +// 16 x 8, rows {256 x 16}; every boundary a DUMMA-k multiple), grid = 40 N +// groups x 16 = 640 one-wave blocks ~ 5.33/CU over 120 CUs, per-block LDS +// 4,352 + 8,192 + 2,048 = 14,592 B -> floor(65,536/14,592) = 4 resident +// blocks/CU, i.e. the iteration-5 co-residency is preserved. Workspace +// need = 16 int32 planes = 1,310,720 B = the caller's full fixed 16-plane +// allocation for this shape, still gated by the launcher's workspace_bytes +// >= kWorkspaceNeeded check (undersized workspaces keep the workspace-less +// fallback). The combine kernel (80 blocks x 64 lanes, second same-stream +// launch in the timed API/Graph) is UNCHANGED code, generic over +// split_count = 16, and sums the sixteen planes ascending per output +// element (exact int32 -> bit-identical to the scalar k-ascending sum), +// staying in the timed Graph. All fallbacks, guards, pack/layout/env/ +// stream/config behavior are unchanged. // -// Iteration 11 (packed col-major B; the validated TP4 gate_up/down_proj -// recipe). The accepted S=16 staged kernel's fresh PMC shows -// lds_bank_conflicts 81,920 == the 8 scalar ds_read_u8 per matrix_b fragment -// (8 steps x 1,280 blocks), i.e. every B fragment read in the zero-barrier K -// loop conflicts at the 20-B k-major pitch, while round 9 falsified -// staging-load MLP batching (58.56 us), so the K-loop B fragment reads are -// the remaining hot-path lever. launch_pack_w8a8_weight now transposes the -// exact (k,n)==(4096,1280) weight to [N,K] n-major once, outside the timed -// region and out of Graph capture (every other (K,N) keeps the identity -// copy); the staged kernel stages its 16-column strip as 16-B int4 straight -// copies into an n-major LDS tile b_s[16][kMaxKlen+16] (row stride -// kMaxKlen+16: 16-B aligned for the int4 stores, 8-B aligned for the -// fragment b64 reads, not a multiple of 128 B -> no bank-phase aliasing), so -// each lane's 8 fragment bytes are consecutive and load with ONE ds_read_b64 -// (load_b_frag_packed) instead of 8 conflicting ds_read_u8. The direct -// S in {2,3,4,5} path reads the packed global layout through the library -// col_major loader, and the scalar fallback decodes the packed layout for -// (k,n)==(4096,1280) (covers the paired M=2 validation shape). Fragment -// contents, the k-ascending int32 du_mma order, the slice-major int32 plane -// layout, the combine+scale kernel and the generic fallback are all -// unchanged, so outputs stay bit-identical (0 mismatches) and Graph -// capture/replay with changed contents must pass. Per-block LDS drops -// 9,280 -> 8,512 B (65,536/8,512 = 7.70, still 7 resident/CU); A/B HBM bytes -// are unchanged (B stays read exactly once: 16-B contiguous per n-row, each -// 64-B sector read once at the uniform S=16 default). +// Iteration 8 (occupancy/resource round, this round): raise the waves-per- +// block occupancy limiter from 1 to 2 wavefronts per block (blockDim 64 -> +// 128). PMC of the accepted iteration-7 code object (source digest +// b64e3209..., grid 640, workgroup 64, arch_vgpr 40, sgpr 48, scratch 0, +// LDS/block 14,592 B -> floor(65,536/14,592) = 4 resident blocks/CU = +// 5.33 waves/CU by grid) shows the residual wall is per-SIMD single-wave +// latency exposure, not issue or HBM traffic: each resident wave is alone +// on its SIMD (4 blocks x 1 wave -> 4 SIMDs), so every per-step +// s_waitcnt lgkmcnt after the staged-fragment du_load_matrix_sync and every +// du_mma accumulator dependency stalls that SIMD with no second independent +// wave to issue (VALU issue utilization ~3-6% of the 26.56 us profiled +// kernel; only 12 one-time staging global loads per block, vmem_read +// 7,680). The fix keeps the iteration-7 geometry byte-for-byte in every +// respect that affects partials (SPLIT_K=16, 40 x 16 = 640 blocks, 8-step +// slices rows {256 x 16}, staged exact byte mirrors, k-ascending int32 +// order, 16-plane workspace) and splits the one-time staging ACROSS the two +// waves (wave 0 copies the A slice, wave 1 copies the B strip; exactly the +// same global bytes read once -- no repeated HBM reads traded for +// occupancy), then each wave runs its OWN 8-step accumulator chain for one +// n16 tile of the block's 32 columns (one du_mma per step per wave; total +// MMA count unchanged). LDS per block stays 14,592 B -> still 4 resident +// blocks/CU, but now 8 resident waves/CU = 2 per SIMD, so a stalling wave +// is covered by its co-resident sibling. The only deliberately added cost +// is LDS re-reads: the shared A fragment is now loaded by both waves +// (predicted lds_instructions 99,840 -> ~125-135k with conflicts scaling at +// the measured ~2.56/instr); LDS is not the wall at ~3-6% issue +// utilization. The combine kernel, guards, workspace-less fallback, scalar +// fallback and pack identity are all unchanged. // -// Iteration 12 (staging-load MLP batching, alignment-safe). The fresh -// exact-source ISA of the accepted iteration-11 object shows both staging -// loops still compile to a fully serialized load -> s_waitcnt vmcnt(0) -> -// store chain (A: 16 sequential global round trips per block; B: 4), i.e. -// every block pays ~20 serial global round trips before the single barrier, -// while the K loop is already lean (2 LDS reads + 2 lgkmcnt waits + 1 v_mmac -// per step). Iteration 11 already validated the B-side fix (16 serialized -// dword loads -> 4 aligned int4 loads as part of the accepted 40.38 -> -// 30.66 us round). This round batches both staging loops 4-deep in source: -// 4 independent loads are issued before their 4 stores, so one vmcnt drain -// covers 4 in-flight round trips (A 16 -> 4, B 4 -> 1; ~20 -> ~5 per block). -// A stays dword-wide because kAStride = kMaxKlen+4 is not 16-B aligned -- -// round 9's 58.56 us regression was misaligned int4 staging on the old -// layouts, not an MLP falsification; B stays int4 (aligned). No layout, -// stride, fragment, slice, plane, combine, guard or fallback changes: -// outputs stay bit-identical (0 mismatches) and Graph capture/replay with -// changed contents must pass. vmem_read/lds instruction counts are -// unchanged (25,600 / 46,080); the wait structure is what changes. +// Iteration 9 (combine restructure round, this round): attack the OTHER half +// of the operator wall. PMC of the accepted iteration-8 code object (source +// digest beb25b76..., official Graph median/P90 28.884/28.933 us) profiles +// BOTH timed dispatches of the operator replay: the 640-block partial kernel +// at 19.361 us AND the 80-block combine kernel at 16.960 us -- the combine is +// 46.7% of the 36.321 us profiled operator aggregate, not a small tail. Its +// structure explains the cost: 80 one-wave 64-lane blocks (<= 80 of 480 +// SIMDs busy, ~40 CUs fully idle) whose lanes each run FOUR output elements +// through a DEPENDENT per-plane chain -- one 4-B global load + s_waitcnt +// vmcnt + int32 add per plane, 16 serial DRAM round trips per element (the +// exact code object's static body is one load/wait/add per plane iteration) +// -- against plane data that is L2-cold at combine start (combine +// l2_hit_rate 6.9%, l2_misses 21,596 ~= the whole 1.31 MB of partial planes, +// evicted by the partial kernel's 5.37 MB streaming read): DRAM-latency +// serialization on only 1.34 MB of reads, plus underfill. This round +// restructures ONLY the combine kernel (the partial kernel, SPLIT_K=16 +// 640-block geometry, staged byte mirrors, k-ascending int32 partials, +// 16-plane 1,310,720 B workspace contract, workspace-less fallback, exact +// guard and scalar fallback are byte-for-byte unchanged): the guard launcher +// switches the combine grid from 80 x 1 one-wave blocks to a 2-D grid of 80 +// n-tiles x 4 four-row m-groups = 320 one-wave blocks, and the kernel's +// split_count == 16 specialization gives every lane exactly ONE output +// element that issues ALL sixteen plane loads as INDEPENDENT reads into a +// 16-int register array (one vmcnt drain, 16-way memory-level parallelism +// per lane), then sums the registers in ascending split order s = 0..15 -- +// the identical int32 sequence as the old ascending loop -- and applies the +// unchanged float scale order and f32_to_bf16_rne, so every output element +// stays bit-identical (this mirrors the validated gate_up S=16 combine +// recipe: issue all 16 plane loads into a register array before the +// ascending sum). Same stream, same two launches, no allocation or sync +// inside the timed API/Graph. // -// The ZTH_W8A8_QKV_SPLIT_K environment override keeps its semantics but now -// selects the data path too: S in {6,8,9,12,16} run the staged kernel (their -// slices fit the 64 KiB LDS), S in {2,3,4,5} keep the iteration-3 direct-load -// kernel (S=2's 2,048-K slice alone would need 64 KiB for B, and S=3's -// 1,408-K slice does not fit either). +// Iterations 10-15 (five-wave occupancy arc on this guard, all archived at +// iterations/hy3_tp8_qkv_proj_m16/iterationN): iteration 10 (direct-register +// epilogue on the two-wave mapping: acc tiles + epilogue barrier removed, +// LDS 14,592 -> 12,544 B = 5 resident blocks/CU = 10 waves/CU) measured +// 20.827/21.585 us and was NOT accepted; iteration 11 raised waves-per-block +// 2 -> 5 (blockDim 320, grid 16 x 16 = 256 blocks, LDS 29,952 B = 2 resident +// blocks/CU = 10 waves/CU) but FAILED exact correctness (its source declared +// no __launch_bounds__ while the launcher requests 320 threads, so the +// compiled descriptor capped the workgroup at 256 and the fifth wavefront's +// tile was absent: 4,091/20,480 mismatches, first at (m,n)=(0,64)); the +// iteration-11 repair was completed in iteration 13 by declaring +// __launch_bounds__(320) (exact correctness passed, 0 mismatches; measured +// 23.764/24.086 us); iteration 14 (ping/pong fragment double-buffer) passed +// correctness but the toolchain spilled the second fragment set to the +// private segment on every K step (.private_segment_fixed_size 48) -> 39.472 +// us, closing the register-pipelining axis; iteration 15 replaced the +// five-wave epilogue __syncthreads() with __threadfence_block() (the acc +// tiles are wave-private and the waves' plane stores are disjoint, so the +// cross-wave join was pure skew cost) and measured 20.613/20.641 us, correct +// but still not below the accepted two-wave 20.179/20.216. // -// Round-3 geometry (direct-load path, still selectable for S in -// {2,3,4,5}): split-K=2 -> 160 -// one-wave zero-barrier blocks -// (160 independent wavefronts = the same total parallelism as iteration 2's -// 80 blocks x 2 waves, but with no in-block barrier, no LDS partial planes, -// and 160 independent block-level load streams that the block scheduler can -// co-resident 2-deep on 40 of the 120 CUs). The trusted occupancy-probe -// sweep SPLIT_K in {2,3,4,5,6,8,9,12,16} is implemented as template -// instantiations and selectable at launch time with the -// ZTH_W8A8_QKV_SPLIT_K environment variable, so the control plane can -// measure split-K=2 and the CU-aligned candidates without source edits. -// For the direct path only, the CU-aligned (integer blocks/CU) candidates -// for N/16=80 tiles were S=3 -> 240 blocks = exactly 2 blocks/CU -// (non-power-of-two), S=6 -> 480 = 4/CU, S=9 -> 720 = 6/CU, S=12 -> 960 = -// 8/CU; with the iteration-7/8 routing those S values now run the staged -// kernel instead (block counts 480/720/960/1,280 for S=6/9/12/16), and -// S in {2,3,4,5} keep the direct path. Every candidate fits the contract -// workspace (1,310,720 bytes = capacity 16 planes of 81,920 B; S=16 uses -// exactly all 16 planes) and keeps -// every K slice a multiple of the DUMMA K tile 32 (non-uniform slices differ -// by at most one 32-K step; S=16 is uniform: 128/16 = 8 steps). +// Iteration 16 (this round): close the five-wave occupancy axis with the +// measured occupancy evidence and apply the repair arc's validated epilogue +// discovery to the accepted mapping. Occupancy ledger of the partial +// kernel, all correct and bit-identical: 8 resident waves/CU (4 x 2-wave +// blocks, LDS 14,592 B) = 20.179/20.216 us (iteration 9, accepted); the two +// independent 10 resident-waves/CU configurations -- iteration 10 (5 x +// 2-wave blocks) 20.827/21.585 and iteration 15 (2 x 5-wave blocks) +// 20.613/20.641 -- BOTH trail the 8-wave configuration, i.e. per-SIMD wave +// density 2.5 (some SIMDs carry 3 waves while others idle) is worse than a +// clean 2.0 on this kernel, and the five-wave mapping (2 or 1 resident +// 5-wave blocks/CU = 10 or 5 waves/CU only) can never reach the 8-wave +// sweet spot. The accepted iteration-9 source below still carries the SAME +// provably-unnecessary epilogue __syncthreads() cross-wave join whose +// five-wave form cost ~3.15 us (iteration 13 23.764 -> iteration 15 20.613 +// with the join removed). Fix (the ONLY code change of this round; the +// two-wave geometry, blockDim 128, grid 40 x 16 = 640 blocks, staging +// partition, stage byte mirrors, 8-step K loop, combine kernel and its +// 320-block launch, workspace, guards and fallbacks are byte-for-byte the +// accepted iteration-9 candidate, verified by diff): replace the partial +// kernel's epilogue __syncthreads() with __threadfence_block(), exactly as +// iteration 15 validated on the five-wave form of this same kernel. +// acc_tile0/acc_tile1 are each written by du_store_matrix_sync and read +// back only by the SAME wave (w == 0 -> acc_tile0, w == 1 -> acc_tile1), and +// the two waves' plane stores touch disjoint 16-column ranges, so no +// epilogue stage has any cross-wave dependency; __threadfence_block() still +// orders each wave's own outstanding LDS writes before its read-back +// ds_reads (wavefront-wide lgkmcnt drain) while costing zero cross-wave +// wait. The staging __syncthreads() above stays (wave 0 stages the A slice +// that wave 1 also reads). Every partial byte and combine output byte is +// unchanged, so the output is bit-identical to iteration 9. // -// Exact int32 accumulation is preserved: each slice accumulates int32 over an -// ascending 32-aligned K range; slices tile [0, K) exactly once; the combine -// sums the partials in int32 ascending split order (per-tile max |dot| = -// 4096*127*127 = 66,064,384 << 2^31), so every per-tile dot total is -// bit-identical to the scalar fallback and to iterations 1-2 regardless of -// SPLIT_K. The scale + bf16 (RN-even) epilogue runs once per output element -// in the combine kernel. +// Iteration 18 (this round): compile-time 8-step K-loop specialization of +// the partial kernel, applied to the current checked-in source (the +// iteration-16 candidate; working-tree digest cbc40932... = the iteration-16 +// archive, verified by diff). Round 17 drafted exactly this edit and +// verified its codegen locally, but the agent session was killed by the +// round infrastructure before a valid proposal landed -- no source change +// was ever applied, no build/benchmark ran, and per policy that +// infrastructure failure does not count toward the valid-round ledger, so +// this round re-submits the same HIP-only change as a fresh valid round. +// The iteration-11 repair item is closed since iteration 13 (archived +// iteration-11 five-wave source + __launch_bounds__(320) ONLY, exact-correct +// with 0 mismatches; iterations 13-16 completed that arc and iteration 16 +// re-based on the accepted two-wave mapping). Status of the current source: +// the iteration-16 epilogue barrier-decoupled two-wave partial kernel +// (blockDim 128, grid 40 x 16 = 640 blocks, LDS 14,592 B = 4 resident +// blocks/CU = 8 resident waves/CU) measured official Graph median/P90 +// 19.9885/20.0133 us (min 19.9557 / max 21.3734, correctness 0 mismatches, +// Graph capture and paired M=2 fallback passed), improving the accepted +// 20.179/20.216 by 0.95%; it is kept as the experimental base ("official +// best is unchanged until cumulative improvement reaches the normal +// acceptance threshold"), and this round keeps building on that source. +// The partial kernel dominates the operator wall, and the exact-source code +// object shows its K loop is ROLLED: the generic loop is runtime-bounded +// (my_steps from the runtime k and gridDim.y args), so the backend emits one +// static body of [ds_read2_b32 (A) + 8 x ds_read_u8 (B) + ~25-VALU byte +// reassembly + 4 lgkmcnt drains + one v_mmac_i32_16x16x32_i8] with a +// loop-carried WAR on the single fragment register set: step s+1's LDS +// reads cannot issue until step s's v_mmac has consumed the same fragment +// registers, so every step's full LDS latency sits serially on the +// accumulator chain with only 2 resident waves/SIMD to cover it (reused PMC +// of this exact geometry: lds_instructions 107,520, lds_wait_instructions +// 19,597, lds_bank_conflicts 276,480). The guard launch is the kernel's +// ONLY call site and always runs the uniform 8-step slice (k == 4096, +// gridDim.y == 16 -> my_steps == 8 == kStageMaxSteps, rows == 256 == +// kStageMaxRows, staged capacity exactly filled), so this round specializes +// that case: when my_steps == kStageMaxSteps the K loop runs with a +// COMPILE-TIME trip count of 8 under #pragma unroll; the backend then fully +// unrolls the chain, renames the fragment registers per step and issues +// each step's LDS reads after the previous step's v_mmac has consumed its +// registers, overlapping the next step's LDS round trip with the current +// v_mmac's execution (no loop-carried WAR; verified in the local exact-flag +// code object of this edit: 8 static v_mmac + 72 static ds_reads in the +// specialized path vs 1 static v_mmac / 10 static ds_reads in the rolled +// predecessor, vgpr 34 / sgpr 39 / private_segment 0 / group_segment 14,592 +// B and exactly ONE s_barrier unchanged -> the 4-blocks/CU = 8-waves/CU +// occupancy is preserved). Both branches issue the identical ascending-s +// du_mma sequence into the same accumulator (v_mmac acc operand v[1:4] +// throughout) over the same staged byte mirrors, so every partial byte and +// every combine output byte is bit-identical to iteration 16 and exact +// correctness must pass with 0 mismatches; the generic rolled loop stays +// verbatim as the else path for any non-guard (my_steps != 8) +// configuration. Geometry, staging partition, stage byte mirrors, the +// blockDim 128 / grid 40 x 16 = 640 blocks launch, the epilogue +// (acc-tile store + __threadfence_block + read-back plane store), the +// combine kernel and its 320-block launch, the 16-plane 1,310,720 B +// workspace need, the workspace-less iteration-1 fallback, the exact +// (m,n,k)==(16,1280,4096) guard and the scalar fallback for every unmatched +// shape (including paired M=2) are all unchanged. // -// Launch is guarded by the exact (m,n,k) = (16,1280,4096) shape; every other -// shape (including the paired M=2 shape with the same (N,K)) falls back to -// the generic scalar kernel below. launch_pack_w8a8_weight is an identity -// device-to-device copy for every (K,N) except the exact pair -// (k,n)==(4096,1280), which is transposed once to an [N,K] n-major packed -// layout (iteration 11) so DUMMA col_major B fragments read 8 consecutive -// bytes per lane; the scalar fallback decodes that packed layout for the -// same pair, so every path stays correct against the packed buffer. +// Iteration 20 (DS-side staging-latency round, accepted): the exact-guard +// launch always stages the full 512-chunk B strip (chunks == 512 == 8 x 64 +// lanes), so wave 1's B copy now runs as ONE straight-line 8-deep batched +// group (all eight 16-B source loads issue before the first LDS store, one +// vmcnt drain covers the whole strip) instead of two serialized 4-deep +// iterations -- the second global round trip is removed from every block's +// staging join critical path (official Graph median/P90 19.967 -> 19.541/ +// 19.558 us, accepted). The generic rolled loop stays verbatim as the else +// path for any non-guard (chunks != 512) configuration. // -// Graph-safety: this file only launches kernels on the caller-provided stream -// (PyTorch's current HIP stream). It performs no allocation, compilation, -// autotuning, packing, host synchronization, or device synchronization, and -// touches no buffer other than the caller-provided out/workspace pointers. -// The split-K choice is a host-side static dispatch decision (env override or -// default 16) made once per launch call, never inside the Graph replay, so -// capture/replay determinism holds. +// Iteration 21 (this round): DS-side bank-skew round for the B-fragment +// ds_read_u8 structure (pre-registered in iterations 18 and 20: the next +// increment must attack the PMC-documented DS ledger of this exact object -- +// lds_bank_conflicts 276,480 = ~2.57 replays per LDS instruction, dominated +// by the 8 byte-granular B-fragment reads per step per wave, 128 per block). +// The row-major B fragment loader reads each lane's 8 bytes at addresses +// base + (8*(lane>>4) + i) * ldm for i = 0..7, so the four row groups +// g = lane>>4 of the fragment's 32 k rows sit 8*ldm bytes apart; with the +// stage_b rows packed at the old ldm == 32, every group sat 256 B (== 0 +// mod 128 B) apart and each ds_read_u8 replayed the SAME LDS bank 4x per +// wavefront. This round pads each staged B row to kStageBRowStride = 36 B +// (32 data + 4 pad): the groups then sit 288 B (== 32 mod 128 B) apart and +// hit banks {b, b+8, b+16, b+24} -- the B-fragment reads become +// conflict-free while every staged data byte, the fragment contents, the +// k-ascending int32 order and the partial planes stay bit-identical (the +// pads are never written or read; only the loader's ldm argument and the +// wave-1 store offsets change). Wave-1 staging stores become four 4-B +// words per 16-B chunk (36*r is only 4-B aligned) at int32 offsets +// 9*r + 4*h + {0,1,2,3}; LDS store bandwidth is unchanged. LDS per block +// grows 14,592 -> 15,616 B (stage_b 8,192 -> 9,216): 4 x 15,616 = 62,464 +// <= 65,536, so the 4-resident-blocks/CU = 8-resident-waves/CU occupancy +// sweet spot is preserved. The workspace-less fallback, the exact +// (m,n,k)==(16,1280,4096) guard, the scalar fallback for every unmatched +// shape (including paired M=2), the combine kernel and its 320-block +// launch, the pack identity and the 16-plane 1,310,720 B workspace need are +// all unchanged. The iteration-11 repair item stays closed since +// iteration 13 (five-wave archive + __launch_bounds__(320) ONLY, exact- +// correct; the five-wave axis closed on measurement: best 20.613/20.641 us +// vs the accepted two-wave mapping 19.541/19.558 us). // -// Header order is deliberate: hip_runtime.h, then hip_bfloat16.h, then -// du_mma.h (du_mma.h on this DTK is not self-contained before the HIP runtime -// headers). +// Iterations 22-25 (rounds 22 and 23 measured and rejected, tree reverted to +// the iteration-21 source e7ae27d0...; round 24 was killed mid-round before +// any measurement -- its archive draft was never compiled or run): round 22 +// removed the acc-tile LDS round trip from the partial epilogue at pinned +// occupancy (direct-register plane stores) and REGRESSED (+0.82 us -> +// 20.276/20.298 us) because the register-distributed acc values force 16-B +// store granules per wavefront instruction instead of the read-back path's +// four full 64-B lines -- rejected and reverted; round 23 widened each +// combine block to a 32-column n-strip with dwordx2 plane loads (160 blocks) +// and measured FLAT (median 19.4504 us vs the 19.4528 shadow, within its own +// +/-0.05 us falsification band, plus a 41.1 us max outlier) -- rejected and +// reverted. Round 23's pre-registered falsification branch (a) therefore +// committed THIS round to a combine block mapping whose per-plane reads run +// in longer contiguous runs per row (or a partial-plane store-order change), +// never occupancy / fragment duplication / raw asm / a scalar-fallback +// change. +// +// Iteration 25 (this round): combine mapping v3 -- one output row per +// block. The combine byte stream is fixed (every plane word is read +// exactly once, 64-B sectors fully consumed), so v1 (iteration 9: blocks of +// 4 rows x 16 cols) and v2 (iteration 23: 4 rows x 32 cols) both issue +// every per-plane wavefront load as FOUR 64-B (v1) or 128-B (v2) runs at +// the 5,120-B m-row stride -- four different DRAM rows, i.e. up to four row +// activations per load instruction -- which is exactly why widening the +// per-row run 64 -> 128 B in v2 could not move the combine (it never +// reduced the rows-per-instruction scatter). The guard combine launch now +// uses one-wave blocks that each own ONE m-row x 64 consecutive columns +// (kDummaCombineColsPerBlock = 64; grid 20 x 16 = 320 blocks, the same +// 320-block count and the same one-output-element-per-lane, 16-independent- +// plane-load structure as iteration 9): every per-plane wavefront load then +// reads 64 lanes x 4 B = 256 B CONTIGUOUS = four consecutive 64-B lines of +// ONE DRAM row (single run, ~4x fewer DRAM row activations for the same +// byte stream). Each element's ascending-s int32 sum, the float scale +// order and every output bf16 word are unchanged element-for-element, so +// the output is bit-identical to the iteration-21 shadow; the partial +// kernel, the 16-plane workspace contract, the exact guard, the +// workspace-less fallback and the scalar fallback (paired M=2 included) are +// byte-for-byte untouched. +// +// launch_pack_w8a8_weight (optional, out of the timed region) is an +// identity device-to-device copy: the logical [K, N] int8 weight and the +// [N, 1] fp32 scales are copied unchanged, so the DUMMA direct loader +// reads the exact logical [K, N] row-major layout. Later rounds may +// change the packed layout here together with the matching GEMM +// interpretation. +// +// This file is compiled with hipcc by torch.utils.cpp_extension; it must +// stay free of torch/ATen dependencies. Only the two extern "C" launchers +// below are the host contract consumed by csrc/bindings.cpp. +// +// Include order is the known-good DTK order: HIP runtime first, then the +// bfloat16 header, then du_mma.h (consumed by the DUMMA kernel below). #include #include #include #include -#include namespace { -// gfx928 native wavefront is 64 lanes; every block size below is a multiple -// of 64. -constexpr int kScalarBlockThreads = 128; -constexpr int kCopyBlockThreads = 256; +// IEEE round-to-nearest-even conversion of an fp32 value to a bf16 bit +// pattern, implemented with pure integer ops. This is exactly the +// conversion applied by the CPU reference (`scaled.to(torch.bfloat16)`), +// independent of which bf16 type family the DTK headers provide. +__device__ __forceinline__ uint16_t f32_to_bf16_rne(float f) { + constexpr uint32_t kRoundBias = 0x7FFFu; + uint32_t u; + __builtin_memcpy(&u, &f, sizeof(u)); + const uint32_t lsb = (u >> 16) & 1u; // tie goes to even + u += kRoundBias + lsb; + return static_cast(u >> 16); +} -// Minimal DUMMA tile constants (gfx928 INT8 support is m16n16k32). -constexpr int kDummaM = 16; -constexpr int kDummaN = 16; -constexpr int kDummaK = 32; -constexpr int kDummaThreads = 64; // one wavefront (64 lanes) per block -constexpr int kCombineThreads = 256; // 16x16 = 256 output elements/block - -// Exact (k,n) pair whose weight launch_pack_w8a8_weight transposes to the -// [N,K] n-major packed layout and that the scalar fallback decodes as -// n-major (iteration 11; covers the paired M=2 validation shape). -constexpr int kPackedK = 4096; -constexpr int kPackedN = 1280; - -// Trusted occupancy-probe split-K candidates (control plane): all fit the -// contract workspace for this shape (1,310,720 B = 16 partial planes of -// 81,920 B each) and produce 32-aligned K slices. Block counts = 80*S: -// S=3 -> 240 blocks = exactly 2 blocks/CU (CU-aligned, non-power-of-two), -// S=6 -> 480 = 4/CU, S=9 -> 720 = 6/CU, S=12 -> 960 = 8/CU, -// S=16 -> 1,280 = 10.67/CU (16 is outside the old probe list but is the -// largest split that fits the 16-plane workspace, divides 128 K-steps -// exactly into uniform 256-K slices, and is the iteration-8 default: it -// lowers the binding occupancy limiter -- per-block LDS footprint -- from -// 12,736 B (5 resident/CU) to 9,280 B (7 resident/CU)). -constexpr int kTrustedSplitK[] = {2, 3, 4, 5, 6, 8, 9, 12, 16}; +} // namespace // --------------------------------------------------------------------------- -// Scalar GEMM kernel: one thread computes one output element (generic -// fallback for every unmatched shape, including the paired M=2 shape). +// Scalar generic GEMM: out[m, n] = bf16(int32(sum_k a[m,k]*b[k,n]) +// * x_scale[m] * weight_scale[n]). +// One thread per output element. Adjacent lanes walk adjacent columns n +// (the fastest-changing output/weight dimension). gfx928 wavefront = 64; +// the launcher always uses blockDim 128 (a multiple of 64). // --------------------------------------------------------------------------- -// Adjacent lanes own adjacent N columns (the fastest-changing dimension), so -// output stores are coalesced across each wavefront. -__global__ __launch_bounds__(kScalarBlockThreads) void -w8a8_scalar_gemm_kernel( - const int8_t* __restrict__ a, - const int8_t* __restrict__ b, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - hip_bfloat16* __restrict__ out, +__global__ void w8a8_scalar_gemm_kernel( + const int8_t* __restrict__ a, // x_q [M, K] row-major + const int8_t* __restrict__ b, // packed_weight [K, N] row-major + const float* __restrict__ x_scale, // [M] (logical [M, 1]) + const float* __restrict__ w_scale, // [N] (logical [N, 1]) + uint16_t* __restrict__ out, // [M, N] bf16 payload int m, int n, int k) { - const int64_t linear = - static_cast(blockIdx.x) * kScalarBlockThreads + threadIdx.x; - const int64_t total = static_cast(m) * n; - if (linear >= total) { + const int tid = + static_cast(blockIdx.x) * static_cast(blockDim.x) + + static_cast(threadIdx.x); + const int total = m * n; + if (tid >= total) { return; } - const int row = static_cast(linear / n); - const int col = static_cast(linear - static_cast(row) * n); + const int row = tid / n; // M is small for decode; N is the fast axis + const int col = tid - row * n; - // Complete K loop, k ascending, accumulated exactly in int32. The maximum - // API K (<= 6144) keeps the int8 dot well inside int32 range. const int8_t* a_row = a + static_cast(row) * k; - // Iteration 11: the exact (k,n)==(4096,1280) pair is packed [N,K] n-major - // by launch_pack_w8a8_weight; the fallback decodes that layout so it stays - // correct against the packed buffer (paired M=2 validation shape included). - // Every other (K,N) keeps the raw [K,N] row-major decode. - const bool packed = (k == kPackedK && n == kPackedN); - const int8_t* b_col = - b + (packed ? static_cast(col) * k : col); - const int64_t b_stride = packed ? 1 : static_cast(n); + const int8_t* b_col = b + col; // stride n along k + int32_t acc = 0; +#pragma unroll 4 for (int kk = 0; kk < k; ++kk) { acc += static_cast(a_row[kk]) * - static_cast(b_col[static_cast(kk) * b_stride]); + static_cast(b_col[static_cast(kk) * n]); } - // Same float32 evaluation order as the exact reference: - // (float(dot) * x_scale[m]) * weight_scale[n], then bf16 (RN). - const float scaled = static_cast(acc) * x_scale[row] * - weight_scale[col]; - out[linear] = __float2bfloat16(scaled); -} - -void launch_scalar_gemm( - const int8_t* a, - const int8_t* b, - const float* x_scale, - const float* weight_scale, - hip_bfloat16* out, - int m, - int n, - int k, - hipStream_t stream) { - const int64_t total = static_cast(m) * n; - const unsigned grid = static_cast( - (total + kScalarBlockThreads - 1) / kScalarBlockThreads); - hipLaunchKernelGGL( - w8a8_scalar_gemm_kernel, - dim3(grid), - dim3(kScalarBlockThreads), - 0, - stream, - a, - b, - x_scale, - weight_scale, - out, - m, - n, - k); + const float scaled = static_cast(acc) * x_scale[row] * w_scale[col]; + out[tid] = f32_to_bf16_rne(scaled); } // --------------------------------------------------------------------------- -// Workspace split-K partial kernel (template over SPLIT_K), M == 16: -// partial[tile, slice][m, n] = int32_dot(x_q[m, k0:k1], weight[k0:k1, n]) +// Exact-guard DUMMA path for hy3_tp8_qkv_proj_m16 (M=16, N=1280, K=4096). +// +// Minimal native gfx928 INT8 tile: DUFragment fragments with a DUFragment +// accumulator, one 64-thread wavefront per block, one 16x16 output tile per +// block (grid = N/16 = 80 blocks). Every K step is an explicit int32 +// du_mma_sync over a 32-aligned ascending k chunk, so the int32 dot is +// identical to the scalar kernel's sum (exact integer addition, no +// overflow: |out| <= K*127*127 ~= 6.6e7 << 2^31). +// +// The gfx928 int8 m16n16k32 accelerator fragment ownership is opaque, so +// the epilogue materializes the tile with du_store_matrix_sync into a +// single-wave LDS staging array (library-defined mem_row_major placement, +// 256 x int32 = 1,024 B) and only afterwards applies the per-element float +// scaling and the bf16 RNE store. BlockDim is one wavefront, so the single +// __syncthreads() below is a same-wave barrier (all 64 threads reach it) +// and there is no cross-wave synchronization anywhere in the kernel. // --------------------------------------------------------------------------- -// Layouts (all contiguous): -// x_q [16, K] int8 row-major (ldm = K) -// weight [K, N] int8 row-major (ldm = N, identity pack layout) -// partials [(SPLIT_K * num_tiles) * 256] int32, plane = slice*num_tiles + -// tile, row-major 16x16 tile per plane (caller workspace) -// One block = one wavefront = one (tile, slice). The K slice is 32-aligned -// (non-uniform: slices differ by at most one 32-K DUMMA step) and covers -// [k0, k0 + steps*32) in ascending order. No LDS, no __syncthreads, no -// atomics: each block writes exactly one private int32 partial plane. -template -__global__ __launch_bounds__(kDummaThreads) void -w8a8_dumma_m16_splitk_partial_kernel( - const int8_t* __restrict__ a, - const int8_t* __restrict__ b, - int32_t* __restrict__ partials, +constexpr int kDummaM = 16; +constexpr int kDummaN = 16; +constexpr int kDummaK = 32; +constexpr int kDummaWaveSize = 64; + +// Combine constants. Iteration 9: the exact-guard launcher (SPLIT_K == 16) +// launched the combine kernel as 80 n-tiles x four-row m-groups of one-wave +// blocks and the kernel specialized split_count == kDummaCombineSplits with +// one output element per lane. Iteration 25 (mapping v3): the guard launch +// becomes 20 column-groups x 16 m-rows of one-wave blocks -- each block owns +// ONE output row x kDummaCombineColsPerBlock = 64 consecutive columns (4 x +// kDummaN), so every per-plane wavefront load reads a single 256-B +// contiguous run (four consecutive 64-B lines inside one DRAM row) instead +// of four 64/128-B runs at the 5,120-B row stride (see the kernel comment +// below for the DRAM row-activation mechanism). +constexpr int kDummaCombineSplits = 16; // == the guard launcher kSplitK +constexpr int kDummaCombineColsPerBlock = 64; // n-columns per combine block + +__global__ void w8a8_dumma_m16_tile_kernel( + const int8_t* __restrict__ a, // x_q [M, K] row-major + const int8_t* __restrict__ b, // packed_weight [K, N] row-major + const float* __restrict__ x_scale, // [M] + const float* __restrict__ w_scale, // [N] + uint16_t* __restrict__ out, // [M, N] bf16 payload int n, int k) { - const int tile = static_cast(blockIdx.x); - const int slice = static_cast(blockIdx.y); - const int n0 = tile * kDummaN; + const int lane = static_cast(threadIdx.x); // blockDim == 64 + const int n0 = static_cast(blockIdx.x) * kDummaN; - // 32-aligned non-uniform slice bounds over the K dimension. - const int total_steps = k / kDummaK; // 4096 / 32 = 128 - const int base = total_steps / SPLIT_K; - const int rem = total_steps - base * SPLIT_K; - const int steps = base + (slice < rem ? 1 : 0); - const int k0 = (slice * base + (slice < rem ? slice : rem)) * kDummaK; + __shared__ __align__(16) int32_t acc_tile[kDummaM * kDummaN]; du::dumma::DUFragment a_frag; - // Iteration 11: b is the packed [N,K] n-major weight for this shape, so - // the B fragment is col_major (each lane's 8 fragment bytes are consecutive - // in memory -> one 8-B read instead of 8 byte reads). du::dumma::DUFragment + signed char, du::dumma::row_major> b_frag; du::dumma::DUFragment acc_frag; du::dumma::du_fill_fragment(acc_frag, 0); - for (int i = 0; i < steps; ++i) { - const int kk = k0 + i * kDummaK; - du::dumma::du_load_matrix_sync(a_frag, a + kk, k); - // Packed n-major row n starts at b[n*k + k0]; the library col_major - // loader addresses p[(lane&15)*ldm + ((lane>>4)<<3) + i], i.e. n-major - // rows of k bytes, so the tile's 16 columns live at n0..n0+15 n-rows - // with row stride k. + // A is [16, k] (row stride k) and B is [k, 16] (row stride n); both stay + // row-major in their global tensors, so each 32-K step loads straight + // from global with the library row_major loaders. + for (int k0 = 0; k0 < k; k0 += kDummaK) { + du::dumma::du_load_matrix_sync(a_frag, a + k0, k); du::dumma::du_load_matrix_sync( - b_frag, b + static_cast(n0) * k + kk, k); + b_frag, b + static_cast(k0) * n + n0, n); du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); } - // Private int32 partial plane: slice-major so a tile's S planes are - // stride-uniform in the combine kernel. - const int num_tiles = static_cast(gridDim.x); - du::dumma::du_store_matrix_sync( - partials + (slice * num_tiles + tile) * (kDummaM * kDummaN), - acc_frag, - kDummaN, - du::dumma::mem_row_major); -} - -// Template launch helper for the partial kernel (one wavefront per block, -// grid = (num_tiles, SPLIT_K)). -template -void launch_splitk_partial( - const int8_t* a, - const int8_t* b, - int32_t* partials, - int n, - int k, - hipStream_t stream, - int num_tiles) { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_m16_splitk_partial_kernel), - dim3(static_cast(num_tiles), static_cast(SPLIT_K)), - dim3(static_cast(kDummaThreads)), - 0, - stream, - a, - b, - partials, - n, - k); -} + du::dumma::du_store_matrix_sync(acc_tile, acc_frag, kDummaN, + du::dumma::mem_row_major); + __syncthreads(); -// Iteration 11 packed B-fragment loader (validated on the TP4 gate_up -// (4096,768) lineage): with n-major [N,K] storage and a col_major fragment, -// du_mma.h's matrix_b col_major address rule is x[i] = p[(lane&15)*ldm + -// ((lane>>4)<<3) + i] -- lane l's 8 fragment bytes are consecutive in -// memory, so one 8-byte LDS read per fragment replaces the 8 ds_read_u8 at -// stride-64 addresses that the row-major loader emits. du_mma_sync consumes -// the fragment as one 64-bit value per lane, so writing the packed bytes -// straight into x[0..7] is exactly the library's own storage convention. -// p must be 8 B aligned (guaranteed: b_s is __align__(256), kBStride % 8 == -// 0, step offsets are multiples of 32, and (lane&15)*kBStride % 8 == 0). -__device__ __forceinline__ void load_b_frag_packed( - du::dumma::DUFragment& b_frag, - const int8_t* p, int ldm, int lane) { - const int row = lane & 15; - const int col = (lane >> 4) << 3; - const uint64_t v = *reinterpret_cast( - p + static_cast(row) * ldm + col); - *reinterpret_cast(b_frag.x) = v; + // 256 elements / 64 lanes: each lane scales and stores four outputs. +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; linear += kDummaWaveSize) { + const int row = linear / kDummaN; + const int col = linear - row * kDummaN; + const int out_col = n0 + col; + const float scaled = static_cast(acc_tile[linear]) * + x_scale[row] * w_scale[out_col]; + out[static_cast(row) * n + out_col] = f32_to_bf16_rne(scaled); + } } // --------------------------------------------------------------------------- -// Whole-K-slice LDS staging partial kernel (M == 16, iteration 7): -// partial[tile, slice][m, n] = int32_dot(x_q[m, k0:k1], weight[k0:k1, n]) +// Split-K partial kernel (iteration 2 probe, widened in iteration 4 to +// multi-N-tile reuse, B-staged in iteration 5): same m16n16k32 int8 DUMMA +// tile as iteration 1, but gridDim = (N/(16*kTilesPerWave), SPLIT_K) so +// every one-wave block owns kTilesPerWave adjacent 16-wide N tiles AND one +// ascending 32-aligned K slice, and writes its exact int32 partials into +// workspace planes partials[split][16][n]. Iteration 4: kTilesPerWave = 2, +// so per 32-K step the wave loads the A fragment once and shares it between +// the two du_mma_sync calls (A bytes of the step reused across both output +// tiles; every B byte is still read exactly once). Iteration 5 (staging +// round, B-only): each block first stages its whole 32-column B slice +// (my_steps*32 rows x 32 B, <= 352 x 32 = 11,264 B under the SPLIT_K=12 +// guard launch) into an LDS tile stage_b that is an EXACT byte mirror of +// the block's global B strip (16-B int4 cooperative loads, 4-deep source +// batching so one vmcnt drain covers 4 in-flight round trips), then after +// one same-wave __syncthreads() the K loop reads the B0/B1 fragments from +// LDS with du_load_matrix_sync (row stride 32) -- same fragment bytes as +// the iteration-4 direct global loads, so the int32 partials stay +// bit-identical -- while A stays on the direct-global path (this round +// isolates the B side only). BlockDim is one 64-lane wavefront: the two +// barriers are the staging barrier above and the single same-wave +// __syncthreads() after the LDS materialization of both partial tiles. +// LDS per block: 11,264 B (stage_b) + 2 x 1,024 B (acc tiles) = 13,312 B +// -> 65,536 / 13,312 = 4.9, so the 4-blocks-per-CU launch co-residency is +// preserved (LDS is not the occupancy limiter). +// +// Iteration 7 (this round): A+B whole-slice LDS staging. Each one-wave +// block now stages BOTH operands of its whole K slice once: stage_a holds +// the 16 m-rows x 256 slice bytes at row stride 272 (256 data + 16 pad for +// bank skew, 4,352 B) and stage_b holds the 256 x 32 B strip (8,192 B); +// both use 16-B int4 cooperative loads with the same 4-deep batching, one +// same-wave __syncthreads(), and then the K loop reads every fragment from +// LDS (A at row stride 272, B0/B1 at row stride 32) -- zero per-step global +// loads and zero vmcnt waits on the MMA path. Geometry moved to SPLIT_K=16 +// (uniform 8-step slices, 40 x 16 = 640 blocks): per-block LDS 4,352 + +// 8,192 + 2,048 = 14,592 B -> floor(65,536/14,592) = 4, so the 4-blocks/CU +// co-residency of iteration 5 is preserved. Barriers per block stay two +// (staging, epilogue). Staged tiles are exact byte mirrors, so partials +// are bit-identical to iteration 5 and the combine kernel is unchanged. +// +// Iteration 8 (occupancy/resource round, this round): tune the waves-per- +// block occupancy limiter. PMC of the accepted iteration-7 code object +// (source digest b64e3209..., grid 640, workgroup 64, arch_vgpr 40, sgpr 48, +// scratch 0, LDS/block 14,592 B -> floor(65,536/14,592) = 4 resident blocks/ +// CU, 5.33 waves/CU by grid) shows a per-SIMD latency-exposure wall: every +// block is ONE wavefront, so at 4 resident blocks/CU the four waves land one +// per SIMD and each SIMD has no second independent wave to issue while the +// resident wave stalls on its per-step s_waitcnt lgkmcnt (lds_wait 8,359) +// after the staged-fragment du_load_matrix_sync and on its du_mma accum +// chain (VALU issue utilization ~3-6% of the 26.56 us profiled kernel; +// vmem_read 7,680 = only the 12 one-time staging loads per block). Fix: +// blockDim 64 -> 128, i.e. TWO 64-lane wavefronts per block, and the block's +// 32-column strip is split by ownership: wave w computes ONLY n-tile w (cols +// [n0+16w, n0+16w+16)) over the shared 8-step slice with its own independent +// accumulator chain, one du_mma per step per wave (same total MMA count). +// LDS per block is UNCHANGED (stage_a 4,352 + stage_b 8,192 + 2 x 1,024 acc +// tiles = 14,592 B), so 4 blocks/CU still fit but now carry 8 resident waves +// = 2 per SIMD: one wave's LDS->MMA stall is filled by the other. The +// one-time staging is partitioned, NOT duplicated: wave 0 copies the shared +// A slice and wave 1 copies the B strip (their iteration-7 loops unchanged), +// so exactly the same global bytes are read once (no repeated HBM reads for +// occupancy) and the staging barrier becomes a cross-wave __syncthreads(). +// Deliberate added cost: the A fragment is re-read from LDS by each wave +// (lds_instructions 99,840 -> ~125-135k predicted; LDS pipe/conflicts are +// not the wall at ~3-6% issue utilization -- the trade is LDS re-reads, not +// HBM traffic). Geometry, planes, combine kernel, guards and fallbacks are +// otherwise unchanged. // --------------------------------------------------------------------------- -// Same 32-aligned non-uniform K slices and the same k-ascending int32 -// du_mma accumulation as w8a8_dumma_m16_splitk_partial_kernel, so the -// combine+scale kernel below sums bit-identical partials and the outputs are -// bit-identical to the accepted iteration-3 kernel. Differences: -// - A (16 x slice) is staged ONCE into LDS with coalesced dword global -// loads (4 B/lane, 256 B per wavefront instruction); B (16 x slice) is -// staged as 16-B int4 loads straight from the packed [N,K] n-major -// weight (iteration 11) into an n-major LDS tile. -// - One __syncthreads() after staging, then a zero-barrier LDS-only K -// loop: du_load_matrix_sync reads the A fragment from LDS and -// load_b_frag_packed reads each B fragment with one ds_read_b64 (~30-cycle -// LDS latency) instead of a global round trip per K step. -// Layouts (all contiguous): -// x_q [16, K] int8 row-major (ldm = K) -// weight [K, N] int8 row-major for every (K,N) except the exact pair -// (k,n)==(4096,1280), which is packed [N,K] n-major (element -// (k,n) at P[n*K+k]; iteration 11) -// partials [(SPLIT_K * num_tiles) * 256] int32, plane = slice*num_tiles + -// tile, row-major 16x16 tile per plane (caller workspace) -// LDS: a_s[16][kMaxKlen + 4] (compile-time row stride kMaxKlen+4, an odd -// dword count -> the library's matrix_a row-major lane pattern (row = -// lane&15, 8 consecutive bytes per lane) is bank-conflict-free); -// b_s[16][kMaxKlen + 16] (iteration 11: n-major, row stride kMaxKlen+16 = -// 16-B aligned for the int4 staging stores, 8-B aligned for the col_major -// fragment b64 reads, not a multiple of 128 B -> the 16 staged n-rows do not -// alias onto one LDS bank phase; each lane's 8 fragment bytes are -// consecutive, so every B fragment loads with ONE ds_read_b64 instead of 8 -// ds_read_u8). Per-template LDS fits: S=6 22,848 B, S=8 16,704 B, S=9 -// 15,680 B, S=12 11,584 B, S=16 8,512 B (all <= 64 KiB; resident blocks/CU -// by LDS: 2/3/4/5/7). Only launched -// under the exact (m,n,k) == (16,1280,4096) guard, so total_steps = k/32 = -// 128 is exact for the compile-time LDS sizing below. -template -__global__ __launch_bounds__(kDummaThreads) void -w8a8_dumma_m16_slicestage_partial_kernel( - const int8_t* __restrict__ a, - const int8_t* __restrict__ b, - int32_t* __restrict__ partials, +__global__ void w8a8_dumma_m16_splitk_partial_kernel( + const int8_t* __restrict__ a, // x_q [16, k] row-major + const int8_t* __restrict__ b, // packed_weight [k, n] row-major + int32_t* __restrict__ partials, // [gridDim.y][16][n] int32 int n, int k) { - const int tile = static_cast(blockIdx.x); - const int slice = static_cast(blockIdx.y); - const int n0 = tile * kDummaN; - - // 32-aligned non-uniform slice bounds (identical formula to the direct - // kernel: slices tile [0, K) exactly once in ascending order). - const int total_steps = k / kDummaK; // 4096 / 32 = 128 under the guard - const int base = total_steps / SPLIT_K; - const int rem = total_steps - base * SPLIT_K; - const int steps = base + (slice < rem ? 1 : 0); - const int k0 = (slice * base + (slice < rem ? slice : rem)) * kDummaK; - const int klen = steps * kDummaK; - - // Compile-time LDS sizing for the longest slice this instantiation can - // receive (base+1 steps, or base when SPLIT_K divides 128 evenly). - constexpr int kMaxSteps = 128 / SPLIT_K + (128 % SPLIT_K == 0 ? 0 : 1); - constexpr int kMaxKlen = kMaxSteps * kDummaK; - __shared__ int8_t a_s[16][kMaxKlen + 4]; - // Iteration 11: B is staged n-major (16 columns x kMaxKlen+16 bytes, row - // stride kMaxKlen+16) so col_major B fragments read 8 consecutive bytes - // per lane (one ds_read_b64). kMaxKlen is a multiple of 32, so the row - // stride is 16-B aligned (int4 staging stores stay aligned), 8-B aligned - // (fragment b64 reads stay aligned) and not a multiple of 128 B (no LDS - // bank-phase aliasing across the 16 staged n-rows). - __shared__ __align__(256) int8_t b_s[16][kMaxKlen + 16]; - - // Row strides are compile-time constants equal to the array row strides - // (a_s: kMaxKlen+4 bytes = odd dword count; b_s: kMaxKlen+16 bytes), so - // the staging stores and the K-loop fragment reads always agree even for - // short slices (klen < kMaxKlen). The A stride stays an odd dword count - // (bank-conflict-free row-major lane pattern); the B n-major stride keeps - // every fragment b64 read 8 B aligned and off the single bank phase. - constexpr int kAStride = kMaxKlen + 4; - constexpr int kBStride = kMaxKlen + 16; - - // Stage A: 16 rows x klen bytes = 4*klen dwords; consecutive lanes copy - // consecutive dwords of one row (coalesced 256 B per wavefront load). - // Iteration 12 (staging MLP, alignment-safe): the fresh exact-source ISA - // (iteration-12 profile, digest ab7fc108...) shows both staging loops - // compile to a fully serialized load -> s_waitcnt vmcnt(0) -> store chain - // (A: 16 sequential global round trips per block, B: 4), and the accepted - // iteration-11 round already validated this fix on the B side (16 - // serialized dword loads -> 4 aligned int4 loads). Both loops below batch - // 4-deep in source: 4 independent loads are issued before their 4 stores, - // so one vmcnt drain covers 4 in-flight round trips (A 16 -> 4, B 4 -> 1). - // A stays dword-wide: 4-B loads AND stores are aligned at every a_s row, - // while kAStride = kMaxKlen+4 is not 16-B aligned, so int4 A staging would - // repeat round 9's misaligned-store regression (58.56 us) -- that round - // falsified misaligned int4 staging, not the MLP idea. B stays int4 (its - // global rows are 16-B aligned and kBStride % 16 == 0). Fragment reads, - // LDS strides, the K loop and the int32 accumulation order are untouched, - // so outputs stay bit-identical (0 mismatches) and Graph capture/replay - // with changed contents must pass. klen % 32 == 0 => a_dwords % 128 == 0, - // so the batch loop is exact whenever klen % 128 == 0 (S=16: klen = 256 -> - // 1024 dwords = 4 batches exactly; S=6/9/12's odd slices fall through to - // the scalar tail, which preserves correctness for every S in the sweep). - const int a_dwords = 4 * klen; - const int a_dw_per_row = klen / 4; - constexpr int kABatch = 4 * kDummaThreads; - { - const auto ld_a = [&](int e) { - const int row = e / a_dw_per_row; - return *reinterpret_cast( - a + static_cast(row) * k + k0 + (e - row * a_dw_per_row) * 4); - }; - const auto st_a = [&](int e, int32_t v) { - const int row = e / a_dw_per_row; - *reinterpret_cast(&a_s[row][(e - row * a_dw_per_row) * 4]) = v; - }; - int e = static_cast(threadIdx.x); - for (; e + 3 * kDummaThreads < a_dwords; e += kABatch) { - const int32_t v0 = ld_a(e); - const int32_t v1 = ld_a(e + kDummaThreads); - const int32_t v2 = ld_a(e + 2 * kDummaThreads); - const int32_t v3 = ld_a(e + 3 * kDummaThreads); - st_a(e, v0); - st_a(e + kDummaThreads, v1); - st_a(e + 2 * kDummaThreads, v2); - st_a(e + 3 * kDummaThreads, v3); - } - for (; e < a_dwords; e += kDummaThreads) { - st_a(e, ld_a(e)); - } - } + constexpr int kTilesPerWave = 2; // two adjacent 16-wide N tiles per block + // Iteration 8 (occupancy round): each block is TWO 64-lane wavefronts + // (blockDim 128). Wave w owns N tile w of the block's 32 columns (cols + // [n0 + 16*w, n0 + 16*w + 16)) and runs an independent 8-step accumulator + // chain over the shared K slice for that tile alone. + const int tid = static_cast(threadIdx.x); // blockDim == 128 + const int w = tid >> 6; // wave id inside block 0/1 + const int lane = tid & (kDummaWaveSize - 1); // lane id within the wave + const int group = static_cast(blockIdx.x); // 32-wide n-group id + const int split = static_cast(blockIdx.y); // k-slice id + const int n0 = group * (kDummaN * kTilesPerWave); - // Stage B (iteration 11): the exact (k,n)==(4096,1280) weight is packed - // [N,K] n-major (element (k,n) at P[n*K+k]) by launch_pack_w8a8_weight, so - // the block's 16-column strip is 16 n-rows of klen contiguous bytes. Each - // lane stages 16-B int4s straight into the n-major LDS tile b_s[16][ - // kMaxKlen+16] (4 loads per lane at the uniform S=16 default: rows 0..15 x - // 16 int4 columns). Iteration 12: the same 4-deep batching as A, so the 4 - // aligned int4 round trips collapse to one vmcnt drain per block. - const int k16_per_row = klen / 16; - const int b_int4s = kDummaN * k16_per_row; - const int4* __restrict__ b4 = reinterpret_cast(b); - int4* __restrict__ lds_b4 = reinterpret_cast(b_s); - constexpr int kBBatch = 4 * kDummaThreads; - { - const auto ld_b = [&](int e) { - const int row = e / k16_per_row; - return b4[(static_cast(n0) + row) * (k >> 4) + (k0 >> 4) + - (e - row * k16_per_row)]; - }; - const auto st_b = [&](int e, const int4& v) { - const int row = e / k16_per_row; - lds_b4[row * (kBStride >> 4) + (e - row * k16_per_row)] = v; - }; - int e = static_cast(threadIdx.x); - for (; e + 3 * kDummaThreads < b_int4s; e += kBBatch) { - const int4 v0 = ld_b(e); - const int4 v1 = ld_b(e + kDummaThreads); - const int4 v2 = ld_b(e + 2 * kDummaThreads); - const int4 v3 = ld_b(e + 3 * kDummaThreads); - st_b(e, v0); - st_b(e + kDummaThreads, v1); - st_b(e + 2 * kDummaThreads, v2); - st_b(e + 3 * kDummaThreads, v3); - } - for (; e < b_int4s; e += kDummaThreads) { - st_b(e, ld_b(e)); - } - } + // Stage-buffer sizing is tied to the ONLY launch site of this kernel: the + // exact (m,n,k)==(16,1280,4096) guard with gridDim.y == 16 (iteration 7), + // so total_steps == 128 and every split owns exactly 8 steps = 256 rows + // (128 == 16 * 8; the generic slice formula below degenerates to uniform + // 8-step slices). Compile-time capacities: 256 rows of B per block and + // 16 m-rows x (256 + 16 pad) bytes of A per block. + constexpr int kStageSplitCount = 16; // gridDim.y of the guard launch + constexpr int kStageMaxSteps = + (128 + kStageSplitCount - 1) / kStageSplitCount; // 8 (uniform) + constexpr int kStageMaxRows = kStageMaxSteps * kDummaK; // 256 rows x 32 + constexpr int kStageARowStride = + kStageMaxRows + 16; // 272 B per A m-row: 256 data + 16 pad for banks + // Iteration 21 (DS-side bank-skew round): each staged B k-row now occupies + // kStageBRowStride = 36 B (32 data bytes + 4 pad bytes). The row-major B + // fragment loader reads a lane's 8 fragment bytes at addresses + // base + (8*(lane>>4) + i) * ldm, i = 0..7: the four lane groups + // g = lane>>4 sit 8*ldm bytes apart, so with the old ldm == 32 every group + // landed 256 B (== 0 mod 128 B) apart and each of the 8 ds_read_u8 per + // fragment replayed the SAME LDS bank 4x (the PMC-documented + // lds_bank_conflicts 276,480 = ~2.57 replays/instr on this object, ~89% of + // it from the B-fragment byte reads: 128 per block). ldm == 36 puts the + // groups 288 B (== 32 mod 128 B, 8 banks) apart, so the four groups hit + // banks {b, b+8, b+16, b+24} -- fully conflict-free -- while the 4-B pad + // per row keeps every row start 4-B aligned for the staging stores. + // stage_b grows 8,192 -> 9,216 B and the block total becomes 4,352 + // (stage_a) + 9,216 (stage_b) + 2,048 (acc tiles) = 15,616 B/block: + // 4 x 15,616 = 62,464 <= 65,536, so the 4-resident-blocks/CU = + // 8-resident-waves/CU occupancy sweet spot is preserved unchanged. + constexpr int kStageBRowStride = kDummaN * 2 + 4; // 36 B per staged B row - __syncthreads(); + __shared__ __align__(16) int32_t acc_tile0[kDummaM * kDummaN]; + __shared__ __align__(16) int32_t acc_tile1[kDummaM * kDummaN]; + __shared__ __align__(16) int8_t stage_a[kDummaM * kStageARowStride]; + __shared__ __align__(16) int8_t stage_b[kStageMaxRows * kStageBRowStride]; du::dumma::DUFragment a_frag; - // Iteration 11: col_major B fragments match the n-major staged tile, so - // each lane's 8 fragment bytes are consecutive (one ds_read_b64 per - // fragment in the K loop instead of 8 ds_read_u8). du::dumma::DUFragment - b_frag; + signed char, du::dumma::row_major> + b_frag; // this wave's B half of the strip (b_col == 16*w bytes) du::dumma::DUFragment - acc_frag; + acc_frag; // this wave's accumulator for its n16 tile du::dumma::du_fill_fragment(acc_frag, 0); - // Zero-barrier LDS-only K loop: every fragment read comes from LDS, so no - // global latency sits on the MMA critical path after the single staging - // barrier above. - const int lane = static_cast(threadIdx.x) & 63; - for (int i = 0; i < steps; ++i) { - du::dumma::du_load_matrix_sync(a_frag, &a_s[0][i * kDummaK], kAStride); - // One 8-B ds_read_b64 per B fragment (col_major: each lane's 8 fragment - // bytes are consecutive in the n-major tile) instead of 8 ds_read_u8 at - // the 20-B k-major pitch. - load_b_frag_packed(b_frag, &b_s[0][i * kDummaK], kBStride, lane); - du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + // 32-aligned non-uniform K slicing over the gridDim.y splits: split s owns + // my_steps ascending 32-K steps beginning at step_begin, slices are + // disjoint and cover [0, k) exactly once in ascending order, so the int32 + // partial sums and the ascending combine below reproduce the scalar + // k-ascending sum bit-for-bit. + const int total_steps = k / kDummaK; + const int split_count = static_cast(gridDim.y); + const int base_steps = total_steps / split_count; + const int rem_steps = total_steps - base_steps * split_count; + const int my_steps = base_steps + (split < rem_steps ? 1 : 0); + const int step_begin = + split * base_steps + (split < rem_steps ? split : rem_steps); + const int k_base = step_begin * kDummaK; // first global k row of slice + const int rows = my_steps * kDummaK; // slice rows: my_steps x 32 + + // Stage the block's operands (iteration 8: PARTITIONED across the two + // wavefronts, not duplicated -- the union reads exactly the same global + // bytes as the iteration-7 single wave, so no byte is fetched twice and + // HBM traffic is unchanged). Wave 0 copies the shared A slice: the 16 + // m-rows x rows bytes of A (global columns [k_base, k_base + rows)), + // stored m-major in stage_a at row stride kStageARowStride (an exact byte + // mirror of global a[m][k_base .. k_base + rows)). Wave 1 copies the + // block's 32-column B strip: my_steps*32 rows x 32 B, stored row-major in + // stage_b at the 36-B row pitch kStageBRowStride (32 data + 4 pad bytes + // per row; iteration 21 bank skew -- an exact byte mirror of the global + // data rows [k_base, k_base + rows) x columns [n0, n0 + 32), pads never + // written or read). Each copy uses the + // unchanged iteration-7 cooperative code: 64 lanes x 16 B per wavefront + // instruction, chunks in LDS int4 units, loads batched 4-deep in source so + // one vmcnt drain covers 4 in-flight global round trips (the two waves' + // round trips overlap further across the SIMDs). Global m-rows are + // k-major at stride k and global B rows are n-major at stride n (k % 16 == + // 0 and n % 16 == 0 under the guard; k_base and n0 are 32-aligned), so + // every 16-B global load is 16-B aligned (the wave-1 LDS stores are + // 4-B-granular since iteration 21: 36*row is only 4-B aligned). + const int a_chunks = (rows * kDummaM) / 16; // rows/16 int4 per m-row + const int chunks = rows * (kDummaN * 2) / 16; // rows*2 int4 chunks + if (w == 0) { + const int per_row = rows / 16; + const int4* __restrict__ a4 = reinterpret_cast(a); + int4* __restrict__ ald4 = reinterpret_cast(stage_a); + const int kb4 = k_base >> 4; // global int4 column of the slice start + const auto ld_chunk = [&](int e) { + const int m = e / per_row; + const int t = e - m * per_row; + return a4[static_cast(m) * (k >> 4) + kb4 + t]; + }; + const auto st_chunk = [&](int e, const int4& v) { + const int m = e / per_row; + const int t = e - m * per_row; + ald4[m * (kStageARowStride / 16) + t] = v; + }; + int e = lane; + for (; e + 3 * kDummaWaveSize < a_chunks; e += 4 * kDummaWaveSize) { + const int4 v0 = ld_chunk(e); + const int4 v1 = ld_chunk(e + kDummaWaveSize); + const int4 v2 = ld_chunk(e + 2 * kDummaWaveSize); + const int4 v3 = ld_chunk(e + 3 * kDummaWaveSize); + st_chunk(e, v0); + st_chunk(e + kDummaWaveSize, v1); + st_chunk(e + 2 * kDummaWaveSize, v2); + st_chunk(e + 3 * kDummaWaveSize, v3); + } + for (; e < a_chunks; e += kDummaWaveSize) { + st_chunk(e, ld_chunk(e)); + } + } else { + const int4* __restrict__ b4 = reinterpret_cast(b); + const auto ld_chunk = [&](int e) { + const int r = e >> 1; + // Global int4 index of (global row k_base + r, col half e&1). + return b4[static_cast(k_base + r) * (n >> 4) + (n0 >> 4) + + (e & 1)]; + }; + // Iteration 21: stage_b rows now live at the 36-B pitch + // kStageBRowStride (32 data bytes + 4 pad per row), so a 16-B chunk + // (row r = e >> 1, col half h = e & 1 of that row) can no longer be + // stored with one 16-B-aligned ds_write_b128 (36*r is 16-B aligned only + // for r % 4 == 0). Each chunk is stored as four consecutive int32 words + // at word offset 9*r + 4*h + {0,1,2,3} == (36*r + 16*h)/4 + {0,1,2,3}: + // every data byte lands at exactly the same byte offset (row*36 + ncol) + // the loader reads back, so the staged data byte mirror of the global B + // strip is unchanged -- only the pad bytes (row*36 + 32..35) are never + // written and never read (the fragment loader only touches ncol < 32). + // LDS store bandwidth is unchanged (4 x 4-B words per chunk); the four + // words of one chunk are consecutive, so a lane's store stream stays + // dense per row half. + int32_t* __restrict__ lds32 = reinterpret_cast(stage_b); + const auto st_chunk = [&](int e, const int4& v) { + const int word0 = (e >> 1) * (kStageBRowStride >> 2) + (e & 1) * 4; + lds32[word0 + 0] = v.x; + lds32[word0 + 1] = v.y; + lds32[word0 + 2] = v.z; + lds32[word0 + 3] = v.w; + }; + // Iteration 20 (guard-slice staging latency specialization): the guard + // launch is the ONLY call site of this kernel and always stages the + // full 512-chunk B strip (chunks == rows*2 == 512 == 8 x 64 lanes, rows + // == 256 == kStageMaxRows, staged capacity exactly filled). The + // runtime-bounded 4-deep loop below then compiles to TWO serialized + // 4-deep iterations (each iteration issues four 16-B source loads and + // drains them with vmcnt waits before the next iteration's loads can + // issue -- two full global round trips on the block's staging critical + // path, one per iteration), and wave 0's A copy finishes in a single + // 4-deep group and then idles at the __syncthreads() join while wave 1 + // completes its second round trip. When chunks == 8*kDummaWaveSize the + // loop now runs as ONE straight-line 8-deep batched group: all eight + // 16-B source loads per lane are issued before the first LDS store, so + // one drain window covers the whole strip and the second round trip is + // removed from the join's critical path. Every chunk id [0, 512) is + // still loaded and stored exactly once by exactly one lane to the same + // stage_b byte (the chunk set, the store order and the byte mirror are + // identical to the rolled loop), so every staged byte, every partial + // and every combine output byte is bit-identical; the generic rolled + // loop stays verbatim as the else path for any non-guard (chunks != + // 512) configuration, which never executes at the only launch site. + if (chunks == 8 * kDummaWaveSize) { + const int4 v0 = ld_chunk(lane); + const int4 v1 = ld_chunk(lane + kDummaWaveSize); + const int4 v2 = ld_chunk(lane + 2 * kDummaWaveSize); + const int4 v3 = ld_chunk(lane + 3 * kDummaWaveSize); + const int4 v4 = ld_chunk(lane + 4 * kDummaWaveSize); + const int4 v5 = ld_chunk(lane + 5 * kDummaWaveSize); + const int4 v6 = ld_chunk(lane + 6 * kDummaWaveSize); + const int4 v7 = ld_chunk(lane + 7 * kDummaWaveSize); + st_chunk(lane, v0); + st_chunk(lane + kDummaWaveSize, v1); + st_chunk(lane + 2 * kDummaWaveSize, v2); + st_chunk(lane + 3 * kDummaWaveSize, v3); + st_chunk(lane + 4 * kDummaWaveSize, v4); + st_chunk(lane + 5 * kDummaWaveSize, v5); + st_chunk(lane + 6 * kDummaWaveSize, v6); + st_chunk(lane + 7 * kDummaWaveSize, v7); + } else { + int e = lane; + for (; e + 3 * kDummaWaveSize < chunks; e += 4 * kDummaWaveSize) { + const int4 v0 = ld_chunk(e); + const int4 v1 = ld_chunk(e + kDummaWaveSize); + const int4 v2 = ld_chunk(e + 2 * kDummaWaveSize); + const int4 v3 = ld_chunk(e + 3 * kDummaWaveSize); + st_chunk(e, v0); + st_chunk(e + kDummaWaveSize, v1); + st_chunk(e + 2 * kDummaWaveSize, v2); + st_chunk(e + 3 * kDummaWaveSize, v3); + } + for (; e < chunks; e += kDummaWaveSize) { + st_chunk(e, ld_chunk(e)); + } + } } + __syncthreads(); // stage_a and stage_b ready for every lane of both waves - // Private int32 partial plane: same slice-major layout as the direct - // kernel (plane = slice*num_tiles + tile). - const int num_tiles = static_cast(gridDim.x); - du::dumma::du_store_matrix_sync( - partials + (slice * num_tiles + tile) * (kDummaM * kDummaN), - acc_frag, - kDummaN, - du::dumma::mem_row_major); -} + // LDS-only K loop (iteration 8): every 32-K step wave w reads the A + // fragment from the staged A mirror at row stride kStageARowStride and its + // own B half from the staged B strip at the 36-B row pitch + // kStageBRowStride (byte column b_col; iteration 21 bank skew) + // -- the same fragment bytes the iteration-7 loads produced, with no + // global round trip and no vmcnt wait on the MMA path -- then issues ONE + // du_mma on its private accumulator. The two resident wavefronts carry + // independent 8-step chains, so a SIMD holding both can issue one wave's + // instructions while the other stalls on its s_waitcnt lgkmcnt / du_mma + // dependency (per-SIMD latency exposure was the iteration-7 wall). + const int b_col = w * kDummaN; // byte offset of this wave's n16 tile -// Template launch helper for the staged partial kernel (one wavefront per -// block, grid = (num_tiles, SPLIT_K)); only instantiated for the SPLIT_K -// values whose longest slice fits the 64 KiB LDS (6, 8, 9, 12, 16). -template -void launch_slicestage_partial( - const int8_t* a, - const int8_t* b, - int32_t* partials, - int n, - int k, - hipStream_t stream, - int num_tiles) { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_m16_slicestage_partial_kernel), - dim3(static_cast(num_tiles), static_cast(SPLIT_K)), - dim3(static_cast(kDummaThreads)), - 0, - stream, - a, - b, - partials, - n, - k); + // Iteration 18 (compile-time 8-step K-loop specialization): the guard + // launch is the ONLY call site of this kernel and always runs the uniform + // 8-step slice (k == 4096, gridDim.y == 16 -> my_steps == 8 == + // kStageMaxSteps, rows == 256 == kStageMaxRows, staged capacity exactly + // filled). The runtime-bounded generic loop below compiles to a ROLLED + // body (verified in the exact-source code object: one static v_mmac per + // step with per-step ds_read2_b32 + 8 x ds_read_u8 + ~25-VALU byte + // reassembly + 4 lgkmcnt drains on the accumulator chain and a + // loop-carried WAR on the single fragment set), so every step's full LDS + // latency sits serially on the MMA path with only 2 resident waves/SIMD + // to cover it (PMC of the exact object: lds_instructions 107,520, + // lds_wait_instructions 19,597). Making the guard trip count + // compile-time 8 lets the backend fully unroll the chain, rename the + // fragment registers per step and hoist later steps' LDS reads across the + // current step's v_mmac (local exact-flag codegen check of this edit: 8 + // static v_mmac in the specialized path with the next steps' LDS reads + // issued before the current v_mmac; vgpr 34 / private_segment 0 / + // group_segment 14,592 B unchanged -> the 4-blocks/CU = 8-waves/CU + // occupancy is preserved). Both branches issue the identical + // ascending-s du_mma sequence into the same accumulator over the same + // staged byte mirrors, so every partial byte and every combine output + // byte is bit-identical to iteration 16; the generic rolled loop stays + // verbatim as the else path for any non-guard (my_steps != 8) + // configuration. + if (my_steps == kStageMaxSteps) { +#pragma unroll + for (int s = 0; s < kStageMaxSteps; ++s) { + du::dumma::du_load_matrix_sync( + a_frag, stage_a + static_cast(s) * kDummaK, + kStageARowStride); + du::dumma::du_load_matrix_sync( + b_frag, + stage_b + + static_cast(s) * kDummaK * kStageBRowStride + b_col, + kStageBRowStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + } else { + for (int s = 0; s < my_steps; ++s) { + du::dumma::du_load_matrix_sync( + a_frag, stage_a + static_cast(s) * kDummaK, + kStageARowStride); + du::dumma::du_load_matrix_sync( + b_frag, + stage_b + + static_cast(s) * kDummaK * kStageBRowStride + b_col, + kStageBRowStride); + du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + } + + // Epilogue (iteration 16, barrier-decoupled): each wave materializes its + // own tile into its own LDS acc buffer (wave 0 -> acc_tile0, wave 1 -> + // acc_tile1, 1,024 B each), then its 64 lanes read back that same buffer + // and store the 256 int32 partials (4 stores per lane) to + // plane[split][m][n0 + b_col + col]. acc_tile{w} is written and read ONLY + // by wave w, and the two waves' plane stores touch disjoint 16-column + // ranges, so nothing in the epilogue needs cross-wave convergence. The + // former epilogue __syncthreads() was a 128-thread cross-wave join whose + // release was gated by the slower of the block's two waves on a kernel + // that runs only 8 resident waves/CU (4 x 2-wave blocks) to hide the skew + // -- the same join whose five-wave form measured ~3.15 us of pure cost + // (iteration 13 23.764/24.086 -> iteration 15 20.613/20.641 us with the + // join removed, exact correctness passing in both). It is replaced by + // __threadfence_block(), which only orders THIS wave's own outstanding LDS + // writes (the du_store_matrix_sync ds_writes of its private acc tile) + // before its read-back ds_reads and costs no cross-wave wait. The staging + // __syncthreads() above stays (wave 0 stages the A slice and wave 1 stages + // the B strip; every wave reads bytes the other wave staged). Every + // partial byte, the k-ascending int32 order and every plane store address + // are unchanged, so the output is bit-identical to iteration 9 and exact + // correctness must pass with 0 mismatches. + if (w == 0) { + du::dumma::du_store_matrix_sync(acc_tile0, acc_frag, kDummaN, + du::dumma::mem_row_major); + } else { + du::dumma::du_store_matrix_sync(acc_tile1, acc_frag, kDummaN, + du::dumma::mem_row_major); + } + __threadfence_block(); // iteration 16: wave-local store->read ordering + // only; no cross-wave join (was __syncthreads()) + + int32_t* plane = partials + static_cast(split) * kDummaM * n; + const int32_t* tile = (w == 0) ? acc_tile0 : acc_tile1; + const int out_col0 = n0 + b_col; +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; + linear += kDummaWaveSize) { + const int row = linear / kDummaN; + const int col = linear - row * kDummaN; + plane[static_cast(row) * n + out_col0 + col] = tile[linear]; + } } // --------------------------------------------------------------------------- -// Combine + scale kernel (M == 16): one 16x16 tile per block, one thread per -// output element. Exact int32 sum of the SPLIT_K workspace partials -// (ascending split order), then x_scale/weight_scale scaling and the bf16 -// (RN-even) store. No barrier is needed: every thread reads only its own -// 16x16 element's S partial values (written by the previous kernel on the -// same stream). +// Split-K combine kernel (iteration 2, restructured in iteration 9, mapping +// v3 in iteration 25): sums the gridDim.y int32 partial planes ascending per +// output element (exact int32, no overflow, bit-identical to the scalar +// k-ascending dot), applies x_scale[row] * weight_scale[col] in the scalar +// float order and stores bf16 through f32_to_bf16_rne. Iteration 9 added +// the guard specialization (split_count == 16): one output element per lane, +// all sixteen plane loads issued as INDEPENDENT reads into a 16-int register +// array (single vmcnt drain instead of a dependent load->wait->add chain of +// 16 DRAM round trips per element), then the ascending-s register sum and +// the scaled bf16 store. Iteration 25 (this round) keeps that per-lane +// structure and changes ONLY the block-to-output mapping: every one-wave +// block owns ONE m-row x kDummaCombineColsPerBlock = 64 consecutive columns +// (guard grid = 20 x 16 = 320 one-wave blocks, same block count as +// iteration 9). In the iteration-9/23 mappings each block covered 4 rows x +// 16/32 columns, so every per-plane wavefront load issued four 64/128-B runs +// at the 5,120-B m-row stride -- four different DRAM rows, i.e. up to four +// row activations per load instruction -- which is why widening the per-row +// run 64 -> 128 B in iteration 23 could not move the combine (it never +// reduced the rows-per-instruction scatter). With ONE row per block every +// per-plane load instruction is 64 lanes x 4 B = 256 B CONTIGUOUS = four +// consecutive 64-B lines of one DRAM row (a single run: ~4x fewer DRAM row +// activations for the identical 1.31 MB byte stream, identical 64-B sectors, +// identical wavefront-load count). The element set is unchanged, so every +// element's ascending-s int32 sum, the float scale order and every output +// bf16 word are bit-identical to the pre-iteration-25 combine; the generic +// loop below stays verbatim for every other call site / split count / grid +// shape. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(kCombineThreads) void -w8a8_dumma_m16_combine_scale_kernel( - const int32_t* __restrict__ partials, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - hip_bfloat16* __restrict__ out, - int split_k, - int num_tiles, - int n) { - const int tile = static_cast(blockIdx.x); - const int linear = static_cast(threadIdx.x); // 0..255 - const int row = linear >> 4; - const int col = linear & 15; - const int out_col = tile * kDummaN + col; - - int32_t sum = 0; - for (int s = 0; s < split_k; ++s) { - sum += partials[(s * num_tiles + tile) * (kDummaM * kDummaN) + linear]; +__global__ void w8a8_dumma_m16_combine_kernel( + const int32_t* __restrict__ partials, // [split_count][16][n] int32 + const float* __restrict__ x_scale, // [16] + const float* __restrict__ w_scale, // [n] + uint16_t* __restrict__ out, // [16, n] bf16 payload + int n, + int split_count) { + const int lane = static_cast(threadIdx.x); // blockDim == 64 + const int n0 = static_cast(blockIdx.x) * kDummaN; + + // --------------------------------------------------------------------- + // Iteration 25 specialization (mapping v3; guard launch only: + // split_count == 16 and gridDim.y == kDummaM == 16 -> blockIdx.x = + // 64-column group in [0, 20), blockIdx.y = one m-row in [0, 16)). Each + // lane owns output element (row, col) with row = blockIdx.y and col = + // blockIdx.x * 64 + lane: the sixteen independent 4-B plane reads (plane + // stride kDummaM * n int32) batch into one vmcnt drain, and because the + // 64 lanes of the block are 64 CONSECUTIVE columns of ONE row, every + // per-plane wavefront load covers 256 contiguous bytes (4 consecutive + // 64-B lines of one DRAM row) instead of four 5,120-B-strided runs; the + // ascending-s register sum then reproduces the exact int32 sequence for + // this element. + // --------------------------------------------------------------------- + if (split_count == kDummaCombineSplits && gridDim.y == kDummaM) { + const int row = static_cast(blockIdx.y); + const int col = + static_cast(blockIdx.x) * kDummaCombineColsPerBlock + lane; + const int64_t elem = static_cast(row) * n + col; + int32_t v[kDummaCombineSplits]; +#pragma unroll + for (int s = 0; s < kDummaCombineSplits; ++s) { + v[s] = partials[static_cast(s) * (kDummaM * n) + elem]; + } + int32_t acc = 0; +#pragma unroll + for (int s = 0; s < kDummaCombineSplits; ++s) { + acc += v[s]; // ascending split order: identical int32 sum + } + const float scaled = + static_cast(acc) * x_scale[row] * w_scale[col]; + out[elem] = f32_to_bf16_rne(scaled); + return; } - const float scaled = - static_cast(sum) * x_scale[row] * weight_scale[out_col]; - out[row * n + out_col] = __float2bfloat16(scaled); +#pragma unroll + for (int linear = lane; linear < kDummaM * kDummaN; + linear += kDummaWaveSize) { + const int row = linear / kDummaN; + const int col = linear - row * kDummaN; + const int out_col = n0 + col; + int32_t acc = 0; + for (int s = 0; s < split_count; ++s) { + acc += partials[(static_cast(s) * kDummaM + row) * n + + out_col]; + } + const float scaled = + static_cast(acc) * x_scale[row] * w_scale[out_col]; + out[static_cast(row) * n + out_col] = f32_to_bf16_rne(scaled); + } } // --------------------------------------------------------------------------- -// Weight pack op (outside the timed region). +// Identity device-to-device packing (bootstrap): grid-stride byte copies on +// the caller stream. Valid for every (K, N), including unmatched shapes. // --------------------------------------------------------------------------- -// For every (K,N) except the exact pair (k,n)==(4096,1280) the packed weight -// stays a contiguous [K, N] int8 buffer with an [N, 1] fp32 scale, -// byte-identical to the raw inputs. For (k,n)==(4096,1280) (iteration 11) -// launch_pack_w8a8_weight transposes the raw [K,N] weight to an [N,K] -// n-major packed layout (element (kk,nn) -> dst[nn*k + kk]) so DUMMA -// col_major B fragments read 8 consecutive bytes per lane (one ds_read_b64 -// instead of 8 ds_read_u8); the staged/direct DUMMA kernels and the scalar -// fallback (for the paired M=2 validation shape) all decode that layout. -__global__ __launch_bounds__(kCopyBlockThreads) void -w8a8_identity_copy_i8_kernel( +__global__ void w8a8_identity_copy_i8_kernel( const int8_t* __restrict__ src, int8_t* __restrict__ dst, int64_t count) { - const int64_t linear = - static_cast(blockIdx.x) * kCopyBlockThreads + threadIdx.x; - if (linear < count) { - dst[linear] = src[linear]; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + i < count; i += stride) { + dst[i] = src[i]; } } -__global__ __launch_bounds__(kCopyBlockThreads) void -w8a8_identity_copy_f32_kernel( +__global__ void w8a8_identity_copy_f32_kernel( const float* __restrict__ src, float* __restrict__ dst, int64_t count) { - const int64_t linear = - static_cast(blockIdx.x) * kCopyBlockThreads + threadIdx.x; - if (linear < count) { - dst[linear] = src[linear]; - } -} - -// Iteration 11 one-time pack for the exact pair (k,n)==(4096,1280): -// transpose the raw [K,N] int8 weight to [N,K] n-major (element (kk,nn) -> -// dst[nn*k + kk]). Runs out of the timed region and out of Graph capture; -// performance is irrelevant (single weight prep per layer load). -__global__ __launch_bounds__(kCopyBlockThreads) void -w8a8_transpose_i8_kernel( - const int8_t* __restrict__ src, - int8_t* __restrict__ dst, - int k, - int n) { - const int64_t linear = - static_cast(blockIdx.x) * kCopyBlockThreads + threadIdx.x; - const int64_t total = static_cast(k) * n; - if (linear >= total) { - return; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + static_cast(threadIdx.x); + i < count; i += stride) { + dst[i] = src[i]; } - const int kk = static_cast(linear / n); - const int nn = static_cast(linear - static_cast(kk) * n); - dst[static_cast(nn) * k + kk] = src[linear]; } -} // namespace - // --------------------------------------------------------------------------- -// Stable host launch symbols (declared in csrc/bindings.cpp). +// Host launcher consumed by csrc/bindings.cpp (gemm_out, timed region). +// The exact (16, 1280, 4096) guard dispatches the split-K DUMMA partial + +// combine kernels (two wavefronts per partial block since iteration 8); +// the scalar fallback below launches 128 threads per 128 outputs. // --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, @@ -765,128 +1131,118 @@ extern "C" void launch_w8a8_gemm( int n, int k, hipStream_t stream) { - auto* out_bf16 = static_cast(out); + (void)workspace; + (void)workspace_bytes; + if (m <= 0 || n <= 0 || k <= 0 || a == nullptr || b == nullptr || + x_scale == nullptr || weight_scale == nullptr || out == nullptr) { + return; + } - // Exact-shape guard for the assigned shape hy3_tp8_qkv_proj_m16 - // (M=16, N=1280, K=4096). The DUMMA specialization can never capture - // other shapes (in particular the paired M=2 shape with the same (N, K)). + // --------------------------------------------------------------------- + // Exact-shape DUMMA guard: hy3_tp8_qkv_proj_m16 (M=16, N=1280, K=4096). + // Lineage: iteration 4 widened each one-wave block to TWO adjacent + // 16-wide N tiles (32 columns) sharing one A-fragment load per 32-K step + // at SPLIT_K=12 (grid 40 x 12 = 480 blocks = 4/CU over 120 CUs); + // iteration 5 staged each block's 32-column B slice into LDS once and + // read the B0/B1 fragments from LDS per step, A direct-global + // (39.404/39.485 us official Graph median/P90, from 41.154/41.233). + // Iteration 7: BOTH operands are staged per K slice -- the A fragment + // also comes from LDS, so the K loop is global-load-free -- which forces + // SPLIT_K=16: sixteen UNIFORM 8-step slices (128 == 16 x 8, rows + // {256 x 16}, exact ascending cover of [0, 4096), every boundary a + // DUMMA-k multiple), grid = 40 N groups x 16 = 640 one-wave blocks + // (~5.33 waves/CU, 4 resident blocks/CU by LDS: 14,592 B/block -> + // 65,536/14,592 = 4.4), each block writing exact int32 partial planes + // (cols [n0, n0+16) and [n0+16, n0+32) of plane[split]), then one + // UNCHANGED 80-block combine kernel (generic over split_count = 16) that + // sums the planes ascending and emits the scaled bf16 output. Staged + // tiles are exact byte mirrors of the global operands, so fragment + // bytes, K-slice bounds, the k-ascending int32 order and the partial + // plane layout are unchanged: the combine+scale output is bit-identical + // to iteration 5. + // Iteration 8 (this round): occupancy/resource change only -- the same + // 640-block SPLIT_K=16 geometry now runs TWO 64-lane wavefronts per block + // (blockDim 128): wave w computes tile w (cols [n0+16w, n0+16w+16)) of + // the block's 32-column strip with its own independent accumulator chain + // over the shared slice. LDS per block is UNCHANGED (14,592 B -> 4 + // resident blocks/CU), so co-residency rises from 4 to 8 resident waves + // per CU = 2 per SIMD, hiding each wave's per-step LDS->MMA stall behind + // the other wave. Staging is partitioned across the two waves (wave 0 = + // A slice, wave 1 = B strip; identical global bytes read exactly once). + // The combine kernel, the 16-plane workspace need (1,310,720 B), the + // k-ascending partial order and the bit-identical output are unchanged. + // Both launches stay on the caller stream inside the timed API/Graph; + // the combine cost is part of the operator wall. Requires 16 int32 + // partial planes (1,310,720 B) of the caller's workspace -- the full + // fixed 16-plane 1,310,720 B allocation for this shape; when the + // workspace is missing/undersized the exact guard falls back to the + // iteration-1 one-wave direct geometry (still no workspace). Every + // unmatched (m, n, k) -- including the paired M=2 API shapes with the + // same (N, K) -- keeps reaching the scalar fallback below. + // Iteration 9 (this round): combine-side restructure only. PMC of the + // accepted iteration-8 code object profiles the combine kernel at 16.960 + // us vs 19.361 us for the partial kernel (46.7% of the 36.321 us profiled + // operator aggregate) with a 6.9% L2 hit rate and only 80 one-wave blocks, + // so the guard now launches the SAME combine kernel on a 2-D grid (80 + // n-tiles x 4 four-row m-groups = 320 one-wave blocks) and the kernel's + // split_count==16 specialization gives each lane one output element that + // batches all sixteen plane reads into a register array before the + // ascending sum. Partials, output bytes, bit-exactness and the partial + // kernel launch are unchanged. + // --------------------------------------------------------------------- if (m == 16 && n == 1280 && k == 4096) { - const int num_tiles = n / kDummaN; // 1280 / 16 = 80 - const int64_t plane_bytes = - static_cast(num_tiles) * kDummaM * kDummaN * - static_cast(sizeof(int32_t)); // 80 * 256 * 4 = 81,920 - - // Split-K selection: default 16 (the iteration-8 staged geometry: 1,280 - // one-wave blocks = 10.67 blocks/CU; per-block LDS drops from 12,736 B - // (5 resident/CU, the binding occupancy limiter at S=12 per PMC: grid - // 960 = 8 waves/CU available, VGPR 24 / 0 scratch not limiting) to - // 9,280 B = 7 resident/CU, with uniform 256-K slices = exactly 8 MMAs - // per block; A/B HBM bytes unchanged). The ZTH_W8A8_QKV_SPLIT_K - // environment variable selects any trusted probe candidate - // {2,3,4,5,6,8,9,12,16} so the sweep can be measured without source - // edits (S=6 -> 480 blocks = 4/CU, S=9 -> 720 = 6/CU, S=12 -> 960 = 8/CU, - // S=16 -> 1,280 = 10.67/CU, S=3 -> 240 direct-load blocks, ...). The - // read happens once per launch call on the host, never during Graph - // replay, so capture/replay determinism is preserved. - int split_k = 16; - if (const char* env = std::getenv("ZTH_W8A8_QKV_SPLIT_K")) { - char* end = nullptr; - const long parsed = std::strtol(env, &end, 10); - if (end != env && *end == '\0') { - for (int t : kTrustedSplitK) { - if (static_cast(parsed) == t) { - split_k = t; - break; - } - } - } + constexpr int kDummaThreads = 64; // one wavefront (tile/combine) + constexpr int kPartialThreads = 2 * kDummaThreads; // 2 waves per block + constexpr int kTilesPerWave = 2; // two 16-wide N tiles per block (n32) + constexpr int kSplitK = 16; // 40 groups x 16 = 640 blocks, 4/CU + const int ntiles = n / kDummaN; // 1280 / 16 = 80 combine tiles + const int ngroups = ntiles / kTilesPerWave; // 80 / 2 = 40 n32 groups + const int64_t kWorkspaceNeeded = + static_cast(kSplitK) * kDummaM * n * sizeof(int32_t); + auto* out_bf16 = reinterpret_cast(out); + if (workspace != nullptr && workspace_bytes >= kWorkspaceNeeded) { + auto* partials = reinterpret_cast(workspace); + hipLaunchKernelGGL(w8a8_dumma_m16_splitk_partial_kernel, + dim3(ngroups, kSplitK), dim3(kPartialThreads), 0, + stream, a, b, partials, n, k); + // Iteration 25 (mapping v3): 2-D combine grid = 20 column-groups x 16 + // one-row blocks = 320 one-wave blocks (one output row x 64 consecutive + // columns per block, one output element per lane, 16-deep independent + // plane-load register batch in the kernel's split_count==16 path; every + // per-plane wavefront load is one 256-B contiguous run of one DRAM row + // instead of four 5,120-B-strided runs). + hipLaunchKernelGGL(w8a8_dumma_m16_combine_kernel, + dim3(n / kDummaCombineColsPerBlock, kDummaM), + dim3(kDummaThreads), 0, stream, + partials, x_scale, weight_scale, out_bf16, n, + kSplitK); + } else { + // Workspace-less fallback: iteration-1 one-wave direct geometry. + hipLaunchKernelGGL(w8a8_dumma_m16_tile_kernel, + dim3(ntiles), dim3(kDummaThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, n, k); } - - // Workspace fit: clamp to the largest trusted candidate that fits the - // caller workspace (contract guarantees 16 planes = 1,310,720 B for this - // shape, so the clamp never fires; kept defensive). - int best_fit = 0; - for (int t : kTrustedSplitK) { - if (static_cast(t) * plane_bytes <= workspace_bytes) { - best_fit = t; - } - } - if (best_fit == 0) { - launch_scalar_gemm(a, b, x_scale, weight_scale, out_bf16, m, n, k, - stream); - return; - } - if (split_k > best_fit) { - split_k = best_fit; - } - - // Int32 partial planes in the caller workspace, then the combine+scale - // kernel, both on the caller's stream inside the timed Graph. - // Iteration-7/8 staged kernels: S in {6,8,9,12,16} (their longest slice - // fits the 64 KiB LDS; see w8a8_dumma_m16_slicestage_partial_kernel). - // Iteration-3 direct-load kernels: S in {2,3,4,5} (their slices exceed - // the staging capacity). Both families produce the same slice-major - // int32 planes, so the combine kernel is shared and unchanged. - int32_t* partials = static_cast(workspace); - switch (split_k) { - case 6: - launch_slicestage_partial<6>(a, b, partials, n, k, stream, - num_tiles); - break; - case 8: - launch_slicestage_partial<8>(a, b, partials, n, k, stream, - num_tiles); - break; - case 9: - launch_slicestage_partial<9>(a, b, partials, n, k, stream, - num_tiles); - break; - case 12: - launch_slicestage_partial<12>(a, b, partials, n, k, stream, - num_tiles); - break; - case 16: - launch_slicestage_partial<16>(a, b, partials, n, k, stream, - num_tiles); - break; - case 2: - launch_splitk_partial<2>(a, b, partials, n, k, stream, num_tiles); - break; - case 3: - launch_splitk_partial<3>(a, b, partials, n, k, stream, num_tiles); - break; - case 4: - launch_splitk_partial<4>(a, b, partials, n, k, stream, num_tiles); - break; - case 5: - launch_splitk_partial<5>(a, b, partials, n, k, stream, num_tiles); - break; - default: - launch_scalar_gemm(a, b, x_scale, weight_scale, out_bf16, m, n, k, - stream); - return; - } - - hipLaunchKernelGGL( - w8a8_dumma_m16_combine_scale_kernel, - dim3(static_cast(num_tiles)), - dim3(static_cast(kCombineThreads)), - 0, - stream, - partials, - x_scale, - weight_scale, - out_bf16, - split_k, - num_tiles, - n); return; } - // Generic scalar fallback for every unmatched (m, n, k). - launch_scalar_gemm(a, b, x_scale, weight_scale, out_bf16, m, n, k, stream); + // --------------------------------------------------------------------- + // Scalar generic fallback: every shape outside the exact guard above. + // --------------------------------------------------------------------- + + constexpr int kBlockThreads = 128; // multiple of the gfx928 wavefront 64 + const int total = m * n; + const int grid = (total + kBlockThreads - 1) / kBlockThreads; + auto* out_bf16 = reinterpret_cast(out); + hipLaunchKernelGGL(w8a8_scalar_gemm_kernel, + dim3(grid), dim3(kBlockThreads), 0, stream, + a, b, x_scale, weight_scale, out_bf16, m, n, k); } +// --------------------------------------------------------------------------- +// Host launcher for the optional out-of-timed-region weight packing. +// Bootstrap: identity device-to-device copy of the int8 [K, N] weight and +// the fp32 [N, 1] scales; valid for every (K, N), unmatched shapes included. +// --------------------------------------------------------------------------- extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -895,47 +1251,27 @@ extern "C" void launch_pack_w8a8_weight( int k, int n, hipStream_t stream) { - // Out-of-timed-region weight prep: [N,K] n-major transpose pack for the - // exact pair (k,n)==(4096,1280) (iteration 11), identity copy otherwise. - const int64_t weight_count = static_cast(k) * n; - const unsigned weight_grid = static_cast( - (weight_count + kCopyBlockThreads - 1) / kCopyBlockThreads); - if (k == kPackedK && n == kPackedN) { - // Iteration 11: [N,K] n-major transpose pack for the exact pair so DUMMA - // col_major B fragments read 8 consecutive bytes per lane (one - // ds_read_b64 instead of 8 ds_read_u8). - hipLaunchKernelGGL( - w8a8_transpose_i8_kernel, - dim3(weight_grid), - dim3(kCopyBlockThreads), - 0, - stream, - raw_weight, - packed_weight, - k, - n); - } else { - hipLaunchKernelGGL( - w8a8_identity_copy_i8_kernel, - dim3(weight_grid), - dim3(kCopyBlockThreads), - 0, - stream, - raw_weight, - packed_weight, - weight_count); + constexpr int kPackBlockThreads = 256; + + const int64_t i8_count = static_cast(k) * n; + if (i8_count > 0 && raw_weight != nullptr && packed_weight != nullptr) { + const int i8_grid = + static_cast((i8_count + kPackBlockThreads - 1) / + kPackBlockThreads); + hipLaunchKernelGGL(w8a8_identity_copy_i8_kernel, + dim3(i8_grid), dim3(kPackBlockThreads), 0, stream, + raw_weight, packed_weight, i8_count); } - const unsigned scale_grid = static_cast( - (static_cast(n) + kCopyBlockThreads - 1) / kCopyBlockThreads); - hipLaunchKernelGGL( - w8a8_identity_copy_f32_kernel, - dim3(scale_grid), - dim3(kCopyBlockThreads), - 0, - stream, - weight_scale, - packed_weight_scale, - static_cast(n)); + const int64_t f32_count = static_cast(n); + if (f32_count > 0 && weight_scale != nullptr && + packed_weight_scale != nullptr) { + const int f32_grid = + static_cast((f32_count + kPackBlockThreads - 1) / + kPackBlockThreads); + hipLaunchKernelGGL(w8a8_identity_copy_f32_kernel, + dim3(f32_grid), dim3(kPackBlockThreads), 0, stream, + weight_scale, packed_weight_scale, f32_count); + } } // @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_down_proj.hip index 01194376..5d133848 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_down_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_down_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=hy3_tp8_shared_down_proj_m16 commit=6f8e396e4eb7b316acecbb928ea5db9f2f3ffe4d added=2026-08-27 +// baseline_us=17.312 speedup=2.24948 // median_us=7.696 p90_us=8.174 // source=hy3-dsh-tp8-m16-2-dc53295a // INT8 W8A8 GEMM implementation for gfx928 (K500SM_AI). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_gate_up_proj.hip index 3b0c1c88..1f40f800 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_gate_up_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_gate_up_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=hy3_tp8_shared_gate_up_proj_m16 commit=2321c42b8618aae23700bc2cf3532b59c8f6c8d8 added=2026-08-27 +// baseline_us=66.527 speedup=6.40915 // median_us=10.38 p90_us=10.42 // source=hy3-dsh-tp8-m16-2-dc53295a // MetaInfer W8A8 INT8 GEMM - HIP implementation (gfx928 / K500SM_AI). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/o_proj.hip index 8cf04a9f..5be78d23 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/o_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/o_proj.hip @@ -1,132 +1,85 @@ -// @@variant shape=hy3_tp8_o_proj_m4096 commit=710edb90ef1f462cd0751356a542a14e8bbba8c6 added=2026-08-27 -// median_us=323 p90_us=325.4 speedup=43.82 baseline_us=1.415e+04 -// source=hy3-dsh-tp8-m4096-1-0ebb994d -// @@variant shape=hy3_tp8_o_proj_m4096 added=2026-08-27 (bootstrap iteration 1) +// @@variant shape=hy3_tp8_o_proj_m4096 commit=5a89acbee264bee5d61f74c281509b86c3fcd4e3 added=2026-09-06 +// median_us=266.8 p90_us=267.2 speedup=53.05 baseline_us=1.415e+04 +// source=hy3-tp8-dsh-m16-and-m4096-2-f7fa9503 // INT8 W8A8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. // -// Worker: worker_1 (physical GPU 1), assigned shape: +// Worker: worker_1 (physical GPU 1). Assigned shapes: +// hy3_tp8_o_proj_m16 : M=16, N=4096, K=1024 // hy3_tp8_o_proj_m4096 : M=4096, N=4096, K=1024 // // Bootstrap strategy (iteration 1, correctness-first but a usable profiling // baseline - a large-Prefill scalar K loop would not be): -// * Large-prefill path (exact (m, n, k) == (4096, 4096, 1024)): +// +// * Large-prefill path (exact (m, n, k) == (4096, 4096, 1024), m4096 +// iteration-1 architecture round - DUMMA throughput baseline with the +// accepted lineage steady state): // native INT8 DUMMA m16n16k32 with int32 accumulation; one 128x64 -// output tile per block; four wavefronts (256 threads); each wave owns -// a 64x32 quadrant built from eight m16n16k32 accumulator fragments; -// the block cooperatively vector-loads A[128,64] and B[64,64] into one -// single-buffered 64-K LDS stage (15,360 B total: 128*80 + 64*80 with -// 16 B padding per row for bank skew); two __syncthreads per stage; -// fused dot * x_scale[m] * weight_scale[n] epilogue stored directly as -// bf16 from the accumulator fragments. Grid dim3(64, 32) = 2048 blocks -// dwarfs the 120 CUs, so no split-K is needed. -// The weight is read in the packed n-major [N, K] int8 layout -// (packed[n*K + kk] == raw[kk*N + n], a full transpose; identity -// [K, N] layout is retained for every other (k, n)). Every B tile of -// one (n-tile, K-stage) is one contiguous 4 KiB vector stream in both -// global and LDS, so each thread's B staging load/store is a single -// 16-byte-aligned int4. launch_pack_w8a8_weight performs the -// permutation once, outside the timed region and outside Graph -// capture, into the same byte-count buffer (graph-stable). -// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), -// including all small-M API cases (M=2, M=16) and M=3072 with the same -// (K, N) = (1024, 4096); the fallback decodes the packed n-major [N, K] -// layout for (k, n) == (1024, 4096) and the identity [K, N] layout for -// every other (k, n). -// * launch_pack_w8a8_weight: for (k, n) == (1024, 4096) transposes the raw -// [K, N] weight into the [N, K] int8 layout; for any other (k, n) it -// is an identity device-to-device copy. Scales are copied identity. -// Packing never happens inside the timed GEMM. +// output tile per block; four wavefronts (256 threads, gfx928 wavefront +// = 64); each wave owns a 64x32 quadrant built from eight m16n16k32 +// int32 accumulator fragments. The block cooperatively vector-loads +// A[128,64] (x_q rows m0..m0+127, row-major stride k) and B[64,64] +// (packed [N, K] weight rows n0..n0+63, n-major stride k) into one +// single-buffered LDS stage (16,384 B total: a_tile 128*88 + b_tile +// 64*80; A rows padded 64 -> 88 B and staged as two 8-B halves per +// 16-B vector, B rows padded 64 -> 80 B); two __syncthreads per stage +// (one after the cooperative load, one before the next stage overwrites +// LDS). Each DUMMA operand fragment is filled by the direct 8-byte +// load_fragment8 LDS read (one contiguous 8-B read per lane per +// fragment on the m-major A tile / n-major B tile - no library loader +// byte-reassembly VALU and no per-byte ds_read_u8), two A fragments +// live per MMAC burst, and the epilogue is fused: int32 accumulator * +// x_scale[m] * weight_scale[n] -> bf16 through the coalesced fragment +// store (4x4 lane transpose + one 8-byte store per lane, 100% store +// sector efficiency). +// Grid dim3(64, 32) = 2048 output-tile blocks dwarfs the 120 CUs, so no +// split-K is used (nothing is ever written to the workspace). // -// Iteration 6 (epilogue round): the scales and bf16 conversion were already -// fused in the compute kernel and the launcher already ignores the workspace -// (no split-K combine pass exists), so the remaining epilogue inefficiency -// is the store pattern: the m16n16k32 lane mapping (row = lane & 15, -// c4 = lane >> 4, frag.x[i] -> column c4 + 4*i) makes the direct store four -// 2-byte scalar stores per lane whose wavefront addresses touch each 32-B -// sector at 25% utilization (32 global_store_short_d16_hi per wavefront per -// block in the gfx928 ISA; vmem_write_instructions 262,144). The new -// store_prefill_fragment_coalesced transposes the 4-element groups within -// each 4-lane column group (two __shfl_xor steps, 16 then 32; this DTK -// lowers them to ds_bpermute at the block tail where the LDS pipe is idle), -// so lane (r, c4) holds the four CONTIGUOUS columns 4*c4..4*c4+3 and issues -// ONE 8-byte store per lane: 64 lanes x 8 B = 512 B per fragment per -// wavefront in 16 fully-used 32-B sectors (100% sector efficiency; -// vmem_write_instructions -> 65,536). The MMAC loop, LDS staging, barriers -// and int32 accumulation order are untouched, and the per-element scale -// multiply order plus bf16 rounding are unchanged => stored bf16 bits are -// identical (mismatch 0, max_abs_error 0.0 expected). This is the exact -// mechanism validated on the worker-29 TP4 qkv_proj lineage (accepted -// iteration 6, 770.14 -> 721.87 us). +// * Decode path (exact (m, n, k) == (16, 4096, 1024), iteration 5 +// architecture round): fused single-kernel split-K=2 - one 16x16 N tile +// per block, grid 256 x 128 threads (2.13 blocks/CU on the measured 120 +// CUs, two co-resident 64-lane wavefronts per block), wave w +// accumulating the ascending-k half w (K/2 = 512, 16 m16n16k32 steps) +// of the tile with the iteration-3 per-wave body (direct global +// fragment loads: matrix_a row_major off x_q, matrix_b col_major off +// the packed [N, K] weight, depth-1 double-buffered prefetch, no LDS in +// the K loop). The split combine is fused into this one kernel: wave 0 +// publishes its int32 accumulator fragment to a 1-KiB shared plane, one +// END-of-K __syncthreads makes it visible, and wave 1 adds plane 0 + +// its own upper-K-half fragment (exact int32, ascending split order, +// bit-identical to the unsplit ascending-k dot) and applies the fused +// x_scale * weight_scale -> bf16 epilogue. One launch on the caller +// stream inside the timed Graph, no caller workspace and no second +// combine kernel: the exact-shape Graph has a single node and no +// ~1-MiB workspace write+read round trip. Iteration 3's two-launch +// workspace partial+combine variant and iteration 2's unsplit two-wave +// kernel are superseded (the fused kernel needs no workspace, so no +// workspace-fit fallback branch remains). // -// Iteration 10 (B-fragment load round): the gfx928 code object of the -// accepted kernel (419.05 us median, 81.995 TOPS, arch_vgpr 80, LDS 15,360 B, -// grid 2048, workgroup 256) shows the DUMMA B-fragment path compiled as ~32 -// per-wave ds_read_u8 byte loads (the row_major int8 loader assigns -// x[i] = p[col*ldm + row + i], i.e. 8 k-rows strided by ldm per lane) plus a -// mask/OR byte-reassembly VALU chain (~6 ops per operand dword) with long -// lgkmcnt wait chains; PMC: lds_bank_conflicts 15,728,640 = 3 x (4,194,304 B -// ds_read_u8 + 1,048,576 A ds_read2_b32) and lds_wait_instructions -// 11,718,821. This round stages B n-major in LDS (b_tile[n][k], 64 rows x 80 -// B - same 5,120 B) and reads each B fragment as ONE 8-byte LDS load written -// directly into the fragment storage (the same 8 bytes in the same x[0..7] -// order; validated load_fragment8 pattern from the worker-29 TP4 gate_up -// lineage iteration 13, VALU 16.08M -> 6.25M), so the per-byte loads, the -// pack VALU and their waits disappear; du_mma_sync casts b.x straight to the -// v_mmac operand, so the operand bit pattern and the int32 accumulation are -// bit-identical (mismatch 0, max_abs_error 0.0 expected). The global pack for -// (k, n) == (1024, 4096) becomes n-major [N, K] so staging stays one -// coalesced int4 per thread (same trade as the worker-29 TP4 down_proj -// iteration 14), and the scalar fallback decodes the same n-major layout. -// Tile, A staging/A fragments, barriers, staging bytes, per-accumulator MMAC -// order and the iteration-6 epilogue are untouched. +// * Generic scalar int8/int32 fallback for every unmatched (m, n, k), +// including the paired M=2 API shape with the same (K, N) = (1024, +// 4096) and any other M tail: one thread per output element, exact +// int32 dot over K, packed [N, K] weight decode (column n = one +// contiguous k-run), fused x_scale * weight_scale, bf16 store. // -// Iteration 11 (repair round): the iteration-10 candidate measured the fast -// mapping (327.36 us median, 104.96 TOPS, ~43.23x vs the fixed Triton -// baseline) but failed correctness (mismatch 16,775,923, max_abs_error 86.5) -// because its pack kernel vectorized along the wrong axis: it copied -// raw[kk][n0..n0+15] (16 consecutive n at fixed k) into packed[n0][kk..kk+15] -// (which must hold 16 consecutive k at fixed n). That wrote only the packed -// rows with n0 % 16 == 0 (93.66% of the buffer unwritten) and did misaligned -// int4 stores whenever kk % 16 != 0. The DUMMA kernel itself is verified -// correct against the gfx928 code object: the B-fragment path compiles to -// four ds_read2_b64 at exactly b_tile[(local_col + lane&15)*80 + kk + -// (lane>>4)*8] for (kk, frag) = (0,0), (32,0), (1280,1), (1312,1) (8-byte -// offset units), the A-fragment ds_read2_b32 instructions are byte-identical -// to the accepted kernel, and the v_mmac operand order (D, A, B) is -// unchanged, so with a correct n-major pack the int32 accumulation and the -// stored bf16 bits are the reference bits (mismatch 0, max_abs_error 0.0 -// expected). This round changes ONLY w8a8_pack_o_proj_panels_kernel: each -// thread now gathers its 16-byte destination vector from 16 raw rows -// (aligned byte loads strided by N) and stores one aligned int4; the pack is -// one-time, out-of-timed-region and out-of-Graph, so the timed GEMM kernel -// (staging, load_fragment8 fragments, MMAC order, epilogue, grid 2048, -// workgroup 256, LDS 15,360 B, arch_vgpr 82, zero scratch) is bit-identical -// to iteration 10. The exact (m, n, k) guard, the scalar fallback decode -// (packed[col*K + kk] == raw[kk*N + col]) and the identity packs are -// untouched. +// * launch_pack_w8a8_weight: out-of-Graph, out-of-timed device transpose +// of the raw [K, N] weight into the packed [N, K] buffer +// (packed[n][k] = weight[k][n]) and an identity copy of weight_scale +// into packed_weight_scale for every (k, n). The [N, K] packed layout +// is what every GEMM consumer below reads: for the gfx928 INT8 +// m16n16k32 matrix_b fragment each lane owns 8 consecutive k bytes of +// one n row, so the [K, N] identity layout would force 8 scattered +// 1-byte global loads per lane per k32 step while the [N, K] layout +// makes them one contiguous 8-B vector load per lane per step. +// Packing never happens inside the timed GEMM and stays outside Graph +// capture; the packed buffer keeps the same byte count (K*N) and the +// same allocated address as raw_weight. // -// Iteration 12 (A-fragment load round): the iteration-11 repair was accepted -// (348.27 us median, 98.66 TOPS, mismatch 0). The exact gfx928 code object of -// the accepted kernel still shows the library du_load_matrix_sync path for the -// four matrix_a fragments emitting, per wave per stage, 8 ds_read2_b32 (A is -// m-major in LDS, so each lane's eight k-bytes are contiguous and the compiler -// already vectorized the byte loads) followed by a ~52-instruction per-dword -// mask/OR byte-reassembly VALU chain (v_and 0xff00/0xff0000/0xff000000 + -// v_or_b32_sdwa src0_sel:BYTE_0 + v_or3_b32) plus 7 lgkmcnt waitcnts that is -// arithmetic IDENTITY on every dword (the loaded bytes reach the v_mmac A -// operand unchanged): dead issue and dead latency on the A-load -> v_mmac -// dependency path. This round applies the exact iteration-10 mechanism -// (load_fragment8, validated on B; the worker-29 TP4 gate_up lineage measured -// VALU 16.08M -> 6.25M, and the TP4 qkv lineage accepted iteration 18 replaced -// the library row_major byte-reassembly with an explicit A-fragment loader) to -// the four A du_load_matrix_sync calls: each A fragment is filled by one -// 8-byte LDS read at a_tile[(local_row + (lane&15))*80 + kk + ((lane>>4)<<3)] -// writing the same 8 bytes to the same x[0..7] slots the row_major loader -// produced => the v_mmac A operand bit pattern and the int32 accumulation are -// bit-identical (mismatch 0, max_abs_error 0.0 expected). Tile, staging, the -// B load_fragment8 fragments, barriers, per-accumulator MMAC order, the -// iteration-6 coalesced epilogue, the pack layout, the exact (m, n, k) guard -// and the scalar fallback are untouched. +// Timed operator rules honored: no allocation, compilation, autotuning, +// weight packing, host/device synchronization, or default-stream launch; the +// GEMM runs on the caller-provided PyTorch current HIP stream; only `out` is +// written by the timed kernels (the fused decode split-K=2 kernel combines +// in block-local LDS, so the caller workspace is left untouched by every +// path). #include #include @@ -145,81 +98,109 @@ constexpr int kWaveSize = 64; // gfx928 native wavefront constexpr int kTargetM = 4096; constexpr int kTargetN = 4096; constexpr int kTargetK = 1024; +constexpr int kDecodeM = 16; // assigned decode shape M constexpr int kBlockM = 128; constexpr int kBlockN = 64; constexpr int kStageK = 64; -constexpr int kBPad = 16; // 64 -> 80-byte LDS row stride (five bank phases) -constexpr int kBStride = kBlockN + kBPad; -constexpr int kAStride = kStageK + kBPad; // 64 -> 80-byte A row stride -constexpr int kBlockThreads = 4 * kWaveSize; +constexpr int kPad = 16; // B: 64 -> 80-byte LDS row stride (bank skew) +constexpr int kAPad = 24; // A: 64 -> 88-byte LDS row stride (22 dwords/row); + // 88 % 16 == 8, so odd rows are only 8-B aligned + // and the A stage is written as two 8-B halves + // (ds_write_b64 class). 22 dwords/row maps the 16 + // rows of a fragment window onto 16 distinct 4-B + // word phases within each fragment q phase (the + // 20-dword 80-B rows of the accepted port alias + // rows r and r+8, doubling the LDS bank pressure + // of every A-fragment read). +constexpr int kAStride = kStageK + kAPad; // A tile row stride (m-major) +constexpr int kBStride = kBlockN + kPad; // B tile row stride (n-major: row + // = n, k bytes along the padded row) +constexpr int kBlockThreads = 4 * kWaveSize; // 256 threads = 4 wavefronts +// Decode launch geometry (iteration 3 grid, iteration 5 fused combine): grid +// 256 = one 16x16 N tile per block reaches 2.13 blocks/CU on the measured 120 +// CUs - the finest legal grid at the two-blocks-per-CU latency-hiding target, +// since the 256 whole 16-col N tiles (4096/16) admit no finer whole-tile +// split - and keeps two co-resident 64-lane wavefronts per 128-thread block +// (the iteration-2 latency-hiding mechanism: one wave's global-load stalls +// overlap the other wave's MMAC/address work). Wave w of block b accumulates +// the ascending-k segment [w*512, (w+1)*512) of tile b. Iteration 3 combined +// the two per-tile halves through two caller-workspace int32 planes plus a +// second combine+scale kernel launch; iteration 5 fuses that combine into the +// same kernel through a 1-KiB block-local shared plane and one END-of-K +// __syncthreads (wave 0 publishes, wave 1 sums plane 0 + its own upper half +// and stores scaled bf16), removing the second Graph node and the workspace +// round trip entirely. +constexpr int kDecodeThreads = kWaveSize * 2; // 128 = 2 wavefronts +constexpr int kDecodeSplitK = 2; // K segments per tile +constexpr int kDecodeSegK = kTargetK / kDecodeSplitK; // 512 +constexpr int kDecodeSegSteps = kDecodeSegK / kTileK; // 16 x k32 steps +constexpr int kNumDecodeTiles = kTargetN / kTileN; // 256 16-col tiles +constexpr int kDecodeLaneInts = 4; // int32 accumulator values per lane +constexpr int kDecodeTileInts = + kWaveSize * kDecodeLaneInts; // 256 int32 = 1 KiB shared plane per block using bf16_t = hip_bfloat16; // --------------------------------------------------------------------------- // Large-M prefill: 128x64 output tile per block, four wavefronts of 64 lanes. // Each wave owns a 64x32 quadrant (eight m16n16k32 int32 accumulators); the -// block cooperatively stages A[128,64] from x_q (row-major, stride k) and -// B[64,64] from the packed n-major [N, K] weight into a single -// 64-K LDS buffer. Two barriers per stage: one after the cooperative load, -// one before the next stage overwrites LDS. Direct fragment epilogue. +// block cooperatively stages A[128,64] from x_q (m-major, row stride k) and +// B[64,64] from the packed n-major [N, K] weight into a single-buffered LDS +// stage. Two barriers per K stage. DUMMA operand fragments are filled by the +// direct 8-byte LDS fill (load_fragment8) and the fused scale -> bf16 +// epilogue stores through the coalesced fragment store. // --------------------------------------------------------------------------- -template -__device__ __forceinline__ void store_prefill_fragment( - const AccFragment& frag, - const float* __restrict__ x_scale, - const float* __restrict__ weight_scale, - bf16_t* __restrict__ out, - int base_row, - int base_col, - int m, - int n, - int lane) { - // Verified gfx928 int8 m16n16k32 accumulator ownership (matches - // du_store_matrix_sync): lane & 15 selects the row, lane >> 4 selects - // col % 4, and x[i] maps to columns col%4 + 4*i. - const int row = base_row + (lane & 15); - if (row >= m) { - return; // tail-M masking: padded rows never write - } - const int col_mod4 = lane >> 4; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int col = base_col + col_mod4 + 4 * i; - const float scaled = static_cast(frag.x[i]) * - x_scale[row] * weight_scale[col]; - out[static_cast(row) * n + col] = __float2bfloat16(scaled); - } -} -// Iteration 6 (epilogue round): coalesced fragment store. The m16n16k32 -// accumulator lane mapping (row = lane & 15, column group c4 = lane >> 4, -// frag.x[i] -> column c4 + 4*i) gives each lane four elements strided by 4 -// columns, so the direct store (store_prefill_fragment) is four 2-byte -// scalar stores per lane whose wavefront addresses touch each 32-B sector at -// 25% utilization (8 B used of 32 B; 4 store instructions per fragment per -// wave, i.e. 32 global_store_short_d16_hi per wavefront per block in the -// current gfx928 ISA). This epilogue transposes the 4-element groups within -// each 4-lane column group (lanes r, r+16, r+32, r+48 -- a 4x4 transpose, -// two 2x2 steps with shfl_xor 16 then 32, one v_cndmask per element per -// step; no staging tile round trip), so lane (r, c4) ends up holding the +// Coalesced fragment store. The m16n16k32 accumulator lane mapping +// (row = lane & 15, column group c4 = lane >> 4, frag.x[i] -> column +// c4 + 4*i) gives each lane four elements strided by 4 columns, so a direct +// per-element store is four 2-byte scalar stores per lane whose wavefront +// addresses touch each 32-B sector at 25% utilization (vmem_write 262,144 +// for this shape). This epilogue transposes the 4-element groups within each +// 4-lane column group (lanes r, r+16, r+32, r+48 -- a 4x4 transpose, two 2x2 +// steps with __shfl_xor 16 then 32), so lane (r, c4) ends up holding the // four CONTIGUOUS columns 4*c4 .. 4*c4+3, converts them to bf16, packs 4 -// bf16 (8 B), and writes ONE 8-byte store per lane: 64 lanes x 8 B = 512 B -// per fragment per wavefront in 16 fully-used 32-B sectors (100% store -// sector efficiency; vmem_write_instructions drops 262,144 -> 65,536 for -// this shape). This is the exact mechanism validated on the worker-29 TP4 -// qkv_proj lineage (accepted iteration 6, 770.14 -> 721.87 us): this DTK's -// __shfl_xor lowers to ds_bpermute (one LDS permute per shuffle, +64 LDS -// instructions per wavefront per block -- accepted because it runs at the -// block tail where the LDS pipe is otherwise idle; there is no global round -// trip and no staging tile). Only the int32 values are re-routed between -// lanes -- the per-element scale multiply order +// bf16 (8 B), and writes ONE 8-byte store per lane (vmem_write 65,536: 100% +// store sector efficiency). This DTK lowers __shfl_xor to ds_bpermute at the +// block tail where the LDS pipe is idle. Only the int32 values are re-routed +// between lanes; the per-element scale multiply order // (float(dot) * x_scale[row] * weight_scale[col]) and the bf16 rounding are -// unchanged, so the stored bits are identical to store_prefill_fragment. -// The row>=m guard is wavefront-uniform (all 64 lanes of a wave share the -// same 16-row window base_row + (lane & 15)), so the shuffles never mix -// active and inactive lanes; base_col is a multiple of 16 and n*2 a multiple -// of 8, so the float4 weight_scale load (col0 % 4 == 0) and the 8-byte store +// unchanged, so the stored bits are identical to the direct per-element +// store. The row >= m guard is wavefront-uniform (all 64 lanes of a wave +// share the same 16-row window), so the shuffles never mix active and +// inactive lanes; base_col is a multiple of 16 and n*2 a multiple of 8, so +// the float4 weight_scale load (col0 % 4 == 0) and the 8-byte store // ((row*n + col0) % 4 == 0 elements -> byte offset % 8 == 0) are aligned. +// This is the exact mechanism validated on the accepted hy3 TP8 o_proj M4096 +// lineage (iteration 6, vmem_write_instructions 262,144 -> 65,536). + +// Plain round-to-nearest-even fp32 -> bf16 conversion (bit 16 of the fp32 +// pattern), without the NaN/Inf signaling-bit fixup of hip's +// __float2bfloat16. The arithmetic (u + 0x7fff + ((u >> 16) & 1)) >> 16 is +// exactly the finite-domain path of DTK's float_2_bfloatraw (exponent bits +// != all-ones), so every finite fp32 input - including +/-0, denormals, and +// values whose RNE carries into/out of the bf16 exponent field - produces +// the identical 16-bit pattern as __float2bfloat16; the outputs can differ +// only for NaN/Inf inputs, which this kernel's epilogue can never produce +// (|int32 dot| <= K*127^2 and both scale inputs are finite, so the scaled +// accumulator is always finite). Replacing the library conversion drops the +// exec-masked special-case chain the DTK otherwise emits per element in the +// epilogue tail (v_and exponent mask + v_cmp_ne + s_and_saveexec + s_xor + +// v_bfe + v_add3 + s_andn2_saveexec + v_mov/v_or + v_cmp_eq_sdwa + +// v_cndmask + s_or in the accepted iteration-5 ISA: ~12 extra SALU/VALU per +// element x 32 elements per wave, all executed once per block tail) for 3 +// VALU per element, with stored bits unchanged (verified against the +// float_2_bfloatraw source: the else-if branch it removes is reachable only +// from Inf/NaN inputs). +__device__ __forceinline__ unsigned short bf16_rne_finite_bits(float f) { + union { + float fp32; + unsigned int u32; + } u = {f}; + u.u32 += 0x7fffu + ((u.u32 >> 16) & 1u); // round half to even at bit 16 + return static_cast(u.u32 >> 16); +} + template __device__ __forceinline__ void store_prefill_fragment_coalesced( const AccFragment& frag, @@ -266,37 +247,37 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced( const int col0 = base_col + 4 * c4; const float xs = x_scale[row]; const float4 ws = *reinterpret_cast(weight_scale + col0); + // Left-associative scale order, matching the logical reference + // (A.float() @ B.float()) * x_scale[m] * weight_scale[n].T. const float v0 = static_cast(f0) * xs * ws.x; const float v1 = static_cast(f1) * xs * ws.y; const float v2 = static_cast(f2) * xs * ws.z; const float v3 = static_cast(f3) * xs * ws.w; const uint64_t packed = - static_cast( - static_cast(__float2bfloat16(v0))) | - (static_cast( - static_cast(__float2bfloat16(v1))) - << 16) | - (static_cast( - static_cast(__float2bfloat16(v2))) - << 32) | - (static_cast( - static_cast(__float2bfloat16(v3))) - << 48); + static_cast(bf16_rne_finite_bits(v0)) | + (static_cast(bf16_rne_finite_bits(v1)) << 16) | + (static_cast(bf16_rne_finite_bits(v2)) << 32) | + (static_cast(bf16_rne_finite_bits(v3)) << 48); *reinterpret_cast(out + static_cast(row) * n + col0) = packed; } -// Iteration 10: direct 8-byte LDS fragment fill for matrix_b on the n-major -// [N, K] B tile. du_load_matrix_sync's int8 matrix_b loaders assign -// x[i] = p[col*ldm + row + i] (row = lane & 15, col = (lane >> 4) << 3), i.e. -// eight consecutive k-values per lane on an [N, K] tile, and du_mma_sync -// passes b.x straight to the v_mmac builtin as one packed 8-byte operand; -// the library path makes the compiler emit eight per-byte ds_read_u8 plus a -// mask/OR reassembly chain per loaded dword. Writing the same 8 bytes -// directly into the fragment storage keeps the operand bit pattern identical -// (exact int32 accumulation unchanged) and lets the compiler feed a single -// 8-byte LDS read straight to the v_mmac. This is the exact pattern -// validated on the worker-29 TP4 gate_up lineage (accepted iteration 13). +// Direct 8-byte LDS fragment fill for the DUMMA m16n16k32 signed char +// operands on a padded tile whose row direction carries the fragment's +// 16-row index and whose columns hold the K bytes (A m-major tile row = m, +// B n-major tile row = n): the fragment slot of lane (r = lane & 15, +// kq = lane >> 4) is the eight CONSECUTIVE k-bytes at (lane & 15) * stride + +// (lane >> 4) * 8. The library du_load_matrix_sync int8 loaders make the +// compiler emit per-byte ds_read_u8 loads plus a mask/OR byte-reassembly +// chain per loaded dword that is an arithmetic identity on the v_mmac +// operand (dead issue and dead latency on the load -> v_mmac dependency +// path). Writing the same 8 bytes straight into the fragment storage (x[0..7] +// of the DUFragmentBase) keeps the operand bit pattern and +// the int32 accumulation identical and lets the compiler feed a single 8-byte +// LDS read to the v_mmac. This is the exact pattern validated on the worker- +// 29 TP4 gate_up lineage (accepted iteration 13) and the TP8 hy3 o_proj M4096 +// lineage (accepted iterations 10/12: 348.27 -> 322.99 us median, VALU +// ~22.56M -> ~15.7M, SQ_WAIT_INST_LDS drop, arch_vgpr 82 -> <= 64). template __device__ __forceinline__ void load_fragment8( Frag& frag, @@ -308,7 +289,7 @@ __device__ __forceinline__ void load_fragment8( *reinterpret_cast(base + off); } -__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( +__global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_prefill_kernel( const int8_t* __restrict__ x_q, const int8_t* __restrict__ weight, const float* __restrict__ x_scale, @@ -320,23 +301,29 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; const int lane = tid % kWaveSize; - const int wave_row = wave >> 1; - const int wave_col = wave & 1; + const int wave_row = wave >> 1; // quadrant rows [0,1] x 64 + const int wave_col = wave & 1; // quadrant cols [0,1] x 32 const int m0 = static_cast(blockIdx.y) * kBlockM; const int n0 = static_cast(blockIdx.x) * kBlockN; const int local_row = wave_row * 64; const int local_col = wave_col * 32; - // Single-buffered 64-K stage: A[128, 80] + B[64, 80] = 15,360 B/block - // (4 blocks/CU fit the 64 KiB LDS budget; the padded 80-byte strides are - // 16-byte-aligned and break the 64-byte LDS bank periodicity). - // Iteration 10: B is staged n-major b_tile[n][k] (64 n-rows x 80 B - the - // same 5,120 B as the old [k][n] tile), so each lane's eight B-fragment - // k-values are contiguous and load as ONE 8-byte LDS read (load_fragment8) - // instead of eight ds_read_u8 + mask/OR pack VALU. + // Single-buffered K-stage: A[128, 88] + B[64, 80] = 16,384 B/block. + // Round-8 residency: the one-stage-ahead staging payload (+12 VGPR) moves + // the code object from the accepted exact-64-VGPR 4-blocks/CU point (64 x + // 1024 threads = 65,536 VGPR exactly) to 80 VGPR => 3 blocks/CU (80 x 256 + // x 3 = 61,440 <= 65,536 VGPR; LDS 16,384 x 3 = 49,152 B; <= 85 VGPR keeps + // 3 blocks, no spills at the offline-verified 80). __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; - __shared__ __align__(16) int8_t b_tile[kBlockN * kBStride]; + __shared__ __align__(16) int8_t b_tile[kStageK * kBStride]; + // DUMMA m16n16k32 signed char operand fragments are filled by + // load_fragment8 with the eight consecutive k-bytes each lane owns on the + // padded tile (A m-major rows, B n-major rows of the [N, K] packed + // weight): no library loader, no per-byte ds_read_u8, no byte-reassembly + // VALU, one 8-byte LDS read straight into the v_mmac operand. Only two A + // fragments are live per MMAC burst, keeping the register footprint at the + // 4-blocks/CU occupancy point. DUFragment a_frag0, a_frag1; DUFragment @@ -352,55 +339,112 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( du_fill_fragment(acc30, 0); du_fill_fragment(acc31, 0); - // Cooperative staging: each thread owns one int4 in A rows [0,64), one in - // A rows [64,128), and one int4 of B. B comes from the packed n-major - // [N, K] layout (iteration 10): B[k0+stage_kv, n0+b_n] is at - // packed[(blockIdx.x*64 + b_n)*K + k0 + stage_kv], so the whole - // (n-tile, K-stage) B tile is one contiguous 4 KiB vector stream and the - // thread's global load and LDS store are each a single 16-byte-aligned - // int4. A[128,64] is 512 int4s; B[64,64] is 256 int4s. - const int vector_byte_offset = tid * static_cast(sizeof(int4)); - const int stage_row = vector_byte_offset / kStageK; - const int stage_col = vector_byte_offset - stage_row * kStageK; - const int b_n = tid >> 2; // n column of this thread's B int4 - const int b_kv = (tid & 3) * 16; // k offset (stage-local) of the B int4 - - for (int k0 = 0; k0 < k; k0 += kStageK) { - *reinterpret_cast(a_tile + stage_row * kAStride + stage_col) = - *reinterpret_cast( - x_q + static_cast(m0 + stage_row) * k + k0 + stage_col); - *reinterpret_cast(a_tile + - (stage_row + kBlockM / 2) * kAStride + stage_col) = - *reinterpret_cast( - x_q + static_cast(m0 + stage_row + kBlockM / 2) * k + - k0 + stage_col); - *reinterpret_cast(b_tile + b_n * kBStride + b_kv) = - *reinterpret_cast( - weight + (static_cast(blockIdx.x) * kBlockN + b_n) * k + - k0 + b_kv); + // Cooperative staging (identity input layouts, [N, K] packed output): + // * A[128,64] is 512 16-B chunks -> two per thread: rows [0,64) and + // [64,128). Each 16-B chunk is loaded from global as two 8-B halves + // and written to LDS as the adjacent 8-B pair (ds_write2_b64 class) + // because the 88-B A row stride is 8 (mod 16): odd rows start at an + // 8-B (not 16-B) aligned LDS offset, so a 16-B aligned LDS store would + // be misaligned. The 88-B stride (22 dwords/row) also maps the 16 rows + // of a fragment window onto 16 distinct word phases within each q + // phase, removing the r/r+8 alias of the old 80-B rows from the + // A-fragment reads. + // * B[64,64] of the packed [N, K] weight (packed[n][k]) is 256 int4s -> + // one int4 per thread; thread (tid>>2) selects the stage-local n row + // and (tid & 3)*16 the k offset, so global reads are 64-B contiguous + // per n row of the packed buffer and LDS stores are 16-B aligned along + // the padded n-major b_tile row (stride 80 unchanged). + const int row_a = tid >> 2; // 0..63 (first A row half) + const int col_a = (tid & 3) << 4; // 0,16,32,48 (8-B half at +8) + const int row_b = tid >> 2; // 0..63 (stage-local n row) + const int col_b = (tid & 3) << 4; // 0,16,32,48 (k offset) + + // Round-8 change: the cooperative staging below is software-pipelined one + // k64 stage ahead. Each thread keeps the 48 B/lane it stages for stage + // s+1 (two 16-B A row chunks + one 16-B B chunk = 12 payload VGPR) in + // registers and issues those global loads right after stage s's LDS + // stores/visibility barrier, so the full global round trip (~500+ cycles + // on gfx928) overlaps stage s's LDS fragment reads and MMACs instead of + // being exposed between the stage barriers and the LDS stores of the next + // stage (in the accepted single-buffered code the stage-s loads are issued + // at the top of the same stage, immediately followed by the vmcnt waits, + // so ~one full global latency is on the per-stage critical path of every + // wave, 16 times per block). The single LDS buffer layout, the exact + // staged bytes, the store order, the two __syncthreads per stage and the + // fragment-read/MMAC stream are unchanged - only the issue TIME of the + // global loads moves one stage earlier (their results wait in registers), + // so every int32 partial and every stored bf16 bit stays identical. The + // 12 payload VGPR take the code object past the exact 64-VGPR + // four-blocks/CU point to the 3-blocks/CU point (predicted arch_vgpr <= + // 80: 80 * 256 threads * 3 blocks = 61,440 <= 65,536 VGPR, LDS 16,384 B * + // 3 = 49,152 B): 12 resident waves per CU instead of 16, each no longer + // stalling per stage on the global round trip - the predicted occupancy + // arm of the experiment (no spill/scratch allowed; > 80 arch_vgpr would + // drop residency to 2 blocks/CU and falsify the mechanism). + const int64_t a_g0_base = + static_cast(m0 + row_a) * k + col_a; // A row 0..63, bytes + const int64_t a_g1_base = + static_cast(m0 + row_a + kBlockM / 2) * k + col_a; // row 64.. + const int64_t b_g_base = + static_cast(n0 + row_b) * k + col_b; // packed B n row, bytes + const int a_lds0 = row_a * kAStride + col_a; + const int a_lds1 = (row_a + kBlockM / 2) * kAStride + col_a; + const int b_lds = row_b * kBStride + col_b; + + // Loop-carried staging payload (48 B/lane): holds the stage-k0 chunks + // until they are stored to LDS at the top of the iteration, then is + // refilled with stage (k0 + kStageK) so the round trip overlaps the MMACs + // below. Same addresses and byte patterns as the accepted staging loads. + int k0 = 0; + unsigned long long a0_lo = 0, a0_hi = 0, a1_lo = 0, a1_hi = 0; + int4 b_chunk = make_int4(0, 0, 0, 0); + if (k0 < k) { // prologue: issue the stage-0 loads (only these are exposed) + a0_lo = *reinterpret_cast(x_q + a_g0_base); + a0_hi = *reinterpret_cast(x_q + a_g0_base + 8); + a1_lo = *reinterpret_cast(x_q + a_g1_base); + a1_hi = *reinterpret_cast(x_q + a_g1_base + 8); + b_chunk = *reinterpret_cast(weight + b_g_base); + } + for (; k0 < k; k0 += kStageK) { + // Barrier A (the accepted code's post-compute barrier, moved to the top + // of the next stage): every wave has finished reading the previous + // stage from the single LDS buffer, so the stores below may overwrite + // it. The stage-k0 global loads were issued one full stage earlier and + // landed during the previous stage's MMACs. __syncthreads(); + *reinterpret_cast(a_tile + a_lds0) = a0_lo; + *reinterpret_cast(a_tile + a_lds0 + 8) = a0_hi; + *reinterpret_cast(a_tile + a_lds1) = a1_lo; + *reinterpret_cast(a_tile + a_lds1 + 8) = a1_hi; + *reinterpret_cast(b_tile + b_lds) = b_chunk; + __syncthreads(); // Barrier B: stage k0 is visible to every wave + + if (k0 + kStageK < k) { // issue stage k0 + kStageK (overlaps the MMACs + const int64_t kn = k0 + kStageK; + a0_lo = + *reinterpret_cast(x_q + a_g0_base + kn); + a0_hi = *reinterpret_cast(x_q + a_g0_base + + kn + 8); + a1_lo = + *reinterpret_cast(x_q + a_g1_base + kn); + a1_hi = *reinterpret_cast(x_q + a_g1_base + + kn + 8); + b_chunk = *reinterpret_cast(weight + b_g_base + kn); + } - // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per kk. + // Each wave consumes its 64x32 quadrant: eight m16n16k32 MMACs per 32-k + // chunk, operands filled by the direct 8-byte loads (A m-major stride + // kAStride, B n-major stride kBStride). The kk order (0 then 32), the + // per-accumulator MMAC order and the ascending-k int32 accumulation are + // unchanged, so every accumulator partial is bit-identical to the + // library-loader path. #pragma unroll for (int kk = 0; kk < kStageK; kk += kTileK) { - // Iteration 10: B is n-major in LDS, so each lane's eight fragment - // k-values are contiguous; one 8-byte load fills the fragment (the - // same 8 bytes the old row_major loader gathered with eight ds_read_u8 - // plus mask/OR pack VALU), keeping the v_mmac operand bit pattern and - // the int32 accumulation identical. load_fragment8(b_frag0, b_tile + local_col * kBStride + kk, kBStride, lane); - load_fragment8(b_frag1, b_tile + (local_col + kTileN) * kBStride + kk, - kBStride, lane); - // Iteration 12: A gets the same direct 8-byte fill (load_fragment8). - // A is m-major in LDS, so lane (r, c4) owns the eight CONSECUTIVE - // k-values a_tile[(local_row + r)*80 + kk + c4*8 .. +7]; the code - // object shows du_load_matrix_sync already vectorizes these into - // ds_read2_b32 but still executes the per-dword mask/OR reconstruction - // (~52 identity VALU + 7 lgkmcnt waits per wave per stage) before the - // v_mmac A operands. Writing the same 8 bytes straight into the - // fragment storage removes that dead chain; the operand bit pattern - // and the int32 accumulation are unchanged. + load_fragment8(b_frag1, + b_tile + (local_col + kTileN) * kBStride + kk, kBStride, + lane); load_fragment8(a_frag0, a_tile + local_row * kAStride + kk, kAStride, lane); load_fragment8(a_frag1, a_tile + (local_row + kTileM) * kAStride + kk, @@ -420,45 +464,196 @@ __global__ __launch_bounds__(kBlockThreads) void w8a8_dumma_128x64x64_kernel( du_mma_sync(acc30, a_frag1, b_frag0, acc30); du_mma_sync(acc31, a_frag1, b_frag1, acc31); } - - // Protect the LDS buffer from the next stage's cooperative overwrite. - __syncthreads(); } const int base_row = m0 + local_row; const int base_col = n0 + local_col; - // Iteration 6: coalesced epilogue store (one 8-byte store per lane per - // fragment; scales + bf16 conversion fused in-kernel, no workspace pass). - store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, - base_row, base_col, m, n, lane); - store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, - base_row, base_col + kTileN, m, n, lane); + store_prefill_fragment_coalesced(acc00, x_scale, weight_scale, out, base_row, + base_col, m, n, lane); + store_prefill_fragment_coalesced(acc01, x_scale, weight_scale, out, base_row, + base_col + kTileN, m, n, lane); store_prefill_fragment_coalesced(acc10, x_scale, weight_scale, out, - base_row + kTileM, base_col, m, n, lane); + base_row + kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced(acc11, x_scale, weight_scale, out, - base_row + kTileM, base_col + kTileN, m, n, - lane); + base_row + kTileM, base_col + kTileN, m, n, lane); store_prefill_fragment_coalesced(acc20, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col, m, n, - lane); + base_row + 2 * kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced(acc21, x_scale, weight_scale, out, - base_row + 2 * kTileM, base_col + kTileN, - m, n, lane); + base_row + 2 * kTileM, base_col + kTileN, m, n, + lane); store_prefill_fragment_coalesced(acc30, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col, m, n, - lane); + base_row + 3 * kTileM, base_col, m, n, lane); store_prefill_fragment_coalesced(acc31, x_scale, weight_scale, out, - base_row + 3 * kTileM, base_col + kTileN, - m, n, lane); + base_row + 3 * kTileM, base_col + kTileN, m, n, + lane); +} + +// Device helper for the fused split-K combine: load a lane's 4 int32 +// accumulator values as one aligned 16-B vector (LDS plane offsets are +// 16-B aligned: 1-KiB plane, 16-B lane stride). +__device__ __forceinline__ void load4_ints(const int32_t* src, + int32_t out[kDecodeLaneInts]) { + const int4 v = *reinterpret_cast(src); + out[0] = v.x; + out[1] = v.y; + out[2] = v.z; + out[3] = v.w; +} + +// --------------------------------------------------------------------------- +// Decode fused split-K=2 kernel (iteration 5, architecture round): one 16x16 +// N tile per block, grid 256 = 2.13 blocks/CU on the measured 120 CUs (the +// finest legal grid at the two-blocks-per-CU target; the 256 whole 16-col N +// tiles admit no finer whole-tile split), block 128 threads = two co-resident +// 64-lane wavefronts. Wave w of block b accumulates the k segment +// [w*512, (w+1)*512) of tile b in ascending k order - 16 m16n16k32 steps with +// the exact iteration-3 per-wave body: same lane mapping, direct global +// fragment loads (matrix_a row_major off x_q, matrix_b col_major off the +// packed [N, K] weight, so each lane's 8 B bytes are one contiguous 8-B +// vector load per k32 step instead of 8 scattered byte loads), depth-1 +// double-buffered fragment prefetch, no LDS and no barrier inside the K loop. +// The per-block split-K combine is fused into this single kernel: wave 0 +// publishes its int32 accumulator fragment (4 int32/lane = one 16-B vector, +// lane-contiguous, 1 KiB total) to a block-local shared plane, one END-of-K +// __syncthreads (executed unconditionally by both wavefronts after the +// symmetric K loop) makes it visible, and wave 1 reads the plane back, adds +// plane 0 + its own upper-K-half fragment with exact int32 arithmetic - the +// ascending-split sum is bit-identical to the unsplit ascending-k int32 dot, +// since K=1024 keeps every partial and the total far below int32 overflow - +// and applies the fused reference epilogue chain +// ((float)acc * x_scale[row]) * weight_scale[col] -> bf16 with the verified +// fragment->out mapping. One launch on the caller stream inside the timed +// Graph: the exact-shape decode is now a single Graph node with no caller +// workspace write, no workspace read and no second combine launch (iteration +// 3's workspace partial+combine pair and iteration 2's unsplit two-wave +// kernel are superseded). 1-KiB LDS/block does not dent occupancy, and the +// one END-of-K barrier is the accepted down_proj lineage pattern at the same +// 256-block / two-wave geometry (w8a8_dumma_m16n16k32_sk2_kernel). M=16 fills +// the m16 tile exactly so every row guard is trivially satisfied; every other +// (m, n, k) including the paired M=2 API shape stays on the scalar fallback. +// --------------------------------------------------------------------------- +__global__ __launch_bounds__(kDecodeThreads) void +w8a8_dumma_m16n16_sk2_fused_kernel( + const int8_t* __restrict__ x_q, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + bf16_t* __restrict__ out, + int m, + int n, + int k) { + const int wave = + static_cast(threadIdx.x) >> 6; // k segment 0..kDecodeSplitK-1 + const int lane = + static_cast(threadIdx.x) & (kWaveSize - 1); // lane 0..63 + const int tile = static_cast(blockIdx.x); // one tile per block + const int n0 = tile * kTileN; + const int k0 = wave * kDecodeSegK; // segment base (0 or 512) + + // 1-KiB in-block int32 partial plane for the fused split-K combine: wave 0 + // publishes, wave 1 reads after the single END-of-K barrier. + __shared__ __align__(16) int32_t s_part[kDecodeTileInts]; + + DUFragment + a_frag[2]; + // matrix_b col_major: the packed weight is [N, K] (n-major rows of length + // k), so the library's col_major matrix_b loader addresses fragment slot + // (k, n) at packed[n * ldm + k]. Each lane's 8 fragment bytes (8 + // consecutive k of one n row) are then one contiguous, 8-B-aligned vector + // load per k32 step instead of 8 scattered byte loads with ldm = n over + // the identity [K, N] layout. + DUFragment + b_frag[2]; + DUFragment acc; + du_fill_fragment(acc, 0); + + // Epilogue scale prefetch (iteration 17): the fused combine/scale epilogue + // lives on the post-barrier critical tail of wave 1, where each lane + // performs 4 unique weight_scale loads + 1 x_scale load before its FMA + // chain. weight_scale is touched exactly once per kernel replay (each + // 16-col tile block reads its own 16 floats of the 4096-entry array), so + // those loads miss L1/L2 on the tail. Issue them before the K loop of the + // segment instead: identical values and the identical left-associative + // multiply chain keep the scaled-bf16 output bit-exact, while the global + // latency now hides under the 16 MMAC steps instead of stalling the + // END-of-K combine. Wave 0 never executes the branch and never reads the + // prefetched registers. + float xs; + float ws[kDecodeLaneInts]; + if (wave == 1) { + const int row = lane & 15; + const int col_mod4 = lane >> 4; + xs = x_scale[row]; +#pragma unroll + for (int i = 0; i < kDecodeLaneInts; ++i) { + ws[i] = weight_scale[n0 + col_mod4 + 4 * i]; + } + } + + // Depth-1 software pipeline over the segment (unchanged iteration-3 body: + // the global loads for step s+1 are issued just before the mmac of step s, + // so one full global round trip overlaps each mmac; kDecodeSegSteps = 16 + // constant-trip iterations let the scheduler hoist loads while registers + // allow). + du_load_matrix_sync(a_frag[0], x_q + k0, kTargetK); + du_load_matrix_sync(b_frag[0], + weight + static_cast(n0) * kTargetK + k0, + kTargetK); + int buf = 0; +#pragma unroll + for (int s = 0; s < kDecodeSegSteps; ++s) { + const int nxt = buf ^ 1; + if (s + 1 < kDecodeSegSteps) { + du_load_matrix_sync(a_frag[nxt], x_q + k0 + (s + 1) * kTileK, + kTargetK); + du_load_matrix_sync(b_frag[nxt], + weight + static_cast(n0) * kTargetK + k0 + + (s + 1) * kTileK, + kTargetK); + } + du_mma_sync(acc, a_frag[buf], b_frag[buf], acc); + buf = nxt; + } + + if (wave == 0) { + // Publish the lower-K-half partial: lane-contiguous fragment order + // acc.x[0..3] = one 16-B vector at s_part[lane * 4] (1 KiB plane). + *reinterpret_cast(&s_part[lane * kDecodeLaneInts]) = + make_int4(acc.x[0], acc.x[1], acc.x[2], acc.x[3]); + } + // One END-of-K barrier: both wavefronts execute it unconditionally after + // their symmetric K loops; makes wave 0's plane visible to wave 1. + __syncthreads(); + + if (wave == 1) { + // Wave 1 owns the upper-K-half partial in `acc`; add the lower-K-half + // plane in ascending split order (plane 0 + plane 1 = full ascending-k + // int32 dot, bit-identical to the unsplit dot by construction). The + // epilogue scales were prefetched into xs/ws before the K loop above. + int32_t p0[kDecodeLaneInts]; + load4_ints(s_part + lane * kDecodeLaneInts, p0); + const int row = lane & 15; + if (row >= m) { + return; // tail-M guard (the exact shape always has m == 16) + } + const int col_mod4 = lane >> 4; +#pragma unroll + for (int i = 0; i < kDecodeLaneInts; ++i) { + const int col = n0 + col_mod4 + 4 * i; + const float scaled = static_cast(p0[i] + acc.x[i]) * xs * ws[i]; + out[static_cast(row) * n + col] = __float2bfloat16(scaled); + } + } } // --------------------------------------------------------------------------- // Generic scalar fallback: one thread per output element. Exact int32 dot -// over K, then fused x_scale * weight_scale, then bf16 store. Handles every -// unmatched (m, n, k), including all small-M API cases (M=2, M=16) and any -// M in (0, 4096] with the same (K, N). For (k, n) == (1024, 4096) it decodes -// the packed n-major [N, K] layout (raw[kk, col] == packed[col*K + kk]); -// every other (k, n) reads the identity [K, N] layout. +// over K (k ascending), then fused x_scale * weight_scale, then bf16 store. +// Handles every unmatched (m, n, k), including the paired M=2 API shape +// (with the same (K, N) as the assigned decode/prefill shapes) and any M +// tail. Reads the packed [N, K] weight layout for every (k, n) - column n +// is one contiguous k-run at weight + n*k - exactly as launch_pack_w8a8_weight +// fills it. // --------------------------------------------------------------------------- __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( const int8_t* __restrict__ x_q, @@ -479,20 +674,11 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( const int col = static_cast(linear - static_cast(row) * n); int32_t acc = 0; const int8_t* a_row = x_q + static_cast(row) * k; - if (k == kTargetK && n == kTargetN) { - // Iteration 10: the packed layout for (k, n) == (1024, 4096) is n-major - // [N, K] (packed[col*K + kk] == raw[kk*N + col]). - for (int kk = 0; kk < k; ++kk) { - acc += static_cast(a_row[kk]) * - static_cast( - weight[static_cast(col) * k + kk]); - } - } else { - const int8_t* b_col = weight + col; - for (int kk = 0; kk < k; ++kk) { - acc += static_cast(a_row[kk]) * - static_cast(b_col[static_cast(kk) * n]); - } + // Packed weight is [N, K] (packed[n][k]), so column n is one contiguous + // k-run starting at packed + n * k (element (k, n) at [n * k + k]). + const int8_t* b_col = weight + static_cast(col) * k; + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * static_cast(b_col[kk]); } const float scaled = static_cast(acc) * x_scale[row] * weight_scale[col]; @@ -500,74 +686,26 @@ __global__ __launch_bounds__(256) void w8a8_gemm_scalar_fallback_kernel( } // --------------------------------------------------------------------------- -// Weight packing (outside the timed region and outside Graph capture; the -// packed buffer keeps the same byte count K*N and the same allocated -// address, so the layout is graph-stable). For the target (k, n) == -// (1024, 4096) the raw [K, N] row-major weight is transposed into the -// n-major [N, K] layout: -// packed[n*K + kk] == raw[kk*N + n] -// so every DUMMA B tile (one 64-column n-tile x one 64-K stage) is a single -// contiguous 4 KiB stream of 16-byte vector loads and each lane's eight -// B-fragment k-values are contiguous in LDS (one 8-byte fragment read). -// Iteration 11 repair: the n-major transpose's 16-byte destination vectors -// run along K, so each thread GATHERS its 16 bytes from 16 raw rows (byte -// loads strided by N) instead of copying a contiguous raw vector (the -// iteration-10 pack copied 16 consecutive n at fixed k, leaving 93.66% of -// the packed buffer unwritten and issuing misaligned int4 stores). The -// pack is one-time and out-of-timed-region, so the gather cost is free. -// For every other (k, n) the pack is an identity device-to-device copy. -// Scales are copied identity in both cases. +// Bootstrap weight packing (outside the timed region and outside Graph +// capture): out-of-Graph device transpose raw [K, N] -> packed [N, K] +// (packed[n][k] = raw[k][n]) for every (k, n), plus an identity scale copy. +// The packed buffer keeps the same byte count (K*N) and the same allocated +// address as raw_weight, so the layout is graph-stable. Every GEMM consumer +// in this file (decode DUMMA, prefill DUMMA, scalar fallback) decodes the +// [N, K] layout; the scale buffer stays n-indexed. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(256) void w8a8_pack_o_proj_panels_kernel( +__global__ __launch_bounds__(256) void w8a8_pack_transpose_kernel( const int8_t* __restrict__ raw, int8_t* __restrict__ packed, int k, - int n) { - // Iteration 11 repair: one thread per 16-byte vector of the packed [N, K] - // layout, but the vector is a GATHER, not a memcpy. The destination vector - // packed[n*K + kk0 .. kk0+15] holds 16 CONSECUTIVE k-values of raw column n - // (packed[n*K + kk0 + j] == raw[(kk0 + j)*N + n], j = 0..15), so its 16 - // source bytes sit in 16 DIFFERENT raw rows (strided by N). The - // iteration-10 pack instead vectorized along the source axis (it copied - // raw[kk][n0..n0+15] - 16 consecutive n at fixed k - into - // packed[n0][kk..kk+15], which must hold 16 consecutive k at fixed n); - // that wrote only the packed rows with n0 % 16 == 0 (93.66% of the buffer - // left unwritten, rows 1..15 of every 16 never touched) and issued - // misaligned int4 stores whenever kk % 16 != 0. Here each thread gathers - // its 16 bytes with aligned byte loads and stores one aligned int4. - const int kk_vectors = k / 16; // 16-k vectors per packed row (k % 16 == 0) - const int64_t total4 = static_cast(n) * kk_vectors; - const int64_t linear = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (linear >= total4) { - return; - } - const int kk0 = static_cast(linear % kk_vectors) * 16; // multiple of 16 - const int nn = static_cast(linear / kk_vectors); // packed row = raw col - uint32_t d[4] = {0, 0, 0, 0}; -#pragma unroll - for (int j = 0; j < 16; ++j) { - const uint8_t b = raw[static_cast(kk0 + j) * n + nn]; - d[j >> 2] |= static_cast(b) << (8 * (j & 3)); - } - int4 v; - v.x = static_cast(d[0]); - v.y = static_cast(d[1]); - v.z = static_cast(d[2]); - v.w = static_cast(d[3]); - // Destination offset nn*K + kk0 is 16-byte aligned (K % 16 == 0 and - // kk0 % 16 == 0), so the int4 store is aligned. - *reinterpret_cast(packed + static_cast(nn) * k + kk0) = v; -} - -__global__ __launch_bounds__(256) void w8a8_pack_identity_kernel( - const int8_t* __restrict__ raw, - int8_t* __restrict__ packed, + int n, int64_t count) { const int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (linear < count) { - packed[linear] = raw[linear]; + const int64_t k_row = linear / n; + const int64_t n_col = linear - k_row * n; + packed[n_col * k + k_row] = raw[linear]; } } @@ -585,9 +723,9 @@ __global__ __launch_bounds__(256) void w8a8_pack_scale_identity_kernel( // --------------------------------------------------------------------------- // Stable host launch symbols consumed by csrc/bindings.cpp. -// Both launchers are pure dispatch: no allocation, no packing, no -// synchronization, no default-stream launch; they run on the caller-provided -// PyTorch stream and are CUDA-Graph safe. +// Both launchers are pure dispatch on the caller-provided HIP stream: no +// allocation, no packing, no synchronization, no default-stream launch; they +// are CUDA/HIP-Graph safe. // --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, @@ -601,29 +739,41 @@ extern "C" void launch_w8a8_gemm( int n, int k, hipStream_t stream) { - (void)workspace; // no split-K: the GEMM does not use the workspace - (void)workspace_bytes; auto* out_bf16 = reinterpret_cast(out); - // Explicit dispatch. The assigned shape (M=4096, N=4096, K=1024) takes the - // DUMMA 128x64 single-buffered path; every other (m, n, k) - including - // small-M API cases (M=2, M=16) and M=3072 with the same (K, N) - takes the - // scalar fallback, which decodes the packed n-major [N, K] layout for - // (k, n) == (1024, 4096) and the identity layout otherwise. - if (m == kTargetM && n == kTargetN && k == kTargetK) { + // Explicit dispatch: + // * assigned decode shape (M=16, N=4096, K=1024, iteration 5): native + // INT8 DUMMA fused split-K=2 - one kernel, 256 blocks (one 16-col N + // tile per block) x 128 threads (wave w = ascending-k half w of the + // tile). Wave 0 publishes its int32 partial to a 1-KiB block-local + // shared plane, one END-of-K __syncthreads, wave 1 sums plane 0 + its + // own upper-K-half fragment (exact int32) and stores scaled bf16. A + // single launch inside the timed Graph at 2.13 blocks/CU; no caller + // workspace and no second combine kernel (workspace left untouched, so + // the iteration-3 workspace-fit fallback no longer exists). + // * assigned large-prefill shape (M=4096, N=4096, K=1024): the DUMMA + // 128x64 single-buffered path (workspace untouched). + // * every other (m, n, k) - including the paired M=2 API shape with the + // same (K, N) - takes the scalar fallback on the packed [N, K] weight + // layout (column n contiguous over k). + if (m == kDecodeM && n == kTargetN && k == kTargetK) { + const dim3 grid(kNumDecodeTiles); // 256 blocks, one tile per block + const dim3 block(kDecodeThreads); // 128 threads = 2 wavefronts + hipLaunchKernelGGL(w8a8_dumma_m16n16_sk2_fused_kernel, grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); + } else if (m == kTargetM && n == kTargetN && k == kTargetK) { const dim3 grid(kTargetN / kBlockN, kTargetM / kBlockM); const dim3 block(kBlockThreads); - hipLaunchKernelGGL(w8a8_dumma_128x64x64_kernel, - grid, block, 0, stream, + hipLaunchKernelGGL(w8a8_dumma_prefill_kernel, grid, block, 0, stream, a, b, x_scale, weight_scale, out_bf16, m, n, k); } else { constexpr int kBlock = 256; const int64_t total = static_cast(m) * n; const dim3 grid(static_cast((total + kBlock - 1) / kBlock)); const dim3 block(kBlock); - hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, grid, block, 0, + stream, a, b, x_scale, weight_scale, out_bf16, m, n, + k); } (void)hipGetLastError(); } @@ -638,29 +788,21 @@ extern "C" void launch_pack_w8a8_weight( hipStream_t stream) { constexpr int kBlock = 256; const dim3 block(kBlock); - if (k == kTargetK && n == kTargetN) { - const int64_t total4 = static_cast(k) * n / 16; - const dim3 pack_grid( - static_cast((total4 + kBlock - 1) / kBlock)); - hipLaunchKernelGGL(w8a8_pack_o_proj_panels_kernel, - pack_grid, block, 0, stream, - raw_weight, packed_weight, k, n); - } else { - const int64_t weight_count = static_cast(k) * n; - const dim3 weight_grid( - static_cast((weight_count + kBlock - 1) / kBlock)); - hipLaunchKernelGGL(w8a8_pack_identity_kernel, - weight_grid, block, 0, stream, - raw_weight, packed_weight, weight_count); - } + + // Out-of-timed device transpose of the raw [K, N] weight into the packed + // [N, K] buffer (same byte count, same address: graph-stable). All GEMM + // consumers in this file decode [N, K] (see launch_w8a8_gemm). + const int64_t weight_count = static_cast(k) * n; + const dim3 weight_grid( + static_cast((weight_count + kBlock - 1) / kBlock)); + hipLaunchKernelGGL(w8a8_pack_transpose_kernel, weight_grid, block, 0, + stream, raw_weight, packed_weight, k, n, weight_count); const int64_t scale_count = n; const dim3 scale_grid( static_cast((scale_count + kBlock - 1) / kBlock)); - hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, - scale_grid, block, 0, stream, - weight_scale, packed_weight_scale, n); + hipLaunchKernelGGL(w8a8_pack_scale_identity_kernel, scale_grid, block, 0, + stream, weight_scale, packed_weight_scale, n); (void)hipGetLastError(); } // @@end -// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_down_proj.hip index 6c5d9375..f02ed9ff 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_down_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M4096/shared_down_proj.hip @@ -1,65 +1,72 @@ -// @@variant shape=hy3_tp8_shared_down_proj_m4096 commit=426ea470e7ee02bcdbe1428e583651b7e518219e added=2026-08-27 -// median_us=171.2 p90_us=171.5 speedup=62.87 baseline_us=1.076e+04 -// source=hy3-dsh-tp8-m4096-1-0ebb994d -// MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). -// worker_3 iteration 1 (architecture round), assigned shape: -// hy3_tp8_shared_down_proj_m4096 : (M, N, K) = (4096, 4096, 192) +// @@variant shape=hy3_tp8_shared_down_proj_m4096 commit=303c7ee331a658a2b081df0b08f962d51259464c added=2026-09-06 +// median_us=89.46 p90_us=89.63 speedup=120.3 baseline_us=1.076e+04 +// source=hy3-tp8-dsh-m16-and-m4096-2-f7fa9503 +// MetaInfer W8A8 INT8 GEMM backend for Hygon K500SM_AI / gfx928 (worker_3). // -// This file provides the two stable host launch symbols consumed by -// csrc/bindings.cpp: -// * launch_w8a8_gemm(...) - timed, graph-safe GEMM on the caller's -// HIP stream; no allocation, packing, -// autotuning, or synchronization. -// * launch_pack_w8a8_weight(...) - out-of-timed-region packing. For the -// exact assigned pair (k, n) == (192, -// 4096) this round emits an n-major -// transpose pack (packed[n][k], one -// contiguous K run per output column); -// every other (k, n) pair keeps the -// identity device-to-device copy so the -// generic fallback and paired small-M -// shapes stay on the raw [K, N] layout. +// Assigned shapes (physical GPU 3): +// hy3_tp8_shared_down_proj_m16 M=16, N=4096, K=192 +// hy3_tp8_shared_down_proj_m4096 M=4096, N=4096, K=192 // -// Round 1 strategy (DUMMA throughput baseline, 2-D macro-tile family): -// * Native INT8 DUMMA m16n16k32 tiled kernel with int32 accumulation. -// * 2-D macro-tile family (64x64, 64x128, 128x64); the active tile is -// selected at compile time (kActiveTileIndex = 1 -> 64x128 (round 4 -// aspect-ratio flip; rounds 1-3 ran 128x64), 512 threads -// = 8 wavefronts, one 32x32 quadrant per wave with four 16x16 int32 -// accumulators kept resident over the whole K loop). All three tiles -// share the templated double-buffered staging pipeline below. -// * K = 192 is pipelined in kPrefillStageK = 64 stages (3 stages) into a -// double-buffered LDS pair with software prefetch: the global int4 -// loads of stage s+1 are issued before the MMAC sequence of stage s and -// the alternate LDS buffer is written right after it, so the vmcnt(0) -// wait and ds_write overlap compute. One barrier per stage. -// * A stays row-major in LDS (stride 88, 8-byte-aligned); B is packed -// n-major at pack time (packed[n][k]) and staged n-major into LDS -// (b_tile[n][k], row stride 72, K-contiguous) with two 8-byte -// ds_write_b64 per thread, consumed through col_major matrix_b -// fragments; both operand paths use the round-8 explicit load_frag8 -// 8-byte ds_read_b64 reads (the bootstrap's row_major matrix_b path was -// byte-granular: 16 ds_read_u8 + ~40 VALU of byte assembly per wave per -// k-step, the measured LDS-issue bottleneck of the 634 us bootstrap). -// * kk body fully unrolled into two independent fragment register sets -// (kStageK == 2 * kTileK) so the second step's LDS latency overlaps the -// first step's swizzle + MMACs. -// * Fused direct fragment -> x_scale -> weight_scale -> bf16 epilogue, -// using the verified gfx928 accumulator ownership (row = lane & 15, -// col_mod4 = lane >> 4, x[i] -> columns col_mod4 + 4*i). The int32 dot -// is exact (max |dot| = 192*127*127 << 2^31), the int32 -> float -// conversion is exact, and the single fp32 multiply followed by one -// bfloat16 rounding matches the CPU int64 reference bit-for-bit. The -// col_major loader places the same B[k][n] values in the same fragment -// registers as row_major, so int32 accumulation is bit-identical to the -// bootstrap path. -// * All other (m, n, k) - including every M < 128 API case and the paired -// M=2 shape with the same (N, K) - go to a scalar int8/int32 fallback -// that decodes the packed layout only for (k, n) == (192, 4096) and the -// raw [K, N] layout for every other pair. +// Logical operator (see int8_w8a8_gemm_api.py): +// out[m, n] = bf16( int32_dot(x_q[m, :], weight[:, n]) +// * x_scale[m, 0] * weight_scale[n, 0] ) // -// Header order is fixed by the control plane: hip_runtime first (du_mma.h -// is not self-contained before it), hip_bfloat16 second, du_mma.h last. +// Pack layout: identity for every weight except the assigned TP8 down-proj +// weight ((k, n) == (192, 4096)), which is packed n-major +// packed[n*k + kk] == raw[kk*n + n_local] (each logical [K, N] column +// becomes a contiguous 192-byte run; same byte count, same buffer). +// launch_pack_w8a8_weight performs the transpose once, out of the timed +// region and Graph capture. Every consumer of the (k, n) == (192, 4096) +// packed buffer (the M=16 kernel below, the M >= 128 tiled DUMMA path, and +// the scalar fallback including the paired M=2 shape) decodes the matching +// n-major layout; all other (k, n) keep the identity pack. +// +// Strategy: +// * exact assigned shape (M=16, N=4096, K=192) -> native INT8 DUMMA +// m16n16k32 kernel with a bounded whole-window LDS prefetch: one +// 64-lane wavefront per block, one 16x16 output tile per block, the wave +// prefetches the full K=192 A/B operand window into LDS once with six +// 16-byte cooperative vector loads issued back to back before any LDS +// store (one __syncthreads; the prologue then pays one in-flight memory +// round trip per block instead of six serialized load->wait->write +// chains), then the 6 m16n16k32 steps consume +// A fragments row_major from a_tile and B fragments col_major from the +// packed n-major stage b_tile via du_load_matrix_sync (no vmem inside +// the K loop; every per-lane fragment byte run is one contiguous 8-byte +// LDS read), explicit int32 accumulation, direct fragment -> +// ((acc * x_scale) * weight_scale) -> bf16 epilogue, no barriers inside +// the K loop. +// * down-proj family (M >= 128, K == 192, N % kActiveBlockN == 0) -> the +// 2-D macro-tile DUMMA throughput kernel: a templated staged pipeline +// over a (kBlockM x kBlockN) output tile with the three macro-tiles of +// the family (64x64, 64x128, 128x64) compiled and a compile-time active +// selection (kActiveTileIndex, see below); the tile-1 64x128 path also +// carries a compile-time pipeline selector (kActivePipeline) between the +// double-buffered form (w8a8_dumma_prefill_kernel) and the round-3 +// single-buffered control (w8a8_dumma_prefill_sb_kernel), both staying +// compiled. K = 192 +// is pipelined in kPrefillStageK = 64 chunks (3 stages): the global +// int4 loads of stage s+1 issue before the MMAC sequence of stage s and +// land in the alternate LDS buffer right after it (one barrier per +// stage), so the vmcnt(0) wait and the ds_write overlap compute. The +// packed n-major weight ((k,n)==(192,4096), P[n*K+kk]) is staged +// n-major into LDS (b_tile[n][k], padded 72-byte rows) and consumed +// through col_major matrix_b fragments; A is staged row-major with +// padded 88-byte rows and consumed row_major. Both operand fragment +// paths use explicit 8-byte LDS reads (load_frag8 -> ds_read_b64, the +// 4-phase conflict-free floor at these even dword strides) instead of +// the library ds_read2_b32 lowering (2-way conflict in every phase at +// ANY even dword row stride). int32 accumulation stays resident over +// the whole K loop in four 16x16 accumulator fragments per wave and the +// epilogue is a direct fragment -> ((acc * x_scale) * weight_scale) -> +// bf16 store. No split-K, no raw asm, no speculative pipeline. +// * Every other shape (including the paired M=2 API shapes with the same +// (N, K)) -> generic scalar int8/int32 fallback. The M=16 kernel keeps +// an exact (m, n, k) == (16, 4096, 192) guard so paired M=2 API shapes +// and unmatched shapes still reach the scalar fallback. +// +// Headers are included in the mandated order for this DTK: HIP runtime first +// (du_mma.h is not self-contained), then hip_bfloat16.h, then du_mma.h. #include #include @@ -69,131 +76,165 @@ namespace { -constexpr int kTileM = 16; -constexpr int kTileN = 16; -constexpr int kTileK = 32; -constexpr int kWaveSize = 64; - -// Staging stage length for the prefill pipeline. 64 = 2x the DUMMA K unit: -// with the 128x64 macro-tile the double-buffered footprint is -// 2*(128*88 + 64*72) = 31,744 B; with the round-4 active 64x128 tile it is -// 2*(64*88 + 128*72) = 29,696 B (A stride 88, transposed-B stride 72; both -// footprints leave 2 blocks/CU LDS-fit at 512 threads (8 waves per block, -// 16 resident wavefronts): 59,392 B <= 64 KiB for 64x128, the same 29,696 B -// per block that round 1 ran resident at). Stage 64 keeps -// the barrier -// count low for the short K = 192 (3 stages -> 4 barriers per block -// including the prologue) while keeping one full 64-k stage of global loads -// in flight. Padded row strides stay aligned: kAStride 88 % 8 == 0 for the -// two-half int64 A staging, kBTStride 72 % 8 == 0 for the col_major 8-byte -// fragment reads and the two-half ds_write_b64 B staging. The A pad is 24 -// (stride 22 dwords): the row_major m16k32 fragment maps lane l to row -// l&15, k-quarter l/16 (8 contiguous bytes per lane), so the 16 fragment -// rows read banks (22*r mod 32) - 16 distinct even bank-pairs - and every -// A ds_read2 hits the 4-phase 32-bank floor with zero conflicts. Stride 20 -// dwords (pad 16) aliases rows r and r+8 onto identical bank pairs -// (20*r mod 32 has period 8), doubling each A fragment read to 8 phases; -// stride 22 is the smallest even dword stride >= 16 whose 16 multiples are -// pairwise distinct mod 32. -constexpr int kPrefillStageK = 64; - -// The only assigned shape for this worker has K = 192, so the optimized -// launch guard is shape-exact on this value (k == kAssignedK). 192 = 3 -// stage-64 chunks. -constexpr int kAssignedK = 192; -constexpr int kAssignedN = 4096; - -// Compile-time selection of the active 2-D macro-tile: +constexpr int kInt8TileM = 16; // DUMMA m16n16k32 (int8 -> int32) tiles +constexpr int kInt8TileN = 16; +constexpr int kInt8TileK = 32; + +constexpr int kWaveSize = 64; // gfx928 native wavefront size + +// Large-M tiled kernel geometry (K=192 = 3 x kPrefillStageK stages). +// +// Compile-time selection of the active 2-D macro-tile of the prefill +// family: // 0 -> 64x64, 1 -> 64x128, 2 -> 128x64. -// All three instantiations are compiled so the family stays measurable; only -// the active one is launched inside the timed region. Round 4 (tile-shape -// round) flips the active tile from 128x64 to 64x128: the first constant- -// occupancy aspect-ratio A/B in this session (same 512 threads / 8 waves / -// 2 blocks/CU / grid 2048; only the (M,N) quadrant partition and grid -// geometry (N/128, M/64) swap). The flip halves per-A-tile global re-reads -// (64 -> 32) and doubles B re-reads (32 -> 64), both 768 KiB L2-resident; -// the active footprint is 29,696 B/block (59,392 B/CU at 2 blocks <= 64 KiB, -// the same per-block footprint round 1 ran resident at). -constexpr int kActiveTileIndex = 1; +// All three instantiations of the templated kernel are compiled so the +// family stays measurable; only the active one is launched inside the timed +// region. The active tile is 64x128: with the (kBlockM x kBlockN) tile, +// 8 wavefronts (512 threads) and the double-buffered stage-64 pipeline the +// per-block LDS footprints are +// 64x64 : 2*(64*88 + 64*72) = 20,480 B (256 threads / 4 waves) +// 64x128 : 2*(64*88 + 128*72) = 29,696 B (512 threads / 8 waves) +// 128x64 : 2*(128*88 + 64*72) = 31,744 B (512 threads / 8 waves) +// and 2 blocks/CU stay LDS-resident for all three (40,960 / 59,392 / 63,488 +// B per CU <= 64 KiB), i.e. 16 resident wavefronts per CU for the 512-thread +// tiles. For the assigned M=4096 x N=4096 grid the 64x128 tile yields +// grid (32, 64) = 2048 blocks and re-reads each A 64-row tile across 32 +// N-blocks and each B 128-column tile across 64 M-blocks (both operand +// streams stay L2-resident). +constexpr int kPrefillStageK = 64; // K depth staged per barrier pair +constexpr int kAssignedK = 192; // the only K the optimized guard allows +constexpr int kActiveTileIndex = 1; // 0 -> 64x64, 1 -> 64x128, 2 -> 128x64 constexpr int kActiveBlockM = kActiveTileIndex == 2 ? 128 : 64; constexpr int kActiveBlockN = kActiveTileIndex == 1 ? 128 : 64; -// LDS padding in bytes added to each staged A row. 24 keeps every row start -// 8-byte aligned for the two-half int64 staging while making the 16 -// m16k32-fragment row bases pairwise distinct mod the 32 LDS banks -// (stride 22 dwords: 22*r mod 32 is distinct for r = 0..15), so the A -// fragment ds_read2 instructions run at the conflict-free 4-phase floor. -// Strides of 20 dwords (pad 16) alias rows r and r+8 onto identical bank -// pairs (8 two-way conflicts per 16-lane group -> 8 phases per read); -// strides >= 16 that are multiples of 128 B alias every row onto one bank -// phase (up to 8-16-way conflicts). -constexpr int kLdsPad = 24; - -// LDS row stride of the transposed (n-major) B tile: each row holds one -// output column n with kStageK contiguous K bytes. kLdsPadB = 8 keeps every -// row start 8-byte aligned so the col_major matrix_b fragment loader can -// lower its 8 contiguous bytes per lane to a single vectorized ds_read2_b32 -// and the B staging can land each n row's k-contiguous 16-byte run as two -// 8-byte ds_write_b64 (72 % 8 == 0; 72 % 16 == 8, so a single ds_write_b128 -// would be 16-byte misaligned for odd n rows - the explicit two-half store -// keeps the lowering legal). +// Pipeline form of the tile-1 (64x128) prefill path, compared in round 3: +// 0 -> single-buffered control (w8a8_dumma_prefill_sb_kernel, ACTIVE this +// round): LDS 14,848 B/block, prologue + 2 barriers/stage boundary = +// 5 barriers/block for K=192, one-stage-ahead VGPR prefetch kept, +// occupancy unchanged at 2 blocks/CU (VGPR-bound); +// 1 -> accepted double-buffered pipeline (w8a8_dumma_prefill_kernel, +// retained and still compiled): LDS 29,696 B/block, one barrier per +// stage (4/block incl. its dead last-stage barrier). +// Tiles 0 and 2 always use the double-buffered form (compile-time retention +// instantiations; only tile 1 is launched for the assigned M=4096 grid). +constexpr int kActivePipeline = 0; + +// LDS row padding of the staged A tile (row-major, one row per output row). +// 24 keeps every row start 8-byte aligned for the two-half int64 staging +// while making the 16 m16k32-fragment row bases pairwise distinct mod the +// 32 LDS banks: row stride kAStride = 88 B = 22 dwords, and {22*r mod 32} +// for r = 0..15 is a permutation of the 16 even bank-pairs, so an 8-byte +// ds_read_b64 fragment read runs at the conflict-free 4-phase floor. +// Stride 20 dwords (pad 16) aliases rows r and r+8 onto identical bank +// pairs at any even dword stride, which is exactly why the ds_read2_b32 +// lowering doubles every phase. +constexpr int kLdsPadA = 24; +constexpr int kAStride = kPrefillStageK + kLdsPadA; // 88, % 8 == 0 + +// LDS row padding of the transposed (n-major) B tile: one row per output +// column, kPrefillStageK contiguous K bytes + 8 B pad -> stride 72 +// (% 8 == 0 for the col_major 8-byte fragment reads and the two-half +// ds_write_b64 B staging; % 16 == 8, so the staging store is the explicit +// two-int64 form, never a 16-byte-misaligned int4 store). {18*r mod 32} +// for r = 0..15 is likewise a permutation of the 16 even bank-pairs. constexpr int kLdsPadB = 8; -constexpr int kBTStride = kPrefillStageK + kLdsPadB; +constexpr int kBTStride = kPrefillStageK + kLdsPadB; // 72, % 8 == 0 + +constexpr int kScalarThreads = 256; // scalar fallback threads per block +constexpr int kPackThreads = 256; // pack copy threads per block using namespace du::dumma; -// gfx928 int8 m16n16k32 accumulator ownership, established against -// du_store_matrix_sync: lane % 16 selects the row, lane / 16 selects the -// column mod 4, and x[i] selects columns separated by four. +// float32 -> bfloat16 with round-to-nearest-even, stored as raw uint16 bits. +// The output tensor is at::kBFloat16 (2 bytes/element); this conversion is +// bit-identical to the hardware RNE cvt used by the reference path for all +// finite inputs and keeps the store independent of bf16 header type details. +__device__ __forceinline__ uint16_t bf16_bits_from_f32(float f) { + uint32_t bits; + __builtin_memcpy(&bits, &f, sizeof(bits)); + const uint32_t rounding_bias = 0x7FFFu + ((bits >> 16) & 1u); + return static_cast((bits + rounding_bias) >> 16); +} + +// Store one 16x16 int32 accumulator fragment scaled to bf16 with the three +// tail-gated weight_scale operands (wv1..wv3) arriving as already-loaded +// registers. +// +// Exact-source ISA of the accepted kernel (round-6 code object, +// compile_cache_key 1df09b96..., bundle +// profiles/.../iteration7/current-best-isa/isa.txt): the compiler already +// hoists x_scale[row] into the staging prologue (global_load_dword at 0x4710, +// before the first s_waitcnt vmcnt) and issues wv0 immediately after the last +// ds_write_b128 (0x4780), but wv1..wv3 are only issued between v_mmac 1..3 +// (0x4844 / 0x484C / 0x4888) and are still in flight when the epilogue +// consumes them: real s_waitcnt vmcnt waits (0x4ACC vmcnt(4), 0x4AD4 vmcnt(3), +// 0x4B20 / 0x4B48 / 0x4B6C vmcnt(3)) sit between the four +// global_store_short_d16_hi, i.e. the per-block store tail can serialize on +// ~L2-latency arrivals of per-lane weight_scale dwords issued only ~3 v_mmacs +// earlier. The exact-shape kernel below therefore loads wv1..wv3 in its +// staging batch (same three vmem read instructions as the generic epilogue, +// moved to the front of the wave's memory FIFO so they land during the +// staging round trip plus the six-step K chain) and passes them into this +// helper, which converts, scales and stores each element purely from +// registers. Math stays byte-identical to the generic fragment stores +// (store_prefill_fragment / the reference formula): row = lane & 15 is +// always < m == 16, and each element keeps the same left-associative +// ((float)acc_i * xv) * wv_i fp32 evaluation order with the same RNE bf16 +// rounding; xv and wv0 keep their proven round-6 issue points (loaded here, +// where the round-6 compiler scheduled them), so the delta is exactly the +// three loads whose issue the round-6 code object placed inside the mmac +// chain. template -__device__ __forceinline__ void store_prefill_fragment( +__device__ __forceinline__ void store_acc_fragment_tail_preloaded( const AccFragment& frag, const float* __restrict__ x_scale, const float* __restrict__ weight_scale, - __hip_bfloat16* __restrict__ out, - int m, + float wv1, + float wv2, + float wv3, + uint16_t* __restrict__ out, int n, - int base_row, int base_col, int lane) { - const int row = base_row + (lane & 15); - if (row >= m) { - return; - } + const int row = lane & 15; // m == 16 guard: always in range const int col_mod4 = lane >> 4; - const float xs = x_scale[row]; -#pragma unroll - for (int i = 0; i < 4; ++i) { - const int col = base_col + col_mod4 + 4 * i; - const float scaled = - static_cast(frag.x[i]) * xs * weight_scale[col]; - out[row * n + col] = __float2bfloat16(scaled); - } + const float xv = x_scale[row]; + const float wv0 = weight_scale[base_col + col_mod4]; + const int64_t out_base = + static_cast(row) * n + base_col + col_mod4; + // Elements i = 0..3 sit at columns base_col + col_mod4 + 4*i; each keeps the + // same left-associative ((float)acc_i * xv) * wv_i evaluation order and RNE + // bf16 rounding as the generic fragment stores. + const float s0 = static_cast(frag.x[0]) * xv; + out[out_base] = bf16_bits_from_f32(s0 * wv0); + const float s1 = static_cast(frag.x[1]) * xv; + out[out_base + 4] = bf16_bits_from_f32(s1 * wv1); + const float s2 = static_cast(frag.x[2]) * xv; + out[out_base + 8] = bf16_bits_from_f32(s2 * wv2); + const float s3 = static_cast(frag.x[3]) * xv; + out[out_base + 12] = bf16_bits_from_f32(s3 * wv3); } // --------------------------------------------------------------------------- -// Round 8 (LDS read-width round): explicit 8-byte fragment loader replacing -// the library du_load_matrix_sync for the int8 m16k32 fragments. The library -// row_major matrix_a and col_major matrix_b loaders both fetch the same -// per-lane bytes - lane l -> row (l & 15), k-quarter (l >> 4), eight -// contiguous bytes at p[row*ldm + (l>>4)*8 .. +7] - and the compiler lowers -// them to ds_read2_b32 (two 4-byte dword phases per half-wave). A -// ds_read2_b32 phase groups 32 lanes (rows 0-15 of k-quarter q with rows -// 0-15 of k-quarter q+1) into one 32-bank cycle, so the kq and kq+1 row sets -// share every phase: {22*r mod 32} and {22*r+2 mod 32} (A stride 22 dwords; -// 18*r and 18*r+2 for B stride 18) are the same 16 even banks, giving a -// 2-way conflict in every phase - 8 phases per read at ANY even dword row -// stride. This is why the exact-source PMC lds_bank_conflicts (2,162,688) -// is byte-identical between the round-5 stride-80 and round-6 stride-88 -// sources. Loading the same 8 bytes as one int64 forces the 8-byte -// ds_read_b64 form, whose phases are 16-lane groups (16 lanes x 8 B = 128 B -// = one bank cycle): within each group the 16 rows land on 16 distinct -// bank-pairs (A: 22*r mod 32 distinct for r = 0..15, B: 18*r mod 32 -// distinct), so every fragment read runs at the 4-phase conflict-free floor. -// The element-to-slot mapping is unchanged (f.x[0..7] = the same 8 bytes in -// the same little-endian order), the MMACs and the exact int32 k-ascending -// accumulation sequence are untouched, and the staging stores, LDS strides, -// barriers, tile, occupancy and epilogue are byte-identical. +// Explicit 8-byte fragment loader for the int8 m16k32 operand fragments. +// +// Both library loaders (row_major matrix_a, col_major matrix_b) place the +// same 8 contiguous bytes per lane into the same x[0..7] slots - lane l -> +// row (l & 15), k-quarter (l >> 4) - and the compiler lowers them to +// ds_read2_b32 (two 4-byte dword phases per half-wave). A ds_read2_b32 +// phase groups 32 lanes (rows 0-15 of one k-quarter with rows 0-15 of the +// other) into one 32-bank cycle, so the two k-quarter row sets share every +// phase: {22*r mod 32} and {22*r+2 mod 32} (A stride 22 dwords; 18*r and +// 18*r+2 for B stride 18) are the same 16 even banks - a 2-way conflict in +// every phase (8 phases per read) at ANY even dword row stride. Loading +// the same 8 bytes as one int64 forces the 8-byte ds_read_b64 form, whose +// phases are 16-lane groups (16 lanes x 8 B = 128 B = one bank cycle): +// within each group the 16 rows land on 16 distinct bank-pairs (A: +// 22*r mod 32 distinct for r = 0..15, B: 18*r mod 32 distinct), so every +// fragment read runs at the 4-phase conflict-free floor. The element-to- +// slot mapping is unchanged, so fragment register contents (and therefore +// the exact int32 k-ascending MMAC accumulation) are identical to the +// library loaders. // --------------------------------------------------------------------------- template __device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, @@ -205,54 +246,124 @@ __device__ __forceinline__ void load_frag8(Frag& f, const int8_t* p, reinterpret_cast(f.x)[0] = v; } +// Store one 16x16 int32 accumulator fragment scaled to bf16 (prefill +// epilogue). gfx928 int8 m16n16k32 accumulator ownership (validated +// against du_store_matrix_sync in this DTK): lane&15 selects the row, +// lane>>4 selects col%4, and frag.x[i] holds the columns col%4 + 4*i. +// Per element the fp32 scaling keeps the same left-associative +// ((float)acc_i * x_scale) * weight_scale order and the same RNE bf16 +// rounding as the reference path and as the M=16 kernel. +template +__device__ __forceinline__ void store_prefill_fragment( + const AccFragment& frag, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + uint16_t* __restrict__ out, + int m, + int n, + int base_row, + int base_col, + int lane) { + const int row = base_row + (lane & 15); + if (row >= m) { + return; // tail-M rows: never touch x_scale/out out of bounds + } + const int col_mod4 = lane >> 4; + const float xs = x_scale[row]; + const int64_t out_base = + static_cast(row) * n + base_col + col_mod4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float scaled = static_cast(frag.x[i]) * xs * + weight_scale[base_col + col_mod4 + 4 * i]; + out[out_base + 4 * i] = bf16_bits_from_f32(scaled); + } +} + +// Round-6 fused-epilogue transpose write, iteration-14 register-batched +// form: scale one 16x16 accumulator fragment with its row scale (xv) and +// its four column scales (wv0..wv3, arriving as already-loaded registers) +// and store the RNE bf16 result into the block's row-major out_tile at the +// element's final (row, col) position. The caller reads the four column +// scales of each 16-column fragment base (tile_col0 and +// tile_col0 + kInt8TileN) ONCE per wave and shares them between the two +// fragments of that base (acc00/acc10 and acc01/acc11), so the per-wave +// epilogue issues 8 ds_read_b32 instead of the round-8 form's 16 in-loop +// reads and no LDS load sits between the last v_mmac and the element store +// tail. Per element the fp32 math keeps the same left-associative +// ((float)acc_i * xv) * wv_i evaluation order and the same RNE bf16 +// rounding as store_prefill_fragment, and every wv_i is the bit-identical +// float the in-loop form read (same buffer s_wscale, same index +// base + c4 + 4*i, no write to s_wscale between the prologue landing and +// the epilogue), so the bytes this lane publishes are identical; only the +// destination differs from a scattered 2-byte global store (an LDS +// transpose slot that the coalesced sweep below reads back). +template +__device__ __forceinline__ void transpose_bf16_fragment_lane( + const AccFragment& frag, float xv, int tile_row0, int tile_col0, + int r_lane, int c4, float wv0, float wv1, float wv2, float wv3, + uint16_t* __restrict__ out_tile, int out_stride) { + const float s0 = static_cast(frag.x[0]) * xv; + out_tile[(tile_row0 + r_lane) * out_stride + tile_col0 + c4] = + bf16_bits_from_f32(s0 * wv0); + const float s1 = static_cast(frag.x[1]) * xv; + out_tile[(tile_row0 + r_lane) * out_stride + tile_col0 + c4 + 4] = + bf16_bits_from_f32(s1 * wv1); + const float s2 = static_cast(frag.x[2]) * xv; + out_tile[(tile_row0 + r_lane) * out_stride + tile_col0 + c4 + 8] = + bf16_bits_from_f32(s2 * wv2); + const float s3 = static_cast(frag.x[3]) * xv; + out_tile[(tile_row0 + r_lane) * out_stride + tile_col0 + c4 + 12] = + bf16_bits_from_f32(s3 * wv3); +} + // --------------------------------------------------------------------------- -// Large-M prefill path. One block computes a kBlockM x kBlockN output tile -// with (kBlockM/32) x (kBlockN/32) wavefronts (one 32x32 quadrant per wave). -// A[kBlockM, K] and B[K, kBlockN] are cooperatively staged into a -// double-buffered, bank-padded LDS in kStageK chunks with VGPR prefetch: -// the global int4 loads for stage s+1 are issued before the MMAC sequence -// of stage s, and the alternate LDS buffer is written right after it, so -// the vmcnt(0) wait and ds_write overlap compute. One barrier per stage. -// A stays row-major in LDS (stride 88); B is packed n-major by -// launch_pack_w8a8_weight (packed[n][k], one contiguous K run per output -// column) and staged n-major (stride 72, K-contiguous) with two 8-byte -// ds_write_b64 per thread, loaded through the round-8 explicit load_frag8 -// 8-byte ds_read_b64 reads (row_major A and col_major B fragments place the -// same 8 contiguous bytes per lane in the same x[0..7] slots, so both -// operand paths are vectorized 8-byte LDS reads). Each wave keeps four -// 16x16 int32 accumulators resident over the whole K loop, then a fused -// scale/bf16 store. The kk ordering (k0 outer, kk inner step kTileK) and -// the acc00/acc01/acc10/acc11 update order are unchanged from the -// bootstrap, and the loaders place the same B[k][n] values in the same -// fragment registers as row_major, so int32 accumulation is bit-identical. +// Large-M prefill path (down-proj family: M >= 128, K == kAssignedK, N a +// multiple of the active macro-tile N). One block computes a kBlockM x +// kBlockN output tile with (kBlockM/32) x (kBlockN/32) wavefronts (one +// 32x32 quadrant per wave, four 16x16 int32 accumulators kept resident +// over the whole K loop). A and the packed n-major B are cooperatively +// staged into a double-buffered, bank-padded LDS in kPrefillStageK = 64 +// chunks with VGPR prefetch: the global int4 loads of stage s+1 are issued +// before the MMAC sequence of stage s and the alternate LDS buffer is +// written right after it, so the vmcnt(0) wait and the ds_write overlap +// compute (one barrier per stage; prologue + 3 stages = 4 barriers per +// block for K = 192). A stays row-major in LDS (stride kAStride = 88); B +// is n-major at pack time (packed[n][k], one contiguous K run per output +// column) and is staged n-major (stride kBTStride = 72, K-contiguous), +// consumed through col_major matrix_b fragments; both operand paths use +// the explicit 8-byte load_frag8 ds_read_b64 reads above. Each stage is +// consumed as kStageK = 2 x kInt8TileK fully unrolled steps into two +// independent fragment register sets, so the second step's LDS latency +// overlaps the first step's MMACs. The epilogue is the fused direct +// fragment -> ((acc * x_scale) * weight_scale) -> bf16 store. // --------------------------------------------------------------------------- template __global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize, 2) void w8a8_dumma_prefill_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, const float* __restrict__ x_scale, const float* __restrict__ weight_scale, - __hip_bfloat16* __restrict__ out, + uint16_t* __restrict__ out, int m, int n, int k) { constexpr int kWavesN = kBlockN / 32; - constexpr int kAStride = kStageK + kLdsPad; static_assert(kAStride % 8 == 0, "A LDS row stride must stay 8-byte aligned"); static_assert(kBTStride % 8 == 0, "transposed-B LDS row stride must stay 8-byte aligned"); - static_assert(kStageK % kTileK == 0, + static_assert(kStageK % kInt8TileK == 0, "K stage must be a multiple of the DUMMA K unit"); static_assert(kStageK % sizeof(int4) == 0, "K stage must keep int4 staging aligned"); - static_assert(kStageK == 2 * kTileK, - "unrolled kk body assumes kStageK == 2 * kTileK"); + static_assert(kStageK == 2 * kInt8TileK, + "unrolled kk body assumes kStageK == 2 * kInt8TileK"); const int tid = static_cast(threadIdx.x); const int wave = tid / kWaveSize; - const int lane = tid % kWaveSize; + const int lane = tid - wave * kWaveSize; const int wave_row = wave / kWavesN; const int wave_col = wave - wave_row * kWavesN; const int m0 = static_cast(blockIdx.y) * kBlockM; @@ -268,28 +379,30 @@ w8a8_dumma_prefill_kernel( const int kStages = k / kStageK; __shared__ __align__(16) int8_t a_tile[2][kBlockM * kAStride]; - // Transposed (n-major) B tile: element B[k0+kk][n] lives at - // b_tile[n][kk], so col_major fragment loads read 8 contiguous K bytes - // per lane (vectorized ds_read2_b32). The staged runs are written n-major - // (packed[n][k]) as two 8-byte ds_write_b64 per thread. + // Transposed (n-major) B tile: logical B[k0+kk][n0+n_local] lives at + // b_tile[n_local][kk], so col_major fragment loads read 8 contiguous K + // bytes per lane (one ds_read_b64). The staged runs are written from the + // n-major packed weight as two 8-byte halves per thread. __shared__ __align__(16) int8_t b_tile[2][kBlockN * kBTStride]; - DUFragment + DUFragment a_frag0, a_frag1, a_frag0_1, a_frag1_1; - DUFragment + DUFragment b_frag0, b_frag1, b_frag0_1, b_frag1_1; - DUFragment + DUFragment acc00, acc01, acc10, acc11; du_fill_fragment(acc00, 0); du_fill_fragment(acc01, 0); du_fill_fragment(acc10, 0); du_fill_fragment(acc11, 0); - // Prologue: prefetch stage 0 into VGPRs and land it in buffer 0. Every + // Prologue: prefetch stage 0 into VGPRs and land it in buffer 0. Every // source address is 16-byte aligned (k % 16 == 0 by the exact k guard and // n % 16 == 0 by the launch guard); A destination rows are 8-byte aligned // (stride 88) and B destination rows 8-byte aligned (stride 72), so all - // staging stores are the two-half int64 ds_write2_b64 form. + // staging stores are the two-half int64 form. int4 a_pref{0, 0, 0, 0}; int4 b_pref{0, 0, 0, 0}; { @@ -299,21 +412,21 @@ w8a8_dumma_prefill_kernel( const int global_row = m0 + row; if (global_row < m) { a_pref = *reinterpret_cast( - x_q + global_row * k + seg * sizeof(int4)); + a + static_cast(global_row) * k + seg * sizeof(int4)); } } if (tid < kBInt4) { const int n_local = tid / kChunksPerBRow; const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); b_pref = *reinterpret_cast( - weight + (n0 + n_local) * k + kchunk); + b + (static_cast(n0) + n_local) * k + kchunk); } if (tid < kAInt4) { const int row = tid / kAInt4PerRow; const int seg = tid - row * kAInt4PerRow; // Two 8-byte halves: kAStride 88 % 16 == 8 makes odd rows 16-byte // misaligned, so the single-int4 store is illegal; the int64 pair - // lowers to one ds_write2_b64 per thread, exactly like the B staging. + // lowers to one ds_write2_b64 per thread. const int64_t* src64 = reinterpret_cast(&a_pref); int64_t* dst64 = reinterpret_cast( a_tile[0] + row * kAStride + seg * sizeof(int4)); @@ -337,7 +450,7 @@ w8a8_dumma_prefill_kernel( const int buf = s & 1; // Issue the global loads for stage s+1 now so their latency hides - // behind the fragment loads and MMACs below. Out-of-range A rows are + // behind the fragment loads and MMACs below. Out-of-range A rows are // zero-filled (tail-M guard). if (s + 1 < kStages) { const int k1 = k0 + kStageK; @@ -347,7 +460,8 @@ w8a8_dumma_prefill_kernel( const int global_row = m0 + row; if (global_row < m) { a_pref = *reinterpret_cast( - x_q + global_row * k + k1 + seg * sizeof(int4)); + a + static_cast(global_row) * k + k1 + + seg * sizeof(int4)); } else { a_pref = int4{0, 0, 0, 0}; } @@ -356,41 +470,35 @@ w8a8_dumma_prefill_kernel( const int n_local = tid / kChunksPerBRow; const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); b_pref = *reinterpret_cast( - weight + (n0 + n_local) * k + k1 + kchunk); + b + (static_cast(n0) + n_local) * k + k1 + kchunk); } } // Consume stage s from buffer buf: four fragment loads and four MMACs - // per wave per kTileK step. kStageK = 64 = 2 kTileK steps, fully - // unrolled into two independent fragment register sets: all eight + // per wave per kInt8TileK step. kStageK = 64 = 2 kInt8TileK steps, + // fully unrolled into two independent fragment register sets: all eight // 8-byte fragment reads of a stage issue before the MMACs, so the // second step's LDS latency overlaps the first step's MMACs. const int local_row = wave_row * 32; const int local_col = wave_col * 32; - // Round 8: custom load_frag8 replaces du_load_matrix_sync for all eight - // fragment reads. Both library loaders (row_major A and col_major B) - // place the same 8 contiguous bytes per lane into the same x[0..7] - // slots; only the LDS instruction changes, from ds_read2_b32 (2-way - // conflict in every 32-lane dword phase at any even row stride) to the - // 8-byte ds_read_b64 (16-lane phases, distinct bank-pairs per group -> - // 4-phase conflict-free floor). Exact int32 accumulation is unchanged. load_frag8(a_frag0, a_tile[buf] + local_row * kAStride, kAStride); load_frag8(a_frag1, a_tile[buf] + (local_row + 16) * kAStride, kAStride); load_frag8(b_frag0, b_tile[buf] + local_col * kBTStride, kBTStride); load_frag8(b_frag1, b_tile[buf] + (local_col + 16) * kBTStride, kBTStride); - // Second 32-k step (kk = kTileK): independent registers, issued while - // the first step's data is still in flight. - load_frag8(a_frag0_1, a_tile[buf] + local_row * kAStride + kTileK, + // Second 32-k step (kk = kInt8TileK): independent registers, issued + // while the first step's data is still in flight. + load_frag8(a_frag0_1, a_tile[buf] + local_row * kAStride + kInt8TileK, kAStride); load_frag8(a_frag1_1, - a_tile[buf] + (local_row + 16) * kAStride + kTileK, kAStride); - // col_major: element B[k0+kk+k][n] at b_tile[n][kk+k], so the loader's - // p[row*ldm + col + i] reads 8 contiguous bytes per lane. - load_frag8(b_frag0_1, b_tile[buf] + local_col * kBTStride + kTileK, + a_tile[buf] + (local_row + 16) * kAStride + kInt8TileK, + kAStride); + // col_major: element B[k0+kk+k][n0+n] lives at b_tile[n][kk+k], so the + // loader's p[row*ldm + k + i] reads 8 contiguous bytes per lane. + load_frag8(b_frag0_1, b_tile[buf] + local_col * kBTStride + kInt8TileK, kBTStride); load_frag8(b_frag1_1, - b_tile[buf] + (local_col + 16) * kBTStride + kTileK, + b_tile[buf] + (local_col + 16) * kBTStride + kInt8TileK, kBTStride); du_mma_sync(acc00, a_frag0, b_frag0, acc00); du_mma_sync(acc01, a_frag0, b_frag1, acc01); @@ -401,7 +509,7 @@ w8a8_dumma_prefill_kernel( du_mma_sync(acc10, a_frag1_1, b_frag0_1, acc10); du_mma_sync(acc11, a_frag1_1, b_frag1_1, acc11); - // Land the prefetched stage s+1 in the alternate buffer. The compiler + // Land the prefetched stage s+1 in the alternate buffer. The compiler // inserts the vmcnt(0) wait here via the data dependency, so the wait // and the ds_write overlap the MMACs above instead of stalling before // them. @@ -427,121 +535,707 @@ w8a8_dumma_prefill_kernel( } // One barrier per stage: makes the alternate-buffer writes visible to // the next stage's reads and retires this stage's reads before the - // parity-flipped buffer is overwritten two stages later. + // parity-flipped buffer is overwritten two stages later. Every thread + // reaches it. __syncthreads(); } const int base_row = m0 + wave_row * 32; const int base_col = n0 + wave_col * 32; - store_prefill_fragment( - acc00, x_scale, weight_scale, out, m, n, - base_row, base_col, lane); - store_prefill_fragment( - acc01, x_scale, weight_scale, out, m, n, - base_row, base_col + 16, lane); - store_prefill_fragment( - acc10, x_scale, weight_scale, out, m, n, - base_row + 16, base_col, lane); - store_prefill_fragment( - acc11, x_scale, weight_scale, out, m, n, - base_row + 16, base_col + 16, lane); + store_prefill_fragment(acc00, x_scale, weight_scale, out, m, n, base_row, + base_col, lane); + store_prefill_fragment(acc01, x_scale, weight_scale, out, m, n, base_row, + base_col + kInt8TileN, lane); + store_prefill_fragment(acc10, x_scale, weight_scale, out, m, n, + base_row + kInt8TileM, base_col, lane); + store_prefill_fragment(acc11, x_scale, weight_scale, out, m, n, + base_row + kInt8TileM, base_col + kInt8TileN, lane); } // --------------------------------------------------------------------------- -// Generic scalar int8 x int8 -> int32 fallback for unmatched (m, n, k) and -// small-M API cases. One thread per output element, coalesced along N. -// For the exact packed pair (k, n) == (192, 4096) the weight is n-major -// (packed[n][k], column stride k); every other pair keeps the raw logical -// [K, N] row-major layout (column stride n). The int32 summation order is -// identical either way (k-ascending), keeping the fallback bit-identical. +// Round-3 pipeline control of the same down-proj family tile: the SINGLE- +// buffered staged pipeline (one LDS buffer per operand instead of the +// parity pair of w8a8_dumma_prefill_kernel above), kept bit-identical and +// structurally identical otherwise: same (kBlockM x kBlockN) tile and wave +// -> quadrant mapping, same one-stage-ahead int4 VGPR prefetch (stage s+1 +// global loads issued before the stage-s fragment loads/MMACs), same +// 8-byte load_frag8 fragment consumption with the identical k-ascending +// int32 MMAC order. The single buffer is overwritten in place, so every +// stage boundary needs a WAR __syncthreads (retire the stage-s fragment +// reads) before the ds_write of stage s+1 and a publish __syncthreads after +// it: prologue + 2 x (kStages - 1) = 5 barriers per block for K = 192 vs +// the double-buffered prologue + kStages = 4 (of which the last-stage loop +// barrier is dead); the final stage's MMACs run straight into the epilogue. +// The SB epilogue is the round-6 FUSED form (scale windows staged in LDS in +// the prologue, per-fragment RNE bf16 transpose through a padded out_tile, +// one publish barrier, fully-coalesced 8-byte sweep stores) - same per- +// element fp32 math and RNE bf16 bytes as the direct round-3/4 fragment +// stores of the DB arm, only the store form differs. Static LDS grows to +// 32,512 B/block for 64x128 (5,632 A + 9,216 B + 256 s_xscale + 512 +// s_wscale + 16,896 out_tile; see the in-body comment), so residency stays +// 2 blocks/CU = 16 waves/CU (2 x 32,512 = 65,024 <= 65,536; the kernel +// stays VGPR-bound at ~56-64 vgpr x 512 threads x 2 blocks <= 65,536 +// VGPRs), which keeps occupancy from confounding the pipeline comparison. // --------------------------------------------------------------------------- -__global__ __launch_bounds__(256) void w8a8_scalar_fallback_kernel( - const int8_t* __restrict__ x_q, - const int8_t* __restrict__ weight, +template +__global__ __launch_bounds__(kBlockM * kBlockN / 1024 * kWaveSize, 2) void +w8a8_dumma_prefill_sb_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, const float* __restrict__ x_scale, const float* __restrict__ weight_scale, - __hip_bfloat16* __restrict__ out, + uint16_t* __restrict__ out, int m, int n, - int k, - int b_col_stride) { - const int linear = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const int total = m * n; - if (linear >= total) { + int k) { + constexpr int kWavesN = kBlockN / 32; + static_assert(kAStride % 8 == 0, + "A LDS row stride must stay 8-byte aligned"); + static_assert(kBTStride % 8 == 0, + "transposed-B LDS row stride must stay 8-byte aligned"); + static_assert(kStageK % kInt8TileK == 0, + "K stage must be a multiple of the DUMMA K unit"); + static_assert(kStageK % sizeof(int4) == 0, + "K stage must keep int4 staging aligned"); + static_assert(kStageK == 2 * kInt8TileK, + "unrolled kk body assumes kStageK == 2 * kInt8TileK"); + + const int tid = static_cast(threadIdx.x); + const int wave = tid / kWaveSize; + const int lane = tid - wave * kWaveSize; + const int wave_row = wave / kWavesN; + const int wave_col = wave - wave_row * kWavesN; + const int m0 = static_cast(blockIdx.y) * kBlockM; + const int n0 = static_cast(blockIdx.x) * kBlockN; + + constexpr int kThreads = kBlockM * kBlockN / 1024 * kWaveSize; + constexpr int kAInt4PerRow = kStageK / sizeof(int4); + constexpr int kAInt4 = kBlockM * kAInt4PerRow; + constexpr int kChunksPerBRow = kStageK / sizeof(int4); + constexpr int kBInt4 = kBlockN * kChunksPerBRow; + static_assert(kAInt4 <= kThreads && kBInt4 <= kThreads, + "one int4 chunk per thread requires chunks <= blockDim"); + // No runtime stage count: the launch guard fixes k == kAssignedK == 192, + // so the kernel body below is the compile-time 3-stage form (k is still + // read at runtime inside run_stage for the global A/B row stride of the + // stage-prefetch addresses; its value is always 192 under the guard). + + // Single-buffered stage: A rows on the padded 88-B stride and the packed + // n-major B rows on the padded 72-B stride, exactly as the double-buffered + // kernel, minus the parity [2] dimension. + __shared__ __align__(16) int8_t a_tile[kBlockM * kAStride]; + __shared__ __align__(16) int8_t b_tile[kBlockN * kBTStride]; + + // Round-6 fused-epilogue staging: + // * s_xscale/s_wscale - the block's per-row x_scale[kBlockM] and + // per-column weight_scale[kBlockN] windows, copied cooperatively into + // LDS inside the stage-0 staging batch (192 dword loads issued with the + // operand vectors, latency hidden under the prologue barrier and the + // whole K chain). The epilogue then scales and converts every + // accumulator element from LDS-resident scales: no global load and no + // vmcnt wait can sit between the last v_mmac and the final store + // (the round-5 exact-source ISA shows the per-fragment scale dword + // loads issued only ~3 v_mmacs before the store tail, with real + // s_waitcnt vmcnt(n) waits interleaved between the four + // global_store_short_d16_hi groups of each fragment). + // * out_tile - the block's (kBlockM x kBlockN) bf16 output tile in + // row-major layout with row stride kOutStride = 132 elements (264 B, + // % 16 == 8 -> every row start stays 8-byte aligned for the coalesced + // ds_read_b64 sweep, while the +4-element pad spreads the 16 fragment + // rows over 16 distinct LDS bank words so the fragment-layout 2-byte + // transpose writes run at the ~1-phase floor). Each wave converts its + // four fragments to RNE bf16 and transposes them through out_tile; one + // extra barrier later every thread stores one contiguous 32-B run of + // one output row as four 8-byte vector stores (fully coalesced lines, + // ~8x fewer L2 write line-touches than the old per-fragment + // global_store_short_d16_hi tail, which scattered each 64-lane store + // instruction over 16 rows' 8-B segments at the 8,192-B row stride). + // Static LDS stays 32,512 B/block (5,632 A + 9,216 B + 256 + 512 scales + + // 16,896 out), so 2 blocks/CU = 16 waves/CU remain resident + // (2 x 32,512 = 65,024 <= 65,536); occupancy remains VGPR-bound. + __shared__ __align__(16) float s_xscale[kBlockM]; + __shared__ __align__(16) float s_wscale[kBlockN]; + constexpr int kOutStride = kBlockN + 4; // 132 bf16 per row (264 B), % 8 == 0 + static_assert((kOutStride * 2) % 8 == 0, + "out_tile rows must stay 8-byte aligned for the b64 sweep"); + __shared__ __align__(16) uint16_t out_tile[kBlockM * kOutStride]; + + DUFragment + a_frag0, a_frag1, a_frag0_1, a_frag1_1; + DUFragment + b_frag0, b_frag1, b_frag0_1, b_frag1_1; + DUFragment + acc00, acc01, acc10, acc11; + du_fill_fragment(acc00, 0); + du_fill_fragment(acc01, 0); + du_fill_fragment(acc10, 0); + du_fill_fragment(acc11, 0); + + int4 a_pref{0, 0, 0, 0}; + int4 b_pref{0, 0, 0, 0}; + + // Prologue: fetch stage 0 into VGPRs, land it in the single buffer with + // the two-half int64 form (kAStride/kBTStride are 8 mod 16, so a single + // 16-byte store would be misaligned on odd rows), publish with one + // barrier. Identical source bytes and store addresses as stage 0 of the + // double-buffered kernel. + { + // Round-6 fused-epilogue staging: the block's x_scale[kBlockM] and + // weight_scale[kBlockN] windows are loaded here, at the FRONT of the + // prologue batch, so their vmcnt results land while the A/B staging + // stores and the publish barrier below drain (the fused epilogue then + // scales/converts every accumulator element purely from LDS: no global + // load and no vmcnt wait can sit between the last v_mmac and the final + // store, replacing the round-3/4 tail that serialized the four + // store_short_d16_hi groups on per-lane weight_scale dword arrivals). + float xscale_v = 0.0f; + float wscale_v = 0.0f; + if (tid < kBlockM) { + const int g_row = m0 + tid; + if (g_row < m) { + xscale_v = x_scale[g_row]; // tail-M rows keep 0, never used + } + } + if (tid < kBlockN) { + wscale_v = weight_scale[n0 + tid]; + } + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + const int global_row = m0 + row; + if (global_row < m) { + a_pref = *reinterpret_cast( + a + static_cast(global_row) * k + seg * sizeof(int4)); + } + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + b_pref = *reinterpret_cast( + b + (static_cast(n0) + n_local) * k + kchunk); + } + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + const int64_t* src64 = reinterpret_cast(&a_pref); + int64_t* dst64 = reinterpret_cast( + a_tile + row * kAStride + seg * sizeof(int4)); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + const int64_t* src64 = reinterpret_cast(&b_pref); + int64_t* dst64 = reinterpret_cast( + b_tile + n_local * kBTStride + kchunk); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + // Land the staged scale windows (loaded at the front of this batch, + // above) so the prologue publish barrier makes them visible to every + // wave's fused epilogue. + if (tid < kBlockM) { + s_xscale[tid] = xscale_v; + } + if (tid < kBlockN) { + s_wscale[tid] = wscale_v; + } + __syncthreads(); + } + + // Iteration-15 exact-K stage unroll: the launch guard fixes + // k == kAssignedK == 192 == 3 x kPrefillStageK for every launch of this + // kernel, so the three stage passes below are written out at compile time + // for that exact K. Each stage keeps the byte-identical round-3/14 + // structure - stage s+1 global int4 prefetch issued before the stage-s + // fragment loads/MMACs, WAR __syncthreads, in-place landing of the + // prefetched stage, publish __syncthreads at the first two boundaries, + // straight into the epilogue after the last stage - with the exact int32 + // k-ascending MMAC order, operand staging bytes, fragment slots, the 6 + // barriers per block and the a_pref/b_pref VGPR prefetch registers reused + // per stage unchanged, so the published bytes are bit-identical. The + // delta vs the rolled loop is only the per-stage control the round-8 code + // object carries (trip-count increment s19, k offset s22, per-pass + // compare/branch and the loop back-edge): with the count compile-time + // known, all of it disappears and the stage-2 MMACs flow straight into the + // fused epilogue. + // + // Iteration-16 exact-K specialization (terminal step of the round-15 + // unroll): the runtime dispatch below ran run_stage(0..2) for the guarded + // K and kept a rolled fallback loop for any other k. That fallback can + // never execute - the host guard use_tiled_dumma requires + // k == kAssignedK == 192 for every launch of this kernel - yet the timed + // code object still carries its full extra stage copy (8 du_mma_sync, 8 + // ds_read_b64 fragment loads, 2 ds_write2_b64 landings, 2 s_barrier and + // the loop control: v_mmac 32 total vs the 24 the executed path runs, + // 8 s_barrier in .text vs the 6 executed) plus the uniform + // s_cmp_eq_u32/s_cbranch on k == kAssignedK, and the register allocator + // still has to allocate one merged CFG containing the loop's live ranges + // and phi copies. This round removes the dispatch and the dead rolled + // copy: the three compile-time stages below run unconditionally (k is + // still read at runtime for the global A/B row stride k of the stage + // prefetch addresses, always 192 under the guard). The executed path - + // stage bodies, prefetch/land pairs, the 6 barriers, the fused epilogue + // and the sweep - is source-identical, so the published bytes stay + // bit-identical; the code object loses only dead code, dead barriers and + // the k compare, and the allocator's problem loses the fallback merge. + // Round-15's linearization raised the object from 57 to 65 VGPRs (sgpr + // 28 -> 32) at the 64-VGPR 2-blocks/CU register boundary every prior + // guard priced; the falsifiable question of this round is how much of + // that +8 was the never-taken fallback path's register/phi cost (then the + // specialized object returns to <= 64 VGPRs) versus inherent ILP of the + // straight-line three-stage schedule (then it stays at 65 and this round + // is the honest null at parity). + const auto run_stage = [&](int s, bool has_next) { + // Issue the global loads for stage s+1 now, at the same point as the + // double-buffered kernel, so their latency hides behind the fragment + // loads and MMACs below. Out-of-range A rows are zero-filled (tail-M + // guard), byte-identical to the rolled zero fill. + if (has_next) { + const int k1 = (s + 1) * kStageK; + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + const int global_row = m0 + row; + if (global_row < m) { + a_pref = *reinterpret_cast( + a + static_cast(global_row) * k + k1 + + seg * sizeof(int4)); + } else { + a_pref = int4{0, 0, 0, 0}; + } + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + b_pref = *reinterpret_cast( + b + (static_cast(n0) + n_local) * k + k1 + kchunk); + } + } + + // Consume stage s from the single buffer: identical fragment loads and + // MMACs to the rolled loop (same padded strides, same element-to-slot + // mapping, same k-ascending int32 order). The second kInt8TileK step + // uses an independent register set so its LDS latency overlaps the first + // step's MMACs. + const int local_row = wave_row * 32; + const int local_col = wave_col * 32; + load_frag8(a_frag0, a_tile + local_row * kAStride, kAStride); + load_frag8(a_frag1, a_tile + (local_row + 16) * kAStride, kAStride); + load_frag8(b_frag0, b_tile + local_col * kBTStride, kBTStride); + load_frag8(b_frag1, b_tile + (local_col + 16) * kBTStride, kBTStride); + load_frag8(a_frag0_1, a_tile + local_row * kAStride + kInt8TileK, + kAStride); + load_frag8(a_frag1_1, + a_tile + (local_row + 16) * kAStride + kInt8TileK, kAStride); + load_frag8(b_frag0_1, b_tile + local_col * kBTStride + kInt8TileK, + kBTStride); + load_frag8(b_frag1_1, + b_tile + (local_col + 16) * kBTStride + kInt8TileK, + kBTStride); + du_mma_sync(acc00, a_frag0, b_frag0, acc00); + du_mma_sync(acc01, a_frag0, b_frag1, acc01); + du_mma_sync(acc10, a_frag1, b_frag0, acc10); + du_mma_sync(acc11, a_frag1, b_frag1, acc11); + du_mma_sync(acc00, a_frag0_1, b_frag0_1, acc00); + du_mma_sync(acc01, a_frag0_1, b_frag1_1, acc01); + du_mma_sync(acc10, a_frag1_1, b_frag0_1, acc10); + du_mma_sync(acc11, a_frag1_1, b_frag1_1, acc11); + + if (has_next) { + // WAR barrier: every stage-s fragment read in the whole block must + // retire before the single buffer is overwritten in place with stage + // s+1. + __syncthreads(); + // Land the prefetched stage s+1. The compiler inserts the vmcnt(0) + // wait here via the a_pref/b_pref data dependency, i.e. after the + // MMACs and the WAR barrier, never before the stage-s burst. + if (tid < kAInt4) { + const int row = tid / kAInt4PerRow; + const int seg = tid - row * kAInt4PerRow; + const int64_t* src64 = reinterpret_cast(&a_pref); + int64_t* dst64 = reinterpret_cast( + a_tile + row * kAStride + seg * sizeof(int4)); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + if (tid < kBInt4) { + const int n_local = tid / kChunksPerBRow; + const int kchunk = (tid - n_local * kChunksPerBRow) * sizeof(int4); + const int64_t* src64 = reinterpret_cast(&b_pref); + int64_t* dst64 = reinterpret_cast( + b_tile + n_local * kBTStride + kchunk); + dst64[0] = src64[0]; + dst64[1] = src64[1]; + } + // Publish: make the stage s+1 writes visible before the next stage's + // fragment reads. + __syncthreads(); + } + }; + // Iteration-16 exact-K specialization: the launch guard fixes + // k == kAssignedK == 192 == 3 x kPrefillStageK for every launch of this + // kernel, so the three compile-time stages run unconditionally below + // (the round-15 rolled fallback for any other k, which could never + // execute and whose dead stage copy, dead barriers and loop control + // stayed in the timed code object, is removed). + run_stage(0, true); + run_stage(1, true); + run_stage(2, false); + + // Round-6 fused epilogue (SB arm only): every accumulator element is + // scaled with the prologue-staged LDS windows (s_xscale/s_wscale) and + // converted to RNE bf16 with the identical left-associative fp32 math as + // the direct fragment stores of rounds 3/4, but each element is transposed + // into the block's padded row-major out_tile instead of a scattered + // 2-byte global store. One extra barrier later a fully-coalesced sweep + // stores the whole tile: each thread copies one contiguous 32-B run of one + // output row to global as four 8-byte vector stores, so the per-block + // epilogue global-store tail collapses from 256 scattered + // store_short_d16_hi instructions (each 64-lane instruction touching 16 + // rows' 8-B segments at the 8,192-B row stride, with vmcnt waits for + // in-flight per-lane scale dwords between the four groups) into 32 fully + // coalesced 8-byte store instructions with no global-load dependency at + // all. Every published byte is bit-identical to the round-3/4 store form. + const int tile_row0 = wave_row * 32; // n0/m0-relative quadrant origin + const int tile_col0 = wave_col * 32; + const int r_lane = lane & 15; + const int c4 = lane >> 4; + // Iteration-14 register-batched column scales: the four s_wscale dwords + // of each 16-column fragment base are read ONCE per wave, back to back, + // and shared by the two fragments of that base (acc00/acc10 use base + // tile_col0, acc01/acc11 use base tile_col0 + kInt8TileN). The round-8 + // form re-read them inside every fragment call's store loop (16 + // ds_read_b32/wave with one load->use dependency per element); this form + // issues 8 ds_read_b32/wave ahead of the element math. Every read is + // unguarded: the index stays inside s_wscale[kBlockN] for every lane + // (tile_col0 + kInt8TileN + c4 + 12 <= 96 + 16 + 15 <= 127) and tail-M + // blocks simply discard the values under the row guards below, so the + // published bytes stay bit-identical to the round-8 form. + const int g_row0 = m0 + tile_row0 + r_lane; // rows of acc00/acc01 + const int g_row1 = m0 + tile_row0 + kInt8TileM + r_lane; // acc10/acc11 + const float wv00 = s_wscale[tile_col0 + c4]; + const float wv01 = s_wscale[tile_col0 + c4 + 4]; + const float wv02 = s_wscale[tile_col0 + c4 + 8]; + const float wv03 = s_wscale[tile_col0 + c4 + 12]; + const float xv0 = s_xscale[tile_row0 + r_lane]; + if (g_row0 < m) { // row tail guard, as in the round-3/4 fragment stores + transpose_bf16_fragment_lane(acc00, xv0, tile_row0, tile_col0, r_lane, + c4, wv00, wv01, wv02, wv03, out_tile, + kOutStride); + } + const float xv1 = s_xscale[tile_row0 + kInt8TileM + r_lane]; + if (g_row1 < m) { + transpose_bf16_fragment_lane(acc10, xv1, tile_row0 + kInt8TileM, + tile_col0, r_lane, c4, wv00, wv01, wv02, + wv03, out_tile, kOutStride); + } + // The tile_col0 batch is consumed; only now read the second 4-dword batch + // (column base tile_col0 + kInt8TileN, shared by acc11/acc01) so at most + // four scale registers are live at once in the tail. + const float wv10 = s_wscale[tile_col0 + kInt8TileN + c4]; + const float wv11 = s_wscale[tile_col0 + kInt8TileN + c4 + 4]; + const float wv12 = s_wscale[tile_col0 + kInt8TileN + c4 + 8]; + const float wv13 = s_wscale[tile_col0 + kInt8TileN + c4 + 12]; + if (g_row1 < m) { + transpose_bf16_fragment_lane(acc11, xv1, tile_row0 + kInt8TileM, + tile_col0 + kInt8TileN, r_lane, c4, wv10, + wv11, wv12, wv13, out_tile, kOutStride); + } + if (g_row0 < m) { + transpose_bf16_fragment_lane(acc01, xv0, tile_row0, + tile_col0 + kInt8TileN, r_lane, c4, wv10, + wv11, wv12, wv13, out_tile, kOutStride); + } + // Publish: every wave's out_tile transpose writes must be visible to the + // coalesced sweep below (the 6th barrier of the block). + __syncthreads(); + // Coalesced sweep: thread tid stores one contiguous 32-B run (16 bf16 = + // four 8-byte vectors) of one output row. kThreadsPerRow consecutive + // threads cover one row, so every 8-byte store instruction in the sweep + // is fully coalesced. Global row stride n is a multiple of kBlockN by the + // launch guard and c is a multiple of 16 elements, so both the LDS source + // (row stride kOutStride * 2 B, 8-byte aligned by the static_assert above) + // and the global destination stay 8-byte aligned. Tail-M rows are skipped + // (their out_tile slots were never transposed). + { + constexpr int kElemsPerThread = (kBlockM * kBlockN) / kThreads; // 16 + constexpr int kThreadsPerRow = kBlockN / kElemsPerThread; // 8 + static_assert(kBlockN % kElemsPerThread == 0, + "sweep requires an integral number of threads per row"); + static_assert((kElemsPerThread * 2) % 8 == 0, + "sweep copy must stay 8-byte aligned"); + const int row_local = tid / kThreadsPerRow; + const int seg = tid - row_local * kThreadsPerRow; + const int g_row = m0 + row_local; + if (g_row < m) { + const int c = seg * kElemsPerThread; + const int64_t* __restrict__ src64 = + reinterpret_cast(out_tile + + row_local * kOutStride + c); + int64_t* __restrict__ dst64 = reinterpret_cast( + out + static_cast(g_row) * n + n0 + c); +#pragma unroll + for (int j = 0; j < (kElemsPerThread * 2) / 8; ++j) { + dst64[j] = src64[j]; + } + } + } +} + +// Exact assigned shape hy3_tp8_shared_down_proj_m16 (M=16, N=4096, K=192): +// one-wave INT8 DUMMA m16n16k32 kernel with a bounded whole-window LDS +// prefetch of the block's operand tiles and the packed n-major weight. +// +// The down-proj weight ((k,n)==(192,4096)) is packed once, out of the timed +// region, as P[n*K+kk] = W[kk*N+n]: each logical n column becomes a +// contiguous 192-byte run, so this block's private B window (n0..n0+15) is +// one contiguous 3,072-byte global region. One 64-lane wavefront per block +// computes one 16x16 output tile (all 16 rows of A, columns n0..n0+15). +// Before the K loop the wave cooperatively prefetches the whole K=192 window +// into LDS with wide 16-byte vector loads -- A[16,192] (3,072 B; the same +// L2-hot bytes every block reads) into a bank-skewed 16x208 padded +// row-major stage, and the packed B strip (3,072 B, fetched from global +// exactly once, fully coalesced) into a 16x208 padded n-major stage. The +// six 16-byte cooperative loads (A vectors lane/lane+64/lane+128 and the B +// strip vectors lane/lane+64/lane+128) are issued back to back into distinct +// registers before any ds_write, so all six global reads run in flight and +// the wave pays one memory round trip for the whole window instead of the +// six serialized global_load -> s_waitcnt vmcnt(0) -> ds_write_b128 chains +// the round-5 exact-source ISA shows (register reuse forced a full wait per +// chunk) -- then +// the six m16n16k32 steps load their fragments with du_load_matrix_sync / +// du_mma_sync straight from LDS (int32 accumulation). A is consumed +// row_major and B is consumed col_major on its n-major stage, so each +// lane's 8 fragment bytes are one contiguous k-run: every per-lane fragment +// read is a single 8-byte LDS read (ds_read_b64) with no byte reassembly, +// the pattern the validated TP4 down-proj lineage (same N=4096 grid=256 +// K-half=192 geometry, 208-byte strides) measured. No vmem instruction or +// global-latency wait remains inside the K loop; the only __syncthreads is +// the single one after the staging stores (one wavefront, none in the +// loop), and the epilogue goes straight from the verified gfx928 +// accumulator fragment ownership to the scaled bf16 store. P is a pure +// transpose of the identity pack (P[n*K+kk] == W[kk*N+n]), so every +// fragment byte equals the byte the round-4 k-major stage delivered for the +// same (kk, n) and the k-ascending int32 accumulation order is unchanged. +__global__ __launch_bounds__(kWaveSize) void +w8a8_gemm_dumma_m16n16k32_packedb_lds_kernel( + const int8_t* __restrict__ a, // [16, K] row-major + const int8_t* __restrict__ b, // packed [N, K] n-major down-proj + const float* __restrict__ x_scale, // [16] + const float* __restrict__ weight_scale, // [N] + uint16_t* __restrict__ out, // [16, N] bf16 bits + int n, + int k) { + const int lane = static_cast(threadIdx.x); // wavefront lane 0..63 + const int n0 = static_cast(blockIdx.x) * kInt8TileN; + (void)k; // exact-shape guard fixes k == 192 == kLdsK (compile-time bound) + + // Whole-window LDS stage (guard guarantees k == 192). Both strides are + // 16-byte multiples so every staging store and fragment load stays + // 16-byte aligned; 208 = 192 + 16 is the validated TP4 down_proj skew, + // used here for the A rows and for the B n-columns alike. + constexpr int kLdsK = 192; // K of the guarded shape + constexpr int kLdsAStride = 208; // A row stride in bytes (192 + 16 pad) + constexpr int kLdsBStride = 208; // B n-column stride in bytes (192 + 16 pad) + __shared__ __align__(16) int8_t a_tile[kInt8TileM * kLdsAStride]; + __shared__ __align__(16) int8_t b_tile[kInt8TileN * kLdsBStride]; + + // A: all 16 rows x 192 bytes are contiguous in global memory (3,072 B = + // 192 int4 vectors) and B: this block's private strip is the 16 packed n + // rows n0..n0+15, one contiguous 3,072 B region (P row n is 192 contiguous + // k bytes). In both windows lane l owns vectors l, l+64 and l+128, i.e. + // six fully coalesced 16-byte loads per lane cover the whole prologue. + // All six loads are issued into distinct registers before any LDS store: + // the round-5 exact-source ISA shows the previous per-vector loops + // degraded into six serialized global_load -> s_waitcnt vmcnt(0) -> + // ds_write_b128 round trips per block (three for A, three for B), each + // paying the full L2/HBM latency on the block's critical path. Batching + // keeps every source byte and every padded tile slot identical while + // collapsing the prologue to one in-flight memory round trip. + const int4 a_v0 = *reinterpret_cast(a + lane * 16); + const int4 a_v1 = + *reinterpret_cast(a + (lane + kWaveSize) * 16); + const int4 a_v2 = + *reinterpret_cast(a + (lane + 2 * kWaveSize) * 16); + const int8_t* b_strip = b + static_cast(n0) * kLdsK; + const int4 b_v0 = *reinterpret_cast(b_strip + lane * 16); + const int4 b_v1 = + *reinterpret_cast(b_strip + (lane + kWaveSize) * 16); + const int4 b_v2 = + *reinterpret_cast(b_strip + (lane + 2 * kWaveSize) * 16); + // wv1..wv3 scale preload, issued inside the staging batch right behind the + // six operand vectors (same three vmem read instructions per block as the + // generic epilogue performs; only the issue point moves to the front of the + // wave's vmem FIFO). col_mod4 = lane >> 4; M == 16 fills the whole tile, so + // these are the exact columns the epilogue scales, and every index stays + // in range (n0 + col_mod4 + 12 <= 4095 < N). See + // store_acc_fragment_tail_preloaded for the round-6 ISA evidence that + // wv1..wv3 were issued between v_mmac 1..3 and waited on between the four + // bf16 stores. + const int col_mod4 = lane >> 4; + const float wv1 = weight_scale[n0 + col_mod4 + 4]; + const float wv2 = weight_scale[n0 + col_mod4 + 8]; + const float wv3 = weight_scale[n0 + col_mod4 + 12]; + constexpr int kChunksPerRow = kLdsK / 16; // 12 x 16-B chunks per 192-B row + // Chunk vec -> padded-tile int4 slot (row = vec / 12, k16 = (vec % 12)*16); + // both tiles use the same 208-byte skew so the mapping is shared. + const auto stage_slot = [](int vec) -> int { + const int row = vec / kChunksPerRow; + return (row * kLdsAStride + (vec - row * kChunksPerRow) * 16) / 16; + }; + reinterpret_cast(a_tile)[stage_slot(lane)] = a_v0; + reinterpret_cast(a_tile)[stage_slot(lane + kWaveSize)] = a_v1; + reinterpret_cast(a_tile)[stage_slot(lane + 2 * kWaveSize)] = a_v2; + reinterpret_cast(b_tile)[stage_slot(lane)] = b_v0; + reinterpret_cast(b_tile)[stage_slot(lane + kWaveSize)] = b_v1; + reinterpret_cast(b_tile)[stage_slot(lane + 2 * kWaveSize)] = b_v2; + __syncthreads(); + + DUFragment + a_frag; + DUFragment + b_frag; + DUFragment acc_frag; + du_fill_fragment(acc_frag, 0); + + // K = 192 = 6 x kInt8TileK DUMMA steps; k ascending, int32 accumulation. + // Fragment loads hit LDS (du_load_matrix_sync on a __shared__ pointer, + // the same pattern the staged prefill kernel above validates), so the K + // loop carries no vmem and no global-latency wait. The bound is the + // compile-time kLdsK (the exact-shape guard fixes k == 192), so the six + // steps fully unroll and the scheduler can issue step i+1's ds_reads while + // step i's v_mmac is still in the pipe, instead of a rolled loop in which + // each iteration's reads can only issue after the previous step retires. + for (int k0 = 0; k0 < kLdsK; k0 += kInt8TileK) { + du_load_matrix_sync(a_frag, a_tile + k0, kLdsAStride); + // col_major on the n-major stage: byte (k, n) sits at n*208 + k, so the + // k0 step's fragment starts at b_tile + k0 and each lane's 8 bytes are + // one contiguous run in its column (one ds_read_b64 per lane). + du_load_matrix_sync(b_frag, b_tile + k0, kLdsBStride); + du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + + // Direct fragment -> ((acc * x_scale) * weight_scale) -> bf16 store with + // wv1..wv3 supplied from the prologue preload: no vmem load issued inside + // the mmac chain remains to gate the store tail (xv and wv0 load here at + // their proven round-6 issue points). M == 16 fills the whole tile; row is + // always in range. + store_acc_fragment_tail_preloaded(acc_frag, x_scale, weight_scale, wv1, wv2, + wv3, out, n, n0, lane); +} + +// Generic scalar int8/int32 fallback: one thread per output element, exact +// int32 accumulation over K, fp32 scaling, bf16 store. Serves every +// unmatched (m, n, k) and all small-M API cases (including the assigned +// M=16 shape at bootstrap). Requires only m >= 1 and the API-validated +// alignment properties; it never reads out of bounds. +__global__ __launch_bounds__(kScalarThreads) +void w8a8_gemm_scalar_fallback_kernel( + const int8_t* __restrict__ a, + const int8_t* __restrict__ b, + const float* __restrict__ x_scale, + const float* __restrict__ weight_scale, + uint16_t* __restrict__ out, + int m, + int n, + int k) { + const int row = static_cast(blockIdx.y); + const int col = + static_cast(blockIdx.x) * kScalarThreads + + static_cast(threadIdx.x); + if (row >= m || col >= n) { return; } - const int row = linear / n; - const int col = linear - row * n; - const int8_t* a_row = x_q + row * k; - // b_col_stride is 1 for the n-major packed (192, 4096) pair (each output - // column has a contiguous K run) and n for the raw [K, N] layout. - const int8_t* b_col = weight + col * b_col_stride; + const int8_t* a_row = a + static_cast(row) * k; int32_t acc = 0; - for (int kk = 0; kk < k; ++kk) { - acc += static_cast(a_row[kk]) * - static_cast(b_col[kk * b_col_stride]); + if (k == 192 && n == 4096) { + // The (k, n) == (192, 4096) down-proj packed buffer is n-major + // P[n*K+kk] = W[kk*N+n]: column col is one contiguous 192-byte run. + const int8_t* b_col = b + static_cast(col) * k; +#pragma unroll 4 + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[kk]); + } + } else { + const int8_t* b_col = b + col; +#pragma unroll 4 + for (int kk = 0; kk < k; ++kk) { + acc += static_cast(a_row[kk]) * + static_cast(b_col[static_cast(kk) * n]); + } } - const float scaled = - static_cast(acc) * x_scale[row] * weight_scale[col]; - out[linear] = __float2bfloat16(scaled); + const float scaled = static_cast(acc) * x_scale[row] * + weight_scale[col]; + out[static_cast(row) * n + col] = bf16_bits_from_f32(scaled); } -// --------------------------------------------------------------------------- -// Pack kernels for launch_pack_w8a8_weight (run once per weight tensor -// outside the timed region and outside CUDA/HIP Graph capture, so plain -// kernels are fine): -// * w8a8_transpose_i8_kernel: n-major transpose pack for the exact -// assigned pair (k, n) == (192, 4096): packed[col * k + row] = -// raw[row * n + col] gives every output column a contiguous K run, -// which the timed GEMM stages with two vectorized ds_write_b64 per -// thread instead of byte-granular scatters. -// * w8a8_identity_copy_i8_kernel: identity device-to-device copy for every -// unmatched (k, n) pair (raw [K, N] layout preserved). -// --------------------------------------------------------------------------- -__global__ __launch_bounds__(256) void w8a8_transpose_i8_kernel( - const int8_t* __restrict__ src, - int8_t* __restrict__ dst, - int k, - int n) { - const int64_t linear = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t total = static_cast(k) * n; - if (linear >= total) { - return; +// Generic element-wise device-to-device copy used by the identity pack. +template +__global__ void w8a8_element_copy_kernel(const T* __restrict__ src, + T* __restrict__ dst, + int64_t count) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + i < count; i += stride) { + dst[i] = src[i]; } - const int row = static_cast(linear / n); - const int col = static_cast(linear - static_cast(row) * n); - dst[static_cast(col) * k + row] = src[linear]; } -__global__ __launch_bounds__(256) void w8a8_identity_copy_i8_kernel( - const int8_t* __restrict__ src, - int8_t* __restrict__ dst, - int64_t numel) { - const int64_t linear = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (linear < numel) { - dst[linear] = src[linear]; +// One-time n-major transpose pack for the TP8 down-proj weight +// ((k, n) == (192, 4096)): packed[n*k + kk] = raw[kk*n + n_local]. Every +// logical [K, N] column of the weight becomes a contiguous k-run of the +// packed buffer (same byte count, same graph-stable buffer). Weight +// preprocessing; runs outside the timed region and Graph capture. +__global__ __launch_bounds__(kPackThreads) +void w8a8_pack_nmajor_downproj_kernel(const int8_t* __restrict__ raw, + int8_t* __restrict__ packed, int k, + int n) { + const int64_t total = static_cast(n) * k; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + idx < total; idx += stride) { + const int n_local = static_cast(idx / k); + const int kk = static_cast(idx - static_cast(n_local) * k); + packed[idx] = raw[static_cast(kk) * n + n_local]; } } -__global__ __launch_bounds__(256) void w8a8_identity_copy_f32_kernel( - const float* __restrict__ src, - float* __restrict__ dst, - int64_t numel) { - const int64_t linear = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (linear < numel) { - dst[linear] = src[linear]; +void launch_pack_identity_impl(const int8_t* raw_weight, + const float* weight_scale, + int8_t* packed_weight, + float* packed_weight_scale, int k, int n, + hipStream_t stream) { + const int64_t weight_count = static_cast(k) * n; + const int64_t scale_count = n; + const dim3 block(kPackThreads); + const unsigned weight_blocks = static_cast( + (weight_count + kPackThreads - 1) / kPackThreads); + const unsigned scale_blocks = static_cast( + (scale_count + kPackThreads - 1) / kPackThreads); + if (weight_count > 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_element_copy_kernel), + dim3(weight_blocks), block, 0, stream, raw_weight, packed_weight, + weight_count); + } + if (scale_count > 0) { + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_element_copy_kernel), dim3(scale_blocks), + block, 0, stream, weight_scale, packed_weight_scale, scale_count); } } } // namespace -// --------------------------------------------------------------------------- -// Host launch symbols (extern "C", consumed by csrc/bindings.cpp). -// --------------------------------------------------------------------------- extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -556,15 +1250,35 @@ extern "C" void launch_w8a8_gemm( hipStream_t stream) { (void)workspace; (void)workspace_bytes; - auto* out_bf16 = reinterpret_cast<__hip_bfloat16*>(out); - - // Native INT8 DUMMA m16n16k32 tiled path for the assigned shape family. - // The guard is shape-exact (k == kAssignedK == 192, the only assigned K), - // keeps M >= 128 (small-M API cases fall through to the scalar path), and - // requires N to be a multiple of the active macro-tile N (4096 always - // is). The active macro-tile is fixed at compile time, so the timed - // region never tunes. All three tile instantiations share the same guard. - if (m >= 128 && k == kAssignedK && (n % kActiveBlockN) == 0) { + auto* out_ptr = static_cast(out); + + // Exact assigned decode shape hy3_tp8_shared_down_proj_m16. The exact + // (m, n, k) == (16, 4096, 192) guard keeps the paired M=2 API shape with + // the same (N, K) on the generic scalar fallback below. + if (m == 16 && n == 4096 && k == 192) { + // One 64-lane wavefront per block, one 16x16 output tile per block: + // grid = 4096 / 16 = 256 blocks, whole-window LDS prefetch of A and of + // the packed n-major B strip, one staging barrier per block, col_major + // B fragment consumption (one 8-byte LDS read per lane per step). + const dim3 block(kWaveSize); + const dim3 grid(static_cast(n / kInt8TileN)); + hipLaunchKernelGGL( + HIP_KERNEL_NAME(w8a8_gemm_dumma_m16n16k32_packedb_lds_kernel), grid, + block, 0, stream, a, b, x_scale, weight_scale, out_ptr, n, k); + return; + } + + // Down-proj family guard (assigned M=4096 shape). The M < 128 shapes - + // including the paired M=2 API shape with the same (N, K) - deliberately + // fall through to the scalar kernel below. The guard is shape-exact on K + // (k == kAssignedK == 192, the only K whose weight is packed n-major) and + // requires N to be a multiple of the active macro-tile N. The active + // macro-tile is fixed at compile time (kActiveTileIndex), so the timed + // region never tunes; the other two family instantiations stay compiled + // for compile-time A/B. + const bool use_tiled_dumma = + m >= 128 && k == kAssignedK && (n % kActiveBlockN) == 0; + if (use_tiled_dumma) { const dim3 grid( static_cast(n / kActiveBlockN), static_cast((m + kActiveBlockM - 1) / kActiveBlockM)); @@ -573,39 +1287,44 @@ extern "C" void launch_w8a8_gemm( if (kActiveTileIndex == 0) { hipLaunchKernelGGL( HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 64, kPrefillStageK>), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); + grid, block, 0, stream, a, b, x_scale, weight_scale, out_ptr, m, n, + k); } else if (kActiveTileIndex == 1) { - hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<64, 128, kPrefillStageK>), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); + if (kActivePipeline == 0) { + // Round-3 single-buffered pipeline control (64x128 tile): same + // bytes, MMAC order and occupancy as the double-buffered arm, LDS + // halved to 14,848 B/block, one extra barrier per stage boundary. + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_sb_kernel<64, 128, kPrefillStageK>), + grid, block, 0, stream, a, b, x_scale, weight_scale, out_ptr, m, + n, k); + } else { + // Accepted double-buffered pipeline (retention arm). + hipLaunchKernelGGL( + HIP_KERNEL_NAME( + w8a8_dumma_prefill_kernel<64, 128, kPrefillStageK>), + grid, block, 0, stream, a, b, x_scale, weight_scale, out_ptr, m, + n, k); + } } else { hipLaunchKernelGGL( - HIP_KERNEL_NAME(w8a8_dumma_prefill_kernel<128, 64, kPrefillStageK>), - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); + HIP_KERNEL_NAME( + w8a8_dumma_prefill_kernel<128, 64, kPrefillStageK>), + grid, block, 0, stream, a, b, x_scale, weight_scale, out_ptr, m, n, + k); } return; } - // Generic scalar fallback for every unmatched (m, n, k) and M < 128. - // The packed weight is n-major only for the exact assigned pair - // (k, n) == (192, 4096); every other pair is identity-packed, so the - // fallback decodes b_col with column stride k for that pair and with - // column stride n otherwise (the paired M=2 shape with the same (N, K) - // is covered by the k-stride branch). - const bool packed_nmajor = (k == kAssignedK && n == kAssignedN); - const int b_col_stride = packed_nmajor ? 1 : n; - const int total = m * n; - constexpr int kFallbackThreads = 256; - const dim3 grid(static_cast( - (total + kFallbackThreads - 1) / kFallbackThreads)); - const dim3 block(static_cast(kFallbackThreads)); - hipLaunchKernelGGL( - w8a8_scalar_fallback_kernel, - grid, block, 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k, b_col_stride); + // Generic scalar fallback (one grid row per M, columns chunked). + const dim3 scalar_grid( + static_cast((n + kScalarThreads - 1) / kScalarThreads), + static_cast(m)); + const dim3 scalar_block(kScalarThreads); + hipLaunchKernelGGL(w8a8_gemm_scalar_fallback_kernel, scalar_grid, + scalar_block, 0, stream, a, b, x_scale, weight_scale, + out_ptr, m, n, k); } extern "C" void launch_pack_w8a8_weight( @@ -616,39 +1335,33 @@ extern "C" void launch_pack_w8a8_weight( int k, int n, hipStream_t stream) { - constexpr int kCopyThreads = 256; - const int64_t weight_numel = static_cast(k) * n; - const dim3 weight_grid(static_cast( - (weight_numel + kCopyThreads - 1) / kCopyThreads)); - - // Round 1: n-major transpose pack (packed[n][k], K-contiguous per output - // column) for the exact assigned pair (k, n) == (192, 4096) only. The API - // contract (int8_w8a8_gemm_api.py) allows any contiguous opaque packed - // layout; the timed GEMM and the (192, 4096) fallback consume exactly - // this layout, so the timed B staging drops from byte-granular scatters - // to two 8-byte ds_write_b64 per thread. Every other (k, n) pair keeps - // the identity device-to-device copy (raw [K, N] layout), so unmatched - // shapes and paired small-M shapes on other pairs stay on the layout - // their fallback decodes. - if (k == kAssignedK && n == kAssignedN) { - hipLaunchKernelGGL( - w8a8_transpose_i8_kernel, - weight_grid, dim3(kCopyThreads), 0, stream, - raw_weight, packed_weight, k, n); - } else { - hipLaunchKernelGGL( - w8a8_identity_copy_i8_kernel, - weight_grid, dim3(kCopyThreads), 0, stream, - raw_weight, packed_weight, weight_numel); + // The assigned TP8 down-proj weight ((k, n) == (192, 4096)) is packed + // n-major (P[n*k + kk] = raw[kk*n + n_local], one contiguous k-run per n + // column) so the M=16 kernel and the M>=128 tiled DUMMA path can consume + // B fragments col_major with one vector LDS read per lane per step. Every + // other (k, n) keeps the logical [K, N] row-major identity pack (weight + // and scale copies), which the generic fallback decodes. + if (k == 192 && n == 4096) { + const int64_t weight_count = static_cast(k) * n; + const int64_t scale_count = n; + const dim3 block(kPackThreads); + if (weight_count > 0) { + const unsigned weight_blocks = static_cast( + (weight_count + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL(HIP_KERNEL_NAME(w8a8_pack_nmajor_downproj_kernel), + dim3(weight_blocks), block, 0, stream, raw_weight, + packed_weight, k, n); + } + if (scale_count > 0) { + const unsigned scale_blocks = static_cast( + (scale_count + kPackThreads - 1) / kPackThreads); + hipLaunchKernelGGL(HIP_KERNEL_NAME(w8a8_element_copy_kernel), + dim3(scale_blocks), block, 0, stream, weight_scale, + packed_weight_scale, scale_count); + } + return; } - - // The scale copy is unchanged (N-length, order-independent). - const dim3 scale_grid(static_cast( - (n + kCopyThreads - 1) / kCopyThreads)); - hipLaunchKernelGGL( - w8a8_identity_copy_f32_kernel, - scale_grid, dim3(kCopyThreads), 0, stream, - weight_scale, packed_weight_scale, static_cast(n)); + launch_pack_identity_impl(raw_weight, weight_scale, packed_weight, + packed_weight_scale, k, n, stream); } // @@end -// @@end diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj.hip index 81b98bb3..4d72a3d5 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp4_qkv_proj_m4096 commit=8e400dcec14217e9b03ccc7512c4d99149820b06 added=2026-08-24 +// baseline_us=31556.2 speedup=20.2283 // median_us=1560 p90_us=1575 // source=minimaxm3-dsh-tp4-m4096-1-0c2f84a9 // MetaInfer W8A8 INT8 GEMM bootstrap for gfx928 (K500SM_AI). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj_and_indexer_qk.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj_and_indexer_qk.hip index eca651cf..0df41638 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj_and_indexer_qk.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/qkv_proj_and_indexer_qk.hip @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp4_qkv_proj_and_indexer_qk_m4096 commit=d6cbef3829c738228f5dab5255b89c384c66c6ef added=2026-08-24 +// baseline_us=35918.8 speedup=33.1966 // median_us=1082 p90_us=1089 // source=minimaxm3-dsh-tp4-m4096-1-0c2f84a9 // MetaInfer W8A8 INT8 GEMM bootstrap for gfx928 (K500SM_AI). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/shared_gate_up_proj.hip index f9b6e0ad..4db5a3fa 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/shared_gate_up_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP4/M4096/shared_gate_up_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp4_shared_gate_up_proj_m4096 commit=5c8bfa17fe45a29749afdfc9124782e9068e415c added=2026-08-24 +// baseline_us=19776.5 speedup=20.0471 // median_us=986.5 p90_us=987.7 // source=minimaxm3-dsh-tp4-m4096-1-0c2f84a9 // MetaInfer W8A8 INT8 GEMM backend for Hygon gfx928 (K500SM_AI). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/o_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/o_proj.hip index f771caa3..e41630bf 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/o_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/o_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_o_proj_m16 commit=35384b7737eacb9fb8946dc1c5f984e8bf364aae added=2026-08-28 +// baseline_us=38.655 speedup=3.39974 // median_us=11.37 p90_us=11.39 // source=minimax-dsh-tp8-m16-1-78402260 // W8A8 INT8 GEMM - gfx928 (K500SM_AI) - worker_2 bootstrap (iteration 1). diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj.hip index bf0cbf03..3434210f 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_qkv_proj_m16 commit=b86aaad82e2a39632cd692ec7fe88dcaeeadaaed added=2026-08-28 +// baseline_us=110.478 speedup=3.32765 // median_us=33.2 p90_us=33.31 // source=minimax-dsh-tp8-m16-1-78402260 // MetaInfer W8A8 INT8 GEMM — gfx928 (K500SM_AI) bootstrap. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj_and_indexer_qk.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj_and_indexer_qk.hip index 3f4a48e3..ed05e4cb 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj_and_indexer_qk.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/qkv_proj_and_indexer_qk.hip @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_qkv_proj_and_indexer_qk_m16 commit=e72025dfcfa50ff3bf3fcb00b87663e65f4e0fc4 added=2026-08-28 +// baseline_us=111.486 speedup=2.80609 // median_us=39.73 p90_us=39.81 // source=minimax-dsh-tp8-m16-1-78402260 // csrc/w8a8_gemm_hip.hip diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_down_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_down_proj.hip index f5930cb7..fe0eade0 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_down_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_down_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_shared_down_proj_m16 commit=2d8b403746ea59e59c41eabce07da2e6b2447072 added=2026-08-28 +// baseline_us=23.152 speedup=2.58595 // median_us=8.953 p90_us=9.069 // source=minimax-dsh-tp8-m16-1-78402260 // MetaInfer W8A8 INT8 GEMM - HIP implementation for Hygon K500SM_AI / gfx928. diff --git a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_gate_up_proj.hip index 4400c8bd..8f28c5a6 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_gate_up_proj.hip +++ b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/minimax-m3/TP8/M16/shared_gate_up_proj.hip @@ -1,4 +1,5 @@ // @@variant shape=minimax_tp8_shared_gate_up_proj_m16 commit=4358dea23da7468d3868ce258b971667a50c139e added=2026-08-28 +// baseline_us=104.527 speedup=6.10912 // median_us=17.11 p90_us=17.13 // source=minimax-dsh-tp8-m16-1-78402260 // MetaInfer W8A8 INT8 GEMM - HIP implementation for Hygon K500SM_AI / gfx928. From 7f1414db948a76058c7cd2c26daed2dc75629c84 Mon Sep 17 00:00:00 2001 From: fms-zth <2579013254@qq.com> Date: Wed, 23 Sep 2026 20:05:45 +0800 Subject: [PATCH 5/7] feat(harness-evolve): import HE (AHE outer loop) from the DCU dev container MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Import metainfer/tasks/harness_evolve/ verbatim from the DCU development container, so the AHE outer loop lives in the repository alongside the DKAO task it drives. Source: 10.8.160.17:/root/zth_agent/MetaInfer/metainfer/tasks/harness_evolve 73 files, imported without local modification. What this adds: - FLOW.md / README.md: the fixed variant-round protocol (HE only launches DKAO tasks; measurement, GPU gate and contamination stay inside DKAO). - orchestrator/: pipeline, decision_engine, evolve, rounds, promotion, champion_selection, supervisor, process_control, variant, pool. - server/ + static/: the HE web surface (routes, plugin, he-api/detail JS). - tests/: 33 tests including test_flow_contract.py, which keeps FLOW.md and the code in sync. Scope note: this commit is material-only. The package is NOT wired into the framework yet and is not expected to import cleanly on this branch — it depends on host APIs that are not part of the DKAO-only PR lineage: - metainfer.orchestrator.gpu_broker (not in this tree) - metainfer.server.* (StateStore / forms / registry, next-app lineage) - metainfer.tasks.dcu_kernel_auto_opt.orchestrator.{harness_io,variant_promote} plus 11 hardcoded /root/zth_agent/... paths carried over from the container. Follow-up: reconcile the host API differences and drop the skipif guards that currently stand in for this module (see the DKAO sync guide). --- metainfer/tasks/harness_evolve/FLOW.md | 252 ++ metainfer/tasks/harness_evolve/README.md | 197 ++ metainfer/tasks/harness_evolve/__init__.py | 13 + metainfer/tasks/harness_evolve/form.yaml | 209 ++ .../harness_evolve/orchestrator/__init__.py | 7 + .../orchestrator/adapters/__init__.py | 1 + .../orchestrator/adapters/eval.py | 834 ++++++ .../orchestrator/attribution.py | 134 + .../orchestrator/champion_selection.py | 211 ++ .../tasks/harness_evolve/orchestrator/cli.py | 462 ++++ .../harness_evolve/orchestrator/config.py | 170 ++ .../orchestrator/decision_engine.py | 588 ++++ .../harness_evolve/orchestrator/evolve.py | 272 ++ .../orchestrator/mechanism_checks.py | 352 +++ .../harness_evolve/orchestrator/phases.py | 40 + .../harness_evolve/orchestrator/pipeline.py | 2422 +++++++++++++++++ .../harness_evolve/orchestrator/plugin.py | 11 + .../tasks/harness_evolve/orchestrator/pool.py | 241 ++ .../orchestrator/pool_builder.py | 105 + .../orchestrator/process_control.py | 133 + .../harness_evolve/orchestrator/promotion.py | 416 +++ .../harness_evolve/orchestrator/publish.py | 200 ++ .../orchestrator/question_roles.py | 428 +++ .../orchestrator/repo_cleanup.py | 300 ++ .../harness_evolve/orchestrator/round_plan.py | 209 ++ .../harness_evolve/orchestrator/rounds.py | 481 ++++ .../harness_evolve/orchestrator/state.py | 195 ++ .../harness_evolve/orchestrator/supervisor.py | 737 +++++ .../harness_evolve/orchestrator/variant.py | 324 +++ .../tasks/harness_evolve/server/__init__.py | 1 + .../tasks/harness_evolve/server/plugin.py | 34 + .../tasks/harness_evolve/server/routes.py | 1553 +++++++++++ .../tasks/harness_evolve/static/dkao.css | 909 +++++++ .../tasks/harness_evolve/static/he-api.js | 142 + .../tasks/harness_evolve/static/he-detail.js | 1346 +++++++++ metainfer/tasks/harness_evolve/static/he.css | 166 ++ .../tasks/harness_evolve/tests/__init__.py | 0 .../tasks/harness_evolve/tests/conftest.py | 19 + .../tests/test_champion_record_failure.py | 135 + .../tests/test_champion_selection.py | 74 + .../tests/test_config_req_contract.py | 100 + .../tests/test_decision_engine.py | 196 ++ .../tests/test_evolve_failure.py | 87 + .../tests/test_flow_contract.py | 220 ++ .../harness_evolve/tests/test_form_rounds.py | 137 + .../tests/test_four_card_overview.py | 123 + .../tests/test_freeze_and_review.py | 125 + .../tests/test_gate_display_contract.py | 431 +++ .../harness_evolve/tests/test_gpu_cards.py | 361 +++ .../tests/test_legacy_protocol_warning.py | 167 ++ .../tests/test_live_dkao_mirror.py | 125 + .../tests/test_live_mirror_layouts.py | 64 + .../harness_evolve/tests/test_liveness.py | 277 ++ .../tests/test_mechanism_gate.py | 293 ++ .../harness_evolve/tests/test_one_driver.py | 240 ++ .../tests/test_operator_stop.py | 437 +++ .../harness_evolve/tests/test_pipeline_dry.py | 132 + .../tests/test_pool_protocol.py | 333 +++ .../tests/test_promotion_flow.py | 698 +++++ .../tests/test_promotion_pipeline.py | 129 + .../tests/test_promotion_routes.py | 99 + .../harness_evolve/tests/test_publish.py | 92 + .../tests/test_question_roles.py | 187 ++ .../tests/test_real_adapters.py | 466 ++++ .../harness_evolve/tests/test_repo_cleanup.py | 152 ++ .../tasks/harness_evolve/tests/test_resume.py | 356 +++ .../tests/test_round_incomplete.py | 144 + .../tests/test_round_selection_evidence.py | 141 + .../harness_evolve/tests/test_state_routes.py | 97 + .../harness_evolve/tests/test_supervisor.py | 541 ++++ .../tests/test_variant_rounds.py | 414 +++ .../tasks/harness_evolve/tools/__init__.py | 0 .../harness_evolve/tools/recover_evolve.py | 130 + 73 files changed, 21817 insertions(+) create mode 100644 metainfer/tasks/harness_evolve/FLOW.md create mode 100644 metainfer/tasks/harness_evolve/README.md create mode 100644 metainfer/tasks/harness_evolve/__init__.py create mode 100644 metainfer/tasks/harness_evolve/form.yaml create mode 100644 metainfer/tasks/harness_evolve/orchestrator/__init__.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/adapters/__init__.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/adapters/eval.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/attribution.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/champion_selection.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/cli.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/config.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/decision_engine.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/evolve.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/mechanism_checks.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/phases.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/pipeline.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/plugin.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/pool.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/pool_builder.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/process_control.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/promotion.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/publish.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/question_roles.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/repo_cleanup.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/round_plan.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/rounds.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/state.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/supervisor.py create mode 100644 metainfer/tasks/harness_evolve/orchestrator/variant.py create mode 100644 metainfer/tasks/harness_evolve/server/__init__.py create mode 100644 metainfer/tasks/harness_evolve/server/plugin.py create mode 100644 metainfer/tasks/harness_evolve/server/routes.py create mode 100644 metainfer/tasks/harness_evolve/static/dkao.css create mode 100644 metainfer/tasks/harness_evolve/static/he-api.js create mode 100644 metainfer/tasks/harness_evolve/static/he-detail.js create mode 100644 metainfer/tasks/harness_evolve/static/he.css create mode 100644 metainfer/tasks/harness_evolve/tests/__init__.py create mode 100644 metainfer/tasks/harness_evolve/tests/conftest.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_champion_record_failure.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_champion_selection.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_config_req_contract.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_decision_engine.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_evolve_failure.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_flow_contract.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_form_rounds.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_four_card_overview.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_freeze_and_review.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_gate_display_contract.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_gpu_cards.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_legacy_protocol_warning.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_live_dkao_mirror.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_live_mirror_layouts.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_liveness.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_mechanism_gate.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_one_driver.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_operator_stop.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_pipeline_dry.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_pool_protocol.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_promotion_flow.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_promotion_pipeline.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_promotion_routes.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_publish.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_question_roles.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_real_adapters.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_repo_cleanup.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_resume.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_round_incomplete.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_round_selection_evidence.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_state_routes.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_supervisor.py create mode 100644 metainfer/tasks/harness_evolve/tests/test_variant_rounds.py create mode 100644 metainfer/tasks/harness_evolve/tools/__init__.py create mode 100644 metainfer/tasks/harness_evolve/tools/recover_evolve.py 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()) From 80c508ec5f3319716ed5e947c13ca09e8c5861f7 Mon Sep 17 00:00:00 2001 From: 6eanut Date: Sun, 20 Sep 2026 10:16:57 +0800 Subject: [PATCH 6/7] feat(evolve-kernel): per-iteration progress table in detail page MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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, and the per-iteration "goal" (what the optimizer set out to do) was not visible anywhere. Add an "Iteration Progress" panel rendered from the already-fetched `/iterations` payload: one row per iteration with status, exec time, speedup, library size, duration and goal. Rows that set a new best exec time are marked green with a "▲" tag — the "is it actually improving?" signal. - export `formatMs` from ok-kernel-library.js so both tables share one formatter instead of duplicating it - coerce `perf.speedup` with Number() before toFixed: the value is stored straight from an LLM-generated perf harness report with no numeric validation, so a malformed row now degrades to "—" instead of throwing - add tests/test_iteration_progress.py pinning the reader fields the table renders, plus that a partially written iteration record does not take the whole list down --- .../tasks/evolve_kernel/static/ok-detail.js | 82 ++++++++++++++++++- .../evolve_kernel/static/ok-kernel-library.js | 2 +- metainfer/tasks/evolve_kernel/static/ok.css | 27 ++++++ .../tests/test_iteration_progress.py | 63 ++++++++++++++ 4 files changed, 172 insertions(+), 2 deletions(-) create mode 100644 metainfer/tasks/evolve_kernel/tests/test_iteration_progress.py diff --git a/metainfer/tasks/evolve_kernel/static/ok-detail.js b/metainfer/tasks/evolve_kernel/static/ok-detail.js index 0770a0a3..b6c3b015 100644 --- a/metainfer/tasks/evolve_kernel/static/ok-detail.js +++ b/metainfer/tasks/evolve_kernel/static/ok-detail.js @@ -9,9 +9,10 @@ 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, @@ -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 }) { @@ -262,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

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.css b/metainfer/tasks/evolve_kernel/static/ok.css index 2ce837b0..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; 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) == [] From 988845d8b819a530dbd0f83ddba24bbb7bd116f7 Mon Sep 17 00:00:00 2001 From: fms-zth <2579013254@qq.com> Date: Sun, 20 Sep 2026 21:22:47 +0800 Subject: [PATCH 7/7] feat(dcu-kernel-auto-opt): refresh DKAO harness from the DCU dev container MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Updates the task plugin to the latest container content (/root/zth_agent/MetaInfer/metainfer/tasks/dcu_kernel_auto_opt), applied as an incremental 3-way merge so the PR-side alignment fixes stay intact (init_or_resume(task_id), CPU-only-CI torch guards, hermetic dsh tests). Added: - orchestrator/contention.py, orchestrator/round_strategy.py and harness_default/systemprompt/round_strategy.yaml - tools/retest_variant.py, tools/baseline/tp8_m8_* (Triton m8 baselines) - reports/: down-proj diagnostic scripts and tooling checks - tests: test_contention, test_execution_profile, test_gpu_mode, test_retest_tool, test_round_strategy, test_shape_catalog_panel, test_ahe_child_embed Changed: planner/gate_policy/gpu_preflight/harness_io/config/prompts, gen_and_opt_pipeline, w8a8_pipeline, w8a8_baselines, server/routes, static detail/shape panels, form.yaml, WORKFLOW.md, bridge/dsh. Dependency guards (tests only): Three features measured by this plugin live outside this PR's scope for now — metainfer.orchestrator.gpu_broker (GPU broker leases), the DCU toolchain (/opt/dtk/bin/hipprof PMC profiling) and the harness_evolve sibling task. Eleven tests that require them are marked with pytest.mark.skipif via importlib.util.find_spec / path checks (same convention as the torch guards), so the suite stays green in CI and they re-enable automatically once those modules land. --- .../tasks/dcu_kernel_auto_opt/WORKFLOW.md | 23 +- .../api/int8w8a8gemm/int8_w8a8_gemm_api.py | 19 +- .../dcu_kernel_auto_opt/assets/w8a8_bench.py | 161 ++ .../dcu_kernel_auto_opt/bridge/dsh/README.md | 10 +- .../bridge/dsh/dsh_agent.py | 36 +- .../bridge/dsh/run_dsh_task.sh | 3 +- .../bridge/dsh/tests/smoke_sdk.py | 2 +- .../bridge/dsh/tests/test_dsh_agent.py | 37 +- metainfer/tasks/dcu_kernel_auto_opt/form.yaml | 21 +- .../harness_default/README.md | 73 +- .../harness_default/manifest.yaml | 20 +- .../harness_default/planner_catalog.yaml | 6 +- .../systemprompt/round_strategy.yaml | 204 ++ .../orchestrator/config.py | 180 +- .../orchestrator/contention.py | 187 ++ .../orchestrator/gen_and_opt_pipeline.py | 109 +- .../orchestrator/gpu_preflight.py | 341 ++- .../orchestrator/harness_io.py | 64 +- .../orchestrator/planner.py | 24 +- .../orchestrator/prompts.py | 264 +- .../orchestrator/round_strategy.py | 454 ++++ .../orchestrator/w8a8_baselines.py | 29 + .../orchestrator/w8a8_pipeline.py | 293 ++- .../reports/_diag_down_proj.sh | 40 + .../reports/_diag_down_proj_profiler.py | 99 + .../reports/_diag_down_proj_sweep.sh | 51 + .../dcu_kernel_auto_opt/reports/_make_xlsx.py | 14 +- .../reports/_tool_check_decompose.py | 109 + .../reports/_tool_check_floor.py | 103 + .../reports/_tool_check_timing.py | 124 + .../dkao_optimized_operators_bestvariant.csv | 72 +- .../dkao_optimized_operators_bestvariant.json | 191 +- .../dkao_optimized_operators_bestvariant.xlsx | Bin 16612 -> 18133 bytes .../dcu_kernel_auto_opt/server/routes.py | 183 +- .../static/dkao-agent-fields.js | 10 +- .../dcu_kernel_auto_opt/static/dkao-detail.js | 135 +- .../static/dkao-shape-input.js | 32 +- .../dcu_kernel_auto_opt/tests/conftest.py | 14 + .../tests/test_ahe_child_embed.py | 175 ++ .../dcu_kernel_auto_opt/tests/test_config.py | 39 +- .../tests/test_contention.py | 91 + .../tests/test_contract_snapshot.py | 74 +- .../tests/test_execution_profile.py | 266 ++ .../tests/test_gpu_mode.py | 108 + .../tests/test_gpu_preflight.py | 407 ++- .../dcu_kernel_auto_opt/tests/test_plugin.py | 2 +- .../tests/test_retest_tool.py | 52 + .../tests/test_round_strategy.py | 490 ++++ .../tests/test_shape_catalog_panel.py | 126 + .../tests/test_w8a8_baselines.py | 31 + .../dcu_kernel_auto_opt/tools/__init__.py | 0 .../tools/baseline/bench_triton_tp8_m8.py | 238 ++ .../tools/baseline/tp8_m8_graph_merged.json | 197 ++ .../tools/baseline/tp8_m8_graph_run1.json | 662 +++++ .../tools/baseline/tp8_m8_graph_run2.json | 662 +++++ .../tools/baseline/tp8_m8_graph_run3.json | 662 +++++ .../tools/retest_variant.py | 331 +++ .../glm52/TP8/M4096/fused_qkv_a_proj.hip | 1455 ++++------- .../int8w8a8-gemm/glm52/TP8/M4096/o_proj.hip | 1370 +++------- .../glm52/TP8/M4096/shared_down_proj.hip | 1900 ++++++-------- .../glm52/TP8/M4096/shared_gate_up_proj.hip | 1018 +++----- .../int8w8a8-gemm/hy3/TP4/M4096/o_proj.hip | 842 ++++-- .../int8w8a8-gemm/hy3/TP4/M4096/qkv_proj.hip | 1171 ++++----- .../int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip | 2263 +++++++++-------- .../hy3/TP8/M16/shared_gate_up_proj.hip | 1419 +++++------ 65 files changed, 12664 insertions(+), 7124 deletions(-) create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/harness_default/systemprompt/round_strategy.yaml create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/orchestrator/contention.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/orchestrator/round_strategy.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj.sh create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj_profiler.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/_diag_down_proj_sweep.sh create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_decompose.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_floor.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/reports/_tool_check_timing.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_ahe_child_embed.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_contention.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_execution_profile.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_mode.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_retest_tool.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_round_strategy.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tests/test_shape_catalog_panel.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/__init__.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/bench_triton_tp8_m8.py create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_merged.json create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run1.json create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run2.json create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/baseline/tp8_m8_graph_run3.json create mode 100644 metainfer/tasks/dcu_kernel_auto_opt/tools/retest_variant.py diff --git a/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md b/metainfer/tasks/dcu_kernel_auto_opt/WORKFLOW.md index 53f8912c..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`。 @@ -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,…)` 已实测 + `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。 @@ -257,7 +271,8 @@ TP8 另有 wq_b/wo_b (1024,4096)、gate_up (4096,512)、down (256,4096)、indexe **规范 skill 以 `metainfer/tasks/dcu_kernel_auto_opt/skills/` 内的副本为种子**(int8-w8a8 家族, 2026-08-27 起随插件一起维护):新机器上 `sync_skill_libraries()`(或 WebUI 同步按钮)会自动 把缺失的 skill 补种进 `~/.dsh/skills/`,再镜像到 `~/.claude/skills/`;已有的库 skill 不被覆盖。 - baseline 测量工具在 `tools/baseline/`(`bench_triton_tp4/tp8_m4096.py` + `int8_utils.py`)。 + 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 里)。 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 6e4b54a5..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 @@ -331,6 +487,11 @@ def _env_default(name: str, fallback: int) -> 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 0bfbd0fc..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() @@ -294,6 +301,13 @@ def _headless_backend_main(args: Any, prompt: str) -> int: 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 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 fb0e2302..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 @@ -166,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 3f1066a4..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." @@ -224,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 index a5a0a41b..67372060 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/harness_default/README.md +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/README.md @@ -1,38 +1,61 @@ # harness_default —— DKAO 可演化 harness 的种子目录 -这是 AHE「组件可观测性」的第一步落地(M1 切片 1,纯新增、不改运行时行为)。 +这是 AHE「组件可观测性」的落地目录:DKAO 的 harness(验收门阈值、轮次指令文案、 +planner 策略、以及后续的 skills/工具/中间件/记忆)**文件化**在这里,成为可版本化、 +可校验、可演化的组件。AHE 外循环(`harness_evolve`)以本目录为种子,拷贝出可写 +workspace(git 仓库),每轮评测 pin 一个快照;Evolve Agent 只改 workspace。 -## 现状与用途 - -DKAO 的 harness(gates 阈值、prompts/菜单、planner、skills、工具集、中间件策略、记忆) -目前散在 Python 代码与静态文件里。AHE 只能演化"文件化、可版本化、可校验"的组件, -因此我们先把它们**播种成文件**,放在本目录: +## 目录内容与接线状态 ``` harness_default/ -├── manifest.yaml # 组件清单(AHE Evolve 唯一可写空间的入口描述) -├── gates.yaml # 验收/plateau/ISA 门的规范值(seed 阶段) -├── planner_catalog.yaml # 优化方案选择机制:方案目录(★ 见 orchestrator/planner.py v0) -├── README.md # 本文件 +├── 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/ # prompts 模板(coordinator/bootstrap/worker/synthesis) - ├── tools/ # cordis 组合与 per-role 工具白名单 - ├── middleware/ # resume/compaction/重试策略参数 - ├── skills/ # SKILL.md 种子 - └── memory/ # LongTermMEMORY(measured 事实) + ├── systemprompt/ # 其余 prompt 模板(coordinator/bootstrap/worker/synthesis) + ├── tools/ # cordis 组合与 per-role 工具白名单 + ├── middleware/ # resume/compaction/重试策略参数 + ├── skills/ # SKILL.md 种子 + └── memory/ # LongTermMEMORY(measured 事实) ``` -## seed 阶段语义(重要) +## 语义(重要) + +**`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 试什么"可演化, +而选择状态机本身不可被改写。 -- **wired: false**:运行时管线仍读 Python 常量(config.py:105、w8a8_pipeline.py:78-80 等), - 本目录文件**尚未被加载器接线**——因此默认行为零变化。 -- 一致性由测试守卫:`tests/test_harness_io.py` 断言 gates.yaml 的值与 Python 常量相等, - 防止两份来源漂移;接线(renderer/loader)在后续切片落地,落地时删掉该守卫并让 - 运行时以本目录为唯一来源。 -- AHE 外循环(harness_evolve)将以本目录为种子,拷贝出可写 workspace(git 仓库), - 每轮评测 pin 一个 workspace 快照;Evolve Agent 只改 workspace。 +**一致性由测试守卫**(drift guard,见 manifest 每个组件的 `drift_guard`): +每个 YAML 种子的取值都被断言等于 Python 内置默认值,因此两条来源不会无声漂移。 +`tests/test_round_strategy.py` 另外用 5400 个渲染用例的 sha256 指纹把文案钉死 —— 任何 +意外改词都会让该测试失败。**有意改文案时**:同步改 `BUILTIN_STRATEGY` 与种子文件, +并有意重算指纹。 ## 读取方式 -见 `orchestrator/harness_io.py`(harness_root 解析、load_manifest/load_gates、seed_workspace)。 -可用环境变量 `METAINFER_HARNESS_ROOT` 覆盖根目录(默认本目录)。 +- 根目录解析:`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/manifest.yaml b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/manifest.yaml index 3512f59a..955405d9 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/harness_default/manifest.yaml +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/manifest.yaml @@ -28,9 +28,19 @@ components: 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/, 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). + 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 index 7eb8302c..2a8868ee 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_catalog.yaml +++ b/metainfer/tasks/dcu_kernel_auto_opt/harness_default/planner_catalog.yaml @@ -3,8 +3,10 @@ # DKAO 优化方案选择机制的「方案目录」:把"每一轮可以试什么方案"枚举成稳定 id, # 供 orchestrator/planner.py(状态条件化选择器 v0)与后续 harness 演化使用。 # -# wired: false —— 运行时仍由 prompts.py::w8a8_round_strategy 按轮数菜单出文字, -# 本目录与 planner.py 当前是纯新增的"选择机制 v0",用于沉淀/测试/演进,尚未接线。 +# 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 菜单近似)。 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/gen_and_opt_pipeline.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gen_and_opt_pipeline.py index de293088..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,7 @@ 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 @@ -55,7 +56,9 @@ RealW8A8OptimizationPipeline, W8A8Runner, _SOURCE_ONLY_AGENT_ARGS, + _gate_preflight_enabled, _sha256_file, + ensure_any_measurement_gate, evaluate_final_target, snapshot_accepted_kernel_artifact, ) @@ -209,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 @@ -253,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), @@ -927,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 } @@ -981,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) @@ -2287,6 +2341,46 @@ 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). @@ -2317,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) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_preflight.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_preflight.py index 723a8a5c..d5a92e4b 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_preflight.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/gpu_preflight.py @@ -1,21 +1,24 @@ """GPU admission pre-check shared by production DKAO tasks and AHE children. -Two different questions are answered separately, because they need different -rules: - -* **Can we run here at all?** Only VRAM pressure decides that. Our W8A8 - kernels need a few GB at most, so the operating rule is simple: a device is - usable while its used-VRAM percentage is at or below - ``vram_limit_percent`` (default 85). Above that the caller prefers another - device, and only then waits. -* **Will the measurement be trustworthy?** Busy compute (high HCU%, high - board power) or foreign processes hold the device (KFD lists pids invisible - inside this container) do NOT block the run — they only mark the question as - ``measurement_suspect`` so a polluted reading cannot silently decide an - A/B comparison. +The 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 listing, so it works inside the task containers without extra tools. +KFD/sysfs listings, so it works inside the task containers without extra tools. """ from __future__ import annotations @@ -26,9 +29,13 @@ from pathlib import Path from typing import Any, Dict, List, Optional -#: A device is usable while at least this much VRAM stays free. Absolute free -#: memory is the right unit for us (one W8A8 worker needs ~1-3GB), whereas a -#: percentage punishes big cards that hold idle reservations. +#: 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 @@ -38,34 +45,131 @@ 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+)\s+" # device index r"(?P[\d.]+)C\s+" - r"(?P[\d.]+)W\s+" - r"\S+\s+" # perf mode - r"[\d.]+W\s+" # power cap - r"(?P[\d.]+)%\s+" - r"(?P[\d.]+)%" + 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) -> Dict[int, Dict[str, float]]: - """Parse ``hy-smi`` output into ``{index: {vram_percent, util_percent, power_w, temp_c}}``.""" +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")) - out[index] = { - "vram_percent": float(match.group("vram")), - "util_percent": float(match.group("util")), - "power_w": float(match.group("power")), - "temp_c": float(match.group("temp")), - } + 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 @@ -104,13 +208,27 @@ def foreign_kfd_pids(kfd_root: Path = KFD_PROC_ROOT) -> List[int]: def sample_gpu_state(gpu_ids: Optional[List[int]] = None, *, samples: int = 3, interval_s: float = 3.0, - smi_reader=read_hy_smi) -> Dict[int, Dict[str, float]]: - """Average a few hy-smi samples per device (drops idle flicker).""" + 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)): - parsed = parse_hy_smi(smi_reader()) - for index, row in parsed.items(): - if gpu_ids is not None and index not in gpu_ids: + 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": []}) @@ -118,12 +236,71 @@ def sample_gpu_state(gpu_ids: Optional[List[int]] = None, *, 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(): - out[index] = { - key: (sum(values) / len(values) if values else 0.0) - for key, values in bucket.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 @@ -142,44 +319,66 @@ def total_vram_gb(card_index: int = 0) -> float: 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]: - """Decide usability (free VRAM >= ``min_free_gb``) and suspicion (busy).""" - vram = float(state.get("vram_percent") or 0.0) + """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 - usable = free_gb >= float(min_free_gb) + 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 util > float(util_suspect_percent): - reasons.append(f"busy: HCU {util:.0f}% > {util_suspect_percent:.0f}%") - if power > float(power_suspect_w): - reasons.append(f"high power draw {power:.0f}W > {power_suspect_w:.0f}W") - # Foreign KFD entries alone are NOT evidence of contention: idle holders - # (long-lived runtimes in other containers) keep their handle while the - # device sits at 0% util / idle power, which is harmless for us. They are - # reported as notes, and the busy/power signals above decide suspicion. + 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 usable: + if not vram_ok: + reasons.append( + f"VRAM {vram:.0f}% > {float(vram_limit_percent):.0f}% limit") + if not idle_ok: reasons.append( - f"only {free_gb:.1f}GB VRAM free (< {float(min_free_gb):.1f}GB)") + 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), - "min_free_gb": float(min_free_gb), + "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, @@ -189,36 +388,48 @@ def check_gpu(gpu_id: int, state: Dict[str, float], *, def preflight_gpus(gpu_ids: List[int], *, samples: int = 3, interval_s: float = 3.0, enabled: bool = True, - min_free_gb: float = DEFAULT_MIN_FREE_GB, + 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).""" + """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": []} + "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 {}, min_free_gb=min_free_gb, + 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"]] - busy_but_usable = [g for g in gpu_ids if not checks[g]["usable"]] - # Devices over the VRAM limit stay usable in principle (configuration - # choice), but when any device is within the limit they are pushed to the - # back so concurrent children do not collide with a full device. - preferred = (clean + suspect + busy_but_usable) if not usable \ - else (clean + suspect) + 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, - "over_limit_ids": busy_but_usable, + "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, diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/harness_io.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/harness_io.py index 82b09fc6..04833baa 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/harness_io.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/harness_io.py @@ -1,17 +1,19 @@ """Harness workspace IO: locate the evolvable-harness seed, read components, and copy a seed workspace to a destination. -This is the first piece of AHE "component observability" for dcu_kernel_auto_opt: -harness components (gates, and later prompts/planner/skills/tools/middleware/ -memory) are seeded as files under ``harness_default/`` so they can become a -versioned, evolvable workspace driven by the ``harness_evolve`` outer loop. - -Slice-1 semantics (important): this module is **read-only/additive**. -Runtime pipelines still read the current Python constants (``config.py``, -``w8a8_pipeline.py``); ``harness_io`` only exposes the seed for inspection and -copying. Consistency tests guard the YAML seed against the Python constants so -the two sources cannot drift before the loader/renderer wiring lands in a later -slice (wiring will make the harness workspace the single source of truth). +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 @@ -63,6 +65,46 @@ def _load_yaml(path: Path) -> Dict[str, Any]: 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()) diff --git a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/planner.py b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/planner.py index e4ef12c1..372853d8 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/planner.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/orchestrator/planner.py @@ -1,17 +1,21 @@ -"""State-conditioned optimization plan selector (v0, not yet wired). +"""State-conditioned optimization plan selector. -DKAO currently chooses the next optimization direction with the round menus in -``prompts.py::w8a8_round_strategy`` (keyed by iteration number). This module is -the first step toward a *state-conditioned* plan selector that decides "what to -try this round" from the current measured state (history + PMC + budget), per -the design in ``docs/dkao_harness_eval_protocol.md`` / M1 slice-2. +DKAO chooses the next optimization direction in one of two ways: -Slice-2 semantics: -- **Additive and inert**: nothing in the runtime pipelines calls this module yet. - ``prompts.py``/``w8a8_pipeline.py`` keep their current behaviour unchanged. +* ``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 - already an evolvable harness component; this module only consumes it. + 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. 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/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 4ec2e402..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 @@ -33,6 +33,232 @@ _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 { @@ -132,6 +358,7 @@ def _round_strategy_text( 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 @@ -858,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 = ( @@ -880,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"] @@ -980,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, @@ -1049,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", @@ -2601,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 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/_make_xlsx.py b/metainfer/tasks/dcu_kernel_auto_opt/reports/_make_xlsx.py index 72e4d538..e9fb3e98 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/reports/_make_xlsx.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/reports/_make_xlsx.py @@ -126,7 +126,8 @@ def main(): key = (r["model"], r["kernel_repo"], r["tp"]) groups.setdefault(key, []).append(r) head2 = ["模型(任务)", "kernel_repo", "TP", "算子数", "平均加速比×", - "最小加速比×", "最大加速比×", "平均最终µs", "平均TOPS", "平均提升%"] + "最小加速比×", "最大加速比×", "平均最终µs", "平均TOPS", "平均提升%", + "验收状态"] ws2.append(head2) for c in range(1, len(head2) + 1): cell = ws2.cell(row=1, column=c) @@ -142,19 +143,21 @@ def main(): sum(sp) / len(sp), min(sp), max(sp), sum(fin) / len(fin), sum(tops) / len(tops), - sum(imp) / len(imp)] + 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 j >= 5 and j != 4: + 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], start=1): + 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}" @@ -232,7 +235,8 @@ def main(): ["计算性能(TOPS)", "= 2×M×N×K / final_us,纯算法逻辑 INT8 运算速率(非实测带宽)"], ["达标/正确性", "target_met: 最终验收 ≥3% 提升门槛;passed: 最终串行验证正确性通过(CPU int64 exact reference)"], ["验收状态", "'最终验收通过' = 该任务串行最终验收全部 shape 通过并写出 final_report.json (status=success);'仅worker验收(未过最终验收)' = 该任务未产出成功 final_report,本行取的是 worker 并行探索阶段已验收最优 variant 的 median(未经最终门限确认)"], - ["覆盖", "Hy3: TP4 M16/M4096, TP8 M4096; MiniMax-M3: TP4 M4096, TP8 M16/M4096; GLM5.2: TP8 M4096 —— 共 33 个优化 shape,全部为各任务自身 plan 的优化算子(不含默认42-shape回归项)"], + ["覆盖", "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): 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 index 10929acb..14e8cab0 100644 --- 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 @@ -1,34 +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 -hy3-dsh-tp4-m16-2-368c654c,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP4-m16-2,4,qkv_proj,16,2560,4096,INT8 W8A8,80.19,27.359989881515503,worker_0 accepted/hy3_tp4_qkv_proj_m16/kernel.hip,27.722550630569458,27.860960960388184,2.89259098373059,189.259098373059,True,True,12.103659741538994,2026-08-26 14:59 -hy3-dsh-tp4-m16-2-368c654c,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP4-m16-2,4,o_proj,16,4096,2048,INT8 W8A8,54.617,15.513530373573303,worker_1 accepted/hy3_tp4_o_proj_m16/kernel.hip,15.558480024337769,15.603289604187012,3.510432890267166,251.0432890267166,True,True,17.25332137715848,2026-08-26 14:59 -hy3-dsh-tp4-m16-2-368c654c,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP4-m16-2,4,shared_gate_up_proj,16,768,4096,INT8 W8A8,73.262,11.581599712371826,worker_2 accepted/hy3_tp4_shared_gate_up_proj_m16/kernel.hip,11.451255083084106,11.476860046386719,6.397726665631898,539.7726665631898,True,True,8.790590661865588,2026-08-26 14:59 -hy3-dsh-tp4-m16-2-368c654c,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP4-m16-2,4,shared_down_proj,16,4096,384,INT8 W8A8,23.216,9.17439490556717,worker_3 accepted/hy3_tp4_shared_down_proj_m16/kernel.hip,9.128850102424622,9.246439933776855,2.5431461508863897,154.31461508863896,True,True,5.513470747715742,2026-08-26 14:59 -hy3-tp4-m4096-dsh-test1-bfdf3002,Hy3 (gfx928 INT8 W8A8),hy3_TP4_M4096_dsh-test1,4,qkv_proj,4096,2560,4096,INT8 W8A8,24045.12,756.7258453369141,worker_0 accepted/hy3_tp4_qkv_proj_m4096/kernel.hip,756.7258453369141,760.15869140625,31.77520650070368,3077.520650070368,True,True,113.51448671844342,2026-08-22 05:44 -hy3-tp4-m4096-dsh-test1-bfdf3002,Hy3 (gfx928 INT8 W8A8),hy3_TP4_M4096_dsh-test1,4,o_proj,4096,4096,2048,INT8 W8A8,19881.949,803.8980102539062,worker_1 accepted/hy3_tp4_o_proj_m4096/kernel.hip,803.8980102539062,805.1364135742188,24.731929605996175,2373.1929605996174,True,True,85.48282973644304,2026-08-22 05:44 -hy3-tp4-m4096-dsh-test1-bfdf3002,Hy3 (gfx928 INT8 W8A8),hy3_TP4_M4096_dsh-test1,4,shared_gate_up_proj,4096,768,4096,INT8 W8A8,6678.282,220.71203231811523,worker_2 accepted/hy3_tp4_shared_gate_up_proj_m4096/kernel.hip,220.71203231811523,221.53764724731445,30.25789726939083,2925.7897269390833,True,True,116.75758455641257,2026-08-22 05:44 -hy3-tp4-m4096-dsh-test1-bfdf3002,Hy3 (gfx928 INT8 W8A8),hy3_TP4_M4096_dsh-test1,4,shared_down_proj,4096,4096,384,INT8 W8A8,4381.289,214.83811378479004,worker_3 accepted/hy3_tp4_shared_down_proj_m4096/kernel.hip,214.83811378479004,215.11651992797852,20.393443802010253,1939.3443802010254,True,True,59.97493489868936,2026-08-22 05:44 -hy3-dsh-tp8-m4096-1-0ebb994d,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M4096-1,8,qkv_proj,4096,1280,4096,INT8 W8A8,19381.383,418.2191276550293,worker_0 accepted/hy3_tp8_qkv_proj_m4096/kernel.hip,318.4573459625244,326.0519790649414,60.86021643313191,5986.021643313191,True,True,134.86789833717404,2026-08-27 07:23 -hy3-dsh-tp8-m4096-1-0ebb994d,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M4096-1,8,o_proj,4096,4096,1024,INT8 W8A8,14151.792,322.98728942871094,worker_1 accepted/hy3_tp8_o_proj_m4096/kernel.hip,282.8514766693115,283.3587074279785,50.0325901304917,4903.25901304917,True,True,121.47625592271805,2026-08-27 07:23 -hy3-dsh-tp8-m4096-1-0ebb994d,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M4096-1,8,shared_gate_up_proj,4096,384,4096,INT8 W8A8,16712.231,115.95122337341309,worker_2 accepted/hy3_tp8_shared_gate_up_proj_m4096/kernel.hip,115.34959316253662,115.53359985351562,144.8833111743286,14388.331117432861,True,True,111.70305446889755,2026-08-27 07:23 -hy3-dsh-tp8-m4096-1-0ebb994d,Hy3 (gfx928 INT8 W8A8),hy3-dsh-TP8-M4096-1,8,shared_down_proj,4096,4096,192,INT8 W8A8,10762.616,171.18735313415527,worker_3 accepted/hy3_tp8_shared_down_proj_m4096/kernel.hip,171.81203842163086,172.1600341796875,62.64180379251588,6164.180379251588,True,True,37.497086951439755,2026-08-27 07:23 -minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,qkv_proj,16,1280,6144,INT8 W8A8,110.478,33.20169925689697,worker_0 accepted/minimax_tp8_qkv_proj_m16/kernel.hip,33.19578528404236,33.328590393066406,3.328073098879459,232.8073098879459,True,True,7.581029876132359,2026-08-28 17:17 -minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,qkv_proj_and_indexer_qk,16,1536,6144,INT8 W8A8,111.486,39.72720980644226,worker_1 accepted/minimax_tp8_qkv_proj_and_indexer_qk_m16/kernel.hip,39.73748564720154,39.8534893989563,2.8055625106680915,180.55625106680915,True,True,7.599622449219237,2026-08-28 17:17 -minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,o_proj,16,6144,1024,INT8 W8A8,38.655,11.368839740753174,worker_2 accepted/minimax_tp8_o_proj_m16/kernel.hip,11.77299976348877,12.421009540557861,3.2833602969975013,228.33602969975013,True,True,17.10070466699301,2026-08-28 17:17 -minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,shared_gate_up_proj,16,768,6144,INT8 W8A8,104.527,17.114474773406982,worker_3 accepted/minimax_tp8_shared_gate_up_proj_m16/kernel.hip,17.31948971748352,17.501089572906494,6.0352239993816355,503.5223999381635,True,True,8.718209743071991,2026-08-28 17:17 -minimax-dsh-tp8-m16-1-78402260,MiniMax-M3 (gfx928 INT8 W8A8),minimax-dsh-TP8-m16-1,8,shared_down_proj,16,6144,384,INT8 W8A8,23.152,8.95284503698349,worker_3 accepted/minimax_tp8_shared_down_proj_m16/kernel.hip,9.341430068016052,10.175809860229492,2.478421379962978,147.8421379962978,True,True,8.082003660070677,2026-08-28 17:17 -minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,qkv_proj,4096,2304,6144,INT8 W8A8,31556.192,1559.944839477539,worker_0 accepted/minimax_tp4_qkv_proj_m4096/kernel.hip,1561.9083404541016,1577.305908203125,20.20361322279994,1920.361322279994,True,True,74.24514869949742,2026-08-23 18:45 -minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,qkv_proj_and_indexer_qk,4096,2560,6144,INT8 W8A8,35918.768,1081.603660583496,worker_1 accepted/minimax_tp4_qkv_proj_and_indexer_qk_m4096/kernel.hip,1089.0154266357422,1102.6435089111328,32.9827908048673,3198.27908048673,True,True,118.31698222866211,2026-08-23 18:45 -minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,o_proj,4096,6144,2048,INT8 W8A8,29702.473,765.9097671508789,worker_2 accepted/minimax_tp4_o_proj_m4096/kernel.hip,752.4849700927734,753.6545562744141,39.47251331324,3847.251331324,True,True,136.98508169710212,2026-08-23 18:45 -minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,shared_gate_up_proj,4096,1536,6144,INT8 W8A8,19776.488,986.5248489379883,worker_3 accepted/minimax_tp4_shared_gate_up_proj_m4096/kernel.hip,962.3474502563477,965.7874298095703,20.55025759639306,1955.025759639306,True,True,80.33419874225937,2026-08-23 18:45 -minimaxm3-dsh-tp4-m4096-1-0c2f84a9,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh_TP4_M4096-1,4,shared_down_proj,4096,6144,768,INT8 W8A8,11759.473,636.2568092346191,worker_3 accepted/minimax_tp4_shared_down_proj_m4096/kernel.hip,634.1561698913574,634.6337509155273,18.543496946524407,1754.3496946524406,True,True,60.95455267843923,2026-08-23 18:45 -minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,qkv_proj,4096,1280,6144,INT8 W8A8,41000.618,475.5491065979004,worker_0 accepted/minimax_tp8_qkv_proj_m4096/kernel.hip,474.8045539855957,475.41736602783203,86.35262163311907,8535.262163311907,True,True,135.68637642417067,2026-08-29 03:28 -minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,qkv_proj_and_indexer_qk,4096,1536,6144,INT8 W8A8,35604.698,531.2225532531738,worker_1 accepted/minimax_tp8_qkv_proj_and_indexer_qk_m4096/kernel.hip,519.466609954834,520.4890060424805,68.54087888939718,6754.087888939718,True,True,148.82460170966874,2026-08-29 03:28 -minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,o_proj,4096,6144,1024,INT8 W8A8,19897.517,397.44619369506836,worker_2 accepted/minimax_tp8_o_proj_m4096/kernel.hip,394.50483322143555,394.9032211303711,50.43668752426038,4943.668752426038,True,True,130.64379245024566,2026-08-29 03:28 -minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,shared_gate_up_proj,4096,768,6144,INT8 W8A8,31882.257,319.16600227355957,worker_3 accepted/minimax_tp8_shared_gate_up_proj_m4096/kernel.hip,311.9345474243164,312.3681640625,102.20816278048036,10120.816278048036,True,True,123.91928365478229,2026-08-29 03:28 -minimaxm3-dsh-tp8-m4096-1-b0482833,MiniMax-M3 (gfx928 INT8 W8A8),minimaxM3-dsh-TP8-M4096-1,8,shared_down_proj,4096,6144,384,INT8 W8A8,22389.162,312.4720001220703,worker_3 accepted/minimax_tp8_shared_down_proj_m4096/kernel.hip,311.00934982299805,311.1349678039551,71.98871034823276,7098.871034823276,True,True,62.14396076195009,2026-08-29 03:28 -glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,fused_qkv_a_proj,4096,2624,6144,INT8 W8A8,68661.417,937.5793075561523,worker_0 accepted/glm_tp8_fused_qkv_a_proj_m4096/kernel.hip,938.0712509155273,939.3240356445312,73.19424503521313,7219.4245035213125,True,True,140.78913965554716,2026-08-31 11:23 -glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,q_b_proj,4096,2048,2048,INT8 W8A8,14385.571,220.13651847839355,worker_1 accepted/glm_tp8_q_b_proj_m4096/kernel.hip,215.435791015625,215.74857711791992,66.77428542482365,6577.428542482365,True,True,159.48946182998893,2026-08-31 11:23 -glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,kv_b_proj,4096,3584,512,INT8 W8A8,6939.939,137.40396976470947,worker_1 accepted/glm_tp8_kv_b_proj_m4096/kernel.hip,134.75635528564453,134.93956565856934,51.499901324055074,5049.990132405507,True,True,111.55233090221004,2026-08-31 11:23 -glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,o_proj,4096,6144,2048,INT8 W8A8,54775.797,759.9214935302734,worker_2 accepted/glm_tp8_o_proj_m4096/kernel.hip,755.8016204833984,756.3575744628906,72.47377554571302,7147.377554571302,True,True,136.3839562001365,2026-08-31 11:23 -glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,shared_gate_up_proj,4096,512,6144,INT8 W8A8,21254.589,273.8345146179199,worker_3 accepted/glm_tp8_shared_gate_up_proj_m4096/kernel.hip,272.7110004425049,272.99739837646484,77.93814318275388,7693.814318275387,True,True,94.49491855548743,2026-08-31 11:23 -glm5-2-dsh-tp8-m4096-1-e6a280a2,GLM5.2 (gfx928 INT8 W8A8),glm5.2-dsh-TP8-M4096-1,8,shared_down_proj,4096,6144,256,INT8 W8A8,17700.262,278.44139099121094,worker_3 accepted/glm_tp8_shared_down_proj_m4096/kernel.hip,273.6263179779053,274.4647216796875,64.6877176537867,6368.77176537867,True,True,47.089410050974806,2026-08-31 11:23 +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 index 7d4d11bb..989c0d54 100644 --- 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 @@ -19,7 +19,8 @@ "target_met": true, "passed": true, "logical_tops": 12.103659741538994, - "finished_at": "2026-08-26 14:59" + "finished_at": "2026-08-26 14:59", + "validation_state": "最终验收通过" }, { "task_id": "hy3-dsh-tp4-m16-2-368c654c", @@ -41,7 +42,8 @@ "target_met": true, "passed": true, "logical_tops": 17.25332137715848, - "finished_at": "2026-08-26 14:59" + "finished_at": "2026-08-26 14:59", + "validation_state": "最终验收通过" }, { "task_id": "hy3-dsh-tp4-m16-2-368c654c", @@ -63,7 +65,8 @@ "target_met": true, "passed": true, "logical_tops": 8.790590661865588, - "finished_at": "2026-08-26 14:59" + "finished_at": "2026-08-26 14:59", + "validation_state": "最终验收通过" }, { "task_id": "hy3-dsh-tp4-m16-2-368c654c", @@ -85,7 +88,8 @@ "target_met": true, "passed": true, "logical_tops": 5.513470747715742, - "finished_at": "2026-08-26 14:59" + "finished_at": "2026-08-26 14:59", + "validation_state": "最终验收通过" }, { "task_id": "hy3-tp4-m4096-dsh-test1-bfdf3002", @@ -107,7 +111,8 @@ "target_met": true, "passed": true, "logical_tops": 113.51448671844342, - "finished_at": "2026-08-22 05:44" + "finished_at": "2026-08-22 05:44", + "validation_state": "最终验收通过" }, { "task_id": "hy3-tp4-m4096-dsh-test1-bfdf3002", @@ -129,7 +134,8 @@ "target_met": true, "passed": true, "logical_tops": 85.48282973644304, - "finished_at": "2026-08-22 05:44" + "finished_at": "2026-08-22 05:44", + "validation_state": "最终验收通过" }, { "task_id": "hy3-tp4-m4096-dsh-test1-bfdf3002", @@ -151,7 +157,8 @@ "target_met": true, "passed": true, "logical_tops": 116.75758455641257, - "finished_at": "2026-08-22 05:44" + "finished_at": "2026-08-22 05:44", + "validation_state": "最终验收通过" }, { "task_id": "hy3-tp4-m4096-dsh-test1-bfdf3002", @@ -173,7 +180,8 @@ "target_met": true, "passed": true, "logical_tops": 59.97493489868936, - "finished_at": "2026-08-22 05:44" + "finished_at": "2026-08-22 05:44", + "validation_state": "最终验收通过" }, { "task_id": "hy3-dsh-tp8-m4096-1-0ebb994d", @@ -195,7 +203,8 @@ "target_met": true, "passed": true, "logical_tops": 134.86789833717404, - "finished_at": "2026-08-27 07:23" + "finished_at": "2026-08-27 07:23", + "validation_state": "最终验收通过" }, { "task_id": "hy3-dsh-tp8-m4096-1-0ebb994d", @@ -217,7 +226,8 @@ "target_met": true, "passed": true, "logical_tops": 121.47625592271805, - "finished_at": "2026-08-27 07:23" + "finished_at": "2026-08-27 07:23", + "validation_state": "最终验收通过" }, { "task_id": "hy3-dsh-tp8-m4096-1-0ebb994d", @@ -239,7 +249,8 @@ "target_met": true, "passed": true, "logical_tops": 111.70305446889755, - "finished_at": "2026-08-27 07:23" + "finished_at": "2026-08-27 07:23", + "validation_state": "最终验收通过" }, { "task_id": "hy3-dsh-tp8-m4096-1-0ebb994d", @@ -261,7 +272,100 @@ "target_met": true, "passed": true, "logical_tops": 37.497086951439755, - "finished_at": "2026-08-27 07:23" + "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", @@ -283,7 +387,8 @@ "target_met": true, "passed": true, "logical_tops": 7.581029876132359, - "finished_at": "2026-08-28 17:17" + "finished_at": "2026-08-28 17:17", + "validation_state": "最终验收通过" }, { "task_id": "minimax-dsh-tp8-m16-1-78402260", @@ -305,7 +410,8 @@ "target_met": true, "passed": true, "logical_tops": 7.599622449219237, - "finished_at": "2026-08-28 17:17" + "finished_at": "2026-08-28 17:17", + "validation_state": "最终验收通过" }, { "task_id": "minimax-dsh-tp8-m16-1-78402260", @@ -327,7 +433,8 @@ "target_met": true, "passed": true, "logical_tops": 17.10070466699301, - "finished_at": "2026-08-28 17:17" + "finished_at": "2026-08-28 17:17", + "validation_state": "最终验收通过" }, { "task_id": "minimax-dsh-tp8-m16-1-78402260", @@ -349,7 +456,8 @@ "target_met": true, "passed": true, "logical_tops": 8.718209743071991, - "finished_at": "2026-08-28 17:17" + "finished_at": "2026-08-28 17:17", + "validation_state": "最终验收通过" }, { "task_id": "minimax-dsh-tp8-m16-1-78402260", @@ -371,7 +479,8 @@ "target_met": true, "passed": true, "logical_tops": 8.082003660070677, - "finished_at": "2026-08-28 17:17" + "finished_at": "2026-08-28 17:17", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", @@ -393,7 +502,8 @@ "target_met": true, "passed": true, "logical_tops": 74.24514869949742, - "finished_at": "2026-08-23 18:45" + "finished_at": "2026-08-23 18:45", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", @@ -415,7 +525,8 @@ "target_met": true, "passed": true, "logical_tops": 118.31698222866211, - "finished_at": "2026-08-23 18:45" + "finished_at": "2026-08-23 18:45", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", @@ -437,7 +548,8 @@ "target_met": true, "passed": true, "logical_tops": 136.98508169710212, - "finished_at": "2026-08-23 18:45" + "finished_at": "2026-08-23 18:45", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", @@ -459,7 +571,8 @@ "target_met": true, "passed": true, "logical_tops": 80.33419874225937, - "finished_at": "2026-08-23 18:45" + "finished_at": "2026-08-23 18:45", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp4-m4096-1-0c2f84a9", @@ -481,7 +594,8 @@ "target_met": true, "passed": true, "logical_tops": 60.95455267843923, - "finished_at": "2026-08-23 18:45" + "finished_at": "2026-08-23 18:45", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", @@ -503,7 +617,8 @@ "target_met": true, "passed": true, "logical_tops": 135.68637642417067, - "finished_at": "2026-08-29 03:28" + "finished_at": "2026-08-29 03:28", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", @@ -525,7 +640,8 @@ "target_met": true, "passed": true, "logical_tops": 148.82460170966874, - "finished_at": "2026-08-29 03:28" + "finished_at": "2026-08-29 03:28", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", @@ -547,7 +663,8 @@ "target_met": true, "passed": true, "logical_tops": 130.64379245024566, - "finished_at": "2026-08-29 03:28" + "finished_at": "2026-08-29 03:28", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", @@ -569,7 +686,8 @@ "target_met": true, "passed": true, "logical_tops": 123.91928365478229, - "finished_at": "2026-08-29 03:28" + "finished_at": "2026-08-29 03:28", + "validation_state": "最终验收通过" }, { "task_id": "minimaxm3-dsh-tp8-m4096-1-b0482833", @@ -591,7 +709,8 @@ "target_met": true, "passed": true, "logical_tops": 62.14396076195009, - "finished_at": "2026-08-29 03:28" + "finished_at": "2026-08-29 03:28", + "validation_state": "最终验收通过" }, { "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", @@ -613,7 +732,8 @@ "target_met": true, "passed": true, "logical_tops": 140.78913965554716, - "finished_at": "2026-08-31 11:23" + "finished_at": "2026-08-31 11:23", + "validation_state": "最终验收通过" }, { "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", @@ -635,7 +755,8 @@ "target_met": true, "passed": true, "logical_tops": 159.48946182998893, - "finished_at": "2026-08-31 11:23" + "finished_at": "2026-08-31 11:23", + "validation_state": "最终验收通过" }, { "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", @@ -657,7 +778,8 @@ "target_met": true, "passed": true, "logical_tops": 111.55233090221004, - "finished_at": "2026-08-31 11:23" + "finished_at": "2026-08-31 11:23", + "validation_state": "最终验收通过" }, { "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", @@ -679,7 +801,8 @@ "target_met": true, "passed": true, "logical_tops": 136.3839562001365, - "finished_at": "2026-08-31 11:23" + "finished_at": "2026-08-31 11:23", + "validation_state": "最终验收通过" }, { "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", @@ -701,7 +824,8 @@ "target_met": true, "passed": true, "logical_tops": 94.49491855548743, - "finished_at": "2026-08-31 11:23" + "finished_at": "2026-08-31 11:23", + "validation_state": "最终验收通过" }, { "task_id": "glm5-2-dsh-tp8-m4096-1-e6a280a2", @@ -723,6 +847,7 @@ "target_met": true, "passed": true, "logical_tops": 47.089410050974806, - "finished_at": "2026-08-31 11:23" + "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 index f97cf0db2c976a0ac6e277f77a96a236065ee91a..dde8c8a40ba55a5e864b4ee6a4ec2e6ead9fcdfd 100644 GIT binary patch delta 15039 zcmZ8|V|ZoFvu|wMwrx#p+qUf;+nm@=Cbm6sCf3BZjhpwq_ndS8XMfmh@AdpvS9NuD zRae!j{XpQwSYUWX8Bj1ZARr(pAd@_XMEE}B?~C50Cab~K_f1+5ARy%A9ey;xRFaa+ zAQRGPZ`6Kwck%py@G|BwC3t_}u8yR=j%IS$e(~Q{&KjNYW5{b}kFWPj564&pbm!Fm_S z6>e^d&f+j5s_@0^30^2qfWLRnGJUl|ao5N@t(H_-0Iqw+drJpkMWCT&S55EQjo7|b z_;+XX{=?}0>);j5xwnbOBO$H<*gAy&U$Hmv{}78_M&mUw9|XcbqM^#qQ)b2j0$OQL zz9Pj2Ff(}C*>>yMxZt)XebCa`*^NI16Lk@$T~uc6Z%Wcek%*YzDT(Q%)Wnj}pxIHr zonhubCv^A#IsR6wylghZXPbuO3g_pCMR&g{iE&Pc=~HjX_5teKB5A0?gb<1a@zpCAc0=|5`(->&Un z-TZxA1W#so;tj4YM;dzp*&Fh`y@E)PRefdPuHe6 zXEztlc0gCgriXd**%Ra!ngc*Z0A7>m^UeJB<7Rb>_e4=Zz6?-IC||otcpZ%P>OTi~ z@p*RiY5ZBunkV2&xqO#L(b~$K&vzUnIOWhI_lo@ZxxMEqkhfW~C6Re|!}mG=VEfve z{5G7Oye>5)&VmT_j%JF@hQF65GXQzECsmyLeVt=kWv`wlpILsdTso-V76Pn>dgGDz3Lhrlc+vIGmSq0L#-q)Fc#U_+moCR~l7)3<2P=UK1q0E`Lf^7?H( z46f_^$CIu~4AXHjE4+G1uP@Yrph4R^aH;3!RfoPusA}u}C8=sL-SGjJ5xHI4?}dz0 z;k6DaJ#X~xCfuX(ytV}X|znRmp5}1dfovLFF`}7~J%0ap~GOSK66jU%l z`=4AMRNWWm=a9M4?n*HjwTr{X+Jv&ZC$@a34n6z>Zt{vE`U@XXYWkDJp;52bhk_ME z!MmZ>dYajiHjBb`T}Oui&BT6xdYN;PD`@5#oW-dxBqsIH0h{?zUe3Q#}X3nOFM-Gv^`iL%Q|y z_iH|Ri^(lp{DYkG)aQ^oL_0*1Q)Elv*h69~7f`iN`EnQUkdD}R@};;%!PQd782lAP zJg51I;wtR84=lv{EHb<32e zi$5&Z5kr_cvFHWpilKG8PrC(LMoi8KXxF4>|Xg&?rRLYtP44S5eZksz{BmpY#ht4#BC9CG{) zOs2JMriTn{4+jE|E^*3BsPVChkcaFU8Z#An4)*}FT-@&#`LSHUyGFn}-s>MmUcVWs zjAuCH2fRoT^u|~uS53_9z3ktfi4>vjSS6qtMKlVT%0snqfsUh2(2dKzuL#IFu zh8T_viku#1(4|Tmlat6>htefhyP`=F$*@gpB^>J+W}DG)5l#5TIwKT;xX6QS3$@Be zXvgCjJLk6d^$W+|!m1TXzCih?k~ON6HIDjH0SK$~6HmF!(8tj33f{>kqK~X!KYEnn zuJI(}pm4ZY2Z#=U*XQ}xtMRZ?uzpiXB%ET0fm38dh<139$gnHNS~m$4peu{ABoR=_ z#s^8+0pGyxvN08*x9yzbl6*x(7Ng)~*9BVT5jOuL5T%zhs$O|^KBQPClB z>pNHoWU<0b$uh4(ExCp^*}`^?n_Gy=kicop!(w15Jx+zvQq`2Wx_rrITW_#gXNnle zlSyz08svEb&}osaLNrmQqo0X*6p}3fx1a^@eWXTIlD~e|<4~A;v z5-UYC?aegp#ykx*GzX+hRmzE4HQ#?bQt~R>eMLn_GF3r&`iTkyQpz$G$--gj`+fR@ z!iH^P?sb_7^C=F-QfwsAqg16zBz6inXi_MQ@)stlHT$l33*(A+{eD5!#q9&uj-)Qa z5up7ynJ&n=Sc~3C7A2Jh1_;QfJW~-KR|Y0lIal(?kBVC= z-ozg`hrR*r$qqTZ%>>*-=OKUy?vlW8-tk(ESj7X#BcIM9;WNeQNe3cd_G7w>h~rw|Z~aT%ja2qP^ZUR5=Fo5z1`YJUx z_8vD3HT_;)o$C)EdJDj&vu2<+hyhb7Vzwa##Wy;Z~vY3Y>%13UOrH-46iuow$$&YNV>{4h(l7 zTeeduU#iS@G@KaZucR6jm=w!A;R_~rNhLJroveVz(g9~RD^-CZ0e0)XbByMtOX=9P z^v=#C4XL6!9~>)H|HAPGG-hOSO{~z%L-|VK91P2&2bj*>P&Z)~Xu{e&rCHBCG&mXs zbSU=4GT*aTRZswX$xizzJ5g*9#3L+Bg^B!bK~6m*)UsQb+oLQsu-r~zH*MFqIEEF6 zuE-EDy;LQ3a=iI|;;G8}nJ^K+?@)A9 z%5vJ>QQ<2vBkTJhRvY3^@L~{ieC=-*QD7min zQ!v#1&>p@Ax_sp$btJhU6tPlFRna}qXLJ+wW|P;XV@wg_y9W9I3%ZSfoE9Du4zzOh z7F7AkVlF9{zspVXGwWViu99+f@jL!4DP!do%07~aQmcvgZh-HO=nFj}EW4;JX zee3Caw8k%4mC~dojG%yhVr`}t4qNbzq4iYZwC3BJX)3i`_GhnIeQp<9o$XFtI-i*O zhc5dHWo*+oU9}9znV`_p@g+I20kK>FzVF{4CU@S_S#?wp%z4kz#97tVjzt7v6^UYr zZ4gxLs=60MX#^(*<4xl~jcfw~_#lmrMH7bLlN4NLlL?tjYrz#3e)U!cE6@n`OsMvc- zLA4gs;3AX1#7v@N5JlAiTYt*HE4_C5A{nZ<_e&IOc zpnGi2g@_8hZCK_f5vZj?8xHhev<5NRjOtSAM>y(4S~@J*3%W{PyHr?#s#Zuhwp>9y zanQ3W$@#Olwydl)i*0(Z3YXjhfLk>xyo0<_2LvpOm0v2B(&wK)Pj3%Dt#38D?mypy z(dhZGG3Y^pVbUnUqEKq^;Wj{U)8pS;V)$F1H5c{z6wr;Gn zn{DQ`cPF5O*wdTpAWgz#(D8P{s<5ZmMy|9Sk&PIN@f48k+`0u?_FRw|dB%CjLFf)y zfH92gyj%9KcY$E*DtI&KEb47r@*uplQ)ioM7k(xWSsBh1T~MH8m6XVulRihtwmdnv zejfvlhkYI+B>1UlYc{X}el#H~Wde598!T5-wgn0=0qbWf{QX@5sd|4yD`;f4hD0^1 znIIhY$%RDY)QEDHX$H2Y*p=Q&18!J^{Mb9dIoc>5bgS3cql(%qJd>eHqQP(w>|!y& zW?6nIGNMVfmBVY?y!Sh#1a9;ncMq-*i~V7x%-1tf)oi2r@Mg~fe9yh>d+w{ky|=)v z%eCK2@j=7e1N-gr<|DQ0JdNN|t<*FJ6kucScEy_7)Mfjs?qmAAFyPL6|Ij;DB>l4M z7)S1!W6%M%0|k?Jq`<2V)KjD+Lw`{m7vt{2zoUx8{f*z@e>MAvdv_P(GsgMV&!S;; zQMquo7_;eYat}!JDm%PIb$_aUIvBtFkka?zFMU1u(1hyf4pK|xGf3FDr?FJ5cLrP! zxUOoTs(ST6r-&k^DnOy7`dt{lriqr>8pW6$*Y9Bi$NugM3X3i402gncdry-)mQ+@+ z>)ZJ}Ab@sbmf3TZ*#HMl&6YxNfO<fw^9RBHk~iRZ9AI)rF%_S%Fs_h{?(tTlvKCYYG3y>>Y$f?S}+Z1GD3Se zM2{&fwu4M=BZHf}DMG*E+t$0p21ehAzBdPcUa0S*(&$^B0)NLTA?1b`0dMksqAK6E z+Jd3kQvpy?U;GTqure?<@aXD(kmY&LOnbB&TSvr40Ta`;1<%8choRLhv8!Fm(zVY% zVsTKjtFU{?U$@agACtDDeWGWSO|Ljy#+m{UD=z3!h$ByN!K&UVi~p<--aDCwaknDQ zpU&Lh5SP+UrOc^5aac@EHa1PtKwBK4$~Ru>vj^;9f+>f3vJ?kaifIm~E}!7ja0#v# zguft|gri9m*rFK)+cVSRW*iltE%Ibc#1`ls=G4-c&=R`(oK%u%QR^K{ABLV-!+xmQ zO+%yYlo2VaO1A`0>7m$PZ)XpW+^~XV;?QA7I&f5VU7Ai+`+dnr<)+G7?UxfJR_su0 z&;tB^5Sg(^J01t7i$>oV!tIqmy$ZEl1#J2sClDR)@z>9l?;5Ov@>5#YE`jOU>)x?k zRoYQoeoFvW`&$A(XgS3gWi`}ZYONFG^UJH@ie_WCtO8y1C!yZj94Qo~u){_W&18jV z#rHu9XqeZ3S%8emu^Mgjc6h@}R-t*r0EDt}^7NU4Z&?RU2xbc8W6ihybHClLe~ZYS zYgEwzU4E16KX!uMr}^{5*<>_b&HjmZyR(bzIr8?zk9B(n$&P`_<7_z5} zuYjQo%0X5v0!tpt5D`4O!T|hiZn|EU#P3%da-AQ*kLHKLb|7cB|Na11-|4?vD2=c%hGzGTDu4-mfm+Pu98q%{TW|G!yS>-WpIH&6|sf z)$_b+kB{Hd=vwNowBQXSx#a!Uj!QpJ8;V&H`)uZcL zB<7?6Lum{iEeAC;z;WfaIhZ*p5BhITMKHM zwK#wOF`)K#U7wZX(!DkI7pqtQNTGE+_iMm4H6A1)k_x%X=>wIBM2vQJ+6L=SrJg6_ zJZ)8en9^(?*d)Na16+XK_lwNKKyr2p?A3pvI{@s z?5A{d@E0JwUL>hll`emoG!8~av&6c#1ysn+9jKfp_@kw=7 z$QVOID&#?rCNBj>+X~t;2_&{98WZqr+ew%&K0{?4^CVoOys;#uZtgGr0l)XZI0y+X4c*4OCQS#A6cii}#0Q)86c1r(9u z+`!ui02M^c%$)ZeN>Oye#i2$r)S1cz3rDpJ-c6CV16=soOGtMxG_3zlr_2bAai%Re zx0WDM_9AVE6d0tK0CA~tH(ny?3(@2Nh^n0j$E|{r#8N-6T>N-bSV3k5r zQ4z{#;RrvN6?4*nabDMTC{2M~cMx@R4QfNcLE3d_`)Akt3Z1T+KL8~3qLxyvd>&-A z%MtmPM@?6xH<&?}h|v65dT5XX-312Jj`5dtoBC0+d>xd)08@=(d_6?52z?|2RB`hp z_ex}UZJpC*sh^UzaQLNAv&Z8<{uV0AwDHE#1c5rP z!Y-Q`WwG-V?nA5Jn_Z-L4eM2H(-p8TUU!S!idxOmDcw)gTFE$>>UF4jutBM)Y6n|X zzHbeA)I+k?XmezI2<(?>0mjyD92|1R|BoOn;%NKV4!uZ+N0?@eWWIAUvMr3^3sFK0 zkH-CS$uMu*de$xlcro0H7qPGZC{B>o88S`Gx7MnWtJaz?Xy6w0jK2Y_)}hqCS3_9s zdo@%K1fD|MtqEDwx$mLT#&hMNkkrQ4`U`67;Pqgv*jQ(w+{Az67!&Z1oPEsRjQ;@YaghnG|}SQBzdDeDSfPZocP<5Smz=?$*6d7^3~Nk zGX>6PXFhd2I~bP$qOx17_n)wq=`dykG9_iv>!Fe6@#_k(DCqspk}(M&m>DJ9jP0E+ z`|4@WQ83S0-RF5&0c26(`FTDuee!V=&IoL62)f%-E+@iUFyscUP8~6x5gyKmQ(Zyk z#LvI@dbd8GrH?&Zv&aNIznz`9_Pi30@JSPH51$i$BK~JJ0;Kh&rUwEDXcHGG`I;09 z!17NGq9(&ifm8x1>?`Vk>dztQg+OKPLBbrS2$~6Jc$>^&)pQG{3d9ti&#J49iZXaH(Nr%h7WEsi6ho#SBT#OV_ zk@3TS4Fza*opmB9Rs&QjA(F{*9^(-Jl^&SGe|~%_nx8D_p>@9unkJ#mTuKmf`voqAli z?nqzKjo zx%DynaDW}a_>ih?gxDTQ>RN38=LPlQeYuhBAgfUOEg*|f3oRgjpsC64t2=a4=XSKUW`!<@RJ?iKhFP6=o8v10r)H}Y`O#uN|paG@1}!& z=MppWx?gA1Wi@-c84s-g;Pe#XuF_|s82-G13`W7dT?)<`Ie9}U+tOnOxUIVFn0rOP ziiE$O7K+{1yWj;;E4k@F)x4`|n?362`km&Y7zh;BJR~#Up3;_(x?zfJ2^G~6SE#V4h~*|lQ;uitqtVff;uem(?8N`@9v16v9GR^p{5r`s1@luaKOPXC6D z`!IlR#-_&IbkpSkf`P{>nBRrb=VoUIlh2V*&6D28a>;eC}<$FADY&ymAhQ(lHz^ZGT6qVCB@!qPAIfZ6zW(2h#>4=B}_kh6->Z0 zQDT{cT+C8Ndew%R)R`cLXvjm2S2!yUDtD{R_K(x(Vtv(^PBq&+IYYN!>^5{y#IJ`y z_>@|qZPxhwFm#lccaA@`*nanKhE(uxLTp6En`?n=Tcb8^?yA@srPF~i3v0*n=F!HX zh4u}@$dlayMM+r$;11&hS~c?{>Tv#&BEV!c5C8!a4oUA}@{QI!16p3T;c^is0ax z=3VCc2&1-#W6&1Ttu`DYK?n;vT!1PO3(&1j5PjxI87Tn=ph4^$HI=~2$3U2bnn5)O zFbxgEh|x)oQv+F!g&<@^qw6|180{CuncbODs5o!s9q1yX991cTpHn*PMSS)IO0^p* zU(!Ihb7>nn`>mxqp}}RO8%d^{aR(>z^t52kN&#g_U zg2yNy{%{vij6z4rZL&Bs-?{WuQwvE5qq-x=8c;3-`=N=w+iY~QW&x9=G@tw~@UR6x z^5ISzfP-sC+q0+s5ft70(x|%c6r7(mA4dTm#}w;cdo3Z%7VD+lWXByUPpi*@1N6)K zndhDTI>#Z^RSxe%7)TE7pD$m8VVXQ?!{6aUsg3;u_@gkg9YiTCSTjI2e#jtLv`nHX zbU5|fbETDFpg<-M5+R%W%1v(E zq4M>BEIsf(c;0J(o>}2rK`r;7uu2Wqv4QRIAg@hnDuxyxx2%M~466-C2>ttMU_AxE zq2yyJ>xDJpGVx|02LL}%XEC!(^2wGQXFdBy25r$kwb6ur&i#~RHr@r`)VVIP*iu`N zX$hVyjgDUg_4LmzUyX@m5SD43L~W4km)X9sxJWn|}cR>Fr!C3Nn_#fPmhUFG&dj7n(8-dmKog zv_TjvD)f;K&dVILu_jGI!8ln>{5aCG>(-`Yw11rD47Li2E$#laWL_48Uy3nKa}Hwr zvx+#i%|wqn*EHYca$2d!y*D$jBN}N6R6kHK1p8L8Vrd-?Q*c~sR}VPr9D6DSa#%$a zMTBQ@jlJL0y>ifw4EkCC=E&w72WUOA|$uO@;mX1tyAj^ZB=j)s57`&Mz!Z%$lr z$kFKJjD11p`nDgJ3T+!dNY?d6l^n6Ki)*GvD ziR7j&F0L;({CwRU{1vjwX85z8V(irCePieuD4>uPXamR*Q)n7U5|ro*O$Zn00^}MS z$Ofny6zRx!8$Lh-P+uFYCgzQj<;h2P+pVE#HET{-c~E&F5|TXG@VJD=lnpc+v(-x) zf%c6Yf%% zLAE~h;HyL7ScjPwiGt!%JL(5G#o$38AJ09aN%NEm9O2 zjl`871O;4d0+LuB{LK;yq@?m#V08$!>7XOp0zo+7E$B74)t`8Z3Evip&=9M52c8E9 zgbc1iXi3uIx_FJ7DpRA@h1Sy5O0lL3BTyh!DBikR^X#}P(ERzZFoU$976KA7jNv-0 zxyvZM{pCQIrD4<5RKC!W0Ls!o>BnkN&wy4q_rJUksS1@8;z?Dfq*M-O6QuVN<9p$% z1N<{`T70r$X+2OHQ~b4UeW`#PSSNCqexc`by!2U_Lhi5!=YWYdIDlB*g(t=ZFNs@hj*7QZYV+v7_Ad_7v ziBy@7<#E;$z0G3CmLx2UgFI1#Ag3D=1dyj9#W4y&PpY%HoGqpV`T8&(bFbew6Hh;D z3&`0f*|M-m)_Qc;@T%-|KK9HO(=}dMcSAzOBvu4+lh2P>L4oGIo2JRB>j!~o1{s{` zL`m@|XT7+2_i^bKC0az0)I0PR$yc+z_ig2$AU-hlP8S&;6RzA|7J?2#s%Chv0JyzX zM+>DuD6g3VA+0RAaqpAZ+9`(g10Y9d^VM2s9zDNuL@A}v2=LTbT3-Y3 zQdA(I?{|)XJd+I1`f0YwXVv!FHk_IM^4&p^Y#85wWWK`0Lq4@?Yr>`4=(}@Mod21 zCsfRiT15YnvsgvM_Y+Ut+82CGfWuL0*Bx)FTIOA${KPaAtS zN=08;33 zIiYSa*K^o5xxQ{NZ?fE`n_^7FKZzmva!{|(=aL|H094Cf=+Pde7U%-=Mb}>xF6hyo zJio@fRVpkiN*UNTN9xFuFe=BdM9B@xSd?VH-a&!Eid?+HnEH-dHB1a7jUJW^T1YV6 z6OokZgP=l=AR=X=he&H-a`{z#b#NML=~kayMCP7|Ije8;+j<9VLT{|rg6Eo%{Y{!UgY zqHLL*fQ~4K&wUKWrZNSSFtT;B3lT&_7%Mz##(CqbZ84mGP{(A9ja*Ji)>VK?tdk7ypeN z-7Ud2C;%NGbtw|tRrPP$leloJJ`(#Ykx5nyevdIm(+`BD!5zKt80PMP=6mckhtKqM0zXL{;UEI(=qvlJ%gg5nvR>7?CQ8;`nm_OU0Wb z!sT&rA~YioBYQ zZQ5mG7p8oP1KV^dw62Ti8*H3~8>fEZ=4X-y_MTn4sW1sw$qi_=oYt&`>zP^qJM9q6 zh~9(&4Fm*YhyTy3kIC_Z;DA&uF9#eMlrI|kLM~0_*{6G~^(2u^ia!imXTz)Skt4HM z`T;vrZ~{Fa>Vz&VkIzzRN2kAjb^5F6@{NpLxt6$v<+fz$d9Ui2t$zBlNgvU6$Icmb zE%QD%#3b9L^qV*fNNA{SZt@^d_Q^=g z2jzlU;*H|qUxT@N+5jl?wk$~hW0d@XE$F|fPh@mFwI^r&RV2}xc1w^+I7henI@Q`S2^A*`Gx*6v4uwY6$DYk<*d!K%wl1 z>6IIaw7w&D+sR+@o6H5}IC)(~Ns8P(SY|>?zPMBTkqS9v<@N|J)E|SFw#M|r6m?<7Js%C#7hBm#4NnW?*nJc+VGR_bB{+hpsqhIc{WRojy&U~(i$+femr1HvPaG7tF z6en-#8P+!cw-s2evhw?Iv--`)P=-(~-+5Npy*87(uSMOMRV#*Qd%%MmDev#;N?537YGnn?fKh85$ z(no(rs8k{W3GsjaCP{*Ew@=ku-)4~c7aOP_r+5E~UN_Ad=H0UlQa%FW!wQB(5gR$N z%5vwW!qHZhCj7BYb_sP)9{EVSz&F;8(x?$uQ{^J=Th4mqty*#8{85Za6@4irMyeV) z-A9Jz>JTX+zt9 zq5DN*$N+A{F3!+Yor=vZKQvMqD6>G4z;WX!spKyBxZ0YVKtu=)l~9?x$Q8uaGt);X z?cF{G&^ru#0!Gj?0$SIA1s&nA;q5r&CANdE8{-@vxJ~UuGfP@`11>7O9yCc0f}7M6 zHrsyqNDEGx*vVH6m9tlZp&`UJ-f<3IKF*PEs z;-+>+hYIhYuxcxHc@e52feuHd$0@2-@+8y&Qq#hRc*2a^%6$prW%5~7fzi5TaDtW2 zP9B5eoKTn!uZ6L^5nTy>QSqx=m9TNngg*?G$WT)DU$;+@ik3=n9p7A><=tG!n7WZcx%;R$q$s0M9T7V?AG=sGO>G@pvu{j}Es zfKF2mz37l`eb12*GXYiypE%$3i5jz|A*2p!bVvgPh8y^eqr1Fm+~e!%%3*v z!i^>Tj2xFom_jC#jGRj`T)LSSbQ^qVIV6JG8sm0vE=QxA0~b7pLpd~7l@bCCxE6Rz z+xFh!%jd3W(YDf@<213M)Awlw(f3a?05rTw_G&Ue;WRKX=}{{zrJV($)TwR;B77M2 z<&;EIxEy3}jzAXND_sSck;B4BL$wVB-&-xm8y>bhEXOB zV==QA^K6iN9|B~lVouqgdBaSVO2x6rxIgJ%vcn0v8WDthk`dT-g0QoZ0K(R6HghF2 zJg;;rU07Pd4wf2ruTr9mfAPD-^n5&t1Xkf|e{ z%C_9t0~WB%GW5lcMXW{c08n*vgzYjbj9y2&yEQ}M!1*7)l!OE(n;q~vX@anAJdEcq zh7OvdU1LgVtc5-t%j3k$0X0lz(bmaO+Wd|EBzo`--8i<9cp;$dKzgX~Q?41%o;zNa z-q3KVNJ*W!X@;D|DJ$as5^E+%*?r-7+en{C2-D10Ll8@7UFzjFADS5+RBDqySO?y9 zF|w(0fh9jDZi1k&LE_j}TISTUT8WT&+Gl_wF#SASiW6MqnP{>U0TdZI`gpvh$d6ZB zh?n=t9^0O|Xu5kN>mUSSlj>@M^DG<2SvF*jiiu5*K6NElqzPi|aRnn>h1l`v8QU-r za`(pWgl69ED@VHQtA-*VLQ8I7u~wlk;+A3{6VwAD1DomLKgeCK)tNLn{3zi*2Qpv{ zIA35vc1yKWIv0q704z0L=9o8o3`b4nz4~ju&n^o=>#woaRuUmSFdgl2RXwTJ>pWII zE^1Sq%W1xU_`Apt?$cZmb^0SU*`K=_`jrinmTdq0nD_A}hd=5-Auxj4rQs`L@vKj2 z>T5}MK8JW=BfRVZLeR&~j-%-f%^)Nxo3*X>pk6yBuO1#00^BR^`mehnYIU^={`xzt zh2d~5}e>LqK)|kFCT!{c7&|;Y}FiAc2VySgitB<W36K9i?r30~XSl8L92b5*8B`7B1WbM-FJ* zYuluWhDKy$N~Sc+{vaKiQnHdD03#Ztsyr-%yF4FAzcNe(h}~~iB15u=92G=8#9ixi zw62u#S9gMqr$X#aEtRU!0`8ydGa|`U9ZD;hL&Kq#dpwl#Oj)X`Hl^e3Cz>(n0*c^r zZ{`{qedDcTbQY%KB`4aFQfKB9!!D}}iz-7R8XOyPMKChh&RBcq8> zyi0T3D2C#fUWAd3!9IcP@X590sJoC@FByc@b#|LNdZ3y*X6tMj_L=vC4 zkGop#!BJb%A66#9RCZs>0C8*g8qMj}mfk5Hke?4m7Jdt@_WMnhLzl|eYFbQ~?|%?1 zP&-D!Z7QB~ZXhOHt-xdxx6`Pn()po_EwczyLyoWJLzMlV$SzK=U2wWZDXTLJsYXP-^kbsd8ARmoGEiR5lcb|>JyNVRpSf-H!&07Ji!JQ98Cg1<9i z++m5^t%Zy5>!%;9@7(0irTYViic7pN_}OpTRzRFZ)Gr7gG=JJ*F>o2!7_&ArEwE_` z@xEfb}4h_BbA`N2ttWjPb_H zq%MMAZcI?_ZocC!CA3jZ@%e=bqTySN(d(nuFegy3Lr6Y;`sHWaJ@i{od_$L4%}~R{ z-P8!Ns3RP<&n{pRsrSa%7ho_BSj~$7KrdlAwID^5p*pj@#0ty^yHt=>KE?+-&Zwlj zvFH0;a7_}GG)gaot3QNNh?ODxE=~-vV;`H6=~KVYn^N0W*}x@XQR9E^$YeyWL#Thc zyIQo$+mb)H@BuUS^Fl}?2E)~K5287}=&`$!5isRKNPp&Rn(Id&%Fpy{Jd#s#y{q!U zAA;sUi5>!mN`T4QknFppO+4Y2I3QYoAPbLM@v|;HeVDU#UKAvB@)Nt)bVoQHBq^S? z9dM1q5*J!)4fV;9Be4{&g*cBY)q|Z@Aqh!LR;ba?$^%3bljS>SPo>N7px|@G`{h{V zY>D9(^Sp7iK~+~o!d@d~8XO>oq7qC(-nKi7h3U|vcJzpy3>@WF!zjT`7%C1!=@W3e zGw#J<+NFGGCM?%{MByW3)M><-geqJgB}x_G_kbp6ycU*sIauJGW3Agr#RSEBlE;5C zWGmyY;sK&q35{hwC$Pkj@*!;o&LsNW1SFv8!1rl}G?Am{p>&gc`D76&mXkFgF%1+T zRY?wb>}H;8-n^>Tn9?99bu3|=|IS!qO1i(Zd3kX7x!l_|txhq^7&dboHiriL5yqDu zQL-vNa4SB5i*&yy@DWnhZc!M>Gf*U-G&Z7Wm;ewoOzL&#@03_R-}gSs_V2|a%elkL zGyVn{Klni<1b_9F>&5wF+g>pf&^vtf$Tvx(Dw&XU*M;V^yAZz$T)P(W$^dj;OKcwH zI@cNq*B@~0%J%QJN*8HM8{c$2Hybr-HF?t4IU4jev`-SLmvN>}E4@r#WB#(Vhrjpa z`=j0Xqd6EHQT~VBPNv~~2mRKPMruAC;D3H3hx7CPt2Y1iOD2!=bN+i1OhDj&cM2*B za1s5ZUjJX^{@;Awx9$J;6A&0llzdHvp4=^f{4emoiyr?k@%KMKKgxfBpo0Gw_`frR z{{xJo`v&49%L~H&JK}#u?{@hAHt+iCWx)~U)u)!yH^wd4g`ewUk`j=1AsRK{&jC24$d~l4h}Z|I5aN1r1{{q;>+?d;9poWq&=92%o(c49GuEl<&h(V+Ken;&Wp z=aQ0b<#(}BBYM&ih)V8Eq;^TD!ej`%q;h*WnHP!BAZmVj1tt!8J_KxfP?%9<)n&>) zYaTH+qZb8s=O=<#I1MZ}AWuXN^`rr6X&U|A8;`dY)B?wZ`MM*}hUD<%=gJ1Chm4{M zQ@Jz2c>QEslNh3-Qpb-@hkSc-6N(Zo6<1zg{l?S%Wslp|t%mOMObr!*qQ+Xi0Hf>X zfS(3{r$3=_YA<=8_17A8m)pxde|tdw!xRP?`q_;)puA8>{}?1L`E*kU69nWc6PQno z3t(aNvbU=?uyw_6Nq(oJw?Ef>2qEqas<+enL;k3%tKt{BF)4*3X>wnPq)*Ylv|X=&2>j zJ&cLiZ0LS=xs%M37A~Y zxz~&t3Gi=xDemro2R%@6d@$$uG)w?H8<;&Bvia4T(1Xmb)u{1QnRT%Hc@p9N^WK24 zE@*-bObB`vB5wKVDEw;n!iTi{2Vj%+bJ@puK-fQ5V6lh(BH*V{t%DH3bK$1DKijKE z%ZvT>UezW-TGVwn{YCE5=lQ$WFd&@e*#G$fW5o^+8tU;j;^=ewV-4YE9I7)vXoEd5 zFUamh#2$}&rk8X39;(&^S^;sQfp0ooHKSZ|dTB*I;&(!Ok>KO0XF*}w!a9K4_SgnE_|;p0wM<@>iaKFPsy(^I(G99M_7RDo9E^Ks!7$b*!O7NAW#Ipd0jH0lW zOWCO|4=IsXwQiwu1*4^IK+xHJ`ak=@RxY_YL~1%jOp{iHN_)jVryjQw;%~JRAWTP< zLt=K0MX}dKON>(;SG#vnnHMqq5(DSF7+OgM2tbvbk8eviK5ZX><7cqe1;JGnq4n8H zde9U)(Q9i{Ya&hzaP_CJwFmr6w}gQriHb{ zc=v~vllQy(!&=#w&1d5i-=0^1mx}=P9y_5uw@PKdE6gtd66zz#me9!~)s|5W=&|C* z+urNqp38Ml%ER7De>c(xtXNU~8^_<+?ca1L$NUPeT=6MRLR^bXyS`K6_w6C@vB~i? z!i7IeuFa|^wfOCPyB88FUtY-N=QN&X^zBEkX(tv&Z)K+(juKMJlEt%A0>=Uez3B`G z!XaXx;Uw_^0s|A{qj#7wqVIUz0Z?nhzG_a2wOBpw4m>2kWhh%+YTFd%8`0GVzPSSE z?Unv|mdO4MUONufjc{`;fIrqgXjdMlDTkjX@e{?K%b~~90F{!g5O#}qcw$3H*7wbW zvrd-L@N&@YigaXe1oJCmWR}OV2xS$mtzc|CRk1{ zGfDb2;41mWni4-JQJX4KgaHLHITh&!Qaogs2R^xoWLT$!7Fyi3Xx7JGeJC3U!o8FO zjGAJgO3${e*)=XGIgYZ)$F?ZNK16hl`Vyn-A2Hy8lN56~c7kimF+|95ra`a5rXn-F zI*?E_$Sb*}tCU0ic@i2>^?~Y4%vGlS)&Sw6pOg)=aEoP_9iR*-aLGw`+aE;Ab#>*? zgAs$w%!QaMvatyD9ueg7u!vB0cI`)4|Jp(Qf=gjq4ZggNue@&6+>ZhM8vTPF?|xpE z0MOmmt&*+gc|3g_c@Ugr?twYDd&0`vvBpyPeZXQq)=jH?4gwY5x{nHhG#VX1ku)_g z4vaRre$aQd@1KNb2M5oopO67ZD+tJxmx&@ZG?8SIT}6P&JCMggRt+M7qD57O0v|T{ zK6_K$>N+8f1cYDbY5Tk361)nF`lx|tBk7aokxAma(HR3TR%p?hW6^r1x~a)FhdRzd zHeRL9rq8)3?YXGmOiXGvL=6haF+ovjm7bTU4IjXN8}sd>MU>0LIc_uU4vtQu1lAg0 z{FNX{1r38YEmMvT3Q}b&HP-Rmk9-M(FNz2sQNh6E>o z3g=@I80$9Z()XK{4in+00_&Us<`Qe9L!FN?N{vMd1h0~sVsyj^zSk*$ANZ{VABUR4 z5~+mZo&87JeJmIW%|1t-Q!ws9uQ3vfvNL5|i0JO@7$co^`Yq-JpWjhuq3bxT1(o>k z5robSYoBi2-_pN;Nhg!H@n2KojB#6(^i0W7Bq~>{ zN%ASY7z~e)BFaY8D+~~y`@V+33x`pZvKR;!t}m{gIF1ZUY+0dtr6Vh`z(=dWQfzW; zsqd~N+0K=N0--QO3(Yoe?bEMY7P-xb8}$cEK=B%%0d=JFCFJE`km_Nz`yY zNvPZJP-5+tTX(f@?Q>dQs@w+}H3tP5jG`z4Goe8H2)987bQRFFSbG^lOfnWy^cEaw zgL;Dr>JR&lkYCas;i&aU{8ZK}g%3?st}Jbh8vz+d0@{@AJ#KPf2JUySU3^1APRpFY zHN(gw{aAt?oTn_J%OGD;k}+wl5F{(1 zaoD$#!syUZ>he1cG(qgy&S40K_}i6323q~5~#ODf|6 z&jj)c79ez~i5(BX8TM1rQO+90!ICe^>TpEe#LHnN0r>`FIx1zi2om||<=P9AZ(wHOi)ke` zRlz6##urYxJK36;{b8dysS~73Fp%!Su$Os+eduN$$jvIXsz7D zCVbEbp;vC<&hsYnc>?a8LmsjkNOd`?1gU22F4D?yaix&K4W7BtO?>A8AIv>J-+y6a zqEZZbPZq}_L5r%+fij8pj-y;PXzPr@O1kR+(8+&6uJa`|7j?*{DZufUVT4m>8mC~^ zRe0v+XQp9{6^@EIoaN^k=fUX+Ml3x9gd^6Z+gTv0D>Q7fW+s*;YK2149w()s^hSn2 zu?2PIs)iwAN38eM|Gf}D-*gkba-C&IpeUM&Iq~i9ikaV_C-G&gknULGm6ql;_JySZ zP^G7#oIQ)9C12P)>iiVYH<~1>Iqm89zdj-=D!{M%g*wK_^eI`@qxs3a|!WbfE#gvaO%^0;UWc4Mglz!a+@RrGc36e@fm9B|vhao`=yr5Od9jdP{Ge$^@ zVWUM;$rk9HriCFYwOqm}qrt&$lF~t*w8mv^O}kN5i{7d!Jm!8=kLZ;tq3}CR!-Jv} z@FAzKGf^MY<`z+MYwytM_}aC1p=O}STb+{=s`%F{__|-4a!G`L4s<<6iU-iwmzy`m4f<9d*Z`-HR#35oDi<$K$O<;wQUs2s)d!w2D{ zLRc5_0w$UjchyfM*gSK6lB`bR*y8X?<^bi35k0HYiG7e$I*C~}syuAx%#2pf-3G3n z=qB`!Pd@wxvBtiMhVIlH*c)R;8aMnG6KlLI)8=%?rQo_E_D`A$9#CO zBOUWRwvL}Wszr46n~P*#637`d%4*DtPLu>y9SG&1QDGD*RVa*V=jl?%)7oy2y&nJa z@o?AvD?r)0cm`RIn88=9N_u^da~|-y>%Mg^_LJ#m_*6$?OpezI;Ok98>fJ`!)-uW| zEx$>Wb18q2WV z2@19df>l{BCAuTKRy7sYko@an9J1MOv|`j6rNxO&5YRoq^1Oy8^VkeCN zeY8Z;W=)B}U*<9s*d*K6z|6ren7bFwihqN$PT5(u2;&xI_+py6<|TVo*e*11y+oWS>F`$LhC%fxekj~+x?)d3Mb zW`R_9A#10#cnut*zw0(qR<`U{h}iU;LZNlPOtXR?l3|u7y-IVP++_AKP|4 zj@r|WW4(qj6#1b$($Mjv#djW7L>(e`X`ryO1dx%j4lBR7can2Q6STVSZJwmpT;1DQ zU)&P=J->I~zANKmEhShE>BqZxUfCbDq^>_ zrkPKzZWv}sf-7O7GJi=&zVyws7l}>@i#U`! z`?Z*=IYFv)r|rQ7?)k+JFmOYAW}IG{<^YTCwaBqeMEaPCW_yW)IKn=2a(D@fOSmaF zEnR#@-nTSi6C-BOg5)f4zdkg4z5z^h zzW3SF@i?BLGkz(|<6Y}i(oq2BJnP){^|Y;IH~Mj>th@3RlQjg#VD4^FXS{ye+~s}` zSJE4Cids=#fY>A0ln#L+e_4tl0d*7U?rfR#rLvoG?}U3r^Rnsx(ii0)eV=qjXCfRj zGkQ2p8$A)no5Ze8|8()BNm~PK_<*_7zs=m(P=)T95P`+CK+BQuNrT4FjLGS4MqFFD z{YOHYXcm)bYWuj9%5(kaGhfcoB@>)~0IkEUx#fX9pMkuB6IhxaUC$xf3>f!e$821G_OgIZoXB{P4zIU zCdaQR%HEv86_W$D9ZG8Jf+?Gde&5UG^$UQoVSt4p74uIgD#*cdwym{VMF8@XQ6g@K zLuR+LwBHeOEJ#-4Y%-FS{n%Zh1k(uFi3cc1DM+R^5ZmxJ0hiSMingstdzyVCHfHXD z2akcx=+S2UP``1ax<%V2ZuAN$eaCBM`Ci-0WOz;aI7S=k)PZlu9U0wdMs0@s=M>7~ z0%5``)}x@ChGBmyCHb(I?5wali1TJMwPd`EJ-&<}M6y;P6H_Wijz7R3!PJf{{k^Fr z&ylDg@l=ZjunjK*$7gQf_pHl<<670O=lxH>0I`<^UWUUpDX)Q#xvJ-_vDUOKAu)%# zUM`($kRQFDur_(8@gyDjH!>L`R@8e07nF323U|#~qEAwTrf7B|ua>+ZPmv|zTl?MD zf{62|q&D~oowPvidIgg%0$x{uw|`e`P(faWc$hR zr`fXSMa8oFP6>`?^apTu#_t%1=}BH ztgiYWX7r69OFY|qANi;9Pu)D<6uD|?#I$s6-n>jr&r#4|1t4w7$GD_vI_Ygbq{YmY za&cItJ#rVfbICP&mjd*KvQpmnb^+aEjcq?uMsET*<;F`N1)07n@>boxZM%B-YK1{} zeO!}v-D7vZ*s{><)U-Uf04Lx4GB9t8?&g70d=DPqL;UTpJP(P|@RwzN@7t>Q!E+5{ z0g=ujrd@yX8_?fW#iyzd?E-&l68YtZ2yM#DozlKuP#Js7J?pA#u)uE(&;lx(f~Dpg zo|dOi(*+-m;ca!!aB8#J6~&@I*-y`_=3!&jIs*)!b-HTnl#J$!Kb7~58nR9vrmo?j9jL#>WWs@nqscseSUG)Hzq+RoxBhLF2jPUe5&&5S*~kZx{5c)B z^rl1THJO%Ro%)b4zDq(ABJc!{&2gZK6hyJguf8dv(Z7rfQu*Q7<22C37p%mH%jW}* zoznOMY9>qY9T(o{oleW;*BQ3J>2p%Y%St@sOWKX?!4TI#iMb7Kf=+Z2nJO$4t`WC47A;gN) z84!~h7%YeB@l(g6Nt_?q==k|&dK7Zk;_d~KR^t-H5^UxDw0Zhl%4EC8)7MEGxq`m` zHm{AE+u4>Q$m;+E>=6Cnd^X+EY5K0~_veu9YB6kTwJUq&LR+=81EDgPK`OsK;BtL$ z^KnvFC&NznuYJt?aMT7n%b*oy(SjnJ6ybV!HVMj5q}!AVRH?{HWxM?I0G{U?#rl zf#^BUbRUpA5r-lEB3FC>ZZ(JH&jAiqvmbj0mb0n3>7!mYr{Z0B6tO0?WnN0_q04KQp_Ki&Fd#?8L0F6>5MU;7@mkJw*(I@X*b=y=D6#}rc_q3| z6$avqsA*A{00^ak7&I3J*k_k0BE={yDIx0X0_EQx`H6g@s6bItiZd8?Wb5)eeKs)( zaX5$q>U=OfR%*!9_#xgP;`Ee3T@oBKNjQhX0wQ9RDIQSU8Rla0Vs=VMtT0y6suD@= zr}6^%ly{i%ZxgVY?^s{_hIwKMD-lCxJcJoJs4*D{sO(J`W1>sJtuSCzuhge5GG%93 z3OaP5!qnoofP%_be$Q*~S4AqGAS)J>BF-noft0AC4trm&HzY>+s~Rc`3W7Tfigmy$ zPOKrR z6M6~}pxuDyF5X3)%Co}L7W!_%28yhlbpnUVG=LaIv8Wy`oO6$$tS)NFIbN9EDlWJl znG!vny(tolv$w#hKMsaGG@SHTEU_5GuJqL;IR`U6PPyrEZvnO3)m_Vz$M@<+`wE}phwgAAZl8-TZ&y#t#kYe zht!$UQX&>WkKlRM(J6k#fH3Mc)_xH+1_LcNHky_@?#fm8!U-aX8UIlA*tk8SPi4#s zkSriY+#mGdG-M2e%1;SnitHnneK6O12XA&5A8Cvtlg>HAEbO%{6jzdsn%039o!3?i zGq>)Ao3g1jq?4xA9c!2z{-a{AsdMshoc!@C#T^x-%BaPJuEnKMUuMK#y=k%Z!&9*9 z+MGG?vQ$DvU@O#9<+Nq!{qz4-^3eXLs%a1hi+eBv0Rj>Q%%{f#u>DiKJJ4QtSeHb3 z#|#F&QBU&iAx{n*#%>ahgR5T`glqbhtgTigTEo!sAwv93sWhTZ=z?A8q@|tIqBuGYveu`SQr-p;`*PJ>Twq}6! zDOs;*KBgEqvm085n>_Y$nJ2Ojr}CUHzu;@*vAF5K&DzEg?fZCjYl>Dd<2|Kt<<~Vi z$E#@^bIxR8JuROWB!mVvnxjB8i7-LAV{c+#@4Uw>mj%%F`b8hpBX2muo&lj?vR}i`e0|QS@+WWD%fq3&;P5GL?s}L1-?_KuHu63cgYP8|!P^%KLZSn$}xS zt?wG|_-A-}{r%e)`N6AA#_4^OlUcSm3of0|y?n?QEK|DuA}gWmlF(}p(|3?u-( z!s-HSb-KCc9T?EyztkLy6IzeM7J# ze!qU_6D`Le8mWhDaiH(EHSoi86G_zUZ<_&I)*6NJ)T#;N-M_J0p%`e$#mtXwryI`h zO@%AAT~;O5>^X`6%SMJq*{(p!``rWMbUlpV5Rd~fAaZi$>EzEQ9-)zlP4qF;W_Q3H zK7Vttn7lu655*arp$T@2CRUL4A>Y5rXN`IO@jBwox|KHcGXi6_)r!Bh`-yvK@5vbu zgb&|rNn#z!K^99U4XY+^L^ib*+xnM`Su9B%PmO{5NIrQvay)23Yc*WzJ47|WjZLkP zF1!hVvx9T7_woVJ)l~R{ljSK1r}y++#Sp)Wx6jrr|B;?MK79$#L>6b4Kqg_Ci;VFF zCf0e7#|@?km1HxU5t{V_N>gWqmaofYNJUmhslu?`vWkjqjC3evUtd^BS(+tbPEa&C zQic|^zjYA9yT9Wv8GBS6m^)cPRN9QNfGRWq?HlQRgz>xlx(Vl+F$*d4?_Ih#o4+_N zNX@q|&KC&P<@63$aj4 zHoa=GLLq#(SP}7iq)vs$ zK5wIw|ICdCMce!c$lZ_e&SaVZ((Lp)Ny9d18E;YsQ2OpdQ%KX`NG7u)GA!y9-wKMB zMY0tY1yhFu%U<<<`O}DB7Z111am;?rW2Gpy!XrDA$Rl2M8dpoF{-4xze zKs`#kml;x|(7o{uJB81z*=8U4hYq4Pu|}b~8=y=rZ_?hQI%S5iv775l4~5>OI`JgF zc=kx0D3ZT4q1hfHuNU&pcG*0*7raI&d}ASjMVEFhLAiPLAqv7J=t()XKh|cI zO{oP6ijGN}`tJ_1TtcScP-%%~BY%8n|qI(qW`+J94Kz zCExST!Qd(L;&<=vzeZHM=H1>3y*v%E&$DQtkCyu7+L6oKU6|sEWOy;gKD!K7K;3hz zt$)Lw>hl|PR+e1Jk7S@^ssW|o_k#Das<`O4=*nHL>3a_EG4rZB)eqE-$R7pjm3{EXWLgHxo=w0v{*WJK6Q_#G`s zch4?7cuBvaawQzNyMnsuyQI##EMbms8@-O7!g`St?s6q|)dNHTBTRK#Z5Mo9l9}YS zkZwP&d560lJ7aXoBiNdJeG~sY<~<9zZORjX`x32vsSv;$5A)QQj)E=~GgwYwJ(3il zPkghbLHeoFbx*3EgTuN@Chm_e5wooiH&wjtM?Ha#P*Jh~r9fTYMS(~nf?PbL5RnXu z6K`TAYMsQ5X}$u;M-Zw>@+fWyiP$}ja}P6>B%KG}?kDwOoid?3THDwAk(m_!aKRpW zm!wPJF3l2XCl7|M`Zc+$cJnE$d!padi!PTSr!i~H-H8gHu! zC_kU)DelA%Zn(C6z(igIbWT4WwzY5C!0RYg!!VHt&#jp@pH%$w^o<(1q8NSh9 z{^I;=-8D8gKikEU1tg|Edi2Rj8F`kpyt2r^J{|0}cg9Z15~YIYvhIST9%j|1vgXW; z{nI#Z$~x6zk!k^WX%If}=u(V6Qu0X^%8mbB7I4MT)xq^3!i?*f%=d1PbvkJd>M3Go zo>b$mS>hcI*TqkmlLVx3M&F18sJ?4Ae~_Jy+2?*XaT#HilG%9{Y56!*ML{qk`&4Mr z*}^_t8=1Q;R9RCse=aY5I!%SJwaAp0KA%XnGe4G+afkxoyX2ecZW#b?Rdf-dzqqEg zieDuX0_YFO*tgOsjkJI!(5LLT%lqYE1so!~T9h+3VCV8MMp?ZTuoFBdo`)N}O7g6@ z*w19NXI6C1+)R5_^F{|UIaMIwp)<5V}qq4 z=;|(_eGq^>)S)GHoaEnzZQn^HW1b%lOBq?a+GUz4dq%mS`sC!_iA5no=%NZ{{=k!# zts}#WSwg^DrZ`i+`%e6EE!Zefl~^%y{(B^Iid?AQdX0nbIjyljEG63hOiIcvNUdPn zd|=&{GwrVPE#EwoqB2X=@ur{PQ}1wI!+v zNbr}0W{$k^x;$rd04Udk2JV;xoY?E`n7&Za}(?m6rI1V zwZx*kQXPz7IdK5ZDPS_H#dyh|!f1c;=L}gjVw}Ajh#+$d5Xj2Or<)n<6rS&#iIB1~ z$p!tz&Ib2Ijur%#;lCZg%KvPgY zf$bf@!I8#@8XpjJgQnNjQy-%Z2_z4IbyFeF@?9_(Rf}q)8;C#2PNH4RDE+2EYgnS> z6D02fj{;einv;)-;dT@&CR&?pDMVqTIzNl9UX-kxz;o?AQ1%B5JzsfJ75m4sw9sF` z1oyX){E;GUN3MGVUz(;R_}@i7L*3f-ecIQQntmucBz0~pBAYKwVQz0bPaBgTTxQ&z zstlbPlAc-?6~jiM6LKh;EH~u`1rxu!QP#`C4a+2`bFQ2cskw->PCDqM^Q6GgisBlh z(G#!ZKm|KjD#b2Rp>rd+Z^A%Lw-^xsIDCr_9_PNIFA#Mdj?rQ3S_(8{hXUtW! zWR7U|UZaW!-_}IITL)vC16!sCG)IK%cbUYkXcryGc4!$ABFVx{Tg&~39^~=>$m;Xd z>M2d2h4IH1!4I$GE9sZm1!bmrmaPSDM01+f7d6OxQ9tA^srs(k6HrYPQ5#G`2^>`X z3N4&(`F{vcgHKOxMsiV7?HxeeFJm+9I7+S!@FArEuffv+R=(8G0tooX14c9IThVCv zr6HbOieY?;gm=D}Z<$JItR(UPvHiN5Pl!FwYRbxalO;6JcdD{30g|HA^A}13$LKea z0q^k*J2#eyQyRX9@G4o|Bs795vEng9wC&XR+I7;jGQHoAa-IhFFQyns!_piqsv!b- zFwcs69YB0&Er3zLC1TbZEPgewIpA6MKA=|uK(La=bTeg#J!iZb5z;D&STjpp!y5ZmuN>UbXKS+p zeI95J#PD3(sQSf&EAMV}CqA*6@4)LZYD6_kNh8PCHE7yZZSL48p*`g`JEw-I+Q_ z%BeE(6~7Qy`ZvAz3wNc|FUbXTQ@877d4hT1k0^m5qzZazG_x{6^zKW7^scKWywaTZ zvr`@S@S%GaN2#~GjSjD1P)%zcxgp+&+l$fGDSltDvZ-0cHtsY42c5BCDN&_C@Gi?% zv-8<%2yV>Q$)zXREYa&yMCOhoSe@E(-jb1U(K6I;@T;O9@Zg-FWY4Hl)eXdN);ZC# z;SkQ{Aj;!ci^^vJc9CjMJSS4eJ#oVk8Y}c8xoLo)?u7FEA_lB7;LXgWJ-VP`DuaMQxro7KY|LHVoCFb32!9hTFfB;@Bz>dQj zGfF_}Z{!(}V^&sB)gPGFbDSDBfU{Z>ByE@72^PY;i z{U4cf&)`2bxbvHw=-xaOw$(98-A0Vd^o0ju^x6*WvDOvhTqjgx`fQxGmmj0_`uPJ?cbWVzp7|?4=^d7Kj%JK<(d{AdF zvnY#+j#5)uQ`#p{XW?S@J=3E>G8Fpa#IgBsesC~JRYC6inS0`zF~wom@+!xpzIVoo zoj8tPx;U93ZSMqG;W#(`!!-qcNYAw8N~wBdH*PPY zoI>=1G!TKc6zIEHv2XIo0L>Ln`|29ds_IUMOnEH@%VGCkmG#WVKb-hoziYIve9MI# z+KfkKV4j6LJOAjs93LuSwc2x;$OL5s)VJil(-D7Nr!I$SwocX@lEosq0IF5Dl zKO}oGt}im?{8p2NqnoX3$sqtgDN{~5sI0Z5xJqaF_SMhDE@Ky00T!3f6dX_3Cm)E? zf*lvry$>?T z$gczFbT{g&H@D~v8~ip?vW*C-^-7mQ*lZc@w>7yagti5#@jN$82f_}$7xiw*O`d0G z7t@#BOfG7nA10qd0mW`)cLvqj$tCYO6n;R7`!alK5snILK#1oea}`t2F!nThv$>4E z`{d61_S&8xh*y<0WBdpn!k*h7ycv@KW+0mRaeN$R7;KU(NQ=c$ z=h(ONaZwYcd~?)r!OVlKcq2Ioru2>`Fn~Wx%npIQZ-J#pfCIdr_MM`+HQ;uNKPyn^ zRXfq@rLXlvO86WDBMUB8hu=WJvLLHGdI2p&W#ta?D7{st8+DkmwgoWr!VDhERyd@q zusiDlLMybe!PF;9fi+noD8_i&ajj5hAOy?`C@cO~vjO{;k2F2cr59)a6~F$727(Lj z{oBxBpHUwSybvS=C^*gwqlAA(>*$wFl}sa=7f(yQk1>KX4ZAdR=8@P5V~<2y3+VyHk>L) zTK&o}q7neXqET6ZyrC{iVpgF*?Mb-!Ll_(r^lv)PFIn?RKopJ3oHe$t!Ep+Be%~^6 zt4v9heQ45)nv$4pB%HwaGNKOhR}0}P9G+n>kRUnQLDeLjKj>)45X&j{&#KIa$u@>7 zQ022f1y0K83(H$%5F_Q5XeF7&%iSKNN|X@yP$s6m*=KfG zQdtbRi#Z#302r2ngK2L+SfZK5iI&2IB7;u-;?cC>ertsKOVBHtImz5yQ}bjSDj5q8t}+SEaFT-8>prqB1WUy_YJ^Zcyf!wy*+ zWFNvac^FNglKzxxm@UJe9Mm$kH(0_gay~Ozbw+{x`A;$IZpZ8*d|J@5B zBF;_xPkYG!6&^uAaK6&c|8oQe)e{1@#4&+2;;8=u|2Gqj{@=i5vM(Uzzo7ri-uyRc zf%@OO@k;!^yZzq_+J6I?8U6(Z{{O)LbxL0?EQtRj+}AY_)Q%lkFMuH_=9C)*#LJG!)6vD&)X~xQpLDvi94HtD`2Y418HDjj s=Zt&};{QrQ;r+*8LOvj^2mvra65-!T<0SFH2>5{&k`z#00{;O2KfudN`v3p{ diff --git a/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py b/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py index 10404739..fb56ba0d 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/server/routes.py @@ -71,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: @@ -473,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) @@ -488,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)) @@ -500,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")) @@ -521,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) @@ -604,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: @@ -625,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) @@ -643,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) @@ -660,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}: @@ -685,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)} @@ -697,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 {} @@ -752,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)} @@ -766,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()} @@ -778,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: 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 43b942bb..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`
+ ${ctx.readOnly ? null : html`
${saving ? "Saving…" : "Use next round"} -
+ `} ${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) { @@ -940,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 57fffc1a..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) { @@ -468,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) { @@ -488,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`
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_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 index f3f0b80c..faa21706 100644 --- a/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py +++ b/metainfer/tasks/dcu_kernel_auto_opt/tests/test_gpu_preflight.py @@ -2,10 +2,28 @@ 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 ================================== ================================================================================================ @@ -17,6 +35,20 @@ ======================================== 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) @@ -27,28 +59,72 @@ def test_parse_hy_smi_rows(): assert parsed[1]["temp_c"] == 49.0 -def test_usability_is_free_vram_not_percentage(): - """The rule is `free >= 4GB`, so a 91%-used card is still usable.""" - ok = gp.check_gpu(0, {"vram_percent": 91.0, "util_percent": 0.0}, +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) - too_full = gp.check_gpu(1, {"vram_percent": 94.0, "util_percent": 0.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 ok["usable"] is True and ok["free_gb"] == 5.76 - assert ok["measurement_suspect"] is False - assert too_full["usable"] is False and too_full["free_gb"] == 3.84 - assert any("VRAM free" in r for r in too_full["reasons"]) - # a full device is a capacity problem, not a pollution problem - assert too_full["measurement_suspect"] is False - assert gp.DEFAULT_MIN_FREE_GB == 4.0 + 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_marks_suspect_without_blocking(): +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 True # usage is allowed (VRAM ok) + 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"]) - assert any("power" 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"]) @@ -63,28 +139,41 @@ def test_idle_foreign_holders_are_not_suspicion(): assert any("foreign" in n for n in idle["notes"]) -def test_preflight_prefers_clean_then_suspect_then_over_limit(monkeypatch): +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": 55.0, + 1: {"vram_percent": 5.0, "util_percent": 0.0, "power_w": 250.0}, - 2: {"vram_percent": 30.0, "util_percent": 3.0}, - 3: {"vram_percent": 97.0, "util_percent": 10.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, 2] - assert plan["suspect_ids"] == [1] + assert plan["clean_ids"] == [0, 1] + assert plan["suspect_ids"] == [] + assert plan["blocked_ids"] == [2, 3] assert plan["over_limit_ids"] == [3] - # usable devices first (clean before suspect); the over-limit device drops - # out while any device is within the limit - assert plan["preferred"] == [0, 2, 1] - assert 3 not in plan["preferred"] + # only gate-passing devices are ever offered + assert plan["preferred"] == [0, 1] + assert "HCU 4.0%" in plan["gate_reasons"]["2"][0] -def test_over_limit_devices_return_when_nothing_else_is_usable(monkeypatch): +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}, @@ -93,7 +182,20 @@ def test_over_limit_devices_return_when_nothing_else_is_usable(monkeypatch): monkeypatch.setattr(gp, "foreign_kfd_pids", lambda *a, **k: []) monkeypatch.setattr(gp, "total_vram_gb", lambda *a, **k: 64.0) plan = gp.preflight_gpus([0, 1]) - assert plan["preferred"] == [0, 1] + 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): @@ -119,3 +221,256 @@ def test_disabled_preflight_returns_default_order(): 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_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_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_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_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/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/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/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 f5387ae1..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,457 +1,78 @@ -// @@variant shape=glm_tp8_fused_qkv_a_proj_m4096 commit=99b9a733eddfc7196f3d186b80affe4324351ac4 added=2026-09-11 -// median_us=883.7 p90_us=885.5 speedup=77.7 baseline_us=6.866e+04 -// source=test-9-11-1-b01955b3/iteration_001/glm_tp8_fused_qkv_a_proj_m4096 +// @@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] ) // -// Round 1 (mandatory establish_arch) -- packed_weight architecture family: +// 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. // -// Bottleneck hypothesis, from the trusted PMC of the current best source -// (the correctness-first bootstrap: w8a8_dumma_prefill_tiled_kernel -// <128, 64, 128>, identity [K, N] weight, library du_load_matrix_sync -// loaders, LDS row strides 128 (A) / 64 (B), grid 41 x 32 = 1312 blocks, -// 256 threads, 2 blocks/CU): -// median 2920.88 us / p90 2962.82 us, lds_instructions 21,999,616, -// lds_bank_conflicts 274,071,552 (~12.5 per LDS instruction), -// lds_wait_instructions 50,181,819, valu_instructions 86,560,512, -// vmem_read_instructions 1,574,400, and a counter-derived operator HBM -// bandwidth of only 50.6 GB/s (147.8 MB per replay at a 90.7% L2 hit -// rate). The exact-shape stage loop is LDS-issue/conflict bound, not -// MMA- or HBM-bound: both LDS row strides alias every row onto one bank -// phase (128 mod 32 == 0, 64 mod 32 == 0), and the row_major matrix_b -// loader of du_load_matrix_sync lowers each m16n16k32 int8 B fragment to -// 8 ds_read_u8 at 16-byte spacing plus a byte-reassembly chain. +// (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. // -// Focused change (ONE architecture family -- packed_weight -- entirely -// inside csrc/w8a8_gemm_hip.hip): -// (1) launch_pack_w8a8_weight packs the exact (k, n) == (6144, 2624) -// weight ONCE, outside the timed region and out of Graph capture, -// into the n-major layout packed[n * k + kk] = raw[kk * n + n] -// (same byte count, same buffer allocation, so every captured -// address is unchanged). Every other (K, N) keeps the identity -// device-to-device copy. -// (2) the exact-shape guard dispatches the new -// w8a8_dumma_prefill_packedb_kernel<128, 64, 128>. It keeps the -// EXACT bootstrap geometry -- grid (N/64) x (M/128) = 41 x 32 = 1312 -// blocks, 256 threads = 4 gfx928 wavefronts, one 64 x 32 quadrant of -// eight m16n16k32 int32 accumulator fragments per wave, the same -// single-buffered 128-K LDS stage with two __syncthreads per stage, -// the same k0-outer / kk-inner ascending int32 accumulation order and -// the same coalesced bf16 epilogue -- and changes only the operand -// path to the packed-weight architecture: -// * B is staged n-major (b_tile[n_row][k]) from the packed [N, K] -// weight, so a col_major matrix_b fragment is 8 contiguous K -// bytes per lane and every staging vector is one coalesced -// 16-byte global read; -// * both LDS tiles use the bank-safe row stride StageK + 8 = 136 -// bytes (34 dwords): 8-byte aligned with an odd 8-byte count -// (136/8 = 17), so the 16 rows of each 16-lane ds_read_b64 phase -// land on 16 distinct bank pairs (the conflict-free 4-phase -// floor), and 136 % 16 == 8 keeps the two-half int64 staging -// form; -// * both fragment loaders are the explicit load_frag8: ONE aligned -// 8-byte load per lane per fragment (row = lane & 15, k-quarter = -// lane >> 4), replacing the byte-reassembly library path with a -// single ds_read_b64 while producing identical fragment values. -// (3) LDS/block grows 24,576 -> 26,112 B (A[128, 136] + B[64, 136]); -// 2 x 26,112 = 52,224 <= 65,536 B keeps TWO resident blocks/CU = -// 8 waves/CU, so grid_blocks (1312), waves_per_block (4) and the -// active CU count (120) are unchanged. -// (4) the scalar fallback decodes the n-major layout when -// (k, n) == (6144, 2624) with ascending kk, so the paired M=2 / -// M=16 API shapes with the same (N, K) stay byte-exact; the generic -// <64, 64, 128> identity arm is untouched. +// 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. // -// Correctness: the same weight elements reach the same fragment slots in -// the same ascending k-chunk sequence (k0 outer over 48 stages of 128 K, -// kk inner ascending over four m16n16k32 steps), so every int32 -// accumulator -- and therefore every bf16 output byte -- is identical to -// the bootstrap kernel and to the scalar reference. -// -// Round 2 (mandated operand-reuse round: cooperative A+B LDS staging vs direct -// loads, quantified A/B reuse per macro-tile, vectorized coalesced global loads -// with a bank-safe LDS layout): -// -// Operand-reuse accounting for the accepted <128, 64, 128> macro-tile -// (M=4096, N=2624, K=6144): -// * per block the unique operand bytes are A = 128 x 6144 = 786,432 and -// B = 6144 x 64 = 393,216; the 41 N-blocks re-read every A row-slab -// (M*K*41 = 1.03 GB/replay, 41x amplification of the 25.2 MB unique A) -// and the 32 M-blocks re-read every B column-slab (K*N*32 = 0.52 -// GB/replay, 32x of the 16.1 MB unique B): 1.55 GB/replay of staged -// global traffic; -// * inside one block each staged 16-B vector is re-read from LDS by the two -// wave columns (A) / the two wave rows (B), i.e. exactly 2x wave-level -// reuse per staged byte, so the LDS read traffic is 2x the staged bytes -// (4.72 GB/replay, ~23 B/cycle/CU at the measured 1.271 GHz -- far below -// the LDS bandwidth), and the counter-derived HBM rate is only -// 108.4 GB/s (147.8 MB/replay at a 90.7% L2 hit rate). The kernel is -// therefore NOT HBM- or LDS-bandwidth bound, so a direct per-wave global -// fragment path is rejected up front: it would issue 4x more vmem -// instructions (4 waves x 48 stages of 8-B fragment loads) and every 8-B -// fragment load would be a row-strided multi-line access instead of the -// coalesced 16-B staging vectors of the cooperative path. Cooperative -// A+B LDS staging is RETAINED -- the round improves it instead. -// -// Remaining exposed cost of the retained path: the VMEM round trip in front -// of every stage consume. The accepted source (PMC digest daf403b3..., -// profiles/.../iteration2/pmc.json: profiled 1368.793 us, grid 1312 blocks, -// 2 blocks/CU, 26,112 B LDS, lds_instructions 7,892,992 with -// lds_wait 4,535,391, lds_bank_conflicts 0, vmem_read 1,574,400, -// grbm_gui_active == grbm_count) issues per 128-K stage 24 wave-level -// global_load_dwordx4 + 24 ds_write2_b64 and then, per the exact-source ISA, -// strictly serializes global_load_dwordx4 -> s_waitcnt vmcnt(0) -> -// ds_write2_b64 -> s_barrier -> 24x ds_read_b64 -> 32x v_mmac -> s_barrier: -// with 8 resident waves/CU the GPU is 100% GRBM GUI-active yet issues only -// ~4.6% of its slots, i.e. the stage boundary -- not HBM, LDS bandwidth or -// MMA throughput -- sets the pace. -// -// Focused change (ONE mechanism -- the K stage pipeline -- entirely inside -// csrc/w8a8_gemm_hip.hip): the exact-shape guard now dispatches -// w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64>. It keeps the SAME -// macro-tile (grid 41 x 32 = 1312 blocks, 256 threads = 4 wavefronts, 64 x 32 -// quadrant of eight m16n16k32 int32 accumulators per wave), the SAME packed -// n-major B (16-B coalesced staging vectors), the SAME explicit load_frag8 -// fragment reads, the SAME bank-safe LDS stride form and the SAME coalesced -// epilogue -- and changes only how the K stages are staged: -// * the single-buffered 128-K stage (TWO __syncthreads per stage, full -// VMEM latency exposed in front of every consume) becomes a -// DOUBLE-buffered 64-K stage with ONE __syncthreads per stage: the -// stage-(s+1) global loads (2 int4 A + 1 int4 B per thread) are issued -// into staging VGPRs BEFORE the stage-s consume and flushed into the idle -// LDS buffer after it, so the stage-(s+1) VMEM latency overlaps the -// stage-s LDS/v_mmac work. -// -// Resources: LDS = 2 x (A[128, 72] + B[64, 72]) = 2 x 13,824 = 27,648 B/block -// (26,112 -> 27,648 B, +5.9%); 2 x 27,648 = 55,296 <= 65,536 B keeps TWO -// resident blocks/CU = 8 waves/CU, occupancy UNCHANGED. kLdsStride = 72 = -// StageK + 8: 8-aligned with an odd 8-byte count (72/8 = 9) so the 16 rows of -// each 16-lane ds_read_b64 phase land on 16 distinct bank pairs (18r mod 32, -// r = 0..15) -- the same conflict-free 4-phase floor as the accepted kernel -- -// and 72 % 16 == 8 keeps the two-half ds_write2_b64 staging form. 12 extra -// VGPR hold the prefetch payload (~72 -> ~84, still <= the 128/thread budget -// of 2 blocks/CU; 0 spills expected). -// -// Correctness: k0-outer over 96 ASCENDING 64-K stages with kk-inner ascending -// over two m16n16k32 steps gives the identical (k0 + kk) k-chunk sequence -// 0, 32, 64, 96, ... as the accepted 48 x 128-K stages, with the identical -// element-to-slot fragment mapping, so every int32 accumulator -- and every -// bf16 output byte -- is bit-identical to the accepted kernel and to the -// scalar reference. The n-major pack, the scalar-fallback n-major decode, the -// generic <64, 64, 128> identity arm and every shape guard are untouched. -// -// Round 3 (mandated pipeline round: compare single and double buffering across -// K tiles; retain double buffering only on ISA/PMC evidence of reduced VMEM -// stalls without harmful LDS or occupancy growth): -// -// * Mandatory comparison (exact trusted PMC/ISA of the two accepted arms): -// arm median us LDS B/blk blocks/CU VGPR -// single 128-K <128,64,128> 1363.078 26,112 2/8 waves 72 -// double 64-K <128,64,64> 1153.970 27,648 2/8 waves 96 -// A single 64-K arm is not built: with NO prefetch it can only add -// barriers (2 per 64-K stage = 192 per block vs the 128-K arm's 96) while -// exposing the same VMEM round trip, so it is strictly dominated and -// rejected by construction. -// RETAINED -- double buffering. Evidence: median 1363.078 -> 1153.970 us -// (-15.3%) and p90 1364.825 -> 1164.522 us (-14.7%); the exact-source PMC -// of the double-buffered arm (digest ee609c3f..., iteration3/pmc.json) -// shows lds_wait_instructions 4,562,210 -> 3,590,796 (-21.3%) with -// grbm_gui_active == grbm_count (100% GUI-active), LDS 27,648 B/block -// (2 x 27,648 = 55,296 <= 65,536 so the trusted occupancy-probe split 2 -// is unchanged = 8 waves/CU), arch_vgpr 96 with 0 spills (<= the -// 128/thread budget of 2 blocks/CU) and sgpr 32. The ISA confirms the -// mechanism: the stage-(s+1) global_load_dwordx4s are issued BEFORE the -// stage-s ds_read_b64 / v_mmac burst and the vmcnt(0) wait lands on the -// ds_write2_b64 flush AFTER the MMACs, while the 128-K arm serializes -// global_load_dwordx4 -> s_waitcnt vmcnt(0) -> ds_write2_b64 -> s_barrier -// with zero compute overlap. So the prefetch converted exposed VMEM -// latency into overlapped work without trading away occupancy. -// LDS-growth verdict: the +1,536 B/block stays inside the 2-block budget -// and lds_wait FELL (the LDS path is not harmed in the net); the new -// lds_bank_conflicts 12,091,392 is exactly 6 x 4 x 4 x 96 x 1312 = one -// conflict per 16-lane phase of every merged ds_read2_b64, i.e. the -// bank-preserving alias of the +16-row second fragment offset -// (16 x 72 = 1152 B = 9 x 128 B; for ANY 8-aligned uniform row stride S, -// 16 rows = 4S dwords with 4S % 32 == 0). The unmerged ds_read_b64 -// control at the same 72/136-byte strides measured 0 conflicts, so this -// is an alias of the merged read form, not a layout defect. -// Every other K-tile pipeline is rejected inside the 65,536 B budget: -// StageK=96 double buffer = 39,936 B/block, StageK=128 double buffer = -// 52,224 B/block and 3 x 64-K = 41,472 B/block all exceed the 32,768 -// B/block that 2 resident blocks/CU permit (each would drop to 1 -// block/CU = 4 waves/CU, a 2x occupancy loss, explicitly excluded), and a -// 32-K stage leaves only 128 B staging vectors per block (half the 256 -// threads idle) while carrying a single m16n16k32 step per stage (too -// little consume work to amortize the barrier). StageK=64 x 2 buffers is -// therefore the unique 2-resident-block K pipeline and is retained. -// -// * Focused change (ONE mechanism -- the B operand's global-memory layout -- -// entirely inside csrc/w8a8_gemm_hip.hip). The pipeline round's own PMC -// identifies the next exposed cost: the accepted source reads -// 170,221,184 B/replay at only an 86.938% L2 hit rate (3,250,848 misses) -// and 166.147 GB/s counter-derived HBM, while the unique operands are only -// 25.2 MB (A) + 16.1 MB (B). The B side causes it: the n-major pack -// packed[n*K + kk] lays a [64, 64] B tile out as 64 separate 64-byte -// column chunks at a 6,144-byte stride, so every one of the 8 warp-level -// 16-byte staging loads touches 32 half-used 128-byte lines. -// launch_pack_w8a8_weight now packs the exact (k, n) == (6144, 2624) -// weight once, outside the timed region and out of Graph capture, into 41 -// B panels of [K, 64] int8 with 16-byte k-runs per column: -// packed[(n>>6)*(k*64) + (kk>>4)*1024 + (n&63)*16 + (kk&15)] -// = raw[kk*n + n] -// (byte-exact bijection 41 x 384 x 64 x 16 = k*n = 16,121,856; same byte -// count, same buffer and same captured addresses; every other (K, N) -// keeps the identity copy). Each [64, 64] B tile of a 64-K stage is then -// FOUR contiguous 1-KiB slabs (one per 16-k-run) and every staging vector -// is a 16-byte-aligned contiguous k-run of one column, so a warp's 32 -// x 16-byte B staging loads touch 4 FULLY-USED 128-byte lines instead of -// 32 half-used ones (8x fewer line requests per warp-load; the tile's line -// footprint is 32 fully-used lines instead of 64 half-used ones). -// The ONLY kernel change is the two B global-load sites (prologue + -// steady-state prefetch) of the dispatched -// w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64>: byte -// (n0>>6)*(k*64) + ((k0>>4)+k16)*1024 + n_row*16 replaces -// (n0+n_row)*k + k0 + k16*16. The LDS staging layout b_tile[2][64][72] -// (two-half ds_write2_b64), the col_major load_frag8 consumes, the -// k0-outer/kk-inner ascending int32 accumulation order, the coalesced -// epilogue, the grid (41 x 32 = 1312 blocks), the block size (256 = 4 -// wavefronts), the retained double-buffered 64-K pipeline (one -// __syncthreads per stage), LDS/block 27,648 B (2 blocks/CU = 8 waves/CU) -// and the VGPR budget are ALL unchanged, so this is a pure global-layout -// A/B. -// * Correctness: the same weight elements reach the same fragment slots in -// the same ascending k-chunk sequence, so every int32 accumulator -- and -// every bf16 output byte -- is bit-identical to the accepted kernel and to -// the scalar reference. The scalar fallback for (k, n) == (6144, 2624) -// decodes the panel layout in ascending 16-byte k-runs (the paired M=2 / -// M=16 API shapes with the same (N, K) stay byte-exact); the generic -// identity arm, the round-2 n-major pack kernel (kept compiled as the -// layout control) and every shape guard are untouched. Packing stays -// outside the timed region and out of Graph capture. -// * Predicted deltas: vmem_read_instructions unchanged (1,574,400), LDS -// counters unchanged (lds_instructions 4,901,632, lds_wait 3,590,796, -// conflicts at the merged-read alias floor 12,091,392), L2 misses strictly -// below 3,250,848 with the hit rate up, operator read bytes strictly below -// 170,221,184 (B line fetches halve), median < 1153.970 us. -// * Superseded by Round 5: the B-panel pack and the b_tile[2][64][72] LDS -// form below are replaced by the fragment-interleaved swizzle. -// -// Round 5 (mandated packing round: one weight packing/swizzle that makes each -// DUMMA B tile vector-loadable and LDS-bank-safe): -// -// * Baseline: the accepted double-buffered <128, 64, 64> kernel (official -// median 937.431 us / p90 939.452 us; exact-source PMC -// profiles/.../iteration4/pmc.json: profiled 889.453 us, grid 41x32 = -// 1312 blocks, 2 blocks/CU = 8 waves/CU, 27,648 B LDS, arch_vgpr 96 with -// 0 spills, lds_instructions 4,901,632 = 9.7 per wave-stage (6 merged -// ds_read2_b64 + 3 ds_write2_b64), lds_wait 3,828,676, -// lds_bank_conflicts 12,091,392 = 6 x 4 x 4 x 96 x 1312, vmem_read -// 1,574,400, L2 hit 90.39%). The surviving exposed LDS cost is the -// b_tile row-major layout: each col_major B fragment is one ds_read_b64 -// at a 72-byte row stride, and the compiler merges the fragment pairs -// into ds_read2_b64 whose +16-row second offset (16 x 72 = 1152 B = -// 9 x 128 B) aliases the same bank pairs -- one conflict per 16-lane -// phase, exactly the measured 12,091,392 conflicts. -// * Mechanism (ONE packing/swizzle, entirely in this file): replace the -// Round-3 B-panel pack with a fragment-interleaved swizzle. For each -// 32-k x 16-n DUMMA B fragment block the 4 k8 groups x 16 n-rows x 8 B -// are stored n-interleaved as [k8][n][8] (512 B per fragment block), -// 32-k major inside the 64-K stage: -// packed[((kk>>5)*n16g + (col>>4))*512 + -// (((kk>>3)&3)*16 + (col&15))*8 + (kk&7)] = raw[kk*n + col], -// with n16g = n>>4 (164 for the exact shape). Byte-exact bijection -// 192 x 164 x 512 = k*n = 16,121,856: same byte count, same buffer and -// same captured addresses (graph-stable), and packing still runs once, -// outside the timed region and out of Graph capture. -// * Kernel-side effect: every 16-byte B staging vector is now 2 n-rows x -// 8 B of one k8 of one fragment block, so the staging store is ONE -// 16-byte-aligned int4 (ds_write_b128) flat copy per thread (256 chunks -// = one per thread per 64-K stage; a warp reads 8 fully-used 128-byte -// lines) and every col_major B fragment is load_frag8_swz: lane -// (n = lane&15, k8 = lane>>4) reads p + k8*128 + n*8, i.e. one -// ds_read_b64 whose 16-lane phase spans exactly one 128-byte k8 block -- -// all 32 banks exactly once per phase, conflict-free by construction. -// The A tile keeps its bank-safe 72-byte stride, and the double-buffered -// 64-K pipeline, the k0-outer/kk-inner ascending int32 order, the -// coalesced epilogue, the grid (41x32 = 1312 blocks), the block -// (256 threads = 4 wavefronts) and the occupancy (2 blocks/CU = 8 -// waves/CU) are all unchanged, so this is a pure packing/swizzle A/B. -// * LDS: b_tile shrinks from 2 x [64, 72] = 9,216 B to 2 x 4,096 B; the -// block total is 18,432 (A) + 8,192 (B) = 26,624 B, so 2 x 26,624 = -// 53,248 <= 65,536 B keeps 2 resident blocks/CU = 8 waves/CU. -// * Correctness: the same weight elements reach the same fragment slots in -// the same k0-outer/kk-inner ascending k-chunk sequence, so every int32 -// accumulator -- and every bf16 output byte -- is bit-identical to the -// accepted kernel and to the scalar reference. The scalar fallback for -// (k, n) == (6144, 2624) decodes the swizzle element-wise in ascending -// kk (the paired M=2 / M=16 API shapes stay byte-exact); the generic -// identity arm, the Round-2 n-major pack control and every shape guard -// are untouched. -// * Predicted deltas: vmem_read_instructions unchanged (~1,574,400), -// LDS/block 27,648 -> 26,624 B, lds_bank_conflicts strictly below -// 12,091,392, lds_wait below 3,828,676, median < 937.431 us. -// Falsifiable: retained only if median_us < 937.4308013916016, -// mismatch_count == 0, max_abs_error == 0.0, -// graph_capture_passed == true, the exact code object shows 0 spills / -// LDS == 26,624 B / VGPR <= 128 (2 blocks/CU), and the PMC shows -// lds_bank_conflicts and lds_wait strictly below the baseline with -// vmem_read not growing; any regression restores the iteration-3 source. -// -// Round 6 (mandated epilogue round: fuse per-row/per-column scales, bf16 -// conversion and the final coalesced store into the compute kernel; remove any -// unnecessary workspace/combine pass): -// -// * Fusion audit (nothing to remove, so the fused epilogue is made cheaper). -// The accepted source ALREADY fuses everything: store_prefill_fragment_ -// coalesced applies x_scale[row](per-row) * weight_scale[col](per-column), -// converts to bf16 and writes one coalesced 8-byte store per lane, and the -// exact-shape operator replay is a SINGLE dispatch -- the trusted PMC -// (profiles/glm_tp8_fused_qkv_a_proj_m4096/iteration6/pmc.json, digest -// 7e5731dc..., compile key 120c807e...) reports operator_aggregate -// kernel_count == 1, scratch_bytes == 0, vmem_write_instructions 41,984 = -// 8 x 4 x 1312 and write_bytes_per_operator_replay 21,500,384 = m*n*2, i.e. -// there is no workspace, no split-K combine pass and no second Graph node -// to remove; workspace_bytes is unused and workspace_split_k_capacity is -// zero for this shape. The remaining epilogue cost is therefore the -// conversion itself. -// * Measured cost of the fused epilogue (exact gfx928 code object of the -// current best, iteration6/current-best-isa/isa.txt). The dispatched -// kernel is ~1,920 instructions; the steady-state K-stage loop body is 107 -// instructions (one `s_branch 65428` back-edge at 0x6BB8 -> 0x6A0C) and -// the prologue is ~259, so the fused epilogue is ~1,490 static -// instructions -- roughly 13% of the dynamic instruction stream over -// 96 stages. The dominant term is that this DTK lowers EACH of the 32 -// per-lane __float2bfloat16 calls to a ~13-instruction EXEC-masked -// Inf/NaN sequence: -// v_and_b32 (b & 0x7f800000) / v_cmp_ne_u32 / s_and_saveexec_b64 / -// s_xor_b64 / v_bfe_u32 (bit 16) / v_add3_u32 (b + lsb + 0x7fff) / -// s_andn2_saveexec_b64 / v_mov_b32 0 / v_or_b32 (b | 0x10000) / -// v_cmp_eq_u32_sdwa / v_cndmask_b64 / s_or_b64 (restore EXEC) -// so ~420 instructions per lane plus 4 x (3 shifts + 3 ors) of packing, -// with 32 dependent SALU EXEC save/restore pairs that cannot be software- -// pipelined and serialise the wavefront (the epilogue is the only place -// VALU issues drop while the exact-shape kernel is 100% GRBM GUI-active). -// * Focused change (ONE mechanism -- the bf16 conversion of the fused -// epilogue -- entirely inside csrc/w8a8_gemm_hip.hip). Replace the library -// conversion + shift/or pack in store_prefill_fragment_coalesced with the -// branch-free round-to-nearest-even integer identity -// bf16_rne_high16(v) = (b + 0x7fff + ((b >> 16) & 1)) >> 16, -// b = __float_as_uint(v), -// which the disassembly proves is exactly what the normal (finite) path of -// the library conversion computes (v_bfe bit 16, v_add3 b + lsb + 0x7fff). -// Adding 0x7fff + the retained bit-16 LSB cannot carry out of the sign, the -// +((b>>16)&1) term rounds exact ties to an even LSB (RNE), max-float still -// rounds to +/-Inf and 0.0/-0.0 and subnormal encodings are preserved, so -// every 16-bit payload is identical. Every value converted here is finite -// by construction (|int32 dot| < 2^31 times finite dequant scales), so the -// removed Inf/NaN special case is unreachable. Four values now pack into -// one 8-byte word with three VALU ops (v_bfe/v_and + v_add3 + v_lshrrev) -// instead of ~13 EXEC-bracketed instructions each, and the 32 EXEC -// save/restore chains disappear. The epilogue structure, the per-row -// x_scale gather, the per-column float4 weight_scale gather, the register -// 4x4 shuffle transpose, the multiply order float(dot) * x_scale * -// weight_scale, the 8-byte-aligned coalesced store address, the grid -// (41 x 32 = 1312), the block (256 = 4 wavefronts), the double-buffered -// 64-K pipeline, the swizzled packed B, LDS/block 26,624 B (2 blocks/CU = -// 8 waves/CU) and the 65,536 B workspace budget are ALL unchanged, so this -// is a pure conversion A/B. The scalar fallback keeps __float2bfloat16 (it -// is not on the timed path); the generic identity arm and every shape -// guard are untouched. -// * Predicted deltas: grid_blocks 1312, LDS/block 26,624 B, 2 blocks/CU = -// 8 waves/CU, vmem_read 1,574,400 and vmem_write 41,984 unchanged; -// valu_instructions strictly below 23,920,384 and the epilogue's static -// EXEC/SALU count down by >1,000; median < 918.191 us with no p90 -// regression. Falsifiable: retained only if median_us <= 909.0 us -// (>= 1% round acceptance) and p90_us <= 919.530 us, mismatch_count == 0, -// max_abs_error == 0.0, graph_capture_passed == true, the exact code -// object shows 0 spills / LDS == 26,624 B / VGPR <= 128 (2 blocks/CU) and -// no __float2bfloat16 EXEC sequence remains in the epilogue; otherwise the -// iteration-5 source is restored. -// -// Round 9 (late ISA-diagnosis round; the control-plane plateau gate is CLOSED -// for this round -- phase hip_only, plateau false, raw inline asm and ISA -// Skills forbidden -- so this stays a HIP-only round): -// -// * Bottleneck evidence (exact trusted gfx928 ISA of the current best, -// profiles/glm_tp8_fused_qkv_a_proj_m4096/iteration6/current-best-isa/ -// isa.txt, digest 7e5731dc...; the round-6 source differs from that -// profiled official best only in the bf16 conversion, so the steady-state -// loop analysed here is the dispatched loop). The steady-state stage body -// is 107 instructions (0x6A0C..0x6BB8, one back-edge `s_branch 65428` at -// 0x6BB8 -> 0x6A0C) and roughly 14 of them are pure A-row tail -// predication: -// 0x6A0C..0x6A20 six unconditional `v_mov_b32` zero-fills of the two -// int4 staging vectors (v33/v34 and v37..v40), -// 0x6A24..0x6A34 `s_and_saveexec_b64` / `s_cbranch_execz` / -// `global_load_dwordx4` / `s_or_b64` for A vector 0, -// 0x6A38..0x6A50 two more `v_mov_b32` zero-fills (v35/v36) plus the -// same EXEC save/restore + uniform branch for A -// vector 1. -// The guard source `a_pre[v] = (g_row < m) ? load : int4{0,0,0,0}` is -// never false for the dispatched M=4096 (M / BM = 4096 / 128 = 32 exact), -// yet its zero-fills execute unconditionally in every one of the 96 -// stages and its two EXEC read-modify-writes serialize the wavefront in -// front of the very global loads that prefetch the next tile. That is -// ~14/107 = 13% of the hot loop's instruction stream, on top of the -// epilogue's per-fragment `row >= m` EXEC bracket (s_and_saveexec / -// s_xor_b64 / s_andn2_saveexec_b64 / s_cbranch_execz / s_or_b64 around -// each of the eight unrolled fragment stores), for a kernel the PMC shows -// 100% GRBM GUI-active yet issuing only ~0.42 instructions per cycle per -// CU (lds_wait_instructions 5,302,705 against lds_instructions -// 4,901,632). -// * Focused change (ONE mechanism -- exact-shape full-tile specialization, -// entirely inside csrc/w8a8_gemm_hip.hip): the dispatched kernel gains a -// compile-time `bool kFullTile` template parameter. The exact-shape guard -// keeps `(m >= 128, n == 2624, k == 6144)` and the same grid (N/64 x -// ceil(M/128)) and block (256 = 4 wavefronts), and selects -// `kFullTile == true` exactly when the caller's m is a multiple of -// BM = 128 -- i.e. when every block tile lies fully inside the output (the -// assigned M=4096 case). In that instantiation the A staging loads -// (prologue AND steady state) drop the `(g_row < m)` select / zero-fill -// and the epilogue drops the per-fragment `row >= m` early return, so the -// loop body keeps only the address math and the three coalesced -// `global_load_dwordx4`. Every other m keeps the guarded instantiation -// byte-for-byte (zero-fill, epilogue row mask, tail-block coverage), and -// the paired M=2/M=16 API shapes with the same (N, K) still take the -// scalar fallback because the `m >= 128` guard is untouched. -// * Why this is the remaining lever and not a new pipeline: rounds 7 -// (deeper register prefetch / stage pairs) and 8 (smaller macro-tile, more -// resident waves) both regressed, so the double-buffered 64-K stage (ONE -// __syncthreads per stage, LDS 26,624 B/block, 2 blocks/CU = 8 waves/CU), -// the 128 x 64 macro-tile, the fragment-interleaved packed B swizzle, the -// load_frag8 / load_frag8_swz consumes, the k0-outer / kk-inner ascending -// m16n16k32 int32 accumulation order and the fused branch-free bf16 -// epilogue are ALL retained unchanged. The round removes instructions that -// are provably dead for this shape without touching any resource, any -// memory access order or any MMA input. -// * Correctness: for m % 128 == 0 no block tile crosses m, so the removed -// guard predicate is never false and the removed zero-fill never reaches an -// accumulator; the surviving assignments, the ascending k-chunk sequence -// and the exact `float(dot) * x_scale * weight_scale` multiply order are -// bit-identical to the guarded path (mismatch 0 / max_abs_error 0.0 -// expected). The guarded instantiation preserves the exact current -// semantics for every m % 128 != 0. -// * Predicted deltas: steady-state loop 107 -> ~93 instructions (the eight -// zero-fills and the two EXEC save/restore + uniform-branch pairs removed), -// valu_instructions strictly below 23,920,384, grid_blocks 1312, -// block_threads 256, LDS/block 26,624 B, arch_vgpr <= 128 with 0 spills, -// vmem_read_instructions 1,574,400, vmem_write_instructions 41,984 and the -// epilogue store count all unchanged. Falsifiable: retained only if -// median_us < 912.7147674560547 (the current shadow base) with no p90 -// regression, mismatch_count == 0, max_abs_error == 0.0, -// graph_capture_passed == true, the exact code object shows 0 spills / -// LDS == 26,624 B / VGPR <= 128 (2 blocks/CU) and the kFullTile symbol's -// loop body no longer contains the A-guard -// `s_and_saveexec`/`s_cbranch_execz`/zero-fill sequence; any median/p90 -// regression, mismatch, spill or occupancy loss restores the round-6 -// source. +// 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). @@ -470,6 +91,12 @@ 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; @@ -477,53 +104,38 @@ constexpr int kExactK = 6144; using namespace du::dumma; -// Coalesced accumulator epilogue for gfx928 INT8 m16n16k32: the accumulator -// lane mapping is row = lane & 15, col_mod4 = lane >> 4, frag.x[i] maps to -// columns col_mod4 + 4*i. A register 4x4 transpose (two 2x2 steps with -// shfl_xor 16 then 32, one v_cndmask per element per step) re-routes the -// 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 -> offset a multiple of 8) are -// aligned. -// -// Round 6: the fused bf16 conversion itself is made branch-free. The exact -// gfx928 code object of the previous best lowers each of the 32 per-lane -// __float2bfloat16 calls to a ~13-instruction sequence that brackets an -// Inf/NaN special case with EXEC save/restore: -// v_and_b32 (b & 0x7f800000) / v_cmp_ne_u32 / s_and_saveexec_b64 / s_xor_b64 -// / v_bfe_u32 (bit 16) / v_add3_u32 (b + lsb + 0x7fff) / s_andn2_saveexec_b64 -// / v_mov_b32 0 / v_or_b32 (b | 0x10000) / v_cmp_eq_u32_sdwa / v_cndmask_b64 -// / s_or_b64 (restore EXEC) -// The EXEC manipulation cannot be software-pipelined and serialises the -// wavefront. For every FINITE fp32 the library conversion is exactly -// round-to-nearest-even at bit 16, which is the branch-free integer identity -// bf16 = (b + 0x7fff + ((b >> 16) & 1)) >> 16, b = __float_as_uint(v) -// (adding 0x7fff + the retained bit-16 LSB never carries out of the sign, and -// the +((b>>16)&1) term breaks exact ties toward an even LSB, so max-float -// still rounds to +/-Inf and -0.0/0.0 are preserved). Every value converted -// here is finite by construction: |int32 dot| < 2^31 and the dequant scales -// are finite, so no Inf/NaN can reach the conversion. This yields the exact -// same 16-bit payloads as __float2bfloat16 while dropping the EXEC sequence -// and the separate shift/or pack (four bf16 pairs become three VALU ops per -// 8-byte word instead of 4 x 13 + 3). +// 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; } -// Round 9 adds the compile-time `kGuarded` flag: kGuarded == false is the -// exact-shape full-tile specialization (m % BM == 0), for which every row of -// every block tile is strictly inside the output, so the `row >= m` early -// return -- which the exact code object lowers to an EXEC save/restore pair -// plus a uniform branch around each of the eight unrolled fragment stores -- -// is removed at compile time. The default kGuarded == true keeps the exact -// current semantics for every other m (including the generic arm and the -// single-buffered siblings, which never pass the flag). +// 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, @@ -567,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); @@ -588,15 +200,14 @@ __device__ __forceinline__ void store_prefill_fragment_coalesced( packed; } -// Explicit 8-byte int8 fragment loader (lineage-validated load_frag8): for -// BOTH the row_major matrix_a and the col_major matrix_b m16n16k32 fragments, -// lane l -> row (l & 15), k-quarter (l >> 4) holds the EIGHT CONTIGUOUS bytes -// at p[row * ldm + (l >> 4) * 8 .. +7]. Loading the 8 bytes as one int64 -// forces the ds_read_b64 form (16-lane phases, distinct bank pairs per group -// -> the 4-phase conflict-free floor with an odd 8-byte LDS row stride), -// identical fragment values to the library loader, so the v_mmac inputs and -// the exact int32 accumulation order are unchanged. `lane` is -// threadIdx.x % kWaveSize, so no wave-id builtin is required. +// 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, @@ -609,17 +220,17 @@ __device__ __forceinline__ void load_frag8( reinterpret_cast(f.x)[0] = v; } -// Round-5 fragment-interleaved B swizzle helpers. +// Fragment-interleaved B swizzle helpers. // // Each 32-k x 16-n DUMMA B fragment block is 512 B = [k8 0..3][n 0..15][8 B] -// (n-interleaved k8 groups). The 16-byte chunk (k32, n16, k8, n_pair) holds -// the two 8-byte k-runs of n-rows 2*n_pair and 2*n_pair+1 of one k8 group, -// and sits at byte (k32 * n16g + n16) * 512 + k8 * 128 + n_pair * 16 of the -// packed [n>>4][k>>5][512 B] array (n16g = n>>4). Global elements satisfy +// (n-interleaved k8 groups). The 16-byte chunk (k32, n16, k8, n_pair) holds the +// two 8-byte k-runs of n-rows 2*n_pair and 2*n_pair+1 of one k8 group, and sits +// at byte (k32 * n16g + n16) * 512 + k8 * 128 + n_pair * 16 of the packed +// [n>>4][k>>5][512 B] array (n16g = n>>4). Global elements satisfy // packed element (col, kk) at ((kk>>5)*n16g + (col>>4))*512 + // (((kk>>3)&3)*16 + (col&15))*8 + (kk&7) -- the exact inverse of this chunk -// map, so staging is a flat 16-byte copy and the fragment load reads the -// same bytes the Round-3 layout delivered. +// 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, @@ -630,12 +241,11 @@ __device__ __forceinline__ int64_t packed_b_swz_offset( static_cast(k8) * 128 + (n_pair << 4); } -// Round-5 col_major B fragment loader on the swizzled tile: lane (n = lane&15, +// 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. +// load_frag8, so the v_mmac operands and the int32 accumulation stay bit-exact. template __device__ __forceinline__ void load_frag8_swz( Frag& f, @@ -648,27 +258,50 @@ __device__ __forceinline__ void load_frag8_swz( 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 = +// 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)). // // 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. 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 above. +// 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( @@ -787,282 +420,146 @@ w8a8_dumma_prefill_tiled_kernel( } } -// Round-1 packed-weight architecture kernel, kept as the single-buffered -// 128-K sibling of the dispatched <128, 64, 64> double-buffered kernel below -// (Round 2 moved the exact-shape dispatch to that sibling; Round 3 keeps this -// kernel compiled but undispatched as the n-major layout / single-buffer -// control, and only the double-buffered kernel reads the packed B; Round 5 -// replaces that packed B layout with the fragment-interleaved swizzle, so this -// n-major control no longer matches the pack either). It is defined for -// the n-major packed_weight[n * K + kk] = raw[kk * N + n] layout produced by -// launch_pack_w8a8_weight. SAME geometry as the bootstrap -// w8a8_dumma_prefill_tiled_kernel<128, 64, 128>: -// grid (N/64) x (M/128) = 41 x 32 = 1312 blocks, 256 threads = 4 wavefronts, -// one 64 x 32 quadrant of eight m16n16k32 int32 accumulators per wave, -// single-buffered 128-K LDS stage with TWO __syncthreads per stage, same -// k0-outer/kk-inner ascending int32 accumulation order, same coalesced -// fragment epilogue. The two differences, both inside the operand path: -// * B is staged n-major (b_tile[n_row][kk] from the packed [N, K] weight) -// so col_major B fragments are 8 contiguous K bytes per lane and staging -// is one coalesced 16-byte vector per (n_row, k16) pair; -// * A and B LDS row strides are 136 B (StageK + 8): 8-byte aligned with an -// odd 8-byte count (136/8 = 17) so the 16 rows of each 16-lane ds_read_b64 -// phase land on 16 distinct bank pairs (the conflict-free 4-phase floor), -// and 136 % 16 == 8 keeps the two-half int64 staging form. -// Both fragment loaders are the explicit load_frag8 above; the staged bytes -// are the same matrix elements in the same fragment slots as the library -// loaders produced from the identity layout, so every v_mmac input and the -// exact int32 accumulation (k0-outer over 48 stages, kk-inner ascending over -// four m16n16k32 steps) are bit-identical to the bootstrap kernel and to the -// scalar reference. -template -__global__ __launch_bounds__(kThreadsPerBlock) void -w8a8_dumma_prefill_packedb_kernel( - const int8_t* __restrict__ x_q, // [M, K] row-major - const int8_t* __restrict__ packed_w, // [N, K] n-major (exact shape) - const float* __restrict__ x_scale, // [M] - const float* __restrict__ weight_scale, // [N] - hip_bfloat16* __restrict__ out, // [M, N] - int m, - int n, - int k) { - constexpr int kWaveM16 = BM / 32; - constexpr int kWaveN16 = BN / 32; - static_assert(BM % 32 == 0 && BN % 32 == 0, - "block tile must be a multiple of the wave quadrant"); - static_assert(kThreadsPerBlock % kWaveSize == 0, - "blockDim must be a multiple of the gfx928 wavefront size"); - // Bank-safe LDS row stride: StageK + 8 bytes = 34 dwords for StageK = 128. - // 8-aligned with an odd 8-byte count -> every lane's 8-byte fragment is a - // ds_read_b64 at the 4-phase conflict-free floor (see the kernel comment - // above); 8 mod 16 keeps the two-half staging store form. - constexpr int kLdsStride = StageK + 8; - static_assert(kLdsStride % 8 == 0 && (kLdsStride / 8) % 2 == 1, - "LDS row stride must be 8-aligned with an odd 8-byte count"); - 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 (the dispatch guarantees - // k % StageK == 0 and n % BN == 0). 8 vectors per StageK row for 128. - constexpr int kVecsPerRow = StageK / static_cast(sizeof(int4)); - constexpr int kAVectors = BM * StageK / static_cast(sizeof(int4)); - constexpr int kBVectors = StageK * BN / static_cast(sizeof(int4)); - - for (int k0 = 0; k0 < k; k0 += StageK) { - // Stage A[BM, StageK] row-major with the bank-safe row stride: thread - // `vec` owns the 16-byte chunk at (local_row, k16); rows past M are - // zero-filled and masked in the epilogue. - for (int vec = tid; vec < kAVectors; vec += kThreadsPerBlock) { - 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: the row stride is 8 mod 16, so odd rows are not - // 16-byte aligned and a single int4 LDS store is illegal. The int64 - // pair lowers to one ds_write2_b64 per thread. - const int64_t* src64 = reinterpret_cast(&v); - int64_t* dst64 = reinterpret_cast( - a_tile + local_row * kLdsStride + k16 * 16); - dst64[0] = src64[0]; - dst64[1] = src64[1]; - } - // Stage B[BN, StageK] n-major from the packed [N, K] weight: thread `vec` - // owns the 16-byte chunk at (n_row, k16); the K axis is always in range. - 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, - static_cast(kLdsStride), - lane); - } -#pragma unroll - for (int i = 0; i < kWaveN16; ++i) { - load_frag8( - b_frag[i], - b_tile + (local_col + i * kTileN) * kLdsStride + kk, - static_cast(kLdsStride), - 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]); - } - } - } - // Single buffer: every wavefront must finish reading LDS before the next - // stage's stores overwrite it. - __syncthreads(); - } - - const int base_row = m0 + wave_row * (BM / 2); - const int base_col = n0 + wave_col * (BN / 2); -#pragma unroll - for (int j = 0; j < kWaveM16; ++j) { -#pragma unroll - for (int i = 0; i < kWaveN16; ++i) { - store_prefill_fragment_coalesced( - acc[j][i], x_scale, weight_scale, out, m, n, - base_row + j * kTileM, base_col + i * kTileN, lane); - } - } -} - -// Round 9 exact-shape A-staging load. For the dispatched M=4096 every 128-row -// block tile lies fully inside the output, so the `(g_row < m)` row guard is -// dead: the exact code object lowers it to eight unconditional zero-fill -// `v_mov_b32` plus two EXEC save/restore (`s_and_saveexec_b64` / -// `s_cbranch_execz` / `s_or_b64`) uniform-branch pairs per K stage, i.e. ~14 -// of the 107 steady-state loop instructions and an EXEC serialization in -// front of the very `global_load_dwordx4`s that prefetch the next tile. The -// kFullTile instantiation removes the select (one coalesced load per staging -// vector, no zero-fill), while kFullTile == false keeps the zero-fill + row -// mask semantics identical to the current kernel for every m % 128 != 0. +// 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__ x_q, - int64_t offset, + const int8_t* __restrict__ p, int m, int g_row) { if (kFullTile) { - return *reinterpret_cast(x_q + offset); + return *reinterpret_cast(p); } - return (g_row < m) ? *reinterpret_cast(x_q + offset) - : int4{0, 0, 0, 0}; + return (g_row < m) ? *reinterpret_cast(p) : int4{0, 0, 0, 0}; +} + +// 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); } -// Round-2 double-buffered packed-B prefill kernel, dispatched ONLY for -// (m >= 128, n == 2624, k == 6144). SAME macro-tile and SAME operand path as -// the accepted w8a8_dumma_prefill_packedb_kernel<128, 64, 128>: grid -// (N/64) x (M/128) = 41 x 32 = 1312 blocks, 256 threads = 4 wavefronts, one -// 64 x 32 quadrant of eight m16n16k32 int32 accumulators per wave, packed -// n-major B, explicit load_frag8 reads, coalesced epilogue. The ONLY change is -// the K pipeline (see the Round-2 section of the file header): -// * 64-K DOUBLE-buffered stages (2 x 13,824 = 27,648 B LDS/block, still -// 2 blocks/CU = 8 waves/CU): the stage-(s+1) global loads are issued into -// 3 int4 staging registers per thread BEFORE the stage-s consume and -// flushed into the idle LDS buffer after it, so the VMEM latency of stage -// s+1 overlaps the LDS/v_mmac work of stage s (one __syncthreads per -// stage instead of two); -// * k0-outer over 96 ASCENDING 64-K stages with kk-inner ascending over two -// m16n16k32 steps, i.e. the same (k0 + kk) k-chunk sequence -// 0, 32, 64, 96, ... as the accepted 48 x 128-K stages, so every int32 -// accumulator is bit-identical. -// kLdsStride = StageK + 8 = 72 B: 8-aligned with an odd 8-byte count (9) so -// the 16 rows of each 16-lane ds_read_b64 phase land on 16 distinct bank pairs +// 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. // -// Round 5 additionally reads the packed B operand from the fragment- -// interleaved swizzle (see w8a8_pack_bswz_i8_kernel and the Round-5 header -// section): the LDS b_tile is now the SAME [k32][n16][k8][n][8] layout as the -// pack (2 x 256 16-byte chunks per 64-K stage), so staging is one flat -// 16-byte copy per thread and each col_major B fragment is one conflict-free -// ds_read_b64 (load_frag8_swz). The int32 accumulation order is untouched. +// 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. // -// Round 9 adds the compile-time `kFullTile` flag: it is true exactly when the -// caller's m is a multiple of BM, so every block tile is fully inside the -// output. In that instantiation the A staging loads drop the row guard (see -// load_a_stage_vec) and the epilogue drops the per-fragment row mask (see -// store_prefill_fragment_coalesced); false keeps the guarded form -// byte-for-byte. Geometry, resources, operand path and accumulation order are -// identical in both instantiations. -template -__global__ __launch_bounds__(kThreadsPerBlock) void +// 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>>4][k>>5][512 B] swizzled B - // fragment blocks (round-5 layout) + // fragment blocks const float* __restrict__ x_scale, // [M] const float* __restrict__ weight_scale, // [N] hip_bfloat16* __restrict__ out, // [M, N] int m, int n, int k) { - constexpr int kWaveM16 = BM / 32; - constexpr int kWaveN16 = BN / 32; - static_assert(BM % 32 == 0 && BN % 32 == 0, + 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"); 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"); - // Round 5: the packed B weight is the fragment-interleaved swizzle (see - // w8a8_pack_bswz_i8_kernel), defined for BN == 64 and StageK == 64 (one - // 32-k x 16-n fragment block per 512 B; 2 k32 blocks x 4 n16 groups x 4 k8 - // groups x 8 n-pairs = 256 16-byte chunks per stage, one per thread). Only - // <128, 64, 64> is ever dispatched by the exact-shape guard, so the chunk - // index math below is exact. - static_assert(BN == 64 && StageK == 64, - "swizzled B pack requires BN == 64 and StageK == 64"); + // 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]; - // Round-5 swizzled B tile: flat [k32][n16][k8][n_pair][16 B] = 4,096 B per - // 64-K stage, byte-identical to the packed global layout of that stage. + // 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 @@ -1079,56 +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"); - // Round 5: the swizzled B stage is exactly one 16-byte chunk per thread, so - // thread `vec` = tid is also the flat [k32][n16][k8][n_pair] chunk index. - static_assert(kBVecPerThread == 1 && kBVectors == kThreadsPerBlock, + // 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]; + // 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] = load_a_stage_vec( - x_q, static_cast(g_row) * k + k0 + k16 * 16, m, g_row); + 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; - // Round 5: the 16-byte chunk (k32 = vec>>7, n16 = (vec>>5)&3, + // 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; - // consecutive lanes read consecutive 16-byte chunks, so a warp's + // (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( - packed_w + - packed_b_swz_offset((k0 >> 5) + (vec >> 7), - (n0 >> 4) + ((vec >> 5) & 3), n16g, - (vec >> 3) & 3, vec & 7)); + 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]); @@ -1139,112 +667,157 @@ w8a8_dumma_prefill_packedb_db_kernel( } #pragma unroll for (int v = 0; v < kBVecPerThread; ++v) { - const int vec = tid + v * kThreadsPerBlock; - // Round 5: the swizzled tile is a flat copy of the packed stage, so the - // commit is ONE 16-byte-aligned int4 (ds_write_b128) per thread. + 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 stage s, issue the stage-(s+1) global loads into the - // staging registers BEFORE consuming stage s from buffer s&1, then flush the - // prefetched vectors into the idle buffer (s+1)&1 and close the stage with - // ONE __syncthreads (so the vmcnt wait for stage s+1 lands after the MMAC - // burst of stage s, overlapping the VMEM latency of s+1 with the stage-s - // LDS/v_mmac work). The last stage is consumed after the loop without a - // trailing barrier. - for (int s = 0; s < kStages - 1; ++s) { - // Prefetch stage s+1: issue the global loads (held in VGPRs 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] = load_a_stage_vec( - x_q, static_cast(g_row) * k + k0 + k16 * 16, m, g_row); - } + 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; - // Round 5: the 16-byte chunk (k32 = vec>>7, n16 = (vec>>5)&3, - // k8 = (vec>>3)&3, n_pair = vec&7) of the swizzled pack sits at byte - // ((k0>>5)+k32)*n16g*512 + ((n0>>4)+n16)*512 + k8*128 + n_pair*16; - // consecutive lanes read consecutive 16-byte chunks, so a warp's - // staging loads touch 8 fully-used 128-byte lines. - b_pre[v] = *reinterpret_cast( - packed_w + - packed_b_swz_offset((k0 >> 5) + (vec >> 7), - (n0 >> 4) + ((vec >> 5) & 3), n16g, - (vec >> 3) & 3, vec & 7)); - } + 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); + // ---- 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, - static_cast(kLdsStride), - lane); - } + 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) { - // Round 5: fragment block (k32 = kk>>5, n16 = wave_col*kWaveN16 + i) - // of the swizzled tile -- 512 B per 16-n x 32-k DUMMA B fragment, one - // conflict-free ds_read_b64 per lane. - load_frag8_swz( - b_frag[i], - b_tile[s & 1] + - ((kk >> 5) * (BN / 16) + wave_col * kWaveN16 + i) * 512, - lane); - } + 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 < kAVecPerThread; ++v) { + a_src[v] += StageK; + } #pragma unroll - for (int v = 0; v < kBVecPerThread; ++v) { - const int vec = tid + v * kThreadsPerBlock; - // Round 5: flat 16-byte swizzled-tile commit (ds_write_b128). - *reinterpret_cast(b_tile[buf] + vec * 16) = b_pre[v]; + 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_row = wave_row * kRowsPerWave; #pragma unroll for (int kk = 0; kk < StageK; kk += kTileK) { #pragma unroll @@ -1257,9 +830,6 @@ w8a8_dumma_prefill_packedb_db_kernel( } #pragma unroll for (int i = 0; i < kWaveN16; ++i) { - // Round 5: fragment block (k32 = kk>>5, n16 = wave_col*kWaveN16 + i) - // of the swizzled tile -- 512 B per 16-n x 32-k DUMMA B fragment, one - // conflict-free ds_read_b64 per lane. load_frag8_swz( b_frag[i], b_tile[s & 1] + @@ -1276,14 +846,14 @@ 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) { - // Round 9: kGuarded == !kFullTile -- the full-tile instantiation skips - // the dead per-fragment `row >= m` EXEC bracket (see the helper above). + // 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); @@ -1292,26 +862,21 @@ 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/packed paths). One output element per +// 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. // -// Round 1: for the exact (k, n) == (6144, 2624) the packed_weight buffer -// holds a packed layout (produced once, outside the timed region), so the -// fallback decodes that layout there with ascending kk (every other (K, N) -// keeps the identity [K, N] row-major copy). The int32 accumulation stays -// kk-ascending 0..k-1 in both arms; integer addition is exact and commutative, -// so the int32 dot -- and the paired M=2 / M=16 API shapes -- remain -// byte-exact against the identity-layout reference. Round 5 replaces the -// exact-shape pack with the fragment-interleaved swizzle (see -// w8a8_pack_bswz_i8_kernel), so the exact-shape decode below walks +// 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. +// 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, // packed weight (swizzled/identity) @@ -1333,8 +898,8 @@ __global__ __launch_bounds__(256) void w8a8_scalar_gemm_kernel( const int8_t* a_row = a + static_cast(row) * k; int32_t acc = 0; if (packed_swz) { - // Round-5 swizzle: one 8-byte k-run per (n, k8) inside a 512-byte - // 32-k x 16-n fragment block; walk it in ascending kk. + // 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 = @@ -1372,46 +937,19 @@ __global__ __launch_bounds__(256) void w8a8_pack_identity_i8_kernel( } } -// Round-1 exact-shape n-major pack: for (k, n) == (6144, 2624) the weight is -// transposed once, outside the timed region and out of Graph capture, into -// packed[n * k + kk] = raw[kk * n + n] so every col_major B fragment is 8 -// contiguous K bytes in the packed buffer and every staging vector is one -// coalesced 16-byte read. Same byte count and buffer as the identity pack, so -// the captured addresses are unchanged. One thread per output byte; runs once -// during weight preparation. -__global__ __launch_bounds__(256) void w8a8_pack_nmajor_i8_kernel( - const int8_t* __restrict__ raw, - int8_t* __restrict__ packed, - 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]; - } -} - -// Round-5 exact-shape fragment-interleaved B swizzle: for (k, n) == -// (6144, 2624) the weight is packed once, outside the timed region and out of -// Graph capture, into 512-byte 32-k x 16-n DUMMA B fragment blocks laid out -// [k32][n16][k8][n][8]: +// Exact-shape fragment-interleaved B swizzle: for (k, n) == (6144, 2624) the +// weight is packed once, outside the timed region and out of Graph capture, +// into 512-byte 32-k x 16-n DUMMA B fragment blocks laid out [k32][n16][k8][n][8]: // packed[((kk >> 5) * n16g + (n >> 4)) * 512 + // (((kk >> 3) & 3) * 16 + (n & 15)) * 8 + (kk & 7)] = raw[kk * n + n], -// where n16g = n >> 4 (164 for the exact shape). This is a byte-exact -// bijection 192 x 164 x 512 = k*n = 16,121,856: same byte count and buffer as -// the identity/n-major/panel packs, so captured addresses are unchanged. Each -// thread packs one 16-byte chunk = the 8-byte k-runs of two n-rows of one k8 -// group (a strided 2x8 gather on the read side; packing is outside timing, so -// the gather is off the critical path). Consecutive chunks are consecutive -// packed bytes, so the kernel's B staging is a flat 16-byte copy and every -// col_major B fragment read is one conflict-free ds_read_b64 (see -// load_frag8_swz). The round-2 n-major pack kernel above stays compiled as the -// layout control. +// 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, @@ -1437,8 +975,8 @@ __global__ __launch_bounds__(256) void w8a8_pack_bswz_i8_kernel( 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). + // 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]; } @@ -1462,10 +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. The workspace is unused for this shape -// (single-buffered, no split-K): m*n*4 = 42,990,080 B exceeds the 16 MiB -// workspace budget, so workspace_split_k_capacity is zero. +// 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, @@ -1486,41 +1024,23 @@ 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). The m >= 128 guard keeps the paired M=2/M=16 - // API shapes with the same (N, K) on the scalar fallback below. Round 2 - // keeps the accepted packed-weight architecture (the exact-fit 2-D macro-tile - // <128, 64, ...> with the packed weight produced once, out of the - // timed region, by launch_pack_w8a8_weight, the bank-safe LDS row strides and - // the explicit 8-byte load_frag8 fragment loads) and changes ONLY the K-stage - // pipeline: the dispatch goes to the 64-K double-buffered prefetch kernel - // <128, 64, 64> instead of the single-buffered <128, 64, 128> one, so the - // stage-(s+1) VMEM latency overlaps the stage-s LDS/v_mmac work. Round 3 - // retains that double-buffered 64-K pipeline unchanged (it is the unique - // 2-resident-block K pipeline inside the 65,536 B LDS budget: 2 x 27,648 = - // 55,296 <= 65,536) and changes only the packed B layout consumed by the two - // B global-load sites -- the weight is now the B-panel layout, so a warp's B - // staging loads touch 4 fully-used 128-byte lines instead of 32 half-used - // ones. Round 5 keeps the kernel, geometry and occupancy identical and - // changes only the packed B layout again: the B-panel pack is replaced by the - // fragment-interleaved swizzle, so the B stage is a flat 16-byte copy and - // every col_major B fragment is one conflict-free ds_read_b64 - // (load_frag8_swz). M = 32*128 and N = 41*64 are both exact, so no tail - // predication is introduced. Round 9 keeps that kernel, geometry and - // occupancy and only selects a compile-time tail-predication specialization: - // when m is a multiple of BM = 128 every block tile is fully inside the - // output, so the instantiation with kFullTile == true drops the dead A-row - // zero-fill/select in both staging sites and the dead epilogue row mask - // (see the Round-9 header section); every other m keeps the guarded - // instantiation byte-for-byte. + // (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)); - if ((m & 127) == 0) { + 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<128, 64, 64, true>), + w8a8_dumma_prefill_packedb_db_kernel< + kPackedBM, kPackedBN, 64, kPackedThreads, true>), grid, block, 0, @@ -1536,7 +1056,8 @@ extern "C" void launch_w8a8_gemm( } else { hipLaunchKernelGGL( HIP_KERNEL_NAME( - w8a8_dumma_prefill_packedb_db_kernel<128, 64, 64, false>), + w8a8_dumma_prefill_packedb_db_kernel< + kPackedBM, kPackedBN, 64, kPackedThreads, false>), grid, block, 0, @@ -1555,8 +1076,8 @@ extern "C" void launch_w8a8_gemm( // 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. + // 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), @@ -1608,7 +1129,7 @@ extern "C" void launch_w8a8_gemm( k); } -// Stable host launch symbol consumed by csrc/bindings.cpp. Round 5: the exact +// 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)] @@ -1628,7 +1149,7 @@ extern "C" void launch_pack_w8a8_weight( constexpr int kPackThreads = 256; const int64_t weight_count = static_cast(k) * n; if (k == kExactK && n == kExactN) { - // Round 5: one thread packs one 16-byte chunk (16 bytes = 2 n-rows x 8 B). + // 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) { 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/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/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/qkv_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/qkv_proj.hip index e0c06dd2..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,459 +1,467 @@ -// @@variant shape=hy3_tp8_qkv_proj_m16 commit=f2b771e350c9fea3b13d59b5fec4219786976bcc added=2026-09-09 -// median_us=19.3 p90_us=19.32 speedup=3.875 baseline_us=74.78 -// source=hy3-dsh-tp8-m16-9-8-0161e718 -// W8A8 INT8 GEMM HIP implementation for Hygon K500SM_AI / gfx928. +// @@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). // -// Worker: worker_0 (physical GPU 0), shape hy3_tp8_qkv_proj_m16 -// M = 16, N = 1280, K = 4096 (logical weight layout [K, N]). +// 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 0 (bootstrap): correctness-first scalar kernel (one thread per -// output element) served every shape and measured 336 us on this shape. +// 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 1 (DUMMA bootstrap, this round): replace the scalar kernel for -// the exact guard (m, n, k) == (16, 1280, 4096) with the minimal native -// gfx928 DUMMA INT8 m16n16k32 tile per the round mandate: -// - one 64-lane wavefront per block, one 16x16 output tile per block -// (grid = N/16 = 80 blocks), no cross-wave barrier; -// - A and B fragments are loaded DIRECTLY from their row-major global -// tensors every 32-K step (no LDS staging, no split-K, no workspace, -// no combine kernel in the timed region); -// - explicit int32 accumulation: du_fill_fragment(acc, 0) then one -// du_mma_sync per ascending 32-K step (exact int32 dot, same sum as -// the scalar kernel since int32 addition is exact and order-free); -// - single-wave LDS epilogue: du_store_matrix_sync materializes the -// 16x16 int32 tile in LDS (library-defined ownership), one -// __syncthreads() reached by all 64 threads, then the 64 lanes scale -// with x_scale[row] * weight_scale[n0+col] (scalar-path float order) -// and store bf16 through the manual round-to-nearest-even conversion. -// This establishes the exact DUMMA API family below so later rounds can -// tune geometry (wavefronts/block, LDS staging, split-K planes). +// 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. // -// Iteration 2 (architecture round, this round): measure grid parallelism -// before polishing. The iteration-1 geometry (80 one-wave blocks over 120 -// CUs, 0.67 blocks/CU) is far below the two-blocks-per-CU latency-hiding -// target, and the measured 130.43 us median (vs 74.782 us Triton Graph -// baseline) is dominated by underfilled CUs and per-block serial load->mma -// chains (128 direct-load 32-K steps per block). This round replaces the -// exact-guard dispatch with a grid-level split-K=3 geometry, still one -// 64-lane wavefront per block and one 16x16 N tile per wave (zero in-block -// barriers): -// - grid = (N/16) * SPLIT_K = 80 * 3 = 240 blocks = exactly 2 blocks per -// device CU (120 CUs), split id on gridDim.y, tile id on gridDim.x; -// - the 4,096-long K axis is partitioned into three legal 32-aligned, -// non-uniform, ascending slices of {1376, 1376, 1344} rows (43/43/42 -// m16n16k32 steps), each block du_fill_fragment(acc,0) then runs its -// own ascending 32-K DUMMA chain over its slice only (bit-exact int32 -// partial dot: |partial| <= 6.6e7 << 2^31, order-free); -// - each block materializes its 16x16 int32 partial (LDS epilogue, same -// single-wave store/barrier pattern as iteration 1) into its own plane -// of the caller's int32 workspace: partials[split][16][N]; every launch -// overwrites every plane, so no workspace clear is needed; -// - a small combine kernel (80 blocks x 64 lanes, second launch on the -// same stream inside the timed API) sums the SPLIT_K planes ascending -// per output element (exact int32, identical to the scalar k-ascending -// sum) and applies x_scale[row] * weight_scale[col] then bf16 RNE, -// keeping every output bit-identical to the CPU reference. -// Split-K combine is part of the operator wall and stays inside the timed -// region/Graph. Workspace need = SPLIT_K*M*N*4 = 245,760 B, well under the -// contract's 16-plane 1,310,720 B allocation for this shape. +// 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. // -// Iteration 3 (architecture round, this round): continue the mandated -// split-K occupancy sweep. Iteration 2's falsifiable claim held (240 -// one-wave blocks = 2 blocks/CU cut the wall from 130.43 to 68.65 us -// median, 1.089x over the 74.782 us Triton Graph baseline), but the -// measured logical rate (86 GB/s, 2.44 TOPS) is still far below the -// headroom of this decode family, so the sweep goes finer inside the -// trusted CU-aligned probe set {2,3,4,5,6,8,9,12}: SPLIT_K=6 -> grid = -// 80 N tiles x 6 K-slices = 480 one-wave zero-in-block-barrier blocks = -// exactly 4 blocks per device CU, and each block's serial 32-K chain -// halves to {22,22,21,21,21,21} steps = rows {704,704,672,672,672,672} -// (the existing generic slice formula; disjoint, ascending, 32-aligned, -// exact cover of [0,4096)). The partial and combine kernels are -// UNCHANGED code -- both were already generic over gridDim.y / -// split_count -- so this round only moves the launcher geometry constant -// kSplitK from 3 to 6. Workspace need = 6*M*N*4 = 491,520 B, still -// inside the caller's 16-plane 1,310,720 B allocation; the combine kernel -// (80 blocks x 64 lanes) sums the six planes ascending per output element -// (exact int32 -> bit-identical to the scalar k-ascending sum) and stays -// in the timed Graph. +// Iteration 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 (architecture/pipeline round, this round): multi-N-tile -// reuse. The iteration-3 code object (PMC digest 26d6df45...) shows the -// poison direct-load pattern per 32-K step: 16 global_load_ubyte (8 per -// A/B fragment) with ~10 scattered s_waitcnt vmcnt before byte reassembly -// and one v_mmac_i32_16x16x32_i8, and no cross-iteration load overlap, so -// at 4 resident one-wave blocks/CU every block's serial chain is one -// ~load-latency step per m16n16k32 mma (52.19 us median; the wall is the -// per-step latency of a one-mma step, not raw issue or HBM traffic). -// This round keeps the same split-K partial + combine architecture and -// the same 4 blocks/CU co-residency but widens each one-wave block to TWO -// adjacent 16-wide N tiles (32 columns): per 32-K step the wave loads the -// A fragment ONCE and issues two du_mma_sync (independent accumulators), -// so the x_q A bytes [16 x 32] of the shared step are reused across the -// two output tiles (per-output A global loads/reassembly halve) while -// each B byte is still read exactly once (no B reuse). To keep 480 -// one-wave blocks = exactly 4/CU over 120 CUs the grid becomes 40 N -// groups x SPLIT_K=12 K-slices ({11x8, 10x4} 32-K steps = rows -// {352x8, 320x4}, disjoint ascending exact cover of [0,4096), every -// boundary a DUMMA-k multiple; existing generic slice formula). Per-block -// serial chain drops to <= 11 steps of TWO mmas, so if the iteration-3 -// wall is per-step load->mma latency the per-output step latency halves -// and the unprofiled Graph median/P90 must drop materially below -// 52.194/52.269 us. Workspace need = 12*M*N*4 = 983,040 B, still inside -// the caller's 16-plane 1,310,720 B allocation; the combine kernel (80 -// blocks x 64 lanes) is UNCHANGED code and sums the twelve planes -// ascending per output element (exact int32 -> bit-identical to the -// scalar k-ascending sum), staying in the timed Graph. The only edits -// are the widened per-wave N tile in the partial kernel plus the launcher -// geometry; the workspace-less fallback keeps the iteration-1 one-wave -// direct kernel byte-for-byte and the scalar fallback still serves every -// unmatched shape including the paired M=2 API shape. +// Iteration 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 5 (packed-weight/staging round, this round): B-only LDS slice -// staging. The iteration-4 code object (PMC digest 1a4e30d0..., vmem_read -// 122,880 = 24 byte-granular global_load_ubyte per 32-K step x 5120 -// block-steps) shows the remaining wall is the poison direct-global B half -// of the per-step load->wait->reassembly->mma chain: per step the wave -// still pays 16 global byte loads (B0/B1) plus ~10 scattered s_waitcnt -// vmcnt on the MMA critical path even though every byte is L2-hot in the -// timed Graph replay. Per the lineage rule, direct-global fragment loads -// are poison: this round removes the B half by staging each block's whole -// 32-column B slice (<= 352 rows x 32 B = 11,264 B) into LDS ONCE at block -// start with 16-B int4 cooperative loads (64 lanes x 16 B per wavefront -// instruction, 4-deep source batching so one vmcnt drain covers 4 -// in-flight round trips), one same-wave __syncthreads(), then every -// 32-K step reads the B0/B1 fragments from LDS with du_load_matrix_sync -// (~30-cycle LDS latency, no global round trip) at row stride 32 -- the -// staged tile is an exact byte mirror of the block's global B strip, so -// fragment contents, the k-ascending int32 order and the partial planes -// are bit-identical to iteration 4. A stays direct-global on purpose: -// this round isolates the B-side hypothesis (A-only staging and the -// packed [N,K] layout remain untested axes). Geometry is UNCHANGED (40 -// N groups x SPLIT_K=12 = 480 one-wave blocks = exactly 4 blocks/CU over -// 120 CUs; per-block LDS 11,264 B + 2 x 1,024 B acc tiles = 13,312 B, -// 65,536/13,312 = 4.9, so the 4/CU co-residency is preserved and LDS does -// not become the occupancy limiter). If the per-step wall is the B-load -// half of the global wait ladder, the official unprofiled Graph median/P90 -// must drop materially below 41.154/41.233 us (partial-kernel profiled -// 36.32 us); if instead the median stays flat or regresses, the wall at 4 -// waves/CU is A-side latency or per-CU issue (16 vs 24 loads immaterial), -// and the next round must stage A too (LDS-only K loop) or raise -// co-residency -- never a scalar-fallback change. Exactness is preserved -// for every shape: the (m,n,k)==(16,1280,4096) guard is unchanged; the -// staged bytes are a copy of the same global bytes the direct loader -// read, the workspace-less fallback keeps the iteration-1 direct kernel -// byte-for-byte (still correct for this shape), the scalar fallback still -// serves every unmatched shape including the paired M=2 API shape, and no -// pack/layout/env/stream/config change is made; all launches stay on the -// caller stream inside the timed API/Graph. +// 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. // -// Iteration 6 (infrastructure failure, no valid artifact): the round-6 -// attempt was killed by the agent infrastructure before it produced a -// candidate result; no round-6 source change was accepted and the session -// was reset to the iteration-5 accepted best source below. The failure does -// not count as a completed optimization round. +// 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. // -// Iteration 7 (staging round, this round): A+B whole-slice LDS staging -// (LDS-only K loop). The iteration-5 result (39.404/39.485 us official -// Graph median/P90 vs iteration-4 41.154/41.233) falsified its own claim: -// the predicted 25-33 us for the B-half hypothesis did not appear, so per -// the iteration-5 contingency text this round stages the shared A operand -// too. The iteration-5 PMC (source digest 5fb42aff..., grid 480 = 40 N -// groups x SPLIT_K=12, vmem_read 46,080, l2_misses 98,879, counter-derived -// HBM ~187.9 GB/s << peak) still shows 8 byte-granular global_load_ubyte of -// the shared A fragment plus ~8 scattered s_waitcnt vmcnt before shift/or -// reassembly on the MMA path every 32-K step (~40,960 of the 46,080 -// vmem_read over 5,120 block-steps): pure per-step A-side global latency -// serialization, not traffic. This round stages BOTH operands once per -// block with 16-B int4 cooperative loads (A slice: 16 m-rows x 256 B + 16 B -// pad/row = 4,352 B at row stride 272 for bank skew; B strip: 256 x 32 = -// 8,192 B), one same-wave __syncthreads(), then an LDS-only K loop (A -// fragment at row stride 272, B0/B1 at row stride 32 via -// du_load_matrix_sync) with zero per-step global loads/waits. Staged tiles -// are exact byte mirrors of the global operands, so fragment bytes, the -// k-ascending int32 order and the partial planes are bit-identical to -// iteration 5. Geometry forced by LDS capacity: a whole-slice A stage does -// not fit at SPLIT_K=12 with 4 resident blocks/CU (A 16 x 352 = 5,632 B + -// B 11,264 B + acc 2,048 B = 18,944 B -> floor(65,536/18,944) = 3), so the -// guard launch moves to SPLIT_K=16: sixteen UNIFORM 8-step slices (128 == -// 16 x 8, rows {256 x 16}; every boundary a DUMMA-k multiple), grid = 40 N -// groups x 16 = 640 one-wave blocks ~ 5.33/CU over 120 CUs, per-block LDS -// 4,352 + 8,192 + 2,048 = 14,592 B -> floor(65,536/14,592) = 4 resident -// blocks/CU, i.e. the iteration-5 co-residency is preserved. Workspace -// need = 16 int32 planes = 1,310,720 B = the caller's full fixed 16-plane -// allocation for this shape, still gated by the launcher's workspace_bytes -// >= kWorkspaceNeeded check (undersized workspaces keep the workspace-less -// fallback). The combine kernel (80 blocks x 64 lanes, second same-stream -// launch in the timed API/Graph) is UNCHANGED code, generic over -// split_count = 16, and sums the sixteen planes ascending per output -// element (exact int32 -> bit-identical to the scalar k-ascending sum), -// staying in the timed Graph. All fallbacks, guards, pack/layout/env/ -// stream/config behavior are unchanged. +// Iteration 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 8 (occupancy/resource round, this round): raise the waves-per- -// block occupancy limiter from 1 to 2 wavefronts per block (blockDim 64 -> -// 128). PMC of the accepted iteration-7 code object (source digest -// b64e3209..., grid 640, workgroup 64, arch_vgpr 40, sgpr 48, scratch 0, -// LDS/block 14,592 B -> floor(65,536/14,592) = 4 resident blocks/CU = -// 5.33 waves/CU by grid) shows the residual wall is per-SIMD single-wave -// latency exposure, not issue or HBM traffic: each resident wave is alone -// on its SIMD (4 blocks x 1 wave -> 4 SIMDs), so every per-step -// s_waitcnt lgkmcnt after the staged-fragment du_load_matrix_sync and every -// du_mma accumulator dependency stalls that SIMD with no second independent -// wave to issue (VALU issue utilization ~3-6% of the 26.56 us profiled -// kernel; only 12 one-time staging global loads per block, vmem_read -// 7,680). The fix keeps the iteration-7 geometry byte-for-byte in every -// respect that affects partials (SPLIT_K=16, 40 x 16 = 640 blocks, 8-step -// slices rows {256 x 16}, staged exact byte mirrors, k-ascending int32 -// order, 16-plane workspace) and splits the one-time staging ACROSS the two -// waves (wave 0 copies the A slice, wave 1 copies the B strip; exactly the -// same global bytes read once -- no repeated HBM reads traded for -// occupancy), then each wave runs its OWN 8-step accumulator chain for one -// n16 tile of the block's 32 columns (one du_mma per step per wave; total -// MMA count unchanged). LDS per block stays 14,592 B -> still 4 resident -// blocks/CU, but now 8 resident waves/CU = 2 per SIMD, so a stalling wave -// is covered by its co-resident sibling. The only deliberately added cost -// is LDS re-reads: the shared A fragment is now loaded by both waves -// (predicted lds_instructions 99,840 -> ~125-135k with conflicts scaling at -// the measured ~2.56/instr); LDS is not the wall at ~3-6% issue -// utilization. The combine kernel, guards, workspace-less fallback, scalar -// fallback and pack identity are all unchanged. +// Iteration 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 9 (combine restructure round, this round): attack the OTHER half -// of the operator wall. PMC of the accepted iteration-8 code object (source -// digest beb25b76..., official Graph median/P90 28.884/28.933 us) profiles -// BOTH timed dispatches of the operator replay: the 640-block partial kernel -// at 19.361 us AND the 80-block combine kernel at 16.960 us -- the combine is -// 46.7% of the 36.321 us profiled operator aggregate, not a small tail. Its -// structure explains the cost: 80 one-wave 64-lane blocks (<= 80 of 480 -// SIMDs busy, ~40 CUs fully idle) whose lanes each run FOUR output elements -// through a DEPENDENT per-plane chain -- one 4-B global load + s_waitcnt -// vmcnt + int32 add per plane, 16 serial DRAM round trips per element (the -// exact code object's static body is one load/wait/add per plane iteration) -// -- against plane data that is L2-cold at combine start (combine -// l2_hit_rate 6.9%, l2_misses 21,596 ~= the whole 1.31 MB of partial planes, -// evicted by the partial kernel's 5.37 MB streaming read): DRAM-latency -// serialization on only 1.34 MB of reads, plus underfill. This round -// restructures ONLY the combine kernel (the partial kernel, SPLIT_K=16 -// 640-block geometry, staged byte mirrors, k-ascending int32 partials, -// 16-plane 1,310,720 B workspace contract, workspace-less fallback, exact -// guard and scalar fallback are byte-for-byte unchanged): the guard launcher -// switches the combine grid from 80 x 1 one-wave blocks to a 2-D grid of 80 -// n-tiles x 4 four-row m-groups = 320 one-wave blocks, and the kernel's -// split_count == 16 specialization gives every lane exactly ONE output -// element that issues ALL sixteen plane loads as INDEPENDENT reads into a -// 16-int register array (one vmcnt drain, 16-way memory-level parallelism -// per lane), then sums the registers in ascending split order s = 0..15 -- -// the identical int32 sequence as the old ascending loop -- and applies the -// unchanged float scale order and f32_to_bf16_rne, so every output element -// stays bit-identical (this mirrors the validated gate_up S=16 combine -// recipe: issue all 16 plane loads into a register array before the -// ascending sum). Same stream, same two launches, no allocation or sync -// inside the timed API/Graph. +// 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. // -// Iterations 10-15 (five-wave occupancy arc on this guard, all archived at -// iterations/hy3_tp8_qkv_proj_m16/iterationN): iteration 10 (direct-register -// epilogue on the two-wave mapping: acc tiles + epilogue barrier removed, -// LDS 14,592 -> 12,544 B = 5 resident blocks/CU = 10 waves/CU) measured -// 20.827/21.585 us and was NOT accepted; iteration 11 raised waves-per-block -// 2 -> 5 (blockDim 320, grid 16 x 16 = 256 blocks, LDS 29,952 B = 2 resident -// blocks/CU = 10 waves/CU) but FAILED exact correctness (its source declared -// no __launch_bounds__ while the launcher requests 320 threads, so the -// compiled descriptor capped the workgroup at 256 and the fifth wavefront's -// tile was absent: 4,091/20,480 mismatches, first at (m,n)=(0,64)); the -// iteration-11 repair was completed in iteration 13 by declaring -// __launch_bounds__(320) (exact correctness passed, 0 mismatches; measured -// 23.764/24.086 us); iteration 14 (ping/pong fragment double-buffer) passed -// correctness but the toolchain spilled the second fragment set to the -// private segment on every K step (.private_segment_fixed_size 48) -> 39.472 -// us, closing the register-pipelining axis; iteration 15 replaced the -// five-wave epilogue __syncthreads() with __threadfence_block() (the acc -// tiles are wave-private and the waves' plane stores are disjoint, so the -// cross-wave join was pure skew cost) and measured 20.613/20.641 us, correct -// but still not below the accepted two-wave 20.179/20.216. +// 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. // -// Iteration 16 (this round): close the five-wave occupancy axis with the -// measured occupancy evidence and apply the repair arc's validated epilogue -// discovery to the accepted mapping. Occupancy ledger of the partial -// kernel, all correct and bit-identical: 8 resident waves/CU (4 x 2-wave -// blocks, LDS 14,592 B) = 20.179/20.216 us (iteration 9, accepted); the two -// independent 10 resident-waves/CU configurations -- iteration 10 (5 x -// 2-wave blocks) 20.827/21.585 and iteration 15 (2 x 5-wave blocks) -// 20.613/20.641 -- BOTH trail the 8-wave configuration, i.e. per-SIMD wave -// density 2.5 (some SIMDs carry 3 waves while others idle) is worse than a -// clean 2.0 on this kernel, and the five-wave mapping (2 or 1 resident -// 5-wave blocks/CU = 10 or 5 waves/CU only) can never reach the 8-wave -// sweet spot. The accepted iteration-9 source below still carries the SAME -// provably-unnecessary epilogue __syncthreads() cross-wave join whose -// five-wave form cost ~3.15 us (iteration 13 23.764 -> iteration 15 20.613 -// with the join removed). Fix (the ONLY code change of this round; the -// two-wave geometry, blockDim 128, grid 40 x 16 = 640 blocks, staging -// partition, stage byte mirrors, 8-step K loop, combine kernel and its -// 320-block launch, workspace, guards and fallbacks are byte-for-byte the -// accepted iteration-9 candidate, verified by diff): replace the partial -// kernel's epilogue __syncthreads() with __threadfence_block(), exactly as -// iteration 15 validated on the five-wave form of this same kernel. -// acc_tile0/acc_tile1 are each written by du_store_matrix_sync and read -// back only by the SAME wave (w == 0 -> acc_tile0, w == 1 -> acc_tile1), and -// the two waves' plane stores touch disjoint 16-column ranges, so no -// epilogue stage has any cross-wave dependency; __threadfence_block() still -// orders each wave's own outstanding LDS writes before its read-back -// ds_reads (wavefront-wide lgkmcnt drain) while costing zero cross-wave -// wait. The staging __syncthreads() above stays (wave 0 stages the A slice -// that wave 1 also reads). Every partial byte and combine output byte is -// unchanged, so the output is bit-identical to iteration 9. +// 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. // -// Iteration 18 (this round): compile-time 8-step K-loop specialization of -// the partial kernel, applied to the current checked-in source (the -// iteration-16 candidate; working-tree digest cbc40932... = the iteration-16 -// archive, verified by diff). Round 17 drafted exactly this edit and -// verified its codegen locally, but the agent session was killed by the -// round infrastructure before a valid proposal landed -- no source change -// was ever applied, no build/benchmark ran, and per policy that -// infrastructure failure does not count toward the valid-round ledger, so -// this round re-submits the same HIP-only change as a fresh valid round. -// The iteration-11 repair item is closed since iteration 13 (archived -// iteration-11 five-wave source + __launch_bounds__(320) ONLY, exact-correct -// with 0 mismatches; iterations 13-16 completed that arc and iteration 16 -// re-based on the accepted two-wave mapping). Status of the current source: -// the iteration-16 epilogue barrier-decoupled two-wave partial kernel -// (blockDim 128, grid 40 x 16 = 640 blocks, LDS 14,592 B = 4 resident -// blocks/CU = 8 resident waves/CU) measured official Graph median/P90 -// 19.9885/20.0133 us (min 19.9557 / max 21.3734, correctness 0 mismatches, -// Graph capture and paired M=2 fallback passed), improving the accepted -// 20.179/20.216 by 0.95%; it is kept as the experimental base ("official -// best is unchanged until cumulative improvement reaches the normal -// acceptance threshold"), and this round keeps building on that source. -// The partial kernel dominates the operator wall, and the exact-source code -// object shows its K loop is ROLLED: the generic loop is runtime-bounded -// (my_steps from the runtime k and gridDim.y args), so the backend emits one -// static body of [ds_read2_b32 (A) + 8 x ds_read_u8 (B) + ~25-VALU byte -// reassembly + 4 lgkmcnt drains + one v_mmac_i32_16x16x32_i8] with a -// loop-carried WAR on the single fragment register set: step s+1's LDS -// reads cannot issue until step s's v_mmac has consumed the same fragment -// registers, so every step's full LDS latency sits serially on the -// accumulator chain with only 2 resident waves/SIMD to cover it (reused PMC -// of this exact geometry: lds_instructions 107,520, lds_wait_instructions -// 19,597, lds_bank_conflicts 276,480). The guard launch is the kernel's -// ONLY call site and always runs the uniform 8-step slice (k == 4096, -// gridDim.y == 16 -> my_steps == 8 == kStageMaxSteps, rows == 256 == -// kStageMaxRows, staged capacity exactly filled), so this round specializes -// that case: when my_steps == kStageMaxSteps the K loop runs with a -// COMPILE-TIME trip count of 8 under #pragma unroll; the backend then fully -// unrolls the chain, renames the fragment registers per step and issues -// each step's LDS reads after the previous step's v_mmac has consumed its -// registers, overlapping the next step's LDS round trip with the current -// v_mmac's execution (no loop-carried WAR; verified in the local exact-flag -// code object of this edit: 8 static v_mmac + 72 static ds_reads in the -// specialized path vs 1 static v_mmac / 10 static ds_reads in the rolled -// predecessor, vgpr 34 / sgpr 39 / private_segment 0 / group_segment 14,592 -// B and exactly ONE s_barrier unchanged -> the 4-blocks/CU = 8-waves/CU -// occupancy is preserved). Both branches issue the identical ascending-s -// du_mma sequence into the same accumulator (v_mmac acc operand v[1:4] -// throughout) over the same staged byte mirrors, so every partial byte and -// every combine output byte is bit-identical to iteration 16 and exact -// correctness must pass with 0 mismatches; the generic rolled loop stays -// verbatim as the else path for any non-guard (my_steps != 8) -// configuration. Geometry, staging partition, stage byte mirrors, the -// blockDim 128 / grid 40 x 16 = 640 blocks launch, the epilogue -// (acc-tile store + __threadfence_block + read-back plane store), the -// combine kernel and its 320-block launch, the 16-plane 1,310,720 B -// workspace need, the workspace-less iteration-1 fallback, the exact -// (m,n,k)==(16,1280,4096) guard and the scalar fallback for every unmatched -// shape (including paired M=2) are all unchanged. +// 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. // -// Iteration 20 (DS-side staging-latency round, accepted): the exact-guard -// launch always stages the full 512-chunk B strip (chunks == 512 == 8 x 64 -// lanes), so wave 1's B copy now runs as ONE straight-line 8-deep batched -// group (all eight 16-B source loads issue before the first LDS store, one -// vmcnt drain covers the whole strip) instead of two serialized 4-deep -// iterations -- the second global round trip is removed from every block's -// staging join critical path (official Graph median/P90 19.967 -> 19.541/ -// 19.558 us, accepted). The generic rolled loop stays verbatim as the else -// path for any non-guard (chunks != 512) configuration. +// 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. // -// Iteration 21 (this round): DS-side bank-skew round for the B-fragment -// ds_read_u8 structure (pre-registered in iterations 18 and 20: the next -// increment must attack the PMC-documented DS ledger of this exact object -- -// lds_bank_conflicts 276,480 = ~2.57 replays per LDS instruction, dominated -// by the 8 byte-granular B-fragment reads per step per wave, 128 per block). -// The row-major B fragment loader reads each lane's 8 bytes at addresses -// base + (8*(lane>>4) + i) * ldm for i = 0..7, so the four row groups -// g = lane>>4 of the fragment's 32 k rows sit 8*ldm bytes apart; with the -// stage_b rows packed at the old ldm == 32, every group sat 256 B (== 0 -// mod 128 B) apart and each ds_read_u8 replayed the SAME LDS bank 4x per -// wavefront. This round pads each staged B row to kStageBRowStride = 36 B -// (32 data + 4 pad): the groups then sit 288 B (== 32 mod 128 B) apart and -// hit banks {b, b+8, b+16, b+24} -- the B-fragment reads become -// conflict-free while every staged data byte, the fragment contents, the -// k-ascending int32 order and the partial planes stay bit-identical (the -// pads are never written or read; only the loader's ldm argument and the -// wave-1 store offsets change). Wave-1 staging stores become four 4-B -// words per 16-B chunk (36*r is only 4-B aligned) at int32 offsets -// 9*r + 4*h + {0,1,2,3}; LDS store bandwidth is unchanged. LDS per block -// grows 14,592 -> 15,616 B (stage_b 8,192 -> 9,216): 4 x 15,616 = 62,464 -// <= 65,536, so the 4-resident-blocks/CU = 8-resident-waves/CU occupancy -// sweet spot is preserved. The workspace-less fallback, the exact -// (m,n,k)==(16,1280,4096) guard, the scalar fallback for every unmatched -// shape (including paired M=2), the combine kernel and its 320-block -// launch, the pack identity and the 16-plane 1,310,720 B workspace need are -// all unchanged. The iteration-11 repair item stays closed since -// iteration 13 (five-wave archive + __launch_bounds__(320) ONLY, exact- -// correct; the five-wave axis closed on measurement: best 20.613/20.641 us -// vs the accepted two-wave mapping 19.541/19.558 us). -// -// Iterations 22-25 (rounds 22 and 23 measured and rejected, tree reverted to -// the iteration-21 source e7ae27d0...; round 24 was killed mid-round before -// any measurement -- its archive draft was never compiled or run): round 22 -// removed the acc-tile LDS round trip from the partial epilogue at pinned -// occupancy (direct-register plane stores) and REGRESSED (+0.82 us -> -// 20.276/20.298 us) because the register-distributed acc values force 16-B -// store granules per wavefront instruction instead of the read-back path's -// four full 64-B lines -- rejected and reverted; round 23 widened each -// combine block to a 32-column n-strip with dwordx2 plane loads (160 blocks) -// and measured FLAT (median 19.4504 us vs the 19.4528 shadow, within its own -// +/-0.05 us falsification band, plus a 41.1 us max outlier) -- rejected and -// reverted. Round 23's pre-registered falsification branch (a) therefore -// committed THIS round to a combine block mapping whose per-plane reads run -// in longer contiguous runs per row (or a partial-plane store-order change), -// never occupancy / fragment duplication / raw asm / a scalar-fallback -// change. -// -// Iteration 25 (this round): combine mapping v3 -- one output row per -// block. The combine byte stream is fixed (every plane word is read -// exactly once, 64-B sectors fully consumed), so v1 (iteration 9: blocks of -// 4 rows x 16 cols) and v2 (iteration 23: 4 rows x 32 cols) both issue -// every per-plane wavefront load as FOUR 64-B (v1) or 128-B (v2) runs at -// the 5,120-B m-row stride -- four different DRAM rows, i.e. up to four row -// activations per load instruction -- which is exactly why widening the -// per-row run 64 -> 128 B in v2 could not move the combine (it never -// reduced the rows-per-instruction scatter). The guard combine launch now -// uses one-wave blocks that each own ONE m-row x 64 consecutive columns -// (kDummaCombineColsPerBlock = 64; grid 20 x 16 = 320 blocks, the same -// 320-block count and the same one-output-element-per-lane, 16-independent- -// plane-load structure as iteration 9): every per-plane wavefront load then -// reads 64 lanes x 4 B = 256 B CONTIGUOUS = four consecutive 64-B lines of -// ONE DRAM row (single run, ~4x fewer DRAM row activations for the same -// byte stream). Each element's ascending-s int32 sum, the float scale -// order and every output bf16 word are unchanged element-for-element, so -// the output is bit-identical to the iteration-21 shadow; the partial -// kernel, the 16-plane workspace contract, the exact guard, the -// workspace-less fallback and the scalar fallback (paired M=2 included) are -// byte-for-byte untouched. -// -// launch_pack_w8a8_weight (optional, out of the timed region) is an -// identity device-to-device copy: the logical [K, N] int8 weight and the -// [N, 1] fp32 scales are copied unchanged, so the DUMMA direct loader -// reads the exact logical [K, N] row-major layout. Later rounds may -// change the packed layout here together with the matching GEMM -// interpretation. -// -// This file is compiled with hipcc by torch.utils.cpp_extension; it must -// stay free of torch/ATen dependencies. Only the two extern "C" launchers -// below are the host contract consumed by csrc/bindings.cpp. -// -// Include order is the known-good DTK order: HIP runtime first, then the -// bfloat16 header, then du_mma.h (consumed by the DUMMA kernel below). +// 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 @@ -463,662 +471,629 @@ namespace { -// IEEE round-to-nearest-even conversion of an fp32 value to a bf16 bit -// pattern, implemented with pure integer ops. This is exactly the -// conversion applied by the CPU reference (`scaled.to(torch.bfloat16)`), -// independent of which bf16 type family the DTK headers provide. -__device__ __forceinline__ uint16_t f32_to_bf16_rne(float f) { - constexpr uint32_t kRoundBias = 0x7FFFu; - uint32_t u; - __builtin_memcpy(&u, &f, sizeof(u)); - const uint32_t lsb = (u >> 16) & 1u; // tie goes to even - u += kRoundBias + lsb; - return static_cast(u >> 16); -} +// 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; -} // namespace +// 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 generic GEMM: out[m, n] = bf16(int32(sum_k a[m,k]*b[k,n]) -// * x_scale[m] * weight_scale[n]). -// One thread per output element. Adjacent lanes walk adjacent columns n -// (the fastest-changing output/weight dimension). gfx928 wavefront = 64; -// the launcher always uses blockDim 128 (a multiple of 64). -// --------------------------------------------------------------------------- -__global__ void w8a8_scalar_gemm_kernel( - const int8_t* __restrict__ a, // x_q [M, K] row-major - const int8_t* __restrict__ b, // packed_weight [K, N] row-major - const float* __restrict__ x_scale, // [M] (logical [M, 1]) - const float* __restrict__ w_scale, // [N] (logical [N, 1]) - uint16_t* __restrict__ out, // [M, N] bf16 payload +// 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, int m, int n, int k) { - const int tid = - static_cast(blockIdx.x) * static_cast(blockDim.x) + - static_cast(threadIdx.x); - const int total = m * n; - if (tid >= 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 = tid / n; // M is small for decode; N is the fast axis - const int col = tid - row * n; - const int8_t* a_row = a + static_cast(row) * k; - const int8_t* b_col = b + col; // stride n along k + 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; -#pragma unroll 4 - for (int kk = 0; kk < k; ++kk) { - acc += static_cast(a_row[kk]) * - static_cast(b_col[static_cast(kk) * n]); + 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; + } } - const float scaled = static_cast(acc) * x_scale[row] * w_scale[col]; - out[tid] = f32_to_bf16_rne(scaled); + // 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[idx] = __float2bfloat16(scaled); } // --------------------------------------------------------------------------- -// Exact-guard DUMMA path for hy3_tp8_qkv_proj_m16 (M=16, N=1280, K=4096). +// 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-thread wavefront per block, one 16x16 output tile per -// block (grid = N/16 = 80 blocks). Every K step is an explicit int32 -// du_mma_sync over a 32-aligned ascending k chunk, so the int32 dot is -// identical to the scalar kernel's sum (exact integer addition, no -// overflow: |out| <= K*127*127 ~= 6.6e7 << 2^31). +// 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). // -// The gfx928 int8 m16n16k32 accelerator fragment ownership is opaque, so -// the epilogue materializes the tile with du_store_matrix_sync into a -// single-wave LDS staging array (library-defined mem_row_major placement, -// 256 x int32 = 1,024 B) and only afterwards applies the per-element float -// scaling and the bf16 RNE store. BlockDim is one wavefront, so the single -// __syncthreads() below is a same-wave barrier (all 64 threads reach it) -// and there is no cross-wave synchronization anywhere in the kernel. +// 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. // --------------------------------------------------------------------------- constexpr int kDummaM = 16; constexpr int kDummaN = 16; constexpr int kDummaK = 32; constexpr int kDummaWaveSize = 64; -// Combine constants. Iteration 9: the exact-guard launcher (SPLIT_K == 16) -// launched the combine kernel as 80 n-tiles x four-row m-groups of one-wave -// blocks and the kernel specialized split_count == kDummaCombineSplits with -// one output element per lane. Iteration 25 (mapping v3): the guard launch -// becomes 20 column-groups x 16 m-rows of one-wave blocks -- each block owns -// ONE output row x kDummaCombineColsPerBlock = 64 consecutive columns (4 x -// kDummaN), so every per-plane wavefront load reads a single 256-B -// contiguous run (four consecutive 64-B lines inside one DRAM row) instead -// of four 64/128-B runs at the 5,120-B row stride (see the kernel comment -// below for the DRAM row-activation mechanism). -constexpr int kDummaCombineSplits = 16; // == the guard launcher kSplitK -constexpr int kDummaCombineColsPerBlock = 64; // n-columns per combine block - -__global__ void w8a8_dumma_m16_tile_kernel( - const int8_t* __restrict__ a, // x_q [M, K] row-major - const int8_t* __restrict__ b, // packed_weight [K, N] row-major - const float* __restrict__ x_scale, // [M] - const float* __restrict__ w_scale, // [N] - uint16_t* __restrict__ out, // [M, N] bf16 payload - int n, - int k) { - const int lane = static_cast(threadIdx.x); // blockDim == 64 - const int n0 = static_cast(blockIdx.x) * kDummaN; - - __shared__ __align__(16) int32_t acc_tile[kDummaM * kDummaN]; - - du::dumma::DUFragment - a_frag; - du::dumma::DUFragment - b_frag; - du::dumma::DUFragment - acc_frag; - du::dumma::du_fill_fragment(acc_frag, 0); - - // A is [16, k] (row stride k) and B is [k, 16] (row stride n); both stay - // row-major in their global tensors, so each 32-K step loads straight - // from global with the library row_major loaders. - for (int k0 = 0; k0 < k; k0 += kDummaK) { - du::dumma::du_load_matrix_sync(a_frag, a + k0, k); - du::dumma::du_load_matrix_sync( - b_frag, b + static_cast(k0) * n + n0, n); - du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } +// 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; - du::dumma::du_store_matrix_sync(acc_tile, acc_frag, kDummaN, - du::dumma::mem_row_major); - __syncthreads(); +// 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; - // 256 elements / 64 lanes: each lane scales and stores four outputs. -#pragma unroll - for (int linear = lane; linear < kDummaM * kDummaN; linear += kDummaWaveSize) { - const int row = linear / kDummaN; - const int col = linear - row * kDummaN; - const int out_col = n0 + col; - const float scaled = static_cast(acc_tile[linear]) * - x_scale[row] * w_scale[out_col]; - out[static_cast(row) * n + out_col] = f32_to_bf16_rne(scaled); - } -} +// 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 -// --------------------------------------------------------------------------- -// Split-K partial kernel (iteration 2 probe, widened in iteration 4 to -// multi-N-tile reuse, B-staged in iteration 5): same m16n16k32 int8 DUMMA -// tile as iteration 1, but gridDim = (N/(16*kTilesPerWave), SPLIT_K) so -// every one-wave block owns kTilesPerWave adjacent 16-wide N tiles AND one -// ascending 32-aligned K slice, and writes its exact int32 partials into -// workspace planes partials[split][16][n]. Iteration 4: kTilesPerWave = 2, -// so per 32-K step the wave loads the A fragment once and shares it between -// the two du_mma_sync calls (A bytes of the step reused across both output -// tiles; every B byte is still read exactly once). Iteration 5 (staging -// round, B-only): each block first stages its whole 32-column B slice -// (my_steps*32 rows x 32 B, <= 352 x 32 = 11,264 B under the SPLIT_K=12 -// guard launch) into an LDS tile stage_b that is an EXACT byte mirror of -// the block's global B strip (16-B int4 cooperative loads, 4-deep source -// batching so one vmcnt drain covers 4 in-flight round trips), then after -// one same-wave __syncthreads() the K loop reads the B0/B1 fragments from -// LDS with du_load_matrix_sync (row stride 32) -- same fragment bytes as -// the iteration-4 direct global loads, so the int32 partials stay -// bit-identical -- while A stays on the direct-global path (this round -// isolates the B side only). BlockDim is one 64-lane wavefront: the two -// barriers are the staging barrier above and the single same-wave -// __syncthreads() after the LDS materialization of both partial tiles. -// LDS per block: 11,264 B (stage_b) + 2 x 1,024 B (acc tiles) = 13,312 B -// -> 65,536 / 13,312 = 4.9, so the 4-blocks-per-CU launch co-residency is -// preserved (LDS is not the occupancy limiter). -// -// Iteration 7 (this round): A+B whole-slice LDS staging. Each one-wave -// block now stages BOTH operands of its whole K slice once: stage_a holds -// the 16 m-rows x 256 slice bytes at row stride 272 (256 data + 16 pad for -// bank skew, 4,352 B) and stage_b holds the 256 x 32 B strip (8,192 B); -// both use 16-B int4 cooperative loads with the same 4-deep batching, one -// same-wave __syncthreads(), and then the K loop reads every fragment from -// LDS (A at row stride 272, B0/B1 at row stride 32) -- zero per-step global -// loads and zero vmcnt waits on the MMA path. Geometry moved to SPLIT_K=16 -// (uniform 8-step slices, 40 x 16 = 640 blocks): per-block LDS 4,352 + -// 8,192 + 2,048 = 14,592 B -> floor(65,536/14,592) = 4, so the 4-blocks/CU -// co-residency of iteration 5 is preserved. Barriers per block stay two -// (staging, epilogue). Staged tiles are exact byte mirrors, so partials -// are bit-identical to iteration 5 and the combine kernel is unchanged. -// -// Iteration 8 (occupancy/resource round, this round): tune the waves-per- -// block occupancy limiter. PMC of the accepted iteration-7 code object -// (source digest b64e3209..., grid 640, workgroup 64, arch_vgpr 40, sgpr 48, -// scratch 0, LDS/block 14,592 B -> floor(65,536/14,592) = 4 resident blocks/ -// CU, 5.33 waves/CU by grid) shows a per-SIMD latency-exposure wall: every -// block is ONE wavefront, so at 4 resident blocks/CU the four waves land one -// per SIMD and each SIMD has no second independent wave to issue while the -// resident wave stalls on its per-step s_waitcnt lgkmcnt (lds_wait 8,359) -// after the staged-fragment du_load_matrix_sync and on its du_mma accum -// chain (VALU issue utilization ~3-6% of the 26.56 us profiled kernel; -// vmem_read 7,680 = only the 12 one-time staging loads per block). Fix: -// blockDim 64 -> 128, i.e. TWO 64-lane wavefronts per block, and the block's -// 32-column strip is split by ownership: wave w computes ONLY n-tile w (cols -// [n0+16w, n0+16w+16)) over the shared 8-step slice with its own independent -// accumulator chain, one du_mma per step per wave (same total MMA count). -// LDS per block is UNCHANGED (stage_a 4,352 + stage_b 8,192 + 2 x 1,024 acc -// tiles = 14,592 B), so 4 blocks/CU still fit but now carry 8 resident waves -// = 2 per SIMD: one wave's LDS->MMA stall is filled by the other. The -// one-time staging is partitioned, NOT duplicated: wave 0 copies the shared -// A slice and wave 1 copies the B strip (their iteration-7 loops unchanged), -// so exactly the same global bytes are read once (no repeated HBM reads for -// occupancy) and the staging barrier becomes a cross-wave __syncthreads(). -// Deliberate added cost: the A fragment is re-read from LDS by each wave -// (lds_instructions 99,840 -> ~125-135k predicted; LDS pipe/conflicts are -// not the wall at ~3-6% issue utilization -- the trade is LDS re-reads, not -// HBM traffic). Geometry, planes, combine kernel, guards and fallbacks are -// otherwise unchanged. -// --------------------------------------------------------------------------- -__global__ void w8a8_dumma_m16_splitk_partial_kernel( - const int8_t* __restrict__ a, // x_q [16, k] row-major - const int8_t* __restrict__ b, // packed_weight [k, n] row-major - int32_t* __restrict__ partials, // [gridDim.y][16][n] int32 +__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) { - constexpr int kTilesPerWave = 2; // two adjacent 16-wide N tiles per block - // Iteration 8 (occupancy round): each block is TWO 64-lane wavefronts - // (blockDim 128). Wave w owns N tile w of the block's 32 columns (cols - // [n0 + 16*w, n0 + 16*w + 16)) and runs an independent 8-step accumulator - // chain over the shared K slice for that tile alone. - const int tid = static_cast(threadIdx.x); // blockDim == 128 - const int w = tid >> 6; // wave id inside block 0/1 - const int lane = tid & (kDummaWaveSize - 1); // lane id within the wave - const int group = static_cast(blockIdx.x); // 32-wide n-group id - const int split = static_cast(blockIdx.y); // k-slice id + 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); - // Stage-buffer sizing is tied to the ONLY launch site of this kernel: the - // exact (m,n,k)==(16,1280,4096) guard with gridDim.y == 16 (iteration 7), - // so total_steps == 128 and every split owns exactly 8 steps = 256 rows - // (128 == 16 * 8; the generic slice formula below degenerates to uniform - // 8-step slices). Compile-time capacities: 256 rows of B per block and - // 16 m-rows x (256 + 16 pad) bytes of A per block. - constexpr int kStageSplitCount = 16; // gridDim.y of the guard launch - constexpr int kStageMaxSteps = - (128 + kStageSplitCount - 1) / kStageSplitCount; // 8 (uniform) - constexpr int kStageMaxRows = kStageMaxSteps * kDummaK; // 256 rows x 32 - constexpr int kStageARowStride = - kStageMaxRows + 16; // 272 B per A m-row: 256 data + 16 pad for banks - // Iteration 21 (DS-side bank-skew round): each staged B k-row now occupies - // kStageBRowStride = 36 B (32 data bytes + 4 pad bytes). The row-major B - // fragment loader reads a lane's 8 fragment bytes at addresses - // base + (8*(lane>>4) + i) * ldm, i = 0..7: the four lane groups - // g = lane>>4 sit 8*ldm bytes apart, so with the old ldm == 32 every group - // landed 256 B (== 0 mod 128 B) apart and each of the 8 ds_read_u8 per - // fragment replayed the SAME LDS bank 4x (the PMC-documented - // lds_bank_conflicts 276,480 = ~2.57 replays/instr on this object, ~89% of - // it from the B-fragment byte reads: 128 per block). ldm == 36 puts the - // groups 288 B (== 32 mod 128 B, 8 banks) apart, so the four groups hit - // banks {b, b+8, b+16, b+24} -- fully conflict-free -- while the 4-B pad - // per row keeps every row start 4-B aligned for the staging stores. - // stage_b grows 8,192 -> 9,216 B and the block total becomes 4,352 - // (stage_a) + 9,216 (stage_b) + 2,048 (acc tiles) = 15,616 B/block: - // 4 x 15,616 = 62,464 <= 65,536, so the 4-resident-blocks/CU = - // 8-resident-waves/CU occupancy sweet spot is preserved unchanged. - constexpr int kStageBRowStride = kDummaN * 2 + 4; // 36 B per staged B row - + // 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[kDummaM * kStageARowStride]; - __shared__ __align__(16) int8_t stage_b[kStageMaxRows * kStageBRowStride]; + __shared__ __align__(16) int8_t stage_a[kStageABytes]; + __shared__ __align__(16) int8_t stage_b[kStageBBytes]; du::dumma::DUFragment a_frag; + // 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; // this wave's B half of the strip (b_col == 16*w bytes) + signed char, du::dumma::col_major> + b_frag; // this wave's own 16-wide B half of the block's 32-column strip du::dumma::DUFragment acc_frag; // this wave's accumulator for its n16 tile du::dumma::du_fill_fragment(acc_frag, 0); - // 32-aligned non-uniform K slicing over the gridDim.y splits: split s owns - // my_steps ascending 32-K steps beginning at step_begin, slices are - // disjoint and cover [0, k) exactly once in ascending order, so the int32 - // partial sums and the ascending combine below reproduce the scalar - // k-ascending sum bit-for-bit. - const int total_steps = k / kDummaK; - const int split_count = static_cast(gridDim.y); - const int base_steps = total_steps / split_count; - const int rem_steps = total_steps - base_steps * split_count; - const int my_steps = base_steps + (split < rem_steps ? 1 : 0); - const int step_begin = - split * base_steps + (split < rem_steps ? split : rem_steps); - const int k_base = step_begin * kDummaK; // first global k row of slice - const int rows = my_steps * kDummaK; // slice rows: my_steps x 32 + // Stage-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; - // Stage the block's operands (iteration 8: PARTITIONED across the two - // wavefronts, not duplicated -- the union reads exactly the same global - // bytes as the iteration-7 single wave, so no byte is fetched twice and - // HBM traffic is unchanged). Wave 0 copies the shared A slice: the 16 - // m-rows x rows bytes of A (global columns [k_base, k_base + rows)), - // stored m-major in stage_a at row stride kStageARowStride (an exact byte - // mirror of global a[m][k_base .. k_base + rows)). Wave 1 copies the - // block's 32-column B strip: my_steps*32 rows x 32 B, stored row-major in - // stage_b at the 36-B row pitch kStageBRowStride (32 data + 4 pad bytes - // per row; iteration 21 bank skew -- an exact byte mirror of the global - // data rows [k_base, k_base + rows) x columns [n0, n0 + 32), pads never - // written or read). Each copy uses the - // unchanged iteration-7 cooperative code: 64 lanes x 16 B per wavefront - // instruction, chunks in LDS int4 units, loads batched 4-deep in source so - // one vmcnt drain covers 4 in-flight global round trips (the two waves' - // round trips overlap further across the SIMDs). Global m-rows are - // k-major at stride k and global B rows are n-major at stride n (k % 16 == - // 0 and n % 16 == 0 under the guard; k_base and n0 are 32-aligned), so - // every 16-B global load is 16-B aligned (the wave-1 LDS stores are - // 4-B-granular since iteration 21: 36*row is only 4-B aligned). - const int a_chunks = (rows * kDummaM) / 16; // rows/16 int4 per m-row - const int chunks = rows * (kDummaN * 2) / 16; // rows*2 int4 chunks - if (w == 0) { - const int per_row = rows / 16; - const int4* __restrict__ a4 = reinterpret_cast(a); - int4* __restrict__ ald4 = reinterpret_cast(stage_a); - const int kb4 = k_base >> 4; // global int4 column of the slice start - const auto ld_chunk = [&](int e) { - const int m = e / per_row; - const int t = e - m * per_row; - return a4[static_cast(m) * (k >> 4) + kb4 + t]; - }; - const auto st_chunk = [&](int e, const int4& v) { - const int m = e / per_row; - const int t = e - m * per_row; - ald4[m * (kStageARowStride / 16) + t] = v; - }; - int e = lane; - for (; e + 3 * kDummaWaveSize < a_chunks; e += 4 * kDummaWaveSize) { - const int4 v0 = ld_chunk(e); - const int4 v1 = ld_chunk(e + kDummaWaveSize); - const int4 v2 = ld_chunk(e + 2 * kDummaWaveSize); - const int4 v3 = ld_chunk(e + 3 * kDummaWaveSize); - st_chunk(e, v0); - st_chunk(e + kDummaWaveSize, v1); - st_chunk(e + 2 * kDummaWaveSize, v2); - st_chunk(e + 3 * kDummaWaveSize, v3); - } - for (; e < a_chunks; e += kDummaWaveSize) { - st_chunk(e, ld_chunk(e)); - } - } else { - const int4* __restrict__ b4 = reinterpret_cast(b); - const auto ld_chunk = [&](int e) { - const int r = e >> 1; - // Global int4 index of (global row k_base + r, col half e&1). - return b4[static_cast(k_base + r) * (n >> 4) + (n0 >> 4) + - (e & 1)]; - }; - // Iteration 21: stage_b rows now live at the 36-B pitch - // kStageBRowStride (32 data bytes + 4 pad per row), so a 16-B chunk - // (row r = e >> 1, col half h = e & 1 of that row) can no longer be - // stored with one 16-B-aligned ds_write_b128 (36*r is 16-B aligned only - // for r % 4 == 0). Each chunk is stored as four consecutive int32 words - // at word offset 9*r + 4*h + {0,1,2,3} == (36*r + 16*h)/4 + {0,1,2,3}: - // every data byte lands at exactly the same byte offset (row*36 + ncol) - // the loader reads back, so the staged data byte mirror of the global B - // strip is unchanged -- only the pad bytes (row*36 + 32..35) are never - // written and never read (the fragment loader only touches ncol < 32). - // LDS store bandwidth is unchanged (4 x 4-B words per chunk); the four - // words of one chunk are consecutive, so a lane's store stream stays - // dense per row half. - int32_t* __restrict__ lds32 = reinterpret_cast(stage_b); - const auto st_chunk = [&](int e, const int4& v) { - const int word0 = (e >> 1) * (kStageBRowStride >> 2) + (e & 1) * 4; - lds32[word0 + 0] = v.x; - lds32[word0 + 1] = v.y; - lds32[word0 + 2] = v.z; - lds32[word0 + 3] = v.w; - }; - // Iteration 20 (guard-slice staging latency specialization): the guard - // launch is the ONLY call site of this kernel and always stages the - // full 512-chunk B strip (chunks == rows*2 == 512 == 8 x 64 lanes, rows - // == 256 == kStageMaxRows, staged capacity exactly filled). The - // runtime-bounded 4-deep loop below then compiles to TWO serialized - // 4-deep iterations (each iteration issues four 16-B source loads and - // drains them with vmcnt waits before the next iteration's loads can - // issue -- two full global round trips on the block's staging critical - // path, one per iteration), and wave 0's A copy finishes in a single - // 4-deep group and then idles at the __syncthreads() join while wave 1 - // completes its second round trip. When chunks == 8*kDummaWaveSize the - // loop now runs as ONE straight-line 8-deep batched group: all eight - // 16-B source loads per lane are issued before the first LDS store, so - // one drain window covers the whole strip and the second round trip is - // removed from the join's critical path. Every chunk id [0, 512) is - // still loaded and stored exactly once by exactly one lane to the same - // stage_b byte (the chunk set, the store order and the byte mirror are - // identical to the rolled loop), so every staged byte, every partial - // and every combine output byte is bit-identical; the generic rolled - // loop stays verbatim as the else path for any non-guard (chunks != - // 512) configuration, which never executes at the only launch site. - if (chunks == 8 * kDummaWaveSize) { - const int4 v0 = ld_chunk(lane); - const int4 v1 = ld_chunk(lane + kDummaWaveSize); - const int4 v2 = ld_chunk(lane + 2 * kDummaWaveSize); - const int4 v3 = ld_chunk(lane + 3 * kDummaWaveSize); - const int4 v4 = ld_chunk(lane + 4 * kDummaWaveSize); - const int4 v5 = ld_chunk(lane + 5 * kDummaWaveSize); - const int4 v6 = ld_chunk(lane + 6 * kDummaWaveSize); - const int4 v7 = ld_chunk(lane + 7 * kDummaWaveSize); - st_chunk(lane, v0); - st_chunk(lane + kDummaWaveSize, v1); - st_chunk(lane + 2 * kDummaWaveSize, v2); - st_chunk(lane + 3 * kDummaWaveSize, v3); - st_chunk(lane + 4 * kDummaWaveSize, v4); - st_chunk(lane + 5 * kDummaWaveSize, v5); - st_chunk(lane + 6 * kDummaWaveSize, v6); - st_chunk(lane + 7 * kDummaWaveSize, v7); - } else { - int e = lane; - for (; e + 3 * kDummaWaveSize < chunks; e += 4 * kDummaWaveSize) { - const int4 v0 = ld_chunk(e); - const int4 v1 = ld_chunk(e + kDummaWaveSize); - const int4 v2 = ld_chunk(e + 2 * kDummaWaveSize); - const int4 v3 = ld_chunk(e + 3 * kDummaWaveSize); - st_chunk(e, v0); - st_chunk(e + kDummaWaveSize, v1); - st_chunk(e + 2 * kDummaWaveSize, v2); - st_chunk(e + 3 * kDummaWaveSize, v3); - } - for (; e < chunks; e += kDummaWaveSize) { - st_chunk(e, ld_chunk(e)); + // 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; + + // 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)); + } } } - } - __syncthreads(); // stage_a and stage_b ready for every lane of both waves - - // LDS-only K loop (iteration 8): every 32-K step wave w reads the A - // fragment from the staged A mirror at row stride kStageARowStride and its - // own B half from the staged B strip at the 36-B row pitch - // kStageBRowStride (byte column b_col; iteration 21 bank skew) - // -- the same fragment bytes the iteration-7 loads produced, with no - // global round trip and no vmcnt wait on the MMA path -- then issues ONE - // du_mma on its private accumulator. The two resident wavefronts carry - // independent 8-step chains, so a SIMD holding both can issue one wave's - // instructions while the other stalls on its s_waitcnt lgkmcnt / du_mma - // dependency (per-SIMD latency exposure was the iteration-7 wall). - const int b_col = w * kDummaN; // byte offset of this wave's n16 tile + // 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 - // Iteration 18 (compile-time 8-step K-loop specialization): the guard - // launch is the ONLY call site of this kernel and always runs the uniform - // 8-step slice (k == 4096, gridDim.y == 16 -> my_steps == 8 == - // kStageMaxSteps, rows == 256 == kStageMaxRows, staged capacity exactly - // filled). The runtime-bounded generic loop below compiles to a ROLLED - // body (verified in the exact-source code object: one static v_mmac per - // step with per-step ds_read2_b32 + 8 x ds_read_u8 + ~25-VALU byte - // reassembly + 4 lgkmcnt drains on the accumulator chain and a - // loop-carried WAR on the single fragment set), so every step's full LDS - // latency sits serially on the MMA path with only 2 resident waves/SIMD - // to cover it (PMC of the exact object: lds_instructions 107,520, - // lds_wait_instructions 19,597). Making the guard trip count - // compile-time 8 lets the backend fully unroll the chain, rename the - // fragment registers per step and hoist later steps' LDS reads across the - // current step's v_mmac (local exact-flag codegen check of this edit: 8 - // static v_mmac in the specialized path with the next steps' LDS reads - // issued before the current v_mmac; vgpr 34 / private_segment 0 / - // group_segment 14,592 B unchanged -> the 4-blocks/CU = 8-waves/CU - // occupancy is preserved). Both branches issue the identical - // ascending-s du_mma sequence into the same accumulator over the same - // staged byte mirrors, so every partial byte and every combine output - // byte is bit-identical to iteration 16; the generic rolled loop stays - // verbatim as the else path for any non-guard (my_steps != 8) - // configuration. - if (my_steps == kStageMaxSteps) { -#pragma unroll - for (int s = 0; s < kStageMaxSteps; ++s) { - du::dumma::du_load_matrix_sync( - a_frag, stage_a + static_cast(s) * kDummaK, - kStageARowStride); - du::dumma::du_load_matrix_sync( - b_frag, - stage_b + - static_cast(s) * kDummaK * kStageBRowStride + b_col, - kStageBRowStride); - du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - } else { - for (int s = 0; s < my_steps; ++s) { - du::dumma::du_load_matrix_sync( - a_frag, stage_a + static_cast(s) * kDummaK, - kStageARowStride); - du::dumma::du_load_matrix_sync( - b_frag, - stage_b + - static_cast(s) * kDummaK * kStageBRowStride + b_col, - kStageBRowStride); - du::dumma::du_mma_sync(acc_frag, a_frag, b_frag, acc_frag); + // 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(); } - // Epilogue (iteration 16, barrier-decoupled): each wave materializes its - // own tile into its own LDS acc buffer (wave 0 -> acc_tile0, wave 1 -> - // acc_tile1, 1,024 B each), then its 64 lanes read back that same buffer - // and store the 256 int32 partials (4 stores per lane) to - // plane[split][m][n0 + b_col + col]. acc_tile{w} is written and read ONLY - // by wave w, and the two waves' plane stores touch disjoint 16-column - // ranges, so nothing in the epilogue needs cross-wave convergence. The - // former epilogue __syncthreads() was a 128-thread cross-wave join whose - // release was gated by the slower of the block's two waves on a kernel - // that runs only 8 resident waves/CU (4 x 2-wave blocks) to hide the skew - // -- the same join whose five-wave form measured ~3.15 us of pure cost - // (iteration 13 23.764/24.086 -> iteration 15 20.613/20.641 us with the - // join removed, exact correctness passing in both). It is replaced by - // __threadfence_block(), which only orders THIS wave's own outstanding LDS - // writes (the du_store_matrix_sync ds_writes of its private acc tile) - // before its read-back ds_reads and costs no cross-wave wait. The staging - // __syncthreads() above stays (wave 0 stages the A slice and wave 1 stages - // the B strip; every wave reads bytes the other wave staged). Every - // partial byte, the k-ascending int32 order and every plane store address - // are unchanged, so the output is bit-identical to iteration 9 and exact - // correctness must pass with 0 mismatches. - if (w == 0) { - du::dumma::du_store_matrix_sync(acc_tile0, acc_frag, kDummaN, - du::dumma::mem_row_major); - } else { - du::dumma::du_store_matrix_sync(acc_tile1, acc_frag, kDummaN, - du::dumma::mem_row_major); - } - __threadfence_block(); // iteration 16: wave-local store->read ordering - // only; no cross-wave join (was __syncthreads()) + // 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(); - int32_t* plane = partials + static_cast(split) * kDummaM * n; - const int32_t* tile = (w == 0) ? acc_tile0 : acc_tile1; - const int out_col0 = n0 + b_col; + // 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; - plane[static_cast(row) * n + out_col0 + col] = tile[linear]; + const int64_t dst = + static_cast(row) * n + n0 + w * kDummaN + col; + plane[dst] = my_tile[linear]; } } -// --------------------------------------------------------------------------- -// Split-K combine kernel (iteration 2, restructured in iteration 9, mapping -// v3 in iteration 25): sums the gridDim.y int32 partial planes ascending per -// output element (exact int32, no overflow, bit-identical to the scalar -// k-ascending dot), applies x_scale[row] * weight_scale[col] in the scalar -// float order and stores bf16 through f32_to_bf16_rne. Iteration 9 added -// the guard specialization (split_count == 16): one output element per lane, -// all sixteen plane loads issued as INDEPENDENT reads into a 16-int register -// array (single vmcnt drain instead of a dependent load->wait->add chain of -// 16 DRAM round trips per element), then the ascending-s register sum and -// the scaled bf16 store. Iteration 25 (this round) keeps that per-lane -// structure and changes ONLY the block-to-output mapping: every one-wave -// block owns ONE m-row x kDummaCombineColsPerBlock = 64 consecutive columns -// (guard grid = 20 x 16 = 320 one-wave blocks, same block count as -// iteration 9). In the iteration-9/23 mappings each block covered 4 rows x -// 16/32 columns, so every per-plane wavefront load issued four 64/128-B runs -// at the 5,120-B m-row stride -- four different DRAM rows, i.e. up to four -// row activations per load instruction -- which is why widening the per-row -// run 64 -> 128 B in iteration 23 could not move the combine (it never -// reduced the rows-per-instruction scatter). With ONE row per block every -// per-plane load instruction is 64 lanes x 4 B = 256 B CONTIGUOUS = four -// consecutive 64-B lines of one DRAM row (a single run: ~4x fewer DRAM row -// activations for the identical 1.31 MB byte stream, identical 64-B sectors, -// identical wavefront-load count). The element set is unchanged, so every -// element's ascending-s int32 sum, the float scale order and every output -// bf16 word are bit-identical to the pre-iteration-25 combine; the generic -// loop below stays verbatim for every other call site / split count / grid -// shape. -// --------------------------------------------------------------------------- -__global__ void w8a8_dumma_m16_combine_kernel( - const int32_t* __restrict__ partials, // [split_count][16][n] int32 - const float* __restrict__ x_scale, // [16] - const float* __restrict__ w_scale, // [n] - uint16_t* __restrict__ out, // [16, n] bf16 payload +// 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 split_count) { - const int lane = static_cast(threadIdx.x); // blockDim == 64 - const int n0 = static_cast(blockIdx.x) * kDummaN; + int split_k) { + const int lane = static_cast(threadIdx.x); - // --------------------------------------------------------------------- - // Iteration 25 specialization (mapping v3; guard launch only: - // split_count == 16 and gridDim.y == kDummaM == 16 -> blockIdx.x = - // 64-column group in [0, 20), blockIdx.y = one m-row in [0, 16)). Each - // lane owns output element (row, col) with row = blockIdx.y and col = - // blockIdx.x * 64 + lane: the sixteen independent 4-B plane reads (plane - // stride kDummaM * n int32) batch into one vmcnt drain, and because the - // 64 lanes of the block are 64 CONSECUTIVE columns of ONE row, every - // per-plane wavefront load covers 256 contiguous bytes (4 consecutive - // 64-B lines of one DRAM row) instead of four 5,120-B-strided runs; the - // ascending-s register sum then reproduces the exact int32 sequence for - // this element. - // --------------------------------------------------------------------- - if (split_count == kDummaCombineSplits && gridDim.y == kDummaM) { + // 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) * kDummaCombineColsPerBlock + lane; + static_cast(blockIdx.x) * kCombineColsPerBlock + lane; const int64_t elem = static_cast(row) * n + col; - int32_t v[kDummaCombineSplits]; + int32_t plane[kCombineSplits]; #pragma unroll - for (int s = 0; s < kDummaCombineSplits; ++s) { - v[s] = partials[static_cast(s) * (kDummaM * n) + elem]; + 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 < kDummaCombineSplits; ++s) { - acc += v[s]; // ascending split order: identical int32 sum + 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] * w_scale[col]; - out[elem] = f32_to_bf16_rne(scaled); + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[elem] = __float2bfloat16(scaled); return; } -#pragma unroll - for (int linear = lane; linear < kDummaM * kDummaN; - linear += kDummaWaveSize) { - const int row = linear / kDummaN; - const int col = linear - row * kDummaN; - const int out_col = n0 + col; - int32_t acc = 0; - for (int s = 0; s < split_count; ++s) { - acc += partials[(static_cast(s) * kDummaM + row) * n + - out_col]; - } - const float scaled = - static_cast(acc) * x_scale[row] * w_scale[out_col]; - out[static_cast(row) * n + out_col] = f32_to_bf16_rne(scaled); + const int i = static_cast(blockIdx.x) * kCombineThreads + lane; + if (i >= total) { + return; + } + + int32_t acc = 0; + for (int s = 0; s < split_k; ++s) { + acc += partials[static_cast(s) * total + i]; } + + const int row = i / n; + const int col = i - row * n; + const float scaled = + static_cast(acc) * x_scale[row] * weight_scale[col]; + out[i] = __float2bfloat16(scaled); } -// --------------------------------------------------------------------------- -// Identity device-to-device packing (bootstrap): grid-stride byte copies on -// the caller stream. Valid for every (K, N), including unmatched shapes. -// --------------------------------------------------------------------------- -__global__ 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 stride = static_cast(gridDim.x) * blockDim.x; - for (int64_t i = static_cast(blockIdx.x) * blockDim.x + - static_cast(threadIdx.x); - i < count; i += stride) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { dst[i] = src[i]; } } -__global__ 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 stride = static_cast(gridDim.x) * blockDim.x; - for (int64_t i = static_cast(blockIdx.x) * blockDim.x + - static_cast(threadIdx.x); - i < count; i += stride) { + const int64_t i = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < count) { dst[i] = src[i]; } } -// --------------------------------------------------------------------------- -// Host launcher consumed by csrc/bindings.cpp (gemm_out, timed region). -// The exact (16, 1280, 4096) guard dispatches the split-K DUMMA partial + -// combine kernels (two wavefronts per partial block since iteration 8); -// the scalar fallback below launches 128 threads per 128 outputs. -// --------------------------------------------------------------------------- +// 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) { + __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 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 + extern "C" void launch_w8a8_gemm( const int8_t* a, const int8_t* b, @@ -1131,118 +1106,147 @@ 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 || a == nullptr || b == nullptr || - x_scale == nullptr || weight_scale == nullptr || out == nullptr) { + // 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 DUMMA guard: hy3_tp8_qkv_proj_m16 (M=16, N=1280, K=4096). - // Lineage: iteration 4 widened each one-wave block to TWO adjacent - // 16-wide N tiles (32 columns) sharing one A-fragment load per 32-K step - // at SPLIT_K=12 (grid 40 x 12 = 480 blocks = 4/CU over 120 CUs); - // iteration 5 staged each block's 32-column B slice into LDS once and - // read the B0/B1 fragments from LDS per step, A direct-global - // (39.404/39.485 us official Graph median/P90, from 41.154/41.233). - // Iteration 7: BOTH operands are staged per K slice -- the A fragment - // also comes from LDS, so the K loop is global-load-free -- which forces - // SPLIT_K=16: sixteen UNIFORM 8-step slices (128 == 16 x 8, rows - // {256 x 16}, exact ascending cover of [0, 4096), every boundary a - // DUMMA-k multiple), grid = 40 N groups x 16 = 640 one-wave blocks - // (~5.33 waves/CU, 4 resident blocks/CU by LDS: 14,592 B/block -> - // 65,536/14,592 = 4.4), each block writing exact int32 partial planes - // (cols [n0, n0+16) and [n0+16, n0+32) of plane[split]), then one - // UNCHANGED 80-block combine kernel (generic over split_count = 16) that - // sums the planes ascending and emits the scaled bf16 output. Staged - // tiles are exact byte mirrors of the global operands, so fragment - // bytes, K-slice bounds, the k-ascending int32 order and the partial - // plane layout are unchanged: the combine+scale output is bit-identical - // to iteration 5. - // Iteration 8 (this round): occupancy/resource change only -- the same - // 640-block SPLIT_K=16 geometry now runs TWO 64-lane wavefronts per block - // (blockDim 128): wave w computes tile w (cols [n0+16w, n0+16w+16)) of - // the block's 32-column strip with its own independent accumulator chain - // over the shared slice. LDS per block is UNCHANGED (14,592 B -> 4 - // resident blocks/CU), so co-residency rises from 4 to 8 resident waves - // per CU = 2 per SIMD, hiding each wave's per-step LDS->MMA stall behind - // the other wave. Staging is partitioned across the two waves (wave 0 = - // A slice, wave 1 = B strip; identical global bytes read exactly once). - // The combine kernel, the 16-plane workspace need (1,310,720 B), the - // k-ascending partial order and the bit-identical output are unchanged. - // Both launches stay on the caller stream inside the timed API/Graph; - // the combine cost is part of the operator wall. Requires 16 int32 - // partial planes (1,310,720 B) of the caller's workspace -- the full - // fixed 16-plane 1,310,720 B allocation for this shape; when the - // workspace is missing/undersized the exact guard falls back to the - // iteration-1 one-wave direct geometry (still no workspace). Every - // unmatched (m, n, k) -- including the paired M=2 API shapes with the - // same (N, K) -- keeps reaching the scalar fallback below. - // Iteration 9 (this round): combine-side restructure only. PMC of the - // accepted iteration-8 code object profiles the combine kernel at 16.960 - // us vs 19.361 us for the partial kernel (46.7% of the 36.321 us profiled - // operator aggregate) with a 6.9% L2 hit rate and only 80 one-wave blocks, - // so the guard now launches the SAME combine kernel on a 2-D grid (80 - // n-tiles x 4 four-row m-groups = 320 one-wave blocks) and the kernel's - // split_count==16 specialization gives each lane one output element that - // batches all sixteen plane reads into a register array before the - // ascending sum. Partials, output bytes, bit-exactness and the partial - // kernel launch are unchanged. - // --------------------------------------------------------------------- - if (m == 16 && n == 1280 && k == 4096) { - constexpr int kDummaThreads = 64; // one wavefront (tile/combine) - constexpr int kPartialThreads = 2 * kDummaThreads; // 2 waves per block - constexpr int kTilesPerWave = 2; // two 16-wide N tiles per block (n32) - constexpr int kSplitK = 16; // 40 groups x 16 = 640 blocks, 4/CU - const int ntiles = n / kDummaN; // 1280 / 16 = 80 combine tiles - const int ngroups = ntiles / kTilesPerWave; // 80 / 2 = 40 n32 groups - const int64_t kWorkspaceNeeded = - static_cast(kSplitK) * kDummaM * n * sizeof(int32_t); - auto* out_bf16 = reinterpret_cast(out); - if (workspace != nullptr && workspace_bytes >= kWorkspaceNeeded) { - auto* partials = reinterpret_cast(workspace); - hipLaunchKernelGGL(w8a8_dumma_m16_splitk_partial_kernel, - dim3(ngroups, kSplitK), dim3(kPartialThreads), 0, - stream, a, b, partials, n, k); - // Iteration 25 (mapping v3): 2-D combine grid = 20 column-groups x 16 - // one-row blocks = 320 one-wave blocks (one output row x 64 consecutive - // columns per block, one output element per lane, 16-deep independent - // plane-load register batch in the kernel's split_count==16 path; every - // per-plane wavefront load is one 256-B contiguous run of one DRAM row - // instead of four 5,120-B-strided runs). - hipLaunchKernelGGL(w8a8_dumma_m16_combine_kernel, - dim3(n / kDummaCombineColsPerBlock, kDummaM), - dim3(kDummaThreads), 0, stream, - partials, x_scale, weight_scale, out_bf16, n, - kSplitK); + __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 = + 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 { - // Workspace-less fallback: iteration-1 one-wave direct geometry. - hipLaunchKernelGGL(w8a8_dumma_m16_tile_kernel, - dim3(ntiles), dim3(kDummaThreads), 0, stream, - a, b, x_scale, weight_scale, out_bf16, n, k); + split_k = 0; + } + if (split_k > total_stages) { + split_k = total_stages; + } + + 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); + } + return; } - return; } - // --------------------------------------------------------------------- - // Scalar generic fallback: every shape outside the exact guard above. - // --------------------------------------------------------------------- + const int64_t total = + static_cast(m) * static_cast(n); + const int64_t blocks = (total + kBlockThreads - 1) / kBlockThreads; - constexpr int kBlockThreads = 128; // multiple of the gfx928 wavefront 64 - const int total = m * n; - const int grid = (total + kBlockThreads - 1) / kBlockThreads; - auto* out_bf16 = reinterpret_cast(out); - hipLaunchKernelGGL(w8a8_scalar_gemm_kernel, - dim3(grid), dim3(kBlockThreads), 0, stream, - a, b, x_scale, weight_scale, out_bf16, m, n, k); + // 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); } -// --------------------------------------------------------------------------- -// Host launcher for the optional out-of-timed-region weight packing. -// Bootstrap: identity device-to-device copy of the int8 [K, N] weight and -// the fp32 [N, 1] scales; valid for every (K, N), unmatched shapes included. -// --------------------------------------------------------------------------- extern "C" void launch_pack_w8a8_weight( const int8_t* raw_weight, const float* weight_scale, @@ -1251,27 +1255,28 @@ extern "C" void launch_pack_w8a8_weight( int k, int n, hipStream_t stream) { - constexpr int kPackBlockThreads = 256; - - const int64_t i8_count = static_cast(k) * n; - if (i8_count > 0 && raw_weight != nullptr && packed_weight != nullptr) { - const int i8_grid = - static_cast((i8_count + kPackBlockThreads - 1) / - kPackBlockThreads); - hipLaunchKernelGGL(w8a8_identity_copy_i8_kernel, - dim3(i8_grid), dim3(kPackBlockThreads), 0, stream, - raw_weight, packed_weight, i8_count); + if (raw_weight == nullptr || weight_scale == nullptr || + packed_weight == nullptr || packed_weight_scale == nullptr) { + return; + } + if (k <= 0 || n <= 0) { + return; } - const int64_t f32_count = static_cast(n); - if (f32_count > 0 && weight_scale != nullptr && - packed_weight_scale != nullptr) { - const int f32_grid = - static_cast((f32_count + kPackBlockThreads - 1) / - kPackBlockThreads); - hipLaunchKernelGGL(w8a8_identity_copy_f32_kernel, - dim3(f32_grid), dim3(kPackBlockThreads), 0, stream, - weight_scale, packed_weight_scale, f32_count); + 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_gate_up_proj.hip b/metainfer/tasks/dcu_kernel_auto_opt/variant/int8w8a8-gemm/hy3/TP8/M16/shared_gate_up_proj.hip index 1f40f800..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,393 +1,206 @@ -// @@variant shape=hy3_tp8_shared_gate_up_proj_m16 commit=2321c42b8618aae23700bc2cf3532b59c8f6c8d8 added=2026-08-27 -// baseline_us=66.527 speedup=6.40915 -// median_us=10.38 p90_us=10.42 -// source=hy3-dsh-tp8-m16-2-dc53295a +// @@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. @@ -400,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) { @@ -467,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) + @@ -486,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 @@ -536,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( @@ -587,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) @@ -633,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 = @@ -707,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 @@ -823,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(); @@ -848,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, @@ -946,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, @@ -1012,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, @@ -1038,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, @@ -1094,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, @@ -1147,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) { @@ -1166,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(