ref(upstream): 搬运 3 大 GDN 上游仓库 — FLA naive ops + vllm GDN 子树 + xllm C++ 参考
来源:
1. fla-org/flash-linear-attention (5538 stars)
→ upstream_ref/fla/ops/gated_delta_rule/naive.py (正确的纯 PyTorch GDN)
→ upstream_ref/fla/ops/gated_delta_rule/chunk.py (Triton chunk kernel)
→ upstream_ref/fla/layers/gated_deltanet.py (层集成)
2. vllm-project/vllm main (88717 stars)
→ upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py (1751行, Qwen3.5 原生 GDN)
→ upstream_ref/vllm_gdn/ops/causal_conv1d.py (1289行, 正确的 Conv1d)
→ upstream_ref/vllm_gdn/third_party/ops/ (FLA Triton ops vendored)
→ upstream_ref/vllm_gdn/models/qwen3_5.py (vllm 最新 Qwen3.5 模型)
3. Deep-Spark/xllm (BI-V100 硬件厂商)
→ upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp (576行)
→ upstream_ref/xllm_latest/core/kernels/npu/npu_causal_conv1d.cpp
→ upstream_ref/xllm_latest/core/kernels/npu/npu_recurrent_gated_delta_rule.cpp
目的: 修复 corex_gdn.py Conv1d groups 接口不匹配问题
错误: conv1d_weight shape (2560,1,4) 被当成 (num_k_heads,1,4) 索引
conv_dim = key_dim*2 + value_dim = 10240, TP=4 后 2560
FLA naive.py 和 vllm qwen_gdn_linear_attn.py 有正确的实现可直接对接
This commit is contained in:
Claude
@@ -0,0 +1,59 @@
|
||||
/* 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
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp"
|
||||
#include "core/kernels/npu/utils.h"
|
||||
#include "core/kernels/npu/xllm_ops/xllm_ops_api.h"
|
||||
|
||||
namespace xllm::kernel::npu {
|
||||
|
||||
torch::Tensor causal_conv1d(const torch::Tensor& x,
|
||||
const torch::Tensor& weight,
|
||||
const torch::Tensor& conv_state,
|
||||
const std::optional<torch::Tensor>& bias_opt,
|
||||
const torch::IntArrayRef query_start_loc_opt,
|
||||
const torch::IntArrayRef cache_indices_opt,
|
||||
const torch::IntArrayRef initial_state_mode_opt,
|
||||
const torch::IntArrayRef num_accepted_tokens_opt,
|
||||
int64_t activation_mode,
|
||||
int64_t pad_slot_id,
|
||||
int64_t run_mode) {
|
||||
check_tensor(x, "x", "causal_conv1d");
|
||||
check_tensor(weight, "weight", "causal_conv1d");
|
||||
check_tensor(conv_state, "conv_state", "causal_conv1d");
|
||||
|
||||
c10::optional<torch::Tensor> bias_tensor = c10::nullopt;
|
||||
if (bias_opt.has_value() && bias_opt.value().defined()) {
|
||||
bias_tensor = bias_opt.value();
|
||||
}
|
||||
|
||||
torch::Tensor output = torch::empty(x.sizes(), x.options());
|
||||
EXEC_NPU_CMD(aclnnCausalConv1d,
|
||||
x,
|
||||
weight,
|
||||
bias_tensor,
|
||||
conv_state,
|
||||
query_start_loc_opt,
|
||||
cache_indices_opt,
|
||||
initial_state_mode_opt,
|
||||
num_accepted_tokens_opt,
|
||||
activation_mode,
|
||||
pad_slot_id,
|
||||
run_mode,
|
||||
output);
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::npu
|
||||
@@ -0,0 +1,83 @@
|
||||
/* 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
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp"
|
||||
#include "core/kernels/npu/npu_ops_api.h"
|
||||
#include "core/kernels/npu/utils.h"
|
||||
|
||||
namespace {
|
||||
|
||||
c10::optional<torch::Tensor> to_c10_optional_tensor(
|
||||
const std::optional<torch::Tensor>& tensor_opt) {
|
||||
if (tensor_opt.has_value() && tensor_opt.value().defined()) {
|
||||
return tensor_opt.value();
|
||||
}
|
||||
return c10::nullopt;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace xllm::kernel::npu {
|
||||
|
||||
torch::Tensor npu_recurrent_gated_delta_rule(
|
||||
const torch::Tensor& query,
|
||||
const torch::Tensor& key,
|
||||
const torch::Tensor& value,
|
||||
torch::Tensor& state,
|
||||
const std::optional<torch::Tensor>& beta,
|
||||
const std::optional<double> scale,
|
||||
const std::optional<torch::Tensor>& actual_seq_lengths,
|
||||
const std::optional<torch::Tensor>& ssm_state_indices,
|
||||
const std::optional<torch::Tensor>& num_accepted_tokens,
|
||||
const std::optional<torch::Tensor>& g,
|
||||
const std::optional<torch::Tensor>& gk) {
|
||||
check_tensor(query, "query", "recurrent_gated_delta_rule");
|
||||
check_tensor(key, "key", "recurrent_gated_delta_rule");
|
||||
check_tensor(value, "value", "recurrent_gated_delta_rule");
|
||||
check_tensor(state, "state", "recurrent_gated_delta_rule");
|
||||
CHECK(scale.has_value())
|
||||
<< "recurrent_gated_delta_rule requires a valid scale value";
|
||||
|
||||
c10::optional<torch::Tensor> beta_tensor = to_c10_optional_tensor(beta);
|
||||
c10::optional<torch::Tensor> actual_seq_lengths_tensor =
|
||||
to_c10_optional_tensor(actual_seq_lengths);
|
||||
c10::optional<torch::Tensor> ssm_state_indices_tensor =
|
||||
to_c10_optional_tensor(ssm_state_indices);
|
||||
c10::optional<torch::Tensor> num_accepted_tokens_tensor =
|
||||
to_c10_optional_tensor(num_accepted_tokens);
|
||||
c10::optional<torch::Tensor> g_tensor = to_c10_optional_tensor(g);
|
||||
c10::optional<torch::Tensor> gk_tensor = to_c10_optional_tensor(gk);
|
||||
float scale_value = static_cast<float>(scale.value());
|
||||
torch::Tensor output = torch::empty_like(value);
|
||||
|
||||
EXEC_NPU_CMD(aclnnRecurrentGatedDeltaRule,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
beta_tensor,
|
||||
state,
|
||||
actual_seq_lengths_tensor,
|
||||
ssm_state_indices_tensor,
|
||||
g_tensor,
|
||||
gk_tensor,
|
||||
num_accepted_tokens_tensor,
|
||||
scale_value,
|
||||
output);
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::npu
|
||||
236
upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.cpp
Normal file
236
upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.cpp
Normal file
@@ -0,0 +1,236 @@
|
||||
/* 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
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "qwen3_5_attention.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include <tuple>
|
||||
|
||||
#include "kernels/ops_api.h"
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
Qwen3_5AttentionImpl::Qwen3_5AttentionImpl(const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options,
|
||||
int32_t layer_id) {
|
||||
const int64_t tp_size = parallel_args.tp_group_->world_size();
|
||||
const int64_t total_num_heads = args.n_heads();
|
||||
const int64_t total_num_kv_heads = args.n_kv_heads().value_or(args.n_heads());
|
||||
layer_id_ = layer_id;
|
||||
rank_ = parallel_args.tp_group_->rank();
|
||||
CHECK(total_num_heads % tp_size == 0);
|
||||
num_heads_ = total_num_heads / tp_size;
|
||||
|
||||
if (total_num_kv_heads >= tp_size) {
|
||||
CHECK(total_num_kv_heads % tp_size == 0);
|
||||
num_kv_heads_ = total_num_kv_heads / tp_size;
|
||||
num_kv_head_replicas_ = 1;
|
||||
} else {
|
||||
CHECK(tp_size % total_num_kv_heads == 0);
|
||||
num_kv_heads_ = 1;
|
||||
num_kv_head_replicas_ = tp_size / total_num_kv_heads;
|
||||
}
|
||||
|
||||
head_dim_ = args.head_dim();
|
||||
q_size_ = num_heads_ * head_dim_;
|
||||
kv_size_ = num_kv_heads_ * head_dim_;
|
||||
scaling_ = 1.0f / std::sqrt(static_cast<float>(head_dim_));
|
||||
attn_output_gate_ = args.attn_output_gate();
|
||||
mrope_cu_seq_lens_ = torch::zeros(2, torch::kInt32).to(options.device());
|
||||
// 1. QKV linear
|
||||
qkv_proj_ = register_module(
|
||||
"qkv_proj",
|
||||
QKVParallelLinear(args.hidden_size(),
|
||||
attn_output_gate_ ? num_heads_ * 2 : num_heads_,
|
||||
num_kv_heads_,
|
||||
args.head_dim(),
|
||||
num_kv_head_replicas_,
|
||||
/*bias=*/args.attention_bias(),
|
||||
/*gather_output=*/false,
|
||||
parallel_args,
|
||||
options));
|
||||
|
||||
// 2. O proj
|
||||
o_proj_ = register_module("o_proj",
|
||||
RowParallelLinear(total_num_heads * head_dim_,
|
||||
args.hidden_size(),
|
||||
/*bias=*/false,
|
||||
/*input_is_parallelized=*/true,
|
||||
/*if_reduce_results=*/true,
|
||||
quant_args,
|
||||
parallel_args.tp_group_,
|
||||
options));
|
||||
|
||||
// 3. Q norm
|
||||
q_norm_ = register_module(
|
||||
"q_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options));
|
||||
|
||||
// 4. K norm
|
||||
k_norm_ = register_module(
|
||||
"k_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options));
|
||||
|
||||
// 5. Attention
|
||||
attn_ = register_module("attn",
|
||||
Attention(num_heads_,
|
||||
head_dim_,
|
||||
scaling_,
|
||||
num_kv_heads_,
|
||||
args.sliding_window()));
|
||||
|
||||
// 6. Rotary embedding
|
||||
const int32_t rotary_dim =
|
||||
static_cast<int32_t>(head_dim_ * args.partial_rotary_factor());
|
||||
rotary_emb_ =
|
||||
register_module("rope",
|
||||
MRotaryEmbedding(rotary_dim,
|
||||
args.max_position_embeddings(),
|
||||
args.rope_theta(),
|
||||
/*interleaved=*/false,
|
||||
args.rope_scaling_mrope_section(),
|
||||
options));
|
||||
}
|
||||
|
||||
void Qwen3_5AttentionImpl::rotary_emb_forward(
|
||||
torch::Tensor& q,
|
||||
torch::Tensor& k,
|
||||
const torch::Tensor& positions,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
auto q_shape = q.sizes();
|
||||
auto k_shape = k.sizes();
|
||||
auto num_tokens = positions.size(-1);
|
||||
mrope_cu_seq_lens_[1] = num_tokens;
|
||||
|
||||
xllm::kernel::RotaryParams rotary_params;
|
||||
bool only_prefill =
|
||||
(attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
|
||||
if (only_prefill) {
|
||||
rotary_params.sin = attn_metadata.mrope_sin;
|
||||
rotary_params.cos = attn_metadata.mrope_cos;
|
||||
rotary_params.position_ids = std::nullopt;
|
||||
rotary_params.cu_query_lens = mrope_cu_seq_lens_;
|
||||
rotary_params.interleaved = false;
|
||||
rotary_params.discrete = false;
|
||||
rotary_params.max_query_len = num_tokens;
|
||||
|
||||
rotary_params.q = q.view({num_tokens, -1, head_dim_});
|
||||
xllm::kernel::apply_rotary(rotary_params);
|
||||
q = rotary_params.q.reshape(q_shape);
|
||||
|
||||
rotary_params.q = k.view({num_tokens, -1, head_dim_});
|
||||
xllm::kernel::apply_rotary(rotary_params);
|
||||
k = rotary_params.q.reshape(k_shape);
|
||||
} else {
|
||||
if (positions.dim() == 2) {
|
||||
rotary_params.position_ids = positions[0];
|
||||
} else {
|
||||
rotary_params.position_ids = positions;
|
||||
}
|
||||
rotary_params.sin = rotary_emb_->get_sin_cache();
|
||||
rotary_params.cos = rotary_emb_->get_cos_cache();
|
||||
|
||||
rotary_params.interleaved = false;
|
||||
rotary_params.discrete = true;
|
||||
rotary_params.max_query_len = num_tokens;
|
||||
rotary_params.q = q.view({1, num_tokens, -1, head_dim_});
|
||||
xllm::kernel::apply_rotary(rotary_params);
|
||||
q = rotary_params.q.reshape(q_shape);
|
||||
|
||||
rotary_params.q = k.view({1, num_tokens, -1, head_dim_});
|
||||
xllm::kernel::apply_rotary(rotary_params);
|
||||
k = rotary_params.q.reshape(k_shape);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3_5AttentionImpl::forward(
|
||||
const torch::Tensor& positions,
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata,
|
||||
KVCache& kv_cache) {
|
||||
// 1. qkv projection
|
||||
auto qkv = qkv_proj_->forward(hidden_states);
|
||||
torch::Tensor q, k, v;
|
||||
torch::Tensor gate;
|
||||
|
||||
if (attn_output_gate_) {
|
||||
// Split qkv for attn_output_gate case: [q_size*2, kv_size, kv_size]
|
||||
auto q_gate = qkv.slice(/*dim=*/-1, 0, q_size_ * 2);
|
||||
k = qkv.slice(/*dim=*/-1, q_size_ * 2, q_size_ * 2 + kv_size_);
|
||||
v = qkv.slice(
|
||||
/*dim=*/-1, q_size_ * 2 + kv_size_, q_size_ * 2 + kv_size_ * 2);
|
||||
v = v.contiguous();
|
||||
|
||||
std::vector<int64_t> orig_shape;
|
||||
for (int64_t i = 0; i < q_gate.dim() - 1; i++) {
|
||||
orig_shape.push_back(q_gate.size(i));
|
||||
}
|
||||
std::vector<int64_t> new_shape = orig_shape;
|
||||
new_shape.push_back(num_heads_);
|
||||
new_shape.push_back(-1);
|
||||
torch::Tensor q_gate_reshaped = q_gate.reshape(new_shape);
|
||||
auto chunks = torch::chunk(q_gate_reshaped, 2, /*dim=*/-1);
|
||||
q = chunks[0];
|
||||
gate = chunks[1];
|
||||
|
||||
std::vector<int64_t> q_new_shape = orig_shape;
|
||||
q_new_shape.push_back(-1);
|
||||
q = q.reshape(q_new_shape);
|
||||
|
||||
std::vector<int64_t> gate_new_shape = orig_shape;
|
||||
gate_new_shape.push_back(-1);
|
||||
gate = gate.reshape(gate_new_shape);
|
||||
} else {
|
||||
// Normal case: [q_size, kv_size, kv_size]
|
||||
q = qkv.slice(/*dim=*/-1, 0, q_size_);
|
||||
k = qkv.slice(/*dim=*/-1, q_size_, q_size_ + kv_size_);
|
||||
v = qkv.slice(/*dim=*/-1, q_size_ + kv_size_, q_size_ + 2 * kv_size_);
|
||||
}
|
||||
|
||||
const int64_t T = q.size(0);
|
||||
|
||||
auto q_reshaped = q.reshape({T, num_heads_, head_dim_});
|
||||
auto q_normed = std::get<0>(q_norm_->forward(q_reshaped));
|
||||
auto k_reshaped = k.reshape({T, num_kv_heads_, head_dim_});
|
||||
auto k_normed = std::get<0>(k_norm_->forward(k_reshaped));
|
||||
|
||||
q = q_normed.view({T, q_size_});
|
||||
k = k_normed.view({T, kv_size_});
|
||||
rotary_emb_forward(q, k, positions, attn_metadata);
|
||||
auto out = std::get<0>(attn_->forward(attn_metadata, q, k, v, kv_cache));
|
||||
|
||||
if (attn_output_gate_) {
|
||||
gate = torch::sigmoid(gate);
|
||||
out = out * gate;
|
||||
}
|
||||
|
||||
out = o_proj_->forward(out);
|
||||
return out;
|
||||
}
|
||||
|
||||
void Qwen3_5AttentionImpl::load_state_dict(const StateDict& state_dict) {
|
||||
qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."});
|
||||
o_proj_->load_state_dict(state_dict.get_dict_with_prefix("o_proj."));
|
||||
if (auto w = state_dict.get_tensor("q_norm.weight"); w.defined()) {
|
||||
q_norm_->load_state_dict(StateDict({{"weight", w}}));
|
||||
}
|
||||
if (auto w = state_dict.get_tensor("k_norm.weight"); w.defined()) {
|
||||
k_norm_->load_state_dict(StateDict({{"weight", w}}));
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
79
upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.h
Normal file
79
upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.h
Normal file
@@ -0,0 +1,79 @@
|
||||
/* 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
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include "attention.h"
|
||||
#include "framework/kv_cache/kv_cache.h"
|
||||
#include "framework/model/model_args.h"
|
||||
#include "framework/parallel_state/parallel_args.h"
|
||||
#include "framework/quant_args.h"
|
||||
#include "framework/state_dict/state_dict.h"
|
||||
#include "layers/common/linear.h"
|
||||
#include "layers/common/partial_rotary_embedding.h"
|
||||
#include "layers/common/qwen3_next_rms_norm.h"
|
||||
#include "layers/common/rotary_embedding.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3_5AttentionImpl : public torch::nn::Module {
|
||||
public:
|
||||
Qwen3_5AttentionImpl() = default;
|
||||
Qwen3_5AttentionImpl(const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options,
|
||||
int32_t layer_id);
|
||||
|
||||
torch::Tensor forward(const torch::Tensor& positions,
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata,
|
||||
KVCache& kv_cache);
|
||||
|
||||
void load_state_dict(const StateDict& state_dict);
|
||||
void rotary_emb_forward(torch::Tensor& q,
|
||||
torch::Tensor& k,
|
||||
const torch::Tensor& positions,
|
||||
const AttentionMetadata& attn_metadata);
|
||||
|
||||
private:
|
||||
int64_t num_heads_;
|
||||
int64_t num_kv_heads_;
|
||||
int64_t num_kv_head_replicas_;
|
||||
int64_t head_dim_;
|
||||
int64_t q_size_;
|
||||
int64_t kv_size_;
|
||||
float scaling_;
|
||||
bool attn_output_gate_;
|
||||
int32_t layer_id_;
|
||||
int32_t rank_;
|
||||
|
||||
QKVParallelLinear qkv_proj_{nullptr};
|
||||
RowParallelLinear o_proj_{nullptr};
|
||||
|
||||
Qwen3NextRMSNorm q_norm_{nullptr};
|
||||
Qwen3NextRMSNorm k_norm_{nullptr};
|
||||
|
||||
Attention attn_{nullptr};
|
||||
MRotaryEmbedding rotary_emb_{nullptr};
|
||||
torch::Tensor mrope_cu_seq_lens_;
|
||||
};
|
||||
TORCH_MODULE(Qwen3_5Attention);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
@@ -0,0 +1,193 @@
|
||||
/* 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
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "qwen3_5_decoder_layer.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include "common/global_flags.h"
|
||||
#include "layers/common/dp_utils.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
namespace {
|
||||
bool use_moe_all2all(bool enable_deep_ep,
|
||||
const ModelInputParams& input_params) {
|
||||
return enable_deep_ep && all_dp_ranks_are_decode(input_params);
|
||||
}
|
||||
|
||||
bool is_moe_layer(const ModelArgs& model_args, int32_t layer_id) {
|
||||
const auto& mlp_only_layers = model_args.mlp_only_layers();
|
||||
return std::count(mlp_only_layers.begin(), mlp_only_layers.end(), layer_id) ==
|
||||
0 &&
|
||||
model_args.n_routed_experts() > 0 &&
|
||||
(layer_id + 1) % model_args.decoder_sparse_step() == 0;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
Qwen3_5DecoderLayerImpl::Qwen3_5DecoderLayerImpl(const ModelContext& context,
|
||||
int32_t layer_id)
|
||||
: parallel_args_(context.get_parallel_args()) {
|
||||
const auto& model_args = context.get_model_args();
|
||||
const auto& quant_args = context.get_quant_args();
|
||||
const auto& options = context.get_tensor_options();
|
||||
|
||||
const bool use_moe = is_moe_layer(model_args, layer_id);
|
||||
|
||||
enable_deep_ep_ = use_moe && FLAGS_expert_parallel_degree == 2;
|
||||
if (enable_deep_ep_) {
|
||||
CHECK_EQ(parallel_args_.dp_size(), parallel_args_.world_size())
|
||||
<< "Qwen3.5 MoE only support deep ep all2all when dp_size == "
|
||||
"world_size";
|
||||
CHECK_EQ(parallel_args_.dp_size(), parallel_args_.ep_size())
|
||||
<< "Qwen3.5 MoE only support deep ep all2all when dp_size == ep_size";
|
||||
}
|
||||
|
||||
auto layer_types = model_args.layer_types();
|
||||
if (layer_types.empty()) {
|
||||
int32_t interval = model_args.full_attention_interval();
|
||||
for (int32_t i = 0; i < model_args.n_layers(); i++) {
|
||||
layer_types.push_back((i + 1) % interval == 0 ? "full_attention"
|
||||
: "linear_attention");
|
||||
}
|
||||
}
|
||||
|
||||
if (layer_id >= 0 && layer_id < static_cast<int32_t>(layer_types.size())) {
|
||||
layer_type_ = layer_types[layer_id];
|
||||
} else {
|
||||
layer_type_ = "full_attention";
|
||||
}
|
||||
|
||||
if (layer_type_ == "linear_attention") {
|
||||
// TODO: support linear attention
|
||||
} else {
|
||||
full_attention_ = register_module(
|
||||
"self_attn",
|
||||
Qwen3_5Attention(
|
||||
model_args, quant_args, parallel_args_, options, layer_id));
|
||||
}
|
||||
|
||||
input_norm_ = register_module(
|
||||
"input_layernorm",
|
||||
Qwen3NextRMSNorm(
|
||||
model_args.hidden_size(), model_args.rms_norm_eps(), options));
|
||||
|
||||
post_norm_ = register_module(
|
||||
"post_attention_layernorm",
|
||||
Qwen3NextRMSNorm(
|
||||
model_args.hidden_size(), model_args.rms_norm_eps(), options));
|
||||
|
||||
if (use_moe) {
|
||||
moe_mlp_ = register_module("mlp",
|
||||
Qwen3_5FusedMoE(model_args,
|
||||
FusedMoEArgs{.is_gated = true},
|
||||
quant_args,
|
||||
parallel_args_,
|
||||
options));
|
||||
} else {
|
||||
mlp_ = register_module("mlp",
|
||||
DenseMLP(model_args.hidden_size(),
|
||||
model_args.intermediate_size(),
|
||||
true,
|
||||
false,
|
||||
model_args.hidden_act(),
|
||||
/*enable_result_reduction=*/true,
|
||||
quant_args,
|
||||
parallel_args_.tp_group_,
|
||||
options));
|
||||
}
|
||||
}
|
||||
|
||||
void Qwen3_5DecoderLayerImpl::load_state_dict(const StateDict& state_dict) {
|
||||
if (layer_type_ == "linear_attention") {
|
||||
// TODO: support linear attention
|
||||
} else {
|
||||
full_attention_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("self_attn."));
|
||||
}
|
||||
input_norm_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("input_layernorm."));
|
||||
post_norm_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("post_attention_layernorm."));
|
||||
if (moe_mlp_) {
|
||||
moe_mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp."));
|
||||
} else {
|
||||
mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp."));
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3_5DecoderLayerImpl::run_moe(
|
||||
torch::Tensor x,
|
||||
const ModelInputParams& input_params) {
|
||||
const bool enable_moe_all2all =
|
||||
use_moe_all2all(enable_deep_ep_, input_params);
|
||||
if (need_dp_moe_gather(parallel_args_, enable_moe_all2all)) {
|
||||
x = gather_dp_tokens(x, input_params, parallel_args_);
|
||||
x = moe_mlp_->forward_experts(x, enable_moe_all2all);
|
||||
return get_dp_local_slice(x, input_params, parallel_args_);
|
||||
}
|
||||
return moe_mlp_->forward_experts(x, enable_moe_all2all);
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>>
|
||||
Qwen3_5DecoderLayerImpl::apply_norm(Qwen3NextRMSNorm& norm,
|
||||
torch::Tensor& input,
|
||||
std::optional<torch::Tensor>& residual) {
|
||||
if (!residual.has_value()) {
|
||||
auto new_residual = input;
|
||||
auto output = std::get<0>(norm->forward(input));
|
||||
return {output, new_residual};
|
||||
}
|
||||
auto orig_dtype = input.dtype();
|
||||
input = input + residual.value();
|
||||
auto new_residual = input;
|
||||
input = input.to(orig_dtype);
|
||||
auto output = std::get<0>(norm->forward(input));
|
||||
return {output, new_residual};
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3_5DecoderLayerImpl::forward(
|
||||
torch::Tensor& x,
|
||||
std::optional<torch::Tensor>& residual,
|
||||
torch::Tensor& positions,
|
||||
const AttentionMetadata& attn_metadata,
|
||||
KVCache& kv_cache,
|
||||
const ModelInputParams& input_params) {
|
||||
// Pre-attention norm
|
||||
std::tie(x, residual) = apply_norm(input_norm_, x, residual);
|
||||
|
||||
// Attention
|
||||
if (full_attention_) {
|
||||
x = full_attention_->forward(positions, x, attn_metadata, kv_cache);
|
||||
} else {
|
||||
// TODO: support linear attention
|
||||
}
|
||||
|
||||
auto orig_dtype = x.dtype();
|
||||
// Post-attention norm
|
||||
std::tie(x, residual) = apply_norm(post_norm_, x, residual);
|
||||
|
||||
// MLP/MoE
|
||||
if (moe_mlp_) {
|
||||
x = run_moe(x, input_params);
|
||||
} else {
|
||||
x = mlp_->forward(x);
|
||||
}
|
||||
x = x.to(orig_dtype);
|
||||
return x;
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
@@ -0,0 +1,73 @@
|
||||
/* 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
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <optional>
|
||||
#include <string>
|
||||
|
||||
#include "framework/kv_cache/kv_cache.h"
|
||||
#include "framework/model/model_args.h"
|
||||
#include "framework/model/model_input_params.h"
|
||||
#include "framework/model_context.h"
|
||||
#include "framework/parallel_state/parallel_args.h"
|
||||
#include "framework/state_dict/state_dict.h"
|
||||
#include "layers/common/dense_mlp.h"
|
||||
#include "layers/common/qwen3_next_rms_norm.h"
|
||||
#include "layers/mlu/qwen3_5_attention.h"
|
||||
#include "layers/mlu/qwen3_5_fused_moe.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3_5DecoderLayerImpl final : public torch::nn::Module {
|
||||
public:
|
||||
Qwen3_5DecoderLayerImpl(const ModelContext& context, int32_t layer_id);
|
||||
|
||||
void load_state_dict(const StateDict& state_dict);
|
||||
|
||||
torch::Tensor forward(torch::Tensor& x,
|
||||
std::optional<torch::Tensor>& residual,
|
||||
torch::Tensor& positions,
|
||||
const AttentionMetadata& attn_metadata,
|
||||
KVCache& kv_cache,
|
||||
const ModelInputParams& input_params);
|
||||
|
||||
private:
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> apply_norm(
|
||||
Qwen3NextRMSNorm& norm,
|
||||
torch::Tensor& input,
|
||||
std::optional<torch::Tensor>& residual);
|
||||
|
||||
torch::Tensor run_moe(torch::Tensor x, const ModelInputParams& input_params);
|
||||
|
||||
std::string layer_type_;
|
||||
Qwen3_5Attention full_attention_{nullptr};
|
||||
// TODO: support linear attention
|
||||
// Qwen3_5GatedDeltaNet linear_attention_{nullptr};
|
||||
DenseMLP mlp_{nullptr};
|
||||
Qwen3_5FusedMoE moe_mlp_{nullptr};
|
||||
Qwen3NextRMSNorm input_norm_{nullptr};
|
||||
Qwen3NextRMSNorm post_norm_{nullptr};
|
||||
ParallelArgs parallel_args_;
|
||||
bool enable_deep_ep_ = false;
|
||||
};
|
||||
|
||||
TORCH_MODULE(Qwen3_5DecoderLayer);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
209
upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.cpp
Normal file
209
upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.cpp
Normal file
@@ -0,0 +1,209 @@
|
||||
/* 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
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "qwen3_5_fused_moe.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include "framework/parallel_state/parallel_state.h"
|
||||
#include "framework/state_dict/utils.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
namespace {
|
||||
torch::Tensor get_tensor_with_weight_suffix(const StateDict& state_dict,
|
||||
const std::string& tensor_name) {
|
||||
auto tensor = state_dict.get_tensor(tensor_name);
|
||||
if (!tensor.defined()) {
|
||||
tensor = state_dict.get_tensor(tensor_name + ".weight");
|
||||
}
|
||||
return tensor;
|
||||
}
|
||||
|
||||
torch::Tensor slice_expert_weights(const torch::Tensor& weight,
|
||||
int64_t start_expert_id,
|
||||
int64_t num_experts_per_rank) {
|
||||
return weight
|
||||
.slice(0, start_expert_id, start_expert_id + num_experts_per_rank)
|
||||
.contiguous();
|
||||
}
|
||||
|
||||
bool load_fused_gate_up_fallback(const StateDict& state_dict,
|
||||
int64_t rank,
|
||||
int64_t world_size,
|
||||
int64_t start_expert_id,
|
||||
int64_t num_experts_per_rank,
|
||||
torch::Tensor& w13) {
|
||||
auto fused_gate_up =
|
||||
get_tensor_with_weight_suffix(state_dict, "gate_up_proj");
|
||||
if (!fused_gate_up.defined()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (world_size > 1) {
|
||||
CHECK_EQ(fused_gate_up.size(1) % 2, 0)
|
||||
<< "gate_up_proj dim1 must be even, got " << fused_gate_up.size(1);
|
||||
const int64_t full_intermediate = fused_gate_up.size(1) / 2;
|
||||
CHECK_EQ(full_intermediate % world_size, 0)
|
||||
<< "gate_up_proj intermediate dim is not divisible by world_size";
|
||||
const int64_t inter_shard = full_intermediate / world_size;
|
||||
|
||||
auto gate_full = fused_gate_up.slice(1, 0, full_intermediate);
|
||||
auto up_full =
|
||||
fused_gate_up.slice(1, full_intermediate, full_intermediate * 2);
|
||||
auto gate_shard =
|
||||
gate_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard);
|
||||
auto up_shard =
|
||||
up_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard);
|
||||
fused_gate_up = torch::cat({gate_shard, up_shard}, 1);
|
||||
}
|
||||
|
||||
auto gate_up_slice = slice_expert_weights(
|
||||
fused_gate_up, start_expert_id, num_experts_per_rank);
|
||||
CHECK_EQ(w13.sizes(), gate_up_slice.sizes())
|
||||
<< "weight size mismatch for " << state_dict.prefix()
|
||||
<< "experts.gate_up_proj";
|
||||
w13.copy_(gate_up_slice);
|
||||
return true;
|
||||
}
|
||||
|
||||
bool load_fused_down_fallback(const StateDict& state_dict,
|
||||
int64_t rank,
|
||||
int64_t world_size,
|
||||
int64_t start_expert_id,
|
||||
int64_t num_experts_per_rank,
|
||||
torch::Tensor& w2) {
|
||||
auto fused_down = get_tensor_with_weight_suffix(state_dict, "down_proj");
|
||||
if (!fused_down.defined()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (world_size > 1) {
|
||||
CHECK_EQ(fused_down.size(2) % world_size, 0)
|
||||
<< "down_proj dim2 is not divisible by world_size";
|
||||
const int64_t down_shard = fused_down.size(2) / world_size;
|
||||
fused_down =
|
||||
fused_down.slice(2, rank * down_shard, (rank + 1) * down_shard);
|
||||
}
|
||||
|
||||
auto down_slice =
|
||||
slice_expert_weights(fused_down, start_expert_id, num_experts_per_rank);
|
||||
CHECK_EQ(w2.sizes(), down_slice.sizes())
|
||||
<< "weight size mismatch for " << state_dict.prefix()
|
||||
<< "experts.down_proj";
|
||||
w2.copy_(down_slice);
|
||||
return true;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
Qwen3_5FusedMoEImpl::Qwen3_5FusedMoEImpl(const ModelArgs& model_args,
|
||||
const FusedMoEArgs& moe_args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options)
|
||||
: FusedMoEImpl(model_args, moe_args, quant_args, parallel_args, options) {
|
||||
if (n_shared_experts_ > 0) {
|
||||
shared_expert_gate_ = register_module(
|
||||
"shared_expert_gate",
|
||||
torch::nn::Linear(
|
||||
torch::nn::LinearOptions(hidden_size_, 1).bias(false)));
|
||||
shared_expert_gate_->weight.set_data(
|
||||
shared_expert_gate_->weight.to(options));
|
||||
}
|
||||
}
|
||||
|
||||
void Qwen3_5FusedMoEImpl::load_experts(const StateDict& state_dict) {
|
||||
FusedMoEImpl::load_experts(state_dict);
|
||||
|
||||
if (!is_smoothquant_) {
|
||||
if (!w13_is_loaded_) {
|
||||
w13_is_loaded_ = load_fused_gate_up_fallback(state_dict,
|
||||
tp_pg_->rank(),
|
||||
tp_pg_->world_size(),
|
||||
start_expert_id_,
|
||||
num_experts_per_rank_,
|
||||
w13_);
|
||||
}
|
||||
|
||||
if (!w2_is_loaded_) {
|
||||
w2_is_loaded_ = load_fused_down_fallback(state_dict,
|
||||
tp_pg_->rank(),
|
||||
tp_pg_->world_size(),
|
||||
start_expert_id_,
|
||||
num_experts_per_rank_,
|
||||
w2_);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void Qwen3_5FusedMoEImpl::load_state_dict(const StateDict& state_dict) {
|
||||
if (state_dict.size() == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (n_shared_experts_ > 0) {
|
||||
shared_experts_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("shared_expert."));
|
||||
auto weight = state_dict.get_tensor("shared_expert_gate.weight");
|
||||
if (weight.defined()) {
|
||||
weight = weight.reshape({weight.size(0), -1});
|
||||
DCHECK_EQ(shared_expert_gate_->weight.sizes(), weight.sizes())
|
||||
<< "proj weight size mismatch for " << name();
|
||||
shared_expert_gate_->weight.data().copy_(weight);
|
||||
}
|
||||
}
|
||||
gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
|
||||
load_experts(state_dict.get_dict_with_prefix("experts."));
|
||||
}
|
||||
|
||||
void Qwen3_5FusedMoEImpl::final_comm_allreduce(
|
||||
torch::Tensor& final_hidden_states,
|
||||
const torch::Tensor& hidden_states,
|
||||
torch::Tensor& shared_expert_output) {
|
||||
auto current_stream = device_.current_stream();
|
||||
routed_stream_->wait_stream(*current_stream);
|
||||
{
|
||||
torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
|
||||
if (tp_pg_->world_size() > 1) {
|
||||
final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
|
||||
}
|
||||
if (parallel_args_.ep_size() > 1) {
|
||||
final_hidden_states = parallel_state::reduce(
|
||||
final_hidden_states, parallel_args_.moe_ep_group_);
|
||||
}
|
||||
}
|
||||
|
||||
if (n_shared_experts_ > 0) {
|
||||
shared_stream_->wait_stream(*current_stream);
|
||||
torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
|
||||
shared_expert_output = shared_experts_(hidden_states);
|
||||
if (shared_expert_gate_) {
|
||||
auto gate = torch::sigmoid(shared_expert_gate_->forward(hidden_states));
|
||||
shared_expert_output = gate * shared_expert_output;
|
||||
}
|
||||
shared_expert_output =
|
||||
shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
|
||||
}
|
||||
|
||||
// join for parallelization
|
||||
current_stream->wait_stream(*routed_stream_);
|
||||
if (n_shared_experts_ > 0) {
|
||||
current_stream->wait_stream(*shared_stream_);
|
||||
final_hidden_states += shared_expert_output;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
47
upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.h
Normal file
47
upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.h
Normal file
@@ -0,0 +1,47 @@
|
||||
/* 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
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "layers/mlu/fused_moe.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3_5FusedMoEImpl final : public FusedMoEImpl {
|
||||
public:
|
||||
Qwen3_5FusedMoEImpl() = default;
|
||||
|
||||
Qwen3_5FusedMoEImpl(const ModelArgs& model_args,
|
||||
const FusedMoEArgs& moe_args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
void load_state_dict(const StateDict& state_dict) override;
|
||||
|
||||
protected:
|
||||
void final_comm_allreduce(torch::Tensor& final_hidden_states,
|
||||
const torch::Tensor& hidden_states,
|
||||
torch::Tensor& shared_expert_output) override;
|
||||
|
||||
private:
|
||||
void load_experts(const StateDict& state_dict);
|
||||
torch::nn::Linear shared_expert_gate_{nullptr};
|
||||
};
|
||||
|
||||
TORCH_MODULE(Qwen3_5FusedMoE);
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
@@ -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(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user