diff --git a/CHANGELOG.md b/CHANGELOG.md index 5df2f34..514f94a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,63 @@ deletes its own evidence is not a retraction. No phase shipped yet (`v0.1.0` is still the scaffold version; a tag lands once a phase's exit criteria are actually met). +- **Phase 5a (int8 weight-only quantization) complete**, 2026-09-17 + (`docs/plans/2026-09-16-phase-5a-quantization-plan.md`). Self-computed, + per-output-channel int8 quantization extending the Triton kernel + itself (not sourced from bitsandbytes/AWQ) — `make check` green + throughout (123 tests). Kernel-level correctness gate passed on a real + L40 (15/15 int8, 25/25 bf16). Measured on the real + `deepseek-ai/deepseek-moe-16b-base` model: int8 kernel **+70.9%** over + stock (a throughput tie with the bf16 naive kernel it's built on, + expected since weight-only quantization saves memory bandwidth, not + FLOPs), **49.89%** expert-weight memory reduction, perfect model-level + top-1/mutual-top-k agreement at every tested position. Two real bugs + found on real hardware: quantizing without freeing the original bf16 + expert weights (holding both copies at once, OOMing a 44GB L40; found + live during the rental) and quantization arithmetic computed in the + weight's own bf16/fp16 dtype instead of float32 (widened round-trip + error, and an fp16-specific scale-clamp cliff; found by the final + whole-branch review), both fixed before merge. Cost: **$1.95** of a $5 + cap, including two RunPod Community Cloud pods that hit a real + host-level GPU passthrough bug. Full account: + `docs/findings/2026-09-17-phase-5a-quantization-run.md`. +- **Phase 4 (disaggregated prefill/decode) complete**, 2026-09-17 + (`docs/plans/2026-09-16-phase-4-disaggregated-prefill-decode-plan.md`). + Continuous-batching prefill/decode workers across two real 4-GPU H100 + SXM topologies (co-located 4-rank EP, disaggregated 2+2-rank EP with a + real cross-rank KV-cache handoff), both byte-exact against a + single-GPU reference — `make check` green throughout (102 tests). + Three real bugs found and fixed on real hardware, none caught by + CPU-only tests. Measured result was mixed, not clean: disaggregated + TTFT beats co-located at concurrency 4 (0.46s vs 1.12s) but loses at + concurrency 8 (0.98s vs 0.70s) — most likely a kernel-warmup confound, + reported as genuinely inconclusive. Cost: **$10.94** of a $40 cap. + Full account: + `docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md`. +- **Phase 3 (multi-GPU expert-parallel serving) complete**, 2026-09-16 + (`docs/plans/2026-09-15-phase-3-multi-gpu-expert-parallel-serving-plan.md`). + Real 2x H200 SXM expert-parallel serving over DeepSeek's own DeepEP + (V1 `Buffer`, after V2's `ElasticBuffer` proved unavailable on this + rental's Fabric-Manager-less hardware) — `make check` green + throughout. Correctness gate passed: perfect top-1 and mutual top-5 + agreement vs. a single-GPU reference. The measured run disproved this + project's own hypothesis: Phase 1's naive-vs-persistent crossover + doesn't reproduce under DeepEP's real per-expert token counts (naive + wins at every tested count on H200; real per-local-expert counts, + median 2 max ~20, sit below Phase 1's smallest tested point of 16). + Cost: **$13.03** of a $25 cap. Full account: + `docs/findings/2026-09-16-phase-3-multi-gpu-ep-run.md`. +- **Phase 2 (vLLM benchmark contribution) complete**, 2026-09-15 + (`docs/plans/2026-09-15-phase-2-vllm-benchmark-contribution-plan.md`). + Neither vLLM's nor SGLang's official MoE benchmark modeled skewed + (zipf) expert load; upstreamed an opt-in `--expert-load-distribution` + flag to vLLM's `benchmarks/kernels/benchmark_moe.py`. Measured on one + rented RTX 3090: `--tune` under zipf vs. uniform routing picks a + different winning Triton config at 4 of 5 tested batch sizes. Open + PR: [vllm-project/vllm#57100](https://github.com/vllm-project/vllm/pull/57100). + Cost: **$0.77**. This phase landed no code in `dispatch` itself, only + docs — its actual contribution is the upstreamed PR. Full account: + `docs/findings/2026-09-15-phase-2-vllm-benchmark-run.md`. - **Phase 1 (custom Triton grouped-GEMM kernel) complete**, 2026-09-15 (`docs/plans/2026-09-15-phase-1-grouped-gemm-plan.md`). Built and tested: a CPU-only MoE reference and grouped-GEMM contract, a naive and diff --git a/CLAUDE.md b/CLAUDE.md index 29d8543..bcdc92e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -83,6 +83,63 @@ cost across both paid sessions: **$0.65** ($0.06 kernel correctness + $0.59 the measured run). Full account: `docs/findings/2026-09-15-phase-1-grouped-gemm-run.md`. +**Phase 2 (vLLM benchmark contribution) is complete, 2026-09-15.** +Neither vLLM's nor SGLang's official MoE benchmark modeled skewed (zipf) +expert load; upstreamed an opt-in `--expert-load-distribution` flag to +vLLM's `benchmarks/kernels/benchmark_moe.py`. Measured on one rented RTX +3090: `--tune` under zipf vs. uniform routing picks a **different +winning Triton config at 4 of 5 tested batch sizes**. Open PR: +[vllm-project/vllm#57100](https://github.com/vllm-project/vllm/pull/57100), +awaiting maintainer review. Cost: **$0.77**. This phase landed no code +in `dispatch` itself, only docs — its actual contribution is the +upstreamed PR. Full account: +`docs/findings/2026-09-15-phase-2-vllm-benchmark-run.md`. + +**Phase 3 (multi-GPU expert-parallel serving) is complete, 2026-09-16**, +merged via PR #3 alongside Phase 2's docs. Real 2x H200 SXM EP over +DeepSeek's own DeepEP (V1 `Buffer`, after V2's `ElasticBuffer` proved +unavailable -- this rental had no GPU Fabric Manager). Correctness gate +passed: perfect top-1 and mutual top-5 agreement vs. a single-GPU +reference. **The measured answer disproved this project's own +hypothesis**: Phase 1's naive-vs-persistent crossover doesn't reproduce +under DeepEP's real per-expert token counts (naive wins at every tested +count on H200; real per-local-expert counts, median 2 max ~20, sit +below Phase 1's smallest tested point of 16). Cost: **$13.03** of a $25 +cap. Full account: +`docs/findings/2026-09-16-phase-3-multi-gpu-ep-run.md`. + +**Phase 4 (disaggregated prefill/decode) is complete, 2026-09-17**, +merged via PR #4. Continuous-batching prefill/decode workers across two +real 4-GPU H100 SXM topologies (co-located 4-rank EP, disaggregated +2+2-rank EP with a real cross-rank KV-cache handoff), both proven +byte-exact against a single-GPU reference. Three real bugs found and +fixed on real hardware, none caught by CPU-only tests. **The measured +result was mixed, not clean**: disaggregated TTFT beats co-located at +concurrency 4 (0.46s vs 1.12s) but loses at concurrency 8 (0.98s vs +0.70s) -- most likely a kernel-warmup confound, reported as genuinely +inconclusive. Cost: **$10.94** of a $40 cap. Full account: +`docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md`. + +**Phase 5a (int8 weight-only quantization) is complete, 2026-09-17**, on +its own branch (`phase-5a-quantization`), PR pending. Self-computed, +per-output-channel int8 quantization extending the Triton kernel itself +(not sourced from bitsandbytes/AWQ). Kernel-level correctness gate +passed on a real L40 (15/15 int8, 25/25 bf16, both first-time on this +hardware). A real bug found live during the GPU session -- quantizing +without freeing the original bf16 expert weights, holding both copies +at once and OOMing a 44GB L40 -- was found and fixed mid-session, and +the final whole-branch review caught two more real precision defects +(quantization arithmetic done in the input's own bf16/fp16 dtype +instead of float32, and an fp16-specific scale-clamp cliff), both fixed +before merge. Measured: int8 kernel **+70.9%** over stock (a throughput +tie with the bf16 naive kernel it's built on, as expected -- weight-only +quantization saves memory, not FLOPs), **49.89%** expert-weight memory +reduction, perfect model-level top-1/mutual-top-k agreement at every +tested position. Cost: **$1.95** of a $5 cap, including two RunPod +Community Cloud pods that hit a real host-level GPU passthrough bug +before a Secure Cloud pod worked. Full account: +`docs/findings/2026-09-17-phase-5a-quantization-run.md`. + ## One governing principle **Correctness before speed.** A kernel is not "fast" until it has been diff --git a/README.md b/README.md index 8e99eab..540c9f7 100644 --- a/README.md +++ b/README.md @@ -1,41 +1,278 @@ # dispatch -An inference engine for the token-routing layer of Mixture-of-Experts (MoE) -language models — a custom GPU kernel for expert dispatch and grouped-GEMM -execution, multi-GPU expert-parallel serving, and disaggregated -prefill/decode, benchmarked against production serving engines (vLLM, -SGLang) on real, measured throughput/latency/cost numbers, not estimates. - -**Status: Phase 4 (disaggregated prefill/decode) complete, open as a PR.** -Phases 0-3 are merged: a plain-HF baseline, a naive and a persistent -Triton grouped-GEMM kernel (**~65-67% faster decode throughput** than -DeepSeek's own stock `moe_infer`, perfect mutual top-5/top-1 logit -agreement), a skewed-load benchmark contribution upstreamed to vLLM, and -real 2-GPU expert-parallel serving via DeepSeek's DeepEP. Phase 4 adds a -continuous-batching prefill/decode scheduler and tests it across two real -4-GPU topologies — a co-located 4-rank EP pool and a disaggregated -2+2-rank EP pool connected by a real cross-rank KV-cache handoff — both -proven byte-exact correct against a single-GPU reference on real -hardware. Three real bugs found and fixed along the way, all now covered -by regression tests. See -[`docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md`](docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md) -for the full measured account, and -[`docs/design/2026-09-14-dispatch-system-design.md`](docs/design/2026-09-14-dispatch-system-design.md) -for the architecture — model choice, what's custom vs. reused, the 7-phase -plan, and explicit scope boundaries — and `docs/STATUS.md` for the -authoritative, same-commit-updated record of what's actually built. -Conventions and tooling carried over from -[almanac](https://github.com/bsreecharanreddy/almanac) and -[canopica](https://github.com/bsreecharanreddy/canopica) where applicable. +[![CI](https://github.com/bsreecharanreddy/dispatch/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/bsreecharanreddy/dispatch/actions/workflows/ci.yml) +![python](https://img.shields.io/badge/python-3.12%2B-blue) +![mypy](https://img.shields.io/badge/mypy-strict-blue) + +An **inference engine for the token-routing layer of Mixture-of-Experts +(MoE) language models** — the architecture nearly every frontier model now +uses (DeepSeek, Llama, Mixtral, Grok, Qwen). A router sends each token to a +handful of specialist sub-networks ("experts") out of many; the hard +engineering problem is making that dispatch fast in practice — a custom GPU +kernel, multi-GPU expert-parallel serving, disaggregated prefill/decode — +not just correct on paper. + +This project writes a real custom Triton kernel for that dispatch step, +gets it running across multiple real GPUs, and reports the result against +production serving-engine methodology (vLLM, SGLang) on rented hardware: +throughput, latency, and cost per million tokens, never estimated. + +**No number in this README is quoted unless it was measured.** Where a +result was a null or a mixed one, it's reported that way. + +--- + +> **Status:** Phases 0-4 merged to `main` across four PRs +> ([#1](https://github.com/bsreecharanreddy/dispatch/pull/1) Phase 0, +> [#2](https://github.com/bsreecharanreddy/dispatch/pull/2) Phase 1, +> [#3](https://github.com/bsreecharanreddy/dispatch/pull/3) Phases 2-3, +> [#4](https://github.com/bsreecharanreddy/dispatch/pull/4) Phase 4) — +> Phase 2 landed as docs only, its actual contribution being the +> upstreamed vLLM PR below, not a dispatch-repo code change. **Phase 5a +> (int8 weight-only quantization) is complete** on its own branch, PR +> pending. Full task-by-task record: +> [`docs/STATUS.md`](docs/STATUS.md). Design and phasing: +> [`docs/design/2026-09-14-dispatch-system-design.md`](docs/design/2026-09-14-dispatch-system-design.md). + +## In sixty seconds + +**What it does.** A pure-PyTorch reference MoE forward pass, a from-scratch +Triton grouped-GEMM kernel checked against it on real hardware before any +speed claim, real 2-to-4-GPU expert-parallel serving over DeepSeek's own +DeepEP library, and a disaggregated prefill/decode scheduler — each phase's +thesis stated up front and reported honestly, including the phases where +the answer was a null or mixed result. + +**Built solo, start to finish** — design docs, the kernel, the multi-GPU +serving path, an upstreamed open-source benchmark contribution, every +rented-GPU session, and the write-ups of what broke along the way. + +**Measured across seven rented-GPU sessions (six phases), $27.69 total, +all under their stated caps:** a custom Triton kernel **~65-73% faster** than DeepSeek's own +stock MoE forward pass at perfect logit agreement (Phase 1); an opt-in +skewed-load benchmark flag upstreamed to vLLM, changing which kernel config +its own tuner picks at 4 of 5 tested batch sizes (Phase 2); a real 2-GPU +DeepEP-backed expert-parallel path, byte-exact against a single-GPU +reference, that also *disproved* this project's own kernel-crossover +hypothesis at real scale (Phase 3); a disaggregated prefill/decode +topology across 4 real GPUs that returned a genuinely mixed result rather +than a forced win (Phase 4); and an int8 weight-only quantized kernel that +cut expert-weight memory **49.89%** at perfect model-level agreement, while +also surfacing and fixing a real bug where the naive quantized path was +holding both the bf16 and int8 copies in memory at once (Phase 5a). + +**123 tests, `make check` green throughout Phase 5a** (lint, `mypy --strict`, +and the full non-GPU suite) — GPU-dependent tests are marked and excluded +from CI by design, then run for real on rented hardware every phase. + +## Why this project + +Nearly every frontier MoE model gets its capacity-per-compute win from the +same mechanic: route each token to a few experts out of many, and make that +routing fast under wildly skewed per-expert token counts. That kernel and +the serving system around it is a specific, in-demand skill set — +inference engineering, the kind OpenAI, Anthropic, and NVIDIA hire for +under that title — distinct from training a model or building a +data/ML platform. + +**The router and the cross-GPU communication layer are both solved, +production-grade problems** — DeepSeek's own DeepEP library does the +all-to-all dispatch/combine, and top-k gating is plain PyTorch. The +differentiated work this project actually does is the grouped-GEMM expert +kernel itself and the serving system built around it — not reinventing +either solved piece. Three decisions carry that scoping, each checked +against real current material before being made, not assumed: + +- **[ADR-0001](docs/adr/0001-grouped-gemm-kernel-over-attention-kernel.md)** + — a grouped-GEMM expert kernel over another from-scratch attention-kernel + reimplementation, a pattern already saturated on GitHub. +- **[ADR-0002](docs/adr/0002-deepep-over-hand-rolled-communication.md)** — + DeepEP over hand-rolled cross-GPU communication, and the real NVLink/SXM + hardware constraint that follows it into Phase 3's cost. +- **[ADR-0003](docs/adr/0003-rust-router-over-python-only-server.md)** — a + Rust router in front of the Python model server (Phase 7), mirroring + Hugging Face's own `text-generation-inference` three-tier split, added + after checking real inference-engineer postings against the original + design and finding real gaps it didn't cover. + +Positioned deliberately differently from this author's other portfolio +projects: `almanac` is a data + ML platform, `canopica` is a governed +decision system with an AI capability layer, `dispatch` is the GPU-systems +piece neither of them touches. + +## The centerpiece + +**Correctness before speed.** A kernel is not "fast" until it has been +proven numerically correct against a reference implementation within a +stated tolerance — a kernel that's fast because it's silently wrong is a +bug, not a result, and it's invisible unless the correctness check is +explicit. Every phase below follows the same build order: reference +implementation first, correctness gate against it on real hardware second, +a speed or scaling claim only after that gate passes. Phase 3's own kernel +crossover hypothesis failed to reproduce at real scale — found *because* +the correctness gate ran before the benchmark did, not after. + +## Architecture — what's custom vs. reused + +```mermaid +flowchart LR + subgraph reused["Reused (production tools)"] + R["Router / top-k gating\n[plain PyTorch]"] + D["Cross-GPU dispatch/combine\n[DeepEP, Phase 3+]"] + Q["Quantization\n[self-computed int8, Phase 5a]"] + B["Benchmark methodology\n[vLLM/SGLang-style metrics]"] + end + + subgraph custom["Custom (this project's actual work)"] + K["Grouped-GEMM expert kernel\n[Triton, Phase 1]"] + RT["Rust router\n[HTTP/gRPC, batching, Phase 7]"] + MS["Python model server\n[runs router+kernel, gRPC, Phase 7]"] + end + + R --> K --> MS + RT -->|gRPC| MS + D <-.-> MS + Q -.-> K + B -->|measures| MS + B -->|measures, same config| V["vLLM / SGLang\n(comparison baseline)"] +``` -## Development +The differentiated work is the grouped-GEMM kernel and the serving system +around it. Phases 0-6 drive the kernel and model directly from lightweight +Python test/benchmark harnesses; the Rust-router/Python-model-server split +is Phase 7's contribution, once there's a real service worth +productionizing. + +## What has been measured + +Each phase started with a stated thesis and reported the real answer, +including the ones that came back null or mixed: + +| Phase | Thesis tested | Measured result | +|---|---|---| +| **0 — baseline** | What does `deepseek-ai/deepseek-moe-16b-base` actually cost per token, unpatched, on one rented GPU? | bf16, single L40, unbatched eager decode: **12.75 tok/s, 0.355s mean TTFT** (p50 0.258s, p99 1.588s). **$0.35** total, three attempts included | +| **1 — grouped-GEMM kernel** | Does a custom Triton grouped-GEMM kernel beat DeepSeek's own stock `moe_infer`, and does a persistent/cache-aware variant beat the naive one? | Naive kernel **+67.2%** (12.55 → 20.98 tok/s), persistent **+65.7%** — both at perfect mutual top-5/top-1 logit agreement. Persistent **ties** naive at the 1-token/step granularity that actually drives decode (a null result the plan predicted before the run), wins 3-8% at 16-128 tokens, loses up to 14% at 512-2048 | +| **2 — vLLM contribution** | Do vLLM's or SGLang's official MoE benchmarks model skewed (zipf) expert load? | Neither did. Upstreamed an opt-in `--expert-load-distribution` flag to vLLM; on one rented RTX 3090, `--tune` under zipf vs. uniform routing picks a **different winning Triton config at 4 of 5 tested batch sizes**. Open: [vllm-project/vllm#57100](https://github.com/vllm-project/vllm/pull/57100) | +| **3 — multi-GPU EP** | Does Phase 1's naive-vs-persistent crossover hold under DeepEP's real per-expert token distribution across GPUs? | **It doesn't apply at real scale.** Naive wins at every tested token count (16-2048) on H200, independent of EP. Real DeepEP dispatch across the model's 27 MoE layers produces per-local-expert counts (median 2, max ~20) far below Phase 1's smallest tested point (16) | +| **4 — disaggregated prefill/decode** | Does separating prefill and decode across GPU pools relieve contention, holding GPU count fixed at 4? | **Mixed, not a clean result.** Disaggregated TTFT beats co-located at concurrency 4 (0.46s vs 1.12s) but loses at concurrency 8 (0.98s vs 0.70s) — most likely a kernel-warmup confound between independent process launches, reported as genuinely inconclusive rather than forced | +| **5a — int8 quantization** | Does a self-computed, per-channel int8 weight-only kernel cut memory with no model-level correctness cost? | **Yes on both counts, no speedup.** **49.89%** expert-weight memory reduction, perfect top-1/mutual-top-k agreement vs. the naive bf16 kernel at every tested position. Throughput was a near-exact tie with naive (**-1.2%**), as expected — weight-only quantization saves memory bandwidth, not FLOPs. A real bug (quantizing without freeing the original bf16 weights, OOMing a 44GB L40) was found and fixed mid-session | + +Full method and every number: one findings doc and one cost doc per phase +in [`docs/findings/`](docs/findings/). + +## What it does not do + +Stated here rather than left for a reader to discover — from the design +doc's own scope boundaries: + +- **Does not train or fine-tune the model.** DeepSeekMoE-16B's published + weights are used as-is. +- **Does not generalize the kernel to other MoE architectures** in this + pass — proving the mechanic on one real model is the scope; generalizing + it is a possible later phase, not a Phase 0-7 requirement. +- **Does not run a production autoscaling or fleet-management layer.** + Phase 7's Kubernetes deployment is applied for real once, verified, and + torn down — it does not keep a service running continuously. +- **Does not reimplement DeepEP, quantization algorithms, or the model's + own forward pass from scratch.** Those are solved problems this project + deliberately reuses; Rust (Phase 7) is scoped to the router only. +- **Does not chase the enterprise-ML-integration profile** — that territory + is deliberately `almanac`'s, not this project's. + +## Stack + +| Layer | Choice | +|---|---| +| Model | `deepseek-ai/deepseek-moe-16b-base` — 16.4B params, fine-grained MoE, shared + routed experts | +| Kernel | Custom Triton grouped-GEMM (naive + persistent/cache-aware), int8 weight-only quantization (Phase 5a) | +| Multi-GPU | DeepSeek's DeepEP — real all-to-all expert-parallel dispatch/combine over NVLink/SXM | +| Serving | Continuous-batching prefill/decode workers, both co-located and disaggregated topologies | +| Benchmark methodology | vLLM/SGLang-style metrics — TTFT, inter-token latency, tokens/sec, cost per million tokens | +| Language | Python 3.12+ — `uv`, `ruff`, `mypy --strict`, `pytest` | +| GPU provisioning | RunPod's API via lightweight scripts (`scripts/gpu/`) — not Terraform; see `CLAUDE.md`'s cost discipline | +| CI | GitHub Actions — lint, `mypy --strict`, and the non-GPU suite on every push | + +## Layout + +```text +src/dispatch/ + benchmark/ metrics, generation harness, reference-logit capture/comparison + kernels/ reference MoE, Triton grouped-GEMM (naive + persistent), int8 quantization, + expert-parallel sharding, backend registry + serving/ KV-cache slice/handoff, continuous-batching prefill/decode, co-located worker +scripts/ + run_baseline.py latency/throughput CLI, stock or kernel-patched, --compare-reference gate + run_kernel_bench.py pure kernel micro-benchmark, refuses to time a backend that disagrees + gpu/ RunPod API client, provisioning CLI, pod-live correctness/concurrency scripts +tests/ unit + integration; `gpu`-marked tests need real CUDA, excluded from CI +docs/design/ the authoritative architecture, phase plan, and scope document +docs/adr/ decisions, each checked against real current material before being made +docs/plans/ per-phase implementation plans, written before any code +docs/findings/ measured results and costs — one run doc and one cost doc per phase +docs/runbooks/ the exact GPU-session steps each phase's real run followed +``` + +## Running it ```bash uv sync --all-extras --dev -make check # lint + typecheck + test — the full gate, same as CI -make check-fast # inner loop +make check # lint + mypy --strict + the non-GPU test suite -- same as CI +make check-fast # inner loop, also skips slow-marked tests ``` +GPU-dependent tests (`pytest -m gpu`) need a real CUDA device and are +excluded from both targets by design — this repo's CI has no GPU runner. +They're run for real on rented hardware every phase; each phase's findings +doc in [`docs/findings/`](docs/findings/) records exactly which GPU, for +how long, and at what cost. + +## Design and the paper trail + +[`docs/design/2026-09-14-dispatch-system-design.md`](docs/design/2026-09-14-dispatch-system-design.md) +is the authoritative architecture, phasing, and scope document. + +| | | +|---|---| +| [`docs/STATUS.md`](docs/STATUS.md) | the verification log, at task granularity, updated in the same commit as the work | +| [`docs/adr/`](docs/adr/) | three decisions, each citing the real material that settled it | +| [`docs/findings/`](docs/findings/) | the measurements themselves, one run doc and one cost doc per phase | +| [`docs/plans/`](docs/plans/) | per-phase implementation plans, written before any code | +| [`docs/runbooks/`](docs/runbooks/) | the exact GPU-session steps each phase's real run followed | + +## Who should look at what + +- **Inference / GPU-systems engineering** — the kernel itself + (`src/dispatch/kernels/grouped_gemm.py`), then + [Phase 1's findings doc](docs/findings/2026-09-15-phase-1-grouped-gemm-run.md): + a from-scratch Triton kernel proven correct on real hardware before a + single speed claim was made. +- **Distributed / multi-GPU systems** — `src/dispatch/kernels/expert_parallel.py` + and `src/dispatch/serving/`, then + [Phase 3's](docs/findings/2026-09-16-phase-3-multi-gpu-ep-run.md) and + [Phase 4's](docs/findings/2026-09-17-phase-4-disaggregated-prefill-decode-run.md) + findings docs for two results — one null, one mixed — found by measuring + on real hardware rather than assumed. +- **ML infra / open-source contribution** — + [Phase 2's findings doc](docs/findings/2026-09-15-phase-2-vllm-benchmark-run.md) + and the open [vLLM PR](https://github.com/vllm-project/vllm/pull/57100): a + real gap found by checking what a production project already ships, not + duplicating it. +- **Anyone checking whether the claims hold** — [`docs/STATUS.md`](docs/STATUS.md) + and [`docs/findings/`](docs/findings/), where a null or mixed result + (Phase 1's persistent-kernel tie, Phase 3's crossover disappearing, + Phase 4's inconclusive disaggregation number) is reported as such rather + than smoothed over. + +`CLAUDE.md` is instructions for AI coding assistants working in this repo, +not a document for a human evaluating the project. + +## Why the name + +"Dispatch" is the literal term for the MoE routing step this project +builds a kernel for — each token sent to a handful of expert sub-networks +out of many. It's also, not incidentally, what this project does with +every rented GPU session: dispatch the work, measure it, tear it down. + ## License Dual-licensed under either [MIT](LICENSE-MIT) or diff --git a/docs/STATUS.md b/docs/STATUS.md index b1799a2..8e4e987 100644 --- a/docs/STATUS.md +++ b/docs/STATUS.md @@ -341,11 +341,80 @@ instruction once both correctness gates had passed. Full account: **Total Phase 4 GPU cost: $10.94** of the $40 cap, 4x H100 SXM, 47 minutes. +## Phase 5a progress + +Design: `docs/design/2026-09-16-phase-5a-quantization.md`. Plan: +`docs/plans/2026-09-16-phase-5a-quantization-plan.md`. Branch +`phase-5a-quantization`. Split from Phase 5 (quantization + +speculative decoding) into two sub-phases; 5b (speculative decoding) is +separate, later work. + +- [x] Task 1: `QuantizedTensor`, per-channel int8 quantize/dequantize (`quantization.py`) +- [x] Task 2: Quantized MoE forward path + memory-footprint helpers +- [x] Task 3: Triton int8 grouped-GEMM kernel + backend resolver (`grouped_gemm_int8.py`, `backends.py`) +- [x] Task 4: `patch_moe_infer_quantized`, sharing layer-iteration with `patch_moe_infer` +- [x] Task 5: `--moe-kernel quantized` wired through `run_baseline.py` +- [x] Task 6: GPU rental runbook -- correctness gate, three-way measured run, memory footprint, cost +- [x] Task 7: this update + +**Phase 5a is complete.** `make check` green throughout (119 tests, +lint and `mypy --strict` clean). + +**Kernel-level correctness gate passed** on a real NVIDIA L40 (RunPod +Secure Cloud): the int8 kernel's first real-hardware run, 15/15 GPU +tests passed against `torch_grouped_matmul_dequant` (an independent +quantize-then-dequantize-then-eager-matmul reference); the existing bf16 +kernels re-verified at 25/25. + +**A real bug found and fixed on real hardware, none caught by CPU-only +tests:** `patch_moe_infer_quantized` quantized each layer's stacked bf16 +weights into int8 but never released the bf16 originals -- +`stack_expert_weights` re-points every expert's `Linear.weight` at a +*view* into one shared bf16 tensor (so building that stack costs no +extra memory in the existing bf16 path), and those views kept the whole +bf16 tensor resident even after quantizing, holding both copies at once +and OOMing a 44GB L40 at DeepSeekMoE-16B's real scale. Fixed with +`_free_expert_weights` (frees each expert's original weight right after +quantizing, since the quantized closure never reads `experts` again) and +a new CPU regression test asserting those weights are actually freed. + +**Measured three-way run** (bf16, single L40, unbatched eager decode, 15 +runs per config, 27/27 MoE layers patched): stock **12.13 tok/s** ($18.78 +per 1M tokens), naive bf16 kernel **20.97 tok/s** (**+72.9%**, $10.86 per +1M tokens -- consistent with Phase 1's separately-measured +67.2% on the +same GPU class), int8 quantized kernel **20.72 tok/s** (**+70.9%** over +stock, but **-1.2%** vs. the naive kernel it's built on -- essentially a +throughput tie, expected: this is a weight-only quantization meant to +save memory bandwidth/footprint, not FLOPs, and the naive kernel already +runs `tl.dot` at native tensor-core precision). + +**Memory footprint**, computed directly from the real model's stacked +expert weights across all 27 MoE layers: bf16 **27.84 GiB** -> int8 +**13.95 GiB**, a **49.89%** reduction (just under a clean 50% from the +per-channel fp32 scale overhead). + +**Model-level agreement** (quantized vs. the naive bf16 run, per the +design doc's two-tier bar): **perfect top-1 agreement and mutual top-k +membership at every tested position**, all 3 prompts -- stronger than +the bar required, which explicitly allowed for real divergence here. + +**Two RunPod Community Cloud pods hit the same real host-level bug this +session** (`cuInit()` returning `CUDA_ERROR_UNKNOWN` even against the +base image's own stock torch, confirmed via a raw ctypes test) before a +Secure Cloud L40 worked immediately -- a new environment finding, not +previously documented in this project. `transformers==4.57.6` + +`DynamicCache.get_usable_length` monkeypatch (Phase 0/1/3/4's fix) was +still needed; new this session, `HF_HUB_ENABLE_HF_TRANSFER=1` was set on +the pod with `hf_transfer` not installed. + +**Total Phase 5a GPU cost: $1.9543** of the $5 cap (39.1%), across all +three pods including both abandoned Community Cloud attempts. Full +account: `docs/findings/2026-09-17-phase-5a-quantization-run.md`. + ## Next step -Phase 4 is complete and ready to push as a PR (`make check` green, -correctness gates passed, findings recorded, cost well under cap). -Phase 3's PR (#3) already merged. Phase 2's PR is still open and awaiting -maintainer review; no further work planned on it beyond responding to -review feedback. Phase 5 (quantization + speculative decoding) is not +Phase 4 is merged (PR #4). Phase 5a is complete and ready to push as a +PR. Phase 3's PR (#3) already merged. Phase 2's PR is still open and +awaiting maintainer review; no further work planned on it beyond +responding to review feedback. Phase 5b (speculative decoding) is not yet planned. diff --git a/docs/design/2026-09-16-phase-5a-quantization.md b/docs/design/2026-09-16-phase-5a-quantization.md new file mode 100644 index 0000000..443188f --- /dev/null +++ b/docs/design/2026-09-16-phase-5a-quantization.md @@ -0,0 +1,222 @@ +# Phase 5a design: int8 weight-only quantized grouped-GEMM kernel + +Per `docs/design/2026-09-14-dispatch-system-design.md` §7 (phase plan, row +5: "Quantization + speculative decoding layered on top"). That row bundles +two independent techniques; decided during brainstorming to split it into +two sequential sub-phases, each single-focus like every phase before it. +**This document covers 5a (quantization) only.** Speculative decoding is +5b, planned separately once 5a is complete. + +## 1. Purpose and thesis + +Every phase through Phase 4 has kept the routed-expert weights in bf16. +Weight-only int8 quantization is the standard first lever production +serving engines pull for memory: it roughly halves the routed-expert +weight footprint (the majority of DeepSeekMoE-16B's 32.8GB) with activations +left untouched, at some cost in output fidelity. + +**Phase 5a's thesis: can the custom Triton grouped-GEMM kernel itself +consume int8 quantized expert weights -- not just call an existing +quantization library on the stock HF path -- and if so, at what memory, +throughput, and quality tradeoff, measured on the same real model and +hardware class this project has used throughout?** + +This is a deliberate scope increase over "call bitsandbytes and report the +number": the kernel-engineering work (a new Triton kernel that dequantizes +inside the tile loop) is the differentiated contribution, consistent with +this project's overall thesis that the kernel is the custom part and +routing/communication/quantization *algorithms* are solved problems +(§3 of the system design doc; also its §11 out-of-scope note that +reimplementing bitsandbytes/AWQ's *algorithm* is not this project's job -- +writing a kernel that executes one is a different claim). + +## 2. Scope, decided during brainstorming + +- **Weight-only int8 (W8A16).** Activations stay bf16. Only the routed + expert weight tensors already flowing through Phase 1's kernel + (`gate_proj`, `up_proj`, `down_proj`) are quantized. Shared experts, + attention, the router/gate, and `lm_head` stay bf16 and untouched -- + exactly Phase 1's kernel scope, not expanded. +- **Per-output-channel symmetric round-to-nearest scales, self-computed.** + `scale = max(abs(weight_row)) / 127` per expert per output channel, + computed directly in `dispatch`'s own code rather than sourced from + bitsandbytes/AWQ. Simple enough (a handful of lines) that it is not + "reimplementing" a library's job -- it keeps the entire + quantize -> kernel -> dequant-in-kernel path inside this project's own, + fully inspectable code. +- **Built on the naive kernel, not persistent.** Phase 1's naive kernel won + every benchmark that has run since (Phase 1's own token-count sweep, + Phase 3's real-scale H200 measurement) -- no reason to add quantization + complexity on top of the kernel variant that has never won. +- **Single GPU, reusing Phase 1's shape**, not Phase 3/4's multi-GPU + rentals -- a kernel-correctness session, then one measured run. + **Budget cap: $5**, set now, before any rental. +- **int4, activation quantization (W8A8), and bitsandbytes/AWQ-sourced + scales are explicitly out of scope for 5a** -- see §7 Non-goals. + +## 3. Architecture + +The existing bf16 pluggable-kernel contract +(`GroupedMatmul = Callable[[Tensor, Tensor, TileSchedule], Tensor]` in +`moe_forward.py`, consumed by `torch_grouped_matmul` and the naive/persistent +Triton kernels) is **left untouched**. Quantized weights need a second +tensor (the per-channel scale) that doesn't fit that signature, and forcing +it in would mean every existing bf16 backend carries dead-weight branching +it doesn't need. Instead, Phase 5a adds a fully parallel set of types and +functions, the same pattern Phase 3 used for EP (`patch_moe_infer_ep` +added alongside `patch_moe_infer`, not folded into it): + +- **`QuantizedTensor`** (new dataclass) -- `data: Tensor` (int8, shape + `(E, N, K)`), `scale: Tensor` (fp32, shape `(E, N)`, one scale per + expert per output channel). +- **`QuantizedStackedExpertWeights`** (new dataclass) -- `gate, up, down: + QuantizedTensor`, mirroring `StackedExpertWeights`. +- **`QuantizedGroupedMatmul`** (new type alias) -- `Callable[[Tensor, + QuantizedTensor, TileSchedule], Tensor]`. +- **`grouped_moe_routed_quantized`** (new function) -- the same body as + `grouped_moe_routed` (group tokens, build tile schedule, three matmuls + through SiLU, ungroup) but typed against `QuantizedStackedExpertWeights` + and `QuantizedGroupedMatmul`. No behavior is shared by mutating the + original; the logic is small enough (four lines) that duplicating it + cleanly beats making the original generic over a union type. +- **`quantize_per_channel_int8`** (new function) -- `Tensor -> QuantizedTensor`, + the round-to-nearest scale computation. +- **`quantize_stacked_weights`** (new function) -- `StackedExpertWeights -> + QuantizedStackedExpertWeights`, applying the above to each projection. + This is the hook point called once per layer at patch time (see below), + never per forward pass. +- **`torch_grouped_matmul_dequant`** (new function) -- `(Tensor, + QuantizedTensor, TileSchedule) -> Tensor`: dequantizes the given + `QuantizedTensor` to fp32 and runs the same per-expert-slice matmul loop + `torch_grouped_matmul` does. This is **not** a model-quality reference -- + it is the Triton int8 kernel's correctness oracle (see §5). +- **`grouped_matmul_int8`** (new Triton kernel, in `grouped_gemm.py` or a + new sibling module if that file gets unwieldy -- decided in the + implementation plan) -- built on the naive kernel's grid/launch + structure; each tile loads an int8 weight sub-tile, widens and multiplies + by its output-channel's fp32 scale, then proceeds through the same + `tl.dot` accumulation the naive kernel uses. +- **`patch_moe_infer_quantized`** (new function, `integration.py`) -- + mirrors `patch_moe_infer` but calls `quantize_stacked_weights` on the + freshly-stacked bf16 weights before building the per-layer closure, and + wires `grouped_moe_routed_quantized` + `grouped_matmul_int8` instead of + the bf16 path. Quantization happens once per layer, at patch time, not + per forward pass. +- **CLI**: `scripts/run_baseline.py --moe-kernel` gains a new value + (e.g. `quantized`) that routes to `patch_moe_infer_quantized` instead of + `patch_moe_infer(model, resolve_backend(name))`. `resolve_backend`'s + `BACKENDS` tuple and its bf16 contract are unchanged. + +## 4. Data flow + +1. **Patch time (once per layer, once per run):** `stack_expert_weights` + builds the bf16 `StackedExpertWeights` exactly as today. + `quantize_stacked_weights` converts it to `QuantizedStackedExpertWeights` + -- three `QuantizedTensor`s, each with int8 data and fp32 per-channel + scales. The original bf16 stacked weights are not retained after this + (memory savings are real, not just reported). +2. **Forward pass (every call, unchanged token-routing logic):** + `grouped_moe_routed_quantized` groups tokens by expert (identical + `group_tokens_by_expert`/`build_tile_schedule` machinery, untouched), + then calls `grouped_matmul_int8` three times (gate, up, down) exactly + where `grouped_moe_routed` calls `matmul` today. +3. **Inside the kernel:** each tile dequantizes its int8 weight slice + against that expert/channel's scale before the same tiled accumulation + the naive kernel already does. Activations (`x`) stay bf16 throughout -- + only the weight side is ever int8. + +## 5. Correctness -- two distinct layers, so the tight bar survives contact with lossy quantization + +Quantization is lossy *by design*, so "correctness before speed" needs two +separate claims rather than one, or the project would either fake precision +it doesn't have or quietly drop the correctness discipline for this phase +alone: + +- **Kernel-level (tight, same rigor as every prior phase).** The Triton + `grouped_matmul_int8` kernel's output vs. `torch_grouped_matmul_dequant` + fed the *same* `QuantizedTensor` -- both sides use identical already-quantized + weights, so nothing here should diverge beyond float-precision noise. + Held to `assert_matches_reference`'s existing tolerance discipline. A + mutation (e.g. forcing every tile to use expert 0's scale) must turn this + red, same bar Phase 1 set. +- **Model-level (stated, expected to show real divergence).** Mutual + top-5 / top-1 logit agreement (`compare_top_k_agreement`, unchanged, no + new machinery) between the int8-kernel-patched real model and Phase 1's + bf16-naive-kernel-patched real model, same 3 prompts used throughout. + This *is* expected to disagree somewhat -- quantization changing outputs + is the entire point of measuring it. Whatever agreement percentage comes + back gets reported as-is, not forced toward 100% or hidden if it's worse + than hoped. + +## 6. Testing + +Extends CLAUDE.md's existing testing table: + +| Layer | What's covered | +|---|---| +| Quantization math (CPU) | `quantize_per_channel_int8`: correct scale computation, round-trip error bounds on known tensors, an all-zero channel doesn't divide by zero, clamping at +-127. | +| Kernel contract (CPU) | `quantize_stacked_weights` produces correctly-shaped `QuantizedStackedExpertWeights`; tiling/grouping logic is reused unchanged from Phase 1, so no new tile-boundary tests are needed. | +| Kernel (`gpu`) | `grouped_matmul_int8` vs. `torch_grouped_matmul_dequant` at toy, decode-, and prefill-shaped dims -- the tight, kernel-level bar from §5. A mutation must turn this red. | +| End to end (`gpu`, paid) | Mutual top-5/top-1 agreement between the int8-patched real model and Phase 1's bf16-naive-patched real model, all 3 prompts -- the model-level bar from §5, reported honestly. **Gates the benchmark**: a run that fails to produce a patched count refuses to run, same discipline as every prior phase's `--moe-kernel`. | +| Benchmarks | Three-way `do_bench`/harness comparison (stock / bf16-naive / int8) on the same prompts and hardware; plus a new pure-function memory-footprint measurement (bytes of `StackedExpertWeights` vs. `QuantizedStackedExpertWeights`, CPU-testable, no GPU needed) since memory reduction is the actual point of quantizing. | + +## 7. Non-goals + +- **int4 or any bit-width below int8.** A real follow-on if 5a's int8 result + is promising, not required for 5a to be "done." +- **Activation quantization (W8A8).** Weight-only only; activations stay + bf16 throughout. +- **bitsandbytes/AWQ-sourced quantization scales.** Decided during + brainstorming: self-computed round-to-nearest scales, to keep the whole + path inside this project's own code (see §2). Not a comment on those + libraries' quality -- their calibration-based schemes are the natural + next comparison point, not part of 5a's claim. +- **Quantizing the persistent kernel.** Only the naive kernel gets an int8 + variant (§2) -- the persistent kernel has not won a single benchmark + since Phase 1's own token-count sweep found its win region doesn't + reach this project's real workload scale (Phase 3's finding). + Revisiting it is not blocked by anything in this design if a future + phase finds a reason to. +- **Combining with multi-GPU EP (Phase 3/4).** Single GPU only, matching + the design doc's Phase 5 cost shape ("single/dual GPU"). A quantized + + EP combination is a plausible future extension, not required here. +- **Speculative decoding.** Phase 5b, planned and executed separately once + 5a's PR merges. +- **Beating vLLM/SGLang's own quantized-serving numbers.** Phase 6 owns + the vLLM/SGLang comparison; 5a's claim is limited to "this project's own + kernel, bf16 vs. int8," an internal comparison like Phase 4's + co-located-vs-disaggregated one. + +## 8. Risk, cost, and rollout + +- **Cost discipline is part of the result, not just a constraint.** $5 is + a ceiling; the gap between cap and actual spend gets reported the way + every prior phase has (Phase 1: $0.65 of a combined $8 across two + sessions; Phase 3: $13.03 of $25; Phase 4: $10.94 of $40). +- **Maximize what's proven before any rental.** Every row in §6 above the + `gpu` line is real, CPU-only, $0 to iterate on -- the quantization math, + the dataclasses, the shape/round-trip tests. The paid session should + start from that suite green and spend rented time only on what + genuinely needs a GPU: the Triton kernel's own correctness, then the + measured run. +- **Named risk: int8 dequant-in-kernel may need a numerically different + tolerance than bf16-to-bf16 comparisons have used.** If + `assert_matches_reference`'s existing `CORRECTNESS_RTOL` doesn't hold + for the dequant path, that gets a stated, justified, and reported + tolerance change -- not a silently loosened one. +- **Named risk: per-channel scale storage/broadcast inside the Triton + kernel is new territory this project's kernels haven't needed before** + (naive/persistent both operate on a single dtype with no side-channel + per-tile metadata beyond the tile schedule already built). If this + proves harder than expected inside the tile loop, that is itself a + reportable finding, same as Phase 3's Fabric Manager discovery or + Phase 1's persistent-kernel null result -- not a reason to force a + workaround that muddies the kernel-level correctness gate. +- **GPU class**: RTX 3090 or L40 class (Phase 0/1's tier), marketplace/spot, + checked live against real-time availability at rental time, same + practice as every prior phase. +- **Session order**: CPU-testable suite green (quantization math, dataclass + shapes) -> rent -> kernel correctness gate (§5, kernel-level) -> only if + that passes, patch the real model and run the model-level agreement + check (§5) -> only if a patched count is nonzero, the three-way + benchmark and memory-footprint measurement -> tear down immediately. diff --git a/docs/findings/2026-09-17-phase-5a-community-attempt1-cost.md b/docs/findings/2026-09-17-phase-5a-community-attempt1-cost.md new file mode 100644 index 0000000..50a1ffa --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-community-attempt1-cost.md @@ -0,0 +1,7 @@ +# phase-5a-community-attempt1 -- GPU rental cost + +- Pod: `4u8rfiy5dz5vj6` (NVIDIA L40 (Community Cloud)) +- Rate: $0.7900/hr +- Duration: 125s (0.035hr) +- Cost: $0.0274 +- Note: Phase 5a attempt 1: RunPod Community Cloud L40 out of stock at create time, pod created then immediately unusable; abandoned. Duration/cost derived from RunPod billing API's exact recorded total ($0.0274) at this pod's quoted $0.79/hr rate. diff --git a/docs/findings/2026-09-17-phase-5a-community-attempt2-cost.md b/docs/findings/2026-09-17-phase-5a-community-attempt2-cost.md new file mode 100644 index 0000000..ae99722 --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-community-attempt2-cost.md @@ -0,0 +1,7 @@ +# phase-5a-community-attempt2 -- GPU rental cost + +- Pod: `wmv06rf1l2sm8s` (NVIDIA L40S (Community Cloud)) +- Rate: $0.7900/hr +- Duration: 1024s (0.284hr) +- Cost: $0.2247 +- Note: Phase 5a attempt 2: RunPod Community Cloud L40S. Host-level GPU passthrough bug -- cuInit() returned CUDA_ERROR_UNKNOWN (999) even against the base images own stock torch, confirmed via a raw ctypes test, reproduced on a second Community Cloud pod too. Abandoned Community Cloud entirely after this. Duration/cost derived from RunPod billing APIs exact recorded total ($0.2247) at this pods quoted $0.79/hr rate. diff --git a/docs/findings/2026-09-17-phase-5a-memory-footprint.json b/docs/findings/2026-09-17-phase-5a-memory-footprint.json new file mode 100644 index 0000000..e5f9e65 --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-memory-footprint.json @@ -0,0 +1,9 @@ +{ + "model": "deepseek-ai/deepseek-moe-16b-base", + "dtype": "bfloat16", + "device": "cuda", + "layers": 27, + "bf16_bytes": 29896998912, + "int8_bytes": 14982119424, + "reduction_pct": 49.89 +} diff --git a/docs/findings/2026-09-17-phase-5a-naive-results.json b/docs/findings/2026-09-17-phase-5a-naive-results.json new file mode 100644 index 0000000..3c9f51a --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-naive-results.json @@ -0,0 +1,33 @@ +{ + "model": "deepseek-ai/deepseek-moe-16b-base", + "device": "cuda", + "dtype": "bfloat16", + "moe_kernel": "naive", + "moe_layers_patched": 27, + "run_count": 15, + "mean_ttft": 0.3907895054658487, + "p50_ttft": 0.046567491997848265, + "p99_ttft": 4.490808108804046, + "mean_inter_token_latency": 0.04587934422397859, + "mean_tokens_per_second": 20.966010253117634, + "reference_comparison": { + "prompt_000_logits": { + "positions": 11, + "top1_agreement": 1.0, + "mutual_top_k": true, + "max_abs_diff": 0.78125 + }, + "prompt_001_logits": { + "positions": 11, + "top1_agreement": 1.0, + "mutual_top_k": true, + "max_abs_diff": 0.640625 + }, + "prompt_002_logits": { + "positions": 7, + "top1_agreement": 1.0, + "mutual_top_k": true, + "max_abs_diff": 1.4271240234375 + } + } +} diff --git a/docs/findings/2026-09-17-phase-5a-quantization-cost.md b/docs/findings/2026-09-17-phase-5a-quantization-cost.md new file mode 100644 index 0000000..e61d829 --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-quantization-cost.md @@ -0,0 +1,7 @@ +# phase-5a-quantization -- GPU rental cost + +- Pod: `924l6eft4d8251` (NVIDIA L40 (Secure Cloud)) +- Rate: $0.8200/hr +- Duration: 7473s (2.076hr) +- Cost: $1.7022 +- Note: Phase 5a: int8 quantized Triton kernel correctness gate (15/15 GPU tests) plus bf16 kernel re-verification (25/25), then the three-way measured run (stock/naive/quantized) and memory-footprint measurement on the real deepseek-ai/deepseek-moe-16b-base model. Includes a mid-session fix: patch_moe_infer_quantized was not freeing the original bf16 expert weights (stack_expert_weights re-points each Linear at a *view* into a shared bf16 tensor), so the quantized run OOMed on first attempt; fixed and rerun successfully. diff --git a/docs/findings/2026-09-17-phase-5a-quantization-run.md b/docs/findings/2026-09-17-phase-5a-quantization-run.md new file mode 100644 index 0000000..e609d14 --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-quantization-run.md @@ -0,0 +1,278 @@ +# Phase 5a int8 quantized kernel run -- findings + +Design: `docs/design/2026-09-16-phase-5a-quantization.md`. Plan: +`docs/plans/2026-09-16-phase-5a-quantization-plan.md`. Branch +`phase-5a-quantization`. + +## Infrastructure + +Three RunPod pods rented this session, in order: + +1. `4u8rfiy5dz5vj6` -- L40, Community Cloud. Out of stock the instant it + came up; unusable, abandoned. +2. `wmv06rf1l2sm8s` -- L40S, Community Cloud. Booted, but `cuInit()` (the + raw CUDA driver API, called directly via ctypes to rule out anything + about this project's own torch/venv choice) returned + `CUDA_ERROR_UNKNOWN` (999) -- even against the base image's own + preinstalled torch. A pod restart changed which `/dev/nvidiaN` device + node appeared but did not fix it. **A genuine RunPod Community Cloud + host-level GPU passthrough bug**, not a code or environment issue on + this project's side -- new finding this session, distinct from every + environment issue Phases 0/1/3/4 documented. +3. `924l6eft4d8251` -- L40, **Secure Cloud**. Same raw `cuInit()` test + returned `0` immediately. Everything below ran here. + +Abandoned Community Cloud entirely for this session after two pods hit +the same bug. Whether it's specific to this account, this data center, or +Community Cloud generally is unknown -- worth a quick stock/host check +before assuming Community Cloud is usable again in a later phase. + +## Kernel-level correctness gate: passed + +On the working Secure Cloud L40 (driver 580.178.04, CUDA 13.0): + +- `pytest -m gpu tests/unit/test_grouped_gemm_int8_kernel.py`: **15/15 + passed** -- the int8 kernel's first execution on real hardware. +- `pytest -m gpu tests/unit/test_grouped_gemm_kernel.py`: **25/25 + passed** -- the existing bf16 naive/persistent kernels, re-verified on + this card. +- `pytest -m "not gpu"` (full CPU suite): **118 passed** (later 119, then + 122, after the fixes below). + +**Not run this session, and worth saying plainly rather than leaving +implicit:** the design doc's own testing table calls for a mutation +check on the kernel (e.g. forcing every tile to read the wrong expert's +scale) turning the suite red, the same bar Phase 1 met and recorded +(`docs/STATUS.md`'s Phase 1 section). Phase 5a's GPU session ran out of +scope before this was attempted -- the 15/15 gate proves the kernel +matches its reference on real, varied per-channel-scale data, but it +does not independently prove the suite is capable of catching a broken +per-channel scale index the way Phase 1's explicit mutation run did. +Flagged by the final whole-branch review; deferred to a future GPU +session rather than invented after the fact. + +## A kernel-precision finding from Task 3's fix round, recorded here + +`grouped_matmul_int8`'s first implementation dequantized each int8 +weight tile to fp32 and cast the activation tile to fp32 on every +K-iteration before calling `tl.dot` -- which silently drops `tl.dot` +off the native bf16/fp16 tensor-core path onto the much slower fp32/TF32 +one. Caught by task review before this ever ran on real hardware: **a +kernel comparison that timed this would have measured "quantization is +slower," and the true cause would have been the dot's precision, not +quantization itself.** The design doc named exactly this risk in +advance ("if this proves harder than expected inside the tile loop, +that is itself a reportable finding"). + +Fixed by casting the int8 tile to the activation's own native dtype +before `tl.dot` and hoisting the per-output-channel scale multiply to +run once *after* the K-loop rather than once per iteration -- exact, +because the scale is K-invariant, and strictly more accurate than the +original (one rounding instead of `K / BLOCK_K` of them). No before/after +timing was measured on real hardware for this specific change in +isolation -- the fix landed before Task 6's rental, so the 20.72 tok/s +measured for the quantized kernel already reflects the corrected, +native-precision version; there is no "slow" data point to compare it +against without deliberately re-introducing the bug. + +## A real bug found mid-session, fixed before the measured run + +The first quantized-model run OOMed inside `patch_moe_infer_quantized` +itself (`stack_expert_weights` -> `torch.stack`), before any generation +happened -- a pure static-weight memory problem, not an activation or +batch-size one. + +**Root cause:** `stack_expert_weights` re-points every expert's +`nn.Linear.weight` at a *view* into one shared bf16 tensor rather than +copying it, specifically so building that stack costs no extra memory +(this is exactly what makes the existing bf16 `patch_moe_infer` path +memory-neutral). `patch_moe_infer_quantized` quantized that stack into +int8 but never dropped those views -- so the model held the bf16 +originals *and* their int8 copies simultaneously, for every layer, for +the rest of the process. Quantizing was making memory usage strictly +worse, not better, and it OOMed a 44GB L40 at DeepSeekMoE-16B's real +scale. + +This is a real defect in already-task-reviewed code (Task 4), found only +by a real hardware run at real model scale -- the same pattern Phases 3 +and 4 both hit (see their findings docs). Ruled and fixed directly rather +than re-dispatched through the SDD task loop, per that skill's own +authority for judgment calls mid-session: + +- Added `_free_expert_weights` (`src/dispatch/kernels/integration.py`): + replaces each expert's `gate_proj`/`up_proj`/`down_proj` `.weight` with + a 0-element `Parameter` immediately after quantizing, since the + quantized closure never reads `experts` again after patch time. +- Added `test_patch_quantized_frees_the_original_bf16_expert_weights` + (CPU-only): asserts `.weight.numel() == 0` on every expert/projection + after patching. +- Verified the existing independent-oracle test + (`test_patched_quantized_model_matches_weights_quantized_in_place`, + which exercises a full patched forward pass) still passes -- proof that + `experts` is genuinely unused after patching, not just unused in the + common case. +- `make check` green (119 passed, 2 skipped, 1 deselected; lint and + `mypy --strict` clean). Commit `c75da79`. + +Re-run after the fix: one transient `CUDACachingAllocator` OOM warning +(the allocator retrying, not failing) instead of a crash, and the run +completed cleanly. + +## Two more real bugs, found by the final whole-branch review + +The final review (broader and more architectural than the per-task +reviews above, dispatched after all seven tasks were otherwise +complete) found two real precision defects in `quantize_per_channel_int8` +that no per-task review or CPU test had caught, because the existing +round-trip test only ran fp32 and the zero-channel tests only checked +for NaN, never range utilization: + +1. **Quantization arithmetic ran in the weight's own dtype.** `absmax`, + `scale`, and the quotient were all computed in bf16/fp16 rather than + float32. Measured: bf16 widens round-trip error to up to **1.5x** the + ideal 0.5-quantization-step bound (bf16's 8-bit mantissa is the + cause). The measured run's own perfect model-level agreement shows + this was empirically tolerable at this model's real scale -- but + nothing in the suite would have noticed if it regressed further. +2. **The all-zero-channel guard clamped every small channel, not just + zero ones.** For fp16 specifically, the guard floored `scale` at + `torch.finfo(torch.float16).tiny` (~6.1e-5). Any channel with + `absmax < 127 x 6.1e-5 ≈ 0.00775` got its scale clamped *upward*, + silently under-using the int8 range -- demonstrated on a channel with + absmax ~0.001: ideal scale ~7.9e-6, clamped to ~6.1e-5, cutting the + largest value's int8 code from 127 down to 16 (a 2.4% relative error + where <0.4% was achievable). bf16 was immune (its own `tiny` is + ~1.18e-38), which is why the measured run above was unaffected -- + this is CLI-reachable today via `--dtype float16 --moe-kernel + quantized`. + +Fixed by computing `absmax`/`scale`/the quotient in float32 regardless +of the weight's own dtype, and changing the zero-channel guard to set +`scale = 1.0` only when `absmax == 0` exactly, rather than clamping +toward a dtype-dependent floor. Added a round-trip test parametrized +over float32/bfloat16/float16 (one bound that must now hold for all +three) and a regression test proving a small-but-nonzero fp16 channel +still uses the full int8 range. `make check` green (122 passed, 2 +skipped, 1 deselected). Commit `589834c`. + +**One caveat this doesn't resolve:** computing in float32 means peak +GPU memory during `patch_moe_infer_quantized` is not actually reduced +by quantizing -- `load_model` loads the full bf16 model, then each +layer transiently doubles to a float32 copy for the quantization math +before `_free_expert_weights` releases the bf16 original. The **49.89%** +figure below is the *steady-state* expert-weight footprint after +quantizing, not a claim about peak memory during the process that gets +there; a genuinely lower-peak path would need to load pre-quantized +weights directly rather than quantizing a fully-materialized bf16 model +in place. Out of scope for this phase; worth naming for whoever picks +up quantize-on-load next. + +## Three-way measured run + +`deepseek-ai/deepseek-moe-16b-base`, bf16, single NVIDIA L40 (Secure +Cloud), unbatched eager-mode decode, 15 runs per configuration (3 prompts +x 5 repetitions, 64 max new tokens), `transformers==4.57.6`, 27/27 MoE +layers patched for both kernel configurations: + +| Config | Tokens/sec | Mean TTFT | Cost / 1M tokens | +|---|---|---|---| +| Stock (unpatched) | 12.126 | 0.363s | $18.78 | +| Naive bf16 kernel | 20.966 | 0.391s | $10.86 | +| int8 quantized kernel | 20.718 | 0.392s | $10.99 | + +Naive bf16 kernel: **+72.9%** over stock -- consistent with Phase 1's +separately-measured +67.2% on the same GPU class, a cross-phase sanity +check. int8 quantized kernel: **+70.9%** over stock, but **-1.2%** vs. the +naive bf16 kernel it's built on top of -- essentially a throughput tie, +not a speedup from quantizing. This int8 kernel dequantizes per-channel +in-kernel and was built to match the naive kernel's own `tl.dot` +precision profile exactly (see Task 3's fix round); it was never expected +to be *faster* than bf16 on compute-bound decode -- weight-only int8 +saves memory bandwidth and footprint, not FLOPs, and this project's own +naive kernel already keeps every `tl.dot` at native tensor-core rate. Any +speedup from int8 specifically would need to come from a wider +batch/sequence-length regime where memory bandwidth, not compute, is the +bottleneck -- out of scope for this phase's single-request decode +benchmark. + +## Memory footprint + +Computed directly from the loaded model's real stacked expert weights +(`stack_expert_weights` + `quantize_stacked_weights` against all 27 MoE +layers, `stacked_weights_nbytes`/`quantized_stacked_weights_nbytes` -- +not derived from GPU memory-usage sampling): + +| | Bytes | | +|---|---|---| +| bf16 (all 27 layers' routed-expert weights) | 29,896,998,912 (27.84 GiB) | | +| int8 (data + per-channel fp32 scale) | 14,982,119,424 (13.95 GiB) | | +| **Reduction** | | **49.89%** | + +Just under a clean 50% because the int8 data itself *is* exactly half of +bf16's bytes, plus a small per-output-channel fp32 scale on top (one +`float32` per row rather than per element) -- expected, and consistent +with the design doc's stated shape. + +## Model-level agreement: quantized vs. naive + +Per the design doc's two-tier correctness bar, the quantized run was +compared against the **naive bf16 run's** reference logits (not stock's: +quantization-induced divergence vs. bf16 is expected, and the naive +kernel is this project's own closest apples-to-apples bf16 comparison +point), across all 3 prompts: + +| Prompt | Positions | `top1_agreement` | `mutual_top_k` | `max_abs_diff` | +|---|---|---|---|---| +| 0 | 11 | 1.0 | true | 2.125 | +| 1 | 11 | 1.0 | true | 1.90625 | +| 2 | 7 | 1.0 | true | 1.34375 | + +**Perfect top-1 agreement and mutual top-k membership at every tested +position** -- int8 weight-only quantization did not change which token +would be selected anywhere in this test, on this model, at this prompt +set. `max_abs_diff` is larger than naive-vs-stock's (0.64-1.43 in Phase +5a's own stock/naive comparison) as expected from quantization error, but +never large enough to move an argmax outside the top-k. This is a +stronger result than the design doc's own bar required (it explicitly +allowed for real divergence here) -- reported as measured, not +downgraded to match the expectation. + +## Environment + +- `transformers==5.17.0` (this repo's pinned floor) still breaks + DeepSeek's own remote-code model file the same two ways Phases 0/1/3/4 + documented: `is_torch_fx_available` removed from + `transformers.utils.import_utils`, and `DynamicCache.get_usable_length` + renamed to `get_seq_length`. **Still unresolved upstream** -- fourth + phase in a row needing the same fix: `uv pip install --reinstall + transformers==4.57.6` (not `uv run`, which re-syncs to the lockfile) + plus the same narrow, pod-local, never-committed `_patch_and_run.py` + monkeypatch. +- **New this session:** the pod's environment had + `HF_HUB_ENABLE_HF_TRANSFER=1` set with the `hf_transfer` package not + installed, which makes `huggingface_hub`'s downloader raise instead of + falling back. Fixed with `uv pip install hf_transfer` -- worth + installing proactively in a future phase's runbook rather than + discovering it mid-download. +- Direct SSH (the pod's own public IP/port) rejected this session's SSH + key even though it was in the pod's own registered-keys list; + RunPod's proxy (`ssh.runpod.io`) with the same key worked immediately. + Not investigated further -- the proxy's known quirks (forces an + interactive shell regardless of a trailing command; worked around with + stdin heredoc scripts ending cleanly on EOF) were sufficient for + everything this session needed. + +## Cost + +| Pod | GPU | Duration | Cost | +|---|---|---|---| +| `4u8rfiy5dz5vj6` | L40 (Community, out of stock) | 125s | $0.0274 | +| `wmv06rf1l2sm8s` | L40S (Community, host GPU bug) | 1024s | $0.2247 | +| `924l6eft4d8251` | L40 (Secure Cloud) | 7473s (2.08hr) | $1.7022 | +| **Total** | | | **$1.9543** | + +Against the **$5 cap**: **39.1% used**, all three pods (including both +abandoned attempts) included. Full records: +`docs/findings/2026-09-17-phase-5a-community-attempt1-cost.md`, +`2026-09-17-phase-5a-community-attempt2-cost.md`, +`2026-09-17-phase-5a-quantization-cost.md`. diff --git a/docs/findings/2026-09-17-phase-5a-quantized-results.json b/docs/findings/2026-09-17-phase-5a-quantized-results.json new file mode 100644 index 0000000..74dceb5 --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-quantized-results.json @@ -0,0 +1,33 @@ +{ + "model": "deepseek-ai/deepseek-moe-16b-base", + "device": "cuda", + "dtype": "bfloat16", + "moe_kernel": "quantized", + "moe_layers_patched": 27, + "run_count": 15, + "mean_ttft": 0.3922873901989078, + "p50_ttft": 0.045856651006033644, + "p99_ttft": 4.5111440703252415, + "mean_inter_token_latency": 0.04649968624976402, + "mean_tokens_per_second": 20.71831170686657, + "reference_comparison": { + "prompt_000_logits": { + "positions": 11, + "top1_agreement": 1.0, + "mutual_top_k": true, + "max_abs_diff": 2.125 + }, + "prompt_001_logits": { + "positions": 11, + "top1_agreement": 1.0, + "mutual_top_k": true, + "max_abs_diff": 1.90625 + }, + "prompt_002_logits": { + "positions": 7, + "top1_agreement": 1.0, + "mutual_top_k": true, + "max_abs_diff": 1.34375 + } + } +} diff --git a/docs/findings/2026-09-17-phase-5a-stock-results.json b/docs/findings/2026-09-17-phase-5a-stock-results.json new file mode 100644 index 0000000..fde999a --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-stock-results.json @@ -0,0 +1,14 @@ +{ + "model": "deepseek-ai/deepseek-moe-16b-base", + "device": "cuda", + "dtype": "bfloat16", + "moe_kernel": "none", + "moe_layers_patched": 0, + "run_count": 15, + "mean_ttft": 0.36332343839943254, + "p50_ttft": 0.26675163400068413, + "p99_ttft": 1.5793800269556222, + "mean_inter_token_latency": 0.07844627715449191, + "mean_tokens_per_second": 12.12603356567426, + "reference_comparison": {} +} diff --git a/docs/plans/2026-09-16-phase-5a-quantization-plan.md b/docs/plans/2026-09-16-phase-5a-quantization-plan.md new file mode 100644 index 0000000..27a1fac --- /dev/null +++ b/docs/plans/2026-09-16-phase-5a-quantization-plan.md @@ -0,0 +1,1627 @@ +# Phase 5a: int8 Quantized Grouped-GEMM Kernel Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Extend dispatch's custom Triton grouped-GEMM kernel to consume +int8 weight-only quantized expert weights, wire it through the existing +patch/CLI plumbing as a new path alongside the bf16 backends, and measure +its real memory/throughput/quality tradeoff against Phase 0's stock model +and Phase 1's bf16 naive kernel on the same hardware. + +**Architecture:** A fully parallel set of quantized types and functions +(`QuantizedTensor`, `QuantizedStackedExpertWeights`, +`grouped_moe_routed_quantized`, all new, in a new +`src/dispatch/kernels/quantization.py`) sits alongside the existing bf16 +contract (`StackedExpertWeights`, `grouped_moe_routed` in `moe_forward.py`) +without modifying it -- the same pattern Phase 3 used for +`patch_moe_infer_ep` next to `patch_moe_infer`. A new Triton kernel +(`grouped_matmul_int8`, in a new `grouped_gemm_int8.py`, built on the +naive kernel's grid/launch structure) dequantizes each int8 weight tile +against its per-output-channel scale inside the tile loop. `integration.py` +gains `patch_moe_infer_quantized`, sharing a small extracted +`_iter_validated_moe_layers` helper with the unmodified `patch_moe_infer`. +`scripts/run_baseline.py --moe-kernel quantized` routes through this new +path instead of the existing `backends.BACKENDS` registry. + +**Tech Stack:** Triton (Linux-only, lazily imported everywhere it's +touched, matching the existing convention). PyTorch. No new third-party +dependency -- quantization scales are computed in dispatch's own code, +not sourced from bitsandbytes/AWQ (design doc §2). + +**Spec:** `docs/design/2026-09-16-phase-5a-quantization.md` (all +sections). Also reused, unmodified: `src/dispatch/kernels/moe_forward.py` +(`StackedExpertWeights`, `stack_expert_weights`, `torch_grouped_matmul`, +`assert_matches_reference`), `src/dispatch/kernels/grouping.py` +(`group_tokens_by_expert`, `ungroup_and_combine`), +`src/dispatch/kernels/tile_schedule.py` (`TileSchedule`, +`build_tile_schedule`), `src/dispatch/kernels/reference_moe.py` +(`MoEConfig`, `ReferenceMoE`), `src/dispatch/benchmark/reference.py` +(`compare_top_k_agreement`, `capture_reference_logits`, `save_reference`, +`load_reference`), `scripts/gpu/provision.py` (`write_cost_record`). + +## Global Constraints + +- **Weight-only int8 (W8A16).** Activations never quantized. Only routed + expert weights already flowing through Phase 1's kernel (`gate_proj`, + `up_proj`, `down_proj`) are quantized -- shared experts, attention, the + gate, and `lm_head` stay bf16 (design doc §2). +- **Per-output-channel symmetric round-to-nearest scales, self-computed** + (`scale = max(abs(weight_row)) / 127`) -- no bitsandbytes/AWQ dependency + (design doc §2). +- **Built on the naive kernel's structure, not persistent** -- naive has + won every benchmark since Phase 1 (design doc §2). +- **The existing bf16 contract is not modified.** + `GroupedMatmul`/`StackedExpertWeights`/`grouped_moe_routed` in + `moe_forward.py`, and `BACKENDS`/`resolve_backend` in `backends.py`, + keep their exact current signatures and behavior. Every new quantized + type/function is additive, in a new module (design doc §3). +- **Two distinct correctness bars, never conflated (design doc §5):** + - *Kernel-level* (tight, same `assert_matches_reference` discipline as + every prior phase): the Triton int8 kernel vs. `torch_grouped_matmul_dequant` + fed the *same* already-quantized weights. A mutation must turn this red. + - *Model-level* (stated, expected to show real divergence): mutual + top-5/top-1 agreement (`compare_top_k_agreement`, unmodified) between + the int8-patched real model and Phase 1's bf16-naive-patched real + model. Reported as measured, never forced toward 100%. +- **Budget cap: $5, a ceiling not a target**, single GPU (RTX 3090/L40 + class), one combined session (correctness gate, then the measured run), + checked live against real-time marketplace/spot availability at rental + time. This is a live session with the user, not something to run + unattended -- get explicit go-ahead before renting (Task 6). +- **Never quote a benchmark number that wasn't measured** on this exact + run (repo-wide rule). +- **Out of scope for this plan** (design doc §7): int4, activation + quantization (W8A8), bitsandbytes/AWQ-sourced scales, a quantized + persistent kernel, combining with multi-GPU EP, speculative decoding + (Phase 5b, separate). + +--- + +### Task 1: Int8 per-channel weight quantization primitives + +**Files:** +- Create: `src/dispatch/kernels/quantization.py` +- Test: `tests/unit/test_quantization.py` + +**Interfaces:** +- Produces: `QuantizedTensor` (frozen dataclass: `data: torch.Tensor` + int8 shape `(E, N, K)`; `scale: torch.Tensor` float32 shape `(E, N)`). + `quantize_per_channel_int8(weight: torch.Tensor) -> QuantizedTensor`. + `dequantize_int8(qtensor: QuantizedTensor) -> torch.Tensor` (returns + float32). + +- [ ] **Step 1: Write the failing tests** + +Create `tests/unit/test_quantization.py`: + +```python +"""CPU-only proof of the int8 quantization math: per-channel scale +computation, round-trip error bounds, and the zero-channel edge case. +Per docs/design/2026-09-16-phase-5a-quantization.md section 2.""" + +from __future__ import annotations + +import torch + +from dispatch.kernels.quantization import dequantize_int8, quantize_per_channel_int8 + + +def test_quantize_dequantize_round_trip_is_within_one_quantization_step() -> None: + torch.manual_seed(0) + weight = torch.randn(3, 5, 7) * 10 + + quantized = quantize_per_channel_int8(weight) + dequantized = dequantize_int8(quantized) + + max_step = weight.abs().amax(dim=-1, keepdim=True) / 127 + assert torch.all((dequantized - weight).abs() <= max_step * 1.01) + + +def test_quantized_data_is_int8_and_within_range() -> None: + torch.manual_seed(0) + weight = torch.randn(2, 4, 6) * 100 + + quantized = quantize_per_channel_int8(weight) + + assert quantized.data.dtype == torch.int8 + assert quantized.scale.dtype == torch.float32 + assert quantized.data.abs().max() <= 127 + + +def test_quantized_shapes_match_weight() -> None: + weight = torch.randn(6, 8, 10) + + quantized = quantize_per_channel_int8(weight) + + assert quantized.data.shape == (6, 8, 10) + assert quantized.scale.shape == (6, 8) + + +def test_an_all_zero_channel_does_not_produce_nan() -> None: + weight = torch.zeros(1, 2, 4) + weight[0, 1] = torch.tensor([1.0, -2.0, 3.0, -4.0]) + + quantized = quantize_per_channel_int8(weight) + dequantized = dequantize_int8(quantized) + + assert torch.isfinite(dequantized).all() + assert torch.equal(dequantized[0, 0], torch.zeros(4)) + assert quantized.data[0, 0].abs().sum() == 0 + + +def test_scale_is_per_expert_per_output_channel() -> None: + weight = torch.zeros(2, 2, 3) + weight[0, 0] = torch.tensor([1.0, -1.0, 0.5]) + weight[0, 1] = torch.tensor([10.0, -10.0, 5.0]) + + quantized = quantize_per_channel_int8(weight) + + assert quantized.scale.shape == (2, 2) + assert quantized.scale[0, 1] > quantized.scale[0, 0] +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/unit/test_quantization.py -v` +Expected: FAIL (collection error) -- `dispatch.kernels.quantization` does not exist yet. + +- [ ] **Step 3: Write the minimal implementation** + +Create `src/dispatch/kernels/quantization.py`: + +```python +"""Weight-only int8 quantization for MoE expert weights: per-output-channel, +symmetric, round-to-nearest scales computed directly in dispatch's own +code -- not sourced from bitsandbytes/AWQ, per +docs/design/2026-09-16-phase-5a-quantization.md section 2. Activations are +never quantized; only expert weight tensors are.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +INT8_MAX = 127 + + +@dataclass(frozen=True) +class QuantizedTensor: + """An (E, N, K) weight tensor quantized to int8, one scale per expert + per output channel (E, N).""" + + data: torch.Tensor + scale: torch.Tensor + + +def quantize_per_channel_int8(weight: torch.Tensor) -> QuantizedTensor: + """weight: (..., N, K). scale[..., n] = max(abs(weight[..., n, :])) / 127. + An all-zero channel's scale is clamped away from zero so dividing by it + is a no-op (result: 0) rather than a NaN-producing divide-by-zero.""" + absmax = weight.detach().abs().amax(dim=-1) + scale = (absmax / INT8_MAX).clamp(min=torch.finfo(torch.float32).tiny) + quantized = (weight.detach() / scale.unsqueeze(-1)).round().clamp(-INT8_MAX, INT8_MAX) + return QuantizedTensor(data=quantized.to(torch.int8), scale=scale.to(torch.float32)) + + +def dequantize_int8(qtensor: QuantizedTensor) -> torch.Tensor: + return qtensor.data.to(torch.float32) * qtensor.scale.unsqueeze(-1) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `uv run pytest tests/unit/test_quantization.py -v` +Expected: 5 passed. + +- [ ] **Step 5: Lint and typecheck** + +Run: `make lint && make typecheck` +Expected: both clean. + +- [ ] **Step 6: Commit** + +```bash +git checkout -b phase-5a-quantization +git add src/dispatch/kernels/quantization.py tests/unit/test_quantization.py +git commit -m "feat: add int8 per-channel weight quantization primitives" +``` + +--- + +### Task 2: Quantized MoE forward path and memory-footprint measurement + +**Files:** +- Modify: `src/dispatch/kernels/quantization.py` +- Test: Modify `tests/unit/test_quantization.py` + +**Interfaces:** +- Consumes: `StackedExpertWeights`, `stack_expert_weights`, + `torch_grouped_matmul` (`moe_forward.py`); `group_tokens_by_expert`, + `ungroup_and_combine` (`grouping.py`); `TileSchedule`, + `build_tile_schedule` (`tile_schedule.py`); `MoEConfig`, `ReferenceMoE` + (`reference_moe.py`, tests only). `QuantizedTensor`, + `quantize_per_channel_int8`, `dequantize_int8` (Task 1). +- Produces: `QuantizedStackedExpertWeights` (dataclass: `gate, up, down: + QuantizedTensor`; `num_experts: int` property). + `QuantizedGroupedMatmul = Callable[[torch.Tensor, QuantizedTensor, + TileSchedule], torch.Tensor]`. + `quantize_stacked_weights(weights: StackedExpertWeights) -> + QuantizedStackedExpertWeights`. + `torch_grouped_matmul_dequant(x: torch.Tensor, qweight: QuantizedTensor, + schedule: TileSchedule) -> torch.Tensor`. + `grouped_moe_routed_quantized(x, topk_idx, topk_weight, weights: + QuantizedStackedExpertWeights, matmul: QuantizedGroupedMatmul, *, + block_m: int = 16) -> torch.Tensor`. + `stacked_weights_nbytes(weights: StackedExpertWeights) -> int`. + `quantized_stacked_weights_nbytes(weights: QuantizedStackedExpertWeights) + -> int`. + +- [ ] **Step 1: Write the failing tests** + +At the top of `tests/unit/test_quantization.py`, replace Task 1's single +`from dispatch.kernels.quantization import dequantize_int8, +quantize_per_channel_int8` line with: + +```python +from dispatch.kernels.moe_forward import ( + StackedExpertWeights, + grouped_moe_routed, + stack_expert_weights, + torch_grouped_matmul, +) +from dispatch.kernels.quantization import ( + dequantize_int8, + grouped_moe_routed_quantized, + quantize_per_channel_int8, + quantize_stacked_weights, + quantized_stacked_weights_nbytes, + stacked_weights_nbytes, + torch_grouped_matmul_dequant, +) +from dispatch.kernels.reference_moe import MoEConfig, ReferenceMoE +``` + +(All imports stay consolidated at the top of the file -- `ruff`'s `E402` +rule, enabled in this repo, rejects a module-level import placed after +existing function/class definitions.) + +Then append the following below Task 1's existing tests: + +```python +TOY_CONFIG = MoEConfig( + hidden_size=8, + moe_intermediate_size=16, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, +) + + +def test_grouped_moe_routed_quantized_matches_dequantized_weights_reference() -> None: + torch.manual_seed(0) + moe = ReferenceMoE(TOY_CONFIG) + hidden_states = torch.randn(7, TOY_CONFIG.hidden_size) + topk_idx, topk_weight = moe.route(hidden_states) + weights = stack_expert_weights(moe.experts) + quantized_weights = quantize_stacked_weights(weights) + + dequantized = StackedExpertWeights( + gate=dequantize_int8(quantized_weights.gate), + up=dequantize_int8(quantized_weights.up), + down=dequantize_int8(quantized_weights.down), + ) + expected = grouped_moe_routed( + hidden_states, topk_idx, topk_weight, dequantized, torch_grouped_matmul + ) + + actual = grouped_moe_routed_quantized( + hidden_states, topk_idx, topk_weight, quantized_weights, torch_grouped_matmul_dequant + ) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + + +def test_quantize_stacked_weights_shapes() -> None: + torch.manual_seed(0) + moe = ReferenceMoE(TOY_CONFIG) + weights = stack_expert_weights(moe.experts) + + quantized = quantize_stacked_weights(weights) + + assert quantized.num_experts == 4 + assert quantized.gate.data.shape == weights.gate.shape + assert quantized.down.data.shape == weights.down.shape + + +def test_quantized_stacked_weights_are_smaller_than_bf16() -> None: + torch.manual_seed(0) + moe = ReferenceMoE(TOY_CONFIG) + weights = stack_expert_weights(moe.experts) + bf16_weights = StackedExpertWeights( + gate=weights.gate.to(torch.bfloat16), + up=weights.up.to(torch.bfloat16), + down=weights.down.to(torch.bfloat16), + ) + quantized = quantize_stacked_weights(weights) + + bf16_bytes = stacked_weights_nbytes(bf16_weights) + int8_bytes = quantized_stacked_weights_nbytes(quantized) + + assert int8_bytes < bf16_bytes + + +def test_stacked_weights_nbytes_counts_every_projection() -> None: + weights = StackedExpertWeights( + gate=torch.zeros(2, 3, 4, dtype=torch.bfloat16), + up=torch.zeros(2, 3, 4, dtype=torch.bfloat16), + down=torch.zeros(2, 4, 3, dtype=torch.bfloat16), + ) + + assert stacked_weights_nbytes(weights) == 3 * (2 * 3 * 4 * 2) + + +def test_quantized_stacked_weights_nbytes_counts_data_and_scale() -> None: + from dispatch.kernels.quantization import QuantizedStackedExpertWeights, QuantizedTensor + + tensor = QuantizedTensor( + data=torch.zeros(2, 3, 4, dtype=torch.int8), scale=torch.zeros(2, 3, dtype=torch.float32) + ) + weights = QuantizedStackedExpertWeights(gate=tensor, up=tensor, down=tensor) + + expected = 3 * (2 * 3 * 4 * 1 + 2 * 3 * 4) + assert quantized_stacked_weights_nbytes(weights) == expected +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/unit/test_quantization.py -v` +Expected: FAIL -- `grouped_moe_routed_quantized`, `quantize_stacked_weights`, +etc. don't exist yet. + +- [ ] **Step 3: Write the minimal implementation** + +At the top of `src/dispatch/kernels/quantization.py`, replace Task 1's +import block with (all imports stay consolidated at the top -- `ruff`'s +`E402`, enabled in this repo, rejects a module-level import placed after +existing definitions): + +```python +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +import torch.nn.functional as F # noqa: N812 -- F is the universal PyTorch convention + +from dispatch.kernels.grouping import group_tokens_by_expert, ungroup_and_combine +from dispatch.kernels.moe_forward import StackedExpertWeights, torch_grouped_matmul +from dispatch.kernels.tile_schedule import TileSchedule, build_tile_schedule + +INT8_MAX = 127 +``` + +Then append the following below Task 1's existing `dequantize_int8`: + +```python +@dataclass(frozen=True) +class QuantizedStackedExpertWeights: + """Each projection's int8-quantized weights for every expert.""" + + gate: QuantizedTensor + up: QuantizedTensor + down: QuantizedTensor + + @property + def num_experts(self) -> int: + return int(self.gate.data.shape[0]) + + +QuantizedGroupedMatmul = Callable[[torch.Tensor, QuantizedTensor, TileSchedule], torch.Tensor] + + +def quantize_stacked_weights(weights: StackedExpertWeights) -> QuantizedStackedExpertWeights: + return QuantizedStackedExpertWeights( + gate=quantize_per_channel_int8(weights.gate), + up=quantize_per_channel_int8(weights.up), + down=quantize_per_channel_int8(weights.down), + ) + + +def torch_grouped_matmul_dequant( + x: torch.Tensor, qweight: QuantizedTensor, schedule: TileSchedule +) -> torch.Tensor: + """The Triton int8 kernel's correctness oracle: dequantizes qweight -- + the *same* already-quantized weights the kernel sees -- and runs the + same per-expert-slice matmul torch_grouped_matmul does. Not a + model-quality reference: both sides here use identical quantized + weights, so nothing should diverge beyond float precision.""" + return torch_grouped_matmul(x, dequantize_int8(qweight).to(x.dtype), schedule) + + +def grouped_moe_routed_quantized( + x: torch.Tensor, + topk_idx: torch.Tensor, + topk_weight: torch.Tensor, + weights: QuantizedStackedExpertWeights, + matmul: QuantizedGroupedMatmul, + *, + block_m: int = 16, +) -> torch.Tensor: + grouping = group_tokens_by_expert(topk_idx, topk_weight, weights.num_experts) + schedule = build_tile_schedule(grouping.group_sizes, block_m) + gathered = x[grouping.sorted_token_idx] + gate_out = matmul(gathered, weights.gate, schedule) + up_out = matmul(gathered, weights.up, schedule) + expert_out = matmul(F.silu(gate_out) * up_out, weights.down, schedule) + return ungroup_and_combine(expert_out, grouping, num_tokens=x.shape[0]) + + +def stacked_weights_nbytes(weights: StackedExpertWeights) -> int: + return sum(t.element_size() * t.nelement() for t in (weights.gate, weights.up, weights.down)) + + +def quantized_stacked_weights_nbytes(weights: QuantizedStackedExpertWeights) -> int: + total = 0 + for qtensor in (weights.gate, weights.up, weights.down): + total += qtensor.data.element_size() * qtensor.data.nelement() + total += qtensor.scale.element_size() * qtensor.scale.nelement() + return total +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `uv run pytest tests/unit/test_quantization.py -v` +Expected: 10 passed. + +- [ ] **Step 5: Lint and typecheck** + +Run: `make lint && make typecheck` +Expected: both clean. + +- [ ] **Step 6: Commit** + +```bash +git add src/dispatch/kernels/quantization.py tests/unit/test_quantization.py +git commit -m "feat: add quantized MoE forward path and memory-footprint measurement" +``` + +--- + +### Task 3: Triton int8 grouped-GEMM kernel + +**Files:** +- Create: `src/dispatch/kernels/grouped_gemm_int8.py` +- Modify: `src/dispatch/kernels/backends.py` +- Test: Create `tests/unit/test_grouped_gemm_int8_kernel.py` +- Test: Modify `tests/unit/test_backends.py` + +**Interfaces:** +- Consumes: `QuantizedTensor`, `quantize_per_channel_int8`, + `quantize_stacked_weights`, `torch_grouped_matmul_dequant`, + `grouped_moe_routed_quantized`, `QuantizedGroupedMatmul` (Tasks 1-2). + `TileSchedule`, `build_tile_schedule` (`tile_schedule.py`). + `DEFAULT_BLOCK_N`, `DEFAULT_BLOCK_K`, `MIN_BLOCK_M` + (`grouped_gemm.py`, reused unchanged). `assert_matches_reference` + (`moe_forward.py`). `MoEConfig`, `ReferenceMoE` (`reference_moe.py`, + tests only). `stack_expert_weights` (`moe_forward.py`, tests only). +- Produces: `grouped_matmul_int8(x: torch.Tensor, qweight: QuantizedTensor, + schedule: TileSchedule, *, block_n: int = DEFAULT_BLOCK_N, block_k: int + = DEFAULT_BLOCK_K) -> torch.Tensor`. + `resolve_quantized_backend() -> QuantizedGroupedMatmul` (`backends.py`). + +- [ ] **Step 1: Write the failing GPU-marked kernel tests** + +Create `tests/unit/test_grouped_gemm_int8_kernel.py`: + +```python +"""GPU-only correctness gate for the int8 kernel: it must match +torch_grouped_matmul_dequant fed the *same* already-quantized weights -- +the kernel-level bar in docs/design/2026-09-16-phase-5a-quantization.md +section 5. Excluded from CI by the `gpu` marker.""" + +from __future__ import annotations + +import pytest +import torch + +from dispatch.kernels.moe_forward import assert_matches_reference, stack_expert_weights +from dispatch.kernels.quantization import ( + grouped_moe_routed_quantized, + quantize_per_channel_int8, + quantize_stacked_weights, + torch_grouped_matmul_dequant, +) +from dispatch.kernels.reference_moe import MoEConfig, ReferenceMoE +from dispatch.kernels.tile_schedule import build_tile_schedule + +pytest.importorskip("triton", reason="triton ships Linux wheels only") + +from dispatch.kernels import grouped_gemm_int8 + +pytestmark = pytest.mark.gpu + +SM80 = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 0), + reason="bf16 tensor-core matmul needs compute capability 8.0+ (Ampere or newer)", +) +DTYPES = [torch.float16, pytest.param(torch.bfloat16, marks=SM80)] + +TOY_DIMS = MoEConfig( + hidden_size=8, + moe_intermediate_size=16, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, +) +REAL_DIMS = MoEConfig( + hidden_size=2048, + moe_intermediate_size=1408, + n_routed_experts=64, + n_shared_experts=2, + num_experts_per_tok=6, +) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize( + ("group_sizes", "n", "k"), + [ + ([5, 0, 3, 40], 24, 16), # toy dims, an empty expert, a multi-tile expert + ([1, 1, 1, 1, 1, 1], 1408, 2048), # decode-shaped, gate/up_proj dims + ([70, 0, 33, 129], 2048, 1408), # prefill-shaped, down_proj dims, partial tiles + ], +) +def test_int8_kernel_matches_dequant_reference( + dtype: torch.dtype, group_sizes: list[int], n: int, k: int +) -> None: + torch.manual_seed(0) + sizes = torch.tensor(group_sizes, device="cuda") + schedule = build_tile_schedule(sizes, block_m=16) + x = torch.randn(int(sizes.sum()), k, device="cuda", dtype=dtype) + weight = torch.randn(len(group_sizes), n, k, device="cuda", dtype=torch.float32) / k**0.5 + qweight = quantize_per_channel_int8(weight) + + actual = grouped_gemm_int8.grouped_matmul_int8(x, qweight, schedule) + + assert_matches_reference(actual, torch_grouped_matmul_dequant(x.float(), qweight, schedule)) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize( + ("config", "num_tokens"), [(TOY_DIMS, 37), (REAL_DIMS, 1), (REAL_DIMS, 256)] +) +def test_int8_kernel_backed_moe_layer_matches_dequant_reference( + dtype: torch.dtype, config: MoEConfig, num_tokens: int +) -> None: + torch.manual_seed(0) + moe = ReferenceMoE(config).to(device="cuda", dtype=dtype) + hidden_states = torch.randn(num_tokens, config.hidden_size, device="cuda", dtype=dtype) + topk_idx, topk_weight = moe.route(hidden_states) + weights = stack_expert_weights(moe.experts) + qweights = quantize_stacked_weights(weights) + + expected = grouped_moe_routed_quantized( + hidden_states, topk_idx, topk_weight, qweights, torch_grouped_matmul_dequant + ) + actual = grouped_moe_routed_quantized( + hidden_states, topk_idx, topk_weight, qweights, grouped_gemm_int8.grouped_matmul_int8 + ) + + assert_matches_reference(actual, expected) + + +def test_int8_kernel_rejects_a_schedule_below_tl_dots_minimum_block_m() -> None: + x = torch.randn(4, 16, device="cuda", dtype=torch.float16) + weight = torch.randn(1, 16, 16, device="cuda", dtype=torch.float32) + qweight = quantize_per_channel_int8(weight) + schedule = build_tile_schedule(torch.tensor([4], device="cuda"), block_m=4) + + with pytest.raises(ValueError, match="block_m"): + grouped_gemm_int8.grouped_matmul_int8(x, qweight, schedule) + + +def test_int8_kernel_rejects_non_cuda_tensors() -> None: + x = torch.randn(4, 16, dtype=torch.float16) + weight = torch.randn(1, 16, 16, dtype=torch.float32) + qweight = quantize_per_channel_int8(weight) + schedule = build_tile_schedule(torch.tensor([4]), block_m=16) + + with pytest.raises(ValueError, match="CUDA"): + grouped_gemm_int8.grouped_matmul_int8(x, qweight, schedule) +``` + +- [ ] **Step 2: Write the failing CPU-only backend-registry test** + +At the top of `tests/unit/test_backends.py`, change the existing +`from dispatch.kernels.backends import resolve_backend` line to: + +```python +from dispatch.kernels.backends import resolve_backend, resolve_quantized_backend +``` + +Then append the following test to the end of the file: + +```python +def test_quantized_backend_maps_to_grouped_matmul_int8(monkeypatch: pytest.MonkeyPatch) -> None: + stand_in = types.ModuleType("dispatch.kernels.grouped_gemm_int8") + stand_in.grouped_matmul_int8 = object() # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "dispatch.kernels.grouped_gemm_int8", stand_in) + monkeypatch.setattr(dispatch.kernels, "grouped_gemm_int8", stand_in, raising=False) + + assert resolve_quantized_backend() is stand_in.grouped_matmul_int8 +``` + +- [ ] **Step 3: Run tests to verify they fail** + +Run: `uv run pytest tests/unit/test_backends.py -v` +Expected: FAIL -- `resolve_quantized_backend` doesn't exist yet. + +Run (only if on a machine with triton and a CUDA device; otherwise this +suite is `gpu`-marked and skipped until Task 6's rental): +`uv run pytest tests/unit/test_grouped_gemm_int8_kernel.py -v -m gpu` +Expected: FAIL -- `dispatch.kernels.grouped_gemm_int8` does not exist yet. + +- [ ] **Step 4: Write the minimal implementation** + +Create `src/dispatch/kernels/grouped_gemm_int8.py`: + +```python +"""Triton int8 weight-only grouped-GEMM kernel: dequantizes each weight +tile against its per-output-channel scale inside the tile loop, then +accumulates exactly like grouped_gemm.py's naive kernel. Activations (x) +stay bf16/fp16 throughout -- only the weight side is ever int8. Held to +assert_matches_reference against torch_grouped_matmul_dequant fed the same +QuantizedTensor, per docs/design/2026-09-16-phase-5a-quantization.md +section 5. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from dispatch.kernels.grouped_gemm import DEFAULT_BLOCK_K, DEFAULT_BLOCK_N, MIN_BLOCK_M +from dispatch.kernels.quantization import QuantizedTensor +from dispatch.kernels.tile_schedule import TileSchedule + + +@triton.jit # type: ignore[untyped-decorator] +def _matmul_tile_int8( # type: ignore[no-untyped-def] # noqa: PLR0913 -- one pointer/stride per tensor, matching grouped_gemm.py's own kernel + x_ptr, + w_ptr, + scale_ptr, + out_ptr, + tile_expert_ptr, + tile_row_start_ptr, + tile_valid_rows_ptr, + m_tile, + n_tile, + n, + k, + stride_xm, + stride_xk, + stride_we, + stride_wn, + stride_wk, + stride_se, + stride_sn, + stride_om, + stride_on, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + expert_id = tl.load(tile_expert_ptr + m_tile) + row_start = tl.load(tile_row_start_ptr + m_tile) + valid_rows = tl.load(tile_valid_rows_ptr + m_tile) + + offs_m = tl.arange(0, BLOCK_M) + offs_n = n_tile * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + row_mask = offs_m[:, None] < valid_rows + col_mask_2d = offs_n[None, :] < n + col_mask_1d = offs_n < n + + x_ptrs = x_ptr + (row_start + offs_m)[:, None] * stride_xm + offs_k[None, :] * stride_xk + w_ptrs = ( + w_ptr + expert_id * stride_we + offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk + ) + scale_ptrs = scale_ptr + expert_id * stride_se + offs_n * stride_sn + scale_tile = tl.load(scale_ptrs, mask=col_mask_1d, other=1.0).to(tl.float32) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k_start in range(0, k, BLOCK_K): + k_mask = (k_start + offs_k) < k + x_tile = tl.load(x_ptrs, mask=row_mask & k_mask[None, :], other=0.0) + w_tile = tl.load(w_ptrs, mask=k_mask[:, None] & col_mask_2d, other=0) + w_dequant = w_tile.to(tl.float32) * scale_tile[None, :] + acc += tl.dot(x_tile.to(tl.float32), w_dequant) + x_ptrs += BLOCK_K * stride_xk + w_ptrs += BLOCK_K * stride_wk + + out_ptrs = out_ptr + (row_start + offs_m)[:, None] * stride_om + offs_n[None, :] * stride_on + tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=row_mask & col_mask_2d) + + +@triton.jit # type: ignore[untyped-decorator] +def _grouped_matmul_int8_kernel( # type: ignore[no-untyped-def] # noqa: PLR0913 + x_ptr, + w_ptr, + scale_ptr, + out_ptr, + tile_expert_ptr, + tile_row_start_ptr, + tile_valid_rows_ptr, + n, + k, + stride_xm, + stride_xk, + stride_we, + stride_wn, + stride_wk, + stride_se, + stride_sn, + stride_om, + stride_on, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + _matmul_tile_int8( + x_ptr, + w_ptr, + scale_ptr, + out_ptr, + tile_expert_ptr, + tile_row_start_ptr, + tile_valid_rows_ptr, + tl.program_id(axis=0), + tl.program_id(axis=1), + n, + k, + stride_xm, + stride_xk, + stride_we, + stride_wn, + stride_wk, + stride_se, + stride_sn, + stride_om, + stride_on, + BLOCK_M, + BLOCK_N, + BLOCK_K, + ) + + +def grouped_matmul_int8( + x: torch.Tensor, + qweight: QuantizedTensor, + schedule: TileSchedule, + *, + block_n: int = DEFAULT_BLOCK_N, + block_k: int = DEFAULT_BLOCK_K, +) -> torch.Tensor: + """One CTA per (m_tile, n_tile), dequantizing qweight's int8 tile + against its per-output-channel scale before accumulating -- the int8 + analog of grouped_gemm.grouped_matmul.""" + out = _validated_quantized_output(x, qweight, schedule) + if schedule.num_tiles == 0: + return out + n, k = qweight.data.shape[1], qweight.data.shape[2] + grid = (schedule.num_tiles, triton.cdiv(n, block_n)) + _grouped_matmul_int8_kernel[grid]( + x, + qweight.data, + qweight.scale, + out, + schedule.tile_expert, + schedule.tile_row_start, + schedule.tile_valid_rows, + n, + k, + *x.stride(), + *qweight.data.stride(), + *qweight.scale.stride(), + *out.stride(), + BLOCK_M=schedule.block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + ) + return out + + +def _validated_quantized_output( + x: torch.Tensor, qweight: QuantizedTensor, schedule: TileSchedule +) -> torch.Tensor: + if not (x.is_cuda and qweight.data.is_cuda and qweight.scale.is_cuda): + raise ValueError("the Triton int8 grouped-GEMM kernel needs CUDA tensors") + if qweight.data.dtype != torch.int8: + raise ValueError(f"expected int8 weight data, got {qweight.data.dtype}") + if qweight.scale.dtype != torch.float32: + raise ValueError(f"expected float32 scales, got {qweight.scale.dtype}") + if x.shape[1] != qweight.data.shape[2]: + raise ValueError(f"K mismatch: x has {x.shape[1]}, qweight has {qweight.data.shape[2]}") + if schedule.block_m < MIN_BLOCK_M: + raise ValueError( + f"schedule.block_m={schedule.block_m} is below tl.dot's minimum of {MIN_BLOCK_M}" + ) + return torch.empty((x.shape[0], qweight.data.shape[1]), device=x.device, dtype=x.dtype) +``` + +Add to `src/dispatch/kernels/backends.py`, after the existing imports: + +```python +from dispatch.kernels.quantization import QuantizedGroupedMatmul +``` + +And append at the end of the file: + +```python +def resolve_quantized_backend() -> QuantizedGroupedMatmul: + from dispatch.kernels import grouped_gemm_int8 # noqa: PLC0415 -- triton is Linux-only + + return grouped_gemm_int8.grouped_matmul_int8 +``` + +- [ ] **Step 5: Run tests to verify they pass** + +Run: `uv run pytest tests/unit/test_backends.py -v` +Expected: all passed. + +Run (GPU host only -- this is the `gpu`-marked suite this task's real +verification happens against; on a non-Linux/no-triton dev machine this +step is deferred to Task 6's rental): +`uv run pytest tests/unit/test_grouped_gemm_int8_kernel.py -v -m gpu -rs` +Expected: all passed, 0 skipped (on a compute-capability-8.0+ card). If +`assert_matches_reference`'s default tolerance does not hold here, do not +silently loosen it -- investigate why first (see design doc section 8's +named risk); any justified tolerance change must be a new, explicitly +named constant with a comment explaining the numerical reason, not a +tweaked parameter. + +- [ ] **Step 6: Lint and typecheck** + +Run: `make lint && make typecheck` +Expected: both clean. (`mypy` runs on `grouped_gemm_int8.py` even without +triton installed locally, same as it already does for `grouped_gemm.py` -- +triton ships type stubs.) + +- [ ] **Step 7: Commit** + +```bash +git add src/dispatch/kernels/grouped_gemm_int8.py src/dispatch/kernels/backends.py \ + tests/unit/test_grouped_gemm_int8_kernel.py tests/unit/test_backends.py +git commit -m "feat: add Triton int8 grouped-GEMM kernel and its backend resolver" +``` + +--- + +### Task 4: `patch_moe_infer_quantized` + +**Files:** +- Modify: `src/dispatch/kernels/integration.py` +- Test: Modify `tests/unit/test_integration.py` + +**Interfaces:** +- Consumes: `QuantizedStackedExpertWeights`, `QuantizedGroupedMatmul`, + `quantize_stacked_weights`, `grouped_moe_routed_quantized`, + `torch_grouped_matmul_dequant`, `dequantize_int8`, + `quantize_per_channel_int8` (Tasks 1-2). `stack_expert_weights` + (`moe_forward.py`, unchanged). +- Produces: `patch_moe_infer_quantized(model: torch.nn.Module, matmul: + QuantizedGroupedMatmul, *, block_m: int = 16) -> int` -- same contract + as the existing `patch_moe_infer`. +- Refactors (no behavior change): `_iter_validated_moe_layers(model: + torch.nn.Module) -> Iterator[tuple[torch.nn.Module, torch.nn.ModuleList]]`, + extracted from `patch_moe_infer`'s body and shared by both patch + functions. + +- [ ] **Step 1: Confirm the existing tests are green before refactoring** + +Run: `uv run pytest tests/unit/test_integration.py -v` +Expected: 3 passed (this is the baseline the refactor in Step 3 must not break). + +- [ ] **Step 2: Write the failing tests for the new function** + +At the top of `tests/unit/test_integration.py`: add `import copy` to the +stdlib import group (before `import pytest`), change the existing +`from dispatch.kernels.integration import patch_moe_infer` line to +`from dispatch.kernels.integration import patch_moe_infer, +patch_moe_infer_quantized`, and add a new import block for +`dispatch.kernels.quantization`: + +```python +import copy + +import pytest +import torch + +from dispatch.kernels.integration import patch_moe_infer, patch_moe_infer_quantized +from dispatch.kernels.moe_forward import torch_grouped_matmul +from dispatch.kernels.quantization import ( + dequantize_int8, + quantize_per_channel_int8, + torch_grouped_matmul_dequant, +) +from dispatch.kernels.reference_moe import MoEConfig, ReferenceMoE +``` + +Then append the following tests to the end of the file: + +```python +def test_patch_quantized_counts_only_moe_layers() -> None: + torch.manual_seed(0) + + assert patch_moe_infer_quantized(FakeModel(num_moe_layers=3), torch_grouped_matmul_dequant) == 3 + + +def test_patched_quantized_model_matches_weights_quantized_in_place() -> None: + """Proves patch_moe_infer_quantized's full wiring (quantize at patch + time, route through grouped_moe_routed_quantized) is mathematically + equivalent to independently replacing every expert Linear's weight + with its own quantize-then-dequantize round trip and running + DeepSeek's stock moe_infer -- an independent computation path, not a + call to any of the same helpers.""" + torch.manual_seed(0) + model = FakeModel(num_moe_layers=3) + hidden_states = torch.randn(9, TOY_CONFIG.hidden_size) + + expected_model = copy.deepcopy(model) + for layer in expected_model.moe_layers: + for expert in layer.experts: + for name in ("gate_proj", "up_proj", "down_proj"): + linear = getattr(expert, name) + requantized = dequantize_int8( + quantize_per_channel_int8(linear.weight.detach().unsqueeze(0)) + ).squeeze(0) + linear.weight = torch.nn.Parameter(requantized, requires_grad=False) + with torch.no_grad(): + expected = expected_model(hidden_states) + + patch_moe_infer_quantized(model, torch_grouped_matmul_dequant) + with torch.no_grad(): + actual = model(hidden_states) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + + +def test_patch_quantized_rejects_experts_that_are_not_a_module_list() -> None: + class Odd(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.experts = torch.nn.Linear(2, 2) + + def moe_infer(self) -> None: + raise NotImplementedError + + with pytest.raises(TypeError, match="ModuleList"): + patch_moe_infer_quantized(Odd(), torch_grouped_matmul_dequant) +``` + +- [ ] **Step 3: Run tests to verify they fail** + +Run: `uv run pytest tests/unit/test_integration.py -v` +Expected: FAIL -- `patch_moe_infer_quantized` doesn't exist yet. + +- [ ] **Step 4: Refactor `patch_moe_infer` and add `patch_moe_infer_quantized`** + +At the top of `src/dispatch/kernels/integration.py`, change +`from collections.abc import Callable` to +`from collections.abc import Callable, Iterator`, and add a new import +block right after the existing `from dispatch.kernels.moe_forward import +(...)` block: + +```python +from dispatch.kernels.quantization import ( + QuantizedGroupedMatmul, + QuantizedStackedExpertWeights, + grouped_moe_routed_quantized, + quantize_stacked_weights, +) +``` + +(All imports stay consolidated at the top of the file -- `ruff`'s `E402`, +enabled in this repo, rejects a module-level import placed after existing +definitions.) + +Then replace `patch_moe_infer` and everything below it with: + +```python +def _iter_validated_moe_layers( + model: torch.nn.Module, +) -> Iterator[tuple[torch.nn.Module, torch.nn.ModuleList]]: + for module in model.modules(): + if not hasattr(module, "moe_infer"): + continue + experts = module.experts + if not isinstance(experts, torch.nn.ModuleList): + raise TypeError( + f"expected {type(module).__name__}.experts to be nn.ModuleList, " + f"got {type(experts).__name__}" + ) + yield module, experts + + +def patch_moe_infer(model: torch.nn.Module, matmul: GroupedMatmul, *, block_m: int = 16) -> int: + """Returns how many MoE layers were patched, for the caller to check. + + Forces the model into eval mode first: DeepSeek's own DeepseekMoE.forward + only calls moe_infer on the not-self.training branch, so a model left in + training mode would report a nonzero patched count while never actually + executing the patched path -- silently comparing the stock model against + itself rather than against the kernel. + """ + model.eval() + patched = 0 + for module, experts in _iter_validated_moe_layers(model): + # An instance attribute shadowing the remote-code method; nn.Module's + # __setattr__ stub only admits Tensor | Module values. + module.moe_infer = _grouped_moe_infer( # type: ignore[assignment] + stack_expert_weights(experts), matmul, block_m + ) + patched += 1 + return patched + + +def patch_moe_infer_quantized( + model: torch.nn.Module, matmul: QuantizedGroupedMatmul, *, block_m: int = 16 +) -> int: + """Same contract as patch_moe_infer, but quantizes each layer's stacked + weights to int8 (docs/design/2026-09-16-phase-5a-quantization.md) once + at patch time, before building the per-layer closure -- not per + forward pass.""" + model.eval() + patched = 0 + for module, experts in _iter_validated_moe_layers(model): + quantized_weights = quantize_stacked_weights(stack_expert_weights(experts)) + module.moe_infer = _grouped_moe_infer_quantized( # type: ignore[assignment] + quantized_weights, matmul, block_m + ) + patched += 1 + return patched + + +def _grouped_moe_infer( + weights: StackedExpertWeights, matmul: GroupedMatmul, block_m: int +) -> MoEInfer: + @torch.no_grad() + def moe_infer( + x: torch.Tensor, flat_expert_indices: torch.Tensor, flat_expert_weights: torch.Tensor + ) -> torch.Tensor: + top_k = flat_expert_indices.numel() // x.shape[0] + return grouped_moe_routed( + x, + flat_expert_indices.view(-1, top_k), + flat_expert_weights.view(-1, top_k), + weights, + matmul, + block_m=block_m, + ) + + return moe_infer + + +def _grouped_moe_infer_quantized( + weights: QuantizedStackedExpertWeights, matmul: QuantizedGroupedMatmul, block_m: int +) -> MoEInfer: + @torch.no_grad() + def moe_infer( + x: torch.Tensor, flat_expert_indices: torch.Tensor, flat_expert_weights: torch.Tensor + ) -> torch.Tensor: + top_k = flat_expert_indices.numel() // x.shape[0] + return grouped_moe_routed_quantized( + x, + flat_expert_indices.view(-1, top_k), + flat_expert_weights.view(-1, top_k), + weights, + matmul, + block_m=block_m, + ) + + return moe_infer +``` + +(`Callable` was already imported at the top of the file for `MoEInfer`; +add `Iterator` to that same `from collections.abc import` line rather than +duplicating the import.) + +- [ ] **Step 5: Run tests to verify they pass** + +Run: `uv run pytest tests/unit/test_integration.py -v` +Expected: 6 passed (the original 3, still green after the refactor, plus +the 3 new ones). + +- [ ] **Step 6: Lint and typecheck** + +Run: `make lint && make typecheck` +Expected: both clean. + +- [ ] **Step 7: Commit** + +```bash +git add src/dispatch/kernels/integration.py tests/unit/test_integration.py +git commit -m "feat: add patch_moe_infer_quantized, sharing layer-iteration with patch_moe_infer" +``` + +--- + +### Task 5: CLI wiring -- `--moe-kernel quantized` + +**Files:** +- Modify: `scripts/run_baseline.py` +- Test: Modify `tests/unit/test_run_baseline.py` + +**Interfaces:** +- Consumes: `resolve_quantized_backend` (`backends.py`, Task 3). + `patch_moe_infer_quantized` (`integration.py`, Task 4). +- Produces: `run_baseline(..., moe_kernel="quantized")` routes through + `patch_moe_infer_quantized`/`resolve_quantized_backend` instead of + `patch_moe_infer`/`resolve_backend`. `--moe-kernel` CLI choices gain + `"quantized"`. + +- [ ] **Step 1: Write the failing tests** + +Add to `tests/unit/test_run_baseline.py`: + +```python +def test_run_baseline_routes_the_quantized_kernel_through_its_own_patch_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: dict[str, object] = {} + + def fake_resolve_quantized_backend() -> object: + calls["resolved"] = True + return "the-quantized-matmul" + + def fake_patch_moe_infer_quantized(model: object, matmul: object) -> int: + calls["patched_model"] = model + calls["patched_matmul"] = matmul + return 27 + + monkeypatch.setattr( + run_baseline_module, "load_model", lambda *a, **k: ("the-model", "the-tokenizer") + ) + monkeypatch.setattr( + run_baseline_module, "resolve_quantized_backend", fake_resolve_quantized_backend + ) + monkeypatch.setattr( + run_baseline_module, "patch_moe_infer_quantized", fake_patch_moe_infer_quantized + ) + monkeypatch.setattr(run_baseline_module, "generate_with_timings", lambda *a, **k: object()) + monkeypatch.setattr(run_baseline_module, "capture_reference_logits", lambda *a, **k: {}) + + _, _, moe_layers_patched = run_baseline_module.run_baseline( + "some/model", + device="cpu", + dtype=torch.float32, + trust_remote_code=False, + prompts=["hi"], + repetitions=1, + max_new_tokens=1, + moe_kernel="quantized", + ) + + assert moe_layers_patched == 27 + assert calls == { + "resolved": True, + "patched_model": "the-model", + "patched_matmul": "the-quantized-matmul", + } + + +def test_run_baseline_refuses_a_quantized_run_that_patches_nothing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + run_baseline_module, "load_model", lambda *a, **k: ("the-model", "the-tokenizer") + ) + monkeypatch.setattr(run_baseline_module, "resolve_quantized_backend", lambda: object()) + monkeypatch.setattr(run_baseline_module, "patch_moe_infer_quantized", lambda model, matmul: 0) + + with pytest.raises(RuntimeError, match="patched no MoE layers"): + run_baseline_module.run_baseline( + "some/model", + device="cpu", + dtype=torch.float32, + trust_remote_code=False, + prompts=["hi"], + repetitions=1, + max_new_tokens=1, + moe_kernel="quantized", + ) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run pytest tests/unit/test_run_baseline.py -v` +Expected: FAIL -- `run_baseline_module` has no `resolve_quantized_backend`/ +`patch_moe_infer_quantized` attribute to monkeypatch, and `"quantized"` +isn't a recognized `moe_kernel` value yet. + +- [ ] **Step 3: Write the minimal implementation** + +In `scripts/run_baseline.py`, change the imports: + +```python +from dispatch.kernels.backends import BACKENDS, resolve_backend, resolve_quantized_backend +from dispatch.kernels.integration import patch_moe_infer, patch_moe_infer_quantized +``` + +Replace the patch block inside `run_baseline`: + +```python + moe_layers_patched = 0 + if moe_kernel == "quantized": + moe_layers_patched = patch_moe_infer_quantized(model, resolve_quantized_backend()) + elif moe_kernel != "none": + moe_layers_patched = patch_moe_infer(model, resolve_backend(moe_kernel)) + if moe_kernel != "none" and moe_layers_patched == 0: + raise RuntimeError( + f"--moe-kernel {moe_kernel} patched no MoE layers: {model_name} has no " + "moe_infer to replace, so this run would time the stock model under a kernel's name" + ) +``` + +Change the `--moe-kernel` argument's choices: + +```python + parser.add_argument("--moe-kernel", default="none", choices=["none", *BACKENDS, "quantized"]) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `uv run pytest tests/unit/test_run_baseline.py -v` +Expected: all passed. + +- [ ] **Step 5: Run the full CPU-only suite** + +Run: `make test` +Expected: all pass, nothing else broken by this task's edits. + +- [ ] **Step 6: Lint and typecheck** + +Run: `make lint && make typecheck` +Expected: both clean. + +- [ ] **Step 7: Commit** + +```bash +git add scripts/run_baseline.py tests/unit/test_run_baseline.py +git commit -m "feat: wire --moe-kernel quantized through patch_moe_infer_quantized" +``` + +--- + +### Task 6: GPU rental runbook -- kernel correctness gate, three-way measured run + +**Files:** +- Create: `docs/runbooks/phase-5a-quantization.md` +- Creates live, on the pod / written by the CLIs (not pre-committed as + code, committed as evidence after the session): + `docs/findings/-phase-5a-stock-reference.safetensors`, + `docs/findings/-phase-5a-stock-results.json`, + `docs/findings/-phase-5a-naive-results.json`, + `docs/findings/-phase-5a-quantized-results.json`, + `docs/findings/-phase-5a-memory-footprint.json` +- Create (via `write_cost_record`, reused unmodified): + `docs/findings/-phase-5a-quantization-cost.md` + +**Budget cap: $5, a ceiling not a target -- be surgical (Global +Constraints). This is a live session with the user, not something to run +unattended -- get explicit go-ahead and confirm the cap before renting.** + +- [ ] **Step 1: Write the runbook** + +Create `docs/runbooks/phase-5a-quantization.md`: + +````markdown +# Runbook: Phase 5a int8 quantized kernel (rented GPU) + +One combined session -- kernel correctness, then the measured run -- +matching Phase 3/4's single-session shape rather than Phase 1's two +separate sessions, per docs/design/2026-09-16-phase-5a-quantization.md +section 8. Budget cap: $5, a ceiling not a target. + +1. Pick the cheapest Community Cloud card with compute capability 8.0+ + (Triton's own floor) -- RTX 3090/4090, L4, L40, A40, A100, checked + live against real-time availability, same practice every prior phase + used. +2. Create and wait, with a 50GB+ volume (Phase 0's real trap: the model + is 32.8GB and the container disk alone is often only 30GB): + + uv run python -m scripts.gpu.provision create --name dispatch-phase-5a \ + --gpu-type "" --image "" \ + --cloud COMMUNITY --disk-gb 60 + uv run python -m scripts.gpu.provision wait --pod-id + +3. Confirm the card: + + nvidia-smi --query-gpu=name,compute_cap --format=csv + +4. Transfer and set up, all of Phase 0/1's known environment traps in one + place -- **confirm live whether DeepSeek's `modeling_deepseek.py` still + needs these** (check + https://huggingface.co/deepseek-ai/deepseek-moe-16b-base/commits/main + since Phase 1's session) before assuming either is still required: + + git archive HEAD | ssh "mkdir -p dispatch && tar -x -C dispatch" + ssh + cd dispatch + export HF_HOME=/workspace/hf_cache # NOT the container disk + command -v uv || pip install uv + uv sync --all-extras --dev + uv pip install transformers==4.57.6 # only if step 4's live check still shows the break + + If the `get_usable_length` break (Phase 0/1) is still present, use the + same pod-local, never-committed wrapper Phase 1's Task 9 used: + + cat > _patch_and_run.py <<'EOF' + import sys + from transformers.cache_utils import DynamicCache + + + def _get_usable_length(self, new_seq_length=None, layer_idx=0): + return self.get_seq_length(layer_idx) + + + DynamicCache.get_usable_length = _get_usable_length + + from scripts.run_baseline import main + + main(sys.argv[1:]) + EOF + + From here on invoke `.venv/bin/python` directly (`uv run` re-syncs to + `uv.lock` on every call and silently undoes the `transformers` override). + If the wrapper isn't needed, invoke `scripts/run_baseline.py` directly + instead of `_patch_and_run.py` in every command below. + +5. Full CPU-testable suite, one more time, on the pod itself: + + .venv/bin/python -m pytest -m "not gpu" -v + + Expected: all pass (confirms the pod's environment doesn't disagree + with what already passed locally). + +6. **Kernel correctness gate -- must pass before any timing.** + + .venv/bin/python -m pytest -m gpu tests/unit/test_grouped_gemm_int8_kernel.py -v -rs + + Expected: all passed, 0 skipped (on a compute-capability-8.0+ card). + Also re-verify the existing bf16 kernels still pass on this card + (cheap, and rules out a hardware-specific surprise unrelated to this + phase's own new code): + + .venv/bin/python -m pytest -m gpu tests/unit/test_grouped_gemm_kernel.py -v -rs + + Anything red: stop. Nothing below is meaningful until this is green. + +7. **The stock run** -- this session's own reference, same + `DEFAULT_PROMPTS` every prior phase has used: + + DATE=$(date +%Y-%m-%d) + .venv/bin/python -m scripts.run_baseline --trust-remote-code \ + --run-label $DATE-phase-5a-stock + +8. **The bf16 naive-kernel run** -- Phase 1's own "before" number, + re-verified on this card, and the model-level agreement bar's + comparison point: + + .venv/bin/python -m scripts.run_baseline --trust-remote-code \ + --moe-kernel naive \ + --compare-reference docs/findings/$DATE-phase-5a-stock-reference.safetensors \ + --run-label $DATE-phase-5a-naive + + Must show `"moe_layers_patched": 27` and `mutual_top_k: true` against + the stock reference. + +9. **The int8 quantized run** -- the model-level correctness bar from + design doc section 5, compared against the bf16 naive run (not the + stock run: quantization-induced divergence is expected, and the naive + kernel is this project's own closest apples-to-apples bf16 comparison + point): + + .venv/bin/python -m scripts.run_baseline --trust-remote-code \ + --moe-kernel quantized \ + --compare-reference docs/findings/$DATE-phase-5a-naive-reference.safetensors \ + --run-label $DATE-phase-5a-quantized + + Must show `"moe_layers_patched": 27`. Record the measured + `mutual_top_k`/`top1_agreement` values from the results JSON exactly as + they come back -- this is expected to show some real divergence; do + not treat anything short of perfect agreement as a failure unless + `mutual_top_k` is false (an argmax landing outside the other side's + top-k, the same bar every prior phase has used to distinguish a real + bug from an acceptable near-tie flip). + +10. **Memory footprint.** From the repo root on the pod: + + .venv/bin/python <<'EOF' + import json + import torch + + from dispatch.benchmark.harness import load_model + from dispatch.kernels.moe_forward import stack_expert_weights + from dispatch.kernels.quantization import ( + quantize_stacked_weights, + quantized_stacked_weights_nbytes, + stacked_weights_nbytes, + ) + + model, _ = load_model( + "deepseek-ai/deepseek-moe-16b-base", + device="cuda", + dtype=torch.bfloat16, + trust_remote_code=True, + ) + + total_bf16 = 0 + total_int8 = 0 + layers = 0 + for module in model.modules(): + if not hasattr(module, "moe_infer"): + continue + weights = stack_expert_weights(module.experts) + total_bf16 += stacked_weights_nbytes(weights) + total_int8 += quantized_stacked_weights_nbytes(quantize_stacked_weights(weights)) + layers += 1 + + report = { + "layers": layers, + "bf16_bytes": total_bf16, + "int8_bytes": total_int8, + "reduction_pct": round(100 * (1 - total_int8 / total_bf16), 2), + } + print(json.dumps(report, indent=2)) + with open("docs/findings/memory-footprint.json", "w") as f: + json.dump(report, f, indent=2) + EOF + + mv docs/findings/memory-footprint.json docs/findings/$DATE-phase-5a-memory-footprint.json + +11. **Throughput comparison.** Read `mean_tokens_per_second` out of the + three results JSONs from steps 7-9 (`$DATE-phase-5a-{stock,naive,quantized}-results.json`) + -- no separate benchmark tool needed, `run_baseline`'s own harness + already measured it identically for all three. + +12. **Cost record and teardown:** + + .venv/bin/python -c " + from pathlib import Path + from scripts.gpu.provision import write_cost_record + write_cost_record( + Path('docs/findings'), + pod_id='', + gpu_type_id='', + cost_per_hour=, + duration_s=, + note='Phase 5a int8 quantized kernel: correctness gate + three-way measured run', + run_label='phase-5a-quantization', + ) + " + + Then tear the pod down immediately: + + uv run python -m scripts.gpu.provision delete --pod-id + uv run python -m scripts.gpu.provision get --pod-id # confirm TERMINATED +```` + +- [ ] **Step 2: Get explicit go-ahead and run the session** + +Confirm with the user: budget cap ($5), GPU class/price at real-time +availability, and that this is a live, watched session -- then execute +the runbook above. Copy back every result JSON and the cost/memory-footprint +files into `docs/findings/` on the local machine before tearing the pod down. + +- [ ] **Step 3: Commit the evidence** + +```bash +git add docs/runbooks/phase-5a-quantization.md docs/findings/*phase-5a* +git commit -m "docs: run Phase 5a GPU correctness gate and measured session" +``` + +--- + +### Task 7: Findings doc and STATUS.md + +**Files:** +- Create: `docs/findings/-phase-5a-quantization-run.md` +- Modify: `docs/STATUS.md` + +- [ ] **Step 1: Write the findings doc** + +Cover, in `docs/findings/-phase-5a-quantization-run.md`: whether the +kernel-level correctness gate passed (Task 6 step 6), and on which GPU; +the stock/naive/quantized throughput numbers from the three results +JSONs, each with its full config (model, dtype, hardware, prompts, +repetitions); the memory-footprint numbers (bytes and percent reduction) +from Task 6 step 10; the model-level agreement result (quantized vs. +naive) exactly as measured, including `top1_agreement` and whether +`mutual_top_k` held -- reported honestly even if agreement is well below +100%; cost per 1M generated tokens for all three configurations (same +`$/hr / (tokens/sec * 3600) * 1e6` computation Phase 1's findings doc +used); the actual GPU type, duration, and cost from Task 6's cost record, +**explicitly compared against the $5 cap**; and whether the +`transformers==4.57.6`/`get_usable_length` workarounds were still needed +or had been resolved upstream since Phase 1. If the kernel correctness +gate failed, or the budget ran out before all three runs completed, say +so plainly and report exactly what was measured, matching this project's +practice of writing down a null or partial result rather than a +flattering guess. + +- [ ] **Step 2: Update STATUS.md** + +Add a "## Phase 5a progress" section following the Phase 0-4 pattern: +design and plan links, task checklist, the kernel-level and model-level +correctness results, the throughput/memory numbers, and total GPU cost +against the $5 cap. Update "## Next step" to note Phase 5a is complete +and Phase 5b (speculative decoding) is next, not yet planned. + +- [ ] **Step 3: Commit** + +```bash +git add docs/findings/-phase-5a-quantization-run.md docs/STATUS.md +git commit -m "docs: record Phase 5a int8 quantized kernel outcome" +``` + +Then open the PR for the whole `phase-5a-quantization` branch, per this +repo's one-branch-per-phase convention (README/CLAUDE.md refresh first, +per the standing instruction to update them proactively at natural +stopping points). + +--- + +## Self-Review Notes + +- **Spec coverage:** design doc §2 (scope: weight-only int8, self-computed + scales, naive-kernel-based, single GPU, $5 cap) is a Global Constraint, + restated in Task 6's runbook. §3 (architecture: parallel quantized + types, new kernel file, `patch_moe_infer_quantized`, CLI routing) maps + to Tasks 1-5 one component per task, in dependency order (quantization + math before the forward path that uses it, before the kernel that's + pluggable into it, before the integration wiring, before the CLI). + §4 (data flow: quantize once at patch time, forward pass reuses + unchanged grouping/tiling, dequant inside the kernel) is implemented in + Task 4's `patch_moe_infer_quantized` (quantize once) and Task 3's kernel + (dequant inside the tile loop) respectively. §5 (two correctness bars) + is a Global Constraint and is concretely implemented: kernel-level in + Task 3's tests (`torch_grouped_matmul_dequant` reference), model-level + in Task 6 step 9 (naive-vs-quantized agreement, reported as measured). + §6 (testing table) maps one-to-one onto Tasks 1-5 (CPU rows) and Task 6 + (the two `gpu`/paid rows, plus the new memory-footprint measurement). + §7 (non-goals) is listed in Global Constraints so no task accidentally + reaches for bitsandbytes, int4, or the persistent kernel. §8 (risk/cost) + is the $5 cap (Global Constraints, Task 6) and the named tolerance/scale + risks (Task 3 step 5's explicit contingency instruction, not a silent + workaround). +- **Placeholder scan:** no task step describes an action without showing + the code; every test has real assertions; the GPU runbook's "confirm + live" instructions (transformers workaround, tolerance) are concrete + contingencies tied to a specific, already-documented prior finding + (Phase 0/1's real bugs), not vague "handle appropriately" language. +- **Type consistency:** `QuantizedTensor` (Task 1) is used identically by + every later task -- `.data` (int8, `(E,N,K)`), `.scale` (float32, + `(E,N)`) -- never renamed. `QuantizedGroupedMatmul`'s signature + (`Callable[[Tensor, QuantizedTensor, TileSchedule], Tensor]`, Task 2) is + what `grouped_matmul_int8` (Task 3), `torch_grouped_matmul_dequant` + (Task 2), and `patch_moe_infer_quantized`'s `matmul` parameter (Task 4) + all satisfy -- verified directly in Task 3's and Task 4's tests, which + pass `torch_grouped_matmul_dequant` and (once rented) + `grouped_matmul_int8` into the exact same call sites. `patch_moe_infer`'s + post-refactor behavior is pinned by Task 4 step 1 running its existing + tests *before* the refactor and step 5 confirming they're still green + *after* -- the same green-refactor-green discipline Phase 1's + `_require_same_keys` extraction used. +- **A decision the design doc left to this plan, resolved here:** where + the new Triton kernel lives (design doc §3 said "grouped_gemm.py or a + new sibling module, decided in the implementation plan"). This plan + puts it in a new `grouped_gemm_int8.py` (Task 3) rather than growing + `grouped_gemm.py` further, importing its tuning constants + (`DEFAULT_BLOCK_N`/`DEFAULT_BLOCK_K`/`MIN_BLOCK_M`) rather than + duplicating them -- keeps the existing file's two bf16 kernels and this + phase's one int8 kernel each in a file with one clear dtype contract, + per the design-for-isolation guidance every prior phase's plan has + followed. diff --git a/docs/runbooks/phase-5a-quantization.md b/docs/runbooks/phase-5a-quantization.md new file mode 100644 index 0000000..f4b5055 --- /dev/null +++ b/docs/runbooks/phase-5a-quantization.md @@ -0,0 +1,162 @@ +# Phase 5a runbook: int8 quantized kernel + +Budget cap: $5, a ceiling not a target. + +**What actually happened (2026-09-17), read this before following the +steps below as written:** `RUNPOD_API_KEY` wasn't set in the shell this +session, so the pod lifecycle ran through the RunPod MCP plugin's tools +(`create-pod`, `get-pod`, `delete-pod`, `list-pod-billing`) instead of +`scripts/gpu/provision.py`'s CLI -- same underlying RunPod API, different +caller. Two Community Cloud pods hit a real host-level GPU bug before a +Secure Cloud L40 worked. Total cost: **$1.9543** of the $5 cap. Full +account, including the correctness gate, the three-way measured run, and +a real bug found and fixed mid-session: +`docs/findings/2026-09-17-phase-5a-quantization-run.md`. + +1. **Create the pod.** Quoted L40 on Community Cloud at $0.69/hr, + confirmed with the user before creating anything. + + **What actually happened:** the first L40 (Community) create attempt + failed -- "no longer any instances available" -- and a retry failed + the same way. Fell back to L40S (Community), which succeeded at + $0.79/hr, but that pod's `cuInit()` (the raw CUDA driver API) returned + `CUDA_ERROR_UNKNOWN` (999) even against the base image's own + preinstalled torch -- confirmed via a standalone ctypes-level test + that doesn't touch this project's own torch/venv choice at all. A pod + restart changed which `/dev/nvidiaN` node appeared but did not fix + it. Reproduced identically on a second, separately-created Community + Cloud pod. **Abandoned Community Cloud entirely** and created an L40 + on **Secure Cloud** instead, at $0.82/hr -- the same ctypes test + returned `0` (success) immediately. Everything below ran on that pod + (`924l6eft4d8251`). + +2. **Transfer the repo.** No `RUNPOD_API_KEY` locally means no `scp` + integration either; this pod's SSH was proxy-only + (`ssh.runpod.io`), which forces an interactive shell regardless of a + trailing command and rejects non-PTY connections outright. Worked + around by piping whole scripts via stdin (`ssh -tt ... < script.sh`, + ending on EOF) for every remote command, and by transferring the repo + as a `git archive HEAD | base64` blob written into a heredoc script + (`base64 -d <<'B64EOF' ... | tar -xzf -`) rather than ever putting the + ~465KB blob directly in a shell command. + +3. **Reinstall `transformers==4.57.6` and apply the known monkeypatch.** + The pod's `transformers==5.17.0` (this repo's pinned floor) breaks + DeepSeek's own remote-code model file the same two ways Phases + 0/1/3/4 documented (`is_torch_fx_available` removed; + `DynamicCache.get_usable_length` renamed to `get_seq_length`) -- + fourth phase running into the same, still-unresolved-upstream break. + + ```bash + uv pip install --reinstall "transformers==4.57.6" + ``` + + (`uv run` re-syncs to the lockfile on every invocation and will + silently revert this -- call `.venv/bin/python` directly instead, same + as every prior phase learned.) Applied via a pod-local, never-committed + `_patch_and_run.py` that monkeypatches `DynamicCache.get_usable_length` + then calls `scripts.run_baseline.main()` -- this is a DeepSeek-repo + problem, not a `dispatch` one. + + **New this session:** the pod also had `HF_HUB_ENABLE_HF_TRANSFER=1` + set with the `hf_transfer` package not installed, which makes + `huggingface_hub`'s downloader raise instead of silently falling back. + Fixed with `uv pip install hf_transfer`. + +4. **Kernel correctness gate:** + + ```bash + .venv/bin/python -m pytest -m gpu tests/unit/test_grouped_gemm_int8_kernel.py -v + .venv/bin/python -m pytest -m gpu tests/unit/test_grouped_gemm_kernel.py -v + ``` + + **Result:** 15/15 and 25/25, both passed -- the int8 kernel's first + real-hardware run. + +5. **Three-way measured run**, each via the same `_patch_and_run.py` + wrapper: + + ```bash + .venv/bin/python _patch_and_run.py --trust-remote-code \ + --run-label $DATE-phase-5a-stock + .venv/bin/python _patch_and_run.py --trust-remote-code \ + --run-label $DATE-phase-5a-naive --moe-kernel naive \ + --compare-reference docs/findings/$DATE-phase-5a-stock-reference.safetensors + .venv/bin/python _patch_and_run.py --trust-remote-code \ + --run-label $DATE-phase-5a-quantized --moe-kernel quantized \ + --compare-reference docs/findings/$DATE-phase-5a-naive-reference.safetensors + ``` + + **What actually happened:** the quantized run OOMed on the first + attempt, before any generation -- a real bug in `patch_moe_infer_quantized` + (it quantized each layer's weights into int8 but never freed the bf16 + originals, so the model held both at once). Fixed + (`_free_expert_weights` in `src/dispatch/kernels/integration.py`, + commit `c75da79`), the fix pushed to the pod, and the quantized run + re-executed successfully. Full root-cause and fix details in the + findings doc. + +6. **Memory footprint** -- a standalone script loading the model fresh + and measuring `stack_expert_weights`/`quantize_stacked_weights`' + byte counts directly, not sampled from `nvidia-smi`: + + ```bash + .venv/bin/python <<'PYEOF' + from transformers.cache_utils import DynamicCache + DynamicCache.get_usable_length = lambda self, new_seq_length=None, layer_idx=0: self.get_seq_length(layer_idx) + + import json + import torch + from dispatch.benchmark.harness import load_model + from dispatch.kernels.moe_forward import stack_expert_weights + from dispatch.kernels.quantization import ( + quantize_stacked_weights, + quantized_stacked_weights_nbytes, + stacked_weights_nbytes, + ) + + model, _ = load_model( + "deepseek-ai/deepseek-moe-16b-base", device="cuda", + dtype=torch.bfloat16, trust_remote_code=True, + ) + total_bf16 = total_int8 = layers = 0 + for module in model.modules(): + if not hasattr(module, "moe_infer"): + continue + weights = stack_expert_weights(module.experts) + total_bf16 += stacked_weights_nbytes(weights) + total_int8 += quantized_stacked_weights_nbytes(quantize_stacked_weights(weights)) + layers += 1 + print(json.dumps({ + "model": "deepseek-ai/deepseek-moe-16b-base", "dtype": "bfloat16", "device": "cuda", + "layers": layers, "bf16_bytes": total_bf16, "int8_bytes": total_int8, + "reduction_pct": round(100 * (1 - total_int8 / total_bf16), 2), + }, indent=2)) + PYEOF + ``` + + **Result:** 27 layers, 27.84 GiB bf16 -> 13.95 GiB int8, **49.89%** + reduction. + +7. **Copy evidence back, record cost, tear down.** Result JSONs were + small enough to `cat` back over the same proxy SSH session rather + than needing a file-transfer tool; the `.safetensors` reference files + were not copied back (`.gitignore` excludes them repo-wide, same as + every prior phase -- they're regenerated by `capture_reference_logits` + rather than relied on as checked-in blobs). Cost recorded locally via + `write_cost_record` (a pure function needing no API key) once per pod + this session, including both abandoned Community Cloud attempts: + + ```python + from pathlib import Path + from scripts.gpu.provision import write_cost_record + write_cost_record( + Path("docs/findings"), pod_id="", gpu_type_id="", + cost_per_hour=, duration_s=, + note="...", run_label="