Add Qwen3.6 MoE inference support

crates/xserv-kernels/src/attention.rs

@@ -412,8 +412,8 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens
         "num_q_heads must be divisible by num_kv_heads"
     );
     assert!(
-        head_dim <= 128,
-        "flash_attention supports head_dim up to 128"
+        head_dim <= 256,
+        "flash_attention supports head_dim up to 256"
     );
 
     // Dispatch to specialized decode kernel for single-token generation
@@ -479,7 +479,7 @@ pub fn flash_attention_sinks(
     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!(head_dim <= 256);
     assert_eq!(
         sinks.shape()[0],
         num_q_heads,
@@ -548,7 +548,7 @@ pub fn paged_decode_attention(
         num_q_heads % num_kv_heads == 0,
         "GQA: num_q_heads must be divisible by num_kv_heads"
     );
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
 
     let scale = 1.0 / (head_dim as f32).sqrt();
     let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
@@ -601,7 +601,7 @@ pub fn paged_decode_attention_tree(
     assert_eq!(q.shape()[2], 1);
     assert_eq!(q.dtype(), DType::BF16);
     assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
 
     let scale = 1.0 / (head_dim as f32).sqrt();
     let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
@@ -653,7 +653,7 @@ pub fn paged_decode_attention_sinks(
     assert_eq!(q.shape()[2], 1);
     assert_eq!(q.dtype(), DType::BF16);
     assert!(num_q_heads % num_kv_heads == 0);
-    assert!(head_dim <= 128);
+    assert!(head_dim <= 256);
 
     let scale = 1.0 / (head_dim as f32).sqrt();
     let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());