data: port complete MoE + xllm layer call chains from upstream repos
MoE call chain from ds_vllm (vllm-project/vllm latest): ex_engine/moe/ — 20 files, 8736 lines - modular_kernel.py (1630 lines) — base classes for modular MoE - experts/fused_batched_moe.py (972 lines) — NaiveBatchedExperts - prepare_finalize/batched.py (171 lines) — token grouping by expert - topk_weight_and_reduce.py (176 lines) — scatter-add finalize - fused_moe.py (1740 lines) — main fused_moe dispatch - config.py (1407 lines) — FusedMoEQuantConfig - activation.py, utils.py, layer.py, etc. xllm layer code (jd-opensource/xllm): ex_engine/xllm_layers/ — 39 files, 5859 lines - ilu/fused_moe.cpp (797 lines) — production ixformer 7-step MoE pipeline - ilu/attention.cpp (189 lines) — paged_attention + flash_attn bridge - npu_torch/qwen3_gated_delta_net_base.cpp (576 lines) — GDN reference - common/rms_norm.cpp, rotary_embedding.cpp, activation.cpp, dense_mlp.cpp xllm ILU kernels — synced 10 files to upstream (diffs from prior edits) These are reference implementations, NOT hand-written. Source repos: vllm-project/vllm, jd-opensource/xllm
This commit is contained in:
Claude
38
ex_engine/xllm_layers/common/activation.cpp
Normal file
38
ex_engine/xllm_layers/common/activation.cpp
Normal file
@@ -0,0 +1,38 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "activation.h"
|
||||
|
||||
#include "kernels/ops_api.h"
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
ActivationImpl::ActivationImpl(const std::string& act_mode, bool is_gated)
|
||||
: act_mode_(act_mode), is_gated_(is_gated) {}
|
||||
|
||||
void ActivationImpl::forward(torch::Tensor& input, torch::Tensor& output) {
|
||||
xllm::kernel::ActivationParams activation_params;
|
||||
activation_params.input = input;
|
||||
activation_params.output = output;
|
||||
activation_params.act_mode = act_mode_;
|
||||
activation_params.is_gated = is_gated_;
|
||||
xllm::kernel::active(activation_params);
|
||||
// Unified assignment: NPU returns new tensor, others modify in-place (no-op
|
||||
// assignment)
|
||||
output = activation_params.output;
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
38
ex_engine/xllm_layers/common/activation.h
Normal file
38
ex_engine/xllm_layers/common/activation.h
Normal file
@@ -0,0 +1,38 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class ActivationImpl : public torch::nn::Module {
|
||||
public:
|
||||
ActivationImpl(const std::string& act_mode, bool is_gated);
|
||||
|
||||
void forward(torch::Tensor& input, torch::Tensor& output);
|
||||
|
||||
private:
|
||||
std::string act_mode_;
|
||||
bool is_gated_;
|
||||
};
|
||||
TORCH_MODULE(Activation);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
141
ex_engine/xllm_layers/common/dense_mlp.cpp
Normal file
141
ex_engine/xllm_layers/common/dense_mlp.cpp
Normal file
@@ -0,0 +1,141 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "dense_mlp.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include "kernels/ops_api.h"
|
||||
#include "platform/device.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
DenseMLPImpl::DenseMLPImpl(int64_t hidden_size,
|
||||
int64_t intermediate_size,
|
||||
bool is_gated,
|
||||
bool has_bias,
|
||||
const std::string& hidden_act,
|
||||
bool enable_result_reduction,
|
||||
const QuantArgs& quant_args,
|
||||
ProcessGroup* process_group,
|
||||
const torch::TensorOptions& options,
|
||||
const std::string& module_prefix)
|
||||
: is_gated_(is_gated),
|
||||
intermediate_size_(intermediate_size),
|
||||
process_group_(process_group),
|
||||
hidden_act_(hidden_act) {
|
||||
// Check if using w8a8 smoothquant quantization
|
||||
is_smoothquant_ = quant_args.quant_method() == kQuantMethodSmoothquant;
|
||||
|
||||
if (is_smoothquant_) {
|
||||
// Safety check: only w8a8 smoothquant is supported
|
||||
if (quant_args.bits() != 8 || !quant_args.activation_dynamic()) {
|
||||
LOG(FATAL)
|
||||
<< "DenseMLP w8a8 mode only supports w8a8 smoothquant quantization. "
|
||||
<< "Got bits=" << quant_args.bits()
|
||||
<< ", activation_dynamic=" << quant_args.activation_dynamic();
|
||||
}
|
||||
}
|
||||
|
||||
// Determine extra args based on quantization mode
|
||||
LinearExtraArgs gate_up_proj_extra_args("none", false);
|
||||
LinearExtraArgs down_proj_extra_args("none", false);
|
||||
if (is_smoothquant_) {
|
||||
// For per-token smoothquant, use specific args
|
||||
down_proj_extra_args = LinearExtraArgs(hidden_act_, is_gated_);
|
||||
}
|
||||
|
||||
// 1. gate + up
|
||||
int64_t out_feature = is_gated_ ? intermediate_size_ * 2 : intermediate_size_;
|
||||
gate_up_proj_ =
|
||||
register_module("gate_up_proj",
|
||||
ColumnParallelLinear(hidden_size,
|
||||
out_feature,
|
||||
/*bias=*/has_bias,
|
||||
/*gather_output=*/false,
|
||||
quant_args,
|
||||
process_group_,
|
||||
options,
|
||||
gate_up_proj_extra_args));
|
||||
|
||||
act_ = register_module("act", Activation(hidden_act_, is_gated_));
|
||||
|
||||
// 2. down
|
||||
const auto down_proj_quant_args =
|
||||
module_prefix.empty()
|
||||
? quant_args
|
||||
: quant_args.for_module(module_prefix + ".down_proj");
|
||||
down_proj_ = register_module("down_proj",
|
||||
RowParallelLinear(intermediate_size_,
|
||||
hidden_size,
|
||||
/*bias=*/has_bias,
|
||||
/*input_is_parallelized=*/true,
|
||||
enable_result_reduction,
|
||||
down_proj_quant_args,
|
||||
process_group_,
|
||||
options,
|
||||
down_proj_extra_args));
|
||||
}
|
||||
|
||||
torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) {
|
||||
// input shape: [num_tokens, hidden_size]
|
||||
auto gate_up = gate_up_proj_->forward(hidden_states);
|
||||
|
||||
if (is_smoothquant_) {
|
||||
// For w8a8 quantization, the active operation is fused with the down_proj
|
||||
return down_proj_->forward(gate_up);
|
||||
} else {
|
||||
torch::Tensor output;
|
||||
if (Device::type_str() != "npu") {
|
||||
int64_t batch_size = gate_up.sizes()[0];
|
||||
output = torch::empty(
|
||||
{batch_size, intermediate_size_ / process_group_->world_size()},
|
||||
gate_up.options());
|
||||
}
|
||||
|
||||
act_->forward(gate_up, output);
|
||||
return down_proj_->forward(output);
|
||||
}
|
||||
}
|
||||
|
||||
void DenseMLPImpl::load_state_dict(const StateDict& state_dict) {
|
||||
gate_up_proj_->load_state_dict(state_dict, {"gate_proj.", "up_proj."});
|
||||
down_proj_->load_state_dict(state_dict.get_dict_with_prefix("down_proj."));
|
||||
}
|
||||
|
||||
void DenseMLPImpl::load_state_dict(const StateDict& state_dict,
|
||||
const std::vector<std::string>& gate_up_name,
|
||||
const std::string& down_name) {
|
||||
if (is_gated_) {
|
||||
CHECK_EQ(gate_up_name.size(), 2);
|
||||
gate_up_proj_->load_state_dict(state_dict, gate_up_name);
|
||||
} else {
|
||||
CHECK_EQ(gate_up_name.size(), 1);
|
||||
gate_up_proj_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix(gate_up_name[0]));
|
||||
}
|
||||
down_proj_->load_state_dict(state_dict.get_dict_with_prefix(down_name));
|
||||
}
|
||||
|
||||
std::optional<torch::Tensor> DenseMLPImpl::get_fp8_input_scale() const {
|
||||
if (gate_up_proj_) {
|
||||
return gate_up_proj_->get_input_scale();
|
||||
}
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
67
ex_engine/xllm_layers/common/dense_mlp.h
Normal file
67
ex_engine/xllm_layers/common/dense_mlp.h
Normal file
@@ -0,0 +1,67 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include "activation.h"
|
||||
#include "framework/model/model_args.h"
|
||||
#include "framework/parallel_state/parallel_args.h"
|
||||
#include "framework/quant_args.h"
|
||||
#include "framework/state_dict/state_dict.h"
|
||||
#include "linear.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class DenseMLPImpl : public torch::nn::Module {
|
||||
public:
|
||||
DenseMLPImpl() = default;
|
||||
DenseMLPImpl(int64_t hidden_size,
|
||||
int64_t intermediate_size,
|
||||
bool is_gated,
|
||||
bool has_bias,
|
||||
const std::string& hidden_act,
|
||||
bool enable_result_reduction,
|
||||
const QuantArgs& quant_args,
|
||||
ProcessGroup* process_group,
|
||||
const torch::TensorOptions& options,
|
||||
const std::string& module_prefix = "");
|
||||
|
||||
torch::Tensor forward(const torch::Tensor& hidden_states);
|
||||
|
||||
void load_state_dict(const StateDict& state_dict);
|
||||
void load_state_dict(const StateDict& state_dict,
|
||||
const std::vector<std::string>& gate_up_name,
|
||||
const std::string& down_name);
|
||||
|
||||
// Get FP8 input scale from gate_up_proj for fused RMSNorm+FP8 quantization
|
||||
std::optional<torch::Tensor> get_fp8_input_scale() const;
|
||||
|
||||
private:
|
||||
bool is_gated_;
|
||||
int64_t intermediate_size_;
|
||||
ProcessGroup* process_group_;
|
||||
ColumnParallelLinear gate_up_proj_{nullptr};
|
||||
RowParallelLinear down_proj_{nullptr};
|
||||
Activation act_{nullptr};
|
||||
bool is_smoothquant_;
|
||||
std::string hidden_act_;
|
||||
};
|
||||
TORCH_MODULE(DenseMLP);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
58
ex_engine/xllm_layers/common/fused_moe.cpp
Normal file
58
ex_engine/xllm_layers/common/fused_moe.cpp
Normal file
@@ -0,0 +1,58 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "fused_moe.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
FusedMoEImpl::FusedMoEImpl(const ModelArgs& /*model_args*/,
|
||||
const FusedMoEArgs& /*moe_args*/,
|
||||
const QuantArgs& /*quant_args*/,
|
||||
const ParallelArgs& /*parallel_args*/,
|
||||
const torch::TensorOptions& /*options*/) {
|
||||
NOT_IMPLEMENTED_WITH_MSG(
|
||||
"FusedMoE is not supported for this backend. Please use CUDA, MLU or "
|
||||
"ILU backend for MoE models.");
|
||||
}
|
||||
|
||||
torch::Tensor FusedMoEImpl::forward_experts(
|
||||
const torch::Tensor& /*hidden_states*/,
|
||||
const torch::Tensor& /*router_logits*/,
|
||||
bool /*enable_all2all_communication*/) {
|
||||
NOT_IMPLEMENTED_WITH_MSG(
|
||||
"FusedMoE is not supported for this backend. Please use CUDA, MLU or "
|
||||
"ILU backend for MoE models.");
|
||||
return torch::Tensor();
|
||||
}
|
||||
|
||||
torch::Tensor FusedMoEImpl::forward(const torch::Tensor& /*hidden_states*/,
|
||||
const ModelInputParams& /*input_params*/) {
|
||||
NOT_IMPLEMENTED_WITH_MSG(
|
||||
"FusedMoE is not supported for this backend. Please use CUDA, MLU or "
|
||||
"ILU backend for MoE models.");
|
||||
return torch::Tensor();
|
||||
}
|
||||
|
||||
void FusedMoEImpl::load_state_dict(const StateDict& /*state_dict*/) {
|
||||
NOT_IMPLEMENTED_WITH_MSG(
|
||||
"FusedMoE is not supported for this backend. Please use CUDA, MLU or "
|
||||
"ILU backend for MoE models.");
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
54
ex_engine/xllm_layers/common/fused_moe.h
Normal file
54
ex_engine/xllm_layers/common/fused_moe.h
Normal file
@@ -0,0 +1,54 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include "dense_mlp.h"
|
||||
#include "framework/model/model_args.h"
|
||||
#include "framework/model/model_input_params.h"
|
||||
#include "framework/parallel_state/parallel_args.h"
|
||||
#include "framework/quant_args.h"
|
||||
#include "framework/state_dict/state_dict.h"
|
||||
#include "framework/state_dict/utils.h"
|
||||
#include "fused_moe_base.h"
|
||||
#include "linear.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
// FusedMoE common implementation - placeholder for unsupported backends
|
||||
// Actual implementations are in backend-specific fused_moe.h files.
|
||||
class FusedMoEImpl : public torch::nn::Module {
|
||||
public:
|
||||
FusedMoEImpl() = default;
|
||||
FusedMoEImpl(const ModelArgs& model_args,
|
||||
const FusedMoEArgs& moe_args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
torch::Tensor forward_experts(const torch::Tensor& hidden_states,
|
||||
const torch::Tensor& router_logits,
|
||||
bool enable_all2all_communication);
|
||||
torch::Tensor forward(const torch::Tensor& hidden_states,
|
||||
const ModelInputParams& input_params);
|
||||
void load_state_dict(const StateDict& state_dict);
|
||||
};
|
||||
TORCH_MODULE(FusedMoE);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
144
ex_engine/xllm_layers/common/rms_norm.cpp
Normal file
144
ex_engine/xllm_layers/common/rms_norm.cpp
Normal file
@@ -0,0 +1,144 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "rms_norm.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include "kernels/ops_api.h"
|
||||
#include "platform/device.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
const static std::string kLayerNormMode = "layernorm";
|
||||
const static std::string kRmsNormMode = "rmsnorm";
|
||||
|
||||
RMSNormImpl::RMSNormImpl(int64_t dim,
|
||||
double eps,
|
||||
const torch::TensorOptions& options)
|
||||
: norm_dim_(dim), eps_(eps), mode_(kRmsNormMode) {
|
||||
weight_ = register_parameter("weight",
|
||||
torch::empty({dim}, options),
|
||||
/*requires_grad=*/false);
|
||||
}
|
||||
|
||||
RMSNormImpl::RMSNormImpl(const ModelContext& context)
|
||||
: RMSNormImpl(context.get_model_args().hidden_size(),
|
||||
context.get_model_args().rms_norm_eps(),
|
||||
context.get_tensor_options()) {}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> RMSNormImpl::forward(
|
||||
torch::Tensor& input,
|
||||
std::optional<torch::Tensor> residual,
|
||||
std::optional<torch::Tensor> inplace_output) {
|
||||
auto org_shape = input.sizes().vec();
|
||||
input = input.reshape({-1, norm_dim_});
|
||||
|
||||
torch::Tensor output;
|
||||
if (Device::type_str() != "npu") {
|
||||
if (inplace_output.has_value()) {
|
||||
output = inplace_output.value();
|
||||
output = output.reshape({-1, norm_dim_});
|
||||
} else {
|
||||
output = torch::empty_like(input);
|
||||
}
|
||||
}
|
||||
|
||||
std::optional<torch::Tensor> residual_out;
|
||||
if (residual.has_value()) {
|
||||
residual.value() = residual.value().reshape({-1, norm_dim_});
|
||||
if (Device::type_str() == "mlu" || Device::type_str() == "ilu") {
|
||||
residual_out = residual.value();
|
||||
}
|
||||
}
|
||||
|
||||
xllm::kernel::FusedLayerNormParams fused_layernorm_params;
|
||||
fused_layernorm_params.input = input;
|
||||
fused_layernorm_params.residual = residual;
|
||||
fused_layernorm_params.output = output;
|
||||
fused_layernorm_params.residual_out = residual_out;
|
||||
fused_layernorm_params.weight = weight_;
|
||||
fused_layernorm_params.eps = eps_;
|
||||
fused_layernorm_params.mode = mode_;
|
||||
fused_layernorm_params.store_output_before_norm = residual_out.has_value();
|
||||
if (bias_.defined()) {
|
||||
fused_layernorm_params.beta = bias_;
|
||||
}
|
||||
|
||||
xllm::kernel::fused_layernorm(fused_layernorm_params);
|
||||
|
||||
output = fused_layernorm_params.output;
|
||||
residual_out = fused_layernorm_params.residual_out;
|
||||
|
||||
output = output.view(org_shape);
|
||||
if (residual_out.has_value()) {
|
||||
residual_out.value() = residual_out.value().view(org_shape);
|
||||
}
|
||||
return std::make_tuple(output, residual_out);
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>>
|
||||
RMSNormImpl::forward_fp8(torch::Tensor& input,
|
||||
const torch::Tensor& fp8_scale,
|
||||
std::optional<torch::Tensor> residual) {
|
||||
// Only supported on CUDA for now
|
||||
CHECK(Device::type_str() == "cuda")
|
||||
<< "forward_fp8 is only supported on CUDA";
|
||||
CHECK(mode_ == kRmsNormMode)
|
||||
<< "forward_fp8 only supports RMSNorm mode, not LayerNorm";
|
||||
|
||||
if (residual.has_value()) {
|
||||
// Fused Add + RMSNorm + FP8 Quantization
|
||||
xllm::kernel::FusedAddRmsNormStaticFp8QuantParams params;
|
||||
params.input = input;
|
||||
params.residual = residual.value();
|
||||
params.weight = weight_;
|
||||
params.scale = fp8_scale;
|
||||
params.epsilon = eps_;
|
||||
|
||||
auto [output, updated_residual] =
|
||||
xllm::kernel::fused_add_rms_norm_static_fp8_quant(params);
|
||||
|
||||
return std::make_tuple(output, updated_residual);
|
||||
} else {
|
||||
// RMSNorm + FP8 Quantization (no residual)
|
||||
xllm::kernel::RmsNormStaticFp8QuantParams params;
|
||||
params.input = input;
|
||||
params.weight = weight_;
|
||||
params.scale = fp8_scale;
|
||||
params.epsilon = eps_;
|
||||
|
||||
auto output = xllm::kernel::rms_norm_static_fp8_quant(params);
|
||||
|
||||
return std::make_tuple(output, std::nullopt);
|
||||
}
|
||||
}
|
||||
|
||||
void RMSNormImpl::load_state_dict(const StateDict& state_dict) {
|
||||
LOAD_WEIGHT(weight);
|
||||
if (bias_.defined()) {
|
||||
LOAD_WEIGHT(bias);
|
||||
}
|
||||
}
|
||||
|
||||
void RMSNormImpl::set_layernorm_mode() {
|
||||
mode_ = kLayerNormMode;
|
||||
bias_ = register_parameter(
|
||||
"bias", torch::empty({norm_dim_}, weight_.options()), false);
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
64
ex_engine/xllm_layers/common/rms_norm.h
Normal file
64
ex_engine/xllm_layers/common/rms_norm.h
Normal file
@@ -0,0 +1,64 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include "core/framework/model_context.h"
|
||||
#include "framework/state_dict/state_dict.h"
|
||||
#include "framework/state_dict/utils.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class RMSNormImpl : public torch::nn::Module {
|
||||
public:
|
||||
RMSNormImpl(int64_t dim, double eps, const torch::TensorOptions& options);
|
||||
RMSNormImpl(const ModelContext& context);
|
||||
|
||||
// Standard forward: returns (normalized_output, updated_residual)
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
|
||||
torch::Tensor& input,
|
||||
std::optional<torch::Tensor> residual = std::nullopt,
|
||||
std::optional<torch::Tensor> inplace_output = std::nullopt);
|
||||
|
||||
// Fused forward with FP8 quantization output (for static quantization)
|
||||
// Returns: (fp8_quantized_output, updated_residual)
|
||||
// This combines RMSNorm + FP8 quantization to reduce memory bandwidth
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward_fp8(
|
||||
torch::Tensor& input,
|
||||
const torch::Tensor& fp8_scale,
|
||||
std::optional<torch::Tensor> residual = std::nullopt);
|
||||
|
||||
void set_layernorm_mode();
|
||||
|
||||
void load_state_dict(const StateDict& state_dict);
|
||||
|
||||
torch::Tensor weight() const { return weight_; }
|
||||
torch::Tensor bias() const { return bias_; }
|
||||
double eps() const { return eps_; }
|
||||
|
||||
private:
|
||||
DEFINE_WEIGHT(weight);
|
||||
DEFINE_WEIGHT(bias);
|
||||
int64_t norm_dim_;
|
||||
double eps_;
|
||||
std::string mode_;
|
||||
};
|
||||
TORCH_MODULE(RMSNorm);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
307
ex_engine/xllm_layers/common/rotary_embedding.cpp
Normal file
307
ex_engine/xllm_layers/common/rotary_embedding.cpp
Normal file
@@ -0,0 +1,307 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "rotary_embedding.h"
|
||||
|
||||
#include "kernels/ops_api.h"
|
||||
#include "platform/device.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
RotaryEmbeddingImpl::RotaryEmbeddingImpl(const ModelContext& context) {
|
||||
LOG(FATAL) << "Not implement currently.";
|
||||
}
|
||||
|
||||
RotaryEmbeddingImpl::RotaryEmbeddingImpl(int64_t rotary_dim,
|
||||
int64_t max_position_embeddings,
|
||||
int64_t rope_theta,
|
||||
bool interleaved,
|
||||
const torch::TensorOptions& options)
|
||||
: interleaved_(interleaved) {
|
||||
auto inv_freq = rotary::compute_inv_freq(rotary_dim, rope_theta, options);
|
||||
const auto cos_sin = rotary::compute_cos_sin_cache(
|
||||
rotary_dim, max_position_embeddings, interleaved, inv_freq, options);
|
||||
cos_sin_cache_ = register_buffer("cos_sin_cache", cos_sin);
|
||||
|
||||
auto cos_sin_vec = cos_sin_cache_.chunk(2, /*dim=*/-1);
|
||||
cos_ = cos_sin_vec[0].view({-1, rotary_dim});
|
||||
sin_ = cos_sin_vec[1].view({-1, rotary_dim});
|
||||
|
||||
// Pre-compute [cos_half, sin_half] format used by the CUDA/ILU/MUSA kernels.
|
||||
const auto dev = Device::type_str();
|
||||
if (dev == "cuda" || dev == "ilu" || dev == "musa") {
|
||||
auto chunks = cos_sin_cache_.chunk(4, -1);
|
||||
precomputed_cos_sin_cache_ =
|
||||
torch::cat({chunks[0], chunks[2]}, -1).contiguous();
|
||||
}
|
||||
}
|
||||
|
||||
void RotaryEmbeddingImpl::forward(torch::Tensor& q,
|
||||
torch::Tensor& k,
|
||||
const torch::Tensor& positions,
|
||||
const torch::Tensor& cu_query_lens,
|
||||
int64_t max_query_len,
|
||||
bool is_prompt) {
|
||||
bool discrete;
|
||||
std::optional<torch::Tensor> position_ids;
|
||||
if (is_prompt) {
|
||||
discrete = false;
|
||||
if (Device::type_str() == "cuda" || Device::type_str() == "npu" ||
|
||||
Device::type_str() == "ilu" || Device::type_str() == "musa") {
|
||||
position_ids = positions;
|
||||
}
|
||||
} else {
|
||||
discrete = true;
|
||||
position_ids = positions;
|
||||
}
|
||||
|
||||
xllm::kernel::RotaryParams rotary_params;
|
||||
rotary_params.q = q;
|
||||
rotary_params.k = k;
|
||||
rotary_params.sin = sin_;
|
||||
rotary_params.cos = cos_;
|
||||
rotary_params.cos_sin = cos_sin_cache_;
|
||||
rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_;
|
||||
rotary_params.position_ids = position_ids;
|
||||
rotary_params.cu_query_lens = cu_query_lens;
|
||||
rotary_params.interleaved = interleaved_;
|
||||
rotary_params.discrete = discrete;
|
||||
rotary_params.max_query_len = max_query_len;
|
||||
xllm::kernel::apply_rotary(rotary_params);
|
||||
|
||||
q = rotary_params.q;
|
||||
k = rotary_params.k;
|
||||
}
|
||||
|
||||
// Single tensor forward for MLA architecture
|
||||
void RotaryEmbeddingImpl::forward(torch::Tensor& input,
|
||||
const torch::Tensor& positions,
|
||||
const torch::Tensor& cu_query_lens,
|
||||
int64_t max_query_len,
|
||||
bool is_prompt) {
|
||||
bool discrete;
|
||||
std::optional<torch::Tensor> position_ids;
|
||||
if (is_prompt) {
|
||||
discrete = false;
|
||||
if (Device::type_str() == "cuda" || Device::type_str() == "npu" ||
|
||||
Device::type_str() == "ilu") {
|
||||
position_ids = positions;
|
||||
}
|
||||
} else {
|
||||
discrete = true;
|
||||
position_ids = positions;
|
||||
}
|
||||
|
||||
xllm::kernel::RotaryParams rotary_params;
|
||||
rotary_params.q = input;
|
||||
rotary_params.sin = sin_;
|
||||
rotary_params.cos = cos_;
|
||||
rotary_params.cos_sin = cos_sin_cache_;
|
||||
rotary_params.position_ids = position_ids;
|
||||
rotary_params.cu_query_lens = cu_query_lens;
|
||||
rotary_params.interleaved = interleaved_;
|
||||
rotary_params.discrete = discrete;
|
||||
rotary_params.max_query_len = max_query_len;
|
||||
xllm::kernel::apply_rotary(rotary_params);
|
||||
|
||||
input = rotary_params.q;
|
||||
}
|
||||
|
||||
MRotaryEmbeddingImpl::MRotaryEmbeddingImpl(
|
||||
int64_t rotary_dim,
|
||||
int64_t max_position_embeddings,
|
||||
int64_t rope_theta,
|
||||
bool interleaved,
|
||||
const std::vector<int64_t>& rope_scaling_mrope_section,
|
||||
const torch::TensorOptions& options)
|
||||
: RotaryEmbeddingImpl(rotary_dim,
|
||||
max_position_embeddings,
|
||||
rope_theta,
|
||||
interleaved,
|
||||
options),
|
||||
mrope_section_(rope_scaling_mrope_section) {
|
||||
mrope_cu_seq_lens_ = torch::zeros(2, torch::kInt32).to(options.device());
|
||||
}
|
||||
|
||||
void MRotaryEmbeddingImpl::forward(torch::Tensor& q,
|
||||
torch::Tensor& k,
|
||||
const torch::Tensor& positions,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
bool only_prefill =
|
||||
(attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
|
||||
if (!only_prefill || mrope_section_.empty()) {
|
||||
torch::Tensor position_ids = positions;
|
||||
if (positions.dim() == 2) {
|
||||
position_ids = positions[0];
|
||||
}
|
||||
return RotaryEmbeddingImpl::forward(q,
|
||||
k,
|
||||
position_ids,
|
||||
attn_metadata.q_cu_seq_lens,
|
||||
attn_metadata.max_query_len,
|
||||
attn_metadata.is_prefill);
|
||||
}
|
||||
|
||||
int64_t num_tokens = positions.size(-1);
|
||||
mrope_cu_seq_lens_[1] = num_tokens;
|
||||
CHECK(attn_metadata.mrope_cos.defined() && attn_metadata.mrope_sin.defined());
|
||||
xllm::kernel::RotaryParams rotary_params;
|
||||
rotary_params.q = q;
|
||||
rotary_params.k = k;
|
||||
rotary_params.sin = attn_metadata.mrope_sin;
|
||||
rotary_params.cos = attn_metadata.mrope_cos;
|
||||
rotary_params.cos_sin = cos_sin_cache_;
|
||||
rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_;
|
||||
rotary_params.position_ids = std::nullopt;
|
||||
rotary_params.cu_query_lens = mrope_cu_seq_lens_;
|
||||
rotary_params.interleaved = interleaved_;
|
||||
rotary_params.discrete = false;
|
||||
rotary_params.max_query_len = num_tokens;
|
||||
xllm::kernel::apply_rotary(rotary_params);
|
||||
|
||||
q = rotary_params.q;
|
||||
k = rotary_params.k;
|
||||
}
|
||||
|
||||
DeepseekScalingRotaryEmbeddingImpl::DeepseekScalingRotaryEmbeddingImpl(
|
||||
int64_t head_size,
|
||||
int64_t rotary_dim,
|
||||
int64_t max_position_embeddings,
|
||||
int64_t rope_scaling_original_max_position_embeddings,
|
||||
int64_t rope_theta,
|
||||
bool interleaved,
|
||||
float scaling_factor,
|
||||
float extrapolation_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float mscale,
|
||||
float mscale_all_dim,
|
||||
const torch::TensorOptions& options)
|
||||
: head_size_(head_size),
|
||||
rotary_dim_(rotary_dim),
|
||||
interleaved_(interleaved) {
|
||||
auto inv_freq = rotary::apply_deepseek_yarn_rope_scaling(
|
||||
scaling_factor,
|
||||
extrapolation_factor,
|
||||
beta_fast,
|
||||
beta_slow,
|
||||
rotary_dim,
|
||||
rope_theta,
|
||||
rope_scaling_original_max_position_embeddings);
|
||||
const auto cos_sin = rotary::compute_cos_sin_cache(rotary_dim,
|
||||
max_position_embeddings,
|
||||
interleaved,
|
||||
scaling_factor,
|
||||
attn_factor,
|
||||
mscale,
|
||||
mscale_all_dim,
|
||||
inv_freq,
|
||||
options);
|
||||
cos_sin_cache_ = register_buffer("cos_sin_cache", cos_sin);
|
||||
|
||||
auto cos_sin_vec = cos_sin_cache_.chunk(2, /*dim=*/-1);
|
||||
cos_ = cos_sin_vec[0].view({-1, rotary_dim});
|
||||
sin_ = cos_sin_vec[1].view({-1, rotary_dim});
|
||||
|
||||
// Pre-compute [cos_half, sin_half] format used by the CUDA/ILU/MUSA kernels.
|
||||
const auto dev = Device::type_str();
|
||||
if (dev == "cuda" || dev == "ilu" || dev == "musa") {
|
||||
auto chunks = cos_sin_cache_.chunk(4, -1);
|
||||
precomputed_cos_sin_cache_ =
|
||||
torch::cat({chunks[0], chunks[2]}, -1).contiguous();
|
||||
}
|
||||
}
|
||||
|
||||
void DeepseekScalingRotaryEmbeddingImpl::forward(
|
||||
torch::Tensor& input,
|
||||
const torch::Tensor& positions,
|
||||
const torch::Tensor& cu_query_lens,
|
||||
int64_t max_query_len,
|
||||
bool is_prompt) {
|
||||
const int32_t dim = -1;
|
||||
bool discrete;
|
||||
std::optional<torch::Tensor> position_ids;
|
||||
if (is_prompt) {
|
||||
discrete = false;
|
||||
position_ids = std::nullopt;
|
||||
} else {
|
||||
discrete = true;
|
||||
position_ids = positions;
|
||||
max_query_len = 1;
|
||||
}
|
||||
auto input_rot = input.slice(dim, 0, rotary_dim_);
|
||||
torch::Tensor input_pass;
|
||||
if (rotary_dim_ < head_size_) {
|
||||
input_pass = input.slice(dim, rotary_dim_, head_size_);
|
||||
}
|
||||
|
||||
xllm::kernel::RotaryParams rotary_params;
|
||||
rotary_params.q = input_rot;
|
||||
rotary_params.sin = sin_;
|
||||
rotary_params.cos = cos_;
|
||||
rotary_params.cos_sin = cos_sin_cache_;
|
||||
rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_;
|
||||
rotary_params.position_ids = position_ids;
|
||||
rotary_params.cu_query_lens = cu_query_lens;
|
||||
rotary_params.interleaved = interleaved_;
|
||||
rotary_params.discrete = discrete;
|
||||
rotary_params.max_query_len = max_query_len;
|
||||
xllm::kernel::apply_rotary(rotary_params);
|
||||
input_rot = rotary_params.q;
|
||||
|
||||
if (rotary_dim_ < head_size_) {
|
||||
input = torch::cat({input_rot, input_pass}, dim);
|
||||
} else {
|
||||
input = input_rot;
|
||||
}
|
||||
}
|
||||
|
||||
// Factory function: creates the appropriate RoPE type based on model args
|
||||
std::shared_ptr<RotaryEmbeddingBase> create_mla_rotary_embedding(
|
||||
const ModelArgs& args,
|
||||
int64_t rotary_dim,
|
||||
int64_t max_position_embeddings,
|
||||
bool interleaved,
|
||||
const torch::TensorOptions& options) {
|
||||
if (args.rope_scaling_rope_type() == "deepseek_yarn") {
|
||||
return std::make_shared<DeepseekScalingRotaryEmbeddingImpl>(
|
||||
rotary_dim, // head_size (same as rotary_dim for MLA)
|
||||
rotary_dim,
|
||||
max_position_embeddings,
|
||||
args.rope_scaling_original_max_position_embeddings(),
|
||||
args.rope_theta(),
|
||||
interleaved,
|
||||
args.rope_scaling_factor(),
|
||||
args.rope_extrapolation_factor(),
|
||||
args.rope_scaling_attn_factor(),
|
||||
args.rope_scaling_beta_fast(),
|
||||
args.rope_scaling_beta_slow(),
|
||||
args.rope_scaling_mscale(),
|
||||
args.rope_scaling_mscale_all_dim(),
|
||||
options);
|
||||
} else {
|
||||
// default rope type
|
||||
return std::make_shared<RotaryEmbeddingImpl>(rotary_dim,
|
||||
max_position_embeddings,
|
||||
args.rope_theta(),
|
||||
interleaved,
|
||||
options);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
158
ex_engine/xllm_layers/common/rotary_embedding.h
Normal file
158
ex_engine/xllm_layers/common/rotary_embedding.h
Normal file
@@ -0,0 +1,158 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
#include <torch/types.h>
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "attention_metadata.h"
|
||||
#include "core/framework/model_context.h"
|
||||
#include "framework/model/model_args.h"
|
||||
#include "rotary_embedding_util.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class RotaryEmbeddingBase : public torch::nn::Module {
|
||||
public:
|
||||
~RotaryEmbeddingBase() override = default;
|
||||
|
||||
virtual void forward(torch::Tensor& input,
|
||||
const torch::Tensor& positions,
|
||||
const torch::Tensor& cu_query_lens,
|
||||
int64_t max_query_len,
|
||||
bool is_prompt) = 0;
|
||||
virtual const torch::Tensor& get_sin_cache() const = 0;
|
||||
virtual const torch::Tensor& get_cos_cache() const = 0;
|
||||
virtual const bool get_interleaved() const = 0;
|
||||
};
|
||||
|
||||
class RotaryEmbeddingImpl : public RotaryEmbeddingBase {
|
||||
public:
|
||||
RotaryEmbeddingImpl(int64_t rotary_dim,
|
||||
int64_t max_position_embeddings,
|
||||
int64_t rope_theta,
|
||||
bool interleaved,
|
||||
const torch::TensorOptions& options);
|
||||
RotaryEmbeddingImpl(const ModelContext& context);
|
||||
|
||||
void forward(torch::Tensor& q,
|
||||
torch::Tensor& k,
|
||||
const torch::Tensor& positions,
|
||||
const torch::Tensor& cu_query_lens,
|
||||
int64_t max_query_len,
|
||||
bool is_prompt);
|
||||
// Single tensor forward for MLA architecture
|
||||
void forward(torch::Tensor& input,
|
||||
const torch::Tensor& positions,
|
||||
const torch::Tensor& cu_query_lens,
|
||||
int64_t max_query_len,
|
||||
bool is_prompt) override;
|
||||
|
||||
const torch::Tensor& precomputed_cos_sin_cache() {
|
||||
return precomputed_cos_sin_cache_;
|
||||
}
|
||||
|
||||
torch::Tensor get_cos_sin_cache() { return cos_sin_cache_; }
|
||||
const torch::Tensor& get_sin_cache() const override { return sin_; }
|
||||
const torch::Tensor& get_cos_cache() const override { return cos_; }
|
||||
const bool get_interleaved() const override { return interleaved_; }
|
||||
|
||||
protected:
|
||||
bool interleaved_;
|
||||
torch::Tensor cos_sin_cache_;
|
||||
// Pre-formatted [cos_half, sin_half] cache for CUDA/MUSA/ILU kernels.
|
||||
// Avoids chunk/cat operations on every forward call.
|
||||
torch::Tensor precomputed_cos_sin_cache_;
|
||||
|
||||
private:
|
||||
torch::Tensor sin_;
|
||||
torch::Tensor cos_;
|
||||
};
|
||||
TORCH_MODULE(RotaryEmbedding);
|
||||
|
||||
class MRotaryEmbeddingImpl : public RotaryEmbeddingImpl {
|
||||
public:
|
||||
MRotaryEmbeddingImpl(int64_t rotary_dim,
|
||||
int64_t max_position_embeddings,
|
||||
int64_t rope_theta,
|
||||
bool interleaved,
|
||||
const std::vector<int64_t>& rope_scaling_mrope_section,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
void forward(torch::Tensor& q,
|
||||
torch::Tensor& k,
|
||||
const torch::Tensor& positions,
|
||||
const AttentionMetadata& attn_metadata);
|
||||
|
||||
private:
|
||||
std::vector<int64_t> mrope_section_;
|
||||
torch::Tensor mrope_cu_seq_lens_;
|
||||
};
|
||||
TORCH_MODULE(MRotaryEmbedding);
|
||||
|
||||
class DeepseekScalingRotaryEmbeddingImpl : public RotaryEmbeddingBase {
|
||||
public:
|
||||
DeepseekScalingRotaryEmbeddingImpl(
|
||||
int64_t head_size,
|
||||
int64_t rotary_dim,
|
||||
int64_t max_position_embeddings,
|
||||
int64_t rope_scaling_original_max_position_embeddings,
|
||||
int64_t rope_theta,
|
||||
bool interleaved,
|
||||
float scaling_factor,
|
||||
float extrapolation_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float mscale,
|
||||
float mscale_all_dim,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
void forward(torch::Tensor& input,
|
||||
const torch::Tensor& positions,
|
||||
const torch::Tensor& cu_query_lens,
|
||||
int64_t max_query_len,
|
||||
bool is_prompt) override;
|
||||
const torch::Tensor& get_sin_cache() const override { return sin_; }
|
||||
const torch::Tensor& get_cos_cache() const override { return cos_; }
|
||||
const bool get_interleaved() const override { return interleaved_; }
|
||||
|
||||
private:
|
||||
int64_t head_size_;
|
||||
int64_t rotary_dim_;
|
||||
bool interleaved_;
|
||||
torch::Tensor sin_;
|
||||
torch::Tensor cos_;
|
||||
torch::Tensor cos_sin_cache_;
|
||||
// Pre-formatted [cos_half, sin_half] cache for CUDA/MUSA/ILU kernels.
|
||||
// Avoids chunk/cat operations on every forward call.
|
||||
torch::Tensor precomputed_cos_sin_cache_;
|
||||
};
|
||||
TORCH_MODULE(DeepseekScalingRotaryEmbedding);
|
||||
|
||||
// Factory function: creates the appropriate RoPE type based on model args
|
||||
std::shared_ptr<RotaryEmbeddingBase> create_mla_rotary_embedding(
|
||||
const ModelArgs& args,
|
||||
int64_t rotary_dim,
|
||||
int64_t max_position_embeddings,
|
||||
bool interleaved,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
189
ex_engine/xllm_layers/ilu/attention.cpp
Normal file
189
ex_engine/xllm_layers/ilu/attention.cpp
Normal file
@@ -0,0 +1,189 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "attention.h"
|
||||
|
||||
#include "kernels/ilu/ilu_ops_api.h"
|
||||
#include "kernels/ops_api.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
AttentionImpl::AttentionImpl(int64_t num_heads,
|
||||
int64_t head_size,
|
||||
float scale,
|
||||
int64_t num_kv_heads,
|
||||
int64_t sliding_window)
|
||||
: num_heads_(num_heads),
|
||||
head_size_(head_size),
|
||||
scale_(scale),
|
||||
num_kv_heads_(num_kv_heads),
|
||||
v_head_dim_(head_size),
|
||||
use_fused_mla_qkv_(false),
|
||||
enable_lighting_indexer_(false),
|
||||
enable_mla_(false),
|
||||
sliding_window_(sliding_window) {
|
||||
if (sliding_window_ > -1) {
|
||||
sliding_window_ = sliding_window_ - 1;
|
||||
}
|
||||
}
|
||||
|
||||
AttentionImpl::AttentionImpl(int64_t num_heads,
|
||||
int64_t head_size,
|
||||
int64_t num_kv_heads,
|
||||
int64_t v_head_dim,
|
||||
int64_t sliding_window,
|
||||
float scale,
|
||||
bool use_fused_mla_qkv,
|
||||
bool enable_lighting_indexer,
|
||||
bool enable_mla)
|
||||
: num_heads_(num_heads),
|
||||
head_size_(head_size),
|
||||
scale_(scale),
|
||||
num_kv_heads_(num_kv_heads),
|
||||
v_head_dim_(v_head_dim),
|
||||
use_fused_mla_qkv_(use_fused_mla_qkv),
|
||||
enable_lighting_indexer_(enable_lighting_indexer),
|
||||
enable_mla_(enable_mla),
|
||||
sliding_window_(sliding_window) {
|
||||
if (sliding_window_ > -1) {
|
||||
sliding_window_ = sliding_window_ - 1;
|
||||
}
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> AttentionImpl::forward(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
torch::Tensor& value,
|
||||
KVCache& kv_cache) {
|
||||
std::optional<torch::Tensor> output_lse = std::nullopt;
|
||||
torch::Tensor output;
|
||||
if (enable_mla_) {
|
||||
output = torch::empty({query.size(0), num_heads_ * v_head_dim_},
|
||||
query.options());
|
||||
} else {
|
||||
output = torch::empty_like(query);
|
||||
}
|
||||
if (attn_metadata.is_dummy) {
|
||||
return std::make_tuple(output, output_lse);
|
||||
}
|
||||
|
||||
bool only_prefill =
|
||||
attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
|
||||
int64_t num_kv_heads = (enable_mla_ && !only_prefill) ? 1 : num_kv_heads_;
|
||||
torch::Tensor k_cache = kv_cache.get_k_cache();
|
||||
std::optional<torch::Tensor> v_cache;
|
||||
std::optional<torch::Tensor> v;
|
||||
if (!enable_mla_) {
|
||||
v = value.view({-1, num_kv_heads, head_size_});
|
||||
v_cache = kv_cache.get_v_cache();
|
||||
}
|
||||
|
||||
bool skip_process_cache = enable_mla_ && (only_prefill || use_fused_mla_qkv_);
|
||||
if (!skip_process_cache) {
|
||||
xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params;
|
||||
reshape_paged_cache_params.key = key.view({-1, num_kv_heads, head_size_});
|
||||
reshape_paged_cache_params.value = v;
|
||||
reshape_paged_cache_params.k_cache = k_cache;
|
||||
reshape_paged_cache_params.v_cache = v_cache;
|
||||
reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping;
|
||||
xllm::kernel::reshape_paged_cache(reshape_paged_cache_params);
|
||||
}
|
||||
|
||||
if (enable_lighting_indexer_ || !only_prefill) {
|
||||
decoder_forward(query, output, k_cache, v_cache, attn_metadata);
|
||||
} else {
|
||||
prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata);
|
||||
}
|
||||
|
||||
int64_t head_size = enable_mla_ ? v_head_dim_ : head_size_;
|
||||
output = output.view({-1, num_heads_ * head_size});
|
||||
return {output, output_lse};
|
||||
}
|
||||
|
||||
void AttentionImpl::prefill_forward(torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
torch::Tensor& value,
|
||||
torch::Tensor& output,
|
||||
const torch::Tensor& k_cache,
|
||||
const std::optional<torch::Tensor>& v_cache,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
|
||||
std::optional<torch::Tensor> output_lse = std::nullopt;
|
||||
query = query.view({-1, num_heads_, head_size_});
|
||||
output = output.view({-1, num_heads_, head_size_v});
|
||||
// torch::Tensor k_cache_ = k_cache;
|
||||
// torch::Tensor v_cache_ = v_cache.value();
|
||||
xllm::kernel::ilu::batch_prefill(query,
|
||||
k_cache,
|
||||
v_cache,
|
||||
output,
|
||||
output_lse,
|
||||
attn_metadata.q_cu_seq_lens,
|
||||
attn_metadata.kv_cu_seq_lens,
|
||||
/*alibi_slope=*/std::nullopt,
|
||||
/*attn_bias=*/std::nullopt,
|
||||
/*q_quant_scale=*/std::nullopt,
|
||||
/*k_quant_scale=*/std::nullopt,
|
||||
/*v_quant_scale=*/std::nullopt,
|
||||
attn_metadata.block_table,
|
||||
attn_metadata.max_query_len,
|
||||
attn_metadata.max_seq_len,
|
||||
scale_,
|
||||
attn_metadata.is_causal,
|
||||
sliding_window_,
|
||||
/*window_size_right=*/-1,
|
||||
attn_metadata.compute_dtype,
|
||||
/*return_lse=*/false);
|
||||
}
|
||||
|
||||
void AttentionImpl::decoder_forward(torch::Tensor& query,
|
||||
torch::Tensor& output,
|
||||
const torch::Tensor& k_cache,
|
||||
const std::optional<torch::Tensor>& v_cache,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
|
||||
query = query.view({-1, 1, num_heads_, head_size_});
|
||||
output = output.view({-1, 1, num_heads_, head_size_v});
|
||||
std::optional<torch::Tensor> output_lse = std::nullopt;
|
||||
|
||||
int64_t block_aligned_max_seq_len =
|
||||
attn_metadata.block_table.size(-1) * k_cache.size(2);
|
||||
|
||||
xllm::kernel::ilu::batch_decode(query,
|
||||
k_cache,
|
||||
output,
|
||||
attn_metadata.block_table,
|
||||
attn_metadata.kv_seq_lens,
|
||||
v_cache,
|
||||
output_lse,
|
||||
/*q_quant_scale=*/std::nullopt,
|
||||
/*k_quant_scale=*/std::nullopt,
|
||||
/*v_quant_scale=*/std::nullopt,
|
||||
/*out_quant_scale=*/std::nullopt,
|
||||
/*alibi_slope=*/std::nullopt,
|
||||
attn_metadata.attn_mask,
|
||||
attn_metadata.compute_dtype,
|
||||
block_aligned_max_seq_len,
|
||||
sliding_window_,
|
||||
/*window_size_right=*/-1,
|
||||
scale_,
|
||||
/*return_lse=*/false,
|
||||
attn_metadata.is_causal,
|
||||
/*kv_cache_quant_bit_size=*/-1);
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
82
ex_engine/xllm_layers/ilu/attention.h
Normal file
82
ex_engine/xllm_layers/ilu/attention.h
Normal file
@@ -0,0 +1,82 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <tuple>
|
||||
|
||||
#include "framework/kv_cache/kv_cache.h"
|
||||
#include "framework/model/model_input_params.h"
|
||||
#include "layers/common/attention_metadata.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
class AttentionImpl : public torch::nn::Module {
|
||||
public:
|
||||
AttentionImpl() = default;
|
||||
|
||||
AttentionImpl(int64_t num_heads,
|
||||
int64_t head_size,
|
||||
float scale,
|
||||
int64_t num_kv_heads,
|
||||
int64_t sliding_window);
|
||||
AttentionImpl(int64_t num_heads,
|
||||
int64_t head_size,
|
||||
int64_t num_kv_heads,
|
||||
int64_t v_head_dim,
|
||||
int64_t sliding_window,
|
||||
float scale,
|
||||
bool use_fused_mla_qkv,
|
||||
bool enable_lighting_indexer,
|
||||
bool enable_mla);
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
torch::Tensor& value,
|
||||
KVCache& kv_cache);
|
||||
|
||||
void prefill_forward(torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
torch::Tensor& value,
|
||||
torch::Tensor& output,
|
||||
const torch::Tensor& k_cache,
|
||||
const std::optional<torch::Tensor>& v_cache,
|
||||
const AttentionMetadata& attn_metadata);
|
||||
|
||||
void decoder_forward(torch::Tensor& query,
|
||||
torch::Tensor& output,
|
||||
const torch::Tensor& k_cache,
|
||||
const std::optional<torch::Tensor>& v_cache,
|
||||
const AttentionMetadata& attn_metadata);
|
||||
|
||||
private:
|
||||
int64_t num_heads_;
|
||||
int64_t head_size_;
|
||||
float scale_;
|
||||
int64_t num_kv_heads_;
|
||||
int64_t v_head_dim_;
|
||||
bool use_fused_mla_qkv_;
|
||||
bool enable_lighting_indexer_;
|
||||
bool enable_mla_;
|
||||
int64_t sliding_window_;
|
||||
};
|
||||
TORCH_MODULE(Attention);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
797
ex_engine/xllm_layers/ilu/fused_moe.cpp
Normal file
797
ex_engine/xllm_layers/ilu/fused_moe.cpp
Normal file
@@ -0,0 +1,797 @@
|
||||
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "fused_moe.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include <iomanip>
|
||||
|
||||
#include "common/global_flags.h"
|
||||
#include "framework/parallel_state/parallel_state.h"
|
||||
#include "kernels/ops_api.h"
|
||||
#include "layers/common/dp_utils.h"
|
||||
#include "util/utils.h"
|
||||
|
||||
namespace {
|
||||
|
||||
int32_t get_dtype_size(torch::ScalarType dtype) {
|
||||
return static_cast<int32_t>(torch::elementSize(dtype));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
|
||||
const FusedMoEArgs& moe_args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options)
|
||||
: num_total_experts_(static_cast<int64_t>(model_args.n_routed_experts())),
|
||||
topk_(model_args.num_experts_per_tok()),
|
||||
num_expert_group_(model_args.n_group()),
|
||||
topk_group_(model_args.topk_group()),
|
||||
route_scale_(model_args.routed_scaling_factor()),
|
||||
hidden_size_(model_args.hidden_size()),
|
||||
n_shared_experts_(model_args.n_shared_experts()),
|
||||
is_gated_(moe_args.is_gated),
|
||||
renormalize_(model_args.norm_topk_prob() ? 1 : 0),
|
||||
hidden_act_(model_args.hidden_act()),
|
||||
scoring_func_(model_args.scoring_func()),
|
||||
quant_args_(quant_args),
|
||||
parallel_args_(parallel_args),
|
||||
options_(options),
|
||||
device_(options.device()) {
|
||||
const int64_t num_experts = num_total_experts_;
|
||||
const int64_t intermediate_size =
|
||||
static_cast<int64_t>(model_args.moe_intermediate_size());
|
||||
const std::string& topk_method = model_args.topk_method();
|
||||
int64_t ep_size = parallel_args.ep_size();
|
||||
int64_t ep_rank = 0;
|
||||
tp_pg_ = parallel_args.tp_group_;
|
||||
if (ep_size > 1) {
|
||||
ep_rank = parallel_args.moe_ep_group_->rank();
|
||||
tp_pg_ = parallel_args.moe_tp_group_;
|
||||
}
|
||||
|
||||
// smoothquant check: If quant_method is not empty, only w8a8 smoothquant is
|
||||
// supported
|
||||
if (!quant_args.quant_method().empty()) {
|
||||
if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 ||
|
||||
!quant_args.activation_dynamic()) {
|
||||
LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when "
|
||||
"quant_method is set. "
|
||||
<< "Got quant_method=" << quant_args.quant_method()
|
||||
<< ", bits=" << quant_args.bits()
|
||||
<< ", activation_dynamic=" << quant_args.activation_dynamic();
|
||||
}
|
||||
// If confirmed as smoothquant w8a8, set is_smoothquant_ to true
|
||||
is_smoothquant_ = true;
|
||||
} else {
|
||||
is_smoothquant_ = false;
|
||||
}
|
||||
|
||||
// Deep EP initialization check
|
||||
enable_deep_ep_ = FLAGS_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)
|
||||
// K is the number of speculative tokens.
|
||||
int64_t dispatch_token_size;
|
||||
if (quant_args.quant_method() == "smoothquant") {
|
||||
// float32 is for the scale of the quantized input
|
||||
dispatch_token_size = hidden_size_ * get_dtype_size(torch::kInt8) +
|
||||
get_dtype_size(torch::kFloat32);
|
||||
} else {
|
||||
dispatch_token_size =
|
||||
hidden_size_ * get_dtype_size(options_.dtype().toScalarType());
|
||||
}
|
||||
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 max_num_tokens_per_rank =
|
||||
(1 + FLAGS_num_speculative_tokens) * effective_seqs * topk_;
|
||||
|
||||
// make sure that all layers share the same deep ep instance
|
||||
// so that the memory footprint is minimized
|
||||
deep_ep_ = DeepEPManager::get_instance(dispatch_token_size,
|
||||
combine_token_size,
|
||||
max_num_tokens_per_rank,
|
||||
num_experts,
|
||||
parallel_args,
|
||||
options_);
|
||||
|
||||
// obtain the buffer and parameters of deep ep
|
||||
deep_ep_buffer_ = deep_ep_->get_buffer();
|
||||
deep_ep_params_ = deep_ep_->get_params();
|
||||
|
||||
// intermediate buffer that can be initialized once
|
||||
// we place these tensor here in order to speed up forward pass
|
||||
int64_t n_tokens_recv = deep_ep_params_.max_num_tokens_recv;
|
||||
int64_t token_bytes = is_smoothquant_
|
||||
? get_dtype_size(torch::kInt8)
|
||||
: get_dtype_size(options_.dtype().toScalarType());
|
||||
token_bytes = token_bytes * hidden_size_;
|
||||
int64_t head_size = n_tokens_recv * token_bytes;
|
||||
dispatch_recv_token_tensor_head_ =
|
||||
deep_ep_buffer_.combine_send_token_tensor.narrow(0, 0, head_size)
|
||||
.view({n_tokens_recv, token_bytes});
|
||||
// input scale in smoothquant
|
||||
if (is_smoothquant_) {
|
||||
int64_t tail_size = n_tokens_recv * get_dtype_size(torch::kFloat32);
|
||||
dispatch_recv_token_tensor_tail_ =
|
||||
deep_ep_buffer_.combine_send_token_tensor
|
||||
.narrow(0, head_size, tail_size)
|
||||
.view({n_tokens_recv, -1});
|
||||
}
|
||||
}
|
||||
|
||||
// calculate the number of experts per rank
|
||||
num_experts_per_rank_ = num_experts / ep_size;
|
||||
start_expert_id_ = ep_rank * num_experts_per_rank_;
|
||||
|
||||
if (topk_method == "noaux_tc") {
|
||||
e_score_correction_bias_ = register_parameter(
|
||||
"e_score_correction_bias", torch::empty({num_experts}, options), false);
|
||||
}
|
||||
|
||||
gate_ = register_module(
|
||||
"gate_proj",
|
||||
ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options));
|
||||
if (n_shared_experts_ > 0) {
|
||||
ProcessGroup* shared_expert_pg;
|
||||
if (parallel_args_.ep_size() > 1) {
|
||||
// we use tp=1 for shared experts computation in deep ep mode
|
||||
CHECK(parallel_args_.ep_size() == parallel_args_.world_size())
|
||||
<< "Models with shared experts only support ep_size equal to "
|
||||
"world size for now.";
|
||||
shared_expert_pg = parallel_args.moe_tp_group_;
|
||||
} else {
|
||||
shared_expert_pg = parallel_args.process_group_;
|
||||
}
|
||||
// The shared experts computation can proceed in parallel with the
|
||||
// final communication step during the MoE computation, as long as it
|
||||
// remains independent of any communication operations. For optimal
|
||||
// performance, ensure that the shared experts layer on each rank always
|
||||
// maintains its own unique weights.
|
||||
shared_experts_ =
|
||||
register_module("shared_experts",
|
||||
DenseMLP(hidden_size_,
|
||||
intermediate_size * n_shared_experts_,
|
||||
is_gated_,
|
||||
false,
|
||||
hidden_act_,
|
||||
/*enable_result_reduction=*/true,
|
||||
quant_args,
|
||||
shared_expert_pg,
|
||||
options));
|
||||
}
|
||||
|
||||
// create weight buffer
|
||||
const int64_t world_size = tp_pg_->world_size();
|
||||
int64_t local_intermediate_size = intermediate_size / world_size;
|
||||
if (is_smoothquant_) {
|
||||
auto quant_option = options_.dtype(torch::kInt8);
|
||||
auto fp_option = options_.dtype(torch::kFloat32);
|
||||
w13_ = register_parameter(
|
||||
"w13",
|
||||
torch::empty(
|
||||
{num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
|
||||
quant_option),
|
||||
false);
|
||||
w13_scale_ = register_parameter(
|
||||
"w13_scale",
|
||||
torch::empty({num_experts_per_rank_, local_intermediate_size * 2},
|
||||
fp_option),
|
||||
false);
|
||||
// Note: We do not check enable_deep_ep_ here, since smooth quantization
|
||||
// information may be needed even when deep EP mode is disabled. This allows
|
||||
// retrieving quantization parameters for any subset of experts as required.
|
||||
input_smooth_ = register_parameter(
|
||||
"input_smooth",
|
||||
torch::empty({num_total_experts_, hidden_size_}, fp_option),
|
||||
false);
|
||||
w2_ = register_parameter(
|
||||
"w2",
|
||||
torch::empty(
|
||||
{num_experts_per_rank_, hidden_size_, local_intermediate_size},
|
||||
quant_option),
|
||||
false);
|
||||
w2_scale_ = register_parameter(
|
||||
"w2_scale",
|
||||
torch::empty({num_experts_per_rank_, hidden_size_}, fp_option),
|
||||
false);
|
||||
act_smooth_ = register_parameter(
|
||||
"act_smooth",
|
||||
torch::empty({num_experts_per_rank_, local_intermediate_size},
|
||||
fp_option),
|
||||
false);
|
||||
|
||||
} else {
|
||||
w13_ = register_parameter(
|
||||
"w13",
|
||||
torch::empty(
|
||||
{num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
|
||||
options_),
|
||||
false);
|
||||
w2_ = register_parameter(
|
||||
"w2",
|
||||
torch::empty(
|
||||
{num_experts_per_rank_, hidden_size_, local_intermediate_size},
|
||||
options_),
|
||||
false);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor FusedMoEImpl::create_group_gemm_output(
|
||||
const torch::Tensor& a,
|
||||
const torch::Tensor& b,
|
||||
const torch::Tensor& group_list,
|
||||
torch::ScalarType dtype,
|
||||
torch::Tensor& workspace) {
|
||||
// unify shape logic: define the target shape once.
|
||||
bool is_3d_weight = (b.dim() != 2);
|
||||
int64_t num_tokens = a.size(0);
|
||||
int64_t out_dim = is_3d_weight ? b.size(1) : b.size(0);
|
||||
|
||||
std::vector<int64_t> output_shape;
|
||||
int64_t required_elements = num_tokens * out_dim;
|
||||
|
||||
if (is_3d_weight) {
|
||||
output_shape = {num_tokens, out_dim};
|
||||
} else {
|
||||
output_shape = {group_list.size(0), num_tokens, out_dim};
|
||||
required_elements *= group_list.size(0);
|
||||
}
|
||||
|
||||
auto options = a.options().dtype(dtype);
|
||||
|
||||
// non-smoothquant: direct allocation
|
||||
if (!is_smoothquant_) {
|
||||
return torch::empty(output_shape, options);
|
||||
}
|
||||
|
||||
// smoothquant: managed workspace logic
|
||||
if (!workspace.defined()) {
|
||||
// Lazy initialization: allocate max buffer for the lifecycle
|
||||
// Note: accessing class members w13_ and w2_ directly for context
|
||||
int64_t max_width = std::max(w13_.size(1), w2_.size(1));
|
||||
workspace = torch::empty({num_tokens * max_width}, options);
|
||||
}
|
||||
|
||||
// view construction
|
||||
CHECK(workspace.numel() >= required_elements)
|
||||
<< "FusedMoE Workspace too small! Alloc: " << workspace.numel()
|
||||
<< ", Req: " << required_elements;
|
||||
|
||||
// utilize the pre-calculated output_shape
|
||||
return workspace.slice(0, 0, required_elements).view(output_shape);
|
||||
}
|
||||
|
||||
torch::Tensor FusedMoEImpl::select_experts(
|
||||
const torch::Tensor& hidden_states_2d,
|
||||
const torch::Tensor& router_logits_2d,
|
||||
SelectedExpertInfo& selected_expert_info,
|
||||
bool enable_all2all_communication) {
|
||||
// prepare the parameters for select_experts
|
||||
std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
|
||||
if (e_score_correction_bias_.defined()) {
|
||||
e_score_correction_bias = e_score_correction_bias_;
|
||||
}
|
||||
int64_t expert_size = w13_.size(0);
|
||||
|
||||
// Step 1: apply softmax topk or sigmoid topk / routing logic
|
||||
torch::Tensor reduce_weight;
|
||||
torch::Tensor expert_id;
|
||||
{
|
||||
xllm::kernel::MoeFusedTopkParams moe_active_topk_params;
|
||||
moe_active_topk_params.input = router_logits_2d;
|
||||
moe_active_topk_params.topk = topk_;
|
||||
moe_active_topk_params.num_expert_group = num_expert_group_;
|
||||
moe_active_topk_params.topk_group = topk_group_;
|
||||
moe_active_topk_params.normalize = renormalize_;
|
||||
moe_active_topk_params.normed_by = "topk_logit";
|
||||
moe_active_topk_params.scoring_func = scoring_func_;
|
||||
moe_active_topk_params.route_scale = route_scale_;
|
||||
moe_active_topk_params.e_score_correction_bias = e_score_correction_bias;
|
||||
std::tie(reduce_weight, expert_id) =
|
||||
xllm::kernel::moe_active_topk(moe_active_topk_params);
|
||||
}
|
||||
|
||||
// Step 2: generate expert ids
|
||||
torch::Tensor gather_idx;
|
||||
torch::Tensor combine_idx;
|
||||
torch::Tensor token_count;
|
||||
std::optional<torch::Tensor> cusum_token_count;
|
||||
{
|
||||
xllm::kernel::MoeGenIdxParams moe_gen_idx_params;
|
||||
moe_gen_idx_params.expert_id = expert_id;
|
||||
moe_gen_idx_params.expert_num = num_total_experts_;
|
||||
std::vector<torch::Tensor> output_vec =
|
||||
xllm::kernel::moe_gen_idx(moe_gen_idx_params);
|
||||
gather_idx = output_vec[0];
|
||||
combine_idx = output_vec[1];
|
||||
token_count = output_vec[2];
|
||||
// during all2all communication, we do not need cusum_token_count in the
|
||||
// following computation
|
||||
if (enable_all2all_communication) {
|
||||
cusum_token_count = std::nullopt;
|
||||
} else {
|
||||
cusum_token_count = output_vec[3];
|
||||
}
|
||||
}
|
||||
|
||||
// Step 3: expand and quantize input if needed
|
||||
torch::Tensor expand_hidden_states;
|
||||
torch::Tensor hidden_states_scale;
|
||||
torch::Tensor token_count_slice;
|
||||
// all2all related variables
|
||||
torch::Tensor dispatch_send_token_tensor;
|
||||
// in all2all, the input is scattered, so there is no need to slice the token
|
||||
// count, and we can use the dispatch buffer directly
|
||||
if (enable_all2all_communication) {
|
||||
token_count_slice = token_count;
|
||||
int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
|
||||
int64_t dispatch_bytes =
|
||||
num_token_expand * deep_ep_params_.dispatch_token_size;
|
||||
dispatch_send_token_tensor =
|
||||
deep_ep_buffer_.dispatch_send_token_tensor.slice(0, 0, dispatch_bytes)
|
||||
.view({num_token_expand, deep_ep_params_.dispatch_token_size});
|
||||
} else {
|
||||
token_count_slice =
|
||||
token_count.slice(0, start_expert_id_, start_expert_id_ + expert_size);
|
||||
}
|
||||
|
||||
if (is_smoothquant_) {
|
||||
xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
|
||||
scaled_quantize_params.x = hidden_states_2d;
|
||||
// use dispatch_send_token_tensor buffer for input
|
||||
// to reduce memory footprint
|
||||
if (enable_all2all_communication) {
|
||||
scaled_quantize_params.smooth = input_smooth_;
|
||||
scaled_quantize_params.output =
|
||||
dispatch_send_token_tensor.slice(1, 0, hidden_size_);
|
||||
} else {
|
||||
scaled_quantize_params.smooth = input_smooth_.slice(
|
||||
0, start_expert_id_, start_expert_id_ + expert_size);
|
||||
scaled_quantize_params.gather_index_start_position =
|
||||
cusum_token_count.value().index({start_expert_id_}).unsqueeze(0);
|
||||
}
|
||||
scaled_quantize_params.token_count = token_count_slice;
|
||||
scaled_quantize_params.gather_index = gather_idx;
|
||||
scaled_quantize_params.act_mode = "none";
|
||||
scaled_quantize_params.active_coef = 1.0;
|
||||
scaled_quantize_params.is_gated = false;
|
||||
scaled_quantize_params.quant_type = torch::kChar;
|
||||
std::tie(expand_hidden_states, hidden_states_scale) =
|
||||
xllm::kernel::scaled_quantize(scaled_quantize_params);
|
||||
if (enable_all2all_communication) {
|
||||
// since view_as_dtype has not supported stride yet,
|
||||
// we need to copy the scale output to the dispatch buffer
|
||||
torch::Tensor dispatch_scale_slice =
|
||||
dispatch_send_token_tensor.slice(1, hidden_size_);
|
||||
torch::Tensor hidden_states_scale_bytes =
|
||||
view_as_dtype(hidden_states_scale, torch::kInt8)
|
||||
.view_as(dispatch_scale_slice);
|
||||
dispatch_scale_slice.copy_(hidden_states_scale_bytes);
|
||||
}
|
||||
} else {
|
||||
xllm::kernel::MoeExpandInputParams moe_expand_input_params;
|
||||
moe_expand_input_params.input = hidden_states_2d;
|
||||
moe_expand_input_params.gather_index = gather_idx;
|
||||
moe_expand_input_params.combine_idx = combine_idx;
|
||||
moe_expand_input_params.topk = topk_;
|
||||
expand_hidden_states =
|
||||
xllm::kernel::moe_expand_input(moe_expand_input_params);
|
||||
if (enable_all2all_communication) {
|
||||
// use copy to place the output inside the dispatch buffer
|
||||
torch::Tensor dispatch_tensor =
|
||||
view_as_dtype(expand_hidden_states, torch::kChar);
|
||||
dispatch_send_token_tensor.copy_(dispatch_tensor);
|
||||
}
|
||||
}
|
||||
|
||||
// collect the selected tensor
|
||||
selected_expert_info.reduce_weight = reduce_weight;
|
||||
selected_expert_info.combine_idx = combine_idx;
|
||||
selected_expert_info.token_count_slice = token_count_slice;
|
||||
selected_expert_info.cusum_token_count = cusum_token_count;
|
||||
if (is_smoothquant_) {
|
||||
selected_expert_info.input_scale = hidden_states_scale;
|
||||
}
|
||||
|
||||
return expand_hidden_states;
|
||||
}
|
||||
|
||||
torch::Tensor FusedMoEImpl::forward_experts(const torch::Tensor& hidden_states,
|
||||
const torch::Tensor& router_logits,
|
||||
bool enable_all2all_communication) {
|
||||
if (!stream_initialized_) {
|
||||
// update device record
|
||||
device_ = xllm::Device(hidden_states.device());
|
||||
|
||||
// acquire streams from the pool again
|
||||
routed_stream_ = device_.get_stream_from_pool();
|
||||
shared_stream_ = device_.get_stream_from_pool();
|
||||
stream_initialized_ = true;
|
||||
}
|
||||
|
||||
std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
|
||||
if (e_score_correction_bias_.defined()) {
|
||||
e_score_correction_bias = e_score_correction_bias_;
|
||||
}
|
||||
|
||||
// prepare the parameters for MoE computation
|
||||
torch::Tensor shared_expert_output;
|
||||
torch::IntArrayRef hidden_states_shape = hidden_states.sizes();
|
||||
torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType();
|
||||
torch::Tensor hidden_states_2d =
|
||||
hidden_states.reshape({-1, hidden_states.size(-1)});
|
||||
torch::Tensor router_logits_2d =
|
||||
router_logits.reshape({-1, router_logits.size(-1)});
|
||||
int64_t group_gemm_max_dim = enable_all2all_communication
|
||||
? deep_ep_params_.max_num_tokens_recv / topk_
|
||||
: hidden_states_2d.size(0);
|
||||
int64_t expert_size = w13_.size(0);
|
||||
|
||||
// Step 1-3: select experts
|
||||
SelectedExpertInfo selected_expert_info;
|
||||
torch::Tensor expand_hidden_states =
|
||||
select_experts(hidden_states_2d,
|
||||
router_logits_2d,
|
||||
selected_expert_info,
|
||||
enable_all2all_communication);
|
||||
|
||||
// Communciation Step 1: Dipatch
|
||||
// intermediate outputs that are used both in dispatch and combine
|
||||
torch::Tensor gather_by_rank_index;
|
||||
torch::Tensor token_sum;
|
||||
if (enable_all2all_communication) {
|
||||
int64_t dispatch_token_num = hidden_states_2d.size(0) * topk_;
|
||||
|
||||
// 1. Dispatch Step: Generate layout and send data
|
||||
deep_ep_->dispatch_step(dispatch_token_num,
|
||||
selected_expert_info.token_count_slice);
|
||||
|
||||
// 2. Process Result: Generate indices and unpack to computation buffer
|
||||
// use the buffer during initialization for the output
|
||||
expand_hidden_states = dispatch_recv_token_tensor_head_;
|
||||
std::optional<torch::Tensor> output_tail = std::nullopt;
|
||||
if (is_smoothquant_) {
|
||||
output_tail = dispatch_recv_token_tensor_tail_;
|
||||
// update selected_expert_info with the tail (input scale)
|
||||
selected_expert_info.input_scale = output_tail;
|
||||
}
|
||||
|
||||
DeepEPMetaResult deep_ep_meta = deep_ep_->process_dispatch_result(
|
||||
num_experts_per_rank_, expand_hidden_states, output_tail);
|
||||
|
||||
// Extract metadata for subsequent steps
|
||||
gather_by_rank_index = deep_ep_meta.gather_rank_index;
|
||||
selected_expert_info.token_count_slice = deep_ep_meta.token_count_slice;
|
||||
token_sum = deep_ep_meta.token_sum;
|
||||
}
|
||||
|
||||
// common gemm workspace for reduce memory footprint
|
||||
torch::Tensor gemm_workspace;
|
||||
|
||||
// Step 4: group gemm 1
|
||||
torch::Tensor gemm1_out =
|
||||
create_group_gemm_output(expand_hidden_states,
|
||||
w13_,
|
||||
selected_expert_info.token_count_slice,
|
||||
hidden_states_dtype,
|
||||
gemm_workspace);
|
||||
// ensure the lifespan of these parameters via brace
|
||||
{
|
||||
xllm::kernel::GroupGemmParams group_gemm_params;
|
||||
torch::ScalarType a_dtype =
|
||||
is_smoothquant_ ? torch::kInt8 : hidden_states_dtype;
|
||||
group_gemm_params.a =
|
||||
view_as_dtype(expand_hidden_states, a_dtype).view({-1, hidden_size_});
|
||||
group_gemm_params.b = w13_;
|
||||
group_gemm_params.token_count =
|
||||
selected_expert_info.token_count_slice.to("cpu");
|
||||
if (is_smoothquant_) {
|
||||
torch::Tensor a_scale =
|
||||
selected_expert_info.input_scale.value().flatten();
|
||||
selected_expert_info.input_scale =
|
||||
view_as_dtype(a_scale, torch::kFloat32);
|
||||
group_gemm_params.a_scale = selected_expert_info.input_scale;
|
||||
group_gemm_params.b_scale = w13_scale_;
|
||||
}
|
||||
group_gemm_params.max_dim = group_gemm_max_dim;
|
||||
group_gemm_params.trans_a = false;
|
||||
group_gemm_params.trans_b = true;
|
||||
group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
|
||||
group_gemm_params.output = gemm1_out;
|
||||
group_gemm_params.combine_idx = std::nullopt;
|
||||
gemm1_out = xllm::kernel::group_gemm(group_gemm_params);
|
||||
}
|
||||
|
||||
// Step 5: activation or scaled quantization(fused with activation)
|
||||
torch::Tensor act_out;
|
||||
torch::Tensor act_out_scale;
|
||||
if (is_smoothquant_) {
|
||||
int64_t slice_dim = gemm1_out.size(1);
|
||||
if (is_gated_) slice_dim /= 2;
|
||||
// slice operation is a view, does not take up extra memory, but points to
|
||||
// the same memory
|
||||
act_out = expand_hidden_states.slice(1, 0, slice_dim);
|
||||
act_out_scale =
|
||||
selected_expert_info.input_scale.value().slice(0, 0, gemm1_out.size(0));
|
||||
// call scaled quantization kernel (also fused with activation)
|
||||
xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
|
||||
scaled_quantize_params.x = gemm1_out;
|
||||
scaled_quantize_params.smooth = act_smooth_;
|
||||
scaled_quantize_params.token_count = selected_expert_info.token_count_slice;
|
||||
scaled_quantize_params.output = act_out;
|
||||
scaled_quantize_params.output_scale = act_out_scale;
|
||||
scaled_quantize_params.act_mode = hidden_act_;
|
||||
scaled_quantize_params.active_coef = 1.0;
|
||||
scaled_quantize_params.is_gated = is_gated_;
|
||||
scaled_quantize_params.quant_type = torch::kChar;
|
||||
std::tie(act_out, act_out_scale) =
|
||||
xllm::kernel::scaled_quantize(scaled_quantize_params);
|
||||
} else {
|
||||
act_out = is_gated_
|
||||
? gemm1_out.slice(1, 0, gemm1_out.size(1) / 2).contiguous()
|
||||
: gemm1_out;
|
||||
// call activation kernel
|
||||
xllm::kernel::ActivationParams activation_params;
|
||||
activation_params.input = gemm1_out;
|
||||
activation_params.output = act_out;
|
||||
activation_params.cusum_token_count =
|
||||
selected_expert_info.cusum_token_count;
|
||||
activation_params.act_mode = hidden_act_;
|
||||
activation_params.is_gated = is_gated_;
|
||||
activation_params.start_expert_id = start_expert_id_;
|
||||
activation_params.expert_size = expert_size;
|
||||
xllm::kernel::active(activation_params);
|
||||
}
|
||||
|
||||
// Step 6: group gemm 2
|
||||
torch::Tensor gemm2_out =
|
||||
create_group_gemm_output(act_out,
|
||||
w2_,
|
||||
selected_expert_info.token_count_slice,
|
||||
hidden_states_dtype,
|
||||
gemm_workspace);
|
||||
// ensure the lifespan of these parameters via brace
|
||||
{
|
||||
xllm::kernel::GroupGemmParams group_gemm_params;
|
||||
group_gemm_params.a = act_out;
|
||||
group_gemm_params.b = w2_;
|
||||
group_gemm_params.token_count =
|
||||
selected_expert_info.token_count_slice.to("cpu");
|
||||
if (is_smoothquant_) {
|
||||
group_gemm_params.a_scale = act_out_scale;
|
||||
group_gemm_params.b_scale = w2_scale_;
|
||||
}
|
||||
group_gemm_params.max_dim = group_gemm_max_dim;
|
||||
group_gemm_params.trans_a = false;
|
||||
group_gemm_params.trans_b = true;
|
||||
group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
|
||||
group_gemm_params.output = gemm2_out;
|
||||
group_gemm_params.combine_idx = selected_expert_info.combine_idx;
|
||||
gemm2_out = xllm::kernel::group_gemm(group_gemm_params);
|
||||
}
|
||||
|
||||
// Communciation Step 2: Combine
|
||||
if (enable_all2all_communication) {
|
||||
int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
|
||||
// Delegate pack, layout generation and combine to DeepEP
|
||||
torch::Tensor combine_send_layout =
|
||||
deep_ep_->combine_step_pack(gemm2_out,
|
||||
gather_by_rank_index,
|
||||
token_sum,
|
||||
hidden_size_,
|
||||
hidden_states_dtype);
|
||||
|
||||
// create a wait event for the current stream to finish computation
|
||||
auto current_stream = device_.current_stream();
|
||||
routed_stream_->wait_stream(*current_stream);
|
||||
// pure communciation kernel: dispatch
|
||||
{
|
||||
torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
|
||||
gemm2_out = deep_ep_->combine_step_comm(combine_send_layout,
|
||||
num_token_expand,
|
||||
hidden_size_,
|
||||
hidden_states_dtype);
|
||||
}
|
||||
|
||||
// pure computation kernel: shared experts
|
||||
if (n_shared_experts_ > 0) {
|
||||
shared_stream_->wait_stream(*current_stream);
|
||||
torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
|
||||
shared_expert_output = shared_experts_(hidden_states);
|
||||
}
|
||||
|
||||
// join for parallelization
|
||||
current_stream->wait_stream(*routed_stream_);
|
||||
if (n_shared_experts_ > 0) {
|
||||
current_stream->wait_stream(*shared_stream_);
|
||||
}
|
||||
}
|
||||
|
||||
// After group gemm is finished, some tensors are no
|
||||
// longer needed. We must explicitly release the memory.
|
||||
expand_hidden_states = torch::Tensor();
|
||||
selected_expert_info.input_scale = std::nullopt;
|
||||
act_out = torch::Tensor();
|
||||
|
||||
// Step 7: combine the intermediate results and get the final hidden states
|
||||
torch::Tensor final_hidden_states;
|
||||
// ensure the lifespan of these parameters via brace
|
||||
{
|
||||
xllm::kernel::MoeCombineResultParams moe_combine_result_params;
|
||||
moe_combine_result_params.input = gemm2_out;
|
||||
moe_combine_result_params.reduce_weight =
|
||||
selected_expert_info.reduce_weight;
|
||||
moe_combine_result_params.gather_ids = selected_expert_info.combine_idx;
|
||||
moe_combine_result_params.cusum_token_count =
|
||||
selected_expert_info.cusum_token_count;
|
||||
moe_combine_result_params.start_expert_id = start_expert_id_;
|
||||
moe_combine_result_params.expert_size = expert_size;
|
||||
moe_combine_result_params.bias = std::nullopt;
|
||||
// if all2all communication is enabled and shared output is provided,
|
||||
// we will fused the add up to combine result
|
||||
if (enable_all2all_communication && n_shared_experts_ > 0) {
|
||||
moe_combine_result_params.residual =
|
||||
shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
|
||||
}
|
||||
|
||||
final_hidden_states =
|
||||
xllm::kernel::moe_combine_result(moe_combine_result_params);
|
||||
}
|
||||
|
||||
// reshape the final hidden states to the original shape
|
||||
final_hidden_states = final_hidden_states.reshape(hidden_states_shape);
|
||||
|
||||
if (enable_all2all_communication) {
|
||||
return final_hidden_states;
|
||||
}
|
||||
|
||||
// Communciation Step 3: AllReduce for non-all2all communication
|
||||
// shared experts can be parallelized with the final communication step
|
||||
// during moe computation.
|
||||
auto current_stream = device_.current_stream();
|
||||
routed_stream_->wait_stream(*current_stream);
|
||||
{
|
||||
torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
|
||||
if (tp_pg_->world_size() > 1) {
|
||||
final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
|
||||
}
|
||||
if (parallel_args_.ep_size() > 1) {
|
||||
final_hidden_states = parallel_state::reduce(
|
||||
final_hidden_states, parallel_args_.moe_ep_group_);
|
||||
}
|
||||
}
|
||||
|
||||
if (n_shared_experts_ > 0) {
|
||||
shared_stream_->wait_stream(*current_stream);
|
||||
torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
|
||||
// for non all2all, we compute the shared experts parallelized with the
|
||||
// final communication step
|
||||
shared_expert_output = shared_experts_(hidden_states);
|
||||
shared_expert_output =
|
||||
shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
|
||||
}
|
||||
|
||||
// join for parallelization
|
||||
current_stream->wait_stream(*routed_stream_);
|
||||
if (n_shared_experts_ > 0) {
|
||||
current_stream->wait_stream(*shared_stream_);
|
||||
final_hidden_states += shared_expert_output;
|
||||
}
|
||||
|
||||
return final_hidden_states;
|
||||
}
|
||||
|
||||
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(),
|
||||
[](int32_t val) { return val == 1; });
|
||||
|
||||
bool is_dp_ep_parallel =
|
||||
parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1;
|
||||
// during all2all communication, the output has been
|
||||
// gathered and sliced by dispatch and combine steps,
|
||||
// so we do not need to gather input and slice output again
|
||||
bool need_gather_and_slice =
|
||||
is_dp_ep_parallel && !enable_all2all_communication;
|
||||
|
||||
auto input = hidden_states;
|
||||
if (need_gather_and_slice) {
|
||||
input = parallel_state::gather(input,
|
||||
parallel_args_.dp_local_process_group_,
|
||||
input_params.dp_global_token_nums);
|
||||
}
|
||||
// MoE Gate
|
||||
auto router_logits = gate_(input);
|
||||
|
||||
// MoE Experts
|
||||
auto output =
|
||||
forward_experts(input, router_logits, enable_all2all_communication);
|
||||
|
||||
if (need_gather_and_slice) {
|
||||
output = get_dp_local_slice(output, input_params, parallel_args_);
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) {
|
||||
if (e_score_correction_bias_.defined() &&
|
||||
!e_score_correction_bias_is_loaded_) {
|
||||
LOAD_WEIGHT(e_score_correction_bias);
|
||||
}
|
||||
}
|
||||
|
||||
void FusedMoEImpl::load_experts(const StateDict& state_dict) {
|
||||
const int64_t rank = tp_pg_->rank();
|
||||
const int64_t world_size = tp_pg_->world_size();
|
||||
const int64_t start_expert_id = start_expert_id_;
|
||||
const int64_t num_experts_per_rank = num_experts_per_rank_;
|
||||
const int64_t num_total_experts = num_total_experts_;
|
||||
std::vector<std::string> prefixes = {"gate_proj.", "up_proj."};
|
||||
if (is_smoothquant_) {
|
||||
LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13);
|
||||
LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale);
|
||||
// When supporting DeepEP All2All mode,
|
||||
// we need to load the complete set of expert weights corresponding to
|
||||
// "up_proj.smooth". Note that even if deep EP mode is not enabled, it
|
||||
// remains possible to retrieve the smooth quantization information for a
|
||||
// subset of experts. Therefore, we intentionally do not check whether
|
||||
// deep_ep_ is enabled in this case.
|
||||
LOAD_MOE_ALL_EXPERT_WEIGHT("up_proj.", "smooth", input_smooth, -1);
|
||||
LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1);
|
||||
LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1);
|
||||
LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0);
|
||||
} else {
|
||||
LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13);
|
||||
LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1);
|
||||
}
|
||||
}
|
||||
|
||||
void FusedMoEImpl::load_state_dict(const StateDict& state_dict) {
|
||||
if (state_dict.size() == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (n_shared_experts_ > 0) {
|
||||
shared_experts_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("shared_experts."));
|
||||
}
|
||||
gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
|
||||
load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate."));
|
||||
load_experts(state_dict.get_dict_with_prefix("experts."));
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
131
ex_engine/xllm_layers/ilu/fused_moe.h
Normal file
131
ex_engine/xllm_layers/ilu/fused_moe.h
Normal file
@@ -0,0 +1,131 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include "framework/model/model_args.h"
|
||||
#include "framework/model/model_input_params.h"
|
||||
#include "framework/parallel_state/parallel_args.h"
|
||||
#include "framework/quant_args.h"
|
||||
#include "framework/state_dict/state_dict.h"
|
||||
#include "framework/state_dict/utils.h"
|
||||
#include "layers/common/deep_ep.h"
|
||||
#include "layers/common/dense_mlp.h"
|
||||
#include "layers/common/fused_moe_base.h"
|
||||
#include "layers/common/linear.h"
|
||||
#include "platform/device.h"
|
||||
#include "util/tensor_helper.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class FusedMoEImpl : public torch::nn::Module {
|
||||
public:
|
||||
FusedMoEImpl() = default;
|
||||
FusedMoEImpl(const ModelArgs& model_args,
|
||||
const FusedMoEArgs& moe_args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
torch::Tensor forward_experts(const torch::Tensor& hidden_states,
|
||||
const torch::Tensor& router_logits,
|
||||
bool enable_all2all_communication);
|
||||
torch::Tensor forward(const torch::Tensor& hidden_states,
|
||||
const ModelInputParams& input_params);
|
||||
void load_state_dict(const StateDict& state_dict);
|
||||
|
||||
private:
|
||||
// struct to store the selected expert info
|
||||
struct SelectedExpertInfo {
|
||||
torch::Tensor reduce_weight;
|
||||
torch::Tensor combine_idx;
|
||||
torch::Tensor token_count_slice;
|
||||
std::optional<torch::Tensor> cusum_token_count;
|
||||
std::optional<torch::Tensor> input_scale;
|
||||
};
|
||||
|
||||
// initial steps for MoE computation, select the experts for each token
|
||||
torch::Tensor select_experts(const torch::Tensor& hidden_states_2d,
|
||||
const torch::Tensor& router_logits_2d,
|
||||
SelectedExpertInfo& selected_expert_info,
|
||||
bool enable_all2all_communication);
|
||||
|
||||
private:
|
||||
int64_t num_total_experts_;
|
||||
int64_t topk_;
|
||||
int64_t num_expert_group_;
|
||||
int64_t topk_group_;
|
||||
double route_scale_;
|
||||
int64_t hidden_size_;
|
||||
int64_t n_shared_experts_;
|
||||
bool is_gated_;
|
||||
int64_t renormalize_;
|
||||
std::string hidden_act_;
|
||||
std::string scoring_func_;
|
||||
bool is_smoothquant_;
|
||||
|
||||
int64_t num_experts_per_rank_;
|
||||
int64_t start_expert_id_;
|
||||
|
||||
// Deep EP related parameters
|
||||
bool enable_deep_ep_;
|
||||
DeepEPBuffer deep_ep_buffer_;
|
||||
DeepEPParams deep_ep_params_;
|
||||
torch::Tensor dispatch_recv_token_tensor_head_;
|
||||
torch::Tensor dispatch_recv_token_tensor_tail_;
|
||||
|
||||
// steams for parallel shared experts
|
||||
std::unique_ptr<Stream> shared_stream_;
|
||||
std::unique_ptr<Stream> routed_stream_;
|
||||
xllm::Device device_;
|
||||
bool stream_initialized_ = false;
|
||||
|
||||
ReplicatedLinear gate_{nullptr};
|
||||
DenseMLP shared_experts_{nullptr};
|
||||
DeepEP deep_ep_{nullptr};
|
||||
|
||||
QuantArgs quant_args_;
|
||||
ParallelArgs parallel_args_;
|
||||
torch::TensorOptions options_;
|
||||
ProcessGroup* tp_pg_;
|
||||
|
||||
DEFINE_WEIGHT(w13);
|
||||
DEFINE_FUSED_WEIGHT(w1);
|
||||
DEFINE_FUSED_WEIGHT(w3);
|
||||
DEFINE_FUSED_WEIGHT(w2);
|
||||
DEFINE_WEIGHT(e_score_correction_bias);
|
||||
DEFINE_WEIGHT(w13_scale);
|
||||
DEFINE_FUSED_WEIGHT(w1_scale);
|
||||
DEFINE_FUSED_WEIGHT(w3_scale);
|
||||
DEFINE_FUSED_WEIGHT(w2_scale);
|
||||
DEFINE_FUSED_WEIGHT(input_smooth);
|
||||
DEFINE_FUSED_WEIGHT(act_smooth);
|
||||
|
||||
void load_e_score_correction_bias(const StateDict& state_dict);
|
||||
void load_experts(const StateDict& state_dict);
|
||||
// create the group gemm output tensor with the workspace
|
||||
torch::Tensor create_group_gemm_output(const torch::Tensor& a,
|
||||
const torch::Tensor& b,
|
||||
const torch::Tensor& group_list,
|
||||
torch::ScalarType dtype,
|
||||
torch::Tensor& workspace);
|
||||
};
|
||||
TORCH_MODULE(FusedMoE);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
@@ -1,4 +1,4 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
@@ -123,53 +123,19 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations(
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor>
|
||||
Qwen3_5GatedDeltaNetImpl::project_decode_inputs(
|
||||
const torch::Tensor& hidden_states) {
|
||||
const auto reshape_projection = [](const torch::Tensor& projection) {
|
||||
return projection.view({projection.size(0), -1, projection.size(-1)});
|
||||
};
|
||||
auto qkv = reshape_projection(in_proj_qkv_->forward(hidden_states));
|
||||
auto z_proj = reshape_projection(in_proj_z_->forward(hidden_states));
|
||||
auto b_proj = reshape_projection(in_proj_b_->forward(hidden_states));
|
||||
auto a_proj = reshape_projection(in_proj_a_->forward(hidden_states));
|
||||
return {merge_qkvz_from_split_activations(qkv, z_proj),
|
||||
merge_ba_from_split_activations(b_proj, a_proj)};
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor>
|
||||
Qwen3_5GatedDeltaNetImpl::project_flat_inputs(
|
||||
const torch::Tensor& hidden_states) {
|
||||
auto qkv = in_proj_qkv_->forward(hidden_states).unsqueeze(0);
|
||||
auto z_proj = in_proj_z_->forward(hidden_states).unsqueeze(0);
|
||||
auto b_proj = in_proj_b_->forward(hidden_states).unsqueeze(0);
|
||||
auto a_proj = in_proj_a_->forward(hidden_states).unsqueeze(0);
|
||||
auto qkvz = merge_qkvz_from_split_activations(qkv, z_proj);
|
||||
auto ba = merge_ba_from_split_activations(b_proj, a_proj);
|
||||
return {qkvz.view({hidden_states.size(0), qkvz.size(-1)}).contiguous(),
|
||||
ba.view({hidden_states.size(0), ba.size(-1)}).contiguous()};
|
||||
}
|
||||
|
||||
std::optional<
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
|
||||
Qwen3_5GatedDeltaNetImpl::project_split_inputs(
|
||||
Qwen3_5GatedDeltaNetImpl::project_padded_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
auto qkv = reshape_projected_tokens_with_pad(
|
||||
attn_metadata, in_proj_qkv_->forward(hidden_states));
|
||||
auto z_proj = reshape_projected_tokens_with_pad(
|
||||
attn_metadata, in_proj_z_->forward(hidden_states));
|
||||
auto b_proj = reshape_projected_tokens_with_pad(
|
||||
attn_metadata, in_proj_b_->forward(hidden_states));
|
||||
auto a_proj = reshape_projected_tokens_with_pad(
|
||||
attn_metadata, in_proj_a_->forward(hidden_states));
|
||||
|
||||
const int64_t batch_size = qkv.size(0);
|
||||
const int64_t seq_len = qkv.size(1);
|
||||
auto z =
|
||||
z_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
|
||||
auto b = b_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
auto a = a_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
return std::make_tuple(qkv, z, b, a);
|
||||
auto qkv = reshape_qkvz_with_pad(attn_metadata,
|
||||
in_proj_qkv_->forward(hidden_states));
|
||||
auto z_proj =
|
||||
reshape_qkvz_with_pad(attn_metadata, in_proj_z_->forward(hidden_states));
|
||||
auto b_proj =
|
||||
reshape_qkvz_with_pad(attn_metadata, in_proj_b_->forward(hidden_states));
|
||||
auto a_proj =
|
||||
reshape_qkvz_with_pad(attn_metadata, in_proj_a_->forward(hidden_states));
|
||||
return {merge_qkvz_from_split_activations(qkv, z_proj),
|
||||
merge_ba_from_split_activations(b_proj, a_proj)};
|
||||
}
|
||||
|
||||
void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict(
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
@@ -17,9 +17,7 @@ limitations under the License.
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <tuple>
|
||||
#include <utility>
|
||||
|
||||
#include "qwen3_next_gated_delta_net.h"
|
||||
@@ -36,15 +34,9 @@ class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl {
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
protected:
|
||||
std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
|
||||
const torch::Tensor& hidden_states) override;
|
||||
std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
|
||||
const torch::Tensor& hidden_states) override;
|
||||
std::optional<
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
|
||||
project_split_inputs(const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) override;
|
||||
bool use_fla_ssm_state_layout() const override { return true; }
|
||||
std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) override;
|
||||
|
||||
void load_projection_state_dict(const StateDict& state_dict) override;
|
||||
void verify_projection_weights(const std::string& prefix) const override;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
@@ -15,12 +15,9 @@ limitations under the License.
|
||||
#include <glog/logging.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <optional>
|
||||
#include <tuple>
|
||||
|
||||
#include "xllm/core/kernels/npu/npu_ops_api.h"
|
||||
#include "xllm/core/kernels/ops_api.h"
|
||||
#include "xllm/core/platform/npu/acl_graph_task_update_context.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
@@ -31,31 +28,6 @@ torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) {
|
||||
return x / norm;
|
||||
}
|
||||
|
||||
torch::Tensor repeat_tensor_heads(const torch::Tensor& tensor,
|
||||
int64_t target_heads,
|
||||
int64_t head_dim) {
|
||||
const int64_t current_heads = tensor.size(head_dim);
|
||||
if (current_heads == target_heads) {
|
||||
return tensor;
|
||||
}
|
||||
CHECK_GT(current_heads, 0) << "current heads must be positive";
|
||||
CHECK_EQ(target_heads % current_heads, 0)
|
||||
<< "target heads must be divisible by current heads, target_heads="
|
||||
<< target_heads << ", current_heads=" << current_heads;
|
||||
|
||||
const int64_t repeats = target_heads / current_heads;
|
||||
std::vector<int64_t> view_shape = tensor.sizes().vec();
|
||||
view_shape.insert(view_shape.begin() + head_dim + 1, 1);
|
||||
std::vector<int64_t> expand_shape = view_shape;
|
||||
expand_shape[head_dim + 1] = repeats;
|
||||
std::vector<int64_t> output_shape = tensor.sizes().vec();
|
||||
output_shape[head_dim] = target_heads;
|
||||
return tensor.unsqueeze(head_dim + 1)
|
||||
.expand(expand_shape)
|
||||
.reshape(output_shape)
|
||||
.contiguous();
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
|
||||
torch::Tensor query,
|
||||
torch::Tensor key,
|
||||
@@ -80,9 +52,6 @@ std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
|
||||
value = to_float32_and_transpose(value);
|
||||
beta = to_float32_and_transpose(beta);
|
||||
g = to_float32_and_transpose(g);
|
||||
const int64_t value_num_heads = value.size(1);
|
||||
query = repeat_tensor_heads(query, value_num_heads, 1);
|
||||
key = repeat_tensor_heads(key, value_num_heads, 1);
|
||||
|
||||
int64_t batch_size = key.size(0);
|
||||
int64_t num_heads = key.size(1);
|
||||
@@ -150,15 +119,12 @@ std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
|
||||
value = to_float32(value);
|
||||
beta = to_float32(beta);
|
||||
g = to_float32(g);
|
||||
const int64_t value_num_heads = value.size(1);
|
||||
query = repeat_tensor_heads(query, value_num_heads, 1);
|
||||
key = repeat_tensor_heads(key, value_num_heads, 1);
|
||||
|
||||
int64_t batch_size = query.size(0);
|
||||
int64_t num_heads = query.size(1);
|
||||
int64_t sequence_length = query.size(2);
|
||||
int64_t k_head_dim = key.size(-1);
|
||||
int64_t v_head_dim = value.size(-1);
|
||||
auto batch_size = query.size(0);
|
||||
auto num_heads = query.size(1);
|
||||
auto sequence_length = query.size(2);
|
||||
auto k_head_dim = key.size(-1);
|
||||
auto v_head_dim = value.size(-1);
|
||||
|
||||
int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size;
|
||||
query = torch::nn::functional::pad(
|
||||
@@ -276,164 +242,6 @@ std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
|
||||
core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
|
||||
return std::make_tuple(core_attn_out, last_recurrent_state);
|
||||
}
|
||||
|
||||
int64_t get_checkpoint_stride(const torch::Tensor& conv_cache,
|
||||
const torch::Tensor& ssm_cache) {
|
||||
if (!conv_cache.defined() || !ssm_cache.defined() ||
|
||||
conv_cache.numel() == 0 || ssm_cache.numel() == 0) {
|
||||
return 1;
|
||||
}
|
||||
CHECK_GT(conv_cache.size(0), 0) << "conv cache must have positive batch dim";
|
||||
CHECK_EQ(ssm_cache.size(0) % conv_cache.size(0), 0)
|
||||
<< "ssm cache checkpoint layout mismatch, ssm_rows=" << ssm_cache.size(0)
|
||||
<< ", conv_rows=" << conv_cache.size(0);
|
||||
return ssm_cache.size(0) / conv_cache.size(0);
|
||||
}
|
||||
|
||||
torch::Tensor build_linear_state_base_indices(
|
||||
const torch::Tensor& logical_state_indices,
|
||||
int64_t checkpoint_stride) {
|
||||
if (checkpoint_stride == 1) {
|
||||
return logical_state_indices;
|
||||
}
|
||||
return logical_state_indices * checkpoint_stride;
|
||||
}
|
||||
|
||||
torch::Tensor expand_sequence_tensor_to_batch(const torch::Tensor& tensor,
|
||||
int64_t target_batch,
|
||||
const char* tensor_name) {
|
||||
CHECK(tensor.defined()) << tensor_name << " must be defined";
|
||||
CHECK_EQ(tensor.dim(), 1) << tensor_name << " must be a 1D tensor.";
|
||||
const int64_t source_batch = tensor.size(0);
|
||||
if (source_batch == target_batch) {
|
||||
return tensor.contiguous();
|
||||
}
|
||||
CHECK_GT(source_batch, 0) << tensor_name << " must not be empty.";
|
||||
CHECK_EQ(target_batch % source_batch, 0)
|
||||
<< tensor_name << " cannot be expanded from " << source_batch << " to "
|
||||
<< target_batch;
|
||||
const int64_t repeat_count = target_batch / source_batch;
|
||||
return tensor.unsqueeze(1)
|
||||
.expand({source_batch, repeat_count})
|
||||
.reshape({target_batch})
|
||||
.contiguous();
|
||||
}
|
||||
|
||||
torch::Tensor run_causal_conv1d_graph_update(
|
||||
const std::shared_ptr<xllm::npu::AclGraphTaskUpdateContext>& graph_context,
|
||||
const torch::Tensor& x,
|
||||
const torch::Tensor& weight,
|
||||
const torch::Tensor& conv_state,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const std::vector<int64_t>& query_start_loc,
|
||||
const std::vector<int64_t>& cache_indices,
|
||||
const std::vector<int64_t>& num_accepted_tokens,
|
||||
xllm::npu::CausalConv1dGraphBranch branch) {
|
||||
CHECK(graph_context != nullptr && graph_context->capturing)
|
||||
<< "causal_conv1d graph update can only be registered during capture";
|
||||
|
||||
c10_npu::NPUStream stream = c10_npu::getCurrentNPUStream();
|
||||
auto event = std::make_shared<c10_npu::NPUEvent>(ACL_EVENT_EXTERNAL);
|
||||
event->block(stream);
|
||||
event->reset(stream);
|
||||
|
||||
torch::Tensor output;
|
||||
c10_npu::graph_task_group_begin(stream);
|
||||
const std::vector<int64_t> empty_host_args;
|
||||
CHECK(!query_start_loc.empty())
|
||||
<< "query_start_loc must be populated for causal_conv1d graph update";
|
||||
CHECK_EQ(query_start_loc.back(), x.size(0))
|
||||
<< "query_start_loc must be padded to x.shape[0] during graph capture";
|
||||
CHECK_EQ(cache_indices.size() + 1, query_start_loc.size())
|
||||
<< "cache_indices must be sequence-scoped";
|
||||
if (branch == xllm::npu::CausalConv1dGraphBranch::kSpecVerify) {
|
||||
CHECK_EQ(num_accepted_tokens.size(), cache_indices.size())
|
||||
<< "num_accepted_tokens must be sequence-scoped for spec verify";
|
||||
}
|
||||
|
||||
output = torch::empty_like(x);
|
||||
xllm::kernel::causal_conv1d_out(output,
|
||||
x,
|
||||
weight,
|
||||
conv_state,
|
||||
bias,
|
||||
torch::IntArrayRef(query_start_loc),
|
||||
torch::IntArrayRef(cache_indices),
|
||||
torch::IntArrayRef(empty_host_args),
|
||||
torch::IntArrayRef(num_accepted_tokens),
|
||||
xllm::npu::kCausalConv1dActivationSilu,
|
||||
xllm::npu::kCausalConv1dGraphPadSlotId,
|
||||
xllm::npu::kCausalConv1dRunModeUpdate);
|
||||
c10_npu::NPUTaskGroupHandle handle = c10_npu::graph_task_group_end(stream);
|
||||
|
||||
xllm::npu::CausalConv1dGraphTask task;
|
||||
task.output = output;
|
||||
task.x = x;
|
||||
task.weight = weight;
|
||||
task.conv_state = conv_state;
|
||||
task.bias = bias;
|
||||
task.activation_mode = xllm::npu::kCausalConv1dActivationSilu;
|
||||
task.pad_slot_id = xllm::npu::kCausalConv1dGraphPadSlotId;
|
||||
task.run_mode = xllm::npu::kCausalConv1dRunModeUpdate;
|
||||
task.branch = branch;
|
||||
task.handle = handle;
|
||||
task.event = std::move(event);
|
||||
graph_context->causal_conv1d_tasks.emplace_back(std::move(task));
|
||||
return output;
|
||||
}
|
||||
|
||||
torch::Tensor run_spec_verify_gated_delta_rule(
|
||||
torch::Tensor query,
|
||||
torch::Tensor key,
|
||||
torch::Tensor value,
|
||||
torch::Tensor g,
|
||||
torch::Tensor beta,
|
||||
torch::Tensor& ssm_cache,
|
||||
const torch::Tensor& checkpoint_indices,
|
||||
const torch::Tensor& num_accepted_tokens,
|
||||
const torch::Tensor& cu_seq_lens,
|
||||
const std::vector<int32_t>& q_seq_lens_vec,
|
||||
double scale) {
|
||||
const auto device = value.device();
|
||||
const int64_t batch_size = value.size(0);
|
||||
const int64_t seq_len = value.size(1);
|
||||
const int64_t total_seq_len = batch_size * seq_len;
|
||||
CHECK_EQ(cu_seq_lens.numel(), batch_size + 1)
|
||||
<< "GDN spec verify cu_seq_lens must be cumulative.";
|
||||
CHECK_EQ(q_seq_lens_vec.size(), static_cast<size_t>(batch_size))
|
||||
<< "GDN spec verify q_seq_lens_vec must be per sequence.";
|
||||
for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
|
||||
CHECK_EQ(q_seq_lens_vec[batch_idx], seq_len)
|
||||
<< "Qwen3.5 spec verify fused recurrent path expects dense "
|
||||
"same-length validate tokens.";
|
||||
}
|
||||
|
||||
xllm::kernel::FusedRecurrentGatedDeltaRuleParams params;
|
||||
params.q = query.reshape({1, total_seq_len, query.size(-2), query.size(-1)})
|
||||
.contiguous();
|
||||
params.k =
|
||||
key.reshape({1, total_seq_len, key.size(-2), key.size(-1)}).contiguous();
|
||||
params.v = value.reshape({1, total_seq_len, value.size(-2), value.size(-1)})
|
||||
.contiguous();
|
||||
params.g = g.to(torch::kFloat32)
|
||||
.reshape({1, total_seq_len, g.size(-1)})
|
||||
.contiguous();
|
||||
params.beta = beta.reshape({1, total_seq_len, beta.size(-1)}).contiguous();
|
||||
params.scale = static_cast<float>(scale);
|
||||
params.initial_state = ssm_cache;
|
||||
params.inplace_final_state = true;
|
||||
params.cu_seqlens = cu_seq_lens.to(torch::kLong).contiguous();
|
||||
params.ssm_state_indices = checkpoint_indices.contiguous();
|
||||
params.num_accepted_tokens =
|
||||
num_accepted_tokens.to(device, torch::kInt32).contiguous();
|
||||
params.use_qk_l2norm_in_kernel = true;
|
||||
|
||||
auto output_and_state =
|
||||
xllm::kernel::fused_recurrent_gated_delta_rule(params);
|
||||
return output_and_state.first.view(
|
||||
{batch_size, seq_len, value.size(-2), value.size(-1)});
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
Qwen3GatedDeltaNetBaseImpl::Qwen3GatedDeltaNetBaseImpl(
|
||||
@@ -495,11 +303,7 @@ void Qwen3GatedDeltaNetBaseImpl::load_common_state_dict(
|
||||
|
||||
if (auto w = state_dict.get_tensor("conv1d.weight"); w.defined()) {
|
||||
conv1d_->load_state_dict(
|
||||
StateDict({{"weight", w.squeeze(1)}},
|
||||
static_cast<std::string>(state_dict.prefix()) + "conv1d."),
|
||||
shard_tensor_count,
|
||||
shard_sizes);
|
||||
conv1d_->weight().set_(conv1d_->weight().transpose(0, 1).contiguous());
|
||||
StateDict({{"weight", w.squeeze(1)}}), shard_tensor_count, shard_sizes);
|
||||
}
|
||||
o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj."));
|
||||
if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) {
|
||||
@@ -518,279 +322,87 @@ void Qwen3GatedDeltaNetBaseImpl::verify_common_loaded_weights(
|
||||
<< prefix << "A_log";
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor>
|
||||
Qwen3GatedDeltaNetBaseImpl::project_padded_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
if (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill) {
|
||||
auto [qkvz_flat, ba_flat] = project_flat_inputs(hidden_states);
|
||||
return {reshape_projected_tokens_with_pad(attn_metadata, qkvz_flat),
|
||||
reshape_projected_tokens_with_pad(attn_metadata, ba_flat)};
|
||||
}
|
||||
return project_decode_inputs(hidden_states);
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata,
|
||||
KVCache& kv_cache,
|
||||
const ModelInputParams& input_params) {
|
||||
// Early-return on dummy shards. Under dp>1, an empty shard is padded with a
|
||||
// fake token by worker_impl but its GDN state tensors (kv_cache_tokens_nums,
|
||||
// linear_state_ids etc.) are left undefined. This mirrors the is_dummy
|
||||
// early-return in Attention::forward (npu_torch/attention.cpp). Uses
|
||||
// zeros_like rather than empty_like so downstream post-norm / mlp do not
|
||||
// read uninitialized data. Placed before FlashComm1 sequence gather so
|
||||
// dummy shards do not enter the collective and waste bandwidth.
|
||||
if (attn_metadata.is_dummy) {
|
||||
return torch::zeros_like(hidden_states);
|
||||
}
|
||||
const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context();
|
||||
torch::Tensor h = hidden_states;
|
||||
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
|
||||
h = gather_sequence(hidden_states, *fc1_ctx);
|
||||
}
|
||||
auto [qkvz_padded, ba_padded] =
|
||||
project_padded_inputs(hidden_states, attn_metadata);
|
||||
int64_t batch_size = qkvz_padded.size(0);
|
||||
int64_t seq_len = qkvz_padded.size(1);
|
||||
|
||||
torch::Tensor qkvz_flat =
|
||||
qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)});
|
||||
torch::Tensor ba_flat =
|
||||
ba_padded.view({batch_size * seq_len, ba_padded.size(-1)});
|
||||
xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params;
|
||||
fused_params.mixed_qkvz = qkvz_flat;
|
||||
fused_params.mixed_ba = ba_flat;
|
||||
fused_params.num_heads_qk = static_cast<int32_t>(num_k_heads_ / tp_size_);
|
||||
fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
|
||||
fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
|
||||
fused_params.head_v = static_cast<int32_t>(head_v_dim_);
|
||||
|
||||
// Save the gathered hidden-state size for potential padding later.
|
||||
const int64_t original_num_tokens = h.size(0);
|
||||
const bool use_spec_verify = input_params.is_spec_verify;
|
||||
const bool is_any_prefill =
|
||||
attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
|
||||
torch::Tensor mixed_qkv, z, b, a;
|
||||
torch::Tensor processed_q, processed_k, processed_v;
|
||||
int64_t batch_size = 0;
|
||||
int64_t seq_len = 0;
|
||||
std::tie(mixed_qkv, z, b, a) =
|
||||
xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
|
||||
|
||||
// Qwen3.5 stores qkv, z, b, and a as separate projection weights, so it can
|
||||
// use their outputs directly in every forward mode. Qwen3Next stores qkvz
|
||||
// and ba as packed weights and uses the fused-split fallback below.
|
||||
auto split_inputs = project_split_inputs(h, attn_metadata);
|
||||
if (split_inputs.has_value()) {
|
||||
std::tie(mixed_qkv, z, b, a) = split_inputs.value();
|
||||
batch_size = mixed_qkv.size(0);
|
||||
seq_len = mixed_qkv.size(1);
|
||||
} else {
|
||||
auto [qkvz_padded, ba_padded] = project_padded_inputs(h, attn_metadata);
|
||||
batch_size = qkvz_padded.size(0);
|
||||
seq_len = qkvz_padded.size(1);
|
||||
|
||||
torch::Tensor qkvz_flat =
|
||||
qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)});
|
||||
torch::Tensor ba_flat =
|
||||
ba_padded.view({batch_size * seq_len, ba_padded.size(-1)});
|
||||
xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params;
|
||||
fused_params.mixed_qkvz = qkvz_flat;
|
||||
fused_params.mixed_ba = ba_flat;
|
||||
fused_params.num_heads_qk = static_cast<int32_t>(num_k_heads_ / tp_size_);
|
||||
fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
|
||||
fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
|
||||
fused_params.head_v = static_cast<int32_t>(head_v_dim_);
|
||||
|
||||
std::tie(mixed_qkv, z, b, a) =
|
||||
xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
|
||||
|
||||
mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)});
|
||||
z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
|
||||
b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
}
|
||||
|
||||
const bool fla_ssm_state_layout = use_fla_ssm_state_layout();
|
||||
const int64_t local_q_heads = num_k_heads_ / tp_size_;
|
||||
const int64_t local_v_heads = num_v_heads_ / tp_size_;
|
||||
const int64_t local_conv_dim =
|
||||
2 * local_q_heads * head_k_dim_ + local_v_heads * head_v_dim_;
|
||||
bool used_direct_prefill_qkv = false;
|
||||
mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)});
|
||||
z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
|
||||
b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
|
||||
torch::Tensor conv_cache = kv_cache.get_conv_cache();
|
||||
torch::Tensor ssm_cache = kv_cache.get_ssm_cache();
|
||||
torch::Device device = mixed_qkv.device();
|
||||
torch::Tensor conv_weight = conv1d_->weight();
|
||||
torch::Tensor logical_state_indices =
|
||||
get_linear_state_indices(input_params, device);
|
||||
const int64_t checkpoint_stride =
|
||||
get_checkpoint_stride(conv_cache, ssm_cache);
|
||||
torch::Tensor linear_state_base_indices =
|
||||
build_linear_state_base_indices(logical_state_indices, checkpoint_stride);
|
||||
auto graph_context = input_params.graph.acl_graph_task_update_context;
|
||||
const bool register_conv1d_graph_update =
|
||||
graph_context != nullptr && graph_context->capturing;
|
||||
torch::Tensor g, beta, core_attn_out, last_recurrent_state;
|
||||
auto device = mixed_qkv.device();
|
||||
auto conv_weight = conv1d_->weight();
|
||||
auto linear_state_indices = get_linear_state_indices(input_params, device);
|
||||
|
||||
if (!use_spec_verify && is_any_prefill) {
|
||||
torch::IntArrayRef num_accepted_tokens_opt;
|
||||
std::vector<int64_t> linear_state_indices_vec(
|
||||
input_params.embedding.linear_state_ids.begin(),
|
||||
input_params.embedding.linear_state_ids.end());
|
||||
torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv);
|
||||
if (attn_metadata.is_prefill) {
|
||||
mixed_qkv = mixed_qkv.transpose(1, 2);
|
||||
torch::Tensor conv_state =
|
||||
(seq_len < conv_kernel_size_ - 1)
|
||||
? torch::pad(mixed_qkv, {0, conv_kernel_size_ - 1 - seq_len})
|
||||
: (seq_len > conv_kernel_size_ - 1)
|
||||
? mixed_qkv.narrow(
|
||||
-1, seq_len - conv_kernel_size_ + 1, conv_kernel_size_ - 1)
|
||||
: mixed_qkv;
|
||||
conv_state = conv_state.transpose(1, 2).contiguous();
|
||||
conv_cache.index_put_({linear_state_indices},
|
||||
conv_state.to(conv_cache.dtype()));
|
||||
torch::Tensor bias;
|
||||
auto conv_output =
|
||||
torch::conv1d(mixed_qkv,
|
||||
conv_weight.unsqueeze(1).to(device),
|
||||
bias,
|
||||
/*stride=*/std::vector<int64_t>{1},
|
||||
/*padding=*/std::vector<int64_t>{3},
|
||||
/*dilation=*/std::vector<int64_t>{1},
|
||||
/*groups=*/static_cast<int64_t>(mixed_qkv.size(1)));
|
||||
mixed_qkv = torch::silu(conv_output.slice(2, 0, seq_len));
|
||||
|
||||
const bool direct_qkv_model_supported =
|
||||
fla_ssm_state_layout && num_k_heads_ % tp_size_ == 0 &&
|
||||
num_v_heads_ % tp_size_ == 0 && local_q_heads > 0 &&
|
||||
local_v_heads > 0 && head_k_dim_ == 128 && head_v_dim_ == 128;
|
||||
const bool direct_qkv_metadata_available =
|
||||
attn_metadata.q_seq_lens_vec.size() ==
|
||||
static_cast<size_t>(batch_size) &&
|
||||
input_params.parallel.query_start_loc.size() ==
|
||||
static_cast<size_t>(batch_size + 1) &&
|
||||
input_params.embedding.linear_state_ids.size() ==
|
||||
static_cast<size_t>(batch_size) &&
|
||||
input_params.linear_state_validity_mask.size() ==
|
||||
static_cast<size_t>(batch_size);
|
||||
int64_t total_valid_tokens = 0;
|
||||
bool direct_qkv_lengths_valid = direct_qkv_metadata_available;
|
||||
if (direct_qkv_metadata_available) {
|
||||
for (const int32_t valid_len : attn_metadata.q_seq_lens_vec) {
|
||||
direct_qkv_lengths_valid =
|
||||
direct_qkv_lengths_valid && valid_len >= 0 && valid_len <= seq_len;
|
||||
total_valid_tokens += valid_len;
|
||||
}
|
||||
}
|
||||
const bool direct_qkv_sequence_supported =
|
||||
direct_qkv_model_supported && direct_qkv_lengths_valid &&
|
||||
conv_input.dim() == 2 && total_valid_tokens == conv_input.size(0);
|
||||
const bool direct_qkv_shape_supported =
|
||||
direct_qkv_sequence_supported && conv_input.size(1) == local_conv_dim &&
|
||||
conv_weight.dim() == 2 && conv_weight.size(0) == 4 &&
|
||||
conv_weight.size(1) == local_conv_dim && conv_cache.dim() == 3 &&
|
||||
conv_cache.size(1) >= 3 && conv_cache.size(2) == local_conv_dim;
|
||||
const bool direct_qkv_dtype_supported =
|
||||
direct_qkv_shape_supported &&
|
||||
conv_input.scalar_type() == torch::kBFloat16 &&
|
||||
conv_weight.scalar_type() == torch::kBFloat16 &&
|
||||
conv_cache.scalar_type() == torch::kBFloat16;
|
||||
const bool use_direct_prefill_qkv =
|
||||
direct_qkv_dtype_supported && conv_input.is_contiguous() &&
|
||||
conv_weight.is_contiguous() && conv_cache.is_contiguous();
|
||||
if (use_direct_prefill_qkv) {
|
||||
std::tie(processed_q, processed_k, processed_v) =
|
||||
xllm::kernel::npu::causal_conv1d_qkv(
|
||||
conv_input,
|
||||
conv_weight,
|
||||
conv_cache,
|
||||
torch::IntArrayRef(input_params.parallel.query_start_loc),
|
||||
torch::IntArrayRef(linear_state_indices_vec),
|
||||
torch::IntArrayRef(input_params.linear_state_validity_mask),
|
||||
local_q_heads,
|
||||
local_v_heads,
|
||||
head_k_dim_,
|
||||
head_v_dim_);
|
||||
used_direct_prefill_qkv = true;
|
||||
} else {
|
||||
mixed_qkv = xllm::kernel::causal_conv1d(
|
||||
conv_input,
|
||||
conv_weight,
|
||||
conv_cache,
|
||||
std::optional<torch::Tensor>(), // bias (no bias for qwen3)
|
||||
torch::IntArrayRef(input_params.parallel.query_start_loc),
|
||||
torch::IntArrayRef(linear_state_indices_vec),
|
||||
torch::IntArrayRef(input_params.linear_state_validity_mask),
|
||||
num_accepted_tokens_opt,
|
||||
xllm::npu::kCausalConv1dActivationSilu,
|
||||
xllm::npu::kCausalConv1dGraphPadSlotId,
|
||||
xllm::npu::kCausalConv1dRunModeForward);
|
||||
|
||||
mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv);
|
||||
mixed_qkv = mixed_qkv.transpose(1, 2);
|
||||
}
|
||||
} else {
|
||||
if (use_spec_verify) {
|
||||
CHECK(input_params.num_accepted_tokens.defined())
|
||||
<< "num_accepted_tokens must be populated for Qwen3.5 spec verify";
|
||||
}
|
||||
torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv);
|
||||
const auto& num_accepted = use_spec_verify
|
||||
? input_params.num_accepted_tokens_host
|
||||
: std::vector<int64_t>();
|
||||
const std::vector<int64_t> linear_state_indices_host(
|
||||
input_params.embedding.linear_state_ids.begin(),
|
||||
input_params.embedding.linear_state_ids.end());
|
||||
if (register_conv1d_graph_update) {
|
||||
if (use_spec_verify) {
|
||||
const auto conv1d_branch =
|
||||
xllm::npu::CausalConv1dGraphBranch::kSpecVerify;
|
||||
mixed_qkv = run_causal_conv1d_graph_update(
|
||||
graph_context,
|
||||
conv_input,
|
||||
conv_weight,
|
||||
conv_cache,
|
||||
std::optional<torch::Tensor>(),
|
||||
input_params.parallel.query_start_loc,
|
||||
linear_state_indices_host,
|
||||
num_accepted,
|
||||
conv1d_branch);
|
||||
} else {
|
||||
auto conv_input_2d = conv_input.dim() == 3
|
||||
? conv_input.reshape({-1, conv_input.size(-1)})
|
||||
: conv_input;
|
||||
xllm::kernel::CausalConv1dUpdateParams conv1d_params;
|
||||
conv1d_params.x = conv_input_2d;
|
||||
conv1d_params.conv_state = conv_cache;
|
||||
conv1d_params.weight = conv_weight;
|
||||
conv1d_params.conv_state_indices = logical_state_indices;
|
||||
conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
|
||||
conv1d_params.max_query_len = attn_metadata.max_query_len;
|
||||
mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
|
||||
if (conv_input.dim() == 3) {
|
||||
mixed_qkv =
|
||||
mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)});
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (use_spec_verify) {
|
||||
torch::Tensor output = torch::empty_like(conv_input);
|
||||
xllm::kernel::causal_conv1d_out(
|
||||
output,
|
||||
conv_input,
|
||||
conv_weight,
|
||||
conv_cache,
|
||||
std::optional<torch::Tensor>(),
|
||||
torch::IntArrayRef(input_params.parallel.query_start_loc),
|
||||
torch::IntArrayRef(linear_state_indices_host),
|
||||
torch::IntArrayRef(std::vector<int64_t>()),
|
||||
torch::IntArrayRef(num_accepted),
|
||||
xllm::npu::kCausalConv1dActivationSilu,
|
||||
xllm::npu::kCausalConv1dGraphPadSlotId,
|
||||
xllm::npu::kCausalConv1dRunModeUpdate);
|
||||
mixed_qkv = output;
|
||||
} else {
|
||||
auto conv_input_2d = conv_input.dim() == 3
|
||||
? conv_input.reshape({-1, conv_input.size(-1)})
|
||||
: conv_input;
|
||||
xllm::kernel::CausalConv1dUpdateParams conv1d_params;
|
||||
conv1d_params.x = conv_input_2d;
|
||||
conv1d_params.conv_state = conv_cache;
|
||||
conv1d_params.weight = conv_weight;
|
||||
conv1d_params.conv_state_indices = logical_state_indices;
|
||||
conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
|
||||
conv1d_params.max_query_len = attn_metadata.max_query_len;
|
||||
mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
|
||||
if (conv_input.dim() == 3) {
|
||||
mixed_qkv =
|
||||
mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)});
|
||||
}
|
||||
}
|
||||
}
|
||||
mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv);
|
||||
xllm::kernel::CausalConv1dUpdateParams conv1d_params;
|
||||
conv1d_params.x = mixed_qkv.reshape({-1, mixed_qkv.size(-1)});
|
||||
conv1d_params.conv_state = conv_cache;
|
||||
conv1d_params.weight = conv_weight;
|
||||
conv1d_params.conv_state_indices = linear_state_indices;
|
||||
conv1d_params.block_idx_last_scheduled_token =
|
||||
std::optional<torch::Tensor>();
|
||||
conv1d_params.initial_state_idx = std::optional<torch::Tensor>();
|
||||
conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
|
||||
conv1d_params.max_query_len = attn_metadata.max_query_len;
|
||||
mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
|
||||
// Reshape back to 3D [batch_size, dim, seq_len]
|
||||
mixed_qkv =
|
||||
mixed_qkv.view({batch_size, -1, mixed_qkv.size(-1)}).contiguous();
|
||||
mixed_qkv = mixed_qkv.transpose(1, 2);
|
||||
}
|
||||
const bool use_fused_sigmoid_gdn_decode =
|
||||
fla_ssm_state_layout && !use_spec_verify && !is_any_prefill &&
|
||||
checkpoint_stride == 1;
|
||||
torch::Tensor g;
|
||||
torch::Tensor beta;
|
||||
|
||||
// Compute gated delta net decay and beta terms.
|
||||
if (use_spec_verify || attn_metadata.is_chunked_prefill ||
|
||||
checkpoint_stride > 1) {
|
||||
beta = torch::sigmoid(b);
|
||||
torch::Tensor A_log_exp = A_log_.exp();
|
||||
torch::Tensor a_float = a.to(torch::kFloat32);
|
||||
torch::Tensor a_plus_dt = a_float + dt_bias_;
|
||||
torch::Tensor softplus_out = torch::nn::functional::softplus(
|
||||
a_plus_dt,
|
||||
torch::nn::functional::SoftplusFuncOptions().beta(1.0).threshold(20.0));
|
||||
g = -A_log_exp * softplus_out;
|
||||
g = g.to(a.dtype()).contiguous();
|
||||
} else if (attn_metadata.is_prefill) {
|
||||
if (attn_metadata.is_prefill) {
|
||||
xllm::kernel::FusedGdnGatingParams gdn_params;
|
||||
gdn_params.A_log = A_log_;
|
||||
gdn_params.a = a.contiguous().view({-1, a.size(-1)});
|
||||
@@ -801,7 +413,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
|
||||
std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
|
||||
g = g.squeeze(0).contiguous().view({batch_size, seq_len, a.size(-1)});
|
||||
beta = beta.squeeze(0).contiguous().view({batch_size, seq_len, b.size(-1)});
|
||||
} else if (!use_fused_sigmoid_gdn_decode) {
|
||||
} else {
|
||||
xllm::kernel::FusedGdnGatingParams gdn_params;
|
||||
gdn_params.A_log = A_log_;
|
||||
gdn_params.a = a.view({-1, a.size(-1)});
|
||||
@@ -811,216 +423,57 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
|
||||
gdn_params.threshold = 20.0f;
|
||||
std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
|
||||
}
|
||||
if (!used_direct_prefill_qkv) {
|
||||
std::tie(processed_q, processed_k, processed_v) =
|
||||
process_mixed_qkv(mixed_qkv);
|
||||
}
|
||||
torch::Tensor core_attn_out;
|
||||
torch::Tensor last_recurrent_state;
|
||||
auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv);
|
||||
// Apply chunked or recurrent gated-delta attention and update caches.
|
||||
if (use_spec_verify) {
|
||||
torch::Tensor spec_num_accepted_tokens = expand_sequence_tensor_to_batch(
|
||||
input_params.num_accepted_tokens.to(device, torch::kInt32),
|
||||
batch_size,
|
||||
"num_accepted_tokens");
|
||||
torch::Tensor spec_linear_state_base_indices =
|
||||
expand_sequence_tensor_to_batch(
|
||||
linear_state_base_indices, batch_size, "linear_state_base_indices");
|
||||
torch::Tensor step_offsets =
|
||||
torch::arange(seq_len,
|
||||
torch::TensorOptions()
|
||||
.dtype(spec_linear_state_base_indices.dtype())
|
||||
.device(device));
|
||||
torch::Tensor checkpoint_indices =
|
||||
spec_linear_state_base_indices.unsqueeze(1) + step_offsets;
|
||||
double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
|
||||
core_attn_out =
|
||||
run_spec_verify_gated_delta_rule(processed_q,
|
||||
processed_k,
|
||||
processed_v,
|
||||
g,
|
||||
beta,
|
||||
ssm_cache,
|
||||
checkpoint_indices,
|
||||
spec_num_accepted_tokens,
|
||||
attn_metadata.q_cu_seq_lens,
|
||||
attn_metadata.q_seq_lens_vec,
|
||||
scale);
|
||||
} else if (is_any_prefill) {
|
||||
CHECK_GE(attn_metadata.q_seq_lens_vec.size(),
|
||||
static_cast<size_t>(batch_size))
|
||||
<< "q_seq_lens_vec must be populated for Qwen3.5 prefill.";
|
||||
const bool use_single_prefill_pack =
|
||||
batch_size == 1 && attn_metadata.q_seq_lens_vec.size() == 1 &&
|
||||
attn_metadata.q_seq_lens_vec[0] == seq_len;
|
||||
torch::Tensor packed_processed_q;
|
||||
torch::Tensor packed_processed_k;
|
||||
torch::Tensor packed_processed_v;
|
||||
torch::Tensor packed_g_tensor;
|
||||
torch::Tensor packed_beta_tensor;
|
||||
if (use_single_prefill_pack) {
|
||||
packed_processed_q = processed_q;
|
||||
packed_processed_k = processed_k;
|
||||
packed_processed_v = processed_v;
|
||||
packed_g_tensor = g;
|
||||
packed_beta_tensor = beta;
|
||||
} else {
|
||||
std::vector<torch::Tensor> packed_q;
|
||||
std::vector<torch::Tensor> packed_k;
|
||||
std::vector<torch::Tensor> packed_v;
|
||||
std::vector<torch::Tensor> packed_g;
|
||||
std::vector<torch::Tensor> packed_beta;
|
||||
packed_q.reserve(batch_size);
|
||||
packed_k.reserve(batch_size);
|
||||
packed_v.reserve(batch_size);
|
||||
packed_g.reserve(batch_size);
|
||||
packed_beta.reserve(batch_size);
|
||||
for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
|
||||
const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx];
|
||||
if (!used_direct_prefill_qkv) {
|
||||
packed_q.emplace_back(processed_q[batch_idx].narrow(
|
||||
/*dim=*/0, /*start=*/0, valid_len));
|
||||
packed_k.emplace_back(processed_k[batch_idx].narrow(
|
||||
/*dim=*/0, /*start=*/0, valid_len));
|
||||
packed_v.emplace_back(processed_v[batch_idx].narrow(
|
||||
/*dim=*/0, /*start=*/0, valid_len));
|
||||
}
|
||||
packed_g.emplace_back(
|
||||
g[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len));
|
||||
packed_beta.emplace_back(
|
||||
beta[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len));
|
||||
}
|
||||
if (used_direct_prefill_qkv) {
|
||||
packed_processed_q = processed_q;
|
||||
packed_processed_k = processed_k;
|
||||
packed_processed_v = processed_v;
|
||||
} else {
|
||||
packed_processed_q = torch::cat(packed_q, 0).unsqueeze(0);
|
||||
packed_processed_k = torch::cat(packed_k, 0).unsqueeze(0);
|
||||
packed_processed_v = torch::cat(packed_v, 0).unsqueeze(0);
|
||||
}
|
||||
packed_g_tensor = torch::cat(packed_g, 0).unsqueeze(0);
|
||||
packed_beta_tensor = torch::cat(packed_beta, 0).unsqueeze(0);
|
||||
}
|
||||
|
||||
xllm::kernel::MegaChunkGdnParams mega_chunk_gdn_params;
|
||||
mega_chunk_gdn_params.q = packed_processed_q;
|
||||
mega_chunk_gdn_params.k = packed_processed_k;
|
||||
mega_chunk_gdn_params.v = packed_processed_v;
|
||||
mega_chunk_gdn_params.g = packed_g_tensor;
|
||||
mega_chunk_gdn_params.beta = packed_beta_tensor;
|
||||
if (attn_metadata.is_prefill) {
|
||||
xllm::kernel::ChunkGatedDeltaRuleParams chunk_gated_delta_params;
|
||||
chunk_gated_delta_params.q = processed_q;
|
||||
chunk_gated_delta_params.k = processed_k;
|
||||
chunk_gated_delta_params.v = processed_v;
|
||||
chunk_gated_delta_params.g = g;
|
||||
chunk_gated_delta_params.beta = beta;
|
||||
// Get initial state from ssm_cache for sequences with previous state
|
||||
// Shape: [batch_size, num_heads, head_k_dim, head_v_dim]
|
||||
torch::Tensor initial_state_tensor =
|
||||
torch::index_select(ssm_cache, 0, linear_state_base_indices);
|
||||
CHECK_EQ(input_params.linear_state_validity_mask.size(),
|
||||
input_params.embedding.linear_state_ids.size())
|
||||
<< "linear state validity mask must be sequence-scoped.";
|
||||
for (size_t i = 0; i < input_params.linear_state_validity_mask.size();
|
||||
++i) {
|
||||
if (input_params.linear_state_validity_mask[i] == 0) {
|
||||
initial_state_tensor.select(0, static_cast<int64_t>(i)).fill_(0.0);
|
||||
}
|
||||
}
|
||||
if (!fla_ssm_state_layout && attn_metadata.is_chunked_prefill) {
|
||||
initial_state_tensor =
|
||||
initial_state_tensor.transpose(-1, -2).contiguous();
|
||||
}
|
||||
mega_chunk_gdn_params.initial_state = initial_state_tensor;
|
||||
mega_chunk_gdn_params.output_final_state = true;
|
||||
mega_chunk_gdn_params.cu_seqlens = attn_metadata.q_cu_seq_lens;
|
||||
mega_chunk_gdn_params.q_seq_lens = c10::ArrayRef<int32_t>(
|
||||
attn_metadata.q_seq_lens_vec.data(), static_cast<size_t>(batch_size));
|
||||
mega_chunk_gdn_params.use_qk_l2norm_in_kernel = !used_direct_prefill_qkv;
|
||||
torch::Tensor packed_core_attn_out;
|
||||
std::tie(packed_core_attn_out, last_recurrent_state) =
|
||||
xllm::kernel::mega_chunk_gdn(mega_chunk_gdn_params);
|
||||
if (use_single_prefill_pack) {
|
||||
core_attn_out = packed_core_attn_out;
|
||||
if (core_attn_out.scalar_type() != processed_v.scalar_type()) {
|
||||
core_attn_out = core_attn_out.to(processed_v.scalar_type());
|
||||
}
|
||||
} else {
|
||||
core_attn_out =
|
||||
used_direct_prefill_qkv
|
||||
? torch::zeros({batch_size, seq_len, local_v_heads, head_v_dim_},
|
||||
z.options())
|
||||
: torch::zeros_like(processed_v);
|
||||
int64_t packed_offset = 0;
|
||||
for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
|
||||
const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx];
|
||||
core_attn_out[batch_idx]
|
||||
.narrow(/*dim=*/0, /*start=*/0, valid_len)
|
||||
.copy_(packed_core_attn_out[0].narrow(
|
||||
/*dim=*/0, packed_offset, valid_len));
|
||||
packed_offset += valid_len;
|
||||
}
|
||||
}
|
||||
torch::Tensor state_to_store = fla_ssm_state_layout
|
||||
? last_recurrent_state
|
||||
: last_recurrent_state.transpose(-1, -2);
|
||||
ssm_cache.index_put_({linear_state_base_indices},
|
||||
state_to_store.to(ssm_cache.dtype()));
|
||||
} else if (checkpoint_stride > 1) {
|
||||
auto ssm_state =
|
||||
torch::index_select(ssm_cache, 0, linear_state_base_indices);
|
||||
if (!fla_ssm_state_layout) {
|
||||
ssm_state = ssm_state.transpose(-1, -2);
|
||||
}
|
||||
ssm_state = ssm_state.contiguous();
|
||||
torch::index_select(ssm_cache, 0, linear_state_indices);
|
||||
// Todo: chunked-prefill/prefix-cache use initial_state
|
||||
initial_state_tensor.fill_(0.0);
|
||||
chunk_gated_delta_params.initial_state = initial_state_tensor;
|
||||
chunk_gated_delta_params.output_final_state = true;
|
||||
chunk_gated_delta_params.cu_seqlens = attn_metadata.q_cu_seq_lens;
|
||||
chunk_gated_delta_params.head_first = false;
|
||||
chunk_gated_delta_params.use_qk_l2norm_in_kernel = true;
|
||||
std::tie(core_attn_out, last_recurrent_state) =
|
||||
torch_recurrent_gated_delta_rule(
|
||||
processed_q, processed_k, processed_v, g, beta, ssm_state);
|
||||
torch::Tensor state_to_store = fla_ssm_state_layout
|
||||
? last_recurrent_state
|
||||
: last_recurrent_state.transpose(-1, -2);
|
||||
ssm_cache.index_put_({linear_state_base_indices},
|
||||
state_to_store.to(ssm_cache.dtype()));
|
||||
xllm::kernel::chunk_gated_delta_rule(chunk_gated_delta_params);
|
||||
ssm_cache.index_put_(
|
||||
{linear_state_indices},
|
||||
last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype()));
|
||||
} else {
|
||||
processed_q = xllm::kernel::l2_norm(processed_q, 1e-6);
|
||||
processed_k = xllm::kernel::l2_norm(processed_k, 1e-6);
|
||||
auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
|
||||
torch::Tensor actual_seq_lengths =
|
||||
torch::cat({zero, attn_metadata.q_seq_lens}, 0);
|
||||
double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
|
||||
if (fla_ssm_state_layout) {
|
||||
xllm::kernel::FusedSigmoidGatingDeltaRuleUpdateParams params;
|
||||
params.A_log = A_log_.contiguous();
|
||||
params.a = a.contiguous();
|
||||
params.dt_bias = dt_bias_.contiguous();
|
||||
params.q = processed_q.contiguous();
|
||||
params.k = processed_k.contiguous();
|
||||
params.v = processed_v.contiguous();
|
||||
params.b = b.contiguous();
|
||||
params.initial_state_source = ssm_cache;
|
||||
params.initial_state_indices = linear_state_base_indices.contiguous();
|
||||
params.cu_seqlens = attn_metadata.q_cu_seq_lens.contiguous();
|
||||
params.scale = static_cast<float>(scale);
|
||||
params.use_qk_l2norm_in_kernel = true;
|
||||
params.softplus_beta = 1.0f;
|
||||
params.softplus_threshold = 20.0f;
|
||||
core_attn_out =
|
||||
xllm::kernel::fused_sigmoid_gating_delta_rule_update(params);
|
||||
} else {
|
||||
processed_q = xllm::kernel::l2_norm(processed_q, /*eps=*/1e-6);
|
||||
processed_k = xllm::kernel::l2_norm(processed_k, /*eps=*/1e-6);
|
||||
auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
|
||||
torch::Tensor actual_seq_lengths =
|
||||
torch::cat({zero, attn_metadata.q_seq_lens}, 0);
|
||||
core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
|
||||
processed_q.reshape(
|
||||
{-1, processed_q.size(-2), processed_q.size(-1)}),
|
||||
processed_k.reshape(
|
||||
{-1, processed_k.size(-2), processed_k.size(-1)}),
|
||||
processed_v.reshape(
|
||||
{-1, processed_v.size(-2), processed_v.size(-1)}),
|
||||
ssm_cache,
|
||||
beta.squeeze(0).contiguous(),
|
||||
scale,
|
||||
actual_seq_lengths,
|
||||
logical_state_indices,
|
||||
c10::nullopt,
|
||||
g.squeeze(0).contiguous(),
|
||||
c10::nullopt)
|
||||
.unsqueeze(0)
|
||||
.contiguous();
|
||||
}
|
||||
core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
|
||||
processed_q.reshape(
|
||||
{-1, processed_q.size(-2), processed_q.size(-1)}),
|
||||
processed_k.reshape(
|
||||
{-1, processed_k.size(-2), processed_k.size(-1)}),
|
||||
processed_v.reshape(
|
||||
{-1, processed_v.size(-2), processed_v.size(-1)}),
|
||||
ssm_cache,
|
||||
beta.squeeze(0).contiguous(),
|
||||
scale,
|
||||
actual_seq_lengths,
|
||||
linear_state_indices,
|
||||
c10::nullopt,
|
||||
g.squeeze(0).contiguous(),
|
||||
c10::nullopt)
|
||||
.unsqueeze(0)
|
||||
.contiguous();
|
||||
}
|
||||
|
||||
auto z_reshaped = z.view({-1, z.size(-1)});
|
||||
auto core_attn_out_reshaped =
|
||||
core_attn_out.view({-1, core_attn_out.size(-1)});
|
||||
@@ -1033,47 +486,25 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
|
||||
auto rearranged_norm =
|
||||
norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)});
|
||||
rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm);
|
||||
// For chunked prefill or spec verify, reshape_projected_tokens_with_pad may
|
||||
// pad each batch to max_len, causing output tokens > original_num_tokens. We
|
||||
// need to slice back to original_num_tokens to match the residual shape.
|
||||
if (rearranged_norm.size(0) > original_num_tokens) {
|
||||
// Slice excess padding tokens
|
||||
rearranged_norm =
|
||||
rearranged_norm.slice(0, 0, original_num_tokens).contiguous();
|
||||
}
|
||||
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
|
||||
return o_proj_->forward(rearranged_norm,
|
||||
row_parallel_reduce_mode_for_fc1(*fc1_ctx));
|
||||
}
|
||||
return o_proj_->forward(rearranged_norm);
|
||||
auto attn_output = o_proj_->forward(rearranged_norm);
|
||||
return attn_output;
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& padded_qkvz) const {
|
||||
const bool has_padded_queries =
|
||||
attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
|
||||
if (!has_padded_queries) {
|
||||
if (!attn_metadata.is_prefill) {
|
||||
return padded_qkvz;
|
||||
}
|
||||
std::vector<torch::Tensor> valid_batches;
|
||||
const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty();
|
||||
int64_t bs = has_host_lens
|
||||
? static_cast<int64_t>(attn_metadata.q_seq_lens_vec.size())
|
||||
: attn_metadata.q_seq_lens.size(0);
|
||||
valid_batches.reserve(bs);
|
||||
int64_t bs = attn_metadata.q_seq_lens.size(0);
|
||||
int64_t max_len = attn_metadata.max_query_len;
|
||||
const auto& ori_seq_lens = attn_metadata.q_seq_lens;
|
||||
auto reshaped_qkvz = padded_qkvz.view({bs, max_len, -1});
|
||||
for (int64_t b = 0; b < bs; ++b) {
|
||||
int64_t ori_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b]
|
||||
: ori_seq_lens[b].template item<int64_t>();
|
||||
torch::Tensor valid_batch =
|
||||
reshaped_qkvz[b].slice(/*dim=*/0, /*start=*/0, ori_len);
|
||||
valid_batches.emplace_back(valid_batch);
|
||||
}
|
||||
if (valid_batches.size() == 1) {
|
||||
return valid_batches[0].contiguous();
|
||||
int64_t ori_len = ori_seq_lens[b].template item<int64_t>();
|
||||
torch::Tensor valid_batch = reshaped_qkvz[b].slice(0, 0, ori_len);
|
||||
valid_batches.push_back(valid_batch);
|
||||
}
|
||||
return torch::cat(valid_batches, 0).contiguous();
|
||||
}
|
||||
@@ -1081,60 +512,41 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
|
||||
torch::Tensor Qwen3GatedDeltaNetBaseImpl::get_linear_state_indices(
|
||||
const ModelInputParams& input_params,
|
||||
const torch::Device& device) const {
|
||||
CHECK(!input_params.embedding.linear_state_ids.empty())
|
||||
CHECK(!input_params.linear_state_ids.empty())
|
||||
<< "linear_state_ids must be populated for gated delta net";
|
||||
if (input_params.embedding.linear_state_indices.defined()) {
|
||||
auto indices = input_params.embedding.linear_state_indices;
|
||||
if (indices.device() != device || indices.scalar_type() != torch::kInt) {
|
||||
indices =
|
||||
indices.to(torch::TensorOptions().dtype(torch::kInt).device(device),
|
||||
/*non_blocking=*/true,
|
||||
/*copy=*/true);
|
||||
}
|
||||
return indices.contiguous();
|
||||
if (input_params.linear_state_indices.defined()) {
|
||||
return input_params.linear_state_indices;
|
||||
}
|
||||
return torch::tensor(
|
||||
input_params.embedding.linear_state_ids,
|
||||
input_params.linear_state_ids,
|
||||
torch::TensorOptions().dtype(torch::kInt).device(device));
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_projected_tokens_with_pad(
|
||||
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& projected_tokens) const {
|
||||
const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty();
|
||||
int64_t bs = has_host_lens
|
||||
? static_cast<int64_t>(attn_metadata.q_seq_lens_vec.size())
|
||||
: attn_metadata.q_seq_lens.size(0);
|
||||
const torch::Tensor& qkvz) const {
|
||||
int64_t bs = attn_metadata.q_seq_lens.size(0);
|
||||
int64_t max_len = attn_metadata.max_query_len;
|
||||
const auto& start_loc = attn_metadata.q_seq_lens;
|
||||
const bool need_padding =
|
||||
attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
|
||||
if (!need_padding) {
|
||||
return projected_tokens.view({bs, -1, projected_tokens.size(-1)});
|
||||
}
|
||||
if (has_host_lens && bs == 1 && attn_metadata.q_seq_lens_vec[0] == max_len &&
|
||||
projected_tokens.dim() == 2 && projected_tokens.size(0) == max_len) {
|
||||
return projected_tokens.view({1, max_len, projected_tokens.size(-1)});
|
||||
if (!attn_metadata.is_prefill) {
|
||||
return qkvz.view({qkvz.size(0), -1, qkvz.size(-1)});
|
||||
}
|
||||
std::vector<torch::Tensor> batches;
|
||||
batches.reserve(bs);
|
||||
int64_t idx = 0;
|
||||
for (int64_t b = 0; b < bs; ++b) {
|
||||
int64_t cur_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b]
|
||||
: start_loc[b].template item<int64_t>();
|
||||
torch::Tensor batch =
|
||||
projected_tokens.slice(/*dim=*/0, idx, idx + cur_len).contiguous();
|
||||
int64_t cur_len = start_loc[b].template item<int64_t>();
|
||||
torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous();
|
||||
idx = idx + cur_len;
|
||||
if (batch.size(0) != max_len) {
|
||||
batch = batch.size(0) > max_len
|
||||
? batch.slice(/*dim=*/0, /*start=*/0, max_len).contiguous()
|
||||
? batch.slice(0, 0, max_len).contiguous()
|
||||
: torch::nn::functional::pad(
|
||||
batch,
|
||||
torch::nn::functional::PadFuncOptions(
|
||||
{0, 0, 0, max_len - batch.size(0)}))
|
||||
.contiguous();
|
||||
}
|
||||
batches.emplace_back(batch);
|
||||
batches.push_back(batch);
|
||||
}
|
||||
auto ret = torch::stack(batches, 0).contiguous();
|
||||
return ret;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
@@ -17,7 +17,6 @@ limitations under the License.
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <tuple>
|
||||
#include <utility>
|
||||
@@ -52,40 +51,19 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module {
|
||||
const ModelInputParams& input_params);
|
||||
|
||||
protected:
|
||||
virtual std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
|
||||
const torch::Tensor& hidden_states) = 0;
|
||||
virtual std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
|
||||
const torch::Tensor& hidden_states) = 0;
|
||||
// Qwen3.5 overrides this to project and reshape its separate qkv/z/b/a
|
||||
// weights in every forward mode. Qwen3Next keeps qkvz/ba packed and returns
|
||||
// nullopt to select the fused-split fallback.
|
||||
virtual std::optional<
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
|
||||
project_split_inputs(const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
return std::nullopt;
|
||||
}
|
||||
virtual bool use_fla_ssm_state_layout() const { return false; }
|
||||
virtual std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) = 0;
|
||||
|
||||
void load_common_state_dict(const StateDict& state_dict);
|
||||
void verify_common_loaded_weights(const std::string& prefix) const;
|
||||
|
||||
torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
|
||||
const torch::Device& device) const;
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata);
|
||||
|
||||
torch::Tensor reshape_qkvz_with_pad(const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& qkvz) const;
|
||||
torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& padded_qkvz) const;
|
||||
|
||||
// Projection outputs are packed as [total_tokens, dim], while GDN kernels
|
||||
// consume dense [batch, max_query_len, dim] tensors. Split the packed tokens
|
||||
// by query length and pad each sequence before entering the kernels.
|
||||
torch::Tensor reshape_projected_tokens_with_pad(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& projected_tokens) const;
|
||||
torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
|
||||
const torch::Device& device) const;
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> process_mixed_qkv(
|
||||
torch::Tensor& mixed_qkv) const;
|
||||
|
||||
Reference in New Issue
Block a user