From e524cd4ede45b264ef7b2c0fd97a52e6356bb54b Mon Sep 17 00:00:00 2001
From: ShareableXue <1006032941@qq.com>
Date: Sat, 4 Jul 2026 15:27:25 +0800
Subject: [PATCH] perf: route Qwen3.5 FSGDR model path to TileLang.

---
 xllm/core/kernels/ops_api.cpp                   |  5 ++++-
 .../npu_torch/qwen3_gated_delta_net_base.cpp    | 17 ++++++++++-------
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/xllm/core/kernels/ops_api.cpp b/xllm/core/kernels/ops_api.cpp
index a7a915c52b..e85b42eaea 100644
--- a/xllm/core/kernels/ops_api.cpp
+++ b/xllm/core/kernels/ops_api.cpp
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "ops_api.h"
 
+#include <tuple>
+
 #if defined(USE_MLU)
 #include "mlu/mlu_ops_api.h"
 #elif defined(USE_NPU)
@@ -1278,7 +1280,7 @@ torch::Tensor hc_pre_inv_rms(HcPreInvRmsParams& params) {
 torch::Tensor fused_sigmoid_gating_delta_rule_update(
     FusedSigmoidGatingDeltaRuleUpdateParams& params) {
 #if defined(USE_NPU)
-  return npu::npu_fused_sigmoid_gating_delta_rule_update(
+  auto result = npu::tilelang::fused_sigmoid_gating_delta_rule(
       params.A_log,
       params.a,
       params.dt_bias,
@@ -1293,6 +1295,7 @@ torch::Tensor fused_sigmoid_gating_delta_rule_update(
       params.use_qk_l2norm_in_kernel,
       params.softplus_beta,
       params.softplus_threshold);
+  return std::get<0>(result);
 #else
   NOT_IMPLEMENTED();
 #endif
diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp
index 7312550dda..3a245675b0 100644
--- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp
+++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp
@@ -848,13 +848,16 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
     double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
     if (fla_ssm_state_layout) {
       xllm::kernel::FusedSigmoidGatingDeltaRuleUpdateParams params;
-      params.A_log = A_log_.contiguous();
-      params.a = a.contiguous();
-      params.dt_bias = dt_bias_.contiguous();
-      params.q = processed_q.contiguous();
-      params.k = processed_k.contiguous();
-      params.v = processed_v.contiguous();
-      params.b = b.contiguous();
+      params.A_log = A_log_;
+      params.a = a.reshape({-1, a.size(-1)});
+      params.dt_bias = dt_bias_;
+      params.q =
+          processed_q.reshape({-1, processed_q.size(-2), processed_q.size(-1)});
+      params.k =
+          processed_k.reshape({-1, processed_k.size(-2), processed_k.size(-1)});
+      params.v =
+          processed_v.reshape({-1, processed_v.size(-2), processed_v.size(-1)});
+      params.b = b.reshape({-1, b.size(-1)});
       params.initial_state_source = ssm_cache;
       params.initial_state_indices = linear_state_base_indices.contiguous();
       params.cu_seqlens = attn_metadata.q_cu_seq_lens.contiguous();