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

gpt-oss: drop debug syncs from forward; GPU broadcast bias-add

Browse Source 
Decode carried three leftover cudaDeviceSynchronize (prefill one) from
NaN debugging — the Qwen3 path has none and the logits D2H in sample()
already orders against the null stream.

add_bias for rows>1 round-tripped the bias through the CPU (D2H + host
tile loop + H2D) on every call — 96 times per prefill across q/k/v/o.
Replace with a bias_add_2d broadcast kernel.

dash5, FP8 TP=2, warm-server: TTFT 35/49/94 -> 29/42/79 ms
(short/medium/long), TPOT 7.19-7.32 -> 6.99-7.21 ms. Greedy tokens
unchanged; GSM8K-50 94%.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-12 17:02:59 +08:00
parent 63f5599717
commit 1897b2e17a
4 changed files with 50 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -450,12 +450,8 @@ impl GptOss {
paged_cache.advance_seq_len(slot, 1);
}
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
let x = Self::norm(&x, &self.norm, &self.norm_bias, eps);
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
let logits = matmul_2d(&x, &self.lm_head_t);
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
logits
matmul_2d(&x, &self.lm_head_t)
}
/// Paged prefill: process full prompt tokens.
@@ -519,9 +515,7 @@ impl GptOss {
}
let x = Self::norm(&x, &self.norm, &self.norm_bias, eps);
let logits = matmul_2d(&x, &self.lm_head_t);
unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); }
logits
matmul_2d(&x, &self.lm_head_t)
}
/// MoE forward pass — fully on GPU via batched GEMM.
@@ -691,31 +685,12 @@ fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
matmul(a, b, GemmBackend::CuBlas)
}
/// Add bias to a 2D tensor: [rows, cols] + [cols] → [rows, cols]
/// Add bias to a 2D tensor: [rows, cols] + [cols] → [rows, cols].
/// Single GPU broadcast kernel — the old rows>1 path tiled the bias on the
/// CPU (D2H + host loop + H2D) on every call, 96×/prefill in the hot path.
fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
assert_eq!(x.ndim(), 2);
assert_eq!(bias.ndim(), 1);
let rows = x.shape()[0];
let cols = x.shape()[1];
assert_eq!(bias.shape()[0], cols, "bias size {} != cols {}", bias.shape()[0], cols);
let x_c = x.contiguous();
if rows == 1 {
// Fast path: reshape bias [cols] → [1, cols] (zero-copy), add directly on GPU
let bias_2d = bias.reshape(&[1, cols]);
return xserv_kernels::add(&x_c, &bias_2d);
}
// General path: tile bias to [rows, cols] via CPU, then add on GPU
let bias_cpu = bias.to_device(Device::Cpu);
let bias_data = bias_cpu.as_slice::<bf16>();
let mut tiled = Vec::with_capacity(rows * cols);
for _ in 0..rows {
tiled.extend_from_slice(bias_data);
}
let bias_tiled = Tensor::from_slice(&tiled, &[rows, cols]).to_device(x.device());
xserv_kernels::add(&x_c, &bias_tiled)
xserv_kernels::bias_add_2d(&x_c, bias)
}
fn shard_rows(t: &Tensor, rank: usize, world: usize) -> Tensor {