Add Triton grouped-FP8 runtime and microbatch server

NOTICE.md

@@ -10,11 +10,17 @@ RivetCoder-9B-A4B combines the following sources:
   endpoint. No Qwen weights are included.
 - Architectural reference: `Akahsizrr/fuse-1-Lite`, revision
   `430c959e47556ae53fed18a9d97f7cf30876e6ff`.
+- Fast-serving kernel reference: `kernels-community/finegrained-fp8`, whose
+  Triton grouped-FP8 and activation-quantization kernels are Copyright 2026 The
+  Hugging Face Inc. team and licensed under Apache License 2.0. RivetCoder's
+  adaptation adds per-output scales, TorchAO expert-bank packing, deterministic
+  Top-4 reduction, and automatic Windows MSVC environment discovery.
 
 Modifications and new work include GLM expert selection, a fixed tied
 identity-Hadamard bridge, expert/router weight folding, per-layer Top-4 routing,
 token gating, bounded residual scaling, Qwen-supervised routing-control training,
-and the custom `fuse_glm` Transformers implementation shipped in this repository.
+the custom `fuse_glm` Transformers implementation, and the inference-only
+grouped-FP8 serving runtime shipped in this repository.
 
 See `provenance/` for the complete local build record with personal filesystem
 paths removed.

README.md

@@ -129,6 +129,61 @@ TorchAO 0.15.0 also reports that its optional C++ extensions are skipped with
 the tested PyTorch 2.12 build. The native CUDA FP8 path used by this checkpoint
 still passed quantization, serialization, clean reload, and forward validation.
 
+## Fast grouped-FP8 serving
+
+The repository includes an inference-only Triton runtime that replaces the
+Python 16-expert loop with two grouped FP8 GEMMs per fused layer: one combined
+gate/up projection and one down projection. It preserves Top-4 routing and the
+reference FP8 logits while releasing the unpacked expert tensors after runtime
+packing.
+
+For direct Transformers use, enable it after loading:
+
+```python
+runtime_report = model.enable_fast_fp8_serving()
+print(runtime_report)
+```
+
+For an OpenAI-compatible, queue-to-completion microbatch server:
+
+```powershell
+pip install -r requirements-serve.txt
+python serve.py `
+  --model HCHs/RivetCoder-9B-A4B-FP8 `
+  --no-local-files-only `
+  --host 0.0.0.0 `
+  --port 8000 `
+  --max-batch-size 16 `
+  --batch-wait-ms 3
+```
+
+On Windows, the first Triton JIT requires Visual Studio 2022 C++ Build Tools.
+The runtime automatically imports the installed Developer environment and
+records the resolved `cl.exe` path in its startup report. The first request for
+a new shape includes autotuning; later calls use the Triton cache.
+
+RTX 5070 Ti validation with a one-token full forward produced:
+
+| Runtime | Latency | Relative throughput |
+|---|---:|---:|
+| Original TorchAO expert loop | 5.100 s | 1.00x |
+| Grouped-FP8 fast path | 1.079 s | 4.73x |
+
+The logits were bit-exact (`MAE=0`, `max error=0`, identical top-1), repeated
+execution was deterministic, and resident VRAM was about 8.43 GiB. With the
+fast path enabled, fixed microbatch throughput scaled as follows:
+
+| Batch | Forward latency | Sequences/s |
+|---:|---:|---:|
+| 1 | 1.080 s | 0.93 |
+| 4 | 1.087 s | 3.68 |
+| 8 | 1.086 s | 7.36 |
+| 16 | 1.117 s | 14.33 |
+
+These are local full-forward measurements, not standardized generation
+benchmarks. Batch 16 increased throughput about 15.5x over batch 1 while adding
+only about 3.4% latency, which is why the bundled server defaults to batch 16.
+
 ## Limitations
 
 - This is an experimental fusion with only 60 routing-control optimizer steps.
@@ -138,13 +193,18 @@ still passed quantization, serialization, clean reload, and forward validation.
   with llama.cpp, LM Studio, or Ollama.
 - Hardware and software combinations other than the tested stack may need
   additional compatibility work.
+- The bundled server does not stream tokens and batches only requests with the
+  same generation parameters.
 
-See `provenance/quantization.json` for the local quantization and placement
-report. Architecture, source-model, expert-selection, and router-training
-provenance are retained from the BF16 repository.
+See `provenance/quantization.json` and `provenance/fast-serving.json` for the
+local quantization, placement, parity, and throughput reports. Architecture,
+source-model, expert-selection, and router-training provenance are retained
+from the BF16 repository.
 
 ## License and attribution
 
 The LFM host remains subject to the included LFM Open License v1.0. GLM-derived
-expert tensors retain the included MIT license and attribution. Review
-`LICENSE`, `NOTICE.md`, and `licenses/` before redistribution or deployment.
+expert tensors retain the included MIT license and attribution. The grouped FP8
+Triton kernels are adapted from Hugging Face's Apache-2.0
+`kernels-community/finegrained-fp8`. Review `LICENSE`, `NOTICE.md`, and
+`licenses/` before redistribution or deployment.

fast_fp8_runtime.py

@@ -0,0 +1,525 @@
+"""Inference-only grouped FP8 expert runtime for RivetCoder.
+
+The Triton kernels in this module are adapted from Hugging Face's
+``kernels-community/finegrained-fp8`` package (Apache-2.0).  The adaptation
+adds per-output weight scales so separately quantized gate/up projections can
+be concatenated without requantizing their checkpoint tensors.
+
+This module is intentionally imported lazily.  Normal BF16 loading, training,
+and CPU execution do not require Triton or TorchAO.
+"""
+
+from __future__ import annotations
+
+import gc
+import os
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def ensure_windows_msvc_environment() -> str | None:
+    """Populate the MSVC environment Triton's Windows launcher JIT needs."""
+
+    if os.name != "nt":
+        return None
+    configured = os.environ.get("CC")
+    if configured and (Path(configured).is_file() or shutil.which(configured)):
+        return configured
+    compiler = shutil.which("cl.exe")
+    if compiler:
+        os.environ["CC"] = compiler
+        return compiler
+
+    candidates = [
+        Path(os.environ.get("ProgramFiles", r"C:\Program Files"))
+        / "Microsoft Visual Studio"
+        / "2022"
+        / edition
+        / "Common7"
+        / "Tools"
+        / "VsDevCmd.bat"
+        for edition in ("Community", "Professional", "Enterprise", "BuildTools")
+    ]
+    vsdevcmd = next((path for path in candidates if path.is_file()), None)
+    if vsdevcmd is None:
+        raise RuntimeError(
+            "Triton needs an MSVC C compiler on Windows. Install Visual Studio 2022 "
+            "C++ Build Tools or launch the server from a Developer PowerShell."
+        )
+
+    command = f'call "{vsdevcmd}" -arch=x64 -host_arch=x64 >nul && set'
+    completed = subprocess.run(
+        command,
+        check=True,
+        capture_output=True,
+        text=True,
+        encoding="utf-8",
+        errors="replace",
+        shell=True,
+        executable=os.environ.get("COMSPEC", "cmd.exe"),
+    )
+    for line in completed.stdout.splitlines():
+        name, separator, value = line.partition("=")
+        if separator and name:
+            os.environ[name] = value
+    compiler = shutil.which("cl.exe")
+    if compiler is None:
+        raise RuntimeError("VsDevCmd completed but cl.exe is still unavailable")
+    os.environ["CC"] = compiler
+    return compiler
+
+
+def _load_triton() -> tuple[Any, Any, Any, Any]:
+    ensure_windows_msvc_environment()
+    try:
+        import triton
+        import triton.language as tl
+        from torch.library import triton_op, wrap_triton
+    except ImportError as error:
+        raise RuntimeError(
+            "Fast FP8 serving requires Triton. Use a PyTorch build that bundles Triton "
+            "or install a Windows-compatible Triton package."
+        ) from error
+    return triton, tl, triton_op, wrap_triton
+
+
+triton, tl, triton_op, wrap_triton = _load_triton()
+
+
+@triton.jit
+def _fp8_per_row_quant_kernel(x_ptr, q_ptr, scale_ptr, K: tl.constexpr):
+    row = tl.program_id(axis=0)
+    offsets = tl.arange(0, K)
+    values = tl.load(x_ptr + row * K + offsets).to(tl.float32)
+    scale = tl.maximum(tl.max(tl.abs(values), axis=0) / 448.0, 1.0e-12)
+    quantized = (values / scale).to(tl.float8e4nv)
+    tl.store(q_ptr + row * K + offsets, quantized)
+    tl.store(scale_ptr + row, scale)
+
+
+@triton_op("rivet_fp8::per_row_quant", mutates_args=())
+def _fp8_per_row_quant(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    if x.ndim != 2 or not x.is_contiguous():
+        raise ValueError("FP8 activation input must be a contiguous 2D tensor")
+    if x.shape[1] <= 0 or x.shape[1] & (x.shape[1] - 1):
+        raise ValueError("FP8 activation width must be a positive power of two")
+    quantized = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    scales = torch.empty(x.shape[0], device=x.device, dtype=torch.float32)
+    wrap_triton(_fp8_per_row_quant_kernel)[(x.shape[0],)](
+        x,
+        quantized,
+        scales,
+        K=x.shape[1],
+    )
+    return quantized, scales
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=warps, num_stages=stages)
+        for warps in (2, 4, 8, 16)
+        for stages in (2, 3, 4, 5)
+    ],
+    key=["N", "K", "BLOCK_M"],
+)
+@triton.jit
+def _grouped_fp8_linear_kernel(
+    A,
+    B,
+    C,
+    AScales,
+    BScales,
+    Offsets,
+    TileOffsets,
+    S,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    stride_am,
+    stride_ak,
+    stride_be,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_bs_e,
+    stride_bs_n,
+    NUM_EXPERTS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    SEARCH_STEPS: tl.constexpr,
+):
+    tile_m = tl.program_id(axis=0)
+    tile_n = tl.program_id(axis=1)
+    total_tiles = tl.load(TileOffsets + NUM_EXPERTS - 1)
+    if tile_m >= total_tiles:
+        return
+
+    low = 0
+    high = NUM_EXPERTS
+    for _ in tl.static_range(SEARCH_STEPS):
+        middle = (low + high) >> 1
+        middle_value = tl.load(TileOffsets + middle)
+        move_right = middle_value <= tile_m
+        low = tl.where(move_right, middle + 1, low)
+        high = tl.where(move_right, high, middle)
+    expert = low.to(tl.int64)
+
+    previous = tl.maximum(expert - 1, 0)
+    expert_start = tl.where(expert == 0, 0, tl.load(Offsets + previous))
+    expert_end = tl.load(Offsets + expert)
+    expert_rows = expert_end - expert_start
+    expert_tile_start = tl.where(expert == 0, 0, tl.load(TileOffsets + previous))
+    local_row_start = (tile_m - expert_tile_start) * BLOCK_M
+
+    row_offsets = local_row_start + tl.arange(0, BLOCK_M)
+    valid_rows = row_offsets < expert_rows
+    global_rows = expert_start + row_offsets
+    output_offsets = tile_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    k_offsets = tl.arange(0, BLOCK_K)
+
+    a_ptrs = A + global_rows[:, None] * stride_am + k_offsets[None, :] * stride_ak
+    b_ptrs = (
+        B
+        + expert * stride_be
+        + output_offsets[None, :] * stride_bn
+        + k_offsets[:, None] * stride_bk
+    )
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for _ in range(0, tl.cdiv(K, BLOCK_K)):
+        a = tl.load(a_ptrs, mask=valid_rows[:, None], other=0.0)
+        b = tl.load(b_ptrs)
+        accumulator += tl.dot(a, b)
+        a_ptrs += BLOCK_K * stride_ak
+        b_ptrs += BLOCK_K * stride_bk
+
+    activation_scale = tl.load(
+        AScales + global_rows,
+        mask=valid_rows,
+        other=0.0,
+    )
+    weight_scale = tl.load(
+        BScales + expert * stride_bs_e + output_offsets * stride_bs_n,
+    )
+    accumulator *= activation_scale[:, None] * weight_scale[None, :]
+
+    if C.dtype.element_ty == tl.bfloat16:
+        result = accumulator.to(tl.bfloat16)
+    elif C.dtype.element_ty == tl.float16:
+        result = accumulator.to(tl.float16)
+    else:
+        result = accumulator
+    c_ptrs = C + global_rows[:, None] * stride_cm + output_offsets[None, :] * stride_cn
+    tl.store(c_ptrs, result, mask=valid_rows[:, None])
+
+
+@triton_op("rivet_fp8::grouped_linear", mutates_args=())
+def _grouped_fp8_linear(
+    activations: torch.Tensor,
+    weights: torch.Tensor,
+    weight_scales: torch.Tensor,
+    offsets: torch.Tensor,
+    tokens_per_expert: torch.Tensor,
+) -> torch.Tensor:
+    if activations.ndim != 2 or not activations.is_contiguous():
+        raise ValueError("activations must be contiguous [routes, hidden]")
+    if weights.ndim != 3 or not weights.is_contiguous():
+        raise ValueError("weights must be contiguous [experts, output, hidden]")
+    if weights.dtype != torch.float8_e4m3fn:
+        raise TypeError("weights must use torch.float8_e4m3fn")
+    experts, output_size, hidden_size = weights.shape
+    if activations.shape[1] != hidden_size:
+        raise ValueError("activation/weight hidden dimensions do not match")
+    if output_size % 128 or hidden_size % 128:
+        raise ValueError("grouped FP8 output and hidden dimensions must be divisible by 128")
+    if weight_scales.shape != (experts, output_size):
+        raise ValueError("weight_scales must have shape [experts, output]")
+    if offsets.shape != (experts,) or tokens_per_expert.shape != (experts,):
+        raise ValueError("offset/count tensors must have one value per expert")
+
+    # TorchAO's checkpoint config uses dynamic PerTensor activation scaling.
+    # Match that per routed expert (rather than per row) so prefill follows the
+    # same quantization semantics as the original expert-by-expert calls.
+    expert_ids = torch.repeat_interleave(
+        torch.arange(experts, device=activations.device),
+        tokens_per_expert.to(torch.long),
+        output_size=activations.shape[0],
+    )
+    row_max = activations.abs().amax(dim=-1)
+    expert_max = torch.zeros(experts, device=activations.device, dtype=activations.dtype)
+    expert_max.scatter_reduce_(0, expert_ids, row_max, reduce="amax", include_self=True)
+    expert_scales = (expert_max / 448.0).float().clamp_min(1.0e-12)
+    activation_scales = expert_scales.index_select(0, expert_ids).contiguous()
+    quantized = (
+        activations.float()
+        .div(activation_scales.unsqueeze(-1))
+        .clamp(min=-448.0, max=448.0)
+        .to(torch.float8_e4m3fn)
+    )
+    output = activations.new_empty((activations.shape[0], output_size))
+    block_m = min(max(triton.next_power_of_2((activations.shape[0] + experts - 1) // experts), 16), 128)
+    tiles_per_expert = (tokens_per_expert + block_m - 1) // block_m
+    tile_offsets = torch.cumsum(tiles_per_expert, dim=0, dtype=torch.int32)
+    max_m_tiles = triton.cdiv(activations.shape[0], block_m) + experts
+    grid = (max_m_tiles, triton.cdiv(output_size, 128))
+    wrap_triton(_grouped_fp8_linear_kernel)[grid](
+        quantized,
+        weights,
+        output,
+        activation_scales,
+        weight_scales,
+        offsets,
+        tile_offsets,
+        activations.shape[0],
+        output_size,
+        hidden_size,
+        quantized.stride(0),
+        quantized.stride(1),
+        weights.stride(0),
+        weights.stride(2),
+        weights.stride(1),
+        output.stride(0),
+        output.stride(1),
+        weight_scales.stride(0),
+        weight_scales.stride(1),
+        NUM_EXPERTS=experts,
+        BLOCK_N=128,
+        BLOCK_K=128,
+        BLOCK_M=block_m,
+        SEARCH_STEPS=experts.bit_length(),
+    )
+    return output
+
+
+def grouped_fp8_linear(
+    activations: torch.Tensor,
+    weights: torch.Tensor,
+    weight_scales: torch.Tensor,
+    offsets: torch.Tensor,
+    tokens_per_expert: torch.Tensor,
+) -> torch.Tensor:
+    return torch.ops.rivet_fp8.grouped_linear(
+        activations,
+        weights,
+        weight_scales,
+        offsets,
+        tokens_per_expert,
+    )
+
+
+def _float8_parts(linear: nn.Linear) -> tuple[torch.Tensor, torch.Tensor]:
+    weight = linear.weight
+    qdata = getattr(weight, "qdata", None)
+    scale = getattr(weight, "scale", None)
+    if qdata is None or scale is None or "Float8" not in type(weight).__name__:
+        raise TypeError("fast serving requires TorchAO Float8Tensor Linear weights")
+    if qdata.dtype != torch.float8_e4m3fn or scale.numel() != 1:
+        raise TypeError("fast serving currently supports per-tensor E4M3 TorchAO weights")
+    return qdata, scale.reshape(())
+
+
+class PackedFp8ExpertBank(nn.Module):
+    """One layer's 16 experts packed into two grouped FP8 projections."""
+
+    def __init__(
+        self,
+        gate_up_qdata: torch.Tensor,
+        gate_up_scales: torch.Tensor,
+        down_qdata: torch.Tensor,
+        down_scales: torch.Tensor,
+        *,
+        gate_clamp_max: float,
+        up_clamp_min: float,
+        up_clamp_max: float,
+    ) -> None:
+        super().__init__()
+        self.register_buffer("gate_up_qdata", gate_up_qdata, persistent=False)
+        self.register_buffer("gate_up_scales", gate_up_scales, persistent=False)
+        self.register_buffer("down_qdata", down_qdata, persistent=False)
+        self.register_buffer("down_scales", down_scales, persistent=False)
+        self.num_experts = int(gate_up_qdata.shape[0])
+        self.intermediate_size = int(gate_up_qdata.shape[1] // 2)
+        self.hidden_size = int(gate_up_qdata.shape[2])
+        self.gate_clamp_max = float(gate_clamp_max)
+        self.up_clamp_min = float(up_clamp_min)
+        self.up_clamp_max = float(up_clamp_max)
+
+    @classmethod
+    @torch.no_grad()
+    def from_experts(
+        cls,
+        experts: nn.ModuleList,
+        *,
+        gate_clamp_max: float,
+        up_clamp_min: float,
+        up_clamp_max: float,
+    ) -> "PackedFp8ExpertBank":
+        if not experts:
+            raise ValueError("cannot pack an empty expert list")
+        first_gate, _ = _float8_parts(experts[0].gate_proj)
+        first_down, _ = _float8_parts(experts[0].down_proj)
+        num_experts = len(experts)
+        intermediate_size, hidden_size = first_gate.shape
+        if tuple(first_down.shape) != (hidden_size, intermediate_size):
+            raise ValueError("unexpected down projection shape")
+        device = first_gate.device
+        gate_up_qdata = torch.empty(
+            (num_experts, 2 * intermediate_size, hidden_size),
+            device=device,
+            dtype=torch.float8_e4m3fn,
+        )
+        gate_up_scales = torch.empty(
+            (num_experts, 2 * intermediate_size), device=device, dtype=torch.float32
+        )
+        down_qdata = torch.empty(
+            (num_experts, hidden_size, intermediate_size),
+            device=device,
+            dtype=torch.float8_e4m3fn,
+        )
+        down_scales = torch.empty((num_experts, hidden_size), device=device, dtype=torch.float32)
+
+        for index, expert in enumerate(experts):
+            gate_qdata, gate_scale = _float8_parts(expert.gate_proj)
+            up_qdata, up_scale = _float8_parts(expert.up_proj)
+            down_expert_qdata, down_scale = _float8_parts(expert.down_proj)
+            if tuple(gate_qdata.shape) != (intermediate_size, hidden_size):
+                raise ValueError("expert gate projection shapes are inconsistent")
+            if tuple(up_qdata.shape) != (intermediate_size, hidden_size):
+                raise ValueError("expert up projection shapes are inconsistent")
+            if tuple(down_expert_qdata.shape) != (hidden_size, intermediate_size):
+                raise ValueError("expert down projection shapes are inconsistent")
+            gate_up_qdata[index, :intermediate_size].copy_(gate_qdata)
+            gate_up_qdata[index, intermediate_size:].copy_(up_qdata)
+            gate_up_scales[index, :intermediate_size].copy_(gate_scale.expand(intermediate_size))
+            gate_up_scales[index, intermediate_size:].copy_(up_scale.expand(intermediate_size))
+            down_qdata[index].copy_(down_expert_qdata)
+            down_scales[index].copy_(down_scale.expand(hidden_size))
+
+        return cls(
+            gate_up_qdata,
+            gate_up_scales,
+            down_qdata,
+            down_scales,
+            gate_clamp_max=gate_clamp_max,
+            up_clamp_min=up_clamp_min,
+            up_clamp_max=up_clamp_max,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        selected_indices: torch.Tensor,
+        selected_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        token_count = hidden_states.shape[0]
+        route_experts = selected_indices.reshape(-1)
+        route_tokens = torch.arange(token_count, device=hidden_states.device).repeat_interleave(
+            selected_indices.shape[-1]
+        )
+        order = torch.argsort(route_experts, stable=True)
+        sorted_experts = route_experts.index_select(0, order)
+        sorted_tokens = route_tokens.index_select(0, order)
+        sorted_hidden = hidden_states.index_select(0, sorted_tokens).contiguous()
+        sorted_route_weights = selected_weights.reshape(-1).index_select(0, order)
+        tokens_per_expert = torch.bincount(
+            sorted_experts, minlength=self.num_experts
+        ).to(dtype=torch.int32)
+        offsets = torch.cumsum(tokens_per_expert, dim=0, dtype=torch.int32)
+
+        gate_up = grouped_fp8_linear(
+            sorted_hidden,
+            self.gate_up_qdata,
+            self.gate_up_scales,
+            offsets,
+            tokens_per_expert,
+        )
+        gate, up = gate_up.split(self.intermediate_size, dim=-1)
+        intermediate = F.silu(gate.clamp(max=self.gate_clamp_max)) * up.clamp(
+            min=self.up_clamp_min,
+            max=self.up_clamp_max,
+        )
+        routed = grouped_fp8_linear(
+            intermediate.contiguous(),
+            self.down_qdata,
+            self.down_scales,
+            offsets,
+            tokens_per_expert,
+        )
+        weighted = routed * sorted_route_weights.to(routed.dtype).unsqueeze(-1)
+        # ``order`` is a permutation, so restoring route order and reducing a
+        # contiguous [token, top_k, hidden] view avoids duplicate-index atomics.
+        # This makes decode deterministic and is friendlier to CUDA graphs.
+        inverse_order = torch.argsort(order)
+        route_outputs = weighted.index_select(0, inverse_order).reshape(
+            token_count, selected_indices.shape[-1], self.hidden_size
+        )
+        # The reference dispatcher visits experts in ascending expert-index
+        # order.  Preserve that BF16 accumulation order to minimize long-stack
+        # drift across 30 residual layers.
+        expert_order = torch.argsort(selected_indices, dim=-1, stable=True)
+        route_outputs = route_outputs.gather(
+            1,
+            expert_order.unsqueeze(-1).expand(-1, -1, self.hidden_size),
+        )
+        output = torch.zeros_like(hidden_states)
+        for route_slot in range(selected_indices.shape[-1]):
+            output = output + route_outputs[:, route_slot]
+        return output
+
+
+@torch.no_grad()
+def install_fast_fp8_runtime(model: nn.Module) -> dict[str, Any]:
+    """Pack every expert layer and enable the inference-only fast path.
+
+    The transformation releases the original per-expert modules to avoid
+    duplicating their FP8 storage.  It is intentionally one-way for the current
+    process; reload the checkpoint to recover trainable/module-list form.
+    """
+
+    if model.training:
+        raise RuntimeError("call model.eval() before enabling fast FP8 serving")
+    if not torch.cuda.is_available():
+        raise RuntimeError("fast FP8 serving requires CUDA")
+    ensure_windows_msvc_environment()
+
+    wrappers = tuple(model.fusion_layers())
+    packed_layers = 0
+    released_experts = 0
+    packed_bytes = 0
+    for wrapper in wrappers:
+        if getattr(wrapper, "fast_expert_bank", None) is not None:
+            continue
+        bank = PackedFp8ExpertBank.from_experts(
+            wrapper.experts,
+            gate_clamp_max=wrapper.experts[0].gate_clamp_max,
+            up_clamp_min=wrapper.experts[0].up_clamp_min,
+            up_clamp_max=wrapper.experts[0].up_clamp_max,
+        )
+        released_experts += len(wrapper.experts)
+        packed_bytes += sum(buffer.numel() * buffer.element_size() for buffer in bank.buffers())
+        wrapper.fast_expert_bank = bank
+        wrapper.experts = nn.ModuleList()
+        wrapper.serving_mode = True
+        wrapper.last_router_state = None
+        wrapper.last_router_diagnostics = None
+        packed_layers += 1
+    # TorchAO tensor subclasses can participate in reference cycles.  Explicit
+    # collection is necessary before the allocator can release the unpacked
+    # per-expert qdata that the packed banks replaced.
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    return {
+        "backend": "triton-grouped-fp8",
+        "packed_layers": packed_layers,
+        "released_experts": released_experts,
+        "packed_bytes": packed_bytes,
+        "cuda_allocated_bytes": torch.cuda.memory_allocated(),
+        "compiler": os.environ.get("CC"),
+    }

licenses/HF-FINEGRAINED-FP8-APACHE-2.0.txt

@@ -0,0 +1,203 @@
+Copyright 2018- The Hugging Face team. All rights reserved.
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

modeling_fuse_glm.py

@@ -214,6 +214,20 @@ class TopKFoldedExpertRouter(nn.Module):
         auxiliary_loss = self.num_experts * torch.sum(probability_fraction * token_fraction)
         return logits, selected_indices, selected_weights, auxiliary_loss
 
+    def forward_for_serving(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Route without training-only softmax, one-hot, or auxiliary loss work."""
+
+        logits = F.linear(hidden_states.float(), self.proj.weight.float())
+        scores = torch.sigmoid(logits)
+        selected_indices = torch.topk(
+            scores + self.e_score_correction_bias.float(), self.top_k, dim=-1
+        ).indices
+        selected_scores = scores.gather(-1, selected_indices)
+        selected_weights = selected_scores / selected_scores.sum(
+            dim=-1, keepdim=True
+        ).clamp_min(1e-12)
+        return selected_indices, selected_weights
+
 
 class FuseGlmFeedForward(nn.Module):
     """Preserve the native LFM FFN and add a sparse expert sidecar in parallel.
@@ -279,6 +293,8 @@ class FuseGlmFeedForward(nn.Module):
         self.raw_residual_scale = nn.Parameter(torch.zeros(()))
         self.last_router_state: RouterState | None = None
         self.last_router_diagnostics: RouterDiagnostics | None = None
+        self.fast_expert_bank: nn.Module | None = None
+        self.serving_mode = False
 
     @property
     def residual_scale(self) -> torch.Tensor:
@@ -309,6 +325,8 @@ class FuseGlmFeedForward(nn.Module):
         selected_indices: torch.Tensor,
         selected_weights: torch.Tensor,
     ) -> torch.Tensor:
+        if self.fast_expert_bank is not None:
+            return self.fast_expert_bank(hidden_states, selected_indices, selected_weights)
         output = torch.zeros_like(hidden_states)
         for expert_index, expert in enumerate(self.experts):
             token_positions, route_slots = torch.where(selected_indices == expert_index)
@@ -320,6 +338,27 @@ class FuseGlmFeedForward(nn.Module):
             output = output.index_add(0, token_positions, expert_output * route_weight)
         return output
 
+    def _forward_for_serving(
+        self, hidden_states: torch.Tensor, base_output: torch.Tensor
+    ) -> torch.Tensor:
+        """Inference fast path without router state or diagnostic construction."""
+
+        original_shape = hidden_states.shape
+        flat_hidden = hidden_states.reshape(-1, original_shape[-1])
+        token_gate = torch.sigmoid(
+            F.linear(
+                flat_hidden.float(),
+                self.token_gate.weight.float(),
+                self.token_gate.bias.float(),
+            )
+        ).to(hidden_states.dtype)
+        selected_indices, selected_weights = self.router.forward_for_serving(flat_hidden)
+        expert_delta = self._dispatch(flat_hidden, selected_indices, selected_weights)
+        expert_delta = expert_delta * token_gate
+        return base_output + self.residual_scale.to(expert_delta.dtype) * expert_delta.reshape(
+            original_shape
+        )
+
     def _make_diagnostics(
         self,
         state: RouterState,
@@ -361,6 +400,8 @@ class FuseGlmFeedForward(nn.Module):
             self.last_router_state = None
             self.last_router_diagnostics = self._disabled_diagnostics(hidden_states)
             return base_output
+        if self.serving_mode and not self.training:
+            return self._forward_for_serving(hidden_states, base_output)
 
         original_shape = hidden_states.shape
         flat_hidden = hidden_states.reshape(-1, original_shape[-1])
@@ -551,6 +592,14 @@ class FuseGlmForCausalLM(Lfm2ForCausalLM):
             wrapper.coding_enabled = bool(enabled)
         return self
 
+    def enable_fast_fp8_serving(self) -> dict[str, Any]:
+        """Pack TorchAO experts into the Triton grouped-FP8 serving runtime."""
+
+        from .fast_fp8_runtime import install_fast_fp8_runtime
+
+        self.eval()
+        return install_fast_fp8_runtime(self)
+
     @property
     def coding_enabled(self) -> bool:
         layers = self.fusion_layers()

provenance/checksums.sha256

@@ -3,9 +3,11 @@
 f3e412383b66fc807c3d0fbe9e6f8a4a58b2638df55d8380e9cab64af30e8675  chat_template.jinja
 1b441cb88813e8c01f8a9ae98452adfe655ba680e1964c787583634ecb8cb473  config.json
 922dfb636b84d09d2f92ba0c5bc226d00d7b41999c7e9aef1939ccb022192427  configuration_fuse_glm.py
+8ced7754937def623087e269d2a0111125ee21046a47b479dd03c8043eb63a78  fast_fp8_runtime.py
 e34e51ccf5a169f0d9cecf8119e9f719132fa0cf5fae7f8b650893821348e62a  generation_config.json
 4d28ca14dedc0b3d0fcc2b3339f0e79931faa33874f3d24f522183a8fc70068c  LICENSE
 30b85b6b9659f2e78aa259f8faf5d920a68dee7c9ced3fa6dba1f19f2bc4fca1  licenses/GLM-MIT.txt
+77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049  licenses/HF-FINEGRAINED-FP8-APACHE-2.0.txt
 4d28ca14dedc0b3d0fcc2b3339f0e79931faa33874f3d24f522183a8fc70068c  licenses/LFM-OPEN-LICENSE-v1.0.txt
 7a3c863c63fb2c52b7f881b483ffc7d1f4b352b2de932b718f62b2b2134de1e8  model-00001-of-00005.safetensors
 e1668117e8a2ad1c7d4ff9a12c1687c9e907a39738dcafd1a48647d38a6efed3  model-00002-of-00005.safetensors
@@ -13,23 +15,26 @@ e1668117e8a2ad1c7d4ff9a12c1687c9e907a39738dcafd1a48647d38a6efed3  model-00002-of
 ebb55bc7f9cd7139d3e21c0677f2b2b7f99ec25925837729b0c07a6efa2b0074  model-00004-of-00005.safetensors
 5aec98e4957c1263db1e79f2014f468620ec1b512831256329d8091f762fc61a  model-00005-of-00005.safetensors
 95b79151fed3d2716210fb3046282a47bf73e2016c0ae59f413b59e221e50fc5  model.safetensors.index.json
-2fe578e98265fa3704aa6971d1234784c4e1ef6ebf9a6df6661c831e340d3f33  modeling_fuse_glm.py
-a1829a85639489b9b4cef7dbb520e34c914bdc4bc1a726f68de5714083fe59b0  NOTICE.md
+200287c00eff2afaa68c9e8355fa8c2c73bcab4c2c79c86fb69c7a12ef885d25  modeling_fuse_glm.py
+e70d7940c5bc65ec7dda47699654521ade102a99163903ef0d67646ffc151af2  NOTICE.md
 b7d698b9414c814a1d78dce6a414c3340f3cd0cd949fc2c6b192bd2e35688710  provenance/assembly.json
 4dbb9025d648edd432e0fcbd6cc1f6dafa96e4aa0e0222171923feb17965aa25  provenance/base-bf16-checksums.sha256
 92eda133f175fa573b88b92fca5995518965117e534f3b732c63e64862488ff0  provenance/bridge.json
 69ae72f0e5add90269c5e949366c10e87d1de64fc815d14d0f330010f0fd1400  provenance/expert-selection.json
+cbcbfc68c1e543df155dbee63fe7eac265ee8b0f0ac4e5f19592da211a3bcae8  provenance/fast-serving.json
 036c4343e5d62e72613b3bb0ad1238d9f8000088d1118439798ac7ade13d7a2f  provenance/folding.json
 2ba3ae47063c1b4a957b8c5983c48d08c0dacb01c5624f856b8414ac7ccfc2c8  provenance/fusion-plan.json
 a0ab71b19384546ea62e09d0c55eee623f59b92fff5140f1baa883d2386ff12f  provenance/pre-repack-index.json
-15d69ac2fd857a1749578273f84f1385b004cd1e3fb9204e5d7684cc6c534e8c  provenance/quantization.json
+dfed0c2864aa6c8b0e187b956d20d03a36f155c4bc1bed5901f5d190b75cb2f7  provenance/quantization.json
 f6c13f0c8a28f1029132af53138b91352db35791b5da4eb0748cde4a42c59e51  provenance/router-training.json
 b8bf61e3cd2abfe2ad4b9b59e71ee48c7b680485963f3bc4e3b19df7849a3016  provenance/selected-expert-tensors.json
 ac437e75d524dd58e1c6934ad0838c62e11c9d0954d24873cb17862f958f64a5  provenance/source-models.json
 88fd032686145809f3dee3e8082a7a3b1cffe2874661289b59472a56ed5fcc1c  provenance/training-data-manifest.json
 f426ea62b798bb2340b8d9cfdea2f53f309044dbe9caa061c208e358ddc39ea9  provenance/training-summary.json
 19f75f7117ca582e393d9f79452af231f93339c55dfcc6bebd2b953f87ed3152  provenance/validation-metrics.json
-3cc74de6c69b37d403f553e376348e1c1addcbae099f90054a4f95640c6a93e3  README.md
+8407f808d4182aefa1316828e855ea6ad8267a1822a939169d6cc94ec6a1f79a  README.md
+da44e453254bbea8d1e9169c37ebd65abcbad816a0889f0ab2d9037ace351e10  requirements-serve.txt
 ce31a666514d30d302472da9387348b4958872f7e307a7ee443ae09f8d1b19b5  requirements.txt
+46da82ff64a9bf9ffa1fcdfcd78a419e0e5497a8eb271233a282c3705db44811  serve.py
 14c60d6814b7f64c69711d5c5d5561d3de9cc3896feebf9613ae8a8523a3497d  tokenizer_config.json
 695be7802a0e4b8a81048f0ff5ebb7fc811a0ba5a6be63dbb24deb5a81096f41  tokenizer.json

provenance/fast-serving.json

@@ -0,0 +1,56 @@
+{
+  "schema": "rivetcoder-fast-fp8-serving",
+  "schema_version": 1,
+  "created_at": "2026-08-27T23:51:49+09:00",
+  "hardware": {
+    "gpu": "NVIDIA GeForce RTX 5070 Ti",
+    "compute_capability": [12, 0],
+    "vram_mib": 16303
+  },
+  "software": {
+    "platform": "Windows 11",
+    "python": "3.12.10",
+    "torch": "2.12.0+cu130",
+    "torchao": "0.15.0",
+    "transformers": "5.16.1",
+    "triton": "3.5.1",
+    "compiler": "MSVC 19.44.35222"
+  },
+  "implementation": {
+    "backend": "Triton grouped FP8",
+    "packed_layers": 30,
+    "packed_experts": 480,
+    "top_k": 4,
+    "expert_gemms_per_layer": 2,
+    "gate_up_shape": [16, 4096, 2048],
+    "down_shape": [16, 2048, 2048],
+    "activation_scaling": "dynamic per routed expert",
+    "weight_scaling": "checkpoint-preserving per projection",
+    "kernel_reference": "kernels-community/finegrained-fp8",
+    "kernel_license": "Apache-2.0"
+  },
+  "full_model_parity": {
+    "baseline_seconds": 5.100,
+    "fast_seconds": 1.079,
+    "speedup": 4.73,
+    "mean_absolute_logit_error": 0.0,
+    "maximum_absolute_logit_error": 0.0,
+    "top1_equal": true,
+    "repeat_maximum_absolute_error": 0.0,
+    "resident_vram_gib": 8.428,
+    "packing_peak_vram_gib": 14.053
+  },
+  "microbatch_forward": [
+    {"batch": 1, "seconds": 1.0800, "sequences_per_second": 0.9259},
+    {"batch": 4, "seconds": 1.0868, "sequences_per_second": 3.6806},
+    {"batch": 8, "seconds": 1.0863, "sequences_per_second": 7.3645},
+    {"batch": 16, "seconds": 1.1169, "sequences_per_second": 14.3259}
+  ],
+  "server": {
+    "protocol": "OpenAI-compatible chat completions",
+    "batching": "queue-to-completion fixed microbatch",
+    "default_max_batch_size": 16,
+    "default_batch_wait_ms": 3.0,
+    "streaming": false
+  }
+}

provenance/quantization.json

@@ -42,6 +42,8 @@
     "clean_reload_seconds": 338.0,
     "clean_reload_fp8_parameters": 1636,
     "clean_reload_non_cuda_parameters": 0,
+    "fast_serving_validated": true,
+    "fast_serving_logits_bit_exact": true,
     "resident_cuda_allocated_bytes": 9039455232,
     "forward_finite": true,
     "forward_logits_shape": [1, 6, 128000]

requirements-serve.txt

@@ -0,0 +1,3 @@
+-r requirements.txt
+fastapi>=0.115
+uvicorn>=0.34

serve.py

@@ -0,0 +1,312 @@
+#!/usr/bin/env python3
+"""OpenAI-compatible, microbatched RivetCoder FP8 serving entry point."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import concurrent.futures
+import json
+import queue
+import threading
+import time
+import uuid
+from collections import deque
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+@dataclass(slots=True)
+class PendingCompletion:
+    messages: list[dict[str, Any]]
+    max_tokens: int
+    temperature: float
+    top_p: float
+    future: concurrent.futures.Future[str]
+
+    @property
+    def batch_key(self) -> tuple[int, float, float]:
+        return self.max_tokens, self.temperature, self.top_p
+
+
+class MicrobatchEngine:
+    """Collect compatible requests briefly, then generate them as one batch."""
+
+    def __init__(
+        self,
+        model: Any,
+        tokenizer: Any,
+        *,
+        max_batch_size: int,
+        batch_wait_ms: float,
+    ) -> None:
+        self.model = model
+        self.tokenizer = tokenizer
+        self.max_batch_size = int(max_batch_size)
+        self.batch_wait_seconds = float(batch_wait_ms) / 1000.0
+        self.incoming: queue.Queue[PendingCompletion | None] = queue.Queue()
+        self.deferred: deque[PendingCompletion] = deque()
+        self.thread = threading.Thread(target=self._worker, name="rivetcoder-gpu", daemon=True)
+        self.thread.start()
+
+    def submit(
+        self,
+        messages: list[dict[str, Any]],
+        *,
+        max_tokens: int,
+        temperature: float,
+        top_p: float,
+    ) -> concurrent.futures.Future[str]:
+        future: concurrent.futures.Future[str] = concurrent.futures.Future()
+        self.incoming.put(
+            PendingCompletion(
+                messages=messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                future=future,
+            )
+        )
+        return future
+
+    def close(self) -> None:
+        self.incoming.put(None)
+
+    def _next_request(self) -> PendingCompletion | None:
+        if self.deferred:
+            return self.deferred.popleft()
+        return self.incoming.get()
+
+    def _collect_batch(self, first: PendingCompletion) -> list[PendingCompletion]:
+        batch = [first]
+        key = first.batch_key
+        deadline = time.perf_counter() + self.batch_wait_seconds
+        while len(batch) < self.max_batch_size:
+            remaining = deadline - time.perf_counter()
+            if remaining <= 0:
+                break
+            try:
+                item = self.incoming.get(timeout=remaining)
+            except queue.Empty:
+                break
+            if item is None:
+                self.incoming.put(None)
+                break
+            if item.batch_key == key:
+                batch.append(item)
+            else:
+                self.deferred.append(item)
+        return batch
+
+    def _worker(self) -> None:
+        while True:
+            first = self._next_request()
+            if first is None:
+                return
+            batch = self._collect_batch(first)
+            try:
+                results = self._generate(batch)
+            except BaseException as error:
+                for item in batch:
+                    item.future.set_exception(error)
+                continue
+            for item, text in zip(batch, results, strict=True):
+                item.future.set_result(text)
+
+    def _generate(self, batch: list[PendingCompletion]) -> list[str]:
+        rendered = [
+            self.tokenizer.apply_chat_template(
+                item.messages,
+                add_generation_prompt=True,
+                tokenize=False,
+            )
+            for item in batch
+        ]
+        encoded = self.tokenizer(
+            rendered,
+            add_special_tokens=False,
+            padding=True,
+            return_tensors="pt",
+        ).to("cuda")
+        prompt_width = encoded["input_ids"].shape[-1]
+        temperature = batch[0].temperature
+        with torch.no_grad():
+            generated = self.model.generate(
+                **encoded,
+                max_new_tokens=batch[0].max_tokens,
+                do_sample=temperature > 0,
+                temperature=max(temperature, 1.0e-5),
+                top_p=batch[0].top_p,
+                use_cache=True,
+                logits_to_keep=1,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+            )
+        return self.tokenizer.batch_decode(
+            generated[:, prompt_width:],
+            skip_special_tokens=True,
+        )
+
+
+def load_runtime(args: argparse.Namespace) -> tuple[Any, Any, dict[str, Any]]:
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model,
+        trust_remote_code=True,
+        local_files_only=args.local_files_only,
+    )
+    tokenizer.padding_side = "left"
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model,
+        trust_remote_code=True,
+        local_files_only=args.local_files_only,
+        dtype=torch.bfloat16,
+        device_map=0,
+        attn_implementation="sdpa",
+    ).eval()
+    model.config.use_cache = True
+    model.set_coding_enabled(True)
+    if not hasattr(model, "enable_fast_fp8_serving"):
+        raise RuntimeError(
+            "The model package does not contain the grouped-FP8 runtime. "
+            "Use the updated RivetCoder FP8 package."
+        )
+    report = model.enable_fast_fp8_serving()
+    return tokenizer, model, report
+
+
+def warmup(model: Any, tokenizer: Any, batch_sizes: list[int]) -> None:
+    text = tokenizer.apply_chat_template(
+        [{"role": "user", "content": "Return the integer 1."}],
+        add_generation_prompt=True,
+        tokenize=False,
+    )
+    for batch_size in batch_sizes:
+        encoded = tokenizer(
+            [text] * batch_size,
+            add_special_tokens=False,
+            padding=True,
+            return_tensors="pt",
+        ).to("cuda")
+        with torch.no_grad():
+            model.generate(
+                **encoded,
+                max_new_tokens=2,
+                do_sample=False,
+                use_cache=True,
+                logits_to_keep=1,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+
+
+def build_app(engine: MicrobatchEngine, runtime_report: dict[str, Any], model_name: str) -> Any:
+    try:
+        from fastapi import FastAPI, HTTPException
+    except ImportError as error:
+        raise RuntimeError("Serving requires fastapi and uvicorn") from error
+
+    app = FastAPI(title="RivetCoder FP8 Server")
+    created = int(time.time())
+
+    @app.get("/health")
+    async def health() -> dict[str, Any]:
+        return {
+            "status": "ok",
+            "model": model_name,
+            "runtime": runtime_report,
+            "max_batch_size": engine.max_batch_size,
+            "batch_wait_ms": engine.batch_wait_seconds * 1000.0,
+            "cuda_allocated_gib": torch.cuda.memory_allocated() / 1024**3,
+        }
+
+    @app.get("/v1/models")
+    async def models() -> dict[str, Any]:
+        return {
+            "object": "list",
+            "data": [{"id": model_name, "object": "model", "created": created, "owned_by": "HCHs"}],
+        }
+
+    @app.post("/v1/chat/completions")
+    async def chat_completions(payload: dict[str, Any]) -> dict[str, Any]:
+        if payload.get("stream", False):
+            raise HTTPException(status_code=400, detail="Streaming is not implemented in the microbatch server")
+        messages = payload.get("messages")
+        if not isinstance(messages, list) or not messages:
+            raise HTTPException(status_code=400, detail="messages must be a non-empty list")
+        max_tokens = int(payload.get("max_tokens", 512))
+        if max_tokens < 1 or max_tokens > 4096:
+            raise HTTPException(status_code=400, detail="max_tokens must be between 1 and 4096")
+        temperature = float(payload.get("temperature", 0.2))
+        top_p = float(payload.get("top_p", 0.95))
+        future = engine.submit(
+            messages,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+        )
+        try:
+            text = await asyncio.wrap_future(future)
+        except Exception as error:
+            raise HTTPException(status_code=500, detail=str(error)) from error
+        completion_id = f"chatcmpl-{uuid.uuid4().hex}"
+        return {
+            "id": completion_id,
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": model_name,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": text},
+                    "finish_reason": "stop",
+                }
+            ],
+        }
+
+    return app
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--model",
+        default=str(Path("RivetCoder-9B-A4B-FP8")),
+        help="Local FP8 model directory or Hugging Face model id",
+    )
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--max-batch-size", type=int, default=16)
+    parser.add_argument("--batch-wait-ms", type=float, default=3.0)
+    parser.add_argument("--warmup-batches", default="1,8,16")
+    parser.add_argument("--local-files-only", action=argparse.BooleanOptionalAction, default=True)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+    torch.set_float32_matmul_precision("high")
+    tokenizer, model, runtime_report = load_runtime(args)
+    warmup_batches = [int(value) for value in args.warmup_batches.split(",") if value]
+    warmup(model, tokenizer, warmup_batches)
+    engine = MicrobatchEngine(
+        model,
+        tokenizer,
+        max_batch_size=args.max_batch_size,
+        batch_wait_ms=args.batch_wait_ms,
+    )
+    app = build_app(engine, runtime_report, str(args.model))
+    print(json.dumps({"runtime": runtime_report, "listen": f"http://{args.host}:{args.port}"}, indent=2))
+    import uvicorn
+
+    uvicorn.run(app, host=args.host, port=args.port, workers=1)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())