Accelerate remaining FP8 linears with direct scaled-mm

NOTICE.md

@@ -14,7 +14,8 @@ RivetCoder-9B-A4B combines the following sources:
   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.
+  Top-4 reduction, direct FP8 Linear dispatch using the checkpoint's existing
+  qdata/scales, 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,

README.md

@@ -133,9 +133,11 @@ still passed quantization, serialization, clean reload, and forward validation.
 
 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.
+gate/up projection and one down projection. It also bypasses TorchAO's
+tensor-subclass dispatch for 196 remaining compatible FP8 Linear modules and
+calls their existing qdata/scales through `_scaled_mm` directly. 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:
 
@@ -166,8 +168,8 @@ 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 |
+| Original TorchAO path | 4.894 s | 1.00x |
+| Grouped/direct-FP8 fast path | 0.304 s | 16.08x |
 
 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
@@ -175,14 +177,15 @@ 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 |
+| 1 | 0.337 s | 2.96 |
+| 4 | 0.316 s | 12.65 |
+| 8 | 0.340 s | 23.50 |
+| 16 | 0.309 s | 51.71 |
 
 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.
+benchmarks. Batch 16 increased throughput about 17.5x over batch 1 without a
+latency increase in this short test, which is why the bundled server defaults
+to batch 16.
 
 ## Limitations
 

fast_fp8_runtime.py

@@ -15,6 +15,7 @@ import gc
 import os
 import shutil
 import subprocess
+import types
 from pathlib import Path
 from typing import Any
 
@@ -324,6 +325,52 @@ def _float8_parts(linear: nn.Linear) -> tuple[torch.Tensor, torch.Tensor]:
     return qdata, scale.reshape(())
 
 
+def _direct_fp8_linear_forward(linear: nn.Linear, hidden_states: torch.Tensor) -> torch.Tensor:
+    """TorchAO-compatible PerTensor FP8 Linear without tensor-subclass dispatch."""
+
+    qdata, weight_scale = _float8_parts(linear)
+    original_shape = hidden_states.shape
+    flattened = hidden_states.reshape(-1, original_shape[-1]).contiguous()
+    activation_scale = (flattened.abs().amax() / 448.0).float().reshape(1, 1)
+    activation_scale = activation_scale.clamp_min(1.0e-12)
+    quantized = (
+        flattened.float()
+        .div(activation_scale)
+        .clamp(min=-448.0, max=448.0)
+        .to(torch.float8_e4m3fn)
+    )
+    output = torch._scaled_mm(
+        quantized,
+        qdata.t(),
+        activation_scale,
+        weight_scale.reshape(1, 1),
+        out_dtype=linear.weight.dtype,
+        use_fast_accum=True,
+    )
+    if linear.bias is not None:
+        output = output + linear.bias
+    return output.reshape(*original_shape[:-1], linear.out_features)
+
+
+def _install_direct_fp8_linears(model: nn.Module) -> int:
+    installed = 0
+    for module in model.modules():
+        if not isinstance(module, nn.Linear) or getattr(module, "_rivet_direct_fp8", False):
+            continue
+        weight = module.weight
+        if (
+            "Float8" not in type(weight).__name__
+            or getattr(weight, "qdata", None) is None
+            or getattr(weight, "scale", None) is None
+            or weight.scale.numel() != 1
+        ):
+            continue
+        module.forward = types.MethodType(_direct_fp8_linear_forward, module)
+        module._rivet_direct_fp8 = True
+        installed += 1
+    return installed
+
+
 class PackedFp8ExpertBank(nn.Module):
     """One layer's 16 experts packed into two grouped FP8 projections."""
 
@@ -514,12 +561,14 @@ def install_fast_fp8_runtime(model: nn.Module) -> dict[str, Any]:
     # per-expert qdata that the packed banks replaced.
     gc.collect()
     torch.cuda.empty_cache()
+    direct_fp8_linears = _install_direct_fp8_linears(model)
 
     return {
         "backend": "triton-grouped-fp8",
         "packed_layers": packed_layers,
         "released_experts": released_experts,
         "packed_bytes": packed_bytes,
+        "direct_fp8_linears": direct_fp8_linears,
         "cuda_allocated_bytes": torch.cuda.memory_allocated(),
         "compiler": os.environ.get("CC"),
     }

provenance/checksums.sha256

@@ -3,7 +3,7 @@
 f3e412383b66fc807c3d0fbe9e6f8a4a58b2638df55d8380e9cab64af30e8675  chat_template.jinja
 1b441cb88813e8c01f8a9ae98452adfe655ba680e1964c787583634ecb8cb473  config.json
 922dfb636b84d09d2f92ba0c5bc226d00d7b41999c7e9aef1939ccb022192427  configuration_fuse_glm.py
-8ced7754937def623087e269d2a0111125ee21046a47b479dd03c8043eb63a78  fast_fp8_runtime.py
+3604561de9be361a595cd4dcbfd592d6fe13d0da5f920939f457d6fd93492bcb  fast_fp8_runtime.py
 e34e51ccf5a169f0d9cecf8119e9f719132fa0cf5fae7f8b650893821348e62a  generation_config.json
 4d28ca14dedc0b3d0fcc2b3339f0e79931faa33874f3d24f522183a8fc70068c  LICENSE
 30b85b6b9659f2e78aa259f8faf5d920a68dee7c9ced3fa6dba1f19f2bc4fca1  licenses/GLM-MIT.txt
@@ -16,12 +16,12 @@ ebb55bc7f9cd7139d3e21c0677f2b2b7f99ec25925837729b0c07a6efa2b0074  model-00004-of
 5aec98e4957c1263db1e79f2014f468620ec1b512831256329d8091f762fc61a  model-00005-of-00005.safetensors
 95b79151fed3d2716210fb3046282a47bf73e2016c0ae59f413b59e221e50fc5  model.safetensors.index.json
 200287c00eff2afaa68c9e8355fa8c2c73bcab4c2c79c86fb69c7a12ef885d25  modeling_fuse_glm.py
-e70d7940c5bc65ec7dda47699654521ade102a99163903ef0d67646ffc151af2  NOTICE.md
+457a3785fb93db67d9ce2161aedbe4de86c240cc7908d01f5693eabbe012776d  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
+40bed73ee90cd26f3da8a0f58d79cd48a73d32af40c269fc56ab592bfba50d4a  provenance/fast-serving.json
 036c4343e5d62e72613b3bb0ad1238d9f8000088d1118439798ac7ade13d7a2f  provenance/folding.json
 2ba3ae47063c1b4a957b8c5983c48d08c0dacb01c5624f856b8414ac7ccfc2c8  provenance/fusion-plan.json
 a0ab71b19384546ea62e09d0c55eee623f59b92fff5140f1baa883d2386ff12f  provenance/pre-repack-index.json
@@ -32,9 +32,9 @@ ac437e75d524dd58e1c6934ad0838c62e11c9d0954d24873cb17862f958f64a5  provenance/sou
 88fd032686145809f3dee3e8082a7a3b1cffe2874661289b59472a56ed5fcc1c  provenance/training-data-manifest.json
 f426ea62b798bb2340b8d9cfdea2f53f309044dbe9caa061c208e358ddc39ea9  provenance/training-summary.json
 19f75f7117ca582e393d9f79452af231f93339c55dfcc6bebd2b953f87ed3152  provenance/validation-metrics.json
-8407f808d4182aefa1316828e855ea6ad8267a1822a939169d6cc94ec6a1f79a  README.md
+81eee8545dadab653bd492dbcee75ed80099dcee73daf92c5c1ebb99512c426b  README.md
 da44e453254bbea8d1e9169c37ebd65abcbad816a0889f0ab2d9037ace351e10  requirements-serve.txt
 ce31a666514d30d302472da9387348b4958872f7e307a7ee443ae09f8d1b19b5  requirements.txt
-46da82ff64a9bf9ffa1fcdfcd78a419e0e5497a8eb271233a282c3705db44811  serve.py
+e3a92f56b19700001d1dee2dd5aaee3f33fb66fd2a8e41eda8fe1da7599bbd78  serve.py
 14c60d6814b7f64c69711d5c5d5561d3de9cc3896feebf9613ae8a8523a3497d  tokenizer_config.json
 695be7802a0e4b8a81048f0ff5ebb7fc811a0ba5a6be63dbb24deb5a81096f41  tokenizer.json

provenance/fast-serving.json

@@ -20,6 +20,7 @@
     "backend": "Triton grouped FP8",
     "packed_layers": 30,
     "packed_experts": 480,
+    "direct_fp8_linears": 196,
     "top_k": 4,
     "expert_gemms_per_layer": 2,
     "gate_up_shape": [16, 4096, 2048],
@@ -30,9 +31,9 @@
     "kernel_license": "Apache-2.0"
   },
   "full_model_parity": {
-    "baseline_seconds": 5.100,
-    "fast_seconds": 1.079,
-    "speedup": 4.73,
+    "baseline_seconds": 4.8937,
+    "fast_seconds": 0.3044,
+    "speedup": 16.08,
     "mean_absolute_logit_error": 0.0,
     "maximum_absolute_logit_error": 0.0,
     "top1_equal": true,
@@ -41,10 +42,10 @@
     "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}
+    {"batch": 1, "seconds": 0.3374, "sequences_per_second": 2.9635},
+    {"batch": 4, "seconds": 0.3162, "sequences_per_second": 12.6488},
+    {"batch": 8, "seconds": 0.3404, "sequences_per_second": 23.5045},
+    {"batch": 16, "seconds": 0.3094, "sequences_per_second": 51.7069}
   ],
   "server": {
     "protocol": "OpenAI-compatible chat completions",

serve.py

@@ -134,18 +134,18 @@ class MicrobatchEngine:
         ).to("cuda")
         prompt_width = encoded["input_ids"].shape[-1]
         temperature = batch[0].temperature
+        generation_kwargs = {
+            "max_new_tokens": batch[0].max_tokens,
+            "do_sample": temperature > 0,
+            "use_cache": True,
+            "logits_to_keep": 1,
+            "pad_token_id": self.tokenizer.pad_token_id,
+            "eos_token_id": self.tokenizer.eos_token_id,
+        }
+        if temperature > 0:
+            generation_kwargs.update(temperature=temperature, top_p=batch[0].top_p)
         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,
-            )
+            generated = self.model.generate(**encoded, **generation_kwargs)
         return self.tokenizer.batch_decode(
             generated[:, prompt_width:],
             skip_special_tokens=True,