From dd5bb5d0e692c4e30568b4d5599167596f293140 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:37:03 -0400 Subject: [PATCH 01/18] docs: add phase 5a quantization design doc Splits the design doc's bundled Phase 5 (quantization + speculative decoding) into two sequential sub-phases per brainstorming decision. 5a scopes an int8 weight-only grouped-GEMM kernel extension; 5b (speculative decoding) is planned separately once 5a merges. --- .../2026-09-16-phase-5a-quantization.md | 222 ++++++++++++++++++ 1 file changed, 222 insertions(+) create mode 100644 docs/design/2026-09-16-phase-5a-quantization.md 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. From 0371da4b068ca9f297749f046ab7cdbde5bb78f5 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:50:22 -0400 Subject: [PATCH 02/18] docs: add phase 5a quantization implementation plan --- .../2026-09-16-phase-5a-quantization-plan.md | 1576 +++++++++++++++++ 1 file changed, 1576 insertions(+) create mode 100644 docs/plans/2026-09-16-phase-5a-quantization-plan.md 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..6a7bc8e --- /dev/null +++ b/docs/plans/2026-09-16-phase-5a-quantization-plan.md @@ -0,0 +1,1576 @@ +# 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** + +Append to `tests/unit/test_quantization.py`: + +```python +from dispatch.kernels.moe_forward import ( + StackedExpertWeights, + grouped_moe_routed, + stack_expert_weights, + torch_grouped_matmul, +) +from dispatch.kernels.quantization import ( + grouped_moe_routed_quantized, + quantize_stacked_weights, + quantized_stacked_weights_nbytes, + stacked_weights_nbytes, + torch_grouped_matmul_dequant, +) +from dispatch.kernels.reference_moe import MoEConfig, ReferenceMoE + +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** + +Append to `src/dispatch/kernels/quantization.py`: + +```python +from collections.abc import Callable + +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 + + +@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** + +Append to `tests/unit/test_backends.py`: + +```python +from dispatch.kernels.backends import resolve_quantized_backend + + +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** + +Add to `tests/unit/test_integration.py` (add `import copy` to the top): + +```python +import copy + +from dispatch.kernels.integration import patch_moe_infer, patch_moe_infer_quantized +from dispatch.kernels.quantization import dequantize_int8, quantize_per_channel_int8, torch_grouped_matmul_dequant + + +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`** + +Replace `src/dispatch/kernels/integration.py`'s `patch_moe_infer` function +and everything below it with: + +```python +from collections.abc import Callable, Iterator + +from dispatch.kernels.quantization import ( + QuantizedGroupedMatmul, + QuantizedStackedExpertWeights, + grouped_moe_routed_quantized, + quantize_stacked_weights, +) + + +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. From 49da8b75f329aa81a3e9be1482d3e387462d6fb8 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:51:14 -0400 Subject: [PATCH 03/18] style: ruff-format the plan doc's embedded code blocks --- docs/plans/2026-09-16-phase-5a-quantization-plan.md | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/plans/2026-09-16-phase-5a-quantization-plan.md b/docs/plans/2026-09-16-phase-5a-quantization-plan.md index 6a7bc8e..c6c0de3 100644 --- a/docs/plans/2026-09-16-phase-5a-quantization-plan.md +++ b/docs/plans/2026-09-16-phase-5a-quantization-plan.md @@ -904,15 +904,17 @@ Add to `tests/unit/test_integration.py` (add `import copy` to the top): import copy from dispatch.kernels.integration import patch_moe_infer, patch_moe_infer_quantized -from dispatch.kernels.quantization import dequantize_int8, quantize_per_channel_int8, torch_grouped_matmul_dequant +from dispatch.kernels.quantization import ( + dequantize_int8, + quantize_per_channel_int8, + torch_grouped_matmul_dequant, +) 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 - ) + 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: From 1acc7c01a64b1a9de351d600c87819f2151c40cb Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:55:24 -0400 Subject: [PATCH 04/18] docs: fix plan's import-placement instructions to avoid E402 --- .../2026-09-16-phase-5a-quantization-plan.md | 67 ++++++++++++++++--- 1 file changed, 58 insertions(+), 9 deletions(-) diff --git a/docs/plans/2026-09-16-phase-5a-quantization-plan.md b/docs/plans/2026-09-16-phase-5a-quantization-plan.md index c6c0de3..27a1fac 100644 --- a/docs/plans/2026-09-16-phase-5a-quantization-plan.md +++ b/docs/plans/2026-09-16-phase-5a-quantization-plan.md @@ -258,7 +258,9 @@ git commit -m "feat: add int8 per-channel weight quantization primitives" - [ ] **Step 1: Write the failing tests** -Append to `tests/unit/test_quantization.py`: +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 ( @@ -268,14 +270,24 @@ from dispatch.kernels.moe_forward import ( 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, @@ -368,18 +380,30 @@ etc. don't exist yet. - [ ] **Step 3: Write the minimal implementation** -Append to `src/dispatch/kernels/quantization.py`: +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.""" @@ -609,12 +633,16 @@ def test_int8_kernel_rejects_non_cuda_tensors() -> None: - [ ] **Step 2: Write the failing CPU-only backend-registry test** -Append to `tests/unit/test_backends.py`: +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_quantized_backend +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] @@ -898,19 +926,32 @@ 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** -Add to `tests/unit/test_integration.py` (add `import copy` to the top): +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) @@ -967,20 +1008,28 @@ Expected: FAIL -- `patch_moe_infer_quantized` doesn't exist yet. - [ ] **Step 4: Refactor `patch_moe_infer` and add `patch_moe_infer_quantized`** -Replace `src/dispatch/kernels/integration.py`'s `patch_moe_infer` function -and everything below it with: +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 collections.abc import Callable, Iterator - 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]]: From 4ce07101dc42b19d0af762f45d0c52595eccb501 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Wed, 16 Sep 2026 23:58:31 -0400 Subject: [PATCH 05/18] feat: add int8 per-channel weight quantization primitives --- src/dispatch/kernels/quantization.py | 36 ++++++++++++++++ tests/unit/test_quantization.py | 63 ++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 src/dispatch/kernels/quantization.py create mode 100644 tests/unit/test_quantization.py diff --git a/src/dispatch/kernels/quantization.py b/src/dispatch/kernels/quantization.py new file mode 100644 index 0000000..fcb3cef --- /dev/null +++ b/src/dispatch/kernels/quantization.py @@ -0,0 +1,36 @@ +"""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) diff --git a/tests/unit/test_quantization.py b/tests/unit/test_quantization.py new file mode 100644 index 0000000..d1a7ce7 --- /dev/null +++ b/tests/unit/test_quantization.py @@ -0,0 +1,63 @@ +"""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] From 1913993afe5fb77dfb51e978b4b4b18f8f9b4898 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Thu, 17 Sep 2026 00:07:09 -0400 Subject: [PATCH 06/18] fix: use input dtype for NaN-guard clamp floor in quantize_per_channel_int8 --- src/dispatch/kernels/quantization.py | 2 +- tests/unit/test_quantization.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/dispatch/kernels/quantization.py b/src/dispatch/kernels/quantization.py index fcb3cef..6fac49d 100644 --- a/src/dispatch/kernels/quantization.py +++ b/src/dispatch/kernels/quantization.py @@ -27,7 +27,7 @@ def quantize_per_channel_int8(weight: torch.Tensor) -> QuantizedTensor: 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) + scale = (absmax / INT8_MAX).clamp(min=torch.finfo(weight.dtype).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)) diff --git a/tests/unit/test_quantization.py b/tests/unit/test_quantization.py index d1a7ce7..aa42f1e 100644 --- a/tests/unit/test_quantization.py +++ b/tests/unit/test_quantization.py @@ -52,6 +52,18 @@ def test_an_all_zero_channel_does_not_produce_nan() -> None: assert quantized.data[0, 0].abs().sum() == 0 +def test_an_all_zero_channel_with_float16_does_not_produce_nan() -> None: + weight = torch.zeros(1, 2, 4, dtype=torch.float16) + weight[0, 1] = torch.tensor([1.0, -2.0, 3.0, -4.0], dtype=torch.float16) + + 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]) From c77e4a9f72359b9eaaa2b36be06a7bb6b366fdb2 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Thu, 17 Sep 2026 00:11:47 -0400 Subject: [PATCH 07/18] feat: add quantized MoE forward path and memory-footprint measurement --- src/dispatch/kernels/quantization.py | 71 +++++++++++++++++++ tests/unit/test_quantization.py | 101 ++++++++++++++++++++++++++- 2 files changed, 171 insertions(+), 1 deletion(-) diff --git a/src/dispatch/kernels/quantization.py b/src/dispatch/kernels/quantization.py index 6fac49d..f5e6444 100644 --- a/src/dispatch/kernels/quantization.py +++ b/src/dispatch/kernels/quantization.py @@ -6,9 +6,15 @@ 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 @@ -34,3 +40,68 @@ def quantize_per_channel_int8(weight: torch.Tensor) -> QuantizedTensor: def dequantize_int8(qtensor: QuantizedTensor) -> torch.Tensor: return qtensor.data.to(torch.float32) * qtensor.scale.unsqueeze(-1) + + +@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( # noqa: PLR0913 -- routing inputs plus a pluggable GEMM and tile size + 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 diff --git a/tests/unit/test_quantization.py b/tests/unit/test_quantization.py index aa42f1e..3c94654 100644 --- a/tests/unit/test_quantization.py +++ b/tests/unit/test_quantization.py @@ -6,7 +6,24 @@ import torch -from dispatch.kernels.quantization import dequantize_int8, quantize_per_channel_int8 +from dispatch.kernels.moe_forward import ( + StackedExpertWeights, + grouped_moe_routed, + stack_expert_weights, + torch_grouped_matmul, +) +from dispatch.kernels.quantization import ( + QuantizedStackedExpertWeights, + QuantizedTensor, + 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 def test_quantize_dequantize_round_trip_is_within_one_quantization_step() -> None: @@ -73,3 +90,85 @@ def test_scale_is_per_expert_per_output_channel() -> None: assert quantized.scale.shape == (2, 2) assert quantized.scale[0, 1] > quantized.scale[0, 0] + + +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: + 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 From b231aa810e2c9fdc461504e7654873bfaa055feb Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Thu, 17 Sep 2026 00:23:15 -0400 Subject: [PATCH 08/18] feat: add Triton int8 grouped-GEMM kernel and its backend resolver --- pyproject.toml | 1 + src/dispatch/kernels/backends.py | 7 + src/dispatch/kernels/grouped_gemm_int8.py | 182 ++++++++++++++++++++ tests/unit/test_backends.py | 11 +- tests/unit/test_grouped_gemm_int8_kernel.py | 114 ++++++++++++ 5 files changed, 314 insertions(+), 1 deletion(-) create mode 100644 src/dispatch/kernels/grouped_gemm_int8.py create mode 100644 tests/unit/test_grouped_gemm_int8_kernel.py diff --git a/pyproject.toml b/pyproject.toml index b28fa60..a4e5f95 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ max-complexity = 10 # Triton kernel signatures are pointer/stride/tile-size lists by construction, # and Triton's convention is UPPER_CASE for tl.constexpr arguments. "src/dispatch/kernels/grouped_gemm.py" = ["N803", "PLR0913", "PLR0917"] +"src/dispatch/kernels/grouped_gemm_int8.py" = ["N803", "PLR0913", "PLR0917"] # Expert-parallel sharding/dispatch functions carry routing inputs plus a # pluggable GEMM and tile size, same shape as moe_forward.py's own functions. "src/dispatch/kernels/expert_parallel.py" = ["PLR0913", "PLR0917"] diff --git a/src/dispatch/kernels/backends.py b/src/dispatch/kernels/backends.py index 7c37b3c..358fb8b 100644 --- a/src/dispatch/kernels/backends.py +++ b/src/dispatch/kernels/backends.py @@ -5,6 +5,7 @@ from __future__ import annotations from dispatch.kernels.moe_forward import GroupedMatmul, torch_grouped_matmul +from dispatch.kernels.quantization import QuantizedGroupedMatmul BACKENDS = ("torch", "naive", "persistent") @@ -19,3 +20,9 @@ def resolve_backend(name: str) -> GroupedMatmul: return ( grouped_gemm.grouped_matmul if name == "naive" else grouped_gemm.grouped_matmul_persistent ) + + +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 diff --git a/src/dispatch/kernels/grouped_gemm_int8.py b/src/dispatch/kernels/grouped_gemm_int8.py new file mode 100644 index 0000000..b1cd34e --- /dev/null +++ b/src/dispatch/kernels/grouped_gemm_int8.py @@ -0,0 +1,182 @@ +"""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] + 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] + 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) diff --git a/tests/unit/test_backends.py b/tests/unit/test_backends.py index 8aa0986..ef05329 100644 --- a/tests/unit/test_backends.py +++ b/tests/unit/test_backends.py @@ -9,7 +9,7 @@ import pytest import dispatch.kernels -from dispatch.kernels.backends import resolve_backend +from dispatch.kernels.backends import resolve_backend, resolve_quantized_backend from dispatch.kernels.moe_forward import torch_grouped_matmul @@ -31,3 +31,12 @@ def test_triton_backend_names_map_to_their_kernels(monkeypatch: pytest.MonkeyPat def test_unknown_backend_raises() -> None: with pytest.raises(ValueError, match="unknown backend"): resolve_backend("cutlass") + + +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 diff --git a/tests/unit/test_grouped_gemm_int8_kernel.py b/tests/unit/test_grouped_gemm_int8_kernel.py new file mode 100644 index 0000000..c81c2e2 --- /dev/null +++ b/tests/unit/test_grouped_gemm_int8_kernel.py @@ -0,0 +1,114 @@ +"""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) From c87c332b0b66bda967d853ffd5358ed3baec945d Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Thu, 17 Sep 2026 00:51:23 -0400 Subject: [PATCH 09/18] fix: keep int8 kernel's tl.dot on native tensor-core precision, validate scale shape Hoist the per-output-channel scale multiply outside the K-loop instead of dequantizing the weight tile every step: tl.dot now runs on x_tile's native bf16/fp16 dtype (int8 -> that dtype is exact) instead of forcing fp32/TF32-rate tl.dot, which would have made the kernel's throughput measurement conflate quantization's real speedup with a dot-precision regression. Also reject a QuantizedTensor whose scale shape doesn't match weight (E, N), since the kernel's masked scale load trusts data's N without checking scale's own shape. --- src/dispatch/kernels/grouped_gemm_int8.py | 29 +++++++++++++++------ tests/unit/test_grouped_gemm_int8_kernel.py | 12 +++++++++ 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/src/dispatch/kernels/grouped_gemm_int8.py b/src/dispatch/kernels/grouped_gemm_int8.py index b1cd34e..cdb841a 100644 --- a/src/dispatch/kernels/grouped_gemm_int8.py +++ b/src/dispatch/kernels/grouped_gemm_int8.py @@ -1,9 +1,14 @@ -"""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 +"""Triton int8 weight-only grouped-GEMM kernel: accumulates the raw +int8-times-activation dot product on native bf16/fp16 tensor cores -- +casting the int8 weight tile up to the activation's dtype, exact since +int8's range is exactly representable in both -- then applies the +per-output-channel scale once to the finished accumulator, exact since +that scale is invariant across K. This keeps tl.dot's precision profile +identical to grouped_gemm.py's bf16 naive kernel instead of falling back +to the slower fp32/TF32 path a per-K-step dequant would force. 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. """ @@ -67,10 +72,16 @@ def _matmul_tile_int8( # type: ignore[no-untyped-def] 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) + # int8 -> x_tile.dtype is exact, so casting the weight tile (rather + # than dequantizing per-K-step) keeps tl.dot on native bf16/fp16 + # tensor cores instead of the ~2x-slower fp32/TF32 path -- the same + # precision profile as grouped_gemm.py's bf16 kernel. The + # per-output-channel scale is invariant across K, so it's exact and + # mathematically equivalent to hoist it outside the loop below. + acc += tl.dot(x_tile, w_tile.to(x_tile.dtype)) x_ptrs += BLOCK_K * stride_xk w_ptrs += BLOCK_K * stride_wk + acc = acc * scale_tile[None, :] 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) @@ -173,6 +184,8 @@ def _validated_quantized_output( 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 qweight.scale.shape != qweight.data.shape[:2]: + raise ValueError(f"scale shape {tuple(qweight.scale.shape)} does not match weight (E, N)") 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: diff --git a/tests/unit/test_grouped_gemm_int8_kernel.py b/tests/unit/test_grouped_gemm_int8_kernel.py index c81c2e2..6b91f39 100644 --- a/tests/unit/test_grouped_gemm_int8_kernel.py +++ b/tests/unit/test_grouped_gemm_int8_kernel.py @@ -10,6 +10,7 @@ from dispatch.kernels.moe_forward import assert_matches_reference, stack_expert_weights from dispatch.kernels.quantization import ( + QuantizedTensor, grouped_moe_routed_quantized, quantize_per_channel_int8, quantize_stacked_weights, @@ -112,3 +113,14 @@ def test_int8_kernel_rejects_non_cuda_tensors() -> None: with pytest.raises(ValueError, match="CUDA"): grouped_gemm_int8.grouped_matmul_int8(x, qweight, schedule) + + +def test_int8_kernel_rejects_a_scale_shape_that_does_not_match_weight_e_n() -> 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) + mismatched = QuantizedTensor(data=qweight.data, scale=qweight.scale.unsqueeze(-1)) + schedule = build_tile_schedule(torch.tensor([4], device="cuda"), block_m=16) + + with pytest.raises(ValueError, match="scale shape"): + grouped_gemm_int8.grouped_matmul_int8(x, mismatched, schedule) From 20d61ffd2d37fb6fda5763c896c74a57abc207b6 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Thu, 17 Sep 2026 08:24:22 -0400 Subject: [PATCH 10/18] feat: add patch_moe_infer_quantized, sharing layer-iteration with patch_moe_infer --- src/dispatch/kernels/integration.py | 71 +++++++++++++++++++++++++---- tests/unit/test_integration.py | 59 +++++++++++++++++++++++- 2 files changed, 119 insertions(+), 11 deletions(-) diff --git a/src/dispatch/kernels/integration.py b/src/dispatch/kernels/integration.py index 9f10c7e..f7ab8da 100644 --- a/src/dispatch/kernels/integration.py +++ b/src/dispatch/kernels/integration.py @@ -6,7 +6,7 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Iterator import torch @@ -16,10 +16,31 @@ grouped_moe_routed, stack_expert_weights, ) +from dispatch.kernels.quantization import ( + QuantizedGroupedMatmul, + QuantizedStackedExpertWeights, + grouped_moe_routed_quantized, + quantize_stacked_weights, +) MoEInfer = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] +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. @@ -31,15 +52,7 @@ def patch_moe_infer(model: torch.nn.Module, matmul: GroupedMatmul, *, block_m: i """ model.eval() patched = 0 - 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__}" - ) + 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] @@ -49,6 +62,24 @@ def patch_moe_infer(model: torch.nn.Module, matmul: GroupedMatmul, *, block_m: i 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: @@ -67,3 +98,23 @@ def moe_infer( ) 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 diff --git a/tests/unit/test_integration.py b/tests/unit/test_integration.py index 2f4de06..16a3b9c 100644 --- a/tests/unit/test_integration.py +++ b/tests/unit/test_integration.py @@ -5,11 +5,19 @@ from __future__ import annotations +import copy +from typing import cast + import pytest import torch -from dispatch.kernels.integration import patch_moe_infer +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 TOY_CONFIG = MoEConfig( @@ -96,3 +104,52 @@ def moe_infer(self) -> None: with pytest.raises(TypeError, match="ModuleList"): patch_moe_infer(Odd(), torch_grouped_matmul) + + +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 cast(torch.nn.ModuleList, 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) From 0fe78bf312db44353a9d3492b22fa419efcf1030 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Thu, 17 Sep 2026 08:31:20 -0400 Subject: [PATCH 11/18] feat: wire --moe-kernel quantized through patch_moe_infer_quantized --- scripts/run_baseline.py | 20 +++++---- tests/unit/test_run_baseline.py | 77 +++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 9 deletions(-) diff --git a/scripts/run_baseline.py b/scripts/run_baseline.py index fbdf433..467746d 100644 --- a/scripts/run_baseline.py +++ b/scripts/run_baseline.py @@ -23,8 +23,8 @@ load_reference, save_reference, ) -from dispatch.kernels.backends import BACKENDS, resolve_backend -from dispatch.kernels.integration import patch_moe_infer +from dispatch.kernels.backends import BACKENDS, resolve_backend, resolve_quantized_backend +from dispatch.kernels.integration import patch_moe_infer, patch_moe_infer_quantized DEFAULT_PROMPTS = [ "The quick brown fox jumps over the lazy dog.", @@ -49,13 +49,15 @@ def run_baseline( # noqa: PLR0913 -- each of these is an independent, user-faci model_name, device=device, dtype=dtype, trust_remote_code=trust_remote_code ) moe_layers_patched = 0 - if moe_kernel != "none": + 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_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" - ) + 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" + ) runs = [ generate_with_timings( @@ -78,7 +80,7 @@ def main(argv: list[str] | None = None) -> None: parser.add_argument("--max-new-tokens", type=int, default=64) parser.add_argument("--output-dir", type=Path, default=Path("docs/findings")) parser.add_argument("--run-label", default=time.strftime("%Y-%m-%d-baseline")) - parser.add_argument("--moe-kernel", default="none", choices=["none", *BACKENDS]) + parser.add_argument("--moe-kernel", default="none", choices=["none", *BACKENDS, "quantized"]) parser.add_argument( "--compare-reference", type=Path, diff --git a/tests/unit/test_run_baseline.py b/tests/unit/test_run_baseline.py index 373b96d..66e18aa 100644 --- a/tests/unit/test_run_baseline.py +++ b/tests/unit/test_run_baseline.py @@ -152,3 +152,80 @@ def test_run_baseline_refuses_a_kernel_run_that_patches_nothing() -> None: max_new_tokens=3, moe_kernel="torch", ) + + +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: + def fake_resolve_quantized_backend() -> object: + return object() + + def fake_patch_moe_infer_quantized(model: object, matmul: object) -> int: + return 0 + + 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 + ) + + 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", + ) From c75da794fc71a39ad001a3d46bc6ddc23a311028 Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Thu, 17 Sep 2026 10:11:24 -0400 Subject: [PATCH 12/18] fix: free original bf16 expert weights in patch_moe_infer_quantized stack_expert_weights re-points every expert's Linear at a view into one shared bf16 tensor rather than copying it. patch_moe_infer_quantized quantized that stack but never dropped those views, so the model held both the bf16 originals and their int8 copies at once for every layer -- defeating quantization's whole point, and OOMing the real Phase 5a measured run on a 44GB L40 at DeepSeekMoE-16B's scale. --- src/dispatch/kernels/integration.py | 18 ++++++++++++++++++ tests/unit/test_integration.py | 15 +++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/src/dispatch/kernels/integration.py b/src/dispatch/kernels/integration.py index f7ab8da..f0211b4 100644 --- a/src/dispatch/kernels/integration.py +++ b/src/dispatch/kernels/integration.py @@ -73,6 +73,7 @@ def patch_moe_infer_quantized( patched = 0 for module, experts in _iter_validated_moe_layers(model): quantized_weights = quantize_stacked_weights(stack_expert_weights(experts)) + _free_expert_weights(experts) module.moe_infer = _grouped_moe_infer_quantized( # type: ignore[assignment] quantized_weights, matmul, block_m ) @@ -80,6 +81,23 @@ def patch_moe_infer_quantized( return patched +def _free_expert_weights(experts: torch.nn.ModuleList) -> None: + """Releases each expert's original bf16 weight now that moe_infer's + quantized closure never reads `experts` again. stack_expert_weights + re-points every expert's Linear at a *view* into one shared bf16 + tensor rather than copying it (so building that stack costs no extra + memory) -- but that means those views keep the whole bf16 tensor + resident for the model's lifetime unless something drops them. Left + alone, a quantized model holds its bf16 weights AND their int8 copies + at once, defeating the point of quantizing (and OOMing on real + hardware at DeepSeekMoE-16B's scale -- found running Phase 5a's + measured session).""" + for expert in experts: + for name in ("gate_proj", "up_proj", "down_proj"): + linear = getattr(expert, name) + linear.weight = torch.nn.Parameter(linear.weight.new_empty(0), requires_grad=False) + + def _grouped_moe_infer( weights: StackedExpertWeights, matmul: GroupedMatmul, block_m: int ) -> MoEInfer: diff --git a/tests/unit/test_integration.py b/tests/unit/test_integration.py index 16a3b9c..db6f5bb 100644 --- a/tests/unit/test_integration.py +++ b/tests/unit/test_integration.py @@ -142,6 +142,21 @@ def test_patched_quantized_model_matches_weights_quantized_in_place() -> None: torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) +def test_patch_quantized_frees_the_original_bf16_expert_weights() -> None: + """A quantized model must not keep both the bf16 originals and their + int8 copies resident -- that defeats quantizing at all, and OOMed a + real rented-GPU run at DeepSeekMoE-16B's scale (Phase 5a).""" + torch.manual_seed(0) + model = FakeModel(num_moe_layers=2) + + patch_moe_infer_quantized(model, torch_grouped_matmul_dequant) + + for layer in model.moe_layers: + for expert in cast(torch.nn.ModuleList, layer.experts): + for name in ("gate_proj", "up_proj", "down_proj"): + assert getattr(expert, name).weight.numel() == 0 + + def test_patch_quantized_rejects_experts_that_are_not_a_module_list() -> None: class Odd(torch.nn.Module): def __init__(self) -> None: From 60c2ebe7be08bff37306911f22883fd6ee91496e Mon Sep 17 00:00:00 2001 From: bsreecharanreddy <283869510+bsreecharanreddy@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:09:31 -0400 Subject: [PATCH 13/18] docs: record Phase 5a int8 quantized kernel outcome Kernel-level correctness gate passed on real hardware (15/15 int8, 25/25 bf16). Measured three-way run: naive bf16 kernel +72.9% over stock, int8 quantized +70.9% over stock but a throughput tie with naive (expected -- weight-only quantization saves memory, not FLOPs). Memory footprint measured directly: 49.89% reduction. Model-level agreement (quantized vs. naive) came back perfect at every tested position, stronger than the design doc's own bar required. Total cost $1.9543 of the $5 cap across all three pods, including two Community Cloud pods that hit a real host-level GPU passthrough bug. --- docs/STATUS.md | 79 ++++++- ...-09-17-phase-5a-community-attempt1-cost.md | 7 + ...-09-17-phase-5a-community-attempt2-cost.md | 7 + .../2026-09-17-phase-5a-memory-footprint.json | 6 + .../2026-09-17-phase-5a-naive-results.json | 33 +++ .../2026-09-17-phase-5a-quantization-cost.md | 7 + .../2026-09-17-phase-5a-quantization-run.md | 193 ++++++++++++++++++ ...2026-09-17-phase-5a-quantized-results.json | 33 +++ .../2026-09-17-phase-5a-stock-results.json | 14 ++ docs/runbooks/phase-5a-quantization.md | 161 +++++++++++++++ 10 files changed, 535 insertions(+), 5 deletions(-) create mode 100644 docs/findings/2026-09-17-phase-5a-community-attempt1-cost.md create mode 100644 docs/findings/2026-09-17-phase-5a-community-attempt2-cost.md create mode 100644 docs/findings/2026-09-17-phase-5a-memory-footprint.json create mode 100644 docs/findings/2026-09-17-phase-5a-naive-results.json create mode 100644 docs/findings/2026-09-17-phase-5a-quantization-cost.md create mode 100644 docs/findings/2026-09-17-phase-5a-quantization-run.md create mode 100644 docs/findings/2026-09-17-phase-5a-quantized-results.json create mode 100644 docs/findings/2026-09-17-phase-5a-stock-results.json create mode 100644 docs/runbooks/phase-5a-quantization.md 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/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..9cba908 --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-memory-footprint.json @@ -0,0 +1,6 @@ +{ + "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..6212bc0 --- /dev/null +++ b/docs/findings/2026-09-17-phase-5a-quantization-run.md @@ -0,0 +1,193 @@ +# 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 after + the fix below). + +## 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. + +## 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/runbooks/phase-5a-quantization.md b/docs/runbooks/phase-5a-quantization.md new file mode 100644 index 0000000..4a0cd6c --- /dev/null +++ b/docs/runbooks/phase-5a-quantization.md @@ -0,0 +1,161 @@ +# 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({ + "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="