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

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

Browse Source 
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:
Gahow Wang
2026-06-07 19:33:07 +08:00
parent e1eb77baa4
commit 9f1fbbb98b
10 changed files with 474 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
crates/xserv-model/src/gpt_oss.rs
View File

@@ -43,10 +43,15 @@ struct GptOssBlock {
router_wt: Tensor,
router_bias: Tensor,
// 3D expert weights for batched GEMM (contiguous on GPU)
expert_gate_up_wt: Tensor, // [local_experts, hidden, 2*inter]
expert_gate_up_wt: Tensor, // [local_experts, hidden, 2*inter] BF16
expert_gate_up_bias: Tensor, // [local_experts, 2*inter]
expert_down_wt: Tensor, // [local_experts, inter, hidden]
expert_down_wt: Tensor, // [local_experts, inter, hidden] BF16
expert_down_bias: Tensor, // [local_experts, hidden]
// FP8 quantized expert weights (Some when running FP8 W8A16)
expert_gate_up_fp8: Option<Tensor>, // [local_experts, hidden, 2*inter] FP8E4M3
expert_gate_up_scale: Option<Tensor>,// [local_experts] F32
expert_down_fp8: Option<Tensor>, // [local_experts, inter, hidden] FP8E4M3
expert_down_scale: Option<Tensor>, // [local_experts] F32
local_experts: usize,
// Activation params
glu_alpha: f32,
@@ -156,17 +161,49 @@ impl GptOss {
let down_3d = take(&mut w, &format!("{p}.mlp.experts.down_proj"));
let down_bias_2d = take(&mut w, &format!("{p}.mlp.experts.down_proj_bias"));
// FP8 scale tensors (present only in FP8-quantized models)
let gate_up_scale = w.remove(&format!("{p}.mlp.experts.gate_up_proj_scale"));
let down_scale = w.remove(&format!("{p}.mlp.experts.down_proj_scale"));
let local_experts = num_experts / world;
let expert_start = rank * local_experts;
let is_fp8 = gate_up_3d.dtype() == xserv_tensor::DType::FP8E4M3;
let inter2 = gate_up_3d.shape()[2]; // 2 * intermediate_size
let hidden = gate_up_3d.shape()[1];
let inter = down_3d.shape()[1]; // intermediate_size
// Slice the rank's range of experts as contiguous 3D tensors on GPU
let expert_gate_up_wt = slice_expert_range_3d(&gate_up_3d, expert_start, local_experts, hidden, inter2).to_device(dev);
let expert_gate_up_wt;
let expert_down_wt;
let expert_gate_up_fp8;
let expert_gate_up_scale_gpu;
let expert_down_fp8;
let expert_down_scale_gpu;
if is_fp8 {
// FP8 path: load quantized weights and scales
expert_gate_up_fp8 = Some(slice_expert_range_3d_raw(&gate_up_3d, expert_start, local_experts, hidden, inter2).to_device(dev));
expert_down_fp8 = Some(slice_expert_range_3d_raw(&down_3d, expert_start, local_experts, inter, hidden).to_device(dev));
// Scales: [num_experts] F32 → slice to [local_experts]
let gu_s = gate_up_scale.expect("FP8 model missing gate_up_proj_scale");
let d_s = down_scale.expect("FP8 model missing down_proj_scale");
expert_gate_up_scale_gpu = Some(slice_scale_range(&gu_s, expert_start, local_experts).to_device(dev));
expert_down_scale_gpu = Some(slice_scale_range(&d_s, expert_start, local_experts).to_device(dev));
// Dummy BF16 tensors (never read in FP8 path)
expert_gate_up_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
expert_down_wt = Tensor::empty(&[1, 1, 1], xserv_tensor::DType::BF16, dev);
} else {
// BF16 path: existing behavior
expert_gate_up_wt = slice_expert_range_3d(&gate_up_3d, expert_start, local_experts, hidden, inter2).to_device(dev);
expert_down_wt = slice_expert_range_3d(&down_3d, expert_start, local_experts, inter, hidden).to_device(dev);
expert_gate_up_fp8 = None;
expert_gate_up_scale_gpu = None;
expert_down_fp8 = None;
expert_down_scale_gpu = None;
}
let expert_gate_up_bias = slice_expert_range_2d(&gate_up_bias_2d, expert_start, local_experts, inter2).to_device(dev);
let expert_down_wt = slice_expert_range_3d(&down_3d, expert_start, local_experts, inter, hidden).to_device(dev);
let expert_down_bias = slice_expert_range_2d(&down_bias_2d, expert_start, local_experts, hidden).to_device(dev);
xserv_cuda::allocator::cached_trim();
@@ -198,6 +235,10 @@ impl GptOss {
expert_gate_up_bias,
expert_down_wt,
expert_down_bias,
expert_gate_up_fp8,
expert_gate_up_scale: expert_gate_up_scale_gpu,
expert_down_fp8,
expert_down_scale: expert_down_scale_gpu,
local_experts,
glu_alpha,
glu_limit,
@@ -208,10 +249,14 @@ impl GptOss {
let local_num_kv_heads = config.num_kv_heads() / world;
let has_norm_bias = norm_bias.is_some();
let is_fp8 = layers.first().map(|l| l.expert_gate_up_fp8.is_some()).unwrap_or(false);
if rank == 0 {
if has_norm_bias {
eprintln!("gpt-oss: detected LayerNorm bias — using LayerNorm instead of RMSNorm");
}
if is_fp8 {
eprintln!("gpt-oss: FP8 E4M3 quantized expert weights detected (W8A16 mode)");
}
}
// Warn about unused weights that the model didn't consume
@@ -470,7 +515,12 @@ impl GptOss {
let x_rep = xserv_kernels::moe::moe_replicate(x, local_experts);
// 4. Batched GEMM gate_up: [E, tokens, hidden] @ [E, hidden, 2*inter] → [E, tokens, 2*inter]
let gate_up = xserv_kernels::moe::batched_gemm_strided(&x_rep, &layer.expert_gate_up_wt);
let gate_up_wt = if let Some(ref fp8) = layer.expert_gate_up_fp8 {
xserv_kernels::quantization::dequant_fp8_to_bf16(fp8, layer.expert_gate_up_scale.as_ref().unwrap())
} else {
layer.expert_gate_up_wt.clone()
};
let gate_up = xserv_kernels::moe::batched_gemm_strided(&x_rep, &gate_up_wt);
// 5. Bias add: gate_up += expert_gate_up_bias (in-place)
xserv_kernels::moe::moe_bias_add_3d(&gate_up, &layer.expert_gate_up_bias);
@@ -484,7 +534,12 @@ impl GptOss {
let activated = activated_flat.reshape(&[local_experts, num_tokens, inter]);
// 7. Batched GEMM down: [E, tokens, inter] @ [E, inter, hidden] → [E, tokens, hidden]
let down = xserv_kernels::moe::batched_gemm_strided(&activated, &layer.expert_down_wt);
let down_wt = if let Some(ref fp8) = layer.expert_down_fp8 {
xserv_kernels::quantization::dequant_fp8_to_bf16(fp8, layer.expert_down_scale.as_ref().unwrap())
} else {
layer.expert_down_wt.clone()
};
let down = xserv_kernels::moe::batched_gemm_strided(&activated, &down_wt);
// 8. Bias add: down += expert_down_bias (in-place)
xserv_kernels::moe::moe_bias_add_3d(&down, &layer.expert_down_bias);
@@ -581,6 +636,28 @@ fn shard_1d(t: &Tensor, rank: usize, world: usize) -> Tensor {
Tensor::from_slice(&shard, &[local])
}
/// Extract experts [start..start+count) from a [num_experts, rows, cols] 3D tensor (any dtype, raw bytes).
fn slice_expert_range_3d_raw(t: &Tensor, start: usize, count: usize, rows: usize, cols: usize) -> Tensor {
assert_eq!(t.ndim(), 3);
let host = t.to_device(Device::Cpu);
let elem_size = t.dtype().size_bytes();
let raw = host.as_raw_bytes();
let stride = rows * cols * elem_size;
let offset = start * stride;
let slice = &raw[offset..offset + count * stride];
Tensor::from_raw_bytes(slice, &[count, rows, cols], t.dtype())
}
/// Slice scale tensor [num_experts] F32 → [count] starting at `start`.
fn slice_scale_range(t: &Tensor, start: usize, count: usize) -> Tensor {
assert_eq!(t.ndim(), 1);
assert_eq!(t.dtype(), xserv_tensor::DType::F32);
let host = t.to_device(Device::Cpu);
let data = host.as_slice::<f32>();
let slice = data[start..start + count].to_vec();
Tensor::from_slice(&slice, &[count])
}
/// Extract experts [start..start+count) from a [num_experts, rows, cols] 3D tensor
fn slice_expert_range_3d(t: &Tensor, start: usize, count: usize, rows: usize, cols: usize) -> Tensor {
assert_eq!(t.ndim(), 3);

4
crates/xserv-model/src/loader.rs
View File

@@ -19,6 +19,7 @@ pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor>
safetensors::Dtype::F32 => DType::F32,
safetensors::Dtype::F16 => DType::F16,
safetensors::Dtype::BF16 => DType::BF16,
safetensors::Dtype::F8_E4M3 => DType::FP8E4M3,
other => {
eprintln!("skipping tensor {name}: unsupported dtype {other:?}");
continue;
@@ -83,5 +84,8 @@ fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
};
Tensor::from_slice(bfs, shape)
}
DType::FP8E4M3 => {
Tensor::from_raw_bytes(raw_bytes, shape, DType::FP8E4M3)
}
}
}