attention: tree-aware paged_decode_attention_tree kernel + wrapper

New CUDA kernel paged_decode_attention_tree_bf16_kernel: same as base
paged_decode_attention but with a per-query mask over the newly-written
K/V region. `tree_mask[i][j] != 0` iff query i attends to newly-written
K/V at slot j. Positions before `tree_start` are always attended.

Motivation: speculative decoding with tree drafting needs siblings at
the same target position to attend to their own branch's history, not
each other's K/V.

Rust binding: paged_decode_attention_tree(...) mirrors
paged_decode_attention plus tree_mask_ptr, tree_start, tree_len.

Forward path: Qwen3::forward_verify_paged_decode_attention_tree_with_hidden
takes explicit positions, kv_lens, and a flattened [N*N] tree_mask.

Sanity check: bench-eagle3's γ_multi path now routes through the tree
kernel with a causal mask (mask[i][j]=1 iff j<=i), producing bit-
equivalent output to the non-tree variant. matched=false pattern +
acceptance rate + speedup all identical to previous run within noise
(11.3% acceptance, 1.00× speedup with the mask-check overhead).

--tree CLI flag is parsed but reserved. Real tree drafting (siblings
sharing a target position) is blocked by KV cache position rigidity:
paged_cache stores K/V at cache-position ≡ target-position, so an
accepted sibling at target position P+1 has its K/V physically at
cache position P+2 (its unique slot in the batched write). Continuing
decode at P+1 would see the WRONG K/V (top-1 sibling's, not accepted
top-2 sibling's). Fix requires either KV-slot remap on acceptance or
a virtual position layer.

Infrastructure is in place, next step is tackling that remap.

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)

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};