From fcd7fa62b76e478fe720b348045cf82069ca7232 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Fri, 29 May 2026 22:29:55 +0800 Subject: [PATCH] =?UTF-8?q?moe(wip):=20decode-path=20debugging=20tools=20?= =?UTF-8?q?=E2=80=94=20decode=20STILL=20broken?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Honest checkpoint. Builds green. The KV-cache decode path (decode_step/ generate) is NOT correct: top-1 diverges from the verified host-attention forward and is non-deterministic run-to-run; generation is garbage. Two attempted fixes are included but did NOT solve it: gemm::matmul_dense (cuBLAS without the m==1 GEMV shortcut) and zeroing the dequant output. What IS verified and correct (unchanged): the host-attention forward (gptoss-logits forward path -> top-1 " Paris"), the MXFP4 GPU dequant kernel (mxfp4-check == numpy), and the sink-attention kernel in isolation (sink-attn-check == CPU ref, max_diff 0.0017). So the decode bug is in how decode_step composes these (KV layout / per-step state), not in the kernels themselves. Root cause still OPEN; see docs/MOE_PROGRESS.md. Co-Authored-By: Claude Opus 4.8 --- crates/xserv-model/src/bin/sink-attn-check.rs | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 crates/xserv-model/src/bin/sink-attn-check.rs diff --git a/crates/xserv-model/src/bin/sink-attn-check.rs b/crates/xserv-model/src/bin/sink-attn-check.rs new file mode 100644 index 0000000..433611f --- /dev/null +++ b/crates/xserv-model/src/bin/sink-attn-check.rs @@ -0,0 +1,57 @@ +//! Unit-test decode_attention_sink (GPU) against a CPU reference on small random +//! inputs. Isolates the sink/window attention kernel from the rest of gpt-oss. +use half::bf16; +use xserv_kernels::decode_attention_sink; +use xserv_tensor::{Device, Tensor}; + +fn main() { + let (n_heads, n_kv, head_dim, kv_len) = (8usize, 2usize, 4usize, 3usize); + let scale = (head_dim as f32).powf(-0.5); + let n_rep = n_heads / n_kv; + + // deterministic pseudo-random fill + let mut seed = 12345u64; + let mut rnd = || { seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); ((seed >> 33) as f32 / (1u64 << 31) as f32) - 1.0 }; + + let q: Vec = (0..n_heads * head_dim).map(|_| rnd()).collect(); // [H,1,D] + let k: Vec = (0..n_kv * kv_len * head_dim).map(|_| rnd()).collect(); // [Hkv,kv,D] + let v: Vec = (0..n_kv * kv_len * head_dim).map(|_| rnd()).collect(); + let sinks: Vec = (0..n_heads).map(|_| rnd()).collect(); + + let qb: Vec = q.iter().map(|&x| bf16::from_f32(x)).collect(); + let kb: Vec = k.iter().map(|&x| bf16::from_f32(x)).collect(); + let vb: Vec = v.iter().map(|&x| bf16::from_f32(x)).collect(); + + let qt = Tensor::from_slice(&qb, &[1, n_heads, 1, head_dim]).to_device(Device::Cuda(0)); + let kt = Tensor::from_slice(&kb, &[1, n_kv, kv_len, head_dim]).to_device(Device::Cuda(0)); + let vt = Tensor::from_slice(&vb, &[1, n_kv, kv_len, head_dim]).to_device(Device::Cuda(0)); + let st = Tensor::from_slice(&sinks, &[n_heads]).to_device(Device::Cuda(0)); // f32 + + let out = decode_attention_sink(&qt, &kt, &vt, &st, scale, 0); // [1,H,1,D] + let outh = out.to_device(Device::Cpu); + let og = outh.as_slice::(); + + // CPU reference: for each q head, softmax over [s_0..s_{kv-1}, sink], drop sink, weight V. + let mut max_diff = 0f32; + for h in 0..n_heads { + let kv = h / n_rep; + let mut s = vec![0f32; kv_len]; + let mut m = sinks[h]; + for j in 0..kv_len { + let mut dot = 0f32; + for d in 0..head_dim { dot += q[h * head_dim + d] * k[(kv * kv_len + j) * head_dim + d]; } + s[j] = dot * scale; + if s[j] > m { m = s[j]; } + } + let mut denom = (sinks[h] - m).exp(); + for j in 0..kv_len { denom += (s[j] - m).exp(); } + for d in 0..head_dim { + let mut acc = 0f32; + for j in 0..kv_len { acc += (s[j] - m).exp() / denom * v[(kv * kv_len + j) * head_dim + d]; } + let got = og[h * head_dim + d].to_f32(); + max_diff = max_diff.max((got - acc).abs()); + } + } + println!("decode_attention_sink vs CPU ref: max_abs_diff = {max_diff:.5}"); + println!("{}", if max_diff < 0.05 { "SINK_KERNEL_OK" } else { "SINK_KERNEL_MISMATCH" }); +}