gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

quantization: single strided-batched FP8 MoE GEMM — cut per-token launches ~768→48

Browse Source 
The plan-cache fix removed the per-expert heuristic churn but still issued one
cublasLtMatmul per expert: ~768 tiny launches per decoded token (16 local
experts × 2 GEMMs × 24 layers), which capped the FP8 decode win at ~1.05× over
BF16. Collapse each MoE GEMM into ONE strided-batched cuBLASLt FP8 matmul
(BATCH_COUNT + strided-batch offsets on all four layouts) → ~48 launches/token.

A single strided call can't carry a per-batch scalar B-scale, so the per-expert
weight scale moves out of the GEMM epilogue into a fused post-scale kernel
(rowwise_scale_moe_bf16) that applies a_scale[token]·b_scale[expert] in one
pass. This is precision-equivalent: BF16's relative error is scale-invariant, so
scaling the unscaled GEMM output afterward loses nothing vs scaling in-epilogue.

Measured on dash5 (gpt-oss-20b, TP=2, 5090), warm-server GSM8K:
  decode TPOT 17.45 → 13.08 ms (FP8 now 1.41× vs BF16 18.39 ms),
  throughput 57.3 → 76.4 tok/s, accuracy unchanged (FP8 91.0% vs BF16 90.0%).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-12 01:23:29 +08:00
parent 24c49c31c2
commit e631a71b68
3 changed files with 150 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
crates/xserv-kernels/src/quantization.rs
View File

@@ -23,10 +23,11 @@ unsafe extern "C" {
num_rows: i32, cols: i32,
stream: *mut c_void,
);
fn launch_rowwise_scale_bf16(
fn launch_rowwise_scale_moe_bf16(
data: *mut c_void,
scales: *const c_void,
num_rows: i32, cols: i32,
a_scales: *const c_void,
b_scales: *const c_void,
num_rows: i32, cols: i32, tokens: i32,
stream: *mut c_void,
);
}
@@ -136,11 +137,11 @@ struct Fp8Plan {
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.
/// Persistent device scalar holding 1.0, used as the A/B scale pointer.
/// Scales are applied post-GEMM, so the in-GEMM scales stay 1.0.
one_buf: GpuBuffer,
/// Cache of prepared matmul plans keyed by (M, N, K).
plans: HashMap<(usize, usize, usize), Fp8Plan>,
/// Cache of prepared matmul plans keyed by (M, N, K, batch).
plans: HashMap<(usize, usize, usize, usize), Fp8Plan>,
}
impl CublasLtContext {
@@ -154,14 +155,15 @@ impl CublasLtContext {
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)) {
/// Get the cached strided-batched plan for (m, n, k, batch), building it on
/// first use.
fn plan(&mut self, m: usize, n: usize, k: usize, batch: usize) -> Fp8Plan {
if let Some(p) = self.plans.get(&(m, n, k, batch)) {
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);
let plan = unsafe { build_fp8_plan(self.handle, one_ptr, m, n, k, batch) };
self.plans.insert((m, n, k, batch), plan);
plan
}
}
@@ -184,16 +186,18 @@ impl Drop for CublasLtContext {
}
}
/// 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.
/// Build a strided-batched FP8 matmul plan for `batch` experts of one
/// (m, n, k) shape. Row-major → cuBLASLt col-major mapping (transA=T,
/// transB=N, m_lt=N, n_lt=M, k_lt=K). A/B scale pointers stay at 1.0 — both
/// the per-expert weight scale and the per-token activation scale are applied
/// post-GEMM in a fused kernel, which lets all experts run in one matmul.
unsafe fn build_fp8_plan(
handle: CublasLtHandle,
one_ptr: *const c_void,
m: usize,
n: usize,
k: usize,
batch: usize,
) -> Fp8Plan {
let m_lt = n as u64;
let n_lt = m as u64;
@@ -209,17 +213,33 @@ unsafe fn build_fp8_plan(
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);
// Per-expert strides in ELEMENTS for the strided-batch layout.
let stride_a = (n * k) as i64; // weights [N, K]
let stride_b = (m * k) as i64; // activations [M, K]
let stride_c = (m * n) as i64; // output [M, N]
let bc = batch as i32;
let set_batch = |layout: CublasLtMatrixLayout, stride: i64| {
cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
&bc as *const i32 as _, 4);
cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
&stride as *const i64 as _, 8);
};
// "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);
set_batch(a_layout, stride_a);
// "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);
set_batch(b_layout, stride_b);
// "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);
set_batch(c_layout, stride_c);
let mut d_layout: CublasLtMatrixLayout = std::ptr::null_mut();
cublasLtMatrixLayoutCreate(&mut d_layout, CUDA_R_16BF, m_lt, n_lt, m_lt as i64);
set_batch(d_layout, stride_c);
let mut pref: CublasLtMatmulPreference = std::ptr::null_mut();
cublasLtMatmulPreferenceCreate(&mut pref);
@@ -233,7 +253,7 @@ unsafe fn build_fp8_plan(
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}");
"cublasLtMatmulAlgoGetHeuristic failed for batched FP8 GEMM (m={m}, n={n}, k={k}, batch={batch}): status={status}, found={found}");
cublasLtMatmulPreferenceDestroy(pref);
Fp8Plan {
@@ -354,71 +374,54 @@ pub fn batched_gemm_fp8(
let c = Tensor::empty(&[batch, m, n], DType::BF16, a_fp8.device());
// 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);
// Cached strided-batched plan: heuristic + descriptor/layout creation
// happen once per (m, n, k, batch). All experts run in ONE matmul.
let plan = ctx.plan(m, n, k, batch);
// 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).
// alpha=1, beta=0, in-GEMM scales=1.0. The unscaled result
// D_raw[e] = A_fp8[e] @ B_fp8[e]^T
// is recovered to the real value by the fused post-scale kernel below.
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).
let b_scale_ptr = unsafe { (b_scales.data_ptr() as *const u8).add(e * 4) as *const c_void };
unsafe {
cublasLtMatmulDescSetAttribute(
plan.desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
&b_scale_ptr as *const _ as _, ptr_sz,
);
let status = cublasLtMatmul(
handle, plan.desc,
&alpha as *const f32 as _,
b_ptr, // cuBLASLt "A" = weights
b_fp8_t.data_ptr() as *const c_void, // cuBLASLt "A" = weights
plan.a_layout,
a_ptr, // cuBLASLt "B" = activations
a_fp8.data_ptr() as *const c_void, // cuBLASLt "B" = activations
plan.b_layout,
&beta as *const f32 as _,
c_ptr, // C (unused with beta=0)
c.data_ptr() as *const c_void, // C (unused with beta=0)
plan.c_layout,
c_ptr, // D = output
c.data_ptr() as *mut c_void, // D = output
plan.d_layout,
&plan.algo,
ws_ptr,
plan.workspace_size,
std::ptr::null_mut(),
);
assert_eq!(status, 0, "cublasLtMatmul FP8 failed for expert {e}: status={status}");
}
assert_eq!(status, 0, "batched cublasLtMatmul FP8 failed: status={status}");
}
});
// Post-GEMM: multiply each row of c by its activation scale.
// c is [batch, M, N] BF16. a_scales is [batch * M] F32.
// This recovers the per-token scale that was divided out during quantization.
// Post-GEMM: recover the real result in one pass.
// c[e, t, :] *= a_scales[e*M + t] * b_scales[e]
// (per-token activation scale × per-expert weight scale). BF16's relative
// error is scale-invariant, so applying the scale here is precision-
// equivalent to folding it into the GEMM epilogue.
let total_rows = (batch * m) as i32;
let cols = n as i32;
unsafe {
launch_rowwise_scale_bf16(
launch_rowwise_scale_moe_bf16(
c.data_ptr() as *mut c_void,
a_scales.data_ptr() as *const c_void,
total_rows, cols,
b_scales.data_ptr() as *const c_void,
total_rows, n as i32, m as i32,
std::ptr::null_mut(),
);
}

37
csrc/quantization/quantize_fp8.cu
View File

@@ -86,6 +86,29 @@ __global__ void rowwise_scale_bf16_kernel(
}
}
// Combined dequant scale for batched MoE FP8 GEMM output.
// data[row, :] *= a_scales[row] * b_scales[row / tokens]
// where row = expert * tokens + token. a_scales is the per-token activation
// scale; b_scales is the per-expert scalar weight scale. Lets a single
// strided-batched FP8 matmul (alpha=1, scales=1) recover the real result in
// one pass instead of folding the weight scale into a per-expert GEMM call.
__global__ void rowwise_scale_moe_bf16_kernel(
__nv_bfloat16* __restrict__ data,
const float* __restrict__ a_scales,
const float* __restrict__ b_scales,
int num_rows, int cols, int tokens
) {
int row = blockIdx.x;
if (row >= num_rows) return;
int tid = threadIdx.x;
float s = a_scales[row] * b_scales[row / tokens];
__nv_bfloat16* row_data = data + (long long)row * cols;
for (int i = tid; i < cols; i += blockDim.x) {
float v = __bfloat162float(row_data[i]) * s;
row_data[i] = __float2bfloat16(v);
}
}
extern "C" {
void launch_rowwise_scale_bf16(
@@ -102,6 +125,20 @@ void launch_rowwise_scale_bf16(
CUDA_CHECK_LAST_ERROR();
}
void launch_rowwise_scale_moe_bf16(
void* data, const void* a_scales, const void* b_scales,
int num_rows, int cols, int tokens,
void* stream
) {
int block = 256;
int grid = num_rows;
rowwise_scale_moe_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
(__nv_bfloat16*)data, (const float*)a_scales, (const float*)b_scales,
num_rows, cols, tokens
);
CUDA_CHECK_LAST_ERROR();
}
void launch_quantize_bf16_to_fp8e4m3_rowwise(
const void* src,
void* dst,

70
docs/benchmarks/fp8-quantization.md
View File

@@ -14,9 +14,10 @@ stay BF16.
- **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`).
runs **one strided-batched cuBLASLt FP8 matmul for all experts** (`alpha=1`,
in-GEMM scales `1.0`); a fused kernel then applies `a_scale[token]·b_scale[expert]`
in a single pass. BF16's relative error is scale-invariant, so applying both
scales post-GEMM is precision-equivalent to folding them into the epilogue.
- Model size: **22 GB** (FP8) vs **39 GB** (BF16). The FP8 model fits on a
single 32 GB 5090; BF16 needs ≥ 2.
@@ -34,34 +35,50 @@ decoded token. This made FP8 **slower than BF16**:
| 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`.
in a thread-local map keyed by `(M, N, K, batch)` so the heuristic runs once per
shape, and allocate the scale buffer once.
## Results — GSM8K (200 problems, greedy, TP=2 on the same 2 GPUs)
## Reducing launches: one strided-batched matmul
The per-expert loop still issued one `cublasLtMatmul` per expert — ~768 tiny
launches per decoded token (16 local experts × 2 GEMMs × 24 layers). Collapsing
each MoE GEMM into a single **strided-batched** cuBLASLt FP8 matmul (BATCH_COUNT
+ strided-batch offsets) drops that to ~48, with a fused post-scale kernel
applying both scales. This required moving the per-expert weight scale out of the
GEMM epilogue (a single strided call can't carry a per-batch scalar) into the
post-scale kernel — precision-equivalent, as noted above.
| (gpt-oss-20b, TP=2) | per-expert FP8 | batched FP8 | BF16 |
|---|---|---|---|
| Decode TPOT | 17.9 ms | **13.8 ms** | 18.8 ms |
| Throughput | 55.8 tok/s | **72.3 tok/s** | 53.2 tok/s |
## Results — GSM8K (greedy, TP=2 on the same 2 GPUs)
200-problem run is the per-expert plan-cache fix; 100-problem run is the
strided-batched version. BF16 is the unchanged baseline in both.
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 |
| metric | FP8 per-expert (n=200) | FP8 batched (n=100) | BF16 |
|---|---|---|---|
| GSM8K accuracy | 93.0 % | 91.0 % | 90.5 / 90.0 % |
| TTFT median | 67.4 ms | 65.0 ms | 68.8 / 69.5 ms |
| TPOT median | 17.45 ms | **13.08 ms** | 18.26 / 18.39 ms |
| TPOT p90 | 17.65 ms | **13.28 ms** | 18.38 / 18.52 ms |
| Throughput | 57.3 tok/s | **76.4 tok/s** | 54.8 / 54.4 tok/s |
| Decode speedup vs BF16 | 1.05× | **1.41×** | 1.00× |
- **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.
- **Accuracy: unchanged.** FP8 is nominally +0.5 … +2.5 pts above BF16, but at
n=100–200 the standard error is ~2–3 pts, so they are statistically
indistinguishable. The takeaway is that neither FP8 quantization nor the
strided-batched rounding degrades accuracy.
- **Decode: FP8 1.41× faster** once batched (TPOT 13.08 vs 18.39 ms), with a
tight p90. The per-expert version was only ~1.05× — the ~768 tiny M=1 launches
per token dominated; batching them into ~48 unlocked most of the FP8
expert-weight-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
@@ -75,9 +92,8 @@ 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.
- Sparse (top-k only) MoE compute instead of dense — currently every token runs
all experts, so only ~top_k/num_experts of the FP8 GEMM work is used.