diff --git a/README.md b/README.md index f2f84779..4f0d0e45 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,73 @@ -# TileFoundry +

+ TileFoundry +

-[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE) ![Status](https://img.shields.io/badge/status-early%20development-orange) +--- + +

+ PyPI + Status: early development + License: MIT +

+ +

+ Documentation · + Installation · + Examples +

**TileFoundry** is a tile-based, agentic platform for automatic high-performance program generation across hardware. -> [!NOTE] -> TileFoundry is in an early design and development stage. APIs and architecture are still evolving, and the project is not yet ready for use. +## Latest News + +- 08/2026 ๐ŸŽ‰: **TileFoundry 0.0.1 is on PyPI** โ€” the first public release. +- 08/2026 ๐Ÿ“ฆ: Four [worked examples](https://github.com/tile-ai/TileFoundry/tree/main/examples) added โ€” Qwen3-1.7B (tilelang), Qwen3.5-35B-A3B (tilelang), MiniCPM3-4B (CuTeDSL) and granite-4.0-h-small (CUDA C) โ€” each one a real agent run kept whole, with the decode throughput it measured. + +## Installation + +TileFoundry needs Python 3.12 or newer. + +```sh +pip install tilefoundry +``` + +Check the install โ€” it prints the commands an agent will ask: + +```sh +tilefoundry +``` + +Running a model an agent generates additionally needs one NVIDIA GPU and the +checkpoint already on disk. + +## Quick Start + +There is **no API to learn** first. Give your coding agent this, with a checkpoint +directory of your own: + +```text +Get real tokens out of Qwen3-1.7B on TileFoundry, and make it fast. +Weights and config: +Backend: tilelang. + +Everything about TileFoundry is to be asked of the `tilefoundry` command -- do not +ask a person, do not go looking elsewhere. The model itself is yours to research. + +Done when this runs from outside, prints the continuation, and reports a +tokens-per-second number measured over the whole generation: + + python run.py --prompt "Write a detailed explanation of how a GPU executes a matrix multiplication." --max-new-tokens 2048 + +Measure over a long generation -- 2048 new tokens, more than 2000 characters of +text. A 32-token sample is too short for the number to mean anything. +``` + +That is the whole input โ€” nothing under it is written by hand. + +Claude Opus 5 at xhigh reasoning effort ran this prompt for 2.1 hours with **no +interaction**, and reached **612 tok/s** on one H200. What it wrote is +[`examples/qwen3_1_7b-tilelang/`](https://github.com/tile-ai/TileFoundry/tree/main/examples/qwen3_1_7b-tilelang). ## License -This project is licensed under the [MIT License](LICENSE). +This project is licensed under the [MIT License](https://github.com/tile-ai/TileFoundry/blob/main/LICENSE). diff --git a/examples/qwen3_1_7b-tilelang/README.md b/examples/qwen3_1_7b-tilelang/README.md new file mode 100644 index 00000000..82e729c2 --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/README.md @@ -0,0 +1,188 @@ +# Qwen3-1.7B on TileFoundry โ€” tilelang kernels + +Verified at **v0.0.1**, 2026-08-04. Not verified since; nothing in CI re-runs it. + +Greedy decoding of the published Qwen3-1.7B checkpoint taken end to end on +TileFoundry: every kernel on the path written here in tilelang, and the whole +decode step replayed from one captured CUDA graph. + + 612.5 tok/s one H200, batch 1, bf16, greedy, 2048 new tokens + +--- + +## 1. Environment + +Nothing here is installed by this directory; it is what the directory was written +and measured against. + +| | | +|---|---| +| GPU | one NVIDIA H200, driver 575.57.08 | +| CUDA | 12.8 (`nvcc`), `CUDA_HOME` must be set | +| Python | 3.12 | +| `tilefoundry` | installed from the wheel, **not** editable โ€” nothing here reads the source tree | +| borrowed from the environment | `torch` 2.9.1+cu128, `transformers`, `tokenizers` | +| extra package | `tilelang` 0.1.12 | +| weights | the published `Qwen/Qwen3-1.7B` checkpoint, 3.8 GB on disk | + +## 2. How this was produced + +One agent, one prompt, no human help after it started: Claude Opus 5 at xhigh +reasoning effort, 2.1 hours, 177 tool calls, no sub-agents. + +The prompt was thirteen lines, where the other three examples here were produced +by one of about eighty. It is the Quick Start in the project's README, run +unedited: + +```text +Get real tokens out of Qwen3-1.7B on TileFoundry, and make it fast. +Weights and config: +Backend: tilelang. + +Everything about TileFoundry is to be asked of the `tilefoundry` command -- do not +ask a person, do not go looking elsewhere. The model itself is yours to research. + +Done when this runs from outside, prints the continuation, and reports a +tokens-per-second number measured over the whole generation: + + python run.py --prompt "Write a detailed explanation of how a GPU executes a matrix multiplication." --max-new-tokens 2048 + +Measure over a long generation -- 2048 new tokens, more than 2000 characters of +text. A 32-token sample is too short for the number to mean anything. +``` + +It names one file, `run.py`. That the work is a reference baseline first, then a +runtime twin, then `tilefoundry check` as the comparison between them โ€” and the +shape of this directory โ€” the agent read out of `tilefoundry tutorial` and decided +for itself. + +## 3. How to use it + + python run.py --ckpt --prompt "..." --max-new-tokens 2048 + +`--ckpt` is required: where the weights live is a fact about the machine, and a +default that exists on only one machine is a guess. + +| flag | | +|---|---| +| `--prompt` | the text to continue; required | +| `--max-new-tokens N` | how many tokens to generate, default 2048 | +| `--device` | pin the runtime device. By default the emptiest visible one is taken, probed through Torch โ€” an exclusive-mode card already has an owner, and only trying it says so | + +The first run compiles the tilelang kernels (once, a few minutes). + + run.py the entry point + ref_src/ verbatim copy of the shipped `qwen3_1_7b` source -- the + reference, never edited + fast/kernels.py the tilelang kernels, one decode step's worth + fast/engine.py weight packing, buffers, and the graph capture + fast/twin.py @runtime_module twins, so `tilefoundry check` can judge + the kernels against the reference + fast/test_kernels.py the torch spelling of every kernel -- the interface they + were written against + fast/arbitrate.py an independent f64 reference, for the one disagreement + in ยง4 that `check` cannot settle + +### Why it is shaped this way + +The authored reference hands each step's key and value back for the caller to +`torch.cat` on. That is right for a reference โ€” it keeps every shape expressed in +`ctx_len` alone โ€” but the cache buffer then moves every step, and a CUDA graph +records addresses. So the engine takes the other form the migrate page names: a +cache of fixed capacity whose write window advances, with the position in a +one-element device tensor. + +Everything a step needs then has a fixed address, so all 254 kernels are captured +once and replayed. The last kernel writes the sampled token back into the input +slot, and while the prompt still has a token left it feeds that one instead โ€” so +one capture walks the prompt and continues past it with **no host round trip +anywhere in the loop**, including no sync to read the token back. + +Decode is one token, so every projection is a GEMV: pure streaming, no reuse. +`q|k|v` and `gate|up` are packed into single matrices at load, because the cost of +a GEMV is the block count it can fill, not its arithmetic. Split-K supplies the +rest of the parallelism, and its f32 partials are reduced by the *consumer* โ€” the +attention merge lands inside `o_proj`, and `silu(gate) * up` inside `down_proj`, so +neither is a launch of its own. + +## 4. Where it stands + +**Measured, on the environment above**, at four levels, each answering something +the level before it cannot. + +1. **Every kernel against a torch statement of the same thing**, at production + dimensions. All exact or within one bf16 rounding. +2. **`tilefoundry check` against the authored HIR** โ€” the comparison the optimize + page asks for. All four decoder-layer functions and all three root functions + pass, at context extents 0 / 1 / 255 / 1024; and the whole model in one decode + step passes on all 57 outputs (logits cosine 0.999951). +3. **An independent f64 reference**, because `check` says outright that a FAIL + "proves disagreement, not which side is closer to truth". Where the twin and the + reference differ on attention, the twin's error against f64 is **3.4e-3 vs the + reference's 8.6e-3** โ€” 2.5x closer. +4. **Against Hugging Face on the real checkpoint** โ€” the L3 bar. Teacher-forced, + **255/256 positions agree**. + +The same authored HIR run through the evaluator instead of this twin decodes at +**14.8 tok/s**, so the twin is about 41x it. + +### The two deliberate departures from the reference + +Both move toward the published model, which is what level 3 and 4 measure: + +- **The `1/sqrt(head_dim)` factor is applied after the dot product, in f32.** The + reference multiplies it onto `q` in bf16 first, rounding every entry a second + time; the exponential downstream turns that into percent-level error on the + attention weights. HF scales after. This is most of the 2.5x in level 3. It is a + fact about kernels that hold the score in bf16, not about the description: the + authored HIR run through the Evaluator is unchanged either way, measured. +- **Attention probabilities are bf16 into the V product**, which is exactly what + HF's attention does. Adopting it moved HF agreement from 253/256 to 255/256. + +### The one remaining disagreement with HF + +At the first generated token HF's bf16 logits for the two candidates are +**exactly equal** (22.625 and 22.625), so `argmax` picks the lower index. This +implementation keeps f32 logits, which resolve a real 0.11 gap, and picks the +other. Neither is wrong; HF's output dtype simply cannot represent the +distinction. Every other position in a 256-token teacher-forced comparison agrees. + +### Three TileLang findings worth keeping + +Each cost real time to locate and each is a cliff, not a gradient: + +- **`from __future__ import annotations` breaks `@T.prim_func`.** Buffers are + declared from parameter annotations, and PEP 563 hands the builder strings + evaluated without the enclosing factory's closure โ€” so the dimensions the + factory exists to bind are exactly what fails to resolve. +- **`T.atomic_min` on shared memory costs 18 ms** where a `T.reduce_max` costs + 2 us. Argmax is written here as two max-reductions instead: the winning value, + then `BN - j` over the entries attaining it, whose max is the lowest winning + index โ€” the same tie-break `torch.argmax` reports. +- **Reducing a GEMV accumulator in place costs two orders of magnitude.** Layout + inference replicates the fragment across all threads to satisfy both uses, which + spills it: 19 ms instead of 165 us. Staging through shared memory and reducing a + fresh fragment fixes it. The same conflict has no workaround when the reduction + is over the full hidden size, which is what stopped the residual-norm folds. + +### Where the time goes + +Marginal in-graph cost per decode step, at 1024 context: + +| | per call | x | step | rate | +|---|---|---|---|---| +| `gate_up` GEMV | 14.1 us | 28 | 396 us | 3.6 TB/s | +| `down` GEMV (+silu) | 7.9 us | 28 | 221 us | 3.2 TB/s | +| `lm_head` (+argmax) | 165 us | 1 | 165 us | 3.8 TB/s | +| attention | 5.6 us | 28 | 156 us | โ€” | +| `qkv` GEMV | 4.9 us | 28 | 138 us | 3.4 TB/s | +| `o` GEMV (+combine) | 4.3 us | 28 | 121 us | 1.9 TB/s | +| norms, rope | ~2 us | 84 | 165 us | โ€” | + +The GEMVs are at the streaming roofline โ€” a sweep over block counts and split-K +factors found nothing better than 1% over the shapes in use, and block count +barely moves them. The remaining gap to the 3.44 GB / step memory floor is +per-kernel ramp and drain, so the only lever left is kernel *count*: the two +residual norms are single-block kernels costing ~1.9 us each of near-pure latency, +and folding them into the following GEMV is worth ~6-10% but is what the +layout-inference cliff above blocks. diff --git a/examples/qwen3_1_7b-tilelang/fast/__init__.py b/examples/qwen3_1_7b-tilelang/fast/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/qwen3_1_7b-tilelang/fast/arbitrate.py b/examples/qwen3_1_7b-tilelang/fast/arbitrate.py new file mode 100644 index 00000000..8eaca382 --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/fast/arbitrate.py @@ -0,0 +1,112 @@ +"""Which side of the attention disagreement is closer to truth? + +`check` reports that the twin and the authored reference differ, and says plainly +that it cannot say which is closer -- "establishing accuracy needs an independent +high-precision reference, which check does not run." This runs one: the same real +activations and the same real weights through the same math in float64. + +The disagreement is not an accident. The reference multiplies `1/sqrt(head_dim)` +onto `q` in bf16 before the dot product, rounding every entry a second time, and +the exponential downstream magnifies it. The kernels apply the factor to the +finished f32 dot instead, which is also what Hugging Face does. This script is the +evidence for calling that an improvement rather than a deviation. +""" +import argparse +from pathlib import Path + +import torch + +from engine import _load_module, default_ref_dir + + +def f64_attention(hidden, w, cos, sin, pos, k_cache, v_cache, scale, cfg): + """Reference math in f64: no intermediate lands in bf16 anywhere.""" + H = cfg.hidden_size + HQ, HKV, D = cfg.num_attention_heads, cfg.num_key_value_heads, cfg.head_dim + G = HQ // HKV + eps = cfg.rms_norm_eps + f = torch.float64 + + x = hidden.reshape(H).to(f) + x = x * torch.rsqrt(x.pow(2).mean() + eps) * w["gamma_in"].to(f) + q = (x @ w["w_q"][0].to(f)).reshape(HQ, D) + k = (x @ w["w_k"][0].to(f)).reshape(HKV, D) + v = (x @ w["w_v"][0].to(f)).reshape(HKV, D) + + def hnorm(t, g): + return t * torch.rsqrt(t.pow(2).mean(-1, keepdim=True) + eps) * g.to(f) + + q = hnorm(q, w["gamma_q"]) + k = hnorm(k, w["gamma_k"]) + + def rope(t): + c, s = cos[pos].to(f), sin[pos].to(f) + half = torch.cat([-t[:, D // 2:], t[:, : D // 2]], dim=-1) + return t * c + half * s + + q, k = rope(q), rope(k) + ctx = int(k_cache.shape[1]) + kk = torch.cat([k_cache[0].to(f), k.unsqueeze(0)], dim=0) # (ctx+1, HKV, D) + vv = torch.cat([v_cache[0].to(f), v.unsqueeze(0)], dim=0) + kk = kk.repeat_interleave(G, dim=1) # (ctx+1, HQ, D) + vv = vv.repeat_interleave(G, dim=1) + sc = (q.unsqueeze(0) * kk).sum(-1) * float(scale.reshape(-1)[0]) + p = torch.softmax(sc, dim=0) + attn = (p.unsqueeze(-1) * vv).sum(0).reshape(HQ * D) + return attn @ w["w_o"][0].to(f) + + +def err(got, truth): + g, t = got.reshape(-1).to(torch.float64), truth.reshape(-1) + return ((g - t).norm() / t.norm()).item() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--real", type=Path, default=Path("real_inputs")) + ap.add_argument("--ckpt", default="../prepared") + ap.add_argument("--layer", type=int, default=0) + args = ap.parse_args() + + from tilefoundry.runtime import SafetensorsResource + import twin + + ref = twin.ref + cfg = ref.config + dev = "cuda:0" + acts = {n: torch.load(args.real / f"{n}.pt").to(dev) + for n in ("hidden", "cos_cache", "sin_cache", "pos_ids", + "k_cache", "v_cache", "scale")} + + loaded = ref.Qwen3_1_7B.load(SafetensorsResource(str(args.ckpt), device=dev)) + lay = getattr(loaded, f"layer{args.layer}") + w = lay.constants + + # both sides take activations alone; each fills its own weights from its own + # reading of the same checkpoint, which is what makes the comparison fair + call = [acts[n] for n in ("hidden", "cos_cache", "sin_cache", "pos_ids", + "k_cache", "v_cache", "scale")] + + truth = f64_attention( + acts["hidden"], w, acts["cos_cache"], acts["sin_cache"], + int(acts["pos_ids"][0]), acts["k_cache"], acts["v_cache"], acts["scale"], cfg, + ) + + ref_out = lay.self_attention(*call)[0] + tw = twin.LayerTwin(ir=ref.Qwen3_1_7B_DecoderLayer) + tw.load(SafetensorsResource(str(args.ckpt), device=dev).subtree(f"layer{args.layer}")) + mine = tw.self_attention(*call)[0] + + e_ref, e_mine = err(ref_out, truth), err(mine, truth) + print(f"independent f64 reference, layer {args.layer}, " + f"ctx_len {int(acts['k_cache'].shape[1])}\n") + print(f" authored HIR (evaluator) rel_l2 vs f64 = {e_ref:.3e}") + print(f" TileLang twin rel_l2 vs f64 = {e_mine:.3e}") + print(f" twin/reference error ratio = {e_mine / e_ref:.3f}") + verdict = ("the twin is CLOSER to truth" if e_mine < e_ref + else "the reference is closer to truth") + print(f"\n -> {verdict}") + + +if __name__ == "__main__": + main() diff --git a/examples/qwen3_1_7b-tilelang/fast/engine.py b/examples/qwen3_1_7b-tilelang/fast/engine.py new file mode 100644 index 00000000..990e0c40 --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/fast/engine.py @@ -0,0 +1,307 @@ +"""One captured decode step, replayed. + +The authored reference hands each step's key and value back for the caller to +``torch.cat`` on. That is the right contract for a reference -- it keeps every +shape expressed in ``ctx_len`` alone -- but it means the cache buffer moves +every step, and a graph records addresses, so nothing built that way can be +replayed. This engine takes the other form the migrate page names: a cache of +fixed capacity whose write window advances, with the position in a one-element +device tensor. + +Everything a step needs then has a fixed address, so the whole step -- 284 +kernels, embedding through the greedy pick -- is captured once and replayed. The +chosen token is written back into the input slot by the last kernel, and while +the prompt still has a token left that kernel feeds that one instead, so the +same capture walks the prompt and continues past it with no host round trip +anywhere in the loop. + +Weights are repacked once at load: ``q|k|v`` become one matrix and ``gate|up`` +another, because a decode GEMV is bandwidth-bound and its cost is the block +count it can fill, not the arithmetic. Two fused reads beat five thin ones. +""" +from __future__ import annotations + +import importlib.util +import json +from dataclasses import dataclass +from pathlib import Path +from time import perf_counter + +import torch + +import kernels as K + +#: Tile shapes per projection, chosen by measuring each under a graph replay: +#: ``(BN, BK, SK, threads)``. ``SK`` is what keeps an H200's 132 SMs busy on the +#: narrow projections -- ``N/BN`` blocks alone leaves most of them idle. +TILES = { + "qkv": (256, 64, 8, 256), + "o": (128, 128, 8, 128), + "gate_up": (256, 64, 2, 256), + "down": (128, 128, 8, 128), + "head": (128, 128, 128), +} +#: Context positions one attention block owns. +SPLIT = 256 + + +def _load_module(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +@dataclass(frozen=True) +class Generated: + """A continuation and the timing of the part that produced it.""" + + tokens: list[int] + seconds: float + prefill_seconds: float + prompt_steps: int + + @property + def tokens_per_second(self) -> float: + return len(self.tokens) / self.seconds + + +class Engine: + """A loaded Qwen3-1.7B and one captured decode step over it.""" + + def __init__(self, ckpt: str | Path, ref_dir: str | Path, *, device: str = "cuda:0", + max_new: int = 2048, prompt_room: int = 512): + from tilefoundry.runtime import SafetensorsResource + + ref_dir = Path(ref_dir) + self.ref = _load_module(ref_dir / "model.py", "ref_model") + alias_mod = _load_module(ref_dir / "hf_alias.py", "ref_alias") + cfg = self.ref.config + self.cfg = cfg + self.device = device + self.dt = torch.bfloat16 + + self.H = cfg.hidden_size + self.HQ = cfg.num_attention_heads + self.HKV = cfg.num_key_value_heads + self.D = cfg.head_dim + self.I = cfg.intermediate_size + self.V = cfg.vocab_size + self.L = cfg.num_hidden_layers + self.eps = cfg.rms_norm_eps + self.scale = self.D ** -0.5 + self.qkv_n = self.HQ * self.D + 2 * self.HKV * self.D + + self.nsteps = max_new + prompt_room + self.cap = ((self.nsteps + SPLIT - 1) // SPLIT) * SPLIT + self.ns = self.cap // SPLIT + + loaded = self.ref.Qwen3_1_7B.load( + SafetensorsResource(str(ckpt), device=device, alias=alias_mod.hf_alias(cfg)) + ) + self._pack(loaded) + del loaded + torch.cuda.empty_cache() + self._buffers() + self._kernels() + self._capture() + + # ---------------------------------------------------------------- weights + def _pack(self, loaded): + """Fuse the projections a single GEMV can serve, and keep the rest as is.""" + self.w_embed = loaded.constants["w_embed"] + self.w_head = loaded.constants["w_head"] + self.gamma_final = loaded.constants["gamma_final"] + self.layers = [] + for i in range(self.L): + c = getattr(loaded, f"layer{i}").constants + self.layers.append({ + "gamma_in": c["gamma_in"], + "gamma_post": c["gamma_post"], + "gamma_q": c["gamma_q"], + "gamma_k": c["gamma_k"], + # [q | k | v] over the output axis: the rope kernel splits it back + "w_qkv": torch.cat([c["w_q"][0], c["w_k"][0], c["w_v"][0]], dim=1) + .contiguous(), + "w_o": c["w_o"][0].contiguous(), + # [gate | up]: silu_mul consumes both halves of one partial + "w_gu": torch.cat([c["w_gate"][0], c["w_up"][0]], dim=1).contiguous(), + "w_down": c["w_down"][0].contiguous(), + }) + cos, sin = self.ref._generation_rope(self.device) + self.cos = cos[: self.cap].contiguous() + self.sin = sin[: self.cap].contiguous() + + # ---------------------------------------------------------------- buffers + def _buffers(self): + dev, dt, f32 = self.device, self.dt, torch.float32 + + def z(*shape, dtype=dt): + return torch.zeros(*shape, device=dev, dtype=dtype) + + self.hid = z(self.H) + self.xn = z(self.H) + self.h1 = z(self.H) + self.xn1 = z(self.H) + self.q = z(self.HQ * self.D) + self.qkv_part = z(TILES["qkv"][2], self.qkv_n, dtype=f32) + self.o_part = z(TILES["o"][2], self.H, dtype=f32) + self.gu_part = z(TILES["gate_up"][2], 2 * self.I, dtype=f32) + self.d_part = z(TILES["down"][2], self.H, dtype=f32) + self.op = z(self.ns, self.HQ, self.D, dtype=f32) + self.mp = z(self.ns, self.HQ, dtype=f32) + self.lp = z(self.ns, self.HQ, dtype=f32) + self.logits = z(self.V, dtype=f32) + # A zeroed cache is what makes masking enough: a position past the end + # contributes exp(-inf) * 0, not exp(-inf) * whatever was there. + self.kc = z(self.L, self.cap, self.HKV * self.D) + self.vc = z(self.L, self.cap, self.HKV * self.D) + self.pos = z(1, dtype=torch.int32) + self.ids = z(1, dtype=torch.int64) + self.inp = z(self.nsteps, dtype=torch.int32) + self.sam = z(self.nsteps, dtype=torch.int32) + self.plen = z(1, dtype=torch.int32) + + # ---------------------------------------------------------------- kernels + def _kernels(self): + H, HQ, HKV, D, I, V = self.H, self.HQ, self.HKV, self.D, self.I, self.V + self.k_embed = K.embed(V, H) + self.k_norm = K.rms_norm(H, self.eps) + self.k_qkv = K.gemv(H, self.qkv_n, *TILES["qkv"]) + self.k_rope = K.qk_rope_cache( + HQ, HKV, D, self.cap, self.cap, TILES["qkv"][2], self.eps + ) + self.k_attn = K.attn_partial(HQ, HKV, D, self.cap, SPLIT, self.scale) + # o_proj merges the attention splits itself; down_proj activates its own + # slice of the gate/up partial. Two fewer launches per layer. + self.k_o = K.gemv_attn_combine(HQ, D, H, *TILES["o"][:3], self.ns, + TILES["o"][3]) + self.k_rn_post = K.resid_rms_norm(H, TILES["o"][2], self.eps) + self.k_gu = K.gemv(H, 2 * I, *TILES["gate_up"]) + self.k_down = K.gemv_silu(I, H, *TILES["down"][:3], + TILES["gate_up"][2], TILES["down"][3]) + self.k_rn_in = K.resid_rms_norm(H, TILES["down"][2], self.eps) + self.k_head, self.nb = K.lm_head(H, V, *TILES["head"]) + self.bv = torch.zeros(self.nb, device=self.device, dtype=torch.float32) + self.bi = torch.zeros(self.nb, device=self.device, dtype=torch.int32) + self.k_sample = K.sample_step(self.nb, self.nsteps) + + def _step(self, record=None): + """One decode step, as the sequence of launches the graph records. + + *record*, when given, collects the hidden state entering the stack and + leaving each layer -- the same series ``output_hidden_states`` returns, + so a disagreement can be pinned to a layer instead of to the model. + """ + self.k_embed(self.w_embed, self.ids, self.hid) + self.k_norm(self.hid, self.layers[0]["gamma_in"], self.xn) + if record is not None: + record.append(self.hid.clone()) + for i, w in enumerate(self.layers): + kc, vc = self.kc[i], self.vc[i] + self.k_qkv(self.xn, w["w_qkv"], self.qkv_part) + self.k_rope(self.qkv_part, w["gamma_q"], w["gamma_k"], self.cos, self.sin, + self.pos, self.pos, kc, vc, self.q) + self.k_attn(self.q, kc, vc, self.pos, self.op, self.mp, self.lp) + self.k_o(self.op, self.mp, self.lp, w["w_o"], self.o_part) + self.k_rn_post(self.hid, self.o_part, w["gamma_post"], self.h1, self.xn1) + self.k_gu(self.xn1, w["w_gu"], self.gu_part) + self.k_down(self.gu_part, w["w_down"], self.d_part) + # The next layer's input norm, or the norm that closes the stack -- + # so the residual add never needs a kernel to itself. + nxt = (self.layers[i + 1]["gamma_in"] if i + 1 < self.L else self.gamma_final) + self.k_rn_in(self.h1, self.d_part, nxt, self.hid, self.xn) + if record is not None: + record.append(self.hid.clone()) + self.k_head(self.xn, self.w_head, self.logits, self.bv, self.bi) + self.k_sample(self.bv, self.bi, self.inp, self.plen, self.ids, self.pos, self.sam) + + def _capture(self): + side = torch.cuda.Stream(device=self.device) + side.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(side): + for _ in range(3): + self._step() + torch.cuda.current_stream(self.device).wait_stream(side) + torch.cuda.synchronize(self.device) + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self._step() + self._reset() + + def _reset(self): + self.pos.zero_() + self.kc.zero_() + self.vc.zero_() + self.sam.zero_() + + # ------------------------------------------------------------- generation + def generate(self, prompt_ids: list[int], max_new: int) -> Generated: + """Walk *prompt_ids*, then continue for *max_new* tokens. + + Timing covers exactly the steps that produce the continuation: the step + at ``pos = len(prompt) - 1`` is the first one whose pick is kept, so the + steps before it are prefill and are not counted. + """ + pl = len(prompt_ids) + if pl < 1: + raise ValueError("decode needs a prompt of at least one token") + if pl + max_new > self.nsteps: + raise ValueError( + f"prompt {pl} + {max_new} new exceeds this engine's {self.nsteps} steps" + ) + self._reset() + self.inp.zero_() + self.inp[:pl] = torch.tensor(prompt_ids, device=self.device, dtype=torch.int32) + self.plen.fill_(pl) + self.ids.fill_(prompt_ids[0]) + + torch.cuda.synchronize(self.device) + t0 = perf_counter() + for _ in range(pl - 1): # prefill: picks discarded + self.graph.replay() + torch.cuda.synchronize(self.device) + t1 = perf_counter() + for _ in range(max_new): # the continuation itself + self.graph.replay() + torch.cuda.synchronize(self.device) + t2 = perf_counter() + + out = self.sam[pl - 1: pl - 1 + max_new].tolist() + return Generated(out, t2 - t1, t1 - t0, pl) + + def trace(self, prompt_ids: list[int]) -> list[torch.Tensor]: + """Per-layer hidden states at the last prompt position.""" + self._reset() + self.inp.zero_() + self.inp[: len(prompt_ids)] = torch.tensor( + prompt_ids, device=self.device, dtype=torch.int32 + ) + self.plen.fill_(len(prompt_ids)) + self.ids.fill_(prompt_ids[0]) + for _ in range(len(prompt_ids) - 1): + self.graph.replay() + rec: list[torch.Tensor] = [] + self._step(record=rec) + torch.cuda.synchronize(self.device) + return rec + + def logits_for(self, prompt_ids: list[int]) -> torch.Tensor: + """The logits after consuming every token of *prompt_ids* -- for checking.""" + self._reset() + self.inp.zero_() + self.inp[: len(prompt_ids)] = torch.tensor( + prompt_ids, device=self.device, dtype=torch.int32 + ) + self.plen.fill_(len(prompt_ids)) + self.ids.fill_(prompt_ids[0]) + for _ in range(len(prompt_ids)): + self.graph.replay() + torch.cuda.synchronize(self.device) + return self.logits.clone() + + +def default_ref_dir() -> Path: + """The pristine copy of the shipped source this engine is measured against.""" + here = Path(__file__).resolve().parent + return here.parent / "ref_src" diff --git a/examples/qwen3_1_7b-tilelang/fast/kernels.py b/examples/qwen3_1_7b-tilelang/fast/kernels.py new file mode 100644 index 00000000..0199e09f --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/fast/kernels.py @@ -0,0 +1,637 @@ +"""TileLang kernels for one Qwen3-1.7B decode step. + +Every kernel takes caller-owned output buffers -- no allocation, and no shape or +trip count that depends on a host value -- because the decode loop replays them +from a captured CUDA graph. The step's position is read from a one-element +device tensor (``Pos``) instead of being baked in, which is what lets a single +capture serve every step of a generation. + +Two shapes recur and set the tiling: + +* **GEMV.** Decode is one token, so every projection is ``(K,) @ (K, N)`` -- + pure memory traffic, no reuse. A block owns ``BN`` output columns and walks + ``K``; weights are stored ``(K, N)``, so a row slice is contiguous and + coalesces. ``N/BN`` blocks alone leave most of an H200's 132 SMs idle on the + small projections, so ``K`` is split ``SK`` ways as well and the partial sums + land in an ``(SK, N)`` f32 buffer that the *consumer* reduces -- the reduce is + never a kernel of its own. + +* **Two-pass attention.** One query row against ``pos+1`` cached keys. A block + owns ``SS`` context positions: pass one reads K and fills the whole score row, + then the row is normalised, then pass two reads V and weights it. Holding the + split's entire score row means no rescale happens inside a block -- only + across blocks, which ``gemv_attn_combine`` folds into ``o_proj``. A block whose + slice starts past the current length writes a neutral partial and exits, so a + fixed grid still costs only the context that exists. + +Two TileLang facts shape the code below and are worth stating once: + +* A fragment lives in the owning thread's registers. Anything one thread must + read that another wrote goes through ``alloc_shared``; ``T.reduce_*`` is the + exception, because its one-element result is replicated and so readable by + every thread. +* A kernel body is source-rewritten, and a plain helper function is not. So + loop-emitting code is written inline here rather than factored out -- a + ``T.Parallel`` inside an ordinary callee would execute as Python and fail. +""" +from functools import lru_cache + +import tilelang +import tilelang.language as T + +DT = "bfloat16" +ACC = "float32" +NEG = -1.0e30 + +tilelang.disable_cache() + + +@lru_cache(maxsize=None) +def _compile(prim): + return tilelang.compile(prim) + + +# --------------------------------------------------------------------------- # +# norms +# --------------------------------------------------------------------------- # + +@lru_cache(maxsize=None) +def rms_norm(H: int, eps: float, threads: int = 256): + """``Xn = bf16(x * rsqrt(mean(x^2) + eps)) * gamma``, one block. + + Qwen3RMSNorm rounds to the input dtype *before* the learned scale, so the + cast sits inside the product rather than after it -- the order the authored + HIR spells out, and not the one a generic rms_norm would take. + """ + @T.prim_func + def main(X: T.Tensor((H,), DT), G: T.Tensor((H,), DT), Xn: T.Tensor((H,), DT)): + with T.Kernel(1, threads=threads): + xs = T.alloc_fragment((H,), ACC) + sq = T.alloc_fragment((H,), ACC) + tot = T.alloc_fragment((1,), ACC) + for i in T.Parallel(H): + xs[i] = X[i].astype(ACC) + sq[i] = xs[i] * xs[i] + T.reduce_sum(sq, tot, dim=0) + for i in T.Parallel(H): + Xn[i] = (xs[i] * T.rsqrt(tot[0] / H + eps)).astype(DT) * G[i] + + return _compile(main) + + +@lru_cache(maxsize=None) +def resid_rms_norm(H: int, SK: int, eps: float, threads: int = 256): + """``Hout = A + reduce(P)``; ``Xn = norm(Hout) * gamma``. One block. + + The residual add and the norm that follows it are the same read of the same + vector, so they are one kernel -- and *P*, the producer's split-K partial, + is reduced here rather than by a kernel of its own. + """ + @T.prim_func + def main( + A: T.Tensor((H,), DT), + P: T.Tensor((SK, H), ACC), + G: T.Tensor((H,), DT), + Hout: T.Tensor((H,), DT), + Xn: T.Tensor((H,), DT), + ): + with T.Kernel(1, threads=threads): + xs = T.alloc_fragment((H,), ACC) + sq = T.alloc_fragment((H,), ACC) + acc = T.alloc_fragment((H,), ACC) + tot = T.alloc_fragment((1,), ACC) + T.clear(acc) + for s in T.serial(SK): + for i in T.Parallel(H): + acc[i] += P[s, i] + for i in T.Parallel(H): + v = A[i] + acc[i].astype(DT) + Hout[i] = v + xs[i] = v.astype(ACC) + sq[i] = xs[i] * xs[i] + T.reduce_sum(sq, tot, dim=0) + for i in T.Parallel(H): + Xn[i] = (xs[i] * T.rsqrt(tot[0] / H + eps)).astype(DT) * G[i] + + return _compile(main) + + +# --------------------------------------------------------------------------- # +# GEMV +# --------------------------------------------------------------------------- # + +@lru_cache(maxsize=None) +def gemv(K: int, N: int, BN: int, BK: int, SK: int, threads: int, stages: int = 3): + """``P[s, n] = sum(X[k] * W[k, n])`` over split *s*'s share of ``K``.""" + KS = K // SK + + @T.prim_func + def main(X: T.Tensor((K,), DT), W: T.Tensor((K, N), DT), P: T.Tensor((SK, N), ACC)): + with T.Kernel(T.ceildiv(N, BN), SK, threads=threads) as (bx, bs): + Ws = T.alloc_shared((BK, BN), DT) + Xs = T.alloc_shared((BK,), DT) + acc = T.alloc_fragment((BN,), ACC) + T.clear(acc) + for ko in T.Pipelined(KS // BK, num_stages=stages): + k0 = bs * KS + ko * BK + T.copy(W[k0:k0 + BK, bx * BN:(bx + 1) * BN], Ws) + T.copy(X[k0:k0 + BK], Xs) + for j in T.Parallel(BN): + for kk in T.serial(BK): + acc[j] += Xs[kk].astype(ACC) * Ws[kk, j].astype(ACC) + for j in T.Parallel(BN): + P[bs, bx * BN + j] = acc[j] + + return _compile(main) + + +#: Two GEMVs below fold their producer in rather than reading a materialised +#: vector. The reason is not the traffic saved -- these vectors are a few KB -- +#: but the kernel saved. A split-K block only ever consumes ``K / SK`` of its +#: input, and for both of these that slice can be built from partials in L2, so +#: the producer stops being a separate launch with its own ramp and drain. + +@lru_cache(maxsize=None) +def gemv_attn_combine( + HQ: int, D: int, N: int, BN: int, BK: int, SK: int, NS: int, + threads: int, stages: int = 3, +): + """``o_proj``, with the attention splits merged into its own input read. + + A block owning ``K / SK`` inputs owns a whole number of attention heads + (``K = HQ * D`` and the split divides the head count), so it can merge just + those heads' partials itself. Each thread redoes the ``NS``-term log-sum-exp + for its own entry -- the partials are a few KB and sit in L2. + """ + K = HQ * D + KS = K // SK + + @T.prim_func + def main( + Op: T.Tensor((NS, HQ, D), ACC), + Mp: T.Tensor((NS, HQ), ACC), + Lp: T.Tensor((NS, HQ), ACC), + W: T.Tensor((K, N), DT), + P: T.Tensor((SK, N), ACC), + ): + with T.Kernel(T.ceildiv(N, BN), SK, threads=threads) as (bx, bs): + Ws = T.alloc_shared((BK, BN), DT) + Xs = T.alloc_shared((KS,), DT) + mx = T.alloc_fragment((KS,), ACC) + den = T.alloc_fragment((KS,), ACC) + num = T.alloc_fragment((KS,), ACC) + acc = T.alloc_fragment((BN,), ACC) + for i in T.Parallel(KS): + h = (bs * KS + i) // D + d = (bs * KS + i) % D + mx[i] = NEG + for s in T.serial(NS): + mx[i] = T.max(mx[i], Mp[s, h]) + den[i] = 0.0 + num[i] = 0.0 + for s in T.serial(NS): + den[i] += Lp[s, h] * T.exp(Mp[s, h] - mx[i]) + num[i] += Op[s, h, d] * T.exp(Mp[s, h] - mx[i]) + Xs[i] = (num[i] / den[i]).astype(DT) + T.sync_threads() + T.clear(acc) + for ko in T.Pipelined(KS // BK, num_stages=stages): + k0 = bs * KS + ko * BK + T.copy(W[k0:k0 + BK, bx * BN:(bx + 1) * BN], Ws) + for j in T.Parallel(BN): + for kk in T.serial(BK): + acc[j] += Xs[ko * BK + kk].astype(ACC) * Ws[kk, j].astype(ACC) + for j in T.Parallel(BN): + P[bs, bx * BN + j] = acc[j] + + return _compile(main) + + +@lru_cache(maxsize=None) +def gemv_silu( + I: int, N: int, BN: int, BK: int, SK: int, SKG: int, threads: int, stages: int = 3 +): + """``down_proj``, with ``silu(gate) * up`` folded into its own input read. + + *GU* is the fused gate/up GEMV's partial, gate in ``[0, I)`` and up in + ``[I, 2I)``. A block reduces and activates only the ``I / SK`` entries it + walks, so the intermediate never reaches HBM in either direction. + """ + KS = I // SK + + @T.prim_func + def main( + GU: T.Tensor((SKG, 2 * I), ACC), + W: T.Tensor((I, N), DT), + P: T.Tensor((SK, N), ACC), + ): + with T.Kernel(T.ceildiv(N, BN), SK, threads=threads) as (bx, bs): + Ws = T.alloc_shared((BK, BN), DT) + Xs = T.alloc_shared((KS,), DT) + g = T.alloc_fragment((KS,), ACC) + u = T.alloc_fragment((KS,), ACC) + acc = T.alloc_fragment((BN,), ACC) + T.clear(g) + T.clear(u) + for s in T.serial(SKG): + for i in T.Parallel(KS): + g[i] += GU[s, bs * KS + i] + u[i] += GU[s, I + bs * KS + i] + for i in T.Parallel(KS): + gb = g[i].astype(DT).astype(ACC) + ub = u[i].astype(DT).astype(ACC) + Xs[i] = ( + (gb / (1.0 + T.exp(-gb))).astype(DT).astype(ACC) * ub + ).astype(DT) + T.sync_threads() + T.clear(acc) + for ko in T.Pipelined(KS // BK, num_stages=stages): + k0 = bs * KS + ko * BK + T.copy(W[k0:k0 + BK, bx * BN:(bx + 1) * BN], Ws) + for j in T.Parallel(BN): + for kk in T.serial(BK): + acc[j] += Xs[ko * BK + kk].astype(ACC) * Ws[kk, j].astype(ACC) + for j in T.Parallel(BN): + P[bs, bx * BN + j] = acc[j] + + return _compile(main) + + +@lru_cache(maxsize=None) +def lm_head(K: int, N: int, BN: int, BK: int, threads: int, stages: int = 3): + """The head, plus each block's own best entry. + + Greedy sampling wants one index out of 152k, and a block already holds its + columns in registers -- so it reduces them here and writes ``(value, + index)``, leaving a 1187-entry reduction instead of a second full pass over + the logits. The logits are written too: they cost one store and are what a + comparison against the reference reads. + """ + NB = (N + BN - 1) // BN + + @T.prim_func + def main( + X: T.Tensor((K,), DT), + W: T.Tensor((K, N), DT), + O: T.Tensor((N,), ACC), + Bv: T.Tensor((NB,), ACC), + Bi: T.Tensor((NB,), "int32"), + ): + with T.Kernel(NB, threads=threads) as bx: + Ws = T.alloc_shared((BK, BN), DT) + Xs = T.alloc_shared((BK,), DT) + acc = T.alloc_fragment((BN,), ACC) + stg = T.alloc_shared((BN,), ACC) + red = T.alloc_fragment((BN,), ACC) + sel = T.alloc_fragment((BN,), ACC) + mx = T.alloc_fragment((1,), ACC) + win = T.alloc_fragment((1,), ACC) + T.clear(acc) + for ko in T.Pipelined(K // BK, num_stages=stages): + T.copy(W[ko * BK:(ko + 1) * BK, bx * BN:(bx + 1) * BN], Ws) + T.copy(X[ko * BK:(ko + 1) * BK], Xs) + for j in T.Parallel(BN): + for kk in T.serial(BK): + acc[j] += Xs[kk].astype(ACC) * Ws[kk, j].astype(ACC) + # The block's own best, found without a second pass over the logits. + # Two things matter about how it is written: + # + # * `acc` is staged to shared and the reductions read a *fresh* + # fragment. Reducing `acc` twice in place makes layout inference + # replicate it across all 128 threads to satisfy both uses, which + # spills the accumulator and costs 19ms instead of 165us. + # * Argmax is two max-reductions, no atomic: the winning value, then + # `BN - j` over the entries attaining it, whose max is the *lowest* + # winning index -- torch.argmax's tie-break. An `atomic_min` on + # shared costs 18ms here; a reduce_max costs 2us. + T.copy(acc, stg) + for j in T.Parallel(BN): + O[bx * BN + j] = stg[j] + T.sync_threads() + for j in T.Parallel(BN): + red[j] = stg[j] + T.reduce_max(red, mx, dim=0) + for j in T.Parallel(BN): + sel[j] = T.if_then_else(red[j] >= mx[0], (BN - j) * 1.0, 0.0) + T.reduce_max(sel, win, dim=0) + if T.get_thread_binding() == 0: + Bv[bx] = mx[0] + Bi[bx] = bx * BN + BN - T.Cast("int32", win[0]) + + return _compile(main), NB + + +# --------------------------------------------------------------------------- # +# q/k norm + rope + cache write +# --------------------------------------------------------------------------- # + +@lru_cache(maxsize=None) +def qk_rope_cache(HQ: int, HKV: int, D: int, MP: int, CAP: int, SK: int, eps: float): + """Reduce the fused QKV partials, then per-head norm, rope, and cache write. + + *P* is laid out ``[q | k | v]`` over its last axis, one block per head of + the three. Query heads keep going, into *Q*. Key and value heads stop here: + this step's entry is written straight into the cache at ``Pos``, so no + caller appends anything and the cache buffer never moves, which is what + makes the step replayable from a fixed graph. + + ``Pos`` is the rotary position and ``Wr`` the slot to write. Decoding passes + the same tensor twice, but the authored reference lets them differ -- it + takes ``pos_ids`` and a prior cache of its own length -- so they are separate + parameters, and the twin that stands for that reference can honour it. + + The ``1/sqrt(head_dim)`` factor is deliberately *not* applied here. The + authored reference multiplies it onto ``q`` in bf16 before the dot, which + rounds every entry once more -- up to 2^-9 relative, and the exponential + downstream turns that into percent-level error on the attention weights. + Hugging Face scales after the dot instead, so ``attn_partial`` folds the + factor into its f32 accumulation. + """ + QN, KN = HQ * D, HKV * D + TOT = HQ + 2 * HKV + + @T.prim_func + def main( + P: T.Tensor((SK, QN + 2 * KN), ACC), + Gq: T.Tensor((D,), DT), + Gk: T.Tensor((D,), DT), + Cos: T.Tensor((MP, D), DT), + Sin: T.Tensor((MP, D), DT), + Pos: T.Tensor((1,), "int32"), + Wr: T.Tensor((1,), "int32"), + Kc: T.Tensor((CAP, KN), DT), + Vc: T.Tensor((CAP, KN), DT), + Q: T.Tensor((HQ * D,), DT), + ): + with T.Kernel(TOT, threads=D) as bh: + xs = T.alloc_fragment((D,), ACC) + sq = T.alloc_fragment((D,), ACC) + acc = T.alloc_fragment((D,), ACC) + tot = T.alloc_fragment((1,), ACC) + nrm = T.alloc_shared((D,), DT) # rope pairs d with d +/- D/2 + head = T.if_then_else( + bh < HQ, bh, T.if_then_else(bh < HQ + HKV, bh - HQ, bh - HQ - HKV) + ) + base = T.if_then_else( + bh < HQ, 0, T.if_then_else(bh < HQ + HKV, QN, QN + KN) + ) + head * D + T.clear(acc) + for s in T.serial(SK): + for d in T.Parallel(D): + acc[d] += P[s, base + d] + for d in T.Parallel(D): + xs[d] = acc[d].astype(DT).astype(ACC) + sq[d] = xs[d] * xs[d] + with T.If(bh >= HQ + HKV): + with T.Then(): + for d in T.Parallel(D): + Vc[Wr[0], head * D + d] = xs[d].astype(DT) + with T.Else(): + T.reduce_sum(sq, tot, dim=0) + for d in T.Parallel(D): + g = T.if_then_else(bh < HQ, Gq[d], Gk[d]) + nrm[d] = (xs[d] * T.rsqrt(tot[0] / D + eps)).astype(DT) * g + T.sync_threads() + for d in T.Parallel(D): + half = T.if_then_else( + d < D // 2, + -nrm[d + D // 2].astype(ACC), + nrm[d - D // 2].astype(ACC), + ) + rot = ( + nrm[d].astype(ACC) * Cos[Pos[0], d].astype(ACC) + + half * Sin[Pos[0], d].astype(ACC) + ).astype(DT) + with T.If(bh < HQ): + with T.Then(): + Q[head * D + d] = rot + with T.Else(): + Kc[Wr[0], head * D + d] = rot + + return _compile(main) + + +# --------------------------------------------------------------------------- # +# attention +# --------------------------------------------------------------------------- # + +@lru_cache(maxsize=None) +def attn_partial( + HQ: int, HKV: int, D: int, CAP: int, SS: int, scale: float, threads: int = 128 +): + """One split's ``(max, sum, weighted values)`` over its slice of the context. + + Grid is ``(splits, kv heads)``: one kv head serves its whole GQA group, so a + key is read once and used by every query head that shares it, rather than + being materialised per head the way a literal reading of the reference's + ``repeat_interleave`` would. + + *scale* is applied to the finished f32 dot product, not to ``q`` beforehand + -- see ``qk_rope_cache`` for why that ordering matters here. + + Both products go through ``T.gemm``, which means padding the group's ``G`` + query rows up to MMA's smallest ``M`` of 16 and leaving the rest zero. The + 8x arithmetic waste costs nothing -- one query row against a few hundred + keys is latency-bound, not flop-bound -- and it buys TileLang's own MMA + shared-memory layouts. Written as a scalar loop instead, consecutive threads + read one row apart, every address lands in the same bank, and the kernel + runs at 0.4 TB/s; this way it is 2.2x quicker. + """ + G = HQ // HKV + NS = (CAP + SS - 1) // SS + M = 16 + + @T.prim_func + def main( + Q: T.Tensor((HQ * D,), DT), + Kc: T.Tensor((CAP, HKV * D), DT), + Vc: T.Tensor((CAP, HKV * D), DT), + Pos: T.Tensor((1,), "int32"), + Op: T.Tensor((NS, HQ, D), ACC), + Mp: T.Tensor((NS, HQ), ACC), + Lp: T.Tensor((NS, HQ), ACC), + ): + with T.Kernel(NS, HKV, threads=threads) as (bs, bh): + Qs = T.alloc_shared((M, D), DT) + Ks = T.alloc_shared((SS, D), DT) + Vs = T.alloc_shared((SS, D), DT) + sc = T.alloc_fragment((M, SS), ACC) + scb = T.alloc_shared((M, SS), DT) + orun = T.alloc_fragment((M, D), ACC) + mx = T.alloc_fragment((M,), ACC) + sm = T.alloc_fragment((M,), ACC) + with T.If(bs * SS < Pos[0] + 1): + with T.Then(): + for m, d in T.Parallel(M, D): + Qs[m, d] = T.if_then_else( + m < G, + Q[(bh * G + T.min(m, G - 1)) * D + d], + T.Cast(DT, 0.0), + ) + T.copy(Kc[bs * SS:(bs + 1) * SS, bh * D:(bh + 1) * D], Ks) + T.copy(Vc[bs * SS:(bs + 1) * SS, bh * D:(bh + 1) * D], Vs) + T.clear(sc) + T.gemm(Qs, Ks, sc, transpose_B=True) + for m, s in T.Parallel(M, SS): + # past the end of the context this position does not exist + sc[m, s] = T.if_then_else( + bs * SS + s < Pos[0] + 1, sc[m, s] * scale, NEG + ) + T.reduce_max(sc, mx, dim=1, clear=True) + for m, s in T.Parallel(M, SS): + sc[m, s] = T.exp(sc[m, s] - mx[m]) + T.reduce_sum(sc, sm, dim=1, clear=True) + # bf16 probabilities into the second product, which is what + # HF's own attention does before it multiplies by V + T.copy(sc, scb) + T.clear(orun) + T.gemm(scb, Vs, orun) + for g, d in T.Parallel(G, D): + Op[bs, bh * G + g, d] = orun[g, d] + for g in T.Parallel(G): + Mp[bs, bh * G + g] = mx[g] + Lp[bs, bh * G + g] = sm[g] + with T.Else(): + for g, d in T.Parallel(G, D): + Op[bs, bh * G + g, d] = 0.0 + for g in T.Parallel(G): + Mp[bs, bh * G + g] = NEG + Lp[bs, bh * G + g] = 0.0 + + return _compile(main) + + +@lru_cache(maxsize=None) +def attn_combine(HQ: int, D: int, CAP: int, SS: int): + """Merge the splits' partials against their joint max; head-major flatten. + + Head-major is not a choice: ``w_o`` was stored expecting attention entry + ``(h, d)`` at ``h * D + d``, matching the authored reshape. + """ + NS = (CAP + SS - 1) // SS + + @T.prim_func + def main( + Op: T.Tensor((NS, HQ, D), ACC), + Mp: T.Tensor((NS, HQ), ACC), + Lp: T.Tensor((NS, HQ), ACC), + O: T.Tensor((HQ * D,), DT), + ): + with T.Kernel(HQ, threads=D) as bh: + mx = T.alloc_fragment((D,), ACC) + den = T.alloc_fragment((D,), ACC) + acc = T.alloc_fragment((D,), ACC) + # Every thread needs the joint max and denominator, and there are + # only NS of each -- so each recomputes them from L2 rather than one + # reducing and broadcasting. A cross-thread reduce over NS values + # has no layout anyway when NS does not divide the thread count. + for d in T.Parallel(D): + mx[d] = NEG + for s in T.serial(NS): + mx[d] = T.max(mx[d], Mp[s, bh]) + den[d] = 0.0 + acc[d] = 0.0 + for s in T.serial(NS): + e = T.exp(Mp[s, bh] - mx[d]) + den[d] += Lp[s, bh] * e + acc[d] += Op[s, bh, d] * e + O[bh * D + d] = (acc[d] / den[d]).astype(DT) + + return _compile(main) + + +# --------------------------------------------------------------------------- # +# MLP activation +# --------------------------------------------------------------------------- # + +@lru_cache(maxsize=None) +def silu_mul(I: int, SK: int, BN: int = 256, threads: int = 256): + """Reduce the fused gate/up partials, then ``silu(gate) * up``. + + *P* holds gate in ``[0, I)`` and up in ``[I, 2I)``: one GEMV produced both, + so one kernel consumes both and neither reaches HBM on its own. + """ + @T.prim_func + def main(P: T.Tensor((SK, 2 * I), ACC), O: T.Tensor((I,), DT)): + with T.Kernel(T.ceildiv(I, BN), threads=threads) as bx: + g = T.alloc_fragment((BN,), ACC) + u = T.alloc_fragment((BN,), ACC) + T.clear(g) + T.clear(u) + for s in T.serial(SK): + for j in T.Parallel(BN): + g[j] += P[s, bx * BN + j] + u[j] += P[s, I + bx * BN + j] + for j in T.Parallel(BN): + gb = g[j].astype(DT).astype(ACC) + ub = u[j].astype(DT).astype(ACC) + O[bx * BN + j] = ( + (gb / (1.0 + T.exp(-gb))).astype(DT).astype(ACC) * ub + ).astype(DT) + + return _compile(main) + + +# --------------------------------------------------------------------------- # +# embedding and sampling +# --------------------------------------------------------------------------- # + +@lru_cache(maxsize=None) +def embed(V: int, H: int, threads: int = 256): + """The decoded token's own row of the table.""" + @T.prim_func + def main(Tbl: T.Tensor((V, H), DT), Ids: T.Tensor((1,), "int64"), O: T.Tensor((H,), DT)): + with T.Kernel(T.ceildiv(H, threads), threads=threads) as bx: + for i in T.Parallel(threads): + O[bx * threads + i] = Tbl[Ids[0], bx * threads + i] + + return _compile(main) + + +@lru_cache(maxsize=None) +def sample_step(NB: int, NSTEPS: int, threads: int = 256): + """Finish the greedy pick, record it, hand the next input on, advance ``Pos``. + + This is the graph's last node and the only one that decides anything: while + the prompt still has a token left it feeds that, otherwise it feeds what it + just sampled. That one device-side choice is what lets a single capture + serve the prompt walk and the continuation with no host round trip between + steps. + """ + PAD = ((NB + threads - 1) // threads) * threads + + @T.prim_func + def main( + Bv: T.Tensor((NB,), ACC), + Bi: T.Tensor((NB,), "int32"), + Inp: T.Tensor((NSTEPS,), "int32"), + PromptLen: T.Tensor((1,), "int32"), + Ids: T.Tensor((1,), "int64"), + Pos: T.Tensor((1,), "int32"), + Sam: T.Tensor((NSTEPS,), "int32"), + ): + with T.Kernel(1, threads=threads): + # Padded to a whole number of thread-rounds: a fragment whose size + # does not divide the thread count has no reduce layout. + f = T.alloc_fragment((PAD,), ACC) + sel = T.alloc_fragment((PAD,), ACC) + mx = T.alloc_fragment((1,), ACC) + win = T.alloc_fragment((1,), ACC) + for k in T.Parallel(PAD): + f[k] = T.if_then_else(k < NB, Bv[T.min(k, NB - 1)], NEG) + T.reduce_max(f, mx, dim=0) + for k in T.Parallel(PAD): + sel[k] = T.if_then_else( + (k < NB) and (f[k] >= mx[0]), (NB - k) * 1.0, 0.0 + ) + T.reduce_max(sel, win, dim=0) + if T.get_thread_binding() == 0: + p = Pos[0] + best = Bi[NB - T.Cast("int32", win[0])] + Sam[p] = best + Ids[0] = T.if_then_else( + p + 1 < PromptLen[0], Inp[p + 1].astype("int64"), best.astype("int64") + ) + Pos[0] = p + 1 + + return _compile(main) diff --git a/examples/qwen3_1_7b-tilelang/fast/test_kernels.py b/examples/qwen3_1_7b-tilelang/fast/test_kernels.py new file mode 100644 index 00000000..c8b85191 --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/fast/test_kernels.py @@ -0,0 +1,204 @@ +"""Each kernel against a torch statement of the same thing, at real dimensions.""" +from __future__ import annotations + +import math +import sys + +import torch + +import kernels as K + +DEV = "cuda:0" +H, HQ, HKV, D, I, V = 2048, 16, 8, 128, 6144, 151936 +EPS = 1e-6 +MP = 4096 +FAIL = [] + + +def report(name, got, ref, rtol=2e-2, atol=2e-2): + got32, ref32 = got.float(), ref.float() + err = (got32 - ref32).abs() + scale = ref32.abs().max().clamp(min=1e-6) + rel = (err.max() / scale).item() + cos = torch.nn.functional.cosine_similarity( + got32.flatten(), ref32.flatten(), dim=0 + ).item() + ok = rel < rtol and cos > 1 - atol + print(f"{'ok ' if ok else 'FAIL'} {name:22s} max_rel={rel:.3g} cosine={cos:.6f}") + if not ok: + FAIL.append(name) + + +def t_rms_norm(): + x = torch.randn(H, device=DEV, dtype=torch.bfloat16) + g = torch.randn(H, device=DEV, dtype=torch.bfloat16) + out = torch.empty(H, device=DEV, dtype=torch.bfloat16) + K.rms_norm(H, EPS)(x, g, out) + x32 = x.float() + ref = (x32 * torch.rsqrt(x32.pow(2).mean() + EPS)).to(torch.bfloat16) * g + report("rms_norm", out, ref) + + +def t_resid_rms_norm(): + SK = 8 + a = torch.randn(H, device=DEV, dtype=torch.bfloat16) + p = torch.randn(SK, H, device=DEV, dtype=torch.float32) + g = torch.randn(H, device=DEV, dtype=torch.bfloat16) + ho = torch.empty(H, device=DEV, dtype=torch.bfloat16) + xn = torch.empty(H, device=DEV, dtype=torch.bfloat16) + K.resid_rms_norm(H, SK, EPS)(a, p, g, ho, xn) + hr = a + p.sum(0).to(torch.bfloat16) + report("resid_rms_norm.h", ho, hr) + h32 = hr.float() + ref = (h32 * torch.rsqrt(h32.pow(2).mean() + EPS)).to(torch.bfloat16) * g + report("resid_rms_norm.xn", xn, ref) + + +def t_gemv(): + for Kd, N, BN, BK, SK, thr in [ + (H, 4096, 256, 64, 8, 256), (H, H, 128, 128, 8, 128), + (H, 2 * I, 256, 64, 2, 256), (I, H, 128, 128, 8, 128), + ]: + x = torch.randn(Kd, device=DEV, dtype=torch.bfloat16) + w = torch.randn(Kd, N, device=DEV, dtype=torch.bfloat16) / Kd**0.5 + p = torch.empty(SK, N, device=DEV, dtype=torch.float32) + K.gemv(Kd, N, BN, BK, SK, thr)(x, w, p) + report(f"gemv {Kd}x{N}", p.sum(0), x.float() @ w.float(), rtol=5e-3) + + +def t_lm_head(): + x = torch.randn(H, device=DEV, dtype=torch.bfloat16) + w = torch.randn(H, V, device=DEV, dtype=torch.bfloat16) / H**0.5 + kern, NB = K.lm_head(H, V, 128, 128, 128) + o = torch.empty(V, device=DEV, dtype=torch.float32) + bv = torch.empty(NB, device=DEV, dtype=torch.float32) + bi = torch.empty(NB, device=DEV, dtype=torch.int32) + kern(x, w, o, bv, bi) + ref = x.float() @ w.float() + report("lm_head.logits", o, ref, rtol=5e-3) + got_idx = int(bi[int(bv.argmax())]) + exp_idx = int(ref.argmax()) + ok = got_idx == exp_idx + print(f"{'ok ' if ok else 'FAIL'} lm_head.argmax got={got_idx} want={exp_idx}") + if not ok: + FAIL.append("lm_head.argmax") + + +def _rope_ref(x, cos, sin, pos): + """x (heads, D) bf16 -> rotated, matching HF apply_rotary_pos_emb.""" + c, s = cos[pos].float(), sin[pos].float() + x1, x2 = x.float()[:, : D // 2], x.float()[:, D // 2:] + half = torch.cat([-x2, x1], dim=-1) + return (x.float() * c + half * s).to(torch.bfloat16) + + +def t_qk_rope_cache(): + SK, CAP, pos = 8, 512, 37 + scale = D**-0.5 + p = torch.randn(SK, HQ * D + 2 * HKV * D, device=DEV, dtype=torch.float32) / SK + gq = torch.randn(D, device=DEV, dtype=torch.bfloat16) + gk = torch.randn(D, device=DEV, dtype=torch.bfloat16) + cos = torch.randn(MP, D, device=DEV, dtype=torch.bfloat16) + sin = torch.randn(MP, D, device=DEV, dtype=torch.bfloat16) + pt = torch.tensor([pos], device=DEV, dtype=torch.int32) + kc = torch.zeros(CAP, HKV * D, device=DEV, dtype=torch.bfloat16) + vc = torch.zeros(CAP, HKV * D, device=DEV, dtype=torch.bfloat16) + q = torch.empty(HQ * D, device=DEV, dtype=torch.bfloat16) + K.qk_rope_cache(HQ, HKV, D, MP, CAP, SK, EPS)(p, gq, gk, cos, sin, pt, pt, kc, vc, q) + + flat = p.sum(0).to(torch.bfloat16).float() + qr = flat[: HQ * D].view(HQ, D) + kr = flat[HQ * D: HQ * D + HKV * D].view(HKV, D) + vr = flat[HQ * D + HKV * D:].view(HKV, D) + + def nrm(t, g): + return (t * torch.rsqrt(t.pow(2).mean(-1, keepdim=True) + EPS)).to(torch.bfloat16) * g + + report("qk_rope.q", q.view(HQ, D), _rope_ref(nrm(qr, gq), cos, sin, pos)) + report("qk_rope.k", kc[pos].view(HKV, D), _rope_ref(nrm(kr, gk), cos, sin, pos)) + report("qk_rope.v", vc[pos].view(HKV, D), vr.to(torch.bfloat16)) + + +def _attn_ref(q, kc, vc, cur): + """q (HQ,D) bf16 scaled; kc/vc (cur,HKV,D) -> (HQ*D,) bf16, f32 softmax.""" + G = HQ // HKV + k = kc[:cur].float().repeat_interleave(G, dim=1) # (cur, HQ, D) + v = vc[:cur].float().repeat_interleave(G, dim=1) + s = (q.float().unsqueeze(0) * k).sum(-1) * D**-0.5 # (cur, HQ) + w = torch.softmax(s, dim=0) + return (w.unsqueeze(-1) * v).sum(0).to(torch.bfloat16).reshape(-1) + + +def t_attn(): + CAP, SS, BS = 2560, 256, 128 + for cur in (1, 5, 256, 257, 700, 2048): + pos = cur - 1 + q = torch.randn(HQ, D, device=DEV, dtype=torch.bfloat16) + kc = torch.zeros(CAP, HKV, D, device=DEV, dtype=torch.bfloat16) + vc = torch.zeros(CAP, HKV, D, device=DEV, dtype=torch.bfloat16) + kc[:cur].normal_() + vc[:cur].normal_() + pt = torch.tensor([pos], device=DEV, dtype=torch.int32) + NS = CAP // SS + op = torch.empty(NS, HQ, D, device=DEV, dtype=torch.float32) + mp = torch.empty(NS, HQ, device=DEV, dtype=torch.float32) + lp = torch.empty(NS, HQ, device=DEV, dtype=torch.float32) + out = torch.empty(HQ * D, device=DEV, dtype=torch.bfloat16) + K.attn_partial(HQ, HKV, D, CAP, SS, D**-0.5)( + q.reshape(-1), kc.view(CAP, -1), vc.view(CAP, -1), pt, op, mp, lp) + K.attn_combine(HQ, D, CAP, SS)(op, mp, lp, out) + report(f"attn cur={cur}", out, _attn_ref(q, kc, vc, cur), rtol=3e-2, atol=3e-2) + + +def t_silu_mul(): + SK = 2 + p = torch.randn(SK, 2 * I, device=DEV, dtype=torch.float32) + o = torch.empty(I, device=DEV, dtype=torch.bfloat16) + K.silu_mul(I, SK)(p, o) + g = p.sum(0)[:I].to(torch.bfloat16) + u = p.sum(0)[I:].to(torch.bfloat16) + report("silu_mul", o, torch.nn.functional.silu(g.float()).to(torch.bfloat16) * u) + + +def t_embed(): + tbl = torch.randn(1024, H, device=DEV, dtype=torch.bfloat16) + ids = torch.tensor([517], device=DEV, dtype=torch.int64) + o = torch.empty(H, device=DEV, dtype=torch.bfloat16) + K.embed(1024, H)(tbl, ids, o) + report("embed", o, tbl[517]) + + +def t_sample_step(): + NB, NSTEPS, PL = 1187, 64, 5 + bv = torch.randn(NB, device=DEV, dtype=torch.float32) + bi = torch.randint(0, V, (NB,), device=DEV, dtype=torch.int32) + inp = torch.arange(NSTEPS, device=DEV, dtype=torch.int32) + 100 + pl = torch.tensor([PL], device=DEV, dtype=torch.int32) + sam = torch.zeros(NSTEPS, device=DEV, dtype=torch.int32) + kern = K.sample_step(NB, NSTEPS) + for pos, expect_prompt in [(2, True), (PL - 1, False), (20, False)]: + ids = torch.zeros(1, device=DEV, dtype=torch.int64) + pt = torch.tensor([pos], device=DEV, dtype=torch.int32) + kern(bv, bi, inp, pl, ids, pt, sam) + best = int(bi[int(bv.argmax())]) + want_next = int(inp[pos + 1]) if expect_prompt else best + ok = int(sam[pos]) == best and int(ids[0]) == want_next and int(pt[0]) == pos + 1 + print(f"{'ok ' if ok else 'FAIL'} sample_step pos={pos:<3d} " + f"sam={int(sam[pos])} want={best} next={int(ids[0])} want={want_next}") + if not ok: + FAIL.append(f"sample_step:{pos}") + + +if __name__ == "__main__": + torch.manual_seed(0) + for fn in [t_rms_norm, t_resid_rms_norm, t_gemv, t_lm_head, + t_qk_rope_cache, t_attn, t_silu_mul, t_embed, t_sample_step]: + try: + fn() + except Exception as exc: + import traceback + print(f"FAIL {fn.__name__}: {type(exc).__name__}: {exc}") + traceback.print_exc(limit=6) + FAIL.append(fn.__name__) + print("\n" + ("ALL PASS" if not FAIL else f"FAILURES: {FAIL}")) + sys.exit(1 if FAIL else 0) diff --git a/examples/qwen3_1_7b-tilelang/fast/twin.py b/examples/qwen3_1_7b-tilelang/fast/twin.py new file mode 100644 index 00000000..4c316522 --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/fast/twin.py @@ -0,0 +1,236 @@ +"""Runtime twins of the shipped modules, so `tilefoundry check` can judge them. + +The engine next door is the thing that runs fast; it deliberately does not have +the reference's shape. Its cache has fixed capacity and its step reads the +position from a device tensor, because a graph records addresses -- and that is a +different contract from the authored one, which hands each step's key and value +back for the caller to append. + +So this file exists to be *comparable*. It wraps the same TileLang kernels in the +reference's own signatures -- prior cache in, this step's entry out, ``ctx_len`` +as a range -- and nothing else. Then + + tilefoundry check twin.py:LayerTwin.mlp --inputs random --out output --fn ... + +runs the evaluator over the authored `@func` and this implementation over the same +draw, and reports whether they agree. The per-function comparison the optimize +page asks for is that command, not a test written here. + +Two deliberate departures from the authored source, both toward the published +model rather than away from it, will show up in the numbers and are worth naming +before the report does: + +* **The scale is applied after the dot, in f32.** The reference multiplies + ``1/sqrt(head_dim)`` onto ``q`` in bf16 first, which rounds every entry a second + time; the exponential downstream turns that into percent-level error on the + attention weights. Hugging Face scales after. See `kernels.qk_rope_cache`. +* **Attention probabilities are bf16 into the V product**, which is what HF's own + attention does, where the reference stays in wider precision through it. + +Measured against Hugging Face on the real checkpoint, the engine built from these +kernels agrees with published greedy decoding on 255 of 256 teacher-forced +positions; the one exception is a position where HF's own bf16 logits are exactly +tied and it picks by index order. +""" +import sys +from functools import lru_cache +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from tilefoundry.runtime.decorator import runtime_func, runtime_module # noqa: E402 + +import kernels as K # noqa: E402 +from engine import _load_module # noqa: E402 + +_REF_DIR = Path(__file__).resolve().parent.parent / "ref_src" +ref = _load_module(_REF_DIR / "model.py", "twin_ref_model") +cfg = ref.config + +H = cfg.hidden_size +HQ = cfg.num_attention_heads +HKV = cfg.num_key_value_heads +D = cfg.head_dim +I = cfg.intermediate_size +V = cfg.vocab_size +EPS = cfg.rms_norm_eps +QKV_N = HQ * D + 2 * HKV * D +SPLIT = 256 + +#: Tiles here are chosen to divide every dimension cleanly, not for speed: this +#: twin answers "is it the same computation", and the engine answers "how fast". +T_QKV = (256, 64, 8, 256) +T_O = (128, 128, 8, 128) +T_GU = (256, 64, 2, 256) +T_DOWN = (128, 128, 8, 128) +T_HEAD = (128, 128, 128) + + +def _cap(n: int) -> int: + return max(SPLIT, ((n + SPLIT - 1) // SPLIT) * SPLIT) + + +@lru_cache(maxsize=None) +def _fuse_key(*keys): + return None + + +class _Fused: + """Concatenated projections, remembered per weight tensor identity. + + The reference declares ``w_q``/``w_k``/``w_v`` separately and a twin is handed + them separately, but one GEMV over ``[q|k|v]`` is the kernel that exists. The + join is a view of the same numbers, so caching it by identity keeps a repeated + check from redoing it. + """ + + def __init__(self): + self._cache = {} + + def __call__(self, *tensors): + key = tuple(t.data_ptr() for t in tensors) + got = self._cache.get(key) + if got is None: + got = torch.cat([t.reshape(t.shape[-2], t.shape[-1]) for t in tensors], + dim=1).contiguous() + self._cache[key] = got + return got + + +_fused = _Fused() + + +def _rms_norm(x, gamma): + """`input_rms_norm` over a (1, 1, H) activation.""" + out = torch.empty(H, device=x.device, dtype=x.dtype) + K.rms_norm(H, EPS)(x.reshape(H), gamma, out) + return out.reshape(1, 1, H) + + +@runtime_module(ref.Qwen3_1_7B_DecoderLayer) +class LayerTwin: + """One decoder layer, in the reference's signatures, on TileLang kernels.""" + + @runtime_func + def input_rms_norm(self, hidden, gamma_in): + return _rms_norm(hidden, gamma_in) + + @runtime_func + def self_attention(self, hidden, gamma_in, w_q, w_k, w_v, gamma_q, gamma_k, + cos_cache, sin_cache, pos_ids, k_cache, v_cache, scale, w_o): + """Fused norm + QKV + q/k norm + rope + GQA attention + o_proj. + + The reference reads a prior cache of ``ctx_len`` positions and returns + this step's own entry. The kernels want one buffer holding both, so a + scratch cache of the next whole split is filled with the prior context, + this step's entry is written at slot ``ctx_len``, attention runs over + ``ctx_len + 1`` positions, and the entry is read back out to return. + """ + dev = hidden.device + ctx = int(k_cache.shape[1]) + cap = _cap(ctx + 1) + mp_rows = int(cos_cache.shape[0]) + sk = T_QKV[2] + + xn = _rms_norm(hidden, gamma_in).reshape(H) + part = torch.empty(sk, QKV_N, device=dev, dtype=torch.float32) + K.gemv(H, QKV_N, *T_QKV)(xn, _fused(w_q, w_k, w_v), part) + + # zeroed: a slot past the end must contribute exp(-inf) * 0, not * junk + kc = torch.zeros(cap, HKV * D, device=dev, dtype=hidden.dtype) + vc = torch.zeros(cap, HKV * D, device=dev, dtype=hidden.dtype) + if ctx: + kc[:ctx] = k_cache[0].reshape(ctx, HKV * D) + vc[:ctx] = v_cache[0].reshape(ctx, HKV * D) + q = torch.empty(HQ * D, device=dev, dtype=hidden.dtype) + write = torch.full((1,), ctx, device=dev, dtype=torch.int32) + K.qk_rope_cache(HQ, HKV, D, mp_rows, cap, sk, EPS)( + part, gamma_q, gamma_k, cos_cache, sin_cache, + pos_ids.to(torch.int32), write, kc, vc, q, + ) + + ns = cap // SPLIT + op = torch.empty(ns, HQ, D, device=dev, dtype=torch.float32) + mpart = torch.empty(ns, HQ, device=dev, dtype=torch.float32) + lpart = torch.empty(ns, HQ, device=dev, dtype=torch.float32) + K.attn_partial(HQ, HKV, D, cap, SPLIT, float(scale.reshape(-1)[0]))( + q, kc, vc, write, op, mpart, lpart + ) + opart = torch.empty(T_O[2], H, device=dev, dtype=torch.float32) + K.gemv_attn_combine(HQ, D, H, *T_O[:3], ns, T_O[3])( + op, mpart, lpart, w_o.reshape(HQ * D, H), opart + ) + out = opart.sum(0).to(hidden.dtype).reshape(1, 1, H) + k_new = kc[ctx].reshape(1, 1, HKV, D).clone() + v_new = vc[ctx].reshape(1, 1, HKV, D).clone() + return out, k_new, v_new + + @runtime_func + def mlp(self, hidden, gamma_post, w_gate, w_up, w_down): + """Fused post-attention norm + dense SwiGLU.""" + dev = hidden.device + xn = _rms_norm(hidden, gamma_post).reshape(H) + gu = torch.empty(T_GU[2], 2 * I, device=dev, dtype=torch.float32) + K.gemv(H, 2 * I, *T_GU)(xn, _fused(w_gate, w_up), gu) + dpart = torch.empty(T_DOWN[2], H, device=dev, dtype=torch.float32) + K.gemv_silu(I, H, *T_DOWN[:3], T_GU[2], T_DOWN[3])( + gu, w_down.reshape(I, H), dpart + ) + return dpart.sum(0).to(hidden.dtype).reshape(1, 1, H) + + @runtime_func + def decoder_layer(self, hidden, gamma_in, w_q, w_k, w_v, gamma_q, gamma_k, + cos_cache, sin_cache, pos_ids, k_cache, v_cache, scale, w_o, + gamma_post, w_gate, w_up, w_down): + """attention + residual, then mlp + residual -- `Qwen3DecoderLayer.forward`. + + The authored `@func` passes every weight down because inside HIR nothing is + bound. On the runtime side they already are: `self.` takes activations + alone and fills its own `ConstTensor` params by name from this loading. So + this body receives the weights (a kernel body's signature includes them) + and does not forward them. + """ + attn_out, k_new, v_new = self.self_attention( + hidden, cos_cache, sin_cache, pos_ids, k_cache, v_cache, scale + ) + h1 = hidden + attn_out + return h1 + self.mlp(h1), k_new, v_new + + +_ROOT_FUNCS = {} + + +@runtime_func +def embed(self, w_embed, token_ids): + out = torch.empty(H, device=w_embed.device, dtype=w_embed.dtype) + K.embed(V, H)(w_embed, token_ids.reshape(1).to(torch.int64), out) + return out.reshape(1, 1, H) + + +@runtime_func +def final_rms_norm(self, hidden, gamma_final): + return _rms_norm(hidden, gamma_final) + + +@runtime_func +def lm_head(self, hidden, w_head): + dev = hidden.device + kern, nb = K.lm_head(H, V, *T_HEAD) + logits = torch.empty(V, device=dev, dtype=torch.float32) + bv = torch.empty(nb, device=dev, dtype=torch.float32) + bi = torch.empty(nb, device=dev, dtype=torch.int32) + kern(hidden.reshape(H), w_head, logits, bv, bi) + return logits.to(hidden.dtype).reshape(1, V) + + +#: The root's children are named `layer0`..`layer27` and the twin's child-attribute +#: name set must equal that exactly, so the class body is built rather than typed. +RootTwin = runtime_module(ref.Qwen3_1_7B)(type("RootTwin", (), { + "__doc__": "The layer stack and the step around it, as a runtime twin.", + "embed": embed, + "final_rms_norm": final_rms_norm, + "lm_head": lm_head, + **{f"layer{i}": LayerTwin for i in range(cfg.num_hidden_layers)}, +})) diff --git a/examples/qwen3_1_7b-tilelang/ref_src/config.json b/examples/qwen3_1_7b-tilelang/ref_src/config.json new file mode 100644 index 00000000..044a86ec --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/ref_src/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "Qwen3ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 6144, + "max_position_embeddings": 40960, + "max_window_layers": 28, + "model_type": "qwen3", + "num_attention_heads": 16, + "num_hidden_layers": 28, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.0", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} \ No newline at end of file diff --git a/examples/qwen3_1_7b-tilelang/ref_src/generation.py b/examples/qwen3_1_7b-tilelang/ref_src/generation.py new file mode 100644 index 00000000..a32dff6f --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/ref_src/generation.py @@ -0,0 +1,88 @@ +"""Autoregressive decode: one token per step; the caller owns the state. + +The orchestrator drives a causal-LM Module's orchestration methods. The model +declares a step's activations, while this loop owns sampling and timing. +""" +from __future__ import annotations + +from dataclasses import dataclass +from time import perf_counter + +import torch + + +@dataclass(frozen=True) +class Decoded: + """The text and decode-only timing from one autoregressive continuation.""" + + text: str + tokens: int + seconds: float + prompt_steps: int + + +def greedy(logits) -> int: + """Choose the vocabulary entry with the largest logit.""" + return int(torch.argmax(logits).item()) + + +def _sync(device) -> None: + backend = getattr(torch, torch.device(device).type, None) + synchronize = getattr(backend, "synchronize", None) + if synchronize is not None: + synchronize(device) + + +def _prompt_ids(tokenizer, prompt: str) -> torch.Tensor: + encoded = tokenizer.encode(prompt) + return torch.tensor(getattr(encoded, "ids", encoded), dtype=torch.int64) + + +def decode(loaded, tokenizer, prompt: str, *, max_new: int, sampler=greedy, eos=(), device=None): + """Decode *prompt* through *loaded* and return its sampled continuation. + + *loaded* may be a ``LoadedModule`` or a runtime twin. Its model supplies the + four orchestration methods this loop drives, including each step's inputs. + """ + device = torch.accelerator.current_accelerator() if device is None else device + prompt_ids = _prompt_ids(tokenizer, prompt).to(device) + if prompt_ids.numel() == 0: + raise ValueError("decode needs a prompt that encodes to at least one token") + prompt_steps = prompt_ids.numel() + input_ids = torch.empty(prompt_steps + max_new, dtype=torch.int64, device=device) + input_ids[:prompt_steps] = prompt_ids + caches = loaded.init_caches(device=device) + + for step in range(prompt_steps): + args = loaded.prepare_inputs_for_generation( + input_ids[: step + 1], step, caches, device=device + ) + logits, fresh = loaded.forward(*args) + caches = loaded.append_cache(caches, fresh) + + sampler(logits) + _sync(device) + started = perf_counter() + output: list[int] = [] + for step in range(max_new): + token = int(sampler(logits)) + if token in eos: + break + output.append(token) + input_ids[prompt_steps + step] = token + args = loaded.prepare_inputs_for_generation( + input_ids[: prompt_steps + step + 1], prompt_steps + step, caches, device=device + ) + logits, fresh = loaded.forward(*args) + caches = loaded.append_cache(caches, fresh) + _sync(device) + elapsed = perf_counter() - started + return Decoded( + text=tokenizer.decode(output), + tokens=len(output), + seconds=elapsed, + prompt_steps=prompt_steps, + ) + + +__all__ = ["Decoded", "decode", "greedy"] diff --git a/examples/qwen3_1_7b-tilelang/ref_src/hf_alias.py b/examples/qwen3_1_7b-tilelang/ref_src/hf_alias.py new file mode 100644 index 00000000..409c277f --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/ref_src/hf_alias.py @@ -0,0 +1,55 @@ +"""The published checkpoint as this model's weights: which key, and how it is stored. + +An entry exists where the published checkpoint and the Module disagree on a +name, stored form, or both. The model declarations remain the contract that the +loaded values must satisfy. +""" +from __future__ import annotations + +import torch +from transformers import Qwen3Config + +from tilefoundry.runtime import Preprocessed + + +def _projection(t: torch.Tensor) -> torch.Tensor: + """HF ``nn.Linear.weight`` ``(out, in)`` -> kernel ``(1, in, out)``.""" + return t.t().unsqueeze(0).contiguous() + + +def _transposed(t: torch.Tensor) -> torch.Tensor: + """HF ``(out, in)`` -> an unbatched kernel ``(in, out)`` weight.""" + return t.t().contiguous() + + +def hf_alias(config: Qwen3Config) -> dict[str, object]: + """Canonical names -> published Qwen3 checkpoint names for *config*.""" + return { + "w_embed": "model.embed_tokens.weight", + "gamma_final": "model.norm.weight", + "w_head": Preprocessed("lm_head.weight", _transposed), + "gamma_in": "input_layernorm.weight", + "gamma_post": "post_attention_layernorm.weight", + "gamma_q": "self_attn.q_norm.weight", + "gamma_k": "self_attn.k_norm.weight", + "w_q": Preprocessed("self_attn.q_proj.weight", _projection), + "w_k": Preprocessed("self_attn.k_proj.weight", _projection), + "w_v": Preprocessed("self_attn.v_proj.weight", _projection), + "w_o": Preprocessed("self_attn.o_proj.weight", _projection), + "w_gate": Preprocessed("mlp.gate_proj.weight", _projection), + "w_up": Preprocessed("mlp.up_proj.weight", _projection), + "w_down": Preprocessed("mlp.down_proj.weight", _projection), + **{f"layer{i}": f"model.layers.{i}" for i in range(config.num_hidden_layers)}, + } + + +def hf_layout_only(config: Qwen3Config) -> dict[str, Preprocessed]: + """The table's layout entries, made relative to an in-memory test mapping.""" + return { + name: Preprocessed(name, value.read) + for name, value in hf_alias(config).items() + if isinstance(value, Preprocessed) + } + + +__all__ = ["hf_alias"] diff --git a/examples/qwen3_1_7b-tilelang/ref_src/model.py b/examples/qwen3_1_7b-tilelang/ref_src/model.py new file mode 100644 index 00000000..18e7d425 --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/ref_src/model.py @@ -0,0 +1,388 @@ +"""Qwen3-1.7B's dense decoder layer and the stack that closes it, as IR Modules. +Companion to ``tests/models/qwen3_5_35b_a3b/model.py``: same +``@module class`` authoring style (each kernel is a named ``@func`` method; the +decorator returns the ``tilefoundry.ir.core.module.Module`` that the class name +binds directly to -- ``Qwen3_1_7B.lookup("self_attention")`` resolves one kernel +to its IR node). What differs from the MoE-30B sibling is the MLP: a single +dense SwiGLU expert (plain gate/up/down projection), with none of the 30B's +runtime top-k expert routing -- no router, no ``topk``, no ``gather``. + +Decode, one token per step. The step's own token count is the literal 1, so the +only dimension carried as a range is the context the step reads: ``ctx_len``, +the length of the KV cache handed in. Everything a caller has to know how to +compute is that one number. + +The cache is explicit tensors in and out, and the two directions are not the +same tensor. What comes in is the context *before* this token -- ``ctx_len`` +positions, read-only. What goes out is this token's own key and value, one +position each. Appending the second to the first is the caller's step, not the +kernel's, and that is what keeps every shape here expressed in ``ctx_len`` +alone: a kernel returning the grown cache would have an axis of ``ctx_len + 1``, +and a sum of a range and a constant cannot feed the matmul that would consume it +(the constraint that makes a step return its own entry rather than the grown +cache). + +That split is also why attention here is an online softmax rather than one +``softmax`` over a concatenated score row. The new token has to attend to itself +as well as to the cache, and the two score groups live in differently shaped +tensors; each is reduced to its own ``(max, sum, weighted values)`` partial and +the partials are merged by the same log-sum-exp rescale +``tests/fixtures/gqa_online.py``'s +combine kernel uses. No mask is needed: a single query at the end of the +context may attend every position there is. + +``self_attention`` and ``mlp`` each fuse their preceding RMSNorm internally +(``input_rms_norm`` / the post-attention norm) -- matching the Qwen3-30B-A3B +sibling's convention (its ``self_attention`` fuses ``input_rms_norm``; its +``moe`` fuses the post-attention norm) so each fused kernel lines up with one +HF pre-norm-then-block composition. ``decoder_layer`` composes +``self_attention`` + residual + ``mlp`` + residual, mirroring +``Qwen3DecoderLayer.forward`` exactly. + +Every RMSNorm here is written out rather than calling ``tf.rms_norm``: +``Qwen3RMSNorm`` rounds the normalised activation to the input dtype *before* the +learned scale multiplies it, where the generic op matches ``torch.nn.RMSNorm`` +and stays in f32 through that multiply. + +``q_norm`` / ``k_norm`` are the same normalisation over the ``head_dim`` axis +alone, applied per head, so they run directly on the ``[1, 1, heads, head_dim]`` +tensor the head split already produces. +""" +from __future__ import annotations + +import json +from functools import lru_cache +from pathlib import Path + +from transformers import Qwen3Config + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf # noqa: F401 โ€” tf used by @func bodies +from tilefoundry.dsl.tf import * # noqa: F401, F403 โ€” bare op bindings for @func bodies +from tilefoundry.ir.types.dim import DimVar +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + + +def published(path: Path | None = None) -> Qwen3Config: + """The checkpoint's own configuration, read by the class Hugging Face uses.""" + path = Path(__file__).parent / "config.json" if path is None else path + return Qwen3Config(**json.loads(path.read_text(encoding="utf-8"))) + + +config = published() + +# The published dtype as the DSL spells it. The checkpoint stores its weights at +# this precision, so it is what a kernel reading them consumes. +_DT = {"bfloat16": "bf16", "float16": "f16", "float32": "f32"}[ + str(config.dtype).removeprefix("torch.") +] +_Q_PROJ = config.num_attention_heads * config.head_dim +_KV_PROJ = config.num_key_value_heads * config.head_dim +_GQA = config.num_attention_heads // config.num_key_value_heads + +#: The variance floor every one of this model's norms adds, read from the +#: checkpoint rather than written down: the norms below are spelled out, so the +#: epsilon `tf.rms_norm` would have carried as an attribute is carried here. +_EPS = config.rms_norm_eps + +# The prior cache this step reads: the only range this model carries. Zero is a +# first step, and the exclusive upper bound is `config.max_position_embeddings` +# because a position beyond +# the rotary cache has no embedding to gather. +C = DimVar("ctx_len", 0, config.max_position_embeddings) + +# One token per step. +S = 1 + +_G = _GQA + + +def _generation_device(device): + import torch # noqa: PLC0415 + + return torch.accelerator.current_accelerator() if device is None else device + + +@lru_cache(maxsize=None) +def _generation_rope(device): + import torch # noqa: PLC0415 + + dim = config.head_dim + inverse = 1.0 / ( + config.rope_parameters["rope_theta"] + ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim) + ) + phases = torch.outer( + torch.arange(config.max_position_embeddings, device=device, dtype=torch.float32), inverse + ) + phases = torch.cat((phases, phases), dim=-1) + return phases.cos().to(config.dtype), phases.sin().to(config.dtype) + + +@module(entry="decoder_layer") +class Qwen3_1_7B_DecoderLayer: + @func + def input_rms_norm( + hidden: Tensor[(1, S, config.hidden_size), _DT], + gamma_in: ConstTensor[(config.hidden_size,), _DT], + ) -> Tensor[(1, S, config.hidden_size), _DT]: + # HF `Qwen3DecoderLayer.input_layernorm`. Written out because + # `Qwen3RMSNorm` ends `self.weight * hidden_states.to(input_dtype)`: it + # rounds to the input dtype before the scale multiplies, where + # `tf.rms_norm` keeps f32 through that multiply. + x32 = tf.cast(hidden, dtype="f32") + variance = tf.reduce(x32 * x32, axes=(-1,), keepdim=True, kind="mean") + normed = tf.cast(x32 * tf.rsqrt(variance + _EPS), dtype=_DT) + return normed * gamma_in + + @func + def self_attention( + hidden: Tensor[(1, S, config.hidden_size), _DT], + gamma_in: ConstTensor[(config.hidden_size,), _DT], + w_q: ConstTensor[(1, config.hidden_size, _Q_PROJ), _DT], + w_k: ConstTensor[(1, config.hidden_size, _KV_PROJ), _DT], + w_v: ConstTensor[(1, config.hidden_size, _KV_PROJ), _DT], + gamma_q: ConstTensor[(config.head_dim,), _DT], + gamma_k: ConstTensor[(config.head_dim,), _DT], + cos_cache: Tensor[(config.max_position_embeddings, config.head_dim), _DT], + sin_cache: Tensor[(config.max_position_embeddings, config.head_dim), _DT], + pos_ids: Tensor[(S,), "i32"], + k_cache: Tensor[(1, C, config.num_key_value_heads, config.head_dim), _DT], + v_cache: Tensor[(1, C, config.num_key_value_heads, config.head_dim), _DT], + scale: Tensor[(1, 1, 1, 1), _DT], + w_o: ConstTensor[(1, _Q_PROJ, config.hidden_size), _DT], + ): + # Fused input_layernorm + self_attn, no residual (the layer owns the + # residual add). Returns the attention output together with this token's + # key and value, which are what the caller appends to the cache. + hidden_norm = input_rms_norm(hidden, gamma_in) + q = tf.reshape( + tf.matmul(hidden_norm, w_q), + new_shape=(1, S, config.num_attention_heads, config.head_dim), + ) + k = tf.reshape( + tf.matmul(hidden_norm, w_k), + new_shape=(1, S, config.num_key_value_heads, config.head_dim), + ) + v = tf.reshape( + tf.matmul(hidden_norm, w_v), + new_shape=(1, S, config.num_key_value_heads, config.head_dim), + ) + q_n32 = tf.cast(q, dtype="f32") + q_n_var = tf.reduce(q_n32 * q_n32, axes=(-1,), keepdim=True, kind="mean") + q_n = tf.cast(q_n32 * tf.rsqrt(q_n_var + _EPS), dtype=_DT) * gamma_q + q_rope, _ = tf.rope(q_n, q_n, cos_cache, sin_cache, pos_ids) + k_n32 = tf.cast(k, dtype="f32") + k_n_var = tf.reduce(k_n32 * k_n32, axes=(-1,), keepdim=True, kind="mean") + k_n = tf.cast(k_n32 * tf.rsqrt(k_n_var + _EPS), dtype=_DT) * gamma_k + _, k_rope = tf.rope(k_n, k_n, cos_cache, sin_cache, pos_ids) + + # Every query head sees its group's key/value head, for the cache and + # for the new token alike. + q_s = tf.reshape(q_rope, new_shape=(1, S, config.num_attention_heads, config.head_dim)) * scale + k_ctx = tf.reshape( + tf.transpose(tf.repeat_interleave(k_cache, repeats=_G, axis=2), perm=(0, 2, 1, 3)), + new_shape=(1, 1, config.num_attention_heads, C, config.head_dim), + ) + v_ctx = tf.reshape( + tf.transpose(tf.repeat_interleave(v_cache, repeats=_G, axis=2), perm=(0, 2, 1, 3)), + new_shape=(1, 1, config.num_attention_heads, C, config.head_dim), + ) + k_new = tf.repeat_interleave(k_rope, repeats=_G, axis=2) + v_new = tf.repeat_interleave(v, repeats=_G, axis=2) + + # Two score groups: one over the cache, one over the token itself. + q_e = tf.reshape(q_s, new_shape=(1, S, config.num_attention_heads, 1, config.head_dim)) + score_ctx = tf.reduce(q_e * k_ctx, axes=(-1,), keepdim=True, kind="sum") + score_new = tf.reduce(q_s * k_new, axes=(-1,), keepdim=True, kind="sum") + + # Log-sum-exp merge of the two groups' partials against their joint max. + peak = tf.max( + tf.reduce(score_ctx, axes=(-2,), keepdim=False, kind="max"), score_new + ) + peak_e = tf.reshape(peak, new_shape=(1, S, config.num_attention_heads, 1, 1)) + p_ctx = tf.exp(score_ctx - peak_e) + p_new = tf.exp(score_new - peak) + total = tf.reduce(p_ctx, axes=(-2,), keepdim=False, kind="sum") + p_new + weighted = ( + tf.reduce(p_ctx * v_ctx, axes=(-2,), keepdim=False, kind="sum") + + p_new * v_new + ) + attn = weighted / total + out = tf.matmul( + tf.reshape(attn, new_shape=(1, S, _Q_PROJ)), w_o + ) + return out, k_rope, v + + @func + def mlp( + hidden: Tensor[(1, S, config.hidden_size), _DT], + gamma_post: ConstTensor[(config.hidden_size,), _DT], + w_gate: ConstTensor[(1, config.hidden_size, config.intermediate_size), _DT], + w_up: ConstTensor[(1, config.hidden_size, config.intermediate_size), _DT], + w_down: ConstTensor[(1, config.intermediate_size, config.hidden_size), _DT], + ) -> Tensor[(1, S, config.hidden_size), _DT]: + # Fused post_attention_layernorm + dense SwiGLU, no residual. + hidden_norm32 = tf.cast(hidden, dtype="f32") + hidden_norm_var = tf.reduce(hidden_norm32 * hidden_norm32, axes=(-1,), keepdim=True, kind="mean") + hidden_norm = tf.cast(hidden_norm32 * tf.rsqrt(hidden_norm_var + _EPS), dtype=_DT) * gamma_post + gate = tf.matmul(hidden_norm, w_gate) + up = tf.matmul(hidden_norm, w_up) + act = tf.silu(gate) + h = act * up + return tf.matmul(h, w_down) + + @func + def decoder_layer( + hidden: Tensor[(1, S, config.hidden_size), _DT], + gamma_in: ConstTensor[(config.hidden_size,), _DT], + w_q: ConstTensor[(1, config.hidden_size, _Q_PROJ), _DT], + w_k: ConstTensor[(1, config.hidden_size, _KV_PROJ), _DT], + w_v: ConstTensor[(1, config.hidden_size, _KV_PROJ), _DT], + gamma_q: ConstTensor[(config.head_dim,), _DT], + gamma_k: ConstTensor[(config.head_dim,), _DT], + cos_cache: Tensor[(config.max_position_embeddings, config.head_dim), _DT], + sin_cache: Tensor[(config.max_position_embeddings, config.head_dim), _DT], + pos_ids: Tensor[(S,), "i32"], + k_cache: Tensor[(1, C, config.num_key_value_heads, config.head_dim), _DT], + v_cache: Tensor[(1, C, config.num_key_value_heads, config.head_dim), _DT], + scale: Tensor[(1, 1, 1, 1), _DT], + w_o: ConstTensor[(1, _Q_PROJ, config.hidden_size), _DT], + gamma_post: ConstTensor[(config.hidden_size,), _DT], + w_gate: ConstTensor[(1, config.hidden_size, config.intermediate_size), _DT], + w_up: ConstTensor[(1, config.hidden_size, config.intermediate_size), _DT], + w_down: ConstTensor[(1, config.intermediate_size, config.hidden_size), _DT], + ): + # One decode step: self_attention + residual, then mlp + residual -- + # mirrors `Qwen3DecoderLayer.forward` exactly -- plus this token's key + # and value passed straight through for the caller to append. + attn_out, k_new, v_new = self_attention( + hidden, gamma_in, w_q, w_k, w_v, gamma_q, gamma_k, + cos_cache, sin_cache, pos_ids, k_cache, v_cache, scale, w_o, + ) + h1 = hidden + attn_out + mlp_out = mlp(h1, gamma_post, w_gate, w_up, w_down) + return h1 + mlp_out, k_new, v_new + + +@module(target=CudaTarget("nvidia.h200_sxm")) +class Qwen3_1_7B: + """The ordered layer stack plus the norm that closes it.""" + + topologies = (Topology("cta", 132), Topology("thread", 512)) + + layers = tuple( + Qwen3_1_7B_DecoderLayer.renamed(f"layer{index}") + for index in range(config.num_hidden_layers) + ) + + @func + def embed( + w_embed: ConstTensor[(config.vocab_size, config.hidden_size), _DT], + token_ids: Tensor[(S,), "i64"], + ) -> Tensor[(1, S, config.hidden_size), _DT]: + # HF `Qwen3Model.embed_tokens`. + return tf.reshape( + tf.gather(w_embed, token_ids, axis=0), new_shape=(1, S, config.hidden_size) + ) + + @func + def final_rms_norm( + hidden: Tensor[(1, S, config.hidden_size), _DT], + gamma_final: ConstTensor[(config.hidden_size,), _DT], + ) -> Tensor[(1, S, config.hidden_size), _DT]: + # HF `Qwen3Model.norm`, applied once after the last layer. + out32 = tf.cast(hidden, dtype="f32") + out_var = tf.reduce(out32 * out32, axes=(-1,), keepdim=True, kind="mean") + out = tf.cast(out32 * tf.rsqrt(out_var + _EPS), dtype=_DT) * gamma_final + return out + + @func + def lm_head( + hidden: Tensor[(1, S, config.hidden_size), _DT], + w_head: ConstTensor[(config.hidden_size, config.vocab_size), _DT], + ) -> Tensor[(1, config.vocab_size), _DT]: + return tf.matmul(tf.reshape(hidden, new_shape=(1, config.hidden_size)), w_head) + + def forward(self, token_ids, cos_cache, sin_cache, pos_ids, scale, caches): + """The whole decode step: this token's row, every layer over it, its logits. + + What comes back is the logits and each layer's own fresh entry; growing the + cache with them is the caller's step, through `append_cache`. + """ + hidden = self.embed(token_ids) + normed, entries = self.decode_hidden( + hidden, cos_cache, sin_cache, pos_ids, scale, caches + ) + return self.lm_head(normed), entries + + def decode_hidden(self, hidden, cos_cache, sin_cache, pos_ids, scale, caches): + """One decode step through every layer, then the final norm. + + *caches* is one layer's context per layer, in layer order. What comes back + is the normalised hidden state and each layer's own cache entry, for the + caller to append -- the same division the single layer makes. + """ + if len(caches) != len(self.modules): + raise ValueError( + f"decoder has {len(self.modules)} layers but was given " + f"{len(caches)} caches" + ) + entries = [] + for layer, (k_cache, v_cache) in zip(self.modules, caches): + hidden, k_new, v_new = layer( + hidden, cos_cache, sin_cache, pos_ids, k_cache, v_cache, scale + ) + entries.append((k_new, v_new)) + return self.final_rms_norm(hidden), tuple(entries) + + def append_cache(self, caches, fresh): + """The cache the next step reads: each layer's context with this step's own + key and value written after it. + + A step hands back its own entry rather than the grown cache, so appending is + the caller's, and the caller of a step is this root -- stated here once so a + caller has none of its own. + """ + import torch # noqa: PLC0415 + + return tuple( + (torch.cat([k_cache, k_new], dim=1), torch.cat([v_cache, v_new], dim=1)) + for (k_cache, v_cache), (k_new, v_new) in zip(caches, fresh) + ) + + def init_caches(self, device=None): + """The per-layer cache container, zero positions long. + + `ctx_len` admits 0, so these are a decode start: the first step of a + sequence attends the one position it brings itself. + """ + import torch # noqa: PLC0415 + + from tilefoundry.evaluator.value import to_torch_dtype # noqa: PLC0415 + from tilefoundry.ir.types import DType # noqa: PLC0415 + + device = _generation_device(device) + empty = (1, 0, config.num_key_value_heads, config.head_dim) + dtype = to_torch_dtype(DType.from_name(_DT)) + return tuple( + ( + torch.zeros(empty, device=device, dtype=dtype), + torch.zeros(empty, device=device, dtype=dtype), + ) + for _ in range(config.num_hidden_layers) + ) + + def prepare_inputs_for_generation(self, input_ids, step, caches, device=None): + """The token and positional activations for one decode step.""" + import torch # noqa: PLC0415 + + device = _generation_device(device) + token_ids = input_ids[step].reshape(1).to(device=device, dtype=torch.int64) + cos_cache, sin_cache = _generation_rope(device) + pos_ids = torch.tensor([step], device=device, dtype=torch.int32) + scale = torch.full( + (1, 1, 1, 1), config.head_dim ** -0.5, device=device, dtype=config.dtype + ) + return token_ids, cos_cache, sin_cache, pos_ids, scale, caches diff --git a/examples/qwen3_1_7b-tilelang/ref_src/run.py b/examples/qwen3_1_7b-tilelang/ref_src/run.py new file mode 100644 index 00000000..45ec8346 --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/ref_src/run.py @@ -0,0 +1,60 @@ +"""Run a shipped causal-LM source directory against its published checkpoint.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from transformers import AutoTokenizer + +from tilefoundry.ir.core.module import Module +from tilefoundry.runtime import SafetensorsResource + + +def _root(namespace: dict[str, object]) -> Module: + roots = [ + value + for value in namespace.values() + if isinstance(value, Module) and value.target is not None + ] + if len(roots) != 1: + raise SystemExit("model.py must declare exactly one root Module with a target") + return roots[0] + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--ckpt", required=True, type=Path, help="published checkpoint directory") + parser.add_argument("--prompt", required=True, help="prompt to continue") + parser.add_argument("--max-new-tokens", required=True, type=int, metavar="TOKENS") + parser.add_argument( + "--device", + help="runtime device (default: Torch's current accelerator)", + ) + parser.epilog = "Context is limited by the kernel declarations in model.py." + args = parser.parse_args(argv) + + import model # noqa: PLC0415 + from generation import decode # noqa: PLC0415 + from hf_alias import hf_alias # noqa: PLC0415 + + device = str(torch.accelerator.current_accelerator()) if args.device is None else args.device + loaded = _root(vars(model)).load( + SafetensorsResource(str(args.ckpt), device=device, alias=hf_alias(model.config)) + ) + tokenizer = AutoTokenizer.from_pretrained(args.ckpt) + decoded = decode( + loaded, + tokenizer, + args.prompt, + max_new=args.max_new_tokens, + device=device, + ) + print(decoded.text) + print(f"{decoded.tokens / decoded.seconds:.1f} tok/s") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/qwen3_1_7b-tilelang/run.py b/examples/qwen3_1_7b-tilelang/run.py new file mode 100644 index 00000000..4dc665ce --- /dev/null +++ b/examples/qwen3_1_7b-tilelang/run.py @@ -0,0 +1,101 @@ +"""Continue a prompt with Qwen3-1.7B on TileFoundry, and report the rate. + + python run.py --prompt "..." --max-new-tokens 2048 + +The model is the shipped `qwen3_1_7b` authored HIR (kept verbatim in `ref_src/`, +which is what everything here is measured against); the implementation that runs +is `fast/`, a TileLang runtime twin whose whole decode step is one captured CUDA +graph. See `fast/kernels.py` for the kernels and `fast/engine.py` for why the +step can be captured at all. + +The rate covers exactly the steps that produce the continuation. Walking the +prompt is reported separately rather than folded in, because averaging a short +prefill into a long generation flatters the number. +""" +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path +from time import perf_counter + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE / "fast")) + +NEED_GIB = 12.0 + + +def pick_device(want: str | None) -> str: + """The emptiest visible accelerator, unless the caller named one. + + Indices come from Torch, not from `nvidia-smi`: this box has several H200s + and some are in exclusive-process mode with an owner already, so the shell + restricts what is visible -- and then a global index names the wrong device + or none at all. Probing through Torch also *tries* each one, which is the + only way to find out that an exclusive device is already taken. + """ + if want: + return want + import torch + + if not torch.cuda.is_available(): + return str(torch.accelerator.current_accelerator()) + best = None + for i in range(torch.cuda.device_count()): + try: + free, _ = torch.cuda.mem_get_info(i) + except Exception: + continue # exclusive mode, already owned + gib = free / 2**30 + if gib >= NEED_GIB and (best is None or gib > best[1]): + best = (i, gib) + if best is None: + raise SystemExit( + f"no visible accelerator has {NEED_GIB:.0f} GiB free " + f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', 'unset')})" + ) + return f"cuda:{best[0]}" + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--prompt", required=True, help="prompt to continue") + ap.add_argument("--max-new-tokens", type=int, default=2048, metavar="TOKENS") + ap.add_argument("--ckpt", type=Path, required=True, + help="published checkpoint directory") + ap.add_argument("--device", default=None, help="runtime device (default: an idle one)") + args = ap.parse_args(argv) + + from transformers import AutoTokenizer + import engine as E + + device = pick_device(args.device) + tokenizer = AutoTokenizer.from_pretrained(str(args.ckpt)) + encoded = tokenizer.encode(args.prompt) + prompt_ids = list(getattr(encoded, "ids", encoded)) + if not prompt_ids: + raise SystemExit("decode needs a prompt that encodes to at least one token") + + t0 = perf_counter() + eng = E.Engine(args.ckpt, HERE / "ref_src", device=device, + max_new=args.max_new_tokens, prompt_room=len(prompt_ids) + 8) + build = perf_counter() - t0 + + result = eng.generate(prompt_ids, args.max_new_tokens) + text = tokenizer.decode(result.tokens) + + print(text) + print() + print(f"{result.tokens_per_second:.1f} tok/s " + f"({len(result.tokens)} tokens in {result.seconds:.3f}s, " + f"{1e3 * result.seconds / len(result.tokens):.3f} ms/token)") + print(f" device {device} | prompt {result.prompt_steps} tokens, " + f"prefill {1e3 * result.prefill_seconds:.1f}ms | " + f"load+compile {build:.1f}s | {len(text)} characters generated") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/no_machine_paths_lint.py b/scripts/no_machine_paths_lint.py index e06f1c52..7c7c6816 100644 --- a/scripts/no_machine_paths_lint.py +++ b/scripts/no_machine_paths_lint.py @@ -30,7 +30,10 @@ #: matched -- pointing inside the repository is what paths here are for. _MACHINE_PATHS = re.compile( r""" - (?:^|[^\w./-]) # not mid-token, so a URL or a longer word is safe + (?:^|(?<=[^\w./])|(?<=-)(?=/)) # not mid-token, so a URL or a longer word is + # safe -- but a `-` immediately before the + # slash is `${VAR:-/path}` or `--ckpt=-`-style + # punctuation, not a token this path is part of ( /(?:home|Users)/[\w.-]+ # a named home directory | /data\d*/(?:shared/)?[\w.-]+/[\w.-]+ # a site-local scratch mount diff --git a/tests/scripts/test_no_machine_paths_lint.py b/tests/scripts/test_no_machine_paths_lint.py index 333b4a8d..231fb32a 100644 --- a/tests/scripts/test_no_machine_paths_lint.py +++ b/tests/scripts/test_no_machine_paths_lint.py @@ -49,6 +49,11 @@ def lint(): f'BASE = "{_USERS}/someone/envs/dev"', # The fallback form: configurable-looking, hardcoded for everyone else. f'os.environ.get("TF_CKPT", "{_SCRATCH}/someone/prepared")', + # The same fallback in shell, which reached a shipped example unreported: the + # `-` of `:-` sat where a token character would, so the whole default was + # invisible to the guard that keeps this off the middle of a URL. + f"CKPT=${{CKPT:-{_SCRATCH}/someone/Qwen3-1.7B}}", + f"python run.py --ckpt={_HOME}/someone/models", ] #: Shapes that resemble the above and are not machine-specific.