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

upstream(xllm): sync to jd-opensource/xllm latest + revert serving_chat.py

Browse Source 
搬运 jd-opensource/xllm 最新代码到 upstream_ref/xllm_latest/:
- core/kernels/ilu/ 10 files (ixformer.h API 不变)
- core/layers/ilu/ 4 files (fused_moe.cpp config 访问从 FLAGS→singleton)
- core/layers/npu_torch/ 14 files (qwen3_gated_delta_net_base.cpp 576→1164行,
  新增 repeat_tensor_heads, checkpoint_stride, spec_verify 等 GDN 功能)
- models/llm/ 5 files (qwen3_5.h 模型注册重构, 新增 qwen3_5_mtp_base.h)
- models/vlm/ 1 file (qwen3_5.h 218→440行)

serving_chat.py: 还原到 8030a11b 原版,删掉 6dcf3590 的语法错误 min(8192,
(缺右括号导致 py_compile 失败)
This commit is contained in:
project6-dev
2026-08-12 04:22:34 +00:00
parent 6dcf3590d5
commit d025b08a95
35 changed files with 1595 additions and 593 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
qwen3_6_scripts/serving_chat.py
View File

@@ -394,7 +394,7 @@ class OpenAIServingChat(OpenAIServing):
assert prompt_inputs is not None
sampling_params: Union[SamplingParams, BeamSearchParams]
default_max_tokens = min(8192, self.max_model_len - len(
default_max_tokens = self.max_model_len - len(
prompt_inputs["prompt_token_ids"])
if request.use_beam_search:
sampling_params = request.to_beam_search_params(

2
upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp
View File

@@ -1,5 +1,5 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/kernels/ilu/group_gemm.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/kernels/ilu/ilu_ops_api.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/kernels/ilu/matmul.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/kernels/ilu/utils.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/layers/ilu/attention.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/layers/ilu/attention.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

37
upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -20,6 +20,9 @@ limitations under the License.
#include <iomanip>
#include "common/global_flags.h"
#include "core/framework/config/eplb_config.h"
#include "core/framework/config/scheduler_config.h"
#include "core/framework/config/speculative_config.h"
#include "framework/parallel_state/parallel_state.h"
#include "kernels/ops_api.h"
#include "layers/common/dp_utils.h"
@@ -86,7 +89,9 @@ FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
}
// Deep EP initialization check
enable_deep_ep_ = FLAGS_expert_parallel_degree == 2 && ep_size > 1;
enable_deep_ep_ =
::xllm::EPLBConfig::get_instance().expert_parallel_degree() == 2 &&
ep_size > 1;
if (enable_deep_ep_) {
// for now, we only implement the deep ep for decode stage.
// so we will assume the max_token_num is limited to max_batch_size * (1+K)
@@ -103,16 +108,20 @@ FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
torch::ScalarType combine_dtype = options_.dtype().toScalarType();
int64_t combine_token_size = hidden_size_ * get_dtype_size(combine_dtype);
// Ensure calculation base is at least ep_size
int64_t effective_seqs =
std::max((int64_t)FLAGS_max_seqs_per_batch, (int64_t)ep_size);
// NOTE: FLAGS_max_seqs_per_batch represents the maximum total batch size,
// regardless of the dp size. To ensure robust scheduling and account
// for the worst-case scenario, we must guarantee that each rank is capable
// of handling the maximum possible number of tokens. Therefore, we define
// max_num_tokens_per_rank as the full maximum value, without dividing by
// either the rank count or the dp size.
int64_t effective_seqs = std::max(
(int64_t)::xllm::SchedulerConfig::get_instance().max_seqs_per_batch(),
(int64_t)ep_size);
// NOTE: ::xllm::SchedulerConfig::get_instance().max_seqs_per_batch()
// represents the maximum total batch size, regardless of the dp size. To
// ensure robust scheduling and account for the worst-case scenario, we must
// guarantee that each rank is capable of handling the maximum possible
// number of tokens. Therefore, we define max_num_tokens_per_rank as the
// full maximum value, without dividing by either the rank count or the dp
// size.
int64_t max_num_tokens_per_rank =
(1 + FLAGS_num_speculative_tokens) * effective_seqs * topk_;
(1 +
::xllm::SpeculativeConfig::get_instance().num_speculative_tokens()) *
effective_seqs * topk_;
// make sure that all layers share the same deep ep instance
// so that the memory footprint is minimized
@@ -714,8 +723,8 @@ torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states,
const ModelInputParams& input_params) {
// we only support all2all communication for decode stage for now
bool enable_all2all_communication =
enable_deep_ep_ && std::all_of(input_params.dp_is_decode.begin(),
input_params.dp_is_decode.end(),
enable_deep_ep_ && std::all_of(input_params.parallel.dp_is_decode.begin(),
input_params.parallel.dp_is_decode.end(),
[](int32_t val) { return val == 1; });
bool is_dp_ep_parallel =
@@ -730,7 +739,7 @@ torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states,
if (need_gather_and_slice) {
input = parallel_state::gather(input,
parallel_args_.dp_local_process_group_,
input_params.dp_global_token_nums);
input_params.parallel.dp_global_token_nums);
}
// MoE Gate
auto router_logits = gate_(input);

2
upstream_ref/xllm_latest/core/layers/ilu/fused_moe.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

58
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
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,21 +123,55 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations(
}
std::pair<torch::Tensor, torch::Tensor>
Qwen3_5GatedDeltaNetImpl::project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) {
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));
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(
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);
}
void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict(
const StateDict& state_dict) {
auto in_proj_qkv_state_dict = state_dict.get_dict_with_prefix("in_proj_qkv.");

16
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
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,9 @@ limitations under the License.
#include <torch/torch.h>
#include <optional>
#include <string>
#include <tuple>
#include <utility>
#include "qwen3_next_gated_delta_net.h"
@@ -34,9 +36,15 @@ class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl {
const torch::TensorOptions& options);
protected:
std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) override;
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; }
void load_projection_state_dict(const StateDict& state_dict) override;
void verify_projection_weights(const std::string& prefix) const override;

858
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
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,9 +15,12 @@ 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 {
@@ -28,6 +31,31 @@ 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,
@@ -52,6 +80,9 @@ 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);
@@ -119,12 +150,15 @@ 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);
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 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);
int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size;
query = torch::nn::functional::pad(
@@ -242,6 +276,164 @@ 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(
@@ -303,7 +495,11 @@ 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)}}), shard_tensor_count, shard_sizes);
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());
}
o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj."));
if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) {
@@ -322,87 +518,279 @@ 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) {
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_);
// 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);
}
// 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;
std::tie(mixed_qkv, z, b, a) =
xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
torch::Tensor processed_q, processed_k, processed_v;
int64_t batch_size = 0;
int64_t seq_len = 0;
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_});
// 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;
torch::Tensor conv_cache = kv_cache.get_conv_cache();
torch::Tensor ssm_cache = kv_cache.get_ssm_cache();
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);
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;
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));
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);
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 {
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();
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);
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 (attn_metadata.is_prefill) {
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) {
xllm::kernel::FusedGdnGatingParams gdn_params;
gdn_params.A_log = A_log_;
gdn_params.a = a.contiguous().view({-1, a.size(-1)});
@@ -413,7 +801,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 {
} else if (!use_fused_sigmoid_gdn_decode) {
xllm::kernel::FusedGdnGatingParams gdn_params;
gdn_params.A_log = A_log_;
gdn_params.a = a.view({-1, a.size(-1)});
@@ -423,57 +811,216 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
gdn_params.threshold = 20.0f;
std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
}
auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv);
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;
// Apply chunked or recurrent gated-delta attention and update caches.
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;
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;
// 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_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;
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();
std::tie(core_attn_out, last_recurrent_state) =
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()));
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()));
} 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)));
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();
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();
}
}
auto z_reshaped = z.view({-1, z.size(-1)});
auto core_attn_out_reshaped =
core_attn_out.view({-1, core_attn_out.size(-1)});
@@ -486,25 +1033,47 @@ 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);
auto attn_output = o_proj_->forward(rearranged_norm);
return attn_output;
// 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);
}
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
const AttentionMetadata& attn_metadata,
const torch::Tensor& padded_qkvz) const {
if (!attn_metadata.is_prefill) {
const bool has_padded_queries =
attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
if (!has_padded_queries) {
return padded_qkvz;
}
std::vector<torch::Tensor> valid_batches;
int64_t bs = attn_metadata.q_seq_lens.size(0);
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 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 = 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);
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();
}
return torch::cat(valid_batches, 0).contiguous();
}
@@ -512,41 +1081,60 @@ 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.linear_state_ids.empty())
CHECK(!input_params.embedding.linear_state_ids.empty())
<< "linear_state_ids must be populated for gated delta net";
if (input_params.linear_state_indices.defined()) {
return input_params.linear_state_indices;
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();
}
return torch::tensor(
input_params.linear_state_ids,
input_params.embedding.linear_state_ids,
torch::TensorOptions().dtype(torch::kInt).device(device));
}
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad(
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_projected_tokens_with_pad(
const AttentionMetadata& attn_metadata,
const torch::Tensor& qkvz) const {
int64_t bs = attn_metadata.q_seq_lens.size(0);
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);
int64_t max_len = attn_metadata.max_query_len;
const auto& start_loc = attn_metadata.q_seq_lens;
if (!attn_metadata.is_prefill) {
return qkvz.view({qkvz.size(0), -1, qkvz.size(-1)});
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)});
}
std::vector<torch::Tensor> batches;
batches.reserve(bs);
int64_t idx = 0;
for (int64_t b = 0; b < bs; ++b) {
int64_t cur_len = start_loc[b].template item<int64_t>();
torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous();
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();
idx = idx + cur_len;
if (batch.size(0) != max_len) {
batch = batch.size(0) > max_len
? batch.slice(0, 0, max_len).contiguous()
? batch.slice(/*dim=*/0, /*start=*/0, max_len).contiguous()
: torch::nn::functional::pad(
batch,
torch::nn::functional::PadFuncOptions(
{0, 0, 0, max_len - batch.size(0)}))
.contiguous();
}
batches.push_back(batch);
batches.emplace_back(batch);
}
auto ret = torch::stack(batches, 0).contiguous();
return ret;

38
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@ limitations under the License.
#include <torch/torch.h>
#include <optional>
#include <string>
#include <tuple>
#include <utility>
@@ -51,20 +52,41 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module {
const ModelInputParams& input_params);
protected:
virtual std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) = 0;
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; }
void load_common_state_dict(const StateDict& state_dict);
void verify_common_loaded_weights(const std::string& prefix) const;
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;
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_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;
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> process_mixed_qkv(
torch::Tensor& mixed_qkv) const;

63
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -20,6 +20,8 @@ limitations under the License.
#include <tuple>
#include <vector>
#include "common/flash_comm1_context.h"
namespace xllm {
namespace layer {
@@ -63,7 +65,8 @@ Qwen3NextAttentionImpl::Qwen3NextAttentionImpl(
/*bias=*/args.attention_bias(),
/*gather_output=*/false,
parallel_args,
options));
options,
quant_args));
// 2. O proj
o_proj_ = register_module("o_proj",
@@ -144,7 +147,14 @@ torch::Tensor Qwen3NextAttentionImpl::forward(
const AttentionMetadata& attn_metadata,
KVCache& kv_cache,
const torch::Tensor& mrope_cos_sin) {
auto qkv = qkv_proj_->forward(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 qkv = qkv_proj_->forward(h);
if (use_fused_qkv_) {
const int64_t T = qkv.size(0);
@@ -168,6 +178,10 @@ torch::Tensor Qwen3NextAttentionImpl::forward(
auto out = std::get<0>(
attn_->forward(attn_metadata, q_flat, k_flat, v_flat, kv_cache));
out = out * torch::sigmoid(gate.view({T, q_size_}));
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
return o_proj_->forward(out, row_parallel_reduce_mode_for_fc1(*fc1_ctx));
}
return o_proj_->forward(out);
}
@@ -198,13 +212,18 @@ torch::Tensor Qwen3NextAttentionImpl::forward(
if (attn_output_gate_) {
out = out * torch::sigmoid(gate);
}
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
return o_proj_->forward(out, row_parallel_reduce_mode_for_fc1(*fc1_ctx));
}
return o_proj_->forward(out);
}
void Qwen3NextAttentionImpl::load_state_dict(const StateDict& state_dict) {
qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."});
if (attn_output_gate_) {
if (attn_output_gate_ && qkv_proj_->is_weight_loaded() &&
!qkv_weight_reordered_) {
// Rearrange q_proj rows from per-head interleaved [q0,g0,q1,g1,...]
// to grouped [q0,q1,...,g0,g1,...] so forward output is [Q|G|K|V].
auto w = qkv_proj_->weight();
@@ -217,6 +236,32 @@ void Qwen3NextAttentionImpl::load_state_dict(const StateDict& state_dict) {
{q_part.reshape({q_size_, hidden}), g_part.reshape({q_size_, hidden})},
0);
qg_rows.copy_(reordered);
// Reorder weight_scale and weight_offset for W8A8 dynamic quantization.
// These are per-channel (per output row) tensors that must match the
// reordered weight layout for correct dequantization.
const int64_t qg_size = q_size_ * 2;
auto reorder_per_channel = [this, qg_size](torch::Tensor tensor) {
if (!tensor.defined() || tensor.numel() == 0) {
return;
}
auto qg_part = tensor.slice(0, 0, qg_size);
auto qg_2d = qg_part.view({num_heads_, 2 * head_dim_});
auto q_scale = qg_2d.slice(1, 0, head_dim_);
auto g_scale = qg_2d.slice(1, head_dim_, 2 * head_dim_);
auto reordered_scale = torch::cat(
{q_scale.reshape({q_size_}), g_scale.reshape({q_size_})}, 0);
qg_part.copy_(reordered_scale);
};
if (qkv_proj_->is_weight_scale_loaded()) {
reorder_per_channel(qkv_proj_->weight_scale());
}
if (qkv_proj_->is_weight_offset_loaded()) {
reorder_per_channel(qkv_proj_->weight_offset());
}
qkv_weight_reordered_ = true;
}
o_proj_->load_state_dict(state_dict.get_dict_with_prefix("o_proj."));
@@ -231,8 +276,14 @@ void Qwen3NextAttentionImpl::load_state_dict(const StateDict& state_dict) {
// uses standard RMSNorm (w only). Pre-add 1 so the fused kernel produces
// the same result as Qwen3NextRMSNorm (gemma_rms_norm).
if (use_fused_qkv_) {
q_norm_->weight().add_(1.0);
k_norm_->weight().add_(1.0);
if (q_norm_->is_weight_loaded() && !q_norm_weight_adjusted_) {
q_norm_->weight().add_(1.0);
q_norm_weight_adjusted_ = true;
}
if (k_norm_->is_weight_loaded() && !k_norm_weight_adjusted_) {
k_norm_->weight().add_(1.0);
k_norm_weight_adjusted_ = true;
}
}
}

5
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -67,6 +67,9 @@ class Qwen3NextAttentionImpl : public torch::nn::Module {
float rms_norm_eps_;
bool use_fused_qkv_;
bool is_interleaved_;
bool qkv_weight_reordered_ = false;
bool q_norm_weight_adjusted_ = false;
bool k_norm_weight_adjusted_ = false;
std::vector<int64_t> mrope_section_;
torch::Tensor mrope_gather_pattern_;

2
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

2
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

17
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
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
@@ -66,13 +66,18 @@ void Qwen3NextGatedDeltaNetImpl::init_next_projections(
}
std::pair<torch::Tensor, torch::Tensor>
Qwen3NextGatedDeltaNetImpl::project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) {
Qwen3NextGatedDeltaNetImpl::project_decode_inputs(
const torch::Tensor& hidden_states) {
auto qkvz = qkvz_proj_->forward(hidden_states);
auto ba = ba_proj_->forward(hidden_states);
return {reshape_qkvz_with_pad(attn_metadata, qkvz),
reshape_qkvz_with_pad(attn_metadata, ba)};
return {qkvz.view({qkvz.size(0), -1, qkvz.size(-1)}),
ba.view({ba.size(0), -1, ba.size(-1)})};
}
std::pair<torch::Tensor, torch::Tensor>
Qwen3NextGatedDeltaNetImpl::project_flat_inputs(
const torch::Tensor& hidden_states) {
return {qkvz_proj_->forward(hidden_states), ba_proj_->forward(hidden_states)};
}
void Qwen3NextGatedDeltaNetImpl::load_state_dict(const StateDict& state_dict) {

9
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -43,9 +43,10 @@ class Qwen3NextGatedDeltaNetImpl : public Qwen3GatedDeltaNetBaseImpl {
const torch::TensorOptions& options,
bool init_projections);
std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) override;
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;
virtual void load_projection_state_dict(const StateDict& state_dict);
virtual void verify_projection_weights(const std::string& prefix) const;

23
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -19,6 +19,8 @@ limitations under the License.
#include <optional>
#include <tuple>
#include "common/flash_comm1_context.h"
namespace xllm {
namespace layer {
@@ -114,11 +116,21 @@ torch::Tensor Qwen3HybridDecoderLayerImplBase::forward(
KVCache& kv_cache,
const ModelInputParams& input_params,
const torch::Tensor& mrope_cos_sin) {
const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context();
// Pre-attention norm
if (!residual.has_value()) {
residual = x;
x = std::get<0>(input_norm_->forward(x));
} else {
if (fc1_ctx && is_sequence_sharded(*fc1_ctx) &&
residual.value().size(0) != x.size(0)) {
residual = maybe_shard_residual(residual.value(), *fc1_ctx);
}
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
CHECK_EQ(residual.value().size(0), x.size(0))
<< "FC1 input residual and hidden states must share the same "
<< "padded local sequence layout.";
}
std::tie(x, residual) = input_norm_->forward(x, residual);
}
@@ -131,6 +143,15 @@ torch::Tensor Qwen3HybridDecoderLayerImplBase::forward(
}
// Post-attention norm
// Ensure the residual layout matches the attention output before post_norm.
if (fc1_ctx && is_sequence_sharded(*fc1_ctx) && residual.has_value() &&
residual.value().size(0) != x.size(0)) {
residual = maybe_shard_residual(residual.value(), *fc1_ctx);
CHECK_EQ(residual.value().size(0), x.size(0))
<< "FC1 post-attention residual and hidden states must share the same "
<< "padded local sequence layout.";
}
std::tie(x, residual) = post_norm_->forward(x, residual);
// MLP forward

2
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

65
upstream_ref/xllm_latest/models/llm/qwen3_5.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -16,16 +16,23 @@ limitations under the License.
#pragma once
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#include "core/layers/npu_torch/qwen3_5_decoder_layer_impl.h"
#include "models/model_registry.h"
#if defined(USE_NPU) || defined(USE_MLU) || defined(USE_MUSA) || \
defined(USE_DCU)
#include "core/layers/qwen3_5_decoder_layer.h"
#include "qwen3_next.h"
#endif
namespace xllm {
#if defined(USE_NPU) || defined(USE_MLU) || defined(USE_MUSA) || \
defined(USE_DCU)
class Qwen3_5ModelImpl : public Qwen3NextModelImpl {
public:
explicit Qwen3_5ModelImpl(const ModelContext& context)
@@ -45,8 +52,24 @@ class Qwen3_5ForCausalLMImpl : public Qwen3NextForCausalLMImpl {
: Qwen3NextForCausalLMImpl(context, /*init_model=*/false) {
set_model_module(std::make_shared<Qwen3_5ModelImpl>(context));
}
torch::Tensor get_input_embeddings(torch::Tensor input_ids) {
return get_word_embedding()(input_ids);
}
void load_model(std::unique_ptr<ModelLoader> loader) {
Qwen3NextForCausalLMImpl::load_model(
std::move(loader), "model.language_model.", "lm_head.");
}
void load_model(std::unique_ptr<ModelLoader> loader,
const std::string& model_prefix) {
Qwen3NextForCausalLMImpl::load_model(
std::move(loader), model_prefix, "lm_head.");
}
};
TORCH_MODULE(Qwen3_5ForCausalLM);
#endif
#define LOAD_ARG_TEXT_OR_ROOT(arg_name, json_key, default_value) \
LOAD_ARG_OR(arg_name, "text_config." json_key, default_value); \
@@ -163,53 +186,43 @@ TORCH_MODULE(Qwen3_5ForCausalLM);
SET_ARG(topk_group, 0); \
SET_ARG(routed_scaling_factor, 1.0f); \
SET_ARG(stop_token_ids, \
std::unordered_set<int32_t>({args->eos_token_id()})); \
std::unordered_set<int32_t>({args->eos_token_id(), 248046})); \
LOAD_ARG_TEXT_OR_ROOT(mamba_ssm_dtype, "mamba_ssm_dtype", "float32")
#define LOAD_QWEN3_5_TYPE_AND_DTYPE(default_model_type) \
LOAD_ARG_OR(model_type, "model_type", default_model_type); \
#define LOAD_QWEN3_5_TEXT_TYPE_AND_DTYPE(default_model_type) \
SET_ARG(model_type, default_model_type); \
LOAD_ARG_OR(dtype, "text_config.dtype", "bfloat16"); \
LOAD_ARG_OR(dtype, "dtype", args->dtype()); \
LOAD_ARG_OR(dtype, "text_config.torch_dtype", args->dtype()); \
LOAD_ARG_OR(dtype, "torch_dtype", args->dtype())
REGISTER_CAUSAL_MODEL(qwen3_5, Qwen3_5ForCausalLM);
REGISTER_MODEL_ARGS(qwen3_5, [&] {
LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5");
LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/0,
/*num_experts=*/0,
/*num_experts_per_tok=*/0,
/*shared_expert_intermediate_size=*/0);
});
REGISTER_MODEL_BACKEND(qwen3_5_text, "llm");
#if defined(USE_NPU) || defined(USE_MLU) || defined(USE_MUSA) || \
defined(USE_DCU)
REGISTER_CAUSAL_MODEL(qwen3_5_text, Qwen3_5ForCausalLM);
#endif
REGISTER_MODEL_ARGS(qwen3_5_text, [&] {
LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_text");
LOAD_QWEN3_5_TEXT_TYPE_AND_DTYPE("qwen3_5_text");
LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/0,
/*num_experts=*/0,
/*num_experts_per_tok=*/0,
/*shared_expert_intermediate_size=*/0);
});
REGISTER_CAUSAL_MODEL(qwen3_5_moe, Qwen3_5ForCausalLM);
REGISTER_MODEL_ARGS(qwen3_5_moe, [&] {
LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_moe");
LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/512,
/*num_experts=*/512,
/*num_experts_per_tok=*/10,
/*shared_expert_intermediate_size=*/512);
});
REGISTER_MODEL_BACKEND(qwen3_5_moe_text, "llm");
#if defined(USE_NPU) || defined(USE_MLU) || defined(USE_MUSA) || \
defined(USE_DCU)
REGISTER_CAUSAL_MODEL(qwen3_5_moe_text, Qwen3_5ForCausalLM);
#endif
REGISTER_MODEL_ARGS(qwen3_5_moe_text, [&] {
LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_moe_text");
LOAD_QWEN3_5_TEXT_TYPE_AND_DTYPE("qwen3_5_moe_text");
LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/512,
/*num_experts=*/512,
/*num_experts_per_tok=*/10,
/*shared_expert_intermediate_size=*/512);
});
#undef LOAD_QWEN3_5_TYPE_AND_DTYPE
#undef LOAD_QWEN3_5_TEXT_TYPE_AND_DTYPE
#undef LOAD_QWEN3_5_NEXT_COMPAT_ARGS
#undef LOAD_QWEN3_5_ROPE_ARG
#undef LOAD_ARG_TEXT_OR_ROOT_CHAIN

253
upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -15,250 +15,26 @@ limitations under the License.
#pragma once
#include <glog/logging.h>
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include "core/framework/model/model_input_params.h"
#include "core/layers/common/linear.h"
#include "models/llm/qwen3_5.h"
#include "models/llm/qwen3_5_mtp_base.h"
#include "models/model_registry.h"
#include "qwen3_5.h"
namespace xllm {
namespace {
StateDict find_lm_head_state_dict(const StateDict& state_dict) {
static const std::vector<std::string> kLmHeadPrefixes = {
"lm_head.",
"model.lm_head.",
"language_model.lm_head.",
"model.language_model.lm_head."};
for (const auto& prefix : kLmHeadPrefixes) {
auto sub_dict = state_dict.get_dict_with_prefix(prefix);
if (sub_dict.get_tensor("weight").defined() ||
sub_dict.get_tensor("qweight").defined()) {
return sub_dict;
}
}
return StateDict({}, "");
}
bool load_qwen3_5_mtp_model_args(const JsonReader& json,
ModelArgs* args,
const std::string& base_model_type,
const std::string& mtp_model_type) {
auto base_loader = ModelRegistry::get_model_args_loader(base_model_type);
if (base_loader == nullptr || base_loader(json, args) == false) {
return false;
}
int32_t mtp_num_layers = args->num_nextn_predict_layers();
if (mtp_num_layers <= 0) {
mtp_num_layers = 1;
}
args->model_type(mtp_model_type);
args->num_nextn_predict_layers(mtp_num_layers);
args->n_layers(mtp_num_layers);
args->layer_types(std::vector<std::string>(
static_cast<size_t>(mtp_num_layers), "full_attention"));
return true;
}
} // namespace
class Qwen3_5MtpModelImpl : public Qwen3HybridModelImplBase {
class Qwen3_5MtpModelImpl final : public Qwen3_5MtpModelImplBase {
public:
explicit Qwen3_5MtpModelImpl(const ModelContext& context)
: Qwen3HybridModelImplBase(context) {
const auto& options = context.get_tensor_options();
const int32_t n_layers =
std::max<int32_t>(static_cast<int32_t>(model_args_.n_layers()), 1);
pre_fc_norm_embedding_ = register_module(
"pre_fc_norm_embedding",
layer::Qwen3NextRMSNorm(
model_args_.hidden_size(), model_args_.rms_norm_eps(), options));
pre_fc_norm_hidden_ = register_module(
"pre_fc_norm_hidden",
layer::Qwen3NextRMSNorm(
model_args_.hidden_size(), model_args_.rms_norm_eps(), options));
fc_ = register_module("fc",
layer::ReplicatedLinear(model_args_.hidden_size() * 2,
model_args_.hidden_size(),
/*bias=*/false,
QuantArgs(),
options));
layers_.reserve(n_layers);
for (int32_t layer_id = 0; layer_id < n_layers; ++layer_id) {
add_decoder_layer(
std::make_shared<layer::Qwen3_5DecoderLayerImpl>(context, layer_id));
}
}
ModelOutput forward(torch::Tensor tokens,
torch::Tensor positions,
std::vector<KVCache>& kv_caches,
const ModelInputParams& input_params) override {
torch::NoGradGuard no_grad;
if (dp_size_ > 1 && tokens.sizes() == 0) {
tokens = torch::tensor({1}).to(torch::kInt32).to(device_);
positions = torch::tensor({0}).to(torch::kInt32).to(device_);
}
auto attn_metadata = layer::AttentionMetadataBuilder::build(
input_params,
model_args_.enable_mla(),
build_attention_mask(input_params));
torch::Tensor embedding = embed_tokens_(tokens);
torch::Tensor hidden = input_params.input_embedding;
if (hidden.defined() == false) {
hidden = embedding;
}
embedding = std::get<0>(pre_fc_norm_embedding_->forward(embedding));
hidden = std::get<0>(pre_fc_norm_hidden_->forward(hidden));
torch::Tensor mtp_hidden = fc_(torch::cat({embedding, hidden}, -1));
CHECK_EQ(kv_caches.size(), layers_.size());
std::optional<torch::Tensor> residual = std::nullopt;
for (size_t i = 0; i < layers_.size(); ++i) {
mtp_hidden = layers_[i]->forward(mtp_hidden,
residual,
positions,
attn_metadata,
kv_caches[i],
input_params);
}
auto [new_mtp_hidden, new_res] = norm_->forward(mtp_hidden, residual);
mtp_hidden = new_mtp_hidden;
return ModelOutput(mtp_hidden);
}
void load_state_dict(const StateDict& state_dict) override {
load_shared_embeddings(state_dict);
load_mtp_state_dict(state_dict);
}
void load_shared_embeddings(const StateDict& state_dict) {
auto embedding_state_dict =
state_dict.get_dict_with_prefix("embed_tokens.");
if (embedding_state_dict.get_tensor("weight").defined()) {
shared_embedding_loaded_ = true;
}
embed_tokens_->load_state_dict(embedding_state_dict);
}
void load_mtp_state_dict(const StateDict& state_dict) {
if (state_dict.get_tensor("pre_fc_norm_embedding.weight").defined()) {
pre_fc_norm_embedding_loaded_ = true;
}
if (state_dict.get_tensor("pre_fc_norm_hidden.weight").defined()) {
pre_fc_norm_hidden_loaded_ = true;
}
if (state_dict.get_tensor("fc.weight").defined() ||
state_dict.get_tensor("fc.qweight").defined()) {
fc_loaded_ = true;
}
if (state_dict.get_tensor("norm.weight").defined()) {
norm_loaded_ = true;
}
pre_fc_norm_embedding_->load_state_dict(
state_dict.get_dict_with_prefix("pre_fc_norm_embedding."));
pre_fc_norm_hidden_->load_state_dict(
state_dict.get_dict_with_prefix("pre_fc_norm_hidden."));
fc_->load_state_dict(state_dict.get_dict_with_prefix("fc."));
for (size_t i = 0; i < layers_.size(); ++i) {
layers_[i]->load_state_dict(
state_dict.get_dict_with_prefix("layers." + std::to_string(i) + "."));
}
norm_->load_state_dict(state_dict.get_dict_with_prefix("norm."));
}
void verify_loaded_weights(const std::string& prefix) const override {
CHECK(shared_embedding_loaded_)
<< "Failed to find shared embedding weights for qwen3.5 mtp draft "
"model";
CHECK(pre_fc_norm_embedding_loaded_)
<< "Failed to find mtp pre_fc_norm_embedding weights for qwen3.5 mtp "
"draft model";
CHECK(pre_fc_norm_hidden_loaded_)
<< "Failed to find mtp pre_fc_norm_hidden weights for qwen3.5 mtp "
"draft model";
CHECK(fc_loaded_) << "Failed to find mtp fc weights for qwen3.5 mtp draft "
"model";
CHECK(norm_loaded_)
<< "Failed to find mtp norm weights for qwen3.5 mtp draft model";
for (size_t i = 0; i < layers_.size(); ++i) {
layers_[i]->verify_loaded_weights(prefix + "layers." + std::to_string(i) +
".");
}
}
private:
layer::Qwen3NextRMSNorm pre_fc_norm_embedding_{nullptr};
layer::Qwen3NextRMSNorm pre_fc_norm_hidden_{nullptr};
layer::ReplicatedLinear fc_{nullptr};
bool shared_embedding_loaded_ = false;
bool pre_fc_norm_embedding_loaded_ = false;
bool pre_fc_norm_hidden_loaded_ = false;
bool fc_loaded_ = false;
bool norm_loaded_ = false;
: Qwen3_5MtpModelImplBase(context) {}
};
class Qwen3_5MtpForCausalLMImpl : public Qwen3HybridForCausalLMImplBase {
class Qwen3_5MtpForCausalLMImpl final : public Qwen3_5MtpForCausalLMImplBase {
public:
explicit Qwen3_5MtpForCausalLMImpl(const ModelContext& context)
: Qwen3HybridForCausalLMImplBase(context) {
mtp_model_ = std::make_shared<Qwen3_5MtpModelImpl>(context);
set_model_module(mtp_model_);
}
void load_model(std::unique_ptr<ModelLoader> loader) {
static const std::vector<std::string> kEmbeddingPrefixes = {
"model.language_model.", "language_model.model.", "model.", ""};
static const std::vector<std::string> kMtpPrefixes = {"mtp.", "model.mtp."};
bool lm_head_loaded = false;
for (const auto& state_dict : loader->get_state_dicts()) {
auto shared_embedding_state_dict =
state_dict->get_dict_with_prefix(kEmbeddingPrefixes);
auto mtp_state_dict = state_dict->get_dict_with_prefix(kMtpPrefixes);
mtp_model_->load_shared_embeddings(shared_embedding_state_dict);
mtp_model_->load_mtp_state_dict(mtp_state_dict);
if (tie_word_embeddings_) {
lm_head_->load_state_dict(
shared_embedding_state_dict.get_dict_with_prefix("embed_tokens."));
if (shared_embedding_state_dict.get_tensor("embed_tokens.weight")
.defined()) {
lm_head_loaded = true;
}
} else {
auto lm_head_state_dict = find_lm_head_state_dict(*state_dict);
lm_head_->load_state_dict(lm_head_state_dict);
if (lm_head_state_dict.get_tensor("weight").defined() ||
lm_head_state_dict.get_tensor("qweight").defined()) {
lm_head_loaded = true;
}
}
}
CHECK(lm_head_loaded)
<< "Failed to find lm_head weights for qwen3.5 mtp draft model";
mtp_model_->verify_loaded_weights("mtp.");
}
private:
std::shared_ptr<Qwen3_5MtpModelImpl> mtp_model_;
: Qwen3_5MtpForCausalLMImplBase(
context,
std::make_shared<Qwen3_5MtpModelImpl>(context)) {}
};
TORCH_MODULE(Qwen3_5MtpForCausalLM);
@@ -267,14 +43,17 @@ REGISTER_CAUSAL_MODEL(qwen3_5_moe_mtp, Qwen3_5MtpForCausalLM);
REGISTER_MODEL_ARGS_LOADER(qwen3_5_mtp,
[](const JsonReader& json, ModelArgs* args) {
return load_qwen3_5_mtp_model_args(
json, args, "qwen3_5", "qwen3_5_mtp");
return qwen3_5_mtp::load_model_args(
json, args, "qwen3_5_text", "qwen3_5_mtp");
});
REGISTER_MODEL_ARGS_LOADER(qwen3_5_moe_mtp,
[](const JsonReader& json, ModelArgs* args) {
return load_qwen3_5_mtp_model_args(
json, args, "qwen3_5_moe", "qwen3_5_moe_mtp");
return qwen3_5_mtp::load_model_args(
json,
args,
"qwen3_5_moe_text",
"qwen3_5_moe_mtp");
});
} // namespace xllm

299
upstream_ref/xllm_latest/models/llm/qwen3_5_mtp_base.h Normal file
View File

@@ -0,0 +1,299 @@
/* 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 <glog/logging.h>
#include <algorithm>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "core/layers/common/linear.h"
#include "core/layers/qwen3_5_decoder_layer.h"
#include "models/llm/qwen3_next_hybrid_base.h"
#include "models/model_registry.h"
namespace xllm {
namespace qwen3_5_mtp {
inline StateDict get_lm_head_dict(const StateDict& state_dict) {
static const std::vector<std::string> kLmHeadPrefixes = {
"lm_head.",
"model.lm_head.",
"language_model.lm_head.",
"model.language_model.lm_head."};
for (const std::string& prefix : kLmHeadPrefixes) {
StateDict sub_dict = state_dict.get_dict_with_prefix(prefix);
if (sub_dict.get_tensor("weight").defined() ||
sub_dict.get_tensor("qweight").defined()) {
return sub_dict;
}
}
return StateDict({}, "");
}
inline bool load_model_args(const JsonReader& json,
ModelArgs* args,
const std::string& base_type,
const std::string& mtp_type) {
ModelArgsLoader base_loader = ModelRegistry::get_model_args_loader(base_type);
if (base_loader == nullptr || base_loader(json, args) == false) {
return false;
}
int32_t mtp_num_layers = args->num_nextn_predict_layers();
if (mtp_num_layers <= 0) {
mtp_num_layers = 1;
}
args->model_type(mtp_type);
args->num_nextn_predict_layers(mtp_num_layers);
args->n_layers(mtp_num_layers);
args->layer_types(std::vector<std::string>(
static_cast<size_t>(mtp_num_layers), "full_attention"));
return true;
}
} // namespace qwen3_5_mtp
class Qwen3_5MtpModelImplBase : public Qwen3HybridModelImplBase {
public:
explicit Qwen3_5MtpModelImplBase(const ModelContext& context)
: Qwen3HybridModelImplBase(context) {
const torch::TensorOptions& options = context.get_tensor_options();
const int32_t n_layers =
std::max<int32_t>(static_cast<int32_t>(model_args_.n_layers()), 1);
pre_fc_norm_embedding_ = register_module(
"pre_fc_norm_embedding",
layer::Qwen3NextRMSNorm(
model_args_.hidden_size(), model_args_.rms_norm_eps(), options));
pre_fc_norm_hidden_ = register_module(
"pre_fc_norm_hidden",
layer::Qwen3NextRMSNorm(
model_args_.hidden_size(), model_args_.rms_norm_eps(), options));
fc_ = register_module("fc",
layer::ReplicatedLinear(model_args_.hidden_size() * 2,
model_args_.hidden_size(),
/*bias=*/false,
QuantArgs(),
options));
layers_.reserve(n_layers);
for (int32_t layer_id = 0; layer_id < n_layers; ++layer_id) {
add_decoder_layer(
std::make_shared<layer::Qwen3_5DecoderLayerImpl>(context, layer_id));
}
}
ModelOutput forward(torch::Tensor tokens,
torch::Tensor positions,
std::vector<KVCache>& kv_caches,
const ModelInputParams& input_params) override {
torch::NoGradGuard no_grad;
if (dp_size_ > 1 && tokens.sizes() == 0) {
tokens = torch::tensor({1}).to(torch::kInt32).to(device_);
positions = torch::tensor({0}).to(torch::kInt32).to(device_);
}
layer::AttentionMetadata attn_metadata =
layer::AttentionMetadataBuilder::build(
input_params,
model_args_.enable_mla(),
build_attention_mask(input_params),
/*device=*/device_);
prepare_mrope(positions, attn_metadata);
torch::Tensor embedding = embed_tokens_(tokens);
torch::Tensor hidden = input_params.embedding.input_embedding;
if (hidden.defined() == false) {
hidden = embedding;
}
embedding = std::get<0>(pre_fc_norm_embedding_->forward(embedding));
hidden = std::get<0>(pre_fc_norm_hidden_->forward(hidden));
torch::Tensor mtp_hidden = fc_(torch::cat({embedding, hidden}, -1));
CHECK_EQ(kv_caches.size(), layers_.size());
torch::Tensor mrope_cos_sin;
for (const layer::Qwen3HybridDecoderLayerModulePtr& layer : layers_) {
mrope_cos_sin = layer->build_mrope_cos_sin(positions);
if (mrope_cos_sin.defined()) {
break;
}
}
std::optional<torch::Tensor> residual = std::nullopt;
for (size_t i = 0; i < layers_.size(); ++i) {
if (!input_params.synchronize_layer(static_cast<uint32_t>(i))) {
return ModelOutput();
}
mtp_hidden = layers_[i]->forward(mtp_hidden,
residual,
positions,
attn_metadata,
kv_caches[i],
input_params,
mrope_cos_sin);
#if defined(USE_NPU)
if (input_params.parallel.layer_synchronizer != nullptr &&
!input_params.parallel.layer_synchronizer->record_event(
static_cast<int64_t>(i), device_.index())) {
return ModelOutput();
}
#endif
}
auto [new_mtp_hidden, new_res] = norm_->forward(mtp_hidden, residual);
mtp_hidden = new_mtp_hidden;
return ModelOutput(mtp_hidden);
}
void load_state_dict(const StateDict& state_dict) override {
load_shared_embeddings(state_dict);
load_mtp_state_dict(state_dict);
}
void load_shared_embeddings(const StateDict& state_dict) {
StateDict embedding_state_dict =
state_dict.get_dict_with_prefix("embed_tokens.");
if (embedding_state_dict.get_tensor("weight").defined()) {
shared_embedding_loaded_ = true;
}
embed_tokens_->load_state_dict(embedding_state_dict);
}
void load_mtp_state_dict(const StateDict& state_dict) {
if (state_dict.get_tensor("pre_fc_norm_embedding.weight").defined()) {
pre_fc_norm_embedding_loaded_ = true;
}
if (state_dict.get_tensor("pre_fc_norm_hidden.weight").defined()) {
pre_fc_norm_hidden_loaded_ = true;
}
if (state_dict.get_tensor("fc.weight").defined() ||
state_dict.get_tensor("fc.qweight").defined()) {
fc_loaded_ = true;
}
if (state_dict.get_tensor("norm.weight").defined()) {
norm_loaded_ = true;
}
pre_fc_norm_embedding_->load_state_dict(
state_dict.get_dict_with_prefix("pre_fc_norm_embedding."));
pre_fc_norm_hidden_->load_state_dict(
state_dict.get_dict_with_prefix("pre_fc_norm_hidden."));
fc_->load_state_dict(state_dict.get_dict_with_prefix("fc."));
for (size_t i = 0; i < layers_.size(); ++i) {
layers_[i]->load_state_dict(
state_dict.get_dict_with_prefix("layers." + std::to_string(i) + "."));
}
norm_->load_state_dict(state_dict.get_dict_with_prefix("norm."));
}
void verify_loaded_weights(const std::string& prefix) const override {
CHECK(shared_embedding_loaded_)
<< "Failed to find shared embedding weights for qwen3.5 mtp draft "
"model";
CHECK(pre_fc_norm_embedding_loaded_)
<< "Failed to find mtp pre_fc_norm_embedding weights for qwen3.5 mtp "
"draft model";
CHECK(pre_fc_norm_hidden_loaded_)
<< "Failed to find mtp pre_fc_norm_hidden weights for qwen3.5 mtp "
"draft model";
CHECK(fc_loaded_) << "Failed to find mtp fc weights for qwen3.5 mtp draft "
"model";
CHECK(norm_loaded_)
<< "Failed to find mtp norm weights for qwen3.5 mtp draft model";
for (size_t i = 0; i < layers_.size(); ++i) {
layers_[i]->verify_loaded_weights(prefix + "layers." + std::to_string(i) +
".");
}
}
protected:
virtual void prepare_mrope(const torch::Tensor& positions,
layer::AttentionMetadata& attn_metadata) const {
UNUSED_PARAMETER(positions);
UNUSED_PARAMETER(attn_metadata);
}
private:
layer::Qwen3NextRMSNorm pre_fc_norm_embedding_{nullptr};
layer::Qwen3NextRMSNorm pre_fc_norm_hidden_{nullptr};
layer::ReplicatedLinear fc_{nullptr};
bool shared_embedding_loaded_ = false;
bool pre_fc_norm_embedding_loaded_ = false;
bool pre_fc_norm_hidden_loaded_ = false;
bool fc_loaded_ = false;
bool norm_loaded_ = false;
};
class Qwen3_5MtpForCausalLMImplBase : public Qwen3HybridForCausalLMImplBase {
public:
void load_model(std::unique_ptr<ModelLoader> loader) {
static const std::vector<std::string> kEmbeddingPrefixes = {
"model.language_model.", "language_model.model.", "model.", ""};
static const std::vector<std::string> kMtpPrefixes = {"mtp.", "model.mtp."};
bool lm_head_loaded = false;
for (const std::unique_ptr<StateDict>& state_dict :
loader->get_state_dicts()) {
StateDict shared_embedding_state_dict =
state_dict->get_dict_with_prefix(kEmbeddingPrefixes);
StateDict mtp_state_dict = state_dict->get_dict_with_prefix(kMtpPrefixes);
mtp_model_->load_shared_embeddings(shared_embedding_state_dict);
mtp_model_->load_mtp_state_dict(mtp_state_dict);
if (tie_word_embeddings_) {
lm_head_->load_state_dict(
shared_embedding_state_dict.get_dict_with_prefix("embed_tokens."));
if (shared_embedding_state_dict.get_tensor("embed_tokens.weight")
.defined()) {
lm_head_loaded = true;
}
} else {
StateDict lm_head_state_dict =
qwen3_5_mtp::get_lm_head_dict(*state_dict);
lm_head_->load_state_dict(lm_head_state_dict);
if (lm_head_state_dict.get_tensor("weight").defined() ||
lm_head_state_dict.get_tensor("qweight").defined()) {
lm_head_loaded = true;
}
}
}
CHECK(lm_head_loaded)
<< "Failed to find lm_head weights for qwen3.5 mtp draft model";
mtp_model_->verify_loaded_weights("mtp.");
}
protected:
Qwen3_5MtpForCausalLMImplBase(
const ModelContext& context,
std::shared_ptr<Qwen3_5MtpModelImplBase> mtp_model)
: Qwen3HybridForCausalLMImplBase(context),
mtp_model_(std::move(mtp_model)) {
set_model_module(mtp_model_);
}
private:
std::shared_ptr<Qwen3_5MtpModelImplBase> mtp_model_;
};
} // namespace xllm

2
upstream_ref/xllm_latest/models/llm/qwen3_next.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.

151
upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -22,18 +22,23 @@ limitations under the License.
#include <string>
#include <vector>
#include "core/common/global_flags.h"
#include "core/common/flash_comm1_context.h"
#include "core/framework/kv_cache/kv_cache.h"
#include "core/framework/model/model_input_params.h"
#include "core/framework/model/model_output.h"
#include "core/framework/model_context.h"
#include "core/framework/model_loader.h"
#include "core/framework/parallel_state/parallel_args.h"
#include "core/layers/common/attention_mask.h"
#include "core/layers/common/attention_metadata_builder.h"
#include "core/layers/common/lm_head.h"
#include "core/layers/common/qwen3_next_rms_norm.h"
#include "core/layers/common/word_embedding.h"
#if defined(USE_NPU)
#include "core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h"
#elif defined(USE_MLU)
#include "core/layers/mlu/qwen3_5/qwen3_5_hybrid_decoder_layer_base.h"
#endif
namespace xllm {
@@ -55,7 +60,14 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
public:
explicit Qwen3HybridModelImplBase(const ModelContext& context)
: device_(context.get_tensor_options().device()),
model_args_(context.get_model_args()) {
model_args_(context.get_model_args()),
parallel_args_(context.get_parallel_args()),
flash_comm1_options_(context.get_flash_comm1_options()) {
if (model_args_.n_routed_experts() > 0) {
flash_comm1_options_.enable_flashcomm1 = false;
flash_comm1_options_.enable_mmrs_fusion = false;
}
auto options = context.get_tensor_options();
auto parallel_args = context.get_parallel_args();
@@ -69,10 +81,12 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
model_args_.hidden_size(), model_args_.rms_norm_eps(), options));
embed_tokens_ =
register_module("embed_tokens", layer::WordEmbedding(context));
int32_t mask_value = FLAGS_enable_chunked_prefill ? -9984 : 1;
attn_mask_ = layer::AttentionMask(options.device(),
options.dtype().toScalarType(),
/*mask_value=*/mask_value);
/*mask_value=*/-9984);
dense_attn_mask_ = layer::AttentionMask(options.device(),
options.dtype().toScalarType(),
/*mask_value=*/1);
dp_size_ = parallel_args.dp_size();
}
@@ -95,8 +109,25 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
layer::AttentionMetadataBuilder::build(
input_params,
model_args_.enable_mla(),
build_attention_mask(input_params));
torch::Tensor h = embed_tokens_(tokens);
build_attention_mask(input_params),
/*device=*/device_);
const int32_t num_tokens = static_cast<int32_t>(tokens.size(0));
const auto& batch_forward_type = input_params.meta.batch_forward_type;
const bool is_prefill_side = batch_forward_type.no_decode();
FlashComm1Context fc1_ctx = build_flash_comm1_context(
num_tokens, is_prefill_side, parallel_args_, flash_comm1_options_);
FlashComm1ContextScope fc1_scope(&fc1_ctx);
torch::Tensor h;
if (input_params.embedding.input_embedding.defined()) {
h = input_params.embedding.input_embedding;
} else {
h = embed_tokens_(tokens);
}
if (is_sequence_sharded(fc1_ctx)) {
h = shard_sequence(h, fc1_ctx);
}
torch::Tensor mrope_cos_sin;
for (const auto& layer : layers_) {
@@ -114,9 +145,19 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
kv_caches[i],
input_params,
mrope_cos_sin);
#if defined(USE_NPU)
if (input_params.parallel.layer_synchronizer != nullptr &&
!input_params.parallel.layer_synchronizer->record_event(
static_cast<int64_t>(i), device_.index())) {
return ModelOutput();
}
#endif
}
auto [hidden_states, residual_out] = norm_->forward(h, residual);
h = hidden_states;
if (is_sequence_sharded(fc1_ctx)) {
h = gather_sequence(h, fc1_ctx);
}
return ModelOutput(h);
}
@@ -155,27 +196,48 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
protected:
torch::Tensor build_attention_mask(const ModelInputParams& input_params) {
max_seq_len_ = std::max(input_params.kv_max_seq_len, max_seq_len_);
if (!FLAGS_enable_chunked_prefill) {
return attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_);
#if defined(USE_NPU)
// On NPU the hybrid path never consumes attn_metadata.attn_mask: full
// attention runs through the fused-infer / paged-attention kernels (which
// carry their own fixed fia_attn_mask or need no mask at all) and linear
// attention is mask-free by construction. Materializing a dense
// [seq_len, seq_len] mask here is pure waste and, for long sequences,
// triggers an NPU OOM. Hand the kernels an empty mask unless a graph buffer
// already supplies one.
if (input_params.graph.attn_mask.defined()) {
return input_params.graph.attn_mask;
}
return torch::Tensor();
#else
if (input_params.graph.attn_mask.defined()) {
return input_params.graph.attn_mask;
}
max_seq_len_ = std::max(input_params.meta.kv_max_seq_len, max_seq_len_);
const bool use_append_mask =
input_params.is_spec_verify ||
input_params.meta.batch_forward_type.is_mixed() ||
input_params.meta.batch_forward_type.is_chunked_prefill();
if (!use_append_mask) {
return dense_attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_);
}
const int32_t num_sequences = input_params.num_sequences;
const int32_t num_sequences = input_params.meta.num_sequences;
if (num_sequences <= 0) {
return attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_);
return dense_attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_);
}
std::vector<torch::Tensor> req_mask_vec;
req_mask_vec.reserve(num_sequences);
for (int32_t j = 0; j < num_sequences; ++j) {
req_mask_vec.emplace_back(
attn_mask_.gen_append_mask(input_params.q_seq_lens_vec[j],
input_params.kv_seq_lens_vec[j],
attn_mask_.gen_append_mask(input_params.attention.host.q_seq_lens[j],
input_params.attention.host.kv_seq_lens[j],
max_seq_len_,
dtype_,
device_));
}
return torch::cat(req_mask_vec, 0);
#endif
}
ModelArgs model_args_;
@@ -183,10 +245,13 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
std::vector<layer::Qwen3HybridDecoderLayerModulePtr> layers_;
int32_t max_seq_len_ = 0;
int32_t dp_size_ = 1;
ParallelArgs parallel_args_;
FlashComm1Options flash_comm1_options_;
torch::Device device_;
torch::ScalarType dtype_ = torch::kFloat;
layer::Qwen3NextRMSNorm norm_{nullptr};
layer::AttentionMask attn_mask_;
layer::AttentionMask dense_attn_mask_;
layer::WordEmbedding embed_tokens_{nullptr};
};
@@ -232,53 +297,28 @@ class Qwen3HybridForCausalLMImplBase : public torch::nn::Module {
}
void load_model(std::unique_ptr<ModelLoader> loader) {
auto has_model_weights = [](const StateDict& dict) {
return dict.get_tensor("embed_tokens.weight").defined() ||
dict.get_dict_with_prefix("layers.").size() > 0 ||
dict.get_tensor("norm.weight").defined();
};
load_model(std::move(loader), "model.", "lm_head.");
}
void load_model(std::unique_ptr<ModelLoader> loader,
const std::string& model_prefix) {
load_model(std::move(loader), model_prefix, "lm_head.");
}
void load_model(std::unique_ptr<ModelLoader> loader,
const std::string& model_prefix,
const std::string& lm_head_prefix) {
auto has_lm_head_weights = [](const StateDict& dict) {
return dict.get_tensor("weight").defined() ||
dict.get_tensor("qweight").defined();
};
for (const auto& state_dict : loader->get_state_dicts()) {
auto model_state_dict = state_dict->get_dict_with_prefix("model.");
if (!has_model_weights(model_state_dict)) {
auto language_model_state_dict =
state_dict->get_dict_with_prefix("language_model.model.");
if (has_model_weights(language_model_state_dict)) {
model_state_dict = language_model_state_dict;
} else {
auto wrapped_language_model_state_dict =
state_dict->get_dict_with_prefix("model.language_model.");
if (has_model_weights(wrapped_language_model_state_dict)) {
model_state_dict = wrapped_language_model_state_dict;
}
}
}
auto model_state_dict = state_dict->get_dict_with_prefix(model_prefix);
model_->load_state_dict(model_state_dict);
auto lm_head_state_dict = state_dict->get_dict_with_prefix("lm_head.");
if (!has_lm_head_weights(lm_head_state_dict)) {
auto language_model_lm_head_state_dict =
state_dict->get_dict_with_prefix("language_model.lm_head.");
if (has_lm_head_weights(language_model_lm_head_state_dict)) {
lm_head_state_dict = language_model_lm_head_state_dict;
} else {
auto wrapped_language_model_lm_head_state_dict =
state_dict->get_dict_with_prefix("model.language_model.lm_head.");
if (has_lm_head_weights(wrapped_language_model_lm_head_state_dict)) {
lm_head_state_dict = wrapped_language_model_lm_head_state_dict;
} else {
auto wrapped_lm_head_state_dict =
state_dict->get_dict_with_prefix("model.lm_head.");
if (has_lm_head_weights(wrapped_lm_head_state_dict)) {
lm_head_state_dict = wrapped_lm_head_state_dict;
}
}
}
}
auto lm_head_state_dict =
state_dict->get_dict_with_prefix(lm_head_prefix);
if (!has_lm_head_weights(lm_head_state_dict) && tie_word_embeddings_) {
auto tied_lm_head_state_dict =
model_state_dict.get_dict_with_prefix("embed_tokens.");
@@ -288,8 +328,7 @@ class Qwen3HybridForCausalLMImplBase : public torch::nn::Module {
}
lm_head_->load_state_dict(lm_head_state_dict);
}
model_->verify_loaded_weights("model.");
model_->verify_loaded_weights(model_prefix);
}
virtual void prepare_expert_weight(int32_t layer_id,
@@ -298,6 +337,8 @@ class Qwen3HybridForCausalLMImplBase : public torch::nn::Module {
}
virtual void update_expert_weight(int32_t layer_id) { return; }
bool is_hybrid_linear_attention() { return true; }
layer::LmHead get_lm_head() { return lm_head_; }
void set_lm_head(layer::LmHead& head) { lm_head_ = head; }

256
upstream_ref/xllm_latest/models/vlm/qwen3_5.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -17,18 +17,30 @@ limitations under the License.
#include "core/framework/model/model_output.h"
#include "core/layers/common/lm_head.h"
#include "core/layers/common/rotary_embedding_util.h"
#include "models/model_registry.h"
#include "models/vlm/mposition/mposition.h"
#include "models/vlm/qwen3_vl_base.h"
#include "processors/multimodal_processor.h"
#include "processors/qwen2_vl_image_processor.h"
#include "processors/qwen3_vl_prompt_processor.h"
#include "processors/qwen3_vl_video_processor.h"
#if defined(USE_NPU)
#include "models/llm/qwen3_5.h"
#include "models/vlm/npu/qwen3_vl.h"
#elif defined(USE_MLU) || defined(USE_DCU)
#include "core/layers/common/qwen3_next_rms_norm.h"
#include "core/layers/common/rms_norm.h"
#include "core/layers/mlu/qwen3_5_decoder_layer.h"
#include "core/layers/qwen3_5_decoder_layer.h"
#include "core/layers/qwen3_vision_layer.h"
#include "models/llm/llm_model_base.h"
#include "models/model_registry.h"
#include "models/vlm/qwen3_vl_base.h"
#include "processors/input_processor.h"
#include "processors/qwen2_vl_image_processor.h"
#include "qwen3_vl.h"
#endif
namespace xllm {
#if !defined(USE_NPU)
class Qwen3_5ModelImpl final
: public LlmModelImplBase<layer::Qwen3_5DecoderLayer> {
public:
@@ -78,33 +90,7 @@ class Qwen3_5ModelImpl final
std::pair<torch::Tensor, torch::Tensor> apply_mrope(
const torch::Tensor positions) override {
auto target_cos_sin = cos_sin_.index({positions});
auto target_cos_sin_chunks = target_cos_sin.chunk(/*chunks=*/2, /*dim=*/-1);
auto cos_pos = target_cos_sin_chunks[0].contiguous();
auto sin_pos = target_cos_sin_chunks[1].contiguous();
auto apply = [this](torch::Tensor x) {
auto freqs_t = x[0].clone();
int64_t mrop_length = static_cast<int64_t>(freqs_t.size(-1) / 2);
for (int32_t dim_idx = 1; dim_idx <= 2; ++dim_idx) {
int64_t offset = dim_idx;
int64_t section_len = mrope_section_[dim_idx];
int64_t length = section_len * 3;
auto idx_first_half = torch::arange(offset, length, 3, torch::kLong);
auto idx_second_half = torch::arange(
offset + mrop_length, length + mrop_length, 3, torch::kLong);
auto idx_tensor =
torch::cat({idx_first_half, idx_second_half}, 0).to(x.device());
auto src = x[dim_idx].index_select(-1, idx_tensor);
freqs_t.index_copy_(-1, idx_tensor, src);
}
return freqs_t;
};
cos_pos = apply(cos_pos.reshape({positions.size(0), -1, cos_pos.size(-1)}));
sin_pos = apply(sin_pos.reshape({positions.size(0), -1, sin_pos.size(-1)}));
return std::make_pair(cos_pos, sin_pos);
return layer::rotary::apply_mrope(cos_sin_, positions, mrope_section_);
}
virtual ModelOutput forward(torch::Tensor tokens,
@@ -120,11 +106,11 @@ class Qwen3_5ModelImpl final
tokens = torch::tensor({1}).to(torch::kInt32).to(tokens.device());
positions = torch::tensor({1}).to(torch::kInt32).to(positions.device());
}
auto& dp_token_nums = input_params_new.dp_global_token_nums;
auto& dp_token_nums = input_params_new.parallel.dp_global_token_nums;
std::replace(dp_token_nums.begin(), dp_token_nums.end(), 0, 1);
}
auto inputs_embeds = input_params.input_embedding;
auto inputs_embeds = input_params.embedding.input_embedding;
torch::Tensor h;
if (inputs_embeds.defined()) {
h = inputs_embeds;
@@ -139,12 +125,8 @@ class Qwen3_5ModelImpl final
}
auto& attn_metadata = *(input_params_new.attn_metadata);
bool only_prefill =
(attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
if (positions.dim() == 2 && only_prefill && !mrope_section_.empty()) {
std::tie(attn_metadata.mrope_cos, attn_metadata.mrope_sin) =
apply_mrope(positions);
}
std::tie(attn_metadata.mrope_cos, attn_metadata.mrope_sin) =
apply_mrope(positions);
std::optional<torch::Tensor> residual;
for (size_t i = 0; i < layers_.size(); i++) {
@@ -169,8 +151,70 @@ class Qwen3_5ModelImpl final
layer::AttentionMetadata get_attention_metadata(
const ModelInputParams& params,
const torch::Tensor& h) {
auto attn_metadata = layer::AttentionMetadataBuilder::build(params, false);
// TODO: support linear attention
auto attn_metadata =
layer::AttentionMetadataBuilder::build(params,
/*enable_mla=*/false,
/*attn_mask=*/{},
h.device());
// Init batch and token_block_offset for GDN attention
if (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill) {
constexpr int32_t kBlockM = 64;
constexpr int64_t pad_slot_id = -1;
constexpr int64_t default_max_num_programs = 1024;
constexpr int64_t chunk_size = 64;
auto seqlens = attn_metadata.q_cu_seq_lens.diff();
auto nums = (seqlens + kBlockM - 1) / kBlockM;
nums = nums.to(torch::kLong);
int32_t tot = nums.sum().item<int32_t>();
torch::Tensor range_batch = torch::arange(nums.size(0), nums.options());
torch::Tensor mlist_tensor = torch::repeat_interleave(range_batch, nums);
int64_t mlist_len = mlist_tensor.size(0);
int64_t max_num_programs =
std::max(default_max_num_programs, mlist_len) * 2;
torch::Tensor batch_ptr =
torch::full({max_num_programs},
pad_slot_id,
torch::dtype(torch::kInt32).device(seqlens.device()));
torch::Tensor token_block_offset_ptr =
torch::full({max_num_programs},
pad_slot_id,
torch::dtype(torch::kInt32).device(seqlens.device()));
std::vector<torch::Tensor> vec;
vec.reserve(nums.size(0));
for (int64_t i = 0; i < nums.size(0); ++i) {
vec.emplace_back(
torch::arange(nums[i].item<int64_t>(), nums.options()));
}
torch::Tensor offsetlist_tensor = torch::cat(vec, -1).to(torch::kInt32);
batch_ptr.narrow(0, 0, mlist_len).copy_(mlist_tensor);
token_block_offset_ptr.narrow(0, 0, mlist_len).copy_(offsetlist_tensor);
// Compute chunk indices for the chunked GDN kernel
{
torch::Tensor lengths = seqlens;
torch::Tensor num_chunks = (lengths + chunk_size - 1) / chunk_size;
num_chunks = num_chunks.to(torch::kLong);
torch::Tensor cumsum = torch::cumsum(num_chunks, 0);
int64_t total_chunks = cumsum[-1].item<int64_t>();
torch::Tensor arange_total =
torch::arange(total_chunks, attn_metadata.q_cu_seq_lens.options());
torch::Tensor zeros = torch::zeros({1}, cumsum.options());
torch::Tensor prefix = torch::cat(
{zeros, cumsum.slice(/*dim=*/0, /*start=*/0, /*end=*/-1)});
torch::Tensor repeats_prefix =
torch::repeat_interleave(prefix, num_chunks);
torch::Tensor indices = arange_total - repeats_prefix;
torch::Tensor mask = indices == 0;
torch::Tensor col0 = mask.cumsum(0) - 1;
attn_metadata.chunk_indices = torch::stack({col0, indices}, /*dim=*/1)
.to(attn_metadata.q_cu_seq_lens)
.to(torch::kInt32);
}
attn_metadata.tot = tot;
attn_metadata.batch = batch_ptr;
attn_metadata.token_block_offset = token_block_offset_ptr;
}
return attn_metadata;
}
};
@@ -193,8 +237,16 @@ class Qwen3_5ForCausalLMImpl : public LlmForCausalLMImplBase<Qwen3_5Model> {
};
TORCH_MODULE(Qwen3_5ForCausalLM);
#endif // !defined(USE_NPU)
#if defined(USE_NPU)
using Qwen3_5_VisionTransformer = npu::model::Qwen3_VisionTransformer;
#else
using Qwen3_5_VisionTransformer = Qwen3_VisionTransformer;
#endif
using Qwen3_5ForConditionalGenerationImpl =
Qwen3VLForConditionalGenerationBase<Qwen3_VisionTransformer,
Qwen3VLForConditionalGenerationBase<Qwen3_5_VisionTransformer,
Qwen3_5ForCausalLM>;
TORCH_MODULE(Qwen3_5ForConditionalGeneration);
@@ -211,6 +263,7 @@ TORCH_MODULE(Qwen3_5ForConditionalGeneration);
LOAD_ARG_OR( \
max_position_embeddings, "text_config.max_position_embeddings", 262144); \
LOAD_ARG_OR(rms_norm_eps, "text_config.rms_norm_eps", 1e-6); \
LOAD_ARG_OR(bos_token_id, "text_config.bos_token_id", 151643); \
LOAD_ARG_OR(eos_token_id, "text_config.eos_token_id", 248044); \
LOAD_ARG_OR( \
rope_theta, "text_config.rope_parameters.rope_theta", 10000000.0f); \
@@ -223,34 +276,50 @@ TORCH_MODULE(Qwen3_5ForConditionalGeneration);
LOAD_ARG_OR( \
linear_value_head_dim, "text_config.linear_value_head_dim", 128); \
LOAD_ARG_OR(linear_num_key_heads, "text_config.linear_num_key_heads", 16); \
LOAD_ARG_OR( \
linear_num_value_heads, "text_config.linear_num_value_heads", 48); \
LOAD_ARG_OR(linear_num_value_heads, \
"text_config.linear_num_value_heads", \
static_cast<int32_t>(args->n_heads() * 2)); \
LOAD_ARG_OR( \
full_attention_interval, "text_config.full_attention_interval", 4); \
LOAD_ARG_OR(attn_output_gate, "text_config.attn_output_gate", false); \
LOAD_ARG_OR(attn_output_gate, "text_config.attn_output_gate", true); \
LOAD_ARG_OR( \
num_nextn_predict_layers, "text_config.mtp_num_hidden_layers", 0); \
LOAD_ARG_OR(num_nextn_predict_layers, \
"text_config.num_nextn_predict_layers", \
args->num_nextn_predict_layers()); \
LOAD_ARG_OR(attention_bias, "text_config.attention_bias", false); \
LOAD_ARG_OR(attention_dropout, "text_config.attention_dropout", 0.0f); \
LOAD_ARG_OR(initializer_range, "text_config.initializer_range", 0.02f); \
LOAD_ARG_OR( \
mlp_only_layers, "text_config.mlp_only_layers", std::vector<int32_t>()); \
LOAD_ARG(rope_scaling_mrope_section, \
"text_config.rope_parameters.mrope_section"); \
LOAD_ARG_OR(rope_scaling_mrope_section, \
"text_config.rope_parameters.mrope_section", \
std::vector<int64_t>({11, 11, 10})); \
LOAD_ARG_OR(rope_scaling_mrope_interleaved, \
"text_config.rope_parameters.mrope_interleaved", \
true); \
LOAD_ARG_OR(rope_scaling_rope_type, \
"text_config.rope_parameters.rope_type", \
"default"); \
if (args->rope_scaling_rope_type() == "default") { \
args->rope_scaling_rope_type() = "mrope"; \
} \
LOAD_ARG_OR(partial_rotary_factor, \
"text_config.rope_parameters.partial_rotary_factor", \
0.25f)
0.25f); \
LOAD_ARG_OR(mamba_ssm_dtype, "text_config.mamba_ssm_dtype", "float32")
#define LOAD_QWEN3_5_VISION_ARGS() \
LOAD_ARG_OR(image_token_id, "image_token_id", 248056); \
LOAD_ARG_OR(video_token_id, "video_token_id", 248057); \
LOAD_ARG_OR(vision_start_token_id, "vision_start_token_id", 248053); \
LOAD_ARG_OR(vision_end_token_id, "vision_end_token_id", 248054); \
LOAD_ARG(mm_deepstack_visual_indexes, \
"vision_config.deepstack_visual_indexes"); \
LOAD_ARG_OR(mm_deepstack_visual_indexes, \
"vision_config.deepstack_visual_indexes", \
std::vector<int64_t>()); \
if (!args->mm_deepstack_visual_indexes().empty()) { \
LOG(FATAL) << "qwen3_5 VLM does not support DeepStack visual indexes"; \
} \
LOAD_ARG_OR(mm_num_hidden_layers, "vision_config.depth", 27); \
LOAD_ARG_OR(mm_hidden_act, "vision_config.hidden_act", "gelu_pytorch_tanh"); \
LOAD_ARG_OR(mm_hidden_size, "vision_config.hidden_size", 1152); \
@@ -261,32 +330,43 @@ TORCH_MODULE(Qwen3_5ForConditionalGeneration);
LOAD_ARG_OR(mm_num_position_embeddings, \
"vision_config.num_position_embeddings", \
2304); \
LOAD_ARG_OR(mm_projection_dim, "vision_config.out_hidden_size", 5120); \
LOAD_ARG_OR(mm_projection_dim, \
"vision_config.out_hidden_size", \
args->hidden_size()); \
LOAD_ARG_OR(mm_patch_size, "vision_config.patch_size", 16); \
LOAD_ARG_OR(mm_spatial_merge_size, "vision_config.spatial_merge_size", 2); \
LOAD_ARG_OR(mm_temporal_patch_size, "vision_config.temporal_patch_size", 2); \
LOAD_ARG_OR_FUNC(mm_head_dim, "head_dim", [&] { \
return args->mm_hidden_size() / args->mm_num_attention_heads(); \
}); \
LOAD_ARG_OR( \
rope_scaling_rope_type, "vision_config.rope_scaling.type", "mrope")
})
REGISTER_INPUT_PROCESSOR(qwen3_5, Qwen2_5_VLInputProcessor);
// qwen3_5/qwen3_5_moe are multimodal entry points. On NPU, text-only serving
// uses qwen3_5_text/qwen3_5_moe_text from llm/qwen3_5.h because the VLM
// request protocol currently requires array-form chat content.
REGISTER_CAUSAL_VLM_MODEL(qwen3_5, Qwen3_5ForConditionalGeneration);
REGISTER_IMAGE_PROCESSOR(qwen3_5, Qwen2VLImageProcessor);
REGISTER_MPOSITION_GENERATOR(qwen3_5, Qwen3VLMPositionGenerator);
using Qwen35MultimodalProcessor = MultimodalProcessor<Qwen3VLPromptProcessor,
Qwen2VLImageProcessor,
Qwen3VLVideoProcessor>;
REGISTER_MULTIMODAL_PROCESSOR(qwen3_5, Qwen35MultimodalProcessor);
REGISTER_MODEL_ARGS(qwen3_5, [&] {
LOAD_QWEN3_5_COMMON_ARGS();
LOAD_QWEN3_5_VISION_ARGS();
SET_ARG(stop_token_ids, std::unordered_set<int32_t>({args->eos_token_id()}));
SET_ARG(num_experts, 0);
SET_ARG(n_routed_experts, 0);
SET_ARG(n_shared_experts, 0);
SET_ARG(stop_token_ids,
std::unordered_set<int32_t>({args->eos_token_id(), 248046}));
});
REGISTER_INPUT_PROCESSOR(qwen3_5_moe, Qwen2_5_VLInputProcessor);
REGISTER_CAUSAL_VLM_MODEL(qwen3_5_moe, Qwen3_5ForConditionalGeneration);
REGISTER_IMAGE_PROCESSOR(qwen3_5_moe, Qwen2VLImageProcessor);
REGISTER_MPOSITION_GENERATOR(qwen3_5_moe, Qwen3VLMPositionGenerator);
REGISTER_MULTIMODAL_PROCESSOR(qwen3_5_moe, Qwen35MultimodalProcessor);
REGISTER_MODEL_ARGS(qwen3_5_moe, [&] {
LOAD_QWEN3_5_COMMON_ARGS();
LOAD_QWEN3_5_VISION_ARGS();
LOAD_ARG_OR(decoder_sparse_step, "text_config.decoder_sparse_step", 1);
LOAD_ARG_OR(moe_intermediate_size, "text_config.moe_intermediate_size", 512);
LOAD_ARG_OR(num_experts, "text_config.num_experts", 512);
@@ -295,7 +375,6 @@ REGISTER_MODEL_ARGS(qwen3_5_moe, [&] {
"text_config.shared_expert_intermediate_size",
512);
LOAD_ARG_OR(norm_topk_prob, "text_config.norm_topk_prob", true);
LOAD_ARG_OR(
n_routed_experts, "text_config.n_routed_experts", args->num_experts());
SET_ARG(n_shared_experts,
@@ -306,7 +385,56 @@ REGISTER_MODEL_ARGS(qwen3_5_moe, [&] {
SET_ARG(topk_group, 0);
SET_ARG(routed_scaling_factor, 1.0f);
SET_ARG(stop_token_ids, std::unordered_set<int32_t>({args->eos_token_id()}));
SET_ARG(stop_token_ids,
std::unordered_set<int32_t>({args->eos_token_id(), 248046}));
});
// Text-only model registrations. On NPU these are handled by llm/qwen3_5.h.
#if !defined(USE_NPU)
// qwen3_5 without vision config (text-only serving).
// Model args are already registered by the VLM registration above.
REGISTER_CAUSAL_MODEL_WITH_VARNAME(qwen3_5_lm, qwen3_5, Qwen3_5ForCausalLM);
REGISTER_CAUSAL_MODEL_WITH_VARNAME(qwen3_5_moe_lm,
qwen3_5_moe,
Qwen3_5ForCausalLM);
REGISTER_CAUSAL_MODEL(qwen3_5_text, Qwen3_5ForCausalLM);
REGISTER_MODEL_ARGS(qwen3_5_text, [&] {
LOAD_QWEN3_5_COMMON_ARGS();
SET_ARG(num_experts, 0);
SET_ARG(n_routed_experts, 0);
SET_ARG(n_shared_experts, 0);
SET_ARG(decoder_sparse_step, 1);
SET_ARG(stop_token_ids,
std::unordered_set<int32_t>({args->eos_token_id(), 248046}));
});
REGISTER_CAUSAL_MODEL(qwen3_5_moe_text, Qwen3_5ForCausalLM);
REGISTER_MODEL_ARGS(qwen3_5_moe_text, [&] {
LOAD_QWEN3_5_COMMON_ARGS();
LOAD_ARG_OR(decoder_sparse_step, "text_config.decoder_sparse_step", 1);
LOAD_ARG_OR(moe_intermediate_size, "text_config.moe_intermediate_size", 512);
LOAD_ARG_OR(num_experts, "text_config.num_experts", 512);
LOAD_ARG_OR(num_experts_per_tok, "text_config.num_experts_per_tok", 10);
LOAD_ARG_OR(shared_expert_intermediate_size,
"text_config.shared_expert_intermediate_size",
512);
LOAD_ARG_OR(norm_topk_prob, "text_config.norm_topk_prob", true);
LOAD_ARG_OR(
n_routed_experts, "text_config.n_routed_experts", args->num_experts());
SET_ARG(n_shared_experts,
args->shared_expert_intermediate_size() > 0 ? 1 : 0);
SET_ARG(scoring_func, "softmax");
SET_ARG(topk_method, "");
SET_ARG(n_group, -1);
SET_ARG(topk_group, 0);
SET_ARG(routed_scaling_factor, 1.0f);
SET_ARG(stop_token_ids,
std::unordered_set<int32_t>({args->eos_token_id(), 248046}));
});
#endif // !defined(USE_NPU)
#undef LOAD_QWEN3_5_VISION_ARGS
#undef LOAD_QWEN3_5_COMMON_ARGS
} // namespace xllm