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:
@@ -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;
|
||||
|
||||
46
crates/xserv-kernels/src/quantization.rs
Normal file
46
crates/xserv-kernels/src/quantization.rs
Normal 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
|
||||
}
|
||||
Reference in New Issue
Block a user