tools: warm-server FP8 vs BF16 benchmark + results doc

fp8_compare.py launches one xserv-server per model (same GPUs / TP for a
fair comparison), gates readiness on a real generation (not /health),
and streams GSM8K through /v1/chat/completions measuring per-request
TTFT (time to first token) and TPOT (mean inter-token latency) plus
exact-match accuracy. docs/benchmarks/fp8-quantization.md records the
quantization scheme, the perf-bug fix, and the dash5 results.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xserv-kernels/src/quantization.rs

@@ -1,4 +1,5 @@
 use std::cell::RefCell;
+use std::collections::HashMap;
 use std::ffi::c_void;
 use xserv_cuda::GpuBuffer;
 use xserv_tensor::{DType, Tensor};
@@ -113,9 +114,33 @@ const CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES: i32 = 1;
 
 const WORKSPACE_BYTES: usize = 32 * 1024 * 1024;
 
+const CUBLASLT_MATMUL_DESC_TRANSA: i32 = 3;
+
+/// A fully-prepared FP8 matmul plan for one (M, N, K) shape: the matmul
+/// descriptor, the four matrix layouts, and the heuristically-chosen algo.
+/// Built once per shape and reused across every expert and every forward
+/// pass — the heuristic search and descriptor/layout creation are the
+/// expensive parts, so doing them once instead of per-expert-per-layer is
+/// the difference between FP8 being faster or slower than BF16.
+#[derive(Clone, Copy)]
+struct Fp8Plan {
+    desc: CublasLtMatmulDesc,
+    a_layout: CublasLtMatrixLayout,
+    b_layout: CublasLtMatrixLayout,
+    c_layout: CublasLtMatrixLayout,
+    d_layout: CublasLtMatrixLayout,
+    algo: CublasLtMatmulAlgo,
+    workspace_size: usize,
+}
+
 struct CublasLtContext {
     handle: CublasLtHandle,
     workspace: GpuBuffer,
+    /// Persistent device scalar holding 1.0, used as the A/B scale pointer
+    /// placeholder. Allocated once instead of per-expert.
+    one_buf: GpuBuffer,
+    /// Cache of prepared matmul plans keyed by (M, N, K).
+    plans: HashMap<(usize, usize, usize), Fp8Plan>,
 }
 
 impl CublasLtContext {
@@ -124,18 +149,100 @@ impl CublasLtContext {
         let status = unsafe { cublasLtCreate(&mut handle) };
         assert_eq!(status, 0, "cublasLtCreate failed: {status}");
         let workspace = GpuBuffer::alloc(WORKSPACE_BYTES).expect("alloc cublasLt workspace");
-        Self { handle, workspace }
+        let mut one_buf = GpuBuffer::alloc(4).expect("alloc cublasLt fp8 scale");
+        one_buf.copy_from_host(&1.0f32.to_le_bytes()).expect("init fp8 scale");
+        Self { handle, workspace, one_buf, plans: HashMap::new() }
+    }
+
+    /// Get the cached plan for (m, n, k), building (and caching) it on first use.
+    fn plan(&mut self, m: usize, n: usize, k: usize) -> Fp8Plan {
+        if let Some(p) = self.plans.get(&(m, n, k)) {
+            return *p;
+        }
+        let one_ptr = self.one_buf.as_ptr() as *const c_void;
+        let plan = unsafe { build_fp8_plan(self.handle, one_ptr, m, n, k) };
+        self.plans.insert((m, n, k), plan);
+        plan
     }
 }
 
 impl Drop for CublasLtContext {
     fn drop(&mut self) {
+        // Tear down cached plans before destroying the handle.
+        for (_, p) in self.plans.drain() {
+            unsafe {
+                cublasLtMatrixLayoutDestroy(p.a_layout);
+                cublasLtMatrixLayoutDestroy(p.b_layout);
+                cublasLtMatrixLayoutDestroy(p.c_layout);
+                cublasLtMatrixLayoutDestroy(p.d_layout);
+                cublasLtMatmulDescDestroy(p.desc);
+            }
+        }
         if !self.handle.is_null() {
             unsafe { cublasLtDestroy(self.handle) };
         }
     }
 }
 
+/// Build an FP8 matmul plan for one (m, n, k) shape. See `batched_gemm_fp8`
+/// for the row-major → cuBLASLt col-major layout mapping (transA=T, transB=N,
+/// m_lt=N, n_lt=M, k_lt=K). The B-scale pointer is initialised to `one_ptr`
+/// and overwritten per-expert at call time.
+unsafe fn build_fp8_plan(
+    handle: CublasLtHandle,
+    one_ptr: *const c_void,
+    m: usize,
+    n: usize,
+    k: usize,
+) -> Fp8Plan {
+    let m_lt = n as u64;
+    let n_lt = m as u64;
+    let k_lt = k as u64;
+
+    let mut desc: CublasLtMatmulDesc = std::ptr::null_mut();
+    cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+
+    // transA=T (required for FP8 on Blackwell)
+    let trans_a: i32 = 1;
+    cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &trans_a as *const i32 as _, 4);
+    let ptr_sz = std::mem::size_of::<*const c_void>();
+    cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &one_ptr as *const _ as _, ptr_sz);
+    cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &one_ptr as *const _ as _, ptr_sz);
+
+    // "A" layout (weights, transposed): physical (K, N) col-major, ld=K
+    let mut a_layout: CublasLtMatrixLayout = std::ptr::null_mut();
+    cublasLtMatrixLayoutCreate(&mut a_layout, CUDA_R_8F_E4M3, k_lt, m_lt, k as i64);
+    // "B" layout (activations): physical (K, M) col-major, ld=K
+    let mut b_layout: CublasLtMatrixLayout = std::ptr::null_mut();
+    cublasLtMatrixLayoutCreate(&mut b_layout, CUDA_R_8F_E4M3, k_lt, n_lt, k as i64);
+    // "C"/"D" layout (output): physical (N, M) col-major, ld=N
+    let mut c_layout: CublasLtMatrixLayout = std::ptr::null_mut();
+    cublasLtMatrixLayoutCreate(&mut c_layout, CUDA_R_16BF, m_lt, n_lt, m_lt as i64);
+    let mut d_layout: CublasLtMatrixLayout = std::ptr::null_mut();
+    cublasLtMatrixLayoutCreate(&mut d_layout, CUDA_R_16BF, m_lt, n_lt, m_lt as i64);
+
+    let mut pref: CublasLtMatmulPreference = std::ptr::null_mut();
+    cublasLtMatmulPreferenceCreate(&mut pref);
+    let ws_bytes = WORKSPACE_BYTES as u64;
+    cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_bytes as *const u64 as _, 8);
+
+    let mut heuristic = std::mem::zeroed::<CublasLtMatmulHeuristicResult>();
+    let mut found: i32 = 0;
+    let status = cublasLtMatmulAlgoGetHeuristic(
+        handle, desc, a_layout, b_layout, c_layout, d_layout,
+        pref, 1, &mut heuristic, &mut found,
+    );
+    assert!(status == 0 && found > 0,
+        "cublasLtMatmulAlgoGetHeuristic failed for FP8 GEMM (m={m}, n={n}, k={k}): status={status}, found={found}");
+    cublasLtMatmulPreferenceDestroy(pref);
+
+    Fp8Plan {
+        desc, a_layout, b_layout, c_layout, d_layout,
+        algo: heuristic.algo,
+        workspace_size: heuristic.workspace_size,
+    }
+}
+
 thread_local! {
     static CUBLASLT_CTX: RefCell<CublasLtContext> = RefCell::new(CublasLtContext::new());
 }
@@ -215,9 +322,9 @@ pub fn quantize_bf16_to_fp8_rowwise(src: &Tensor) -> (Tensor, Tensor) {
 /// as [b, N, K] for cuBLASLt FP8 compatibility.
 ///
 /// a_fp8:    [batch, M, K] FP8E4M3  (activations, quantized per-row)
-/// a_scales: [batch * M] F32        (per-token scales, collapsed to per-batch max)
+/// a_scales: [batch * M] F32        (per-token activation scales, applied post-GEMM)
 /// b_fp8_t:  [batch, N, K] FP8E4M3  (weights, TRANSPOSED for cuBLASLt)
-/// b_scales: [batch] F32            (per-expert scalar scales)
+/// b_scales: [batch] F32            (per-expert scalar weight scales, applied in-GEMM)
 ///
 /// Returns:  [batch, M, N] BF16
 pub fn batched_gemm_fp8(
@@ -240,127 +347,64 @@ pub fn batched_gemm_fp8(
     let k = a_fp8.shape()[2]; // hidden
     let n = b_fp8_t.shape()[1]; // out_dim (from transposed weight)
 
-    // Per-token scales → per-expert scales (max over tokens within each expert batch)
+    // a_scales: [batch * M] per-token activation scales (applied post-GEMM, per row).
-    // a_scales: [batch * M] → we take the max per expert to get [batch] scalar scales
+    // b_scales: [batch] per-expert scalar weight scales (applied in-GEMM via B-scale ptr).
-    // This is a slight accuracy tradeoff vs per-token, but allows scalar GEMM scale mode.
     assert_eq!(a_scales.shape()[0], batch * m);
     assert_eq!(b_scales.shape()[0], batch);
 
     let c = Tensor::empty(&[batch, m, n], DType::BF16, a_fp8.device());
 
-    // Read weight scales to host for the per-expert loop
-    let b_scales_cpu = b_scales.to_device(xserv_tensor::Device::Cpu);
-    let b_s_data = b_scales_cpu.as_slice::<f32>();
-
-    CUBLASLT_CTX.with(|cell| {
-        let ctx = cell.borrow();
-        let handle = ctx.handle;
-
     // Strides (in bytes) for one expert slice
     let stride_a = m * k;     // FP8: 1 byte per elem
     let stride_b = n * k;     // FP8: 1 byte per elem (transposed: [N, K])
     let stride_c = m * n * 2; // BF16: 2 bytes per elem
 
+    CUBLASLT_CTX.with(|cell| {
+        let mut ctx = cell.borrow_mut();
+        let handle = ctx.handle;
+        let ws_ptr = ctx.workspace.as_ptr() as *mut c_void;
+        // Build (or fetch) the cached plan for this shape — heuristic search and
+        // descriptor/layout creation happen once per (m, n, k), not per-expert.
+        let plan = ctx.plan(m, n, k);
+
+        // alpha=1, beta=0. Per-expert weight scale is supplied via the cuBLASLt
+        // B-scale pointer (device, scalar): cuBLASLt computes in the FP32 epilogue
+        //   D = (1.0 * A_fp8) @ (b_scale[e] * B_fp8)^T = b_scale[e] * (A_fp8 @ B_fp8^T)
+        // Per-token activation scale (a_scale) is applied post-GEMM (per row).
+        let alpha: f32 = 1.0;
+        let beta: f32 = 0.0;
+        let ptr_sz = std::mem::size_of::<*const c_void>();
+
         for e in 0..batch {
             let a_ptr = unsafe { (a_fp8.data_ptr() as *const u8).add(e * stride_a) as *const c_void };
             let b_ptr = unsafe { (b_fp8_t.data_ptr() as *const u8).add(e * stride_b) as *const c_void };
             let c_ptr = unsafe { (c.data_ptr() as *mut u8).add(e * stride_c) as *mut c_void };
-
+            // Device pointer to this expert's scalar weight scale (FP32, 4 bytes).
-            // alpha = b_scale (weight scale). Per-row activation scale applied post-GEMM.
+            let b_scale_ptr = unsafe { (b_scales.data_ptr() as *const u8).add(e * 4) as *const c_void };
-            // GEMM computes: D = alpha * (A_fp8 @ B_fp8_T)
-            //   = b_scale * ((A_real / a_scale_row) @ (B_real / b_scale))
-            //   = (A_real / a_scale_row) @ B_real
-            // Post-multiply row i by a_scale[i] to recover the correct result.
-            let alpha: f32 = b_s_data[e];
-            let beta: f32 = 0.0;
 
             unsafe {
-                // cuBLASLt FP8 on Blackwell requires transA=T, transB=N.
+                cublasLtMatmulDescSetAttribute(
-                // cuBLASLt computes: D(m,n) = op(A)(m,k) * B(k,n) with transA=T
+                    plan.desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
-                //
+                    &b_scale_ptr as *const _ as _, ptr_sz,
-                // We want: D_row[M,N] = A_act_row[M,K] @ B_wt_row[K,N]
-                // Map to cuBLASLt with m_lt=N, n_lt=M, k_lt=K:
-                //   "A" (transA=T): stored as (K, N) col-major ld=K → transposed to (N, K)
-                //     Our weights are stored TRANSPOSED as [E, N, K] row-major = col-major (K, N) ld=K ✓
-                //   "B" (transB=N): stored as (K, M) col-major ld=K
-                //     Our activations A_act_row[M,K] = col-major (K, M) ld=K ✓
-                //   "D": stored as (N, M) col-major ld=N
-                //     Our output D_row[M,N] = col-major (N, M) ld=N ✓
-                let m_lt = n as u64;
-                let n_lt = m as u64;
-                let k_lt = k as u64;
-
-                let mut matmul_desc: CublasLtMatmulDesc = std::ptr::null_mut();
-                cublasLtMatmulDescCreate(&mut matmul_desc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
-
-                // Set transA=T (required for FP8 on Blackwell)
-                let trans_a: i32 = 1; // CUBLAS_OP_T
-                cublasLtMatmulDescSetAttribute(matmul_desc, 3 /*TRANSA*/, &trans_a as *const i32 as _, 4);
-
-                // FP8 requires scale pointers. We fold the actual scales into alpha,
-                // so set dummy 1.0 scale pointers on device.
-                let one_val: f32 = 1.0;
-                let mut one_buf = xserv_cuda::allocator::cached_alloc(4).unwrap();
-                one_buf.copy_from_host(&one_val.to_le_bytes()).unwrap();
-                let one_ptr = one_buf.as_ptr() as *const c_void;
-                cublasLtMatmulDescSetAttribute(matmul_desc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &one_ptr as *const _ as _, std::mem::size_of::<*const c_void>());
-                cublasLtMatmulDescSetAttribute(matmul_desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &one_ptr as *const _ as _, std::mem::size_of::<*const c_void>());
-
-                // "A" layout (weights, transposed): physical (K, N) col-major, ld=K
-                let mut a_layout: CublasLtMatrixLayout = std::ptr::null_mut();
-                cublasLtMatrixLayoutCreate(&mut a_layout, CUDA_R_8F_E4M3, k_lt, m_lt, k as i64);
-
-                // "B" layout (activations): physical (K, M) col-major, ld=K
-                let mut b_layout: CublasLtMatrixLayout = std::ptr::null_mut();
-                cublasLtMatrixLayoutCreate(&mut b_layout, CUDA_R_8F_E4M3, k_lt, n_lt, k as i64);
-
-                // "C"/"D" layout (output): physical (N, M) col-major, ld=N
-                let mut c_layout: CublasLtMatrixLayout = std::ptr::null_mut();
-                cublasLtMatrixLayoutCreate(&mut c_layout, CUDA_R_16BF, m_lt, n_lt, m_lt as i64);
-                let mut d_layout: CublasLtMatrixLayout = std::ptr::null_mut();
-                cublasLtMatrixLayoutCreate(&mut d_layout, CUDA_R_16BF, m_lt, n_lt, m_lt as i64);
-
-                // Get algo heuristic
-                let mut pref: CublasLtMatmulPreference = std::ptr::null_mut();
-                cublasLtMatmulPreferenceCreate(&mut pref);
-                let ws_bytes = WORKSPACE_BYTES as u64;
-                cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_bytes as *const u64 as _, 8);
-
-                let mut heuristic = std::mem::zeroed::<CublasLtMatmulHeuristicResult>();
-                let mut found: i32 = 0;
-                let status = cublasLtMatmulAlgoGetHeuristic(
-                    handle, matmul_desc,
-                    a_layout, b_layout, c_layout, d_layout,
-                    pref, 1, &mut heuristic, &mut found,
                 );
-                assert!(status == 0 && found > 0,
-                    "cublasLtMatmulAlgoGetHeuristic failed for FP8 GEMM: status={status}, found={found}");
-
                 let status = cublasLtMatmul(
-                    handle, matmul_desc,
+                    handle, plan.desc,
                     &alpha as *const f32 as _,
                     b_ptr,    // cuBLASLt "A" = weights
-                    a_layout,
+                    plan.a_layout,
                     a_ptr,    // cuBLASLt "B" = activations
-                    b_layout,
+                    plan.b_layout,
                     &beta as *const f32 as _,
                     c_ptr,    // C (unused with beta=0)
-                    c_layout,
+                    plan.c_layout,
                     c_ptr,    // D = output
-                    d_layout,
+                    plan.d_layout,
-                    &heuristic.algo,
+                    &plan.algo,
-                    ctx.workspace.as_ptr() as *mut c_void,
+                    ws_ptr,
-                    heuristic.workspace_size,
+                    plan.workspace_size,
                     std::ptr::null_mut(),
                 );
                 assert_eq!(status, 0, "cublasLtMatmul FP8 failed for expert {e}: status={status}");
-
-                cublasLtMatmulPreferenceDestroy(pref);
-                cublasLtMatrixLayoutDestroy(a_layout);
-                cublasLtMatrixLayoutDestroy(b_layout);
-                cublasLtMatrixLayoutDestroy(c_layout);
-                cublasLtMatrixLayoutDestroy(d_layout);
-                cublasLtMatmulDescDestroy(matmul_desc);
             }
         }
     });

docs/benchmarks/fp8-quantization.md

@@ -0,0 +1,83 @@
+# FP8 W8A8 quantization — gpt-oss-20b (dash5, 8× RTX 5090)
+
+Operator-level FP8 E4M3 quantization of the MoE expert weights, with real
+cuBLASLt FP8 tensor-core GEMM (W8A8: FP8 weights × dynamically-quantized FP8
+activations). All other tensors (attention, router, embeddings, norms, biases)
+stay BF16.
+
+## Scheme
+
+- **Weights** (`tools/quantize_fp8.py`): expert `gate_up_proj` / `down_proj`
+  quantized BF16 → FP8 E4M3 with a **per-expert scalar** scale (`absmax/448`).
+  Stored transposed `[E, N, K]` because cuBLASLt FP8 on Blackwell (sm120)
+  requires `transA=T`.
+- **Activations**: quantized dynamically at runtime, **per-token** (per-row
+  absmax), recovered by a post-GEMM row scale.
+- **Compute**: `batched_gemm_fp8` (`crates/xserv-kernels/src/quantization.rs`)
+  runs one cuBLASLt FP8 matmul per expert; the per-expert weight scale is
+  supplied via the cuBLASLt B-scale device pointer (FP32 epilogue, so precision
+  matches folding it into `alpha`).
+- Model size: **22 GB** (FP8) vs **39 GB** (BF16). The FP8 model fits on a
+  single 32 GB 5090; BF16 needs ≥ 2.
+
+## The performance bug that was fixed
+
+`batched_gemm_fp8` originally rebuilt the entire cuBLASLt plan **per expert,
+per GEMM, per layer, on every forward pass** — running the algo heuristic
+search, creating/destroying the descriptor + 4 layouts + preference, and
+`cudaMalloc`-ing a 4-byte scale buffer — roughly 1500 heuristic searches per
+decoded token. This made FP8 **slower than BF16**:
+
+| | FP8 (buggy) | FP8 (fixed) | BF16 |
+|---|---|---|---|
+| Decode TPOT | 27.0 ms | **17.9 ms** | 18.8 ms |
+| Throughput | 37 tok/s | **55.8 tok/s** | 53.2 tok/s |
+
+Fix: cache the cuBLASLt plan (descriptor + layouts + heuristically-chosen algo)
+in a thread-local map keyed by `(M, N, K)` so the heuristic runs once per shape;
+allocate the scale buffer once; pass per-expert weight scales by device pointer.
+The per-expert loop now issues only `cublasLtMatmul`.
+
+## Results — GSM8K (200 problems, greedy, TP=2 on the same 2 GPUs)
+
+Harness: `tools/fp8_compare.py` — a warm `xserv-server` per model, GSM8K streamed
+through `/v1/chat/completions`; TTFT = time to first token, TPOT = mean
+inter-token latency, per request.
+
+| metric | FP8 W8A8 | BF16 |
+|---|---|---|
+| GSM8K accuracy | **93.0 %** | 90.5 % |
+| TTFT median | 67.4 ms | 68.8 ms |
+| TTFT p90 | 90.4 ms | 96.7 ms |
+| TPOT median | **17.45 ms** | 18.26 ms |
+| TPOT p90 | 17.65 ms | 18.38 ms |
+| Throughput | **57.3 tok/s** | 54.8 tok/s |
+| Mean output tokens | 288 | 293 |
+
+- **Accuracy: unchanged.** FP8 is nominally +2.5 pts, but with n=200 the
+  standard error is ~2.1 pts, so the two are statistically indistinguishable.
+  The takeaway is that FP8 did **not** degrade accuracy.
+- **Decode: FP8 ~5 % faster** (TPOT 17.45 vs 18.26 ms), reproducible across
+  runs, with a tighter p90. Modest because the dense-MoE path loads *all*
+  experts every token and FP8 only halves the *expert* bytes; the per-expert
+  M=1 launches and M=1 tensor-core inefficiency absorb much of the bandwidth
+  saving.
+- **Prefill (TTFT): comparable.** A multi-length sweep (113 / 561 / 1681 tokens)
+  gave FP8 480 / 362 / 2451 ms vs BF16 558 / 282 / 2287 ms — non-monotonic, i.e.
+  dominated by fixed overhead (cuBLAS lazy init + FP8's one-time per-shape
+  heuristic), not prefill compute, at these lengths.
+
+## Single-GPU (TP=1)
+
+FP8 runs gpt-oss-20b on **one** 5090 (`bench-gpt-oss --tp 1`, GPU6): TTFT 538 ms,
+TPOT 29.0 ms, 34.5 tok/s. BF16 cannot (39 GB > 32 GB). This — fitting a model
+that otherwise needs two GPUs onto one — is the largest practical win.
+
+## Follow-ups (not done)
+
+- Strided-batched FP8 (one call instead of ~768 per-expert launches per token) —
+  requires folding the per-expert weight scale into the post-scale kernel, at a
+  BF16-intermediate precision cost.
+- Per-channel (per-output-row) weight scales for better accuracy headroom than
+  per-tensor.
+- Warm common prefill shapes at load to hide the first-request heuristic stall.

tools/fp8_compare.py

@@ -0,0 +1,283 @@
+#!/usr/bin/env python3
+"""Compare FP8-W8A8 vs BF16 gpt-oss on one box: GSM8K accuracy + TTFT/TPOT.
+
+For each model it launches a warm xserv-server (same GPUs / same TP for a fair
+compute comparison), waits for a *real* generation to succeed (not /health),
+then streams N GSM8K problems through /v1/chat/completions measuring per-request
+TTFT (time to first token) and TPOT (mean inter-token latency). Accuracy is the
+exact-match rate on the extracted final number.
+
+Run it ON the GPU box (it manages the servers itself):
+
+    python3 tools/fp8_compare.py \
+        --fp8 /opt/wjh/models/gpt-oss-20b-fp8 \
+        --bf16 /opt/wjh/models/gpt-oss-20b-bf16 \
+        --gpus 0,1 --tp 2 --limit 150 --max-tokens 512
+"""
+
+import argparse
+import json
+import os
+import re
+import signal
+import subprocess
+import sys
+import time
+import urllib.request
+import urllib.error
+from pathlib import Path
+
+SCRIPT_DIR = Path(__file__).parent
+GSM8K = SCRIPT_DIR / "bench" / "data" / "gsm8k.json"
+SERVER_BIN = SCRIPT_DIR.parent / "target" / "release" / "xserv-server"
+SYSTEM = ("You are a careful math problem solver. Solve the problem step by step. "
+          "Put your final numeric answer inside \\boxed{}.")
+
+_BOXED_RE = re.compile(r"\\boxed\s*\{([^{}]*)\}")
+_NUM_RE = re.compile(r"-?\d+(?:,\d{3})*(?:\.\d+)?")
+
+
+def normalize_num(s):
+    s = s.replace(",", "").strip()
+    try:
+        f = float(s)
+    except ValueError:
+        return None
+    return str(int(f)) if f == int(f) else f"{f:g}"
+
+
+def extract_answer(text):
+    if not text:
+        return None
+    boxed = _BOXED_RE.findall(text)
+    if boxed:
+        nums = _NUM_RE.findall(boxed[-1])
+        if nums:
+            return normalize_num(nums[-1])
+    nums = _NUM_RE.findall(text)
+    if nums:
+        return normalize_num(nums[-1])
+    return None
+
+
+def pct(vals, p):
+    if not vals:
+        return 0.0
+    s = sorted(vals)
+    i = max(0, min(len(s) - 1, int(round((p / 100.0) * (len(s) - 1)))))
+    return s[i]
+
+
+# ---------- server lifecycle ----------
+
+def gpu_mem_used_mb(gpus):
+    out = subprocess.check_output(
+        ["nvidia-smi", "--query-gpu=index,memory.used", "--format=csv,noheader,nounits"],
+        text=True)
+    used = {}
+    for line in out.strip().splitlines():
+        idx, mem = [x.strip() for x in line.split(",")]
+        used[int(idx)] = int(mem)
+    return max(used.get(g, 0) for g in gpus)
+
+
+def start_server(model_dir, port, tp, gpus, log_path):
+    env = dict(os.environ)
+    env["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpus)
+    cmd = [str(SERVER_BIN), str(model_dir), "--port", str(port),
+           "--tp", str(tp), "--max-seq-len", "2048", "--max-batch", "8"]
+    logf = open(log_path, "wb")
+    # New session so we can kill the whole process tree without touching ours.
+    p = subprocess.Popen(cmd, stdout=logf, stderr=subprocess.STDOUT,
+                         env=env, start_new_session=True)
+    return p
+
+
+def stop_server(p, gpus, drain_to_mb=2000, timeout=120):
+    if p.poll() is None:
+        try:
+            os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+        except ProcessLookupError:
+            pass
+    try:
+        p.wait(timeout=30)
+    except subprocess.TimeoutExpired:
+        try:
+            os.killpg(os.getpgid(p.pid), signal.SIGKILL)
+        except ProcessLookupError:
+            pass
+    # Wait for VRAM to drain so the next server starts clean.
+    t0 = time.time()
+    while time.time() - t0 < timeout:
+        if gpu_mem_used_mb(gpus) < drain_to_mb:
+            return
+        time.sleep(2)
+
+
+def wait_ready(base, model_id, timeout=900):
+    """Gate on a real 1-token generation, not /health (which lies during load)."""
+    t0 = time.time()
+    body = json.dumps({
+        "model": model_id,
+        "messages": [{"role": "user", "content": "hi"}],
+        "max_tokens": 1, "temperature": 0.0, "stream": False,
+    }).encode()
+    while time.time() - t0 < timeout:
+        try:
+            req = urllib.request.Request(base + "/v1/chat/completions", data=body,
+                                         headers={"Content-Type": "application/json"})
+            with urllib.request.urlopen(req, timeout=120) as r:
+                if r.status == 200:
+                    json.loads(r.read())
+                    return True
+        except Exception:
+            time.sleep(3)
+    return False
+
+
+# ---------- one streamed request ----------
+
+def stream_chat(base, model_id, user, max_tokens):
+    body = json.dumps({
+        "model": model_id,
+        "messages": [{"role": "system", "content": SYSTEM},
+                     {"role": "user", "content": user}],
+        "max_tokens": max_tokens, "temperature": 0.0, "stream": True,
+    }).encode()
+    req = urllib.request.Request(base + "/v1/chat/completions", data=body,
+                                 headers={"Content-Type": "application/json"})
+    t0 = time.perf_counter()
+    ttft = None
+    t_last = t0
+    n = 0
+    parts = []
+    with urllib.request.urlopen(req, timeout=300) as resp:
+        for raw in resp:
+            line = raw.decode("utf-8", "ignore").strip()
+            if not line.startswith("data:"):
+                continue
+            data = line[5:].strip()
+            if data == "[DONE]":
+                break
+            try:
+                obj = json.loads(data)
+            except json.JSONDecodeError:
+                continue
+            delta = obj["choices"][0].get("delta", {})
+            content = delta.get("content")
+            if content:
+                now = time.perf_counter()
+                if ttft is None:
+                    ttft = now - t0
+                n += 1
+                t_last = now
+                parts.append(content)
+    ttft = ttft if ttft is not None else (time.perf_counter() - t0)
+    decode_span = t_last - t0 - ttft
+    tpot = decode_span / (n - 1) if n > 1 else 0.0
+    return "".join(parts), ttft, tpot, n
+
+
+def run_eval(base, model_id, problems, max_tokens):
+    correct = 0
+    ttfts, tpots, toks = [], [], []
+    n_scored = 0
+    for i, prob in enumerate(problems):
+        q = prob["problem"].replace("\n", " ")
+        try:
+            text, ttft, tpot, n = stream_chat(base, model_id, q, max_tokens)
+        except Exception as e:
+            print(f"  [E] {i+1}/{len(problems)} {e}", flush=True)
+            continue
+        pred = extract_answer(text)
+        gold = normalize_num(prob["answer"])
+        ok = pred is not None and gold is not None and pred == gold
+        correct += int(ok)
+        n_scored += 1
+        ttfts.append(ttft * 1000.0)
+        if tpot > 0:
+            tpots.append(tpot * 1000.0)
+        toks.append(n)
+        mark = "✓" if ok else "✗"
+        print(f"  [{mark}] {i+1:3d}/{len(problems)} gold={prob['answer']:>7s} "
+              f"pred={str(pred):>7s} ttft={ttft*1000:6.1f}ms tpot={tpot*1000:5.1f}ms tok={n}",
+              flush=True)
+    return {
+        "accuracy": correct / max(n_scored, 1),
+        "correct": correct, "scored": n_scored,
+        "ttft_ms_median": pct(ttfts, 50), "ttft_ms_p90": pct(ttfts, 90),
+        "tpot_ms_median": pct(tpots, 50), "tpot_ms_p90": pct(tpots, 90),
+        "tok_per_s_median": (1000.0 / pct(tpots, 50)) if pct(tpots, 50) > 0 else 0.0,
+        "mean_tokens": sum(toks) / max(len(toks), 1),
+    }
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--fp8", required=True)
+    ap.add_argument("--bf16", required=True)
+    ap.add_argument("--gpus", default="0,1")
+    ap.add_argument("--tp", type=int, default=2)
+    ap.add_argument("--limit", type=int, default=150)
+    ap.add_argument("--max-tokens", type=int, default=512)
+    ap.add_argument("--port", type=int, default=18080)
+    ap.add_argument("--out", default=None)
+    args = ap.parse_args()
+
+    gpus = [int(g) for g in args.gpus.split(",")]
+    with open(GSM8K) as f:
+        problems = json.load(f)[:args.limit]
+
+    base = f"http://127.0.0.1:{args.port}"
+    results = {}
+    for label, model_dir in [("FP8_W8A8", args.fp8), ("BF16", args.bf16)]:
+        model_id = Path(model_dir).name
+        log_path = f"/tmp/xserv_{label}.log"
+        print(f"\n{'='*72}\n  {label}  ({model_dir}, tp={args.tp}, gpus={gpus})\n{'='*72}", flush=True)
+        print(f"  starting server (log: {log_path}) ...", flush=True)
+        p = start_server(model_dir, args.port, args.tp, gpus, log_path)
+        try:
+            if not wait_ready(base, model_id):
+                print(f"  SERVER NOT READY — tail of log:", flush=True)
+                print(subprocess.run(["tail", "-30", log_path], capture_output=True, text=True).stdout)
+                stop_server(p, gpus)
+                continue
+            print(f"  ready. running {len(problems)} GSM8K problems...", flush=True)
+            t0 = time.time()
+            r = run_eval(base, model_id, problems, args.max_tokens)
+            r["wall_s"] = time.time() - t0
+            results[label] = r
+            print(f"  -> acc={r['accuracy']*100:.1f}%  ttft_med={r['ttft_ms_median']:.1f}ms  "
+                  f"tpot_med={r['tpot_ms_median']:.1f}ms  ({r['tok_per_s_median']:.1f} tok/s)", flush=True)
+        finally:
+            print(f"  stopping server...", flush=True)
+            stop_server(p, gpus)
+
+    print(f"\n{'='*72}\n  SUMMARY (gpt-oss-20b, tp={args.tp}, GSM8K n={args.limit})\n{'='*72}")
+    print(f"{'metric':<26s} {'FP8_W8A8':>14s} {'BF16':>14s}")
+    print("-" * 56)
+    f8, b6 = results.get("FP8_W8A8", {}), results.get("BF16", {})
+    def row(name, key, fmt, scale=1.0):
+        a = f8.get(key); b = b6.get(key)
+        if a is None or b is None:
+            return
+        print(f"{name:<26s} {fmt.format(a*scale):>14s} {fmt.format(b*scale):>14s}")
+    row("GSM8K accuracy (%)", "accuracy", "{:.1f}", 100.0)
+    row("TTFT median (ms)", "ttft_ms_median", "{:.1f}")
+    row("TTFT p90 (ms)", "ttft_ms_p90", "{:.1f}")
+    row("TPOT median (ms)", "tpot_ms_median", "{:.2f}")
+    row("TPOT p90 (ms)", "tpot_ms_p90", "{:.2f}")
+    row("Throughput (tok/s)", "tok_per_s_median", "{:.1f}")
+    row("Mean output tokens", "mean_tokens", "{:.0f}")
+    if f8 and b6 and b6.get("tpot_ms_median"):
+        sp = b6["tpot_ms_median"] / f8["tpot_ms_median"] if f8.get("tpot_ms_median") else 0
+        print(f"\n  FP8 decode speedup vs BF16: {sp:.2f}x")
+
+    out = args.out or f"/tmp/fp8_compare_{int(time.time())}.json"
+    with open(out, "w") as f:
+        json.dump({"args": vars(args), "results": results}, f, indent=2)
+    print(f"\n  saved: {out}")
+
+
+if __name__ == "__main__":
+    main()