dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

data: port complete MoE + xllm layer call chains from upstream repos

Browse Source 
MoE call chain from ds_vllm (vllm-project/vllm latest):
  ex_engine/moe/ — 20 files, 8736 lines
  - modular_kernel.py (1630 lines) — base classes for modular MoE
  - experts/fused_batched_moe.py (972 lines) — NaiveBatchedExperts
  - prepare_finalize/batched.py (171 lines) — token grouping by expert
  - topk_weight_and_reduce.py (176 lines) — scatter-add finalize
  - fused_moe.py (1740 lines) — main fused_moe dispatch
  - config.py (1407 lines) — FusedMoEQuantConfig
  - activation.py, utils.py, layer.py, etc.

xllm layer code (jd-opensource/xllm):
  ex_engine/xllm_layers/ — 39 files, 5859 lines
  - ilu/fused_moe.cpp (797 lines) — production ixformer 7-step MoE pipeline
  - ilu/attention.cpp (189 lines) — paged_attention + flash_attn bridge
  - npu_torch/qwen3_gated_delta_net_base.cpp (576 lines) — GDN reference
  - common/rms_norm.cpp, rotary_embedding.cpp, activation.cpp, dense_mlp.cpp

xllm ILU kernels — synced 10 files to upstream (diffs from prior edits)

These are reference implementations, NOT hand-written.
Source repos: vllm-project/vllm, jd-opensource/xllm
This commit is contained in:
Claude
2026-08-15 14:26:18 +00:00
parent 7aa5054574
commit 6415249693
47 changed files with 10894 additions and 851 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2025-2026 The xLLM Authors.
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
@@ -123,53 +123,19 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations(
}
std::pair<torch::Tensor, torch::Tensor>
Qwen3_5GatedDeltaNetImpl::project_decode_inputs(
const torch::Tensor& hidden_states) {
const auto reshape_projection = [](const torch::Tensor& projection) {
return projection.view({projection.size(0), -1, projection.size(-1)});
};
auto qkv = reshape_projection(in_proj_qkv_->forward(hidden_states));
auto z_proj = reshape_projection(in_proj_z_->forward(hidden_states));
auto b_proj = reshape_projection(in_proj_b_->forward(hidden_states));
auto a_proj = reshape_projection(in_proj_a_->forward(hidden_states));
return {merge_qkvz_from_split_activations(qkv, z_proj),
merge_ba_from_split_activations(b_proj, a_proj)};
}
std::pair<torch::Tensor, torch::Tensor>
Qwen3_5GatedDeltaNetImpl::project_flat_inputs(
const torch::Tensor& hidden_states) {
auto qkv = in_proj_qkv_->forward(hidden_states).unsqueeze(0);
auto z_proj = in_proj_z_->forward(hidden_states).unsqueeze(0);
auto b_proj = in_proj_b_->forward(hidden_states).unsqueeze(0);
auto a_proj = in_proj_a_->forward(hidden_states).unsqueeze(0);
auto qkvz = merge_qkvz_from_split_activations(qkv, z_proj);
auto ba = merge_ba_from_split_activations(b_proj, a_proj);
return {qkvz.view({hidden_states.size(0), qkvz.size(-1)}).contiguous(),
ba.view({hidden_states.size(0), ba.size(-1)}).contiguous()};
}
std::optional<
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
Qwen3_5GatedDeltaNetImpl::project_split_inputs(
Qwen3_5GatedDeltaNetImpl::project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) {
auto qkv = reshape_projected_tokens_with_pad(
attn_metadata, in_proj_qkv_->forward(hidden_states));
auto z_proj = reshape_projected_tokens_with_pad(
attn_metadata, in_proj_z_->forward(hidden_states));
auto b_proj = reshape_projected_tokens_with_pad(
attn_metadata, in_proj_b_->forward(hidden_states));
auto a_proj = reshape_projected_tokens_with_pad(
attn_metadata, in_proj_a_->forward(hidden_states));
const int64_t batch_size = qkv.size(0);
const int64_t seq_len = qkv.size(1);
auto z =
z_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
auto b = b_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
auto a = a_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
return std::make_tuple(qkv, z, b, a);
auto qkv = reshape_qkvz_with_pad(attn_metadata,
in_proj_qkv_->forward(hidden_states));
auto z_proj =
reshape_qkvz_with_pad(attn_metadata, in_proj_z_->forward(hidden_states));
auto b_proj =
reshape_qkvz_with_pad(attn_metadata, in_proj_b_->forward(hidden_states));
auto a_proj =
reshape_qkvz_with_pad(attn_metadata, in_proj_a_->forward(hidden_states));
return {merge_qkvz_from_split_activations(qkv, z_proj),
merge_ba_from_split_activations(b_proj, a_proj)};
}
void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict(

16
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2025-2026 The xLLM Authors.
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -17,9 +17,7 @@ limitations under the License.
#include <torch/torch.h>
#include <optional>
#include <string>
#include <tuple>
#include <utility>
#include "qwen3_next_gated_delta_net.h"
@@ -36,15 +34,9 @@ class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl {
const torch::TensorOptions& options);
protected:
std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
const torch::Tensor& hidden_states) override;
std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
const torch::Tensor& hidden_states) override;
std::optional<
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
project_split_inputs(const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) override;
bool use_fla_ssm_state_layout() const override { return true; }
std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) override;
void load_projection_state_dict(const StateDict& state_dict) override;
void verify_projection_weights(const std::string& prefix) const override;

858
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2025-2026 The xLLM Authors.
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
@@ -15,12 +15,9 @@ limitations under the License.
#include <glog/logging.h>
#include <torch/torch.h>
#include <optional>
#include <tuple>
#include "xllm/core/kernels/npu/npu_ops_api.h"
#include "xllm/core/kernels/ops_api.h"
#include "xllm/core/platform/npu/acl_graph_task_update_context.h"
namespace xllm {
namespace layer {
@@ -31,31 +28,6 @@ torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) {
return x / norm;
}
torch::Tensor repeat_tensor_heads(const torch::Tensor& tensor,
int64_t target_heads,
int64_t head_dim) {
const int64_t current_heads = tensor.size(head_dim);
if (current_heads == target_heads) {
return tensor;
}
CHECK_GT(current_heads, 0) << "current heads must be positive";
CHECK_EQ(target_heads % current_heads, 0)
<< "target heads must be divisible by current heads, target_heads="
<< target_heads << ", current_heads=" << current_heads;
const int64_t repeats = target_heads / current_heads;
std::vector<int64_t> view_shape = tensor.sizes().vec();
view_shape.insert(view_shape.begin() + head_dim + 1, 1);
std::vector<int64_t> expand_shape = view_shape;
expand_shape[head_dim + 1] = repeats;
std::vector<int64_t> output_shape = tensor.sizes().vec();
output_shape[head_dim] = target_heads;
return tensor.unsqueeze(head_dim + 1)
.expand(expand_shape)
.reshape(output_shape)
.contiguous();
}
std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
torch::Tensor query,
torch::Tensor key,
@@ -80,9 +52,6 @@ std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
value = to_float32_and_transpose(value);
beta = to_float32_and_transpose(beta);
g = to_float32_and_transpose(g);
const int64_t value_num_heads = value.size(1);
query = repeat_tensor_heads(query, value_num_heads, 1);
key = repeat_tensor_heads(key, value_num_heads, 1);
int64_t batch_size = key.size(0);
int64_t num_heads = key.size(1);
@@ -150,15 +119,12 @@ std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
value = to_float32(value);
beta = to_float32(beta);
g = to_float32(g);
const int64_t value_num_heads = value.size(1);
query = repeat_tensor_heads(query, value_num_heads, 1);
key = repeat_tensor_heads(key, value_num_heads, 1);
int64_t batch_size = query.size(0);
int64_t num_heads = query.size(1);
int64_t sequence_length = query.size(2);
int64_t k_head_dim = key.size(-1);
int64_t v_head_dim = value.size(-1);
auto batch_size = query.size(0);
auto num_heads = query.size(1);
auto sequence_length = query.size(2);
auto k_head_dim = key.size(-1);
auto v_head_dim = value.size(-1);
int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size;
query = torch::nn::functional::pad(
@@ -276,164 +242,6 @@ std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
return std::make_tuple(core_attn_out, last_recurrent_state);
}
int64_t get_checkpoint_stride(const torch::Tensor& conv_cache,
const torch::Tensor& ssm_cache) {
if (!conv_cache.defined() || !ssm_cache.defined() ||
conv_cache.numel() == 0 || ssm_cache.numel() == 0) {
return 1;
}
CHECK_GT(conv_cache.size(0), 0) << "conv cache must have positive batch dim";
CHECK_EQ(ssm_cache.size(0) % conv_cache.size(0), 0)
<< "ssm cache checkpoint layout mismatch, ssm_rows=" << ssm_cache.size(0)
<< ", conv_rows=" << conv_cache.size(0);
return ssm_cache.size(0) / conv_cache.size(0);
}
torch::Tensor build_linear_state_base_indices(
const torch::Tensor& logical_state_indices,
int64_t checkpoint_stride) {
if (checkpoint_stride == 1) {
return logical_state_indices;
}
return logical_state_indices * checkpoint_stride;
}
torch::Tensor expand_sequence_tensor_to_batch(const torch::Tensor& tensor,
int64_t target_batch,
const char* tensor_name) {
CHECK(tensor.defined()) << tensor_name << " must be defined";
CHECK_EQ(tensor.dim(), 1) << tensor_name << " must be a 1D tensor.";
const int64_t source_batch = tensor.size(0);
if (source_batch == target_batch) {
return tensor.contiguous();
}
CHECK_GT(source_batch, 0) << tensor_name << " must not be empty.";
CHECK_EQ(target_batch % source_batch, 0)
<< tensor_name << " cannot be expanded from " << source_batch << " to "
<< target_batch;
const int64_t repeat_count = target_batch / source_batch;
return tensor.unsqueeze(1)
.expand({source_batch, repeat_count})
.reshape({target_batch})
.contiguous();
}
torch::Tensor run_causal_conv1d_graph_update(
const std::shared_ptr<xllm::npu::AclGraphTaskUpdateContext>& graph_context,
const torch::Tensor& x,
const torch::Tensor& weight,
const torch::Tensor& conv_state,
const std::optional<torch::Tensor>& bias,
const std::vector<int64_t>& query_start_loc,
const std::vector<int64_t>& cache_indices,
const std::vector<int64_t>& num_accepted_tokens,
xllm::npu::CausalConv1dGraphBranch branch) {
CHECK(graph_context != nullptr && graph_context->capturing)
<< "causal_conv1d graph update can only be registered during capture";
c10_npu::NPUStream stream = c10_npu::getCurrentNPUStream();
auto event = std::make_shared<c10_npu::NPUEvent>(ACL_EVENT_EXTERNAL);
event->block(stream);
event->reset(stream);
torch::Tensor output;
c10_npu::graph_task_group_begin(stream);
const std::vector<int64_t> empty_host_args;
CHECK(!query_start_loc.empty())
<< "query_start_loc must be populated for causal_conv1d graph update";
CHECK_EQ(query_start_loc.back(), x.size(0))
<< "query_start_loc must be padded to x.shape[0] during graph capture";
CHECK_EQ(cache_indices.size() + 1, query_start_loc.size())
<< "cache_indices must be sequence-scoped";
if (branch == xllm::npu::CausalConv1dGraphBranch::kSpecVerify) {
CHECK_EQ(num_accepted_tokens.size(), cache_indices.size())
<< "num_accepted_tokens must be sequence-scoped for spec verify";
}
output = torch::empty_like(x);
xllm::kernel::causal_conv1d_out(output,
x,
weight,
conv_state,
bias,
torch::IntArrayRef(query_start_loc),
torch::IntArrayRef(cache_indices),
torch::IntArrayRef(empty_host_args),
torch::IntArrayRef(num_accepted_tokens),
xllm::npu::kCausalConv1dActivationSilu,
xllm::npu::kCausalConv1dGraphPadSlotId,
xllm::npu::kCausalConv1dRunModeUpdate);
c10_npu::NPUTaskGroupHandle handle = c10_npu::graph_task_group_end(stream);
xllm::npu::CausalConv1dGraphTask task;
task.output = output;
task.x = x;
task.weight = weight;
task.conv_state = conv_state;
task.bias = bias;
task.activation_mode = xllm::npu::kCausalConv1dActivationSilu;
task.pad_slot_id = xllm::npu::kCausalConv1dGraphPadSlotId;
task.run_mode = xllm::npu::kCausalConv1dRunModeUpdate;
task.branch = branch;
task.handle = handle;
task.event = std::move(event);
graph_context->causal_conv1d_tasks.emplace_back(std::move(task));
return output;
}
torch::Tensor run_spec_verify_gated_delta_rule(
torch::Tensor query,
torch::Tensor key,
torch::Tensor value,
torch::Tensor g,
torch::Tensor beta,
torch::Tensor& ssm_cache,
const torch::Tensor& checkpoint_indices,
const torch::Tensor& num_accepted_tokens,
const torch::Tensor& cu_seq_lens,
const std::vector<int32_t>& q_seq_lens_vec,
double scale) {
const auto device = value.device();
const int64_t batch_size = value.size(0);
const int64_t seq_len = value.size(1);
const int64_t total_seq_len = batch_size * seq_len;
CHECK_EQ(cu_seq_lens.numel(), batch_size + 1)
<< "GDN spec verify cu_seq_lens must be cumulative.";
CHECK_EQ(q_seq_lens_vec.size(), static_cast<size_t>(batch_size))
<< "GDN spec verify q_seq_lens_vec must be per sequence.";
for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
CHECK_EQ(q_seq_lens_vec[batch_idx], seq_len)
<< "Qwen3.5 spec verify fused recurrent path expects dense "
"same-length validate tokens.";
}
xllm::kernel::FusedRecurrentGatedDeltaRuleParams params;
params.q = query.reshape({1, total_seq_len, query.size(-2), query.size(-1)})
.contiguous();
params.k =
key.reshape({1, total_seq_len, key.size(-2), key.size(-1)}).contiguous();
params.v = value.reshape({1, total_seq_len, value.size(-2), value.size(-1)})
.contiguous();
params.g = g.to(torch::kFloat32)
.reshape({1, total_seq_len, g.size(-1)})
.contiguous();
params.beta = beta.reshape({1, total_seq_len, beta.size(-1)}).contiguous();
params.scale = static_cast<float>(scale);
params.initial_state = ssm_cache;
params.inplace_final_state = true;
params.cu_seqlens = cu_seq_lens.to(torch::kLong).contiguous();
params.ssm_state_indices = checkpoint_indices.contiguous();
params.num_accepted_tokens =
num_accepted_tokens.to(device, torch::kInt32).contiguous();
params.use_qk_l2norm_in_kernel = true;
auto output_and_state =
xllm::kernel::fused_recurrent_gated_delta_rule(params);
return output_and_state.first.view(
{batch_size, seq_len, value.size(-2), value.size(-1)});
}
} // namespace
Qwen3GatedDeltaNetBaseImpl::Qwen3GatedDeltaNetBaseImpl(
@@ -495,11 +303,7 @@ void Qwen3GatedDeltaNetBaseImpl::load_common_state_dict(
if (auto w = state_dict.get_tensor("conv1d.weight"); w.defined()) {
conv1d_->load_state_dict(
StateDict({{"weight", w.squeeze(1)}},
static_cast<std::string>(state_dict.prefix()) + "conv1d."),
shard_tensor_count,
shard_sizes);
conv1d_->weight().set_(conv1d_->weight().transpose(0, 1).contiguous());
StateDict({{"weight", w.squeeze(1)}}), shard_tensor_count, shard_sizes);
}
o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj."));
if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) {
@@ -518,279 +322,87 @@ void Qwen3GatedDeltaNetBaseImpl::verify_common_loaded_weights(
<< prefix << "A_log";
}
std::pair<torch::Tensor, torch::Tensor>
Qwen3GatedDeltaNetBaseImpl::project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) {
if (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill) {
auto [qkvz_flat, ba_flat] = project_flat_inputs(hidden_states);
return {reshape_projected_tokens_with_pad(attn_metadata, qkvz_flat),
reshape_projected_tokens_with_pad(attn_metadata, ba_flat)};
}
return project_decode_inputs(hidden_states);
}
torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata,
KVCache& kv_cache,
const ModelInputParams& input_params) {
// Early-return on dummy shards. Under dp>1, an empty shard is padded with a
// fake token by worker_impl but its GDN state tensors (kv_cache_tokens_nums,
// linear_state_ids etc.) are left undefined. This mirrors the is_dummy
// early-return in Attention::forward (npu_torch/attention.cpp). Uses
// zeros_like rather than empty_like so downstream post-norm / mlp do not
// read uninitialized data. Placed before FlashComm1 sequence gather so
// dummy shards do not enter the collective and waste bandwidth.
if (attn_metadata.is_dummy) {
return torch::zeros_like(hidden_states);
}
const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context();
torch::Tensor h = hidden_states;
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
h = gather_sequence(hidden_states, *fc1_ctx);
}
auto [qkvz_padded, ba_padded] =
project_padded_inputs(hidden_states, attn_metadata);
int64_t batch_size = qkvz_padded.size(0);
int64_t seq_len = qkvz_padded.size(1);
torch::Tensor qkvz_flat =
qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)});
torch::Tensor ba_flat =
ba_padded.view({batch_size * seq_len, ba_padded.size(-1)});
xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params;
fused_params.mixed_qkvz = qkvz_flat;
fused_params.mixed_ba = ba_flat;
fused_params.num_heads_qk = static_cast<int32_t>(num_k_heads_ / tp_size_);
fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
fused_params.head_v = static_cast<int32_t>(head_v_dim_);
// Save the gathered hidden-state size for potential padding later.
const int64_t original_num_tokens = h.size(0);
const bool use_spec_verify = input_params.is_spec_verify;
const bool is_any_prefill =
attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
torch::Tensor mixed_qkv, z, b, a;
torch::Tensor processed_q, processed_k, processed_v;
int64_t batch_size = 0;
int64_t seq_len = 0;
std::tie(mixed_qkv, z, b, a) =
xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
// Qwen3.5 stores qkv, z, b, and a as separate projection weights, so it can
// use their outputs directly in every forward mode. Qwen3Next stores qkvz
// and ba as packed weights and uses the fused-split fallback below.
auto split_inputs = project_split_inputs(h, attn_metadata);
if (split_inputs.has_value()) {
std::tie(mixed_qkv, z, b, a) = split_inputs.value();
batch_size = mixed_qkv.size(0);
seq_len = mixed_qkv.size(1);
} else {
auto [qkvz_padded, ba_padded] = project_padded_inputs(h, attn_metadata);
batch_size = qkvz_padded.size(0);
seq_len = qkvz_padded.size(1);
torch::Tensor qkvz_flat =
qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)});
torch::Tensor ba_flat =
ba_padded.view({batch_size * seq_len, ba_padded.size(-1)});
xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params;
fused_params.mixed_qkvz = qkvz_flat;
fused_params.mixed_ba = ba_flat;
fused_params.num_heads_qk = static_cast<int32_t>(num_k_heads_ / tp_size_);
fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
fused_params.head_v = static_cast<int32_t>(head_v_dim_);
std::tie(mixed_qkv, z, b, a) =
xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)});
z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_});
a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_});
}
const bool fla_ssm_state_layout = use_fla_ssm_state_layout();
const int64_t local_q_heads = num_k_heads_ / tp_size_;
const int64_t local_v_heads = num_v_heads_ / tp_size_;
const int64_t local_conv_dim =
2 * local_q_heads * head_k_dim_ + local_v_heads * head_v_dim_;
bool used_direct_prefill_qkv = false;
mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)});
z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_});
a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_});
torch::Tensor conv_cache = kv_cache.get_conv_cache();
torch::Tensor ssm_cache = kv_cache.get_ssm_cache();
torch::Device device = mixed_qkv.device();
torch::Tensor conv_weight = conv1d_->weight();
torch::Tensor logical_state_indices =
get_linear_state_indices(input_params, device);
const int64_t checkpoint_stride =
get_checkpoint_stride(conv_cache, ssm_cache);
torch::Tensor linear_state_base_indices =
build_linear_state_base_indices(logical_state_indices, checkpoint_stride);
auto graph_context = input_params.graph.acl_graph_task_update_context;
const bool register_conv1d_graph_update =
graph_context != nullptr && graph_context->capturing;
torch::Tensor g, beta, core_attn_out, last_recurrent_state;
auto device = mixed_qkv.device();
auto conv_weight = conv1d_->weight();
auto linear_state_indices = get_linear_state_indices(input_params, device);
if (!use_spec_verify && is_any_prefill) {
torch::IntArrayRef num_accepted_tokens_opt;
std::vector<int64_t> linear_state_indices_vec(
input_params.embedding.linear_state_ids.begin(),
input_params.embedding.linear_state_ids.end());
torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv);
if (attn_metadata.is_prefill) {
mixed_qkv = mixed_qkv.transpose(1, 2);
torch::Tensor conv_state =
(seq_len < conv_kernel_size_ - 1)
? torch::pad(mixed_qkv, {0, conv_kernel_size_ - 1 - seq_len})
: (seq_len > conv_kernel_size_ - 1)
? mixed_qkv.narrow(
-1, seq_len - conv_kernel_size_ + 1, conv_kernel_size_ - 1)
: mixed_qkv;
conv_state = conv_state.transpose(1, 2).contiguous();
conv_cache.index_put_({linear_state_indices},
conv_state.to(conv_cache.dtype()));
torch::Tensor bias;
auto conv_output =
torch::conv1d(mixed_qkv,
conv_weight.unsqueeze(1).to(device),
bias,
/*stride=*/std::vector<int64_t>{1},
/*padding=*/std::vector<int64_t>{3},
/*dilation=*/std::vector<int64_t>{1},
/*groups=*/static_cast<int64_t>(mixed_qkv.size(1)));
mixed_qkv = torch::silu(conv_output.slice(2, 0, seq_len));
const bool direct_qkv_model_supported =
fla_ssm_state_layout && num_k_heads_ % tp_size_ == 0 &&
num_v_heads_ % tp_size_ == 0 && local_q_heads > 0 &&
local_v_heads > 0 && head_k_dim_ == 128 && head_v_dim_ == 128;
const bool direct_qkv_metadata_available =
attn_metadata.q_seq_lens_vec.size() ==
static_cast<size_t>(batch_size) &&
input_params.parallel.query_start_loc.size() ==
static_cast<size_t>(batch_size + 1) &&
input_params.embedding.linear_state_ids.size() ==
static_cast<size_t>(batch_size) &&
input_params.linear_state_validity_mask.size() ==
static_cast<size_t>(batch_size);
int64_t total_valid_tokens = 0;
bool direct_qkv_lengths_valid = direct_qkv_metadata_available;
if (direct_qkv_metadata_available) {
for (const int32_t valid_len : attn_metadata.q_seq_lens_vec) {
direct_qkv_lengths_valid =
direct_qkv_lengths_valid && valid_len >= 0 && valid_len <= seq_len;
total_valid_tokens += valid_len;
}
}
const bool direct_qkv_sequence_supported =
direct_qkv_model_supported && direct_qkv_lengths_valid &&
conv_input.dim() == 2 && total_valid_tokens == conv_input.size(0);
const bool direct_qkv_shape_supported =
direct_qkv_sequence_supported && conv_input.size(1) == local_conv_dim &&
conv_weight.dim() == 2 && conv_weight.size(0) == 4 &&
conv_weight.size(1) == local_conv_dim && conv_cache.dim() == 3 &&
conv_cache.size(1) >= 3 && conv_cache.size(2) == local_conv_dim;
const bool direct_qkv_dtype_supported =
direct_qkv_shape_supported &&
conv_input.scalar_type() == torch::kBFloat16 &&
conv_weight.scalar_type() == torch::kBFloat16 &&
conv_cache.scalar_type() == torch::kBFloat16;
const bool use_direct_prefill_qkv =
direct_qkv_dtype_supported && conv_input.is_contiguous() &&
conv_weight.is_contiguous() && conv_cache.is_contiguous();
if (use_direct_prefill_qkv) {
std::tie(processed_q, processed_k, processed_v) =
xllm::kernel::npu::causal_conv1d_qkv(
conv_input,
conv_weight,
conv_cache,
torch::IntArrayRef(input_params.parallel.query_start_loc),
torch::IntArrayRef(linear_state_indices_vec),
torch::IntArrayRef(input_params.linear_state_validity_mask),
local_q_heads,
local_v_heads,
head_k_dim_,
head_v_dim_);
used_direct_prefill_qkv = true;
} else {
mixed_qkv = xllm::kernel::causal_conv1d(
conv_input,
conv_weight,
conv_cache,
std::optional<torch::Tensor>(), // bias (no bias for qwen3)
torch::IntArrayRef(input_params.parallel.query_start_loc),
torch::IntArrayRef(linear_state_indices_vec),
torch::IntArrayRef(input_params.linear_state_validity_mask),
num_accepted_tokens_opt,
xllm::npu::kCausalConv1dActivationSilu,
xllm::npu::kCausalConv1dGraphPadSlotId,
xllm::npu::kCausalConv1dRunModeForward);
mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv);
mixed_qkv = mixed_qkv.transpose(1, 2);
}
} else {
if (use_spec_verify) {
CHECK(input_params.num_accepted_tokens.defined())
<< "num_accepted_tokens must be populated for Qwen3.5 spec verify";
}
torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv);
const auto& num_accepted = use_spec_verify
? input_params.num_accepted_tokens_host
: std::vector<int64_t>();
const std::vector<int64_t> linear_state_indices_host(
input_params.embedding.linear_state_ids.begin(),
input_params.embedding.linear_state_ids.end());
if (register_conv1d_graph_update) {
if (use_spec_verify) {
const auto conv1d_branch =
xllm::npu::CausalConv1dGraphBranch::kSpecVerify;
mixed_qkv = run_causal_conv1d_graph_update(
graph_context,
conv_input,
conv_weight,
conv_cache,
std::optional<torch::Tensor>(),
input_params.parallel.query_start_loc,
linear_state_indices_host,
num_accepted,
conv1d_branch);
} else {
auto conv_input_2d = conv_input.dim() == 3
? conv_input.reshape({-1, conv_input.size(-1)})
: conv_input;
xllm::kernel::CausalConv1dUpdateParams conv1d_params;
conv1d_params.x = conv_input_2d;
conv1d_params.conv_state = conv_cache;
conv1d_params.weight = conv_weight;
conv1d_params.conv_state_indices = logical_state_indices;
conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
conv1d_params.max_query_len = attn_metadata.max_query_len;
mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
if (conv_input.dim() == 3) {
mixed_qkv =
mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)});
}
}
} else {
if (use_spec_verify) {
torch::Tensor output = torch::empty_like(conv_input);
xllm::kernel::causal_conv1d_out(
output,
conv_input,
conv_weight,
conv_cache,
std::optional<torch::Tensor>(),
torch::IntArrayRef(input_params.parallel.query_start_loc),
torch::IntArrayRef(linear_state_indices_host),
torch::IntArrayRef(std::vector<int64_t>()),
torch::IntArrayRef(num_accepted),
xllm::npu::kCausalConv1dActivationSilu,
xllm::npu::kCausalConv1dGraphPadSlotId,
xllm::npu::kCausalConv1dRunModeUpdate);
mixed_qkv = output;
} else {
auto conv_input_2d = conv_input.dim() == 3
? conv_input.reshape({-1, conv_input.size(-1)})
: conv_input;
xllm::kernel::CausalConv1dUpdateParams conv1d_params;
conv1d_params.x = conv_input_2d;
conv1d_params.conv_state = conv_cache;
conv1d_params.weight = conv_weight;
conv1d_params.conv_state_indices = logical_state_indices;
conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
conv1d_params.max_query_len = attn_metadata.max_query_len;
mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
if (conv_input.dim() == 3) {
mixed_qkv =
mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)});
}
}
}
mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv);
xllm::kernel::CausalConv1dUpdateParams conv1d_params;
conv1d_params.x = mixed_qkv.reshape({-1, mixed_qkv.size(-1)});
conv1d_params.conv_state = conv_cache;
conv1d_params.weight = conv_weight;
conv1d_params.conv_state_indices = linear_state_indices;
conv1d_params.block_idx_last_scheduled_token =
std::optional<torch::Tensor>();
conv1d_params.initial_state_idx = std::optional<torch::Tensor>();
conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
conv1d_params.max_query_len = attn_metadata.max_query_len;
mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
// Reshape back to 3D [batch_size, dim, seq_len]
mixed_qkv =
mixed_qkv.view({batch_size, -1, mixed_qkv.size(-1)}).contiguous();
mixed_qkv = mixed_qkv.transpose(1, 2);
}
const bool use_fused_sigmoid_gdn_decode =
fla_ssm_state_layout && !use_spec_verify && !is_any_prefill &&
checkpoint_stride == 1;
torch::Tensor g;
torch::Tensor beta;
// Compute gated delta net decay and beta terms.
if (use_spec_verify || attn_metadata.is_chunked_prefill ||
checkpoint_stride > 1) {
beta = torch::sigmoid(b);
torch::Tensor A_log_exp = A_log_.exp();
torch::Tensor a_float = a.to(torch::kFloat32);
torch::Tensor a_plus_dt = a_float + dt_bias_;
torch::Tensor softplus_out = torch::nn::functional::softplus(
a_plus_dt,
torch::nn::functional::SoftplusFuncOptions().beta(1.0).threshold(20.0));
g = -A_log_exp * softplus_out;
g = g.to(a.dtype()).contiguous();
} else if (attn_metadata.is_prefill) {
if (attn_metadata.is_prefill) {
xllm::kernel::FusedGdnGatingParams gdn_params;
gdn_params.A_log = A_log_;
gdn_params.a = a.contiguous().view({-1, a.size(-1)});
@@ -801,7 +413,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
g = g.squeeze(0).contiguous().view({batch_size, seq_len, a.size(-1)});
beta = beta.squeeze(0).contiguous().view({batch_size, seq_len, b.size(-1)});
} else if (!use_fused_sigmoid_gdn_decode) {
} else {
xllm::kernel::FusedGdnGatingParams gdn_params;
gdn_params.A_log = A_log_;
gdn_params.a = a.view({-1, a.size(-1)});
@@ -811,216 +423,57 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
gdn_params.threshold = 20.0f;
std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
}
if (!used_direct_prefill_qkv) {
std::tie(processed_q, processed_k, processed_v) =
process_mixed_qkv(mixed_qkv);
}
torch::Tensor core_attn_out;
torch::Tensor last_recurrent_state;
auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv);
// Apply chunked or recurrent gated-delta attention and update caches.
if (use_spec_verify) {
torch::Tensor spec_num_accepted_tokens = expand_sequence_tensor_to_batch(
input_params.num_accepted_tokens.to(device, torch::kInt32),
batch_size,
"num_accepted_tokens");
torch::Tensor spec_linear_state_base_indices =
expand_sequence_tensor_to_batch(
linear_state_base_indices, batch_size, "linear_state_base_indices");
torch::Tensor step_offsets =
torch::arange(seq_len,
torch::TensorOptions()
.dtype(spec_linear_state_base_indices.dtype())
.device(device));
torch::Tensor checkpoint_indices =
spec_linear_state_base_indices.unsqueeze(1) + step_offsets;
double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
core_attn_out =
run_spec_verify_gated_delta_rule(processed_q,
processed_k,
processed_v,
g,
beta,
ssm_cache,
checkpoint_indices,
spec_num_accepted_tokens,
attn_metadata.q_cu_seq_lens,
attn_metadata.q_seq_lens_vec,
scale);
} else if (is_any_prefill) {
CHECK_GE(attn_metadata.q_seq_lens_vec.size(),
static_cast<size_t>(batch_size))
<< "q_seq_lens_vec must be populated for Qwen3.5 prefill.";
const bool use_single_prefill_pack =
batch_size == 1 && attn_metadata.q_seq_lens_vec.size() == 1 &&
attn_metadata.q_seq_lens_vec[0] == seq_len;
torch::Tensor packed_processed_q;
torch::Tensor packed_processed_k;
torch::Tensor packed_processed_v;
torch::Tensor packed_g_tensor;
torch::Tensor packed_beta_tensor;
if (use_single_prefill_pack) {
packed_processed_q = processed_q;
packed_processed_k = processed_k;
packed_processed_v = processed_v;
packed_g_tensor = g;
packed_beta_tensor = beta;
} else {
std::vector<torch::Tensor> packed_q;
std::vector<torch::Tensor> packed_k;
std::vector<torch::Tensor> packed_v;
std::vector<torch::Tensor> packed_g;
std::vector<torch::Tensor> packed_beta;
packed_q.reserve(batch_size);
packed_k.reserve(batch_size);
packed_v.reserve(batch_size);
packed_g.reserve(batch_size);
packed_beta.reserve(batch_size);
for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx];
if (!used_direct_prefill_qkv) {
packed_q.emplace_back(processed_q[batch_idx].narrow(
/*dim=*/0, /*start=*/0, valid_len));
packed_k.emplace_back(processed_k[batch_idx].narrow(
/*dim=*/0, /*start=*/0, valid_len));
packed_v.emplace_back(processed_v[batch_idx].narrow(
/*dim=*/0, /*start=*/0, valid_len));
}
packed_g.emplace_back(
g[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len));
packed_beta.emplace_back(
beta[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len));
}
if (used_direct_prefill_qkv) {
packed_processed_q = processed_q;
packed_processed_k = processed_k;
packed_processed_v = processed_v;
} else {
packed_processed_q = torch::cat(packed_q, 0).unsqueeze(0);
packed_processed_k = torch::cat(packed_k, 0).unsqueeze(0);
packed_processed_v = torch::cat(packed_v, 0).unsqueeze(0);
}
packed_g_tensor = torch::cat(packed_g, 0).unsqueeze(0);
packed_beta_tensor = torch::cat(packed_beta, 0).unsqueeze(0);
}
xllm::kernel::MegaChunkGdnParams mega_chunk_gdn_params;
mega_chunk_gdn_params.q = packed_processed_q;
mega_chunk_gdn_params.k = packed_processed_k;
mega_chunk_gdn_params.v = packed_processed_v;
mega_chunk_gdn_params.g = packed_g_tensor;
mega_chunk_gdn_params.beta = packed_beta_tensor;
if (attn_metadata.is_prefill) {
xllm::kernel::ChunkGatedDeltaRuleParams chunk_gated_delta_params;
chunk_gated_delta_params.q = processed_q;
chunk_gated_delta_params.k = processed_k;
chunk_gated_delta_params.v = processed_v;
chunk_gated_delta_params.g = g;
chunk_gated_delta_params.beta = beta;
// Get initial state from ssm_cache for sequences with previous state
// Shape: [batch_size, num_heads, head_k_dim, head_v_dim]
torch::Tensor initial_state_tensor =
torch::index_select(ssm_cache, 0, linear_state_base_indices);
CHECK_EQ(input_params.linear_state_validity_mask.size(),
input_params.embedding.linear_state_ids.size())
<< "linear state validity mask must be sequence-scoped.";
for (size_t i = 0; i < input_params.linear_state_validity_mask.size();
++i) {
if (input_params.linear_state_validity_mask[i] == 0) {
initial_state_tensor.select(0, static_cast<int64_t>(i)).fill_(0.0);
}
}
if (!fla_ssm_state_layout && attn_metadata.is_chunked_prefill) {
initial_state_tensor =
initial_state_tensor.transpose(-1, -2).contiguous();
}
mega_chunk_gdn_params.initial_state = initial_state_tensor;
mega_chunk_gdn_params.output_final_state = true;
mega_chunk_gdn_params.cu_seqlens = attn_metadata.q_cu_seq_lens;
mega_chunk_gdn_params.q_seq_lens = c10::ArrayRef<int32_t>(
attn_metadata.q_seq_lens_vec.data(), static_cast<size_t>(batch_size));
mega_chunk_gdn_params.use_qk_l2norm_in_kernel = !used_direct_prefill_qkv;
torch::Tensor packed_core_attn_out;
std::tie(packed_core_attn_out, last_recurrent_state) =
xllm::kernel::mega_chunk_gdn(mega_chunk_gdn_params);
if (use_single_prefill_pack) {
core_attn_out = packed_core_attn_out;
if (core_attn_out.scalar_type() != processed_v.scalar_type()) {
core_attn_out = core_attn_out.to(processed_v.scalar_type());
}
} else {
core_attn_out =
used_direct_prefill_qkv
? torch::zeros({batch_size, seq_len, local_v_heads, head_v_dim_},
z.options())
: torch::zeros_like(processed_v);
int64_t packed_offset = 0;
for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx];
core_attn_out[batch_idx]
.narrow(/*dim=*/0, /*start=*/0, valid_len)
.copy_(packed_core_attn_out[0].narrow(
/*dim=*/0, packed_offset, valid_len));
packed_offset += valid_len;
}
}
torch::Tensor state_to_store = fla_ssm_state_layout
? last_recurrent_state
: last_recurrent_state.transpose(-1, -2);
ssm_cache.index_put_({linear_state_base_indices},
state_to_store.to(ssm_cache.dtype()));
} else if (checkpoint_stride > 1) {
auto ssm_state =
torch::index_select(ssm_cache, 0, linear_state_base_indices);
if (!fla_ssm_state_layout) {
ssm_state = ssm_state.transpose(-1, -2);
}
ssm_state = ssm_state.contiguous();
torch::index_select(ssm_cache, 0, linear_state_indices);
// Todo: chunked-prefill/prefix-cache use initial_state
initial_state_tensor.fill_(0.0);
chunk_gated_delta_params.initial_state = initial_state_tensor;
chunk_gated_delta_params.output_final_state = true;
chunk_gated_delta_params.cu_seqlens = attn_metadata.q_cu_seq_lens;
chunk_gated_delta_params.head_first = false;
chunk_gated_delta_params.use_qk_l2norm_in_kernel = true;
std::tie(core_attn_out, last_recurrent_state) =
torch_recurrent_gated_delta_rule(
processed_q, processed_k, processed_v, g, beta, ssm_state);
torch::Tensor state_to_store = fla_ssm_state_layout
? last_recurrent_state
: last_recurrent_state.transpose(-1, -2);
ssm_cache.index_put_({linear_state_base_indices},
state_to_store.to(ssm_cache.dtype()));
xllm::kernel::chunk_gated_delta_rule(chunk_gated_delta_params);
ssm_cache.index_put_(
{linear_state_indices},
last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype()));
} else {
processed_q = xllm::kernel::l2_norm(processed_q, 1e-6);
processed_k = xllm::kernel::l2_norm(processed_k, 1e-6);
auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
torch::Tensor actual_seq_lengths =
torch::cat({zero, attn_metadata.q_seq_lens}, 0);
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.initial_state_source = ssm_cache;
params.initial_state_indices = linear_state_base_indices.contiguous();
params.cu_seqlens = attn_metadata.q_cu_seq_lens.contiguous();
params.scale = static_cast<float>(scale);
params.use_qk_l2norm_in_kernel = true;
params.softplus_beta = 1.0f;
params.softplus_threshold = 20.0f;
core_attn_out =
xllm::kernel::fused_sigmoid_gating_delta_rule_update(params);
} else {
processed_q = xllm::kernel::l2_norm(processed_q, /*eps=*/1e-6);
processed_k = xllm::kernel::l2_norm(processed_k, /*eps=*/1e-6);
auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
torch::Tensor actual_seq_lengths =
torch::cat({zero, attn_metadata.q_seq_lens}, 0);
core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
processed_q.reshape(
{-1, processed_q.size(-2), processed_q.size(-1)}),
processed_k.reshape(
{-1, processed_k.size(-2), processed_k.size(-1)}),
processed_v.reshape(
{-1, processed_v.size(-2), processed_v.size(-1)}),
ssm_cache,
beta.squeeze(0).contiguous(),
scale,
actual_seq_lengths,
logical_state_indices,
c10::nullopt,
g.squeeze(0).contiguous(),
c10::nullopt)
.unsqueeze(0)
.contiguous();
}
core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
processed_q.reshape(
{-1, processed_q.size(-2), processed_q.size(-1)}),
processed_k.reshape(
{-1, processed_k.size(-2), processed_k.size(-1)}),
processed_v.reshape(
{-1, processed_v.size(-2), processed_v.size(-1)}),
ssm_cache,
beta.squeeze(0).contiguous(),
scale,
actual_seq_lengths,
linear_state_indices,
c10::nullopt,
g.squeeze(0).contiguous(),
c10::nullopt)
.unsqueeze(0)
.contiguous();
}
auto z_reshaped = z.view({-1, z.size(-1)});
auto core_attn_out_reshaped =
core_attn_out.view({-1, core_attn_out.size(-1)});
@@ -1033,47 +486,25 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
auto rearranged_norm =
norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)});
rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm);
// For chunked prefill or spec verify, reshape_projected_tokens_with_pad may
// pad each batch to max_len, causing output tokens > original_num_tokens. We
// need to slice back to original_num_tokens to match the residual shape.
if (rearranged_norm.size(0) > original_num_tokens) {
// Slice excess padding tokens
rearranged_norm =
rearranged_norm.slice(0, 0, original_num_tokens).contiguous();
}
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
return o_proj_->forward(rearranged_norm,
row_parallel_reduce_mode_for_fc1(*fc1_ctx));
}
return o_proj_->forward(rearranged_norm);
auto attn_output = o_proj_->forward(rearranged_norm);
return attn_output;
}
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
const AttentionMetadata& attn_metadata,
const torch::Tensor& padded_qkvz) const {
const bool has_padded_queries =
attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
if (!has_padded_queries) {
if (!attn_metadata.is_prefill) {
return padded_qkvz;
}
std::vector<torch::Tensor> valid_batches;
const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty();
int64_t bs = has_host_lens
? static_cast<int64_t>(attn_metadata.q_seq_lens_vec.size())
: attn_metadata.q_seq_lens.size(0);
valid_batches.reserve(bs);
int64_t bs = attn_metadata.q_seq_lens.size(0);
int64_t max_len = attn_metadata.max_query_len;
const auto& ori_seq_lens = attn_metadata.q_seq_lens;
auto reshaped_qkvz = padded_qkvz.view({bs, max_len, -1});
for (int64_t b = 0; b < bs; ++b) {
int64_t ori_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b]
: ori_seq_lens[b].template item<int64_t>();
torch::Tensor valid_batch =
reshaped_qkvz[b].slice(/*dim=*/0, /*start=*/0, ori_len);
valid_batches.emplace_back(valid_batch);
}
if (valid_batches.size() == 1) {
return valid_batches[0].contiguous();
int64_t ori_len = ori_seq_lens[b].template item<int64_t>();
torch::Tensor valid_batch = reshaped_qkvz[b].slice(0, 0, ori_len);
valid_batches.push_back(valid_batch);
}
return torch::cat(valid_batches, 0).contiguous();
}
@@ -1081,60 +512,41 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
torch::Tensor Qwen3GatedDeltaNetBaseImpl::get_linear_state_indices(
const ModelInputParams& input_params,
const torch::Device& device) const {
CHECK(!input_params.embedding.linear_state_ids.empty())
CHECK(!input_params.linear_state_ids.empty())
<< "linear_state_ids must be populated for gated delta net";
if (input_params.embedding.linear_state_indices.defined()) {
auto indices = input_params.embedding.linear_state_indices;
if (indices.device() != device || indices.scalar_type() != torch::kInt) {
indices =
indices.to(torch::TensorOptions().dtype(torch::kInt).device(device),
/*non_blocking=*/true,
/*copy=*/true);
}
return indices.contiguous();
if (input_params.linear_state_indices.defined()) {
return input_params.linear_state_indices;
}
return torch::tensor(
input_params.embedding.linear_state_ids,
input_params.linear_state_ids,
torch::TensorOptions().dtype(torch::kInt).device(device));
}
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_projected_tokens_with_pad(
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad(
const AttentionMetadata& attn_metadata,
const torch::Tensor& projected_tokens) const {
const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty();
int64_t bs = has_host_lens
? static_cast<int64_t>(attn_metadata.q_seq_lens_vec.size())
: attn_metadata.q_seq_lens.size(0);
const torch::Tensor& qkvz) const {
int64_t bs = attn_metadata.q_seq_lens.size(0);
int64_t max_len = attn_metadata.max_query_len;
const auto& start_loc = attn_metadata.q_seq_lens;
const bool need_padding =
attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
if (!need_padding) {
return projected_tokens.view({bs, -1, projected_tokens.size(-1)});
}
if (has_host_lens && bs == 1 && attn_metadata.q_seq_lens_vec[0] == max_len &&
projected_tokens.dim() == 2 && projected_tokens.size(0) == max_len) {
return projected_tokens.view({1, max_len, projected_tokens.size(-1)});
if (!attn_metadata.is_prefill) {
return qkvz.view({qkvz.size(0), -1, qkvz.size(-1)});
}
std::vector<torch::Tensor> batches;
batches.reserve(bs);
int64_t idx = 0;
for (int64_t b = 0; b < bs; ++b) {
int64_t cur_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b]
: start_loc[b].template item<int64_t>();
torch::Tensor batch =
projected_tokens.slice(/*dim=*/0, idx, idx + cur_len).contiguous();
int64_t cur_len = start_loc[b].template item<int64_t>();
torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous();
idx = idx + cur_len;
if (batch.size(0) != max_len) {
batch = batch.size(0) > max_len
? batch.slice(/*dim=*/0, /*start=*/0, max_len).contiguous()
? batch.slice(0, 0, max_len).contiguous()
: torch::nn::functional::pad(
batch,
torch::nn::functional::PadFuncOptions(
{0, 0, 0, max_len - batch.size(0)}))
.contiguous();
}
batches.emplace_back(batch);
batches.push_back(batch);
}
auto ret = torch::stack(batches, 0).contiguous();
return ret;

38
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2025-2026 The xLLM Authors.
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -17,7 +17,6 @@ limitations under the License.
#include <torch/torch.h>
#include <optional>
#include <string>
#include <tuple>
#include <utility>
@@ -52,40 +51,19 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module {
const ModelInputParams& input_params);
protected:
virtual std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
const torch::Tensor& hidden_states) = 0;
virtual std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
const torch::Tensor& hidden_states) = 0;
// Qwen3.5 overrides this to project and reshape its separate qkv/z/b/a
// weights in every forward mode. Qwen3Next keeps qkvz/ba packed and returns
// nullopt to select the fused-split fallback.
virtual std::optional<
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
project_split_inputs(const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) {
return std::nullopt;
}
virtual bool use_fla_ssm_state_layout() const { return false; }
virtual std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) = 0;
void load_common_state_dict(const StateDict& state_dict);
void verify_common_loaded_weights(const std::string& prefix) const;
torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
const torch::Device& device) const;
std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata);
torch::Tensor reshape_qkvz_with_pad(const AttentionMetadata& attn_metadata,
const torch::Tensor& qkvz) const;
torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata,
const torch::Tensor& padded_qkvz) const;
// Projection outputs are packed as [total_tokens, dim], while GDN kernels
// consume dense [batch, max_query_len, dim] tensors. Split the packed tokens
// by query length and pad each sequence before entering the kernels.
torch::Tensor reshape_projected_tokens_with_pad(
const AttentionMetadata& attn_metadata,
const torch::Tensor& projected_tokens) const;
torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
const torch::Device& device) const;
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> process_mixed_qkv(
torch::Tensor& mixed_qkv) const;