[CPU] Add Qwen3.5 model optimization for CPU (#19484)

Co-authored-by: Zheng, Beilei <beilei.zheng@intel.com>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>

sgl-kernel/csrc/cpu/model/qwen3.cpp

@@ -61,6 +61,41 @@ void fused_qkvzba_split_reshape_cat_impl(
     }
   });
 }
+
+template <typename scalar_t>
+void fused_qkvzba_split_reshape_cat_contiguous_impl(
+    const scalar_t* __restrict__ mixed_qkvz,
+    const scalar_t* __restrict__ mixed_ba,
+    scalar_t* __restrict__ mixed_qkv,
+    scalar_t* __restrict__ z,
+    scalar_t* __restrict__ b,
+    scalar_t* __restrict__ a,
+    int64_t batch,
+    int64_t k_tp,
+    int64_t v_tp,
+    int64_t num_heads_v,
+    int64_t qkv_dim,
+    int64_t qkv_strideB,
+    int64_t qkvz_strideB,
+    int64_t ba_strideB) {
+  at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t bi = begin; bi < end; ++bi) {
+      scalar_t* __restrict__ qkv_out_ptr = mixed_qkv + bi * qkv_strideB;
+      const scalar_t* __restrict__ qkv_in_ptr = mixed_qkvz + bi * qkvz_strideB;
+      scalar_t* __restrict__ z_out_ptr = z + bi * v_tp;
+      const scalar_t* __restrict__ z_in_ptr = qkv_in_ptr + qkv_dim;
+      copy_stub(qkv_out_ptr, qkv_in_ptr, qkv_dim);
+      copy_stub(z_out_ptr, z_in_ptr, v_tp);
+      scalar_t* __restrict__ b_out_ptr = b + bi * num_heads_v;
+      const scalar_t* __restrict__ b_in_ptr = mixed_ba + bi * ba_strideB;
+      scalar_t* __restrict__ a_out_ptr = a + bi * num_heads_v;
+      const scalar_t* __restrict__ a_in_ptr = b_in_ptr + num_heads_v;
+      copy_stub(b_out_ptr, b_in_ptr, num_heads_v);
+      copy_stub(a_out_ptr, a_in_ptr, num_heads_v);
+    }
+  });
+}
+
 }  // anonymous namespace
 
 // mixed_qkvz: [batch, num_heads_qk * head_qk * 2 + num_heads_v * head_v * 2]
@@ -83,6 +118,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_re
   CHECK_EQ(mixed_qkvz.size(1), expected_dim);
   CHECK_EQ(mixed_ba.size(0), batch);
   CHECK_EQ(mixed_ba.size(1), ba_dim);
+  TORCH_CHECK(mixed_ba.scalar_type() == mixed_qkvz.scalar_type(), "mixed_ba and mixed_qkvz must share same dtype");
   CHECK_EQ(num_heads_v % num_heads_qk, 0);
   at::Tensor mixed_qkv = at::empty({batch, qkv_dim}, mixed_qkvz.options());
   at::Tensor z = at::empty({batch, num_heads_v, head_v}, mixed_qkvz.options());
@@ -112,3 +148,53 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_re
   });
   return std::make_tuple(mixed_qkv, z, b, a);
 }
+
+// mixed_qkvz: [batch, num_heads_qk * head_qk * 2 + num_heads_v * head_v * 2]
+// mixed_ba: [batch, num_heads_v * 2]
+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_reshape_cat_contiguous_cpu(
+    const at::Tensor& mixed_qkvz,
+    const at::Tensor& mixed_ba,
+    int64_t num_heads_qk,
+    int64_t num_heads_v,
+    int64_t head_qk,
+    int64_t head_v) {
+  CHECK_DIM(2, mixed_qkvz);
+  CHECK_DIM(2, mixed_ba);
+  CHECK_INPUT(mixed_qkvz);
+  CHECK_INPUT(mixed_ba);
+  int64_t batch = mixed_qkvz.size(0);
+  int64_t k_tp = num_heads_qk * head_qk;
+  int64_t v_tp = num_heads_v * head_v;
+  int64_t qkv_dim = k_tp * 2 + v_tp;
+  int64_t ba_dim = num_heads_v * 2;
+  int64_t expected_dim = qkv_dim + v_tp;
+  CHECK_EQ(mixed_qkvz.size(1), expected_dim);
+  CHECK_EQ(mixed_ba.size(0), batch);
+  CHECK_EQ(mixed_ba.size(1), ba_dim);
+  TORCH_CHECK(mixed_ba.scalar_type() == mixed_qkvz.scalar_type(), "mixed_ba and mixed_qkvz must share same dtype");
+  at::Tensor mixed_qkv = at::empty({batch, qkv_dim}, mixed_qkvz.options());
+  at::Tensor z = at::empty({batch, num_heads_v, head_v}, mixed_qkvz.options());
+  at::Tensor b = at::empty({batch, num_heads_v}, mixed_ba.options());
+  at::Tensor a = at::empty({batch, num_heads_v}, mixed_ba.options());
+  int64_t qkvz_strideB = mixed_qkvz.size(1);
+  int64_t qkv_strideB = mixed_qkv.size(1);
+  int64_t ba_strideB = mixed_ba.size(1);
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(mixed_qkvz.scalar_type(), "fused_qkvzba_split_reshape_cat_contiguous_impl", [&] {
+    fused_qkvzba_split_reshape_cat_contiguous_impl<scalar_t>(
+        mixed_qkvz.data_ptr<scalar_t>(),
+        mixed_ba.data_ptr<scalar_t>(),
+        mixed_qkv.data_ptr<scalar_t>(),
+        z.data_ptr<scalar_t>(),
+        b.data_ptr<scalar_t>(),
+        a.data_ptr<scalar_t>(),
+        batch,
+        k_tp,
+        v_tp,
+        num_heads_v,
+        qkv_dim,
+        qkv_strideB,
+        qkvz_strideB,
+        ba_strideB);
+  });
+  return std::make_tuple(mixed_qkv, z, b, a);
+}