hotkernel 0.1.0__tar.gz → 0.1.1__tar.gz

{hotkernel-0.1.0 → hotkernel-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: hotkernel
-Version: 0.1.0
+Version: 0.1.1
 Summary: Find the hot kernel, cool it down: an Nsight/SASS MCP server with a verify-and-record contract and a kernel-speedrun loop for coding agents.
 Project-URL: Homepage, https://github.com/xinhuagu/hotkernel
 Project-URL: Issues, https://github.com/xinhuagu/hotkernel/issues
@@ -36,7 +36,7 @@ English | [中文](https://github.com/xinhuagu/hotkernel/blob/main/README_zh.md)
 
 </div>
 
-Ask Claude to "make this kernel faster" today and it guesses: it cannot see the Nsight report, does not know which kernel eats 60% of the step, has never read the SASS the compiler produced, and cannot benchmark its own rewrite on your card. hotkernel is the missing pair of hands. Nine tools wrap `nsys`, `ncu`, `cuobjdump`, `torch.profiler` and a strict verify-and-record contract; one skill turns them into a loop: profile, pick the hot kernel, read its SASS, check it against the roofline, rewrite it in Triton, verify against the baseline, record the verdict, keep or revert, repeat.
+Ask Claude to "make this kernel faster" today and it guesses: it cannot see the Nsight report, does not know which kernel eats 60% of the step, has never read the SASS the compiler produced, and cannot benchmark its own rewrite on your card. hotkernel is the missing pair of hands. Ten tools wrap `nsys`, `ncu`, `cuobjdump`, `torch.profiler` and a strict verify-and-record contract; one skill turns them into a loop: profile, pick the hot kernel, read its SASS, check it against the roofline, rewrite it in Triton, verify against the baseline, record the verdict, keep or revert, repeat.
 
 **Three targets, three autonomous agent runs, one RTX 5090: rmsnorm 3.2× faster than eager PyTorch, softmax 1.2× faster than `torch.softmax`, gemv level with cuBLAS. Every kernel the agent kept runs at 95–105% of the card's measured streaming bandwidth.**¹
 
@@ -44,14 +44,31 @@ Ask Claude to "make this kernel faster" today and it guesses: it cannot see the
 
 ## Install
 
+Claude Code, one command (Claude Code 2.1.275 or later): it adds this repository as a plugin marketplace and installs the plugin, which is the MCP server plus the kernel-speedrun skill.
+
+```bash
+claude plugin install hotkernel --marketplace xinhuagu/hotkernel
+```
+
+Inside a session the same is `/plugin marketplace add xinhuagu/hotkernel` followed by `/plugin install hotkernel@hotkernel`. The plugin has one option, the Python that has your PyTorch and Triton: set it with `/plugin configure hotkernel@hotkernel`, or at install time with `claude plugin install hotkernel@hotkernel --config hotkernel_python=$HOME/venvs/mytorch/bin/python`. Left empty, hotkernel uses the first `python` on PATH that can import torch, and says so if there is none. The server itself stays torch-free and runs as `uvx hotkernel`.
+
+Codex and other MCP clients (Cursor, DSH) keep the plain MCP setup:
+
 ```bash
-claude mcp add -e HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel -- uvx hotkernel   # Claude Code
-codex mcp add --env HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel -- uvx hotkernel # Codex
+codex mcp add --env HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel -- uvx hotkernel
 ```
 
-Any other MCP client: `{"command": "uvx", "args": ["hotkernel"], "env": {"HOTKERNEL_PYTHON": "..."}}`. `HOTKERNEL_PYTHON` is the interpreter that has your PyTorch and Triton; the server itself stays torch-free.
+or `{"command": "uvx", "args": ["hotkernel"], "env": {"HOTKERNEL_PYTHON": "..."}}` in the client's MCP config, plus a copy of [`skills/kernel-speedrun/`](https://github.com/xinhuagu/hotkernel/tree/main/skills/kernel-speedrun) in the agent's skills directory. Claude Code without the plugin works the same way: `claude mcp add -e HOTKERNEL_PYTHON=... hotkernel -- uvx hotkernel`.
+
+## Point it at your own kernel
+
+`hotkernel new` turns any PyTorch function into a target with the same contract as the examples: the function becomes the baseline, its fp32 result the reference, tolerances default by dtype, and the protocol-3 timer is vendored next to the scripts, so they run in your torch interpreter whether or not hotkernel is installed there.
+
+```bash
+hotkernel new gelu --from torch.nn.functional:gelu --shape 4096x8192 --dtype bf16 --out targets
+```
 
-The loop is a skill: copy [`skills/kernel-speedrun/`](https://github.com/xinhuagu/hotkernel/tree/main/skills/kernel-speedrun) into your agent's skills directory (Claude Code: `~/.claude/skills/`).
+writes `targets/gelu_baseline.py`, `targets/gelu_candidate.py` (starts as the baseline; this is the file the loop edits) and `targets/_hotkernel_contract.py`. A function of your own works the same way, with extra inputs in call order: `--from ./mykernels.py:fused --extra w:8192`. Then `bench` the baseline for the number to beat, and `verify` the untouched candidate against it, which must come back `inconclusive` before anything changed. Agents get the same through the `new_target` tool. On the reference card the untouched gelu candidate benched at 0.0952 ms and verified `inconclusive`; a deliberately wrong one came back `incorrect`.
 
 ## Before / after
 
@@ -80,6 +97,7 @@ Every row on [LEADERBOARD.md](https://github.com/xinhuagu/hotkernel/blob/main/LE
 | `bench` | run a script n times, collect its `{"ms": …, "ok": …}` line, return mean/std/min | your torch |
 | `verify` | run baseline and candidate interleaved, n times each; speedup and a verdict: `faster`, `slower`, `inconclusive`, `incorrect` (any `ok=false`, no speedup reported), `invalid` (shape, dtype, seed or timing protocol changed) | your torch |
 | `record` | write the verify result to `results/<target>/`, regenerate the leaderboard from the logs, and say whether the change beats the run's best so far | nothing |
+| `new_target` | scaffold a target from your own function: baseline, blank candidate and the vendored contract, the same as `hotkernel new` | nothing |
 
 `torch_trace`, `bench`, `verify` and `record` work on any machine with PyTorch. Nsight is optional and only makes the picture sharper.
 

{hotkernel-0.1.0 → hotkernel-0.1.1}/README.md

@@ -14,7 +14,7 @@ English | [中文](https://github.com/xinhuagu/hotkernel/blob/main/README_zh.md)
 
 </div>
 
-Ask Claude to "make this kernel faster" today and it guesses: it cannot see the Nsight report, does not know which kernel eats 60% of the step, has never read the SASS the compiler produced, and cannot benchmark its own rewrite on your card. hotkernel is the missing pair of hands. Nine tools wrap `nsys`, `ncu`, `cuobjdump`, `torch.profiler` and a strict verify-and-record contract; one skill turns them into a loop: profile, pick the hot kernel, read its SASS, check it against the roofline, rewrite it in Triton, verify against the baseline, record the verdict, keep or revert, repeat.
+Ask Claude to "make this kernel faster" today and it guesses: it cannot see the Nsight report, does not know which kernel eats 60% of the step, has never read the SASS the compiler produced, and cannot benchmark its own rewrite on your card. hotkernel is the missing pair of hands. Ten tools wrap `nsys`, `ncu`, `cuobjdump`, `torch.profiler` and a strict verify-and-record contract; one skill turns them into a loop: profile, pick the hot kernel, read its SASS, check it against the roofline, rewrite it in Triton, verify against the baseline, record the verdict, keep or revert, repeat.
 
 **Three targets, three autonomous agent runs, one RTX 5090: rmsnorm 3.2× faster than eager PyTorch, softmax 1.2× faster than `torch.softmax`, gemv level with cuBLAS. Every kernel the agent kept runs at 95–105% of the card's measured streaming bandwidth.**¹
 
@@ -22,14 +22,31 @@ Ask Claude to "make this kernel faster" today and it guesses: it cannot see the
 
 ## Install
 
+Claude Code, one command (Claude Code 2.1.275 or later): it adds this repository as a plugin marketplace and installs the plugin, which is the MCP server plus the kernel-speedrun skill.
+
+```bash
+claude plugin install hotkernel --marketplace xinhuagu/hotkernel
+```
+
+Inside a session the same is `/plugin marketplace add xinhuagu/hotkernel` followed by `/plugin install hotkernel@hotkernel`. The plugin has one option, the Python that has your PyTorch and Triton: set it with `/plugin configure hotkernel@hotkernel`, or at install time with `claude plugin install hotkernel@hotkernel --config hotkernel_python=$HOME/venvs/mytorch/bin/python`. Left empty, hotkernel uses the first `python` on PATH that can import torch, and says so if there is none. The server itself stays torch-free and runs as `uvx hotkernel`.
+
+Codex and other MCP clients (Cursor, DSH) keep the plain MCP setup:
+
 ```bash
-claude mcp add -e HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel -- uvx hotkernel   # Claude Code
-codex mcp add --env HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel -- uvx hotkernel # Codex
+codex mcp add --env HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel -- uvx hotkernel
 ```
 
-Any other MCP client: `{"command": "uvx", "args": ["hotkernel"], "env": {"HOTKERNEL_PYTHON": "..."}}`. `HOTKERNEL_PYTHON` is the interpreter that has your PyTorch and Triton; the server itself stays torch-free.
+or `{"command": "uvx", "args": ["hotkernel"], "env": {"HOTKERNEL_PYTHON": "..."}}` in the client's MCP config, plus a copy of [`skills/kernel-speedrun/`](https://github.com/xinhuagu/hotkernel/tree/main/skills/kernel-speedrun) in the agent's skills directory. Claude Code without the plugin works the same way: `claude mcp add -e HOTKERNEL_PYTHON=... hotkernel -- uvx hotkernel`.
+
+## Point it at your own kernel
+
+`hotkernel new` turns any PyTorch function into a target with the same contract as the examples: the function becomes the baseline, its fp32 result the reference, tolerances default by dtype, and the protocol-3 timer is vendored next to the scripts, so they run in your torch interpreter whether or not hotkernel is installed there.
+
+```bash
+hotkernel new gelu --from torch.nn.functional:gelu --shape 4096x8192 --dtype bf16 --out targets
+```
 
-The loop is a skill: copy [`skills/kernel-speedrun/`](https://github.com/xinhuagu/hotkernel/tree/main/skills/kernel-speedrun) into your agent's skills directory (Claude Code: `~/.claude/skills/`).
+writes `targets/gelu_baseline.py`, `targets/gelu_candidate.py` (starts as the baseline; this is the file the loop edits) and `targets/_hotkernel_contract.py`. A function of your own works the same way, with extra inputs in call order: `--from ./mykernels.py:fused --extra w:8192`. Then `bench` the baseline for the number to beat, and `verify` the untouched candidate against it, which must come back `inconclusive` before anything changed. Agents get the same through the `new_target` tool. On the reference card the untouched gelu candidate benched at 0.0952 ms and verified `inconclusive`; a deliberately wrong one came back `incorrect`.
 
 ## Before / after
 
@@ -58,6 +75,7 @@ Every row on [LEADERBOARD.md](https://github.com/xinhuagu/hotkernel/blob/main/LE
 | `bench` | run a script n times, collect its `{"ms": …, "ok": …}` line, return mean/std/min | your torch |
 | `verify` | run baseline and candidate interleaved, n times each; speedup and a verdict: `faster`, `slower`, `inconclusive`, `incorrect` (any `ok=false`, no speedup reported), `invalid` (shape, dtype, seed or timing protocol changed) | your torch |
 | `record` | write the verify result to `results/<target>/`, regenerate the leaderboard from the logs, and say whether the change beats the run's best so far | nothing |
+| `new_target` | scaffold a target from your own function: baseline, blank candidate and the vendored contract, the same as `hotkernel new` | nothing |
 
 `torch_trace`, `bench`, `verify` and `record` work on any machine with PyTorch. Nsight is optional and only makes the picture sharper.
 

{hotkernel-0.1.0 → hotkernel-0.1.1}/README_zh.md

@@ -14,7 +14,7 @@
 
 </div>
 
-今天你让 Claude"把这个 kernel 跑快点"，它只能靠猜：它看不到 Nsight 的报告，不知道哪个 kernel 吃掉了一步的 60%，没读过编译器生成的 SASS，也没法在你的卡上测自己改的版本。hotkernel 就是它缺的那双手。九个工具包住 `nsys`、`ncu`、`cuobjdump`、`torch.profiler` 和一套严格的验证与记录协议；一个 skill 把它们串成循环：剖析、选最耗时的 kernel、读它的 SASS、对照 roofline、用 Triton 重写、对着基线验证、记录裁决、保留或回滚、再来。
+今天你让 Claude"把这个 kernel 跑快点"，它只能靠猜：它看不到 Nsight 的报告，不知道哪个 kernel 吃掉了一步的 60%，没读过编译器生成的 SASS，也没法在你的卡上测自己改的版本。hotkernel 就是它缺的那双手。十个工具包住 `nsys`、`ncu`、`cuobjdump`、`torch.profiler` 和一套严格的验证与记录协议；一个 skill 把它们串成循环：剖析、选最耗时的 kernel、读它的 SASS、对照 roofline、用 Triton 重写、对着基线验证、记录裁决、保留或回滚、再来。
 
 **三个目标，三次 agent 自主运行，一张 RTX 5090：rmsnorm 比 PyTorch 朴素实现快 3.2 倍，softmax 比 `torch.softmax` 快 1.2 倍，gemv 与 cuBLAS 持平。agent 保留下来的每个 kernel 都跑到了这张卡实测流式带宽的 95% 到 105%。**¹
 
@@ -22,14 +22,31 @@
 
 ## 安装
 
+Claude Code 一条命令（Claude Code 2.1.275 及以上）：把本仓库加为插件市场并装上插件，插件 = MCP server + kernel-speedrun skill。
+
+```bash
+claude plugin install hotkernel --marketplace xinhuagu/hotkernel
+```
+
+在会话里等价于先 `/plugin marketplace add xinhuagu/hotkernel` 再 `/plugin install hotkernel@hotkernel`。插件只有一个选项：装了你 PyTorch 和 Triton 的那个 Python。用 `/plugin configure hotkernel@hotkernel` 设置，或者安装时直接给：`claude plugin install hotkernel@hotkernel --config hotkernel_python=$HOME/venvs/mytorch/bin/python`。留空时 hotkernel 用 PATH 上第一个能 `import torch` 的 `python`，一个都没有就明确报错。server 本身不依赖 torch，以 `uvx hotkernel` 运行。
+
+Codex 和其他 MCP 客户端（Cursor、DSH）仍按普通 MCP 方式配置：
+
 ```bash
-claude mcp add -e HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel -- uvx hotkernel   # Claude Code
-codex mcp add --env HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel -- uvx hotkernel # Codex
+codex mcp add --env HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel -- uvx hotkernel
 ```
 
-其他 MCP 客户端：`{"command": "uvx", "args": ["hotkernel"], "env": {"HOTKERNEL_PYTHON": "..."}}`。`HOTKERNEL_PYTHON` 指向装了你 PyTorch 和 Triton 的解释器；server 本身不依赖 torch。
+或者在客户端的 MCP 配置里写 `{"command": "uvx", "args": ["hotkernel"], "env": {"HOTKERNEL_PYTHON": "..."}}`，再把 [`skills/kernel-speedrun/`](https://github.com/xinhuagu/hotkernel/tree/main/skills/kernel-speedrun) 复制到 agent 的 skills 目录。不装插件的 Claude Code 也一样：`claude mcp add -e HOTKERNEL_PYTHON=... hotkernel -- uvx hotkernel`。
+
+## 对准你自己的 kernel
+
+`hotkernel new` 把任意一个 PyTorch 函数变成和示例同一契约的目标：这个函数就是基线，它的 fp32 结果是参考实现，容差按 dtype 给默认值，协议 3 的计时器以文件形式放在脚本旁边，所以脚本在你的 torch 解释器里就能跑，那里装没装 hotkernel 都行。
+
+```bash
+hotkernel new gelu --from torch.nn.functional:gelu --shape 4096x8192 --dtype bf16 --out targets
+```
 
-循环是一个 skill：把 [`skills/kernel-speedrun/`](https://github.com/xinhuagu/hotkernel/tree/main/skills/kernel-speedrun) 复制到你 agent 的 skills 目录（Claude Code 是 `~/.claude/skills/`）。
+会生成 `targets/gelu_baseline.py`、`targets/gelu_candidate.py`（一开始就是基线本身，循环改的就是这个文件）和 `targets/_hotkernel_contract.py`。自己写的函数同样用法，额外输入按调用顺序给：`--from ./mykernels.py:fused --extra w:8192`。然后对基线 `bench` 得到要打败的数字，再把没改过的候选 `verify` 一次，改动之前它必须是 `inconclusive`。agent 用 `new_target` 工具做同样的事。在参考卡上，没改过的 gelu 候选 bench 为 0.0952 ms、verify 判 `inconclusive`，故意写错的候选判 `incorrect`。
 
 ## 改前 / 改后
 
@@ -58,6 +75,7 @@ codex mcp add --env HOTKERNEL_PYTHON=$HOME/venvs/mytorch/bin/python hotkernel --
 | `bench` | 把脚本跑 n 次，收集它打印的 `{"ms": …, "ok": …}`，返回均值、标准差、最小值 | 你的 torch |
 | `verify` | 基线和候选交替各跑 n 次；给出加速比和裁决：`faster`、`slower`、`inconclusive`、`incorrect`（任一次 `ok=false`，不报加速比）、`invalid`（shape、dtype、种子或计时协议被改） | 你的 torch |
 | `record` | 把 verify 结果写到 `results/<目标>/`，从日志重新生成排行榜，并告诉 agent 这次改动是否超过本轮最好成绩 | 无 |
+| `new_target` | 用你自己的函数生成目标：基线、空白候选和随附的契约文件，与 `hotkernel new` 相同 | 无 |
 
 `torch_trace`、`bench`、`verify`、`record` 在任何装了 PyTorch 的机器上都能用。Nsight 是可选的，只是让图更清楚。
 

{hotkernel-0.1.0 → hotkernel-0.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hotkernel"
-version = "0.1.0"
+version = "0.1.1"
 description = "Find the hot kernel, cool it down: an Nsight/SASS MCP server with a verify-and-record contract and a kernel-speedrun loop for coding agents."
 readme = "README.md"
 license = "Apache-2.0"

{hotkernel-0.1.0 → hotkernel-0.1.1}/skills/kernel-speedrun/SKILL.md

@@ -6,8 +6,11 @@ description: Overnight loop that makes one GPU kernel faster on the GPU in this
 # kernel-speedrun
 
 You have the `hotkernel` MCP tools: `gpu_info`, `roofline`, `bench`, `torch_trace`, `profile`,
-`top_kernels`, `disasm`, `verify`, `record`. This skill is the loop that turns them into a
-speedup that someone else can reproduce.
+`top_kernels`, `disasm`, `verify`, `record`, and `new_target` (scaffold a target from the
+user's own function). This skill is the loop that turns them into a speedup that someone else
+can reproduce. If the user's kernel has no target yet, start with
+`new_target(target, from_="module:function", shape="4096x8192", dtype="bf16")` and bench the
+baseline it writes; the generated candidate is the file you edit.
 
 ## The contract
 

{hotkernel-0.1.0 → hotkernel-0.1.1}/src/hotkernel/__init__.py

@@ -5,4 +5,4 @@ hands on NVIDIA's profiling and disassembly tools, plus the data needed to run a
 overnight kernel-speedrun loop on the GPU you own.
 """
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"

hotkernel-0.1.1/src/hotkernel/newtarget.py

@@ -0,0 +1,263 @@
+"""Scaffold a hotkernel target from the user's own PyTorch function.
+
+Used by the `hotkernel new` command and the `new_target` MCP tool. It writes three files into
+the output directory, none of which import hotkernel, so they run in whatever interpreter has
+the user's PyTorch:
+
+  _hotkernel_contract.py   the vendored contract (protocol 3 timer, Target, tolerances)
+  <target>_baseline.py     the user's function as the baseline, plus the Target definition
+  <target>_candidate.py    the file the loop edits; starts as the baseline
+
+This module imports no torch: the generated scripts do.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.resources
+import json
+import re
+import sys
+import time
+from pathlib import Path
+
+CONTRACT_FILE = "_hotkernel_contract.py"
+_NAME_RE = re.compile(r"^[a-z][a-z0-9_]*$")
+_INPUT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+def contract_source() -> str:
+    """The contract file as shipped in this hotkernel version."""
+    return (importlib.resources.files("hotkernel").joinpath("templates/contract.py")
+            .read_text(encoding="utf-8"))
+
+
+def contract_protocol() -> int:
+    m = re.search(r"^PROTOCOL\s*=\s*(\d+)", contract_source(), re.MULTILINE)
+    return int(m.group(1)) if m else 0
+
+
+def parse_shape(spec: str) -> tuple[int, ...]:
+    """'4096x8192', '4096,8192' or '8192' -> tuple of ints."""
+    parts = [p for p in re.split(r"[x,*]", str(spec).strip().lower()) if p]
+    try:
+        dims = tuple(int(p) for p in parts)
+    except ValueError:
+        raise ValueError(f"bad shape {spec!r}; write it like 4096x8192") from None
+    if not dims or any(d <= 0 for d in dims):
+        raise ValueError(f"bad shape {spec!r}; every dimension must be a positive integer")
+    return dims
+
+
+def parse_inputs(shape: str, extra: list[str] | None) -> dict[str, tuple[int, ...]]:
+    """The primary input is `x` with `shape`; extras are 'name:shape' items in call order."""
+    shapes = {"x": parse_shape(shape)}
+    for item in extra or []:
+        if ":" not in item:
+            raise ValueError(f"extra input {item!r} must look like name:shape, e.g. w:8192")
+        name, s = item.split(":", 1)
+        if not _INPUT_RE.match(name) or name in shapes:
+            raise ValueError(f"bad or duplicate input name {name!r}")
+        shapes[name] = parse_shape(s)
+    return shapes
+
+
+def parse_from(spec: str) -> tuple[str, str, bool]:
+    """'torch.nn.functional:gelu' -> (module, function, False);
+    'mykernels.py:fused' or './dir/k.py:f' -> (file path, function, True)."""
+    if ":" not in spec:
+        raise ValueError(f"--from must be module:function or file.py:function, got {spec!r}")
+    module, function = spec.rsplit(":", 1)
+    if not _INPUT_RE.match(function):
+        raise ValueError(f"bad function name {function!r}")
+    is_file = module.endswith(".py")
+    if not is_file and not re.match(r"^[A-Za-z_][\w.]*$", module):
+        raise ValueError(f"bad module name {module!r}")
+    return module, function, is_file
+
+
+_BASELINE = '''"""hotkernel target "{target}": {spec} is the baseline.
+
+Generated by `hotkernel new` on {date} (contract protocol {protocol}). Everything that makes
+the number honest is here and in {contract}: the inputs (shapes, dtype, seed), the reference
+and its tolerances, the input rotation and the timer. A candidate defines a function with the
+same signature and calls run(candidate, impl="..."); it cannot change any of this without
+changing the contract fields that verify() checks. The loop never edits this file.
+
+The reference is the same function run in fp32 and cast back to DTYPE. If your function
+needs a different reference (a different formula, a library call), edit `reference` below
+before the first verify(), not after.
+"""
+
+{binding}
+from _hotkernel_contract import PROTOCOL, Target  # noqa: E402
+
+DTYPE = {dtype!r}
+SHAPES = {shapes}  # inputs in call order
+ROTATE = {rotate}  # streamed inputs that get K distinct copies (> 2x L2 in total)
+ITERS = {iters}
+SEED = {seed}
+RTOL, ATOL = {rtol}, {atol}  # None = default for DTYPE (see DEFAULT_TOLERANCES in the contract)
+FLOPS = {flops}  # nominal flops per call for roofline(); None = not stated
+
+assert PROTOCOL == {protocol}, "contract file does not match the protocol this target was generated for"
+
+
+def reference({params}):
+    return baseline_fn({float_args}).to({primary}.dtype)
+
+
+TARGET = Target(name={target!r}, dtype=DTYPE, shapes=SHAPES, rotate=ROTATE, reference=reference,
+                rtol=RTOL, atol=ATOL, iters=ITERS, seed=SEED, flops=FLOPS,
+                reference_name={refname!r})
+
+
+def run(fn, impl: str, extra: dict | None = None) -> None:
+    TARGET.run(fn, impl, extra)
+
+
+if __name__ == "__main__":
+    run(baseline_fn, impl={spec!r})
+'''
+
+_CANDIDATE = '''"""Candidate for the "{target}" target: the file the kernel-speedrun loop edits.
+
+It starts as the baseline itself, so the first verify() is a fair "inconclusive" and every
+later verdict measures only what changed. Define {target}_candidate({params}) and leave the
+run() call alone: inputs, seed, reference, tolerances, rotation and timer are the baseline's,
+and verify() reports `invalid` if the contract changes. Triton and torch.compile are fair game.
+"""
+
+from {target}_baseline import baseline_fn, run
+
+
+def {target}_candidate({params}):
+    # Starting point: the baseline. Replace with something faster; keep the signature.
+    return baseline_fn({params})
+
+
+if __name__ == "__main__":
+    run({target}_candidate, impl="candidate (starts as the baseline)")
+'''
+
+_BIND_MODULE = '''import importlib
+
+import torch  # noqa: F401  (the generated scripts run in the interpreter that has your torch)
+
+baseline_fn = getattr(importlib.import_module({module!r}), {function!r})'''
+
+_BIND_FILE = '''import importlib.util
+import sys
+
+import torch  # noqa: F401  (the generated scripts run in the interpreter that has your torch)
+
+_spec = importlib.util.spec_from_file_location({modname!r}, {path!r})
+_mod = importlib.util.module_from_spec(_spec)
+sys.modules[{modname!r}] = _mod
+_spec.loader.exec_module(_mod)
+baseline_fn = getattr(_mod, {function!r})'''
+
+
+def new_target(target: str, from_: str, shape: str, dtype: str = "bf16",
+               extra: list[str] | None = None, rotate: list[str] | None = None,
+               out_dir: str = ".", flops: int | None = None, iters: int = 50, seed: int = 0,
+               rtol: float | None = None, atol: float | None = None, force: bool = False) -> dict:
+    """Write the three files and return what was written plus the next commands. Errors come
+    back as {"error": ...} so an agent can read them."""
+    try:
+        if not _NAME_RE.match(target):
+            raise ValueError(f"target name {target!r} must be lowercase letters, digits and _")
+        shapes = parse_inputs(shape, extra)
+        module, function, is_file = parse_from(from_)
+        rotate_names = tuple(rotate) if rotate else ("x",)
+        if any(r not in shapes for r in rotate_names):
+            raise ValueError(f"--rotate names must be inputs: {list(shapes)}")
+    except ValueError as e:
+        return {"error": str(e)}
+
+    out = Path(out_dir).expanduser().resolve()
+    out.mkdir(parents=True, exist_ok=True)
+    contract = out / CONTRACT_FILE
+    baseline = out / f"{target}_baseline.py"
+    candidate = out / f"{target}_candidate.py"
+    src = contract_source()
+    if contract.exists() and contract.read_text(encoding="utf-8") != src and not force:
+        return {"error": f"{contract} exists and differs from this hotkernel's contract "
+                         f"(protocol {contract_protocol()}); pass force=True / --force to replace "
+                         "it, which re-times every target in that directory under the new contract"}
+    clashes = [str(p) for p in (baseline, candidate) if p.exists()]
+    if clashes and not force:
+        return {"error": f"already exists: {clashes}; pass force=True / --force to overwrite"}
+
+    params = ", ".join(shapes)
+    primary = rotate_names[0]
+    if is_file:
+        path = str(Path(module).expanduser().resolve())
+        modname = "_hk_user_" + re.sub(r"\W", "_", Path(module).stem)
+        binding = _BIND_FILE.format(modname=modname, path=path, function=function)
+        spec = f"{Path(module).name}:{function}"
+    else:
+        binding = _BIND_MODULE.format(module=module, function=function)
+        spec = f"{module}:{function}"
+    protocol = contract_protocol()
+    baseline_text = _BASELINE.format(
+        target=target, spec=spec, date=time.strftime("%Y-%m-%d"), protocol=protocol,
+        contract=CONTRACT_FILE, binding=binding, dtype=dtype,
+        shapes="{" + ", ".join(f"{n!r}: {s!r}" for n, s in shapes.items()) + "}",
+        rotate=repr(tuple(rotate_names)), iters=int(iters), seed=int(seed),
+        rtol=repr(rtol), atol=repr(atol), flops=repr(None if flops is None else int(flops)),
+        params=params, float_args=", ".join(f"{n}.float()" for n in shapes), primary=primary,
+        refname=f"{spec} in fp32")
+    candidate_text = _CANDIDATE.format(target=target, params=params)
+
+    contract.write_text(src, encoding="utf-8", newline="\n")
+    baseline.write_text(baseline_text, encoding="utf-8", newline="\n")
+    candidate.write_text(candidate_text, encoding="utf-8", newline="\n")
+    rel = out.as_posix()
+    return {
+        "target": target, "protocol": protocol,
+        "files": {"contract": str(contract), "baseline": str(baseline), "candidate": str(candidate)},
+        "inputs": {n: list(s) for n, s in shapes.items()}, "dtype": dtype, "rotate": list(rotate_names),
+        "next": [
+            (f"python {baseline.name} (run from {rel}, in the interpreter that has your torch) "
+             "to see the baseline's JSON line"),
+            f"bench(script='{rel}/{baseline.name}', n=5): the number to beat",
+            (f"verify(candidate='{rel}/{candidate.name}', baseline='{rel}/{baseline.name}', n=7): "
+             "must be 'inconclusive' before you change anything"),
+            f"edit {candidate.name}, then verify() and record(target='{target}', ...) after every change",
+        ],
+    }
+
+
+def build_parser(prog: str = "hotkernel new") -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(prog=prog, description="Scaffold a hotkernel target from your "
+                                "own PyTorch function: baseline, blank candidate and the vendored contract.")
+    p.add_argument("target", help="target name, e.g. gelu (lowercase letters, digits, _)")
+    p.add_argument("--from", dest="from_", required=True, metavar="MODULE:FUNCTION",
+                   help="the function to speed up: torch.nn.functional:gelu or ./mykernels.py:fused")
+    p.add_argument("--shape", required=True, help="shape of the primary input x, e.g. 4096x8192")
+    p.add_argument("--dtype", default="bf16", help="bf16 (default), fp16, fp32, fp64")
+    p.add_argument("--extra", nargs="*", default=[], metavar="NAME:SHAPE",
+                   help="more inputs in call order, e.g. w:8192 b:8192 (random, same dtype)")
+    p.add_argument("--rotate", nargs="*", default=None, metavar="NAME",
+                   help="inputs that stream through memory and get K copies (default: x)")
+    p.add_argument("--out", default=".", help="directory for the generated files (default: .)")
+    p.add_argument("--flops", type=int, default=None, help="nominal flops per call for roofline()")
+    p.add_argument("--iters", type=int, default=50)
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--rtol", type=float, default=None, help="override the dtype default")
+    p.add_argument("--atol", type=float, default=None, help="override the dtype default")
+    p.add_argument("--force", action="store_true", help="overwrite existing files")
+    return p
+
+
+def main(argv: list[str]) -> int:
+    args = build_parser().parse_args(argv)
+    result = new_target(args.target, args.from_, args.shape, args.dtype, args.extra, args.rotate,
+                        args.out, args.flops, args.iters, args.seed, args.rtol, args.atol, args.force)
+    print(json.dumps(result, indent=1))
+    return 1 if "error" in result else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

{hotkernel-0.1.0 → hotkernel-0.1.1}/src/hotkernel/server.py

@@ -13,6 +13,7 @@ can read them instead of crashing):
   verify        baseline vs candidate, interleaved, n runs each: speedup + verdict
                 (faster / slower / inconclusive / incorrect / invalid)
   record        write a verify result to results/<target>/ and append the LEADERBOARD row
+  new_target    scaffold baseline + blank candidate + vendored contract from the user's function
 
 Design rules:
   * hotkernel's own environment stays light (no torch). Anything that needs the user's
@@ -64,8 +65,46 @@ _DEFAULT_TIMEOUT = 600
 
 # --------------------------------------------------------------------------- helpers
 
+_PYTHON_HELP = (
+    "Set HOTKERNEL_PYTHON to the interpreter that has your PyTorch (and Triton), for example "
+    "~/venvs/mytorch/bin/python: `claude mcp add -e HOTKERNEL_PYTHON=... hotkernel -- uvx hotkernel`, "
+    "or for the Claude Code plugin the 'Python with your PyTorch' option "
+    "(`claude plugin configure hotkernel@hotkernel`), or an env entry in any other MCP client's "
+    "config. hotkernel itself stays torch-free."
+)
+
+
+@functools.lru_cache(maxsize=1)
+def _probe_torch_python() -> str | None:
+    """The first python/python3 on PATH that can import torch, or None. Probed once per server."""
+    for name in ("python3", "python"):
+        exe = shutil.which(name)
+        if not exe:
+            continue
+        r = _run([exe, "-c", "import torch, sys; print(sys.executable)"], timeout=120)
+        if "error" not in r and r["returncode"] == 0:
+            return r["stdout"].strip().splitlines()[-1] if r["stdout"].strip() else exe
+    return None
+
+
+def _python_or_error() -> tuple[str | None, dict | None]:
+    """HOTKERNEL_PYTHON if set (an empty value or an unexpanded ${...} placeholder counts as
+    unset), else the first python on PATH that imports torch, else a clear error."""
+    env = (os.environ.get("HOTKERNEL_PYTHON") or "").strip()
+    if env and not env.startswith("${"):
+        return env, None
+    found = _probe_torch_python()
+    if found:
+        return found, None
+    return None, {"error": "HOTKERNEL_PYTHON is not set and no python/python3 on PATH can import "
+                           "torch, so this tool cannot run the target. " + _PYTHON_HELP}
+
+
 def _python() -> str:
-    return os.environ.get("HOTKERNEL_PYTHON") or shutil.which("python") or sys.executable
+    """Best guess for places that only need an interpreter path (Triton's bundled cuobjdump,
+    profile commands); tools that run a target use _python_or_error() and fail loudly."""
+    py, _ = _python_or_error()
+    return py or sys.executable
 
 
 def _local_path(p: str | None) -> str | None:
@@ -236,7 +275,9 @@ def bench(script: str, n: int = 5, args: list[str] | None = None, cwd: str | Non
     Contract: the script prints one JSON line to stdout containing at least {"ms": <float>}
     and optionally {"ok": <bool>} for a correctness check. Returns mean/std/min over the n
     timed runs plus every raw sample, so a LEADERBOARD row can cite n and seed."""
-    py = _python()
+    py, err = _python_or_error()
+    if err:
+        return err
     cmd = [py, script, *(args or [])]
     samples: list[float] = []
     oks: list[bool] = []
@@ -328,7 +369,9 @@ def torch_trace(script: str, top: int = 15, args: list[str] | None = None,
     """Run a Python script under torch.profiler in the user's interpreter. Returns `kernels`
     (real GPU kernels by device time, pct over total kernel time) and `ops` (aten ops with
     input shapes, inclusive time). Works without Nsight installed; use this first."""
-    py = _python()
+    py, err = _python_or_error()
+    if err:
+        return err
     with tempfile.NamedTemporaryFile("w", suffix="_hotkernel_torchprof.py", delete=False) as f:
         f.write(_TORCH_PROF_RUNNER)
         runner = f.name
@@ -741,7 +784,9 @@ def verify(candidate: str, baseline: str, n: int = 7, warmup: int = 2,
     speedup is reported), 'invalid' if the two scripts disagree on shape/dtype/iters/seed or the
     baseline fails its own check. Also records gpu, driver, power limit, timestamp and the exact
     commands. Pass the result to record()."""
-    py = _python()
+    py, err = _python_or_error()
+    if err:
+        return err
     b_cmd = [py, baseline, *(args or [])]
     c_cmd = [py, candidate, *(args or [])]
     t0 = time.perf_counter()
@@ -967,7 +1012,42 @@ def record(target: str, verify_result: dict | str, note: str = "", who: str = ""
             "advice": advice, **counts}
 
 
+@mcp.tool()
+def new_target(target: str, from_: str, shape: str, dtype: str = "bf16",
+               extra: list[str] | None = None, rotate: list[str] | None = None,
+               out_dir: str = ".", flops: int | None = None, iters: int = 50, seed: int = 0,
+               rtol: float | None = None, atol: float | None = None, force: bool = False) -> dict:
+    """Scaffold a target from the user's own PyTorch function so the loop can run on it: writes
+    <out_dir>/<target>_baseline.py (the function as the baseline plus the Target definition:
+    inputs, seed, fp32 reference, dtype-default tolerances, rotation, timer),
+    <out_dir>/<target>_candidate.py (the file to edit; starts as the baseline) and the vendored
+    _hotkernel_contract.py (protocol 3) they import, so they run in HOTKERNEL_PYTHON without
+    hotkernel installed there. from_: 'torch.nn.functional:gelu' or './mykernels.py:fused';
+    shape: the primary input x, e.g. '4096x8192'; extra: more inputs in call order like
+    ['w:8192', 'b:8192'] (random tensors of the same dtype); rotate: inputs that stream through
+    memory (default ['x']); flops: nominal flops per call for roofline(). Same as the CLI
+    `hotkernel new`. Then bench() the baseline and verify() the blank candidate (expect
+    'inconclusive') before changing anything."""
+    from .newtarget import new_target as _new
+    out = _local_path(out_dir) or out_dir
+    return _new(target, from_, shape, dtype, extra, rotate, out, flops, iters, seed, rtol, atol, force)
+
+
 def main() -> None:
+    """`hotkernel` starts the MCP server over stdio (what `uvx hotkernel` in an MCP config
+    runs); `hotkernel new ...` scaffolds a target; `hotkernel --version` prints the version."""
+    argv = sys.argv[1:]
+    if argv and argv[0] == "new":
+        from .newtarget import main as new_main
+        sys.exit(new_main(argv[1:]))
+    if argv and argv[0] in ("--version", "-V"):
+        print(f"hotkernel {__version__}")
+        return
+    if argv and argv[0] in ("-h", "--help"):
+        print("usage: hotkernel            start the MCP server (stdio)\n"
+              "       hotkernel new ...    scaffold a target from your function (see `hotkernel new -h`)\n"
+              "       hotkernel --version")
+        return
     mcp.run()
 
 

hotkernel-0.1.1/src/hotkernel/templates/contract.py

@@ -0,0 +1,172 @@
+"""hotkernel target contract, protocol 3. This file is vendored next to every target by
+`hotkernel new` (as `_hotkernel_contract.py`), so a target runs in the interpreter that has
+your PyTorch whether or not hotkernel is installed there. Do not edit the copy; regenerate it.
+
+What it owns, so that a candidate cannot change it: the inputs (shapes, dtype, seed), the
+reference and its tolerances, the input rotation, and the timer. A candidate is one function
+with the baseline's signature handed to the baseline's run().
+
+Timing (protocol 3, steady-state rotation, no flush): K distinct sets of the rotated inputs,
+K > 2 * L2 / bytes per call, cycled call after call; the last K outputs are held so the
+allocator gives every call a different output block (K + 1 blocks in a ring). Each call reads
+a cold input and evicts the previous call's dirty output: exactly its own bytes of DRAM
+traffic, the steady-state cost a real pipeline pays. Earlier protocols flushed L2 with a 256 MB
+zero fill (protocol 2) or did nothing (protocol 1); both reported wrong bandwidths on a card
+with a 96 MB L2. See docs/platform-matrix.md in the hotkernel repository for the calibration.
+"""
+
+from __future__ import annotations
+
+import collections
+import json
+import math
+import sys
+
+import torch
+
+PROTOCOL = 3
+TIMING = "per-call CUDA events, steady-state rotation over K buffer pairs > 2x L2, no flush"
+
+# Defaults for allclose(out.float(), ref.float(), rtol, atol) by output dtype: the bf16 and
+# fp16 values allow one rounding of the result, fp32 allows a different accumulation order.
+DEFAULT_TOLERANCES = {
+    torch.bfloat16: (2e-2, 2e-2),
+    torch.float16: (1e-2, 1e-3),
+    torch.float32: (1e-4, 1e-5),
+    torch.float64: (1e-8, 1e-10),
+}
+_DTYPES = {"bf16": torch.bfloat16, "bfloat16": torch.bfloat16, "fp16": torch.float16,
+           "half": torch.float16, "float16": torch.float16, "fp32": torch.float32,
+           "float32": torch.float32, "float": torch.float32, "fp64": torch.float64,
+           "float64": torch.float64}
+
+
+def parse_dtype(name: str | torch.dtype) -> torch.dtype:
+    if isinstance(name, torch.dtype):
+        return name
+    try:
+        return _DTYPES[str(name).lower().removeprefix("torch.")]
+    except KeyError:
+        raise ValueError(f"unknown dtype {name!r}; use one of {sorted(_DTYPES)}") from None
+
+
+def default_tolerances(dtype: torch.dtype) -> tuple[float, float]:
+    return DEFAULT_TOLERANCES.get(dtype, (1e-4, 1e-5))
+
+
+def l2_bytes() -> int:
+    return int(torch.cuda.get_device_properties(0).L2_cache_size)
+
+
+def rotation_k(pair_bytes: int) -> int:
+    """Smallest K with K * pair_bytes > 2 * L2, and at least 2."""
+    return max(2, math.floor(2 * l2_bytes() / max(pair_bytes, 1)) + 1)
+
+
+def time_kernel(fn, inputs: list[tuple], iters: int = 50, warmup: int | None = None) -> float:
+    """Mean milliseconds per call of fn(*inputs[i % K]) over `iters` timed calls, cycling the
+    K input tuples. fn must return a freshly allocated tensor (its output); the last K outputs
+    are held so outputs rotate too. Warmup covers every input at least twice."""
+    k = len(inputs)
+    held: collections.deque = collections.deque(maxlen=k)
+    for i in range(2 * k + 3 if warmup is None else warmup):
+        held.append(fn(*inputs[i % k]))
+    torch.cuda.synchronize()
+    events = []
+    for i in range(iters):
+        start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+        start.record()
+        out = fn(*inputs[i % k])
+        end.record()
+        held.append(out)  # frees the oldest output after the event; the ring keeps rotating
+        events.append((start, end))
+    torch.cuda.synchronize()
+    return sum(s.elapsed_time(e) for s, e in events) / iters
+
+
+def _nbytes(t) -> int:
+    return int(t.numel() * t.element_size()) if isinstance(t, torch.Tensor) else 0
+
+
+class Target:
+    """One kernel target: its inputs, reference, tolerances and timing, all in one place.
+
+    shapes: name -> shape of every input, in call order. rotate: the names that get K distinct
+    copies (the big streamed tensors); the others are shared across calls. reference(*inputs)
+    returns the correct output (usually the same math in fp32, cast back). flops is the nominal
+    flop count per call for roofline() and may be None.
+    """
+
+    def __init__(self, name: str, dtype, shapes: dict[str, tuple[int, ...]], rotate: tuple[str, ...],
+                 reference, rtol: float | None = None, atol: float | None = None, iters: int = 50,
+                 seed: int = 0, flops: int | None = None, min_bytes: int | None = None,
+                 reference_name: str = "reference"):
+        self.name = name
+        self.dtype = parse_dtype(dtype)
+        self.shapes = {k: tuple(int(d) for d in v) for k, v in shapes.items()}
+        self.rotate = tuple(rotate)
+        unknown = set(self.rotate) - set(self.shapes)
+        if unknown or not self.rotate:
+            raise ValueError(f"rotate must name inputs from {list(self.shapes)}; got {self.rotate}")
+        self.reference = reference
+        d_rtol, d_atol = default_tolerances(self.dtype)
+        self.rtol = d_rtol if rtol is None else float(rtol)
+        self.atol = d_atol if atol is None else float(atol)
+        self.iters = int(iters)
+        self.seed = int(seed)
+        self.flops = None if flops is None else int(flops)
+        self.min_bytes = None if min_bytes is None else int(min_bytes)
+        self.reference_name = reference_name
+
+    @property
+    def primary(self) -> str:
+        return self.rotate[0]
+
+    def make_inputs(self, k: int) -> list[tuple]:
+        """K tuples in call order: rotated inputs are fresh per set, the others shared."""
+        torch.manual_seed(self.seed)
+        rotated = {n: [torch.randn(*self.shapes[n], device="cuda", dtype=self.dtype) for _ in range(k)]
+                   for n in self.rotate}
+        shared = {n: torch.randn(*self.shapes[n], device="cuda", dtype=self.dtype)
+                  for n in self.shapes if n not in self.rotate}
+        return [tuple(rotated[n][i] if n in rotated else shared[n] for n in self.shapes)
+                for i in range(k)]
+
+    def run(self, fn, impl: str, extra: dict | None = None) -> None:
+        """Check fn against the reference, time it under protocol 3, print the JSON line, exit
+        0 if correct and 1 if not. The candidate's only job is to pass fn."""
+        rotated_bytes = sum(math.prod(self.shapes[n]) * torch.finfo(self.dtype).bits // 8
+                            for n in self.rotate)
+        k = rotation_k(2 * rotated_bytes)  # assume output ~ input until the first call tells
+        inputs = self.make_inputs(k)
+
+        out = fn(*inputs[0])
+        ref = self.reference(*inputs[0])
+        ok = bool(torch.allclose(out.float(), ref.float(), rtol=self.rtol, atol=self.atol))
+        pair_bytes = rotated_bytes + _nbytes(out)
+        k_needed = rotation_k(pair_bytes)
+        if k_needed > k:
+            inputs = self.make_inputs(k_needed)
+            k = k_needed
+        del out, ref
+
+        ms = time_kernel(fn, inputs, iters=self.iters)
+        min_bytes = pair_bytes if self.min_bytes is None else self.min_bytes
+        line = {
+            "ms": round(ms, 5), "ok": ok,
+            "shape": list(self.shapes[self.primary]), "dtype": str(self.dtype),
+            "iters": self.iters, "seed": self.seed, "protocol": PROTOCOL, "timing": TIMING,
+            "target": self.name, "shapes": {n: list(s) for n, s in self.shapes.items()},
+            "rotate": list(self.rotate), "rtol": self.rtol, "atol": self.atol,
+            "reference": self.reference_name,
+            "rotation_k": k, "rotation_bytes": k * pair_bytes, "l2_bytes": l2_bytes(),
+            "min_bytes": min_bytes,
+            "achieved_gbs_lower_bound": round(min_bytes / (ms / 1e3) / 1e9, 1),
+            "gpu": torch.cuda.get_device_name(0), "torch": torch.__version__,
+            "impl": impl,
+        }
+        if self.flops is not None:
+            line["flops"] = self.flops
+        line.update(extra or {})
+        print(json.dumps(line))
+        sys.exit(0 if ok else 1)