Skip to content

perf: route Qwen3.5 FSGDR model path to TileLang. #1889

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ShareableXue wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf: route Qwen3.5 FSGDR model path to TileLang. #1889

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion xllm/core/kernels/ops_api.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
17 changes: 10 additions & 7 deletions xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)});

@gemini-code-assist gemini-code-assist Bot Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The attn_metadata.q_cu_seq_lens tensor (assigned to params.cu_seqlens on line 863) is not guaranteed to be of type torch::kInt32 and can often be torch::kInt64/torch::kLong. However, the TileLang wrapper fused_sigmoid_gating_delta_rule has an explicit assertion CHECK_EQ(cu_seqlens.scalar_type(), torch::kInt32) which will fail and cause a runtime crash if a 64-bit integer tensor is passed. Please explicitly cast attn_metadata.q_cu_seq_lens to torch::kInt32 before calling .contiguous(), similar to how it is handled in other kernels (e.g., causal_conv1d_update or chunk_gated_delta_rule).

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();
Expand Down
Loading