quantization: add FP8 E4M3 W8A16 for gpt-oss MoE expert weights

Store expert gate_up_proj and down_proj weights in FP8 E4M3 (1 byte/elem)
with per-expert FP32 scale factors. At inference, a fused CUDA kernel
dequantizes to BF16 before the existing cuBLAS batched GEMM.

Results on gpt-oss-20b (50-problem GSM8K subset):
  - FP8 TP=1: 47/50 = 94.0% (single RTX 5090, ~25 GB VRAM)
  - BF16 TP=2: 47/50 = 94.0% (requires 2× RTX 5090, ~39 GB total)

No measurable accuracy degradation. Model size: 41.8 GB → 22.7 GB (−46%).

New files:
  - tools/quantize_fp8.py: offline BF16→FP8 conversion script
  - csrc/quantization/dequant_fp8.cu: per-expert-scale dequant kernel
  - crates/xserv-kernels/src/quantization.rs: Rust FFI wrapper
  - tools/eval_gsm8k_batch.sh: GSM8K accuracy evaluation harness

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-07 19:33:07 +08:00
parent e1eb77baa4
commit 9f1fbbb98b
10 changed files with 474 additions and 6 deletions

View File

@@ -6,6 +6,7 @@ pub mod embedding;
pub mod gemm;
pub mod layernorm;
pub mod moe;
pub mod quantization;
pub mod rmsnorm;
pub mod rope;
pub mod softmax;

View File

@@ -0,0 +1,46 @@
use std::ffi::c_void;
use xserv_tensor::{DType, Tensor};
unsafe extern "C" {
fn launch_dequant_fp8e4m3_to_bf16(
src: *const c_void,
scales: *const c_void,
dst: *mut c_void,
num_experts: i32, rows: i32, cols: i32,
stream: *mut c_void,
);
}
/// Dequantize a 3D FP8 E4M3 tensor to BF16 using per-expert FP32 scales.
///
/// src: [num_experts, rows, cols] FP8E4M3, contiguous, GPU
/// scales: [num_experts] F32, contiguous, GPU
///
/// Returns: [num_experts, rows, cols] BF16
pub fn dequant_fp8_to_bf16(src: &Tensor, scales: &Tensor) -> Tensor {
assert_eq!(src.ndim(), 3, "dequant_fp8_to_bf16: src must be 3D");
assert_eq!(src.dtype(), DType::FP8E4M3);
assert!(src.is_contiguous());
assert_eq!(scales.ndim(), 1);
assert_eq!(scales.dtype(), DType::F32);
assert!(scales.is_contiguous());
let num_experts = src.shape()[0];
let rows = src.shape()[1];
let cols = src.shape()[2];
assert_eq!(scales.shape()[0], num_experts);
let out = Tensor::empty(&[num_experts, rows, cols], DType::BF16, src.device());
unsafe {
launch_dequant_fp8e4m3_to_bf16(
src.data_ptr() as *const c_void,
scales.data_ptr() as *const c_void,
out.data_ptr() as *mut c_void,
num_experts as i32, rows as i32, cols as i32,
std::ptr::null_mut(),
);
}
out
}