diff --git a/crates/xserv-kernels/src/attention.rs b/crates/xserv-kernels/src/attention.rs index fca9b66..9d1f340 100644 --- a/crates/xserv-kernels/src/attention.rs +++ b/crates/xserv-kernels/src/attention.rs @@ -83,6 +83,24 @@ unsafe extern "C" { 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, @@ -515,6 +533,62 @@ pub fn paged_decode_attention( 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) diff --git a/crates/xserv-kernels/src/lib.rs b/crates/xserv-kernels/src/lib.rs index 7e62694..73ba21d 100644 --- a/crates/xserv-kernels/src/lib.rs +++ b/crates/xserv-kernels/src/lib.rs @@ -16,7 +16,8 @@ pub use activation::{add, bias_add_2d, gelu, gpt_oss_glu, mul, scale, silu, silu pub use argmax::{argmax_bf16_single, argmax_bf16_to_host}; pub use attention::{ attention, decode_attention, flash_attention, flash_attention_sinks, paged_decode_attention, - paged_decode_attention_sinks, reshape_and_cache_batched_bf16, reshape_and_cache_bf16, + 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}; diff --git a/crates/xserv-model/src/bin/bench-eagle3.rs b/crates/xserv-model/src/bin/bench-eagle3.rs index 30a2e15..bd1bae7 100644 --- a/crates/xserv-model/src/bin/bench-eagle3.rs +++ b/crates/xserv-model/src/bin/bench-eagle3.rs @@ -94,6 +94,7 @@ fn main() { let max_seq_len = arg_usize(&args, "--max-seq-len", DEFAULT_MAX_SEQ_LEN); let device = arg_usize(&args, "--device", 0) as u32; let gamma = arg_usize(&args, "--gamma", 2).max(1); + let use_tree = args.iter().any(|a| a == "--tree"); xserv_cuda::device::set_device(device).unwrap(); let info = xserv_cuda::device::device_info(device).unwrap(); @@ -150,7 +151,11 @@ fn main() { baseline_tokens += baseline.ids.len(); drop(baseline_cache); - // Speculative with EAGLE, γ from CLI. + // Speculative with EAGLE, γ from CLI. Verify uses the tree kernel with + // a causal mask (equivalent to non-tree behavior); a real tree + // (siblings sharing target positions) would require KV cache slot + // remap after acceptance, which is out of scope for this iteration. + let _ = use_tree; // reserved for future tree drafting let mut target_cache = new_cache(&target_config, max_seq_len, device); let spec = if gamma == 1 { run_eagle_gamma1( @@ -441,9 +446,23 @@ fn run_eagle_gamma_multi( for &d in drafts.iter() { verify_input.push(d); } + let n = verify_input.len(); + let pos_offset = cache.seq_len(slot); + let positions_v: Vec = (0..n).map(|i| (pos_offset + i) as u32).collect(); + let kv_lens_v: Vec = (0..n).map(|i| (pos_offset + i + 1) as i32).collect(); + // Causal mask over new-writes: mask[i][j] = 1 iff j <= i. + let mut tree_mask: Vec = vec![0; n * n]; + for i in 0..n { + for j in 0..=i { + tree_mask[i * n + j] = 1; + } + } let (verify_logits, verify_hooks) = target - .forward_verify_paged_decode_attention_with_hidden( + .forward_verify_paged_decode_attention_tree_with_hidden( &verify_input, + &positions_v, + &kv_lens_v, + &tree_mask, slot, cache, &EAGLE_HOOK_LAYERS, diff --git a/crates/xserv-model/src/qwen3.rs b/crates/xserv-model/src/qwen3.rs index 21e9156..c20a0a6 100644 --- a/crates/xserv-model/src/qwen3.rs +++ b/crates/xserv-model/src/qwen3.rs @@ -1230,6 +1230,138 @@ impl Qwen3 { (logits, hidden_arr) } + /// Tree-aware verify: like `_with_hidden` but supports sibling candidates + /// sharing the same target position. Caller supplies per-token positions + /// (for RoPE), kv_lens (attention context length), and a flattened + /// `tree_mask` (`[new_tokens, new_tokens]` i32; `mask[i, j]!=0` iff query i + /// attends to newly-written K/V at slot j). Positions in the paged cache + /// before pos_offset are always attended (regular history). + #[allow(clippy::too_many_arguments)] + pub fn forward_verify_paged_decode_attention_tree_with_hidden( + &self, + token_ids: &[u32], + positions: &[u32], + kv_lens: &[i32], + tree_mask: &[i32], + slot: usize, + paged_cache: &mut PagedKVCache, + hook_layers: &[usize; 3], + ) -> (Tensor, [Tensor; 3]) { + let new_tokens = token_ids.len(); + assert_eq!(positions.len(), new_tokens); + assert_eq!(kv_lens.len(), new_tokens); + assert_eq!(tree_mask.len(), new_tokens * new_tokens); + + let pos_offset = paged_cache.seq_len(slot); + let num_heads = self.local_num_heads; + let num_kv_heads = self.local_num_kv_heads; + let head_dim = self.config.head_dim(); + let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32; + + paged_cache.ensure_capacity(slot, pos_offset + new_tokens); + paged_cache.advance_seq_len(slot, new_tokens); + + let slots = vec![slot; new_tokens]; + paged_cache.sync_active_batch_with_lens(&slots, kv_lens); + let bt_ptr = paged_cache.block_table_gpu().as_ptr() as *const i32; + let cl_ptr = paged_cache.context_lens_gpu().as_ptr() as *const i32; + let max_blocks = paged_cache.max_blocks_per_seq(); + + // Upload tree_mask [new_tokens, new_tokens] i32 to GPU. + let mask_bytes: &[u8] = unsafe { + std::slice::from_raw_parts(tree_mask.as_ptr() as *const u8, tree_mask.len() * 4) + }; + let mut mask_buf = + xserv_cuda::allocator::cached_alloc(mask_bytes.len()).expect("alloc tree_mask"); + mask_buf.copy_from_host(mask_bytes).unwrap(); + let mask_ptr = mask_buf.as_ptr() as *const i32; + + let mut x = embedding(&self.embed_tokens, token_ids); + let mut hooks: [Option; 3] = [None, None, None]; + + for (layer_idx, layer) in self.layers.iter().enumerate() { + let residual = x.clone(); + let normed = rmsnorm(&x, &layer.input_norm, eps); + + let qkv = matmul_2d(&normed, &layer.qkv_proj_wt); + let q_dim = num_heads * head_dim; + let kv_dim = num_kv_heads * head_dim; + let q_all = qkv.narrow(1, 0, q_dim); + let k_all = qkv.narrow(1, q_dim, kv_dim); + let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim); + + let q_flat = q_all + .contiguous() + .reshape(&[new_tokens * num_heads, head_dim]); + let k_flat = k_all + .contiguous() + .reshape(&[new_tokens * num_kv_heads, head_dim]); + let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps); + let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps); + + let q_3d = q_normed.reshape(&[new_tokens, num_heads, head_dim]); + let k_3d = k_normed.reshape(&[new_tokens, num_kv_heads, head_dim]); + rope_inplace(&q_3d, &self.rope_cache, positions); + rope_inplace(&k_3d, &self.rope_cache, positions); + + let v_3d = v_all + .contiguous() + .reshape(&[new_tokens, num_kv_heads, head_dim]); + paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, new_tokens); + + let q_decode = q_3d.reshape(&[new_tokens, num_heads, 1, head_dim]); + let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void; + let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void; + let attn_out = xserv_kernels::paged_decode_attention_tree( + &q_decode, + k_pool_ptr, + v_pool_ptr, + bt_ptr, + cl_ptr, + mask_ptr, + new_tokens, + num_heads, + num_kv_heads, + head_dim, + max_blocks, + pos_offset, + new_tokens, + ); + + let attn_merged = attn_out.reshape(&[new_tokens, num_heads * head_dim]); + let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt); + self.all_reduce(&attn_proj); + + let (normed, x_new) = + xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps); + let residual = x_new.clone(); + + let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt); + let ffn_dim = gate_up.shape()[1] / 2; + let gate = gate_up.narrow(1, 0, ffn_dim).contiguous(); + let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous(); + let hidden_states = xserv_kernels::silu_mul(&gate, &up); + let down = matmul_2d(&hidden_states, &layer.down_proj_wt); + self.all_reduce(&down); + x = add_any(&residual, &down); + + for (h_idx, &h_layer) in hook_layers.iter().enumerate() { + if layer_idx == h_layer { + hooks[h_idx] = Some(x.clone()); + } + } + } + + let x = rmsnorm(&x, &self.norm, eps); + let logits = matmul_2d(&x, &self.lm_head_t); + let hidden_arr = [ + hooks[0].take().expect("hook layer 0 not reached"), + hooks[1].take().expect("hook layer 1 not reached"), + hooks[2].take().expect("hook layer 2 not reached"), + ]; + (logits, hidden_arr) + } + /// Forward with GPU-resident KV cache and GPU transpose/reshape kernels. pub fn forward_gpu_cache(&self, token_ids: &[u32], cache: &mut GpuKVCache) -> Tensor { let new_tokens = token_ids.len(); diff --git a/csrc/attention/paged_attention.cu b/csrc/attention/paged_attention.cu index 29c3066..0859485 100644 --- a/csrc/attention/paged_attention.cu +++ b/csrc/attention/paged_attention.cu @@ -189,6 +189,169 @@ __global__ void paged_decode_attention_bf16_kernel( } } +// 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(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. @@ -389,6 +552,36 @@ 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<<>>( + (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,