diff --git a/qwen3_6_scripts/serving_chat.py b/qwen3_6_scripts/serving_chat.py index 8d62b6e4..4ed841b5 100644 --- a/qwen3_6_scripts/serving_chat.py +++ b/qwen3_6_scripts/serving_chat.py @@ -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( diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp index ae2a16ba..1ad364a4 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp +++ b/upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp index aa257bf1..ad3cd295 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp +++ b/upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp index 794f9bd9..21c15d8c 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp +++ b/upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/group_gemm.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/group_gemm.cpp index 38743e66..290299a0 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/group_gemm.cpp +++ b/upstream_ref/xllm_latest/core/kernels/ilu/group_gemm.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/ilu_ops_api.h b/upstream_ref/xllm_latest/core/kernels/ilu/ilu_ops_api.h index e4fd7853..3dedd7da 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/ilu_ops_api.h +++ b/upstream_ref/xllm_latest/core/kernels/ilu/ilu_ops_api.h @@ -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. diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h b/upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h index 57ce66dc..83bad88e 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h +++ b/upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h @@ -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. diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/matmul.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/matmul.cpp index 91b6868f..f90c0d47 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/matmul.cpp +++ b/upstream_ref/xllm_latest/core/kernels/ilu/matmul.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp index c5a98595..e451bc36 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp +++ b/upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp index 89370b79..45af7656 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp +++ b/upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/utils.h b/upstream_ref/xllm_latest/core/kernels/ilu/utils.h index e8af0c3c..9fd15298 100644 --- a/upstream_ref/xllm_latest/core/kernels/ilu/utils.h +++ b/upstream_ref/xllm_latest/core/kernels/ilu/utils.h @@ -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. diff --git a/upstream_ref/xllm_latest/core/layers/ilu/attention.cpp b/upstream_ref/xllm_latest/core/layers/ilu/attention.cpp index b66f28a4..62d36a62 100644 --- a/upstream_ref/xllm_latest/core/layers/ilu/attention.cpp +++ b/upstream_ref/xllm_latest/core/layers/ilu/attention.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/layers/ilu/attention.h b/upstream_ref/xllm_latest/core/layers/ilu/attention.h index a971835f..bf4b59ba 100644 --- a/upstream_ref/xllm_latest/core/layers/ilu/attention.h +++ b/upstream_ref/xllm_latest/core/layers/ilu/attention.h @@ -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. diff --git a/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp b/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp index 4238012e..7c829ad3 100644 --- a/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp +++ b/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp @@ -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 #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); diff --git a/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.h b/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.h index 3e477064..8d4e9da9 100644 --- a/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.h +++ b/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.h @@ -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. diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.cpp index c5c9dc7f..a0bd62cf 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.cpp +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.h index bb3929d6..6d6881a4 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.h +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.h @@ -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. diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp index 7d572476..0b97be6d 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp @@ -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 -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 +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> +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."); diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h index bec6c1c6..7c782e3f 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h @@ -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 +#include #include +#include #include #include "qwen3_next_gated_delta_net.h" @@ -34,9 +36,15 @@ class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl { const torch::TensorOptions& options); protected: - std::pair project_padded_inputs( - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata) override; + std::pair project_decode_inputs( + const torch::Tensor& hidden_states) override; + std::pair project_flat_inputs( + const torch::Tensor& hidden_states) override; + std::optional< + std::tuple> + 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; diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp index cec9a95e..7f8b4b5c 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp @@ -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 #include +#include #include +#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 view_shape = tensor.sizes().vec(); + view_shape.insert(view_shape.begin() + head_dim + 1, 1); + std::vector expand_shape = view_shape; + expand_shape[head_dim + 1] = repeats; + std::vector 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_recurrent_gated_delta_rule( torch::Tensor query, torch::Tensor key, @@ -52,6 +80,9 @@ std::tuple 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_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_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& graph_context, + const torch::Tensor& x, + const torch::Tensor& weight, + const torch::Tensor& conv_state, + const std::optional& bias, + const std::vector& query_start_loc, + const std::vector& cache_indices, + const std::vector& 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(ACL_EVENT_EXTERNAL); + event->block(stream); + event->reset(stream); + + torch::Tensor output; + c10_npu::graph_task_group_begin(stream); + const std::vector 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& 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(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(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(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 +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(num_k_heads_ / tp_size_); - fused_params.num_heads_v = static_cast(num_v_heads_ / tp_size_); - fused_params.head_qk = static_cast(head_k_dim_); - fused_params.head_v = static_cast(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(num_k_heads_ / tp_size_); + fused_params.num_heads_v = static_cast(num_v_heads_ / tp_size_); + fused_params.head_qk = static_cast(head_k_dim_); + fused_params.head_v = static_cast(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{1}, - /*padding=*/std::vector{3}, - /*dilation=*/std::vector{1}, - /*groups=*/static_cast(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 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(batch_size) && + input_params.parallel.query_start_loc.size() == + static_cast(batch_size + 1) && + input_params.embedding.linear_state_ids.size() == + static_cast(batch_size) && + input_params.linear_state_validity_mask.size() == + static_cast(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(), // 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(); - conv1d_params.initial_state_idx = std::optional(); - 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(); + const std::vector 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(), + 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::IntArrayRef(input_params.parallel.query_start_loc), + torch::IntArrayRef(linear_state_indices_host), + torch::IntArrayRef(std::vector()), + 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(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(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 packed_q; + std::vector packed_k; + std::vector packed_v; + std::vector packed_g; + std::vector 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(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( + attn_metadata.q_seq_lens_vec.data(), static_cast(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(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(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 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(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(); - 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(); + 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(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 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(); - 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(); + 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; diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h index 2994f329..fdc82b4d 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h @@ -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 +#include #include #include #include @@ -51,20 +52,41 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module { const ModelInputParams& input_params); protected: - virtual std::pair project_padded_inputs( - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata) = 0; + virtual std::pair project_decode_inputs( + const torch::Tensor& hidden_states) = 0; + virtual std::pair 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> + 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 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 process_mixed_qkv( torch::Tensor& mixed_qkv) const; diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.cpp index 9961170a..c1dec2e9 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.cpp +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.cpp @@ -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 #include +#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; + } } } diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.h index 4a4149d1..45347fb9 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.h +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.h @@ -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 mrope_section_; torch::Tensor mrope_gather_pattern_; diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.cpp index 7bb7235b..de56dc2f 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.cpp +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.cpp @@ -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. diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.h index b790e1d8..658b8d23 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.h +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.h @@ -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. diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.cpp index eeeb49f1..f9b394c3 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.cpp +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.cpp @@ -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 -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 +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) { diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.h index 6c8acb9b..ebf39e8f 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.h +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.h @@ -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 project_padded_inputs( - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata) override; + std::pair project_decode_inputs( + const torch::Tensor& hidden_states) override; + std::pair 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; diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp index 904561d6..543b37fb 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp @@ -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 #include +#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 diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h index 098f8f2e..fb6d3a6c 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h @@ -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. diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_5.h b/upstream_ref/xllm_latest/models/llm/qwen3_5.h index 7c712d39..7e4ed875 100644 --- a/upstream_ref/xllm_latest/models/llm/qwen3_5.h +++ b/upstream_ref/xllm_latest/models/llm/qwen3_5.h @@ -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 +#include #include #include +#include #include -#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(context)); } + + torch::Tensor get_input_embeddings(torch::Tensor input_ids) { + return get_word_embedding()(input_ids); + } + + void load_model(std::unique_ptr loader) { + Qwen3NextForCausalLMImpl::load_model( + std::move(loader), "model.language_model.", "lm_head."); + } + + void load_model(std::unique_ptr 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({args->eos_token_id()})); \ + std::unordered_set({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 diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h b/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h index b4b06b90..8a379c34 100644 --- a/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h +++ b/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h @@ -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 - -#include #include -#include -#include -#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 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( - static_cast(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(static_cast(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(context, layer_id)); - } - } - - ModelOutput forward(torch::Tensor tokens, - torch::Tensor positions, - std::vector& 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 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(context); - set_model_module(mtp_model_); - } - - void load_model(std::unique_ptr loader) { - static const std::vector kEmbeddingPrefixes = { - "model.language_model.", "language_model.model.", "model.", ""}; - static const std::vector 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 mtp_model_; + : Qwen3_5MtpForCausalLMImplBase( + context, + std::make_shared(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 diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp_base.h b/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp_base.h new file mode 100644 index 00000000..64b66816 --- /dev/null +++ b/upstream_ref/xllm_latest/models/llm/qwen3_5_mtp_base.h @@ -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 + +#include +#include +#include +#include +#include + +#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 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( + static_cast(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(static_cast(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(context, layer_id)); + } + } + + ModelOutput forward(torch::Tensor tokens, + torch::Tensor positions, + std::vector& 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 residual = std::nullopt; + for (size_t i = 0; i < layers_.size(); ++i) { + if (!input_params.synchronize_layer(static_cast(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(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 loader) { + static const std::vector kEmbeddingPrefixes = { + "model.language_model.", "language_model.model.", "model.", ""}; + static const std::vector kMtpPrefixes = {"mtp.", "model.mtp."}; + bool lm_head_loaded = false; + + for (const std::unique_ptr& 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 mtp_model) + : Qwen3HybridForCausalLMImplBase(context), + mtp_model_(std::move(mtp_model)) { + set_model_module(mtp_model_); + } + + private: + std::shared_ptr mtp_model_; +}; + +} // namespace xllm diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_next.h b/upstream_ref/xllm_latest/models/llm/qwen3_next.h index 9acf0efc..2c19e872 100644 --- a/upstream_ref/xllm_latest/models/llm/qwen3_next.h +++ b/upstream_ref/xllm_latest/models/llm/qwen3_next.h @@ -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. diff --git a/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h b/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h index bbf60bc8..83e42e59 100644 --- a/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h +++ b/upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h @@ -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 #include -#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(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(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 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 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 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 loader, + const std::string& model_prefix) { + load_model(std::move(loader), model_prefix, "lm_head."); + } + + void load_model(std::unique_ptr 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; } diff --git a/upstream_ref/xllm_latest/models/vlm/qwen3_5.h b/upstream_ref/xllm_latest/models/vlm/qwen3_5.h index 9a9a9c02..291bc6a7 100644 --- a/upstream_ref/xllm_latest/models/vlm/qwen3_5.h +++ b/upstream_ref/xllm_latest/models/vlm/qwen3_5.h @@ -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 { public: @@ -78,33 +90,7 @@ class Qwen3_5ModelImpl final std::pair 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(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 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(); + 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 vec; + vec.reserve(nums.size(0)); + for (int64_t i = 0; i < nums.size(0); ++i) { + vec.emplace_back( + torch::arange(nums[i].item(), 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(); + 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 { }; 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; 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(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()); \ - 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({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()); \ + 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; +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({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({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({args->eos_token_id()})); + SET_ARG(stop_token_ids, + std::unordered_set({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({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({args->eos_token_id(), 248046})); +}); +#endif // !defined(USE_NPU) + +#undef LOAD_QWEN3_5_VISION_ARGS +#undef LOAD_QWEN3_5_COMMON_ARGS + } // namespace xllm