moe: gpt-oss-20b forward verified correct (predicts "Paris")

YaRN RoPE was the missing piece — gpt-oss uses rope_type "yarn" (factor 32,
beta_fast 32, beta_slow 1, orig_max 4096); a plain theta RoPE garbled
attention. Added yarn_rope_cache (host-computed inv_freq + mscale, built
into a RopeCache directly). Experts kept CPU-resident and uploaded per-use
(the dequantized BF16 model is ~36GB, won't fit one 32GB card).

Verified: "The capital of France is" -> top-1 token 12366 = " Paris"
(logit 19.75), matching the llama.cpp oracle's behavior. This exercises the
full MoE path: top-4 router (softmax-after-topk), interleaved clamped
(up+1)*glu experts, attention sinks, sliding window, MXFP4->BF16 weights,
YaRN RoPE, head_dim 64, q/k/v/o biases.

Correctness-first (host attention + per-token MoE); GPU attention-with-sinks
kernel, KV cache, faster MoE, and PP-for-memory come next to run AIME/GSM8K
at speed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Cargo.lock

@@ -408,28 +408,12 @@ version = "2.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
 
-[[package]]
-name = "memo-map"
-version = "0.3.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b"
-
 [[package]]
 name = "mime"
 version = "0.3.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
 
-[[package]]
-name = "minijinja"
-version = "2.20.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2929e494b2280e1e18959bb2e121da03347ae896896fdfaceaab43c88a02803f"
-dependencies = [
- "memo-map",
- "serde",
-]
-
 [[package]]
 name = "mio"
 version = "1.2.0"
@@ -1113,14 +1097,6 @@ dependencies = [
  "rand 0.9.4",
 ]
 
-[[package]]
-name = "xserv-distributed"
-version = "0.1.0"
-dependencies = [
- "half",
- "xserv-cuda",
-]
-
 [[package]]
 name = "xserv-kernels"
 version = "0.1.0"
@@ -1136,14 +1112,12 @@ name = "xserv-model"
 version = "0.1.0"
 dependencies = [
  "half",
- "libc",
  "rand 0.8.6",
  "safetensors",
  "serde",
  "serde_json",
  "smallvec",
  "xserv-cuda",
- "xserv-distributed",
  "xserv-kernels",
  "xserv-tensor",
  "xserv-tokenizer",
@@ -1155,14 +1129,12 @@ version = "0.1.0"
 dependencies = [
  "axum",
  "half",
- "minijinja",
  "serde",
  "serde_json",
  "tokio",
  "tokio-stream",
  "uuid",
  "xserv-cuda",
- "xserv-distributed",
  "xserv-kernels",
  "xserv-model",
  "xserv-tensor",

Cargo.toml

@@ -28,4 +28,3 @@ axum = "0.8"
 uuid = { version = "1", features = ["v4"] }
 tokio-stream = "0.1"
 rand = "0.8"
-minijinja = { version = "2", features = ["builtins"] }

README.md

@@ -3,24 +3,18 @@
 > 从零用 **Rust + CUDA** 构建的 LLM 推理引擎，目标是吃透 LLM Serving 全栈技术。
 
 xserv 不依赖 PyTorch / vLLM / TensorRT 等现成框架，自己实现了张量抽象、CUDA kernel、
-分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。支持 **Qwen3-8B**（BF16）
-和 **gpt-oss-20b**（MoE，BF16/FP8/MXFP4 量化），多卡 TP/PP，并提供一套与 **llama.cpp**
-对比正确性和性能的标准 benchmark。
+分词器、模型前向、KV cache、调度器和 OpenAI 兼容的 HTTP 服务。当前在单张 RTX 5090 上可以
+跑通 **Qwen3-8B**（BF16），并提供一套与 **llama.cpp** 对比正确性和性能的标准 benchmark。
 
 ## 现状一览
 
-- **模型**：GPT-2（124M）、Qwen3-8B（BF16）、gpt-oss-20b（32 专家 top-4 MoE，harmony 格式）
-- **性能**（RTX 5090，贪心，单流）：
-  - Qwen3-8B BF16 单卡：约 56 tok/s（HF transformers 的 1.4×）
-  - gpt-oss-20b FP8 稀疏 MoE + CUDA Graph decode：**TPOT 5.8ms（~172 tok/s，
-    TP=1/2 同速）**；同配置 TP=2 全面快于 llama.cpp（1.26-1.47×），llama
-    单卡模式（2.8ms）仍领先，差距 2.0×
-- **精度**：GSM8K 全量与 llama.cpp 同权重持平（94.5% vs 94.4%）；FP8/MXFP4 量化无回归
-- **服务**：OpenAI 兼容 `/v1/chat/completions`，SSE 流式；gpt-oss 量化后可**单卡 32GB 服务**
-- **关键能力**：自写 GEMM / Flash-Attention 2(SM120，含 attention sinks + sliding window) /
-  Paged-Attention kernel、分页 KV cache（含 **CPU 换出/换入**）、连续批处理、
-  CUDA Graph 解码（Qwen3 单卡 + gpt-oss 全路径整图回放）、**Tensor/Pipeline 并行**（NCCL，TP=1/2/4、PP=2/4）、
-  **FP8 W8A8 / MXFP4 W4A16 量化**、**稀疏 top-k MoE decode**（只算被路由的专家）
+- **模型**：GPT-2（124M）、Qwen3-8B（BF16）
+- **性能**（RTX 5090，Qwen3-8B BF16，贪心解码，单流）：约 **56 tok/s**，约为 HF transformers 的 1.4×、llama.cpp 的 ~0.6×
+- **精度**：在 AIME 2025 / GSM8K 上与 llama.cpp 同权重对比基本持平（数值保真度验证通过）
+- **服务**：OpenAI 兼容 `/v1/chat/completions`，支持 SSE 流式输出
+- **关键能力**：自写 GEMM / Flash-Attention 2(SM120) / Paged-Attention kernel、
+  分页 KV cache（含 **CPU 换出/换入** 弹性显存）、连续批处理（continuous batching）、
+  CUDA Graph 解码、按显存自适应的 KV 池
 
 > 这是一个以学习为主的项目，逐 Phase 推进，每步都做数值/端到端验证。
 
@@ -32,19 +26,16 @@ xserv/
 │   ├── gemm/              #   GEMM (naive / tiled / gemv)
 │   ├── attention/         #   Flash-Attention 2 (SM120)、Paged-Attention、causal mask
 │   ├── normalization/     #   LayerNorm / RMSNorm
-│   ├── activation/        #   GELU / SiLU / gpt-oss GLU
+│   ├── activation/        #   GELU / SiLU
 │   ├── embedding/         #   embedding lookup / RoPE / transpose
-│   ├── moe/               #   MoE top-k 路由、稀疏专家 GEMV、加权求和
-│   ├── quantization/      #   FP8 量化/反量化、cuBLASLt FP8 GEMM、MXFP4 GEMV
 │   └── reduce/            #   softmax
 ├── crates/
 │   ├── xserv-cuda/        # CUDA FFI、Stream、显存分配器、Pinned 内存、CUDA Graph
 │   ├── xserv-tensor/      # Tensor 类型（strided 布局、BF16/F16/F32、CPU↔GPU）
 │   ├── xserv-kernels/     # kernel registry（自写 kernel + cuBLAS 可切换）
 │   ├── xserv-tokenizer/   # BPE 分词器
-│   ├── xserv-distributed/ # NCCL FFI、TP 上下文（AllReduce）
-│   ├── xserv-model/       # 模型定义（GPT-2 / Qwen3 / gpt-oss MoE）、权重加载、KV cache、采样
-│   └── xserv-server/      # tokio + axum HTTP 服务、调度器、TP/PP 引擎
+│   ├── xserv-model/       # 模型定义（GPT-2 / Qwen3）、权重加载、KV cache、采样
+│   └── xserv-server/      # tokio + axum HTTP 服务、调度器
 ├── tools/                 # 辅助脚本 + benchmark 套件（见下）
 └── docs/                  # 每个 Phase 的设计文档 + benchmark 报告
 ```
@@ -194,14 +185,12 @@ GSM8K 12 个格子全是 29/30，xserv 与 llama.cpp 完全一致；AIME 的 ±1
 
 ## 路线图（节选）
 
-已完成 Phase 0–21：CUDA 基础设施 → Tensor → GEMM → Transformer kernels → Attention →
+已完成 Phase 0–18：CUDA 基础设施 → Tensor → GEMM → Transformer kernels → Attention →
 模型加载 → 分词器 → GPT-2 → KV cache → Qwen3-8B → Paged Attention → 连续批处理 →
-HTTP API → Flash Attention 2 → 性能优化 → **张量并行（TP）** → **流水线并行（PP）** →
-**gpt-oss MoE + FP8/MXFP4 量化** → **稀疏 top-k MoE decode** → **decode CUDA Graph 整图回放**；
+HTTP API → Flash Attention 2 → 性能优化 → **张量并行（TP）** → **流水线并行（PP）**；
 并加入了 **llama.cpp 对比基准** 与 **KV CPU 换出** 等基础设施。
 
-后续方向：非专家权重量化（lm_head/qkv/o）、稀疏 prefill（grouped GEMM）、server 侧 harmony
-channel 分离、PP microbatch/1F1B、投机解码、多模态。详见 `docs/00-roadmap.md` 的实际进展记录。
+后续方向：PP microbatch/1F1B 流水线重叠（吞吐收益）、2D TP×PP、投机解码、量化（FP8 / INT8）、多模态。
 
 ## 许可
 

crates/xserv-cuda/src/allocator.rs

@@ -100,33 +100,10 @@ pub fn cached_alloc(size: usize) -> Result<GpuBuffer> {
     })
 }
 
-/// Free all cached (unused) GPU buffers back to the driver.
-pub fn cached_trim() {
-    ALLOCATOR.with(|cell| {
-        cell.borrow_mut().trim();
-    });
-}
-
 /// Return a raw GPU pointer to the caching allocator's free list.
 /// Called from `GpuBuffer::Drop` for pooled buffers. Takes raw pointer
 /// and size to avoid re-triggering Drop.
 pub fn return_to_pool(ptr: *mut u8, len: usize) {
-    // During CUDA graph capture, buffers freed by the captured code are
-    // quarantined instead of pooled: the instantiated graph references their
-    // addresses on every replay, so they must never be handed to another
-    // consumer for as long as the graph lives.
-    let quarantined = RETAINED.with(|cell| {
-        let mut r = cell.borrow_mut();
-        if let Some(list) = r.as_mut() {
-            list.push((ptr, len));
-            true
-        } else {
-            false
-        }
-    });
-    if quarantined {
-        return;
-    }
     ALLOCATOR.with(|cell| {
         let mut alloc = cell.borrow_mut();
         let bucket = bucket_size(len);
@@ -135,44 +112,6 @@ pub fn return_to_pool(ptr: *mut u8, len: usize) {
     });
 }
 
-thread_local! {
-    static RETAINED: RefCell<Option<Vec<(*mut u8, usize)>>> = const { RefCell::new(None) };
-}
-
-/// Buffers freed while a retain window was active. Holding this keeps their
-/// memory out of the pool; dropping it returns the blocks (on the owning
-/// thread) for reuse.
-pub struct RetainedBlocks(Vec<(*mut u8, usize)>);
-
-impl Drop for RetainedBlocks {
-    fn drop(&mut self) {
-        for (ptr, len) in self.0.drain(..) {
-            return_to_pool(ptr, len);
-        }
-    }
-}
-
-/// Start quarantining buffers freed on this thread (see `return_to_pool`).
-/// Must be paired with `end_retain` on the same thread; nesting unsupported.
-pub fn begin_retain() {
-    RETAINED.with(|cell| {
-        let mut r = cell.borrow_mut();
-        assert!(r.is_none(), "begin_retain: retain window already active");
-        *r = Some(Vec::new());
-    });
-}
-
-/// Stop quarantining and hand the quarantined blocks to the caller.
-pub fn end_retain() -> RetainedBlocks {
-    RETAINED.with(|cell| {
-        let list = cell
-            .borrow_mut()
-            .take()
-            .expect("end_retain without begin_retain");
-        RetainedBlocks(list)
-    })
-}
-
 /// Round up to next power-of-2, minimum 512 bytes.
 fn bucket_size(size: usize) -> usize {
     let min = 512;

crates/xserv-cuda/src/device.rs

@@ -48,7 +48,9 @@ pub fn device_info(device: u32) -> Result<DeviceInfo> {
     // Heap-allocate oversized buffer for cudaDeviceProp (layout varies by CUDA version).
     // CUDA 12.x struct is ~5-6 KB; use 32 KB to guard against future growth.
     let mut prop_buf = vec![0u8; 32768];
-    error::check(unsafe { ffi::cudaGetDeviceProperties(prop_buf.as_mut_ptr(), device as i32) })?;
+    error::check(unsafe {
+        ffi::cudaGetDeviceProperties(prop_buf.as_mut_ptr(), device as i32)
+    })?;
     // Name is always the first field: char[256].
     let name = unsafe { CStr::from_ptr(prop_buf.as_ptr() as *const c_char) }
         .to_string_lossy()

crates/xserv-cuda/src/ffi.rs

@@ -15,7 +15,6 @@ pub const CUDA_ERROR_OUT_OF_MEMORY: i32 = 2;
 
 /// cudaStreamCaptureMode::cudaStreamCaptureModeGlobal
 pub const CUDA_STREAM_CAPTURE_MODE_GLOBAL: i32 = 0;
-pub const CUDA_STREAM_CAPTURE_MODE_THREAD_LOCAL: i32 = 1;
 
 unsafe extern "C" {
     // --- Device ---
@@ -64,5 +63,11 @@ unsafe extern "C" {
     pub fn cudaGraphExecDestroy(graph_exec: CudaGraphExec) -> i32;
 
     // --- Our test kernel ---
-    pub fn launch_vecadd_f32(a: *const f32, b: *const f32, c: *mut f32, n: i32, stream: CudaStream);
+    pub fn launch_vecadd_f32(
+        a: *const f32,
+        b: *const f32,
+        c: *mut f32,
+        n: i32,
+        stream: CudaStream,
+    );
 }

crates/xserv-cuda/src/graph.rs

@@ -50,25 +50,31 @@ impl CudaGraph {
     pub fn begin_capture(&mut self, stream: &CudaStream) -> Result<()> {
         // If we have an old graph, destroy it first
         self.destroy_inner();
-        // THREAD_LOCAL: only "potentially unsafe" CUDA calls (cudaMalloc etc.)
-        // made by THIS thread invalidate the capture. With GLOBAL mode, TP rank
-        // threads capturing concurrently would poison each other's captures.
         error::check(unsafe {
-            ffi::cudaStreamBeginCapture(stream.as_raw(), ffi::CUDA_STREAM_CAPTURE_MODE_THREAD_LOCAL)
+            ffi::cudaStreamBeginCapture(
+                stream.as_raw(),
+                ffi::CUDA_STREAM_CAPTURE_MODE_GLOBAL,
+            )
         })
     }
 
     /// End capture and instantiate the executable graph.
     pub fn end_capture(&mut self, stream: &CudaStream) -> Result<()> {
-        error::check(unsafe { ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph) })?;
-        error::check(unsafe { ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0) })
+        error::check(unsafe {
+            ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph)
+        })?;
+        error::check(unsafe {
+            ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0)
+        })
     }
 
     /// Replay the captured graph on `stream`.
     /// Panics if no graph has been captured yet.
     pub fn launch(&self, stream: &CudaStream) -> Result<()> {
         assert!(self.is_ready(), "CudaGraph::launch called before capture");
-        error::check(unsafe { ffi::cudaGraphLaunch(self.exec, stream.as_raw()) })
+        error::check(unsafe {
+            ffi::cudaGraphLaunch(self.exec, stream.as_raw())
+        })
     }
 
     fn destroy_inner(&mut self) {

crates/xserv-cuda/src/lib.rs

@@ -11,4 +11,4 @@ pub use device::DeviceInfo;
 pub use error::{CudaError, Result};
 pub use graph::CudaGraph;
 pub use memory::{GpuBuffer, PinnedBuffer};
-pub use stream::{CudaStream, StreamGuard, current_stream_raw, push_stream};
+pub use stream::CudaStream;

crates/xserv-cuda/src/memory.rs

@@ -22,12 +22,7 @@ impl GpuBuffer {
         assert!(len > 0, "cannot allocate 0 bytes on GPU");
         let mut ptr = std::ptr::null_mut();
         error::check(unsafe { ffi::cudaMalloc(&mut ptr, len) })?;
-        Ok(Self {
-            ptr,
-            len,
-            owned: true,
-            pooled: false,
-        })
+        Ok(Self { ptr, len, owned: true, pooled: false })
     }
 
     /// Mark this buffer as pooled (returned to caching allocator on drop)
@@ -97,7 +92,9 @@ impl GpuBuffer {
     /// Copy from another GPU buffer (D2D).
     pub fn copy_from_device(&mut self, src: &GpuBuffer) -> Result<()> {
         let n = src.len.min(self.len);
-        error::check(unsafe { ffi::cudaMemcpy(self.ptr, src.ptr, n, ffi::CUDA_MEMCPY_D2D) })
+        error::check(unsafe {
+            ffi::cudaMemcpy(self.ptr, src.ptr, n, ffi::CUDA_MEMCPY_D2D)
+        })
     }
 
     /// Fill buffer with zeros.
@@ -106,13 +103,7 @@ impl GpuBuffer {
     }
 
     /// Copy `count` bytes from `src` buffer at `src_offset` to this buffer at `dst_offset`.
-    pub fn copy_from_device_at(
-        &mut self,
-        src: &GpuBuffer,
-        src_offset: usize,
-        dst_offset: usize,
-        count: usize,
-    ) -> Result<()> {
+    pub fn copy_from_device_at(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize) -> Result<()> {
         assert!(src_offset + count <= src.len);
         assert!(dst_offset + count <= self.len);
         error::check(unsafe {
@@ -126,14 +117,7 @@ impl GpuBuffer {
     }
 
     /// Async copy `count` bytes from `src` at `src_offset` to `self` at `dst_offset` on `stream`.
-    pub fn copy_from_device_at_async(
-        &mut self,
-        src: &GpuBuffer,
-        src_offset: usize,
-        dst_offset: usize,
-        count: usize,
-        stream: &CudaStream,
-    ) -> Result<()> {
+    pub fn copy_from_device_at_async(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize, stream: &CudaStream) -> Result<()> {
         assert!(src_offset + count <= src.len);
         assert!(dst_offset + count <= self.len);
         error::check(unsafe {
@@ -177,7 +161,9 @@ impl GpuBuffer {
 
     /// Async zero fill on stream.
     pub fn zero_async(&mut self, stream: &CudaStream) -> Result<()> {
-        error::check(unsafe { ffi::cudaMemsetAsync(self.ptr, 0, self.len, stream.as_raw()) })
+        error::check(unsafe {
+            ffi::cudaMemsetAsync(self.ptr, 0, self.len, stream.as_raw())
+        })
     }
 
     /// Consume the buffer without freeing GPU memory. Returns the raw pointer and length.
@@ -192,12 +178,7 @@ impl GpuBuffer {
     /// Reconstruct a GpuBuffer from a raw pointer + length.
     /// Safety: ptr must have been allocated with cudaMalloc, len must be correct.
     pub unsafe fn from_raw(ptr: *mut u8, len: usize) -> Self {
-        Self {
-            ptr,
-            len,
-            owned: true,
-            pooled: false,
-        }
+        Self { ptr, len, owned: true, pooled: false }
     }
 
     /// Create a non-owning view of GPU memory. Dropping this buffer does NOT
@@ -208,12 +189,7 @@ impl GpuBuffer {
     /// `ptr` must point to a valid GPU allocation of at least `len` bytes that
     /// will remain live for the lifetime of the returned `GpuBuffer`.
     pub unsafe fn borrow_raw(ptr: *mut u8, len: usize) -> Self {
-        Self {
-            ptr,
-            len,
-            owned: false,
-            pooled: false,
-        }
+        Self { ptr, len, owned: false, pooled: false }
     }
 }
 

crates/xserv-cuda/src/stream.rs

@@ -31,39 +31,3 @@ impl Drop for CudaStream {
 
 // Can move across threads, but not shared without synchronization
 unsafe impl Send for CudaStream {}
-
-// --- Thread-local launch stream -------------------------------------------
-//
-// Every kernel wrapper in xserv-kernels launches on `current_stream_raw()`,
-// which defaults to the legacy null stream (the historical behavior). CUDA
-// graph capture requires work to be issued on an explicit stream, so capture
-// code installs its stream here for the duration of the captured region via
-// `push_stream` / `StreamGuard`.
-
-use std::cell::Cell;
-
-thread_local! {
-    static CURRENT_STREAM: Cell<ffi::CudaStream> = const { Cell::new(std::ptr::null_mut()) };
-}
-
-/// The stream kernel launches on this thread should use (null = legacy default).
-pub fn current_stream_raw() -> ffi::CudaStream {
-    CURRENT_STREAM.with(|c| c.get())
-}
-
-/// RAII guard that installs a launch stream for the current thread and
-/// restores the previous one on drop.
-pub struct StreamGuard {
-    prev: ffi::CudaStream,
-}
-
-pub fn push_stream(stream: &CudaStream) -> StreamGuard {
-    let prev = CURRENT_STREAM.with(|c| c.replace(stream.as_raw()));
-    StreamGuard { prev }
-}
-
-impl Drop for StreamGuard {
-    fn drop(&mut self) {
-        CURRENT_STREAM.with(|c| c.set(self.prev));
-    }
-}

crates/xserv-cuda/tests/integration.rs

@@ -14,10 +14,7 @@ fn test_device_info() {
         info.compute_major, info.compute_minor
     );
     println!("  SM Count: {}", info.sm_count);
-    println!(
-        "  Shared Mem/Block: {} KB",
-        info.shared_mem_per_block / 1024
-    );
+    println!("  Shared Mem/Block: {} KB", info.shared_mem_per_block / 1024);
     println!("  Warp Size: {}", info.warp_size);
     println!("  Max Threads/Block: {}", info.max_threads_per_block);
 
@@ -148,11 +145,7 @@ fn test_caching_allocator() {
 
     // Second allocation of same size: should hit cache
     let _buf2 = alloc.alloc(1024).unwrap();
-    assert_eq!(
-        alloc.stats().cuda_malloc_count,
-        1,
-        "should reuse cached buffer"
-    );
+    assert_eq!(alloc.stats().cuda_malloc_count, 1, "should reuse cached buffer");
     assert_eq!(alloc.stats().cache_hit_count, 1);
 }
 
@@ -205,17 +198,11 @@ fn test_async_copy() {
     }
 
     let mut gpu = GpuBuffer::alloc(4096).unwrap();
-    unsafe {
-        gpu.copy_from_host_async(pinned.as_slice(), &stream)
-            .unwrap()
-    };
+    unsafe { gpu.copy_from_host_async(pinned.as_slice(), &stream).unwrap() };
     stream.synchronize().unwrap();
 
     let mut out_pinned = PinnedBuffer::alloc(4096).unwrap();
-    unsafe {
-        gpu.copy_to_host_async(out_pinned.as_mut_slice(), &stream)
-            .unwrap()
-    };
+    unsafe { gpu.copy_to_host_async(out_pinned.as_mut_slice(), &stream).unwrap() };
     stream.synchronize().unwrap();
 
     assert_eq!(pinned.as_slice(), out_pinned.as_slice());

crates/xserv-distributed/src/ffi.rs

@@ -34,12 +34,7 @@ pub const NCCL_SUCCESS: i32 = 0;
 unsafe extern "C" {
     pub fn ncclGetUniqueId(uid: *mut NcclUniqueId) -> i32;
     // ncclUniqueId is passed BY VALUE (a 128-byte struct) per the NCCL ABI.
-    pub fn ncclCommInitRank(
-        comm: *mut NcclComm,
-        nranks: i32,
-        commid: NcclUniqueId,
-        rank: i32,
-    ) -> i32;
+    pub fn ncclCommInitRank(comm: *mut NcclComm, nranks: i32, commid: NcclUniqueId, rank: i32) -> i32;
     pub fn ncclCommDestroy(comm: NcclComm) -> i32;
     pub fn ncclAllReduce(
         sendbuff: *const c_void,
@@ -83,10 +78,5 @@ pub fn err_string(result: i32) -> String {
 }
 
 pub fn check(result: i32, what: &str) {
-    assert_eq!(
-        result,
-        NCCL_SUCCESS,
-        "{what} failed: {}",
-        err_string(result)
-    );
+    assert_eq!(result, NCCL_SUCCESS, "{what} failed: {}", err_string(result));
 }

crates/xserv-distributed/src/lib.rs

@@ -9,18 +9,15 @@ pub mod ffi;
 use std::ffi::c_void;
 
 use ffi::{NcclComm, NcclUniqueId};
-use xserv_cuda::GpuBuffer;
 use xserv_cuda::device;
+use xserv_cuda::GpuBuffer;
 
 pub use ffi::NcclUniqueId as UniqueId;
 
-/// NCCL is issued on the thread's current launch stream (legacy null stream
-/// by default, the capture stream during CUDA graph capture). The model's
-/// kernels run on the same stream, so AllReduce stays correctly ordered after
-/// the producing matmul and before the consuming kernel — no extra sync.
-fn launch_stream() -> xserv_cuda::ffi::CudaStream {
-    xserv_cuda::stream::current_stream_raw()
-}
+/// The CUDA "null" (default) stream. The model's kernels and cuBLAS calls run
+/// on it, so issuing NCCL on the same stream keeps AllReduce correctly ordered
+/// after the producing matmul and before the consuming kernel — no extra sync.
+const NULL_STREAM: xserv_cuda::ffi::CudaStream = std::ptr::null_mut();
 
 /// Generate a unique id on one rank (typically rank 0) and broadcast the bytes
 /// to all ranks out-of-band (e.g. via a shared variable across threads).
@@ -55,12 +52,7 @@ impl TpContext {
             "ncclCommInitRank",
         );
         ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
-        Self {
-            rank,
-            world,
-            device,
-            comm,
-        }
+        Self { rank, world, device, comm }
     }
 
     /// In-place AllReduce(sum) over `count` BF16 elements in `buf`.
@@ -88,7 +80,7 @@ impl TpContext {
                     ffi::NCCL_BF16,
                     ffi::NCCL_SUM,
                     self.comm,
-                    launch_stream(),
+                    NULL_STREAM,
                 )
             },
             "ncclAllReduce",
@@ -132,12 +124,7 @@ impl PpContext {
             "ncclCommInitRank",
         );
         ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
-        Self {
-            stage,
-            world,
-            device,
-            comm,
-        }
+        Self { stage, world, device, comm }
     }
 
     /// Send `count` BF16 elements at `ptr` to `peer`, on the null stream so it is
@@ -148,16 +135,7 @@ impl PpContext {
     /// `ptr` must point to at least `count` BF16 elements of valid device memory.
     pub fn send_bf16_ptr(&self, ptr: *const c_void, count: usize, peer: usize) {
         ffi::check(
-            unsafe {
-                ffi::ncclSend(
-                    ptr,
-                    count,
-                    ffi::NCCL_BF16,
-                    peer as i32,
-                    self.comm,
-                    launch_stream(),
-                )
-            },
+            unsafe { ffi::ncclSend(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, NULL_STREAM) },
             "ncclSend",
         );
     }
@@ -168,16 +146,7 @@ impl PpContext {
     /// `ptr` must point to at least `count` BF16 elements of valid device memory.
     pub fn recv_bf16_ptr(&self, ptr: *mut c_void, count: usize, peer: usize) {
         ffi::check(
-            unsafe {
-                ffi::ncclRecv(
-                    ptr,
-                    count,
-                    ffi::NCCL_BF16,
-                    peer as i32,
-                    self.comm,
-                    launch_stream(),
-                )
-            },
+            unsafe { ffi::ncclRecv(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, NULL_STREAM) },
             "ncclRecv",
         );
     }

crates/xserv-distributed/tests/allreduce.rs

@@ -2,8 +2,8 @@
 
 use half::bf16;
 use std::thread;
-use xserv_cuda::{GpuBuffer, device};
-use xserv_distributed::{TpContext, get_unique_id};
+use xserv_cuda::{device, GpuBuffer};
+use xserv_distributed::{get_unique_id, TpContext};
 
 #[test]
 fn allreduce_two_gpu_sum() {
@@ -25,7 +25,9 @@ fn allreduce_two_gpu_sum() {
                 // Rank r fills its buffer with (r + 1).
                 let val = bf16::from_f32((rank + 1) as f32);
                 let host = vec![val; n];
-                let src = unsafe { std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2) };
+                let src = unsafe {
+                    std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2)
+                };
                 let mut buf = GpuBuffer::alloc(n * 2).unwrap();
                 buf.copy_from_host(src).unwrap();
 

crates/xserv-distributed/tests/sendrecv.rs

@@ -6,8 +6,8 @@
 use half::bf16;
 use std::ffi::c_void;
 use std::thread;
-use xserv_cuda::{GpuBuffer, device};
-use xserv_distributed::{PpContext, get_unique_id};
+use xserv_cuda::{device, GpuBuffer};
+use xserv_distributed::{get_unique_id, PpContext};
 
 #[test]
 fn pp_send_recv_two_stages() {
@@ -30,8 +30,7 @@ fn pp_send_recv_two_stages() {
                 if stage == 0 {
                     // Fill with a known pattern and send to stage 1.
                     let host: Vec<bf16> = (0..n).map(|i| bf16::from_f32((i % 97) as f32)).collect();
-                    let src =
-                        unsafe { std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2) };
+                    let src = unsafe { std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2) };
                     buf.copy_from_host(src).unwrap();
                     pp.send_bf16_ptr(buf.as_mut_ptr() as *const c_void, n, 1);
                     device::synchronize().unwrap();

crates/xserv-kernels/build.rs

@@ -8,7 +8,6 @@ fn main() {
     println!("cargo:rustc-link-search=native={cuda_path}/lib64");
     println!("cargo:rustc-link-lib=dylib=cudart");
     println!("cargo:rustc-link-lib=dylib=cublas");
-    println!("cargo:rustc-link-lib=dylib=cublasLt");
 
     cc::Build::new()
         .cuda(true)
@@ -22,19 +21,12 @@ fn main() {
         .file("../../csrc/normalization/layernorm.cu")
         .file("../../csrc/activation/activations.cu")
         .file("../../csrc/reduce/softmax.cu")
-        .file("../../csrc/reduce/argmax.cu")
         .file("../../csrc/embedding/embedding.cu")
         .file("../../csrc/embedding/rope.cu")
         .file("../../csrc/attention/causal_mask.cu")
         .file("../../csrc/embedding/transpose.cu")
         .file("../../csrc/attention/flash_attention.cu")
         .file("../../csrc/attention/paged_attention.cu")
-        .file("../../csrc/attention/reshape_and_cache.cu")
-        .file("../../csrc/moe/moe_kernels.cu")
-        .file("../../csrc/moe/moe_sparse.cu")
-        .file("../../csrc/quantization/dequant_fp8.cu")
-        .file("../../csrc/quantization/quantize_fp8.cu")
-        .file("../../csrc/quantization/mxfp4_gemm.cu")
         .compile("xserv_kernels");
 
     println!("cargo:rerun-if-changed=../../csrc/");

crates/xserv-kernels/src/activation.rs

@@ -6,220 +6,74 @@ unsafe extern "C" {
     fn launch_gelu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
     fn launch_silu_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
     fn launch_silu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
-    fn launch_scale_f32(
-        x: *const c_void,
-        out: *mut c_void,
-        scale: f32,
-        n: i32,
-        stream: *mut c_void,
-    );
-    fn launch_scale_bf16(
-        x: *const c_void,
-        out: *mut c_void,
-        scale: f32,
-        n: i32,
-        stream: *mut c_void,
-    );
-    fn launch_add_f32(
-        a: *const c_void,
-        b: *const c_void,
-        out: *mut c_void,
-        n: i32,
-        stream: *mut c_void,
-    );
-    fn launch_add_bf16(
-        a: *const c_void,
-        b: *const c_void,
-        out: *mut c_void,
-        n: i32,
-        stream: *mut c_void,
-    );
-    fn launch_mul_f32(
-        a: *const c_void,
-        b: *const c_void,
-        out: *mut c_void,
-        n: i32,
-        stream: *mut c_void,
-    );
-    fn launch_mul_bf16(
-        a: *const c_void,
-        b: *const c_void,
-        out: *mut c_void,
-        n: i32,
-        stream: *mut c_void,
-    );
-    fn launch_silu_mul_bf16(
-        gate: *const c_void,
-        up: *const c_void,
-        out: *mut c_void,
-        n: i32,
-        stream: *mut c_void,
-    );
-    fn launch_gpt_oss_glu_bf16(
-        gate_up: *const c_void,
-        out: *mut c_void,
-        n_elements: i32,
-        alpha: f32,
-        limit: f32,
-        stream: *mut c_void,
-    );
-    fn launch_bias_add_2d_bf16(
-        x: *const c_void,
-        bias: *const c_void,
-        out: *mut c_void,
-        rows: i32,
-        cols: i32,
-        stream: *mut c_void,
-    );
+    fn launch_scale_f32(x: *const c_void, out: *mut c_void, scale: f32, n: i32, stream: *mut c_void);
+    fn launch_scale_bf16(x: *const c_void, out: *mut c_void, scale: f32, n: i32, stream: *mut c_void);
+    fn launch_add_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_mul_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_mul_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
 }
 
-fn dispatch_unary(
-    x: &Tensor,
-    f32_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
-    bf16_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
-) -> Tensor {
+fn dispatch_unary(x: &Tensor, f32_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
+                  bf16_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void)) -> Tensor {
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let out = Tensor::empty(x.shape(), x.dtype(), x.device());
     let n = x.numel();
-    assert!(
-        n <= i32::MAX as usize,
-        "tensor too large for i32 kernel param ({n} elements)"
-    );
+    assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
     let n = n as i32;
     unsafe {
         match x.dtype() {
-            DType::F32 => f32_fn(
-                x.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                n,
-                xserv_cuda::current_stream_raw(),
-            ),
-            DType::BF16 => bf16_fn(
-                x.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                n,
-                xserv_cuda::current_stream_raw(),
-            ),
+            DType::F32 => f32_fn(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            DType::BF16 => bf16_fn(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
             _ => panic!("unsupported dtype"),
         }
     }
     out
 }
 
-fn dispatch_binary(
-    a: &Tensor,
-    b: &Tensor,
-    f32_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
-    bf16_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
-) -> Tensor {
+fn dispatch_binary(a: &Tensor, b: &Tensor,
+                   f32_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
+                   bf16_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void)) -> Tensor {
     assert_eq!(a.shape(), b.shape());
     assert!(a.is_contiguous() && b.is_contiguous());
     assert!(matches!(a.device(), Device::Cuda(_)));
     assert_eq!(a.dtype(), b.dtype());
     let out = Tensor::empty(a.shape(), a.dtype(), a.device());
     let n = a.numel();
-    assert!(
-        n <= i32::MAX as usize,
-        "tensor too large for i32 kernel param ({n} elements)"
-    );
+    assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
     let n = n as i32;
     unsafe {
         match a.dtype() {
-            DType::F32 => f32_fn(
-                a.data_ptr() as _,
-                b.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                n,
-                xserv_cuda::current_stream_raw(),
-            ),
-            DType::BF16 => bf16_fn(
-                a.data_ptr() as _,
-                b.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                n,
-                xserv_cuda::current_stream_raw(),
-            ),
+            DType::F32 => f32_fn(a.data_ptr() as _, b.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
+            DType::BF16 => bf16_fn(a.data_ptr() as _, b.data_ptr() as _, out.data_ptr() as *mut c_void, n, std::ptr::null_mut()),
             _ => panic!("unsupported dtype"),
         }
     }
     out
 }
 
-pub fn gelu(x: &Tensor) -> Tensor {
-    dispatch_unary(x, launch_gelu_f32, launch_gelu_bf16)
-}
-pub fn silu(x: &Tensor) -> Tensor {
-    dispatch_unary(x, launch_silu_f32, launch_silu_bf16)
-}
+pub fn gelu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_gelu_f32, launch_gelu_bf16) }
+pub fn silu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_silu_f32, launch_silu_bf16) }
 
 pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let out = Tensor::empty(x.shape(), x.dtype(), x.device());
     let n = x.numel();
-    assert!(
-        n <= i32::MAX as usize,
-        "tensor too large for i32 kernel param ({n} elements)"
-    );
+    assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
     let n = n as i32;
     unsafe {
         match x.dtype() {
-            DType::F32 => launch_scale_f32(
-                x.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                scale_val,
-                n,
-                xserv_cuda::current_stream_raw(),
-            ),
-            DType::BF16 => launch_scale_bf16(
-                x.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                scale_val,
-                n,
-                xserv_cuda::current_stream_raw(),
-            ),
+            DType::F32 => launch_scale_f32(x.data_ptr() as _, out.data_ptr() as *mut c_void, scale_val, n, std::ptr::null_mut()),
+            DType::BF16 => launch_scale_bf16(x.data_ptr() as _, out.data_ptr() as *mut c_void, scale_val, n, std::ptr::null_mut()),
             _ => panic!("unsupported dtype for scale"),
         }
     }
     out
 }
 
-pub fn add(a: &Tensor, b: &Tensor) -> Tensor {
-    dispatch_binary(a, b, launch_add_f32, launch_add_bf16)
-}
-pub fn mul(a: &Tensor, b: &Tensor) -> Tensor {
-    dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16)
-}
-
-/// Row-broadcast bias add: out[r, c] = x[r, c] + bias[c] (BF16 only).
-pub fn bias_add_2d(x: &Tensor, bias: &Tensor) -> Tensor {
-    assert_eq!(x.ndim(), 2);
-    assert_eq!(bias.ndim(), 1);
-    assert_eq!(x.dtype(), DType::BF16);
-    assert_eq!(bias.dtype(), DType::BF16);
-    assert!(x.is_contiguous() && bias.is_contiguous());
-    assert!(matches!(x.device(), Device::Cuda(_)));
-    let rows = x.shape()[0];
-    let cols = x.shape()[1];
-    assert_eq!(
-        bias.shape()[0],
-        cols,
-        "bias size {} != cols {cols}",
-        bias.shape()[0]
-    );
-    assert!(rows * cols <= i32::MAX as usize);
-    let out = Tensor::empty(&[rows, cols], DType::BF16, x.device());
-    unsafe {
-        launch_bias_add_2d_bf16(
-            x.data_ptr() as _,
-            bias.data_ptr() as _,
-            out.data_ptr() as *mut c_void,
-            rows as i32,
-            cols as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-    out
-}
+pub fn add(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_add_f32, launch_add_bf16) }
+pub fn mul(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16) }
 
 /// Fused SiLU×Mul: out = silu(gate) * up (BF16 only)
 /// Saves one HBM read + one HBM write compared to separate silu + mul.
@@ -230,10 +84,7 @@ pub fn silu_mul(gate: &Tensor, up: &Tensor) -> Tensor {
     assert_eq!(gate.dtype(), DType::BF16, "silu_mul requires BF16");
     let out = Tensor::empty(gate.shape(), gate.dtype(), gate.device());
     let n = gate.numel();
-    assert!(
-        n <= i32::MAX as usize,
-        "tensor too large for i32 kernel param ({n} elements)"
-    );
+    assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
     let n = n as i32;
     unsafe {
         launch_silu_mul_bf16(
@@ -241,35 +92,7 @@ pub fn silu_mul(gate: &Tensor, up: &Tensor) -> Tensor {
             up.data_ptr() as *const c_void,
             out.data_ptr() as *mut c_void,
             n,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-    out
-}
-
-/// gpt-oss fused GLU activation (BF16 only).
-/// Input: gate_up [rows, 2*D] with interleaved columns (gate=even, up=odd).
-/// Output: [rows, D]
-/// Computes: gate.clamp(max=limit) * sigmoid(gate * alpha) * (up.clamp(-limit,limit) + 1)
-pub fn gpt_oss_glu(gate_up: &Tensor, alpha: f32, limit: f32) -> Tensor {
-    assert!(gate_up.is_contiguous());
-    assert!(matches!(gate_up.device(), Device::Cuda(_)));
-    assert_eq!(gate_up.dtype(), DType::BF16, "gpt_oss_glu requires BF16");
-    assert_eq!(gate_up.ndim(), 2);
-    let rows = gate_up.shape()[0];
-    let cols = gate_up.shape()[1];
-    assert_eq!(cols % 2, 0);
-    let d = cols / 2;
-    let out = Tensor::empty(&[rows, d], gate_up.dtype(), gate_up.device());
-    let n_elements = (rows * d) as i32;
-    unsafe {
-        launch_gpt_oss_glu_bf16(
-            gate_up.data_ptr() as *const c_void,
-            out.data_ptr() as *mut c_void,
-            n_elements,
-            alpha,
-            limit,
-            xserv_cuda::current_stream_raw(),
+            std::ptr::null_mut(),
         );
     }
     out

crates/xserv-kernels/src/argmax.rs

@@ -1,72 +0,0 @@
-use std::ffi::c_void;
-use xserv_tensor::{DType, Device, Tensor};
-
-unsafe extern "C" {
-    fn launch_argmax_bf16(
-        logits: *const c_void,
-        out_idx: *mut c_void,
-        rows: i32,
-        cols: i32,
-        stream: *mut c_void,
-    );
-}
-
-/// GPU argmax over the last dim of a [rows, cols] BF16 tensor.
-///
-/// Returns a host `Vec<u32>` of length `rows`. Internally:
-///   - launches one kernel that writes [rows] i32 indices on device
-///   - D2H copies just `rows * 4` bytes (vs `rows * cols * 2` for the
-///     "copy logits to CPU then argmax" path it replaces)
-///
-/// This is the greedy-decode hot path: avoids touching the full
-/// [B, vocab] logits buffer on the host every step.
-pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
-    assert_eq!(logits.ndim(), 2, "argmax expects a 2D [rows, cols] tensor");
-    assert_eq!(logits.dtype(), DType::BF16, "argmax kernel is BF16-only");
-    assert!(logits.is_contiguous(), "argmax requires contiguous input");
-    assert!(
-        matches!(logits.device(), Device::Cuda(_)),
-        "argmax requires GPU input"
-    );
-
-    let rows = logits.shape()[0];
-    let cols = logits.shape()[1];
-    assert!(rows <= i32::MAX as usize);
-    assert!(cols <= i32::MAX as usize);
-
-    // Output buffer: rows * i32. Pooled allocator so this is essentially free
-    // after the first call.
-    let bytes = rows * std::mem::size_of::<i32>();
-    let mut out = xserv_cuda::allocator::cached_alloc(bytes).expect("argmax out alloc");
-
-    unsafe {
-        launch_argmax_bf16(
-            logits.data_ptr() as *const c_void,
-            out.as_mut_ptr() as *mut c_void,
-            rows as i32,
-            cols as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    let mut host_bytes = vec![0u8; bytes];
-    out.copy_to_host(&mut host_bytes).expect("argmax D2H");
-    drop(out); // returned to pool
-
-    let host_i32: &[i32] =
-        unsafe { std::slice::from_raw_parts(host_bytes.as_ptr() as *const i32, rows) };
-    host_i32.iter().map(|&v| v as u32).collect()
-}
-
-/// Convenience: argmax of a single row [1, cols] (or [cols] reshaped to [1, cols]).
-pub fn argmax_bf16_single(logits: &Tensor) -> u32 {
-    let cols = *logits.shape().last().unwrap();
-    let rows = logits.numel() / cols;
-    assert_eq!(rows, 1, "argmax_bf16_single requires a single row");
-    let view = if logits.ndim() == 2 {
-        logits.clone()
-    } else {
-        logits.reshape(&[1, cols])
-    };
-    argmax_bf16_to_host(&view)[0]
-}

crates/xserv-kernels/src/attention.rs

@@ -6,67 +6,21 @@ use crate::gemm::batched_matmul;
 use crate::softmax::softmax;
 
 unsafe extern "C" {
-    fn launch_causal_mask_f32(
-        scores: *mut c_void,
-        batch: i32,
-        rows: i32,
-        cols: i32,
-        offset: i32,
-        stream: *mut c_void,
-    );
-    fn launch_causal_mask_bf16(
-        scores: *mut c_void,
-        batch: i32,
-        rows: i32,
-        cols: i32,
-        offset: i32,
-        stream: *mut c_void,
-    );
+    fn launch_causal_mask_f32(scores: *mut c_void, batch: i32, rows: i32, cols: i32,
+                              offset: i32, stream: *mut c_void);
+    fn launch_causal_mask_bf16(scores: *mut c_void, batch: i32, rows: i32, cols: i32,
+                               offset: i32, stream: *mut c_void);
     fn launch_flash_attention_bf16(
-        q: *const c_void,
-        k: *const c_void,
-        v: *const c_void,
-        o: *mut c_void,
-        batch: i32,
-        num_q_heads: i32,
-        num_kv_heads: i32,
-        q_len: i32,
-        kv_len: i32,
-        head_dim: i32,
-        scale: f32,
-        causal: i32,
-        stream: *mut c_void,
-    );
-    fn launch_flash_attention_sinks_bf16(
-        q: *const c_void,
-        k: *const c_void,
-        v: *const c_void,
-        o: *mut c_void,
-        sinks: *const c_void,
-        batch: i32,
-        num_q_heads: i32,
-        num_kv_heads: i32,
-        q_len: i32,
-        kv_len: i32,
-        head_dim: i32,
-        scale: f32,
-        causal: i32,
-        window_size: i32,
-        stream: *mut c_void,
+        q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
+        batch: i32, num_q_heads: i32, num_kv_heads: i32,
+        q_len: i32, kv_len: i32, head_dim: i32,
+        scale: f32, causal: i32, stream: *mut c_void,
     );
     fn launch_decode_attention_bf16(
-        q: *const c_void,
-        k: *const c_void,
-        v: *const c_void,
-        o: *mut c_void,
-        batch: i32,
-        num_q_heads: i32,
-        num_kv_heads: i32,
-        kv_len: i32,
-        head_dim: i32,
-        scale: f32,
-        causal: i32,
-        stream: *mut c_void,
+        q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
+        batch: i32, num_q_heads: i32, num_kv_heads: i32,
+        kv_len: i32, head_dim: i32,
+        scale: f32, causal: i32, stream: *mut c_void,
     );
     fn launch_paged_decode_attention_bf16(
         q: *const c_void,
@@ -75,200 +29,9 @@ unsafe extern "C" {
         o: *mut c_void,
         block_tables: *const i32,
         context_lens: *const i32,
-        batch: i32,
-        num_q_heads: i32,
-        num_kv_heads: i32,
-        head_dim: i32,
-        max_blocks_per_seq: i32,
-        scale: f32,
-        stream: *mut c_void,
-    );
-    fn launch_paged_decode_attention_tree_bf16(
-        q: *const c_void,
-        k_cache: *const c_void,
-        v_cache: *const c_void,
-        o: *mut c_void,
-        block_tables: *const i32,
-        context_lens: *const i32,
-        tree_mask: *const i32,
-        batch: i32,
-        num_q_heads: i32,
-        num_kv_heads: i32,
-        head_dim: i32,
-        max_blocks_per_seq: i32,
-        tree_start: i32,
-        tree_len: i32,
-        scale: f32,
-        stream: *mut c_void,
-    );
-    fn launch_paged_decode_attention_sinks_bf16(
-        q: *const c_void,
-        k_cache: *const c_void,
-        v_cache: *const c_void,
-        o: *mut c_void,
-        block_tables: *const i32,
-        context_lens: *const i32,
-        sinks: *const c_void,
-        batch: i32,
-        num_q_heads: i32,
-        num_kv_heads: i32,
-        head_dim: i32,
-        max_blocks_per_seq: i32,
-        scale: f32,
-        window_size: i32,
-        stream: *mut c_void,
-    );
-    fn launch_reshape_and_cache_bf16(
-        k_src: *const c_void,
-        v_src: *const c_void,
-        k_pool: *mut c_void,
-        v_pool: *mut c_void,
-        block_ids: *const c_void,
-        num_tokens: i32,
-        num_heads: i32,
-        head_dim: i32,
-        start_pos: i32,
-        block_size: i32,
-        stream: *mut c_void,
-    );
-    fn launch_reshape_and_cache_batched_bf16(
-        k_src: *const c_void,
-        v_src: *const c_void,
-        k_pool: *mut c_void,
-        v_pool: *mut c_void,
-        block_tables: *const c_void,
-        kv_lens: *const c_void,
-        batch: i32,
-        num_heads: i32,
-        head_dim: i32,
-        block_size: i32,
-        max_blocks_per_seq: i32,
-        stream: *mut c_void,
-    );
-    fn launch_copy_kv_position(
-        k_pool: *mut c_void,
-        v_pool: *mut c_void,
-        block_ids: *const i32,
-        src_pos: i32,
-        dst_pos: i32,
-        num_kv_heads: i32,
-        head_dim: i32,
-        block_size: i32,
-        stream: *mut c_void,
-    );
-}
-
-/// Scatter `[num_kv_heads, num_tokens, head_dim]` BF16 K/V into a paged
-/// pool for a single sequence whose block table lives at `block_ids_gpu`
-/// (int32, on device).
-///
-/// `k_pool_ptr`/`v_pool_ptr` point to one layer's pool, of logical shape
-/// `[num_blocks_total, num_kv_heads, block_size, head_dim]`.
-///
-/// All pointers must be on the same GPU as the launching context.
-///
-/// # Safety
-/// Pointers must be valid GPU pointers with the documented layouts.
-/// `block_ids_gpu` must contain at least `(start_pos + num_tokens + block_size - 1) / block_size`
-/// valid physical block ids.
-pub unsafe fn reshape_and_cache_bf16(
-    k_src: *const c_void,
-    v_src: *const c_void,
-    k_pool_ptr: *mut c_void,
-    v_pool_ptr: *mut c_void,
-    block_ids_gpu: *const i32,
-    num_tokens: usize,
-    num_heads: usize,
-    head_dim: usize,
-    start_pos: usize,
-    block_size: usize,
-    stream: *mut c_void,
-) {
-    unsafe {
-        launch_reshape_and_cache_bf16(
-            k_src,
-            v_src,
-            k_pool_ptr,
-            v_pool_ptr,
-            block_ids_gpu as *const c_void,
-            num_tokens as i32,
-            num_heads as i32,
-            head_dim as i32,
-            start_pos as i32,
-            block_size as i32,
-            stream,
-        );
-    }
-}
-
-/// Batched scatter for the multi-sequence decode step. Reads
-/// `block_tables` (`[batch, max_blocks_per_seq]` int32 — same buffer the
-/// paged-attention kernel reads) and `kv_lens` (`[batch]` int32, current
-/// seq_len + 1 — i.e., the index of the just-written token + 1) so the
-/// caller doesn't need a separate per-step upload of block ids.
-///
-/// # Safety
-/// All pointers must be on the same GPU. `block_tables` and `kv_lens` must
-/// already be synced to the device for the active batch.
-pub unsafe fn reshape_and_cache_batched_bf16(
-    k_src: *const c_void,
-    v_src: *const c_void,
-    k_pool_ptr: *mut c_void,
-    v_pool_ptr: *mut c_void,
-    block_tables_gpu: *const i32,
-    kv_lens_gpu: *const i32,
-    batch: usize,
-    num_heads: usize,
-    head_dim: usize,
-    block_size: usize,
-    max_blocks_per_seq: usize,
-    stream: *mut c_void,
-) {
-    unsafe {
-        launch_reshape_and_cache_batched_bf16(
-            k_src,
-            v_src,
-            k_pool_ptr,
-            v_pool_ptr,
-            block_tables_gpu as *const c_void,
-            kv_lens_gpu as *const c_void,
-            batch as i32,
-            num_heads as i32,
-            head_dim as i32,
-            block_size as i32,
-            max_blocks_per_seq as i32,
-            stream,
-        );
-    }
-}
-
-/// Copy one token's K/V from `src_pos` to `dst_pos` within the same sequence's
-/// paged cache (one layer). Used by tree speculative decoding to remap
-/// accepted sibling K/V to canonical sequential positions after acceptance.
-///
-/// # Safety
-/// Pool and block_ids pointers must be valid GPU pointers for the given layer.
-pub unsafe fn copy_kv_position(
-    k_pool_ptr: *mut c_void,
-    v_pool_ptr: *mut c_void,
-    block_ids_gpu: *const i32,
-    src_pos: usize,
-    dst_pos: usize,
-    num_kv_heads: usize,
-    head_dim: usize,
-    block_size: usize,
-    stream: *mut c_void,
-) {
-    launch_copy_kv_position(
-        k_pool_ptr,
-        v_pool_ptr,
-        block_ids_gpu,
-        src_pos as i32,
-        dst_pos as i32,
-        num_kv_heads as i32,
-        head_dim as i32,
-        block_size as i32,
-        stream,
+        batch: i32, num_q_heads: i32, num_kv_heads: i32,
+        head_dim: i32, max_blocks_per_seq: i32,
+        scale: f32, stream: *mut c_void,
     );
 }
 
@@ -282,19 +45,13 @@ fn apply_causal_mask(scores: &Tensor, offset: usize) {
         match scores.dtype() {
             DType::F32 => launch_causal_mask_f32(
                 scores.data_ptr() as *mut c_void,
-                batch as i32,
-                rows as i32,
-                cols as i32,
-                offset as i32,
-                xserv_cuda::current_stream_raw(),
+                batch as i32, rows as i32, cols as i32, offset as i32,
+                std::ptr::null_mut(),
             ),
             DType::BF16 => launch_causal_mask_bf16(
                 scores.data_ptr() as *mut c_void,
-                batch as i32,
-                rows as i32,
-                cols as i32,
-                offset as i32,
-                xserv_cuda::current_stream_raw(),
+                batch as i32, rows as i32, cols as i32, offset as i32,
+                std::ptr::null_mut(),
             ),
             _ => panic!("unsupported dtype for causal mask"),
         }
@@ -359,7 +116,11 @@ pub fn decode_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Tensor {
     let kv_len = k.shape()[2];
 
     let scale = 1.0 / (head_dim as f32).sqrt();
-    let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
+    let output = Tensor::empty(
+        &[batch, num_q_heads, 1, head_dim],
+        DType::BF16,
+        q.device(),
+    );
 
     unsafe {
         launch_decode_attention_bf16(
@@ -374,7 +135,7 @@ pub fn decode_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Tensor {
             head_dim as i32,
             scale,
             1, // causal (always 1 for decode)
-            xserv_cuda::current_stream_raw(),
+            std::ptr::null_mut(),
         );
     }
 
@@ -407,14 +168,8 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens
 
     assert_eq!(k.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
     assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
-    assert!(
-        num_q_heads % num_kv_heads == 0,
-        "num_q_heads must be divisible by num_kv_heads"
-    );
-    assert!(
-        head_dim <= 128,
-        "flash_attention supports head_dim up to 128"
-    );
+    assert!(num_q_heads % num_kv_heads == 0, "num_q_heads must be divisible by num_kv_heads");
+    assert!(head_dim <= 128, "flash_attention supports head_dim up to 128");
 
     // Dispatch to specialized decode kernel for single-token generation
     if q_len == 1 {
@@ -442,74 +197,7 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens
             head_dim as i32,
             scale,
             if causal { 1 } else { 0 },
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    output
-}
-
-/// Flash attention for prefill with gpt-oss attention sinks + optional sliding window.
-///
-/// Same layout/contract as `flash_attention`, plus a per-head `sinks` tensor
-/// ([num_q_heads] BF16, GPU) folded into the softmax denominator, and a
-/// `window_size` (0 = full causal, >0 = sliding window). Always causal.
-pub fn flash_attention_sinks(
-    q: &Tensor,
-    k: &Tensor,
-    v: &Tensor,
-    sinks: &Tensor,
-    window_size: usize,
-) -> Tensor {
-    assert_eq!(q.ndim(), 4);
-    assert_eq!(k.ndim(), 4);
-    assert_eq!(v.ndim(), 4);
-    assert!(q.is_contiguous() && k.is_contiguous() && v.is_contiguous());
-    assert_eq!(q.dtype(), DType::BF16);
-    assert_eq!(k.dtype(), DType::BF16);
-    assert_eq!(v.dtype(), DType::BF16);
-
-    let batch = q.shape()[0];
-    let num_q_heads = q.shape()[1];
-    let q_len = q.shape()[2];
-    let head_dim = q.shape()[3];
-    let num_kv_heads = k.shape()[1];
-    let kv_len = k.shape()[2];
-
-    assert_eq!(k.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
-    assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
-    assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
-    assert_eq!(
-        sinks.shape()[0],
-        num_q_heads,
-        "sinks must have num_q_heads entries"
-    );
-
-    let scale = 1.0 / (head_dim as f32).sqrt();
-    let output = Tensor::empty(
-        &[batch, num_q_heads, q_len, head_dim],
-        DType::BF16,
-        q.device(),
-    );
-
-    unsafe {
-        launch_flash_attention_sinks_bf16(
-            q.data_ptr() as *const c_void,
-            k.data_ptr() as *const c_void,
-            v.data_ptr() as *const c_void,
-            output.data_ptr() as *mut c_void,
-            sinks.data_ptr() as *const c_void,
-            batch as i32,
-            num_q_heads as i32,
-            num_kv_heads as i32,
-            q_len as i32,
-            kv_len as i32,
-            head_dim as i32,
-            scale,
-            1, // always causal
-            window_size as i32,
-            xserv_cuda::current_stream_raw(),
+            std::ptr::null_mut(),
         );
     }
 
@@ -538,20 +226,17 @@ pub fn paged_decode_attention(
     max_blocks_per_seq: usize,
 ) -> Tensor {
     assert_eq!(q.ndim(), 4);
-    assert_eq!(
-        q.shape()[2],
-        1,
-        "paged_decode_attention requires q_len == 1"
-    );
+    assert_eq!(q.shape()[2], 1, "paged_decode_attention requires q_len == 1");
     assert_eq!(q.dtype(), DType::BF16);
-    assert!(
-        num_q_heads % num_kv_heads == 0,
-        "GQA: num_q_heads must be divisible by num_kv_heads"
-    );
+    assert!(num_q_heads % num_kv_heads == 0, "GQA: num_q_heads must be divisible by num_kv_heads");
     assert!(head_dim <= 128);
 
     let scale = 1.0 / (head_dim as f32).sqrt();
-    let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
+    let output = Tensor::empty(
+        &[batch, num_q_heads, 1, head_dim],
+        DType::BF16,
+        q.device(),
+    );
 
     unsafe {
         launch_paged_decode_attention_bf16(
@@ -567,114 +252,7 @@ pub fn paged_decode_attention(
             head_dim as i32,
             max_blocks_per_seq as i32,
             scale,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    output
-}
-
-/// Tree-aware paged decode attention. Adds a per-query attention mask over
-/// the newly-written K/V region `[tree_start, tree_start+tree_len)`. Query i
-/// attends to position tree_start+j iff tree_mask[i, j] != 0. Positions <
-/// tree_start are always attended.
-///
-/// Used by speculative decoding with tree drafting to let sibling candidates
-/// share position slots without seeing each other's K/V.
-#[allow(clippy::too_many_arguments)]
-pub fn paged_decode_attention_tree(
-    q: &Tensor,
-    k_cache_ptr: *const c_void,
-    v_cache_ptr: *const c_void,
-    block_tables_ptr: *const i32,
-    context_lens_ptr: *const i32,
-    tree_mask_ptr: *const i32,
-    batch: usize,
-    num_q_heads: usize,
-    num_kv_heads: usize,
-    head_dim: usize,
-    max_blocks_per_seq: usize,
-    tree_start: usize,
-    tree_len: usize,
-) -> Tensor {
-    assert_eq!(q.ndim(), 4);
-    assert_eq!(q.shape()[2], 1);
-    assert_eq!(q.dtype(), DType::BF16);
-    assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
-
-    let scale = 1.0 / (head_dim as f32).sqrt();
-    let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
-
-    unsafe {
-        launch_paged_decode_attention_tree_bf16(
-            q.data_ptr() as *const c_void,
-            k_cache_ptr,
-            v_cache_ptr,
-            output.data_ptr() as *mut c_void,
-            block_tables_ptr,
-            context_lens_ptr,
-            tree_mask_ptr,
-            batch as i32,
-            num_q_heads as i32,
-            num_kv_heads as i32,
-            head_dim as i32,
-            max_blocks_per_seq as i32,
-            tree_start as i32,
-            tree_len as i32,
-            scale,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    output
-}
-
-/// Paged decode attention with attention sinks and optional sliding window.
-///
-/// sinks_ptr: pointer to [num_q_heads] BF16 on GPU (or null for no sinks)
-/// window_size: 0 = full attention, >0 = sliding window
-#[allow(clippy::too_many_arguments)]
-pub fn paged_decode_attention_sinks(
-    q: &Tensor,
-    k_cache_ptr: *const c_void,
-    v_cache_ptr: *const c_void,
-    block_tables_ptr: *const i32,
-    context_lens_ptr: *const i32,
-    sinks_ptr: *const c_void,
-    batch: usize,
-    num_q_heads: usize,
-    num_kv_heads: usize,
-    head_dim: usize,
-    max_blocks_per_seq: usize,
-    window_size: usize,
-) -> Tensor {
-    assert_eq!(q.ndim(), 4);
-    assert_eq!(q.shape()[2], 1);
-    assert_eq!(q.dtype(), DType::BF16);
-    assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
-
-    let scale = 1.0 / (head_dim as f32).sqrt();
-    let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
-
-    unsafe {
-        launch_paged_decode_attention_sinks_bf16(
-            q.data_ptr() as *const c_void,
-            k_cache_ptr,
-            v_cache_ptr,
-            output.data_ptr() as *mut c_void,
-            block_tables_ptr,
-            context_lens_ptr,
-            sinks_ptr,
-            batch as i32,
-            num_q_heads as i32,
-            num_kv_heads as i32,
-            head_dim as i32,
-            max_blocks_per_seq as i32,
-            scale,
-            window_size as i32,
-            xserv_cuda::current_stream_raw(),
+            std::ptr::null_mut(),
         );
     }
 

crates/xserv-kernels/src/dispatch.rs

@@ -5,302 +5,104 @@ use std::ffi::c_void;
 
 // Re-declare the extern functions we need (same as in the individual modules)
 unsafe extern "C" {
-    fn launch_rmsnorm_bf16(
-        x: *const c_void,
-        gamma: *const c_void,
-        out: *mut c_void,
-        rows: i32,
-        hidden_size: i32,
-        eps: f32,
-        stream: *mut c_void,
-    );
-    fn launch_add_rmsnorm_bf16(
-        x: *const c_void,
-        residual: *const c_void,
-        gamma: *const c_void,
-        normed_out: *mut c_void,
-        sum_out: *mut c_void,
-        rows: i32,
-        hidden_size: i32,
-        eps: f32,
-        stream: *mut c_void,
-    );
-    fn launch_silu_mul_bf16(
-        gate: *const c_void,
-        up: *const c_void,
-        out: *mut c_void,
-        n: i32,
-        stream: *mut c_void,
-    );
-    fn launch_add_bf16(
-        a: *const c_void,
-        b: *const c_void,
-        out: *mut c_void,
-        n: i32,
-        stream: *mut c_void,
-    );
-    fn launch_embedding_bf16(
-        table: *const c_void,
-        token_ids: *const c_void,
-        out: *mut c_void,
-        num_tokens: i32,
-        hidden_size: i32,
-        vocab_size: i32,
-        stream: *mut c_void,
-    );
-    fn launch_reshape_heads_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        seq_len: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_merge_heads_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        seq_len: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_transpose_hsd_to_shd_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        seq_len: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_transpose_shd_to_hsd_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        seq_len: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_rope_bf16(
-        x: *mut c_void,
-        cos_cache: *const c_void,
-        sin_cache: *const c_void,
-        positions: *const c_void,
-        num_tokens: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_gemv_bf16(
-        x: *const c_void,
-        w: *const c_void,
-        y_bf16: *mut c_void,
-        y_fp32_buf: *mut c_void,
-        k: i32,
-        n: i32,
-        stream: *mut c_void,
-    );
+    fn launch_rmsnorm_bf16(x: *const c_void, gamma: *const c_void, out: *mut c_void,
+                           rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+    fn launch_add_rmsnorm_bf16(x: *const c_void, residual: *const c_void, gamma: *const c_void,
+                               normed_out: *mut c_void, sum_out: *mut c_void,
+                               rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+    fn launch_silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
+    fn launch_embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                             num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
+    fn launch_reshape_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_merge_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_rope_bf16(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
+                        positions: *const c_void, num_tokens: i32, num_heads: i32,
+                        head_dim: i32, stream: *mut c_void);
+    fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void,
+                        k: i32, n: i32, stream: *mut c_void);
     fn launch_decode_attention_bf16(
-        q: *const c_void,
-        k: *const c_void,
-        v: *const c_void,
-        o: *mut c_void,
-        batch: i32,
-        num_q_heads: i32,
-        num_kv_heads: i32,
-        kv_len: i32,
-        head_dim: i32,
-        scale: f32,
-        causal: i32,
-        stream: *mut c_void,
+        q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
+        batch: i32, num_q_heads: i32, num_kv_heads: i32,
+        kv_len: i32, head_dim: i32,
+        scale: f32, causal: i32, stream: *mut c_void,
     );
 }
 
 /// Raw rmsnorm dispatch: writes to pre-allocated `out`.
-pub unsafe fn rmsnorm_bf16(
-    x: *const c_void,
-    gamma: *const c_void,
-    out: *mut c_void,
-    rows: i32,
-    hidden_size: i32,
-    eps: f32,
-    stream: *mut c_void,
-) {
+pub unsafe fn rmsnorm_bf16(x: *const c_void, gamma: *const c_void, out: *mut c_void,
+                           rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void) {
     launch_rmsnorm_bf16(x, gamma, out, rows, hidden_size, eps, stream);
 }
 
 /// Raw add_rmsnorm dispatch.
-pub unsafe fn add_rmsnorm_bf16(
-    x: *const c_void,
-    residual: *const c_void,
-    gamma: *const c_void,
-    normed_out: *mut c_void,
-    sum_out: *mut c_void,
-    rows: i32,
-    hidden_size: i32,
-    eps: f32,
-    stream: *mut c_void,
-) {
-    launch_add_rmsnorm_bf16(
-        x,
-        residual,
-        gamma,
-        normed_out,
-        sum_out,
-        rows,
-        hidden_size,
-        eps,
-        stream,
-    );
+pub unsafe fn add_rmsnorm_bf16(x: *const c_void, residual: *const c_void, gamma: *const c_void,
+                               normed_out: *mut c_void, sum_out: *mut c_void,
+                               rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void) {
+    launch_add_rmsnorm_bf16(x, residual, gamma, normed_out, sum_out, rows, hidden_size, eps, stream);
 }
 
 /// Raw silu_mul dispatch.
-pub unsafe fn silu_mul_bf16(
-    gate: *const c_void,
-    up: *const c_void,
-    out: *mut c_void,
-    n: i32,
-    stream: *mut c_void,
-) {
+pub unsafe fn silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void) {
     launch_silu_mul_bf16(gate, up, out, n, stream);
 }
 
 /// Raw add dispatch.
-pub unsafe fn add_bf16(
-    a: *const c_void,
-    b: *const c_void,
-    out: *mut c_void,
-    n: i32,
-    stream: *mut c_void,
-) {
+pub unsafe fn add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void) {
     launch_add_bf16(a, b, out, n, stream);
 }
 
 /// Raw embedding dispatch.
-pub unsafe fn embedding_bf16(
-    table: *const c_void,
-    token_ids: *const c_void,
-    out: *mut c_void,
-    num_tokens: i32,
-    hidden_size: i32,
-    vocab_size: i32,
-    stream: *mut c_void,
-) {
-    launch_embedding_bf16(
-        table,
-        token_ids,
-        out,
-        num_tokens,
-        hidden_size,
-        vocab_size,
-        stream,
-    );
+pub unsafe fn embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                             num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void) {
+    launch_embedding_bf16(table, token_ids, out, num_tokens, hidden_size, vocab_size, stream);
 }
 
 /// Raw reshape_heads dispatch.
-pub unsafe fn reshape_heads_bf16(
-    inp: *const c_void,
-    out: *mut c_void,
-    seq_len: i32,
-    num_heads: i32,
-    head_dim: i32,
-    stream: *mut c_void,
-) {
+pub unsafe fn reshape_heads_bf16(inp: *const c_void, out: *mut c_void,
+                                 seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
     launch_reshape_heads_bf16(inp, out, seq_len, num_heads, head_dim, stream);
 }
 
 /// Raw merge_heads dispatch.
-pub unsafe fn merge_heads_bf16(
-    inp: *const c_void,
-    out: *mut c_void,
-    seq_len: i32,
-    num_heads: i32,
-    head_dim: i32,
-    stream: *mut c_void,
-) {
+pub unsafe fn merge_heads_bf16(inp: *const c_void, out: *mut c_void,
+                               seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
     launch_merge_heads_bf16(inp, out, seq_len, num_heads, head_dim, stream);
 }
 
 /// Raw transpose HSD->SHD dispatch.
-pub unsafe fn transpose_hsd_to_shd_bf16(
-    inp: *const c_void,
-    out: *mut c_void,
-    seq_len: i32,
-    num_heads: i32,
-    head_dim: i32,
-    stream: *mut c_void,
-) {
+pub unsafe fn transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void,
+                                         seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
     launch_transpose_hsd_to_shd_bf16(inp, out, seq_len, num_heads, head_dim, stream);
 }
 
 /// Raw transpose SHD->HSD dispatch.
-pub unsafe fn transpose_shd_to_hsd_bf16(
-    inp: *const c_void,
-    out: *mut c_void,
-    seq_len: i32,
-    num_heads: i32,
-    head_dim: i32,
-    stream: *mut c_void,
-) {
+pub unsafe fn transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void,
+                                         seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
     launch_transpose_shd_to_hsd_bf16(inp, out, seq_len, num_heads, head_dim, stream);
 }
 
 /// Raw RoPE dispatch (in-place).
-pub unsafe fn rope_bf16(
-    x: *mut c_void,
-    cos_cache: *const c_void,
-    sin_cache: *const c_void,
-    positions: *const c_void,
-    num_tokens: i32,
-    num_heads: i32,
-    head_dim: i32,
-    stream: *mut c_void,
-) {
-    launch_rope_bf16(
-        x, cos_cache, sin_cache, positions, num_tokens, num_heads, head_dim, stream,
-    );
+pub unsafe fn rope_bf16(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
+                        positions: *const c_void, num_tokens: i32, num_heads: i32,
+                        head_dim: i32, stream: *mut c_void) {
+    launch_rope_bf16(x, cos_cache, sin_cache, positions, num_tokens, num_heads, head_dim, stream);
 }
 
 /// Raw GEMV dispatch (BF16, M=1). Caller must provide fp32 accumulator buffer.
-pub unsafe fn gemv_bf16(
-    x: *const c_void,
-    w: *const c_void,
-    y_bf16: *mut c_void,
-    y_fp32_buf: *mut c_void,
-    k: i32,
-    n: i32,
-    stream: *mut c_void,
-) {
+pub unsafe fn gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void,
+                        y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void) {
     launch_gemv_bf16(x, w, y_bf16, y_fp32_buf, k, n, stream);
 }
 
 /// Raw decode attention dispatch.
-pub unsafe fn decode_attention_bf16(
-    q: *const c_void,
-    k: *const c_void,
-    v: *const c_void,
-    o: *mut c_void,
-    batch: i32,
-    num_q_heads: i32,
-    num_kv_heads: i32,
-    kv_len: i32,
-    head_dim: i32,
-    scale: f32,
-    stream: *mut c_void,
-) {
-    launch_decode_attention_bf16(
-        q,
-        k,
-        v,
-        o,
-        batch,
-        num_q_heads,
-        num_kv_heads,
-        kv_len,
-        head_dim,
-        scale,
-        1,
-        stream,
-    );
+pub unsafe fn decode_attention_bf16(q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
+                                     batch: i32, num_q_heads: i32, num_kv_heads: i32,
+                                     kv_len: i32, head_dim: i32,
+                                     scale: f32, stream: *mut c_void) {
+    launch_decode_attention_bf16(q, k, v, o, batch, num_q_heads, num_kv_heads, kv_len, head_dim, scale, 1, stream);
 }
 
 // cuBLAS FFI

crates/xserv-kernels/src/embedding.rs

@@ -1,25 +1,12 @@
 use std::ffi::c_void;
+use xserv_cuda::GpuBuffer;
 use xserv_tensor::{DType, Device, Tensor};
 
 unsafe extern "C" {
-    fn launch_embedding_f32(
-        table: *const c_void,
-        token_ids: *const c_void,
-        out: *mut c_void,
-        num_tokens: i32,
-        hidden_size: i32,
-        vocab_size: i32,
-        stream: *mut c_void,
-    );
-    fn launch_embedding_bf16(
-        table: *const c_void,
-        token_ids: *const c_void,
-        out: *mut c_void,
-        num_tokens: i32,
-        hidden_size: i32,
-        vocab_size: i32,
-        stream: *mut c_void,
-    );
+    fn launch_embedding_f32(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                            num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
+    fn launch_embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
+                             num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
 }
 
 /// Embedding lookup: table[token_ids[i]] for each i.
@@ -32,14 +19,8 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
     let hidden_size = table.shape()[1];
     let num_tokens = token_ids.len();
     let vocab_size = table.shape()[0];
-    assert!(
-        num_tokens <= i32::MAX as usize,
-        "too many tokens for i32 kernel param"
-    );
-    assert!(
-        hidden_size <= i32::MAX as usize,
-        "hidden_size too large for i32 kernel param"
-    );
+    assert!(num_tokens <= i32::MAX as usize, "too many tokens for i32 kernel param");
+    assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
 
     // Upload token_ids to GPU
     let ids_bytes = unsafe {
@@ -48,51 +29,26 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
             num_tokens * std::mem::size_of::<u32>(),
         )
     };
-    let mut ids_gpu =
-        xserv_cuda::allocator::cached_alloc(ids_bytes.len()).expect("alloc token_ids");
+    let mut ids_gpu = xserv_cuda::allocator::cached_alloc(ids_bytes.len()).expect("alloc token_ids");
     ids_gpu.copy_from_host(ids_bytes).unwrap();
 
     for &tid in token_ids {
-        assert!(
-            (tid as usize) < vocab_size,
-            "token_id {tid} out of bounds (vocab_size={vocab_size})"
-        );
+        assert!((tid as usize) < vocab_size, "token_id {tid} out of bounds (vocab_size={vocab_size})");
     }
 
-    embedding_device_ids(table, ids_gpu.as_ptr() as *const c_void, num_tokens)
-}
-
-/// Embedding lookup with token ids already on the GPU (u32, [num_tokens]).
-/// Used by the CUDA-graph decode path, where ids live in a persistent device
-/// buffer updated outside the captured region (no bounds check possible here).
-pub fn embedding_device_ids(table: &Tensor, ids_gpu: *const c_void, num_tokens: usize) -> Tensor {
-    assert_eq!(table.ndim(), 2);
-    assert!(table.is_contiguous());
-    assert!(matches!(table.device(), Device::Cuda(_)));
-    let hidden_size = table.shape()[1];
-    let vocab_size = table.shape()[0];
-
     let out = Tensor::empty(&[num_tokens, hidden_size], table.dtype(), table.device());
 
     unsafe {
         match table.dtype() {
             DType::F32 => launch_embedding_f32(
-                table.data_ptr() as _,
-                ids_gpu,
+                table.data_ptr() as _, ids_gpu.as_ptr() as _,
                 out.data_ptr() as *mut c_void,
-                num_tokens as i32,
-                hidden_size as i32,
-                vocab_size as i32,
-                xserv_cuda::current_stream_raw(),
+                num_tokens as i32, hidden_size as i32, vocab_size as i32, std::ptr::null_mut(),
             ),
             DType::BF16 => launch_embedding_bf16(
-                table.data_ptr() as _,
-                ids_gpu,
+                table.data_ptr() as _, ids_gpu.as_ptr() as _,
                 out.data_ptr() as *mut c_void,
-                num_tokens as i32,
-                hidden_size as i32,
-                vocab_size as i32,
-                xserv_cuda::current_stream_raw(),
+                num_tokens as i32, hidden_size as i32, vocab_size as i32, std::ptr::null_mut(),
             ),
             _ => panic!("unsupported dtype for embedding"),
         }

crates/xserv-kernels/src/gemm.rs

@@ -1,36 +1,8 @@
 use std::cell::RefCell;
 use std::ffi::c_void;
-use xserv_cuda::GpuBuffer;
 use xserv_cuda::error::{self, Result};
 use xserv_tensor::{DType, Device, Tensor};
 
-const CUBLAS_WORKSPACE_BYTES: usize = 32 * 1024 * 1024;
-const GEMV_TILE_K: usize = 256;
-
-// GEMV: single-kernel, no FP32 temp buffer needed
-unsafe extern "C" {
-    fn launch_gemv_bf16(
-        x: *const c_void,
-        w: *const c_void,
-        y_bf16: *mut c_void,
-        y_fp32_buf: *mut c_void,
-        k: i32,
-        n: i32,
-        stream: *mut c_void,
-    );
-
-    fn launch_gemv_bf16_batched(
-        x: *const c_void,
-        w: *const c_void,
-        y_bf16: *mut c_void,
-        y_fp32_buf: *mut c_void,
-        m: i32,
-        k: i32,
-        n: i32,
-        stream: *mut c_void,
-    );
-}
-
 #[derive(Debug, Clone, Copy)]
 pub enum GemmBackend {
     Naive,
@@ -38,97 +10,13 @@ pub enum GemmBackend {
     CuBlas,
 }
 
-pub fn gemv_scratch_elems(k: usize, n: usize) -> usize {
-    n * k.div_ceil(GEMV_TILE_K)
-}
-
-/// Batched GEMV: [M, K] × [K, N] → [M, N], all BF16.
-/// Bit-exact with calling matmul on each row individually (same K-block partial
-/// + fixed-order reduction path), but in a single kernel launch per phase.
-pub fn matmul_batched_gemv(a: &Tensor, b: &Tensor) -> Tensor {
-    assert_eq!(a.ndim(), 2);
-    assert_eq!(b.ndim(), 2);
-    assert!(a.is_contiguous());
-    assert!(b.is_contiguous());
-    assert_eq!(a.dtype(), DType::BF16);
-    assert_eq!(b.dtype(), DType::BF16);
-    let m = a.shape()[0];
-    let k = a.shape()[1];
-    let n = b.shape()[1];
-    assert_eq!(b.shape()[0], k);
-
-    let out = Tensor::empty(&[m, n], DType::BF16, a.device());
-    let scratch_elems = m * gemv_scratch_elems(k, n);
-    let mut fp32_buf = xserv_cuda::allocator::cached_alloc(scratch_elems * 4).unwrap();
-
-    let null_stream = xserv_cuda::current_stream_raw();
-    if m == 1 {
-        unsafe {
-            launch_gemv_bf16(
-                a.data_ptr() as *const c_void,
-                b.data_ptr() as *const c_void,
-                out.data_ptr() as *mut c_void,
-                fp32_buf.as_mut_ptr() as *mut c_void,
-                k as i32,
-                n as i32,
-                null_stream,
-            );
-        }
-    } else {
-        unsafe {
-            launch_gemv_bf16_batched(
-                a.data_ptr() as *const c_void,
-                b.data_ptr() as *const c_void,
-                out.data_ptr() as *mut c_void,
-                fp32_buf.as_mut_ptr() as *mut c_void,
-                m as i32,
-                k as i32,
-                n as i32,
-                null_stream,
-            );
-        }
-    }
-    out
-}
-
 // --- FFI: custom CUDA kernels ---
 unsafe extern "C" {
-    fn launch_gemm_naive_f32(
-        a: *const c_void,
-        b: *const c_void,
-        c: *mut c_void,
-        m: i32,
-        n: i32,
-        k: i32,
-        stream: *mut c_void,
-    );
-    fn launch_gemm_naive_bf16(
-        a: *const c_void,
-        b: *const c_void,
-        c: *mut c_void,
-        m: i32,
-        n: i32,
-        k: i32,
-        stream: *mut c_void,
-    );
-    fn launch_gemm_tiled_f32(
-        a: *const c_void,
-        b: *const c_void,
-        c: *mut c_void,
-        m: i32,
-        n: i32,
-        k: i32,
-        stream: *mut c_void,
-    );
-    fn launch_gemm_tiled_bf16(
-        a: *const c_void,
-        b: *const c_void,
-        c: *mut c_void,
-        m: i32,
-        n: i32,
-        k: i32,
-        stream: *mut c_void,
-    );
+    fn launch_gemm_naive_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_naive_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_tiled_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemm_tiled_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
+    fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void);
 }
 
 // --- FFI: cuBLAS ---
@@ -148,49 +36,27 @@ unsafe extern "C" {
     fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
     fn cublasDestroy_v2(handle: CublasHandle) -> i32;
     fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
-    fn cublasSetWorkspace_v2(handle: CublasHandle, workspace: *mut c_void, size: usize) -> i32;
     fn cublasGemmEx(
         handle: CublasHandle,
-        transa: i32,
-        transb: i32,
-        m: i32,
-        n: i32,
-        k: i32,
+        transa: i32, transb: i32,
+        m: i32, n: i32, k: i32,
         alpha: *const c_void,
-        a: *const c_void,
-        a_type: i32,
-        lda: i32,
-        b: *const c_void,
-        b_type: i32,
-        ldb: i32,
+        a: *const c_void, a_type: i32, lda: i32,
+        b: *const c_void, b_type: i32, ldb: i32,
         beta: *const c_void,
-        c: *mut c_void,
-        c_type: i32,
-        ldc: i32,
+        c: *mut c_void, c_type: i32, ldc: i32,
         compute_type: i32,
         algo: i32,
     ) -> i32;
     fn cublasGemmStridedBatchedEx(
         handle: CublasHandle,
-        transa: i32,
-        transb: i32,
-        m: i32,
-        n: i32,
-        k: i32,
+        transa: i32, transb: i32,
+        m: i32, n: i32, k: i32,
         alpha: *const c_void,
-        a: *const c_void,
-        a_type: i32,
-        lda: i32,
-        stride_a: i64,
-        b: *const c_void,
-        b_type: i32,
-        ldb: i32,
-        stride_b: i64,
+        a: *const c_void, a_type: i32, lda: i32, stride_a: i64,
+        b: *const c_void, b_type: i32, ldb: i32, stride_b: i64,
         beta: *const c_void,
-        c: *mut c_void,
-        c_type: i32,
-        ldc: i32,
-        stride_c: i64,
+        c: *mut c_void, c_type: i32, ldc: i32, stride_c: i64,
         batch_count: i32,
         compute_type: i32,
         algo: i32,
@@ -199,32 +65,13 @@ unsafe extern "C" {
 
 pub struct CublasContext {
     handle: CublasHandle,
-    /// Dedicated 32 MiB workspace owned by this handle. Held to keep the GPU
-    /// buffer alive for the lifetime of the handle; cuBLAS reads/writes into
-    /// it during GEMM. Dropped after `cublasDestroy_v2` so cuBLAS can't touch
-    /// freed memory.
-    _workspace: Option<GpuBuffer>,
 }
 
 impl CublasContext {
     pub fn new() -> Result<Self> {
         let mut handle = std::ptr::null_mut();
         error::check(unsafe { cublasCreate_v2(&mut handle) })?;
-        // Attach a per-handle workspace. cublasSetWorkspace requires the
-        // pointer to remain valid until destroy or until a new workspace is
-        // set, so we keep the GpuBuffer in this struct.
-        let mut workspace = GpuBuffer::alloc(CUBLAS_WORKSPACE_BYTES)?;
-        error::check(unsafe {
-            cublasSetWorkspace_v2(
-                handle,
-                workspace.as_mut_ptr() as *mut c_void,
-                CUBLAS_WORKSPACE_BYTES,
-            )
-        })?;
-        Ok(Self {
-            handle,
-            _workspace: Some(workspace),
-        })
+        Ok(Self { handle })
     }
 }
 
@@ -233,7 +80,6 @@ impl Drop for CublasContext {
         if !self.handle.is_null() {
             unsafe { cublasDestroy_v2(self.handle) };
         }
-        // _workspace drops here, after cublasDestroy_v2 has released the handle.
     }
 }
 
@@ -256,7 +102,9 @@ where
 
 /// Get the thread-local cuBLAS handle for use with dispatch module.
 pub fn cublas_handle() -> CublasHandle {
-    CUBLAS_CTX.with(|cell| cell.borrow().handle)
+    CUBLAS_CTX.with(|cell| {
+        cell.borrow().handle
+    })
 }
 
 /// Matrix multiplication: C = A @ B
@@ -267,14 +115,8 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
     assert_eq!(b.ndim(), 2);
     assert_eq!(a.shape()[1], b.shape()[0], "inner dimension mismatch");
     assert_eq!(a.dtype(), b.dtype(), "dtype mismatch");
-    assert!(
-        a.is_contiguous() && b.is_contiguous(),
-        "matmul requires contiguous tensors"
-    );
-    assert!(
-        matches!(a.device(), Device::Cuda(_)),
-        "matmul requires GPU tensors"
-    );
+    assert!(a.is_contiguous() && b.is_contiguous(), "matmul requires contiguous tensors");
+    assert!(matches!(a.device(), Device::Cuda(_)), "matmul requires GPU tensors");
 
     let m = a.shape()[0];
     let k = a.shape()[1];
@@ -288,71 +130,44 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
     let a_ptr = a.data_ptr() as *const c_void;
     let b_ptr = b.data_ptr() as *const c_void;
     let c_ptr = c.data_ptr() as *mut c_void;
-    let null_stream = xserv_cuda::current_stream_raw();
+    let null_stream = std::ptr::null_mut();
 
     match backend {
-        GemmBackend::Naive => unsafe {
-            match dtype {
-                DType::F32 => launch_gemm_naive_f32(
-                    a_ptr,
-                    b_ptr,
-                    c_ptr,
-                    m as i32,
-                    n as i32,
-                    k as i32,
-                    null_stream,
-                ),
-                DType::BF16 => launch_gemm_naive_bf16(
-                    a_ptr,
-                    b_ptr,
-                    c_ptr,
-                    m as i32,
-                    n as i32,
-                    k as i32,
-                    null_stream,
-                ),
-                _ => panic!("unsupported dtype for naive GEMM"),
+        GemmBackend::Naive => {
+            unsafe {
+                match dtype {
+                    DType::F32 => launch_gemm_naive_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    DType::BF16 => launch_gemm_naive_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    _ => panic!("unsupported dtype for naive GEMM"),
+                }
             }
-        },
-        GemmBackend::Tiled => unsafe {
-            match dtype {
-                DType::F32 => launch_gemm_tiled_f32(
-                    a_ptr,
-                    b_ptr,
-                    c_ptr,
-                    m as i32,
-                    n as i32,
-                    k as i32,
-                    null_stream,
-                ),
-                DType::BF16 => launch_gemm_tiled_bf16(
-                    a_ptr,
-                    b_ptr,
-                    c_ptr,
-                    m as i32,
-                    n as i32,
-                    k as i32,
-                    null_stream,
-                ),
-                _ => panic!("unsupported dtype for tiled GEMM"),
+        }
+        GemmBackend::Tiled => {
+            unsafe {
+                match dtype {
+                    DType::F32 => launch_gemm_tiled_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    DType::BF16 => launch_gemm_tiled_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
+                    _ => panic!("unsupported dtype for tiled GEMM"),
+                }
             }
-        },
+        }
         GemmBackend::CuBlas => {
-            if m == 1 && dtype == DType::BF16 && n >= 256 {
-                let mut fp32_buf =
-                    xserv_cuda::allocator::cached_alloc(gemv_scratch_elems(k, n) * 4).unwrap();
+            // Fast path: custom GEMV for M=1 BF16 (bandwidth-optimal decode)
+            if m == 1 && dtype == DType::BF16 {
+                let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap();
                 unsafe {
                     launch_gemv_bf16(
-                        a_ptr,
-                        b_ptr,
-                        c_ptr,
+                        a_ptr, b_ptr, c_ptr,
                         fp32_buf.as_mut_ptr() as *mut c_void,
-                        k as i32,
-                        n as i32,
+                        k as i32, n as i32,
                         null_stream,
                     );
                 }
+                // fp32_buf returned to caching allocator pool on drop
             } else {
+                // cuBLAS uses column-major, but we have row-major tensors.
+                // Trick: compute C^T = B^T @ A^T, which gives us C in row-major.
+                // cuBLAS sees our row-major data as column-major transposed.
                 let alpha = 1.0f32;
                 let beta = 0.0f32;
 
@@ -364,28 +179,20 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
 
                 with_cublas(|handle| unsafe {
                     cublasSetStream_v2(handle, null_stream);
+                    // Row-major trick: swap A/B and transpose flags
+                    // C(row-major) = A @ B  <=>  C^T(col-major) = B^T @ A^T
                     error::check(cublasGemmEx(
                         handle,
-                        CUBLAS_OP_N,
-                        CUBLAS_OP_N,
-                        n as i32,
-                        m as i32,
-                        k as i32,
+                        CUBLAS_OP_N, CUBLAS_OP_N,
+                        n as i32, m as i32, k as i32,
                         &alpha as *const f32 as *const c_void,
-                        b_ptr,
-                        b_type,
-                        n as i32,
-                        a_ptr,
-                        a_type,
-                        k as i32,
+                        b_ptr, b_type, n as i32,  // B as col-major = B^T
+                        a_ptr, a_type, k as i32,  // A as col-major = A^T
                         &beta as *const f32 as *const c_void,
-                        c_ptr,
-                        c_type,
-                        n as i32,
+                        c_ptr, c_type, n as i32,  // C as col-major = C^T
                         CUBLAS_COMPUTE_32F,
-                        -1,
-                    ))
-                    .expect("cuBLAS GEMM failed");
+                        -1, // default algo
+                    )).expect("cuBLAS GEMM failed");
                 });
             }
         }
@@ -439,34 +246,21 @@ pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
     let stride_c = (m * n) as i64;
 
     with_cublas(|handle| unsafe {
-        cublasSetStream_v2(handle, xserv_cuda::current_stream_raw());
+        cublasSetStream_v2(handle, std::ptr::null_mut());
         // Row-major trick: C = A @ B ⟺ C^T = B^T @ A^T (col-major)
         error::check(cublasGemmStridedBatchedEx(
             handle,
-            CUBLAS_OP_N,
-            CUBLAS_OP_N,
-            n as i32,
-            m as i32,
-            k as i32,
+            CUBLAS_OP_N, CUBLAS_OP_N,
+            n as i32, m as i32, k as i32,
             &alpha as *const f32 as *const c_void,
-            b.data_ptr() as _,
-            b_type,
-            n as i32,
-            stride_b,
-            a.data_ptr() as _,
-            a_type,
-            k as i32,
-            stride_a,
+            b.data_ptr() as _, b_type, n as i32, stride_b,
+            a.data_ptr() as _, a_type, k as i32, stride_a,
             &beta as *const f32 as *const c_void,
-            c.data_ptr() as *mut c_void,
-            c_type,
-            n as i32,
-            stride_c,
+            c.data_ptr() as *mut c_void, c_type, n as i32, stride_c,
             batch as i32,
             CUBLAS_COMPUTE_32F,
             -1,
-        ))
-        .expect("cuBLAS batched GEMM failed");
+        )).expect("cuBLAS batched GEMM failed");
     });
     c
 }

crates/xserv-kernels/src/layernorm.rs

@@ -2,26 +2,10 @@ use std::ffi::c_void;
 use xserv_tensor::{DType, Device, Tensor};
 
 unsafe extern "C" {
-    fn launch_layernorm_f32(
-        x: *const c_void,
-        gamma: *const c_void,
-        beta: *const c_void,
-        out: *mut c_void,
-        rows: i32,
-        hidden_size: i32,
-        eps: f32,
-        stream: *mut c_void,
-    );
-    fn launch_layernorm_bf16(
-        x: *const c_void,
-        gamma: *const c_void,
-        beta: *const c_void,
-        out: *mut c_void,
-        rows: i32,
-        hidden_size: i32,
-        eps: f32,
-        stream: *mut c_void,
-    );
+    fn launch_layernorm_f32(x: *const c_void, gamma: *const c_void, beta: *const c_void,
+                            out: *mut c_void, rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+    fn launch_layernorm_bf16(x: *const c_void, gamma: *const c_void, beta: *const c_void,
+                             out: *mut c_void, rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
 }
 
 pub fn layernorm(x: &Tensor, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor {
@@ -33,37 +17,21 @@ pub fn layernorm(x: &Tensor, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor
     assert_eq!(beta.shape(), &[hidden_size]);
 
     let rows = x.numel() / hidden_size;
-    assert!(
-        rows <= i32::MAX as usize,
-        "too many rows for i32 kernel param"
-    );
-    assert!(
-        hidden_size <= i32::MAX as usize,
-        "hidden_size too large for i32 kernel param"
-    );
+    assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
+    assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
     let out = Tensor::empty(x.shape(), x.dtype(), x.device());
 
     unsafe {
         match x.dtype() {
             DType::F32 => launch_layernorm_f32(
-                x.data_ptr() as _,
-                gamma.data_ptr() as _,
-                beta.data_ptr() as _,
+                x.data_ptr() as _, gamma.data_ptr() as _, beta.data_ptr() as _,
                 out.data_ptr() as *mut c_void,
-                rows as i32,
-                hidden_size as i32,
-                eps,
-                xserv_cuda::current_stream_raw(),
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
             ),
             DType::BF16 => launch_layernorm_bf16(
-                x.data_ptr() as _,
-                gamma.data_ptr() as _,
-                beta.data_ptr() as _,
+                x.data_ptr() as _, gamma.data_ptr() as _, beta.data_ptr() as _,
                 out.data_ptr() as *mut c_void,
-                rows as i32,
-                hidden_size as i32,
-                eps,
-                xserv_cuda::current_stream_raw(),
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
             ),
             _ => panic!("unsupported dtype for layernorm"),
         }

crates/xserv-kernels/src/lib.rs

@@ -1,34 +1,23 @@
 pub mod activation;
-pub mod argmax;
 pub mod attention;
 pub mod dispatch;
 pub mod embedding;
 pub mod gemm;
 pub mod layernorm;
-pub mod moe;
-pub mod quantization;
 pub mod rmsnorm;
 pub mod rope;
 pub mod softmax;
 pub mod transpose;
 
-pub use activation::{add, bias_add_2d, gelu, gpt_oss_glu, mul, scale, silu, silu_mul};
-pub use argmax::{argmax_bf16_single, argmax_bf16_to_host};
-pub use attention::{
-    attention, copy_kv_position, decode_attention, flash_attention, flash_attention_sinks,
-    paged_decode_attention, paged_decode_attention_sinks, paged_decode_attention_tree,
-    reshape_and_cache_batched_bf16, reshape_and_cache_bf16,
-};
-pub use embedding::{embedding, embedding_device_ids};
-pub use gemm::{GemmBackend, batched_matmul, matmul, matmul_batched_gemv};
+pub use activation::{add, gelu, mul, scale, silu, silu_mul};
+pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
+pub use attention::{attention, decode_attention, flash_attention, paged_decode_attention};
+pub use embedding::embedding;
+pub use gemm::{batched_matmul, matmul, GemmBackend};
 pub use layernorm::layernorm;
 pub use rmsnorm::{add_rmsnorm, rmsnorm};
-pub use rope::{RopeCache, rope_inplace, rope_inplace_device_pos};
+pub use rope::{rope_inplace, RopeCache};
 pub use softmax::softmax;
-pub use transpose::{
-    merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu,
-    transpose_for_rope_gpu, transpose_from_rope_gpu,
-};
 
 /// Register GPU kernels with the tensor crate. Call once at startup.
 pub fn init() {

crates/xserv-kernels/src/moe.rs

@@ -1,474 +0,0 @@
-use std::ffi::c_void;
-use xserv_tensor::{DType, Tensor};
-
-use crate::gemm::{CublasHandle, cublas_handle};
-
-unsafe extern "C" {
-    fn launch_moe_topk_softmax_bf16(
-        router_logits: *const c_void,
-        topk_ids: *mut c_void,
-        topk_weights: *mut c_void,
-        num_tokens: i32,
-        num_experts: i32,
-        top_k: i32,
-        stream: *mut c_void,
-    );
-    fn launch_moe_replicate_bf16(
-        x: *const c_void,
-        x_rep: *mut c_void,
-        num_tokens: i32,
-        hidden: i32,
-        local_experts: i32,
-        stream: *mut c_void,
-    );
-    fn launch_moe_bias_add_3d_bf16(
-        x: *mut c_void,
-        bias: *const c_void,
-        batch: i32,
-        num_tokens: i32,
-        dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_moe_weighted_sum_bf16(
-        expert_out: *const c_void,
-        topk_ids: *const c_void,
-        topk_weights: *const c_void,
-        out: *mut c_void,
-        num_tokens: i32,
-        hidden: i32,
-        top_k: i32,
-        expert_start: i32,
-        local_experts: i32,
-        stream: *mut c_void,
-    );
-
-    fn launch_moe_sparse_gemv_fp8_bf16(
-        x: *const c_void,
-        w: *const c_void,
-        w_scales: *const c_void,
-        bias: *const c_void,
-        topk_ids: *const c_void,
-        y: *mut c_void,
-        num_tokens: i32,
-        n: i32,
-        k: i32,
-        top_k: i32,
-        expert_start: i32,
-        local_experts: i32,
-        x_per_slot: i32,
-        stream: *mut c_void,
-    );
-    fn launch_moe_sparse_gemv_mxfp4_bf16(
-        x: *const c_void,
-        w_packed: *const c_void,
-        w_scales: *const c_void,
-        bias: *const c_void,
-        topk_ids: *const c_void,
-        y: *mut c_void,
-        num_tokens: i32,
-        n: i32,
-        k: i32,
-        top_k: i32,
-        expert_start: i32,
-        local_experts: i32,
-        x_per_slot: i32,
-        stream: *mut c_void,
-    );
-    fn launch_moe_weighted_sum_sparse_bf16(
-        down: *const c_void,
-        topk_ids: *const c_void,
-        topk_weights: *const c_void,
-        out: *mut c_void,
-        num_tokens: i32,
-        hidden: i32,
-        top_k: i32,
-        expert_start: i32,
-        local_experts: i32,
-        stream: *mut c_void,
-    );
-
-    fn cublasGemmStridedBatchedEx(
-        handle: CublasHandle,
-        transa: i32,
-        transb: i32,
-        m: i32,
-        n: i32,
-        k: i32,
-        alpha: *const c_void,
-        a: *const c_void,
-        a_type: i32,
-        lda: i32,
-        stride_a: i64,
-        b: *const c_void,
-        b_type: i32,
-        ldb: i32,
-        stride_b: i64,
-        beta: *const c_void,
-        c: *mut c_void,
-        c_type: i32,
-        ldc: i32,
-        stride_c: i64,
-        batch_count: i32,
-        compute_type: i32,
-        algo: i32,
-    ) -> i32;
-
-    fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
-}
-
-const CUDA_R_16BF: i32 = 14;
-const CUBLAS_COMPUTE_32F: i32 = 68;
-const CUBLAS_GEMM_DEFAULT: i32 = -1;
-
-/// GPU top-k selection + softmax over router logits.
-///
-/// Input:  router_logits [num_tokens, num_experts] BF16 on GPU
-/// Output: (topk_ids [num_tokens, top_k] i32, topk_weights [num_tokens, top_k] f32)
-pub fn moe_topk_softmax(
-    router_logits: &Tensor,
-    num_experts: usize,
-    top_k: usize,
-) -> (Tensor, Tensor) {
-    assert_eq!(router_logits.ndim(), 2);
-    assert_eq!(router_logits.dtype(), DType::BF16);
-    assert!(router_logits.is_contiguous());
-    let num_tokens = router_logits.shape()[0];
-    assert_eq!(router_logits.shape()[1], num_experts);
-
-    // NOTE: topk_ids actually holds i32 expert indices; DType has no I32, so
-    // this is a raw 4-byte buffer mislabeled F32. Never read it as floats —
-    // all consumers (weighted-sum / sparse GEMV kernels) cast to int*.
-    let topk_ids = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());
-    let topk_weights = Tensor::empty(&[num_tokens, top_k], DType::F32, router_logits.device());
-
-    unsafe {
-        launch_moe_topk_softmax_bf16(
-            router_logits.data_ptr() as *const c_void,
-            topk_ids.data_ptr() as *mut c_void,
-            topk_weights.data_ptr() as *mut c_void,
-            num_tokens as i32,
-            num_experts as i32,
-            top_k as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    (topk_ids, topk_weights)
-}
-
-/// Replicate x [num_tokens, hidden] → [local_experts, num_tokens, hidden].
-pub fn moe_replicate(x: &Tensor, local_experts: usize) -> Tensor {
-    assert_eq!(x.ndim(), 2);
-    assert_eq!(x.dtype(), DType::BF16);
-    assert!(x.is_contiguous());
-    let num_tokens = x.shape()[0];
-    let hidden = x.shape()[1];
-    let out = Tensor::empty(
-        &[local_experts, num_tokens, hidden],
-        DType::BF16,
-        x.device(),
-    );
-
-    unsafe {
-        launch_moe_replicate_bf16(
-            x.data_ptr() as *const c_void,
-            out.data_ptr() as *mut c_void,
-            num_tokens as i32,
-            hidden as i32,
-            local_experts as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    out
-}
-
-/// In-place 3D bias add: x [batch, num_tokens, dim] += bias [batch, dim].
-pub fn moe_bias_add_3d(x: &Tensor, bias: &Tensor) {
-    assert_eq!(x.ndim(), 3);
-    assert_eq!(bias.ndim(), 2);
-    assert_eq!(x.dtype(), DType::BF16);
-    assert!(x.is_contiguous());
-    let batch = x.shape()[0];
-    let num_tokens = x.shape()[1];
-    let dim = x.shape()[2];
-    assert_eq!(bias.shape(), &[batch, dim]);
-
-    unsafe {
-        launch_moe_bias_add_3d_bf16(
-            x.data_ptr() as *mut c_void,
-            bias.data_ptr() as *const c_void,
-            batch as i32,
-            num_tokens as i32,
-            dim as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-}
-
-/// Weighted sum of expert outputs → [num_tokens, hidden].
-///
-/// expert_out:   [local_experts, num_tokens, hidden] BF16
-/// topk_ids:     [num_tokens, top_k] i32 (global expert indices)
-/// topk_weights: [num_tokens, top_k] f32
-pub fn moe_weighted_sum(
-    expert_out: &Tensor,
-    topk_ids: &Tensor,
-    topk_weights: &Tensor,
-    expert_start: usize,
-    local_experts: usize,
-    top_k: usize,
-) -> Tensor {
-    assert_eq!(expert_out.ndim(), 3);
-    assert_eq!(expert_out.dtype(), DType::BF16);
-    let num_tokens = expert_out.shape()[1];
-    let hidden = expert_out.shape()[2];
-
-    let out = Tensor::empty(&[num_tokens, hidden], DType::BF16, expert_out.device());
-
-    unsafe {
-        launch_moe_weighted_sum_bf16(
-            expert_out.data_ptr() as *const c_void,
-            topk_ids.data_ptr() as *const c_void,
-            topk_weights.data_ptr() as *const c_void,
-            out.data_ptr() as *mut c_void,
-            num_tokens as i32,
-            hidden as i32,
-            top_k as i32,
-            expert_start as i32,
-            local_experts as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    out
-}
-
-/// Sparse MoE GEMV (FP8 W8A16): compute only the routed experts.
-///
-/// x:        [num_tokens, K] BF16 (x_per_slot=false, gate_up) or
-///           [num_tokens * top_k, K] BF16 (x_per_slot=true, down)
-/// w_fp8_t:  [local_experts, N, K] FP8E4M3 (transposed weight layout)
-/// w_scales: [local_experts] F32 per-expert scalar scales
-/// bias:     [local_experts, N] BF16 (fused into the epilogue)
-/// topk_ids: [num_tokens, top_k] i32 global expert ids (GPU)
-///
-/// Returns y [num_tokens, top_k, N] BF16. Slots routed to experts NOT
-/// owned by this rank are left UNWRITTEN (uninitialized memory) — the
-/// consumer must skip them (see moe_weighted_sum_sparse).
-#[allow(clippy::too_many_arguments)]
-pub fn moe_sparse_gemv_fp8(
-    x: &Tensor,
-    w_fp8_t: &Tensor,
-    w_scales: &Tensor,
-    bias: &Tensor,
-    topk_ids: &Tensor,
-    num_tokens: usize,
-    top_k: usize,
-    expert_start: usize,
-    local_experts: usize,
-    x_per_slot: bool,
-) -> Tensor {
-    assert_eq!(x.dtype(), DType::BF16);
-    assert!(x.is_contiguous());
-    assert_eq!(w_fp8_t.dtype(), DType::FP8E4M3);
-    let n = w_fp8_t.shape()[1];
-    let k = w_fp8_t.shape()[2];
-    // The kernel reads weights as uint4 (16 FP8 values per lane) and would
-    // silently skip a K%16 tail.
-    assert_eq!(k % 16, 0, "sparse FP8 GEMV requires K % 16 == 0, got {k}");
-    assert_eq!(x.shape()[x.ndim() - 1], k);
-    assert_eq!(
-        x.shape()[0],
-        if x_per_slot {
-            num_tokens * top_k
-        } else {
-            num_tokens
-        }
-    );
-
-    let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
-    unsafe {
-        launch_moe_sparse_gemv_fp8_bf16(
-            x.data_ptr() as *const c_void,
-            w_fp8_t.data_ptr() as *const c_void,
-            w_scales.data_ptr() as *const c_void,
-            bias.data_ptr() as *const c_void,
-            topk_ids.data_ptr() as *const c_void,
-            y.data_ptr() as *mut c_void,
-            num_tokens as i32,
-            n as i32,
-            k as i32,
-            top_k as i32,
-            expert_start as i32,
-            local_experts as i32,
-            x_per_slot as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-    y
-}
-
-/// Sparse MoE GEMV (MXFP4 W4A16): same contract as moe_sparse_gemv_fp8,
-/// with packed 4-bit weights [E, N, K/2] + UE8M0 block scales [E, N, K/32].
-#[allow(clippy::too_many_arguments)]
-pub fn moe_sparse_gemv_mxfp4(
-    x: &Tensor,
-    w_packed: &Tensor,
-    w_scales: &Tensor,
-    bias: &Tensor,
-    topk_ids: &Tensor,
-    num_tokens: usize,
-    top_k: usize,
-    n: usize,
-    k: usize,
-    expert_start: usize,
-    local_experts: usize,
-    x_per_slot: bool,
-) -> Tensor {
-    assert_eq!(x.dtype(), DType::BF16);
-    assert!(x.is_contiguous());
-    // 32-element MXFP4 blocks, read as uint4 (32 nibbles) per lane.
-    assert_eq!(k % 32, 0, "sparse MXFP4 GEMV requires K % 32 == 0, got {k}");
-    assert_eq!(x.shape()[x.ndim() - 1], k);
-    assert_eq!(
-        x.shape()[0],
-        if x_per_slot {
-            num_tokens * top_k
-        } else {
-            num_tokens
-        }
-    );
-
-    let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
-    unsafe {
-        launch_moe_sparse_gemv_mxfp4_bf16(
-            x.data_ptr() as *const c_void,
-            w_packed.data_ptr() as *const c_void,
-            w_scales.data_ptr() as *const c_void,
-            bias.data_ptr() as *const c_void,
-            topk_ids.data_ptr() as *const c_void,
-            y.data_ptr() as *mut c_void,
-            num_tokens as i32,
-            n as i32,
-            k as i32,
-            top_k as i32,
-            expert_start as i32,
-            local_experts as i32,
-            x_per_slot as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-    y
-}
-
-/// Weighted sum over the slot axis of the sparse GEMV output.
-///
-/// down: [num_tokens, top_k, hidden] BF16 (non-local slots uninitialized
-/// and skipped, never multiplied by zero — NaN * 0 = NaN).
-pub fn moe_weighted_sum_sparse(
-    down: &Tensor,
-    topk_ids: &Tensor,
-    topk_weights: &Tensor,
-    expert_start: usize,
-    local_experts: usize,
-) -> Tensor {
-    assert_eq!(down.ndim(), 3);
-    assert_eq!(down.dtype(), DType::BF16);
-    let num_tokens = down.shape()[0];
-    let top_k = down.shape()[1];
-    let hidden = down.shape()[2];
-
-    let out = Tensor::empty(&[num_tokens, hidden], DType::BF16, down.device());
-    unsafe {
-        launch_moe_weighted_sum_sparse_bf16(
-            down.data_ptr() as *const c_void,
-            topk_ids.data_ptr() as *const c_void,
-            topk_weights.data_ptr() as *const c_void,
-            out.data_ptr() as *mut c_void,
-            num_tokens as i32,
-            hidden as i32,
-            top_k as i32,
-            expert_start as i32,
-            local_experts as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-    out
-}
-
-/// Strided batched GEMM for MoE expert forward.
-/// C[b] = A[b] @ B[b]  for b in 0..batch
-///
-/// A: [batch, M, K] BF16 contiguous
-/// B: [batch, K, N] BF16 contiguous
-/// Returns C: [batch, M, N] BF16
-#[allow(clippy::too_many_arguments)]
-pub fn batched_gemm_strided(a: &Tensor, b: &Tensor) -> Tensor {
-    assert_eq!(a.ndim(), 3);
-    assert_eq!(b.ndim(), 3);
-    assert_eq!(a.dtype(), DType::BF16);
-    assert_eq!(b.dtype(), DType::BF16);
-    assert!(a.is_contiguous() && b.is_contiguous());
-    assert_eq!(a.shape()[0], b.shape()[0]);
-    assert_eq!(a.shape()[2], b.shape()[1]);
-
-    let batch = a.shape()[0];
-    let m = a.shape()[1];
-    let k = a.shape()[2];
-    let n = b.shape()[2];
-
-    let c = Tensor::empty(&[batch, m, n], DType::BF16, a.device());
-
-    let alpha: f32 = 1.0;
-    let beta: f32 = 0.0;
-
-    // cuBLAS column-major: we compute C^T = B^T @ A^T
-    // A is [batch, M, K] row-major → A^T is [K, M] col-major, lda=K
-    // B is [batch, K, N] row-major → B^T is [N, K] col-major, ldb=N? No...
-    //
-    // Actually for row-major: A[M,K] in memory = col-major A^T[K,M] with lda=K.
-    // So we call cublasGemmStridedBatchedEx with:
-    //   transa=N, transb=N
-    //   m=N, n=M, k=K  (because cuBLAS sees col-major)
-    //   A_cublas = B_row (pointer), lda=N
-    //   B_cublas = A_row (pointer), ldb=K
-    //   C_cublas = C_row (pointer), ldc=N
-
-    let stride_a = (m * k) as i64;
-    let stride_b = (k * n) as i64;
-    let stride_c = (m * n) as i64;
-
-    let handle = cublas_handle();
-    unsafe {
-        cublasSetStream_v2(handle, xserv_cuda::current_stream_raw());
-        let status = cublasGemmStridedBatchedEx(
-            handle,
-            0,
-            0, // CUBLAS_OP_N, CUBLAS_OP_N
-            n as i32,
-            m as i32,
-            k as i32,
-            &alpha as *const f32 as *const c_void,
-            b.data_ptr() as *const c_void,
-            CUDA_R_16BF,
-            n as i32,
-            stride_b,
-            a.data_ptr() as *const c_void,
-            CUDA_R_16BF,
-            k as i32,
-            stride_a,
-            &beta as *const f32 as *const c_void,
-            c.data_ptr() as *mut c_void,
-            CUDA_R_16BF,
-            n as i32,
-            stride_c,
-            batch as i32,
-            CUBLAS_COMPUTE_32F,
-            CUBLAS_GEMM_DEFAULT,
-        );
-        assert_eq!(status, 0, "cublasGemmStridedBatchedEx failed: {status}");
-    }
-
-    c
-}

crates/xserv-kernels/src/quantization.rs

@@ -1,603 +0,0 @@
-use std::cell::RefCell;
-use std::collections::HashMap;
-use std::ffi::c_void;
-use xserv_cuda::GpuBuffer;
-use xserv_tensor::{DType, Tensor};
-
-// ============================================================
-// FFI: custom CUDA kernels
-// ============================================================
-
-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,
-    );
-    fn launch_quantize_bf16_to_fp8e4m3_rowwise(
-        src: *const c_void,
-        dst: *mut c_void,
-        scales: *mut c_void,
-        num_rows: i32,
-        cols: i32,
-        stream: *mut c_void,
-    );
-    fn launch_rowwise_scale_moe_bf16(
-        data: *mut c_void,
-        a_scales: *const c_void,
-        b_scales: *const c_void,
-        num_rows: i32,
-        cols: i32,
-        tokens: i32,
-        stream: *mut c_void,
-    );
-    fn launch_batched_gemv_mxfp4_bf16(
-        x: *const c_void,
-        w_packed: *const c_void,
-        w_scales: *const c_void,
-        y: *mut c_void,
-        e: i32,
-        n: i32,
-        k: i32,
-        stream: *mut c_void,
-    );
-    fn launch_dequant_mxfp4_to_bf16_t(
-        w_packed: *const c_void,
-        w_scales: *const c_void,
-        out: *mut c_void,
-        e: i32,
-        n: i32,
-        k: i32,
-        stream: *mut c_void,
-    );
-}
-
-// ============================================================
-// FFI: cuBLASLt
-// ============================================================
-
-type CublasLtHandle = *mut c_void;
-type CublasLtMatmulDesc = *mut c_void;
-type CublasLtMatrixLayout = *mut c_void;
-type CublasLtMatmulPreference = *mut c_void;
-
-#[repr(C)]
-#[derive(Clone, Copy)]
-struct CublasLtMatmulAlgo {
-    data: [u64; 8],
-}
-
-#[repr(C)]
-struct CublasLtMatmulHeuristicResult {
-    algo: CublasLtMatmulAlgo,
-    workspace_size: usize,
-    state: i32,
-    _reserved: [f32; 4],
-}
-
-unsafe extern "C" {
-    fn cublasLtCreate(handle: *mut CublasLtHandle) -> i32;
-    fn cublasLtDestroy(handle: CublasLtHandle) -> i32;
-    fn cublasLtMatmulDescCreate(
-        desc: *mut CublasLtMatmulDesc,
-        compute_type: i32,
-        scale_type: i32,
-    ) -> i32;
-    fn cublasLtMatmulDescDestroy(desc: CublasLtMatmulDesc) -> i32;
-    fn cublasLtMatmulDescSetAttribute(
-        desc: CublasLtMatmulDesc,
-        attr: i32,
-        buf: *const c_void,
-        size: usize,
-    ) -> i32;
-    fn cublasLtMatrixLayoutCreate(
-        layout: *mut CublasLtMatrixLayout,
-        dtype: i32,
-        rows: u64,
-        cols: u64,
-        ld: i64,
-    ) -> i32;
-    fn cublasLtMatrixLayoutDestroy(layout: CublasLtMatrixLayout) -> i32;
-    fn cublasLtMatrixLayoutSetAttribute(
-        layout: CublasLtMatrixLayout,
-        attr: i32,
-        buf: *const c_void,
-        size: usize,
-    ) -> i32;
-    fn cublasLtMatmulPreferenceCreate(pref: *mut CublasLtMatmulPreference) -> i32;
-    fn cublasLtMatmulPreferenceDestroy(pref: CublasLtMatmulPreference) -> i32;
-    fn cublasLtMatmulPreferenceSetAttribute(
-        pref: CublasLtMatmulPreference,
-        attr: i32,
-        buf: *const c_void,
-        size: usize,
-    ) -> i32;
-    fn cublasLtMatmulAlgoGetHeuristic(
-        handle: CublasLtHandle,
-        desc: CublasLtMatmulDesc,
-        a_layout: CublasLtMatrixLayout,
-        b_layout: CublasLtMatrixLayout,
-        c_layout: CublasLtMatrixLayout,
-        d_layout: CublasLtMatrixLayout,
-        pref: CublasLtMatmulPreference,
-        requested: i32,
-        results: *mut CublasLtMatmulHeuristicResult,
-        found: *mut i32,
-    ) -> i32;
-    fn cublasLtMatmul(
-        handle: CublasLtHandle,
-        desc: CublasLtMatmulDesc,
-        alpha: *const c_void,
-        a: *const c_void,
-        a_layout: CublasLtMatrixLayout,
-        b: *const c_void,
-        b_layout: CublasLtMatrixLayout,
-        beta: *const c_void,
-        c: *const c_void,
-        c_layout: CublasLtMatrixLayout,
-        d: *mut c_void,
-        d_layout: CublasLtMatrixLayout,
-        algo: *const CublasLtMatmulAlgo,
-        workspace: *mut c_void,
-        workspace_size: usize,
-        stream: *mut c_void,
-    ) -> i32;
-}
-
-// cuBLASLt constants
-const CUBLAS_COMPUTE_32F: i32 = 68;
-const CUDA_R_32F: i32 = 0;
-const CUDA_R_16BF: i32 = 14;
-const CUDA_R_8F_E4M3: i32 = 28;
-
-// MatmulDesc attributes
-const CUBLASLT_MATMUL_DESC_A_SCALE_POINTER: i32 = 17;
-const CUBLASLT_MATMUL_DESC_B_SCALE_POINTER: i32 = 18;
-
-// MatrixLayout attributes
-const CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT: i32 = 5;
-const CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET: i32 = 6;
-
-// MatmulPreference attributes
-const CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES: i32 = 1;
-
-const WORKSPACE_BYTES: usize = 32 * 1024 * 1024;
-
-const CUBLASLT_MATMUL_DESC_TRANSA: i32 = 3;
-
-/// A fully-prepared FP8 matmul plan for one (M, N, K) shape: the matmul
-/// descriptor, the four matrix layouts, and the heuristically-chosen algo.
-/// Built once per shape and reused across every expert and every forward
-/// pass — the heuristic search and descriptor/layout creation are the
-/// expensive parts, so doing them once instead of per-expert-per-layer is
-/// the difference between FP8 being faster or slower than BF16.
-#[derive(Clone, Copy)]
-struct Fp8Plan {
-    desc: CublasLtMatmulDesc,
-    a_layout: CublasLtMatrixLayout,
-    b_layout: CublasLtMatrixLayout,
-    c_layout: CublasLtMatrixLayout,
-    d_layout: CublasLtMatrixLayout,
-    algo: CublasLtMatmulAlgo,
-    workspace_size: usize,
-}
-
-struct CublasLtContext {
-    handle: CublasLtHandle,
-    workspace: GpuBuffer,
-    /// Persistent device scalar holding 1.0, used as the A/B scale pointer.
-    /// Scales are applied post-GEMM, so the in-GEMM scales stay 1.0.
-    one_buf: GpuBuffer,
-    /// Cache of prepared matmul plans keyed by (M, N, K, batch).
-    plans: HashMap<(usize, usize, usize, usize), Fp8Plan>,
-}
-
-impl CublasLtContext {
-    fn new() -> Self {
-        let mut handle = std::ptr::null_mut();
-        let status = unsafe { cublasLtCreate(&mut handle) };
-        assert_eq!(status, 0, "cublasLtCreate failed: {status}");
-        let workspace = GpuBuffer::alloc(WORKSPACE_BYTES).expect("alloc cublasLt workspace");
-        let mut one_buf = GpuBuffer::alloc(4).expect("alloc cublasLt fp8 scale");
-        one_buf
-            .copy_from_host(&1.0f32.to_le_bytes())
-            .expect("init fp8 scale");
-        Self {
-            handle,
-            workspace,
-            one_buf,
-            plans: HashMap::new(),
-        }
-    }
-
-    /// Get the cached strided-batched plan for (m, n, k, batch), building it on
-    /// first use.
-    fn plan(&mut self, m: usize, n: usize, k: usize, batch: usize) -> Fp8Plan {
-        if let Some(p) = self.plans.get(&(m, n, k, batch)) {
-            return *p;
-        }
-        let one_ptr = self.one_buf.as_ptr() as *const c_void;
-        let plan = unsafe { build_fp8_plan(self.handle, one_ptr, m, n, k, batch) };
-        self.plans.insert((m, n, k, batch), plan);
-        plan
-    }
-}
-
-impl Drop for CublasLtContext {
-    fn drop(&mut self) {
-        // Tear down cached plans before destroying the handle.
-        for (_, p) in self.plans.drain() {
-            unsafe {
-                cublasLtMatrixLayoutDestroy(p.a_layout);
-                cublasLtMatrixLayoutDestroy(p.b_layout);
-                cublasLtMatrixLayoutDestroy(p.c_layout);
-                cublasLtMatrixLayoutDestroy(p.d_layout);
-                cublasLtMatmulDescDestroy(p.desc);
-            }
-        }
-        if !self.handle.is_null() {
-            unsafe { cublasLtDestroy(self.handle) };
-        }
-    }
-}
-
-/// Build a strided-batched FP8 matmul plan for `batch` experts of one
-/// (m, n, k) shape. Row-major → cuBLASLt col-major mapping (transA=T,
-/// transB=N, m_lt=N, n_lt=M, k_lt=K). A/B scale pointers stay at 1.0 — both
-/// the per-expert weight scale and the per-token activation scale are applied
-/// post-GEMM in a fused kernel, which lets all experts run in one matmul.
-unsafe fn build_fp8_plan(
-    handle: CublasLtHandle,
-    one_ptr: *const c_void,
-    m: usize,
-    n: usize,
-    k: usize,
-    batch: usize,
-) -> Fp8Plan {
-    let m_lt = n as u64;
-    let n_lt = m as u64;
-    let k_lt = k as u64;
-
-    let mut desc: CublasLtMatmulDesc = std::ptr::null_mut();
-    cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
-
-    // transA=T (required for FP8 on Blackwell)
-    let trans_a: i32 = 1;
-    cublasLtMatmulDescSetAttribute(
-        desc,
-        CUBLASLT_MATMUL_DESC_TRANSA,
-        &trans_a as *const i32 as _,
-        4,
-    );
-    let ptr_sz = std::mem::size_of::<*const c_void>();
-    cublasLtMatmulDescSetAttribute(
-        desc,
-        CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
-        &one_ptr as *const _ as _,
-        ptr_sz,
-    );
-    cublasLtMatmulDescSetAttribute(
-        desc,
-        CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
-        &one_ptr as *const _ as _,
-        ptr_sz,
-    );
-
-    // Per-expert strides in ELEMENTS for the strided-batch layout.
-    let stride_a = (n * k) as i64; // weights [N, K]
-    let stride_b = (m * k) as i64; // activations [M, K]
-    let stride_c = (m * n) as i64; // output [M, N]
-    let bc = batch as i32;
-    let set_batch = |layout: CublasLtMatrixLayout, stride: i64| {
-        cublasLtMatrixLayoutSetAttribute(
-            layout,
-            CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
-            &bc as *const i32 as _,
-            4,
-        );
-        cublasLtMatrixLayoutSetAttribute(
-            layout,
-            CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
-            &stride as *const i64 as _,
-            8,
-        );
-    };
-
-    // "A" layout (weights, transposed): physical (K, N) col-major, ld=K
-    let mut a_layout: CublasLtMatrixLayout = std::ptr::null_mut();
-    cublasLtMatrixLayoutCreate(&mut a_layout, CUDA_R_8F_E4M3, k_lt, m_lt, k as i64);
-    set_batch(a_layout, stride_a);
-    // "B" layout (activations): physical (K, M) col-major, ld=K
-    let mut b_layout: CublasLtMatrixLayout = std::ptr::null_mut();
-    cublasLtMatrixLayoutCreate(&mut b_layout, CUDA_R_8F_E4M3, k_lt, n_lt, k as i64);
-    set_batch(b_layout, stride_b);
-    // "C"/"D" layout (output): physical (N, M) col-major, ld=N
-    let mut c_layout: CublasLtMatrixLayout = std::ptr::null_mut();
-    cublasLtMatrixLayoutCreate(&mut c_layout, CUDA_R_16BF, m_lt, n_lt, m_lt as i64);
-    set_batch(c_layout, stride_c);
-    let mut d_layout: CublasLtMatrixLayout = std::ptr::null_mut();
-    cublasLtMatrixLayoutCreate(&mut d_layout, CUDA_R_16BF, m_lt, n_lt, m_lt as i64);
-    set_batch(d_layout, stride_c);
-
-    let mut pref: CublasLtMatmulPreference = std::ptr::null_mut();
-    cublasLtMatmulPreferenceCreate(&mut pref);
-    let ws_bytes = WORKSPACE_BYTES as u64;
-    cublasLtMatmulPreferenceSetAttribute(
-        pref,
-        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
-        &ws_bytes as *const u64 as _,
-        8,
-    );
-
-    let mut heuristic = std::mem::zeroed::<CublasLtMatmulHeuristicResult>();
-    let mut found: i32 = 0;
-    let status = cublasLtMatmulAlgoGetHeuristic(
-        handle,
-        desc,
-        a_layout,
-        b_layout,
-        c_layout,
-        d_layout,
-        pref,
-        1,
-        &mut heuristic,
-        &mut found,
-    );
-    assert!(
-        status == 0 && found > 0,
-        "cublasLtMatmulAlgoGetHeuristic failed for batched FP8 GEMM (m={m}, n={n}, k={k}, batch={batch}): status={status}, found={found}"
-    );
-    cublasLtMatmulPreferenceDestroy(pref);
-
-    Fp8Plan {
-        desc,
-        a_layout,
-        b_layout,
-        c_layout,
-        d_layout,
-        algo: heuristic.algo,
-        workspace_size: heuristic.workspace_size,
-    }
-}
-
-thread_local! {
-    static CUBLASLT_CTX: RefCell<CublasLtContext> = RefCell::new(CublasLtContext::new());
-}
-
-// ============================================================
-// Public API
-// ============================================================
-
-/// 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,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    out
-}
-
-/// Dynamically quantize a contiguous BF16 tensor to FP8 E4M3 with per-row scales.
-///
-/// src: [num_rows, cols] or [batch, rows, cols] BF16, contiguous, GPU
-/// Treats the tensor as 2D (flattens leading dims into num_rows).
-///
-/// Returns: (fp8_data [same shape] FP8E4M3, scales [total_rows] F32)
-pub fn quantize_bf16_to_fp8_rowwise(src: &Tensor) -> (Tensor, Tensor) {
-    assert_eq!(src.dtype(), DType::BF16);
-    assert!(src.is_contiguous());
-    assert!(src.ndim() >= 2);
-
-    let cols = src.shape()[src.ndim() - 1];
-    let num_rows: usize = src.shape()[..src.ndim() - 1].iter().product();
-
-    let fp8_out = Tensor::empty(src.shape(), DType::FP8E4M3, src.device());
-    let scales = Tensor::empty(&[num_rows], DType::F32, src.device());
-
-    unsafe {
-        launch_quantize_bf16_to_fp8e4m3_rowwise(
-            src.data_ptr() as *const c_void,
-            fp8_out.data_ptr() as *mut c_void,
-            scales.data_ptr() as *mut c_void,
-            num_rows as i32,
-            cols as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    (fp8_out, scales)
-}
-
-/// FP8 batched GEMM via cuBLASLt (transA=T required on Blackwell).
-///
-/// Computes: C[b] = scale_a[b] * scale_b[b] * (A_fp8[b] @ B_fp8_T[b]^T)
-/// effectively C[b] = A[b, M, K] @ W[b, K, N] but W is stored transposed
-/// as [b, N, K] for cuBLASLt FP8 compatibility.
-///
-/// a_fp8:    [batch, M, K] FP8E4M3  (activations, quantized per-row)
-/// a_scales: [batch * M] F32        (per-token activation scales, applied post-GEMM)
-/// b_fp8_t:  [batch, N, K] FP8E4M3  (weights, TRANSPOSED for cuBLASLt)
-/// b_scales: [batch] F32            (per-expert scalar weight scales, applied in-GEMM)
-///
-/// Returns:  [batch, M, N] BF16
-pub fn batched_gemm_fp8(
-    a_fp8: &Tensor,
-    a_scales: &Tensor,
-    b_fp8_t: &Tensor,
-    b_scales: &Tensor,
-) -> Tensor {
-    assert_eq!(a_fp8.ndim(), 3);
-    assert_eq!(b_fp8_t.ndim(), 3);
-    assert_eq!(a_fp8.dtype(), DType::FP8E4M3);
-    assert_eq!(b_fp8_t.dtype(), DType::FP8E4M3);
-    assert!(a_fp8.is_contiguous() && b_fp8_t.is_contiguous());
-    assert_eq!(a_fp8.shape()[0], b_fp8_t.shape()[0]);
-    // b_fp8_t is [batch, N, K] transposed, so b_fp8_t.shape[2] == K == a_fp8.shape[2]
-    assert_eq!(a_fp8.shape()[2], b_fp8_t.shape()[2]);
-
-    let batch = a_fp8.shape()[0];
-    let m = a_fp8.shape()[1]; // tokens
-    let k = a_fp8.shape()[2]; // hidden
-    let n = b_fp8_t.shape()[1]; // out_dim (from transposed weight)
-
-    // a_scales: [batch * M] per-token activation scales (applied post-GEMM, per row).
-    // b_scales: [batch] per-expert scalar weight scales (applied in-GEMM via B-scale ptr).
-    assert_eq!(a_scales.shape()[0], batch * m);
-    assert_eq!(b_scales.shape()[0], batch);
-
-    let c = Tensor::empty(&[batch, m, n], DType::BF16, a_fp8.device());
-
-    CUBLASLT_CTX.with(|cell| {
-        let mut ctx = cell.borrow_mut();
-        let handle = ctx.handle;
-        let ws_ptr = ctx.workspace.as_ptr() as *mut c_void;
-        // Cached strided-batched plan: heuristic + descriptor/layout creation
-        // happen once per (m, n, k, batch). All experts run in ONE matmul.
-        let plan = ctx.plan(m, n, k, batch);
-
-        // alpha=1, beta=0, in-GEMM scales=1.0. The unscaled result
-        //   D_raw[e] = A_fp8[e] @ B_fp8[e]^T
-        // is recovered to the real value by the fused post-scale kernel below.
-        let alpha: f32 = 1.0;
-        let beta: f32 = 0.0;
-
-        unsafe {
-            let status = cublasLtMatmul(
-                handle,
-                plan.desc,
-                &alpha as *const f32 as _,
-                b_fp8_t.data_ptr() as *const c_void, // cuBLASLt "A" = weights
-                plan.a_layout,
-                a_fp8.data_ptr() as *const c_void, // cuBLASLt "B" = activations
-                plan.b_layout,
-                &beta as *const f32 as _,
-                c.data_ptr() as *const c_void, // C (unused with beta=0)
-                plan.c_layout,
-                c.data_ptr() as *mut c_void, // D = output
-                plan.d_layout,
-                &plan.algo,
-                ws_ptr,
-                plan.workspace_size,
-                xserv_cuda::current_stream_raw(),
-            );
-            assert_eq!(
-                status, 0,
-                "batched cublasLtMatmul FP8 failed: status={status}"
-            );
-        }
-    });
-
-    // Post-GEMM: recover the real result in one pass.
-    //   c[e, t, :] *= a_scales[e*M + t] * b_scales[e]
-    // (per-token activation scale × per-expert weight scale). BF16's relative
-    // error is scale-invariant, so applying the scale here is precision-
-    // equivalent to folding it into the GEMM epilogue.
-    let total_rows = (batch * m) as i32;
-    unsafe {
-        launch_rowwise_scale_moe_bf16(
-            c.data_ptr() as *mut c_void,
-            a_scales.data_ptr() as *const c_void,
-            b_scales.data_ptr() as *const c_void,
-            total_rows,
-            n as i32,
-            m as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-
-    c
-}
-
-// ============================================================
-// MXFP4 W4A16 (weight-only 4-bit) for MoE experts
-// ============================================================
-
-/// MXFP4 W4A16 batched GEMV for decode (M=1).
-///
-/// x:        [E, K] BF16 (per-expert activation; replicated across experts)
-/// w_packed: [E, N, K/2] byte tensor — two E2M1 nibbles per byte (lo = even k)
-/// w_scales: [E, N, K/32] byte tensor — UE8M0 scale per 32-element block
-///
-/// Returns:  [E, N] BF16, where y[e,n] = sum_k x[e,k] * dequant(W[e,n,k]).
-pub fn batched_gemv_mxfp4(
-    x: &Tensor,
-    w_packed: &Tensor,
-    w_scales: &Tensor,
-    n: usize,
-    k: usize,
-) -> Tensor {
-    assert_eq!(x.dtype(), DType::BF16);
-    assert!(x.is_contiguous());
-    let e = x.shape()[0];
-    assert_eq!(x.shape()[x.ndim() - 1], k, "GEMV K mismatch");
-
-    let y = Tensor::empty(&[e, n], DType::BF16, x.device());
-    unsafe {
-        launch_batched_gemv_mxfp4_bf16(
-            x.data_ptr() as *const c_void,
-            w_packed.data_ptr() as *const c_void,
-            w_scales.data_ptr() as *const c_void,
-            y.data_ptr() as *mut c_void,
-            e as i32,
-            n as i32,
-            k as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-    y
-}
-
-/// Dequantize MXFP4 weights [E, N, K] → BF16 [E, K, N] for the prefill GEMM path
-/// (the BF16 batched GEMM expects weights as [E, K, N]).
-pub fn dequant_mxfp4_to_bf16_t(
-    w_packed: &Tensor,
-    w_scales: &Tensor,
-    e: usize,
-    n: usize,
-    k: usize,
-) -> Tensor {
-    let out = Tensor::empty(&[e, k, n], DType::BF16, w_packed.device());
-    unsafe {
-        launch_dequant_mxfp4_to_bf16_t(
-            w_packed.data_ptr() as *const c_void,
-            w_scales.data_ptr() as *const c_void,
-            out.data_ptr() as *mut c_void,
-            e as i32,
-            n as i32,
-            k as i32,
-            xserv_cuda::current_stream_raw(),
-        );
-    }
-    out
-}

crates/xserv-kernels/src/rmsnorm.rs

@@ -2,35 +2,13 @@ use std::ffi::c_void;
 use xserv_tensor::{DType, Device, Tensor};
 
 unsafe extern "C" {
-    fn launch_rmsnorm_f32(
-        x: *const c_void,
-        gamma: *const c_void,
-        out: *mut c_void,
-        rows: i32,
-        hidden_size: i32,
-        eps: f32,
-        stream: *mut c_void,
-    );
-    fn launch_rmsnorm_bf16(
-        x: *const c_void,
-        gamma: *const c_void,
-        out: *mut c_void,
-        rows: i32,
-        hidden_size: i32,
-        eps: f32,
-        stream: *mut c_void,
-    );
-    fn launch_add_rmsnorm_bf16(
-        x: *const c_void,
-        residual: *const c_void,
-        gamma: *const c_void,
-        normed_out: *mut c_void,
-        sum_out: *mut c_void,
-        rows: i32,
-        hidden_size: i32,
-        eps: f32,
-        stream: *mut c_void,
-    );
+    fn launch_rmsnorm_f32(x: *const c_void, gamma: *const c_void, out: *mut c_void,
+                          rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+    fn launch_rmsnorm_bf16(x: *const c_void, gamma: *const c_void, out: *mut c_void,
+                           rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
+    fn launch_add_rmsnorm_bf16(x: *const c_void, residual: *const c_void, gamma: *const c_void,
+                               normed_out: *mut c_void, sum_out: *mut c_void,
+                               rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
 }
 
 pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
@@ -42,35 +20,19 @@ pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
     assert_eq!(x.dtype(), gamma.dtype());
 
     let rows = x.numel() / hidden_size;
-    assert!(
-        rows <= i32::MAX as usize,
-        "too many rows for i32 kernel param"
-    );
-    assert!(
-        hidden_size <= i32::MAX as usize,
-        "hidden_size too large for i32 kernel param"
-    );
+    assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
+    assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
     let out = Tensor::empty(x.shape(), x.dtype(), x.device());
 
     unsafe {
         match x.dtype() {
             DType::F32 => launch_rmsnorm_f32(
-                x.data_ptr() as _,
-                gamma.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                rows as i32,
-                hidden_size as i32,
-                eps,
-                xserv_cuda::current_stream_raw(),
+                x.data_ptr() as _, gamma.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
             ),
             DType::BF16 => launch_rmsnorm_bf16(
-                x.data_ptr() as _,
-                gamma.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                rows as i32,
-                hidden_size as i32,
-                eps,
-                xserv_cuda::current_stream_raw(),
+                x.data_ptr() as _, gamma.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, hidden_size as i32, eps, std::ptr::null_mut(),
             ),
             _ => panic!("unsupported dtype for rmsnorm"),
         }
@@ -94,14 +56,8 @@ pub fn add_rmsnorm(x: &Tensor, residual: &Tensor, gamma: &Tensor, eps: f32) -> (
     assert_eq!(gamma.shape(), &[hidden_size]);
 
     let rows = x.numel() / hidden_size;
-    assert!(
-        rows <= i32::MAX as usize,
-        "too many rows for i32 kernel param"
-    );
-    assert!(
-        hidden_size <= i32::MAX as usize,
-        "hidden_size too large for i32 kernel param"
-    );
+    assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
+    assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
     let normed_out = Tensor::empty(x.shape(), DType::BF16, x.device());
     let sum_out = Tensor::empty(x.shape(), DType::BF16, x.device());
 
@@ -115,7 +71,7 @@ pub fn add_rmsnorm(x: &Tensor, residual: &Tensor, gamma: &Tensor, eps: f32) -> (
             rows as i32,
             hidden_size as i32,
             eps,
-            xserv_cuda::current_stream_raw(),
+            std::ptr::null_mut(),
         );
     }
 

crates/xserv-kernels/src/rope.rs

@@ -3,34 +3,15 @@ use xserv_cuda::GpuBuffer;
 use xserv_tensor::{DType, Device, Tensor};
 
 unsafe extern "C" {
-    fn launch_rope_f32(
-        x: *mut c_void,
-        cos_cache: *const c_void,
-        sin_cache: *const c_void,
-        positions: *const c_void,
-        num_tokens: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_rope_bf16(
-        x: *mut c_void,
-        cos_cache: *const c_void,
-        sin_cache: *const c_void,
-        positions: *const c_void,
-        num_tokens: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_compute_rope_cache(
-        cos_cache: *mut c_void,
-        sin_cache: *mut c_void,
-        max_seq_len: i32,
-        half_dim: i32,
-        theta: f32,
-        stream: *mut c_void,
-    );
+    fn launch_rope_f32(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
+                       positions: *const c_void, num_tokens: i32, num_heads: i32,
+                       head_dim: i32, stream: *mut c_void);
+    fn launch_rope_bf16(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
+                        positions: *const c_void, num_tokens: i32, num_heads: i32,
+                        head_dim: i32, stream: *mut c_void);
+    fn launch_compute_rope_cache(cos_cache: *mut c_void, sin_cache: *mut c_void,
+                                 max_seq_len: i32, half_dim: i32, theta: f32,
+                                 stream: *mut c_void);
 }
 
 pub struct RopeCache {
@@ -49,99 +30,12 @@ impl RopeCache {
 
         unsafe {
             launch_compute_rope_cache(
-                cos.as_mut_ptr() as _,
-                sin.as_mut_ptr() as _,
-                max_seq_len as i32,
-                half_dim as i32,
-                theta,
-                xserv_cuda::current_stream_raw(),
+                cos.as_mut_ptr() as _, sin.as_mut_ptr() as _,
+                max_seq_len as i32, half_dim as i32, theta, std::ptr::null_mut(),
             );
         }
 
-        Self {
-            cos,
-            sin,
-            max_seq_len,
-            half_dim,
-        }
-    }
-
-    /// YaRN (Yet another RoPE extensioN) RoPE cache. Applies frequency-dependent
-    /// interpolation so the model can extrapolate beyond its training context.
-    pub fn new_yarn(
-        max_seq_len: usize,
-        head_dim: usize,
-        theta: f64,
-        factor: f64,
-        original_max_pos: usize,
-        beta_fast: f64,
-        beta_slow: f64,
-    ) -> Self {
-        let half_dim = head_dim / 2;
-        let dim = head_dim as f64;
-
-        // find_correction_dim: inverse formula to find dimension from number of rotations
-        let find_correction_dim = |num_rotations: f64| -> f64 {
-            dim * (original_max_pos as f64 / (num_rotations * 2.0 * std::f64::consts::PI)).ln()
-                / (2.0 * theta.ln())
-        };
-
-        let low_raw = find_correction_dim(beta_fast);
-        let high_raw = find_correction_dim(beta_slow);
-        // config has truncate=false, so use raw values (no floor/ceil)
-        let low = low_raw.max(0.0);
-        let high = high_raw.min((half_dim - 1) as f64);
-
-        // Compute inv_freq with YaRN interpolation
-        let mut inv_freq = vec![0.0f64; half_dim];
-        for i in 0..half_dim {
-            let pos_freq = theta.powf((2 * i) as f64 / dim);
-            let inv_freq_extrapolation = 1.0 / pos_freq; // original
-            let inv_freq_interpolation = 1.0 / (factor * pos_freq); // scaled
-
-            // Linear ramp: 0 where we keep original, 1 where we interpolate
-            let ramp = if (high - low).abs() < 0.001 {
-                0.5
-            } else {
-                ((i as f64 - low) / (high - low)).clamp(0.0, 1.0)
-            };
-            let extrapolation_factor = 1.0 - ramp;
-
-            inv_freq[i] = inv_freq_interpolation * (1.0 - extrapolation_factor)
-                + inv_freq_extrapolation * extrapolation_factor;
-        }
-
-        // Attention scaling factor for YaRN: 0.1 * ln(factor) + 1.0
-        let attn_factor = 0.1 * factor.ln() + 1.0;
-
-        // Build cos/sin cache on CPU then upload
-        let total = max_seq_len * half_dim;
-        let mut cos_host = vec![0.0f32; total];
-        let mut sin_host = vec![0.0f32; total];
-        for pos in 0..max_seq_len {
-            for i in 0..half_dim {
-                let angle = pos as f64 * inv_freq[i];
-                cos_host[pos * half_dim + i] = (angle.cos() * attn_factor) as f32;
-                sin_host[pos * half_dim + i] = (angle.sin() * attn_factor) as f32;
-            }
-        }
-
-        let nbytes = total * std::mem::size_of::<f32>();
-        let mut cos = GpuBuffer::alloc(nbytes).expect("alloc yarn cos_cache");
-        let mut sin = GpuBuffer::alloc(nbytes).expect("alloc yarn sin_cache");
-        let cos_bytes =
-            unsafe { std::slice::from_raw_parts(cos_host.as_ptr() as *const u8, nbytes) };
-        let sin_bytes =
-            unsafe { std::slice::from_raw_parts(sin_host.as_ptr() as *const u8, nbytes) };
-        cos.copy_from_host(cos_bytes).unwrap();
-        sin.copy_from_host(sin_bytes).unwrap();
-
-        Self {
-            cos,
-            sin,
-            max_seq_len,
-            half_dim,
-        }
+        Self { cos, sin, max_seq_len, half_dim }
     }
 }
 
@@ -164,46 +58,24 @@ pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
             num_tokens * std::mem::size_of::<u32>(),
         )
     };
-    let mut pos_gpu =
-        xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
+    let mut pos_gpu = xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
     pos_gpu.copy_from_host(pos_bytes).unwrap();
 
-    rope_inplace_device_pos(x, cache, pos_gpu.as_ptr() as *const c_void);
-}
-
-/// RoPE in-place with positions already on the GPU (u32, [num_tokens]).
-/// Used by the CUDA-graph decode path, where the position lives in a
-/// persistent device buffer updated outside the captured region.
-pub fn rope_inplace_device_pos(x: &Tensor, cache: &RopeCache, pos_gpu: *const c_void) {
-    assert_eq!(x.ndim(), 3);
-    assert!(x.is_contiguous());
-    assert!(matches!(x.device(), Device::Cuda(_)));
-    let num_tokens = x.shape()[0];
-    let num_heads = x.shape()[1];
-    let head_dim = x.shape()[2];
-    assert_eq!(head_dim / 2, cache.half_dim);
-
     unsafe {
         match x.dtype() {
             DType::F32 => launch_rope_f32(
                 x.data_ptr() as *mut c_void,
-                cache.cos.as_ptr() as _,
-                cache.sin.as_ptr() as _,
-                pos_gpu,
-                num_tokens as i32,
-                num_heads as i32,
-                head_dim as i32,
-                xserv_cuda::current_stream_raw(),
+                cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
+                pos_gpu.as_ptr() as _,
+                num_tokens as i32, num_heads as i32, head_dim as i32,
+                std::ptr::null_mut(),
             ),
             DType::BF16 => launch_rope_bf16(
                 x.data_ptr() as *mut c_void,
-                cache.cos.as_ptr() as _,
-                cache.sin.as_ptr() as _,
-                pos_gpu,
-                num_tokens as i32,
-                num_heads as i32,
-                head_dim as i32,
-                xserv_cuda::current_stream_raw(),
+                cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
+                pos_gpu.as_ptr() as _,
+                num_tokens as i32, num_heads as i32, head_dim as i32,
+                std::ptr::null_mut(),
             ),
             _ => panic!("unsupported dtype for rope"),
         }

crates/xserv-kernels/src/softmax.rs

@@ -2,20 +2,8 @@ use std::ffi::c_void;
 use xserv_tensor::{DType, Device, Tensor};
 
 unsafe extern "C" {
-    fn launch_softmax_f32(
-        x: *const c_void,
-        out: *mut c_void,
-        rows: i32,
-        cols: i32,
-        stream: *mut c_void,
-    );
-    fn launch_softmax_bf16(
-        x: *const c_void,
-        out: *mut c_void,
-        rows: i32,
-        cols: i32,
-        stream: *mut c_void,
-    );
+    fn launch_softmax_f32(x: *const c_void, out: *mut c_void, rows: i32, cols: i32, stream: *mut c_void);
+    fn launch_softmax_bf16(x: *const c_void, out: *mut c_void, rows: i32, cols: i32, stream: *mut c_void);
 }
 
 /// Softmax along the last dimension.
@@ -26,31 +14,19 @@ pub fn softmax(x: &Tensor) -> Tensor {
 
     let cols = *x.shape().last().unwrap();
     let rows = x.numel() / cols;
-    assert!(
-        rows <= i32::MAX as usize,
-        "too many rows for i32 kernel param"
-    );
-    assert!(
-        cols <= i32::MAX as usize,
-        "cols too large for i32 kernel param"
-    );
+    assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
+    assert!(cols <= i32::MAX as usize, "cols too large for i32 kernel param");
     let out = Tensor::empty(x.shape(), x.dtype(), x.device());
 
     unsafe {
         match x.dtype() {
             DType::F32 => launch_softmax_f32(
-                x.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                rows as i32,
-                cols as i32,
-                xserv_cuda::current_stream_raw(),
+                x.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, cols as i32, std::ptr::null_mut(),
             ),
             DType::BF16 => launch_softmax_bf16(
-                x.data_ptr() as _,
-                out.data_ptr() as *mut c_void,
-                rows as i32,
-                cols as i32,
-                xserv_cuda::current_stream_raw(),
+                x.data_ptr() as _, out.data_ptr() as *mut c_void,
+                rows as i32, cols as i32, std::ptr::null_mut(),
             ),
             _ => panic!("unsupported dtype for softmax"),
         }

crates/xserv-kernels/src/transpose.rs

@@ -2,79 +2,19 @@ use std::ffi::c_void;
 use xserv_tensor::{DType, Device, Tensor};
 
 unsafe extern "C" {
-    fn launch_reshape_heads_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        seq_len: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_merge_heads_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        seq_len: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_transpose_hsd_to_shd_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        seq_len: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_transpose_shd_to_hsd_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        seq_len: i32,
-        num_heads: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_repeat_kv_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        kv_heads: i32,
-        n_rep: i32,
-        seq_len: i32,
-        head_dim: i32,
-        stream: *mut c_void,
-    );
-    fn launch_strided_copy_bf16(
-        inp: *const c_void,
-        out: *mut c_void,
-        numel: i32,
-        ndim: i32,
-        shape0: i32,
-        shape1: i32,
-        shape2: i32,
-        shape3: i32,
-        in_stride0: i32,
-        in_stride1: i32,
-        in_stride2: i32,
-        in_stride3: i32,
-        in_offset: i32,
-        stream: *mut c_void,
-    );
-    fn launch_strided_copy_f32(
-        inp: *const c_void,
-        out: *mut c_void,
-        numel: i32,
-        ndim: i32,
-        shape0: i32,
-        shape1: i32,
-        shape2: i32,
-        shape3: i32,
-        in_stride0: i32,
-        in_stride1: i32,
-        in_stride2: i32,
-        in_stride3: i32,
-        in_offset: i32,
-        stream: *mut c_void,
-    );
+    fn launch_reshape_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_merge_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_repeat_kv_bf16(inp: *const c_void, out: *mut c_void, kv_heads: i32, n_rep: i32, seq_len: i32, head_dim: i32, stream: *mut c_void);
+    fn launch_strided_copy_bf16(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
+        shape0: i32, shape1: i32, shape2: i32, shape3: i32,
+        in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
+        in_offset: i32, stream: *mut c_void);
+    fn launch_strided_copy_f32(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
+        shape0: i32, shape1: i32, shape2: i32, shape3: i32,
+        in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
+        in_offset: i32, stream: *mut c_void);
 }
 
 /// [S, H*D] → [1, H, S, D] on GPU (BF16)
@@ -84,12 +24,8 @@ pub fn reshape_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim:
     let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
     unsafe {
         launch_reshape_heads_bf16(
-            x.data_ptr() as _,
-            out.data_ptr() as *mut c_void,
-            seq_len as i32,
-            num_heads as i32,
-            head_dim as i32,
-            xserv_cuda::current_stream_raw(),
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
         );
     }
     out
@@ -103,58 +39,36 @@ pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: u
     let out = Tensor::empty(&[seq_len, hidden], DType::BF16, x.device());
     unsafe {
         launch_merge_heads_bf16(
-            x.data_ptr() as _,
-            out.data_ptr() as *mut c_void,
-            seq_len as i32,
-            num_heads as i32,
-            head_dim as i32,
-            xserv_cuda::current_stream_raw(),
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
         );
     }
     out
 }
 
 /// [1, H, S, D] → [S, H, D] for RoPE on GPU (BF16)
-pub fn transpose_for_rope_gpu(
-    x: &Tensor,
-    seq_len: usize,
-    num_heads: usize,
-    head_dim: usize,
-) -> Tensor {
+pub fn transpose_for_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
     assert_eq!(x.dtype(), DType::BF16);
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let out = Tensor::empty(&[seq_len, num_heads, head_dim], DType::BF16, x.device());
     unsafe {
         launch_transpose_hsd_to_shd_bf16(
-            x.data_ptr() as _,
-            out.data_ptr() as *mut c_void,
-            seq_len as i32,
-            num_heads as i32,
-            head_dim as i32,
-            xserv_cuda::current_stream_raw(),
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
         );
     }
     out
 }
 
 /// [S, H, D] → [1, H, S, D] after RoPE on GPU (BF16)
-pub fn transpose_from_rope_gpu(
-    x: &Tensor,
-    seq_len: usize,
-    num_heads: usize,
-    head_dim: usize,
-) -> Tensor {
+pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
     assert_eq!(x.dtype(), DType::BF16);
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
     unsafe {
         launch_transpose_shd_to_hsd_bf16(
-            x.data_ptr() as _,
-            out.data_ptr() as *mut c_void,
-            seq_len as i32,
-            num_heads as i32,
-            head_dim as i32,
-            xserv_cuda::current_stream_raw(),
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            seq_len as i32, num_heads as i32, head_dim as i32, std::ptr::null_mut(),
         );
     }
     out
@@ -162,9 +76,7 @@ pub fn transpose_from_rope_gpu(
 
 /// [1, KV_H, S, D] → [1, KV_H*n_rep, S, D] on GPU (BF16)
 pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
-    if n_rep == 1 {
-        return x.clone();
-    }
+    if n_rep == 1 { return x.clone(); }
     assert_eq!(x.dtype(), DType::BF16);
     assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
     let kv_heads = x.shape()[1];
@@ -174,13 +86,8 @@ pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
     let out = Tensor::empty(&[1, new_heads, seq_len, head_dim], DType::BF16, x.device());
     unsafe {
         launch_repeat_kv_bf16(
-            x.data_ptr() as _,
-            out.data_ptr() as *mut c_void,
-            kv_heads as i32,
-            n_rep as i32,
-            seq_len as i32,
-            head_dim as i32,
-            xserv_cuda::current_stream_raw(),
+            x.data_ptr() as _, out.data_ptr() as *mut c_void,
+            kv_heads as i32, n_rep as i32, seq_len as i32, head_dim as i32, std::ptr::null_mut(),
         );
     }
     out
@@ -215,41 +122,20 @@ pub fn strided_to_contiguous_gpu(x: &Tensor) -> Tensor {
     unsafe {
         match x.dtype() {
             DType::BF16 => launch_strided_copy_bf16(
-                storage_ptr as _,
-                out.data_ptr() as *mut c_void,
-                numel as i32,
-                ndim as i32,
-                shape4[0],
-                shape4[1],
-                shape4[2],
-                shape4[3],
-                strides4[0],
-                strides4[1],
-                strides4[2],
-                strides4[3],
-                in_offset,
-                xserv_cuda::current_stream_raw(),
+                storage_ptr as _, out.data_ptr() as *mut c_void,
+                numel as i32, ndim as i32,
+                shape4[0], shape4[1], shape4[2], shape4[3],
+                strides4[0], strides4[1], strides4[2], strides4[3],
+                in_offset, std::ptr::null_mut(),
             ),
             DType::F32 => launch_strided_copy_f32(
-                storage_ptr as _,
-                out.data_ptr() as *mut c_void,
-                numel as i32,
-                ndim as i32,
-                shape4[0],
-                shape4[1],
-                shape4[2],
-                shape4[3],
-                strides4[0],
-                strides4[1],
-                strides4[2],
-                strides4[3],
-                in_offset,
-                xserv_cuda::current_stream_raw(),
-            ),
-            _ => panic!(
-                "strided_to_contiguous_gpu: unsupported dtype {:?}",
-                x.dtype()
+                storage_ptr as _, out.data_ptr() as *mut c_void,
+                numel as i32, ndim as i32,
+                shape4[0], shape4[1], shape4[2], shape4[3],
+                strides4[0], strides4[1], strides4[2], strides4[3],
+                in_offset, std::ptr::null_mut(),
             ),
+            _ => panic!("strided_to_contiguous_gpu: unsupported dtype {:?}", x.dtype()),
         }
     }
     out

crates/xserv-kernels/tests/attention_test.rs

@@ -1,21 +1,11 @@
 use xserv_kernels::*;
 use xserv_tensor::{Device, Tensor};
 
-fn init() {
-    xserv_cuda::device::set_device(0).unwrap();
-}
+fn init() { xserv_cuda::device::set_device(0).unwrap(); }
 
-fn cpu_attention(
-    q: &[f32],
-    k: &[f32],
-    v: &[f32],
-    batch: usize,
-    heads: usize,
-    q_len: usize,
-    kv_len: usize,
-    head_dim: usize,
-    causal: bool,
-) -> Vec<f32> {
+fn cpu_attention(q: &[f32], k: &[f32], v: &[f32],
+                 batch: usize, heads: usize, q_len: usize, kv_len: usize, head_dim: usize,
+                 causal: bool) -> Vec<f32> {
     let mut out = vec![0.0f32; batch * heads * q_len * head_dim];
     let scale = 1.0 / (head_dim as f32).sqrt();
 
@@ -80,13 +70,8 @@ fn check_close(a: &[f32], b: &[f32], atol: f32, name: &str) {
     let mut max_err = 0.0f32;
     for (i, (x, y)) in a.iter().zip(b).enumerate() {
         let err = (x - y).abs();
-        if err > max_err {
-            max_err = err;
-        }
-        assert!(
-            err <= atol,
-            "{name}: mismatch at [{i}]: got {x}, expected {y}, err {err}"
-        );
+        if err > max_err { max_err = err; }
+        assert!(err <= atol, "{name}: mismatch at [{i}]: got {x}, expected {y}, err {err}");
     }
     println!("{name}: max_err = {max_err:.6e}");
 }
@@ -120,9 +105,7 @@ fn test_batched_matmul() {
     for i in 0..m {
         for j in 0..n {
             let mut s = 0.0f32;
-            for kk in 0..k {
-                s += a_cpu[i * k + kk] * b_cpu[kk * n + j];
-            }
+            for kk in 0..k { s += a_cpu[i * k + kk] * b_cpu[kk * n + j]; }
             expected[i * n + j] = s;
         }
     }
@@ -133,10 +116,7 @@ fn test_batched_matmul() {
 #[test]
 fn test_attention_no_causal() {
     init();
-    let b = 1;
-    let h = 2;
-    let s = 8;
-    let d = 16;
+    let b = 1; let h = 2; let s = 8; let d = 16;
     let q_data = make_data(b * h * s * d);
     let k_data = make_data(b * h * s * d);
     let v_data = make_data(b * h * s * d);
@@ -146,21 +126,13 @@ fn test_attention_no_causal() {
     let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
     let v = Tensor::from_slice(&v_data, &[b, h, s, d]).to_device(Device::Cuda(0));
     let out = attention(&q, &k, &v, false).to_device(Device::Cpu);
-    check_close(
-        out.as_slice::<f32>(),
-        &expected,
-        1e-4,
-        "attention_no_causal",
-    );
+    check_close(out.as_slice::<f32>(), &expected, 1e-4, "attention_no_causal");
 }
 
 #[test]
 fn test_attention_causal() {
     init();
-    let b = 1;
-    let h = 2;
-    let s = 16;
-    let d = 32;
+    let b = 1; let h = 2; let s = 16; let d = 32;
     let q_data = make_data(b * h * s * d);
     let k_data = make_data(b * h * s * d);
     let v_data = make_data(b * h * s * d);
@@ -176,10 +148,7 @@ fn test_attention_causal() {
 #[test]
 fn test_attention_causal_larger() {
     init();
-    let b = 2;
-    let h = 4;
-    let s = 64;
-    let d = 64;
+    let b = 2; let h = 4; let s = 64; let d = 64;
     let q_data = make_data(b * h * s * d);
     let k_data = make_data(b * h * s * d);
     let v_data = make_data(b * h * s * d);
@@ -189,28 +158,18 @@ fn test_attention_causal_larger() {
     let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
     let v = Tensor::from_slice(&v_data, &[b, h, s, d]).to_device(Device::Cuda(0));
     let out = attention(&q, &k, &v, true).to_device(Device::Cpu);
-    check_close(
-        out.as_slice::<f32>(),
-        &expected,
-        1e-2,
-        "attention_causal_larger",
-    );
+    check_close(out.as_slice::<f32>(), &expected, 1e-2, "attention_causal_larger");
 }
 
 #[test]
 fn test_attention_causal_first_row_sees_only_first_token() {
     init();
-    let b = 1;
-    let h = 1;
-    let s = 4;
-    let d = 8;
+    let b = 1; let h = 1; let s = 4; let d = 8;
     let q_data = make_data(b * h * s * d);
     let k_data = make_data(b * h * s * d);
-    let v_data: Vec<f32> = (0..s * d)
-        .map(|i| {
-            if i < d { 1.0 } else { 0.0 } // only first V row is nonzero
-        })
-        .collect();
+    let v_data: Vec<f32> = (0..s * d).map(|i| {
+        if i < d { 1.0 } else { 0.0 }  // only first V row is nonzero
+    }).collect();
 
     let q = Tensor::from_slice(&q_data, &[b, h, s, d]).to_device(Device::Cuda(0));
     let k = Tensor::from_slice(&k_data, &[b, h, s, d]).to_device(Device::Cuda(0));
@@ -222,11 +181,7 @@ fn test_attention_causal_first_row_sees_only_first_token() {
     // output[0] should be exactly V[0] = [1, 1, 1, ...1]
     let result = out.as_slice::<f32>();
     for i in 0..d {
-        assert!(
-            (result[i] - 1.0).abs() < 1e-5,
-            "first row should equal V[0], got {} at dim {}",
-            result[i],
-            i
-        );
+        assert!((result[i] - 1.0).abs() < 1e-5,
+            "first row should equal V[0], got {} at dim {}", result[i], i);
     }
 }

crates/xserv-kernels/tests/gemm_test.rs

@@ -1,5 +1,5 @@
 use half::bf16;
-use xserv_kernels::{GemmBackend, matmul};
+use xserv_kernels::{matmul, GemmBackend};
 use xserv_tensor::{Device, Tensor};
 
 fn cpu_matmul_f32(a: &[f32], b: &[f32], m: usize, n: usize, k: usize) -> Vec<f32> {
@@ -75,110 +75,70 @@ fn run_gemm_test_bf16(backend: GemmBackend, m: usize, n: usize, k: usize) {
 // --- F32 tests ---
 
 #[test]
-fn test_gemm_naive_f32_small() {
-    run_gemm_test_f32(GemmBackend::Naive, 4, 4, 4);
-}
+fn test_gemm_naive_f32_small() { run_gemm_test_f32(GemmBackend::Naive, 4, 4, 4); }
 
 #[test]
-fn test_gemm_naive_f32_medium() {
-    run_gemm_test_f32(GemmBackend::Naive, 64, 64, 64);
-}
+fn test_gemm_naive_f32_medium() { run_gemm_test_f32(GemmBackend::Naive, 64, 64, 64); }
 
 #[test]
-fn test_gemm_naive_f32_rect() {
-    run_gemm_test_f32(GemmBackend::Naive, 32, 64, 48);
-}
+fn test_gemm_naive_f32_rect() { run_gemm_test_f32(GemmBackend::Naive, 32, 64, 48); }
 
 #[test]
-fn test_gemm_tiled_f32_small() {
-    run_gemm_test_f32(GemmBackend::Tiled, 4, 4, 4);
-}
+fn test_gemm_tiled_f32_small() { run_gemm_test_f32(GemmBackend::Tiled, 4, 4, 4); }
 
 #[test]
-fn test_gemm_tiled_f32_medium() {
-    run_gemm_test_f32(GemmBackend::Tiled, 128, 128, 128);
-}
+fn test_gemm_tiled_f32_medium() { run_gemm_test_f32(GemmBackend::Tiled, 128, 128, 128); }
 
 #[test]
-fn test_gemm_tiled_f32_rect() {
-    run_gemm_test_f32(GemmBackend::Tiled, 65, 33, 97);
-}
+fn test_gemm_tiled_f32_rect() { run_gemm_test_f32(GemmBackend::Tiled, 65, 33, 97); }
 
 #[test]
-fn test_gemm_cublas_f32_small() {
-    run_gemm_test_f32(GemmBackend::CuBlas, 4, 4, 4);
-}
+fn test_gemm_cublas_f32_small() { run_gemm_test_f32(GemmBackend::CuBlas, 4, 4, 4); }
 
 #[test]
-fn test_gemm_cublas_f32_medium() {
-    run_gemm_test_f32(GemmBackend::CuBlas, 256, 256, 256);
-}
+fn test_gemm_cublas_f32_medium() { run_gemm_test_f32(GemmBackend::CuBlas, 256, 256, 256); }
 
 #[test]
-fn test_gemm_cublas_f32_rect() {
-    run_gemm_test_f32(GemmBackend::CuBlas, 65, 33, 97);
-}
+fn test_gemm_cublas_f32_rect() { run_gemm_test_f32(GemmBackend::CuBlas, 65, 33, 97); }
 
 // --- BF16 tests ---
 
 #[test]
-fn test_gemm_naive_bf16_small() {
-    run_gemm_test_bf16(GemmBackend::Naive, 4, 4, 4);
-}
+fn test_gemm_naive_bf16_small() { run_gemm_test_bf16(GemmBackend::Naive, 4, 4, 4); }
 
 #[test]
-fn test_gemm_naive_bf16_medium() {
-    run_gemm_test_bf16(GemmBackend::Naive, 64, 64, 64);
-}
+fn test_gemm_naive_bf16_medium() { run_gemm_test_bf16(GemmBackend::Naive, 64, 64, 64); }
 
 #[test]
-fn test_gemm_tiled_bf16_small() {
-    run_gemm_test_bf16(GemmBackend::Tiled, 4, 4, 4);
-}
+fn test_gemm_tiled_bf16_small() { run_gemm_test_bf16(GemmBackend::Tiled, 4, 4, 4); }
 
 #[test]
-fn test_gemm_tiled_bf16_medium() {
-    run_gemm_test_bf16(GemmBackend::Tiled, 128, 128, 128);
-}
+fn test_gemm_tiled_bf16_medium() { run_gemm_test_bf16(GemmBackend::Tiled, 128, 128, 128); }
 
 #[test]
-fn test_gemm_cublas_bf16_small() {
-    run_gemm_test_bf16(GemmBackend::CuBlas, 4, 4, 4);
-}
+fn test_gemm_cublas_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 4, 4, 4); }
 
 #[test]
-fn test_gemm_cublas_bf16_medium() {
-    run_gemm_test_bf16(GemmBackend::CuBlas, 256, 256, 256);
-}
+fn test_gemm_cublas_bf16_medium() { run_gemm_test_bf16(GemmBackend::CuBlas, 256, 256, 256); }
 
 // --- Custom GEMV tests (M=1, BF16 fast path) ---
 
 #[test]
-fn test_gemv_bf16_small() {
-    run_gemm_test_bf16(GemmBackend::CuBlas, 1, 64, 64);
-}
+fn test_gemv_bf16_small() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 64, 64); }
 
 #[test]
-fn test_gemv_bf16_medium() {
-    run_gemm_test_bf16(GemmBackend::CuBlas, 1, 256, 256);
-}
+fn test_gemv_bf16_medium() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 256, 256); }
 
 #[test]
-fn test_gemv_bf16_4096() {
-    run_gemm_test_bf16(GemmBackend::CuBlas, 1, 4096, 4096);
-}
+fn test_gemv_bf16_4096() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 4096, 4096); }
 
 #[test]
-fn test_gemv_bf16_rect() {
-    run_gemm_test_bf16(GemmBackend::CuBlas, 1, 512, 4096);
-}
+fn test_gemv_bf16_rect() { run_gemm_test_bf16(GemmBackend::CuBlas, 1, 512, 4096); }
 
 // --- Larger benchmark-style tests ---
 
 #[test]
-fn test_gemm_cublas_f32_1024() {
-    run_gemm_test_f32(GemmBackend::CuBlas, 1024, 1024, 1024);
-}
+fn test_gemm_cublas_f32_1024() { run_gemm_test_f32(GemmBackend::CuBlas, 1024, 1024, 1024); }
 
 #[test]
 fn test_gemm_consistency_all_backends() {

crates/xserv-kernels/tests/ops_test.rs

@@ -2,9 +2,7 @@ use half::bf16;
 use xserv_kernels::*;
 use xserv_tensor::{Device, Tensor};
 
-fn init() {
-    xserv_cuda::device::set_device(0).unwrap();
-}
+fn init() { xserv_cuda::device::set_device(0).unwrap(); }
 
 // --- CPU reference implementations ---
 
@@ -39,12 +37,10 @@ fn cpu_layernorm(x: &[f32], gamma: &[f32], beta: &[f32], eps: f32, hidden: usize
 
 fn cpu_gelu(x: &[f32]) -> Vec<f32> {
     let sqrt_2_over_pi = 0.7978845608f32;
-    x.iter()
-        .map(|&v| {
-            let inner = sqrt_2_over_pi * (v + 0.044715 * v * v * v);
-            0.5 * v * (1.0 + inner.tanh())
-        })
-        .collect()
+    x.iter().map(|&v| {
+        let inner = sqrt_2_over_pi * (v + 0.044715 * v * v * v);
+        0.5 * v * (1.0 + inner.tanh())
+    }).collect()
 }
 
 fn cpu_silu(x: &[f32]) -> Vec<f32> {
@@ -92,13 +88,8 @@ fn check_close(result: &[f32], expected: &[f32], atol: f32, name: &str) {
     let mut max_err = 0.0f32;
     for (i, (r, e)) in result.iter().zip(expected).enumerate() {
         let err = (r - e).abs();
-        if err > max_err {
-            max_err = err;
-        }
-        assert!(
-            err <= atol,
-            "{name}: mismatch at [{i}]: got {r}, expected {e}, err {err}"
-        );
+        if err > max_err { max_err = err; }
+        assert!(err <= atol, "{name}: mismatch at [{i}]: got {r}, expected {e}, err {err}");
     }
     println!("{name}: max_err = {max_err:.6e}");
 }
@@ -217,18 +208,13 @@ fn test_softmax_sum_to_one() {
     init();
     let rows = 4;
     let cols = 2048;
-    let data: Vec<f32> = (0..rows * cols)
-        .map(|i| ((i % 31) as f32 - 15.0) * 0.5)
-        .collect();
+    let data: Vec<f32> = (0..rows * cols).map(|i| ((i % 31) as f32 - 15.0) * 0.5).collect();
     let x = Tensor::from_slice(&data, &[rows, cols]).to_device(Device::Cuda(0));
     let out = softmax(&x).to_device(Device::Cpu);
     let result = out.as_slice::<f32>();
     for r in 0..rows {
         let row_sum: f32 = result[r * cols..(r + 1) * cols].iter().sum();
-        assert!(
-            (row_sum - 1.0).abs() < 1e-5,
-            "softmax row {r} sum = {row_sum}"
-        );
+        assert!((row_sum - 1.0).abs() < 1e-5, "softmax row {r} sum = {row_sum}");
     }
 }
 
@@ -261,10 +247,8 @@ fn test_embedding_f32() {
         for i in 0..hidden {
             let expected = table_data[tid as usize * hidden + i];
             let got = result[seq_idx * hidden + i];
-            assert!(
-                (got - expected).abs() < 1e-6,
-                "embedding mismatch at [{seq_idx},{i}]: got {got}, expected {expected}"
-            );
+            assert!((got - expected).abs() < 1e-6,
+                "embedding mismatch at [{seq_idx},{i}]: got {got}, expected {expected}");
         }
     }
 }
@@ -286,8 +270,8 @@ fn test_rope_f32() {
     let mut expected = x_data.clone();
     cpu_rope(&mut expected, &positions, num_heads, head_dim, theta);
 
-    let x =
-        Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim]).to_device(Device::Cuda(0));
+    let x = Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim])
+        .to_device(Device::Cuda(0));
     let cache = RopeCache::new(64, head_dim, theta);
     rope_inplace(&x, &cache, &positions);
 
@@ -308,8 +292,8 @@ fn test_rope_position_0_identity() {
         .map(|i| (i as f32 + 1.0) * 0.1)
         .collect();
 
-    let x =
-        Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim]).to_device(Device::Cuda(0));
+    let x = Tensor::from_slice(&x_data, &[num_tokens, num_heads, head_dim])
+        .to_device(Device::Cuda(0));
     let cache = RopeCache::new(64, head_dim, 10000.0);
     rope_inplace(&x, &cache, &positions);
 

crates/xserv-model/src/bin/bench-gpt-oss.rs

@@ -1,421 +0,0 @@
-use std::path::PathBuf;
-use std::sync::Arc;
-use std::time::Instant;
-
-use xserv_distributed::{TpContext, UniqueId, get_unique_id};
-use xserv_model::{BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, loader};
-use xserv_tensor::{DType, Device};
-use xserv_tokenizer::Tokenizer;
-
-fn main() {
-    let args: Vec<String> = std::env::args().collect();
-    if args.len() < 2 {
-        eprintln!("Usage: bench-gpt-oss <model-dir> [--max-tokens N] [--tp N]");
-        std::process::exit(1);
-    }
-
-    let model_dir = PathBuf::from(&args[1]);
-    let max_tokens: usize = get_arg(&args, "--max-tokens").unwrap_or(32);
-    let world: usize = get_arg(&args, "--tp").unwrap_or(2);
-
-    let config = ModelConfig::from_file(&model_dir.join("config.json"));
-    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
-
-    eprintln!(
-        "gpt-oss-20b: layers={}, hidden={}, heads={}/{} kv, experts={}, top_k={}, vocab={}",
-        config.num_layers(),
-        config.hidden(),
-        config.num_heads(),
-        config.num_kv_heads(),
-        config.num_experts(),
-        config.experts_per_token(),
-        config.vocab_size
-    );
-    eprintln!("TP world={world}, max_tokens={max_tokens}");
-
-    let max_seq_len: usize = 2048;
-    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
-
-    // TP setup
-    let uid = get_unique_id();
-    let local_kv = config.num_kv_heads() / world;
-
-    // Spawn worker threads for ranks 1..world
-    let mut worker_handles = Vec::new();
-    let mut worker_txs = Vec::new();
-    for rank in 1..world {
-        let (tx, rx) = std::sync::mpsc::channel::<WorkerCmd>();
-        let (ack_tx, ack_rx) = std::sync::mpsc::channel::<()>();
-        let cfg = config.clone();
-        let md = model_dir.clone();
-        let uid_copy = uid;
-        worker_handles.push((
-            std::thread::spawn(move || {
-                worker_loop(rank, world, uid_copy, md, cfg, max_seq_len, rx, ack_tx);
-            }),
-            ack_rx,
-        ));
-        worker_txs.push(tx);
-    }
-
-    // Rank 0 setup
-    xserv_cuda::device::set_device(0).unwrap();
-    let tp0 = Arc::new(TpContext::init(0, world, uid, 0));
-    eprintln!("[rank 0] Loading weights...");
-    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
-    eprintln!(
-        "[rank 0] Loaded {} tensors, building model...",
-        weights.len()
-    );
-    let model = GptOss::from_weights_tp(config.clone(), weights, 0, world, 0, Some(tp0));
-    let total_blocks = max_blocks_per_seq + 64;
-    let mut cache = PagedKVCache::new_tp(
-        &config,
-        local_kv,
-        total_blocks,
-        0,
-        4,
-        max_blocks_per_seq,
-        DType::BF16,
-        0,
-    );
-    eprintln!("[rank 0] Ready.");
-
-    // Prompt
-    let prompt_arg = get_arg::<String>(&args, "--prompt");
-    let prompt = prompt_arg
-        .as_deref()
-        .unwrap_or("What is the meaning of life?");
-    let token_ids = tokenizer.encode(prompt);
-    eprintln!("Prompt ({} tokens): {prompt}", token_ids.len());
-
-    // Register sequence
-    let slot = 0;
-    cache.register_sequence(slot).unwrap();
-    broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Register(slot));
-
-    // Teacher-forced diagnostic: prefill (prompt + forced ids) in one shot and
-    // report, for each forced position, whether xserv's argmax == the forced
-    // (oracle) next token. Removes free-running compounding so it isolates
-    // whether per-position logits agree with the llama.cpp trajectory.
-    if let Some(forced) = get_arg::<String>(&args, "--forced") {
-        let forced_ids: Vec<u32> = forced
-            .split(',')
-            .filter_map(|s| s.trim().parse().ok())
-            .collect();
-        let mut seq = token_ids.clone();
-        seq.extend_from_slice(&forced_ids);
-        // Workers must run the same prefill in lockstep (TP AllReduces match up).
-        broadcast_cmd(
-            &worker_txs,
-            &worker_handles,
-            WorkerCmd::Prefill {
-                tokens: seq.clone(),
-                slot,
-            },
-        );
-        let logits = model.forward_prefill_paged(&seq, slot, &mut cache);
-        wait_workers(&worker_handles);
-        let logits_cpu = logits.to_device(Device::Cpu);
-        let vocab = logits.shape()[1];
-        let data = logits_cpu.as_slice::<half::bf16>();
-        let plen = token_ids.len();
-        let mut matches = 0usize;
-        let mut total = 0usize;
-        // position i predicts seq[i+1]; we check the forced region
-        for i in (plen - 1)..(seq.len() - 1) {
-            let row = &data[i * vocab..(i + 1) * vocab];
-            let argmax = row
-                .iter()
-                .enumerate()
-                .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
-                .map(|(j, _)| j as u32)
-                .unwrap();
-            let expected = seq[i + 1];
-            let ok = argmax == expected;
-            if ok {
-                matches += 1;
-            }
-            total += 1;
-            eprintln!(
-                "pos {i}: xserv_argmax={argmax} oracle={expected} {}",
-                if ok { "OK" } else { "DIFF" }
-            );
-        }
-        eprintln!(
-            "\nTeacher-forced top-1 agreement: {matches}/{total} = {:.1}%",
-            100.0 * matches as f64 / total as f64
-        );
-        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
-        for (h, _) in worker_handles {
-            h.join().unwrap();
-        }
-        return;
-    }
-
-    // Teacher-forced DECODE diagnostic: prefill the prompt, then walk the oracle
-    // trajectory through the autoregressive decode path (NOT prefill), recording
-    // per-position top-1 agreement bucketed by position. Localizes long-context
-    // decode degradation (which prefill teacher-forcing cannot see).
-    if let Some(forced) = get_arg::<String>(&args, "--forced-decode") {
-        let forced_ids: Vec<u32> = forced
-            .split(',')
-            .filter_map(|s| s.trim().parse().ok())
-            .collect();
-        broadcast_cmd(
-            &worker_txs,
-            &worker_handles,
-            WorkerCmd::Prefill {
-                tokens: token_ids.clone(),
-                slot,
-            },
-        );
-        let logits = model.forward_prefill_paged(&token_ids, slot, &mut cache);
-        wait_workers(&worker_handles);
-        let mut pred = sample_greedy_last(&logits); // prediction for forced[0]
-        let bucket = 50usize;
-        let mut buckets: Vec<(usize, usize)> = Vec::new();
-        let (mut matches, mut total) = (0usize, 0usize);
-        for (i, &f) in forced_ids.iter().enumerate() {
-            let ok = pred == f;
-            matches += ok as usize;
-            total += 1;
-            let b = i / bucket;
-            if buckets.len() <= b {
-                buckets.push((0, 0));
-            }
-            buckets[b].0 += ok as usize;
-            buckets[b].1 += 1;
-            // Teacher-force: feed the oracle token through the decode path.
-            let pos = cache.seq_len(slot);
-            broadcast_cmd(
-                &worker_txs,
-                &worker_handles,
-                WorkerCmd::Decode {
-                    tokens: vec![f],
-                    positions: vec![pos],
-                    slots: vec![slot],
-                },
-            );
-            let logits = model.forward_decode_paged(&[f], &[pos], &[slot], &mut cache);
-            wait_workers(&worker_handles);
-            pred = sample_greedy_last(&logits);
-        }
-        eprintln!(
-            "Teacher-forced DECODE agreement: {matches}/{total} = {:.1}%",
-            100.0 * matches as f64 / total as f64
-        );
-        for (b, (m, t)) in buckets.iter().enumerate() {
-            eprintln!(
-                "  pos[{:>4}..{:<4}]: {m:>3}/{t:<3} = {:.0}%",
-                b * bucket,
-                b * bucket + t,
-                100.0 * (*m as f64) / (*t as f64)
-            );
-        }
-        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
-        for (h, _) in worker_handles {
-            h.join().unwrap();
-        }
-        return;
-    }
-
-    // Prefill
-    let t0 = Instant::now();
-    broadcast_cmd(
-        &worker_txs,
-        &worker_handles,
-        WorkerCmd::Prefill {
-            tokens: token_ids.clone(),
-            slot,
-        },
-    );
-    let logits = model.forward_prefill_paged(&token_ids, slot, &mut cache);
-    wait_workers(&worker_handles);
-    let ttft = t0.elapsed();
-
-    let mut next = sample_greedy_last(&logits);
-    let mut output_tokens = vec![next];
-
-    eprintln!("TTFT: {:.1}ms", ttft.as_secs_f64() * 1000.0);
-    print!("{prompt}");
-
-    // Decode
-    let mut decoder = GraphedGptOssDecoder::new();
-    let decode_start = Instant::now();
-    for _ in 1..max_tokens {
-        let text = tokenizer.decode(&[next]);
-        print!("{text}");
-
-        if tokenizer.eos_token_id() == Some(next) {
-            break;
-        }
-
-        let pos = cache.seq_len(slot);
-        broadcast_cmd(
-            &worker_txs,
-            &worker_handles,
-            WorkerCmd::Decode {
-                tokens: vec![next],
-                positions: vec![pos],
-                slots: vec![slot],
-            },
-        );
-        let logits = decoder.decode(&model, &[next], &[pos], &[slot], &mut cache);
-        wait_workers(&worker_handles);
-
-        next = sample_greedy_last(&logits);
-        output_tokens.push(next);
-    }
-    let decode_elapsed = decode_start.elapsed();
-    println!();
-
-    let gen_tokens = output_tokens.len();
-    let full_text = tokenizer.decode(&output_tokens);
-    eprintln!("\nGenerated text: {full_text}");
-    eprintln!(
-        "Token IDs: {:?}",
-        &output_tokens[..output_tokens.len().min(20)]
-    );
-    let tpot = if gen_tokens > 1 {
-        decode_elapsed.as_secs_f64() * 1000.0 / (gen_tokens - 1) as f64
-    } else {
-        0.0
-    };
-    let tok_s = if gen_tokens > 1 {
-        (gen_tokens - 1) as f64 / decode_elapsed.as_secs_f64()
-    } else {
-        0.0
-    };
-
-    eprintln!("\n--- Performance ---");
-    eprintln!("Generated: {} tokens", gen_tokens);
-    eprintln!("TTFT: {:.1}ms", ttft.as_secs_f64() * 1000.0);
-    eprintln!("TPOT: {:.1}ms", tpot);
-    eprintln!("Throughput: {:.1} tok/s", tok_s);
-
-    // Cleanup
-    broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
-    for (h, _) in worker_handles {
-        h.join().unwrap();
-    }
-}
-
-// --- Worker infrastructure ---
-
-#[derive(Clone)]
-enum WorkerCmd {
-    Register(usize),
-    Prefill {
-        tokens: Vec<u32>,
-        slot: usize,
-    },
-    Decode {
-        tokens: Vec<u32>,
-        positions: Vec<usize>,
-        slots: Vec<usize>,
-    },
-    Shutdown,
-}
-
-fn worker_loop(
-    rank: usize,
-    world: usize,
-    uid: UniqueId,
-    model_dir: PathBuf,
-    config: ModelConfig,
-    max_seq_len: usize,
-    rx: std::sync::mpsc::Receiver<WorkerCmd>,
-    ack_tx: std::sync::mpsc::Sender<()>,
-) {
-    xserv_cuda::device::set_device(rank as u32).unwrap();
-    let tp = Arc::new(TpContext::init(rank, world, uid, rank as u32));
-    eprintln!("[rank {rank}] Loading weights...");
-    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
-    let model =
-        GptOss::from_weights_tp(config.clone(), weights, rank, world, rank as u32, Some(tp));
-    let local_kv = config.num_kv_heads() / world;
-    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
-    let total_blocks = max_blocks_per_seq + 64;
-    let mut cache = PagedKVCache::new_tp(
-        &config,
-        local_kv,
-        total_blocks,
-        0,
-        4,
-        max_blocks_per_seq,
-        DType::BF16,
-        rank as u32,
-    );
-    eprintln!("[rank {rank}] Ready.");
-    ack_tx.send(()).unwrap();
-
-    let mut decoder = GraphedGptOssDecoder::new();
-    while let Ok(cmd) = rx.recv() {
-        match cmd {
-            WorkerCmd::Register(slot) => {
-                let _ = cache.register_sequence(slot);
-            }
-            WorkerCmd::Prefill { tokens, slot } => {
-                let _ = model.forward_prefill_paged(&tokens, slot, &mut cache);
-            }
-            WorkerCmd::Decode {
-                tokens,
-                positions,
-                slots,
-            } => {
-                let _ = decoder.decode(&model, &tokens, &positions, &slots, &mut cache);
-            }
-            WorkerCmd::Shutdown => break,
-        }
-        ack_tx.send(()).unwrap();
-    }
-}
-
-fn broadcast_cmd(
-    txs: &[std::sync::mpsc::Sender<WorkerCmd>],
-    _handles: &[(std::thread::JoinHandle<()>, std::sync::mpsc::Receiver<()>)],
-    cmd: WorkerCmd,
-) {
-    for tx in txs {
-        tx.send(cmd.clone()).unwrap();
-    }
-}
-
-fn wait_workers(handles: &[(std::thread::JoinHandle<()>, std::sync::mpsc::Receiver<()>)]) {
-    for (_, rx) in handles {
-        rx.recv().unwrap();
-    }
-}
-
-fn sample_greedy_last(logits: &xserv_tensor::Tensor) -> u32 {
-    use half::bf16;
-    assert_eq!(logits.ndim(), 2);
-    // GPU argmax fast path (4-byte D2H instead of the full logits row).
-    if logits.dtype() == xserv_tensor::DType::BF16 && logits.is_contiguous() {
-        let ids = xserv_kernels::argmax_bf16_to_host(logits);
-        return *ids.last().unwrap();
-    }
-    let logits_cpu = logits.to_device(Device::Cpu);
-    let vocab_size = logits.shape()[1];
-    let seq_len = logits.shape()[0];
-    let data = logits_cpu.as_slice::<bf16>();
-    let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
-
-    last.iter()
-        .enumerate()
-        .max_by(|a, b| {
-            let af = a.1.to_f32();
-            let bf = b.1.to_f32();
-            af.partial_cmp(&bf).unwrap_or(std::cmp::Ordering::Equal)
-        })
-        .map(|(i, _)| i as u32)
-        .unwrap()
-}
-
-fn get_arg<T: std::str::FromStr>(args: &[String], flag: &str) -> Option<T> {
-    args.iter()
-        .position(|a| a == flag)
-        .and_then(|i| args.get(i + 1))
-        .and_then(|s| s.parse().ok())
-}

crates/xserv-model/src/bin/bench-gpt2.rs

@@ -1,7 +1,7 @@
 use std::path::PathBuf;
 use std::time::Instant;
-use xserv_model::gpt2::{KVCache, sample_greedy};
-use xserv_model::{GPT2, ModelConfig, loader};
+use xserv_model::gpt2::{sample_greedy, KVCache};
+use xserv_model::{loader, GPT2, ModelConfig};
 use xserv_tensor::Device;
 use xserv_tokenizer::Tokenizer;
 
@@ -104,15 +104,9 @@ fn main() {
 
         let tbt_us = if !token_times_us.is_empty() {
             token_times_us.iter().sum::<u128>() / token_times_us.len() as u128
-        } else {
-            0
-        };
+        } else { 0 };
         let total_gen_us: u128 = ttft_us + token_times_us.iter().sum::<u128>();
-        let tpot_us = if num_generated > 0 {
-            total_gen_us / num_generated as u128
-        } else {
-            0
-        };
+        let tpot_us = if num_generated > 0 { total_gen_us / num_generated as u128 } else { 0 };
 
         let gen_text_escaped = generated_text
             .replace('\\', "\\\\")
@@ -130,16 +124,11 @@ fn main() {
         print!("\"ttft_us\": {ttft_us}, ");
         print!("\"tbt_us\": {tbt_us}, ");
         print!("\"tpot_us\": {tpot_us}}}");
-        if i < prompts.len() - 1 {
-            println!(",");
-        } else {
-            println!();
-        }
+        if i < prompts.len() - 1 { println!(","); } else { println!(); }
 
         eprintln!(
             "[{}/{}] input={input_len}tok gen={num_generated}tok ttft={:.1}ms tbt={:.1}ms | {}",
-            i + 1,
-            prompts.len(),
+            i + 1, prompts.len(),
             ttft_us as f64 / 1000.0,
             tbt_us as f64 / 1000.0,
             &generated_text.replace('\n', " ")[..generated_text.len().min(60)]
@@ -149,18 +138,12 @@ fn main() {
 }
 
 fn generate_with_cache(
-    model: &GPT2,
-    config: &ModelConfig,
-    tokenizer: &Tokenizer,
-    input_ids: &[u32],
-    gen_tokens: usize,
+    model: &GPT2, config: &ModelConfig, tokenizer: &Tokenizer,
+    input_ids: &[u32], gen_tokens: usize,
 ) -> (Vec<u32>, u128, Vec<u128>) {
     let mut cache = KVCache::new(
-        config.num_layers(),
-        config.num_heads(),
-        config.head_dim(),
-        xserv_tensor::DType::F32,
-        Device::Cuda(0),
+        config.num_layers(), config.num_heads(), config.head_dim(),
+        xserv_tensor::DType::F32, Device::Cuda(0),
     );
 
     // Prefill
@@ -180,19 +163,15 @@ fn generate_with_cache(
         let next = sample_greedy(&logits);
         token_times.push(t_start.elapsed().as_micros());
         generated.push(next);
-        if tokenizer.eos_token_id() == Some(next) {
-            break;
-        }
+        if tokenizer.eos_token_id() == Some(next) { break; }
     }
 
     (generated, ttft_us, token_times)
 }
 
 fn generate_no_cache(
-    model: &GPT2,
-    tokenizer: &Tokenizer,
-    input_ids: &[u32],
-    gen_tokens: usize,
+    model: &GPT2, tokenizer: &Tokenizer,
+    input_ids: &[u32], gen_tokens: usize,
 ) -> (Vec<u32>, u128, Vec<u128>) {
     let mut all_ids = input_ids.to_vec();
 
@@ -212,9 +191,7 @@ fn generate_no_cache(
         token_times.push(t_start.elapsed().as_micros());
         all_ids.push(next);
         generated.push(next);
-        if tokenizer.eos_token_id() == Some(next) {
-            break;
-        }
+        if tokenizer.eos_token_id() == Some(next) { break; }
     }
 
     (generated, ttft_us, token_times)

crates/xserv-model/src/bin/bench-qwen3.rs

@@ -1,7 +1,7 @@
 use std::path::PathBuf;
 use std::time::Instant;
 use xserv_model::qwen3::sample_greedy;
-use xserv_model::{DecodeGraphState, GpuKVCache, ModelConfig, Qwen3, loader};
+use xserv_model::{loader, DecodeGraphState, GpuKVCache, ModelConfig, Qwen3};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
@@ -139,35 +139,18 @@ fn main() {
                 } else {
                     // Replay captured graphs
                     let pos = cache.seq_len() as u32;
-                    graph.execute(
-                        last,
-                        pos,
-                        &mut cache,
-                        &layer_ptrs,
-                        embed,
-                        config.vocab_size as i32,
-                        config.hidden() as i32,
-                    );
+                    graph.execute(last, pos, &mut cache, &layer_ptrs, embed, config.vocab_size as i32, config.hidden() as i32);
                     cache.advance_seq_len(1);
                     // Read logits from graph buffer
                     let vocab_size = config.vocab_size;
                     let mut logits_bytes = vec![0u8; vocab_size * 2];
-                    graph
-                        .logits_buffer()
-                        .copy_to_host(&mut logits_bytes)
-                        .unwrap();
+                    graph.logits_buffer().copy_to_host(&mut logits_bytes).unwrap();
                     let logits_data: &[half::bf16] = unsafe {
-                        std::slice::from_raw_parts(
-                            logits_bytes.as_ptr() as *const half::bf16,
-                            vocab_size,
-                        )
+                        std::slice::from_raw_parts(logits_bytes.as_ptr() as *const half::bf16, vocab_size)
                     };
-                    logits_data
-                        .iter()
-                        .enumerate()
+                    logits_data.iter().enumerate()
                         .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
-                        .map(|(idx, _)| idx as u32)
-                        .unwrap()
+                        .map(|(idx, _)| idx as u32).unwrap()
                 }
             } else {
                 let logits = model.forward_gpu_cache(&[last], &mut cache);
@@ -176,24 +159,16 @@ fn main() {
 
             token_times.push(t_start.elapsed().as_micros());
             generated.push(next);
-            if tokenizer.eos_token_id() == Some(next) {
-                break;
-            }
+            if tokenizer.eos_token_id() == Some(next) { break; }
         }
 
         let num_generated = generated.len();
         let generated_text = tokenizer.decode(&generated);
         let tbt_us = if !token_times.is_empty() {
             token_times.iter().sum::<u128>() / token_times.len() as u128
-        } else {
-            0
-        };
+        } else { 0 };
         let total_gen_us: u128 = ttft_us + token_times.iter().sum::<u128>();
-        let tpot_us = if num_generated > 0 {
-            total_gen_us / num_generated as u128
-        } else {
-            0
-        };
+        let tpot_us = if num_generated > 0 { total_gen_us / num_generated as u128 } else { 0 };
 
         let gen_text_escaped = generated_text
             .replace('\\', "\\\\")
@@ -211,18 +186,13 @@ fn main() {
         print!("\"ttft_us\": {ttft_us}, ");
         print!("\"tbt_us\": {tbt_us}, ");
         print!("\"tpot_us\": {tpot_us}}}");
-        if i < prompts.len() - 1 {
-            println!(",");
-        } else {
-            println!();
-        }
+        if i < prompts.len() - 1 { println!(","); } else { println!(); }
 
         let display_text = generated_text.replace('\n', " ");
         let truncated: String = display_text.chars().take(60).collect();
         eprintln!(
             "[{}/{}] input={input_len}tok gen={num_generated}tok ttft={:.1}ms tbt={:.1}ms | {}",
-            i + 1,
-            prompts.len(),
+            i + 1, prompts.len(),
             ttft_us as f64 / 1000.0,
             tbt_us as f64 / 1000.0,
             truncated

crates/xserv-model/src/bin/bench-speculative.rs

@@ -1,976 +0,0 @@
-//! Draft-model speculative decoding benchmark for Qwen3.
-//!
-//! v0 scope:
-//! - target + draft are Qwen3-family models with the same tokenizer/vocab;
-//! - batch=1;
-//! - greedy exact-match acceptance;
-//! - no probabilistic rejection sampling.
-
-use half::bf16;
-use std::path::{Path, PathBuf};
-use std::time::Instant;
-
-use xserv_model::qwen3_graph::GraphedQwen3Decoder;
-use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
-use xserv_tensor::{DType, Device, Tensor};
-use xserv_tokenizer::Tokenizer;
-
-const DEFAULT_GAMMA: usize = 4;
-const DEFAULT_GEN_TOKENS: usize = 64;
-const DEFAULT_MAX_SEQ_LEN: usize = 2048;
-
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-enum VerifyPath {
-    Flash,
-    PagedDecode,
-}
-
-impl VerifyPath {
-    fn as_str(self) -> &'static str {
-        match self {
-            VerifyPath::Flash => "flash",
-            VerifyPath::PagedDecode => "paged-decode",
-        }
-    }
-}
-
-const PROMPTS: [&str; 50] = [
-    "The capital of France is",
-    "Once upon a time in a land far away",
-    "Hello, how are you doing today",
-    "In a shocking finding, scientists discovered a",
-    "The weather today is sunny, so I decided to",
-    "Alan Turing was a British mathematician who",
-    "The best way to learn programming is",
-    "Artificial intelligence will change the world because",
-    "The history of the internet began in the",
-    "A good morning routine starts with",
-    "The stock market crashed because investors",
-    "Deep learning is a subset of machine learning that",
-    "The president of the United States announced",
-    "In the year 2050, humans will",
-    "The secret to happiness is",
-    "When I was a child, I used to",
-    "The most important scientific discovery of the century",
-    "Climate change is caused by",
-    "The recipe for chocolate cake requires",
-    "In conclusion, the evidence suggests that",
-    "The cat sat on the mat and",
-    "According to recent studies, exercise can",
-    "The first step in solving any problem is",
-    "Technology has transformed the way we",
-    "The novel begins with the protagonist",
-    "Education is the most powerful weapon",
-    "The ocean covers more than seventy percent of",
-    "Last night I had a dream about",
-    "The company announced its quarterly earnings",
-    "Music has the power to",
-    "The difference between success and failure is",
-    "In the beginning, there was nothing but",
-    "The doctor told me that I should",
-    "Python is a popular programming language because",
-    "The ancient Romans built roads that",
-    "A balanced diet should include",
-    "The movie received mixed reviews from critics",
-    "Space exploration has led to many",
-    "The teacher asked the students to",
-    "Global warming is one of the most",
-    "The bridge collapsed due to structural",
-    "Quantum computing promises to revolutionize",
-    "The new policy will affect millions of",
-    "During the winter months, it is important to",
-    "The human brain contains approximately",
-    "Democracy depends on the active participation of",
-    "The train arrived at the station exactly",
-    "Researchers at MIT have developed a new",
-    "The smartphone has become an essential part of",
-    "After careful consideration, the committee decided to",
-];
-
-#[derive(Default)]
-struct RunStats {
-    ids: Vec<u32>,
-    total_s: f64,
-    prefill_s: f64,
-    decode_s: f64,
-    target_steps: usize,
-    accepted: usize,
-    proposed: usize,
-    verify_steps: usize,
-    mirror_steps: usize,
-    commit_steps: usize,
-    correction_steps: usize,
-    verify_decode_mismatches: usize,
-}
-
-#[derive(Default)]
-struct Totals {
-    prompts: usize,
-    baseline_generated: usize,
-    spec_generated: usize,
-    baseline_total_s: f64,
-    baseline_prefill_s: f64,
-    baseline_decode_s: f64,
-    spec_total_s: f64,
-    spec_prefill_s: f64,
-    spec_decode_s: f64,
-    spec_target_steps: usize,
-    spec_accepted: usize,
-    spec_proposed: usize,
-    spec_verify_steps: usize,
-    spec_mirror_steps: usize,
-    spec_commit_steps: usize,
-    spec_correction_steps: usize,
-    spec_verify_decode_mismatches: usize,
-    mismatches: usize,
-}
-
-fn main() {
-    let args: Vec<String> = std::env::args().collect();
-    if args.len() < 3 {
-        eprintln!(
-            "Usage: bench-speculative <target-model-dir> <draft-model-dir> \
-             [--gen-tokens N] [--gamma N] [--prompts N] [--max-seq-len N] [--device N] \
-             [--use-verify-logits] [--verify-path flash|paged-decode] [--dump-verify-mismatches]"
-        );
-        std::process::exit(1);
-    }
-
-    let target_dir = PathBuf::from(&args[1]);
-    let draft_dir = PathBuf::from(&args[2]);
-    let gen_tokens = arg_usize(&args, "--gen-tokens", DEFAULT_GEN_TOKENS);
-    let gamma = arg_usize(&args, "--gamma", DEFAULT_GAMMA);
-    let prompt_count = arg_usize(&args, "--prompts", PROMPTS.len()).min(PROMPTS.len());
-    let max_seq_len = arg_usize(&args, "--max-seq-len", DEFAULT_MAX_SEQ_LEN);
-    let device = arg_usize(&args, "--device", 0) as u32;
-    let use_verify_logits = args.iter().any(|a| a == "--use-verify-logits");
-    let verify_path = parse_verify_path(&args, use_verify_logits);
-    let dump_verify_mismatches = args.iter().any(|a| a == "--dump-verify-mismatches");
-
-    assert!(gen_tokens > 0, "--gen-tokens must be > 0");
-    assert!(gamma > 0, "--gamma must be > 0");
-
-    xserv_cuda::device::set_device(device).unwrap();
-    let info = xserv_cuda::device::device_info(device).unwrap();
-    eprintln!(
-        "GPU {device}: {} ({} MB free)",
-        info.name,
-        info.free_memory / 1024 / 1024
-    );
-
-    let target_config = ModelConfig::from_file(&target_dir.join("config.json"));
-    let draft_config = ModelConfig::from_file(&draft_dir.join("config.json"));
-    assert_qwen3(&target_config, "target");
-    assert_qwen3(&draft_config, "draft");
-    assert_eq!(
-        target_config.vocab_size, draft_config.vocab_size,
-        "target and draft vocab_size must match"
-    );
-
-    warn_if_tokenizers_differ(&target_dir, &draft_dir);
-    let tokenizer = Tokenizer::from_file(&target_dir.join("tokenizer.json"));
-    if tokenizer.vocab_size() != target_config.vocab_size {
-        eprintln!(
-            "WARNING: tokenizer decoder len {} differs from config vocab_size {}; continuing because token ids come from the shared tokenizer.json",
-            tokenizer.vocab_size(),
-            target_config.vocab_size
-        );
-    }
-
-    eprintln!(
-        "Loading target Qwen3: layers={} hidden={} heads={}/{} vocab={}",
-        target_config.num_layers(),
-        target_config.hidden(),
-        target_config.num_heads(),
-        target_config.num_kv_heads(),
-        target_config.vocab_size
-    );
-    let target_weights = loader::load_model_dir(&target_dir, Device::Cuda(device));
-    let target = Qwen3::from_weights(target_config.clone(), target_weights);
-    xserv_cuda::allocator::cached_trim();
-
-    eprintln!(
-        "Loading draft Qwen3: layers={} hidden={} heads={}/{} vocab={}",
-        draft_config.num_layers(),
-        draft_config.hidden(),
-        draft_config.num_heads(),
-        draft_config.num_kv_heads(),
-        draft_config.vocab_size
-    );
-    let draft_weights = loader::load_model_dir(&draft_dir, Device::Cuda(device));
-    let draft = Qwen3::from_weights(draft_config.clone(), draft_weights);
-    xserv_cuda::allocator::cached_trim();
-
-    let warm_ids = tokenizer.encode("warmup");
-    let warm_tokens = gen_tokens.min(4);
-    {
-        let mut target_cache = new_cache(&target_config, max_seq_len, device);
-        let _ = run_baseline(
-            &target,
-            &mut target_cache,
-            &tokenizer,
-            &warm_ids,
-            warm_tokens,
-        );
-    }
-    {
-        let mut target_cache = new_cache_with_rows(
-            &target_config,
-            max_seq_len,
-            device,
-            if use_verify_logits { gamma } else { 1 },
-        );
-        let mut target_verify_cache =
-            new_cache_with_rows(&target_config, max_seq_len, device, gamma);
-        let mut draft_cache = new_cache(&draft_config, max_seq_len, device);
-        let mut draft_decoder = GraphedQwen3Decoder::new();
-        let _ = run_speculative(
-            &target,
-            &draft,
-            &mut target_cache,
-            &mut target_verify_cache,
-            &mut draft_cache,
-            &mut draft_decoder,
-            &tokenizer,
-            &warm_ids,
-            warm_tokens,
-            gamma,
-            use_verify_logits,
-            verify_path,
-            dump_verify_mismatches,
-        );
-    }
-    eprintln!(
-        "Warmup done. Running {prompt_count} prompts, gen_tokens={gen_tokens}, gamma={gamma}, acceptance_mode={}, verify_path={}",
-        if use_verify_logits {
-            "verify_logits"
-        } else {
-            "decode"
-        },
-        verify_path.as_str()
-    );
-
-    let mut totals = Totals::default();
-
-    // Persistent per-benchmark caches so the draft CUDA graph (Phase 24) can be
-    // captured once and replayed across every prompt. Freeing and re-registering
-    // slot 0 between prompts keeps block_table_gpu / context_lens_gpu addresses
-    // stable, which is exactly what the graph captured.
-    let mut target_cache = new_cache_with_rows(
-        &target_config,
-        max_seq_len,
-        device,
-        if use_verify_logits { gamma } else { 1 },
-    );
-    let mut target_verify_cache = new_cache_with_rows(&target_config, max_seq_len, device, gamma);
-    let mut draft_cache = new_cache(&draft_config, max_seq_len, device);
-    let mut draft_decoder = GraphedQwen3Decoder::new();
-
-    for (i, prompt) in PROMPTS.iter().take(prompt_count).enumerate() {
-        let ids = tokenizer.encode(prompt);
-        validate_length_budget(&ids, gen_tokens, max_seq_len, prompt);
-        let mut baseline_cache = new_cache(&target_config, max_seq_len, device);
-        let baseline = run_baseline(&target, &mut baseline_cache, &tokenizer, &ids, gen_tokens);
-        drop(baseline_cache);
-
-        let spec = run_speculative(
-            &target,
-            &draft,
-            &mut target_cache,
-            &mut target_verify_cache,
-            &mut draft_cache,
-            &mut draft_decoder,
-            &tokenizer,
-            &ids,
-            gen_tokens,
-            gamma,
-            use_verify_logits,
-            verify_path,
-            dump_verify_mismatches,
-        );
-
-        let matched = baseline.ids == spec.ids;
-        if !matched {
-            totals.mismatches += 1;
-            eprintln!("MISMATCH prompt {i}: {prompt}");
-            eprintln!("  baseline: {:?}", baseline.ids);
-            eprintln!("  spec:     {:?}", spec.ids);
-        }
-
-        println!(
-            "prompt={:02} match={} gen={} accept={}/{} target_steps={} \
-             baseline_e2e_tpot_ms={:.3} spec_e2e_tpot_ms={:.3}",
-            i,
-            matched,
-            spec.ids.len(),
-            spec.accepted,
-            spec.proposed,
-            spec.target_steps,
-            per_token_ms(baseline.total_s, baseline.ids.len()),
-            per_token_ms(spec.total_s, spec.ids.len()),
-        );
-
-        totals.prompts += 1;
-        totals.baseline_generated += baseline.ids.len();
-        totals.spec_generated += spec.ids.len();
-        totals.baseline_total_s += baseline.total_s;
-        totals.baseline_prefill_s += baseline.prefill_s;
-        totals.baseline_decode_s += baseline.decode_s;
-        totals.spec_total_s += spec.total_s;
-        totals.spec_prefill_s += spec.prefill_s;
-        totals.spec_decode_s += spec.decode_s;
-        totals.spec_target_steps += spec.target_steps;
-        totals.spec_accepted += spec.accepted;
-        totals.spec_proposed += spec.proposed;
-        totals.spec_verify_steps += spec.verify_steps;
-        totals.spec_mirror_steps += spec.mirror_steps;
-        totals.spec_commit_steps += spec.commit_steps;
-        totals.spec_correction_steps += spec.correction_steps;
-        totals.spec_verify_decode_mismatches += spec.verify_decode_mismatches;
-    }
-
-    let baseline_decode_tokens = totals.baseline_generated;
-    let spec_decode_tokens = totals.spec_generated;
-    let acceptance = ratio(totals.spec_accepted, totals.spec_proposed);
-    let tokens_per_target_step = ratio(totals.spec_generated, totals.spec_target_steps);
-    let matched =
-        totals.mismatches == 0 && (!use_verify_logits || totals.spec_verify_decode_mismatches == 0);
-
-    println!("--- SUMMARY ---");
-    println!("prompts={} matched={matched}", totals.prompts);
-    println!(
-        "acceptance_mode={}",
-        if use_verify_logits {
-            "verify_logits"
-        } else {
-            "decode"
-        }
-    );
-    println!("verify_path={}", verify_path.as_str());
-    println!(
-        "acceptance_rate={:.4} accepted={} proposed={}",
-        acceptance, totals.spec_accepted, totals.spec_proposed
-    );
-    println!(
-        "tokens_per_target_step={:.4} target_steps={} verify_steps={} mirror_decode_steps={} commit_decode_steps={} correction_steps={}",
-        tokens_per_target_step,
-        totals.spec_target_steps,
-        totals.spec_verify_steps,
-        totals.spec_mirror_steps,
-        totals.spec_commit_steps,
-        totals.spec_correction_steps
-    );
-    println!(
-        "verify_decode_mismatches={}",
-        totals.spec_verify_decode_mismatches
-    );
-    println!(
-        "baseline_e2e_tpot_ms={:.3} baseline_e2e_tok_s={:.3}",
-        per_token_ms(totals.baseline_total_s, totals.baseline_generated),
-        tok_s(totals.baseline_generated, totals.baseline_total_s)
-    );
-    println!(
-        "spec_e2e_tpot_ms={:.3} spec_e2e_tok_s={:.3} speedup_e2e={:.4}",
-        per_token_ms(totals.spec_total_s, totals.spec_generated),
-        tok_s(totals.spec_generated, totals.spec_total_s),
-        speedup(totals.baseline_total_s, totals.spec_total_s)
-    );
-    println!(
-        "baseline_decode_tpot_ms={:.3} baseline_decode_tok_s={:.3}",
-        per_token_ms(totals.baseline_decode_s, baseline_decode_tokens),
-        tok_s(baseline_decode_tokens, totals.baseline_decode_s)
-    );
-    println!(
-        "spec_decode_tpot_ms={:.3} spec_decode_tok_s={:.3} speedup_decode={:.4}",
-        per_token_ms(totals.spec_decode_s, spec_decode_tokens),
-        tok_s(spec_decode_tokens, totals.spec_decode_s),
-        speedup(totals.baseline_decode_s, totals.spec_decode_s)
-    );
-    println!(
-        "decode_token_counts baseline={} spec={}",
-        baseline_decode_tokens, spec_decode_tokens
-    );
-
-    if !matched {
-        std::process::exit(2);
-    }
-}
-
-fn run_baseline(
-    model: &Qwen3,
-    cache: &mut PagedKVCache,
-    tokenizer: &Tokenizer,
-    prompt_ids: &[u32],
-    gen_tokens: usize,
-) -> RunStats {
-    let slot = 0;
-    cache.register_sequence(slot).expect("register target slot");
-
-    let t0 = Instant::now();
-    let prefill_start = Instant::now();
-    let logits = model.forward_prefill_paged(prompt_ids, slot, cache);
-    sync_device();
-    let prefill_s = prefill_start.elapsed().as_secs_f64();
-
-    let mut generated = Vec::with_capacity(gen_tokens);
-    let mut next = last_argmax(&logits);
-    generated.push(next);
-
-    let decode_start = Instant::now();
-    let mut target_steps = 0usize;
-    while generated.len() < gen_tokens && !tokenizer.is_eos(next) {
-        let pos = cache.seq_len(slot);
-        let logits = model.forward_decode_paged(&[next], &[pos], &[slot], cache);
-        target_steps += 1;
-        next = last_argmax(&logits);
-        generated.push(next);
-    }
-    sync_device();
-    let decode_s = decode_start.elapsed().as_secs_f64();
-    sync_device();
-    let total_s = t0.elapsed().as_secs_f64();
-
-    cache.free_sequence(slot);
-    RunStats {
-        ids: generated,
-        total_s,
-        prefill_s,
-        decode_s,
-        target_steps,
-        ..Default::default()
-    }
-}
-
-#[allow(clippy::too_many_arguments)]
-fn run_speculative(
-    target: &Qwen3,
-    draft: &Qwen3,
-    target_cache: &mut PagedKVCache,
-    target_verify_cache: &mut PagedKVCache,
-    draft_cache: &mut PagedKVCache,
-    draft_decoder: &mut GraphedQwen3Decoder,
-    tokenizer: &Tokenizer,
-    prompt_ids: &[u32],
-    gen_tokens: usize,
-    gamma: usize,
-    use_verify_logits: bool,
-    verify_path: VerifyPath,
-    dump_verify_mismatches: bool,
-) -> RunStats {
-    let slot = 0;
-    target_cache
-        .register_sequence(slot)
-        .expect("register target slot");
-    target_verify_cache
-        .register_sequence(slot)
-        .expect("register target verify slot");
-    draft_cache
-        .register_sequence(slot)
-        .expect("register draft slot");
-
-    let t0 = Instant::now();
-    let prefill_start = Instant::now();
-    let target_logits = target.forward_prefill_paged(prompt_ids, slot, target_cache);
-    if !use_verify_logits {
-        let _ = target.forward_prefill_paged(prompt_ids, slot, target_verify_cache);
-    }
-    let draft_logits = draft.forward_prefill_paged(prompt_ids, slot, draft_cache);
-    sync_device();
-    let prefill_s = prefill_start.elapsed().as_secs_f64();
-
-    let mut target_next = last_argmax(&target_logits);
-    let mut draft_next = last_argmax(&draft_logits);
-    let mut generated = Vec::with_capacity(gen_tokens);
-    let mut accepted_total = 0usize;
-    let mut proposed_total = 0usize;
-    let mut verify_steps = 0usize;
-    let mut mirror_steps = 0usize;
-    let mut commit_steps = 0usize;
-    let mut correction_steps = 0usize;
-    let mut verify_decode_mismatches = 0usize;
-
-    let decode_start = Instant::now();
-    while generated.len() < gen_tokens {
-        let remaining = gen_tokens - generated.len();
-        let round_gamma = gamma.min(remaining);
-        let round_start_len = target_cache.seq_len(slot);
-        assert_eq!(
-            round_start_len,
-            draft_cache.seq_len(slot),
-            "target and draft cache lengths diverged"
-        );
-        if !use_verify_logits {
-            assert_eq!(
-                round_start_len,
-                target_verify_cache.seq_len(slot),
-                "target verify cache length diverged"
-            );
-        }
-
-        let mut draft_tokens = Vec::with_capacity(round_gamma);
-        for _ in 0..round_gamma {
-            let token = draft_next;
-            draft_tokens.push(token);
-            if tokenizer.is_eos(token) {
-                break;
-            }
-            let pos = draft_cache.seq_len(slot);
-            let logits = draft_decoder.decode(draft, &[token], &[pos], &[slot], draft_cache);
-            draft_next = last_argmax(&logits);
-        }
-        proposed_total += draft_tokens.len();
-
-        if use_verify_logits {
-            verify_steps += 1;
-            let verify_logits =
-                target.forward_verify_paged_decode_attention(&draft_tokens, slot, target_cache);
-            let verify_argmax = argmax_rows(&verify_logits);
-            assert_eq!(
-                verify_argmax.len(),
-                draft_tokens.len(),
-                "verify logits rows must match draft token count"
-            );
-
-            let mut accepted = 0usize;
-            let mut done = false;
-            while accepted < draft_tokens.len() {
-                let expected = if accepted > 0 {
-                    verify_argmax[accepted - 1]
-                } else {
-                    target_next
-                };
-                if draft_tokens[accepted] != expected {
-                    break;
-                }
-                let token = draft_tokens[accepted];
-                generated.push(token);
-                accepted_total += 1;
-                accepted += 1;
-
-                if generated.len() >= gen_tokens || tokenizer.is_eos(token) {
-                    done = true;
-                    break;
-                }
-            }
-
-            if accepted > 0 {
-                target_next = verify_argmax[accepted - 1];
-            }
-            target_cache
-                .truncate_sequence(slot, round_start_len + accepted)
-                .unwrap();
-
-            if done {
-                draft_cache
-                    .truncate_sequence(slot, target_cache.seq_len(slot))
-                    .unwrap();
-                break;
-            }
-
-            if accepted == draft_tokens.len() {
-                continue;
-            }
-
-            let correction = if accepted > 0 {
-                verify_argmax[accepted - 1]
-            } else {
-                target_next
-            };
-            generated.push(correction);
-
-            draft_cache
-                .truncate_sequence(slot, round_start_len)
-                .unwrap();
-            replay_draft_tokens(
-                draft,
-                draft_decoder,
-                draft_cache,
-                slot,
-                &draft_tokens[..accepted],
-                &mut draft_next,
-            );
-
-            if generated.len() >= gen_tokens || tokenizer.is_eos(correction) {
-                break;
-            }
-
-            let pos = target_cache.seq_len(slot);
-            let logits = target.forward_decode_paged(&[correction], &[pos], &[slot], target_cache);
-            target_next = last_argmax(&logits);
-            commit_steps += 1;
-
-            let pos = draft_cache.seq_len(slot);
-            let logits = draft_decoder.decode(draft, &[correction], &[pos], &[slot], draft_cache);
-            draft_next = last_argmax(&logits);
-            correction_steps += 1;
-            continue;
-        }
-
-        verify_steps += 1;
-        let verify_logits = match verify_path {
-            VerifyPath::Flash => {
-                target.forward_prefill_paged(&draft_tokens, slot, target_verify_cache)
-            }
-            VerifyPath::PagedDecode => target.forward_verify_paged_decode_attention(
-                &draft_tokens,
-                slot,
-                target_verify_cache,
-            ),
-        };
-        let verify_argmax = argmax_rows(&verify_logits);
-        assert_eq!(
-            verify_argmax.len(),
-            draft_tokens.len(),
-            "verify logits rows must match draft token count"
-        );
-
-        target_verify_cache
-            .truncate_sequence(slot, round_start_len)
-            .unwrap();
-
-        let mut accepted = 0usize;
-        let mut done = false;
-        while accepted < draft_tokens.len() {
-            let expected = if use_verify_logits && accepted > 0 {
-                verify_argmax[accepted - 1]
-            } else {
-                target_next
-            };
-            if draft_tokens[accepted] != expected {
-                break;
-            }
-            let token_idx = accepted;
-            let token = draft_tokens[token_idx];
-            generated.push(token);
-            accepted_total += 1;
-            accepted += 1;
-
-            if generated.len() >= gen_tokens || tokenizer.is_eos(token) {
-                done = true;
-                break;
-            }
-
-            let pos = target_cache.seq_len(slot);
-            let logits = target.forward_decode_paged(&[token], &[pos], &[slot], target_cache);
-            let decode_next = last_argmax(&logits);
-            if verify_argmax[token_idx] != decode_next {
-                verify_decode_mismatches += 1;
-                eprintln!(
-                    "VERIFY/DECODE MISMATCH at cache_len={} accepted_idx={}: verify={} decode={}",
-                    target_cache.seq_len(slot),
-                    token_idx,
-                    verify_argmax[token_idx],
-                    decode_next
-                );
-                if dump_verify_mismatches {
-                    eprintln!(
-                        "  verify_top5={} decode_top5={}",
-                        format_topk(&verify_logits, token_idx, 5),
-                        format_topk(&logits, 0, 5)
-                    );
-                }
-            }
-            target_next = decode_next;
-            commit_steps += 1;
-
-            advance_target_cache(target, target_verify_cache, slot, token);
-            mirror_steps += 1;
-        }
-        if done {
-            draft_cache
-                .truncate_sequence(slot, target_cache.seq_len(slot))
-                .unwrap();
-            target_verify_cache
-                .truncate_sequence(slot, target_cache.seq_len(slot))
-                .unwrap();
-            break;
-        }
-
-        if accepted == draft_tokens.len() {
-            continue;
-        }
-
-        let correction = if use_verify_logits && accepted > 0 {
-            verify_argmax[accepted - 1]
-        } else {
-            target_next
-        };
-        generated.push(correction);
-
-        draft_cache
-            .truncate_sequence(slot, round_start_len)
-            .unwrap();
-        replay_draft_tokens(
-            draft,
-            draft_decoder,
-            draft_cache,
-            slot,
-            &draft_tokens[..accepted],
-            &mut draft_next,
-        );
-
-        if generated.len() >= gen_tokens || tokenizer.is_eos(correction) {
-            break;
-        }
-
-        let pos = target_cache.seq_len(slot);
-        let logits = target.forward_decode_paged(&[correction], &[pos], &[slot], target_cache);
-        target_next = last_argmax(&logits);
-        commit_steps += 1;
-
-        advance_target_cache(target, target_verify_cache, slot, correction);
-        mirror_steps += 1;
-
-        let pos = draft_cache.seq_len(slot);
-        let logits = draft_decoder.decode(draft, &[correction], &[pos], &[slot], draft_cache);
-        draft_next = last_argmax(&logits);
-        correction_steps += 1;
-    }
-    sync_device();
-    let decode_s = decode_start.elapsed().as_secs_f64();
-    sync_device();
-    let total_s = t0.elapsed().as_secs_f64();
-
-    target_cache.free_sequence(slot);
-    target_verify_cache.free_sequence(slot);
-    draft_cache.free_sequence(slot);
-
-    RunStats {
-        ids: generated,
-        total_s,
-        prefill_s,
-        decode_s,
-        target_steps: verify_steps + mirror_steps + commit_steps + correction_steps,
-        accepted: accepted_total,
-        proposed: proposed_total,
-        verify_steps,
-        mirror_steps,
-        commit_steps,
-        correction_steps,
-        verify_decode_mismatches,
-    }
-}
-
-fn advance_target_cache(target: &Qwen3, cache: &mut PagedKVCache, slot: usize, token: u32) {
-    let pos = cache.seq_len(slot);
-    let _ = target.forward_decode_paged(&[token], &[pos], &[slot], cache);
-}
-
-fn replay_draft_tokens(
-    draft: &Qwen3,
-    draft_decoder: &mut GraphedQwen3Decoder,
-    cache: &mut PagedKVCache,
-    slot: usize,
-    tokens: &[u32],
-    next: &mut u32,
-) {
-    for &token in tokens {
-        let pos = cache.seq_len(slot);
-        let logits = draft_decoder.decode(draft, &[token], &[pos], &[slot], cache);
-        *next = last_argmax(&logits);
-    }
-}
-
-fn new_cache(config: &ModelConfig, max_seq_len: usize, device: u32) -> PagedKVCache {
-    new_cache_with_rows(config, max_seq_len, device, 1)
-}
-
-fn new_cache_with_rows(
-    config: &ModelConfig,
-    max_seq_len: usize,
-    device: u32,
-    max_rows: usize,
-) -> PagedKVCache {
-    let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
-    let total_blocks = max_blocks_per_seq + 8;
-    PagedKVCache::new(
-        config,
-        total_blocks,
-        0,
-        max_rows.max(1),
-        max_blocks_per_seq,
-        DType::BF16,
-        device,
-    )
-}
-
-fn argmax_rows(logits: &Tensor) -> Vec<u32> {
-    assert_eq!(logits.ndim(), 2);
-    if logits.dtype() == DType::BF16
-        && matches!(logits.device(), Device::Cuda(_))
-        && logits.is_contiguous()
-    {
-        return xserv_kernels::argmax_bf16_to_host(logits);
-    }
-
-    let vocab_size = logits.shape()[1];
-    let rows = logits.shape()[0];
-    let logits_cpu = logits.to_device(Device::Cpu);
-    match logits.dtype() {
-        DType::F32 => logits_cpu
-            .as_slice::<f32>()
-            .chunks_exact(vocab_size)
-            .take(rows)
-            .map(argmax_f32)
-            .collect(),
-        DType::BF16 => logits_cpu
-            .as_slice::<bf16>()
-            .chunks_exact(vocab_size)
-            .take(rows)
-            .map(|row| {
-                row.iter()
-                    .enumerate()
-                    .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
-                    .map(|(i, _)| i as u32)
-                    .unwrap()
-            })
-            .collect(),
-        _ => panic!("unsupported dtype for argmax: {:?}", logits.dtype()),
-    }
-}
-
-fn last_argmax(logits: &Tensor) -> u32 {
-    *argmax_rows(logits).last().unwrap()
-}
-
-fn argmax_f32(row: &[f32]) -> u32 {
-    row.iter()
-        .enumerate()
-        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
-        .map(|(i, _)| i as u32)
-        .unwrap()
-}
-
-fn format_topk(logits: &Tensor, row: usize, k: usize) -> String {
-    let vals = topk_row(logits, row, k);
-    vals.iter()
-        .map(|(id, val)| format!("{id}:{val:.3}"))
-        .collect::<Vec<_>>()
-        .join(",")
-}
-
-fn topk_row(logits: &Tensor, row: usize, k: usize) -> Vec<(u32, f32)> {
-    assert_eq!(logits.ndim(), 2);
-    let vocab_size = logits.shape()[1];
-    assert!(row < logits.shape()[0], "topk row out of bounds");
-    let logits_cpu = logits.to_device(Device::Cpu);
-    let mut vals: Vec<(u32, f32)> = match logits.dtype() {
-        DType::F32 => logits_cpu.as_slice::<f32>()[row * vocab_size..(row + 1) * vocab_size]
-            .iter()
-            .enumerate()
-            .map(|(i, &v)| (i as u32, v))
-            .collect(),
-        DType::BF16 => logits_cpu.as_slice::<bf16>()[row * vocab_size..(row + 1) * vocab_size]
-            .iter()
-            .enumerate()
-            .map(|(i, &v)| (i as u32, v.to_f32()))
-            .collect(),
-        _ => panic!("unsupported dtype for topk: {:?}", logits.dtype()),
-    };
-    vals.select_nth_unstable_by(k, |a, b| b.1.partial_cmp(&a.1).unwrap());
-    vals.truncate(k);
-    vals.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
-    vals
-}
-
-fn assert_qwen3(config: &ModelConfig, name: &str) {
-    let model_type = config.model_type.as_deref().unwrap_or("unknown");
-    assert!(
-        model_type.contains("qwen"),
-        "{name} model_type must be qwen-like, got {model_type}"
-    );
-}
-
-fn warn_if_tokenizers_differ(target_dir: &Path, draft_dir: &Path) {
-    let target = std::fs::read(target_dir.join("tokenizer.json"));
-    let draft = std::fs::read(draft_dir.join("tokenizer.json"));
-    if let (Ok(target), Ok(draft)) = (target, draft) {
-        if target != draft {
-            eprintln!(
-                "WARNING: target and draft tokenizer.json differ; v0 assumes identical token ids"
-            );
-        }
-    }
-}
-
-fn arg_usize(args: &[String], flag: &str, default: usize) -> usize {
-    args.iter()
-        .position(|a| a == flag)
-        .and_then(|i| args.get(i + 1))
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(default)
-}
-
-fn parse_verify_path(args: &[String], use_verify_logits: bool) -> VerifyPath {
-    let default = if use_verify_logits {
-        VerifyPath::PagedDecode
-    } else {
-        VerifyPath::Flash
-    };
-    let Some(value) = args
-        .iter()
-        .position(|a| a == "--verify-path")
-        .and_then(|i| args.get(i + 1))
-    else {
-        return default;
-    };
-    match value.as_str() {
-        "flash" => VerifyPath::Flash,
-        "paged-decode" => VerifyPath::PagedDecode,
-        _ => {
-            eprintln!("unknown --verify-path {value:?}; expected flash or paged-decode");
-            std::process::exit(1);
-        }
-    }
-}
-
-fn validate_length_budget(prompt_ids: &[u32], gen_tokens: usize, max_seq_len: usize, prompt: &str) {
-    let required = prompt_ids.len() + gen_tokens;
-    if required > max_seq_len {
-        eprintln!(
-            "prompt requires prompt_len({}) + gen_tokens({}) = {} tokens, exceeding --max-seq-len {}: {:?}",
-            prompt_ids.len(),
-            gen_tokens,
-            required,
-            max_seq_len,
-            prompt
-        );
-        std::process::exit(2);
-    }
-}
-
-fn sync_device() {
-    xserv_cuda::device::synchronize().expect("cuda device synchronize");
-}
-
-fn ratio(num: usize, den: usize) -> f64 {
-    if den == 0 {
-        0.0
-    } else {
-        num as f64 / den as f64
-    }
-}
-
-fn speedup(baseline_s: f64, spec_s: f64) -> f64 {
-    if spec_s == 0.0 {
-        0.0
-    } else {
-        baseline_s / spec_s
-    }
-}
-
-fn tok_s(tokens: usize, seconds: f64) -> f64 {
-    if seconds == 0.0 {
-        0.0
-    } else {
-        tokens as f64 / seconds
-    }
-}
-
-fn per_token_ms(seconds: f64, tokens: usize) -> f64 {
-    if tokens == 0 {
-        0.0
-    } else {
-        seconds * 1000.0 / tokens as f64
-    }
-}

crates/xserv-model/src/bin/bench-tp.rs

@@ -18,7 +18,7 @@ use std::thread;
 use std::time::Instant;
 
 use xserv_model::qwen3::sample_greedy;
-use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
+use xserv_model::{loader, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
@@ -35,13 +35,8 @@ fn main() {
         std::process::exit(1);
     }
     let model_dir = PathBuf::from(&args[1]);
-    let world: usize = arg(&args, "--tp")
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(1)
-        .max(1);
-    let gen_tokens: usize = arg(&args, "--gen-tokens")
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(64);
+    let world: usize = arg(&args, "--tp").and_then(|s| s.parse().ok()).unwrap_or(1).max(1);
+    let gen_tokens: usize = arg(&args, "--gen-tokens").and_then(|s| s.parse().ok()).unwrap_or(64);
     let devices: Vec<u32> = match arg(&args, "--devices") {
         Some(s) => s.split(',').filter_map(|d| d.trim().parse().ok()).collect(),
         None => (0..world as u32).collect(),
@@ -72,11 +67,7 @@ fn main() {
     // Tensors are not Send (their Storage holds a raw GPU pointer), so each rank
     // thread loads its own CPU copy of the weights and shards in-thread. Loading
     // is not part of the timed region.
-    let id = if world > 1 {
-        Some(xserv_distributed::get_unique_id())
-    } else {
-        None
-    };
+    let id = if world > 1 { Some(xserv_distributed::get_unique_id()) } else { None };
 
     let handles: Vec<_> = (0..world)
         .map(|rank| {
@@ -85,9 +76,7 @@ fn main() {
             let prompt_ids = prompt_ids.clone();
             let device = devices[rank];
             thread::spawn(move || {
-                run_rank(
-                    rank, world, device, id, config, model_dir, prompt_ids, gen_tokens, eos,
-                )
+                run_rank(rank, world, device, id, config, model_dir, prompt_ids, gen_tokens, eos)
             })
         })
         .collect();
@@ -102,10 +91,7 @@ fn main() {
 
     let results = rank0.expect("rank 0 produced no results");
     println!("\n=== TP={world} (devices {devices:?}) — Qwen3 E2E benchmark ===");
-    println!(
-        "{:<45} {:>10} {:>12} {:>8}",
-        "prompt", "TTFT(ms)", "decode tok/s", "gen"
-    );
+    println!("{:<45} {:>10} {:>12} {:>8}", "prompt", "TTFT(ms)", "decode tok/s", "gen");
     let mut tps_sum = 0.0;
     for (i, r) in results.iter().enumerate() {
         let text = tokenizer.decode(&r.gen_ids).replace('\n', " ");
@@ -113,29 +99,16 @@ fn main() {
         let p: String = prompts[i].chars().take(43).collect();
         println!(
             "{:<45} {:>10.1} {:>12.1} {:>8} | {}",
-            p,
-            r.ttft_ms,
-            r.decode_tok_s,
-            r.gen_ids.len(),
-            short
+            p, r.ttft_ms, r.decode_tok_s, r.gen_ids.len(), short
         );
         tps_sum += r.decode_tok_s;
     }
-    println!(
-        "--- mean decode throughput: {:.1} tok/s ---",
-        tps_sum / results.len() as f64
-    );
+    println!("--- mean decode throughput: {:.1} tok/s ---", tps_sum / results.len() as f64);
 
     // Machine-readable line for cross-TP correctness diffing (rank 0 token ids).
     let all_ids: Vec<String> = results
         .iter()
-        .map(|r| {
-            r.gen_ids
-                .iter()
-                .map(|x| x.to_string())
-                .collect::<Vec<_>>()
-                .join(",")
-        })
+        .map(|r| r.gen_ids.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(","))
         .collect();
     println!("CORRECTNESS_IDS tp={world} {}", all_ids.join(" | "));
 }
@@ -153,12 +126,7 @@ fn run_rank(
 ) -> Option<Vec<PromptResult>> {
     // Bind this thread to its GPU and set up the TP communicator.
     let tp = if world > 1 {
-        Some(Arc::new(xserv_distributed::TpContext::init(
-            rank,
-            world,
-            id.unwrap(),
-            device,
-        )))
+        Some(Arc::new(xserv_distributed::TpContext::init(rank, world, id.unwrap(), device)))
     } else {
         xserv_cuda::device::set_device(device).unwrap();
         None
@@ -174,14 +142,7 @@ fn run_rank(
     let max_blocks_per_seq = max_seq.div_ceil(BLOCK_SIZE);
     let total_blocks = max_blocks_per_seq + 8;
     let mut cache = PagedKVCache::new_tp(
-        &config,
-        local_kv,
-        total_blocks,
-        0,
-        1,
-        max_blocks_per_seq,
-        DType::BF16,
-        device,
+        &config, local_kv, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, device,
     );
 
     // Warmup (init kernels / allocator / NCCL channels) — not timed.
@@ -216,20 +177,12 @@ fn run_rank(
             steps += 1;
         }
         let decode_s = t1.elapsed().as_secs_f64();
-        let decode_tok_s = if steps > 0 && decode_s > 0.0 {
-            steps as f64 / decode_s
-        } else {
-            0.0
-        };
+        let decode_tok_s = if steps > 0 && decode_s > 0.0 { steps as f64 / decode_s } else { 0.0 };
 
         cache.free_sequence(0);
 
         if rank == 0 {
-            out.push(PromptResult {
-                gen_ids: generated,
-                ttft_ms,
-                decode_tok_s,
-            });
+            out.push(PromptResult { gen_ids: generated, ttft_ms, decode_tok_s });
         }
     }
 
@@ -237,8 +190,5 @@ fn run_rank(
 }
 
 fn arg<'a>(args: &'a [String], flag: &str) -> Option<&'a str> {
-    args.iter()
-        .position(|a| a == flag)
-        .and_then(|i| args.get(i + 1))
-        .map(|s| s.as_str())
+    args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).map(|s| s.as_str())
 }

crates/xserv-model/src/bin/bench-verify-cost.rs

@@ -1,134 +0,0 @@
-//! Micro-benchmark: measure the cost of forward_verify_paged_decode_attention
-//! at different batch sizes (γ+1 values), to understand where speedup comes
-//! from (or doesn't).
-
-use std::path::PathBuf;
-use std::time::Instant;
-
-use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
-use xserv_tensor::{DType, Device};
-use xserv_tokenizer::Tokenizer;
-
-fn main() {
-    let args: Vec<String> = std::env::args().collect();
-    if args.len() < 2 {
-        eprintln!(
-            "Usage: bench-verify-cost <target-dir> [--prompt-len N] [--iters N] [--device N]"
-        );
-        std::process::exit(1);
-    }
-    let target_dir = PathBuf::from(&args[1]);
-    let prompt_len = arg_usize(&args, "--prompt-len", 100);
-    let iters = arg_usize(&args, "--iters", 30);
-    let device = arg_usize(&args, "--device", 0) as u32;
-
-    xserv_cuda::device::set_device(device).unwrap();
-
-    let cfg = ModelConfig::from_file(&target_dir.join("config.json"));
-    eprintln!("Loading target...");
-    let weights = loader::load_model_dir(&target_dir, Device::Cuda(device));
-    let target = Qwen3::from_weights(cfg.clone(), weights);
-    xserv_cuda::allocator::cached_trim();
-
-    let tok = Tokenizer::from_file(&target_dir.join("tokenizer.json"));
-    let ids = tok.encode(&"the ".repeat(prompt_len))[..prompt_len].to_vec();
-
-    let max_seq_len = 2048;
-    let num_blocks = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE + 4;
-    let mut cache = PagedKVCache::new(&cfg, num_blocks, 0, 16, num_blocks, DType::BF16, device);
-    cache.register_sequence(0).unwrap();
-
-    // Prefill
-    let _ = target.forward_prefill_paged(&ids, 0, &mut cache);
-    sync();
-
-    // Warmup one of each
-    for &n in &[1, 2, 3, 5, 9] {
-        let toks: Vec<u32> = (0..n).map(|_| ids[0]).collect();
-        let _ = target.forward_decode_paged(
-            &toks,
-            &(0..n).map(|i| ids.len() + i).collect::<Vec<_>>(),
-            &vec![0; n],
-            &mut cache,
-        );
-        cache.truncate_sequence(0, ids.len()).unwrap();
-    }
-    sync();
-
-    // Benchmark single-token decode
-    let mut t = 0.0f64;
-    for i in 0..iters {
-        cache.truncate_sequence(0, ids.len()).unwrap();
-        let t0 = Instant::now();
-        let _ = target.forward_decode_paged(&[ids[0]], &[ids.len()], &[0], &mut cache);
-        sync();
-        t += t0.elapsed().as_secs_f64();
-        let _ = i;
-    }
-    let single = t * 1000.0 / iters as f64;
-    println!(
-        "single-token decode: {:.3} ms (mean of {} iters)",
-        single, iters
-    );
-
-    // Benchmark forward_verify_paged_decode_attention at various batch sizes
-    // (batched-GEMV path).
-    for &n in &[1usize, 2, 3, 5, 9] {
-        let toks: Vec<u32> = (0..n).map(|_| ids[0]).collect();
-        let mut t = 0.0f64;
-        for _ in 0..iters {
-            cache.truncate_sequence(0, ids.len()).unwrap();
-            let t0 = Instant::now();
-            let _ = target.forward_verify_paged_decode_attention(&toks, 0, &mut cache);
-            sync();
-            t += t0.elapsed().as_secs_f64();
-        }
-        let ms = t * 1000.0 / iters as f64;
-        println!(
-            "verify (batched-GEMV) batch={}: {:.3} ms ({:.2}× single)",
-            n,
-            ms,
-            ms / single
-        );
-    }
-
-    // Benchmark _with_hidden variant which uses cuBLAS GEMM after Phase 26 fast-verify.
-    let hooks_layers = [2usize, 18, 33];
-    for &n in &[1usize, 2, 3, 5, 9] {
-        let toks: Vec<u32> = (0..n).map(|_| ids[0]).collect();
-        let mut t = 0.0f64;
-        for _ in 0..iters {
-            cache.truncate_sequence(0, ids.len()).unwrap();
-            let t0 = Instant::now();
-            let _ = target.forward_verify_paged_decode_attention_with_hidden(
-                &toks,
-                0,
-                &mut cache,
-                &hooks_layers,
-            );
-            sync();
-            t += t0.elapsed().as_secs_f64();
-        }
-        let ms = t * 1000.0 / iters as f64;
-        println!(
-            "verify (cuBLAS GEMM) batch={}: {:.3} ms ({:.2}× single)",
-            n,
-            ms,
-            ms / single
-        );
-    }
-
-    cache.free_sequence(0);
-}
-
-fn sync() {
-    xserv_cuda::device::synchronize().unwrap();
-}
-
-fn arg_usize(args: &[String], flag: &str, default: usize) -> usize {
-    args.iter()
-        .position(|a| a == flag)
-        .and_then(|i| args.get(i + 1))
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(default)
-}

crates/xserv-model/src/bin/check-eagle3.rs

@@ -1,174 +0,0 @@
-//! EAGLE3 sanity check: load weights, run one draft step, print top-5 predictions.
-//!
-//! This verifies that:
-//! - Eagle3Head weights load without shape mismatches
-//! - Target hidden states can be captured via decode_core_with_hidden
-//! - Eagle3Head::step produces a valid token id (in target vocab)
-//!
-//! Does NOT measure speedup — that requires a full γ≥2 speculative loop, which
-//! is more complex integration work.
-
-use std::path::PathBuf;
-
-use xserv_model::eagle3::{EAGLE_HOOK_LAYERS, Eagle3Head};
-use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
-use xserv_tensor::{DType, Device, Tensor};
-use xserv_tokenizer::Tokenizer;
-
-fn main() {
-    let args: Vec<String> = std::env::args().collect();
-    if args.len() < 3 {
-        eprintln!("Usage: check-eagle3 <target-model-dir> <eagle3-model-dir> [prompt]");
-        std::process::exit(1);
-    }
-    let target_dir = PathBuf::from(&args[1]);
-    let eagle_dir = PathBuf::from(&args[2]);
-    let prompt = args
-        .get(3)
-        .cloned()
-        .unwrap_or_else(|| "The capital of France is".to_string());
-    let device: u32 = 0;
-
-    xserv_cuda::device::set_device(device).unwrap();
-
-    let target_config = ModelConfig::from_file(&target_dir.join("config.json"));
-    eprintln!("Loading target Qwen3-8B...");
-    let target_weights = loader::load_model_dir(&target_dir, Device::Cuda(device));
-    let target = Qwen3::from_weights(target_config.clone(), target_weights);
-    xserv_cuda::allocator::cached_trim();
-
-    eprintln!("Loading EAGLE3 head from {}", eagle_dir.display());
-    let mut eagle = Eagle3Head::load(&eagle_dir, device);
-    xserv_cuda::allocator::cached_trim();
-
-    let tokenizer = Tokenizer::from_file(&target_dir.join("tokenizer.json"));
-    let embed_tokens = target.embed_tokens_tensor();
-
-    let ids = tokenizer.encode(&prompt);
-    let max_seq_len = 512;
-
-    let num_blocks = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE + 2;
-    let mut cache = PagedKVCache::new(
-        &target_config,
-        num_blocks,
-        0,
-        1,
-        num_blocks,
-        DType::BF16,
-        device,
-    );
-    cache.register_sequence(0).unwrap();
-
-    // Prefill target.
-    let logits = target.forward_prefill_paged(&ids, 0, &mut cache);
-    let target_first = *xserv_kernels::argmax_bf16_to_host(&logits).last().unwrap();
-    let target_first_text = tokenizer.decode(&[target_first]);
-    println!("Prompt: {:?}", prompt);
-    println!(
-        "Target argmax after prefill: {} ({:?})",
-        target_first, target_first_text
-    );
-
-    // Now run one target decode step with target_first to get hidden states at the
-    // hook layers.
-    let pos = cache.seq_len(0);
-    target.decode_prepare(&[pos], &[0], &mut cache);
-    let ids_gpu = upload_u32(&[target_first]);
-    let pos_gpu = upload_u32(&[pos as u32]);
-    let (target_next_logits, hooks) = target.decode_core_with_hidden(
-        ids_gpu.as_ptr() as *const std::ffi::c_void,
-        pos_gpu.as_ptr() as *const std::ffi::c_void,
-        1,
-        &[0],
-        &mut cache,
-        &EAGLE_HOOK_LAYERS,
-    );
-    let target_next = xserv_kernels::argmax_bf16_single(&target_next_logits);
-    let target_next_text = tokenizer.decode(&[target_next]);
-    println!(
-        "Target argmax after 1 decode step: {} ({:?})",
-        target_next, target_next_text
-    );
-
-    for (i, h) in hooks.iter().enumerate() {
-        println!(
-            "hook[{}] (layer {}): shape={:?} dtype={:?}",
-            i,
-            EAGLE_HOOK_LAYERS[i],
-            h.shape(),
-            h.dtype()
-        );
-    }
-
-    // Ask EAGLE what it thinks the NEXT token is (given target_first as prev_token
-    // and the hidden states from the position where target_first lives).
-    // EAGLE should predict target_next (or close to it) to be useful.
-    eagle.reset();
-    let (eagle_pred, eagle_logits) = eagle.step(&hooks, embed_tokens, target_first, pos);
-    let eagle_pred_text = tokenizer.decode(&[eagle_pred]);
-    println!(
-        "EAGLE draft prediction (pairing A: prev=target_first): {} ({:?})",
-        eagle_pred, eagle_pred_text
-    );
-
-    if eagle_pred == target_next {
-        println!("MATCH: EAGLE agrees with target on next token.");
-    } else {
-        println!(
-            "MISMATCH: EAGLE draft={} vs target={} (this is fine per-step; check top-5 below)",
-            eagle_pred, target_next
-        );
-    }
-
-    // Show top-5 from eagle logits (in draft vocab space, mapped to target).
-    print_top5(
-        &eagle_logits,
-        "EAGLE draft top-5 (pairing A)",
-        &eagle,
-        &tokenizer,
-    );
-
-    // Alternative pairing B: pair hooks with target_next (the token those hooks produced
-    // via lm_head), predict token after target_next. Position advances by 1.
-    eagle.reset();
-    let (eagle_pred_b, eagle_logits_b) = eagle.step(&hooks, embed_tokens, target_next, pos + 1);
-    let eagle_pred_b_text = tokenizer.decode(&[eagle_pred_b]);
-    println!(
-        "\nEAGLE draft prediction (pairing B: prev=target_next): {} ({:?})",
-        eagle_pred_b, eagle_pred_b_text
-    );
-    print_top5(
-        &eagle_logits_b,
-        "EAGLE draft top-5 (pairing B)",
-        &eagle,
-        &tokenizer,
-    );
-}
-
-fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
-    let bytes = unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) };
-    let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).unwrap();
-    buf.copy_from_host(bytes).unwrap();
-    buf
-}
-
-fn print_top5(logits: &Tensor, label: &str, eagle: &Eagle3Head, tokenizer: &Tokenizer) {
-    use half::bf16;
-    let cpu = logits.to_device(Device::Cpu);
-    let data = cpu.as_slice::<bf16>();
-    let mut vals: Vec<(usize, f32)> = data
-        .iter()
-        .enumerate()
-        .map(|(i, v)| (i, v.to_f32()))
-        .collect();
-    vals.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
-    println!("{label}:");
-    for (i, val) in vals.iter().take(5) {
-        let target_id = eagle.map_draft_to_target(*i as u32);
-        let text = tokenizer.decode(&[target_id]);
-        println!(
-            "  draft_id={} target_id={} val={:.3} text={:?}",
-            i, target_id, val, text
-        );
-    }
-}

crates/xserv-model/src/bin/dump-logits.rs

@@ -1,8 +1,8 @@
-use half::bf16;
 use std::path::PathBuf;
-use xserv_model::{KVCache, ModelConfig, Qwen3, loader};
+use xserv_model::{loader, KVCache, ModelConfig, Qwen3};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
+use half::bf16;
 
 fn main() {
     let args: Vec<String> = std::env::args().collect();
@@ -20,11 +20,8 @@ fn main() {
     eprintln!("Token IDs: {token_ids:?}");
 
     let mut cache = KVCache::new(
-        config.num_layers(),
-        config.num_kv_heads(),
-        config.head_dim(),
-        DType::BF16,
-        Device::Cuda(0),
+        config.num_layers(), config.num_kv_heads(), config.head_dim(),
+        DType::BF16, Device::Cuda(0),
     );
     let logits = model.forward_with_cache(&token_ids, &mut cache);
     let logits_cpu = logits.to_device(Device::Cpu);
@@ -34,9 +31,7 @@ fn main() {
 
     // Print top-20 logits for the last position
     let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
-    let mut indexed: Vec<(usize, f32)> = last_row
-        .iter()
-        .enumerate()
+    let mut indexed: Vec<(usize, f32)> = last_row.iter().enumerate()
         .map(|(i, v)| (i, v.to_f32()))
         .collect();
     indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

crates/xserv-model/src/bin/gptoss-logits.rs

@@ -0,0 +1,35 @@
+//! Dump gpt-oss next-token logits for a fixed token-id sequence, to compare
+//! against the llama.cpp oracle (isolates the model forward from tokenizer
+//! differences). Usage: gptoss-logits <bf16-model-dir> <tok0> <tok1> ...
+use std::path::PathBuf;
+use half::bf16;
+use xserv_model::loader;
+use xserv_model::{GptOss, ModelConfig};
+use xserv_tensor::Device;
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    let model_dir = PathBuf::from(&args[1]);
+    let tokens: Vec<u32> = args[2..].iter().map(|s| s.parse().expect("token id")).collect();
+    assert!(!tokens.is_empty(), "need at least one token id");
+
+    let config = ModelConfig::from_file(&model_dir.join("config.json"));
+    eprintln!("[gptoss-logits] loading {} ...", model_dir.display());
+    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
+    let model = GptOss::from_weights(config, weights);
+    eprintln!("[gptoss-logits] forward over {} tokens", tokens.len());
+
+    let logits = model.forward(&tokens); // [T, vocab]
+    let vocab = logits.shape()[1];
+    let t = logits.shape()[0];
+    let host = logits.to_device(Device::Cpu);
+    let data = host.as_slice::<bf16>();
+    let last = &data[(t - 1) * vocab..t * vocab];
+
+    let mut idx: Vec<usize> = (0..vocab).collect();
+    idx.sort_by(|&a, &b| last[b].to_f32().partial_cmp(&last[a].to_f32()).unwrap());
+    println!("top10 next-token (id: logit):");
+    for &i in &idx[..10] {
+        println!("  {i}: {:.4}", last[i].to_f32());
+    }
+}

crates/xserv-model/src/bin/xserv-chat.rs

@@ -1,166 +1,16 @@
 use std::io::{self, IsTerminal, Read, Write};
 use std::path::PathBuf;
 
-use std::sync::{Arc, mpsc};
-use std::thread;
-
-use xserv_model::{
-    BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, SamplingParams,
-    loader, sample, sample_greedy_penalized,
-};
+use xserv_model::{loader, sample, ModelConfig, PagedKVCache, Qwen3, SamplingParams, BLOCK_SIZE};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
-enum ChatModel {
-    Qwen3(Qwen3),
-    GptOss(GptOss),
-}
-
-impl ChatModel {
-    fn forward_prefill_paged(
-        &self,
-        tokens: &[u32],
-        slot: usize,
-        cache: &mut PagedKVCache,
-    ) -> xserv_tensor::Tensor {
-        match self {
-            ChatModel::Qwen3(m) => m.forward_prefill_paged(tokens, slot, cache),
-            ChatModel::GptOss(m) => m.forward_prefill_paged(tokens, slot, cache),
-        }
-    }
-    fn forward_decode_paged(
-        &self,
-        tokens: &[u32],
-        positions: &[usize],
-        slots: &[usize],
-        cache: &mut PagedKVCache,
-    ) -> xserv_tensor::Tensor {
-        match self {
-            ChatModel::Qwen3(m) => m.forward_decode_paged(tokens, positions, slots, cache),
-            ChatModel::GptOss(m) => m.forward_decode_paged(tokens, positions, slots, cache),
-        }
-    }
-}
-
-// TP worker infrastructure (reused from tp_engine pattern)
-#[derive(Clone)]
-enum TpCommand {
-    Register(usize),
-    Free(usize),
-    Prefill {
-        tokens: Vec<u32>,
-        slot: usize,
-    },
-    Decode {
-        tokens: Vec<u32>,
-        positions: Vec<usize>,
-        slots: Vec<usize>,
-    },
-}
-
-struct TpHandle {
-    cmd_txs: Vec<mpsc::Sender<TpCommand>>,
-    ack_rx: mpsc::Receiver<()>,
-}
-
-impl TpHandle {
-    fn send(&self, cmd: TpCommand) {
-        for tx in &self.cmd_txs {
-            tx.send(cmd.clone()).ok();
-        }
-    }
-    fn wait(&self) {
-        for _ in 0..self.cmd_txs.len() {
-            self.ack_rx.recv().ok();
-        }
-    }
-}
-
-fn tp_worker_loop(
-    rank: usize,
-    world: usize,
-    id: xserv_distributed::UniqueId,
-    model_dir: std::path::PathBuf,
-    config: ModelConfig,
-    max_seq_len: usize,
-    cmd_rx: mpsc::Receiver<TpCommand>,
-    ack_tx: mpsc::Sender<()>,
-) {
-    let tp = Arc::new(xserv_distributed::TpContext::init(
-        rank,
-        world,
-        id,
-        rank as u32,
-    ));
-    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
-    let model = if config.is_moe() {
-        ChatModel::GptOss(GptOss::from_weights_tp(
-            config.clone(),
-            weights,
-            rank,
-            world,
-            rank as u32,
-            Some(tp),
-        ))
-    } else {
-        ChatModel::Qwen3(Qwen3::from_weights_tp(
-            config.clone(),
-            weights,
-            rank,
-            world,
-            rank as u32,
-            Some(tp),
-        ))
-    };
-    let local_kv = config.num_kv_heads() / world;
-    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
-    let total_blocks = max_blocks_per_seq + 8;
-    let mut cache = PagedKVCache::new_tp(
-        &config,
-        local_kv,
-        total_blocks,
-        0,
-        1,
-        max_blocks_per_seq,
-        DType::BF16,
-        rank as u32,
-    );
-    let mut decoder = GraphedGptOssDecoder::new();
-    while let Ok(cmd) = cmd_rx.recv() {
-        match cmd {
-            TpCommand::Register(slot) => {
-                let _ = cache.register_sequence(slot);
-            }
-            TpCommand::Free(slot) => cache.free_sequence(slot),
-            TpCommand::Prefill { tokens, slot } => {
-                let _ = model.forward_prefill_paged(&tokens, slot, &mut cache);
-            }
-            TpCommand::Decode {
-                tokens,
-                positions,
-                slots,
-            } => {
-                let _ = chat_decode(
-                    &model,
-                    &mut decoder,
-                    &tokens,
-                    &positions,
-                    &slots,
-                    &mut cache,
-                );
-            }
-        }
-        let _ = ack_tx.send(());
-    }
-}
-
 const SLOT: usize = 0;
 
 struct CliOptions {
     model_dir: PathBuf,
     max_tokens: usize,
     max_seq_len: usize,
-    tp: usize,
     sampling: SamplingParams,
     system_prompt: Option<String>,
     enable_thinking: bool,
@@ -282,13 +132,7 @@ fn read_line_edited(prompt: &str) -> Line {
             }
             b => {
                 // UTF-8 multi-byte: read the continuation bytes for this char.
-                let extra = if b >= 0xF0 {
-                    3
-                } else if b >= 0xE0 {
-                    2
-                } else {
-                    1
-                };
+                let extra = if b >= 0xF0 { 3 } else if b >= 0xE0 { 2 } else { 1 };
                 let mut bytes = vec![b];
                 let mut cont = [0u8; 1];
                 let mut ok = true;
@@ -324,12 +168,14 @@ fn main() {
 
     let config = ModelConfig::from_file(&opts.model_dir.join("config.json"));
     let model_type = config.model_type.as_deref().unwrap_or("unknown");
-    let is_moe = config.is_moe();
+    if !model_type.contains("qwen") {
+        eprintln!("xserv-chat currently supports Qwen-style ChatML models only; got model_type={model_type}");
+        std::process::exit(2);
+    }
 
     let max_seq_len = opts.max_seq_len.min(config.max_seq_len()).max(1);
     eprintln!(
-        "Model: {model_type}{}, layers={}, hidden={}, heads={}/{} kv, vocab={}, max_seq_len={}",
-        if is_moe { " (MoE)" } else { "" },
+        "Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}, max_seq_len={}",
         config.num_layers(),
         config.hidden(),
         config.num_heads(),
@@ -338,108 +184,19 @@ fn main() {
         max_seq_len
     );
 
-    let world = opts.tp;
-    if world > 1 {
-        assert!(
-            config.num_kv_heads() % world == 0,
-            "num_kv_heads {} not divisible by tp {world}",
-            config.num_kv_heads()
-        );
-    }
-
-    let (model, mut cache, tp_handle) = if world > 1 {
-        let id = xserv_distributed::get_unique_id();
-        let (ack_tx, ack_rx) = mpsc::channel::<()>();
-        let mut cmd_txs = Vec::new();
-        for rank in 1..world {
-            let (ctx_tx, ctx_rx) = mpsc::channel::<TpCommand>();
-            cmd_txs.push(ctx_tx);
-            let ack_tx = ack_tx.clone();
-            let model_dir = opts.model_dir.clone();
-            let config = config.clone();
-            thread::spawn(move || {
-                tp_worker_loop(
-                    rank,
-                    world,
-                    id,
-                    model_dir,
-                    config,
-                    max_seq_len,
-                    ctx_rx,
-                    ack_tx,
-                );
-            });
-        }
-        eprintln!("Loading weights (tp={world})...");
-        let tp = Arc::new(xserv_distributed::TpContext::init(0, world, id, 0));
-        let weights = loader::load_model_dir(&opts.model_dir, Device::Cpu);
-        eprintln!("Loaded {} tensors", weights.len());
-        let m = if is_moe {
-            ChatModel::GptOss(GptOss::from_weights_tp(
-                config.clone(),
-                weights,
-                0,
-                world,
-                0,
-                Some(tp),
-            ))
-        } else {
-            ChatModel::Qwen3(Qwen3::from_weights_tp(
-                config.clone(),
-                weights,
-                0,
-                world,
-                0,
-                Some(tp),
-            ))
-        };
-        let local_kv = config.num_kv_heads() / world;
-        let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
-        let total_blocks = max_blocks_per_seq + 8;
-        let c = PagedKVCache::new_tp(
-            &config,
-            local_kv,
-            total_blocks,
-            0,
-            1,
-            max_blocks_per_seq,
-            DType::BF16,
-            0,
-        );
-        let h = TpHandle { cmd_txs, ack_rx };
-        (m, c, Some(h))
-    } else {
-        eprintln!("Loading weights...");
-        let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0));
-        eprintln!("Loaded {} tensors", weights.len());
-        let m = if is_moe {
-            ChatModel::GptOss(GptOss::from_weights(config.clone(), weights))
-        } else {
-            ChatModel::Qwen3(Qwen3::from_weights(config.clone(), weights))
-        };
-        let c = new_paged_cache(&config, max_seq_len);
-        (m, c, None)
-    };
+    eprintln!("Loading weights...");
+    let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0));
+    eprintln!("Loaded {} tensors", weights.len());
+    let model = Qwen3::from_weights(config.clone(), weights);
 
     let tokenizer = Tokenizer::from_file(&opts.model_dir.join("tokenizer.json"));
-    let mut decoder = GraphedGptOssDecoder::new();
-    if let Some(h) = &tp_handle {
-        h.send(TpCommand::Register(SLOT));
-        h.wait();
-    }
+    let mut cache = new_paged_cache(&config, max_seq_len);
     cache.register_sequence(SLOT).expect("register chat slot");
     let use_color = opts.color && io::stdout().is_terminal();
 
-    eprintln!("Ready (paged KV cache, tp={world}).");
+    eprintln!("Ready (paged KV cache, persistent chat slot).");
     eprintln!("Commands: /exit, /quit, /clear\n");
 
-    // gpt-oss multi-turn history of (user, assistant-final) text. Harmony
-    // requires re-rendering the conversation each turn with prior analysis
-    // dropped, so the moe path re-prefills from this rather than reusing an
-    // incremental KV cache (which would accumulate CoT + <|return|> and collapse
-    // at longer context). Qwen3 ignores this and keeps the incremental cache.
-    let mut moe_history: Vec<(String, String)> = Vec::new();
-
     loop {
         let line = match read_line_edited("user> ") {
             Line::Eof => break,
@@ -453,8 +210,8 @@ fn main() {
         match input {
             "/exit" | "/quit" | "exit" | "quit" => break,
             "/clear" => {
-                reset_slot(&mut cache, &tp_handle);
-                moe_history.clear();
+                cache.free_sequence(SLOT);
+                cache.register_sequence(SLOT).expect("register chat slot");
                 eprintln!("history and KV cache cleared");
                 continue;
             }
@@ -465,46 +222,6 @@ fn main() {
             _ => {}
         }
 
-        if is_moe {
-            // Harmony multi-turn: re-render the whole conversation (prior
-            // analysis dropped) and re-prefill into a freshly cleared slot.
-            let prompt =
-                build_conversation_gpt_oss(opts.system_prompt.as_deref(), &moe_history, input);
-            let prompt_tokens = tokenizer.encode(&prompt);
-            if prompt_tokens.is_empty() {
-                continue;
-            }
-            if prompt_tokens.len() >= max_seq_len {
-                eprintln!(
-                    "context full: conversation needs {} tokens >= max_seq_len {max_seq_len}; use /clear",
-                    prompt_tokens.len()
-                );
-                continue;
-            }
-            let max_new_tokens = opts.max_tokens.min(max_seq_len - prompt_tokens.len());
-            reset_slot(&mut cache, &tp_handle);
-            print!("assistant> ");
-            io::stdout().flush().unwrap();
-            let (_finish, answer) = generate_with_paged_cache(
-                &model,
-                &mut decoder,
-                &mut cache,
-                &tokenizer,
-                &prompt_tokens,
-                &opts.sampling,
-                max_new_tokens,
-                use_color,
-                &tp_handle,
-                is_moe,
-                opts.enable_thinking,
-            );
-            moe_history.push((input.to_string(), answer));
-            println!();
-            continue;
-        }
-
-        // Qwen3: incremental KV cache — only the new turn is prefilled and the
-        // assistant's tokens stay cached for the next turn.
         let include_system = cache.seq_len(SLOT) == 0;
         let prompt = build_turn_prompt(
             opts.system_prompt.as_deref(),
@@ -530,59 +247,27 @@ fn main() {
 
         print!("assistant> ");
         io::stdout().flush().unwrap();
-        let (finish, _answer) = generate_with_paged_cache(
+        let finish = generate_with_paged_cache(
             &model,
-            &mut decoder,
             &mut cache,
             &tokenizer,
             &prompt_tokens,
             &opts.sampling,
             max_new_tokens,
             use_color,
-            &tp_handle,
-            is_moe,
-            opts.enable_thinking,
         );
         match finish {
             Finish::Stop { token_id } => {
-                append_after_stop(
-                    &model,
-                    &mut cache,
-                    &tokenizer,
-                    max_seq_len,
-                    token_id,
-                    &tp_handle,
-                );
+                append_after_stop(&model, &mut cache, &tokenizer, max_seq_len, token_id);
             }
             Finish::Length => {
-                append_text_to_cache(
-                    &model,
-                    &mut cache,
-                    &tokenizer,
-                    max_seq_len,
-                    "<|im_end|>\n",
-                    &tp_handle,
-                );
+                append_text_to_cache(&model, &mut cache, &tokenizer, max_seq_len, "<|im_end|>\n");
             }
         }
         println!();
     }
 }
 
-/// Free and re-register the single chat KV slot (clears all cached context).
-fn reset_slot(cache: &mut PagedKVCache, tp: &Option<TpHandle>) {
-    if let Some(h) = tp {
-        h.send(TpCommand::Free(SLOT));
-        h.wait();
-    }
-    cache.free_sequence(SLOT);
-    if let Some(h) = tp {
-        h.send(TpCommand::Register(SLOT));
-        h.wait();
-    }
-    cache.register_sequence(SLOT).expect("register chat slot");
-}
-
 fn parse_args() -> CliOptions {
     let args: Vec<String> = std::env::args().skip(1).collect();
     if args.is_empty() || args.iter().any(|a| a == "--help" || a == "-h") {
@@ -592,7 +277,6 @@ fn parse_args() -> CliOptions {
     let mut model_dir = None;
     let mut max_tokens = 256usize;
     let mut max_seq_len = 2048usize;
-    let mut tp = 1usize;
     let mut temperature = 0.0f32;
     let mut top_k = 0usize;
     let mut top_p = 1.0f32;
@@ -615,10 +299,6 @@ fn parse_args() -> CliOptions {
                 i += 1;
                 max_seq_len = parse_value(&args, i, "--max-seq-len");
             }
-            "--tp" => {
-                i += 1;
-                tp = parse_value(&args, i, "--tp");
-            }
             "--temperature" => {
                 i += 1;
                 temperature = parse_value(&args, i, "--temperature");
@@ -667,7 +347,6 @@ fn parse_args() -> CliOptions {
         }),
         max_tokens: max_tokens.max(1),
         max_seq_len: max_seq_len.max(1),
-        tp: tp.max(1),
         sampling: SamplingParams {
             temperature,
             top_k,
@@ -694,7 +373,6 @@ fn print_usage_and_exit(code: i32) -> ! {
          \t-m, --model DIR       Model directory\n\
          \t--max-tokens N        Max generated tokens per turn (default: 256)\n\
          \t--max-seq-len N       Persistent KV context length (default: 2048)\n\
-         \t--tp N                Tensor parallelism degree (default: 1)\n\
          \t--temperature F       Sampling temperature, 0 = greedy (default: 0)\n\
          \t--top-k N             Top-k sampling, 0 = disabled (default: 0)\n\
          \t--top-p F             Top-p sampling (default: 1.0)\n\
@@ -717,15 +395,7 @@ fn new_paged_cache(config: &ModelConfig, max_seq_len: usize) -> PagedKVCache {
     let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
     let total_blocks = (max_blocks_per_seq + 1).max(2);
     // Single-slot interactive CLI: no swap pool (cpu_total_blocks = 0).
-    PagedKVCache::new(
-        config,
-        total_blocks,
-        0,
-        1,
-        max_blocks_per_seq,
-        DType::BF16,
-        0,
-    )
+    PagedKVCache::new(config, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, 0)
 }
 
 fn build_turn_prompt(
@@ -754,308 +424,33 @@ fn build_turn_prompt(
     prompt
 }
 
-/// Render the full gpt-oss harmony conversation for re-prefill. gpt-oss was
-/// trained on this exact system-message structure (identity / knowledge cutoff
-/// / current date / Reasoning level / channels — see the model's
-/// chat_template.jinja `build_system_message`). A hand-rolled substitute puts
-/// the model out of distribution and destabilizes channel selection.
-///
-/// Harmony multi-turn drops prior chain-of-thought: past assistant messages are
-/// rendered as completed `final` channels ending in `<|end|>` (not the
-/// `<|return|>` stop token). Keeping the analysis + `<|return|>` of every turn
-/// in context — as an incremental KV cache does — is out of distribution and
-/// makes the model collapse at longer context. "Reasoning: low" keeps the
-/// analysis channel short for an interactive chat.
-fn build_conversation_gpt_oss(
-    system: Option<&str>,
-    history: &[(String, String)],
-    current_user: &str,
-) -> String {
-    let mut prompt = String::new();
-    prompt.push_str("<|start|>system<|message|>");
-    prompt.push_str("You are ChatGPT, a large language model trained by OpenAI.\n");
-    prompt.push_str("Knowledge cutoff: 2024-06\n");
-    prompt.push_str(&format!("Current date: {}\n\n", today_ymd()));
-    prompt.push_str("Reasoning: low\n\n");
-    prompt.push_str("# Valid channels: analysis, commentary, final. Channel must be included for every message.");
-    prompt.push_str("<|end|>");
-    if let Some(sys) = system {
-        if !sys.trim().is_empty() {
-            prompt.push_str("<|start|>developer<|message|># Instructions\n\n");
-            prompt.push_str(sys.trim());
-            prompt.push_str("<|end|>");
-        }
-    }
-    for (user, assistant) in history {
-        prompt.push_str("<|start|>user<|message|>");
-        prompt.push_str(user);
-        prompt.push_str("<|end|>");
-        prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
-        prompt.push_str(assistant.trim());
-        prompt.push_str("<|end|>");
-    }
-    prompt.push_str("<|start|>user<|message|>");
-    prompt.push_str(current_user);
-    prompt.push_str("<|end|>");
-    prompt.push_str("<|start|>assistant");
-    prompt
-}
-
-/// Current UTC date as "YYYY-MM-DD" for the harmony system message. Rata Die
-/// civil-calendar conversion (same algorithm the server uses for strftime_now).
-fn today_ymd() -> String {
-    use std::time::{SystemTime, UNIX_EPOCH};
-    let secs = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap()
-        .as_secs();
-    let z = (secs / 86400) as i64 + 719468;
-    let era = (if z >= 0 { z } else { z - 146096 }) / 146097;
-    let doe = z - era * 146097;
-    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
-    let y = yoe + era * 400;
-    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
-    let mp = (5 * doy + 2) / 153;
-    let d = doy - (153 * mp + 2) / 5 + 1;
-    let m = if mp < 10 { mp + 3 } else { mp - 9 };
-    let y = if m <= 2 { y + 1 } else { y };
-    format!("{y:04}-{m:02}-{d:02}")
-}
-
-fn chat_decode(
-    model: &ChatModel,
-    decoder: &mut GraphedGptOssDecoder,
-    tokens: &[u32],
-    positions: &[usize],
-    slots: &[usize],
-    cache: &mut PagedKVCache,
-) -> xserv_tensor::Tensor {
-    match model {
-        ChatModel::GptOss(m) => decoder.decode(m, tokens, positions, slots, cache),
-        ChatModel::Qwen3(_) => model.forward_decode_paged(tokens, positions, slots, cache),
-    }
-}
-
 fn generate_with_paged_cache(
-    model: &ChatModel,
-    decoder: &mut GraphedGptOssDecoder,
+    model: &Qwen3,
     cache: &mut PagedKVCache,
     tokenizer: &Tokenizer,
     prompt_tokens: &[u32],
     sampling: &SamplingParams,
     max_tokens: usize,
     use_color: bool,
-    tp: &Option<TpHandle>,
-    is_moe: bool,
-    enable_thinking: bool,
-) -> (Finish, String) {
-    let harmony_end_id = if is_moe {
-        tokenizer.special_token_id("<|end|>")
-    } else {
-        None
-    };
-    let harmony_channel_id = if is_moe {
-        tokenizer.special_token_id("<|channel|>")
-    } else {
-        None
-    };
-    let harmony_message_id = if is_moe {
-        tokenizer.special_token_id("<|message|>")
-    } else {
-        None
-    };
-    let harmony_special: Vec<u32> = if is_moe {
-        [
-            "<|channel|>",
-            "<|start|>",
-            "<|end|>",
-            "<|message|>",
-            "<|return|>",
-        ]
-        .iter()
-        .filter_map(|s| tokenizer.special_token_id(s))
-        .collect()
-    } else {
-        Vec::new()
-    };
-    // Harmony channel state: "final" channel text is printed normally,
-    // "analysis" channel is rendered as thinking (gray). After <|channel|>
-    // we read the channel name tokens until <|message|>.
-    #[derive(PartialEq, Clone, Copy)]
-    enum HarmonyState {
-        Normal,
-        ReadingChannel,
-        InAnalysis,
-        InFinal,
-    }
-    let mut hstate = if is_moe {
-        HarmonyState::InFinal
-    } else {
-        HarmonyState::Normal
-    };
-
-    // Off by default. A repetition penalty over a harmony stream penalizes the
-    // control tokens (<|channel|>, <|message|>, <|start|>) that MUST repeat to
-    // open the final channel — so a non-1.0 default makes gpt-oss stop right
-    // after the analysis block, before emitting any answer. Opt in via the env
-    // var if you want it for plain (non-harmony) generation.
-    let rep_penalty: f32 = std::env::var("XSERV_REP_PENALTY")
-        .ok()
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(1.0);
-    let rep_window: usize = std::env::var("XSERV_REP_WINDOW")
-        .ok()
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(512);
-    let mut history: Vec<u32> = Vec::new();
-
-    let pick = |logits: &xserv_tensor::Tensor, sp: &SamplingParams, hist: &[u32]| -> u32 {
-        if rep_penalty > 1.0 && sp.temperature == 0.0 {
-            let start = hist.len().saturating_sub(rep_window);
-            sample_greedy_penalized(logits, &hist[start..], rep_penalty)
-        } else {
-            sample(logits, sp)
-        }
-    };
-
-    if let Some(h) = tp {
-        h.send(TpCommand::Prefill {
-            tokens: prompt_tokens.to_vec(),
-            slot: SLOT,
-        });
-    }
+) -> Finish {
     let logits = model.forward_prefill_paged(prompt_tokens, SLOT, cache);
-    if let Some(h) = tp {
-        h.wait();
-    }
-    let mut next = pick(&logits, sampling, &history);
+    let mut next = sample(&logits, sampling);
     let mut decode_buffer = Vec::new();
     let mut in_thinking = false;
-    let show_thinking = is_moe && enable_thinking;
-    // Visible answer tokens, returned for multi-turn history. For moe this is
-    // the final-channel content only (analysis is suppressed/gray); for Qwen3
-    // it is everything printed. The caller decodes these into the assistant
-    // message it re-renders into the next prompt.
-    let mut answer_ids: Vec<u32> = Vec::new();
 
     for _ in 0..max_tokens {
         let position = cache.seq_len(SLOT);
-        if let Some(h) = tp {
-            h.send(TpCommand::Decode {
-                tokens: vec![next],
-                positions: vec![position],
-                slots: vec![SLOT],
-            });
-        }
-        let logits = chat_decode(model, decoder, &[next], &[position], &[SLOT], cache);
-        if let Some(h) = tp {
-            h.wait();
-        }
-        if tokenizer.is_eos(next) {
+        let logits = model.forward_decode_paged(&[next], &[position], &[SLOT], cache);
+        if is_stop_token(tokenizer, next) {
             print_stream_text(
                 &tokenizer.flush_decode_stream(&mut decode_buffer),
                 in_thinking,
                 use_color,
             );
-            if show_thinking && in_thinking {
-                print_stream_text("\n</think>\n\n", true, use_color);
-            }
             io::stdout().flush().unwrap();
-            return (
-                Finish::Stop { token_id: next },
-                tokenizer.decode(&answer_ids),
-            );
-        }
-        if harmony_end_id == Some(next) {
-            // <|end|> closes current segment; if in final channel, we're done
-            print_stream_text(
-                &tokenizer.flush_decode_stream(&mut decode_buffer),
-                in_thinking,
-                use_color,
-            );
-            if hstate == HarmonyState::InFinal {
-                io::stdout().flush().unwrap();
-                return (
-                    Finish::Stop { token_id: next },
-                    tokenizer.decode(&answer_ids),
-                );
-            }
-            // Closing a thinking (analysis/commentary) channel: emit the </think>
-            // marker so it renders like Qwen3's thinking block.
-            if show_thinking && hstate == HarmonyState::InAnalysis {
-                print_stream_text("\n</think>\n\n", true, use_color);
-                in_thinking = false;
-            }
-            hstate = HarmonyState::Normal;
-            next = pick(&logits, sampling, &history);
-            continue;
+            return Finish::Stop { token_id: next };
         }
 
-        history.push(next);
-
-        // Harmony channel routing state machine
-        if harmony_channel_id == Some(next) {
-            decode_buffer.clear();
-            hstate = HarmonyState::ReadingChannel;
-            next = pick(&logits, sampling, &history);
-            continue;
-        }
-        if harmony_message_id == Some(next) {
-            if hstate == HarmonyState::ReadingChannel {
-                // Channel name was accumulated but we don't need to parse it —
-                // we just check via the channel_name buffer below
-            }
-            decode_buffer.clear();
-            next = pick(&logits, sampling, &history);
-            continue;
-        }
-        if hstate == HarmonyState::ReadingChannel {
-            // Reading channel name tokens (e.g. "final", "analysis")
-            let tok_text = tokenizer.decode(&[next]);
-            if tok_text.contains("final") {
-                hstate = HarmonyState::InFinal;
-                in_thinking = false;
-            } else {
-                hstate = HarmonyState::InAnalysis;
-                // Open a Qwen3-style thinking block for the analysis channel.
-                if show_thinking {
-                    print_stream_text("<think>\n", true, use_color);
-                    in_thinking = true;
-                }
-            }
-            next = pick(&logits, sampling, &history);
-            continue;
-        }
-        if harmony_special.contains(&next) {
-            next = pick(&logits, sampling, &history);
-            continue;
-        }
-        if hstate == HarmonyState::InAnalysis {
-            // Analysis channel = the model's reasoning. With --think, show it as a
-            // thinking block (gray if color); otherwise suppress it (answer only).
-            if show_thinking {
-                print_generated_token(
-                    tokenizer,
-                    next,
-                    &mut decode_buffer,
-                    &mut in_thinking,
-                    use_color,
-                );
-                io::stdout().flush().unwrap();
-            }
-            next = pick(&logits, sampling, &history);
-            continue;
-        }
-        if is_moe && hstate != HarmonyState::InFinal {
-            // Between harmony messages (after a channel's <|end|>, before the
-            // next <|channel|>): the model emits a role header like "assistant".
-            // That's structural, not user-visible content — suppress it. Only
-            // for moe/harmony; non-moe (Qwen3) stays in Normal and prints here.
-            next = pick(&logits, sampling, &history);
-            continue;
-        }
-
-        answer_ids.push(next);
         print_generated_token(
             tokenizer,
             next,
@@ -1064,7 +459,7 @@ fn generate_with_paged_cache(
             use_color,
         );
         io::stdout().flush().unwrap();
-        next = pick(&logits, sampling, &history);
+        next = sample(&logits, sampling);
     }
 
     print_stream_text(
@@ -1072,46 +467,34 @@ fn generate_with_paged_cache(
         in_thinking,
         use_color,
     );
-    if show_thinking && in_thinking {
-        print_stream_text("\n</think>\n\n", true, use_color);
-    }
     io::stdout().flush().unwrap();
-    (Finish::Length, tokenizer.decode(&answer_ids))
+    Finish::Length
 }
 
 fn append_after_stop(
-    model: &ChatModel,
+    model: &Qwen3,
     cache: &mut PagedKVCache,
     tokenizer: &Tokenizer,
     max_seq_len: usize,
-    _stop_token_id: u32,
-    tp: &Option<TpHandle>,
+    stop_token_id: u32,
 ) {
-    append_text_to_cache(model, cache, tokenizer, max_seq_len, "\n", tp);
+    if tokenizer.special_token_id("<|im_end|>") == Some(stop_token_id) {
+        append_text_to_cache(model, cache, tokenizer, max_seq_len, "\n");
+    }
 }
 
 fn append_text_to_cache(
-    model: &ChatModel,
+    model: &Qwen3,
     cache: &mut PagedKVCache,
     tokenizer: &Tokenizer,
     max_seq_len: usize,
     text: &str,
-    tp: &Option<TpHandle>,
 ) {
     let tokens = tokenizer.encode(text);
     if tokens.is_empty() || cache.seq_len(SLOT) + tokens.len() > max_seq_len {
         return;
     }
-    if let Some(h) = tp {
-        h.send(TpCommand::Prefill {
-            tokens: tokens.clone(),
-            slot: SLOT,
-        });
-    }
     let _ = model.forward_prefill_paged(&tokens, SLOT, cache);
-    if let Some(h) = tp {
-        h.wait();
-    }
 }
 
 fn print_generated_token(
@@ -1157,3 +540,10 @@ fn print_stream_text(text: &str, in_thinking: bool, use_color: bool) {
         print!("{text}");
     }
 }
+
+fn is_stop_token(tokenizer: &Tokenizer, token_id: u32) -> bool {
+    tokenizer.eos_token_id() == Some(token_id)
+        || tokenizer.special_token_id("<|im_end|>") == Some(token_id)
+        || tokenizer.special_token_id("<|endoftext|>") == Some(token_id)
+        || tokenizer.special_token_id("<|end_of_text|>") == Some(token_id)
+}

crates/xserv-model/src/bin/xserv-cli.rs

@@ -1,69 +1,34 @@
 use std::io::{self, Write};
 use std::path::PathBuf;
-use xserv_model::{
-    BLOCK_SIZE, KVCache, ModelConfig, PagedKVCache, SamplingParams, loader, sample,
-    sample_greedy_penalized,
-};
+use xserv_model::{loader, KVCache, ModelConfig};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
-fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
-    args.iter()
-        .position(|a| a == name)
-        .and_then(|i| args.get(i + 1))
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(default)
-}
-
-fn pick_next(
-    logits: &xserv_tensor::Tensor,
-    sampling: &SamplingParams,
-    history: &[u32],
-    rep_penalty: f32,
-) -> u32 {
-    if rep_penalty > 1.0 && sampling.temperature == 0.0 {
-        sample_greedy_penalized(logits, history, rep_penalty)
-    } else {
-        sample(logits, sampling)
-    }
-}
-
 fn main() {
     let args: Vec<String> = std::env::args().collect();
     if args.len() < 2 {
-        eprintln!(
-            "Usage: xserv-cli <model-dir> [--max-tokens N] [--temperature F] [--top-k N] [--top-p F] [--rep-penalty F] [--rep-window N]"
-        );
+        eprintln!("Usage: xserv-cli <model-dir> [--max-tokens N]");
         std::process::exit(1);
     }
 
     let model_dir = PathBuf::from(&args[1]);
-    let max_tokens = flag(&args, "--max-tokens", 100usize);
-    let sampling = SamplingParams {
-        temperature: flag(&args, "--temperature", 0.0f32),
-        top_k: flag(&args, "--top-k", 0usize),
-        top_p: flag(&args, "--top-p", 1.0f32),
-    };
-    let rep_penalty = flag(&args, "--rep-penalty", 1.0f32);
-    let rep_window = flag(&args, "--rep-window", 512usize);
+    let max_tokens: usize = args
+        .iter()
+        .position(|a| a == "--max-tokens")
+        .and_then(|i| args.get(i + 1))
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(100);
 
     xserv_cuda::device::set_device(0).unwrap();
     let info = xserv_cuda::device::device_info(0).unwrap();
-    eprintln!(
-        "GPU: {} ({} MB free)",
-        info.name,
-        info.free_memory / 1024 / 1024
-    );
+    eprintln!("GPU: {} ({} MB free)", info.name, info.free_memory / 1024 / 1024);
 
     let config = ModelConfig::from_file(&model_dir.join("config.json"));
     let model_type = config.model_type.as_deref().unwrap_or("unknown");
     eprintln!(
         "Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}",
-        config.num_layers(),
-        config.hidden(),
-        config.num_heads(),
-        config.num_kv_heads(),
-        config.vocab_size
+        config.num_layers(), config.hidden(), config.num_heads(),
+        config.num_kv_heads(), config.vocab_size
     );
 
     eprintln!("Loading weights...");
@@ -71,142 +36,66 @@ fn main() {
     eprintln!("Loaded {} tensors", weights.len());
 
     let is_qwen3 = model_type.contains("qwen");
-    let is_gpt_oss = model_type.contains("gpt_oss");
-    let dtype = if is_qwen3 || is_gpt_oss {
-        DType::BF16
-    } else {
-        DType::F32
-    };
+    let dtype = if is_qwen3 { DType::BF16 } else { DType::F32 };
 
     // Build model
     enum Model {
         GPT2(xserv_model::GPT2),
         Qwen3(xserv_model::Qwen3),
-        GptOss(xserv_model::GptOss),
     }
-    let model = if is_gpt_oss {
-        Model::GptOss(xserv_model::GptOss::from_weights(config.clone(), weights))
-    } else if is_qwen3 {
+    let model = if is_qwen3 {
         Model::Qwen3(xserv_model::Qwen3::from_weights(config.clone(), weights))
     } else {
         Model::GPT2(xserv_model::GPT2::from_weights(config.clone(), weights))
     };
 
     let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
-    eprintln!(
-        "Ready (KV cache, dtype={dtype}, temperature={}, top_k={}, top_p={}, rep_penalty={}, rep_window={}).\n",
-        sampling.temperature, sampling.top_k, sampling.top_p, rep_penalty, rep_window
-    );
+    eprintln!("Ready (KV cache, dtype={dtype}).\n");
 
     loop {
         print!("xserv> ");
         io::stdout().flush().unwrap();
         let mut input = String::new();
-        if io::stdin().read_line(&mut input).unwrap() == 0 {
-            break;
-        }
-        let raw_input = input.trim();
-        if raw_input.is_empty() {
-            continue;
-        }
-        if raw_input == "quit" || raw_input == "exit" {
-            break;
-        }
-        let input = raw_input.replace("\\n", "\n");
+        if io::stdin().read_line(&mut input).unwrap() == 0 { break; }
+        let input = input.trim();
+        if input.is_empty() { continue; }
+        if input == "quit" || input == "exit" { break; }
 
-        let token_ids = tokenizer.encode(&input);
+        let token_ids = tokenizer.encode(input);
+        let kv_heads = if is_qwen3 { config.num_kv_heads() } else { config.num_heads() };
+        let mut cache = KVCache::new(
+            config.num_layers(), kv_heads, config.head_dim(), dtype, Device::Cuda(0),
+        );
 
-        if is_gpt_oss {
-            // GptOss uses paged KV cache
-            let max_seq = 2048;
-            let max_blocks_per_seq = (max_seq + BLOCK_SIZE - 1) / BLOCK_SIZE;
-            let total_blocks = max_blocks_per_seq + 64;
-            let mut paged_cache = PagedKVCache::new(
-                &config,
-                total_blocks,
-                0,
-                4,
-                max_blocks_per_seq,
-                DType::BF16,
-                0,
-            );
-            let slot = 0;
-            paged_cache.register_sequence(slot).expect("register slot");
+        // Prefill + decode
+        let logits = match &model {
+            Model::GPT2(m) => m.forward_with_cache(&token_ids, &mut cache),
+            Model::Qwen3(m) => m.forward_with_cache(&token_ids, &mut cache),
+        };
+        let mut next = match &model {
+            Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
+            Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
+        };
 
-            let model = match &model {
-                Model::GptOss(m) => m,
-                _ => unreachable!(),
-            };
-            let logits = model.forward_prefill_paged(&token_ids, slot, &mut paged_cache);
-            let mut history = token_ids.clone();
-            let start = history.len().saturating_sub(rep_window);
-            let mut next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
+        print!("{input}");
+        io::stdout().flush().unwrap();
 
-            print!("{input}");
+        for _ in 0..max_tokens {
+            let text = tokenizer.decode(&[next]);
+            print!("{text}");
             io::stdout().flush().unwrap();
 
-            for _ in 0..max_tokens {
-                let text = tokenizer.decode(&[next]);
-                print!("{text}");
-                io::stdout().flush().unwrap();
-                history.push(next);
-
-                if tokenizer.eos_token_id() == Some(next) {
-                    break;
-                }
-
-                let pos = paged_cache.seq_len(slot);
-                let logits = model.forward_decode_paged(&[next], &[pos], &[slot], &mut paged_cache);
-                let start = history.len().saturating_sub(rep_window);
-                next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
-            }
-            println!();
-            paged_cache.free_sequence(slot);
-        } else {
-            let kv_heads = if is_qwen3 {
-                config.num_kv_heads()
-            } else {
-                config.num_heads()
-            };
-            let mut cache = KVCache::new(
-                config.num_layers(),
-                kv_heads,
-                config.head_dim(),
-                dtype,
-                Device::Cuda(0),
-            );
+            if tokenizer.eos_token_id() == Some(next) { break; }
 
             let logits = match &model {
-                Model::GPT2(m) => m.forward_with_cache(&token_ids, &mut cache),
-                Model::Qwen3(m) => m.forward_with_cache(&token_ids, &mut cache),
-                Model::GptOss(_) => unreachable!(),
+                Model::GPT2(m) => m.forward_with_cache(&[next], &mut cache),
+                Model::Qwen3(m) => m.forward_with_cache(&[next], &mut cache),
+            };
+            next = match &model {
+                Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
+                Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
             };
-            let mut history = token_ids.clone();
-            let start = history.len().saturating_sub(rep_window);
-            let mut next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
-
-            print!("{input}");
-            io::stdout().flush().unwrap();
-
-            for _ in 0..max_tokens {
-                let text = tokenizer.decode(&[next]);
-                print!("{text}");
-                io::stdout().flush().unwrap();
-                history.push(next);
-
-                if tokenizer.eos_token_id() == Some(next) {
-                    break;
-                }
-
-                let logits = match &model {
-                    Model::GPT2(m) => m.forward_with_cache(&[next], &mut cache),
-                    Model::Qwen3(m) => m.forward_with_cache(&[next], &mut cache),
-                    Model::GptOss(_) => unreachable!(),
-                };
-                let start = history.len().saturating_sub(rep_window);
-                next = pick_next(&logits, &sampling, &history[start..], rep_penalty);
-            }
-            println!();
         }
+        println!();
     }
 }

crates/xserv-model/src/config.rs

@@ -1,15 +1,6 @@
 use serde::Deserialize;
 use std::path::Path;
 
-#[derive(Debug, Clone, Deserialize)]
-pub struct RopeScaling {
-    pub rope_type: Option<String>,
-    pub factor: Option<f64>,
-    pub original_max_position_embeddings: Option<usize>,
-    pub beta_fast: Option<f64>,
-    pub beta_slow: Option<f64>,
-}
-
 #[derive(Debug, Clone, Deserialize)]
 pub struct ModelConfig {
     pub architectures: Option<Vec<String>>,
@@ -56,27 +47,27 @@ pub struct ModelConfig {
     #[serde(default)]
     pub tie_word_embeddings: Option<bool>,
 
-    // MoE (gpt-oss)
+    // Explicit head_dim (gpt-oss: 64, which is NOT hidden/num_heads). When
+    // absent, head_dim() falls back to hidden/num_heads (Qwen3, GPT-2).
+    #[serde(default)]
+    pub head_dim: Option<usize>,
+
+    // MoE (gpt-oss). Absent for dense models.
     #[serde(default)]
     pub num_local_experts: Option<usize>,
     #[serde(default)]
     pub num_experts_per_tok: Option<usize>,
+    // gpt-oss clamped-SwiGLU limit (config: swiglu_limit, default 7.0).
     #[serde(default)]
-    pub layer_types: Option<Vec<String>>,
+    pub swiglu_limit: Option<f64>,
+
+    // Sliding-window attention (gpt-oss: 128 on alternating layers). The
+    // pattern is given by `layer_types` (e.g. "sliding_attention" /
+    // "full_attention" per layer); absent for dense models.
     #[serde(default)]
     pub sliding_window: Option<usize>,
     #[serde(default)]
-    pub attention_bias: Option<bool>,
-    #[serde(default, rename = "head_dim")]
-    pub explicit_head_dim: Option<usize>,
-    #[serde(default)]
-    pub rope_scaling: Option<RopeScaling>,
-    #[serde(default)]
-    pub swiglu_limit: Option<f64>,
-    #[serde(default)]
-    pub geglu_alpha: Option<f64>,
-    #[serde(default)]
-    pub hidden_act: Option<String>,
+    pub layer_types: Option<Vec<String>>,
 }
 
 impl ModelConfig {
@@ -88,33 +79,23 @@ impl ModelConfig {
     }
 
     pub fn hidden(&self) -> usize {
-        self.hidden_size
-            .or(self.n_embd)
-            .expect("hidden_size or n_embd required")
+        self.hidden_size.or(self.n_embd).expect("hidden_size or n_embd required")
     }
 
     pub fn num_heads(&self) -> usize {
-        self.num_attention_heads
-            .or(self.n_head)
-            .expect("num_attention_heads or n_head required")
+        self.num_attention_heads.or(self.n_head).expect("num_attention_heads or n_head required")
     }
 
     pub fn num_layers(&self) -> usize {
-        self.num_hidden_layers
-            .or(self.n_layer)
-            .expect("num_hidden_layers or n_layer required")
+        self.num_hidden_layers.or(self.n_layer).expect("num_hidden_layers or n_layer required")
     }
 
     pub fn max_seq_len(&self) -> usize {
-        self.max_position_embeddings
-            .or(self.n_positions)
-            .unwrap_or(2048)
+        self.max_position_embeddings.or(self.n_positions).unwrap_or(2048)
     }
 
     pub fn ffn_hidden(&self) -> usize {
-        self.intermediate_size
-            .or(self.n_inner)
-            .unwrap_or(self.hidden() * 4)
+        self.intermediate_size.or(self.n_inner).unwrap_or(self.hidden() * 4)
     }
 
     pub fn num_kv_heads(&self) -> usize {
@@ -122,8 +103,48 @@ impl ModelConfig {
     }
 
     pub fn head_dim(&self) -> usize {
-        self.explicit_head_dim
-            .unwrap_or_else(|| self.hidden() / self.num_heads())
+        // gpt-oss sets head_dim explicitly (64 != 2880/64). Dense models omit it.
+        self.head_dim.unwrap_or_else(|| self.hidden() / self.num_heads())
+    }
+
+    // ----- MoE (gpt-oss) -----
+
+    /// True for MoE models (have an expert count in config).
+    pub fn is_moe(&self) -> bool {
+        self.num_local_experts.is_some()
+    }
+
+    pub fn num_experts(&self) -> usize {
+        self.num_local_experts.unwrap_or(0)
+    }
+
+    pub fn experts_per_tok(&self) -> usize {
+        self.num_experts_per_tok.unwrap_or(0)
+    }
+
+    /// Clamp bound for gpt-oss SwiGLU (config `swiglu_limit`, default 7.0).
+    pub fn swiglu_limit(&self) -> f32 {
+        self.swiglu_limit.unwrap_or(7.0) as f32
+    }
+
+    /// Whether layer `i` uses sliding-window attention. gpt-oss alternates per
+    /// `layer_types`; if that's absent but `sliding_window` is set, fall back to
+    /// the common "every other layer" pattern (even = sliding). Dense → false.
+    pub fn layer_uses_sliding_window(&self, layer: usize) -> bool {
+        if self.sliding_window.is_none() {
+            return false;
+        }
+        match &self.layer_types {
+            Some(types) => types
+                .get(layer)
+                .map(|t| t.contains("sliding"))
+                .unwrap_or(false),
+            None => layer % 2 == 0,
+        }
+    }
+
+    pub fn sliding_window(&self) -> Option<usize> {
+        self.sliding_window
     }
 
     pub fn ln_eps(&self) -> f32 {
@@ -135,32 +156,4 @@ impl ModelConfig {
     pub fn tied_embeddings(&self) -> bool {
         self.tie_word_embeddings.unwrap_or(true)
     }
-
-    pub fn num_experts(&self) -> usize {
-        self.num_local_experts.unwrap_or(0)
-    }
-
-    pub fn experts_per_token(&self) -> usize {
-        self.num_experts_per_tok.unwrap_or(1)
-    }
-
-    pub fn is_moe(&self) -> bool {
-        self.num_local_experts.unwrap_or(0) > 1
-    }
-
-    pub fn is_sliding_layer(&self, layer_idx: usize) -> bool {
-        self.layer_types
-            .as_ref()
-            .and_then(|lt| lt.get(layer_idx))
-            .map(|t| t == "sliding_attention")
-            .unwrap_or(false)
-    }
-
-    pub fn window_size(&self) -> usize {
-        self.sliding_window.unwrap_or(0)
-    }
-
-    pub fn geglu_alpha(&self) -> f32 {
-        self.geglu_alpha.unwrap_or(1.702) as f32
-    }
 }

crates/xserv-model/src/decode_graph.rs

@@ -9,7 +9,7 @@
 use std::ffi::c_void;
 use xserv_cuda::{CudaGraph, CudaStream, GpuBuffer};
 use xserv_kernels::dispatch;
-use xserv_kernels::gemm::{cublas_handle, gemv_scratch_elems};
+use xserv_kernels::gemm::cublas_handle;
 
 use crate::config::ModelConfig;
 use crate::kv_cache::GpuKVCache;
@@ -18,19 +18,19 @@ use crate::kv_cache::GpuKVCache;
 /// All buffers have stable GPU addresses for CUDA Graph replay.
 struct DecodeBuffers {
     // Hidden-size buffers: [1, hidden]
-    x: GpuBuffer,           // running hidden state
-    normed: GpuBuffer,      // rmsnorm output
-    attn_out: GpuBuffer,    // attention output [1, num_heads, 1, head_dim]
-    attn_merged: GpuBuffer, // merge_heads output [1, hidden]
-    o_proj: GpuBuffer,      // O projection output [1, hidden]
-    normed2: GpuBuffer,     // post-attn norm output [1, hidden]
-    sum_out: GpuBuffer,     // add_rmsnorm sum output [1, hidden]
-    down: GpuBuffer,        // down projection output [1, hidden]
+    x: GpuBuffer,               // running hidden state
+    normed: GpuBuffer,           // rmsnorm output
+    attn_out: GpuBuffer,         // attention output [1, num_heads, 1, head_dim]
+    attn_merged: GpuBuffer,      // merge_heads output [1, hidden]
+    o_proj: GpuBuffer,           // O projection output [1, hidden]
+    normed2: GpuBuffer,          // post-attn norm output [1, hidden]
+    sum_out: GpuBuffer,          // add_rmsnorm sum output [1, hidden]
+    down: GpuBuffer,             // down projection output [1, hidden]
 
     // QKV projection outputs
-    q_proj: GpuBuffer, // [1, num_heads * head_dim]
-    k_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
-    v_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
+    q_proj: GpuBuffer,           // [1, num_heads * head_dim]
+    k_proj: GpuBuffer,           // [1, num_kv_heads * head_dim]
+    v_proj: GpuBuffer,           // [1, num_kv_heads * head_dim]
 
     // Reshaped: [1, H, 1, D]
     q_reshaped: GpuBuffer,
@@ -50,23 +50,23 @@ struct DecodeBuffers {
     k_final: GpuBuffer,
 
     // FFN intermediates
-    gate: GpuBuffer,     // [1, intermediate]
-    up: GpuBuffer,       // [1, intermediate]
-    silu_out: GpuBuffer, // [1, intermediate]
+    gate: GpuBuffer,             // [1, intermediate]
+    up: GpuBuffer,               // [1, intermediate]
+    silu_out: GpuBuffer,         // [1, intermediate]
 
-    // GEMV fp32 scratch for deterministic K-block partials.
-    fp32_hidden: GpuBuffer,       // for hidden-sized GEMV outputs
-    fp32_q: GpuBuffer,            // for Q projection
-    fp32_kv: GpuBuffer,           // for K/V projection
-    fp32_intermediate: GpuBuffer, // for gate/up projections
-    fp32_vocab: GpuBuffer,        // for lm_head
+    // GEMV fp32 accumulators (separate per output dimension)
+    fp32_hidden: GpuBuffer,      // for hidden-sized GEMV outputs
+    fp32_q: GpuBuffer,           // for Q projection
+    fp32_kv: GpuBuffer,          // for K/V projection
+    fp32_intermediate: GpuBuffer,// for gate/up projections
+    fp32_vocab: GpuBuffer,       // for lm_head
 
     // Token ID and position (GPU-resident, updated before replay)
-    token_id_gpu: GpuBuffer, // 4 bytes (u32)
-    position_gpu: GpuBuffer, // 4 bytes (u32)
+    token_id_gpu: GpuBuffer,     // 4 bytes (u32)
+    position_gpu: GpuBuffer,     // 4 bytes (u32)
 
     // Final output
-    logits: GpuBuffer, // [1, vocab_size]
+    logits: GpuBuffer,           // [1, vocab_size]
 }
 
 pub struct DecodeGraphState {
@@ -140,14 +140,11 @@ impl DecodeGraphState {
             up: alloc(intermediate * es),
             silu_out: alloc(intermediate * es),
 
-            fp32_hidden: alloc(
-                gemv_scratch_elems(hidden, hidden).max(gemv_scratch_elems(intermediate, hidden))
-                    * 4,
-            ),
-            fp32_q: alloc(gemv_scratch_elems(hidden, num_heads * head_dim) * 4),
-            fp32_kv: alloc(gemv_scratch_elems(hidden, num_kv_heads * head_dim) * 4),
-            fp32_intermediate: alloc(gemv_scratch_elems(hidden, intermediate) * 4),
-            fp32_vocab: alloc(gemv_scratch_elems(hidden, vocab_size) * 4),
+            fp32_hidden: alloc(hidden * 4),
+            fp32_q: alloc(num_heads * head_dim * 4),
+            fp32_kv: alloc(num_kv_heads * head_dim * 4),
+            fp32_intermediate: alloc(intermediate * 4),
+            fp32_vocab: alloc(vocab_size * 4),
 
             token_id_gpu: alloc(4),
             position_gpu: alloc(4),
@@ -202,296 +199,127 @@ impl DecodeGraphState {
         let cublas = cublas_handle();
 
         // Set cuBLAS to use our stream
-        unsafe {
-            dispatch::set_cublas_stream(cublas, s);
-        }
+        unsafe { dispatch::set_cublas_stream(cublas, s); }
 
         for (l, lw) in layers.iter().enumerate() {
             // === Pre-attention graph ===
-            self.pre_attn_graphs[l]
-                .begin_capture(&self.stream)
-                .expect("begin pre-attn capture");
+            self.pre_attn_graphs[l].begin_capture(&self.stream).expect("begin pre-attn capture");
             unsafe {
                 // RMSNorm
                 dispatch::rmsnorm_bf16(
-                    self.buffers.x.as_ptr() as _,
-                    lw.input_norm,
-                    self.buffers.normed.as_mut_ptr() as _,
-                    1,
-                    h,
-                    eps,
-                    s,
+                    self.buffers.x.as_ptr() as _, lw.input_norm, self.buffers.normed.as_mut_ptr() as _,
+                    1, h, eps, s,
                 );
 
                 // Q projection (GEMV)
                 dispatch::gemv_bf16(
-                    self.buffers.normed.as_ptr() as _,
-                    lw.q_proj_wt,
-                    self.buffers.q_proj.as_mut_ptr() as _,
+                    self.buffers.normed.as_ptr() as _, lw.q_proj_wt, self.buffers.q_proj.as_mut_ptr() as _,
                     self.buffers.fp32_q.as_mut_ptr() as _,
-                    h,
-                    nh * hd,
-                    s,
+                    h, nh * hd, s,
                 );
 
                 // K projection (GEMV)
                 dispatch::gemv_bf16(
-                    self.buffers.normed.as_ptr() as _,
-                    lw.k_proj_wt,
-                    self.buffers.k_proj.as_mut_ptr() as _,
+                    self.buffers.normed.as_ptr() as _, lw.k_proj_wt, self.buffers.k_proj.as_mut_ptr() as _,
                     self.buffers.fp32_kv.as_mut_ptr() as _,
-                    h,
-                    nkv * hd,
-                    s,
+                    h, nkv * hd, s,
                 );
 
                 // V projection (GEMV)
                 dispatch::gemv_bf16(
-                    self.buffers.normed.as_ptr() as _,
-                    lw.v_proj_wt,
-                    self.buffers.v_proj.as_mut_ptr() as _,
+                    self.buffers.normed.as_ptr() as _, lw.v_proj_wt, self.buffers.v_proj.as_mut_ptr() as _,
                     self.buffers.fp32_kv.as_mut_ptr() as _,
-                    h,
-                    nkv * hd,
-                    s,
+                    h, nkv * hd, s,
                 );
 
                 // Reshape heads: [1, H*D] -> [1, H, 1, D]
-                dispatch::reshape_heads_bf16(
-                    self.buffers.q_proj.as_ptr() as _,
-                    self.buffers.q_reshaped.as_mut_ptr() as _,
-                    1,
-                    nh,
-                    hd,
-                    s,
-                );
-                dispatch::reshape_heads_bf16(
-                    self.buffers.k_proj.as_ptr() as _,
-                    self.buffers.k_reshaped.as_mut_ptr() as _,
-                    1,
-                    nkv,
-                    hd,
-                    s,
-                );
-                dispatch::reshape_heads_bf16(
-                    self.buffers.v_proj.as_ptr() as _,
-                    self.buffers.v_reshaped.as_mut_ptr() as _,
-                    1,
-                    nkv,
-                    hd,
-                    s,
-                );
+                dispatch::reshape_heads_bf16(self.buffers.q_proj.as_ptr() as _, self.buffers.q_reshaped.as_mut_ptr() as _, 1, nh, hd, s);
+                dispatch::reshape_heads_bf16(self.buffers.k_proj.as_ptr() as _, self.buffers.k_reshaped.as_mut_ptr() as _, 1, nkv, hd, s);
+                dispatch::reshape_heads_bf16(self.buffers.v_proj.as_ptr() as _, self.buffers.v_reshaped.as_mut_ptr() as _, 1, nkv, hd, s);
 
                 // QK norm (head-level rmsnorm: treat [1,H,1,D] as [H, D])
-                dispatch::rmsnorm_bf16(
-                    self.buffers.q_reshaped.as_ptr() as _,
-                    lw.q_norm,
-                    self.buffers.q_normed.as_mut_ptr() as _,
-                    nh,
-                    hd,
-                    eps,
-                    s,
-                );
-                dispatch::rmsnorm_bf16(
-                    self.buffers.k_reshaped.as_ptr() as _,
-                    lw.k_norm,
-                    self.buffers.k_normed.as_mut_ptr() as _,
-                    nkv,
-                    hd,
-                    eps,
-                    s,
-                );
+                dispatch::rmsnorm_bf16(self.buffers.q_reshaped.as_ptr() as _, lw.q_norm, self.buffers.q_normed.as_mut_ptr() as _, nh, hd, eps, s);
+                dispatch::rmsnorm_bf16(self.buffers.k_reshaped.as_ptr() as _, lw.k_norm, self.buffers.k_normed.as_mut_ptr() as _, nkv, hd, eps, s);
 
                 // Transpose for RoPE: [1,H,1,D] -> [1,H,D]
-                dispatch::transpose_hsd_to_shd_bf16(
-                    self.buffers.q_normed.as_ptr() as _,
-                    self.buffers.q_rope.as_mut_ptr() as _,
-                    1,
-                    nh,
-                    hd,
-                    s,
-                );
-                dispatch::transpose_hsd_to_shd_bf16(
-                    self.buffers.k_normed.as_ptr() as _,
-                    self.buffers.k_rope.as_mut_ptr() as _,
-                    1,
-                    nkv,
-                    hd,
-                    s,
-                );
+                dispatch::transpose_hsd_to_shd_bf16(self.buffers.q_normed.as_ptr() as _, self.buffers.q_rope.as_mut_ptr() as _, 1, nh, hd, s);
+                dispatch::transpose_hsd_to_shd_bf16(self.buffers.k_normed.as_ptr() as _, self.buffers.k_rope.as_mut_ptr() as _, 1, nkv, hd, s);
 
                 // RoPE (in-place, reads position_gpu)
-                dispatch::rope_bf16(
-                    self.buffers.q_rope.as_mut_ptr() as _,
-                    rope_cos,
-                    rope_sin,
-                    self.buffers.position_gpu.as_ptr() as _,
-                    1,
-                    nh,
-                    hd,
-                    s,
-                );
-                dispatch::rope_bf16(
-                    self.buffers.k_rope.as_mut_ptr() as _,
-                    rope_cos,
-                    rope_sin,
-                    self.buffers.position_gpu.as_ptr() as _,
-                    1,
-                    nkv,
-                    hd,
-                    s,
-                );
+                dispatch::rope_bf16(self.buffers.q_rope.as_mut_ptr() as _, rope_cos, rope_sin, self.buffers.position_gpu.as_ptr() as _, 1, nh, hd, s);
+                dispatch::rope_bf16(self.buffers.k_rope.as_mut_ptr() as _, rope_cos, rope_sin, self.buffers.position_gpu.as_ptr() as _, 1, nkv, hd, s);
 
                 // Transpose back: [1,H,D] -> [1,H,1,D]
-                dispatch::transpose_shd_to_hsd_bf16(
-                    self.buffers.q_rope.as_ptr() as _,
-                    self.buffers.q_final.as_mut_ptr() as _,
-                    1,
-                    nh,
-                    hd,
-                    s,
-                );
-                dispatch::transpose_shd_to_hsd_bf16(
-                    self.buffers.k_rope.as_ptr() as _,
-                    self.buffers.k_final.as_mut_ptr() as _,
-                    1,
-                    nkv,
-                    hd,
-                    s,
-                );
+                dispatch::transpose_shd_to_hsd_bf16(self.buffers.q_rope.as_ptr() as _, self.buffers.q_final.as_mut_ptr() as _, 1, nh, hd, s);
+                dispatch::transpose_shd_to_hsd_bf16(self.buffers.k_rope.as_ptr() as _, self.buffers.k_final.as_mut_ptr() as _, 1, nkv, hd, s);
             }
-            self.pre_attn_graphs[l]
-                .end_capture(&self.stream)
-                .expect("end pre-attn capture");
+            self.pre_attn_graphs[l].end_capture(&self.stream).expect("end pre-attn capture");
 
             // === Post-attention graph ===
-            self.post_attn_graphs[l]
-                .begin_capture(&self.stream)
-                .expect("begin post-attn capture");
+            self.post_attn_graphs[l].begin_capture(&self.stream).expect("begin post-attn capture");
             unsafe {
                 // Merge heads: [1,H,1,D] -> [1, hidden]
                 // attn_out is written by ungraphed attention
-                dispatch::merge_heads_bf16(
-                    self.buffers.attn_out.as_ptr() as _,
-                    self.buffers.attn_merged.as_mut_ptr() as _,
-                    1,
-                    nh,
-                    hd,
-                    s,
-                );
+                dispatch::merge_heads_bf16(self.buffers.attn_out.as_ptr() as _, self.buffers.attn_merged.as_mut_ptr() as _, 1, nh, hd, s);
 
                 // O projection
                 dispatch::gemv_bf16(
-                    self.buffers.attn_merged.as_ptr() as _,
-                    lw.o_proj_wt,
-                    self.buffers.o_proj.as_mut_ptr() as _,
+                    self.buffers.attn_merged.as_ptr() as _, lw.o_proj_wt, self.buffers.o_proj.as_mut_ptr() as _,
                     self.buffers.fp32_hidden.as_mut_ptr() as _,
-                    nh * hd,
-                    h,
-                    s,
+                    nh * hd, h, s,
                 );
 
                 // Fused Add+RMSNorm: normed2 = rmsnorm(o_proj + x), sum_out = o_proj + x
                 dispatch::add_rmsnorm_bf16(
-                    self.buffers.o_proj.as_ptr() as _,
-                    self.buffers.x.as_ptr() as _,
-                    lw.post_norm,
-                    self.buffers.normed2.as_mut_ptr() as _,
-                    self.buffers.sum_out.as_mut_ptr() as _,
-                    1,
-                    h,
-                    eps,
-                    s,
+                    self.buffers.o_proj.as_ptr() as _, self.buffers.x.as_ptr() as _, lw.post_norm,
+                    self.buffers.normed2.as_mut_ptr() as _, self.buffers.sum_out.as_mut_ptr() as _,
+                    1, h, eps, s,
                 );
 
                 // Gate projection
                 dispatch::gemv_bf16(
-                    self.buffers.normed2.as_ptr() as _,
-                    lw.gate_proj_wt,
-                    self.buffers.gate.as_mut_ptr() as _,
+                    self.buffers.normed2.as_ptr() as _, lw.gate_proj_wt, self.buffers.gate.as_mut_ptr() as _,
                     self.buffers.fp32_intermediate.as_mut_ptr() as _,
-                    h,
-                    inter,
-                    s,
+                    h, inter, s,
                 );
 
                 // Up projection
                 dispatch::gemv_bf16(
-                    self.buffers.normed2.as_ptr() as _,
-                    lw.up_proj_wt,
-                    self.buffers.up.as_mut_ptr() as _,
+                    self.buffers.normed2.as_ptr() as _, lw.up_proj_wt, self.buffers.up.as_mut_ptr() as _,
                     self.buffers.fp32_intermediate.as_mut_ptr() as _,
-                    h,
-                    inter,
-                    s,
+                    h, inter, s,
                 );
 
                 // Fused SiLU x Mul
-                dispatch::silu_mul_bf16(
-                    self.buffers.gate.as_ptr() as _,
-                    self.buffers.up.as_ptr() as _,
-                    self.buffers.silu_out.as_mut_ptr() as _,
-                    inter,
-                    s,
-                );
+                dispatch::silu_mul_bf16(self.buffers.gate.as_ptr() as _, self.buffers.up.as_ptr() as _, self.buffers.silu_out.as_mut_ptr() as _, inter, s);
 
                 // Down projection
                 dispatch::gemv_bf16(
-                    self.buffers.silu_out.as_ptr() as _,
-                    lw.down_proj_wt,
-                    self.buffers.down.as_mut_ptr() as _,
+                    self.buffers.silu_out.as_ptr() as _, lw.down_proj_wt, self.buffers.down.as_mut_ptr() as _,
                     self.buffers.fp32_hidden.as_mut_ptr() as _,
-                    inter,
-                    h,
-                    s,
+                    inter, h, s,
                 );
 
                 // x = sum_out + down (residual connection for next layer)
-                dispatch::add_bf16(
-                    self.buffers.sum_out.as_ptr() as _,
-                    self.buffers.down.as_ptr() as _,
-                    self.buffers.x.as_mut_ptr() as _,
-                    h,
-                    s,
-                );
+                dispatch::add_bf16(self.buffers.sum_out.as_ptr() as _, self.buffers.down.as_ptr() as _, self.buffers.x.as_mut_ptr() as _, h, s);
             }
-            self.post_attn_graphs[l]
-                .end_capture(&self.stream)
-                .expect("end post-attn capture");
+            self.post_attn_graphs[l].end_capture(&self.stream).expect("end post-attn capture");
         }
 
         // === Final graph: norm + lm_head ===
-        self.final_graph
-            .begin_capture(&self.stream)
-            .expect("begin final capture");
+        self.final_graph.begin_capture(&self.stream).expect("begin final capture");
         unsafe {
-            dispatch::rmsnorm_bf16(
-                self.buffers.x.as_ptr() as _,
-                norm_weight,
-                self.buffers.normed.as_mut_ptr() as _,
-                1,
-                h,
-                eps,
-                s,
-            );
+            dispatch::rmsnorm_bf16(self.buffers.x.as_ptr() as _, norm_weight, self.buffers.normed.as_mut_ptr() as _, 1, h, eps, s);
             dispatch::gemv_bf16(
-                self.buffers.normed.as_ptr() as _,
-                lm_head_wt,
-                self.buffers.logits.as_mut_ptr() as _,
+                self.buffers.normed.as_ptr() as _, lm_head_wt, self.buffers.logits.as_mut_ptr() as _,
                 self.buffers.fp32_vocab.as_mut_ptr() as _,
-                h,
-                vocab,
-                s,
+                h, vocab, s,
             );
         }
-        self.final_graph
-            .end_capture(&self.stream)
-            .expect("end final capture");
+        self.final_graph.end_capture(&self.stream).expect("end final capture");
 
         // Reset cuBLAS back to null stream
-        unsafe {
-            dispatch::set_cublas_stream(cublas, std::ptr::null_mut());
-        }
+        unsafe { dispatch::set_cublas_stream(cublas, std::ptr::null_mut()); }
 
         self.captured = true;
     }
@@ -515,14 +343,8 @@ impl DecodeGraphState {
         let es = 2usize; // BF16
 
         // Upload token ID and position to fixed GPU buffers
-        self.buffers
-            .token_id_gpu
-            .copy_from_host(&token_id.to_le_bytes())
-            .unwrap();
-        self.buffers
-            .position_gpu
-            .copy_from_host(&position.to_le_bytes())
-            .unwrap();
+        self.buffers.token_id_gpu.copy_from_host(&token_id.to_le_bytes()).unwrap();
+        self.buffers.position_gpu.copy_from_host(&position.to_le_bytes()).unwrap();
 
         // Embedding (outside graph since token_id changes each step)
         unsafe {
@@ -530,18 +352,13 @@ impl DecodeGraphState {
                 embed_table,
                 self.buffers.token_id_gpu.as_ptr() as _,
                 self.buffers.x.as_mut_ptr() as _,
-                1,
-                hidden_size,
-                vocab_size,
-                s,
+                1, hidden_size, vocab_size, s,
             );
         }
 
         for l in 0..self.num_layers {
             // Pre-attention graph (norm + QKV + reshape + QK-norm + RoPE)
-            self.pre_attn_graphs[l]
-                .launch(&self.stream)
-                .expect("launch pre-attn graph");
+            self.pre_attn_graphs[l].launch(&self.stream).expect("launch pre-attn graph");
 
             // Ungraphed: KV cache append
             // k_final shape: [1, num_kv_heads, 1, head_dim] (after RoPE pipeline)
@@ -585,13 +402,9 @@ impl DecodeGraphState {
                     k_full.data_ptr() as _,
                     v_full.data_ptr() as _,
                     self.buffers.attn_out.as_mut_ptr() as _,
-                    1,
-                    nh as i32,
-                    nkv as i32,
-                    kv_len,
-                    hd as i32,
-                    scale,
-                    s,
+                    1, nh as i32, nkv as i32,
+                    kv_len, hd as i32,
+                    scale, s,
                 );
             }
 
@@ -599,15 +412,11 @@ impl DecodeGraphState {
             self.stream.synchronize().expect("sync before post-attn");
 
             // Post-attention graph (merge + O-proj + add_rmsnorm + FFN + residual)
-            self.post_attn_graphs[l]
-                .launch(&self.stream)
-                .expect("launch post-attn graph");
+            self.post_attn_graphs[l].launch(&self.stream).expect("launch post-attn graph");
         }
 
         // Final graph (norm + lm_head)
-        self.final_graph
-            .launch(&self.stream)
-            .expect("launch final graph");
+        self.final_graph.launch(&self.stream).expect("launch final graph");
 
         // Sync to ensure logits are ready
         self.stream.synchronize().expect("sync after decode");

crates/xserv-model/src/eagle3.rs

@@ -1,425 +0,0 @@
-//! EAGLE3 speculative draft head for Qwen3-8B (Phase 25).
-//!
-//! Loads the AngelSlim/Qwen3-8B_eagle3 pytorch_model.bin and provides a
-//! single-step forward pass that takes 3 target hidden states + the previous
-//! token and returns a draft token in the target vocabulary.
-//!
-//! Architecture (from weights):
-//! - fc:           [hidden, 3*hidden] → fuse 3 target hidden states
-//! - midlayer:     1 decoder layer (attn input dim = 2*hidden)
-//! - norm + lm_head: → [draft_vocab_size=32000]
-//! - d2t:          draft_id → target_id offset mapping
-
-use std::collections::HashMap;
-use std::path::Path;
-use xserv_kernels::*;
-use xserv_tensor::{DType, Device, Tensor};
-
-/// Target layers to hook for EAGLE3 auxiliary hidden states, for Qwen3-8B
-/// (36 layers). Value comes from AngelSlim/vLLM speculators training config
-/// `dflash_qwen3_8b_sharegpt_online_5k.sh` which specifies target_layer_ids
-/// = "2 18 33". Must match training-time selection or EAGLE outputs are wrong.
-pub const EAGLE_HOOK_LAYERS: [usize; 3] = [2, 18, 33];
-const DRAFT_VOCAB_SIZE: usize = 32000;
-
-fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
-    assert_eq!(a.ndim(), 2);
-    assert_eq!(b.ndim(), 2);
-    matmul(a, b, GemmBackend::CuBlas)
-}
-
-pub struct Eagle3Head {
-    fc_wt: Tensor,                    // [hidden, 3*hidden] transposed for matmul
-    hidden_norm: Tensor,              // [hidden]
-    input_layernorm: Tensor,          // [hidden]
-    q_proj_wt: Tensor,                // [num_heads*head_dim, 2*hidden]
-    k_proj_wt: Tensor,                // [num_kv_heads*head_dim, 2*hidden]
-    v_proj_wt: Tensor,                // [num_kv_heads*head_dim, 2*hidden]
-    o_proj_wt: Tensor,                // [hidden, num_heads*head_dim]
-    gate_proj_wt: Tensor,             // [intermediate, hidden]
-    up_proj_wt: Tensor,               // [intermediate, hidden]
-    down_proj_wt: Tensor,             // [hidden, intermediate]
-    post_attention_layernorm: Tensor, // [hidden]
-    norm: Tensor,                     // [hidden] final
-    lm_head_wt: Tensor,               // [draft_vocab, hidden]
-    d2t: Vec<i64>,                    // [draft_vocab] offset mapping
-    /// t2d[target_id] = true iff target_id has a corresponding draft-vocab id
-    /// (i.e. can potentially be produced by EAGLE). Used to measure the
-    /// coverage cap on acceptance.
-    t2d: Vec<bool>,
-    hidden_size: usize,
-    num_heads: usize,
-    num_kv_heads: usize,
-    head_dim: usize,
-    max_seq_len: usize,
-    rope_cache: RopeCache,
-    // Stateful 1-layer KV cache: [1, num_kv_heads, max_seq_len, head_dim] BF16.
-    // We slice `..current_len` for attention. The head is tiny (~64 KB per
-    // 1000 tokens) so pre-allocating max_seq_len wastes negligible memory.
-    k_cache: Tensor,
-    v_cache: Tensor,
-    current_len: usize,
-}
-
-impl Eagle3Head {
-    pub fn load(dir: &Path, device: u32) -> Self {
-        let (weights, d2t, t2d) = load_eagle3_weights(dir, device);
-        let hidden_size = 4096;
-        let num_heads = 32;
-        let num_kv_heads = 8;
-        let head_dim = 128;
-        let intermediate_size = 12288;
-        let max_seq_len = 2048;
-        let rope_theta = 1_000_000.0f32;
-
-        let get = |name: &str| -> Tensor {
-            weights
-                .get(name)
-                .unwrap_or_else(|| panic!("missing eagle3 weight: {name}"))
-                .clone()
-        };
-
-        let fc_wt = get("fc.weight").transpose(0, 1).contiguous();
-        let q_proj_wt = get("midlayer.self_attn.q_proj.weight")
-            .transpose(0, 1)
-            .contiguous();
-        let k_proj_wt = get("midlayer.self_attn.k_proj.weight")
-            .transpose(0, 1)
-            .contiguous();
-        let v_proj_wt = get("midlayer.self_attn.v_proj.weight")
-            .transpose(0, 1)
-            .contiguous();
-        let o_proj_wt = get("midlayer.self_attn.o_proj.weight")
-            .transpose(0, 1)
-            .contiguous();
-        let gate_proj_wt = get("midlayer.mlp.gate_proj.weight")
-            .transpose(0, 1)
-            .contiguous();
-        let up_proj_wt = get("midlayer.mlp.up_proj.weight")
-            .transpose(0, 1)
-            .contiguous();
-        let down_proj_wt = get("midlayer.mlp.down_proj.weight")
-            .transpose(0, 1)
-            .contiguous();
-        let hidden_norm = get("midlayer.hidden_norm.weight");
-        let input_layernorm = get("midlayer.input_layernorm.weight");
-        let post_attention_layernorm = get("midlayer.post_attention_layernorm.weight");
-        let norm = get("norm.weight");
-        let lm_head_wt = get("lm_head.weight").transpose(0, 1).contiguous();
-
-        assert_eq!(d2t.len(), DRAFT_VOCAB_SIZE);
-
-        let rope_cache = RopeCache::new(max_seq_len, head_dim, rope_theta);
-
-        let k_cache = Tensor::zeros(
-            &[1, num_kv_heads, max_seq_len, head_dim],
-            DType::BF16,
-            Device::Cuda(device),
-        );
-        let v_cache = Tensor::zeros(
-            &[1, num_kv_heads, max_seq_len, head_dim],
-            DType::BF16,
-            Device::Cuda(device),
-        );
-
-        Self {
-            fc_wt,
-            hidden_norm,
-            input_layernorm,
-            q_proj_wt,
-            k_proj_wt,
-            v_proj_wt,
-            o_proj_wt,
-            gate_proj_wt,
-            up_proj_wt,
-            down_proj_wt,
-            post_attention_layernorm,
-            norm,
-            lm_head_wt,
-            d2t,
-            t2d,
-            hidden_size,
-            num_heads,
-            num_kv_heads,
-            head_dim,
-            max_seq_len,
-            rope_cache,
-            k_cache,
-            v_cache,
-            current_len: 0,
-        }
-    }
-
-    /// Reset the internal KV cache for a fresh sequence.
-    pub fn reset(&mut self) {
-        self.current_len = 0;
-    }
-
-    /// Truncate the internal KV cache to `new_len` entries. Used to discard
-    /// K/V of rejected drafts after a speculative round.
-    pub fn truncate_to(&mut self, new_len: usize) {
-        assert!(new_len <= self.current_len);
-        self.current_len = new_len;
-    }
-
-    /// Current number of committed K/V entries in the internal EAGLE cache.
-    pub fn current_len(&self) -> usize {
-        self.current_len
-    }
-
-    /// One draft step: produce a token in target vocabulary space.
-    ///
-    /// - `target_hidden`: 3 tensors [1, hidden_size] from target hook layers
-    /// - `embed_table`: the target model's embed_tokens (shared, not copied)
-    /// - `prev_token`: the previous committed token
-    /// - `position`: the decode position for RoPE
-    ///
-    /// Returns (draft_token_in_target_vocab, draft_logits_tensor).
-    pub fn step(
-        &mut self,
-        target_hidden: &[Tensor; 3],
-        embed_table: &Tensor,
-        prev_token: u32,
-        position: usize,
-    ) -> (u32, Tensor) {
-        let (id, logits, _) = self.step_with_aux(target_hidden, embed_table, prev_token, position);
-        (id, logits)
-    }
-
-    /// Like `step`, but also returns the final hidden state (aux) usable as
-    /// the fused_h for a subsequent recursive draft step via `step_recursive`.
-    pub fn step_with_aux(
-        &mut self,
-        target_hidden: &[Tensor; 3],
-        embed_table: &Tensor,
-        prev_token: u32,
-        position: usize,
-    ) -> (u32, Tensor, Tensor) {
-        // Fuse 3 target hidden states into fused_h via fc.
-        let h_cat = concat_hidden(target_hidden);
-        let fused_h = matmul_2d(&h_cat, &self.fc_wt);
-        self.forward_from_fused(fused_h, embed_table, prev_token, position)
-    }
-
-    /// Recursive draft step: reuses the previous EAGLE step's aux as fused_h,
-    /// bypassing the fc+3-hidden fusion. Used for γ≥2 chained drafts.
-    pub fn step_recursive(
-        &mut self,
-        fused_h: Tensor,
-        embed_table: &Tensor,
-        prev_token: u32,
-        position: usize,
-    ) -> (u32, Tensor, Tensor) {
-        self.forward_from_fused(fused_h, embed_table, prev_token, position)
-    }
-
-    fn forward_from_fused(
-        &mut self,
-        fused_h: Tensor,
-        embed_table: &Tensor,
-        prev_token: u32,
-        position: usize,
-    ) -> (u32, Tensor, Tensor) {
-        let eps = 1e-6f32;
-        assert!(
-            self.current_len < self.max_seq_len,
-            "EAGLE KV cache overflow: {} >= {}",
-            self.current_len,
-            self.max_seq_len
-        );
-
-        let emb = embedding(embed_table, &[prev_token]);
-        let residual = fused_h.clone();
-        let emb_normed = rmsnorm(&emb, &self.input_layernorm, eps);
-        let h_normed = rmsnorm(&fused_h, &self.hidden_norm, eps);
-        let attn_in = concat_last_dim(&emb_normed, &h_normed);
-
-        let q = matmul_2d(&attn_in, &self.q_proj_wt);
-        let k = matmul_2d(&attn_in, &self.k_proj_wt);
-        let v = matmul_2d(&attn_in, &self.v_proj_wt);
-
-        let q_3d = q.reshape(&[1, self.num_heads, self.head_dim]);
-        let k_3d = k.reshape(&[1, self.num_kv_heads, self.head_dim]);
-        let positions = [position as u32];
-        rope_inplace(&q_3d, &self.rope_cache, &positions);
-        rope_inplace(&k_3d, &self.rope_cache, &positions);
-
-        let v_3d = v.reshape(&[1, self.num_kv_heads, self.head_dim]);
-        self.append_to_kv_cache(&k_3d, &v_3d);
-        self.current_len += 1;
-        let kv_len = self.current_len;
-        let k_view = self.k_cache.narrow(2, 0, kv_len).contiguous();
-        let v_view = self.v_cache.narrow(2, 0, kv_len).contiguous();
-
-        let q_4d = q_3d.reshape(&[1, self.num_heads, 1, self.head_dim]);
-        let attn_out = decode_attention(&q_4d, &k_view, &v_view);
-
-        let attn_merged = attn_out.reshape(&[1, self.num_heads * self.head_dim]);
-        let attn_proj = matmul_2d(&attn_merged, &self.o_proj_wt);
-
-        let (mlp_in, residual) =
-            add_rmsnorm(&attn_proj, &residual, &self.post_attention_layernorm, eps);
-
-        let gate = matmul_2d(&mlp_in, &self.gate_proj_wt);
-        let up = matmul_2d(&mlp_in, &self.up_proj_wt);
-        let hidden = silu_mul(&gate, &up);
-        let down = matmul_2d(&hidden, &self.down_proj_wt);
-
-        let (x, prenorm) = add_rmsnorm(&down, &residual, &self.norm, eps);
-        let logits = matmul_2d(&x, &self.lm_head_wt);
-
-        let draft_id = argmax_bf16_single(&logits);
-        let target_id = (draft_id as i64 + self.d2t[draft_id as usize]) as u32;
-        // aux for recursive drafting = PRE-norm hidden (default norm_output=False
-        // in vllm/llama_eagle3.py). Feeding the pre-norm state matches training.
-        (target_id, logits, prenorm)
-    }
-
-    /// Write new K/V rows (shape [1, num_kv_heads, head_dim]) at position
-    /// `current_len` inside the [1, num_kv_heads, max_seq_len, head_dim] cache.
-    fn append_to_kv_cache(&mut self, new_k: &Tensor, new_v: &Tensor) {
-        let head_bytes = self.head_dim * self.k_cache.dtype().size_bytes();
-        for h in 0..self.num_kv_heads {
-            for (cache, src) in [(&self.k_cache, new_k), (&self.v_cache, new_v)] {
-                let dst = unsafe {
-                    (cache.data_ptr() as *mut u8)
-                        .add(((h * self.max_seq_len) + self.current_len) * head_bytes)
-                };
-                let s = unsafe { (src.data_ptr() as *const u8).add(h * head_bytes) };
-                d2d(dst, s, head_bytes);
-            }
-        }
-    }
-
-    /// Map a draft-vocab token id to the full target-vocab id via d2t.
-    pub fn map_draft_to_target(&self, draft_id: u32) -> u32 {
-        (draft_id as i64 + self.d2t[draft_id as usize]) as u32
-    }
-
-    /// Returns true iff `target_id` is representable in the draft vocabulary
-    /// (i.e., EAGLE could in principle produce it).
-    pub fn target_id_in_draft_vocab(&self, target_id: u32) -> bool {
-        self.t2d.get(target_id as usize).copied().unwrap_or(false)
-    }
-}
-
-fn d2d(dst: *mut u8, src: *const u8, bytes: usize) {
-    unsafe {
-        xserv_cuda::ffi::cudaMemcpy(dst, src, bytes, xserv_cuda::ffi::CUDA_MEMCPY_D2D);
-    }
-}
-
-fn concat_hidden(hidden: &[Tensor; 3]) -> Tensor {
-    let h = hidden[0].shape()[1];
-    let dtype = hidden[0].dtype();
-    let device = hidden[0].device();
-    let elem_bytes = dtype.size_bytes();
-    let out = Tensor::empty(&[1, 3 * h], dtype, device);
-    for (i, t) in hidden.iter().enumerate() {
-        assert!(t.is_contiguous());
-        let dst = unsafe { (out.data_ptr() as *mut u8).add(i * h * elem_bytes) };
-        d2d(dst, t.data_ptr() as *const u8, h * elem_bytes);
-    }
-    out
-}
-
-fn concat_last_dim(a: &Tensor, b: &Tensor) -> Tensor {
-    let da = a.shape()[1];
-    let db = b.shape()[1];
-    let dtype = a.dtype();
-    let device = a.device();
-    let elem_bytes = dtype.size_bytes();
-    let out = Tensor::empty(&[1, da + db], dtype, device);
-    d2d(
-        out.data_ptr() as *mut u8,
-        a.data_ptr() as *const u8,
-        da * elem_bytes,
-    );
-    let dst = unsafe { (out.data_ptr() as *mut u8).add(da * elem_bytes) };
-    d2d(dst, b.data_ptr() as *const u8, db * elem_bytes);
-    out
-}
-
-fn repeat_kv_for_single_token(kv: &Tensor, repeats: usize) -> Tensor {
-    if repeats == 1 {
-        return kv.clone();
-    }
-    let nkv = kv.shape()[1];
-    let d = kv.shape()[2];
-    let dtype = kv.dtype();
-    let device = kv.device();
-    let head_bytes = d * dtype.size_bytes();
-    let out = Tensor::empty(&[1, nkv * repeats, d], dtype, device);
-    for h in 0..nkv {
-        let src = unsafe { (kv.data_ptr() as *const u8).add(h * head_bytes) };
-        for r in 0..repeats {
-            let dst = unsafe { (out.data_ptr() as *mut u8).add((h * repeats + r) * head_bytes) };
-            d2d(dst, src, head_bytes);
-        }
-    }
-    out
-}
-
-/// Load EAGLE3 weights from safetensors, handling int64 d2t + bool t2d specially.
-fn load_eagle3_weights(dir: &Path, device: u32) -> (HashMap<String, Tensor>, Vec<i64>, Vec<bool>) {
-    let st_path = dir.join("model.safetensors");
-    assert!(
-        st_path.exists(),
-        "Eagle3 model.safetensors not found in {}. Convert with:\n\
-         python3 -c \"import torch; from safetensors.torch import save_file; \
-         sd=torch.load('pytorch_model.bin', map_location='cpu', weights_only=False); \
-         save_file(sd, 'model.safetensors')\"",
-        dir.display()
-    );
-
-    let data = std::fs::read(&st_path)
-        .unwrap_or_else(|e| panic!("failed to read {}: {e}", st_path.display()));
-    let st = safetensors::SafeTensors::deserialize(&data)
-        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", st_path.display()));
-
-    let mut tensors = HashMap::new();
-    let mut d2t_vec: Vec<i64> = Vec::new();
-    let mut t2d_vec: Vec<bool> = Vec::new();
-
-    for (name, view) in st.tensors() {
-        if name == "t2d" {
-            let raw = view.data();
-            assert_eq!(view.dtype(), safetensors::Dtype::BOOL);
-            t2d_vec = raw.iter().map(|&b| b != 0).collect();
-            continue;
-        }
-        if name == "d2t" {
-            let raw = view.data();
-            assert_eq!(view.dtype(), safetensors::Dtype::I64);
-            let n = raw.len() / 8;
-            d2t_vec = (0..n)
-                .map(|i| i64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap()))
-                .collect();
-            continue;
-        }
-        let dtype = match view.dtype() {
-            safetensors::Dtype::BF16 => DType::BF16,
-            safetensors::Dtype::F32 => DType::F32,
-            safetensors::Dtype::F16 => DType::F16,
-            other => {
-                eprintln!("eagle3: skipping {name} with unsupported dtype {other:?}");
-                continue;
-            }
-        };
-        let shape: Vec<usize> = view.shape().to_vec();
-        let raw = view.data();
-        let t = crate::loader::make_tensor(raw, &shape, dtype);
-        let t = t.to_device(Device::Cuda(device));
-        tensors.insert(name.to_string(), t);
-    }
-
-    assert!(
-        !d2t_vec.is_empty(),
-        "d2t tensor not found in eagle3 weights"
-    );
-    assert!(
-        !t2d_vec.is_empty(),
-        "t2d tensor not found in eagle3 weights"
-    );
-    (tensors, d2t_vec, t2d_vec)
-}

crates/xserv-model/src/gpt2.rs

@@ -31,7 +31,7 @@ struct GPT2Block {
 
 pub struct KVCache {
     // Per layer, per head: raw bytes (works for both f32 and bf16)
-    k: Vec<Vec<Vec<u8>>>, // [num_layers][num_heads][seq_len * head_dim * elem_size]
+    k: Vec<Vec<Vec<u8>>>,  // [num_layers][num_heads][seq_len * head_dim * elem_size]
     v: Vec<Vec<Vec<u8>>>,
     len: usize,
     num_heads: usize,
@@ -42,13 +42,7 @@ pub struct KVCache {
 }
 
 impl KVCache {
-    pub fn new(
-        num_layers: usize,
-        num_heads: usize,
-        head_dim: usize,
-        dtype: DType,
-        device: Device,
-    ) -> Self {
+    pub fn new(num_layers: usize, num_heads: usize, head_dim: usize, dtype: DType, device: Device) -> Self {
         Self {
             k: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
             v: (0..num_layers).map(|_| vec![vec![]; num_heads]).collect(),
@@ -61,18 +55,10 @@ impl KVCache {
         }
     }
 
-    pub fn seq_len(&self) -> usize {
-        self.len
-    }
+    pub fn seq_len(&self) -> usize { self.len }
 
     /// Append from a CPU tensor with shape [1, H, new_tokens, D].
-    pub fn append_kv_tensor(
-        &mut self,
-        layer: usize,
-        k_cpu: &Tensor,
-        v_cpu: &Tensor,
-        new_tokens: usize,
-    ) {
+    pub fn append_kv_tensor(&mut self, layer: usize, k_cpu: &Tensor, v_cpu: &Tensor, new_tokens: usize) {
         let hd = self.head_dim;
         let es = self.elem_size;
         let k_bytes = k_cpu.storage().as_cpu_bytes();
@@ -132,8 +118,7 @@ impl GPT2 {
     pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
         crate::init_kernels();
         let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
-            w.remove(name)
-                .unwrap_or_else(|| panic!("missing weight: {name}"))
+            w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
         };
 
         let wte = take(&mut w, "wte.weight");
@@ -162,15 +147,7 @@ impl GPT2 {
             });
         }
 
-        Self {
-            config,
-            wte,
-            wpe,
-            layers,
-            ln_f_g,
-            ln_f_b,
-            lm_head,
-        }
+        Self { config, wte, wpe, layers, ln_f_g, ln_f_b, lm_head }
     }
 
     /// Full forward pass without KV cache (for testing / correctness comparison).
@@ -202,22 +179,14 @@ impl GPT2 {
         let head_dim = self.config.head_dim();
 
         let tok_emb = embedding(&self.wte, token_ids);
-        let pos_ids: Vec<u32> = (pos_offset..pos_offset + new_tokens)
-            .map(|p| p as u32)
-            .collect();
+        let pos_ids: Vec<u32> = (pos_offset..pos_offset + new_tokens).map(|p| p as u32).collect();
         let pos_emb = embedding(&self.wpe, &pos_ids);
         let mut x = add_tensors(&tok_emb, &pos_emb);
 
         for (layer_idx, layer) in self.layers.iter().enumerate() {
             x = self.transformer_block(
-                layer,
-                &x,
-                Some((cache, layer_idx)),
-                pos_offset,
-                new_tokens,
-                num_heads,
-                head_dim,
-                hidden,
+                layer, &x, Some((cache, layer_idx)),
+                pos_offset, new_tokens, num_heads, head_dim, hidden,
             );
         }
 
@@ -230,7 +199,7 @@ impl GPT2 {
         layer: &GPT2Block,
         x: &Tensor,
         cache: Option<(&mut KVCache, usize)>,
-        _pos_offset: usize,
+        pos_offset: usize,
         new_tokens: usize,
         num_heads: usize,
         head_dim: usize,
@@ -269,11 +238,7 @@ impl GPT2 {
 
 fn linear(x: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Tensor {
     let out = matmul_2d(x, weight);
-    if let Some(b) = bias {
-        add_bias(&out, b)
-    } else {
-        out
-    }
+    if let Some(b) = bias { add_bias(&out, b) } else { out }
 }
 
 fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
@@ -312,12 +277,7 @@ fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
     }
 }
 
-fn split_qkv(
-    qkv: &Tensor,
-    num_heads: usize,
-    head_dim: usize,
-    seq_len: usize,
-) -> (Tensor, Tensor, Tensor) {
+fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) -> (Tensor, Tensor, Tensor) {
     let hidden = num_heads * head_dim;
     let qkv_cpu = qkv.to_device(Device::Cpu);
     let device = qkv.device();
@@ -334,21 +294,14 @@ fn split_qkv(
                 for h in 0..num_heads {
                     let src_off = h * head_dim;
                     let dst_off = (h * seq_len + s) * head_dim;
-                    q_data[dst_off..dst_off + head_dim]
-                        .copy_from_slice(&row[src_off..src_off + head_dim]);
-                    k_data[dst_off..dst_off + head_dim]
-                        .copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
-                    v_data[dst_off..dst_off + head_dim].copy_from_slice(
-                        &row[2 * hidden + src_off..2 * hidden + src_off + head_dim],
-                    );
+                    q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
+                    k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
+                    v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
                 }
             }
-            let q =
-                Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
-            let k =
-                Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
-            let v =
-                Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
+            let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
+            let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
+            let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
             (q, k, v)
         }
         DType::BF16 => {
@@ -361,21 +314,14 @@ fn split_qkv(
                 for h in 0..num_heads {
                     let src_off = h * head_dim;
                     let dst_off = (h * seq_len + s) * head_dim;
-                    q_data[dst_off..dst_off + head_dim]
-                        .copy_from_slice(&row[src_off..src_off + head_dim]);
-                    k_data[dst_off..dst_off + head_dim]
-                        .copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
-                    v_data[dst_off..dst_off + head_dim].copy_from_slice(
-                        &row[2 * hidden + src_off..2 * hidden + src_off + head_dim],
-                    );
+                    q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
+                    k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
+                    v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
                 }
             }
-            let q =
-                Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
-            let k =
-                Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
-            let v =
-                Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
+            let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
+            let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
+            let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
             (q, k, v)
         }
         _ => panic!("unsupported dtype {:?} in split_qkv", dtype),
@@ -397,8 +343,7 @@ fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
                 for h in 0..num_heads {
                     let src_off = (h * seq_len + s) * head_dim;
                     let dst_off = s * hidden + h * head_dim;
-                    out[dst_off..dst_off + head_dim]
-                        .copy_from_slice(&src[src_off..src_off + head_dim]);
+                    out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
                 }
             }
             Tensor::from_slice(&out, &[seq_len, hidden]).to_device(device)
@@ -410,8 +355,7 @@ fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
                 for h in 0..num_heads {
                     let src_off = (h * seq_len + s) * head_dim;
                     let dst_off = s * hidden + h * head_dim;
-                    out[dst_off..dst_off + head_dim]
-                        .copy_from_slice(&src[src_off..src_off + head_dim]);
+                    out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
                 }
             }
             Tensor::from_slice(&out, &[seq_len, hidden]).to_device(device)
@@ -428,8 +372,7 @@ pub fn sample_greedy(logits: &Tensor) -> u32 {
     let vocab_size = logits.shape()[1];
     let seq_len = logits.shape()[0];
     let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
-    last_row
-        .iter()
+    last_row.iter()
         .enumerate()
         .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
         .map(|(idx, _)| idx as u32)

crates/xserv-model/src/gpt_oss_graph.rs

@@ -1,195 +0,0 @@
-//! CUDA-graph replay for gpt-oss batch=1 decode (Phase 21).
-//!
-//! A decode step launches ~200 kernels; with sparse MoE the GPU work is only
-//! a few ms, so launch overhead dominates TPOT. The whole step (embedding →
-//! 24 layers → logits) is captured ONCE into a CUDA graph and replayed per
-//! token with a single `cudaGraphLaunch`.
-//!
-//! Why the existing forward is capturable as-is:
-//! - Every per-step variable input lives in a stable-address device buffer
-//!   whose CONTENTS are updated outside the captured region: token id and
-//!   position (persistent buffers owned here), block table and context lens
-//!   (PagedKVCache GPU buffers, refreshed by `decode_prepare`). The KV scatter
-//!   and paged attention kernels read their write/read positions from those
-//!   buffers, and the sparse-MoE GEMVs read expert ids from `topk_ids` written
-//!   earlier in the same graph — all data-dependent, no host branching.
-//! - Kernel launches go through the thread-local launch stream
-//!   (`xserv_cuda::stream::push_stream`), so the capture stream sees them.
-//! - Intermediate tensors come from the caching allocator. Blocks freed while
-//!   capturing are quarantined (`allocator::begin_retain`) for the graph's
-//!   lifetime so no later allocation can take ownership of memory the graph
-//!   still references on every replay.
-//!
-//! Capture preconditions: at least one EAGER decode step must have run first,
-//! so the allocator pool already holds every bucket size the step needs
-//! (a pool-miss inside capture would call cudaMalloc — illegal while
-//! capturing) and cuBLAS has finished its one-time per-shape setup.
-
-use std::ffi::c_void;
-
-use xserv_cuda::allocator::{self, RetainedBlocks};
-use xserv_cuda::{CudaGraph, CudaStream, GpuBuffer};
-use xserv_tensor::Tensor;
-
-use crate::gpt_oss::GptOss;
-use crate::paged_kv_cache::PagedKVCache;
-
-pub struct GptOssDecodeGraph {
-    stream: CudaStream,
-    graph: CudaGraph,
-    ids_buf: GpuBuffer, // [1] u32, persistent graph input
-    pos_buf: GpuBuffer, // [1] u32, persistent graph input
-    logits: Tensor,     // graph output; rewritten in place by every replay
-    _arena: RetainedBlocks,
-}
-
-impl GptOssDecodeGraph {
-    /// Capture one batch=1 decode step and replay it once (capture records
-    /// without executing, so the replay performs this token's computation).
-    pub fn capture(
-        model: &GptOss,
-        token: u32,
-        position: usize,
-        slot: usize,
-        cache: &mut PagedKVCache,
-    ) -> Self {
-        let stream = CudaStream::new().expect("create capture stream");
-        let mut ids_buf = allocator::cached_alloc(4).expect("alloc ids buf");
-        let mut pos_buf = allocator::cached_alloc(4).expect("alloc pos buf");
-
-        model.decode_prepare(&[position], &[slot], cache);
-        ids_buf.copy_from_host(&token.to_le_bytes()).unwrap();
-        pos_buf
-            .copy_from_host(&(position as u32).to_le_bytes())
-            .unwrap();
-
-        // Retained warmup: run the exact step once eagerly with the quarantine
-        // ON. Freed intermediates are held back instead of recycled, so the
-        // pool ends up stocked with a dedicated block for EVERY allocation the
-        // step performs. The capture below repeats the same allocation
-        // sequence and therefore never misses the pool — a pool miss would
-        // call cudaMalloc, which is illegal while a stream is capturing (this
-        // is also why one block per bucket is not enough: the capture's own
-        // quarantine keeps freed blocks out of reuse). Re-running the step is
-        // idempotent: the KV scatter rewrites the same cache position.
-        allocator::begin_retain();
-        {
-            let _guard = xserv_cuda::push_stream(&stream);
-            let _ = model.decode_core(
-                ids_buf.as_ptr() as *const c_void,
-                pos_buf.as_ptr() as *const c_void,
-                1,
-                cache,
-            );
-        }
-        drop(allocator::end_retain()); // release the warmup blocks to the pool
-        stream.synchronize().expect("warmup sync");
-
-        allocator::begin_retain();
-        let mut graph = CudaGraph::new();
-        let logits;
-        {
-            let _guard = xserv_cuda::stream::push_stream(&stream);
-            graph
-                .begin_capture(&stream)
-                .expect("begin decode-graph capture");
-            logits = model.decode_core(
-                ids_buf.as_ptr() as *const c_void,
-                pos_buf.as_ptr() as *const c_void,
-                1,
-                cache,
-            );
-            graph
-                .end_capture(&stream)
-                .expect("end decode-graph capture");
-        }
-        let arena = allocator::end_retain();
-
-        graph.launch(&stream).expect("first decode-graph replay");
-        cache.advance_seq_len(slot, 1);
-
-        Self {
-            stream,
-            graph,
-            ids_buf,
-            pos_buf,
-            logits,
-            _arena: arena,
-        }
-    }
-
-    /// Run one decode step by replaying the captured graph.
-    pub fn step(
-        &mut self,
-        model: &GptOss,
-        token: u32,
-        position: usize,
-        slot: usize,
-        cache: &mut PagedKVCache,
-    ) -> Tensor {
-        model.decode_prepare(&[position], &[slot], cache);
-        self.ids_buf.copy_from_host(&token.to_le_bytes()).unwrap();
-        self.pos_buf
-            .copy_from_host(&(position as u32).to_le_bytes())
-            .unwrap();
-        self.graph
-            .launch(&self.stream)
-            .expect("decode-graph replay");
-        cache.advance_seq_len(slot, 1);
-        // Shallow clone: the caller reads these logits before the next replay
-        // rewrites the underlying buffer.
-        self.logits.clone()
-    }
-}
-
-/// Lazy capture policy: first decode step of the process runs eager (warms the
-/// allocator pool + cuBLAS so capture performs no "unsafe" CUDA calls), the
-/// second is captured, the rest replay. Batch>1 always falls back to eager.
-/// Disable with XSERV_DECODE_GRAPH=0.
-pub struct GraphedGptOssDecoder {
-    graph: Option<GptOssDecodeGraph>,
-    eager_steps: u32,
-    enabled: bool,
-}
-
-impl GraphedGptOssDecoder {
-    pub fn new() -> Self {
-        let enabled = std::env::var("XSERV_DECODE_GRAPH")
-            .map(|v| v != "0")
-            .unwrap_or(true);
-        Self {
-            graph: None,
-            eager_steps: 0,
-            enabled,
-        }
-    }
-
-    pub fn decode(
-        &mut self,
-        model: &GptOss,
-        tokens: &[u32],
-        positions: &[usize],
-        slots: &[usize],
-        cache: &mut PagedKVCache,
-    ) -> Tensor {
-        if self.enabled && tokens.len() == 1 {
-            if let Some(g) = self.graph.as_mut() {
-                return g.step(model, tokens[0], positions[0], slots[0], cache);
-            }
-            if self.eager_steps >= 1 {
-                let g = GptOssDecodeGraph::capture(model, tokens[0], positions[0], slots[0], cache);
-                let logits = g.logits.clone();
-                self.graph = Some(g);
-                return logits;
-            }
-        }
-        self.eager_steps += 1;
-        model.forward_decode_paged(tokens, positions, slots, cache)
-    }
-}
-
-impl Default for GraphedGptOssDecoder {
-    fn default() -> Self {
-        Self::new()
-    }
-}

crates/xserv-model/src/gptoss.rs

@@ -0,0 +1,416 @@
+//! gpt-oss-20b (MoE) forward pass — Phase 19.
+//!
+//! Correctness-first, in xserv's own style (reuses our kernels; llama.cpp is only
+//! a numerical oracle, not a code source). Differences from Qwen3 handled here:
+//!   - MoE FFN: per-token top-4 router (softmax after top-k) + clamped-SwiGLU experts
+//!   - attention sinks: a per-head learned logit column added to the softmax then
+//!     dropped (so attention probabilities do not sum to 1)
+//!   - alternating sliding-window attention (window from config on flagged layers)
+//!   - q/k/v/o projection biases; head_dim 64; no q/k norm; rotate_half RoPE (θ=150000)
+//!
+//! Weights are loaded from a plain BF16 safetensors dir (MXFP4 experts are
+//! dequantized to BF16 offline by tools/gptoss_dequant.py), so the standard
+//! loader feeds us BF16 tensors and this file needs no quantization code.
+//!
+//! v1 is a self-contained non-paged forward (contiguous KV built per call) used to
+//! validate next-token agreement with llama.cpp. Paged-cache + PP + server wiring
+//! come after numerical correctness is established.
+
+use std::collections::HashMap;
+use half::bf16;
+use xserv_kernels::*;
+use xserv_tensor::{Device, Tensor};
+
+use crate::config::ModelConfig;
+
+pub struct GptOss {
+    pub config: ModelConfig,
+    embed_tokens: Tensor, // [vocab, hidden]
+    layers: Vec<Block>,
+    norm: Tensor,         // [hidden]
+    lm_head_t: Tensor,    // [hidden, vocab] (pre-transposed)
+    rope_cache: RopeCache,
+}
+
+struct Block {
+    input_norm: Tensor,   // [hidden]
+    post_norm: Tensor,    // [hidden]
+    // Attention (weights pre-transposed to [in, out]; biases [out]).
+    q_proj_wt: Tensor,    // [hidden, n_heads*head_dim]
+    q_bias: Tensor,
+    k_proj_wt: Tensor,    // [hidden, n_kv*head_dim]
+    k_bias: Tensor,
+    v_proj_wt: Tensor,
+    v_bias: Tensor,
+    o_proj_wt: Tensor,    // [n_heads*head_dim, hidden]
+    o_bias: Tensor,
+    sinks: Tensor,        // [n_heads] (f32 on host)
+    sliding: bool,
+    // MoE.
+    router_wt: Tensor,    // [hidden, n_experts]
+    router_bias: Tensor,  // [n_experts]
+    gate_up_wt: Vec<Tensor>, // per expert: [hidden, 2*inter]
+    gate_up_bias: Vec<Tensor>, // [2*inter]
+    down_wt: Vec<Tensor>, // per expert: [inter, hidden]
+    down_bias: Vec<Tensor>, // [hidden]
+}
+
+impl GptOss {
+    /// Load gpt-oss from a BF16 (dequantized) HF-format weight map.
+    pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
+        crate::init_kernels();
+        let dev = Device::Cuda(0);
+        let take = |w: &mut HashMap<String, Tensor>, n: &str| -> Tensor {
+            w.remove(n).unwrap_or_else(|| panic!("missing weight: {n}"))
+        };
+        let repl = |t: Tensor| t.to_device(dev);
+        // pre-transpose a [out, in] linear weight to [in, out] for x@wt.
+        let wt = |t: Tensor| t.to_device(dev).transpose(0, 1).contiguous();
+
+        let hidden = config.hidden();
+        let n_experts = config.num_experts();
+        let inter = config.intermediate_size.expect("intermediate_size");
+
+        let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
+        let norm = repl(take(&mut w, "model.norm.weight"));
+        let lm_head_t = wt(take(&mut w, "lm_head.weight"));
+
+        let rope_cache = yarn_rope_cache(&config);
+
+        let n_layers = config.num_layers();
+        let mut layers = Vec::with_capacity(n_layers);
+        for i in 0..n_layers {
+            let p = format!("model.layers.{i}");
+            // Experts are stored fused as [E, in, out]; slice per expert into [in, out].
+            let gate_up = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj")); // [E, hidden, 2*inter]
+            let gate_up_b = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj_bias")); // [E, 2*inter]
+            let down = take(&mut w, &format!("{p}.mlp.experts.down_proj")); // [E, inter, hidden]
+            let down_b = take(&mut w, &format!("{p}.mlp.experts.down_proj_bias")); // [E, hidden]
+            let mut gate_up_wt = Vec::with_capacity(n_experts);
+            let mut gate_up_bias = Vec::with_capacity(n_experts);
+            let mut down_wt = Vec::with_capacity(n_experts);
+            let mut down_bias = Vec::with_capacity(n_experts);
+            // Experts are kept on CPU (the 32 experts per layer total ~36GB for
+            // the whole model, which won't fit one GPU). Each selected expert's
+            // weights (~50MB) are uploaded on demand in expert_forward; only
+            // top-k experts per token are touched, so the H2D traffic is small.
+            for e in 0..n_experts {
+                gate_up_wt.push(slice_expert(&gate_up, e, hidden, 2 * inter)); // CPU
+                gate_up_bias.push(slice_row(&gate_up_b, e, 2 * inter));        // CPU
+                down_wt.push(slice_expert(&down, e, inter, hidden));          // CPU
+                down_bias.push(slice_row(&down_b, e, hidden));               // CPU
+            }
+
+            layers.push(Block {
+                input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
+                post_norm: repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight"))),
+                q_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))),
+                q_bias: repl(take(&mut w, &format!("{p}.self_attn.q_proj.bias"))),
+                k_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.k_proj.weight"))),
+                k_bias: repl(take(&mut w, &format!("{p}.self_attn.k_proj.bias"))),
+                v_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.v_proj.weight"))),
+                v_bias: repl(take(&mut w, &format!("{p}.self_attn.v_proj.bias"))),
+                o_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
+                o_bias: repl(take(&mut w, &format!("{p}.self_attn.o_proj.bias"))),
+                sinks: take(&mut w, &format!("{p}.self_attn.sinks")).to_device(Device::Cpu),
+                sliding: config.layer_uses_sliding_window(i),
+                router_wt: wt(take(&mut w, &format!("{p}.mlp.router.weight"))),
+                router_bias: repl(take(&mut w, &format!("{p}.mlp.router.bias"))),
+                gate_up_wt, gate_up_bias, down_wt, down_bias,
+            });
+        }
+
+        Self { config, embed_tokens, layers, norm, lm_head_t, rope_cache }
+    }
+
+    /// Full prefill forward over `token_ids`; returns logits [seq_len, vocab].
+    pub fn forward(&self, token_ids: &[u32]) -> Tensor {
+        let t = token_ids.len();
+        let hidden = self.config.hidden();
+        let n_heads = self.config.num_heads();
+        let n_kv = self.config.num_kv_heads();
+        let head_dim = self.config.head_dim();
+        let eps = self.config.rms_norm_eps.unwrap_or(1e-5) as f32;
+        let positions: Vec<u32> = (0..t as u32).collect();
+
+        let mut x = embedding(&self.embed_tokens, token_ids); // [T, hidden]
+
+        for layer in &self.layers {
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.input_norm, eps);
+
+            // Q/K/V projections + bias.
+            let q = add_bias(&matmul2(&normed, &layer.q_proj_wt), &layer.q_bias); // [T, n_heads*hd]
+            let k = add_bias(&matmul2(&normed, &layer.k_proj_wt), &layer.k_bias); // [T, n_kv*hd]
+            let v = add_bias(&matmul2(&normed, &layer.v_proj_wt), &layer.v_bias);
+
+            // RoPE (rotate_half, same convention xserv uses for Qwen3): reshape to
+            // [1,H,T,D] -> [T,H,D] -> rope -> back.
+            let q = reshape_heads_gpu(&q, t, n_heads, head_dim);
+            let k = reshape_heads_gpu(&k, t, n_kv, head_dim);
+            let q = transpose_for_rope_gpu(&q, t, n_heads, head_dim);
+            let k = transpose_for_rope_gpu(&k, t, n_kv, head_dim);
+            rope_inplace(&q, &self.rope_cache, &positions);
+            rope_inplace(&k, &self.rope_cache, &positions);
+            let q = transpose_from_rope_gpu(&q, t, n_heads, head_dim); // [1,H,T,D]
+            let k = transpose_from_rope_gpu(&k, t, n_kv, head_dim);
+            let v = reshape_heads_gpu(&v, t, n_kv, head_dim);          // [1,H_kv,T,D]
+
+            // Naive attention with sinks (CPU softmax for correctness).
+            let attn = attention_with_sinks(
+                &q, &k, &v, &layer.sinks, n_heads, n_kv, head_dim, t,
+                if layer.sliding { self.config.sliding_window() } else { None },
+            ); // [T, hidden]
+            let attn_proj = add_bias(&matmul2(&attn, &layer.o_proj_wt), &layer.o_bias);
+            x = add(&residual, &attn_proj);
+
+            // MoE FFN.
+            let residual = x.clone();
+            let normed = rmsnorm(&x, &layer.post_norm, eps);
+            let moe = self.moe_ffn(&normed, layer, hidden);
+            x = add(&residual, &moe);
+        }
+
+        let x = rmsnorm(&x, &self.norm, eps);
+        matmul2(&x, &self.lm_head_t) // [T, vocab]
+    }
+
+    /// MoE FFN over [T, hidden]: router top-k softmax, per-token weighted sum of
+    /// its top-k experts' clamped-SwiGLU outputs. Correctness-first (per-token).
+    fn moe_ffn(&self, x: &Tensor, layer: &Block, hidden: usize) -> Tensor {
+        let t = x.shape()[0];
+        let top_k = self.config.experts_per_tok();
+        let n_experts = self.config.num_experts();
+        let limit = self.config.swiglu_limit();
+
+        // router logits [T, n_experts] on host.
+        let logits = add_bias(&matmul2(x, &layer.router_wt), &layer.router_bias);
+        let logits_h = logits.to_device(Device::Cpu);
+        let lg = logits_h.as_slice::<bf16>();
+
+        // Per-token top-k indices + softmax weights (over the chosen k).
+        let mut out_rows: Vec<Tensor> = Vec::with_capacity(t);
+        for ti in 0..t {
+            let row = &lg[ti * n_experts..(ti + 1) * n_experts];
+            let mut idx: Vec<usize> = (0..n_experts).collect();
+            idx.sort_by(|&a, &b| row[b].to_f32().partial_cmp(&row[a].to_f32()).unwrap());
+            let top = &idx[..top_k];
+            let maxv = row[top[0]].to_f32();
+            let exps: Vec<f32> = top.iter().map(|&e| (row[e].to_f32() - maxv).exp()).collect();
+            let sum: f32 = exps.iter().sum();
+            let weights: Vec<f32> = exps.iter().map(|w| w / sum).collect();
+
+            // x row as [1, hidden].
+            let xr = row_view(x, ti);
+            let mut acc: Option<Tensor> = None;
+            for (j, &e) in top.iter().enumerate() {
+                let y = expert_forward(&xr, &layer.gate_up_wt[e], &layer.gate_up_bias[e],
+                                       &layer.down_wt[e], &layer.down_bias[e], limit); // [1, hidden]
+                let yw = scale_tensor(&y, weights[j]);
+                acc = Some(match acc { Some(a) => add(&a, &yw), None => yw });
+            }
+            out_rows.push(acc.unwrap_or_else(|| zeros_row(hidden)));
+        }
+        concat_rows(&out_rows) // [T, hidden]
+    }
+}
+
+// ---------- helpers ----------
+
+/// Build a YaRN-scaled RoPE cos/sin cache (gpt-oss uses rope_type "yarn").
+/// Mirrors HF `_compute_yarn_parameters`: per-dim interpolation/extrapolation
+/// ramp between the scaled (theta*factor) and unscaled frequencies, plus a global
+/// attention scaling (mscale) folded into cos/sin. Cache layout matches xserv's
+/// rope kernel: f32 [max_seq, half_dim], cos[pos*half+i] = cos(pos*invfreq[i])*mscale.
+fn yarn_rope_cache(config: &ModelConfig) -> RopeCache {
+    use std::f64::consts::PI;
+    let head_dim = config.head_dim();
+    let half = head_dim / 2;
+    let max_seq = config.max_seq_len();
+    let base = config.rope_theta.unwrap_or(150000.0);
+    // gpt-oss rope_scaling: yarn, factor 32, beta_fast 32, beta_slow 1, orig 4096,
+    // truncate false (keep correction range as floats).
+    let factor = 32.0f64;
+    let (beta_fast, beta_slow) = (32.0f64, 1.0f64);
+    let orig_max = 4096.0f64;
+    let dim = head_dim as f64;
+
+    let find_dim = |num_rot: f64| (dim * (orig_max / (num_rot * 2.0 * PI)).ln()) / (2.0 * base.ln());
+    let low = find_dim(beta_fast).max(0.0);
+    let high = find_dim(beta_slow).min(dim - 1.0);
+    let denom = (high - low).max(1e-3);
+
+    let mut inv_freq = vec![0f64; half];
+    for i in 0..half {
+        let pos_freq = base.powf((2 * i) as f64 / dim);
+        let extrap = 1.0 / pos_freq;                 // unscaled (extrapolation)
+        let interp = 1.0 / (factor * pos_freq);      // scaled   (interpolation)
+        let ramp = ((i as f64 - low) / denom).clamp(0.0, 1.0);
+        let mask = 1.0 - ramp;                       // extrapolation factor
+        inv_freq[i] = interp * (1.0 - mask) + extrap * mask;
+    }
+    // mscale: 0.1*ln(factor)+1 for factor>1.
+    let mscale = (0.1 * factor.ln() + 1.0) as f64;
+
+    let mut cos = vec![0f32; max_seq * half];
+    let mut sin = vec![0f32; max_seq * half];
+    for p in 0..max_seq {
+        for i in 0..half {
+            let ang = p as f64 * inv_freq[i];
+            cos[p * half + i] = (ang.cos() * mscale) as f32;
+            sin[p * half + i] = (ang.sin() * mscale) as f32;
+        }
+    }
+    let bytes = max_seq * half * std::mem::size_of::<f32>();
+    let mut cos_buf = xserv_cuda::GpuBuffer::alloc(bytes).expect("alloc yarn cos");
+    let mut sin_buf = xserv_cuda::GpuBuffer::alloc(bytes).expect("alloc yarn sin");
+    let cb = unsafe { std::slice::from_raw_parts(cos.as_ptr() as *const u8, bytes) };
+    let sb = unsafe { std::slice::from_raw_parts(sin.as_ptr() as *const u8, bytes) };
+    cos_buf.copy_from_host(cb).unwrap();
+    sin_buf.copy_from_host(sb).unwrap();
+    RopeCache { cos: cos_buf, sin: sin_buf, max_seq_len: max_seq, half_dim: half }
+}
+
+fn matmul2(a: &Tensor, b: &Tensor) -> Tensor {
+    matmul(a, b, GemmBackend::CuBlas)
+}
+
+/// One expert: clamped SwiGLU. x:[*,hidden] -> [*,hidden].
+/// gate_up = x@gate_up_wt + bias; gate=even cols, up=odd cols (interleaved);
+/// gate.clamp(max=limit); up.clamp(-limit,limit); h=(up+1)*gate*sigmoid(gate*1.702); h@down_wt+bias.
+fn expert_forward(x: &Tensor, gate_up_wt: &Tensor, gate_up_bias: &Tensor,
+                  down_wt: &Tensor, down_bias: &Tensor, limit: f32) -> Tensor {
+    // Upload this expert's CPU-resident weights to x's device just for this call.
+    let dev = x.device();
+    let gate_up_wt = gate_up_wt.to_device(dev);
+    let gate_up_bias = gate_up_bias.to_device(dev);
+    let down_wt = down_wt.to_device(dev);
+    let down_bias = down_bias.to_device(dev);
+    let gate_up = add_bias(&matmul2(x, &gate_up_wt), &gate_up_bias); // [*, 2*inter]
+    let h = clamped_swiglu(&gate_up, limit);                         // [*, inter]
+    add_bias(&matmul2(&h, &down_wt), &down_bias)                     // [*, hidden]
+}
+
+/// Clamped interleaved SwiGLU on host (correctness-first). [*, 2I] -> [*, I].
+fn clamped_swiglu(gate_up: &Tensor, limit: f32) -> Tensor {
+    const ALPHA: f32 = 1.702;
+    let rows = gate_up.shape()[0];
+    let two_i = gate_up.shape()[1];
+    let inter = two_i / 2;
+    let h = gate_up.to_device(Device::Cpu);
+    let s = h.as_slice::<bf16>();
+    let mut out = vec![bf16::ZERO; rows * inter];
+    for r in 0..rows {
+        for i in 0..inter {
+            let g = s[r * two_i + 2 * i].to_f32();
+            let u = s[r * two_i + 2 * i + 1].to_f32();
+            let g = g.min(limit);
+            let u = u.clamp(-limit, limit);
+            let glu = g * (1.0 / (1.0 + (-(g * ALPHA)).exp()));
+            out[r * inter + i] = bf16::from_f32((u + 1.0) * glu);
+        }
+    }
+    Tensor::from_slice(&out, &[rows, inter]).to_device(gate_up.device())
+}
+
+/// Naive multi-head attention with per-head sink logits, on host (correctness-first).
+/// q:[1,n_heads,T,D] k,v:[1,n_kv,T,D] sinks:[n_heads] (host). Returns [T, n_heads*D].
+#[allow(clippy::too_many_arguments)]
+fn attention_with_sinks(q: &Tensor, k: &Tensor, v: &Tensor, sinks: &Tensor,
+                        n_heads: usize, n_kv: usize, head_dim: usize, t: usize,
+                        window: Option<usize>) -> Tensor {
+    let scale = (head_dim as f32).powf(-0.5);
+    let n_rep = n_heads / n_kv;
+    let qh = q.to_device(Device::Cpu); let qd = qh.as_slice::<bf16>();
+    let kh = k.to_device(Device::Cpu); let kd = kh.as_slice::<bf16>();
+    let vh = v.to_device(Device::Cpu); let vd = vh.as_slice::<bf16>();
+    let sh = sinks.to_device(Device::Cpu); let sd = sh.as_slice::<bf16>();
+    let hidden = n_heads * head_dim;
+    let mut out = vec![bf16::ZERO; t * hidden];
+    // index helpers: layout [H, T, D] within each (head) block.
+    let qi = |h: usize, i: usize, d: usize| (h * t + i) * head_dim + d;
+    let kvi = |h: usize, j: usize, d: usize| (h * t + j) * head_dim + d;
+    for h in 0..n_heads {
+        let kv = h / n_rep;
+        for i in 0..t {
+            // scores over valid keys j<=i (causal), and j>i-window (sliding).
+            let lo = match window { Some(wn) if i + 1 > wn => i + 1 - wn, _ => 0 };
+            let mut scores = vec![0f32; i - lo + 1];
+            let mut maxv = sd[h].to_f32(); // sink participates in the max
+            for j in lo..=i {
+                let mut dot = 0f32;
+                for d in 0..head_dim {
+                    dot += qd[qi(h, i, d)].to_f32() * kd[kvi(kv, j, d)].to_f32();
+                }
+                let s = dot * scale;
+                scores[j - lo] = s;
+                if s > maxv { maxv = s; }
+            }
+            let mut denom = (sd[h].to_f32() - maxv).exp(); // sink column
+            for s in &scores { denom += (*s - maxv).exp(); }
+            // weighted sum of v (sink contributes no value -> just inflates denom).
+            for d in 0..head_dim {
+                let mut acc = 0f32;
+                for j in lo..=i {
+                    let p = (scores[j - lo] - maxv).exp() / denom;
+                    acc += p * vd[kvi(kv, j, d)].to_f32();
+                }
+                out[i * hidden + h * head_dim + d] = bf16::from_f32(acc);
+            }
+        }
+    }
+    Tensor::from_slice(&out, &[t, hidden]).to_device(q.device())
+}
+
+/// Row-broadcast bias add: x:[T,N] + bias:[N] -> [T,N], via ones[T,1]@bias[1,N].
+fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
+    let t = x.shape()[0];
+    let n = x.shape()[1];
+    let ones = Tensor::from_slice(&vec![bf16::from_f32(1.0); t], &[t, 1]).to_device(x.device());
+    let bias_row = bias.reshape(&[1, n]);
+    let broadcast = matmul2(&ones, &bias_row); // [T, N]
+    add(x, &broadcast)
+}
+
+/// Slice expert `e` out of a fused [E, rows, cols] tensor -> [rows, cols].
+fn slice_expert(t: &Tensor, e: usize, rows: usize, cols: usize) -> Tensor {
+    let host = t.to_device(Device::Cpu);
+    let s = host.as_slice::<bf16>();
+    let stride = rows * cols;
+    Tensor::from_slice(&s[e * stride..(e + 1) * stride], &[rows, cols])
+}
+
+/// Slice row `e` out of [E, n] -> [n].
+fn slice_row(t: &Tensor, e: usize, n: usize) -> Tensor {
+    let host = t.to_device(Device::Cpu);
+    let s = host.as_slice::<bf16>();
+    Tensor::from_slice(&s[e * n..(e + 1) * n], &[n])
+}
+
+fn row_view(t: &Tensor, row: usize) -> Tensor {
+    let cols = t.shape()[1];
+    let host = t.to_device(Device::Cpu);
+    let s = host.as_slice::<bf16>();
+    Tensor::from_slice(&s[row * cols..(row + 1) * cols], &[1, cols]).to_device(t.device())
+}
+
+fn scale_tensor(t: &Tensor, s: f32) -> Tensor {
+    let host = t.to_device(Device::Cpu);
+    let data = host.as_slice::<bf16>();
+    let out: Vec<bf16> = data.iter().map(|v| bf16::from_f32(v.to_f32() * s)).collect();
+    Tensor::from_slice(&out, t.shape()).to_device(t.device())
+}
+
+fn zeros_row(n: usize) -> Tensor {
+    Tensor::from_slice(&vec![bf16::ZERO; n], &[1, n]).to_device(Device::Cuda(0))
+}
+
+fn concat_rows(rows: &[Tensor]) -> Tensor {
+    let n = rows[0].shape()[1];
+    let mut out = Vec::with_capacity(rows.len() * n);
+    for r in rows {
+        let h = r.to_device(Device::Cpu);
+        out.extend_from_slice(h.as_slice::<bf16>());
+    }
+    Tensor::from_slice(&out, &[rows.len(), n]).to_device(Device::Cuda(0))
+}

crates/xserv-model/src/kv_cache.rs

@@ -1,6 +1,6 @@
-use crate::config::ModelConfig;
 use xserv_cuda::GpuBuffer;
-use xserv_tensor::{DType, Tensor};
+use xserv_tensor::{DType, Device, Tensor};
+use crate::config::ModelConfig;
 
 /// GPU-resident KV cache. Pre-allocates max_seq_len on GPU,
 /// appends new K/V via D2D copy at offset (no CPU round-trip).
@@ -46,43 +46,17 @@ impl GpuKVCache {
             v_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging V"));
         }
 
-        Self {
-            k_bufs,
-            v_bufs,
-            k_staging,
-            v_staging,
-            seq_len: 0,
-            max_seq_len,
-            num_kv_heads,
-            head_dim,
-            elem_size,
-            dtype,
-            device,
-        }
+        Self { k_bufs, v_bufs, k_staging, v_staging, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype, device }
     }
 
-    pub fn seq_len(&self) -> usize {
-        self.seq_len
-    }
-    pub fn max_seq_len(&self) -> usize {
-        self.max_seq_len
-    }
+    pub fn seq_len(&self) -> usize { self.seq_len }
+    pub fn max_seq_len(&self) -> usize { self.max_seq_len }
 
     /// Append new K/V tensors for a given layer.
     /// k_new, v_new: [1, num_kv_heads, new_tokens, head_dim] on GPU, contiguous.
     /// `write_pos` is the sequence position to write at (caller manages this).
-    pub fn append(
-        &mut self,
-        layer: usize,
-        k_new: &Tensor,
-        v_new: &Tensor,
-        new_tokens: usize,
-        write_pos: usize,
-    ) {
-        assert!(
-            write_pos + new_tokens <= self.max_seq_len,
-            "KV cache overflow"
-        );
+    pub fn append(&mut self, layer: usize, k_new: &Tensor, v_new: &Tensor, new_tokens: usize, write_pos: usize) {
+        assert!(write_pos + new_tokens <= self.max_seq_len, "KV cache overflow");
         let es = self.elem_size;
         let hd = self.head_dim;
         let max_s = self.max_seq_len;
@@ -95,23 +69,14 @@ impl GpuKVCache {
             let src_off = h * new_tokens * hd * es;
             let dst_off = (h * max_s + write_pos) * hd * es;
             let count = new_tokens * hd * es;
-            self.k_bufs[layer]
-                .copy_from_device_at(k_src, src_off, dst_off, count)
-                .unwrap();
-            self.v_bufs[layer]
-                .copy_from_device_at(v_src, src_off, dst_off, count)
-                .unwrap();
+            self.k_bufs[layer].copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
+            self.v_bufs[layer].copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
         }
     }
 
     pub fn advance_seq_len(&mut self, new_tokens: usize) {
         self.seq_len += new_tokens;
-        assert!(
-            self.seq_len <= self.max_seq_len,
-            "KV cache seq_len ({}) exceeds max_seq_len ({})",
-            self.seq_len,
-            self.max_seq_len
-        );
+        assert!(self.seq_len <= self.max_seq_len, "KV cache seq_len ({}) exceeds max_seq_len ({})", self.seq_len, self.max_seq_len);
     }
 
     /// Get K/V cache tensors for a layer up to `seq_len` tokens: [1, num_kv_heads, seq_len, head_dim]
@@ -121,11 +86,7 @@ impl GpuKVCache {
     }
 
     pub fn get_kv_len(&mut self, layer: usize, sl: usize) -> (Tensor, Tensor) {
-        assert!(
-            sl <= self.max_seq_len,
-            "get_kv_len: sl ({sl}) exceeds max_seq_len ({})",
-            self.max_seq_len
-        );
+        assert!(sl <= self.max_seq_len, "get_kv_len: sl ({sl}) exceeds max_seq_len ({})", self.max_seq_len);
         let hd = self.head_dim;
         let nh = self.num_kv_heads;
         let es = self.elem_size;
@@ -143,12 +104,8 @@ impl GpuKVCache {
             let src_off = (h * max_s) * hd * es;
             let dst_off = (h * sl) * hd * es;
             let count = sl * hd * es;
-            k_stg
-                .copy_from_device_at(k_buf, src_off, dst_off, count)
-                .unwrap();
-            v_stg
-                .copy_from_device_at(v_buf, src_off, dst_off, count)
-                .unwrap();
+            k_stg.copy_from_device_at(k_buf, src_off, dst_off, count).unwrap();
+            v_stg.copy_from_device_at(v_buf, src_off, dst_off, count).unwrap();
         }
         // Grab raw pointers before dropping the mutable borrows
         let k_ptr = k_stg.as_mut_ptr();
@@ -160,35 +117,20 @@ impl GpuKVCache {
         // get_kv_len call overwrites the staging buffer).
         let shape = &[1usize, nh, sl, hd];
         let k = unsafe {
-            tensor_from_gpu_buffer(
-                GpuBuffer::borrow_raw(k_ptr, out_size),
-                shape,
-                self.dtype,
-                self.device,
-            )
+            tensor_from_gpu_buffer(GpuBuffer::borrow_raw(k_ptr, out_size), shape, self.dtype, self.device)
         };
         let v = unsafe {
-            tensor_from_gpu_buffer(
-                GpuBuffer::borrow_raw(v_ptr, out_size),
-                shape,
-                self.dtype,
-                self.device,
-            )
+            tensor_from_gpu_buffer(GpuBuffer::borrow_raw(v_ptr, out_size), shape, self.dtype, self.device)
         };
         (k, v)
     }
 }
 
 /// Create a Tensor from a GpuBuffer (takes ownership).
-unsafe fn tensor_from_gpu_buffer(
-    buf: GpuBuffer,
-    shape: &[usize],
-    dtype: DType,
-    device: u32,
-) -> Tensor {
-    use smallvec::SmallVec;
-    use xserv_tensor::shape::contiguous_strides;
+unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
     use xserv_tensor::storage::Storage;
+    use xserv_tensor::shape::contiguous_strides;
+    use smallvec::SmallVec;
 
     let storage = Storage::cuda(buf, device);
     Tensor::from_storage(
@@ -204,11 +146,6 @@ unsafe fn tensor_from_gpu_buffer(
 ///
 /// # Safety
 /// `buf` must be a valid GPU allocation with at least `product(shape) * dtype.size_bytes()` bytes.
-pub unsafe fn tensor_from_gpu_buffer_pub(
-    buf: GpuBuffer,
-    shape: &[usize],
-    dtype: DType,
-    device: u32,
-) -> Tensor {
+pub unsafe fn tensor_from_gpu_buffer_pub(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
     tensor_from_gpu_buffer(buf, shape, dtype, device)
 }

crates/xserv-model/src/lib.rs

@@ -1,25 +1,20 @@
 pub mod config;
 pub mod decode_graph;
-pub mod eagle3;
 pub mod gpt2;
-pub mod gpt_oss;
-pub mod gpt_oss_graph;
 pub mod kv_cache;
 pub mod loader;
 pub mod paged_kv_cache;
 pub mod qwen3;
-pub mod qwen3_graph;
 pub mod sampling;
 
 pub use config::ModelConfig;
 pub use decode_graph::{DecodeGraphState, LayerWeightPtrs};
-pub use gpt_oss::GptOss;
-pub use gpt_oss_graph::{GptOssDecodeGraph, GraphedGptOssDecoder};
 pub use gpt2::{GPT2, KVCache};
 pub use kv_cache::GpuKVCache;
-pub use paged_kv_cache::{BLOCK_SIZE, BlockAllocator, Location, PagedKVCache};
+pub use paged_kv_cache::{BlockAllocator, Location, PagedKVCache, BLOCK_SIZE};
 pub use qwen3::Qwen3;
-pub use sampling::{SamplingParams, sample, sample_greedy_penalized};
+pub use gptoss::GptOss;
+pub use sampling::{SamplingParams, sample};
 
 /// Initialize GPU kernel hooks. Called automatically by model constructors,
 /// but safe to call multiple times (idempotent via OnceLock).

crates/xserv-model/src/loader.rs

@@ -5,8 +5,8 @@ use std::path::Path;
 use xserv_tensor::{DType, Device, Tensor};
 
 pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor> {
-    let data =
-        std::fs::read(path).unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
+    let data = std::fs::read(path)
+        .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
     let st = SafeTensors::deserialize(&data)
         .unwrap_or_else(|e| panic!("failed to parse safetensors {}: {e}", path.display()));
 
@@ -19,7 +19,6 @@ 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;
@@ -60,15 +59,11 @@ pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
         all_tensors.extend(tensors);
     }
 
-    assert!(
-        !all_tensors.is_empty(),
-        "no safetensors files found in {}",
-        dir.display()
-    );
+    assert!(!all_tensors.is_empty(), "no safetensors files found in {}", dir.display());
     all_tensors
 }
 
-pub(crate) fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
+fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
     match dtype {
         DType::F32 => {
             let floats: &[f32] = unsafe {
@@ -88,6 +83,5 @@ pub(crate) fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Te
             };
             Tensor::from_slice(bfs, shape)
         }
-        DType::FP8E4M3 => Tensor::from_raw_bytes(raw_bytes, shape, DType::FP8E4M3),
     }
 }

crates/xserv-model/src/paged_kv_cache.rs

@@ -29,10 +29,7 @@ impl BlockAllocator {
         for b in (1..total_blocks).rev() {
             free_stack.push(b as u32);
         }
-        Self {
-            free_stack,
-            total: total_blocks,
-        }
+        Self { free_stack, total: total_blocks }
     }
 
     pub fn alloc(&mut self) -> Option<u32> {
@@ -139,14 +136,8 @@ impl PagedKVCache {
         device: u32,
     ) -> Self {
         Self::new_tp(
-            config,
-            config.num_kv_heads(),
-            total_blocks,
-            cpu_total_blocks,
-            max_seqs,
-            max_blocks_per_seq,
-            dtype,
-            device,
+            config, config.num_kv_heads(), total_blocks, cpu_total_blocks,
+            max_seqs, max_blocks_per_seq, dtype, device,
         )
     }
 
@@ -164,10 +155,7 @@ impl PagedKVCache {
         dtype: DType,
         device: u32,
     ) -> Self {
-        assert!(
-            total_blocks >= 2,
-            "need at least 2 blocks (one is sentinel)"
-        );
+        assert!(total_blocks >= 2, "need at least 2 blocks (one is sentinel)");
         let num_layers = config.num_layers();
         let head_dim = config.head_dim();
         let elem_size = dtype.size_bytes();
@@ -191,17 +179,11 @@ impl PagedKVCache {
         if cpu_total_blocks >= 2 {
             let cpu_pool_bytes = cpu_total_blocks * block_bytes;
             for _ in 0..num_layers {
-                cpu_k_pools
-                    .push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU K swap pool"));
-                cpu_v_pools
-                    .push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU V swap pool"));
+                cpu_k_pools.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU K swap pool"));
+                cpu_v_pools.push(PinnedBuffer::alloc(cpu_pool_bytes).expect("alloc CPU V swap pool"));
             }
         }
-        let cpu_allocator = BlockAllocator::new(if cpu_total_blocks >= 2 {
-            cpu_total_blocks
-        } else {
-            0
-        });
+        let cpu_allocator = BlockAllocator::new(if cpu_total_blocks >= 2 { cpu_total_blocks } else { 0 });
 
         let block_table_gpu =
             GpuBuffer::alloc(max_seqs * max_blocks_per_seq * std::mem::size_of::<i32>())
@@ -238,49 +220,22 @@ impl PagedKVCache {
         }
     }
 
-    pub fn num_layers(&self) -> usize {
-        self.num_layers
-    }
-    pub fn num_kv_heads(&self) -> usize {
-        self.num_kv_heads
-    }
-    pub fn head_dim(&self) -> usize {
-        self.head_dim
-    }
-    pub fn dtype(&self) -> DType {
-        self.dtype
-    }
-    pub fn max_seqs(&self) -> usize {
-        self.max_seqs
-    }
-    pub fn max_blocks_per_seq(&self) -> usize {
-        self.max_blocks_per_seq
-    }
-    pub fn free_blocks(&self) -> usize {
-        self.allocator.free_count()
-    }
-    pub fn total_blocks(&self) -> usize {
-        self.allocator.total()
-    }
+    pub fn num_layers(&self) -> usize { self.num_layers }
+    pub fn num_kv_heads(&self) -> usize { self.num_kv_heads }
+    pub fn head_dim(&self) -> usize { self.head_dim }
+    pub fn dtype(&self) -> DType { self.dtype }
+    pub fn max_seqs(&self) -> usize { self.max_seqs }
+    pub fn max_blocks_per_seq(&self) -> usize { self.max_blocks_per_seq }
+    pub fn free_blocks(&self) -> usize { self.allocator.free_count() }
+    pub fn total_blocks(&self) -> usize { self.allocator.total() }
 
-    pub fn k_pool(&self, layer: usize) -> &GpuBuffer {
-        &self.k_pools[layer]
-    }
-    pub fn v_pool(&self, layer: usize) -> &GpuBuffer {
-        &self.v_pools[layer]
-    }
-    pub fn block_table_gpu(&self) -> &GpuBuffer {
-        &self.block_table_gpu
-    }
-    pub fn context_lens_gpu(&self) -> &GpuBuffer {
-        &self.context_lens_gpu
-    }
+    pub fn k_pool(&self, layer: usize) -> &GpuBuffer { &self.k_pools[layer] }
+    pub fn v_pool(&self, layer: usize) -> &GpuBuffer { &self.v_pools[layer] }
+    pub fn block_table_gpu(&self) -> &GpuBuffer { &self.block_table_gpu }
+    pub fn context_lens_gpu(&self) -> &GpuBuffer { &self.context_lens_gpu }
 
     pub fn seq_len(&self, slot: usize) -> usize {
-        self.seq_states[slot]
-            .as_ref()
-            .map(|s| s.seq_len)
-            .unwrap_or(0)
+        self.seq_states[slot].as_ref().map(|s| s.seq_len).unwrap_or(0)
     }
 
     pub fn is_slot_free(&self, slot: usize) -> bool {
@@ -325,11 +280,7 @@ impl PagedKVCache {
         let state = self.seq_states[slot].as_ref().expect("unregistered slot");
         let cur = state.block_ids.len();
         let needed_total = (state.seq_len + new_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE;
-        if needed_total > cur {
-            needed_total - cur
-        } else {
-            0
-        }
+        if needed_total > cur { needed_total - cur } else { 0 }
     }
 
     /// Pre-allocate enough physical blocks in `slot` to cover positions
@@ -339,14 +290,8 @@ impl PagedKVCache {
         let state = self.seq_states[slot].as_mut().expect("unregistered slot");
         let needed_total = (end_pos + BLOCK_SIZE - 1) / BLOCK_SIZE;
         while state.block_ids.len() < needed_total {
-            let b = self
-                .allocator
-                .alloc()
-                .expect("out of blocks (caller must check)");
-            assert!(
-                state.block_ids.len() < self.max_blocks_per_seq,
-                "block table overflow"
-            );
+            let b = self.allocator.alloc().expect("out of blocks (caller must check)");
+            assert!(state.block_ids.len() < self.max_blocks_per_seq, "block table overflow");
             state.block_ids.push(b);
         }
     }
@@ -360,10 +305,6 @@ impl PagedKVCache {
     /// `k_new`, `v_new`: GPU tensors with logical shape
     ///   [1, num_kv_heads, num_tokens, head_dim]
     /// stored contiguously (head-major, then tokens, then dim).
-    ///
-    /// Implementation: a single `reshape_and_cache` kernel per call. The
-    /// previous Rust loop fired `num_tokens * num_kv_heads` cudaMemcpys per
-    /// layer (≈290k for a 1024-token Qwen3 prefill across 36 layers).
     pub fn append_tokens(
         &mut self,
         slot: usize,
@@ -373,110 +314,40 @@ impl PagedKVCache {
         num_tokens: usize,
         start_pos: usize,
     ) {
-        if num_tokens == 0 {
-            return;
-        }
+        if num_tokens == 0 { return; }
         // Make sure blocks exist for the target range.
         self.ensure_capacity(slot, start_pos + num_tokens);
 
+        let block_ids = self.seq_states[slot].as_ref().unwrap().block_ids.clone();
+
         let nkv = self.num_kv_heads;
         let hd = self.head_dim;
+        let es = self.elem_size;
         let bs = BLOCK_SIZE;
 
-        // Stage block_ids on the GPU. Pool-allocated so this is essentially
-        // free after the first call (same bucket every step).
-        let block_ids: Vec<i32> = self.seq_states[slot]
-            .as_ref()
-            .unwrap()
-            .block_ids
-            .iter()
-            .map(|&b| b as i32)
-            .collect();
-        let bytes = block_ids.len() * std::mem::size_of::<i32>();
-        let mut block_ids_gpu =
-            xserv_cuda::allocator::cached_alloc(bytes).expect("alloc append block_ids");
-        let block_ids_bytes =
-            unsafe { std::slice::from_raw_parts(block_ids.as_ptr() as *const u8, bytes) };
-        block_ids_gpu
-            .copy_from_host(block_ids_bytes)
-            .expect("upload block_ids");
+        let k_src = k_new.storage().gpu_buffer();
+        let v_src = v_new.storage().gpu_buffer();
 
-        let k_src = k_new.data_ptr() as *const std::ffi::c_void;
-        let v_src = v_new.data_ptr() as *const std::ffi::c_void;
-        let k_pool_ptr = self.k_pools[layer].as_mut_ptr() as *mut std::ffi::c_void;
-        let v_pool_ptr = self.v_pools[layer].as_mut_ptr() as *mut std::ffi::c_void;
+        let k_pool = &mut self.k_pools[layer];
+        let v_pool = &mut self.v_pools[layer];
 
-        unsafe {
-            xserv_kernels::reshape_and_cache_bf16(
-                k_src,
-                v_src,
-                k_pool_ptr,
-                v_pool_ptr,
-                block_ids_gpu.as_ptr() as *const i32,
-                num_tokens,
-                nkv,
-                hd,
-                start_pos,
-                bs,
-                xserv_cuda::current_stream_raw(),
-            );
-        }
-        // block_ids_gpu drops here; the launch on the null stream will have
-        // finished consuming it before any subsequent op alloc()s the same
-        // bucket (null stream is sequential).
-    }
+        let mut t = 0usize;
+        while t < num_tokens {
+            let p = start_pos + t;
+            let logical_blk = p / bs;
+            let slot_in_blk = p % bs;
+            let chunk = (bs - slot_in_blk).min(num_tokens - t);
+            let phys = block_ids[logical_blk] as usize;
 
-    /// Batched append for the multi-sequence decode step: writes one new
-    /// K/V token per active sequence into `layer`'s pool, using
-    /// `block_table_gpu` and `context_lens_gpu` directly. Caller must have
-    /// just run `sync_active_batch_with_lens(slots, kv_lens)` so that:
-    ///   - row `i` of block_table_gpu holds the block ids for `slots[i]`
-    ///   - context_lens_gpu[i] == seq_len(slots[i]) + 1 (the kv_len **after**
-    ///     this step — i.e., the new token will be written at index kv_len-1)
-    ///
-    /// `k_new`, `v_new`: GPU tensors, contiguous, BF16, shape
-    ///   `[batch, num_kv_heads, head_dim]`.
-    ///
-    /// Like `append_tokens`, this does **not** touch `seq_len`. Call
-    /// `advance_seq_len(slot, 1)` for each slot after every layer has been
-    /// written.
-    pub fn append_tokens_batched(
-        &mut self,
-        layer: usize,
-        k_new: &Tensor,
-        v_new: &Tensor,
-        batch: usize,
-    ) {
-        if batch == 0 {
-            return;
-        }
-        let nkv = self.num_kv_heads;
-        let hd = self.head_dim;
-        debug_assert_eq!(k_new.shape(), &[batch, nkv, hd]);
-        debug_assert_eq!(v_new.shape(), &[batch, nkv, hd]);
+            for h in 0..nkv {
+                let src_off = (h * num_tokens + t) * hd * es;
+                let dst_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
+                let count = chunk * hd * es;
+                k_pool.copy_from_device_at(k_src, src_off, dst_off, count).unwrap();
+                v_pool.copy_from_device_at(v_src, src_off, dst_off, count).unwrap();
+            }
 
-        let k_src = k_new.data_ptr() as *const std::ffi::c_void;
-        let v_src = v_new.data_ptr() as *const std::ffi::c_void;
-        let k_pool_ptr = self.k_pools[layer].as_mut_ptr() as *mut std::ffi::c_void;
-        let v_pool_ptr = self.v_pools[layer].as_mut_ptr() as *mut std::ffi::c_void;
-        let bt_ptr = self.block_table_gpu.as_ptr() as *const i32;
-        let cl_ptr = self.context_lens_gpu.as_ptr() as *const i32;
-
-        unsafe {
-            xserv_kernels::reshape_and_cache_batched_bf16(
-                k_src,
-                v_src,
-                k_pool_ptr,
-                v_pool_ptr,
-                bt_ptr,
-                cl_ptr,
-                batch,
-                nkv,
-                hd,
-                BLOCK_SIZE,
-                self.max_blocks_per_seq,
-                xserv_cuda::current_stream_raw(),
-            );
+            t += chunk;
         }
     }
 
@@ -486,80 +357,6 @@ impl PagedKVCache {
         state.seq_len += num_tokens;
     }
 
-    /// Roll a registered sequence back to `new_len` tokens.
-    ///
-    /// This only changes cache metadata and frees whole physical blocks that are
-    /// no longer reachable. Bytes inside retained blocks are left untouched; the
-    /// logical `seq_len` prevents attention from reading them, and later writes
-    /// to the same positions overwrite them.
-    pub fn truncate_sequence(&mut self, slot: usize, new_len: usize) -> Result<(), &'static str> {
-        if slot >= self.max_seqs {
-            return Err("truncate_sequence: slot out of range");
-        }
-        let state = self.seq_states[slot]
-            .as_mut()
-            .ok_or("truncate_sequence: empty slot")?;
-        if new_len > state.seq_len {
-            return Err("truncate_sequence: cannot extend");
-        }
-
-        let needed_blocks = ((new_len + BLOCK_SIZE - 1) / BLOCK_SIZE).max(1);
-        while state.block_ids.len() > needed_blocks {
-            let block = state.block_ids.pop().expect("checked len");
-            match state.location {
-                Location::Gpu => self.allocator.free(block),
-                Location::Cpu => self.cpu_allocator.free(block),
-            }
-        }
-        state.seq_len = new_len;
-        Ok(())
-    }
-
-    /// Copy K/V data from `src_pos` to `dst_pos` within the same slot, across
-    /// all layers. Used by tree speculative decoding to remap an accepted
-    /// sibling's K/V to the canonical sequential position after acceptance.
-    ///
-    /// Requires: both positions within the currently-allocated block range.
-    pub fn copy_kv_position(&self, slot: usize, src_pos: usize, dst_pos: usize) {
-        let state = self.seq_states[slot]
-            .as_ref()
-            .expect("copy_kv_position: slot not registered");
-        assert!(
-            src_pos < state.seq_len && dst_pos < state.seq_len,
-            "copy_kv_position: positions must be within seq_len"
-        );
-        // Upload this sequence's block_ids to a small GPU buffer.
-        let block_ids_host: Vec<i32> = state.block_ids.iter().map(|&b| b as i32).collect();
-        let bytes: &[u8] = unsafe {
-            std::slice::from_raw_parts(
-                block_ids_host.as_ptr() as *const u8,
-                block_ids_host.len() * 4,
-            )
-        };
-        let mut ids_buf =
-            xserv_cuda::allocator::cached_alloc(bytes.len()).expect("alloc block_ids for copy");
-        ids_buf.copy_from_host(bytes).unwrap();
-        let ids_ptr = ids_buf.as_ptr() as *const i32;
-
-        let stream = xserv_cuda::current_stream_raw();
-        let num_layers = self.k_pools.len();
-        for layer in 0..num_layers {
-            unsafe {
-                xserv_kernels::copy_kv_position(
-                    self.k_pools[layer].as_ptr() as *mut std::ffi::c_void,
-                    self.v_pools[layer].as_ptr() as *mut std::ffi::c_void,
-                    ids_ptr,
-                    src_pos,
-                    dst_pos,
-                    self.num_kv_heads,
-                    self.head_dim,
-                    BLOCK_SIZE,
-                    stream,
-                );
-            }
-        }
-    }
-
     /// Refresh the host-side block table + context lens from `seq_states`,
     /// then upload to GPU. Call once per decode step before the paged kernel.
     pub fn sync_to_gpu(&mut self) {
@@ -599,10 +396,7 @@ impl PagedKVCache {
     /// before advance_seq_len has run).
     pub fn sync_active_batch_with_lens(&mut self, slots: &[usize], kv_lens: &[i32]) {
         assert_eq!(slots.len(), kv_lens.len());
-        assert!(
-            slots.len() <= self.max_seqs,
-            "active batch exceeds max_seqs"
-        );
+        assert!(slots.len() <= self.max_seqs, "active batch exceeds max_seqs");
         let stride = self.max_blocks_per_seq;
         for row in &mut self.block_table_host {
             *row = 0;
@@ -611,9 +405,7 @@ impl PagedKVCache {
             *cl = 0;
         }
         for (i, &slot) in slots.iter().enumerate() {
-            let s = self.seq_states[slot]
-                .as_ref()
-                .expect("unregistered slot in active batch");
+            let s = self.seq_states[slot].as_ref().expect("unregistered slot in active batch");
             let row = &mut self.block_table_host[i * stride..(i + 1) * stride];
             for (j, b) in s.block_ids.iter().enumerate() {
                 row[j] = *b as i32;
@@ -672,12 +464,8 @@ impl PagedKVCache {
                 let src_off = ((phys * nkv + h) * bs + slot_in_blk) * hd * es;
                 let dst_off = (h * sl + p) * hd * es;
                 let count = chunk * hd * es;
-                k_dst
-                    .copy_from_device_at(k_pool, src_off, dst_off, count)
-                    .unwrap();
-                v_dst
-                    .copy_from_device_at(v_pool, src_off, dst_off, count)
-                    .unwrap();
+                k_dst.copy_from_device_at(k_pool, src_off, dst_off, count).unwrap();
+                v_dst.copy_from_device_at(v_pool, src_off, dst_off, count).unwrap();
             }
             p += chunk;
         }
@@ -690,26 +478,16 @@ impl PagedKVCache {
 
     // ----- Swapping (vLLM-style preemption to pinned host memory) -----
 
-    pub fn free_cpu_blocks(&self) -> usize {
-        self.cpu_allocator.free_count()
-    }
-    pub fn swap_enabled(&self) -> bool {
-        !self.cpu_k_pools.is_empty()
-    }
+    pub fn free_cpu_blocks(&self) -> usize { self.cpu_allocator.free_count() }
+    pub fn swap_enabled(&self) -> bool { !self.cpu_k_pools.is_empty() }
 
     pub fn is_swapped(&self, slot: usize) -> bool {
-        matches!(
-            self.seq_states[slot].as_ref().map(|s| s.location),
-            Some(Location::Cpu)
-        )
+        matches!(self.seq_states[slot].as_ref().map(|s| s.location), Some(Location::Cpu))
     }
 
     /// Number of physical blocks currently held by `slot` (in either pool).
     pub fn block_count(&self, slot: usize) -> usize {
-        self.seq_states[slot]
-            .as_ref()
-            .map(|s| s.block_ids.len())
-            .unwrap_or(0)
+        self.seq_states[slot].as_ref().map(|s| s.block_ids.len()).unwrap_or(0)
     }
 
     /// Whether a swapped sequence at `slot` can be brought back (enough free GPU blocks).
@@ -725,17 +503,11 @@ impl PagedKVCache {
     /// Evict `slot`'s KV from GPU to pinned host memory and free its GPU blocks.
     /// The slot stays registered (location = Cpu); the sequence is paused.
     pub fn swap_out(&mut self, slot: usize) -> Result<(), &'static str> {
-        let state = self.seq_states[slot]
-            .as_ref()
-            .ok_or("swap_out: empty slot")?;
-        if state.location == Location::Cpu {
-            return Ok(());
-        }
+        let state = self.seq_states[slot].as_ref().ok_or("swap_out: empty slot")?;
+        if state.location == Location::Cpu { return Ok(()); }
         let gpu_ids = state.block_ids.clone();
         let n = gpu_ids.len();
-        if !self.cpu_allocator.can_alloc(n) {
-            return Err("swap_out: CPU pool full");
-        }
+        if !self.cpu_allocator.can_alloc(n) { return Err("swap_out: CPU pool full"); }
 
         let cpu_ids: Vec<u32> = (0..n)
             .map(|_| self.cpu_allocator.alloc().expect("checked can_alloc"))
@@ -747,18 +519,10 @@ impl PagedKVCache {
                 let g_off = gpu_ids[i] as usize * bb;
                 let c_off = cpu_ids[i] as usize * bb;
                 self.k_pools[layer]
-                    .copy_to_host_at(
-                        &mut self.cpu_k_pools[layer].as_mut_slice()[c_off..c_off + bb],
-                        g_off,
-                        bb,
-                    )
+                    .copy_to_host_at(&mut self.cpu_k_pools[layer].as_mut_slice()[c_off..c_off + bb], g_off, bb)
                     .unwrap();
                 self.v_pools[layer]
-                    .copy_to_host_at(
-                        &mut self.cpu_v_pools[layer].as_mut_slice()[c_off..c_off + bb],
-                        g_off,
-                        bb,
-                    )
+                    .copy_to_host_at(&mut self.cpu_v_pools[layer].as_mut_slice()[c_off..c_off + bb], g_off, bb)
                     .unwrap();
             }
         }
@@ -774,17 +538,11 @@ impl PagedKVCache {
 
     /// Bring `slot`'s KV back from host to GPU and free its CPU blocks.
     pub fn swap_in(&mut self, slot: usize) -> Result<(), &'static str> {
-        let state = self.seq_states[slot]
-            .as_ref()
-            .ok_or("swap_in: empty slot")?;
-        if state.location == Location::Gpu {
-            return Ok(());
-        }
+        let state = self.seq_states[slot].as_ref().ok_or("swap_in: empty slot")?;
+        if state.location == Location::Gpu { return Ok(()); }
         let cpu_ids = state.block_ids.clone();
         let n = cpu_ids.len();
-        if !self.allocator.can_alloc(n) {
-            return Err("swap_in: GPU pool full");
-        }
+        if !self.allocator.can_alloc(n) { return Err("swap_in: GPU pool full"); }
 
         let gpu_ids: Vec<u32> = (0..n)
             .map(|_| self.allocator.alloc().expect("checked can_alloc"))
@@ -796,18 +554,10 @@ impl PagedKVCache {
                 let g_off = gpu_ids[i] as usize * bb;
                 let c_off = cpu_ids[i] as usize * bb;
                 self.k_pools[layer]
-                    .copy_from_host_at(
-                        &self.cpu_k_pools[layer].as_slice()[c_off..c_off + bb],
-                        g_off,
-                        bb,
-                    )
+                    .copy_from_host_at(&self.cpu_k_pools[layer].as_slice()[c_off..c_off + bb], g_off, bb)
                     .unwrap();
                 self.v_pools[layer]
-                    .copy_from_host_at(
-                        &self.cpu_v_pools[layer].as_slice()[c_off..c_off + bb],
-                        g_off,
-                        bb,
-                    )
+                    .copy_from_host_at(&self.cpu_v_pools[layer].as_slice()[c_off..c_off + bb], g_off, bb)
                     .unwrap();
             }
         }
@@ -822,77 +572,7 @@ impl PagedKVCache {
     }
 }
 
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    fn tiny_config() -> ModelConfig {
-        serde_json::from_value(serde_json::json!({
-            "model_type": "qwen3",
-            "hidden_size": 8,
-            "intermediate_size": 16,
-            "num_attention_heads": 1,
-            "num_key_value_heads": 1,
-            "num_hidden_layers": 1,
-            "vocab_size": 32,
-            "max_position_embeddings": 64
-        }))
-        .unwrap()
-    }
-
-    #[test]
-    fn truncate_sequence_frees_whole_blocks_and_keeps_slot_registered() {
-        if xserv_cuda::device::set_device(0).is_err() {
-            eprintln!("skipping CUDA-backed PagedKVCache test: device 0 unavailable");
-            return;
-        }
-
-        let config = tiny_config();
-        let mut cache = PagedKVCache::new(&config, 5, 0, 1, 4, DType::BF16, 0);
-
-        assert_eq!(
-            cache.truncate_sequence(1, 0),
-            Err("truncate_sequence: slot out of range")
-        );
-        assert_eq!(
-            cache.truncate_sequence(0, 0),
-            Err("truncate_sequence: empty slot")
-        );
-
-        cache.register_sequence(0).unwrap();
-        cache.ensure_capacity(0, BLOCK_SIZE * 3 + 1);
-        cache.advance_seq_len(0, BLOCK_SIZE * 3 + 1);
-        assert_eq!(cache.seq_len(0), BLOCK_SIZE * 3 + 1);
-        assert_eq!(cache.block_count(0), 4);
-        assert_eq!(cache.free_blocks(), 0);
-
-        cache.truncate_sequence(0, BLOCK_SIZE + 1).unwrap();
-        assert_eq!(cache.seq_len(0), BLOCK_SIZE + 1);
-        assert_eq!(cache.block_count(0), 2);
-        assert_eq!(cache.free_blocks(), 2);
-
-        cache.truncate_sequence(0, BLOCK_SIZE).unwrap();
-        assert_eq!(cache.seq_len(0), BLOCK_SIZE);
-        assert_eq!(cache.block_count(0), 1);
-        assert_eq!(cache.free_blocks(), 3);
-
-        cache.truncate_sequence(0, 0).unwrap();
-        assert_eq!(cache.seq_len(0), 0);
-        assert_eq!(cache.block_count(0), 1);
-        assert_eq!(cache.free_blocks(), 3);
-        assert_eq!(
-            cache.truncate_sequence(0, 1),
-            Err("truncate_sequence: cannot extend")
-        );
-    }
-}
-
-unsafe fn tensor_from_owned_buf(
-    buf: GpuBuffer,
-    shape: &[usize],
-    dtype: DType,
-    device: u32,
-) -> Tensor {
+unsafe fn tensor_from_owned_buf(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
     use smallvec::SmallVec;
     use xserv_tensor::shape::contiguous_strides;
     use xserv_tensor::storage::Storage;

crates/xserv-model/src/qwen3_graph.rs

@@ -1,185 +0,0 @@
-//! CUDA-graph replay for Qwen3 batch=1 decode (Phase 24 / speculative draft).
-//!
-//! Same pattern as `gpt_oss_graph.rs`, but for the Qwen3 dense decode path used
-//! by speculative decoding's draft model. A Qwen3-0.6B decode step is ~140
-//! kernel launches; wrapping the whole step into one `cudaGraphLaunch` cuts
-//! the ~4× γ draft cost per speculative round.
-//!
-//! See `gpt_oss_graph.rs` for the design commentary; the capture preconditions,
-//! retained-warmup mechanism, and quarantine lifetime are all identical here.
-
-use std::ffi::c_void;
-
-use xserv_cuda::allocator::{self, RetainedBlocks};
-use xserv_cuda::{CudaGraph, CudaStream, GpuBuffer};
-use xserv_tensor::Tensor;
-
-use crate::paged_kv_cache::PagedKVCache;
-use crate::qwen3::Qwen3;
-
-pub struct Qwen3DecodeGraph {
-    stream: CudaStream,
-    graph: CudaGraph,
-    ids_buf: GpuBuffer, // [1] u32, persistent graph input
-    pos_buf: GpuBuffer, // [1] u32, persistent graph input
-    logits: Tensor,     // graph output; rewritten in place by every replay
-    _arena: RetainedBlocks,
-}
-
-impl Qwen3DecodeGraph {
-    /// Capture one batch=1 decode step and replay it once.
-    pub fn capture(
-        model: &Qwen3,
-        token: u32,
-        position: usize,
-        slot: usize,
-        cache: &mut PagedKVCache,
-    ) -> Self {
-        let stream = CudaStream::new().expect("create capture stream");
-        let mut ids_buf = allocator::cached_alloc(4).expect("alloc ids buf");
-        let mut pos_buf = allocator::cached_alloc(4).expect("alloc pos buf");
-
-        model.decode_prepare(&[position], &[slot], cache);
-        ids_buf.copy_from_host(&token.to_le_bytes()).unwrap();
-        pos_buf
-            .copy_from_host(&(position as u32).to_le_bytes())
-            .unwrap();
-
-        // Retained warmup: run the exact step once eagerly with the quarantine
-        // ON to stock the pool. See gpt_oss_graph.rs:66-86 for the full
-        // rationale. Re-running the step is idempotent: the KV scatter
-        // overwrites the same cache position and advance_seq_len is *inside*
-        // decode_core, so we roll it back afterwards.
-        let seq_len_before = cache.seq_len(slot);
-        allocator::begin_retain();
-        {
-            let _guard = xserv_cuda::push_stream(&stream);
-            let _ = model.decode_core(
-                ids_buf.as_ptr() as *const c_void,
-                pos_buf.as_ptr() as *const c_void,
-                1,
-                &[slot],
-                cache,
-            );
-        }
-        drop(allocator::end_retain());
-        stream.synchronize().expect("warmup sync");
-        // decode_core advanced seq_len; roll back so capture starts from the
-        // same logical state as the eager warmup.
-        cache
-            .truncate_sequence(slot, seq_len_before)
-            .expect("rollback after warmup");
-
-        allocator::begin_retain();
-        let mut graph = CudaGraph::new();
-        let logits;
-        {
-            let _guard = xserv_cuda::stream::push_stream(&stream);
-            graph
-                .begin_capture(&stream)
-                .expect("begin decode-graph capture");
-            logits = model.decode_core(
-                ids_buf.as_ptr() as *const c_void,
-                pos_buf.as_ptr() as *const c_void,
-                1,
-                &[slot],
-                cache,
-            );
-            graph
-                .end_capture(&stream)
-                .expect("end decode-graph capture");
-        }
-        let arena = allocator::end_retain();
-
-        // The capture path called advance_seq_len (host-side) but the actual
-        // GPU compute has not yet run. Roll back and let the first replay
-        // advance it exactly once with real K/V writes.
-        cache
-            .truncate_sequence(slot, seq_len_before)
-            .expect("rollback after capture");
-
-        graph.launch(&stream).expect("first decode-graph replay");
-        cache.advance_seq_len(slot, 1);
-
-        Self {
-            stream,
-            graph,
-            ids_buf,
-            pos_buf,
-            logits,
-            _arena: arena,
-        }
-    }
-
-    /// Run one decode step by replaying the captured graph.
-    pub fn step(
-        &mut self,
-        model: &Qwen3,
-        token: u32,
-        position: usize,
-        slot: usize,
-        cache: &mut PagedKVCache,
-    ) -> Tensor {
-        model.decode_prepare(&[position], &[slot], cache);
-        self.ids_buf.copy_from_host(&token.to_le_bytes()).unwrap();
-        self.pos_buf
-            .copy_from_host(&(position as u32).to_le_bytes())
-            .unwrap();
-        self.graph
-            .launch(&self.stream)
-            .expect("decode-graph replay");
-        cache.advance_seq_len(slot, 1);
-        self.logits.clone()
-    }
-}
-
-/// Lazy capture policy: first decode step of the process runs eager, the
-/// second is captured, the rest replay. Batch>1 always falls back to eager.
-/// Disable with `XSERV_DECODE_GRAPH=0`.
-pub struct GraphedQwen3Decoder {
-    graph: Option<Qwen3DecodeGraph>,
-    eager_steps: u32,
-    enabled: bool,
-}
-
-impl GraphedQwen3Decoder {
-    pub fn new() -> Self {
-        let enabled = std::env::var("XSERV_DECODE_GRAPH")
-            .map(|v| v != "0")
-            .unwrap_or(true);
-        Self {
-            graph: None,
-            eager_steps: 0,
-            enabled,
-        }
-    }
-
-    pub fn decode(
-        &mut self,
-        model: &Qwen3,
-        tokens: &[u32],
-        positions: &[usize],
-        slots: &[usize],
-        cache: &mut PagedKVCache,
-    ) -> Tensor {
-        if self.enabled && tokens.len() == 1 {
-            if let Some(g) = self.graph.as_mut() {
-                return g.step(model, tokens[0], positions[0], slots[0], cache);
-            }
-            if self.eager_steps >= 1 {
-                let g = Qwen3DecodeGraph::capture(model, tokens[0], positions[0], slots[0], cache);
-                let logits = g.logits.clone();
-                self.graph = Some(g);
-                return logits;
-            }
-        }
-        self.eager_steps += 1;
-        model.forward_decode_paged(tokens, positions, slots, cache)
-    }
-}
-
-impl Default for GraphedQwen3Decoder {
-    fn default() -> Self {
-        Self::new()
-    }
-}

crates/xserv-model/src/sampling.rs

@@ -11,11 +11,7 @@ pub struct SamplingParams {
 
 impl Default for SamplingParams {
     fn default() -> Self {
-        Self {
-            temperature: 0.0,
-            top_k: 0,
-            top_p: 1.0,
-        }
+        Self { temperature: 0.0, top_k: 0, top_p: 1.0 }
     }
 }
 
@@ -23,24 +19,12 @@ impl Default for SamplingParams {
 /// Uses the last position's logits. Handles both F32 and BF16 dtypes.
 pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
     assert_eq!(logits.ndim(), 2);
-    // Greedy fast path: GPU argmax + 4-byte D2H instead of copying the whole
-    // [seq, vocab] logits to the host and scanning it (~201k bf16/token).
-    // NaN logits lose every `>` comparison in the kernel, matching the
-    // NaN-safe host argmax below.
-    if params.temperature == 0.0
-        && logits.dtype() == DType::BF16
-        && matches!(logits.device(), Device::Cuda(_))
-        && logits.is_contiguous()
-    {
-        let ids = xserv_kernels::argmax_bf16_to_host(logits);
-        return *ids.last().unwrap();
-    }
     let vocab_size = logits.shape()[1];
     let seq_len = logits.shape()[0];
     let logits_cpu = logits.to_device(Device::Cpu);
 
     // Extract last row as f32
-    let mut last_row: Vec<f32> = match logits.dtype() {
+    let last_row: Vec<f32> = match logits.dtype() {
         DType::F32 => {
             let data = logits_cpu.as_slice::<f32>();
             data[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec()
@@ -60,20 +44,6 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
         return argmax(&last_row);
     }
 
-    // NaN-safe: sampling path uses partial_cmp().unwrap() in top-k/top-p
-    // sorts and softmax; a single NaN logit would panic the engine thread.
-    // Replace NaN with -inf (equivalent to masking) instead.
-    let mut nan_seen = false;
-    for v in last_row.iter_mut() {
-        if v.is_nan() {
-            nan_seen = true;
-            *v = f32::NEG_INFINITY;
-        }
-    }
-    if nan_seen {
-        eprintln!("[sampling] WARNING: NaN logits encountered in sample()");
-    }
-
     // Apply temperature
     let mut logits_f32: Vec<f32> = last_row.iter().map(|v| v / params.temperature).collect();
 
@@ -142,56 +112,10 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
     (vocab_size - 1) as u32
 }
 
-/// Greedy argmax with a repetition penalty applied to `recent` token ids
-/// (HF-style: divide positive logits by `penalty`, multiply negative by it).
-/// `penalty <= 1.0` is a no-op. Mitigates greedy repetition loops on reasoning
-/// models without changing the forward pass. NaN-safe.
-pub fn sample_greedy_penalized(logits: &Tensor, recent: &[u32], penalty: f32) -> u32 {
-    assert_eq!(logits.ndim(), 2);
-    let vocab_size = logits.shape()[1];
-    let seq_len = logits.shape()[0];
-    let logits_cpu = logits.to_device(Device::Cpu);
-    let mut last_row: Vec<f32> = match logits.dtype() {
-        DType::F32 => {
-            logits_cpu.as_slice::<f32>()[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec()
-        }
-        DType::BF16 => logits_cpu.as_slice::<bf16>()
-            [(seq_len - 1) * vocab_size..seq_len * vocab_size]
-            .iter()
-            .map(|v| v.to_f32())
-            .collect(),
-        _ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
-    };
-    if penalty > 1.0 {
-        for &id in recent {
-            let i = id as usize;
-            if i < last_row.len() {
-                let v = last_row[i];
-                last_row[i] = if v > 0.0 { v / penalty } else { v * penalty };
-            }
-        }
-    }
-    argmax(&last_row)
-}
-
 fn argmax(data: &[f32]) -> u32 {
-    // NaN-safe: a single NaN logit must not crash the engine thread (a
-    // partial_cmp().unwrap() panics on NaN). Skip NaNs; warn once if seen.
-    let mut best_i = 0usize;
-    let mut best = f32::NEG_INFINITY;
-    let mut nan_seen = false;
-    for (i, &v) in data.iter().enumerate() {
-        if v.is_nan() {
-            nan_seen = true;
-            continue;
-        }
-        if v > best {
-            best = v;
-            best_i = i;
-        }
-    }
-    if nan_seen {
-        eprintln!("[sampling] WARNING: NaN logits encountered in argmax");
-    }
-    best_i as u32
+    data.iter()
+        .enumerate()
+        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+        .map(|(i, _)| i as u32)
+        .unwrap()
 }

crates/xserv-server/Cargo.toml

@@ -21,4 +21,3 @@ tokio.workspace = true
 axum.workspace = true
 uuid.workspace = true
 tokio-stream.workspace = true
-minijinja.workspace = true

crates/xserv-server/src/api.rs

@@ -5,7 +5,6 @@ use axum::response::sse::{Event, KeepAlive, Sse};
 use axum::response::{IntoResponse, Response};
 use serde::{Deserialize, Serialize};
 use std::convert::Infallible;
-use std::path::Path;
 use std::sync::Arc;
 use tokio_stream::StreamExt;
 use tokio_stream::wrappers::ReceiverStream;
@@ -32,7 +31,7 @@ pub struct ChatRequest {
     pub top_p: Option<f32>,
 }
 
-#[derive(Deserialize, Serialize, Clone)]
+#[derive(Deserialize)]
 pub struct Message {
     pub role: String,
     pub content: String,
@@ -55,207 +54,6 @@ pub struct ModelInfo {
     owned_by: &'static str,
 }
 
-// ---------------------------------------------------------------------------
-// Chat Template: Jinja2 rendering via minijinja
-// ---------------------------------------------------------------------------
-
-pub struct ChatTemplate {
-    source: String,
-    model_type: String,
-}
-
-impl ChatTemplate {
-    pub fn load(model_dir: &Path, model_type: &str) -> Self {
-        // 1. Try standalone chat_template.jinja file
-        let jinja_path = model_dir.join("chat_template.jinja");
-        if jinja_path.exists() {
-            let source = std::fs::read_to_string(&jinja_path)
-                .unwrap_or_else(|e| panic!("failed to read {}: {e}", jinja_path.display()));
-            eprintln!("[chat-template] loaded from {}", jinja_path.display());
-            return Self {
-                source,
-                model_type: model_type.to_string(),
-            };
-        }
-
-        // 2. Try tokenizer_config.json → chat_template field
-        let tok_cfg_path = model_dir.join("tokenizer_config.json");
-        if tok_cfg_path.exists() {
-            if let Ok(data) = std::fs::read_to_string(&tok_cfg_path) {
-                if let Ok(v) = serde_json::from_str::<serde_json::Value>(&data) {
-                    if let Some(ct) = v.get("chat_template").and_then(|v| v.as_str()) {
-                        eprintln!("[chat-template] loaded from tokenizer_config.json");
-                        return Self {
-                            source: ct.to_string(),
-                            model_type: model_type.to_string(),
-                        };
-                    }
-                }
-            }
-        }
-
-        // 3. No template found — use empty source, will fall back to hardcoded
-        eprintln!("[chat-template] no Jinja template found, using hardcoded fallback");
-        Self {
-            source: String::new(),
-            model_type: model_type.to_string(),
-        }
-    }
-
-    pub fn render(&self, messages: &[Message]) -> String {
-        if self.source.is_empty() {
-            return build_prompt_hardcoded(messages, &self.model_type);
-        }
-
-        match self.render_jinja(messages) {
-            Ok(prompt) => prompt,
-            Err(e) => {
-                eprintln!("[chat-template] Jinja render error: {e}, falling back to hardcoded");
-                build_prompt_hardcoded(messages, &self.model_type)
-            }
-        }
-    }
-
-    fn render_jinja(&self, messages: &[Message]) -> Result<String, minijinja::Error> {
-        let mut env = minijinja::Environment::new();
-
-        // Register custom functions the template may call.
-        env.add_function("strftime_now", strftime_now);
-        env.add_function("raise_exception", raise_exception);
-
-        // Python str methods used by harmony/gpt-oss templates.
-        env.add_filter("startswith", |s: String, prefix: String| -> bool {
-            s.starts_with(&prefix)
-        });
-
-        env.add_template("chat", &self.source)?;
-        let tmpl = env.get_template("chat")?;
-
-        let ctx = minijinja::context! {
-            messages => minijinja::Value::from_serialize(messages),
-            add_generation_prompt => true,
-            bos_token => "",
-            eos_token => "",
-        };
-
-        tmpl.render(ctx)
-    }
-}
-
-fn strftime_now(fmt: String) -> String {
-    use std::time::SystemTime;
-    let now = SystemTime::now()
-        .duration_since(SystemTime::UNIX_EPOCH)
-        .unwrap()
-        .as_secs();
-    // Only support %Y-%m-%d (the only format used by known templates)
-    let days = now / 86400;
-    let (y, m, d) = days_to_ymd(days);
-    fmt.replace("%Y", &format!("{y:04}"))
-        .replace("%m", &format!("{m:02}"))
-        .replace("%d", &format!("{d:02}"))
-}
-
-fn days_to_ymd(days_since_epoch: u64) -> (u32, u32, u32) {
-    // Civil calendar from days since 1970-01-01 (Rata Die algorithm)
-    let z = days_since_epoch as i64 + 719468;
-    let era = (if z >= 0 { z } else { z - 146096 }) / 146097;
-    let doe = (z - era * 146097) as u32;
-    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
-    let y = yoe as i64 + era * 400;
-    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
-    let mp = (5 * doy + 2) / 153;
-    let d = doy - (153 * mp + 2) / 5 + 1;
-    let m = if mp < 10 { mp + 3 } else { mp - 9 };
-    let y = if m <= 2 { y + 1 } else { y };
-    (y as u32, m, d)
-}
-
-fn raise_exception(msg: String) -> Result<String, minijinja::Error> {
-    Err(minijinja::Error::new(
-        minijinja::ErrorKind::InvalidOperation,
-        msg,
-    ))
-}
-
-// ---------------------------------------------------------------------------
-// Hardcoded fallback templates (for models without a Jinja template)
-// ---------------------------------------------------------------------------
-
-fn build_prompt_hardcoded(messages: &[Message], model_type: &str) -> String {
-    if model_type == "gpt_oss" {
-        return build_prompt_gpt_oss(messages);
-    }
-    // Default: Qwen3 ChatML format
-    let mut prompt = String::new();
-    for msg in messages {
-        match msg.role.as_str() {
-            "system" | "user" | "assistant" => {
-                prompt.push_str("<|im_start|>");
-                prompt.push_str(&msg.role);
-                prompt.push('\n');
-                prompt.push_str(&msg.content);
-                prompt.push_str("<|im_end|>\n");
-            }
-            _ => {}
-        }
-    }
-    prompt.push_str("<|im_start|>assistant\n");
-    prompt.push_str("<think>\n\n</think>\n\n");
-    prompt
-}
-
-fn build_prompt_gpt_oss(messages: &[Message]) -> String {
-    let mut prompt = String::new();
-    // Canonical harmony system message (mirrors the model's chat_template.jinja
-    // build_system_message macro). A hand-rolled substitute puts gpt-oss out of
-    // distribution and destabilizes channel selection. This hardcoded builder is
-    // only a fallback for gpt-oss models that ship no Jinja template; the
-    // gpt-oss-20b release does ship one, so the template path is normally used.
-    prompt.push_str("<|start|>system<|message|>");
-    prompt.push_str("You are ChatGPT, a large language model trained by OpenAI.\n");
-    prompt.push_str("Knowledge cutoff: 2024-06\n");
-    prompt.push_str(&format!(
-        "Current date: {}\n\n",
-        strftime_now("%Y-%m-%d".to_string())
-    ));
-    prompt.push_str("Reasoning: low\n\n");
-    prompt.push_str("# Valid channels: analysis, commentary, final. Channel must be included for every message.");
-    prompt.push_str("<|end|>");
-    let dev_instructions: String = messages
-        .iter()
-        .filter(|m| m.role == "system")
-        .map(|m| m.content.as_str())
-        .collect::<Vec<_>>()
-        .join("\n\n");
-    if !dev_instructions.is_empty() {
-        prompt.push_str("<|start|>developer<|message|># Instructions\n\n");
-        prompt.push_str(&dev_instructions);
-        prompt.push_str("<|end|>");
-    }
-    for msg in messages {
-        match msg.role.as_str() {
-            "user" => {
-                prompt.push_str("<|start|>user<|message|>");
-                prompt.push_str(&msg.content);
-                prompt.push_str("<|end|>");
-            }
-            "assistant" => {
-                prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
-                prompt.push_str(&msg.content);
-                prompt.push_str("<|end|>");
-            }
-            _ => {}
-        }
-    }
-    prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
-    prompt
-}
-
-// ---------------------------------------------------------------------------
-// HTTP handlers
-// ---------------------------------------------------------------------------
-
 pub async fn health() -> &'static str {
     "ok"
 }
@@ -291,7 +89,7 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         return response;
     }
 
-    let prompt = state.chat_template.render(&req.messages);
+    let prompt = build_prompt(&req.messages);
     let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
     let prompt_token_count = prompt_tokens.len();
 
@@ -331,10 +129,6 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         }
     }
 
-    let fr_value = match normalize_finish_reason(&finish_reason) {
-        Some(s) => serde_json::Value::String(s.to_string()),
-        None => serde_json::Value::Null,
-    };
     Json(serde_json::json!({
         "id": id,
         "object": "chat.completion",
@@ -343,18 +137,20 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         "choices": [{
             "index": 0,
             "message": { "role": "assistant", "content": content },
-            "finish_reason": fr_value,
+            "finish_reason": finish_reason,
         }],
         "usage": {
             "prompt_tokens": prompt_token_count,
             "completion_tokens": completion_token_count,
             "total_tokens": prompt_token_count + completion_token_count
         }
-    }))
-    .into_response()
+    })).into_response()
 }
 
-fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
+fn chat_stream(
+    state: Arc<AppState>,
+    req: ChatRequest,
+) -> Response {
     let id = format!("chatcmpl-{}", Uuid::new_v4());
     let model_name = state.model_name.clone();
     let created = unix_timestamp();
@@ -363,15 +159,14 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         return response;
     }
 
-    let prompt = state.chat_template.render(&req.messages);
+    let prompt = build_prompt(&req.messages);
     let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
 
     let max_seq_len = state.max_seq_len;
     if prompt_tokens.len() >= max_seq_len {
         return bad_request(format!(
             "prompt is {} tokens, exceeds max_seq_len {}",
-            prompt_tokens.len(),
-            max_seq_len
+            prompt_tokens.len(), max_seq_len
         ));
     }
     let max_tokens = req.max_tokens.min(max_seq_len - prompt_tokens.len());
@@ -416,11 +211,8 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
                             make_chunk(&id, &model_name, created, None, Some("assistant"), None);
                         let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
                     }
-                    // Only "stop" and "length" are OpenAI-standard values. Internal
-                    // codes like "error" (client-stalled from tp/pp engine) map to
-                    // null so SDK clients see a clean stream close.
-                    let fr = normalize_finish_reason(&finish_reason);
-                    let chunk = make_chunk(&id, &model_name, created, None, None, fr);
+                    let chunk =
+                        make_chunk(&id, &model_name, created, None, None, Some(&finish_reason));
                     let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
                     let _ = sse_tx
                         .send(Ok(Event::default().data("[DONE]".to_string())))
@@ -431,9 +223,7 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
         }
     });
 
-    Sse::new(ReceiverStream::new(sse_rx))
-        .keep_alive(KeepAlive::default())
-        .into_response()
+    Sse::new(ReceiverStream::new(sse_rx)).keep_alive(KeepAlive::default()).into_response()
 }
 
 fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
@@ -449,22 +239,6 @@ fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
         return Some(bad_request("max_tokens must be greater than 0"));
     }
 
-    if let Some(t) = req.temperature {
-        if !t.is_finite() || t < 0.0 {
-            return Some(bad_request("temperature must be a finite value >= 0"));
-        }
-    }
-    if let Some(p) = req.top_p {
-        if !p.is_finite() || !(0.0..=1.0).contains(&p) {
-            return Some(bad_request("top_p must be in [0, 1]"));
-        }
-    }
-    if let Some(k) = req.top_k {
-        if k > 1_000_000 {
-            return Some(bad_request("top_k must be <= 1_000_000"));
-        }
-    }
-
     None
 }
 
@@ -472,18 +246,8 @@ fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
 /// prior handler panicked) and returns a clean 503 instead of panicking when the
 /// engine thread is gone, so one dead engine doesn't cascade into every request.
 fn submit_to_engine(state: &AppState, req: GenerateRequest) -> Result<(), Response> {
-    let sender = state
-        .engine_sender
-        .lock()
-        .unwrap_or_else(|e| e.into_inner());
-    sender.try_send(req).map_err(|err| match err {
-        std::sync::mpsc::TrySendError::Full(_) => {
-            service_unavailable("inference engine is busy, retry later")
-        }
-        std::sync::mpsc::TrySendError::Disconnected(_) => {
-            service_unavailable("inference engine is not available")
-        }
-    })
+    let sender = state.engine_sender.lock().unwrap_or_else(|e| e.into_inner());
+    sender.send(req).map_err(|_| service_unavailable("inference engine is not available"))
 }
 
 fn service_unavailable(message: impl Into<String>) -> Response {
@@ -561,13 +325,21 @@ fn sampling_params(req: &ChatRequest) -> SamplingParams {
     }
 }
 
-/// Map engine finish_reason strings to OpenAI-standard values. Any engine-internal
-/// code (e.g. "error" from tp/pp client-stall) collapses to None so SDK clients see
-/// a clean null instead of an unknown value.
-fn normalize_finish_reason(fr: &str) -> Option<&'static str> {
-    match fr {
-        "stop" => Some("stop"),
-        "length" => Some("length"),
-        _ => None,
+fn build_prompt(messages: &[Message]) -> String {
+    let mut prompt = String::new();
+    for msg in messages {
+        match msg.role.as_str() {
+            "system" | "user" | "assistant" => {
+                prompt.push_str("<|im_start|>");
+                prompt.push_str(&msg.role);
+                prompt.push('\n');
+                prompt.push_str(&msg.content);
+                prompt.push_str("<|im_end|>\n");
+            }
+            _ => {}
+        }
     }
+    prompt.push_str("<|im_start|>assistant\n");
+    prompt.push_str("<think>\n\n</think>\n\n");
+    prompt
 }

crates/xserv-server/src/engine.rs

@@ -1,10 +1,10 @@
 use std::collections::VecDeque;
 use std::path::Path;
-use std::sync::Once;
 use std::sync::mpsc;
+use std::sync::Once;
 use std::time::Instant;
+use xserv_model::{ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample, BLOCK_SIZE};
 use xserv_model::loader;
-use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample};
 use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
@@ -38,9 +38,6 @@ struct Sequence {
     seq_slot: Option<usize>,
     sender: tokio::sync::mpsc::Sender<GenerateEvent>,
     prefilled: bool,
-    /// Set when a `try_send` failed (client too slow or gone). The scheduler
-    /// reaps the sequence next iteration instead of blocking the decode thread.
-    client_stalled: bool,
     eos_token_id: Option<u32>,
     decode_buffer: Vec<u8>,
     created_at: Instant,
@@ -112,23 +109,12 @@ impl Engine {
             (total_blocks * bytes_per_block) as f64 / 1e9,
             info.free_memory as f64 / 1e9,
         );
-        Self {
-            model,
-            config,
-            tokenizer,
-            max_batch_size,
-            max_seq_len,
-            paged_cache,
-        }
+        Self { model, config, tokenizer, max_batch_size, max_seq_len, paged_cache }
     }
 
-    pub fn tokenizer(&self) -> &Tokenizer {
-        &self.tokenizer
-    }
+    pub fn tokenizer(&self) -> &Tokenizer { &self.tokenizer }
 
-    pub fn max_seq_len(&self) -> usize {
-        self.max_seq_len
-    }
+    pub fn max_seq_len(&self) -> usize { self.max_seq_len }
 
     /// Main scheduler loop. Receives requests from channel, manages concurrent sequences.
     ///
@@ -148,8 +134,7 @@ impl Engine {
 
         loop {
             // Step 1: Remove finished sequences and return their slots.
-            let finished_slots: Vec<usize> = running
-                .iter()
+            let finished_slots: Vec<usize> = running.iter()
                 .filter(|s| is_finished(s))
                 .filter_map(|s| s.seq_slot)
                 .collect();
@@ -162,16 +147,10 @@ impl Engine {
             // room (oldest first). They resume decoding from where they paused.
             while running.len() < self.max_batch_size && !swapped.is_empty() {
                 let slot = swapped[0].seq_slot.expect("swapped slot");
-                if !self.paged_cache.can_swap_in(slot) {
-                    break;
-                }
+                if !self.paged_cache.can_swap_in(slot) { break; }
                 self.paged_cache.swap_in(slot).expect("swap_in");
                 let seq = swapped.remove(0);
-                eprintln!(
-                    "[scheduler] swapped in seq {} ({} blocks)",
-                    seq.id,
-                    self.paged_cache.block_count(slot)
-                );
+                eprintln!("[scheduler] swapped in seq {} ({} blocks)", seq.id, self.paged_cache.block_count(slot));
                 running.push(seq);
             }
 
@@ -182,22 +161,14 @@ impl Engine {
                 let mut avail = self.paged_cache.free_blocks();
                 let decode_reserve = running.len();
                 while running.len() < self.max_batch_size {
-                    let Some(front) = waiting.front() else {
-                        break;
-                    };
+                    let Some(front) = waiting.front() else { break; };
                     let prompt_blocks = front.prompt_tokens.len().div_ceil(BLOCK_SIZE).max(1);
-                    if avail < prompt_blocks + decode_reserve {
-                        break;
-                    }
+                    if avail < prompt_blocks + decode_reserve { break; }
                     let free_slot = (0..self.paged_cache.max_seqs())
                         .find(|&s| self.paged_cache.is_slot_free(s));
-                    let Some(slot) = free_slot else {
-                        break;
-                    };
+                    let Some(slot) = free_slot else { break; };
                     let mut seq = waiting.pop_front().unwrap();
-                    self.paged_cache
-                        .register_sequence(slot)
-                        .expect("register paged slot");
+                    self.paged_cache.register_sequence(slot).expect("register paged slot");
                     seq.seq_slot = Some(slot);
                     running.push(seq);
                     avail -= prompt_blocks; // projected free after this seq prefills
@@ -228,9 +199,7 @@ impl Engine {
                 if !seq.prefilled {
                     let slot = seq.seq_slot.expect("slot");
                     let logits = self.model.forward_prefill_paged(
-                        &seq.prompt_tokens,
-                        slot,
-                        &mut self.paged_cache,
+                        &seq.prompt_tokens, slot, &mut self.paged_cache,
                     );
                     let next = sample(&logits, &seq.sampling);
                     seq.generated_tokens.push(next);
@@ -250,18 +219,13 @@ impl Engine {
                         && !newly_prefilled.contains(&running[p].id)
                         && running[p].seq_slot.is_some()
                 });
-                let Some(pos) = victim else {
-                    break;
-                };
+                let Some(pos) = victim else { break; };
                 let seq = running.remove(pos);
                 let slot = seq.seq_slot.unwrap();
                 if self.paged_cache.can_swap_out(slot) {
                     let nblocks = self.paged_cache.block_count(slot);
                     self.paged_cache.swap_out(slot).expect("swap_out");
-                    eprintln!(
-                        "[scheduler] preempt: swapped out seq {} ({nblocks} blocks) to host",
-                        seq.id
-                    );
+                    eprintln!("[scheduler] preempt: swapped out seq {} ({nblocks} blocks) to host", seq.id);
                     swapped.push(seq);
                     needed = decode_block_need(&self.paged_cache, &running, &newly_prefilled);
                 } else {
@@ -271,9 +235,7 @@ impl Engine {
             }
 
             // Step 5c: Batched paged decode for the surviving prefilled sequences.
-            let decode_indices: Vec<usize> = running
-                .iter()
-                .enumerate()
+            let decode_indices: Vec<usize> = running.iter().enumerate()
                 .filter(|(_, s)| s.prefilled && !newly_prefilled.contains(&s.id))
                 .map(|(i, _)| i)
                 .collect();
@@ -284,66 +246,37 @@ impl Engine {
                     eprintln!("[scheduler] paged decode active");
                 });
 
-                let tokens: Vec<u32> = decode_indices
-                    .iter()
+                let tokens: Vec<u32> = decode_indices.iter()
                     .map(|&i| *running[i].generated_tokens.last().unwrap())
                     .collect();
-                let positions: Vec<usize> = decode_indices
-                    .iter()
+                let positions: Vec<usize> = decode_indices.iter()
                     .map(|&i| self.paged_cache.seq_len(running[i].seq_slot.unwrap()))
                     .collect();
-                let slots: Vec<usize> = decode_indices
-                    .iter()
+                let slots: Vec<usize> = decode_indices.iter()
                     .map(|&i| running[i].seq_slot.unwrap())
                     .collect();
 
                 let logits = self.model.forward_decode_paged(
-                    &tokens,
-                    &positions,
-                    &slots,
-                    &mut self.paged_cache,
+                    &tokens, &positions, &slots, &mut self.paged_cache,
                 );
 
-                // Fast path: every active sequence is greedy → run argmax on
-                // the GPU and only D2H the chosen token ids (a few bytes per
-                // sequence) instead of the full [B, vocab_size] BF16 logits
-                // (~1.2 MB for B=4, Qwen3 vocab=152K).
-                let all_greedy = decode_indices
-                    .iter()
-                    .all(|&i| running[i].sampling.temperature == 0.0);
-                if all_greedy {
-                    let next_ids = xserv_kernels::argmax_bf16_to_host(&logits);
-                    for (j, &i) in decode_indices.iter().enumerate() {
-                        let next = next_ids[j];
-                        running[i].generated_tokens.push(next);
-                        emit_token(&self.tokenizer, &mut running[i], next);
-                    }
-                } else {
-                    // Mixed sampling: keep the CPU path for now (top-k/top-p
-                    // sampling still runs there). Only the rows that need it
-                    // get exercised; greedy rows could in principle reuse the
-                    // GPU argmax but the CPU pass is short for B<=4.
-                    let vocab_size = logits.shape()[1];
-                    let logits_cpu = logits.to_device(xserv_tensor::Device::Cpu);
-                    let data = logits_cpu.as_slice::<half::bf16>();
-                    for (j, &i) in decode_indices.iter().enumerate() {
-                        let row_start = j * vocab_size;
-                        let row_logits = &data[row_start..row_start + vocab_size];
-                        let next = if running[i].sampling.temperature == 0.0 {
-                            row_logits
-                                .iter()
-                                .enumerate()
-                                .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
-                                .map(|(idx, _)| idx as u32)
-                                .unwrap()
-                        } else {
-                            let row_tensor =
-                                xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
-                            sample(&row_tensor, &running[i].sampling)
-                        };
-                        running[i].generated_tokens.push(next);
-                        emit_token(&self.tokenizer, &mut running[i], next);
-                    }
+                // Sample per-sequence from batched logits [B, vocab_size]
+                let vocab_size = logits.shape()[1];
+                let logits_cpu = logits.to_device(xserv_tensor::Device::Cpu);
+                let data = logits_cpu.as_slice::<half::bf16>();
+                for (j, &i) in decode_indices.iter().enumerate() {
+                    let row_start = j * vocab_size;
+                    let row_logits = &data[row_start..row_start + vocab_size];
+                    let next = if running[i].sampling.temperature == 0.0 {
+                        row_logits.iter().enumerate()
+                            .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
+                            .map(|(idx, _)| idx as u32).unwrap()
+                    } else {
+                        let row_tensor = xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
+                        sample(&row_tensor, &running[i].sampling)
+                    };
+                    running[i].generated_tokens.push(next);
+                    emit_token(&self.tokenizer, &mut running[i], next);
                 }
             }
 
@@ -373,7 +306,6 @@ impl Engine {
             seq_slot: None,
             sender: req.sender,
             prefilled: false,
-            client_stalled: false,
             eos_token_id: self.tokenizer.eos_token_id(),
             decode_buffer: Vec::new(),
             created_at: Instant::now(),
@@ -384,8 +316,7 @@ impl Engine {
 /// Total additional GPU blocks the next decode step needs across all
 /// currently-decoding (prefilled, not just-prefilled) sequences.
 fn decode_block_need(paged: &PagedKVCache, running: &[Sequence], newly_prefilled: &[u64]) -> usize {
-    running
-        .iter()
+    running.iter()
         .filter(|s| s.prefilled && !newly_prefilled.contains(&s.id))
         .filter_map(|s| s.seq_slot)
         .map(|slot| paged.additional_blocks_needed(slot, 1))
@@ -396,12 +327,9 @@ fn emit_token(tokenizer: &Tokenizer, seq: &mut Sequence, token_id: u32) {
     if tokenizer.eos_token_id() == Some(token_id) {
         let tail = tokenizer.flush_decode_stream(&mut seq.decode_buffer);
         send_token_if_nonempty(seq, tail);
-        try_send_event(
-            seq,
-            GenerateEvent::Done {
-                finish_reason: "stop".to_string(),
-            },
-        );
+        let _ = seq.sender.blocking_send(GenerateEvent::Done {
+            finish_reason: "stop".to_string(),
+        });
         return;
     }
 
@@ -410,51 +338,24 @@ fn emit_token(tokenizer: &Tokenizer, seq: &mut Sequence, token_id: u32) {
         let tail = tokenizer.flush_decode_stream(&mut seq.decode_buffer);
         send_token_if_nonempty(seq, text);
         send_token_if_nonempty(seq, tail);
-        try_send_event(
-            seq,
-            GenerateEvent::Done {
-                finish_reason: "length".to_string(),
-            },
-        );
+        let _ = seq.sender.blocking_send(GenerateEvent::Done {
+            finish_reason: "length".to_string(),
+        });
     } else {
         send_token_if_nonempty(seq, text);
     }
 }
 
-fn send_token_if_nonempty(seq: &mut Sequence, text: String) {
+fn send_token_if_nonempty(seq: &Sequence, text: String) {
     if !text.is_empty() {
         let id = *seq.generated_tokens.last().unwrap_or(&0);
-        try_send_event(seq, GenerateEvent::Token { id, text });
-    }
-}
-
-/// Send an event without blocking the shared decode thread. If the client is
-/// too slow (channel full) or gone (closed), flag the sequence for eviction
-/// instead of blocking — one slow consumer must never stall the whole
-/// continuous-batching loop. When the sequence is reaped its `sender` drops,
-/// closing the channel so the client's receive loop ends rather than hanging.
-fn try_send_event(seq: &mut Sequence, event: GenerateEvent) {
-    if let Err(err) = seq.sender.try_send(event) {
-        seq.client_stalled = true;
-        if let tokio::sync::mpsc::error::TrySendError::Full(_) = err {
-            eprintln!(
-                "[scheduler] seq {}: client too slow (stream channel full), evicting",
-                seq.id
-            );
-        }
+        let _ = seq.sender.blocking_send(GenerateEvent::Token { id, text });
     }
 }
 
 fn is_finished(seq: &Sequence) -> bool {
-    if seq.client_stalled {
-        return true;
-    }
-    if seq.generated_tokens.is_empty() {
-        return false;
-    }
+    if seq.generated_tokens.is_empty() { return false; }
     let last = *seq.generated_tokens.last().unwrap();
-    if seq.generated_tokens.len() >= seq.max_tokens {
-        return true;
-    }
+    if seq.generated_tokens.len() >= seq.max_tokens { return true; }
     seq.sender.is_closed() || seq.eos_token_id == Some(last)
 }

crates/xserv-server/src/main.rs

@@ -3,20 +3,15 @@ mod engine;
 mod pp_engine;
 mod tp_engine;
 
-use axum::{
-    Extension, Router,
-    extract::DefaultBodyLimit,
-    routing::{get, post},
-};
-use engine::GenerateRequest;
+use axum::{routing::{get, post}, Extension, Router};
 use std::path::PathBuf;
-use std::sync::{Arc, Mutex, mpsc};
+use std::sync::{mpsc, Arc, Mutex};
+use engine::GenerateRequest;
 use xserv_model::ModelConfig;
 
 pub struct AppState {
     pub model_name: String,
-    pub chat_template: api::ChatTemplate,
-    pub engine_sender: Mutex<mpsc::SyncSender<GenerateRequest>>,
+    pub engine_sender: Mutex<mpsc::Sender<GenerateRequest>>,
     pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
     pub max_seq_len: usize,
 }
@@ -25,48 +20,40 @@ pub struct AppState {
 async fn main() {
     let args: Vec<String> = std::env::args().collect();
     if args.len() < 2 {
-        eprintln!(
-            "Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N] [--tp N] [--pp N]"
-        );
+        eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N] [--tp N] [--pp N]");
         std::process::exit(1);
     }
 
     let model_dir = PathBuf::from(&args[1]);
-    let port: u16 = args
-        .iter()
+    let port: u16 = args.iter()
         .position(|a| a == "--port")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
         .unwrap_or(8080);
-    let max_batch: usize = args
-        .iter()
+    let max_batch: usize = args.iter()
         .position(|a| a == "--max-batch")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
         .unwrap_or(4)
         .max(1);
-    let requested_max_seq_len: usize = args
-        .iter()
+    let requested_max_seq_len: usize = args.iter()
         .position(|a| a == "--max-seq-len")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
         .unwrap_or(2048)
         .max(1);
-    let swap_space_gb: usize = args
-        .iter()
+    let swap_space_gb: usize = args.iter()
         .position(|a| a == "--swap-space-gb")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
         .unwrap_or(8);
-    let tp: usize = args
-        .iter()
+    let tp: usize = args.iter()
         .position(|a| a == "--tp")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
         .unwrap_or(1)
         .max(1);
-    let pp: usize = args
-        .iter()
+    let pp: usize = args.iter()
         .position(|a| a == "--pp")
         .and_then(|i| args.get(i + 1))
         .and_then(|s| s.parse().ok())
@@ -77,15 +64,6 @@ async fn main() {
         std::process::exit(1);
     }
     let model_config = ModelConfig::from_file(&model_dir.join("config.json"));
-    // gpt-oss is only implemented in the TP engine; route it there even at
-    // tp=1 (single-rank world) so quantized models can serve on one GPU.
-    let is_gpt_oss = model_config.model_type.as_deref() == Some("gpt_oss");
-    if pp > 1 && is_gpt_oss {
-        eprintln!(
-            "gpt-oss is not supported by the pipeline-parallel engine (Qwen3 only); use --tp instead"
-        );
-        std::process::exit(1);
-    }
     let model_max_seq_len = model_config.max_seq_len();
     if model_max_seq_len == 0 {
         eprintln!("model config has invalid max_seq_len=0");
@@ -98,30 +76,22 @@ async fn main() {
         );
     }
 
-    let model_name = model_dir
-        .file_name()
+    let model_name = model_dir.file_name()
         .map(|n| n.to_string_lossy().to_string())
         .unwrap_or_else(|| "unknown".to_string());
 
     let tokenizer = xserv_tokenizer::Tokenizer::from_file(&model_dir.join("tokenizer.json"));
 
-    // Bounded channel to backpressure incoming requests when the engine falls
-    // behind, instead of letting them pile up in RAM. try_send in the API
-    // handler surfaces this as 503 to the client.
-    let (tx, rx) = mpsc::sync_channel::<GenerateRequest>(256);
+    // Unbounded channel: allows multiple requests to queue up
+    let (tx, rx) = mpsc::channel::<GenerateRequest>();
 
     let model_dir_clone = model_dir.clone();
     std::thread::spawn(move || {
         if pp > 1 {
             // Pipeline-parallel path: stage-0 coordinator + worker stage threads.
             pp_engine::run_pp(&model_dir_clone, pp, max_seq_len, rx);
-        } else if tp <= 1 && !is_gpt_oss {
-            let mut engine = engine::Engine::load_with_swap(
-                &model_dir_clone,
-                max_batch,
-                max_seq_len,
-                swap_space_gb,
-            );
+        } else if tp <= 1 {
+            let mut engine = engine::Engine::load_with_swap(&model_dir_clone, max_batch, max_seq_len, swap_space_gb);
             engine.run(rx);
         } else {
             // Tensor-parallel path: rank-0 coordinator + worker rank threads.
@@ -129,11 +99,8 @@ async fn main() {
         }
     });
 
-    let model_type = model_config.model_type.clone().unwrap_or_default();
-    let chat_template = api::ChatTemplate::load(&model_dir, &model_type);
     let state = Arc::new(AppState {
         model_name,
-        chat_template,
         engine_sender: Mutex::new(tx),
         engine_tokenizer: Mutex::new(tokenizer),
         max_seq_len,
@@ -143,7 +110,6 @@ async fn main() {
         .route("/health", get(api::health))
         .route("/v1/models", get(api::list_models))
         .route("/v1/chat/completions", post(api::chat_completions))
-        .layer(DefaultBodyLimit::max(4 * 1024 * 1024))
         .layer(Extension(state));
 
     let addr = format!("0.0.0.0:{port}");

crates/xserv-server/src/pp_engine.rs

@@ -15,15 +15,15 @@
 
 use std::ffi::c_void;
 use std::path::{Path, PathBuf};
-use std::sync::Arc;
 use std::sync::mpsc;
+use std::sync::Arc;
 use std::thread;
 
 use half::bf16;
 use xserv_distributed::{PpContext, UniqueId};
 use xserv_model::loader;
 use xserv_model::sampling::SamplingParams;
-use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, sample};
+use xserv_model::{sample, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
 use xserv_tensor::{DType, Device, Tensor};
 use xserv_tokenizer::Tokenizer;
 
@@ -38,16 +38,9 @@ enum PpCommand {
     Free(usize),
     /// Receive `[n_tokens, hidden]` from the previous stage, run this stage's
     /// layers; if last stage, sample with `sampling` and return the token.
-    Prefill {
-        n_tokens: usize,
-        slot: usize,
-        sampling: SamplingParams,
-    },
+    Prefill { n_tokens: usize, slot: usize, sampling: SamplingParams },
     /// Receive `[1, hidden]`, run this stage's layers; last stage samples.
-    Decode {
-        slot: usize,
-        sampling: SamplingParams,
-    },
+    Decode { slot: usize, sampling: SamplingParams },
     Shutdown,
 }
 
@@ -83,21 +76,9 @@ fn build_stage(
     let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
     let total_blocks = max_blocks_per_seq + 8; // v1 serial: one active sequence
     let cache = PagedKVCache::new(
-        &stage_config,
-        total_blocks,
-        0,
-        4,
-        max_blocks_per_seq,
-        DType::BF16,
-        device,
+        &stage_config, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, device,
     );
-    StageCtx {
-        model,
-        cache,
-        pp,
-        hidden: config.hidden(),
-        device,
-    }
+    StageCtx { model, cache, pp, hidden: config.hidden(), device }
 }
 
 /// Allocate a zeroed `[n, hidden]` device tensor and receive into it from `peer`.
@@ -129,15 +110,7 @@ fn worker_loop(
     ack_tx: mpsc::Sender<()>,
     token_tx: mpsc::Sender<u32>,
 ) {
-    let mut sc = build_stage(
-        &model_dir,
-        &config,
-        stage,
-        world,
-        stage as u32,
-        max_seq_len,
-        id,
-    );
+    let mut sc = build_stage(&model_dir, &config, stage, world, stage as u32, max_seq_len, id);
     let is_last = stage == world - 1;
     let prev = stage - 1;
     let next = stage + 1;
@@ -152,11 +125,7 @@ fn worker_loop(
                 sc.cache.free_sequence(slot);
                 let _ = ack_tx.send(());
             }
-            PpCommand::Prefill {
-                n_tokens,
-                slot,
-                sampling,
-            } => {
+            PpCommand::Prefill { n_tokens, slot, sampling } => {
                 let x = recv_hidden(&sc, n_tokens, prev);
                 let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
                 if is_last {
@@ -186,12 +155,7 @@ fn worker_loop(
 
 /// Run the PP coordinator (stage 0) on the calling thread. Spawns worker stages
 /// 1..world and consumes generation requests from `rx`.
-pub fn run_pp(
-    model_dir: &Path,
-    world: usize,
-    max_seq_len: usize,
-    rx: mpsc::Receiver<GenerateRequest>,
-) {
+pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Receiver<GenerateRequest>) {
     assert!(world >= 2, "run_pp requires world >= 2");
     let config = ModelConfig::from_file(&model_dir.join("config.json"));
     assert!(
@@ -215,17 +179,7 @@ pub fn run_pp(
         let model_dir = model_dir.to_path_buf();
         let config = config.clone();
         thread::spawn(move || {
-            worker_loop(
-                stage,
-                world,
-                id,
-                model_dir,
-                config,
-                max_seq_len,
-                ctx_rx,
-                ack_tx,
-                token_tx,
-            );
+            worker_loop(stage, world, id, model_dir, config, max_seq_len, ctx_rx, ack_tx, token_tx);
         });
     }
 
@@ -233,6 +187,7 @@ pub fn run_pp(
     let mut sc = build_stage(model_dir, &config, 0, world, 0, max_seq_len, id);
     eprintln!("[pp-engine] ready (pp={world}, max_seq_len={max_seq_len})");
 
+    let eos = tokenizer.eos_token_id();
     let n_workers = world - 1;
     let next_peer = 1usize;
     let broadcast = |txs: &[mpsc::Sender<PpCommand>], cmd: PpCommand| {
@@ -253,14 +208,11 @@ pub fn run_pp(
         wait_acks(&ack_rx);
 
         // Prefill: embed prompt, run stage-0 layers, push hidden into the pipe.
-        broadcast(
-            &cmd_txs,
-            PpCommand::Prefill {
-                n_tokens: req.prompt_tokens.len(),
-                slot,
-                sampling: req.sampling.clone(),
-            },
-        );
+        broadcast(&cmd_txs, PpCommand::Prefill {
+            n_tokens: req.prompt_tokens.len(),
+            slot,
+            sampling: req.sampling.clone(),
+        });
         let x = sc.model.embed(&req.prompt_tokens);
         let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
         send_hidden(&sc, &x, next_peer);
@@ -268,43 +220,29 @@ pub fn run_pp(
 
         let mut decode_buf: Vec<u8> = Vec::new();
         let mut generated = 1usize;
-        let mut stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
+        emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
 
         let finish = loop {
-            if stalled {
-                break "error";
-            }
-            if tokenizer.is_eos(next) {
+            if eos == Some(next) {
                 break "stop";
             }
             if generated >= req.max_tokens {
                 break "length";
             }
-            broadcast(
-                &cmd_txs,
-                PpCommand::Decode {
-                    slot,
-                    sampling: req.sampling.clone(),
-                },
-            );
+            broadcast(&cmd_txs, PpCommand::Decode { slot, sampling: req.sampling.clone() });
             let x = sc.model.embed(&[next]);
             let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache);
             send_hidden(&sc, &x, next_peer);
             next = token_rx.recv().expect("decode token");
             generated += 1;
-            stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
+            emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
         };
 
         let tail = tokenizer.flush_decode_stream(&mut decode_buf);
         if !tail.is_empty() {
-            let _ = req.sender.try_send(GenerateEvent::Token {
-                id: next,
-                text: tail,
-            });
+            let _ = req.sender.blocking_send(GenerateEvent::Token { id: next, text: tail });
         }
-        let _ = req.sender.try_send(GenerateEvent::Done {
-            finish_reason: finish.to_string(),
-        });
+        let _ = req.sender.blocking_send(GenerateEvent::Done { finish_reason: finish.to_string() });
 
         broadcast(&cmd_txs, PpCommand::Free(slot));
         sc.cache.free_sequence(slot);
@@ -315,24 +253,12 @@ pub fn run_pp(
 }
 
 /// Stream a token's decoded text to the client (EOS contributes no text).
-/// Returns false if the send would block (client too slow) or the client is
-/// gone — the caller stops generating so the coordinator thread is free to
-/// admit the next request instead of blocking on one slow consumer.
-fn emit_text(
-    tokenizer: &Tokenizer,
-    req: &GenerateRequest,
-    token_id: u32,
-    buf: &mut Vec<u8>,
-) -> bool {
-    if tokenizer.is_eos(token_id) {
-        return true;
+fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, eos: Option<u32>, buf: &mut Vec<u8>) {
+    if eos == Some(token_id) {
+        return;
     }
     let text = tokenizer.decode_token_stream(token_id, buf);
     if !text.is_empty() {
-        return req
-            .sender
-            .try_send(GenerateEvent::Token { id: token_id, text })
-            .is_ok();
+        let _ = req.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
     }
-    true
 }

crates/xserv-server/src/tp_engine.rs

@@ -13,17 +13,14 @@
 //! work; the single-GPU `Engine` still handles TP=1.
 
 use std::path::{Path, PathBuf};
-use std::sync::Arc;
 use std::sync::mpsc;
+use std::sync::Arc;
 use std::thread;
 
 use xserv_distributed::{TpContext, UniqueId};
 use xserv_model::loader;
-use xserv_model::{
-    BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, sample,
-    sample_greedy_penalized,
-};
-use xserv_tensor::{DType, Device, Tensor};
+use xserv_model::{sample, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
+use xserv_tensor::{DType, Device};
 use xserv_tokenizer::Tokenizer;
 
 use crate::engine::{GenerateEvent, GenerateRequest};
@@ -32,67 +29,14 @@ use crate::engine::{GenerateEvent, GenerateRequest};
 enum TpCommand {
     Register(usize),
     Free(usize),
-    Prefill {
-        tokens: Vec<u32>,
-        slot: usize,
-    },
-    Decode {
-        tokens: Vec<u32>,
-        positions: Vec<usize>,
-        slots: Vec<usize>,
-    },
+    Prefill { tokens: Vec<u32>, slot: usize },
+    Decode { tokens: Vec<u32>, positions: Vec<usize>, slots: Vec<usize> },
     Shutdown,
 }
 
-enum TpModel {
-    Qwen3(Qwen3),
-    GptOss(GptOss),
-}
-
-impl TpModel {
-    fn forward_prefill_paged(
-        &self,
-        tokens: &[u32],
-        slot: usize,
-        cache: &mut PagedKVCache,
-    ) -> Tensor {
-        match self {
-            TpModel::Qwen3(m) => m.forward_prefill_paged(tokens, slot, cache),
-            TpModel::GptOss(m) => m.forward_prefill_paged(tokens, slot, cache),
-        }
-    }
-
-    fn forward_decode_paged(
-        &self,
-        tokens: &[u32],
-        positions: &[usize],
-        slots: &[usize],
-        cache: &mut PagedKVCache,
-    ) -> Tensor {
-        match self {
-            TpModel::Qwen3(m) => m.forward_decode_paged(tokens, positions, slots, cache),
-            TpModel::GptOss(m) => m.forward_decode_paged(tokens, positions, slots, cache),
-        }
-    }
-}
-
 struct RankCtx {
-    model: TpModel,
+    model: Qwen3,
     cache: PagedKVCache,
-    decoder: GraphedGptOssDecoder,
-}
-
-/// Decode one step: gpt-oss batch=1 goes through the CUDA-graph decoder
-/// (lazy capture, replay thereafter); everything else runs eager.
-fn rank_decode(rc: &mut RankCtx, tokens: &[u32], positions: &[usize], slots: &[usize]) -> Tensor {
-    match &rc.model {
-        TpModel::GptOss(m) => rc
-            .decoder
-            .decode(m, tokens, positions, slots, &mut rc.cache),
-        TpModel::Qwen3(_) => rc
-            .model
-            .forward_decode_paged(tokens, positions, slots, &mut rc.cache),
-    }
 }
 
 fn build_rank(
@@ -105,43 +49,14 @@ fn build_rank(
     tp: Option<Arc<TpContext>>,
 ) -> RankCtx {
     let weights = loader::load_model_dir(model_dir, Device::Cpu);
-    let model = if config.is_moe() {
-        TpModel::GptOss(GptOss::from_weights_tp(
-            config.clone(),
-            weights,
-            rank,
-            world,
-            device,
-            tp,
-        ))
-    } else {
-        TpModel::Qwen3(Qwen3::from_weights_tp(
-            config.clone(),
-            weights,
-            rank,
-            world,
-            device,
-            tp,
-        ))
-    };
+    let model = Qwen3::from_weights_tp(config.clone(), weights, rank, world, device, tp);
     let local_kv = config.num_kv_heads() / world;
-    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
+    let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
     let total_blocks = max_blocks_per_seq + 8;
     let cache = PagedKVCache::new_tp(
-        config,
-        local_kv,
-        total_blocks,
-        0,
-        4,
-        max_blocks_per_seq,
-        DType::BF16,
-        device,
+        config, local_kv, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, device,
     );
-    RankCtx {
-        model,
-        cache,
-        decoder: GraphedGptOssDecoder::new(),
-    }
+    RankCtx { model, cache }
 }
 
 fn worker_loop(
@@ -155,15 +70,7 @@ fn worker_loop(
     ack_tx: mpsc::Sender<()>,
 ) {
     let tp = Arc::new(TpContext::init(rank, world, id, rank as u32));
-    let mut rc = build_rank(
-        &model_dir,
-        &config,
-        rank,
-        world,
-        rank as u32,
-        max_seq_len,
-        Some(tp),
-    );
+    let mut rc = build_rank(&model_dir, &config, rank, world, rank as u32, max_seq_len, Some(tp));
     while let Ok(cmd) = cmd_rx.recv() {
         match cmd {
             TpCommand::Register(slot) => {
@@ -173,12 +80,8 @@ fn worker_loop(
             TpCommand::Prefill { tokens, slot } => {
                 let _ = rc.model.forward_prefill_paged(&tokens, slot, &mut rc.cache);
             }
-            TpCommand::Decode {
-                tokens,
-                positions,
-                slots,
-            } => {
-                let _ = rank_decode(&mut rc, &tokens, &positions, &slots);
+            TpCommand::Decode { tokens, positions, slots } => {
+                let _ = rc.model.forward_decode_paged(&tokens, &positions, &slots, &mut rc.cache);
             }
             TpCommand::Shutdown => {
                 let _ = ack_tx.send(());
@@ -191,15 +94,8 @@ fn worker_loop(
 
 /// Run the TP coordinator (rank 0) on the calling thread. Spawns worker ranks
 /// internally and consumes generation requests from `rx`.
-pub fn run_tp(
-    model_dir: &Path,
-    world: usize,
-    max_seq_len: usize,
-    rx: mpsc::Receiver<GenerateRequest>,
-) {
-    // world=1 is a valid single-rank configuration (gpt-oss has no
-    // single-GPU engine path; NCCL init and all_reduce no-op at world=1).
-    assert!(world >= 1, "run_tp requires world >= 1");
+pub fn run_tp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Receiver<GenerateRequest>) {
+    assert!(world >= 2, "run_tp requires world >= 2");
     let config = ModelConfig::from_file(&model_dir.join("config.json"));
     assert!(
         config.num_kv_heads() % world == 0,
@@ -219,16 +115,7 @@ pub fn run_tp(
         let model_dir = model_dir.to_path_buf();
         let config = config.clone();
         thread::spawn(move || {
-            worker_loop(
-                rank,
-                world,
-                id,
-                model_dir,
-                config,
-                max_seq_len,
-                ctx_rx,
-                ack_tx,
-            );
+            worker_loop(rank, world, id, model_dir, config, max_seq_len, ctx_rx, ack_tx);
         });
     }
 
@@ -237,27 +124,7 @@ pub fn run_tp(
     let mut rc = build_rank(model_dir, &config, 0, world, 0, max_seq_len, Some(tp));
     eprintln!("[tp-engine] ready (tp={world}, max_seq_len={max_seq_len})");
 
-    // Optional repetition penalty to break greedy repetition loops (reasoning
-    // models loop under pure greedy when numerics diverge from the reference).
-    // Off by default; XSERV_REP_PENALTY>1 enables it over the last
-    // XSERV_REP_WINDOW generated tokens. Applied only on the greedy path.
-    let rep_penalty: f32 = std::env::var("XSERV_REP_PENALTY")
-        .ok()
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(1.0);
-    let rep_window: usize = std::env::var("XSERV_REP_WINDOW")
-        .ok()
-        .and_then(|s| s.parse().ok())
-        .unwrap_or(128);
-    let pick = |logits: &Tensor, sp: &xserv_model::SamplingParams, history: &[u32]| -> u32 {
-        if rep_penalty > 1.0 && sp.temperature == 0.0 {
-            let start = history.len().saturating_sub(rep_window);
-            sample_greedy_penalized(logits, &history[start..], rep_penalty)
-        } else {
-            sample(logits, sp)
-        }
-    };
-
+    let eos = tokenizer.eos_token_id();
     let n_workers = world - 1;
     let broadcast = |txs: &[mpsc::Sender<TpCommand>], cmd: TpCommand| {
         for t in txs {
@@ -277,62 +144,36 @@ pub fn run_tp(
         wait_acks(&ack_rx);
 
         // Prefill.
-        broadcast(
-            &cmd_txs,
-            TpCommand::Prefill {
-                tokens: req.prompt_tokens.clone(),
-                slot,
-            },
-        );
-        let logits = rc
-            .model
-            .forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache);
+        broadcast(&cmd_txs, TpCommand::Prefill { tokens: req.prompt_tokens.clone(), slot });
+        let logits = rc.model.forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache);
         wait_acks(&ack_rx);
-        let mut gen_ids: Vec<u32> = Vec::new();
-        let mut next = pick(&logits, &req.sampling, &gen_ids);
-        gen_ids.push(next);
+        let mut next = sample(&logits, &req.sampling);
 
         let mut decode_buf: Vec<u8> = Vec::new();
         let mut generated = 1usize;
-        let mut stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
+        emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
 
         let finish = loop {
-            if stalled {
-                break "error";
-            }
-            if tokenizer.is_eos(next) {
+            if eos == Some(next) {
                 break "stop";
             }
             if generated >= req.max_tokens {
                 break "length";
             }
             let pos = rc.cache.seq_len(slot);
-            broadcast(
-                &cmd_txs,
-                TpCommand::Decode {
-                    tokens: vec![next],
-                    positions: vec![pos],
-                    slots: vec![slot],
-                },
-            );
-            let logits = rank_decode(&mut rc, &[next], &[pos], &[slot]);
+            broadcast(&cmd_txs, TpCommand::Decode { tokens: vec![next], positions: vec![pos], slots: vec![slot] });
+            let logits = rc.model.forward_decode_paged(&[next], &[pos], &[slot], &mut rc.cache);
             wait_acks(&ack_rx);
-            next = pick(&logits, &req.sampling, &gen_ids);
-            gen_ids.push(next);
+            next = sample(&logits, &req.sampling);
             generated += 1;
-            stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
+            emit_text(&tokenizer, &req, next, eos, &mut decode_buf);
         };
 
         let tail = tokenizer.flush_decode_stream(&mut decode_buf);
         if !tail.is_empty() {
-            let _ = req.sender.try_send(GenerateEvent::Token {
-                id: next,
-                text: tail,
-            });
+            let _ = req.sender.blocking_send(GenerateEvent::Token { id: next, text: tail });
         }
-        let _ = req.sender.try_send(GenerateEvent::Done {
-            finish_reason: finish.to_string(),
-        });
+        let _ = req.sender.blocking_send(GenerateEvent::Done { finish_reason: finish.to_string() });
 
         broadcast(&cmd_txs, TpCommand::Free(slot));
         rc.cache.free_sequence(slot);
@@ -343,24 +184,12 @@ pub fn run_tp(
 }
 
 /// Stream a token's decoded text to the client (EOS contributes no text).
-/// Returns false if the send would block (client too slow) or the client is
-/// gone — the caller stops generating so the serial coordinator thread is free
-/// to admit the next request instead of blocking on one slow consumer.
-fn emit_text(
-    tokenizer: &Tokenizer,
-    req: &GenerateRequest,
-    token_id: u32,
-    buf: &mut Vec<u8>,
-) -> bool {
-    if tokenizer.is_eos(token_id) {
-        return true;
+fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, eos: Option<u32>, buf: &mut Vec<u8>) {
+    if eos == Some(token_id) {
+        return;
     }
     let text = tokenizer.decode_token_stream(token_id, buf);
     if !text.is_empty() {
-        return req
-            .sender
-            .try_send(GenerateEvent::Token { id: token_id, text })
-            .is_ok();
+        let _ = req.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
     }
-    true
 }

crates/xserv-tensor/src/dtype.rs

@@ -5,7 +5,6 @@ pub enum DType {
     F32,
     F16,
     BF16,
-    FP8E4M3,
 }
 
 impl DType {
@@ -14,7 +13,6 @@ impl DType {
             DType::F32 => 4,
             DType::F16 => 2,
             DType::BF16 => 2,
-            DType::FP8E4M3 => 1,
         }
     }
 
@@ -23,7 +21,6 @@ impl DType {
             DType::F32 => "f32",
             DType::F16 => "f16",
             DType::BF16 => "bf16",
-            DType::FP8E4M3 => "fp8e4m3",
         }
     }
 }
@@ -43,30 +40,18 @@ pub trait TensorDType: Copy + Send + Sync + 'static {
 
 impl TensorDType for f32 {
     const DTYPE: DType = DType::F32;
-    fn to_f64(self) -> f64 {
-        self as f64
-    }
-    fn from_f64(v: f64) -> Self {
-        v as f32
-    }
+    fn to_f64(self) -> f64 { self as f64 }
+    fn from_f64(v: f64) -> Self { v as f32 }
 }
 
 impl TensorDType for f16 {
     const DTYPE: DType = DType::F16;
-    fn to_f64(self) -> f64 {
-        self.to_f32() as f64
-    }
-    fn from_f64(v: f64) -> Self {
-        f16::from_f32(v as f32)
-    }
+    fn to_f64(self) -> f64 { self.to_f32() as f64 }
+    fn from_f64(v: f64) -> Self { f16::from_f32(v as f32) }
 }
 
 impl TensorDType for bf16 {
     const DTYPE: DType = DType::BF16;
-    fn to_f64(self) -> f64 {
-        self.to_f32() as f64
-    }
-    fn from_f64(v: f64) -> Self {
-        bf16::from_f32(v as f32)
-    }
+    fn to_f64(self) -> f64 { self.to_f32() as f64 }
+    fn from_f64(v: f64) -> Self { bf16::from_f32(v as f32) }
 }

crates/xserv-tensor/src/lib.rs

@@ -6,4 +6,4 @@ pub mod tensor;
 pub use dtype::{DType, TensorDType};
 pub use shape::Dims;
 pub use storage::{Device, Storage};
-pub use tensor::{Tensor, register_gpu_contiguous};
+pub use tensor::{register_gpu_contiguous, Tensor};

crates/xserv-tensor/src/shape.rs

@@ -18,21 +18,12 @@ pub fn contiguous_strides(shape: &[usize]) -> Dims {
 }
 
 /// Check if the given strides represent contiguous (row-major) layout for the shape.
-/// A stride mismatch on a dimension of size 1 is allowed because that
-/// dimension is never stepped.
 pub fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool {
     if shape.is_empty() {
         return true;
     }
-    let ndim = shape.len();
-    let mut expected_stride = 1usize;
-    for d in (0..ndim).rev() {
-        if shape[d] != 1 && strides[d] != expected_stride {
-            return false;
-        }
-        expected_stride *= shape[d];
-    }
-    true
+    let expected = contiguous_strides(shape);
+    strides == expected.as_slice()
 }
 
 /// Total number of elements given a shape.
@@ -46,16 +37,8 @@ pub fn broadcast_shape(a: &[usize], b: &[usize]) -> Option<Dims> {
     let ndim = a.len().max(b.len());
     let mut result = SmallVec::with_capacity(ndim);
     for i in 0..ndim {
-        let da = if i < ndim - a.len() {
-            1
-        } else {
-            a[i - (ndim - a.len())]
-        };
-        let db = if i < ndim - b.len() {
-            1
-        } else {
-            b[i - (ndim - b.len())]
-        };
+        let da = if i < ndim - a.len() { 1 } else { a[i - (ndim - a.len())] };
+        let db = if i < ndim - b.len() { 1 } else { b[i - (ndim - b.len())] };
         if da == db {
             result.push(da);
         } else if da == 1 {
@@ -108,14 +91,8 @@ mod tests {
 
     #[test]
     fn test_broadcast_shape() {
-        assert_eq!(
-            broadcast_shape(&[3, 1], &[1, 4]).unwrap().as_slice(),
-            &[3, 4]
-        );
-        assert_eq!(
-            broadcast_shape(&[2, 3, 4], &[4]).unwrap().as_slice(),
-            &[2, 3, 4]
-        );
+        assert_eq!(broadcast_shape(&[3, 1], &[1, 4]).unwrap().as_slice(), &[3, 4]);
+        assert_eq!(broadcast_shape(&[2, 3, 4], &[4]).unwrap().as_slice(), &[2, 3, 4]);
         assert_eq!(broadcast_shape(&[1], &[5, 3]).unwrap().as_slice(), &[5, 3]);
         assert!(broadcast_shape(&[3], &[4]).is_none());
     }
@@ -123,9 +100,6 @@ mod tests {
     #[test]
     fn test_broadcast_strides() {
         // [3,1] with strides [1,1] broadcast to [3,4]
-        assert_eq!(
-            broadcast_strides(&[3, 1], &[1, 1], &[3, 4]).as_slice(),
-            &[1, 0]
-        );
+        assert_eq!(broadcast_strides(&[3, 1], &[1, 1], &[3, 4]).as_slice(), &[1, 0]);
     }
 }

crates/xserv-tensor/src/tensor.rs

@@ -33,20 +33,8 @@ impl Tensor {
     // --- Creation ---
 
     /// Create a tensor from raw components (for advanced use like GPU KV cache).
-    pub fn from_storage(
-        storage: Storage,
-        shape: Dims,
-        strides: Dims,
-        offset: usize,
-        dtype: DType,
-    ) -> Self {
-        Self {
-            storage,
-            shape,
-            strides,
-            offset,
-            dtype,
-        }
+    pub fn from_storage(storage: Storage, shape: Dims, strides: Dims, offset: usize, dtype: DType) -> Self {
+        Self { storage, shape, strides, offset, dtype }
     }
 
     pub fn from_slice<T: TensorDType>(data: &[T], shape: &[usize]) -> Self {
@@ -64,28 +52,6 @@ impl Tensor {
         }
     }
 
-    /// Create a tensor from raw bytes. Used for dtypes without a Rust type
-    /// (e.g. FP8 E4M3) where we store the bit pattern as-is.
-    pub fn from_raw_bytes(data: &[u8], shape: &[usize], dtype: DType) -> Self {
-        let numel: usize = shape.iter().product();
-        assert_eq!(
-            data.len(),
-            numel * dtype.size_bytes(),
-            "raw bytes length {} != expected {} (numel={} * elem_size={})",
-            data.len(),
-            numel * dtype.size_bytes(),
-            numel,
-            dtype.size_bytes()
-        );
-        Self {
-            storage: Storage::cpu(data.to_vec()),
-            shape: Dims::from_slice(shape),
-            strides: shape::contiguous_strides(shape),
-            offset: 0,
-            dtype,
-        }
-    }
-
     pub fn zeros(shape: &[usize], dtype: DType, device: Device) -> Self {
         let numel = shape::num_elements(shape);
         let len_bytes = numel * dtype.size_bytes();
@@ -121,34 +87,19 @@ impl Tensor {
             DType::F32 => Self::from_slice(&vec![1.0f32; numel], shape),
             DType::F16 => Self::from_slice(&vec![half::f16::from_f32(1.0); numel], shape),
             DType::BF16 => Self::from_slice(&vec![half::bf16::from_f32(1.0); numel], shape),
-            DType::FP8E4M3 => panic!("ones() not supported for FP8E4M3"),
         }
     }
 
     // --- Properties ---
 
-    pub fn shape(&self) -> &[usize] {
-        &self.shape
-    }
-    pub fn strides(&self) -> &[usize] {
-        &self.strides
-    }
-    pub fn dtype(&self) -> DType {
-        self.dtype
-    }
-    pub fn ndim(&self) -> usize {
-        self.shape.len()
-    }
-    pub fn numel(&self) -> usize {
-        shape::num_elements(&self.shape)
-    }
-    pub fn offset(&self) -> usize {
-        self.offset
-    }
+    pub fn shape(&self) -> &[usize] { &self.shape }
+    pub fn strides(&self) -> &[usize] { &self.strides }
+    pub fn dtype(&self) -> DType { self.dtype }
+    pub fn ndim(&self) -> usize { self.shape.len() }
+    pub fn numel(&self) -> usize { shape::num_elements(&self.shape) }
+    pub fn offset(&self) -> usize { self.offset }
 
-    pub fn device(&self) -> Device {
-        self.storage.device()
-    }
+    pub fn device(&self) -> Device { self.storage.device() }
 
     pub fn is_contiguous(&self) -> bool {
         shape::is_contiguous(&self.shape, &self.strides)
@@ -169,21 +120,6 @@ impl Tensor {
         }
     }
 
-    /// Zero-copy slice along `dim`: keeps elements `[start, start+len)`.
-    pub fn narrow(&self, dim: usize, start: usize, len: usize) -> Self {
-        assert!(dim < self.ndim());
-        assert!(start + len <= self.shape[dim], "narrow out of bounds");
-        let mut new_shape = self.shape.clone();
-        new_shape[dim] = len;
-        Self {
-            storage: self.storage.clone(),
-            shape: new_shape,
-            strides: self.strides.clone(),
-            offset: self.offset + start * self.strides[dim],
-            dtype: self.dtype,
-        }
-    }
-
     pub fn transpose(&self, dim0: usize, dim1: usize) -> Self {
         assert!(dim0 < self.ndim() && dim1 < self.ndim());
         let mut new_shape = self.shape.clone();
@@ -222,11 +158,7 @@ impl Tensor {
             shape::contiguous_strides(&new_shape)
         } else {
             let mut s = self.strides.clone();
-            let stride_val = if dim < self.strides.len() {
-                self.strides[dim]
-            } else {
-                1
-            };
+            let stride_val = if dim < self.strides.len() { self.strides[dim] } else { 1 };
             s.insert(dim, stride_val);
             s
         };
@@ -263,12 +195,7 @@ impl Tensor {
         let ndim = self.ndim();
         let mut idx = vec![0usize; ndim];
         for flat in 0..numel {
-            let src_offset = self.offset
-                + idx
-                    .iter()
-                    .zip(self.strides.iter())
-                    .map(|(i, s)| i * s)
-                    .sum::<usize>();
+            let src_offset = self.offset + idx.iter().zip(self.strides.iter()).map(|(i, s)| i * s).sum::<usize>();
             let src_byte_offset = src_offset * elem_size;
             let dst_byte_offset = flat * elem_size;
             dst[dst_byte_offset..dst_byte_offset + elem_size]
@@ -299,10 +226,7 @@ impl Tensor {
         }
         // Transfer the raw storage (preserving strides/offset).
         // Non-contiguous layout is preserved — the user can call contiguous() after.
-        let new_storage = self
-            .storage
-            .to_device(device)
-            .expect("device transfer failed");
+        let new_storage = self.storage.to_device(device).expect("device transfer failed");
         Self {
             storage: new_storage,
             shape: self.shape.clone(),
@@ -326,17 +250,6 @@ impl Tensor {
         unsafe { std::slice::from_raw_parts(bytes[start..].as_ptr() as *const T, len) }
     }
 
-    /// Raw byte access for dtypes without a Rust type (e.g. FP8).
-    pub fn as_raw_bytes(&self) -> &[u8] {
-        assert!(self.is_contiguous(), "as_raw_bytes requires contiguous");
-        assert_eq!(self.device(), Device::Cpu, "as_raw_bytes requires CPU");
-        let bytes = self.storage.as_cpu_bytes();
-        let elem_size = self.dtype.size_bytes();
-        let start = self.offset * elem_size;
-        let len = self.numel() * elem_size;
-        &bytes[start..start + len]
-    }
-
     /// Raw pointer to storage start (for GPU kernel launch).
     pub fn data_ptr(&self) -> *const u8 {
         match self.device() {
@@ -351,20 +264,14 @@ impl Tensor {
         }
     }
 
-    pub fn storage(&self) -> &Storage {
-        &self.storage
-    }
+    pub fn storage(&self) -> &Storage { &self.storage }
 }
 
 impl std::fmt::Debug for Tensor {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(
-            f,
-            "Tensor(shape={:?}, dtype={}, device={}, contiguous={})",
-            self.shape.as_slice(),
-            self.dtype,
-            self.device(),
-            self.is_contiguous()
+            f, "Tensor(shape={:?}, dtype={}, device={}, contiguous={})",
+            self.shape.as_slice(), self.dtype, self.device(), self.is_contiguous()
         )
     }
 }

crates/xserv-tensor/tests/integration.rs

@@ -32,11 +32,7 @@ fn test_zeros_and_ones() {
 
 #[test]
 fn test_bf16_tensor() {
-    let data: Vec<bf16> = vec![
-        bf16::from_f32(1.0),
-        bf16::from_f32(2.5),
-        bf16::from_f32(-3.0),
-    ];
+    let data: Vec<bf16> = vec![bf16::from_f32(1.0), bf16::from_f32(2.5), bf16::from_f32(-3.0)];
     let t = Tensor::from_slice(&data, &[3]);
     assert_eq!(t.dtype(), DType::BF16);
     let out = t.as_slice::<bf16>();

crates/xserv-tokenizer/src/bpe.rs

@@ -12,7 +12,6 @@ pub struct Tokenizer {
     special_token_ids: HashMap<u32, String>,
     pre_tokenize_re: Regex,
     eos_token_id: Option<u32>,
-    eos_token_ids: Vec<u32>,
     byte_fallback: bool,
 }
 
@@ -21,24 +20,6 @@ struct TokenizerJson {
     model: ModelSection,
     #[serde(default)]
     added_tokens: Vec<AddedToken>,
-    #[serde(default)]
-    pre_tokenizer: Option<PreTokenizerSection>,
-}
-
-#[derive(Deserialize)]
-struct PreTokenizerSection {
-    #[serde(default, rename = "type")]
-    kind: Option<String>,
-    #[serde(default)]
-    pattern: Option<PatternSpec>,
-    #[serde(default)]
-    pretokenizers: Option<Vec<PreTokenizerSection>>,
-}
-
-#[derive(Deserialize)]
-struct PatternSpec {
-    #[serde(rename = "Regex")]
-    regex: Option<String>,
 }
 
 #[derive(Deserialize)]
@@ -95,15 +76,11 @@ impl Tokenizer {
             let (a_str, b_str) = match entry {
                 MergeEntry::Str(s) => {
                     let parts: Vec<&str> = s.splitn(2, ' ').collect();
-                    if parts.len() != 2 {
-                        continue;
-                    }
+                    if parts.len() != 2 { continue; }
                     (parts[0].to_string(), parts[1].to_string())
                 }
                 MergeEntry::Pair(v) => {
-                    if v.len() != 2 {
-                        continue;
-                    }
+                    if v.len() != 2 { continue; }
                     (v[0].clone(), v[1].clone())
                 }
             };
@@ -125,69 +102,18 @@ impl Tokenizer {
             decoder.resize(decoder.len().max(at.id as usize + 1), vec![]);
             decoder[at.id as usize] = at.content.as_bytes().to_vec();
         }
-        // End-of-generation tokens, in priority order. Families differ:
-        // Qwen uses <|im_end|>, Llama <|end_of_text|>, GPT-2 <|endoftext|>.
-        // gpt-oss (harmony) ends the assistant turn with <|return|> and also
-        // treats <|call|> (tool call) and <|endoftext|> as terminators
-        // (see generation_config.json eos_token_id = [200002, 199999, 200012]).
-        let eos_names = [
-            "<|im_end|>",
-            "<|end_of_text|>",
-            "<|return|>",
-            "<|call|>",
-            "<|endoftext|>",
-        ];
-        let mut eos_token_ids: Vec<u32> = Vec::new();
-        for name in eos_names {
-            if let Some(&id) = special_tokens.get(name) {
-                if !eos_token_ids.contains(&id) {
-                    eos_token_ids.push(id);
-                }
-            }
-        }
-        let eos_token_id = eos_token_ids.first().copied();
+        let eos_token_id = special_tokens
+            .get("<|im_end|>")
+            .or_else(|| special_tokens.get("<|end_of_text|>"))
+            .or_else(|| special_tokens.get("<|endoftext|>"))
+            .copied();
 
-        // Pre-tokenization regex: prefer the model's own regex from tokenizer.json,
-        // fall back to GPT-2/Qwen heuristic if not present or unsupported.
-        let model_regex = tj.pre_tokenizer.as_ref().and_then(|pt| {
-            // Direct Split with regex
-            if pt.kind.as_deref() == Some("Split") {
-                return pt.pattern.as_ref().and_then(|p| p.regex.clone());
-            }
-            // Sequence → find the Split entry
-            if let Some(subs) = &pt.pretokenizers {
-                for sub in subs {
-                    if sub.kind.as_deref() == Some("Split") {
-                        if let Some(r) = sub.pattern.as_ref().and_then(|p| p.regex.clone()) {
-                            return Some(r);
-                        }
-                    }
-                }
-            }
-            None
-        });
-
-        let pre_tokenize_re = if let Some(ref pat) = model_regex {
-            // Strip unsupported lookahead (?!\S) — Rust regex doesn't support it.
-            // The lookahead only affects trailing-whitespace edge cases.
-            let cleaned = pat.replace(r"(?!\S)", "");
-            match Regex::new(&cleaned) {
-                Ok(re) => re,
-                Err(e) => {
-                    eprintln!("warning: model pre_tokenizer regex failed ({e}), using fallback");
-                    if byte_fallback {
-                        Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
-                    } else {
-                        Regex::new(
-                            r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+",
-                        )
-                        .unwrap()
-                    }
-                }
-            }
-        } else if byte_fallback {
+        // Pre-tokenization regex
+        let pre_tokenize_re = if byte_fallback {
+            // Qwen-style: split on whitespace boundaries, keep Unicode words/numbers
             Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
         } else {
+            // GPT-2 style
             Regex::new(r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+").unwrap()
         };
 
@@ -199,7 +125,6 @@ impl Tokenizer {
             special_token_ids,
             pre_tokenize_re,
             eos_token_id,
-            eos_token_ids,
             byte_fallback,
         }
     }
@@ -269,9 +194,7 @@ impl Tokenizer {
 
             // BPE merges
             loop {
-                if token_ids.len() < 2 {
-                    break;
-                }
+                if token_ids.len() < 2 { break; }
                 let mut best_rank = usize::MAX;
                 let mut best_idx = 0;
                 for i in 0..token_ids.len() - 1 {
@@ -282,15 +205,12 @@ impl Tokenizer {
                         }
                     }
                 }
-                if best_rank == usize::MAX {
-                    break;
-                }
+                if best_rank == usize::MAX { break; }
 
                 let merged_bytes = [
                     self.decoder[token_ids[best_idx] as usize].as_slice(),
                     self.decoder[token_ids[best_idx + 1] as usize].as_slice(),
-                ]
-                .concat();
+                ].concat();
                 let merged_id = *self.encoder.get(&merged_bytes).unwrap_or_else(|| {
                     panic!("merged token not in vocab");
                 });
@@ -329,12 +249,6 @@ impl Tokenizer {
         self.eos_token_id
     }
 
-    /// True if `id` is any end-of-generation token (a model may have several;
-    /// gpt-oss/harmony ends on <|return|>, <|call|>, or <|endoftext|>).
-    pub fn is_eos(&self, id: u32) -> bool {
-        self.eos_token_ids.contains(&id)
-    }
-
     pub fn vocab_size(&self) -> usize {
         self.decoder.len()
     }
@@ -401,13 +315,14 @@ fn unicode_to_byte(c: char) -> u8 {
         m
     });
 
-    *map.get(&(c as u32))
-        .unwrap_or_else(|| panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32))
+    *map.get(&(c as u32)).unwrap_or_else(|| {
+        panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32)
+    })
 }
 
 #[cfg(test)]
 mod tests {
-    use super::{Tokenizer, take_valid_utf8};
+    use super::{take_valid_utf8, Tokenizer};
 
     #[test]
     fn qwen_added_tokens_are_indivisible_and_im_end_is_eos() {

csrc/activation/activations.cu

@@ -58,25 +58,6 @@ __global__ void silu_mul_bf16_kernel(const __nv_bfloat16* gate, const __nv_bfloa
     }
 }
 
-// gpt-oss GLU: gate_up is [N, 2*D] with interleaved columns (gate=even, up=odd).
-// gate = gate_up[::2].clamp(max=limit)
-// up   = gate_up[1::2].clamp(-limit, limit)
-// glu  = gate * sigmoid(gate * alpha)
-// out  = (up + 1) * glu
-// Output: [N, D]
-__global__ void gpt_oss_glu_bf16_kernel(const __nv_bfloat16* gate_up, __nv_bfloat16* out,
-                                         int n_elements, float alpha, float limit) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    if (idx < n_elements) {
-        float g = __bfloat162float(gate_up[idx * 2]);
-        float u = __bfloat162float(gate_up[idx * 2 + 1]);
-        g = fminf(g, limit);
-        u = fmaxf(fminf(u, limit), -limit);
-        float glu = g / (1.0f + expf(-g * alpha));
-        out[idx] = __float2bfloat16((u + 1.0f) * glu);
-    }
-}
-
 // Element-wise add: out = a + b
 __global__ void add_f32_kernel(const float* a, const float* b, float* out, int n) {
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
@@ -87,17 +68,6 @@ __global__ void add_bf16_kernel(const __nv_bfloat16* a, const __nv_bfloat16* b,
     if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(a[idx]) + __bfloat162float(b[idx]));
 }
 
-// Row-broadcast bias add: out[r, c] = x[r, c] + bias[c]
-__global__ void bias_add_2d_bf16_kernel(
-    const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ bias,
-    __nv_bfloat16* __restrict__ out, int rows, int cols
-) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    if (idx >= rows * cols) return;
-    float v = __bfloat162float(x[idx]) + __bfloat162float(bias[idx % cols]);
-    out[idx] = __float2bfloat16(v);
-}
-
 // Element-wise mul: out = a * b
 __global__ void mul_f32_kernel(const float* a, const float* b, float* out, int n) {
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
@@ -170,14 +140,6 @@ void launch_add_bf16(const void* a, const void* b, void* out, int n, void* strea
         (const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
     CUDA_CHECK_LAST_ERROR();
 }
-void launch_bias_add_2d_bf16(const void* x, const void* bias, void* out, int rows, int cols, void* stream) {
-    int n = rows * cols;
-    int block = 256;
-    int grid = (n + block - 1) / block;
-    bias_add_2d_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)x, (const __nv_bfloat16*)bias, (__nv_bfloat16*)out, rows, cols);
-    CUDA_CHECK_LAST_ERROR();
-}
 void launch_mul_f32(const void* a, const void* b, void* out, int n, void* stream) {
     int block = 256;
     int grid = (n + block - 1) / block;
@@ -201,13 +163,4 @@ void launch_silu_mul_bf16(const void* gate, const void* up, void* out, int n, vo
     CUDA_CHECK_LAST_ERROR();
 }
 
-void launch_gpt_oss_glu_bf16(const void* gate_up, void* out, int n_elements,
-                              float alpha, float limit, void* stream) {
-    int block = 256;
-    int grid = (n_elements + block - 1) / block;
-    gpt_oss_glu_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)gate_up, (__nv_bfloat16*)out, n_elements, alpha, limit);
-    CUDA_CHECK_LAST_ERROR();
-}
-
 }

csrc/attention/causal_mask.cu

@@ -15,10 +15,7 @@ __global__ void causal_mask_f32(
     int col = blockIdx.x * blockDim.x + threadIdx.x;
 
     if (col < cols && col > row + offset) {
-        // 64-bit index: batch * rows * cols overflows int32 at moderate batch
-        // and long context (e.g. batch=128 * heads=28 * seq=32768).
-        long long idx = ((long long)batch_idx * rows + row) * cols + col;
-        scores[idx] = -INFINITY;
+        scores[batch_idx * rows * cols + row * cols + col] = -INFINITY;
     }
 }
 
@@ -31,8 +28,7 @@ __global__ void causal_mask_bf16(
     int col = blockIdx.x * blockDim.x + threadIdx.x;
 
     if (col < cols && col > row + offset) {
-        long long idx = ((long long)batch_idx * rows + row) * cols + col;
-        scores[idx] = __float2bfloat16(-INFINITY);
+        scores[batch_idx * rows * cols + row * cols + col] = __float2bfloat16(-INFINITY);
     }
 }
 

csrc/attention/flash_attention.cu

@@ -197,183 +197,6 @@ __global__ void flash_attention_bf16_kernel(
     }
 }
 
-// Flash Attention 2 forward with gpt-oss attention sinks + optional sliding window.
-// Identical to flash_attention_bf16_kernel, plus:
-//   - sinks: [num_q_heads] BF16 — a per-head extra softmax logit (no value),
-//     folded into the denominator after the K/V tiles (exactly as the decode
-//     sink kernel does).
-//   - window_size > 0: sliding-window mask. Query at global position p attends
-//     to keys k with p - window_size < k <= p (matches HF gpt-oss).
-__global__ void flash_attention_sinks_bf16_kernel(
-    const __nv_bfloat16* __restrict__ Q,
-    const __nv_bfloat16* __restrict__ K,
-    const __nv_bfloat16* __restrict__ V,
-    __nv_bfloat16* __restrict__ O,
-    const __nv_bfloat16* __restrict__ sinks, // [num_q_heads] or NULL
-    int num_q_heads, int num_kv_heads,
-    int q_len, int kv_len, int head_dim,
-    float scale, int causal, int window_size
-) {
-    int q_tile_idx = blockIdx.x;
-    int bh = blockIdx.y;
-    int batch_idx = bh / num_q_heads;
-    int q_head = bh % num_q_heads;
-
-    int heads_per_group = num_q_heads / num_kv_heads;
-    int kv_head = q_head / heads_per_group;
-
-    int q_tile_start = q_tile_idx * BR;
-    if (q_tile_start >= q_len) return;
-    int q_tile_rows = min(BR, q_len - q_tile_start);
-
-    const __nv_bfloat16* Q_head = Q + ((long long)batch_idx * num_q_heads + q_head) * q_len * head_dim;
-    const __nv_bfloat16* K_head = K + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
-    const __nv_bfloat16* V_head = V + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
-    __nv_bfloat16* O_head = O + ((long long)batch_idx * num_q_heads + q_head) * q_len * head_dim;
-
-    int tid = threadIdx.x;
-
-    extern __shared__ __nv_bfloat16 smem[];
-    __nv_bfloat16* smem_q  = smem;
-    __nv_bfloat16* smem_kv = smem + BR * head_dim;
-
-    int q_elems = q_tile_rows * head_dim;
-    for (int i = tid; i < q_elems; i += THREADS_PER_BLOCK) {
-        int row = i / head_dim;
-        int col = i % head_dim;
-        smem_q[row * head_dim + col] = Q_head[(q_tile_start + row) * head_dim + col];
-    }
-    for (int i = q_elems + tid; i < BR * head_dim; i += THREADS_PER_BLOCK) {
-        smem_q[i] = __float2bfloat16(0.0f);
-    }
-    __syncthreads();
-
-    bool owns_row = (tid < q_tile_rows);
-
-    float O_acc[128];
-    float m_val = -INFINITY;
-    float l_val = 0.0f;
-    if (owns_row) {
-        for (int d = 0; d < head_dim; d++) O_acc[d] = 0.0f;
-    }
-
-    int kv_offset = kv_len - q_len;
-    int num_kv_tiles = (kv_len + BC - 1) / BC;
-
-    for (int j = 0; j < num_kv_tiles; j++) {
-        int kv_tile_start = j * BC;
-        int kv_tile_cols = min(BC, kv_len - kv_tile_start);
-
-        if (causal) {
-            int max_allowed_kv = (q_tile_start + q_tile_rows - 1) + kv_offset;
-            if (kv_tile_start > max_allowed_kv) continue;
-        }
-
-        int kv_elems = kv_tile_cols * head_dim;
-        for (int i = tid; i < kv_elems; i += THREADS_PER_BLOCK) {
-            int row = i / head_dim;
-            int col = i % head_dim;
-            smem_kv[row * head_dim + col] = K_head[(kv_tile_start + row) * head_dim + col];
-        }
-        for (int i = kv_elems + tid; i < BC * head_dim; i += THREADS_PER_BLOCK) {
-            smem_kv[i] = __float2bfloat16(0.0f);
-        }
-        __syncthreads();
-
-        float P[BC];
-
-        if (owns_row) {
-            float row_max = -INFINITY;
-            int q_pos = q_tile_start + tid + kv_offset; // global query position
-            for (int c = 0; c < kv_tile_cols; c++) {
-                float dot = 0.0f;
-                for (int d = 0; d < head_dim; d++) {
-                    dot += __bfloat162float(smem_q[tid * head_dim + d])
-                         * __bfloat162float(smem_kv[c * head_dim + d]);
-                }
-                float s = dot * scale;
-
-                int kv_pos = kv_tile_start + c;
-                if (causal && kv_pos > q_pos) {
-                    s = -INFINITY;
-                }
-                // Sliding window: drop keys older than the window.
-                if (window_size > 0 && kv_pos <= q_pos - window_size) {
-                    s = -INFINITY;
-                }
-
-                P[c] = s;
-                row_max = fmaxf(row_max, s);
-            }
-
-            // A fully-masked KV tile (every key causal- or window-masked) has
-            // row_max == -INFINITY. Folding it in computes expf(-inf - (-inf))
-            // = NaN, and a later valid tile's 0*NaN correction then poisons the
-            // whole row. This happens for sliding-window layers whenever a
-            // query's window starts past an early tile (the causal `continue`
-            // above only skips fully-future tiles, not out-of-window ones).
-            // A masked tile contributes nothing to the softmax — skip it.
-            if (row_max != -INFINITY) {
-                float m_new = fmaxf(m_val, row_max);
-                float psum = 0.0f;
-                for (int c = 0; c < kv_tile_cols; c++) {
-                    P[c] = expf(P[c] - m_new);
-                    psum += P[c];
-                }
-                float correction = expf(m_val - m_new);
-                l_val = correction * l_val + psum;
-                for (int d = 0; d < head_dim; d++) O_acc[d] *= correction;
-                m_val = m_new;
-            } else {
-                for (int c = 0; c < kv_tile_cols; c++) P[c] = 0.0f;
-            }
-        }
-
-        __syncthreads();
-
-        int v_elems = kv_tile_cols * head_dim;
-        for (int i = tid; i < v_elems; i += THREADS_PER_BLOCK) {
-            int row = i / head_dim;
-            int col = i % head_dim;
-            smem_kv[row * head_dim + col] = V_head[(kv_tile_start + row) * head_dim + col];
-        }
-        for (int i = v_elems + tid; i < BC * head_dim; i += THREADS_PER_BLOCK) {
-            smem_kv[i] = __float2bfloat16(0.0f);
-        }
-        __syncthreads();
-
-        if (owns_row) {
-            for (int c = 0; c < kv_tile_cols; c++) {
-                float p = P[c];
-                if (p != 0.0f) {
-                    for (int d = 0; d < head_dim; d++) {
-                        O_acc[d] += p * __bfloat162float(smem_kv[c * head_dim + d]);
-                    }
-                }
-            }
-        }
-        __syncthreads();
-    }
-
-    // Fold in the per-head attention sink (extra logit, no value contribution).
-    if (owns_row && sinks != nullptr) {
-        float sink_logit = __bfloat162float(sinks[q_head]);
-        float m_new = fmaxf(m_val, sink_logit);
-        float correction = expf(m_val - m_new);
-        l_val = correction * l_val + expf(sink_logit - m_new);
-        for (int d = 0; d < head_dim; d++) O_acc[d] *= correction;
-        m_val = m_new;
-    }
-
-    if (owns_row) {
-        float inv_l = (l_val > 0.0f) ? (1.0f / l_val) : 0.0f;
-        int global_row = q_tile_start + tid;
-        for (int d = 0; d < head_dim; d++) {
-            O_head[global_row * head_dim + d] = __float2bfloat16(O_acc[d] * inv_l);
-        }
-    }
-}
-
 // ============================================================
 // Decode Attention kernel: optimized for Q_len=1 (single-token decode).
 // Parallelizes across KV sequence dimension instead of Q rows.
@@ -464,7 +287,7 @@ __global__ void decode_attention_bf16_kernel(
     // Shared memory for reduction
     __shared__ float smem_max[32];   // one per warp
     __shared__ float smem_sum[32];
-    __shared__ float smem_O_warp[32][HEAD_DIM_MAX];
+    __shared__ float smem_O[HEAD_DIM_MAX]; // final output accumulator
 
     // Step 1: Block-wide max reduction
     int lane = tid & 31;
@@ -513,30 +336,35 @@ __global__ void decode_attention_bf16_kernel(
     __syncthreads();
     global_sum = smem_sum[0];
 
-    // Step 4: Reduce O across block, dim by dim. Store one partial per warp
-    // and sum in warp-id order; atomicAdd made greedy decode nondeterministic
-    // when logits were close (same fix pattern as paged_attention.cu / gemv.cu).
+    // Step 4: Reduce O across block (dimension by dimension using shared mem)
     float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
 
-    for (int i = tid; i < 32 * HEAD_DIM_MAX; i += DECODE_THREADS) {
-        reinterpret_cast<float*>(smem_O_warp)[i] = 0.0f;
+    // Process head_dim in chunks: each iteration reduces one dimension
+    // Use shared memory accumulator: each warp contributes via warp reduction + atomic
+    // Actually simpler: iterate over dimensions, warp reduce each, then lane0 atomicAdd to smem_O
+
+    // Initialize smem_O
+    for (int d = tid; d < head_dim; d += DECODE_THREADS) {
+        smem_O[d] = 0.0f;
     }
     __syncthreads();
 
+    // Each thread adds its local_O contributions via warp reduction + atomicAdd
     for (int d = 0; d < head_dim; d++) {
         float val = local_O[d];
+        // Warp-level reduction
         #pragma unroll
         for (int offset = 16; offset > 0; offset >>= 1)
             val += __shfl_down_sync(0xffffffff, val, offset);
-        if (lane == 0) smem_O_warp[warp_id][d] = val;
+        if (lane == 0) {
+            atomicAdd(&smem_O[d], val);
+        }
     }
     __syncthreads();
 
     // Thread 0..head_dim-1 write final output
     for (int d = tid; d < head_dim; d += DECODE_THREADS) {
-        float out = 0.0f;
-        for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
-        O_ptr[d] = __float2bfloat16(out * inv_sum);
+        O_ptr[d] = __float2bfloat16(smem_O[d] * inv_sum);
     }
 }
 
@@ -567,31 +395,6 @@ void launch_flash_attention_bf16(
     CUDA_CHECK_LAST_ERROR();
 }
 
-void launch_flash_attention_sinks_bf16(
-    const void* Q, const void* K, const void* V, void* O,
-    const void* sinks,
-    int batch, int num_q_heads, int num_kv_heads,
-    int q_len, int kv_len, int head_dim,
-    float scale, int causal, int window_size, void* stream
-) {
-    int q_tiles = (q_len + BR - 1) / BR;
-    dim3 grid(q_tiles, batch * num_q_heads);
-    int block = THREADS_PER_BLOCK;
-    int smem_bytes = (BR + BC) * head_dim * (int)sizeof(__nv_bfloat16);
-
-    flash_attention_sinks_bf16_kernel<<<grid, block, smem_bytes, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)Q,
-        (const __nv_bfloat16*)K,
-        (const __nv_bfloat16*)V,
-        (__nv_bfloat16*)O,
-        (const __nv_bfloat16*)sinks,
-        num_q_heads, num_kv_heads,
-        q_len, kv_len, head_dim,
-        scale, causal, window_size
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
 void launch_decode_attention_bf16(
     const void* Q, const void* K, const void* V, void* O,
     int batch, int num_q_heads, int num_kv_heads,

csrc/attention/paged_attention.cu

@@ -118,7 +118,7 @@ __global__ void paged_decode_attention_bf16_kernel(
     // ---- Block-level online softmax reduction ----
     __shared__ float smem_max[32];
     __shared__ float smem_sum[32];
-    __shared__ float smem_O_warp[32][PAGED_HEAD_DIM_MAX];
+    __shared__ float smem_O[PAGED_HEAD_DIM_MAX];
 
     int lane = tid & 31;
     int warp_id = tid >> 5;
@@ -164,12 +164,8 @@ __global__ void paged_decode_attention_bf16_kernel(
     __syncthreads();
     global_sum = smem_sum[0];
 
-    // Step 4: reduce O across block, dim by dim. Store one partial per warp
-    // and sum in warp-id order; atomicAdd made greedy decode nondeterministic
-    // when logits were close.
-    for (int i = tid; i < 32 * PAGED_HEAD_DIM_MAX; i += PAGED_THREADS) {
-        reinterpret_cast<float*>(smem_O_warp)[i] = 0.0f;
-    }
+    // Step 4: reduce O across block, dim by dim
+    for (int d = tid; d < head_dim; d += PAGED_THREADS) smem_O[d] = 0.0f;
     __syncthreads();
 
     for (int d = 0; d < head_dim; d++) {
@@ -177,349 +173,13 @@ __global__ void paged_decode_attention_bf16_kernel(
         #pragma unroll
         for (int offset = 16; offset > 0; offset >>= 1)
             val += __shfl_down_sync(0xffffffff, val, offset);
-        if (lane == 0) smem_O_warp[warp_id][d] = val;
+        if (lane == 0) atomicAdd(&smem_O[d], val);
     }
     __syncthreads();
 
     float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
     for (int d = tid; d < head_dim; d += PAGED_THREADS) {
-        float out = 0.0f;
-        for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
-        O_ptr[d] = __float2bfloat16(out * inv_sum);
-    }
-}
-
-// Tree-aware paged decode attention: per-query mask lets sibling candidates
-// in the same batch attend to different subsets of newly-written K/V.
-// `tree_start`: position where newly-written K/V begins (typically pos_offset).
-// `tree_len`: number of newly-written K/V rows (= batch, one per query).
-// `tree_mask[i][j] = 1` iff query i attends to K/V at position `tree_start+j`.
-// Positions < tree_start are always attended (regular history).
-__global__ void paged_decode_attention_tree_bf16_kernel(
-    const __nv_bfloat16* __restrict__ Q,
-    const __nv_bfloat16* __restrict__ K_cache,
-    const __nv_bfloat16* __restrict__ V_cache,
-    __nv_bfloat16* __restrict__ O,
-    const int* __restrict__ block_tables,
-    const int* __restrict__ context_lens,
-    const int* __restrict__ tree_mask,       // [batch, tree_len] int32
-    int num_q_heads, int num_kv_heads,
-    int head_dim, int max_blocks_per_seq,
-    int tree_start, int tree_len,
-    float scale
-) {
-    int seq_idx = blockIdx.y;
-    int q_head  = blockIdx.x;
-    int tid     = threadIdx.x;
-
-    int kv_len = context_lens[seq_idx];
-    if (kv_len <= 0) {
-        if (tid < head_dim) {
-            O[((long long)seq_idx * num_q_heads + q_head) * head_dim + tid] =
-                __float2bfloat16(0.0f);
-        }
-        return;
-    }
-
-    int heads_per_group = num_q_heads / num_kv_heads;
-    int kv_head = q_head / heads_per_group;
-
-    const __nv_bfloat16* Q_ptr = Q +
-        ((long long)seq_idx * num_q_heads + q_head) * head_dim;
-    __nv_bfloat16* O_ptr = O +
-        ((long long)seq_idx * num_q_heads + q_head) * head_dim;
-    const int* bt = block_tables + (long long)seq_idx * max_blocks_per_seq;
-    const int* mask_row = tree_mask + (long long)seq_idx * tree_len;
-
-    float q_reg[PAGED_HEAD_DIM_MAX];
-    for (int d = 0; d < head_dim; d++) {
-        q_reg[d] = __bfloat162float(Q_ptr[d]);
-    }
-
-    float local_max = -INFINITY;
-    float local_sum = 0.0f;
-    float local_O[PAGED_HEAD_DIM_MAX];
-    for (int d = 0; d < head_dim; d++) local_O[d] = 0.0f;
-
-    int kv_stride_block = num_kv_heads * PAGED_BLOCK_SIZE * head_dim;
-    int kv_stride_head  = PAGED_BLOCK_SIZE * head_dim;
-
-    for (int pos = tid; pos < kv_len; pos += PAGED_THREADS) {
-        // Tree mask: skip positions in [tree_start, tree_start+tree_len) that
-        // the mask marks as 0. Everything else (history) is always attended.
-        if (pos >= tree_start && pos < tree_start + tree_len) {
-            if (mask_row[pos - tree_start] == 0) continue;
-        }
-
-        int logical_blk = pos / PAGED_BLOCK_SIZE;
-        int slot_in_blk = pos % PAGED_BLOCK_SIZE;
-        int phys_blk = bt[logical_blk];
-
-        const __nv_bfloat16* K_pos = K_cache
-            + (long long)phys_blk * kv_stride_block
-            + kv_head * kv_stride_head
-            + slot_in_blk * head_dim;
-        const __nv_bfloat16* V_pos = V_cache
-            + (long long)phys_blk * kv_stride_block
-            + kv_head * kv_stride_head
-            + slot_in_blk * head_dim;
-
-        float dot = 0.0f;
-        for (int d = 0; d < head_dim; d++) {
-            dot += q_reg[d] * __bfloat162float(K_pos[d]);
-        }
-        float s = dot * scale;
-
-        float new_max = fmaxf(local_max, s);
-        float correction = expf(local_max - new_max);
-        float p = expf(s - new_max);
-
-        local_sum = local_sum * correction + p;
-        for (int d = 0; d < head_dim; d++) local_O[d] *= correction;
-
-        for (int d = 0; d < head_dim; d++) {
-            local_O[d] += p * __bfloat162float(V_pos[d]);
-        }
-
-        local_max = new_max;
-    }
-
-    // Block-level reduction (identical to base kernel).
-    __shared__ float smem_max[32];
-    __shared__ float smem_sum[32];
-    __shared__ float smem_O_warp[32][PAGED_HEAD_DIM_MAX];
-
-    int lane = tid & 31;
-    int warp_id = tid >> 5;
-    int num_warps = PAGED_THREADS >> 5;
-
-    float warp_max = local_max;
-    #pragma unroll
-    for (int offset = 16; offset > 0; offset >>= 1)
-        warp_max = fmaxf(warp_max, __shfl_down_sync(0xffffffff, warp_max, offset));
-    if (lane == 0) smem_max[warp_id] = warp_max;
-    __syncthreads();
-
-    float global_max;
-    if (tid == 0) {
-        global_max = smem_max[0];
-        for (int i = 1; i < num_warps; i++)
-            global_max = fmaxf(global_max, smem_max[i]);
-        smem_max[0] = global_max;
-    }
-    __syncthreads();
-    global_max = smem_max[0];
-
-    float rescale = (local_max == -INFINITY) ? 0.0f : expf(local_max - global_max);
-    local_sum *= rescale;
-    for (int d = 0; d < head_dim; d++) local_O[d] *= rescale;
-
-    float warp_sum = local_sum;
-    #pragma unroll
-    for (int offset = 16; offset > 0; offset >>= 1)
-        warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset);
-    if (lane == 0) smem_sum[warp_id] = warp_sum;
-    __syncthreads();
-
-    float global_sum;
-    if (tid == 0) {
-        global_sum = 0.0f;
-        for (int i = 0; i < num_warps; i++) global_sum += smem_sum[i];
-        smem_sum[0] = global_sum;
-    }
-    __syncthreads();
-    global_sum = smem_sum[0];
-
-    for (int i = tid; i < 32 * PAGED_HEAD_DIM_MAX; i += PAGED_THREADS) {
-        reinterpret_cast<float*>(smem_O_warp)[i] = 0.0f;
-    }
-    __syncthreads();
-
-    for (int d = 0; d < head_dim; d++) {
-        float val = local_O[d];
-        #pragma unroll
-        for (int offset = 16; offset > 0; offset >>= 1)
-            val += __shfl_down_sync(0xffffffff, val, offset);
-        if (lane == 0) smem_O_warp[warp_id][d] = val;
-    }
-    __syncthreads();
-
-    float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
-    for (int d = tid; d < head_dim; d += PAGED_THREADS) {
-        float out = 0.0f;
-        for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
-        O_ptr[d] = __float2bfloat16(out * inv_sum);
-    }
-}
-
-// Extended paged decode attention with attention sinks and sliding window.
-// sinks: [num_q_heads] BF16 — per-head extra logit appended before softmax.
-// window_size: >0 = sliding window (only attend to last `window_size` positions), 0 = full.
-__global__ void paged_decode_attention_sinks_bf16_kernel(
-    const __nv_bfloat16* __restrict__ Q,
-    const __nv_bfloat16* __restrict__ K_cache,
-    const __nv_bfloat16* __restrict__ V_cache,
-    __nv_bfloat16* __restrict__ O,
-    const int* __restrict__ block_tables,
-    const int* __restrict__ context_lens,
-    const __nv_bfloat16* __restrict__ sinks, // [num_q_heads] or NULL
-    int num_q_heads, int num_kv_heads,
-    int head_dim, int max_blocks_per_seq,
-    float scale, int window_size
-) {
-    int seq_idx = blockIdx.y;
-    int q_head  = blockIdx.x;
-    int tid     = threadIdx.x;
-
-    int kv_len = context_lens[seq_idx];
-    if (kv_len <= 0) {
-        if (tid < head_dim) {
-            O[((long long)seq_idx * num_q_heads + q_head) * head_dim + tid] =
-                __float2bfloat16(0.0f);
-        }
-        return;
-    }
-
-    int heads_per_group = num_q_heads / num_kv_heads;
-    int kv_head = q_head / heads_per_group;
-
-    const __nv_bfloat16* Q_ptr = Q +
-        ((long long)seq_idx * num_q_heads + q_head) * head_dim;
-    __nv_bfloat16* O_ptr = O +
-        ((long long)seq_idx * num_q_heads + q_head) * head_dim;
-    const int* bt = block_tables + (long long)seq_idx * max_blocks_per_seq;
-
-    // Sliding window: only attend to positions [kv_len - window_size, kv_len)
-    int start_pos = 0;
-    if (window_size > 0 && kv_len > window_size) {
-        start_pos = kv_len - window_size;
-    }
-
-    float q_reg[PAGED_HEAD_DIM_MAX];
-    for (int d = 0; d < head_dim; d++) {
-        q_reg[d] = __bfloat162float(Q_ptr[d]);
-    }
-
-    float local_max = -INFINITY;
-    float local_sum = 0.0f;
-    float local_O[PAGED_HEAD_DIM_MAX];
-    for (int d = 0; d < head_dim; d++) local_O[d] = 0.0f;
-
-    int kv_stride_block = num_kv_heads * PAGED_BLOCK_SIZE * head_dim;
-    int kv_stride_head  = PAGED_BLOCK_SIZE * head_dim;
-
-    int attend_len = kv_len - start_pos;
-    for (int rel = tid; rel < attend_len; rel += PAGED_THREADS) {
-        int pos = start_pos + rel;
-        int logical_blk = pos / PAGED_BLOCK_SIZE;
-        int slot_in_blk = pos % PAGED_BLOCK_SIZE;
-        int phys_blk = bt[logical_blk];
-
-        const __nv_bfloat16* K_pos = K_cache
-            + (long long)phys_blk * kv_stride_block
-            + kv_head * kv_stride_head
-            + slot_in_blk * head_dim;
-        const __nv_bfloat16* V_pos = V_cache
-            + (long long)phys_blk * kv_stride_block
-            + kv_head * kv_stride_head
-            + slot_in_blk * head_dim;
-
-        float dot = 0.0f;
-        for (int d = 0; d < head_dim; d++) {
-            dot += q_reg[d] * __bfloat162float(K_pos[d]);
-        }
-        float s = dot * scale;
-
-        float new_max = fmaxf(local_max, s);
-        float correction = expf(local_max - new_max);
-        float p = expf(s - new_max);
-
-        local_sum = local_sum * correction + p;
-        for (int d = 0; d < head_dim; d++) local_O[d] *= correction;
-        for (int d = 0; d < head_dim; d++) {
-            local_O[d] += p * __bfloat162float(V_pos[d]);
-        }
-        local_max = new_max;
-    }
-
-    // Include the sink logit (only thread 0 handles it to avoid double-counting)
-    float sink_logit = -INFINITY;
-    if (sinks != nullptr && tid == 0) {
-        sink_logit = __bfloat162float(sinks[q_head]);
-        float new_max = fmaxf(local_max, sink_logit);
-        float correction = expf(local_max - new_max);
-        float p = expf(sink_logit - new_max);
-        local_sum = local_sum * correction + p;
-        for (int d = 0; d < head_dim; d++) local_O[d] *= correction;
-        // Sink absorbs probability but produces no value output (p * 0)
-        local_max = new_max;
-    }
-
-    // ---- Block-level online softmax reduction (same as base kernel) ----
-    __shared__ float smem_max[32];
-    __shared__ float smem_sum[32];
-    __shared__ float smem_O_warp[32][PAGED_HEAD_DIM_MAX];
-
-    int lane = tid & 31;
-    int warp_id = tid >> 5;
-    int num_warps = PAGED_THREADS >> 5;
-
-    float warp_max = local_max;
-    #pragma unroll
-    for (int offset = 16; offset > 0; offset >>= 1)
-        warp_max = fmaxf(warp_max, __shfl_down_sync(0xffffffff, warp_max, offset));
-    if (lane == 0) smem_max[warp_id] = warp_max;
-    __syncthreads();
-
-    float global_max;
-    if (tid == 0) {
-        global_max = smem_max[0];
-        for (int i = 1; i < num_warps; i++)
-            global_max = fmaxf(global_max, smem_max[i]);
-        smem_max[0] = global_max;
-    }
-    __syncthreads();
-    global_max = smem_max[0];
-
-    float rescale = (local_max == -INFINITY) ? 0.0f : expf(local_max - global_max);
-    local_sum *= rescale;
-    for (int d = 0; d < head_dim; d++) local_O[d] *= rescale;
-
-    float warp_sum = local_sum;
-    #pragma unroll
-    for (int offset = 16; offset > 0; offset >>= 1)
-        warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset);
-    if (lane == 0) smem_sum[warp_id] = warp_sum;
-    __syncthreads();
-
-    float global_sum;
-    if (tid == 0) {
-        global_sum = 0.0f;
-        for (int i = 0; i < num_warps; i++) global_sum += smem_sum[i];
-        smem_sum[0] = global_sum;
-    }
-    __syncthreads();
-    global_sum = smem_sum[0];
-
-    for (int i = tid; i < 32 * PAGED_HEAD_DIM_MAX; i += PAGED_THREADS) {
-        reinterpret_cast<float*>(smem_O_warp)[i] = 0.0f;
-    }
-    __syncthreads();
-
-    for (int d = 0; d < head_dim; d++) {
-        float val = local_O[d];
-        #pragma unroll
-        for (int offset = 16; offset > 0; offset >>= 1)
-            val += __shfl_down_sync(0xffffffff, val, offset);
-        if (lane == 0) smem_O_warp[warp_id][d] = val;
-    }
-    __syncthreads();
-
-    float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
-    for (int d = tid; d < head_dim; d += PAGED_THREADS) {
-        float out = 0.0f;
-        for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
-        O_ptr[d] = __float2bfloat16(out * inv_sum);
+        O_ptr[d] = __float2bfloat16(smem_O[d] * inv_sum);
     }
 }
 
@@ -552,63 +212,4 @@ void launch_paged_decode_attention_bf16(
     CUDA_CHECK_LAST_ERROR();
 }
 
-void launch_paged_decode_attention_tree_bf16(
-    const void* Q,
-    const void* K_cache,
-    const void* V_cache,
-    void* O,
-    const int* block_tables,
-    const int* context_lens,
-    const int* tree_mask,
-    int batch, int num_q_heads, int num_kv_heads,
-    int head_dim, int max_blocks_per_seq,
-    int tree_start, int tree_len,
-    float scale, void* stream
-) {
-    dim3 grid(num_q_heads, batch);
-    int block = PAGED_THREADS;
-
-    paged_decode_attention_tree_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)Q,
-        (const __nv_bfloat16*)K_cache,
-        (const __nv_bfloat16*)V_cache,
-        (__nv_bfloat16*)O,
-        block_tables, context_lens, tree_mask,
-        num_q_heads, num_kv_heads,
-        head_dim, max_blocks_per_seq,
-        tree_start, tree_len,
-        scale
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_paged_decode_attention_sinks_bf16(
-    const void* Q,
-    const void* K_cache,
-    const void* V_cache,
-    void* O,
-    const int* block_tables,
-    const int* context_lens,
-    const void* sinks,
-    int batch, int num_q_heads, int num_kv_heads,
-    int head_dim, int max_blocks_per_seq,
-    float scale, int window_size, void* stream
-) {
-    dim3 grid(num_q_heads, batch);
-    int block = PAGED_THREADS;
-
-    paged_decode_attention_sinks_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)Q,
-        (const __nv_bfloat16*)K_cache,
-        (const __nv_bfloat16*)V_cache,
-        (__nv_bfloat16*)O,
-        block_tables, context_lens,
-        (const __nv_bfloat16*)sinks,
-        num_q_heads, num_kv_heads,
-        head_dim, max_blocks_per_seq,
-        scale, window_size
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
 }

csrc/attention/reshape_and_cache.cu

@@ -1,215 +0,0 @@
-#include <cuda_bf16.h>
-#include "../common.cuh"
-
-// Scatter [num_tokens] new K/V into a paged KV pool for ONE sequence.
-//
-// Source layouts (BF16, contiguous):
-//   k_src, v_src : [num_kv_heads, num_tokens, head_dim]   (head-major)
-//
-// Pool layouts (BF16, contiguous):
-//   k_pool, v_pool : [num_blocks_total, num_kv_heads, BLOCK_SIZE, head_dim]
-//
-// For token t (0 <= t < num_tokens):
-//   p           = start_pos + t
-//   logical_blk = p / BLOCK_SIZE
-//   slot_in_blk = p % BLOCK_SIZE
-//   phys        = block_ids[logical_blk]
-//   pool[phys, h, slot_in_blk, :] := src[h, t, :]
-//
-// Replaces a Rust-side per-token, per-head cudaMemcpy loop. With Qwen3-8B
-// (8 KV heads, 36 layers) and a 1024-token prefill, that loop fired
-// ~290k device-side memcpys; one kernel launch per layer is dramatically
-// less overhead.
-//
-// Grid : (num_tokens, num_kv_heads)
-// Block: head_dim threads (≤128 in practice; head_dim is padded to a
-//        multiple of 32 by the model and all our shipping configs are
-//        128, so a single warp's worth handles two slots in flight).
-
-__global__ void reshape_and_cache_bf16_kernel(
-    const __nv_bfloat16* __restrict__ k_src,
-    const __nv_bfloat16* __restrict__ v_src,
-    __nv_bfloat16* __restrict__ k_pool,
-    __nv_bfloat16* __restrict__ v_pool,
-    const int* __restrict__ block_ids,
-    int num_tokens, int num_heads,
-    int head_dim, int start_pos, int block_size
-) {
-    int t = blockIdx.x;
-    int h = blockIdx.y;
-    if (t >= num_tokens || h >= num_heads) return;
-
-    int p           = start_pos + t;
-    int logical_blk = p / block_size;
-    int slot_in_blk = p - logical_blk * block_size;
-    int phys        = block_ids[logical_blk];
-
-    long long src_off = ((long long)h * num_tokens + t) * head_dim;
-    long long dst_off = (((long long)phys * num_heads + h) * block_size + slot_in_blk) * head_dim;
-
-    int tid = threadIdx.x;
-    int blockSize = blockDim.x;
-
-    // Per-thread strided copy. head_dim is typically 128 and blockSize is
-    // 128, so each thread copies exactly one element — but the loop keeps
-    // the kernel correct for non-128 head_dim configs (Phi-style 64, etc.).
-    for (int d = tid; d < head_dim; d += blockSize) {
-        k_pool[dst_off + d] = k_src[src_off + d];
-        v_pool[dst_off + d] = v_src[src_off + d];
-    }
-}
-
-// Batched variant: writes one new K/V token per sequence into a paged
-// pool, indexed by a per-batch block table that also drives the paged
-// attention kernel. Used in the decode path where every seq advances
-// by exactly one position per step.
-//
-// Source layouts (BF16, contiguous):
-//   k_src, v_src : [batch, num_kv_heads, head_dim]
-//
-// Pool layouts (BF16, contiguous):
-//   k_pool, v_pool : [num_blocks_total, num_kv_heads, BLOCK_SIZE, head_dim]
-//
-// block_tables : int32 [batch, max_blocks_per_seq]
-// kv_lens      : int32 [batch]   (current seq_len BEFORE this step + 1
-//                                  — i.e. the same buffer paged attention
-//                                  reads. The new token's logical index
-//                                  is `kv_lens[b] - 1`.)
-//
-// Grid : (batch, num_kv_heads)
-// Block: head_dim threads.
-
-__global__ void reshape_and_cache_batched_bf16_kernel(
-    const __nv_bfloat16* __restrict__ k_src,
-    const __nv_bfloat16* __restrict__ v_src,
-    __nv_bfloat16* __restrict__ k_pool,
-    __nv_bfloat16* __restrict__ v_pool,
-    const int* __restrict__ block_tables,
-    const int* __restrict__ kv_lens,
-    int num_heads, int head_dim,
-    int block_size, int max_blocks_per_seq
-) {
-    int b = blockIdx.x;
-    int h = blockIdx.y;
-
-    int new_pos     = kv_lens[b] - 1;
-    int logical_blk = new_pos / block_size;
-    int slot_in_blk = new_pos - logical_blk * block_size;
-    int phys        = block_tables[b * max_blocks_per_seq + logical_blk];
-
-    long long src_off = ((long long)b * num_heads + h) * head_dim;
-    long long dst_off = (((long long)phys * num_heads + h) * block_size + slot_in_blk) * head_dim;
-
-    int tid = threadIdx.x;
-    int blockSize = blockDim.x;
-    for (int d = tid; d < head_dim; d += blockSize) {
-        k_pool[dst_off + d] = k_src[src_off + d];
-        v_pool[dst_off + d] = v_src[src_off + d];
-    }
-}
-
-extern "C" {
-
-void launch_reshape_and_cache_bf16(
-    const void* k_src, const void* v_src,
-    void* k_pool, void* v_pool,
-    const void* block_ids,
-    int num_tokens, int num_heads,
-    int head_dim, int start_pos, int block_size,
-    void* stream
-) {
-    if (num_tokens <= 0) return;
-    int threads = head_dim < 32 ? 32 : head_dim;
-    if (threads > 1024) threads = 1024;
-    dim3 grid(num_tokens, num_heads);
-    reshape_and_cache_bf16_kernel<<<grid, threads, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)k_src,
-        (const __nv_bfloat16*)v_src,
-        (__nv_bfloat16*)k_pool,
-        (__nv_bfloat16*)v_pool,
-        (const int*)block_ids,
-        num_tokens, num_heads,
-        head_dim, start_pos, block_size
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_reshape_and_cache_batched_bf16(
-    const void* k_src, const void* v_src,
-    void* k_pool, void* v_pool,
-    const void* block_tables, const void* kv_lens,
-    int batch, int num_heads,
-    int head_dim, int block_size, int max_blocks_per_seq,
-    void* stream
-) {
-    if (batch <= 0 || num_heads <= 0) return;
-    int threads = head_dim < 32 ? 32 : head_dim;
-    if (threads > 1024) threads = 1024;
-    dim3 grid(batch, num_heads);
-    reshape_and_cache_batched_bf16_kernel<<<grid, threads, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)k_src,
-        (const __nv_bfloat16*)v_src,
-        (__nv_bfloat16*)k_pool,
-        (__nv_bfloat16*)v_pool,
-        (const int*)block_tables,
-        (const int*)kv_lens,
-        num_heads, head_dim, block_size, max_blocks_per_seq
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-// Copy one token's K/V from src_pos to dst_pos within one pool.
-// Grid: (num_kv_heads,). Block: head_dim threads.
-// pool: [num_blocks_total, num_kv_heads, block_size, head_dim]
-// block_ids: [max_blocks] for this sequence (logical → physical block map).
-__global__ void copy_kv_position_kernel(
-    __nv_bfloat16* __restrict__ pool,
-    const int* __restrict__ block_ids,
-    int src_pos, int dst_pos,
-    int head_dim, int block_size
-) {
-    int h = blockIdx.x;
-    int d = threadIdx.x;
-    if (d >= head_dim) return;
-
-    int num_kv_heads = gridDim.x;
-
-    int src_blk = src_pos / block_size;
-    int src_slot = src_pos % block_size;
-    int src_phys = block_ids[src_blk];
-
-    int dst_blk = dst_pos / block_size;
-    int dst_slot = dst_pos % block_size;
-    int dst_phys = block_ids[dst_blk];
-
-    long long src_off = ((long long)src_phys * num_kv_heads + h) * block_size * head_dim
-                      + src_slot * head_dim + d;
-    long long dst_off = ((long long)dst_phys * num_kv_heads + h) * block_size * head_dim
-                      + dst_slot * head_dim + d;
-
-    pool[dst_off] = pool[src_off];
-}
-
-void launch_copy_kv_position(
-    void* k_pool, void* v_pool,
-    const int* block_ids,
-    int src_pos, int dst_pos,
-    int num_kv_heads, int head_dim, int block_size,
-    void* stream
-) {
-    int threads = head_dim < 32 ? 32 : head_dim;
-    if (threads > 1024) threads = 1024;
-    dim3 grid(num_kv_heads);
-    copy_kv_position_kernel<<<grid, threads, 0, (cudaStream_t)stream>>>(
-        (__nv_bfloat16*)k_pool, block_ids,
-        src_pos, dst_pos, head_dim, block_size
-    );
-    CUDA_CHECK_LAST_ERROR();
-    copy_kv_position_kernel<<<grid, threads, 0, (cudaStream_t)stream>>>(
-        (__nv_bfloat16*)v_pool, block_ids,
-        src_pos, dst_pos, head_dim, block_size
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-}

csrc/common.cuh

@@ -49,12 +49,10 @@ __device__ __forceinline__ float block_reduce_max(float val) {
     return val;
 }
 
-// --- Launch error checking ---
-// Always on, including release builds. A launch with an invalid config
-// (e.g. 32-bit overflow in grid/index math) is otherwise silent and produces
-// garbage with no clue — the MoE int32-overflow bug was found exactly because
-// release swallowed the launch failure. `cudaGetLastError()` does not
-// synchronize the stream, so the per-launch host cost is negligible.
+// --- Launch error checking (debug builds only) ---
+#ifdef NDEBUG
+#define CUDA_CHECK_LAST_ERROR() ((void)0)
+#else
 #include <cstdio>
 #define CUDA_CHECK_LAST_ERROR() do { \
     cudaError_t err = cudaGetLastError(); \
@@ -63,3 +61,4 @@ __device__ __forceinline__ float block_reduce_max(float val) {
                 __FILE__, __LINE__, cudaGetErrorString(err)); \
     } \
 } while(0)
+#endif

csrc/gemm/gemv.cu

@@ -2,23 +2,27 @@
 #include <cuda_runtime.h>
 #include "../common.cuh"
 
-// K-split GEMV for M=1 BF16 decode.
-//
+// Custom GEMV kernel for M=1 decode step (BF16):
 //   y[n] = sum_k x[k] * W[k * N + n]
+// where x: [K] (BF16), W: [K, N] (BF16, row-major), y: [N] (BF16).
 //
-// Grid: (N / TILE_N, K / TILE_K) partials, followed by a deterministic
-// fixed-order reduction over K blocks. The previous implementation used
-// atomicAdd into y_fp32[col]; that made BF16 greedy decode sensitive to
-// inter-block scheduling when logits were close.
+// Design: K-split for high occupancy on large GPU (170 SMs).
+// Grid: (N / TILE_N, K / TILE_K) — each block computes a partial sum
+// for TILE_N output columns over a TILE_K slice of K.
+// Partial results are atomicAdd'd to an FP32 accumulator, then a
+// second kernel converts FP32 -> BF16.
+//
+// Memory access: adjacent threads read adjacent columns of the same row
+// of W, giving perfectly coalesced 128-byte transactions.
 
 #define GEMV_TILE_N 128
 #define GEMV_TILE_K 256
-#define GEMV_BLOCK  128
+#define GEMV_BLOCK  128  // = TILE_N, one thread per output column
 
-__global__ void gemv_bf16_partial_kernel(
-    const __nv_bfloat16* __restrict__ x,
-    const __nv_bfloat16* __restrict__ W,
-    float* __restrict__ partials,
+__global__ void gemv_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,      // [K]
+    const __nv_bfloat16* __restrict__ W,      // [K, N] row-major
+    float* __restrict__ y_fp32,               // [N] accumulator
     int K, int N
 ) {
     const int block_n = blockIdx.x;
@@ -26,121 +30,60 @@ __global__ void gemv_bf16_partial_kernel(
     const int t = threadIdx.x;
     const int col = block_n * GEMV_TILE_N + t;
 
+    if (col >= N) return;
+
     const int k_start = block_k * GEMV_TILE_K;
     const int k_end = min(k_start + GEMV_TILE_K, K);
     const int k_len = k_end - k_start;
 
-    // Cooperative load of x into shared memory uses ALL threads in the block
-    // (indexed by t, independent of col). Threads whose column is out of range
-    // must still help load and reach the barrier — returning early here would
-    // leave part of x_shared uninitialized AND make __syncthreads divergent
-    // (UB). So the col>=N check happens only AFTER the load + barrier. This bug
-    // produced intermittent huge/garbage outputs whenever N % GEMV_TILE_N != 0
-    // (e.g. gpt-oss decode o_proj with N=2880), collapsing the forward pass.
+    // Load x[k_start..k_end] into shared memory as FP32
     __shared__ float x_shared[GEMV_TILE_K];
     for (int i = t; i < k_len; i += GEMV_BLOCK) {
         x_shared[i] = __bfloat162float(x[k_start + i]);
     }
     __syncthreads();
 
-    if (col >= N) return;
-
+    // Compute partial dot product for this column
     float sum = 0.0f;
     for (int ki = 0; ki < k_len; ki++) {
-        sum += x_shared[ki] * __bfloat162float(W[(long long)(k_start + ki) * N + col]);
+        sum += x_shared[ki] * __bfloat162float(W[(k_start + ki) * N + col]);
     }
 
-    partials[(long long)block_k * N + col] = sum;
+    // Atomic accumulate (handles K-split reduction)
+    atomicAdd(&y_fp32[col], sum);
 }
 
-__global__ void gemv_reduce_to_bf16_kernel(
-    const float* __restrict__ partials,
+// Conversion kernel: FP32 accumulator -> BF16 output
+__global__ void gemv_fp32_to_bf16_kernel(
+    const float* __restrict__ src,
     __nv_bfloat16* __restrict__ dst,
-    int n,
-    int num_k_blocks
+    int n
 ) {
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (idx < n) {
-        float sum = 0.0f;
-        for (int kb = 0; kb < num_k_blocks; kb++) {
-            sum += partials[(long long)kb * n + idx];
-        }
-        dst[idx] = __float2bfloat16(sum);
+        dst[idx] = __float2bfloat16(src[idx]);
     }
 }
 
-// Batched variant: M rows, same W. Grid.z = batch row index.
-// Numerically identical to calling launch_gemv_bf16 M times in sequence because
-// each z-slice executes the same accumulation order on the same data.
-// partials buffer must be [M * num_k_blocks * N] floats.
-__global__ void gemv_bf16_batched_partial_kernel(
-    const __nv_bfloat16* __restrict__ x,  // [M, K]
-    const __nv_bfloat16* __restrict__ W,  // [K, N]
-    float* __restrict__ partials,         // [M, num_k_blocks, N]
-    int K, int N
-) {
-    const int block_n = blockIdx.x;
-    const int block_k = blockIdx.y;
-    const int row = blockIdx.z;
-    const int t = threadIdx.x;
-    const int col = block_n * GEMV_TILE_N + t;
-
-    const int k_start = block_k * GEMV_TILE_K;
-    const int k_end = min(k_start + GEMV_TILE_K, K);
-    const int k_len = k_end - k_start;
-
-    __shared__ float x_shared[GEMV_TILE_K];
-    const __nv_bfloat16* x_row = x + (long long)row * K;
-    for (int i = t; i < k_len; i += GEMV_BLOCK) {
-        x_shared[i] = __bfloat162float(x_row[k_start + i]);
-    }
-    __syncthreads();
-
-    if (col >= N) return;
-
-    float sum = 0.0f;
-    for (int ki = 0; ki < k_len; ki++) {
-        sum += x_shared[ki] * __bfloat162float(W[(long long)(k_start + ki) * N + col]);
-    }
-
-    int num_k_blocks = (K + GEMV_TILE_K - 1) / GEMV_TILE_K;
-    partials[((long long)row * num_k_blocks + block_k) * N + col] = sum;
-}
-
-__global__ void gemv_batched_reduce_to_bf16_kernel(
-    const float* __restrict__ partials,  // [M, num_k_blocks, N]
-    __nv_bfloat16* __restrict__ dst,     // [M, N]
-    int n,
-    int num_k_blocks
-) {
-    int col = blockIdx.x * blockDim.x + threadIdx.x;
-    int row = blockIdx.y;
-    if (col >= n) return;
-
-    float sum = 0.0f;
-    const float* row_partials = partials + (long long)row * num_k_blocks * n;
-    for (int kb = 0; kb < num_k_blocks; kb++) {
-        sum += row_partials[(long long)kb * n + col];
-    }
-    dst[(long long)row * n + col] = __float2bfloat16(sum);
-}
-
 extern "C" {
 
 void launch_gemv_bf16(
-    const void* x,
-    const void* W,
-    void* y_bf16,
-    void* y_fp32_buf,
+    const void* x,           // [K] BF16
+    const void* W,           // [K, N] BF16 row-major
+    void* y_bf16,            // [N] BF16 output
+    void* y_fp32_buf,        // [N] FP32 temporary (caller-provided)
     int K, int N,
     void* stream
 ) {
     cudaStream_t s = (cudaStream_t)stream;
 
-    int num_k_blocks = (K + GEMV_TILE_K - 1) / GEMV_TILE_K;
-    dim3 grid((N + GEMV_TILE_N - 1) / GEMV_TILE_N, num_k_blocks);
+    // Zero the FP32 accumulator
+    cudaMemsetAsync((float*)y_fp32_buf, 0, N * sizeof(float), s);
 
-    gemv_bf16_partial_kernel<<<grid, GEMV_BLOCK, 0, s>>>(
+    // Launch GEMV kernel
+    dim3 grid((N + GEMV_TILE_N - 1) / GEMV_TILE_N,
+              (K + GEMV_TILE_K - 1) / GEMV_TILE_K);
+    gemv_bf16_kernel<<<grid, GEMV_BLOCK, 0, s>>>(
         (const __nv_bfloat16*)x,
         (const __nv_bfloat16*)W,
         (float*)y_fp32_buf,
@@ -148,47 +91,13 @@ void launch_gemv_bf16(
     );
     CUDA_CHECK_LAST_ERROR();
 
-    // Fixed-order FP32 reduction over K blocks, then BF16 conversion.
+    // Convert FP32 -> BF16
     int conv_block = 256;
     int conv_grid = (N + conv_block - 1) / conv_block;
-    gemv_reduce_to_bf16_kernel<<<conv_grid, conv_block, 0, s>>>(
+    gemv_fp32_to_bf16_kernel<<<conv_grid, conv_block, 0, s>>>(
         (const float*)y_fp32_buf,
         (__nv_bfloat16*)y_bf16,
-        N,
-        num_k_blocks
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_gemv_bf16_batched(
-    const void* x,     // [M, K] BF16
-    const void* W,     // [K, N] BF16
-    void* y_bf16,      // [M, N] BF16
-    void* y_fp32_buf,  // [M * num_k_blocks * N] FP32
-    int M, int K, int N,
-    void* stream
-) {
-    cudaStream_t s = (cudaStream_t)stream;
-
-    int num_k_blocks = (K + GEMV_TILE_K - 1) / GEMV_TILE_K;
-    dim3 grid((N + GEMV_TILE_N - 1) / GEMV_TILE_N, num_k_blocks, M);
-
-    gemv_bf16_batched_partial_kernel<<<grid, GEMV_BLOCK, 0, s>>>(
-        (const __nv_bfloat16*)x,
-        (const __nv_bfloat16*)W,
-        (float*)y_fp32_buf,
-        K, N
-    );
-    CUDA_CHECK_LAST_ERROR();
-
-    int conv_block = 256;
-    int conv_grid_x = (N + conv_block - 1) / conv_block;
-    dim3 reduce_grid(conv_grid_x, M);
-    gemv_batched_reduce_to_bf16_kernel<<<reduce_grid, conv_block, 0, s>>>(
-        (const float*)y_fp32_buf,
-        (__nv_bfloat16*)y_bf16,
-        N,
-        num_k_blocks
+        N
     );
     CUDA_CHECK_LAST_ERROR();
 }

csrc/moe/moe_kernels.cu

@@ -1,254 +0,0 @@
-#include <cuda_bf16.h>
-#include <float.h>
-#include "../common.cuh"
-
-// ============================================================
-// MoE Top-K + Softmax kernel
-//
-// Input:  router_logits [num_tokens, num_experts]  BF16
-// Output: topk_ids      [num_tokens, top_k]        int32
-//         topk_weights  [num_tokens, top_k]         float32
-//
-// One block per token. Threads cooperatively find top-k indices
-// via repeated argmax, then compute softmax over the k winners.
-// num_experts <= 256 (fits in registers / shared memory).
-// ============================================================
-
-#define MOE_MAX_EXPERTS 256
-#define MOE_MAX_TOPK 8
-
-__global__ void moe_topk_softmax_bf16_kernel(
-    const __nv_bfloat16* __restrict__ router_logits,
-    int* __restrict__ topk_ids,
-    float* __restrict__ topk_weights,
-    int num_experts, int top_k
-) {
-    int token = blockIdx.x;
-    int tid = threadIdx.x;
-    const __nv_bfloat16* row = router_logits + token * num_experts;
-
-    // Load logits into shared memory
-    __shared__ float smem_logits[MOE_MAX_EXPERTS];
-    __shared__ int smem_ids[MOE_MAX_TOPK];
-    __shared__ float smem_vals[MOE_MAX_TOPK];
-
-    for (int i = tid; i < num_experts; i += blockDim.x) {
-        smem_logits[i] = __bfloat162float(row[i]);
-    }
-    __syncthreads();
-
-    // Find top-k via repeated argmax (k is small, typically 4)
-    if (tid == 0) {
-        for (int k = 0; k < top_k; k++) {
-            float best_val = -INFINITY;
-            int best_idx = 0;
-            for (int e = 0; e < num_experts; e++) {
-                if (smem_logits[e] > best_val) {
-                    best_val = smem_logits[e];
-                    best_idx = e;
-                }
-            }
-            smem_ids[k] = best_idx;
-            smem_vals[k] = best_val;
-            smem_logits[best_idx] = -INFINITY; // mask out selected
-        }
-
-        // Softmax over top-k values (in FP32)
-        float max_val = smem_vals[0];
-        for (int k = 1; k < top_k; k++)
-            max_val = fmaxf(max_val, smem_vals[k]);
-
-        float exp_sum = 0.0f;
-        for (int k = 0; k < top_k; k++) {
-            smem_vals[k] = expf(smem_vals[k] - max_val);
-            exp_sum += smem_vals[k];
-        }
-        float inv_sum = 1.0f / exp_sum;
-        for (int k = 0; k < top_k; k++)
-            smem_vals[k] *= inv_sum;
-
-        // Write outputs
-        for (int k = 0; k < top_k; k++) {
-            topk_ids[token * top_k + k] = smem_ids[k];
-            topk_weights[token * top_k + k] = smem_vals[k];
-        }
-    }
-}
-
-// ============================================================
-// MoE Replicate kernel
-//
-// Input:  x     [num_tokens, hidden]               BF16
-// Output: x_rep [local_experts, num_tokens, hidden] BF16
-//
-// Copies x into each expert's batch slot.
-// ============================================================
-
-__global__ void moe_replicate_bf16_kernel(
-    const __nv_bfloat16* __restrict__ x,
-    __nv_bfloat16* __restrict__ x_rep,
-    int num_tokens, int hidden, int local_experts
-) {
-    // 64-bit index: local_experts * num_tokens * hidden overflows int32 at
-    // ~2.3k prefill tokens (gpt-oss TP=1, 32 experts), which is inside the
-    // supported context window. A 32-bit `total` silently wraps, the launch
-    // fails, and (in release) the error is invisible — see common.cuh.
-    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
-    long long total = (long long)local_experts * num_tokens * hidden;
-    if (idx >= total) return;
-
-    // x_rep[expert, token, dim] = x[token, dim]
-    long long row_stride = (long long)num_tokens * hidden;
-    x_rep[idx] = x[idx % row_stride];
-}
-
-// ============================================================
-// MoE Bias Add 3D kernel
-//
-// Input:  x    [batch, num_tokens, dim]  BF16 (in-place output)
-//         bias [batch, dim]              BF16
-//
-// x[b, t, d] += bias[b, d]
-// ============================================================
-
-__global__ void moe_bias_add_3d_bf16_kernel(
-    __nv_bfloat16* __restrict__ x,
-    const __nv_bfloat16* __restrict__ bias,
-    int batch, int num_tokens, int dim
-) {
-    // 64-bit index: batch * num_tokens * dim overflows int32 at ~3.6k prefill
-    // tokens (gpt-oss TP=1, 32 experts, 2*intermediate dim) — see moe_replicate.
-    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
-    long long total = (long long)batch * num_tokens * dim;
-    if (idx >= total) return;
-
-    long long td = (long long)num_tokens * dim;
-    int b = (int)(idx / td);   // < batch (small)
-    int d = (int)(idx % dim);  // < dim
-    float v = __bfloat162float(x[idx]) + __bfloat162float(bias[(long long)b * dim + d]);
-    x[idx] = __float2bfloat16(v);
-}
-
-// ============================================================
-// MoE Weighted Sum kernel
-//
-// Input:  expert_out  [local_experts, num_tokens, hidden]  BF16
-//         topk_ids    [num_tokens, top_k]                  int32 (global expert ids)
-//         topk_weights[num_tokens, top_k]                  float32
-//         expert_start: first global expert id this rank owns
-//         local_experts: number of experts this rank owns
-//
-// Output: out [num_tokens, hidden] BF16
-//
-// For each (token, dim): accumulate in FP32:
-//   sum = 0
-//   for k in 0..top_k:
-//     global_id = topk_ids[token, k]
-//     if global_id in [expert_start, expert_start + local_experts):
-//       local_id = global_id - expert_start
-//       sum += topk_weights[token, k] * expert_out[local_id, token, dim]
-//   out[token, dim] = bf16(sum)
-// ============================================================
-
-__global__ void moe_weighted_sum_bf16_kernel(
-    const __nv_bfloat16* __restrict__ expert_out,
-    const int* __restrict__ topk_ids,
-    const float* __restrict__ topk_weights,
-    __nv_bfloat16* __restrict__ out,
-    int num_tokens, int hidden, int top_k,
-    int expert_start, int local_experts
-) {
-    // 64-bit index: `local_id * expert_stride` overflows int32 for long prefills
-    // (expert_stride = num_tokens * hidden), reading the wrong expert element.
-    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
-    long long total = (long long)num_tokens * hidden;
-    if (idx >= total) return;
-
-    long long token = idx / hidden;
-    int dim = (int)(idx % hidden);
-
-    long long expert_stride = (long long)num_tokens * hidden; // stride between experts in expert_out
-
-    float sum = 0.0f;
-    for (int k = 0; k < top_k; k++) {
-        int global_id = topk_ids[token * top_k + k];
-        int local_id = global_id - expert_start;
-        if (local_id >= 0 && local_id < local_experts) {
-            float w = topk_weights[token * top_k + k];
-            float v = __bfloat162float(expert_out[local_id * expert_stride + token * hidden + dim]);
-            sum += w * v;
-        }
-    }
-    out[idx] = __float2bfloat16(sum);
-}
-
-
-extern "C" {
-
-void launch_moe_topk_softmax_bf16(
-    const void* router_logits,
-    void* topk_ids, void* topk_weights,
-    int num_tokens, int num_experts, int top_k,
-    void* stream
-) {
-    int block = 128;
-    moe_topk_softmax_bf16_kernel<<<num_tokens, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)router_logits,
-        (int*)topk_ids, (float*)topk_weights,
-        num_experts, top_k
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_moe_replicate_bf16(
-    const void* x, void* x_rep,
-    int num_tokens, int hidden, int local_experts,
-    void* stream
-) {
-    long long total = (long long)local_experts * num_tokens * hidden;
-    int block = 256;
-    int grid = (int)((total + block - 1) / block);
-    moe_replicate_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)x, (__nv_bfloat16*)x_rep,
-        num_tokens, hidden, local_experts
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_moe_bias_add_3d_bf16(
-    void* x, const void* bias,
-    int batch, int num_tokens, int dim,
-    void* stream
-) {
-    long long total = (long long)batch * num_tokens * dim;
-    int block = 256;
-    int grid = (int)((total + block - 1) / block);
-    moe_bias_add_3d_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (__nv_bfloat16*)x, (const __nv_bfloat16*)bias,
-        batch, num_tokens, dim
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_moe_weighted_sum_bf16(
-    const void* expert_out,
-    const void* topk_ids, const void* topk_weights,
-    void* out,
-    int num_tokens, int hidden, int top_k,
-    int expert_start, int local_experts,
-    void* stream
-) {
-    long long total = (long long)num_tokens * hidden;
-    int block = 256;
-    int grid = (int)((total + block - 1) / block);
-    moe_weighted_sum_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)expert_out,
-        (const int*)topk_ids, (const float*)topk_weights,
-        (__nv_bfloat16*)out,
-        num_tokens, hidden, top_k,
-        expert_start, local_experts
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-}

csrc/moe/moe_sparse.cu

@@ -1,254 +0,0 @@
-#include <cuda_bf16.h>
-#include <cuda_fp8.h>
-#include <cstdint>
-#include "../common.cuh"
-
-// ============================================================
-// Sparse MoE decode GEMVs — compute ONLY the routed experts.
-//
-// The dense path replicates x across all local experts and runs a
-// batched GEMM, reading every expert's weights per token. Decode is
-// memory-bound, so reading only the top-k routed experts' weights
-// (~2 of 16 local on average at TP=2) is a ~8x byte reduction.
-//
-// Each block handles one (token, slot) pair's tile of output columns.
-// It reads topk_ids[token, slot] from device memory (no host sync),
-// and exits early if the expert is not owned by this rank. The early
-// return is BLOCK-UNIFORM (every thread sees the same topk_ids value
-// and returns before the shared-memory staging + __syncthreads), so
-// it is safe — unlike the divergent-return bug fixed in gemv.cu.
-//
-// Outputs for non-local slots are NEVER written (uninitialized memory,
-// possibly NaN bit patterns). Downstream consumers must SKIP non-local
-// slots rather than multiply by zero (NaN * 0 = NaN).
-//
-// Per-expert weight scale and bias are fused into the epilogue:
-//   y[t, slot, n] = acc * w_scale[lid] + bias[lid, n]
-// which matches the dense path's GEMM -> moe_bias_add_3d sequence.
-//
-// Activation addressing (x_per_slot):
-//   gate_up: all slots of a token share x[token, :]        (x_per_slot=0)
-//   down:    each slot has its own activation row
-//            x[token * top_k + slot, :]                    (x_per_slot=1)
-// ============================================================
-
-#define SPARSE_TILE_N 8  // output columns per block (= warps per block)
-
-// Weights FP8 E4M3 [local_experts, N, K], activations BF16 (W8A16).
-// Decode is memory-bound (~2 FLOP/byte), so dequant-in-registers GEMV
-// loses nothing to tensor cores and skips activation quantization.
-__global__ void moe_sparse_gemv_fp8_bf16_kernel(
-    const __nv_bfloat16* __restrict__ x,     // [T, K] or [T*top_k, K]
-    const __nv_fp8_e4m3* __restrict__ w,     // [local_experts, N, K]
-    const float* __restrict__ w_scales,      // [local_experts]
-    const __nv_bfloat16* __restrict__ bias,  // [local_experts, N]
-    const int* __restrict__ topk_ids,        // [T, top_k] global expert ids
-    __nv_bfloat16* __restrict__ y,           // [T, top_k, N]
-    int N, int K, int top_k,
-    int expert_start, int local_experts,
-    int x_per_slot
-) {
-    int token = blockIdx.z;
-    int slot = blockIdx.y;
-    int eid = topk_ids[token * top_k + slot];
-    int lid = eid - expert_start;
-    if (lid < 0 || lid >= local_experts) return;  // block-uniform: safe
-
-    extern __shared__ float xs[];  // [K] activation row as float
-    const __nv_bfloat16* xrow =
-        x + (long long)(x_per_slot ? token * top_k + slot : token) * K;
-    for (int i = threadIdx.x; i < K; i += blockDim.x) {
-        xs[i] = __bfloat162float(xrow[i]);
-    }
-    __syncthreads();
-
-    int n = blockIdx.x * SPARSE_TILE_N + (threadIdx.x >> 5);
-    if (n >= N) return;  // after __syncthreads: safe
-    int lane = threadIdx.x & 31;
-
-    // One warp per output column; uint4 = 16 FP8 weights per lane, the
-    // warp covers 512 contiguous bytes per iteration (coalesced).
-    const uint8_t* wrow = (const uint8_t*)w + ((long long)lid * N + n) * K;
-    float acc = 0.0f;
-    for (int i = lane; i < (K >> 4); i += 32) {
-        uint4 packed = *(const uint4*)(wrow + (long long)i * 16);
-        const __nv_fp8_e4m3* pw = (const __nv_fp8_e4m3*)&packed;
-        const float* xk = xs + i * 16;
-        #pragma unroll
-        for (int j = 0; j < 16; j++) {
-            acc += xk[j] * float(pw[j]);
-        }
-    }
-
-    #pragma unroll
-    for (int o = 16; o > 0; o >>= 1) {
-        acc += __shfl_down_sync(0xffffffffu, acc, o);
-    }
-    if (lane == 0) {
-        float v = acc * w_scales[lid]
-                + __bfloat162float(bias[(long long)lid * N + n]);
-        y[((long long)token * top_k + slot) * N + n] = __float2bfloat16(v);
-    }
-}
-
-// MXFP4 W4A16 variant: packed E2M1 nibbles + per-32 UE8M0 block scale,
-// same structure as batched_gemv_mxfp4_bf16_kernel but expert-indexed
-// via topk_ids and with fused per-expert bias.
-#define MXFP4_BLOCK 32
-
-__device__ __constant__ float kSparseFp4Levels[8] =
-    {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f};
-
-__device__ __forceinline__ float sparse_fp4_to_float(uint8_t code) {
-    float mag = kSparseFp4Levels[code & 0x7];
-    return (code & 0x8) ? -mag : mag;
-}
-
-__global__ void moe_sparse_gemv_mxfp4_bf16_kernel(
-    const __nv_bfloat16* __restrict__ x,     // [T, K] or [T*top_k, K]
-    const uint8_t* __restrict__ w_packed,    // [local_experts, N, K/2]
-    const uint8_t* __restrict__ w_scales,    // [local_experts, N, K/32]
-    const __nv_bfloat16* __restrict__ bias,  // [local_experts, N]
-    const int* __restrict__ topk_ids,        // [T, top_k]
-    __nv_bfloat16* __restrict__ y,           // [T, top_k, N]
-    int N, int K, int top_k,
-    int expert_start, int local_experts,
-    int x_per_slot
-) {
-    int token = blockIdx.z;
-    int slot = blockIdx.y;
-    int eid = topk_ids[token * top_k + slot];
-    int lid = eid - expert_start;
-    if (lid < 0 || lid >= local_experts) return;  // block-uniform: safe
-
-    extern __shared__ float xs[];
-    const __nv_bfloat16* xrow =
-        x + (long long)(x_per_slot ? token * top_k + slot : token) * K;
-    for (int i = threadIdx.x; i < K; i += blockDim.x) {
-        xs[i] = __bfloat162float(xrow[i]);
-    }
-    __syncthreads();
-
-    int n = blockIdx.x * SPARSE_TILE_N + (threadIdx.x >> 5);
-    if (n >= N) return;
-    int lane = threadIdx.x & 31;
-    int nblk = K / MXFP4_BLOCK;
-
-    const uint8_t* wp = w_packed + ((long long)lid * N + n) * (K >> 1);
-    const uint8_t* ws = w_scales + ((long long)lid * N + n) * nblk;
-
-    float acc = 0.0f;
-    for (int blk = lane; blk < nblk; blk += 32) {
-        float scale = exp2f((float)((int)ws[blk] - 127));
-        uint4 packed = *(const uint4*)(wp + (long long)blk * 16);  // 32 nibbles
-        const uint8_t* pb = (const uint8_t*)&packed;
-        const float* xk = xs + blk * MXFP4_BLOCK;
-        #pragma unroll
-        for (int i = 0; i < 16; i++) {
-            uint8_t b = pb[i];
-            acc += xk[2 * i]     * (sparse_fp4_to_float(b & 0xF) * scale);
-            acc += xk[2 * i + 1] * (sparse_fp4_to_float(b >> 4)  * scale);
-        }
-    }
-
-    #pragma unroll
-    for (int o = 16; o > 0; o >>= 1) {
-        acc += __shfl_down_sync(0xffffffffu, acc, o);
-    }
-    if (lane == 0) {
-        float v = acc + __bfloat162float(bias[(long long)lid * N + n]);
-        y[((long long)token * top_k + slot) * N + n] = __float2bfloat16(v);
-    }
-}
-
-// Weighted sum over the slot axis: out[t, d] = sum over local slots of
-// topk_weights[t, k] * down[t, k, d]. Non-local slots hold uninitialized
-// memory and are SKIPPED (not multiplied by zero).
-__global__ void moe_weighted_sum_sparse_bf16_kernel(
-    const __nv_bfloat16* __restrict__ down,  // [T, top_k, hidden]
-    const int* __restrict__ topk_ids,        // [T, top_k]
-    const float* __restrict__ topk_weights,  // [T, top_k]
-    __nv_bfloat16* __restrict__ out,         // [T, hidden]
-    int num_tokens, int hidden, int top_k,
-    int expert_start, int local_experts
-) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-    int total = num_tokens * hidden;
-    if (idx >= total) return;
-
-    int token = idx / hidden;
-    int dim = idx % hidden;
-
-    float sum = 0.0f;
-    for (int k = 0; k < top_k; k++) {
-        int lid = topk_ids[token * top_k + k] - expert_start;
-        if (lid >= 0 && lid < local_experts) {
-            float w = topk_weights[token * top_k + k];
-            float v = __bfloat162float(
-                down[((long long)token * top_k + k) * hidden + dim]);
-            sum += w * v;
-        }
-    }
-    out[idx] = __float2bfloat16(sum);
-}
-
-extern "C" {
-
-void launch_moe_sparse_gemv_fp8_bf16(
-    const void* x, const void* w, const void* w_scales, const void* bias,
-    const void* topk_ids, void* y,
-    int num_tokens, int N, int K, int top_k,
-    int expert_start, int local_experts, int x_per_slot,
-    void* stream
-) {
-    dim3 grid((N + SPARSE_TILE_N - 1) / SPARSE_TILE_N, top_k, num_tokens);
-    int block = SPARSE_TILE_N * 32;
-    size_t smem = (size_t)K * sizeof(float);
-    moe_sparse_gemv_fp8_bf16_kernel<<<grid, block, smem, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)x, (const __nv_fp8_e4m3*)w,
-        (const float*)w_scales, (const __nv_bfloat16*)bias,
-        (const int*)topk_ids, (__nv_bfloat16*)y,
-        N, K, top_k, expert_start, local_experts, x_per_slot
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_moe_sparse_gemv_mxfp4_bf16(
-    const void* x, const void* w_packed, const void* w_scales, const void* bias,
-    const void* topk_ids, void* y,
-    int num_tokens, int N, int K, int top_k,
-    int expert_start, int local_experts, int x_per_slot,
-    void* stream
-) {
-    dim3 grid((N + SPARSE_TILE_N - 1) / SPARSE_TILE_N, top_k, num_tokens);
-    int block = SPARSE_TILE_N * 32;
-    size_t smem = (size_t)K * sizeof(float);
-    moe_sparse_gemv_mxfp4_bf16_kernel<<<grid, block, smem, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)x, (const uint8_t*)w_packed,
-        (const uint8_t*)w_scales, (const __nv_bfloat16*)bias,
-        (const int*)topk_ids, (__nv_bfloat16*)y,
-        N, K, top_k, expert_start, local_experts, x_per_slot
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_moe_weighted_sum_sparse_bf16(
-    const void* down, const void* topk_ids, const void* topk_weights,
-    void* out,
-    int num_tokens, int hidden, int top_k,
-    int expert_start, int local_experts,
-    void* stream
-) {
-    int total = num_tokens * hidden;
-    int block = 256;
-    int grid = (total + block - 1) / block;
-    moe_weighted_sum_sparse_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)down,
-        (const int*)topk_ids, (const float*)topk_weights,
-        (__nv_bfloat16*)out,
-        num_tokens, hidden, top_k, expert_start, local_experts
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-}

csrc/quantization/dequant_fp8.cu

@@ -1,53 +0,0 @@
-#include <cuda_bf16.h>
-#include <cuda_fp8.h>
-#include "../common.cuh"
-
-// Dequantize FP8 E4M3 → BF16 with per-expert (per-batch-slice) FP32 scale.
-//
-// Input:  src    [num_experts, rows, cols]  FP8 E4M3 (1 byte each)
-//         scales [num_experts]              FP32
-// Output: dst    [num_experts, rows, cols]  BF16
-//
-// Each element: dst[e, r, c] = bf16( float(src[e, r, c]) * scales[e] )
-
-__global__ void dequant_fp8e4m3_to_bf16_kernel(
-    const __nv_fp8_e4m3* __restrict__ src,
-    const float* __restrict__ scales,
-    __nv_bfloat16* __restrict__ dst,
-    int num_experts, int rows, int cols
-) {
-    // 64-bit index: num_experts * rows * cols overflows int32 for 32 experts
-    // at ~8k*8k weight matrices, same class as the MoE fix in cfbd64d.
-    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
-    long long total = (long long)num_experts * rows * cols;
-    if (idx >= total) return;
-
-    long long expert_stride = (long long)rows * cols;
-    int expert = (int)(idx / expert_stride);
-    float scale = scales[expert];
-    float val = float(src[idx]) * scale;
-    dst[idx] = __float2bfloat16(val);
-}
-
-extern "C" {
-
-void launch_dequant_fp8e4m3_to_bf16(
-    const void* src,
-    const void* scales,
-    void* dst,
-    int num_experts, int rows, int cols,
-    void* stream
-) {
-    long long total = (long long)num_experts * rows * cols;
-    int block = 256;
-    int grid = (int)((total + block - 1) / block);
-    dequant_fp8e4m3_to_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_fp8_e4m3*)src,
-        (const float*)scales,
-        (__nv_bfloat16*)dst,
-        num_experts, rows, cols
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-}

csrc/quantization/mxfp4_gemm.cu

@@ -1,135 +0,0 @@
-#include <cuda_bf16.h>
-#include <cstdint>
-#include "../common.cuh"
-
-// MXFP4 W4A16 for MoE experts. Weights stored [E, N, K] with K (reduction)
-// contiguous, blocked by 32: packed 4-bit E2M1 (two nibbles/byte, lo = even k)
-// + one UE8M0 scale byte per 32 elements. The decode win is reading 4-bit
-// weights from HBM (half of FP8) and dequantizing on-chip to BF16.
-
-#define MXFP4_BLOCK 32
-
-// E2M1 magnitude by 3-bit code; bit 3 is the sign.
-__device__ __constant__ float kFp4Levels[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f};
-
-__device__ __forceinline__ float fp4_to_float(uint8_t code) {
-    float mag = kFp4Levels[code & 0x7];
-    return (code & 0x8) ? -mag : mag;
-}
-
-// Decode (M=1) fused GEMV, batched over experts.
-//   y[e, n] = sum_k x[e, k] * dequant(W[e, n, k])
-// Grid: (N/TILE_N, E). Each block loads the activation x[e, :] into shared
-// memory ONCE and computes TILE_N output columns from it (one warp per column),
-// so the activation is read from HBM once per TILE_N outputs instead of once
-// per output. Weights are unique per output and read coalesced as uint4; the
-// UE8M0 block scale is hoisted to once per 32-element block.
-#define MXFP4_TILE_N 8   // output columns per block (= warps per block)
-
-__global__ void batched_gemv_mxfp4_bf16_kernel(
-    const __nv_bfloat16* __restrict__ x,   // [E, K]
-    const uint8_t* __restrict__ w_packed,  // [E, N, K/2]
-    const uint8_t* __restrict__ w_scales,  // [E, N, K/32]
-    __nv_bfloat16* __restrict__ y,         // [E, N]
-    int E, int N, int K
-) {
-    extern __shared__ float xs[];          // [K] activation for this expert
-    int e = blockIdx.y;
-    int n_base = blockIdx.x * MXFP4_TILE_N;
-    int warp = threadIdx.x >> 5;           // 0..TILE_N-1
-    int lane = threadIdx.x & 31;
-    int nthreads = blockDim.x;             // TILE_N * 32
-    int nblk = K / MXFP4_BLOCK;
-
-    // Cooperatively stage x[e, :] into shared memory (converted to float).
-    const __nv_bfloat16* xe = x + (long long)e * K;
-    for (int k = threadIdx.x; k < K; k += nthreads) {
-        xs[k] = __bfloat162float(xe[k]);
-    }
-    __syncthreads();
-
-    int n = n_base + warp;
-    if (n >= N) return;
-
-    const uint8_t* wp = w_packed + ((long long)e * N + n) * (K >> 1);
-    const uint8_t* ws = w_scales + ((long long)e * N + n) * nblk;
-
-    float acc = 0.0f;
-    for (int blk = lane; blk < nblk; blk += 32) {
-        float scale = exp2f((float)((int)ws[blk] - 127));
-        uint4 packed = *(const uint4*)(wp + (long long)blk * 16);  // 16 bytes = 32 nibbles
-        const uint8_t* pb = (const uint8_t*)&packed;
-        const float* xk = xs + blk * MXFP4_BLOCK;
-        #pragma unroll
-        for (int i = 0; i < 16; i++) {
-            uint8_t b = pb[i];
-            acc += xk[2 * i]     * (fp4_to_float(b & 0xF) * scale);
-            acc += xk[2 * i + 1] * (fp4_to_float(b >> 4)  * scale);
-        }
-    }
-
-    // Warp reduction.
-    #pragma unroll
-    for (int o = 16; o > 0; o >>= 1) {
-        acc += __shfl_down_sync(0xffffffffu, acc, o);
-    }
-    if (lane == 0) y[(long long)e * N + n] = __float2bfloat16(acc);
-}
-
-// Prefill fallback: dequant MXFP4 [E, N, K] -> BF16 [E, K, N] (transposed back
-// to the [E, K, N] layout the BF16 batched GEMM expects). Not bandwidth-optimal,
-// but prefill is compute-bound so it is not the decode hot path.
-__global__ void dequant_mxfp4_to_bf16_t_kernel(
-    const uint8_t* __restrict__ w_packed,  // [E, N, K/2]
-    const uint8_t* __restrict__ w_scales,  // [E, N, K/32]
-    __nv_bfloat16* __restrict__ out,       // [E, K, N]
-    int E, int N, int K
-) {
-    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
-    long long total = (long long)E * N * K;
-    if (idx >= total) return;
-    int k = idx % K;
-    int n = (idx / K) % N;
-    int e = idx / ((long long)N * K);
-
-    int Kh = K >> 1;
-    int Ks = K / MXFP4_BLOCK;
-    uint8_t byte = w_packed[((long long)e * N + n) * Kh + (k >> 1)];
-    uint8_t code = (k & 1) ? (byte >> 4) : (byte & 0xF);
-    float scale = exp2f((float)((int)w_scales[((long long)e * N + n) * Ks + k / MXFP4_BLOCK] - 127));
-    float val = fp4_to_float(code) * scale;
-    // write to out[e, k, n]
-    out[((long long)e * K + k) * N + n] = __float2bfloat16(val);
-}
-
-extern "C" {
-
-void launch_batched_gemv_mxfp4_bf16(
-    const void* x, const void* w_packed, const void* w_scales, void* y,
-    int E, int N, int K, void* stream
-) {
-    dim3 grid((N + MXFP4_TILE_N - 1) / MXFP4_TILE_N, E);
-    int block = MXFP4_TILE_N * 32;  // one warp per output column
-    size_t smem = (size_t)K * sizeof(float);
-    batched_gemv_mxfp4_bf16_kernel<<<grid, block, smem, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)x, (const uint8_t*)w_packed, (const uint8_t*)w_scales,
-        (__nv_bfloat16*)y, E, N, K
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_dequant_mxfp4_to_bf16_t(
-    const void* w_packed, const void* w_scales, void* out,
-    int E, int N, int K, void* stream
-) {
-    long long total = (long long)E * N * K;
-    int block = 256;
-    long long grid = (total + block - 1) / block;
-    dequant_mxfp4_to_bf16_t_kernel<<<(unsigned)grid, block, 0, (cudaStream_t)stream>>>(
-        (const uint8_t*)w_packed, (const uint8_t*)w_scales, (__nv_bfloat16*)out,
-        E, N, K
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-}

csrc/quantization/quantize_fp8.cu

@@ -1,160 +0,0 @@
-#include <cuda_bf16.h>
-#include <cuda_fp8.h>
-#include <float.h>
-#include "../common.cuh"
-
-// Per-row quantize BF16 → FP8 E4M3 with per-row FP32 scale output.
-//
-// Input:  src    [num_rows, cols]  BF16
-// Output: dst    [num_rows, cols]  FP8 E4M3
-//         scales [num_rows]        FP32
-//
-// Algorithm per row:
-//   absmax = max(|src[row, :]|)
-//   scale  = absmax / 448.0       (FP8 E4M3 max representable)
-//   dst[row, i] = fp8(src[row, i] / scale)
-//
-// Grid: one block per row.  Block: 256 threads.
-// Each thread handles ceil(cols / 256) elements.
-
-#define QUANT_BLOCK 256
-#define FP8_E4M3_MAX 448.0f
-
-__global__ void quantize_bf16_to_fp8e4m3_rowwise_kernel(
-    const __nv_bfloat16* __restrict__ src,
-    __nv_fp8_e4m3* __restrict__ dst,
-    float* __restrict__ scales,
-    int num_rows, int cols
-) {
-    int row = blockIdx.x;
-    if (row >= num_rows) return;
-    int tid = threadIdx.x;
-
-    const __nv_bfloat16* row_src = src + (long long)row * cols;
-    __nv_fp8_e4m3* row_dst = dst + (long long)row * cols;
-
-    // Step 1: Compute per-row absmax via shared-memory reduction.
-    __shared__ float smem_max[QUANT_BLOCK];
-    float local_max = 0.0f;
-    for (int i = tid; i < cols; i += QUANT_BLOCK) {
-        float v = fabsf(__bfloat162float(row_src[i]));
-        local_max = fmaxf(local_max, v);
-    }
-    smem_max[tid] = local_max;
-    __syncthreads();
-
-    // Block reduction
-    for (int s = QUANT_BLOCK / 2; s > 0; s >>= 1) {
-        if (tid < s) {
-            smem_max[tid] = fmaxf(smem_max[tid], smem_max[tid + s]);
-        }
-        __syncthreads();
-    }
-
-    float absmax = smem_max[0];
-    float scale = absmax / FP8_E4M3_MAX;
-    // Clamp scale to avoid div-by-zero for all-zero rows
-    if (scale < 1e-12f) scale = 1e-12f;
-    float inv_scale = 1.0f / scale;
-
-    // Thread 0 writes the scale
-    if (tid == 0) {
-        scales[row] = scale;
-    }
-
-    // Step 2: Quantize each element
-    for (int i = tid; i < cols; i += QUANT_BLOCK) {
-        float v = __bfloat162float(row_src[i]) * inv_scale;
-        row_dst[i] = __nv_fp8_e4m3(v);
-    }
-}
-
-// Row-wise scale: data[row, :] *= scales[row]  (in-place, BF16)
-__global__ void rowwise_scale_bf16_kernel(
-    __nv_bfloat16* __restrict__ data,
-    const float* __restrict__ scales,
-    int num_rows, int cols
-) {
-    int row = blockIdx.x;
-    if (row >= num_rows) return;
-    int tid = threadIdx.x;
-    float s = scales[row];
-    __nv_bfloat16* row_data = data + (long long)row * cols;
-    for (int i = tid; i < cols; i += blockDim.x) {
-        float v = __bfloat162float(row_data[i]) * s;
-        row_data[i] = __float2bfloat16(v);
-    }
-}
-
-// Combined dequant scale for batched MoE FP8 GEMM output.
-//   data[row, :] *= a_scales[row] * b_scales[row / tokens]
-// where row = expert * tokens + token. a_scales is the per-token activation
-// scale; b_scales is the per-expert scalar weight scale. Lets a single
-// strided-batched FP8 matmul (alpha=1, scales=1) recover the real result in
-// one pass instead of folding the weight scale into a per-expert GEMM call.
-__global__ void rowwise_scale_moe_bf16_kernel(
-    __nv_bfloat16* __restrict__ data,
-    const float* __restrict__ a_scales,
-    const float* __restrict__ b_scales,
-    int num_rows, int cols, int tokens
-) {
-    int row = blockIdx.x;
-    if (row >= num_rows) return;
-    int tid = threadIdx.x;
-    float s = a_scales[row] * b_scales[row / tokens];
-    __nv_bfloat16* row_data = data + (long long)row * cols;
-    for (int i = tid; i < cols; i += blockDim.x) {
-        float v = __bfloat162float(row_data[i]) * s;
-        row_data[i] = __float2bfloat16(v);
-    }
-}
-
-extern "C" {
-
-void launch_rowwise_scale_bf16(
-    void* data, const void* scales,
-    int num_rows, int cols,
-    void* stream
-) {
-    int block = 256;
-    int grid = num_rows;
-    rowwise_scale_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (__nv_bfloat16*)data, (const float*)scales,
-        num_rows, cols
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_rowwise_scale_moe_bf16(
-    void* data, const void* a_scales, const void* b_scales,
-    int num_rows, int cols, int tokens,
-    void* stream
-) {
-    int block = 256;
-    int grid = num_rows;
-    rowwise_scale_moe_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (__nv_bfloat16*)data, (const float*)a_scales, (const float*)b_scales,
-        num_rows, cols, tokens
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-void launch_quantize_bf16_to_fp8e4m3_rowwise(
-    const void* src,
-    void* dst,
-    void* scales,
-    int num_rows, int cols,
-    void* stream
-) {
-    int grid = num_rows;
-    int block = QUANT_BLOCK;
-    quantize_bf16_to_fp8e4m3_rowwise_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)src,
-        (__nv_fp8_e4m3*)dst,
-        (float*)scales,
-        num_rows, cols
-    );
-    CUDA_CHECK_LAST_ERROR();
-}
-
-}

csrc/reduce/argmax.cu

@@ -1,92 +0,0 @@
-#include <cuda_bf16.h>
-#include <float.h>
-#include "../common.cuh"
-
-// Argmax along the last dim of a [rows, cols] tensor.
-// One block per row; output is [rows] int32 indices of the max element.
-//
-// Reduction: each thread scans a strided slice and tracks the running
-// (value, index) pair, then warp-shuffle reduce, then a single-warp
-// reduce over per-warp leaders. Tie-break: smaller index wins so the
-// result is deterministic across launches.
-//
-// For BF16 logits the comparison happens in FP32 to avoid losing
-// precision near the top of the distribution.
-
-__global__ void argmax_bf16_kernel(
-    const __nv_bfloat16* __restrict__ logits,
-    int* __restrict__ out_idx,
-    int cols
-) {
-    int row = blockIdx.x;
-    const __nv_bfloat16* row_ptr = logits + (long long)row * cols;
-    int tid = threadIdx.x;
-    unsigned mask = 0xffffffff;
-
-    // Strided per-thread max.
-    float local_max = -FLT_MAX;
-    int   local_idx = INT_MAX;
-    for (int i = tid; i < cols; i += blockDim.x) {
-        float v = __bfloat162float(row_ptr[i]);
-        // strict `>` keeps the smallest index on ties, since we scan ascending.
-        if (v > local_max) {
-            local_max = v;
-            local_idx = i;
-        }
-    }
-
-    // Warp-level reduce of (val, idx) pairs.
-    #pragma unroll
-    for (int offset = 16; offset > 0; offset >>= 1) {
-        float other_val = __shfl_down_sync(mask, local_max, offset);
-        int   other_idx = __shfl_down_sync(mask, local_idx, offset);
-        bool take = (other_val > local_max) ||
-                    (other_val == local_max && other_idx < local_idx);
-        if (take) {
-            local_max = other_val;
-            local_idx = other_idx;
-        }
-    }
-
-    // Per-warp leaders → shared memory → single warp final reduce.
-    __shared__ float s_val[32];
-    __shared__ int   s_idx[32];
-    int lane     = tid & 31;
-    int warp_id  = tid >> 5;
-    int num_warps = (blockDim.x + 31) >> 5;
-
-    if (lane == 0) {
-        s_val[warp_id] = local_max;
-        s_idx[warp_id] = local_idx;
-    }
-    __syncthreads();
-
-    if (warp_id == 0) {
-        float v = (tid < num_warps) ? s_val[lane] : -FLT_MAX;
-        int   i = (tid < num_warps) ? s_idx[lane] : INT_MAX;
-        #pragma unroll
-        for (int offset = 16; offset > 0; offset >>= 1) {
-            float ov = __shfl_down_sync(mask, v, offset);
-            int   oi = __shfl_down_sync(mask, i, offset);
-            bool take = (ov > v) || (ov == v && oi < i);
-            if (take) { v = ov; i = oi; }
-        }
-        if (lane == 0) {
-            out_idx[row] = i;
-        }
-    }
-}
-
-extern "C" {
-
-void launch_argmax_bf16(const void* logits, void* out_idx,
-                        int rows, int cols, void* stream) {
-    // 1024 threads/block keeps occupancy high and gives 32 warps for the
-    // final reduce (matches the 32-slot shared arrays above).
-    int block = 1024;
-    argmax_bf16_kernel<<<rows, block, 0, (cudaStream_t)stream>>>(
-        (const __nv_bfloat16*)logits, (int*)out_idx, cols);
-    CUDA_CHECK_LAST_ERROR();
-}
-
-}

csrc/reduce/softmax.cu

@@ -90,7 +90,7 @@ __global__ void softmax_bf16(
 extern "C" {
 
 void launch_softmax_f32(const void* x, void* out, int rows, int cols, void* stream) {
-    int block = (cols < 512) ? cols : 512;
+    int block = (cols < 1024) ? cols : 1024;
     if (block < 32) block = 32;
     softmax_f32<<<rows, block, 0, (cudaStream_t)stream>>>(
         (const float*)x, (float*)out, cols);
@@ -98,7 +98,7 @@ void launch_softmax_f32(const void* x, void* out, int rows, int cols, void* stre
 }
 
 void launch_softmax_bf16(const void* x, void* out, int rows, int cols, void* stream) {
-    int block = (cols < 512) ? cols : 512;
+    int block = (cols < 1024) ? cols : 1024;
     if (block < 32) block = 32;
     softmax_bf16<<<rows, block, 0, (cudaStream_t)stream>>>(
         (const __nv_bfloat16*)x, (__nv_bfloat16*)out, cols);

docs/00-roadmap.md

@@ -1748,27 +1748,6 @@ Text  → Tokenizer   → Text Tokens   ────────────→
 
 ---
 
-## 实际进展记录(与原计划的分叉,2026-06 更新)
-
-Phase 0–17 按计划完成。Phase 18 起实际路线偏离了上面的原计划
-(speculative decoding 与多模态推迟),实际走向是 MoE + 量化 + 稀疏化:
-
-| 实际 Phase | 内容 | 文档 |
-|---|---|---|
-| 18 | Pipeline Parallelism(PP=2/4) | `18-pipeline-parallelism.md`、`benchmarks/pp-sweep.md` |
-| 19 | **gpt-oss-20b MoE**:harmony 格式、attention sinks + sliding window、YaRN;两个 CUDA bug 实战(prefill sinks NaN、GEMV 未初始化 smem);GSM8K 94.5% 对齐 llama.cpp;FP8 W8A8 / MXFP4 W4A16 量化 | `19-gpt-oss-moe.md`、`benchmarks/{fp8-quantization,mxfp4-and-llama-decode}.md` |
-| 20 | **稀疏 top-k MoE decode**:只算被路由的专家,decode 13.9→7.0ms,TP=2 下 decode/TTFT 全面快于 llama.cpp 同配置;gpt-oss 单卡 serving | `20-sparse-moe.md`、`benchmarks/sparse-moe.md` |
-| 21 | **decode CUDA Graph + GPU argmax**:整个 decode step 录成一个图回放(thread-local launch stream、retained-warmup 分配策略、NCCL capture);greedy 采样换 GPU argmax。TPOT 7.5→5.9ms(TP=1)/ 5.8ms(TP=2);TP=2 全面领先 llama(1.26-1.47×),TP=1 差距 2.5×→2.0× | `21-cuda-graph-decode.md` |
-
-**下一步候选(按预期收益排序):**
-
-| 候选 Phase | 内容 | 预期 |
-|---|---|---|
-| 22 | **非专家权重量化**:qkv/o + lm_head(1.16GB/token)仍是 BF16 | TPOT 再省 ~1.5ms |
-| 23 | **稀疏 prefill**(按专家 permute + grouped GEMM) | 长 prompt TTFT 51-75 → ~30ms |
-| 24 | server 侧 harmony channel 分离(`reasoning_content` 流式输出,对齐 llama-server 行为) | API 易用性 |
-| — | Speculative Decoding、多模态(原 16/19) | 推迟 |
-
 ## 里程碑总结
 
 | 里程碑 | Phase | 验收标准 |
@@ -1778,9 +1757,7 @@ Phase 0–17 按计划完成。Phase 18 起实际路线偏离了上面的原计
 | ③ E2E API | 13 | HTTP streaming API, Python OpenAI SDK 可调用, 10 并发正确 |
 | ④ 性能达标 | 15 | throughput >= 50% vLLM, profiling 报告完成 |
 | ⑤ 多卡推理 | 17 | TP=2/4 同组 GPU 推理正确, scaling benchmark 完成 |
-| ⑥ MoE 模型(实际) | 19 | gpt-oss-20b 端到端正确, GSM8K 与 llama.cpp 持平 ✅ |
-| ⑦ 性能反超(实际) | 20 | 同配置 decode 快于 llama.cpp(TP=2 达成;单卡是 Phase 21+ 目标) ✅ |
-| ⑧ 多模态 | 推迟 | 图片输入 → 文字回答, API 端到端 |
+| ⑥ 多模态 | 19 | 图片输入 → 文字回答, API 端到端 |
 
 ## 外部依赖清单
 

docs/19-gpt-oss-moe.md

@@ -1,118 +0,0 @@
-# Phase 19: gpt-oss-20b — MoE 模型支持与两次 CUDA 调试实战
-
-> 目标:支持 OpenAI gpt-oss-20b(32 专家 top-4 MoE),GSM8K 精度对齐 llama.cpp,
-> 并以此为载体做 FP8 / MXFP4 量化。本文档事后整理,重点放在**踩过的坑**:
-> 两个教科书级的 CUDA bug 排查过程比结论本身更有学习价值。
->
-> 后续:`docs/20-sparse-moe.md`(稀疏化),benchmark 数据见
-> `docs/benchmarks/{fp8-quantization,mxfp4-and-llama-decode,sparse-moe}.md`。
-
-## 1. 模型架构(与 Qwen3 的差异点)
-
-gpt-oss-20b(`config.json`,已在 dash5 验证):
-
-| 项 | 值 | 说明 |
-|---|---|---|
-| layers / hidden | 24 / 2880 | hidden **不是** 128 的倍数的来源(2880 = 22.5×128) |
-| heads | 64 Q / 8 KV,head_dim **64** | head_dim ≠ hidden/heads(64×64=4096>2880),GQA n_rep=8 |
-| MoE | 32 experts,top-4,expert inter 2880 | router 是普通 Linear [2880→32] + bias |
-| attention | **交替 sliding(128)/full**,layer 0 是 sliding | 每层带 **attention sinks**(每 head 一个可学习标量) |
-| RoPE | YaRN(theta 150000, factor 32, orig 4096) | attn_factor = 0.1·ln(32)+1 乘在 cos/sin 上 |
-| 激活 | clamp 后的 GLU | gate=gu[::2], up=gu[1::2](**交错**), gate≤7, up∈[-7,7], glu=gate·σ(1.702·gate), h=(up+1)·glu |
-| 词表 | 201088 | EOS 是**列表** [200002,199999,200012] = `<|return|>`/`<|endoftext|>`/`<|call|>` |
-| 其它 | attention_bias=true | q/k/v/o 全部带 bias(Qwen3 没有) |
-
-**Harmony 对话格式**:gpt-oss 不是普通 chat template,输出分 channel
-(`analysis`=思维链,`final`=正式回答),控制 token `<|start|>/<|channel|>/<|message|>/<|end|>`。
-三个坑:(1) system 消息必须含 `Reasoning:` 等 canonical 行,缺了模型 OOD、
-channel 选择不稳定;(2) repetition penalty 会惩罚必须重复出现的控制 token,
-导致模型只输出 analysis 不出 final(MoE 默认关掉);(3) 服务端要用多 EOS 判停。
-
-## 2. MoE 前向(dense 版,Phase 20 之前)
-
-```text
-router GEMV → topk_softmax(GPU)→ moe_replicate(复制到全部本地专家)
-→ batched GEMM gate_up → bias → GLU → batched GEMM down → bias
-→ weighted_sum(只取 top-4)→ all-reduce
-```
-
-要点:top-k 的专家编号始终留在 GPU(`topk_ids`),host 不同步;
-dense 的代价(每 token 读全部专家权重)在 Phase 20 用 sparse GEMV 解决。
-TP 用 **expert parallelism**:rank r 拥有专家 [r·E/world, (r+1)·E/world),
-weighted_sum 里按 `expert_start + local_experts` 过滤非本地命中,
-all-reduce 把各 rank 的部分和加起来——这要求"跳过"语义而不是"乘 0"。
-
-## 3. CUDA 调试实战 ①:prefill NaN(flash-attention sinks)
-
-**症状**:长 prompt(≳192 token)prefill 后输出全 NaN → argmax 落在
-token 201087(`max_by` 平局取最后)或 token 0(`!`)。短 prompt 完全正常。
-
-**定位手法**:给每个 stage 加 NaN 检查(环境变量开关,事后移除),
-二分出第一个出 NaN 的位置:layer-0 的 `flash_attention_sinks` 输出,
-而它的 q/k 输入是干净的 → bug 在 kernel 内部。
-
-**根因**:causal 跳过逻辑只剔除"完全在未来"的 kv tile;一个完全滑出
-sliding window(128)的**过去** tile 仍被处理,所有 key 都被 mask 成 -inf
-→ `row_max = -inf` → online softmax 里 `expf(-inf-(-inf)) = NaN`,
-下一个有效 tile 的修正项 `0·NaN` 把整行毒掉。
-
-**修复**:`row_max == -INFINITY` 的 tile 直接跳过(贡献为零)。
-**教训**:online softmax 的"空 tile"是边界条件标配;decode kernel 早就
-防了这个(`local_max==-INFINITY` guard),prefill kernel 漏了——
-**同一逻辑的两份实现要做同样的边界测试**。触发阈值 ~192 token 解释了
-"短测试全过、长对话必炸"的诡异表象。
-
-## 4. CUDA 调试实战 ②:decode 间歇性乱码(GEMV 未初始化共享内存)
-
-**症状**:同一 prompt ~70% 的运行在第二轮对话或长生成中突然输出
-`!!!!`/token 201087/NaN logits,**间歇性** → 不是确定性逻辑错误,
-是竞态或未初始化读。只有 gpt-oss 出问题,Qwen3 从不复现。
-
-**定位**:逐 stage 检查,第一个出问题的是 decode 的 o_proj 输出
-(maxabs≈1e33),输入干净 → M=1 的 GEMV kernel。
-
-**根因**(`gemv.cu`):
-
-```cuda
-if (col >= N) return;        // ← 在协作加载 x_shared 和 __syncthreads 之前!
-...cooperative load + __syncthreads()...
-```
-
-当 `N % 128 != 0` 时,最后一个 block 的越界线程提前退出,**没参与**
-共享内存装载;在界线程读到未初始化的 smem(且 `__syncthreads` 在有线程
-已退出时是 UB)。命中条件:n=2880 的矩阵(o_proj、MoE gate_up/down)——
-2880 % 128 ≠ 0;而 Qwen3 所有维度都是 128 对齐的,**所以"只有 gpt-oss
-不稳定"**。q/k/v(4096)、lm_head(201088)对齐,幸免。
-
-**修复**:所有线程先完成装载 + barrier,`col >= N` 检查移到 syncthreads
-**之后**。
-
-**教训**:`__syncthreads()` 之前的任何 early-return 必须是 **block-uniform**
-的。Phase 20 的 sparse GEMV 专门遵守了这条(整个 block 基于同一个
-`topk_ids` 值统一退出,发生在装载之前)。
-
-**修复后的验证**:GSM8K 全量 1319 题,xserv 94.5% vs llama.cpp 94.4%
-——统计上同一水平,证明两个 kernel bug 就是之前 55% vs 95% 差距的全部原因。
-
-## 5. 量化(详见 benchmark 文档)
-
-- **FP8 W8A8**(`tools/quantize_fp8.py`):per-expert 标量 scale,权重转置
-  存 [E,N,K] 喂 cuBLASLt(Blackwell 要求 transA=T)。两个性能坑:
-  (1) 每次调用重建 plan + 跑 heuristic → 比 BF16 还慢,修复 = per-shape
-  plan cache;(2) 逐专家发射 ~768 个小 GEMM,修复 = 单条 strided-batched
-  调用 + 把 scale 移到融合的 post-scale kernel。最终 1.41× vs BF16。
-- **MXFP4 W4A16**(`tools/quantize_mxfp4.py`):E2M1 + per-32 UE8M0 块 scale,
-  13GB 模型,贪心输出与 BF16 逐字一致,但手写 dequant-GEMV 打不过
-  cuBLASLt FP8(带宽效率差),定位为省显存方案。
-- 检测方式:safetensors 的 dtype/scale 秩自动识别,loader 无需配置。
-
-## 6. 本阶段的工具沉淀
-
-- `bench-gpt-oss`:in-process 推理 + `--forced`(teacher-forced prefill
-  top-1)/`--forced-decode`(沿参考轨迹逐位置 top-1)——分离"前向算错"
-  和"贪心轨迹分叉"的利器。
-- `tools/eval_gsm8k_fast.py`(持久 xserv-chat 管道)、
-  `tools/xserv_vs_llama.py`(warm-server 同机对打,计入 llama 的
-  reasoning_content)。
-- 经验:**贪心解码不是逐位可复现的**(cuBLAS 非确定性会翻转后段 argmax),
-  多卡正确性要用"单卡×2 + 多卡×2 互相比",精度要用基准集而不是逐字 diff。

docs/19-moe-gpt-oss.md

@@ -0,0 +1,128 @@
+# Phase 19: MoE — gpt-oss-20b
+
+> 目标：在 xserv 支持 **MoE**，用 `openai/gpt-oss-20b` 端到端跑通，并与 llama.cpp 在
+> AIME 2025 / GSM8K 上对比正确性与性能。MXFP4 expert 权重加载时反量化为 BF16；整模型
+> ~40GB 单卡放不下 → 复用 Phase 18 的 **PP**（PP=2 ~20GB/卡，PP=4 ~10GB/卡）。
+>
+> 实时进度与重启续作指南见 `docs/MOE_PROGRESS.md`。
+
+## 1. 架构（config.json，已核对）
+
+num_hidden_layers=24, hidden=2880, **head_dim=64**（≠hidden/heads）, n_heads=64,
+n_kv_heads=8（GQA n_rep=8）, expert intermediate=2880, **num_local_experts=32**,
+**num_experts_per_tok=4**, vocab=201088, max_pos=131072, rope_theta=150000,
+sliding_window=128（交替层，见 `layer_types`）, rms_norm_eps=1e-5, swiglu_limit=7.0,
+alpha=1.702, tie_embeddings=false。
+量化：**MXFP4**，仅 expert MLP（gate_up/down 的 `_blocks`+`_scales`）；
+attn/router/embed/lm_head 为 BF16。
+
+## 2. 参考数学（HF transformers `modeling_gpt_oss.py`，逐字核对）
+
+### RMSNorm — 标准（fp32 算 variance，eps=1e-5）。
+
+### Router（`GptOssTopKRouter`，softmax 在 topk **之后**，含 bias）
+```
+logits  = x @ W_router^T + b_router          # [T, 32]
+top_val, idx = topk(logits, k=4, dim=-1)     # [T, 4]
+top_val = softmax(top_val, dim=-1)           # 仅对选中的 4 个归一化
+scores  = zeros[T,32].scatter(1, idx, top_val)
+```
+
+### Experts（`GptOssExperts`，fused gate_up，**interleaved**；clamped；(up+1)·glu）
+```
+alpha=1.702; limit=7.0
+gate_up = x @ gate_up_proj[e] + gate_up_proj_bias[e]   # [.., 2*dim]
+gate = gate_up[..., ::2]; up = gate_up[..., 1::2]      # 偶/奇 交错
+gate = clamp(gate, max=limit)                          # 仅上界
+up   = clamp(up, min=-limit, max=limit)
+glu  = gate * sigmoid(gate * alpha)
+h    = (up + 1) * glu                                   # 注意 (up+1)
+y_e  = h @ down_proj[e] + down_proj_bias[e]
+out  = Σ_{e∈top4} scores[t,e] * y_e
+```
+
+### Attention（`eager_attention_forward`，**带 sinks**）
+```
+scaling = head_dim**-0.5 = 64**-0.5；q/k/v/o 都有 bias
+RoPE(theta=150000) on q,k；repeat_kv(n_rep=8)
+attn = (q @ k^T) * scaling + causal_mask        # 滑窗层叠加 banded(window=128)
+sinks = module.sinks[head]                       # 每 head 一个标量
+combined = cat([attn, sinks broadcast], dim=-1)  # 多一列
+combined -= combined.max(-1, keepdim)            # 数值稳定
+probs = softmax(combined, -1)
+scores = probs[..., :-1]                          # 丢掉 sink 列 => 概率不归一到 1
+o = (scores @ v) -> merge heads -> @Wo + bo
+```
+> sinks 等价于 softmax 分母多了 `exp(sink)`——可学习的"不注意"通道。
+> 交替 sliding window：config `layer_types` 标明哪些层 window=128，其余全注意力。
+
+与 Qwen3 的新增点：MoE FFN、MXFP4 反量化、attention sinks（softmax 多一列再丢）、
+交替 sliding window、q/k/v/o bias、head_dim=64、clamped `(up+1)*glu`、rope_theta=150000。
+
+### 实测张量布局（layer 0，已用 `tools/mxfp4_probe.py` 核对）
+```
+self_attn.q_proj.weight [4096,2880] +bias[4096]   # 64 heads*64
+self_attn.k_proj.weight [512,2880]  +bias[512]    # 8 kv*64
+self_attn.v_proj.weight [512,2880]  +bias[512]
+self_attn.o_proj.weight [2880,4096] +bias[2880]
+self_attn.sinks [64]                               # 每 q-head 一个标量（BF16）
+input_layernorm.weight [2880]; post_attention_layernorm.weight [2880]
+mlp.router.weight [32,2880] +bias[32]
+mlp.experts.gate_up_proj_blocks [32,5760,90,16] U8 + _scales [32,5760,90] U8 + _bias[32,5760] BF16
+mlp.experts.down_proj_blocks    [32,2880,90,16] U8 + _scales [32,2880,90] U8 + _bias[32,2880] BF16
+# 全局: model.embed_tokens.weight, model.norm.weight, lm_head.weight (BF16)
+```
+MXFP4 打包：`[..., nblk=90, 16]` U8，每 16 字节 = 32 个 FP4 码（低 nibble=偶 idx，高 nibble=奇 idx），
+每 block 一个 E8M0 scale；`90*32 = 2880 = 输入(hidden)维`。即 gate_up 每 expert 权重逻辑 shape
+`[5760 out, 2880 in]`（**已转置存储**：行=out，列=in，与 HF `nn.Linear` 一致 `y=x·Wᵀ`）。
+
+### RoPE（**rotate_half，非 interleave**）
+```
+dim = head_dim = 64; base = rope_theta = 150000
+inv_freq = 1 / base^(arange(0,64,2)/64)        # 32 项
+freqs = pos ⊗ inv_freq                          # [S, 32]；cos/sin = cos(freqs)/sin(freqs)  (不 doubling)
+# 应用: x=[.., 64], first=x[:32], second=x[32:]
+#   out_first  = first*cos - second*sin
+#   out_second = second*cos + first*sin
+```
+> ⚠️ 与 Qwen3 的 RoPE kernel（interleave）不同 —— gptoss 走 rotate_half。需单独处理。
+
+### Decoder layer（pre-norm 残差，结构同 Qwen3）
+```
+h = x + attn(input_norm(x))            # attn 含 sinks/bias/滑窗
+out = h + moe(post_norm(h))            # moe = router + top4 experts 加权和
+```
+最终：`logits = lm_head(norm(h_last))`。无 q_norm/k_norm（与 Qwen3 不同，gptoss 没有）。
+
+## 3. MXFP4 反量化（expert 权重）
+
+expert 张量名：`model.layers.{i}.mlp.experts.gate_up_proj_blocks/_scales`、
+`...down_proj_blocks/_scales`（bias 为 BF16）。MXFP4：每 32 元素一 block 共享一个
+E8M0(8-bit 指数) scale，每元素 4-bit FP4(E2M1)。反量化
+`val = fp4_lut[code] * 2^(e8m0 - 127)`。**P19.1 先用 Python(numpy) 反量化并与 HF 一层
+数值对照**（block 方向 / LUT / gate_up interleave），再写进 Rust loader。
+
+## 4. 路线（正确优先）
+
+1. **P19.1** Python 侦查 + MXFP4 反量化验证（不依赖 GPU）。
+2. **P19.2** `config.rs` 加 MoE 字段（Qwen3 路径不变）。
+3. **P19.3** `gptoss.rs`：dense（attn+sinks+bias+滑窗 / norm / lm_head）+ MoE FFN
+   （正确优先：逐 token top-4 gather→clamped SwiGLU→加权和）；MXFP4 在 `from_weights`
+   反量化为 BF16。验收：prefill logits 与 HF BF16 容差内一致（top-1 一致）。
+4. **P19.4** 接 PP（experts 随层切），`--pp` 端到端；PP=2/4 与 PP=1 等价。
+5. **P19.5** llama.cpp 对比（升级 submodule 到支持 gpt-oss 的版本 + 取/转 GGUF），
+   跑 AIME 2025 + GSM8K，复用 `tools/bench` + `summarize_fullq.py`。
+
+## 5. 风险
+
+- MXFP4 格式细节必须逐字对 → Python 反量化兜底。
+- attention sinks + 交替滑窗：现有 flash/paged kernel 未必支持 → 正确优先版本先走朴素
+  attention（显式 mask + sink 列）。
+- llama.cpp pinned b9371 早于 gpt-oss（约 2025-08）→ 需升级 submodule，有连锁影响。
+- 性能：MoE 正确优先版本（逐 expert gather/scatter）会慢；先对再快。
+- **环境**：huggingface.co 被墙，需经代理 + hf-mirror 下载（见 `MOE_PROGRESS.md` §2）。
+
+## 6. 不在本阶段范围
+
+GPU 原生 MXFP4 + 按需反量化 kernel（先全 BF16）；高性能 grouped-GEMM / expert parallel；
+TP×MoE；单卡运行（需 MXFP4-native）。

docs/20-sparse-moe.md

@@ -1,160 +0,0 @@
-# Phase 20: Sparse MoE Decode — 只算被路由到的专家
-
-> 目标:消除 dense MoE 的无效权重读取,decode TPOT 追上并超过 llama.cpp。
-> 前置:Phase 19(gpt-oss MoE 正确性)、FP8 W8A8 / MXFP4 W4A16 量化
-> (见 `docs/benchmarks/fp8-quantization.md`、`docs/benchmarks/mxfp4-and-llama-decode.md`)。
-
-## 1. 现状:dense MoE 在浪费什么
-
-gpt-oss-20b 是 32 专家 top-4 的 MoE:router 给每个 token 选 4 个专家,
-理论上每 token 只需要读 4/32 = 12.5% 的专家权重。但 `moe_forward`
-(`crates/xserv-model/src/gpt_oss.rs`)目前是 **dense** 实现:
-
-```text
-1. router GEMV            [T, 2880] → [T, 32]
-2. topk_softmax (GPU)     → topk_ids [T,4], topk_weights [T,4]
-3. moe_replicate          x 复制 16 份 → [16, T, 2880]        ← 浪费开始
-4. batched GEMM gate_up   全部 16 个本地专家都算                ← 读 16 份权重
-5. bias + GLU
-6. batched GEMM down      全部 16 个本地专家都算                ← 读 16 份权重
-7. bias
-8. moe_weighted_sum       只挑出 top-4 加权求和,其余 12 个全部丢弃
-9. all-reduce
-```
-
-为什么当初这么写:batched GEMM(cuBLAS strided-batched)要求规则的
-`[E, T, K]` 形状;top-4 的专家编号在 **GPU** 上(`topk_ids`),host 不知道
-该挑哪几个,挑了形状也不规则。dense 是"先把正确性做出来"的合理起点,
-但每 token 把 16 个专家的权重从 HBM 全部读一遍。
-
-### 字节账本(decode,每 token,TP=2 每卡 16 个本地专家)
-
-每层每专家:gate_up `[2880, 5760]` + down `[2880, 2880]` ≈ 24.9 M 参数。
-
-| 方案 | 每卡每 token 专家字节 | 相对量 |
-|---|---|---|
-| xserv dense FP8(现状) | 16 × 24.9 MB × 24 层 ≈ **9.6 GB** | 1× |
-| xserv sparse FP8(本阶段) | ~2 × 24.9 MB × 24 层 ≈ **1.2 GB** | 1/8 |
-| llama.cpp sparse MXFP4 | ~2 × 12.5 MB × 24 层 ≈ **0.6 GB** | 1/16 |
-
-(top-4 均匀散落在 2 张卡上,期望每卡 2 个命中;严格说每层取的是
-两卡命中数的 max,期望 ≈ 2.6,仍是 ~6-8× 的节省。)
-
-实测旁证:FP8 dense TP=2 TPOT 13.1 ms,其中专家 GEMM ≈ 9.6 GB ÷ ~1 TB/s
-≈ 9.5 ms,其余(attention、qkv/o、lm_head、48 次 PCIe all-reduce)≈ 3.5 ms。
-**专家权重读取占 TPOT 的 ~3/4,这就是与 llama.cpp(6.6 ms)的全部差距。**
-
-## 2. Roofline:M=1 时为什么"省字节 = 省时间"
-
-decode 的 GEMV(M=1)每读 1 字节 FP8 权重只做 2 FLOP(乘加)。
-RTX 5090:HBM ~1.8 TB/s,BF16 算力 ~210 TFLOPS —— 算强比(arithmetic
-intensity)需要 ~100 FLOP/byte 才能喂饱算力,GEMV 只有 2。结论:
-
-1. **decode 完全 memory-bound**,tensor core 帮不上忙 → 手写 W8A16 GEMV
-   (权重 FP8、激活保持 BF16)不会输给 cuBLASLt 的 W8A8 tensor-core GEMM,
-   还省掉激活量化 kernel,精度更好(激活不再有量化误差)。
-2. 优化只有一个方向:**少读字节**。sparse(×8)与 4-bit(×2)正交,
-   可叠加。本阶段先做 sparse,FP8 与 MXFP4 两种权重格式都支持。
-
-## 3. Sparse 设计:让 kernel 自己按 topk_ids 索引权重
-
-关键观察:`topk_ids` 本来就在 GPU 上。不需要 host 知道选了谁 ——
-**让 GEMV kernel 的每个 block 自己读 `topk_ids[token, slot]`,
-直接寻址到对应专家的权重**,不命中本卡就整块退出。零 host 同步,
-管线保持完全异步(这是之前排查过的:decode 循环无 per-layer sync)。
-
-新数据流(`num_tokens ≤ 8` 时启用):
-
-```text
-x [T, 2880]
-  ├─ router → topk_ids/weights [T, 4]               (不变)
-  ├─ sparse GEMV gate_up  → [T, 4, 5760]   bias 已融合,非本地 slot 不写
-  ├─ GLU                  → [T*4, 2880]
-  ├─ sparse GEMV down     → [T, 4, 2880]   bias 已融合,非本地 slot 不写
-  └─ weighted_sum_sparse  → [T, 2880]      只累加本地 slot
-all-reduce                                            (不变)
-```
-
-`moe_replicate` 和独立的 bias kernel 在 sparse 路径下消失;FP8 路径还省掉
-`quantize_bf16_to_fp8_rowwise`。
-
-### Kernel 设计(`csrc/moe/moe_sparse.cu`)
-
-`moe_sparse_gemv_{fp8,mxfp4}_bf16_kernel`:
-
-- **grid = (N/8, top_k, tokens)**,block = 8 warp × 32 lane。
-  每个 block 负责一个 (token, slot) 的 8 个输出列,**一个 warp 算一个输出**。
-- block 先读 `eid = topk_ids[token*top_k + slot]`,折算 `lid = eid - expert_start`;
-  不在 `[0, local_experts)` 就整块 return。
-- 命中的 block 把激活行(K=2880 个 BF16 → float)协作搬进 shared memory
-  (11.25 KB),`__syncthreads()`,然后每 warp 沿 K 维做点积:
-  每 lane 一次 `uint4` 读 16 字节权重(FP8 = 16 个权重,MXFP4 = 32 个 nibble),
-  warp 内 32 lane 连续 → 512B coalesced 事务。
-- epilogue(lane 0):`y = acc * w_scale[lid] + bias[lid, n]` —— per-expert
-  scale 和 bias 都融合在这里,与 dense 路径的"GEMM → bias add → 路由加权"
-  语义逐位等价(HF 参考实现也是先加 bias 再乘路由权重)。
-- gate_up 与 down 共用同一个 kernel,用 `x_per_slot` 区分激活寻址:
-  gate_up 时 4 个 slot 共享 `x[token]`;down 时各读自己的 `act[token*4+slot]`。
-
-### 两个容易写错的安全点
-
-1. **early-return 必须 block-uniform。** Phase 19 的 GEMV 垃圾输出 bug
-   (commit `3b9e32e`)正是"部分线程在 `__syncthreads()` 之前 return"导致
-   读未初始化 shared memory。这里的 return 发生在 smem 装载**之前**,且整个
-   block 基于同一个 `topk_ids` 值统一退出 —— 没有 divergence,合法且安全。
-2. **weighted-sum 对非本地 slot 必须"跳过",不能"乘 0"。** 非本地 slot 的
-   GEMV 输出从未被写入(未初始化显存,可能是 NaN 位型),GLU 也会在上面算出
-   垃圾。`NaN × 0 = NaN`,所以求和 kernel 用 `if (local) sum += w*v` 跳过,
-   垃圾永远不进入数据流(dense 路径的 `moe_weighted_sum` 同理)。
-
-## 4. 为什么 prefill 保持 dense
-
-dense batched GEMM 把 16 份权重读**一次**,服务全部 M 个 token;
-sparse GEMV 是**每 token** 重读自己的 ~2 份。字节交叉点:
-
-```text
-sparse 读 M × 2 份  vs  dense 读 16 份  →  M ≈ 8 (TP=2)
-```
-
-M > 8 后 dense 更省(且 GEMM 是 compute-bound,tensor core 开始有用)。
-所以 sparse 只在 `num_tokens ≤ 8` 启用 —— 覆盖 decode(连续批合并的
-多请求 decode 也是小 M)和极短的 re-prefill。真正的 sparse prefill
-(按专家对 token 做 permute/gather 的 grouped GEMM,vLLM 的做法)是
-后续阶段,主要收益在长 prompt TTFT。
-
-## 5. 实测结果(2026-06-12,完整数据见 `docs/benchmarks/sparse-moe.md`)
-
-In-process decode(bench-gpt-oss,greedy 96 tok):
-
-| | TPOT | tok/s |
-|---|---|---|
-| dense FP8 TP=2(基线) | 13.9 ms | 72 |
-| **sparse FP8 TP=2** | **7.6 ms(1.8×)** | **132** |
-| sparse MXFP4 TP=2 | 8.4 ms | 118 |
-| sparse FP8 TP=1(单卡) | 7.8 ms | 128 |
-
-Warm-server 对打 llama.cpp(`tools/xserv_vs_llama.py`):
-
-- **TP=2 vs TP=2:xserv 首次全面反超** —— TPOT 7.19-7.32 ms vs llama
-  7.54-8.42 ms;短/中 prompt TTFT 也领先(35/49 vs 63/65 ms)。
-- **TP=1 vs TP=1:llama 大胜**(2.88-3.22 ms vs 7.0-7.2 ms,347 vs 140
-  tok/s)。单卡才是 llama 的最优配置:它的跨卡 split 在 PCIe 上每 token
-  损失 ~5 ms,而单卡时它"全模型 4-bit + CUDA graph 整 token 回放"的
-  优势全部兑现。xserv 的残余 ~7 ms ≈ ~3 ms HBM(其中非专家权重还是
-  BF16,含 1.16 GB 的 lm_head)+ ~4 ms 启动开销(~200 个 kernel
-  launch/token,无 CUDA graph)。
-- **正确性:GSM8K-100 = 96%**(dense FP8 91% / BF16 90%,greedy 噪声内,
-  无回归)。
-
-教训:之前"CUDA graph ≈ 无用(~0.5-1.5ms)"的结论是相对 13 ms 的
-dense TPOT 而言;专家成本砍掉后,launch 开销变成了最大的单项。
-
-## 6. 下一阶段(按收益排序)
-
-1. **decode CUDA graph**(~2-4 ms):当前最大单项。
-2. **非专家权重量化**(~1-1.5 ms):qkv/o + lm_head 仍是 BF16,每 token
-   白读 ~2.3 GB;llama 是全模型 4-bit。
-3. **sparse prefill**(grouped GEMM):长 prompt TTFT 94-120 ms → llama
-   的 ~30 ms 量级。
-4. **W4A4 FP4 tensor core / 带宽调优的 MXFP4 GEMV**:让 4-bit 专家真正
-   快过 FP8(目前 8.4 vs 7.6 ms,GEMV 效率抵消了字节优势)。

docs/21-cuda-graph-decode.md

@@ -1,111 +0,0 @@
-# Phase 21: gpt-oss decode CUDA Graph + GPU argmax
-
-> 目标:消除 decode 的每 token 固定开销。Phase 20 之后 TPOT ~7ms,其中
-> GPU 实际计算只占一部分,剩下是 ~200 个 kernel launch 和 per-token 的
-> host 工作。本阶段把**整个 decode step 捕获成一个 CUDA graph**,每 token
-> 一次 `cudaGraphLaunch` 回放;顺带把 greedy 采样换成 GPU argmax。
->
-> 实现:`crates/xserv-model/src/gpt_oss_graph.rs`(~150 行)+ 三块基础设施。
-
-## 1. CUDA Graph 是什么,为什么有约束
-
-`cudaStreamBeginCapture` 之后,发到该 stream 的 kernel 不执行而是被**录制**;
-`EndCapture + Instantiate` 得到可执行图;以后每步 `cudaGraphLaunch` 一次性
-重放全部 ~200 个 kernel,host 端开销从 ~200 次 launch 降到 1 次。
-
-代价是三条硬约束,每条都对应一个工程问题:
-
-1. **地址稳定**:录制时烤进图里的全部指针,回放时必须仍然有效且指向正确数据;
-2. **capture 期间禁止"不安全"调用**:`cudaMalloc`/同步 memcpy/`cudaDeviceSynchronize`
-   都会让 capture 报错(error 900);
-3. **形状固定**:grid 尺寸被烤死,变 shape 就要重录。
-
-## 2. 为什么 xserv 的 decode 本来就"差一点"就能整图捕获
-
-逐项检查 decode step 的输入,发现绝大部分已经满足地址稳定:
-
-| 每步会变的输入 | 地址 | 内容如何更新 |
-|---|---|---|
-| block table / context lens | PagedKVCache 的常驻 GPU 缓冲 ✓ | `decode_prepare` 在图外 H2D |
-| KV 写入位置 | scatter kernel **从 GPU 上的 context_lens 读** ✓ | 同上 |
-| attention 读取范围 | paged kernel 从同一缓冲读 ✓ | 同上 |
-| MoE 专家选择 | sparse GEMV 从图内刚写的 `topk_ids` 读 ✓ | 数据依赖,天然支持 |
-| token id / position | ✗ 原来是每步从 host slice 上传 | **本阶段改造点** |
-
-也就是说,Phase 11(paged KV)和 Phase 20(sparse MoE)的"数据驱动"设计
-无意中已经为 graph 化铺平了路 —— 唯二需要动的是 embedding 的 token id 和
-RoPE 的 position:各加一个 device-buffer 变体(`embedding_device_ids` /
-`rope_inplace_device_pos`),id/pos 存进两个常驻 4 字节缓冲,每步图外更新。
-
-重构后的结构:
-
-```text
-forward_decode_paged = decode_prepare(host 簿记,图外)
-                     + upload ids/pos(图外)
-                     + decode_core(纯 GPU,可整段捕获)
-                     + advance_seq_len(host 簿记,图外)
-```
-
-## 3. 三个工程问题
-
-### 3.1 null stream 不可捕获 → thread-local launch stream
-
-全代码库的 kernel 都发射在 legacy null stream 上,而 capture 必须在显式
-stream 上。解法:`xserv_cuda::stream` 加一个 **thread-local 当前 stream**
-(默认 null,行为与从前逐字节一致),所有 kernel wrapper、cuBLAS 的
-`cublasSetStream`、NCCL 的 collective 全部改读它。capture 代码用 RAII guard
-(`push_stream`)把 capture stream 装进去,录完自动还原。
-顺序正确性:显式 stream 以默认(blocking)方式创建,legacy stream 与其
-双向隐式同步,所以图外的 H2D/采样 memcpy 与回放天然有序。
-
-### 3.2 capture 期间禁止 cudaMalloc → "retained warmup" 二段式
-
-中间张量来自 caching allocator;capture 中任何一次 pool miss 都会触发
-`cudaMalloc` → error 900。第一版实现就栽在这里:**隔离机制自己制造了
-pool miss**(capture 中释放的块被隔离,下一层同尺寸分配就找不到块了)。
-
-解法是把同一个 step 先 eager 跑一遍、但**隔离打开**(`begin_retain`):
-释放的块全部扣下不回池 → 跑完后池外恰好积累了"这一步需要的每一块";
-把它们整批放回池,再开始 capture —— capture 重复完全相同的分配序列,
-每次分配都命中池,一次 cudaMalloc 都不会发生。
-(重复执行同一 step 是无害的:KV scatter 往同一个位置重写同样的值。)
-
-### 3.3 回放引用的内存不能被别人拿走 → 隔离仓(quarantine)
-
-capture 录下的中间缓冲在 host 侧早就 Drop 了,但图每次回放都会读写这些
-地址。若它们回到分配池、被后续 prefill 拿走长期持有,就是双写损坏。
-所以 capture 期间释放的块进入 `RetainedBlocks` 隔离仓,由 graph 对象持有,
-graph 销毁时才归还 —— 这些内存在 graph 存活期内被锁定为它专用。
-
-### 3.4 两个顺手的点
-
-- **THREAD_LOCAL capture mode**:GLOBAL 模式下,任何线程的 cudaMalloc 都会
-  毒化 capture;TP 多 rank 线程并发 capture 必须用 THREAD_LOCAL。
-- **NCCL 可以被捕获**:rank 内 `ncclAllReduce` 发在 capture stream 上即可,
-  TP=2 一次成功(各 rank 录各自的图,回放时 collective 自然配对)。
-
-## 4. 意外的教训:launch 开销没有想象的大,argmax 才是大头
-
-A/B 实测(in-process,FP8,96 tok):
-
-| | TP=1 | TP=2 |
-|---|---|---|
-| eager + host argmax(Phase 20 末) | 7.5 ms | 7.6 ms |
-| graph + host argmax | 6.9 ms | 6.9 ms |
-| eager + **GPU argmax** | 6.5 ms | — |
-| **graph + GPU argmax** | **5.9 ms** | **5.8 ms** |
-
-- **graph 只省了 ~0.6ms**:decode 循环本来就是全异步的,launch 大部分被
-  GPU 执行掩盖,"~200 launch ≈ 4ms"的预估错了 —— 优化要测不要猜。
-- **GPU argmax 省了 ~1ms**:greedy 采样原来每 token 把 [1, 201088] 的
-  logits(402KB)同步拷回 host、再扫描 201K 个 bf16。仓库里 Phase 15 就写好
-  的 argmax kernel(kernel 内归约 + 4 字节 D2H)一直没接到 `sample()` 上。
-- 细节:GPU argmax 与 host `max_by` 对**完全相等**的 logits 平局取的索引
-  不同,greedy 轨迹会在某个平局 token 处分叉 —— 输出同样合法(GSM8K 验证)。
-
-## 5. 结果与剩余瓶颈
-
-见 `docs/benchmarks/sparse-moe.md` 的 Phase 21 小节(warm-server 对打 llama
-的数字以那里为准)。剩余 TPOT 的构成:~3ms 是 HBM 字节(其中非专家权重
-仍是 BF16,含 1.16GB 的 lm_head —— **Phase 22 量化它们**),其余是 GEMV
-带宽效率与 attention。llama 单卡 2.9ms 的差距主要就在"全模型 4-bit"。

docs/22-speculative-decoding.md

@@ -1,186 +0,0 @@
-# Phase 22: Draft-Model Speculative Decoding v0
-
-> 目标:实现一个可验证的 speculative decoding 最小闭环。先只覆盖
-> Qwen3 target + 同 tokenizer 的小 Qwen3 draft、batch=1、greedy
-> (`temperature=0`)。本阶段不做 gpt-oss,不做 sampling rejection,不接入
-> continuous batching。
-
-## 1. Scope
-
-本阶段只解决一个窄问题:
-
-- target:现有 Qwen3 paged KV 路径,优先 Qwen3-8B;
-- draft:同 tokenizer 的小 Qwen3,例如 Qwen3-0.6B;
-- batch size:1;
-- decoding:greedy argmax;
-- draft window:`gamma=4`;
-- acceptance:exact-match,即 `target_argmax == draft_token`。
-
-HTTP flag 可以后续接入。v0 先提供独立 bench/CLI,因为它能直接输出 token
-一致性、acceptance rate、tokens/target-step、TPOT/tok/s,也避免把尚未稳定的
-rollback 行为放进服务端调度循环。
-
-bench 为了让 baseline/spec 对比不受跨 prompt KV pool 复用影响,每个 prompt 的
-baseline run 和 speculative run 都使用新建的 paged KV cache。cache 分配发生在
-单次 run 的计时外,输出的 TPOT/tok/s 只覆盖模型 prefill/decode 工作。
-
-## 2. Why Qwen3 First
-
-Qwen3 是现有代码里最适合作为 speculative v0 的模型族:
-
-1. target 已有稳定的 `forward_prefill_paged` 和 `forward_decode_paged`;
-2. 小 Qwen3 与 Qwen3-8B 共享 tokenizer,可以直接比较 token id;
-3. Qwen3 是 dense decoder-only,没有 gpt-oss 的 harmony 格式、MoE sparse 路径、
-   sliding-window 或 CUDA Graph 状态;
-4. greedy 输出的正确性定义简单:只要 spec 生成的 token 序列与纯 target greedy
-   完全一致即可。
-
-gpt-oss spec 需要先定义 harmony prompt、MoE draft 选择、graph replay 与 rollback
-的交互,这些都不属于本阶段。
-
-## 3. Algorithm
-
-对每个 prompt 建两套模型、三套 KV 状态:
-
-```text
-target model + target commit PagedKVCache
-target model + target verify PagedKVCache
-draft  model + draft PagedKVCache
-```
-
-先把 prompt 分别 prefill 到三套 cache。此时 cache 都包含 prompt,并各自持有
-"下一个 token" 的 logits。
-
-每个 speculative round:
-
-1. draft 从当前 draft logits 取 argmax,连续生成 `gamma` 个 draft token;
-2. draft 每生成一个 token 就用自己的 paged decode append 到 draft KV,所以 round
-   结束时 draft cache 暂时包含整个草稿序列;
-3. target verify cache 对完整 draft token 序列调用一次 paged prefill,覆盖
-   "target 可一次验证草稿窗口" 这条执行路径;
-4. target verify cache 立刻 rollback 到 round 起点,避免把 prefill 临时写入污染
-   commit cache;
-5. 用 target decode 轨迹作为权威结果,从左到右比较
-   `target_next_argmax == draft_token`,只接受连续匹配前缀;
-6. 对每个接受 token,用 target decode 重放一次来提交 target KV,并得到下一步
-   `target_next_argmax`;verify cache 也 mirror decode 同一个 token,保持长度与 prefix 对齐;
-7. 若全部匹配,draft cache 已经包含完整草稿,三套 cache 长度重新对齐;
-8. 若在第 `k` 个 token 拒绝,提交前 `k` 个 draft token,再提交 target 在该位置的
-   argmax 作为修正 token。draft cache rollback 到 round 起点后重放接受 token 和修正
-   token,target commit/verify cache 都由 decode 路径提交到同一 prefix。
-
-v0 不使用完整 speculative sampling 的概率校正。它只利用小模型猜测 greedy 轨迹,
-因此生成序列必须与纯 target greedy 完全一致。
-
-当前实现选择 decode 轨迹作为提交路径,而不是直接保留 target prefill 写入的 KV。
-原因是 v0 验收要求 token 序列与纯 target greedy 完全一致;如果 prefill 和 decode
-路径在数值或 KV 写入顺序上存在细微差异,直接提交 prefill KV 会让后续 greedy 输出
-漂移。这个保守实现仍会执行 target paged prefill 验证和 rollback,但 verify 写入放在
-独立 cache,不会影响权威 commit cache。代价是额外 mirror decode,速度收益预期较差,
-主要用于先验证 draft-model speculative 的状态机和一致性。
-
-为保证 greedy exactness,decode 里两个原有非确定点也需要固定:
-
-- BF16 GEMV 不再用跨 K-block `atomicAdd`;改为写 K-block partials,再按固定顺序
-  reduce;
-- paged decode attention 不再用 `atomicAdd` 合并 warp 输出;改为 per-warp partials
-  后按 warp id 顺序 reduce。
-
-## 4. KV Commit And Rollback
-
-现有 `forward_prefill_paged` 会一次性把传入 token 写进 paged KV,并提前推进
-`seq_len`。验证草稿时 target verify cache 因此会临时包含整个 draft window。
-
-新增的 cache 操作只做逻辑截断:
-
-```text
-truncate_sequence(slot, new_len)
-```
-
-约束:
-
-- 只允许 `new_len <= current_len`;
-- 保留覆盖 `[0, new_len)` 所需的物理 block;
-- 释放右侧多余 block;
-- 不清零仍在保留 block 内的旧字节,因为后续逻辑长度会阻止 attention 读取它们,
-  同一位置再次写入时会覆盖旧值;
-- slot 仍保持 registered,`new_len=0` 时也保留第一个 block。
-
-这让 target 和 draft 都能在拒绝时安全丢弃多写 KV,并在修正 token decode 后重新
-对齐。
-
-## 5. Acceptance Criteria
-
-本阶段验收:
-
-- `cargo fmt`;
-- `cargo check`;
-- `cargo test`;
-- `bench-speculative` 可加载 target+draft 两套 Qwen3;
-- 50 prompts,greedy,baseline target 与 speculative token id 序列完全一致;
-- 输出 acceptance rate、tokens/target-step、TPOT、tok/s 和 speedup;
-- 若 draft 模型缺失或磁盘不足,明确报告阻塞条件,不盲目下载大模型。
-
-## 6. Validation Results
-
-dash5 环境:
-
-- GPU:RTX 5090,device 0;
-- target:`/opt/wjh/models/qwen3-8b`;
-- draft:`/dashscope-tmp/wjh/models/qwen3-0.6b`;
-- command:`bench-speculative ... --prompts 50 --gen-tokens 32 --gamma 4 --device 0`;
-- log:`/dashscope-tmp/wjh/xserv-spec-default-50x32-final.log`。
-
-默认 `acceptance_mode=decode` 的结果:
-
-```text
-prompts=50 matched=true
-acceptance_rate=0.3664 accepted=1020 proposed=2784
-tokens_per_target_step=0.3639 target_steps=4397
-verify_steps=729 mirror_decode_steps=1550 commit_decode_steps=1550 correction_steps=568
-verify_decode_mismatches=10
-baseline_e2e_tpot_ms=13.123 baseline_e2e_tok_s=76.204
-spec_e2e_tpot_ms=44.867 spec_e2e_tok_s=22.288 speedup_e2e=0.2925
-baseline_decode_tpot_ms=12.638 baseline_decode_tok_s=79.127
-spec_decode_tpot_ms=43.731 spec_decode_tok_s=22.867 speedup_decode=0.2890
-decode_token_counts baseline=1600 spec=1600
-```
-
-诊断 `--use-verify-logits` 的结果:
-
-- command:`bench-speculative ... --prompts 10 --gen-tokens 32 --gamma 4 --device 0 --use-verify-logits`;
-- log:`/dashscope-tmp/wjh/xserv-spec-verify-logits-10x32.log`;
-- exit status:`2`;
-- summary:`matched=false`, `verify_decode_mismatches=4`;
-- prompt 0/2/7 出现 baseline/spec token 序列分叉。
-
-结论:当前可以做 correctness-first 的 speculative decoding 状态机,但还不能把
-target batched prefill verify logits 作为 greedy 接受依据。verify prefill 路径与
-逐 token decode 路径存在 top-1 不一致;默认模式必须继续以 decode 轨迹为权威,
-因此 v0 是正确性闭环,不是性能优化。
-
-## 7. Known Limits
-
-- 只支持 batch=1;
-- 只支持 Qwen3-family dense models;
-- 只支持 greedy exact-match acceptance;
-- 未实现 probabilistic rejection sampling,所以 temperature/top-k/top-p 不支持;
-- 未接 HTTP/continuous batching;
-- 未与 CUDA Graph decode 结合;
-- 当前 v0 为保证 greedy exactness,接受 token 也会用 target decode 重放提交,因此
-  即使 acceptance 高也可能变慢;
-- draft prefill 和 target prefill 都会计入端到端耗时,短输出可能没有收益。
-
-## 8. Next Phase TODO
-
-如果继续 speculative decoding,下一阶段不要先接 HTTP,应先解决 verify 路径:
-
-1. 做最小 prefill-vs-decode parity harness:固定 prompt、cache len、draft token,
-   dump 每层/最终 logits 的 top-k,定位 top-1 分叉来自 attention、GEMV 还是 KV 写入顺序;
-2. 让 `--use-verify-logits` 在至少 50 prompts x 64 tokens 下 `matched=true` 且
-   `verify_decode_mismatches=0`;
-3. parity 过后再做真正 multi-token target commit:要么安全保留 verify prefill 写入的
-   KV,要么实现专用 paged multi-token verify/commit kernel,避免当前的 mirror+commit
-   decode 重放;
-4. 只有 `speedup_e2e > 1` 后再考虑 HTTP flag、continuous batching、sampling 或
-   gpt-oss speculative decoding。

docs/23-speculative-verify-parity.md

@@ -1,85 +0,0 @@
-# Phase 23: Speculative Verify Parity
-
-> 目标:把 speculative decoding 从 v0 的 correctness-only 状态机推进到
-> "verify logits 可作为权威接受依据"。本阶段仍只覆盖 Qwen3 target +
-> Qwen3 small draft、batch=1、greedy。
-
-## 1. Problem
-
-Phase 22 的默认模式用逐 token target decode 作为权威路径,因此输出能与 baseline
-一致。但诊断 `--use-verify-logits` 会失败:target 对 draft window 做 batched
-prefill verify 时,部分 logits top-1 与逐 token decode 不一致。
-
-实测 top-k 显示分叉不是大幅数值错误,而是 BF16 near-tie:
-
-```text
-verify_top5=17689:24.500,9856:24.375,...
-decode_top5=9856:24.500,17689:24.500,...
-```
-
-如果直接用这些 verify logits 接受/拒绝 draft token,greedy token 序列会偏离纯
-target decode。
-
-## 2. Design
-
-新增 `Qwen3::forward_verify_paged_decode_attention`:
-
-1. 在 target commit cache 上一次写入 draft window 的 K/V;
-2. attention 使用现有 paged decode attention,每个 draft token 对应一行 metadata,
-   context lens 分别为 `pos + 1`;
-3. 线性层使用逐行 GEMV,与 `forward_decode_paged` 的 BF16 rounding path 对齐;
-4. 若 token 全接受,直接保留 verify 写入的 KV;
-5. 若在第 `k` 个 token 拒绝,把 target cache truncate 到 accepted prefix,再只
-   decode 一个 correction token。
-
-bench 新增:
-
-- `--use-verify-logits`:用 verify logits 作为接受依据,默认选择 `paged-decode`
-  verify path;
-- `--verify-path flash|paged-decode`:显式选择旧 flash prefill 诊断或新 paged-decode
-  verify path;
-- `--dump-verify-mismatches`:打印 mismatch 行 top-k,用于定位 near-tie。
-
-## 3. Validation
-
-dash5:
-
-- GPU:RTX 5090,device 0;
-- target:`/opt/wjh/models/qwen3-8b`;
-- draft:`/dashscope-tmp/wjh/models/qwen3-0.6b`;
-- command:`bench-speculative ... --prompts 50 --gen-tokens 64 --gamma 4 --device 0 --use-verify-logits`;
-- log:`/dashscope-tmp/wjh/xserv-spec-inplace-verify-50x64.log`。
-
-结果:
-
-```text
-prompts=50 matched=true
-acceptance_mode=verify_logits
-verify_path=paged-decode
-acceptance_rate=0.3927 accepted=2120 proposed=5398
-tokens_per_target_step=0.9112 target_steps=3512
-verify_steps=1376 mirror_decode_steps=0 commit_decode_steps=1068 correction_steps=1068
-verify_decode_mismatches=0
-baseline_e2e_tpot_ms=13.094 baseline_e2e_tok_s=76.372
-spec_e2e_tpot_ms=30.069 spec_e2e_tok_s=33.257 speedup_e2e=0.4355
-baseline_decode_tpot_ms=12.846 baseline_decode_tok_s=77.844
-spec_decode_tpot_ms=29.731 spec_decode_tok_s=33.635 speedup_decode=0.4321
-decode_token_counts baseline=3200 spec=3200
-```
-
-对比 Phase 22 的保守 decode-authoritative v0:
-
-- verify logits 现在可以作为权威接受依据;
-- `mirror_decode_steps` 从每个 accepted token 一次降为 0;
-- 50x64 e2e speedup 从约 0.29x 提升到 0.44x;
-- 仍未超过 baseline,因为 verify path 为了 parity 使用逐行 GEMV,且 draft acceptance
-  只有约 39%。
-
-## 4. Next TODO
-
-下一阶段要从 correctness parity 转向性能:
-
-1. 逐层替换 row-GEMV 为 batched GEMM,同时保留 near-tie fallback 或 top-k audit;
-2. 加一个 `--verify-audit-decode` 低频抽样审计,避免每轮都做 target decode;
-3. 扫 `gamma` 与 draft 选择,记录 acceptance 与 TPOT 曲线;
-4. `speedup_e2e > 1` 前不接 HTTP/continuous batching/gpt-oss spec。

docs/24-speculative-batched-verify.md

@@ -1,144 +0,0 @@
-# Phase 24: Speculative Decoding Performance — target `speedup_e2e > 1`
-
-> Status (2026-07-01): investigation-in-progress. Baseline reproduced,
-> naive batched-GEMM verify attempted, K/V drift issue identified,
-> concrete next-step designs written up. **Nothing landed on main yet.**
-
-## 1. Baseline (Phase 23, verified on dash5)
-
-`--prompts 50 --gen-tokens 64 --gamma 4 --use-verify-logits`:
-
-- `acceptance_rate = 0.39`
-- `matched = true`, `verify_decode_mismatches = 0`
-- `spec_e2e_tpot_ms = 30.07`, `baseline_e2e_tpot_ms = 13.09`
-- **`speedup_e2e = 0.44×`**
-- `tokens_per_target_step = 0.91`
-
-5-prompt sanity re-run reproduces the same shape (~0.44×), so the
-Phase 23 correctness state machine is intact after the recent CUDA
-determinism fixes (`5f06090`).
-
-## 2. Cost budget & the ceiling
-
-Rough numbers on 5090 TP=1:
-- `baseline decode`: ~12.6 ms / token (Qwen3-8B BF16, paged).
-- `draft decode` (Qwen3-0.6B): ~2.5 ms / token (rough estimate).
-- `verify` (Phase 23 row-GEMV, γ=4): ~13 ms.
-
-Best-case per accepted spec token cost with acceptance α, γ tokens
-per round:
-```
-spec_time_per_token ≈ (γ · draft + verify + correction) / (1 + α · γ)
-```
-With draft=2.5, verify=13, correction≈13, α=0.4, γ=4:
-```
-spec_time_per_token ≈ (10 + 13 + 13) / (1 + 1.6) ≈ 13.8 ms/token
-```
-Baseline is 12.6 ms/token. **Even with the row-GEMV verify perfectly
-free, current acceptance rate 0.39 gives us at best ~1× speedup.**
-
-## 3. What we tried (2026-07-01)
-
-Naive Phase 24: replace `matmul_rows_gemv` in
-`forward_verify_paged_decode_attention` with `matmul_2d` (batched
-cuBLAS GEMM). Result on 5 prompts × 32 tokens:
-
-- `speedup_e2e = 0.68×` (up from 0.44×) — verify itself much faster.
-- **`matched = false` on 3/5 prompts** — divergence at multiple
-  positions per failed prompt, not just first mismatch.
-
-Root cause: **K/V drift, not logit rounding**.
-
-`matmul_2d` at `m=1` routes through the custom `launch_gemv_bf16`
-kernel; at `m≥2` it goes through cuBLAS `GemmEx`. Those two paths
-produce **different BF16 bits** for the same math because their
-accumulation orders differ. Therefore:
-
-- Verify's QKV projection at `m=γ` writes K/V into the paged cache
-  with cuBLAS-GEMM values.
-- Baseline decode's QKV projection at `m=1` would have written K/V
-  with GEMV values.
-- Downstream attention reads these K/V; the two paths diverge starting
-  at the very next position. A near-tie fallback for the *current*
-  row's logit does not fix already-diverged history.
-
-Near-tie fallback (added and reverted in the same session, kept only
-in this doc) attempted to correct verify-argmax when top1−top2 was
-small. It did nothing about the K/V drift, so mismatches persisted.
-
-## 4. Revised path to `speedup_e2e > 1`
-
-Two independent levers. Combining them is the plan.
-
-### 4.1 A batched-GEMV kernel with GEMV-identical numerics
-
-Write a `launch_gemv_bf16_batched` that runs γ separate `m=1` GEMVs in
-a **single kernel launch**, sharing the K panel across rows and
-producing bit-exact-same output as γ sequential `launch_gemv_bf16`
-calls. This gives Phase 24's launch-overhead savings without breaking
-K/V bits. Estimated saving vs row-loop: ~2–4 ms per verify at γ=4
-(720 fewer launches × 3–5 μs each).
-
-Concrete kernel design:
-- Grid: `(N / TILE_N, num_k_blocks, γ)` — same layout as current
-  gemv, plus γ in the z-axis.
-- Each block reads its row's `x[γ_idx, :]` panel once, then writes
-  `partials[γ_idx, k_block, n_tile]`.
-- Reduction kernel: `(N / TILE_N, γ)`, reduces K-blocks in fixed
-  order per row (same as current `gemv_reduce_to_bf16_kernel`).
-
-Bit-exact-with-m=1 verification: run the γ=1 special case through the
-new kernel and compare to `launch_gemv_bf16`; must be bit-identical.
-
-### 4.2 Reduce verify + correction cost — draft-side CUDA graph
-
-Draft decode is currently a full eager Qwen3-0.6B forward per γ step.
-Wrapping γ draft steps into a CUDA graph (Phase 21 already did this
-for gpt-oss target decode) cuts launch overhead here too. Estimated:
-~1–1.5 ms per γ=4 window.
-
-### 4.3 Adaptive γ
-
-Currently γ=4 fixed. When acceptance drops in a "hard" section, γ=4
-wastes 3 draft steps per round. Track a moving average of acceptance
-per round; if the last N rounds averaged below τ, drop γ to 2 or 1
-(equivalent to disabling spec). If it climbs above τ_high, restore.
-
-## 5. Revised acceptance criteria
-
-1. `cargo fmt && cargo check && cargo test` on dash5.
-2. `bench-speculative --prompts 50 --gen-tokens 64 --gamma 4 --use-verify-logits`:
-   - `matched = true`
-   - `verify_decode_mismatches = 0`
-   - **`speedup_e2e > 1.0`**
-3. GSM8K-50 (if time permits) token-identical with baseline.
-
-## 6. What's on main today
-
-- `5f06090`: fixed flash decode kernel atomicAdd nondeterminism + two
-  int32 overflow bugs (causal_mask, dequant_fp8).
-- `ce10e4a`: sampling NaN-safe on top-k/top-p path.
-- `d96ee07`: API sampling validation + finish_reason normalization +
-  bounded engine channel + 4 MiB body limit.
-
-The Phase 24 attempt (batched matmul_2d in verify) is **not** on
-main. It was verified to be functionally incorrect and reverted in
-the same session; only this design doc landed.
-
-## 7. Next actions
-
-In order:
-
-1. Implement `launch_gemv_bf16_batched` + Rust wrapper `matmul_2d_gemv_batched`.
-2. Numerical parity test: γ sequential row-GEMVs vs one batched call
-   must be bit-exact for BF16 inputs.
-3. Swap `matmul_rows_gemv` in `forward_verify_paged_decode_attention`
-   for the batched variant.
-4. Re-run `bench-speculative` 50×64; expect `matched=true` and
-   `speedup_e2e` climbing from 0.44× toward the 1.0× ceiling
-   established by 4.1's launch-overhead savings alone.
-5. If still <1×, layer on 4.2 (draft CUDA graph) and 4.3 (adaptive γ).
-6. If still <1× after 4.1–4.3, the arithmetic in §2 suggests this
-   draft/target pair is fundamentally not favourable. At that point
-   Phase 25 should look at (a) smaller draft, or (b) drafting via
-   n-gram / prompt-lookup speculators.

docs/25-speculative-methods-comparison.md

@@ -1,300 +0,0 @@
-# Phase 25: 三种投机解码方法对比 — Small Model / EAGLE / MTP
-
-> 目标：把 speculative decoding 三种主流范式（本项目已试过一种，另两种未实现）
-> 讲清楚，并把 EAGLE3-Qwen3-8B 的实际权重结构展开来看。
-
-## 1. 为什么需要多种范式
-
-Speculative decoding 的核心公式：
-
-```
-speedup = tokens_generated / target_forward_passes
-        ≈ (1 + α·γ) / (1 + draft_cost/verify_cost)
-```
-
-- `α` = acceptance rate（draft 每 token 被接受的概率）
-- `γ` = draft window size（每轮生成的 draft 数）
-- `draft_cost / verify_cost` = draft 一次前向 vs target 一次前向的耗时比
-
-**要 `speedup > 1`，两条路**：把 `α·γ` 做大，或把 `draft_cost/verify_cost` 做小。
-三种范式的本质区别就是**在这两个变量上的取舍**：
-
-| 范式 | draft 模型 | draft cost | α (Qwen3) | 需要训练 | 目标模型是否要改 |
-|------|-----------|-----------|-----------|---------|-------------------|
-| Small-Model | 独立小 LM | 中 (~20% target) | 40% (γ=4) | 无 | 无 |
-| EAGLE (1/2/3) | 1-layer head 读 target hidden | 低 (~10%) | 70%+ (γ=6+) | 蒸馏训练 | 无 (推理路径加 hook) |
-| MTP | target 内嵌多 head | 极低 (∈ target 前向) | 70%+ | 预训练时就要有 | 是（架构层面就是这样的） |
-
-**结论**：
-- Small-Model 是 v0，配置最简单但天花板低。
-- **EAGLE3 是当前性价比最高的落地方案**：draft cost 极低，α 高，需要一次蒸馏训练（约 100k tokens 数据），但对目标模型无侵入。
-- MTP 是 DeepSeek-V3 / DeepSeek-R1 那种"模型天生就懂"的方案，加速比最高但**必须在预训练时就设计进去**，无法事后加装到 Qwen3。
-
----
-
-## 2. Small-Model Speculative（本项目 Phase 22-24 已实现）
-
-### 结构
-
-- **Draft**: 独立的、小得多的同族 LM。要求：**tokenizer 完全一致**（vocab 也一致）。
-- **Verify**: target 用 batched forward 一次算 γ 个位置的 logits，从左往右比较
-  `draft_tokens[i] == target_argmax[i]`，接受最长匹配前缀。
-
-### 算法伪代码
-
-```python
-for _ in gen_tokens:
-    round_start = len(committed)
-    # 1. draft γ steps
-    draft_tokens = [draft.decode(prev) for _ in range(gamma)]
-    # 2. target verify all γ positions in one forward
-    verify_logits = target.forward(committed + draft_tokens[:γ])
-    # 3. accept longest matching prefix
-    accepted = 0
-    while accepted < γ and draft_tokens[accepted] == argmax(verify_logits[accepted-1]):
-        accepted += 1
-    # 4. correction: use target's answer as the next token
-    correction = argmax(verify_logits[accepted-1] if accepted>0 else prev_target_logits)
-    committed.extend(draft_tokens[:accepted] + [correction])
-```
-
-### 优点
-
-- **零训练**。任何同 tokenizer 的两个 LM 组合都能跑。
-- 语义正确性直接保证：只要 accept 逻辑严格，输出等价于纯 target greedy。
-- 代码简单，是理解 speculative decoding 最好的教学入口。
-
-### 缺点（本项目实测在 dash5 上）
-
-Qwen3-0.6B / Qwen3-8B 组合：
-
-| γ | acceptance | speedup_e2e |
-|---|---|---|
-| 1 | 66.5% | 0.57× |
-| 4 | 40.3% | 0.49× |
-| 8 | 25.1% | 0.36× |
-
-即使加上：deterministic gemv/attention、batched GEMV verify kernel、
-Qwen3 whole-step CUDA graph for draft，**仍然 speedup < 1**。
-
-根本原因两点：
-1. **Draft 太贵**：0.6B 一次 decode ~ 2.5 ms，target 8B ~ 12 ms → draft/verify ≈ 20%。
-   γ=4 时，draft 4×2.5=10 ms 单独就占了 verify (13 ms) 的 77%。
-2. **Draft 太蠢**：只用 next-token cross-entropy 训练的独立小模型，
-   跟 target 的 top-1 一致率不高，α 快速衰减（γ=4 → 40%）。
-
-理论上限（假设 verify 免费）：`speedup ≤ (1 + α·γ) ≈ 2.6×`。
-实际上 verify 花掉了绝大部分预算，跑到 0.5× 就到头了。
-
-### 什么时候能赢？
-
-只有当 `draft_cost / verify_cost < acceptance_rate` 时才可能 >1×。
-Qwen3-0.6B 的 draft_cost 太高，需要 draft 是 target 的 **~1/40** 才行
-（8B target 需要 ~200M draft）。Qwen3 没有官方 200M 的成员。
-
----
-
-## 3. EAGLE3（本 Phase 要做的方案）
-
-### 3.1 一句话概括
-
-**EAGLE3 = 用 target 自己的 hidden states 当作 draft 的输入**，
-draft 头只有 1 层 decoder + 1 个 FC 融合层，参数量 ~750M（vs Qwen3-0.6B 的 1.2 GB），
-且更重要的是：draft 前向**不需要重跑 embedding、不需要多层 attention 累积**，
-成本大约是 target 一次 decode 的 **~1/10**。
-
-### 3.2 权重结构（dash5 上下载的 `AngelSlim/Qwen3-8B_eagle3` 实测）
-
-```
-d2t: (32000,)                   int64    # 每个 draft-vocab id → 加多少变成 target-vocab id
-t2d: (151936,)                  bool     # target-vocab id 是否在 draft 频繁词表中
-midlayer.self_attn.q_proj.weight: (4096, 8192)  bf16
-midlayer.self_attn.k_proj.weight: (1024, 8192)  bf16
-midlayer.self_attn.v_proj.weight: (1024, 8192)  bf16
-midlayer.self_attn.o_proj.weight: (4096, 4096)  bf16
-midlayer.mlp.gate_proj.weight:    (12288, 4096) bf16
-midlayer.mlp.up_proj.weight:      (12288, 4096) bf16
-midlayer.mlp.down_proj.weight:    (4096, 12288) bf16
-midlayer.hidden_norm.weight:      (4096,)       bf16   # 融合特征的 pre-attn norm
-midlayer.input_layernorm.weight:  (4096,)       bf16   # draft 嵌入的 pre-attn norm
-midlayer.post_attention_layernorm.weight: (4096,)  bf16
-norm.weight:                      (4096,)       bf16
-fc.weight:                        (4096, 12288) bf16   # 3×hidden → hidden fusion
-lm_head.weight:                   (32000, 4096) bf16   # 输出 draft-vocab
-```
-
-**关键观察**：
-- `fc.weight (4096, 12288)`：**输入是 target 三个不同层的 hidden state 拼起来**
-  （low + mid + high level），一次 FC 融合成 EAGLE 内部的 hidden dim。这是 EAGLE3
-  跟 EAGLE1/2 最大的区别（前两代只用 target 最后一层）。
-- `q_proj.weight (4096, 8192)`：**8192 = 4096 × 2**。attention 输入是
-  `concat(embed(draft_token), fused_target_hidden)`，两个 4096 拼起来。
-  也就是每次预测下一个 token 时，"prompt" 是"上一个 draft token 的 embedding"
-  +"target 对上一个位置的隐状态"。
-- `lm_head.weight (32000, 4096)`：**只输出 32000 个高频 token**（vs target 的 151936）。
-  预测出的 draft-vocab id 用 `d2t` 表查得到真实 target-vocab id：
-  `real_id = draft_id + d2t[draft_id]`。这一步把 lm_head 从 622 MB 压到 131 MB。
-
-### 3.3 推理时的数据流
-
-```
-target 前向（正常执行）：
-    tokens t_0..t_n
-    → embed → layer0 → layer1 → ... → layer35 → norm → logits
-                 ↓        ↓          ↓
-              h_low   h_mid      h_high        （在特定层 hook 出来）
-    logits → sample → t_{n+1}
-
-EAGLE draft γ 步：
-    输入：三个 hidden state h_low[n], h_mid[n], h_high[n] （target 已经算好了）
-    输入：t_{n+1} （target 刚采样出来的下一个 token）
-
-    for k in 0..γ:
-        fused_h = fc(concat(h_low[n+k], h_mid[n+k], h_high[n+k]))     # 4096
-        emb = embed_tokens(t_{n+k+1})                                  # 4096
-        # 这里 embed_tokens 和 target 共享（EAGLE 不重复存 embedding）
-        x_attn_in = concat(embed_norm(emb), hidden_norm(fused_h))     # 8192
-        x = self_attn(x_attn_in) + emb                                 # residual is emb
-        x = mlp(post_norm(x)) + x
-        x = norm(x)
-        draft_logits_small = lm_head(x)                                # 32000
-        draft_id_small = argmax(draft_logits_small)
-        t_{n+k+2} = draft_id_small + d2t[draft_id_small]              # → target vocab
-
-        # 关键：EAGLE 自己会预测下一步的 hidden state 逼近
-        # target 在该位置的 hidden state，供下一 draft 步用。
-        h_low[n+k+1] = h_mid[n+k+1] = h_high[n+k+1] = x
-```
-
-**为什么快？**
-1. 只有 1 层 decoder（vs Qwen3-0.6B 的 28 层）。
-2. 每步计算量 = `attn(hidden=4096, kv_heads=8) + mlp(intermediate=12288) + lm_head(V=32000)`
-   ≈ 1 层 Qwen3-8B decoder + 一个小 lm_head。整个 draft 步 ≈ target 单层 forward + 半个 lm_head，
-   远小于 target 完整 forward。
-3. Draft 的 KV cache 也只有 1 层（vs 28 或 36）。
-4. Embedding 表复用 target 的（不重复算）。
-
-**Acceptance rate 高的原因**：draft 直接使用了 target 的隐状态，
-不是"用另一个小模型独立猜"，α 通常 ≥70%。
-
-### 3.4 与本项目现有 speculative 架构的集成点
-
-保留 Phase 22-24 的所有状态机（verify + accept-reject + correction），
-**只把 draft 换成 EAGLE3 head**。API 契约：
-
-```rust
-// 现在 (Qwen3 draft)
-let draft_logits = draft_decoder.decode(&draft, &[token], &[pos], &[slot], draft_cache);
-let draft_next = last_argmax(&draft_logits);
-
-// EAGLE3 draft
-let draft_logits = eagle.step(&target_hidden_low, &target_hidden_mid, &target_hidden_high, token, pos);
-let draft_next_small = last_argmax(&draft_logits);
-let draft_next = draft_next_small + eagle.d2t[draft_next_small as usize];
-```
-
-**新增到 xserv 的东西**：
-1. Target 侧：改造 `Qwen3::decode_core` 让它在特定 3 层（比如 1/3、2/3、末层的
-   `post_attention_layernorm` 之后）把 hidden state export 出来。
-2. 新模块 `eagle3.rs`：加载 `AngelSlim/Qwen3-8B_eagle3` 权重，暴露 `step()` 方法。
-3. `bench-speculative` 增加 `--drafter eagle3` 分支，draft 改用 EAGLE head。
-
-**不变的东西**：verify path、accept-reject 逻辑、near-tie fallback、CUDA graph
-框架、matched=true 的正确性验证。
-
-### 3.5 Acceptance 上限
-
-按 EAGLE3 paper 的报告，Qwen3-8B 上 γ=6 acceptance ≈ 0.75，speedup 通常 2-3×。
-本项目实测目标：`speedup_e2e > 1` 是保底，`> 2` 是 stretch goal。
-
----
-
-## 4. Multi-Token Prediction (MTP)
-
-### 4.1 一句话概括
-
-**MTP = 在 target 模型的最后加 N 个"预测未来第 k 步"的 head**，
-每个 head 都在预训练阶段和主 head 一起联合训练。推理时这些 head 天然可以并行
-生成 γ 个 draft，然后主 head 一次前向验证。
-
-代表实现：**DeepSeek-V3/R1、Meta MTP 论文（Gloeckle et al., 2024）**。
-
-### 4.2 架构
-
-DeepSeek-V3 的做法：
-
-```
-       [ target 主 decoder，61 层 ]
-                    ↓
-         final hidden h  (2048)
-       /              \
-      main_head        MTP_head_1
-      (predict t_{n+1}) (predict t_{n+2}
-                          given h and t_{n+1})
-```
-
-- 每个 MTP_head 是**一个完整的 transformer block** + linear head（含 embedding
-  proj + attention + MLP）。
-- 训练时：MTP_head_k 的 target 是 `t_{n+k+1}`，loss 加权求和（DeepSeek-V3 训练时权重 0.3）。
-- 推理时：main_head 得到 `t_{n+1}` 后，用 MTP_head_1 得到 `t_{n+2}`（作为 draft），
-  可以级联 MTP_head_2 得到 `t_{n+3}`……然后 target 主前向一次性验证。
-
-**DeepSeek-V3 论文**（arxiv 2412.19437）报告：
-- MTP module 1 层，depth=1，参数占总模型 ~2%。
-- MTP accept rate ≈ 85-90%。
-- 端到端 tps 提升 1.8×。
-
-### 4.3 与 EAGLE 的对比
-
-| 维度 | EAGLE3 | MTP |
-|-----|--------|-----|
-| 加装时机 | 蒸馏训练（一天量级 GPU-hour） | 必须预训练时就设计进去 |
-| Draft 模型独立性 | 独立文件，target 不用改 | 是 target 的一部分 |
-| 深度 | 递归自回归，可 γ=6+ | 通常最多深度 = MTP 头数 (DeepSeek=1) |
-| 训练开销 | 蒸馏，用 target 输出当监督 | 预训练时加多任务 loss |
-| 落地到 Qwen3 | 已有开源权重可直接用 | 需要重新预训练，不可行 |
-
-### 4.4 为什么我们不做 MTP
-
-- Qwen3-8B 没有预训练的 MTP head。要 MTP 就得**自己重新预训练 Qwen3**，不现实。
-- 若要用现成 MTP，只能换到 DeepSeek-V3 这种自带 MTP 的模型；那对整个 xserv 目标
-  (Qwen3 + gpt-oss serving) 是绕道。
-
----
-
-## 5. 三者选型表
-
-给未来的自己或读者一个简明选型：
-
-| 场景 | 选谁 |
-|-----|-----|
-| 已有小同族模型，想快速验证 spec framework | Small-Model（本项目 Phase 22-24） |
-| 已有 target 模型，希望加速但不想改 target 训练 | **EAGLE3**（如有开源 head） |
-| 有充足资源自己预训练一个新 target | MTP（内嵌，加速比最高） |
-| 目标模型是 DeepSeek-V3/R1 | 用它自带的 MTP head |
-| 目标模型是 Qwen3 / LLaMA / GPT-OSS | 找 EAGLE3 蒸馏权重（本 Phase 走这条） |
-
----
-
-## 6. 本 Phase 的实施计划
-
-1. **写这份文档**（正在做）。
-2. **`xserv-model` 新增 `eagle3.rs`**：定义 `Eagle3Head` 结构，加载
-   `AngelSlim/Qwen3-8B_eagle3` 权重。
-3. **修改 `Qwen3::decode_core`**：在 3 个位置 hook hidden state（用 usize const
-   `EAGLE_LOW_LAYER`, `EAGLE_MID_LAYER`, `EAGLE_HIGH_LAYER`；对 36 层默认 12/24/35）。
-   返回值改成 `(Tensor, Option<[Tensor; 3]>)`，第二个 tuple 只在开启 eagle 时填。
-4. **新增 `Eagle3Head::step(hidden_states, token, pos) -> Tensor`**：一层 attention+
-   MLP + lm_head，输出 draft-vocab logits，caller 做 d2t 映射。EAGLE 自己也有
-   一个 1-层的 KV cache（每轮 spec 结束时清空）。
-5. **`bench-speculative` 加 `--drafter [qwen3|eagle3]` 开关**。EAGLE 分支复用现有
-   verify+accept 逻辑，只替换 draft 环节。
-6. **γ 扫**：预期 γ=6 时 acceptance > 0.7、speedup_e2e > 1.5×。
-
-## Sources
-
-- EAGLE-3 paper (arxiv 2503.01840): "Scaling up Inference Acceleration of Large Language Models via Training-time Test"
-- SafeAILab/EAGLE GitHub: reference implementation
-- AngelSlim/Qwen3-8B_eagle3 on ModelScope/HuggingFace: pre-trained head we're using
-- DeepSeek-V3 Technical Report (arxiv 2412.19437): MTP architecture
-- Gloeckle et al. 2024 "Better & Faster Large Language Models via Multi-token Prediction"

docs/26-eagle3-bug-hunt.md

@@ -1,300 +0,0 @@
-# Phase 26: EAGLE3 Implementation Follow-up & Bug Hunt
-
-> Companion to docs/25 (which explains the three speculative paradigms).
-> This doc records the actual EAGLE3 implementation, the bugs we found,
-> the fixes, and why `speedup > 1` remains out of reach.
-
-## Implementation Timeline
-
-Commits are on `main`:
-
-1. **`e04a8ff`** — Eagle3Head module + decode_core_with_hidden hook mechanism +
-   check-eagle3 sanity binary. Weights load; top-5 predictions are
-   thematically coherent (Paris/Tokyo/Madrid for "capital of France is").
-2. **`8f11d6e`** — Fixed EAGLE_HOOK_LAYERS from equally-spaced `[11, 23, 35]`
-   to `[2, 18, 33]` (from vLLM speculators' training config for Qwen3-8B).
-3. **`68b55fa`** — First bench-eagle3 γ=1 loop. matched=true but acceptance
-   only 1.3%.
-4. **`a24621f`** — Residual chain fix + stateful KV cache: acceptance jumps
-   to 20% at γ=1.
-5. **`1492515`** — γ≥2 scaffolding: `step_with_aux` + `step_recursive` +
-   `forward_verify_paged_decode_attention_with_hidden`. matched=false at
-   γ≥2 due to K/V bugs.
-6. **`d2c55c4`** — γ≥2 correctness fixes: matched=true across full sweep.
-
-## Bugs Fixed (γ≥2)
-
-### Bug A: Truncate dropped needed K/V
-
-Old code:
-```rust
-cache.truncate_sequence(slot, round_pos - 1).unwrap();
-let (verify_logits, _) = target.forward_verify_...(&[prev_token, d0, d1], ...);
-```
-
-`round_pos - 1` was the position where the last committed token
-(`pending_prev`) lived. Truncating dropped its K/V. Then verify wrote
-`prev_token` at that slot AGAIN, but this is a DIFFERENT bit pattern —
-the previous single-token decode wrote via `matmul_2d` (m=1 → custom
-GEMV) while verify wrote via `matmul_batched_gemv` (m=γ+1). Same math,
-same output bytes... IN PRINCIPLE. But re-writing K/V that was already
-there introduces a small numerical drift.
-
-**Fix**: Don't truncate. Let verify start at `cache.seq_len` and write
-γ+1 new positions forward. `pending_prev`'s K/V stays intact from the
-previous round's write.
-
-### Bug B: EAGLE cache accumulated rejected drafts
-
-Each EAGLE `step_with_aux` or `step_recursive` writes one K/V entry to
-EAGLE's internal cache. Per round we call it γ times (once with the
-target hooks, γ-1 times recursively). All γ writes happen regardless of
-how many drafts are eventually accepted.
-
-If `k < γ` drafts accepted, EAGLE's cache has γ entries for a round
-that committed only k+1 tokens (pending_prev + k drafts). The extra
-γ-k-1 entries hold K/V for hallucinated drafts that never got
-committed — polluting future rounds.
-
-**Fix**: Add `Eagle3Head::truncate_to(new_len)`. After acceptance,
-truncate to `eagle_len_before + k + 1`.
-
-### Bug C: aux output was normed, should be pre-norm
-
-vLLM's `llama_eagle3.py` (line ~150):
-```python
-hidden_states, hidden_prenorm = self.norm(hidden_states, residual)
-aux_output = hidden_states if self.norm_output else hidden_prenorm
-```
-
-Default `norm_output=False` → aux = hidden_prenorm (pre-RMSNorm
-residual sum). I was returning `hidden_states` (normed).
-
-**Fix**: return the second output of `add_rmsnorm`, which is `x + residual`
-(pre-norm). Small effect on acceptance (~1%).
-
-### Bug D: EAGLE draft position off-by-one
-
-`pending_prev` is at target position `p`. EAGLE step 0 should compute
-RoPE at position `p` (matching pending_prev's target position). I was
-passing `p + 1`.
-
-**Fix**: pass `p + k` for the k-th EAGLE step (k = 0..γ-1).
-
-## Final Measurements
-
-Setup: dash5 (RTX 5090), Qwen3-8B target + AngelSlim/Qwen3-8B_eagle3 head,
-5 prompts × 32 tokens, greedy, matched=true across all runs.
-
-| γ | acceptance | verify_cost (× single decode) | speedup_e2e |
-|---|------------|-------------------------------|-------------|
-| 1 (single-decode verify) | 22.7% | 1.00 | **0.95×** |
-| 1 (batched verify)       | 20.6% | ~1.5 | 0.75× |
-| 2 | 12.6% | ~1.7 | 0.59× |
-| 3 |  9.1% | ~2.1 | 0.48× |
-| 4 |  7.6% | ~2.4 | 0.41× |
-| 6 |  5.2% | ~3.1 | 0.32× |
-| 8 |  4.1% | ~3.7 | 0.27× |
-
-Per-slot diagnostic (γ=8, aggregated over 5 prompts):
-```
-d[0]=12/125(0.10)  d[1]=8/122(0.07)   d[2]=5/119(0.04)
-d[3]=6/116(0.05)   d[4]=8/113(0.07)   d[5]=13/110(0.12)
-d[6]=17/107(0.16)  d[7]=17/104(0.16)
-```
-
-Later positions (d[5..7]) surprisingly show HIGHER acceptance than d[1..3].
-Explanation: once EAGLE hallucinates its own chain, target's `verify_argmax`
-follows that hallucinated context and often converges to plausible common
-tokens (spaces, commas, "the"). This helps per-slot rate but not
-longest-prefix acceptance (first mismatch kills the whole tail).
-
-## Why speedup < 1
-
-The speedup formula:
-```
-speedup ≈ (1 + avg_accepted_per_round) / verify_cost_relative_to_single_decode
-```
-
-Sub-1 across the sweep because:
-
-- **verify_cost grows linearly with γ+1**. Each verify slot is one BF16 GEMV
-  row across all Qwen3-8B layers. Batching gets some memory-bound sharing
-  but not enough to make γ+1 slots free.
-- **avg_accepted per round grows only sub-linearly** because acceptance rate
-  degrades at later chain positions (~half every 2 steps).
-
-To reach `speedup > 1` we need avg_accepted > (verify_cost - 1). With
-verify_cost ≈ 1.7 at γ=2, need avg_accepted > 0.7. Observed 0.25.
-
-## Path Forward
-
-Three levers, all significant work:
-
-### 1. Tree-based drafting (biggest lever, +2-3× acceptance)
-
-EAGLE-3 paper reports 60-70% acceptance using TREE decoding: at each
-recursive step, EAGLE proposes top-k candidates instead of top-1. The
-target's verify then evaluates all tree branches in one forward using
-paged attention with tree-aware masking.
-
-Reference: `SafeAILab/EAGLE` uses trees with depth 6 and 26+ nodes.
-
-Implementation cost: significant. Requires:
-- Tree-aware batched verify (multi-branch attention masking).
-- Tree navigation / longest-accepted-path selection.
-- KV cache management for accepted branch vs discarded branches.
-
-### 2. Cheaper batched verify
-
-Current batched verify at γ+1 tokens uses `matmul_batched_gemv` (per-row
-GEMV) plus `paged_decode_attention` batch=γ+1. Both scale roughly
-linearly with γ+1.
-
-Potential improvements:
-- **Flash Attention** with multi-query: each of the γ+1 queries shares
-  the same K/V cache pointers, so a single kernel can read K/V once and
-  compute γ+1 outputs. Currently they're independent kernel launches per
-  query.
-- **Cheaper QKV projection at m>1**: matmul_batched_gemv is bit-exact
-  per row but doesn't amortize K/V loading across rows. Could use cuBLAS
-  GEMM at m=γ+1 (faster but different BF16 rounding → K/V drift).
-
-### 3. Better draft (smaller EAGLE, different training)
-
-The AngelSlim Qwen3-8B_eagle3 head is 750MB (~1 layer of the 8B model).
-Alternatives:
-- Smaller Qwen3 (0.6B) as draft: already tried, γ=1 gets 40% acceptance
-  but draft cost ~2.5ms (vs EAGLE's ~0.5ms).
-- Different EAGLE weights: `Zjcxy-SmartAI/Eagle3-Qwen3-8B-zh` (Chinese-
-  tuned), or train our own with tree-time supervision.
-
-## Recommendation
-
-Given effort/reward:
-
-**Short-term (1 session)**: implement tree-based drafting with depth=2,
-width=2 (4 candidates per round). Reuse existing batched verify with
-tree-aware masking. Expect acceptance to double (25% → 50%+).
-
-**Medium-term (2-3 sessions)**: fully tree of depth=6, width=varying, +
-flash-attention-2 batched verify kernel. This matches the vLLM
-implementation and should approach 2× speedup.
-
-**Alternative (if EAGLE is a dead-end)**: switch to lookahead decoding
-(Yaniv Leviathan-style) which doesn't require a draft model at all —
-uses n-gram lookup + Jacobi iteration on the target.
-
-The infrastructure to enable this (Eagle3Head, batched verify, cache
-truncation, position management) is now solid on `main`. What's missing
-is the tree-aware acceptance algorithm and possibly a faster verify
-kernel.
-
----
-
-## Epilogue (`06a798c`): cuBLAS GEMM verify → speedup > 1 achieved
-
-Actioned option 2 above: swapped `matmul_batched_gemv` for `matmul_2d`
-(cuBLAS GEMM) inside `forward_verify_paged_decode_attention_with_hidden`.
-
-Micro-benchmark (bench-verify-cost.rs, RTX 5090, prompt_len=100):
-
-| batch | batched-GEMV verify | cuBLAS-GEMM verify |
-|-------|---------------------|--------------------|
-| 1     | 13.14 ms (1.05×)    | 13.04 ms (1.04×)   |
-| 2     | 19.51 ms (1.56×)    | 13.52 ms (1.08×)   |
-| 3     | 26.10 ms (2.09×)    | 13.59 ms (1.09×)   |
-| 5     | 38.72 ms (3.10×)    | 13.88 ms (1.11×)   |
-| 9     | 64.15 ms (5.14×)    | 15.03 ms (1.20×)   |
-
-cuBLAS GEMM at m>1 amortizes K/V load across all queries, giving
-near-flat scaling (compute-bound). GEMV loads K/V per row → linear.
-
-50 prompts × 64 tokens γ sweep with cuBLAS verify:
-
-| γ | acceptance | speedup_e2e |
-|---|------------|-------------|
-| 1 (single-decode) | 29.8% | 0.95× |
-| **2**             | **16.9%** | **1.10×** ← best |
-| 3                 | 11.6% | 1.06× |
-| 4                 |  8.9% | 1.02× |
-| 5                 |  7.2% | 0.96× |
-| 6                 |  6.0% | 0.93× |
-| 8                 |  4.5% | 0.86× |
-
-Tradeoff: `matched=false`. cuBLAS GEMM at m>1 rounds BF16 differently
-from custom GEMV at m=1. K/V bytes written by verify differ from what
-a per-token decode would write, and downstream token choices diverge
-from the strict-baseline path.
-
-The spec output is still a VALID target output (still coherent English,
-still target-model semantics), just via a slightly different numerical
-approximation path. This is the industry norm for "lossless spec
-decoding": distribution preserved modulo BF16 rounding, not bit-exact
-with a specific numerical path.
-
-`speedup_e2e = 1.10×` is a real, measurable win at γ=2 on 50×64 prompts.
-Higher γ gives diminishing returns because acceptance drops faster than
-verify saves (already max at γ=2). To push higher, we'd need better
-draft (tree decoding, larger EAGLE head, or different EAGLE weights).
-
----
-
-## Epilogue 2 (`fd392f7`): Tree attention kernel + why tree drafting is stuck
-
-Wrote the tree-aware paged decode attention kernel:
-`paged_decode_attention_tree_bf16_kernel` takes an extra `[batch, batch]`
-i32 mask that lets each query select which of the newly-written K/V
-rows it attends to. Positions before `tree_start` always attended.
-
-Rust wrapper `paged_decode_attention_tree` + forward variant
-`Qwen3::forward_verify_paged_decode_attention_tree_with_hidden` (takes
-explicit positions, kv_lens, tree_mask) all landed.
-
-Sanity check: bench-eagle3's γ_multi verify path was switched to route
-through the tree kernel with a causal mask. matched=false pattern
-identical, acceptance ~identical, speedup within noise of the non-tree
-version. Kernel is correct.
-
-### The blocker: KV cache position rigidity
-
-Wrote out the top-2 sibling tree structure on paper. Discovered a
-fundamental issue: the paged K/V cache stores K/V at physical positions
-that are 1-to-1 with target positions. If verify writes 4 K/V rows at
-cache positions `[P, P+1, P+2, P+3]` corresponding to
-`[pending_prev, d0_top1, d0_top2, d1_chain_from_top1]`, then:
-
-- If `d0_top1` accepted: its K/V is at physical slot P+1, matching
-  target position P+1. Continuing decode from position P+1 reads the
-  right K/V. ✓
-- If `d0_top2` accepted: its K/V is at physical slot P+2, but its
-  semantic target position is P+1. Continuing decode from target
-  position P+2 would look at physical slot P+2 and read d0_top2's K/V —
-  but semantically, position P+1 should have d0_top2's K/V, and position
-  P+2 should have whatever comes after d0_top2 (unknown). Continuing
-  decode reads the wrong K/V. ✗
-
-Fixing this requires one of:
-1. **KV slot remap on acceptance**: physically copy d0_top2's K/V from
-   slot P+2 to slot P+1 across all layers. Costs one full-layer memcpy
-   per acceptance of a non-top-1 sibling. Doable but adds ~2ms per event.
-2. **Virtual-position paged cache**: introduce a per-slot position
-   translation table so K/V at physical slot X has logical position Y.
-   Requires modifying every attention kernel to consult this table
-   (invasive).
-3. **Restart top-2 branches from a decode**: if top-2 accepted, discard
-   the tree K/V past pending_prev and run a full single-token target
-   decode with d0_top2 to properly write its K/V at target position P+1.
-   Costs ~1 full decode per accepted top-2, which likely eats the win.
-
-Given (1) is the least invasive but still complex, and (3) may not net
-positive speedup, this exceeds a single-session scope.
-
-**Concluding numbers on xserv main**:
-- Best speedup: **1.10×** at γ=2 (cuBLAS-GEMM verify, no tree).
-- Tree kernel + wrapper ready and correctness-verified.
-- Full tree drafting requires KV remap work (Phase 27+ scope).
-
-Everything lands cleanly on `main`. Any future session can start from
-the tree kernel and implement the KV remap; the correctness harness is
-in place (matched=true after remap = success criterion).

docs/27-speculative-quality-gsm8k.md

@@ -1,177 +0,0 @@
-# Phase 27 — Speculative Decoding Quality: Task-Level Correctness at Scale
-
-**Goal**: prove tree-drafting speculative decoding preserves output quality
-**despite** batched-verify BF16 rounding differences (`matched=false` on
-token-by-token comparison).
-
-## TL;DR
-
-| Suite | N | baseline_acc | spec_acc | agreement | tpot base→spec | **speedup** |
-|-------|---|:-----------:|:--------:|:---------:|:--------------:|:-----------:|
-| GSM8K | 1000 | 93.50% | 93.30% | 97.50% | 13.33 → 8.97 ms | **1.486×** |
-| AIME2025 | 30 | 16.67% | 13.33% | 23.33% | 17.18 → 11.64 ms | **1.475×** |
-
-- **Speedup is model+workload driven, not accuracy-driven** — the same
-  1.47-1.49× shows up on high-accuracy chat math (GSM8K) and on saturated
-  long-reasoning math the model can't actually solve (AIME).
-- **GSM8K**: on 1000 problems, spec accuracy is within 0.2 pp of baseline
-  (933 vs 935 correct). Where the two disagree (25 of 1000): baseline wins
-  9 times, spec wins 7 times, they're both wrong 9 times. Net effect on
-  aggregate accuracy is a wash.
-- **AIME**: at 8B params Qwen3 is far below the accuracy floor (16.67% =
-  5/30). Divergences here reflect the fact that both trajectories are
-  wandering through low-probability sequences; agreement drops to 23% but
-  spec is only 1 problem behind baseline.
-
-## Why AIME agreement is low but speedup unchanged
-
-AIME2025 pushes Qwen3-8B way outside its competence. Both baseline and spec
-generate long, meandering, often-wrong reasoning; small BF16 rounding
-differences in tree-verify snowball across ~2000 gen-tokens into completely
-different (still-wrong) answers. This is expected: when the target
-distribution has no dominant mode, top-1 argmax is dictated by noise,
-and any batched-verify rounding will flip it.
-
-Crucially, `speedup_e2e = 1.475×` on AIME matches `1.486×` on GSM8K to
-within ~1%. The wall-clock benefit does not depend on the task being
-solvable — it depends on EAGLE3 draft quality (which stays ~21% on both
-suites) and the batched-verify cost model.
-
-## How the test was run
-
-Extended `bench-eagle3` (from Phase 27) accepts any JSON file with the
-`{id, problem, answer}` schema. Same binary → same code paths.
-
-```bash
-# GSM8K — 1000 problems, gen_tokens=512, max_seq_len=1024
-./target/release/bench-eagle3 \
-    /opt/wjh/models/qwen3-8b \
-    /dashscope-tmp/wjh/models/qwen3-8b-eagle3 \
-    --gsm8k tools/bench/data/gsm8k.json \
-    --tree --prompts 1000 --gen-tokens 512 --max-seq-len 1024
-
-# AIME2025 — 30 problems, gen_tokens=2048, max_seq_len=4096
-./target/release/bench-eagle3 \
-    /opt/wjh/models/qwen3-8b \
-    /dashscope-tmp/wjh/models/qwen3-8b-eagle3 \
-    --gsm8k tools/bench/data/aime2025.json \
-    --tree --prompts 30 --gen-tokens 2048 --max-seq-len 4096
-```
-
-Chat template used (`build_chat_prompt`, math-solver system prompt):
-```
-<|im_start|>system
-You are a careful math problem solver. Solve the problem step by step. Put your final numeric answer inside \boxed{}.
-<|im_end|>
-<|im_start|>user
-{problem}
-<|im_end|>
-<|im_start|>assistant
-<think>
-
-</think>
-
-```
-
-## GSM8K result (1000 problems)
-
-```
---- SUMMARY ---
-prompts=1000 matched=false
-acceptance_rate=0.2120 accepted=125326 proposed=591156 target_steps=149789
-baseline_tpot_ms=13.331 baseline_tok_s=75.013
-spec_tpot_ms=8.971 spec_tok_s=111.474 speedup_e2e=1.4861
-gsm8k: baseline_acc=0.9350 (935/1000) spec_acc=0.9330 (933/1000) agreement=0.9750 (975/1000)
-```
-
-Disagreement analysis (25/1000 questions where extracted answers differ):
-- baseline correct, spec wrong: **9**
-- spec correct, baseline wrong: **7**
-- both wrong (different wrong answers): **9**
-
-The counts are essentially symmetric — spec is not systematically worse.
-
-## AIME2025 result (30 problems, 2048 gen-tokens)
-
-```
---- SUMMARY ---
-prompts=30 matched=false
-acceptance_rate=0.2034 accepted=23511 proposed=115596 target_steps=28959
-baseline_tpot_ms=17.177 baseline_tok_s=58.219
-spec_tpot_ms=11.642 spec_tok_s=85.896 speedup_e2e=1.4754
-gsm8k: baseline_acc=0.1667 (5/30) spec_acc=0.1333 (4/30) agreement=0.2333 (7/30)
-```
-
-Note: the label `gsm8k` in the summary line is a hardcoded label — the
-data is AIME2025, wrapped in the same chat template.
-
-Disagreement analysis (23/30 questions differ):
-- baseline correct, spec wrong: 1
-- spec correct, baseline wrong: 0
-- both wrong (different wrong answers): 22
-
-## Absolute performance
-
-| metric | baseline | tree-spec |
-|--------|----------|-----------|
-| GSM8K tpot | 13.33 ms | 8.97 ms |
-| GSM8K tok/s | 75.0 | 111.5 |
-| AIME tpot | 17.18 ms | 11.64 ms |
-| AIME tok/s | 58.2 | 85.9 |
-
-AIME's absolute tpot is higher than GSM8K because average KV length is
-larger (avg completion ~1500 tokens vs ~350 for GSM8K), which slows the
-paged attention kernel roughly linearly. **Both suites see the same relative
-speedup**, confirming EAGLE3 tree-drafting benefits scale with context
-length rather than depending on it.
-
-## Interpretation
-
-The Phase 26 `matched=false` flag has been fully characterized on 1030
-real problems:
-
-1. **On solvable tasks (GSM8K)**: spec accuracy is within noise (Δacc =
-   -0.2 pp on 1000 samples, 95% CI easily includes zero). This is what
-   vLLM and SGLang call "lossless" speculative decoding.
-
-2. **On hard tasks (AIME)**: both baseline and spec meander through wrong
-   answers; agreement collapses because the argmax distribution is nearly
-   flat. Speedup is preserved.
-
-3. **Draft acceptance is the invariant**: acceptance_rate = 21.2% (GSM8K)
-   vs 20.3% (AIME) — nearly identical, because EAGLE3's draft quality
-   depends on target distribution predictability, which is similar for
-   both math-formatted chat prompts.
-
-Speculative decoding is **correctness-preserving in expectation**, not
-bit-exact. This is the same guarantee production systems ship.
-
-## What was NOT changed
-
-- No changes to kernels, attention, KV cache, EAGLE3 head, or the tree
-  drafting policy (still γ=2 top-3 as in commit `2fe903e`).
-- Bench binary already supported `--gsm8k <path>` from commit `264c004`;
-  we simply pointed it at both `gsm8k.json` and `aime2025.json`.
-
-## Files touched
-
-- `docs/27-speculative-quality-gsm8k.md` — rewritten with 1000-scale
-  GSM8K and 30-problem AIME2025 results.
-
-## Reproduction
-
-```bash
-# on dash5 (5090)
-cd /opt/wjh/projects/xserv
-./target/release/bench-eagle3 /opt/wjh/models/qwen3-8b \
-    /dashscope-tmp/wjh/models/qwen3-8b-eagle3 \
-    --gsm8k tools/bench/data/gsm8k.json \
-    --tree --prompts 1000 --gen-tokens 512 --max-seq-len 1024
-# ~90 minutes wall-clock on 5090
-
-./target/release/bench-eagle3 /opt/wjh/models/qwen3-8b \
-    /dashscope-tmp/wjh/models/qwen3-8b-eagle3 \
-    --gsm8k tools/bench/data/aime2025.json \
-    --tree --prompts 30 --gen-tokens 2048 --max-seq-len 4096
-# ~11 minutes wall-clock on 5090
-```

docs/MOE_PROGRESS.md

@@ -0,0 +1,177 @@
+# MoE (gpt-oss-20b) — 工作进度与续作指南
+
+> **中断原因**：用户要重启 dash5 机器（IP 等可能变），让我先把当前 MoE 支持工作的状态
+> 完整记录到本文件，重启后据此继续。本文件是"重启后从这里接着干"的唯一入口。
+
+最后更新：Phase 18 (PP) 已完成并 push；Phase 19 (MoE/gpt-oss-20b) 刚起步（下载受阻，
+架构与参考数学已侦查清楚）。
+
+---
+
+## 0. 一句话现状
+
+- ✅ **Phase 18 流水线并行 (PP)** 全部完成、验证、benchmark，已 commit 并 push 到
+  `origin/phase18-pipeline-parallelism`（gitea）。
+- 🚧 **Phase 19 MoE (gpt-oss-20b)** 刚开始：架构 + HF 参考数学已核对（见
+  `docs/19-moe-gpt-oss.md`），**模型还没下载完**（HF 被墙，正在解决下载路径），代码未动。
+
+---
+
+## 1. 环境关键事实（重启后很可能变 / 需重新确认）
+
+- **本机**（开发机，非 GPU）：`/home/gahow/projects/xserv`，有公网（走代理）。
+- **dash5**（GPU 机，8×RTX 5090，无 NVLink，0-3/4-7 分组）：通过 `ssh dash5` 访问。
+  - 远端仓库目录：`/opt/wjh/projects/xserv`，模型目录：`/opt/wjh/models/`。
+  - **dash5 无外网、无 rsync**；同步用 `./tools/sync-and-build.sh`（tar over ssh）。
+  - cargo 在 `$HOME/.cargo/bin`；CUDA 12.9 在 `/usr/local/cuda-12.9`。
+  - ⚠️ **重启后 `ssh dash5` 的 IP/可达性可能变** —— 先 `ssh dash5 hostname` 确认；
+    若连不上，检查 `~/.ssh/config` 里 `dash*` 配置 / 让用户给新地址。
+- **HTTP 代理**（本机环境变量,重启后可能还在 `/etc/environment` 或 shell）：
+  `http_proxy=https_proxy=all_proxy=http://ipads:ipads123@202.120.40.82:11235`
+- **huggingface.co 被墙**（`SSL_ERROR_SYSCALL`，即使过代理）。pypi 可过代理。
+- **`huggingface_hub` 不是预装**，已用 `pip install --user --break-system-packages
+  huggingface_hub safetensors` 装好（1.17.0）；venv 不可用（无 ensurepip）。
+
+---
+
+## 2. gpt-oss-20b 下载（**当前卡点**）
+
+目标：下到本机 `~/models/gpt-oss-20b`，再 tar-over-ssh 拷到 dash5
+`/opt/wjh/models/gpt-oss-20b`。
+
+**已验证可行的下载路径**（重启后照此做）：
+- huggingface.co 直连/经代理都失败。
+- hf-mirror.com 的 `/resolve/` 会 **308 跳回 huggingface.co**（也被墙）——所以不能用
+  `curl -L` 跟跳转，`huggingface_hub` 设 `HF_ENDPOINT` 在新版(1.17)上 HEAD 也失败。
+- ✅ **能用的办法**：直接走 **hf-mirror 的 `/raw/`（小文件）和实际 CDN，经代理 curl**。
+  已成功取到 `config.json`（200, 1799 bytes）：
+  ```bash
+  curl -s -x "http://ipads:ipads123@202.120.40.82:11235" \
+    "https://hf-mirror.com/openai/gpt-oss-20b/raw/main/config.json" -o config.json
+  ```
+  大文件（safetensors）要用 `/resolve/main/<file>` 且 **指定 `-x` 代理、不要 `-L`**，
+  若仍 308 跳回 hf.co，则改用 hf-mirror 的 LFS 直链或 `huggingface_hub` 配合
+  `HF_ENDPOINT=https://hf-mirror.com` + 代理（库内部不跟 308）。**下载脚本草稿在
+  `/tmp/dl_shards.sh`（重启后 /tmp 会清空，需重建）。**
+
+**待下载文件**（3 个分片 + 元数据，总 ~13.5GB MXFP4）：
+- `model-00000-of-00002.safetensors`、`model-00001-of-00002.safetensors`、
+  `model-00002-of-00002.safetensors`（注意是 0/1/2 三个，命名 of-00002）
+- `model.safetensors.index.json`、`config.json`、`tokenizer.json`、
+  `tokenizer_config.json`、`special_tokens_map.json`、`generation_config.json`、
+  `chat_template.jinja`
+
+**重启后第一步**：`ls -la ~/models/gpt-oss-20b/` 看已下了哪些、`wc -c` 校验分片大小，
+断点续传用 `curl -C -`。
+
+---
+
+## 3. gpt-oss-20b 架构（config.json 已核对）
+
+| 字段 | 值 |
+|------|----|
+| layers | 24；hidden 2880；**head_dim 64**（≠ hidden/heads！）|
+| heads | 64 q-heads / 8 kv-heads（GQA，n_rep=8）|
+| experts | num_local_experts **32**，num_experts_per_tok **4**（top-4）|
+| expert intermediate | 2880 |
+| vocab | 201088；max_pos 131072；tie_embeddings false |
+| rope_theta | 150000（核对是否有 rope_scaling/YaRN）|
+| sliding_window | 128（**交替层**，见 config `layer_types`）|
+| rms_norm_eps | 1e-5；swiglu_limit 7.0；alpha 1.702 |
+| 量化 | **MXFP4**，仅 expert MLP（gate_up/down 的 `_blocks`+`_scales`）；attn/router/embed/lm_head 为 BF16 |
+
+---
+
+## 4. HF 参考数学（已从 transformers `modeling_gpt_oss.py` 逐字核对，务必照抄）
+
+完整版见 `docs/19-moe-gpt-oss.md` §2。要点：
+
+**Router**（softmax 在 topk **之后**）：
+```
+logits = x @ W_router^T + b_router          # [T,32]
+top_val, idx = topk(logits, 4)
+top_val = softmax(top_val)                  # 只对选中的 4 个归一化
+scores = scatter to [T,32] (其余 0)
+```
+
+**Experts**（fused gate_up，**交错** ::2 / 1::2；clamped；(up+1)·glu）：
+```
+alpha=1.702, limit=7.0
+gate_up = x @ gate_up_proj[e] + bias        # [.., 2*2880]
+gate = gate_up[..., ::2]; up = gate_up[..., 1::2]
+gate = clamp(gate, max=limit)               # 仅上界
+up   = clamp(up, min=-limit, max=limit)
+glu  = gate * sigmoid(gate * alpha)
+h    = (up + 1) * glu                        # 注意 (up+1)
+y_e  = h @ down_proj[e] + bias
+out  = Σ_{e∈top4} scores[t,e] * y_e
+```
+
+**Attention（带 sinks）**：
+```
+scaling = 64 ** -0.5；q/k/v/o 都有 bias
+RoPE(theta=150000) on q,k；repeat_kv(n_rep=8)
+attn = (q@k^T)*scaling + causal(+ 滑窗层叠加 banded window=128)
+combined = cat([attn, sinks_per_head], dim=-1)   # 每 head 一个标量 sink，多一列
+combined -= combined.max(-1, keepdim)            # 数值稳定
+probs = softmax(combined, -1)
+scores = probs[..., :-1]                          # 丢掉 sink 列（概率不归一到 1！）
+o = scores @ v -> merge heads -> @Wo + bo
+```
+
+**RMSNorm**：标准（fp32 算 variance，eps=1e-5）。
+
+参考源码已存（重启后 /tmp 清空需重取）：`pip download transformers --no-deps`
+解 wheel 取 `transformers/models/gpt_oss/modeling_gpt_oss.py`（967 行）。
+
+---
+
+## 5. MXFP4 反量化（expert 权重）
+
+- expert 张量名：`model.layers.{i}.mlp.experts.gate_up_proj_blocks` + `..._scales`，
+  `...down_proj_blocks` + `..._scales`（bias 是 BF16 的 `gate_up_proj_bias`/`down_proj_bias`）。
+- MXFP4：每 **32** 元素一 block，共享一个 **E8M0**（8-bit 指数）scale，每元素 4-bit
+  FP4(E2M1，16 码字)。反量化 `val = fp4_lut[code] * 2^(e8m0 - 127)`。
+- **决策（已定）**：加载时在 CPU 反量化成 BF16（dash5 ~1TB 内存），整模型 ~40GB BF16，
+  单卡放不下 → 走 **Phase 18 的 PP**（PP=2 ~20GB/卡，PP=4 ~10GB/卡）。不写 GPU 原生
+  MXFP4 kernel（风险高、慢），先正确跑通+对比，后续再优化。
+
+---
+
+## 6. 实施路线（Phase 19，逐步可验证）
+
+1. **P19.1** Python(numpy) 读 safetensors + MXFP4 反量化，与 HF 一层数值对照（确认 LUT /
+   block 方向 / gate_up 交错对得上）。**不依赖 GPU，重启后可先做。**
+2. **P19.2** `crates/xserv-model/src/config.rs`：加 MoE 字段
+   （num_local_experts / num_experts_per_tok / sliding_window / swiglu_limit /
+   显式 head_dim / expert intermediate），保持 Qwen3 路径不变。
+3. **P19.3** 新文件 `crates/xserv-model/src/gptoss.rs`：dense（attn+sinks+bias+滑窗 /
+   RMSNorm / lm_head）+ MoE FFN（正确优先：逐 token top-4 gather→clamped SwiGLU→加权和）。
+   MXFP4 在 `from_weights` 反量化为 BF16。验收：prefill logits 与 HF BF16 容差内一致。
+4. **P19.4** `from_weights_pp` 支持 gpt-oss（experts 随层切），`--pp` 端到端；
+   PP=2/4 与 PP=1 等价（沿用 Phase 18 的"单卡×2 vs ppN×2"对照法）。注：~40GB 需 PP≥2。
+5. **P19.5** llama.cpp 对比：**pinned submodule b9371 早于 gpt-oss（约 2025-08 落地），
+   需升级 submodule** 到支持 gpt-oss 的版本 + 取/转 GGUF；跑 AIME 2025 + GSM8K，
+   复用 `tools/bench/` + `tools/bench/summarize_fullq.py`（已有，PP 阶段写的）。
+
+---
+
+## 7. 复用 Phase 18 的资产
+
+- 多卡：`--pp N`（已验证），`crates/xserv-distributed`（NCCL P2P + AllReduce）。
+- bench：`tools/bench/runner.py`（支持 `--pp`/`--tp`）、`summarize_fullq.py`、
+  `tools/pp_quality_full.sh`（xserv 0-3 ‖ llama 4-7 并行跑 AIME+GSM8K 的范式可直接改用）。
+- 教训（见全局 memory）：用对 model 名（不是 "q"）；就绪判定用真实生成不是 /health；
+  贪心 run-to-run 不可复现（cuBLAS）；显存快照要等模型加载完；严格串行避免同组 GPU 互扰；
+  长任务用持久前台 ssh + `run_in_background`，别让一个网络失败 cancel 掉整批命令。
+
+---
+
+## 8. 重启后立即要做（checklist）
+
+1. `ssh dash5 hostname` 确认 GPU 机可达（不行就问用户新地址 / 改 ~/.ssh/config）。
+2. `git -C ~/projects/xserv log --oneline -6` 确认 PP 5 个 commit 还在
+   （`859c0cc..` 那串，分支 `phase18-pipeline-parallelism`）。
+3. `ls -la ~/models/gpt-oss-20b/` 看下载进度，续传缺的分片（§2）。
+4. 重新 `pip download transformers` 取参考源码（/tmp 已清）。
+5. 从 §6 的 P19.1 接着干。

docs/TO-BE-FIXED.md

@@ -0,0 +1,214 @@
+# xserv — To Be Fixed (2026-05-23 审查更新)
+
+> 由全面审查产出的修复清单。每项修复有明确验收标准。
+> 优先级: P0 (阻塞可用性) > P1 (严重bug/性能) > P2 (重要改进) > P3 (设计债务)
+
+---
+
+## 第一批：P0 — 阻塞可用性
+
+### FIX-01: 全局 cuBLAS handle [P0-性能] ❌未修
+
+**问题**: `gemm.rs` 中 `matmul` (line 146) 和 `batched_matmul` (line 224) 每次调用都 `CublasContext::new()` 创建+销毁 handle。Qwen3-8B 一次 forward ~252 次 matmul。
+
+**修复要求**:
+- 使用 thread-local 单例 cuBLAS handle
+- handle 生命周期覆盖整个进程
+- `matmul` / `batched_matmul` 函数体内不再有 `CublasContext::new()`
+
+**验收标准**:
+1. `grep -n "CublasContext::new" crates/xserv-kernels/src/gemm.rs` 只出现 1 次（thread_local 初始化处）
+2. 编译通过，现有 gemm_test 全部通过
+
+---
+
+### FIX-16: EOS token 泄漏到 API 响应 [P0-功能] ❌新发现
+
+**问题**: `engine.rs:218` 中 `emit_token` 先发 `GenerateEvent::Token { text: "<|im_end|>" }` 再发 `Done`。`api.rs:110-111` 把所有 Token text 拼到 content 里，导致最终响应包含 `<|im_end|>` 乱码。
+
+**修复要求**:
+- `emit_token` 中，当 token 是 EOS 时，不发送 Token event（或发送空 text），直接发 Done
+- 或者: API 层收到 Done 时丢弃最后一个 token 的 text（如果 finish_reason == "stop"）
+
+**验收标准**:
+1. 发送请求，响应 content 不包含 `<|im_end|>` 或其他 special token 文本
+2. streaming 模式下最后一个 content chunk 不是 EOS 文本
+3. 编译通过
+
+---
+
+### FIX-17: max_seq_len 硬编码 256 [P0-功能] ❌新发现
+
+**问题**: `engine.rs:53` 硬编码 `let max_seq_len = 256`，超过就 KV cache panic。
+
+**修复要求**:
+- `Engine::load` 接受 `max_seq_len` 参数（或从 config 读取，上限为 config.max_seq_len()）
+- `main.rs` 中通过命令行参数或环境变量传入，默认值改为 2048
+- 同步更新 RoPE cache 上限（当前 `qwen3.rs:45` 限制 8192，应与 max_seq_len 一致）
+
+**验收标准**:
+1. `grep -n "let max_seq_len = 256" crates/xserv-server/` 返回 0 行
+2. 启动 server 时 `--max-seq-len 4096` 可用
+3. 编译通过
+
+---
+
+### FIX-18: max_tokens 无上限校验 [P0-功能] ❌新发现
+
+**问题**: API 不校验 `max_tokens`，客户端可发 `max_tokens: 1000000` 导致 KV cache panic。
+
+**修复要求**:
+- `api.rs` 中 clamp `max_tokens` 到 `engine.max_seq_len - prompt_tokens.len()`
+- 如果 prompt 已超过 max_seq_len，返回 400 错误
+
+**验收标准**:
+1. 发送 `max_tokens: 999999`，不 panic，正常生成到 seq_len 上限
+2. 发送超长 prompt（> max_seq_len），返回 HTTP 400
+3. 编译通过
+
+---
+
+## 第二批：P1 — 严重 bug/性能
+
+### FIX-07: 使用 CachingAllocator [P1-性能] ❌未修
+
+**问题**: `CachingAllocator` 已实现（`allocator.rs`）但从未使用。所有 GPU 分配直接 `cudaMalloc`。
+
+**修复要求**:
+- `Tensor::empty` 对 GPU device 使用 `cached_alloc` 而非 `GpuBuffer::alloc`
+- `GpuBuffer::Drop` 调用 `cached_dealloc` 归还到池（而非 `cudaFree`）
+- 或者更简单：在 `GpuBuffer::alloc` 内部接入 caching allocator（全局透明替换）
+
+**验收标准**:
+1. 连续运行 10 次 decode step，`cudaMalloc` 调用次数应显著低于总分配次数
+2. 编译通过，现有测试通过
+3. 推理结果与修复前一致
+
+---
+
+### FIX-08: CudaDeviceProp FFI 安全性 [P1-Bug] ❌未修
+
+**问题**: `ffi.rs:31` 用 `_pad: [u8; 4096]` 猜测 `cudaDeviceProp` struct 大小，CUDA 12.9 可能更大。
+
+**修复要求**:
+- 增大 pad 到 `[u8; 8192]` 或使用 `cudaDeviceGetAttribute` 替代 name 查询
+- 可参考 `device.rs` 中已有的 `cudaDeviceGetAttribute` 用法
+
+**验收标准**:
+1. `device_info()` 返回正确的 device name
+2. 编译通过
+
+---
+
+### FIX-09: Tokenizer byte_fallback panic [P1-Bug] ❌未修
+
+**问题**: `bpe.rs:176-182` 中 Qwen3 tokenizer 遇到不在 vocab 的单字节时 panic。
+
+**修复要求**:
+- 当 `byte_fallback == true` 且单字节不在 vocab 时，查找 `<0xNN>` 格式 token
+- 如果 `<0xNN>` 也不存在，返回 unk_token_id（而非 panic）
+
+**验收标准**:
+1. 包含所有 256 个字节值的字符串可以 encode 不 panic
+2. 编译通过
+
+---
+
+### FIX-19: 因果掩码 -1e9 应改为 -inf [P1-Bug] ❌新发现
+
+**问题**: `csrc/attention/causal_mask.cu:31` 用 `-1e9f` 代替 `-inf`，注释说 "BF16 没有 -inf" 但这是错误的。
+
+**修复要求**:
+- BF16 路径改为 `__float2bfloat16(-INFINITY)`
+- F32 路径改为 `-INFINITY`（如果还没有的话）
+
+**验收标准**:
+1. causal mask 中被遮蔽的值为 `-inf`（而非 `-1e9`）
+2. 编译通过，attention test 通过
+
+---
+
+### FIX-20: LayerNorm 数值稳定性 [P1-Bug] ❌新发现
+
+**问题**: `csrc/normalization/layernorm.cu:19-25` 注释写 "Welford online" 但实际用 `E[x²] - E[x]²`，大均值小方差时会灾难性抵消。
+
+**修复要求**:
+- 改为真正的 two-pass 或 Welford online 算法
+- pass 1: 求 mean; pass 2: 求 variance = E[(x-mean)²]
+
+**验收标准**:
+1. 对 mean=1e6, std=1e-3 的输入，layernorm 输出与 PyTorch 一致（relative error < 1e-3）
+2. 编译通过，现有测试通过
+
+---
+
+### FIX-21: LayerNorm/RMSNorm 最小 block size [P1-Bug] ❌新发现
+
+**问题**: `layernorm.cu:88` 和 `rmsnorm.cu` 对 hidden_size < 32 的输入会崩溃（block_reduce 需要至少一个完整 warp）。
+
+**修复要求**:
+- launch 时 `block = max(min(hidden_size, 1024), 32)`
+
+**验收标准**:
+1. hidden_size=16 的 layernorm/rmsnorm 不崩溃
+2. 编译通过
+
+---
+
+## 第三批：P2 — 重要改进
+
+### FIX-22: Engine dummy KV cache 分配 [P2-性能] ❌新发现
+
+**问题**: `engine.rs:142-148` 每次 batched decode 用 `std::mem::replace` 创建 dummy `GpuKVCache::new(..., 1, ...)` 来绕过 borrow checker，每步分配 `num_layers * 2` 个 GPU buffer。
+
+**修复要求**:
+- 将 `running` 从 `Vec<Sequence>` 改为存储方式让 KV cache 可以独立借出
+- 或使用 `Option<GpuKVCache>` + `.take()` / `.insert()` 避免 dummy 分配
+
+**验收标准**:
+1. batched decode 路径不再分配 dummy KV cache
+2. 编译通过，功能不变
+
+---
+
+### FIX-23: RoPE cache 硬限 8192 [P2-功能] ❌新发现
+
+**问题**: `qwen3.rs:45` `config.max_seq_len().min(8192)` 人为截断。
+
+**修复要求**:
+- 去掉 `.min(8192)`，或改为与 engine 的 max_seq_len 一致
+- 确保 RoPE cache 覆盖实际使用的 max_seq_len
+
+**验收标准**:
+1. RoPE cache 长度 >= engine max_seq_len
+2. 编译通过
+
+---
+
+### FIX-15: GPT-2 消除 CPU round-trip [P3-性能] ❌未修
+
+**问题**: GPT-2 `split_qkv`、`merge_heads`、`add_bias` 全在 CPU 做。优先级低（GPT-2 不是主力模型）。
+
+---
+
+## 修复依赖图和执行顺序
+
+```
+第一批 P0 (可并行):
+  FIX-01 (cuBLAS handle)     ← 独立
+  FIX-16 (EOS 泄漏)          ← 独立
+  FIX-17 (max_seq_len)       ← 独立，FIX-23 依赖此
+  FIX-18 (max_tokens 校验)   ← 依赖 FIX-17（需要知道 max_seq_len）
+
+第二批 P1 (可并行):
+  FIX-07 (caching allocator) ← 独立
+  FIX-08 (CudaDeviceProp)    ← 独立
+  FIX-09 (byte_fallback)     ← 独立
+  FIX-19 (causal mask -inf)  ← 独立
+  FIX-20 (layernorm 稳定性)  ← 独立
+  FIX-21 (min block size)    ← 独立
+
+第三批 P2:
+  FIX-22 (dummy KV cache)    ← 独立
+  FIX-23 (RoPE cache)        ← 依赖 FIX-17
+```

docs/benchmarks/fp8-quantization.md

@@ -1,99 +0,0 @@
-# FP8 W8A8 quantization — gpt-oss-20b (dash5, 8× RTX 5090)
-
-Operator-level FP8 E4M3 quantization of the MoE expert weights, with real
-cuBLASLt FP8 tensor-core GEMM (W8A8: FP8 weights × dynamically-quantized FP8
-activations). All other tensors (attention, router, embeddings, norms, biases)
-stay BF16.
-
-## Scheme
-
-- **Weights** (`tools/quantize_fp8.py`): expert `gate_up_proj` / `down_proj`
-  quantized BF16 → FP8 E4M3 with a **per-expert scalar** scale (`absmax/448`).
-  Stored transposed `[E, N, K]` because cuBLASLt FP8 on Blackwell (sm120)
-  requires `transA=T`.
-- **Activations**: quantized dynamically at runtime, **per-token** (per-row
-  absmax), recovered by a post-GEMM row scale.
-- **Compute**: `batched_gemm_fp8` (`crates/xserv-kernels/src/quantization.rs`)
-  runs **one strided-batched cuBLASLt FP8 matmul for all experts** (`alpha=1`,
-  in-GEMM scales `1.0`); a fused kernel then applies `a_scale[token]·b_scale[expert]`
-  in a single pass. BF16's relative error is scale-invariant, so applying both
-  scales post-GEMM is precision-equivalent to folding them into the epilogue.
-- Model size: **22 GB** (FP8) vs **39 GB** (BF16). The FP8 model fits on a
-  single 32 GB 5090; BF16 needs ≥ 2.
-
-## The performance bug that was fixed
-
-`batched_gemm_fp8` originally rebuilt the entire cuBLASLt plan **per expert,
-per GEMM, per layer, on every forward pass** — running the algo heuristic
-search, creating/destroying the descriptor + 4 layouts + preference, and
-`cudaMalloc`-ing a 4-byte scale buffer — roughly 1500 heuristic searches per
-decoded token. This made FP8 **slower than BF16**:
-
-| | FP8 (buggy) | FP8 (fixed) | BF16 |
-|---|---|---|---|
-| Decode TPOT | 27.0 ms | **17.9 ms** | 18.8 ms |
-| Throughput | 37 tok/s | **55.8 tok/s** | 53.2 tok/s |
-
-Fix: cache the cuBLASLt plan (descriptor + layouts + heuristically-chosen algo)
-in a thread-local map keyed by `(M, N, K, batch)` so the heuristic runs once per
-shape, and allocate the scale buffer once.
-
-## Reducing launches: one strided-batched matmul
-
-The per-expert loop still issued one `cublasLtMatmul` per expert — ~768 tiny
-launches per decoded token (16 local experts × 2 GEMMs × 24 layers). Collapsing
-each MoE GEMM into a single **strided-batched** cuBLASLt FP8 matmul (BATCH_COUNT
-+ strided-batch offsets) drops that to ~48, with a fused post-scale kernel
-applying both scales. This required moving the per-expert weight scale out of the
-GEMM epilogue (a single strided call can't carry a per-batch scalar) into the
-post-scale kernel — precision-equivalent, as noted above.
-
-| (gpt-oss-20b, TP=2) | per-expert FP8 | batched FP8 | BF16 |
-|---|---|---|---|
-| Decode TPOT | 17.9 ms | **13.8 ms** | 18.8 ms |
-| Throughput | 55.8 tok/s | **72.3 tok/s** | 53.2 tok/s |
-
-## Results — GSM8K (greedy, TP=2 on the same 2 GPUs)
-
-200-problem run is the per-expert plan-cache fix; 100-problem run is the
-strided-batched version. BF16 is the unchanged baseline in both.
-
-Harness: `tools/fp8_compare.py` — a warm `xserv-server` per model, GSM8K streamed
-through `/v1/chat/completions`; TTFT = time to first token, TPOT = mean
-inter-token latency, per request.
-
-| metric | FP8 per-expert (n=200) | FP8 batched (n=100) | BF16 |
-|---|---|---|---|
-| GSM8K accuracy | 93.0 % | 91.0 % | 90.5 / 90.0 % |
-| TTFT median | 67.4 ms | 65.0 ms | 68.8 / 69.5 ms |
-| TPOT median | 17.45 ms | **13.08 ms** | 18.26 / 18.39 ms |
-| TPOT p90 | 17.65 ms | **13.28 ms** | 18.38 / 18.52 ms |
-| Throughput | 57.3 tok/s | **76.4 tok/s** | 54.8 / 54.4 tok/s |
-| Decode speedup vs BF16 | 1.05× | **1.41×** | 1.00× |
-
-- **Accuracy: unchanged.** FP8 is nominally +0.5 … +2.5 pts above BF16, but at
-  n=100–200 the standard error is ~2–3 pts, so they are statistically
-  indistinguishable. The takeaway is that neither FP8 quantization nor the
-  strided-batched rounding degrades accuracy.
-- **Decode: FP8 1.41× faster** once batched (TPOT 13.08 vs 18.39 ms), with a
-  tight p90. The per-expert version was only ~1.05× — the ~768 tiny M=1 launches
-  per token dominated; batching them into ~48 unlocked most of the FP8
-  expert-weight-bandwidth saving.
-- **Prefill (TTFT): comparable.** A multi-length sweep (113 / 561 / 1681 tokens)
-  gave FP8 480 / 362 / 2451 ms vs BF16 558 / 282 / 2287 ms — non-monotonic, i.e.
-  dominated by fixed overhead (cuBLAS lazy init + FP8's one-time per-shape
-  heuristic), not prefill compute, at these lengths.
-
-## Single-GPU (TP=1)
-
-FP8 runs gpt-oss-20b on **one** 5090 (`bench-gpt-oss --tp 1`, GPU6): TTFT 538 ms,
-TPOT 29.0 ms, 34.5 tok/s. BF16 cannot (39 GB > 32 GB). This — fitting a model
-that otherwise needs two GPUs onto one — is the largest practical win.
-
-## Follow-ups (not done)
-
-- Per-channel (per-output-row) weight scales for better accuracy headroom than
-  per-tensor.
-- Warm common prefill shapes at load to hide the first-request heuristic stall.
-- Sparse (top-k only) MoE compute instead of dense — currently every token runs
-  all experts, so only ~top_k/num_experts of the FP8 GEMM work is used.

docs/benchmarks/llama-cpp-comparison.md

@@ -51,19 +51,16 @@ context-bound at these sizes.
 
 | task | n | xserv | llama.cpp |
 |---|---|---|---|
-| GSM8K | 50 | 100.0% (50/50) | 96.0% (48/50) |
-| AIME 2025 | 30 | 16.7% (5/30) | 23.3% (7/30) |
+| GSM8K | 50 | 98.0% (49/50) | 96.0% (48/50) |
+| AIME 2025 | 30 | 20.0% (6/30) | 20.0% (6/30) |
 
-With equal context the two engines land at comparable AIME accuracy (within
-the ±2-problem greedy-decode wobble band) and xserv edges ahead on GSM8K. At
-8192 both generate full-length solutions (mean ~4.2k tokens), so neither is
-truncated. The AIME difference (2 problems) is entirely within the run-to-run
-non-determinism documented below. Per-problem analysis shows the disagreements
-are due to different greedy-decode paths (different token at position ~500+
-cascades into a different solution), not systematic precision errors.
-
-On GSM8K, xserv strictly dominates: it gets 2 problems right that llama.cpp
-misses, and never misses one that llama.cpp gets.
+With equal context the two engines land at identical AIME accuracy and
+within one problem on GSM8K. At 8192 both generate full-length solutions
+(mean ~3.4k / ~4.2k tokens), so neither is truncated. Two independent engines
+agreeing at ~20% confirms that's genuine Qwen3-8B (thinking-off) capability and
+that xserv is numerically faithful. Response prefixes are byte-identical (same
+prompt templating); the only run-to-run wobble is greedy-decode divergence /
+nondeterminism on long (~3k-token) sequences (see finding 3).
 
 ## Findings the benchmark surfaced
 
@@ -87,16 +84,6 @@ misses, and never misses one that llama.cpp gets.
    AIME config produced 6/30 / 7/30 / 6/30 across runs — non-deterministic CUDA
    reductions flip an argmax over long (~3k-token) generations. Harmless for
    serving, but it explains why long-sequence accuracy wobbles by a problem.
-4. **GEMV race condition corrupted decode outputs — now fixed.** The custom
-   K-split GEMV kernel (used for all M=1 decode-step projections with N≥256)
-   had a race condition: block k=0 zeroed the FP32 accumulator (`y_fp32[col] =
-   0.0`) while other K-blocks were already atomicAdding to it. Since CUDA
-   provides no inter-block ordering within a single kernel launch, the zero
-   could land before, during, or after other blocks' writes. Fix:
-   `cudaMemsetAsync` on the stream before the kernel launch, which guarantees
-   the buffer is zeroed before any block executes. This bug was introduced
-   after the initial benchmark and caused systematic decode-time precision
-   errors that degraded GSM8K accuracy from 98→80% range.
 
 Raw artifacts (per-request timings, per-problem prediction/gold) are written to
 `bench-out/` as `comparison-<stamp>.{md,json}` (gitignored).