From 1ea5f6306902e49437667e16259fa4eb7307b6a8 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Tue, 14 Jul 2026 05:37:05 +0000 Subject: [PATCH 1/5] pepe: glm 5.2 trial --- csrc/config/config_factory.cpp | 2 +- csrc/config/quant_config.cpp | 4 + csrc/layers/linear/base_linear.hpp | 1 + csrc/layers/mlp/mlp.hpp | 2 + csrc/layers/moe/legacy/moe_mlp.hpp | 8 + csrc/layers/quantization/glm_w4a8.cpp | 249 ++++++++++++++++++ csrc/layers/quantization/glm_w4a8.hpp | 77 ++++++ csrc/layers/quantization/glm_w8a8.cpp | 92 +++++++ csrc/layers/quantization/glm_w8a8.hpp | 31 +++ csrc/layers/quantization/quantization.hpp | 2 + .../quantization/quantization_scheme.hpp | 3 + csrc/models/glm_moe_dsa/glm_attention.cpp | 63 +++++ csrc/models/glm_moe_dsa/glm_attention.hpp | 28 ++ csrc/models/glm_moe_dsa/glm_model.cpp | 74 ++++++ csrc/models/glm_moe_dsa/glm_model.hpp | 44 ++++ csrc/models/glm_moe_dsa/glm_moe.cpp | 117 ++++++++ csrc/models/glm_moe_dsa/glm_moe.hpp | 49 ++++ .../models/glm_moe_dsa/glm_vocab_parallel.cpp | 42 +++ .../models/glm_moe_dsa/glm_vocab_parallel.hpp | 26 ++ python/infinilm/modeling_utils.py | 28 ++ xmake.lua | 1 + 21 files changed, 942 insertions(+), 1 deletion(-) create mode 100644 csrc/layers/quantization/glm_w4a8.cpp create mode 100644 csrc/layers/quantization/glm_w4a8.hpp create mode 100644 csrc/layers/quantization/glm_w8a8.cpp create mode 100644 csrc/layers/quantization/glm_w8a8.hpp create mode 100644 csrc/models/glm_moe_dsa/glm_attention.cpp create mode 100644 csrc/models/glm_moe_dsa/glm_attention.hpp create mode 100644 csrc/models/glm_moe_dsa/glm_model.cpp create mode 100644 csrc/models/glm_moe_dsa/glm_model.hpp create mode 100644 csrc/models/glm_moe_dsa/glm_moe.cpp create mode 100644 csrc/models/glm_moe_dsa/glm_moe.hpp create mode 100644 csrc/models/glm_moe_dsa/glm_vocab_parallel.cpp create mode 100644 csrc/models/glm_moe_dsa/glm_vocab_parallel.hpp diff --git a/csrc/config/config_factory.cpp b/csrc/config/config_factory.cpp index 0467f4536..5543831a3 100644 --- a/csrc/config/config_factory.cpp +++ b/csrc/config/config_factory.cpp @@ -12,7 +12,7 @@ std::shared_ptr ConfigFactory::createConfig(const const auto &config_map = models::get_model_config_map(); auto it = config_map.find(model_type); if (it != config_map.end()) { - it->second(model_config); + model_config = it->second(model_config); } else { throw std::invalid_argument("infinilm::config::ConfigFactory::createConfig: Unsupported model config type: " + model_type); } diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index da261ce09..a0b739b13 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -20,6 +20,10 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); } else if (quant_method == "gptq") { return std::make_shared(quantization_config); + } else if (quant_method == "glm_w8a8" || quant_method == "w8a8") { + return std::make_shared(quantization_config); + } else if (quant_method == "glm_w4a8" || quant_method == "w4a8") { + return std::make_shared(quantization_config); } else { return std::make_shared(quantization_config); } diff --git a/csrc/layers/linear/base_linear.hpp b/csrc/layers/linear/base_linear.hpp index a36954836..8fe334d4c 100644 --- a/csrc/layers/linear/base_linear.hpp +++ b/csrc/layers/linear/base_linear.hpp @@ -35,6 +35,7 @@ class BaseLinear : public infinicore::nn::Module { void set_alpha(float alpha) { alpha_ = alpha; } // Accessors for parameters (backward compatible) + void release_parameters() { parameters_.clear(); } infinicore::Tensor weight() const; infinicore::Tensor bias() const; infinicore::Tensor weight_scale() const; diff --git a/csrc/layers/mlp/mlp.hpp b/csrc/layers/mlp/mlp.hpp index abd81bd88..ff8be6a77 100644 --- a/csrc/layers/mlp/mlp.hpp +++ b/csrc/layers/mlp/mlp.hpp @@ -38,10 +38,12 @@ class MLP : public infinicore::nn::Module { void process_weights_after_loading() override { gate_up_proj_->process_weights_after_loading(); + down_proj_->process_weights_after_loading(); } void reset_runtime_state() const override { gate_up_proj_->reset_runtime_state(); + down_proj_->reset_runtime_state(); } // Module information diff --git a/csrc/layers/moe/legacy/moe_mlp.hpp b/csrc/layers/moe/legacy/moe_mlp.hpp index 77d39944d..75eda82a5 100644 --- a/csrc/layers/moe/legacy/moe_mlp.hpp +++ b/csrc/layers/moe/legacy/moe_mlp.hpp @@ -18,7 +18,15 @@ class MoeMLP : public infinicore::nn::Module { infinicore::Tensor gate_weight() const { return gate_proj_->weight(); } infinicore::Tensor up_weight() const { return up_proj_->weight(); } infinicore::Tensor down_weight() const { return down_proj_->weight(); } + infinicore::Tensor gate_weight_scale() const { return gate_proj_->weight_scale(); } + infinicore::Tensor up_weight_scale() const { return up_proj_->weight_scale(); } + infinicore::Tensor down_weight_scale() const { return down_proj_->weight_scale(); } void set_alpha(float alpha) { down_proj_->set_alpha(alpha); } + void release_parameters() { + gate_proj_->release_parameters(); + up_proj_->release_parameters(); + down_proj_->release_parameters(); + } protected: std::shared_ptr gate_proj_; diff --git a/csrc/layers/quantization/glm_w4a8.cpp b/csrc/layers/quantization/glm_w4a8.cpp new file mode 100644 index 000000000..8dfdee887 --- /dev/null +++ b/csrc/layers/quantization/glm_w4a8.cpp @@ -0,0 +1,249 @@ +#include "glm_w4a8.hpp" + +#include "infinicore/ops/dynamic_scaled_int8_quant.hpp" +#include "infinicore/ops/mul_scalar.hpp" +#include "infinicore/ops/scaled_mm_w4a8.hpp" + +#include +#include +#include +#include + +namespace infinilm::quantization { +namespace { + +infinicore::Tensor repack_out_khalf_to_k_outhalf_cpu(const infinicore::Tensor &packed_weight, + const infinicore::Device &device) { + if (packed_weight->ndim() != 2 || packed_weight->dtype() != infinicore::DataType::I8) { + throw std::runtime_error("GlmW4A8 expects int8 weight with shape [out, in/2]"); + } + const size_t out_features = packed_weight->size(0); + const size_t in_half = packed_weight->size(1); + const size_t in_features = in_half * 2; + if ((out_features % 2) != 0) { + throw std::runtime_error("GlmW4A8 requires even out_features for packed runtime layout"); + } + + auto src_cpu = packed_weight->to(infinicore::Device::cpu())->contiguous(); + auto dst_cpu = infinicore::Tensor::empty({in_features, out_features / 2}, infinicore::DataType::I8, infinicore::Device::cpu()); + const auto *src = reinterpret_cast(src_cpu->data()); + auto *dst = reinterpret_cast(dst_cpu->data()); + + for (size_t k = 0; k < in_features; ++k) { + const size_t k_byte = k / 2; + const bool k_high = (k & 1) != 0; + for (size_t o = 0; o < out_features; ++o) { + const std::uint8_t src_byte = src[o * in_half + k_byte]; + const std::uint8_t nibble = k_high ? ((src_byte >> 4) & 0x0f) : (src_byte & 0x0f); + std::uint8_t &dst_byte = dst[k * (out_features / 2) + (o / 2)]; + if ((o & 1) == 0) { + dst_byte = static_cast((dst_byte & 0xf0) | nibble); + } else { + dst_byte = static_cast((dst_byte & 0x0f) | (nibble << 4)); + } + } + } + return dst_cpu->to(device); +} + +infinicore::Tensor normalize_scale(const infinicore::Tensor &scale, size_t out_features, const infinicore::Device &device) { + if (scale->dtype() != infinicore::DataType::F32) { + throw std::runtime_error("GlmW4A8 expects float32 weight_scale"); + } + if (scale->ndim() == 2 && scale->size(0) == out_features && scale->size(1) == 1) { + return scale->is_contiguous() ? scale : scale->contiguous(); + } + if (scale->ndim() == 2 && scale->size(0) == 1 && scale->size(1) == out_features) { + auto cpu = scale->to(infinicore::Device::cpu())->contiguous(); + auto fixed_cpu = infinicore::Tensor::empty({out_features, 1}, infinicore::DataType::F32, infinicore::Device::cpu()); + auto *src = reinterpret_cast(cpu->data()); + auto *dst = reinterpret_cast(fixed_cpu->data()); + for (size_t i = 0; i < out_features; ++i) { + dst[i] = src[i]; + } + return fixed_cpu->to(device); + } + throw std::runtime_error("GlmW4A8 expects weight_scale shape [out,1] or [1,out]"); +} + +infinicore::Tensor run_w4a8_forward(const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha) { + auto weight = params.at("weight"); + auto weight_scale = params.at("weight_scale"); + if (weight->ndim() != 2 || weight->dtype() != infinicore::DataType::I8) { + throw std::runtime_error("GlmW4A8Runtime expects int8 runtime weight [in,out/2]"); + } + const size_t in_features = weight->size(0); + const size_t out_features = weight->size(1) * 2; + if (weight_scale->ndim() != 2 || weight_scale->size(0) != out_features || weight_scale->size(1) != 1) { + throw std::runtime_error("GlmW4A8Runtime expects weight_scale [out,1]"); + } + + std::optional bias_opt; + if (has_bias) { + bias_opt = params.at("bias"); + } + + auto effective_weight_scale = weight_scale; + if (std::fabs(alpha - 1.0f) > 1e-7f) { + effective_weight_scale = infinicore::op::mul_scalar(weight_scale, static_cast(alpha)); + } + + auto x = input->is_contiguous() ? input : input->contiguous(); + if (x->size(x->ndim() - 1) != in_features) { + throw std::runtime_error("GlmW4A8Runtime input hidden size mismatch"); + } + std::vector original_shape = x->shape(); + size_t m = 1; + for (size_t i = 0; i + 1 < original_shape.size(); ++i) { + m *= original_shape[i]; + } + auto x2d = (x->ndim() == 2) ? x : x->view({m, in_features}); + + auto x_i8 = infinicore::Tensor::empty({m, in_features}, infinicore::DataType::I8, x->device()); + auto x_scale = infinicore::Tensor::empty({m, 1}, infinicore::DataType::F32, x->device()); + infinicore::op::dynamic_scaled_int8_quant_(x_i8, x2d, x_scale); + + auto out2d = infinicore::Tensor::empty({m, out_features}, x->dtype(), x->device()); + infinicore::op::scaled_mm_w4a8_(out2d, x_i8, weight, x_scale, effective_weight_scale, bias_opt, false); + + if (original_shape.size() == 2) { + return out2d; + } + original_shape.back() = out_features; + return out2d->view(original_shape); +} + +std::vector split_runtime_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int tp_rank, int tp_size, int tp_num_heads) { + std::vector result; + auto w_it = params.find("weight"); + auto s_it = params.find("weight_scale"); + auto b_it = params.find("bias"); + for (const auto &s : splits) { + if ((s.start % 2) != 0 || (s.size % 2) != 0) { + throw std::runtime_error("GlmW4A8Runtime split requires even output offsets/sizes"); + } + result.push_back({s.prefix + ".weight", + infinicore::nn::Parameter( + w_it->second->narrow({{1, s.start / 2, s.size / 2}}), + 1, tp_rank, tp_size, s.num_shards)}); + result.push_back({s.prefix + ".weight_scale", + infinicore::nn::Parameter( + s_it->second->narrow({{0, s.start, s.size}}), + 0, tp_rank, tp_size, s.num_shards)}); + if (b_it != params.end()) { + result.push_back({s.prefix + ".bias", + infinicore::nn::Parameter( + b_it->second->narrow({{0, s.start, s.size}}), + 0, tp_rank, tp_size, s.num_shards)}); + } + } + (void)tp_num_heads; + return result; +} + +} // namespace + +std::vector GlmW4A8::get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int /*tp_num_heads*/, + const infinicore::DataType &dtype, + bool bias) const { + if ((in_features % 2) != 0) { + throw std::runtime_error("GlmW4A8 requires even in_features"); + } + std::vector descs; + // Checkpoint layout. ColumnParallel splits output dim0; RowParallel splits input packed dim1. + int weight_split_dim = split_dim; + descs.push_back({"weight", {out_features, in_features / 2}, infinicore::DataType::I8, weight_split_dim, tp_rank, tp_size}); + int scale_split_dim = (split_dim == 0) ? 0 : -1; + descs.push_back({"weight_scale", {out_features, 1}, infinicore::DataType::F32, scale_split_dim, scale_split_dim == 0 ? tp_rank : 0, scale_split_dim == 0 ? tp_size : 1}); + if (bias) { + descs.push_back({"bias", {out_features}, dtype, -1, 0, 1}); + } + return descs; +} + +infinicore::Tensor GlmW4A8::forward(const ParamsMap &, const infinicore::Tensor &, bool, float) const { + throw std::runtime_error("GlmW4A8 must be converted to runtime layout after loading before forward"); +} + +std::vector GlmW4A8::split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int narrow_dim, + int tp_rank, int tp_size, int tp_num_heads) const { + // Loader/checkpoint layout: output dimension is dim0 for both weight and scale. + std::vector result; + auto w_it = params.find("weight"); + auto s_it = params.find("weight_scale"); + auto b_it = params.find("bias"); + (void)narrow_dim; + for (const auto &s : splits) { + result.push_back({s.prefix + ".weight", + infinicore::nn::Parameter(w_it->second->narrow({{0, s.start, s.size}}), 0, tp_rank, tp_size, s.num_shards)}); + result.push_back({s.prefix + ".weight_scale", + infinicore::nn::Parameter(s_it->second->narrow({{0, s.start, s.size}}), 0, tp_rank, tp_size, s.num_shards)}); + if (b_it != params.end()) { + result.push_back({s.prefix + ".bias", + infinicore::nn::Parameter(b_it->second->narrow({{0, s.start, s.size}}), 0, tp_rank, tp_size, s.num_shards)}); + } + } + (void)tp_num_heads; + return result; +} + +std::shared_ptr GlmW4A8::process_weights_after_loading( + ParamsMap ¶ms, + const infinicore::Device &device, + int /*split_dim*/) const { + auto weight = params.at("weight"); + const size_t out_features = weight->size(0); + params["weight"] = repack_out_khalf_to_k_outhalf_cpu(weight, device); + params["weight_scale"] = normalize_scale(params.at("weight_scale"), out_features, device); + return std::make_shared(get_config()); +} + +std::vector GlmW4A8Runtime::get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int /*tp_num_heads*/, + const infinicore::DataType &dtype, + bool bias) const { + if ((out_features % 2) != 0) { + throw std::runtime_error("GlmW4A8Runtime requires even out_features"); + } + std::vector descs; + int weight_split_dim = split_dim >= 0 ? (1 - split_dim) : -1; + descs.push_back({"weight", {in_features, out_features / 2}, infinicore::DataType::I8, weight_split_dim, tp_rank, tp_size}); + int scale_split_dim = (split_dim == 0) ? 0 : -1; + descs.push_back({"weight_scale", {out_features, 1}, infinicore::DataType::F32, scale_split_dim, scale_split_dim == 0 ? tp_rank : 0, scale_split_dim == 0 ? tp_size : 1}); + if (bias) { + descs.push_back({"bias", {out_features}, dtype, -1, 0, 1}); + } + return descs; +} + +infinicore::Tensor GlmW4A8Runtime::forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha) const { + return run_w4a8_forward(params, input, has_bias, alpha); +} + +std::vector GlmW4A8Runtime::split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int /*narrow_dim*/, + int tp_rank, int tp_size, int tp_num_heads) const { + return split_runtime_params(params, splits, tp_rank, tp_size, tp_num_heads); +} + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/glm_w4a8.hpp b/csrc/layers/quantization/glm_w4a8.hpp new file mode 100644 index 000000000..df19f5f1f --- /dev/null +++ b/csrc/layers/quantization/glm_w4a8.hpp @@ -0,0 +1,77 @@ +#pragma once +#include "base_quantization.hpp" + +namespace infinilm::quantization { + +class GlmW4A8Runtime; + +// Loader-side GLM W4A8 quantization. Checkpoint weight layout is [out, in/2] +// with two 4-bit signed weights packed in one int8 byte along K. +class GlmW4A8 : public BaseQuantization { +public: + explicit GlmW4A8(const nlohmann::json &quant_config) + : BaseQuantization(quant_config) {} + + QuantScheme get_quant_scheme() const override { return QuantScheme::GLM_W4A8; } + + std::vector get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int tp_num_heads, + const infinicore::DataType &dtype, + bool bias) const override; + + // Initial checkpoint layout has output dimension on dim0. + int get_fused_split_dim() const override { return 0; } + + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha = 1.0f) const override; + + std::vector split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int narrow_dim, + int tp_rank, int tp_size, int tp_num_heads) const override; + + std::shared_ptr process_weights_after_loading( + ParamsMap ¶ms, + const infinicore::Device &device, + int split_dim = -1) const override; +}; + +// Runtime GLM W4A8 layout for CUINFER scaled_mm_w4a8: weight [in, out/2]. +class GlmW4A8Runtime : public BaseQuantization { +public: + explicit GlmW4A8Runtime(const nlohmann::json &quant_config) + : BaseQuantization(quant_config) {} + + QuantScheme get_quant_scheme() const override { return QuantScheme::GLM_W4A8_RUNTIME; } + + std::vector get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int tp_num_heads, + const infinicore::DataType &dtype, + bool bias) const override; + + // Runtime packed layout has logical output dimension on dim1, scaled by 2. + int get_fused_split_dim() const override { return 1; } + size_t get_logical_dim_size(size_t raw_size) const override { return raw_size * 2; } + + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha = 1.0f) const override; + + std::vector split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int narrow_dim, + int tp_rank, int tp_size, int tp_num_heads) const override; +}; + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/glm_w8a8.cpp b/csrc/layers/quantization/glm_w8a8.cpp new file mode 100644 index 000000000..6ccc75b6f --- /dev/null +++ b/csrc/layers/quantization/glm_w8a8.cpp @@ -0,0 +1,92 @@ +#include "glm_w8a8.hpp" + +#include "infinicore/ops/dynamic_scaled_int8_quant.hpp" +#include "infinicore/ops/mul_scalar.hpp" +#include "infinicore/ops/scaled_mm_w8a8.hpp" + +#include +#include +#include + +namespace infinilm::quantization { + +std::vector GlmW8A8::get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int /*tp_num_heads*/, const infinicore::DataType &dtype, + bool bias) const { + std::vector result; + result.push_back({"weight", {out_features, in_features}, infinicore::DataType::I8, split_dim, tp_rank, tp_size}); + const int scale_split_dim = split_dim == 0 ? 0 : -1; + result.push_back({"weight_scale", {out_features, 1}, infinicore::DataType::F32, scale_split_dim, scale_split_dim == 0 ? tp_rank : 0, scale_split_dim == 0 ? tp_size : 1}); + if (bias) { + result.push_back({"bias", {out_features}, dtype, split_dim == 0 ? 0 : -1, split_dim == 0 ? tp_rank : 0, split_dim == 0 ? tp_size : 1}); + } + return result; +} + +infinicore::Tensor GlmW8A8::forward( + const ParamsMap ¶ms, const infinicore::Tensor &input, + bool has_bias, float alpha) const { + auto weight = params.at("weight"); + auto weight_scale = params.at("weight_scale"); + if (weight->ndim() != 2 || weight->dtype() != infinicore::DataType::I8) { + throw std::runtime_error("GlmW8A8 expects int8 weight [out,in]"); + } + const size_t out_features = weight->size(0), in_features = weight->size(1); + if (weight_scale->ndim() != 2 || weight_scale->size(0) != out_features || weight_scale->size(1) != 1 || weight_scale->dtype() != infinicore::DataType::F32) { + throw std::runtime_error("GlmW8A8 expects float32 weight_scale [out,1]"); + } + + auto x = input->is_contiguous() ? input : input->contiguous(); + if (x->size(x->ndim() - 1) != in_features) { + throw std::runtime_error("GlmW8A8 input hidden size mismatch"); + } + auto shape = x->shape(); + size_t m = 1; + for (size_t i = 0; i + 1 < shape.size(); ++i) { + m *= shape[i]; + } + auto x2d = x->ndim() == 2 ? x : x->view({m, in_features}); + auto x_i8 = infinicore::Tensor::empty({m, in_features}, infinicore::DataType::I8, x->device()); + auto x_scale = infinicore::Tensor::empty({m, 1}, infinicore::DataType::F32, x->device()); + infinicore::op::dynamic_scaled_int8_quant_(x_i8, x2d, x_scale); + + auto effective_scale = weight_scale; + if (std::fabs(alpha - 1.0f) > 1e-7f) { + effective_scale = infinicore::op::mul_scalar(weight_scale, static_cast(alpha)); + } + std::optional bias; + if (has_bias) { + bias = params.at("bias"); + } + auto out = infinicore::Tensor::empty({m, out_features}, x->dtype(), x->device()); + infinicore::op::scaled_mm_w8a8_(out, x_i8, weight, x_scale, effective_scale, bias, true); + if (shape.size() == 2) { + return out; + } + shape.back() = out_features; + return out->view(shape); +} + +std::vector GlmW8A8::split_params( + const std::unordered_map ¶ms, + const std::vector &splits, int /*narrow_dim*/, + int tp_rank, int tp_size, int /*tp_num_heads*/) const { + std::vector result; + const auto &weight = params.at("weight"); + const auto &scale = params.at("weight_scale"); + auto bias = params.find("bias"); + for (const auto &s : splits) { + result.push_back({s.prefix + ".weight", infinicore::nn::Parameter( + weight->narrow({{0, s.start, s.size}}), 0, tp_rank, tp_size, s.num_shards)}); + result.push_back({s.prefix + ".weight_scale", infinicore::nn::Parameter( + scale->narrow({{0, s.start, s.size}}), 0, tp_rank, tp_size, s.num_shards)}); + if (bias != params.end()) { + result.push_back({s.prefix + ".bias", infinicore::nn::Parameter( + bias->second->narrow({{0, s.start, s.size}}), 0, tp_rank, tp_size, s.num_shards)}); + } + } + return result; +} +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/glm_w8a8.hpp b/csrc/layers/quantization/glm_w8a8.hpp new file mode 100644 index 000000000..d9164a616 --- /dev/null +++ b/csrc/layers/quantization/glm_w8a8.hpp @@ -0,0 +1,31 @@ +#pragma once +#include "base_quantization.hpp" + +namespace infinilm::quantization { + +// Dense GLM W8A8 checkpoint/runtime layout. We retain [out, in] so CUINFER +// scaled_mm can consume checkpoint shards directly with trans_weight=true. +class GlmW8A8 final : public BaseQuantization { +public: + explicit GlmW8A8(const nlohmann::json &quant_config) + : BaseQuantization(quant_config) {} + + QuantScheme get_quant_scheme() const override { return QuantScheme::GLM_W8A8; } + int get_fused_split_dim() const override { return 0; } + + std::vector get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int tp_num_heads, const infinicore::DataType &dtype, + bool bias) const override; + + infinicore::Tensor forward( + const ParamsMap ¶ms, const infinicore::Tensor &input, + bool has_bias, float alpha = 1.0f) const override; + + std::vector split_params( + const std::unordered_map ¶ms, + const std::vector &splits, int narrow_dim, + int tp_rank, int tp_size, int tp_num_heads) const override; +}; +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/quantization.hpp b/csrc/layers/quantization/quantization.hpp index 4fcd9e61f..217468078 100644 --- a/csrc/layers/quantization/quantization.hpp +++ b/csrc/layers/quantization/quantization.hpp @@ -4,6 +4,8 @@ #include "awq_marlin.hpp" #include "base_quantization.hpp" #include "compressed_tensors.hpp" +#include "glm_w4a8.hpp" +#include "glm_w8a8.hpp" #include "gptq.hpp" #include "gptq_marlin.hpp" #include "gptq_qy.hpp" diff --git a/csrc/layers/quantization/quantization_scheme.hpp b/csrc/layers/quantization/quantization_scheme.hpp index 4ab57c71a..b40ad4c65 100644 --- a/csrc/layers/quantization/quantization_scheme.hpp +++ b/csrc/layers/quantization/quantization_scheme.hpp @@ -10,6 +10,9 @@ enum class QuantScheme { GPTQ_W4A16_QY, GPTQ_W4A16, GPTQ_MARLIN_W4A16, + GLM_W8A8, + GLM_W4A8, + GLM_W4A8_RUNTIME, }; enum class KVQuantAlgo { diff --git a/csrc/models/glm_moe_dsa/glm_attention.cpp b/csrc/models/glm_moe_dsa/glm_attention.cpp new file mode 100644 index 000000000..9d010c759 --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_attention.cpp @@ -0,0 +1,63 @@ +#include "glm_attention.hpp" +#include "../../global_state/global_state.hpp" +#include "../../layers/rotary_embedding/rotary_embedding.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/ops/broadcast_to.hpp" +#include "infinicore/ops/cat.hpp" +#include +#include +#include +namespace infinilm::models::glm_moe_dsa { +GlmAttention::GlmAttention(std::shared_ptr c, size_t layer, const infinicore::Device &d) { + qn_ = c->get("qk_nope_head_dim"); + qr_ = c->get("qk_rope_head_dim"); + qh_ = qn_ + qr_; + vh_ = c->get("v_head_dim"); + ql_ = c->get("q_lora_rank"); + kvl_ = c->get("kv_lora_rank"); + auto total = c->get("num_attention_heads"); + auto &r = infinilm::global_state::get_tensor_model_parallel_rank_info(); + heads_ = total / r.tp_size; + scale_ = 1 / std::sqrt(float(qh_)); + auto q = c->get_quantization_method(); + auto dt = c->get_dtype(); + auto h = c->get("hidden_size"); + auto eps = c->get("rms_norm_eps"); + INFINICORE_NN_MODULE_INIT(q_a_proj, h, ql_, q, false, dt, d); + INFINICORE_NN_MODULE_INIT(q_a_layernorm, ql_, eps, dt, d); + INFINICORE_NN_MODULE_INIT(q_b_proj, ql_, total * qh_, q, false, dt, d, r.tp_rank, r.tp_size); + INFINICORE_NN_MODULE_INIT(kv_a_proj_with_mqa, h, kvl_ + qr_, q, false, dt, d); + INFINICORE_NN_MODULE_INIT(kv_a_layernorm, kvl_, eps, dt, d); + INFINICORE_NN_MODULE_INIT(kv_b_proj, kvl_, total * (qn_ + vh_), q, false, dt, d, r.tp_rank, r.tp_size); + INFINICORE_NN_MODULE_INIT(o_proj, total * vh_, h, q, false, dt, d, r.tp_rank, r.tp_size, r.comm); + rope_ = infinilm::layers::rotary_embedding::get_rope(c, d); + auto backend = infinilm::global_state::get_infinilm_config().attention_backend; + attn_ = std::make_shared(heads_, qh_, scale_, heads_, layer, ks_, vs_, backend); + infinilm::layers::attention::init_kv_cache_quant_params([this](const std::string &n, infinicore::nn::Parameter p) { register_parameter(n, std::move(p)); }, d, ks_, vs_); +} +infinicore::Tensor GlmAttention::forward(const infinicore::Tensor &pos, const infinicore::Tensor &x) const { + auto sh = x->shape(); + size_t b = sh[0], s = sh[1], m = b * s; + auto xm = x; + auto qa = q_a_proj_->forward(xm); + auto qan = q_a_layernorm_->forward(qa); + auto q = q_b_proj_->forward(qan)->view({m, heads_, qh_}); + auto qn = q->narrow({{2, 0, qn_}})->contiguous(), qp = q->narrow({{2, qn_, qr_}})->contiguous(); + auto kva = kv_a_proj_with_mqa_->forward(xm)->view({m, kvl_ + qr_}); + auto kvc = kva->narrow({{1, 0, kvl_}})->contiguous(); + auto kp = kva->narrow({{1, kvl_, qr_}})->contiguous(); + auto kvn = kv_a_layernorm_->forward(kvc); + auto kv = kv_b_proj_->forward(kvn)->view({m, heads_, qn_ + vh_}); + auto kn = kv->narrow({{2, 0, qn_}})->contiguous(); + auto v = kv->narrow({{2, qn_, vh_}})->contiguous(); + auto p = pos->contiguous()->view({m}); + qp = rope_->forward(qp, p, true); + auto kpr = rope_->forward(kp->view({m, 1, qr_}), p, true); + kp = infinicore::op::broadcast_to(kpr, {static_cast(m), static_cast(heads_), static_cast(qr_)})->contiguous(); + auto qs = infinicore::op::cat({qn, qp}, 2); + auto k = infinicore::op::cat({kn, kp}, 2); + auto vv = v; + auto out = attn_->forward(qs, k, vv)->view({b, s, heads_ * vh_}); + return o_proj_->forward(out); +} +} // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_attention.hpp b/csrc/models/glm_moe_dsa/glm_attention.hpp new file mode 100644 index 000000000..4e07f8a46 --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_attention.hpp @@ -0,0 +1,28 @@ +#pragma once +#include "../../config/model_config.hpp" +#include "../../layers/attention/attention.hpp" +#include "../../layers/linear/linear.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/nn/rmsnorm.hpp" +#include "infinicore/nn/rope.hpp" +namespace infinilm::models::glm_moe_dsa { +class GlmAttention final : public infinicore::nn::Module { +public: + GlmAttention(std::shared_ptr, size_t, const infinicore::Device &); + infinicore::Tensor forward(const infinicore::Tensor &, const infinicore::Tensor &) const; + +private: + size_t heads_{0}, qh_{0}, qn_{0}, qr_{0}, vh_{0}, ql_{0}, kvl_{0}; + float scale_{1}; + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, q_a_proj); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, q_a_layernorm); + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_b_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, kv_a_proj_with_mqa); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, kv_a_layernorm); + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, kv_b_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, o_proj); + std::shared_ptr rope_; + std::shared_ptr attn_; + infinicore::nn::Parameter ks_, vs_; +}; +} // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_model.cpp b/csrc/models/glm_moe_dsa/glm_model.cpp new file mode 100644 index 000000000..331abae70 --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_model.cpp @@ -0,0 +1,74 @@ +#include "glm_model.hpp" +#include "../../global_state/global_state.hpp" +#include "../models_registry.hpp" +#include "infinicore/ops.hpp" +#include +#include +#include +namespace infinilm::models::glm_moe_dsa { +GlmDecoder::GlmDecoder(std::shared_ptr c, size_t i, const infinicore::Device &d) { + auto h = c->get("hidden_size"); + auto e = c->get("rms_norm_eps"); + auto dt = c->get_dtype(); + INFINICORE_NN_MODULE_INIT(input_layernorm, h, e, dt, d); + INFINICORE_NN_MODULE_INIT(post_attention_layernorm, h, e, dt, d); + INFINICORE_NN_MODULE_INIT(self_attn, c, i, d); + moe_ = i >= c->get_or("first_k_dense_replace", 0); + if (moe_) { + moe_mlp_ = register_module("mlp", c, d); + } else { + dense_mlp_ = register_module("mlp", c, d); + } +} +void GlmDecoder::forward(const infinicore::Tensor &p, infinicore::Tensor &x, infinicore::Tensor &r) const { + input_layernorm_->forward_inplace(x, r); + x = self_attn_->forward(p, x); + post_attention_layernorm_->forward_inplace(x, r); + x = moe_ ? moe_mlp_->forward(x) : dense_mlp_->forward(x); +} +GlmModel::GlmModel(std::shared_ptr c, const infinicore::Device &d) { + auto dt = c->get_dtype(); + auto h = c->get("hidden_size"); + INFINICORE_NN_MODULE_INIT(embed_tokens, c->get("vocab_size"), h, dt, d); + for (size_t i = 0; i < c->get("num_hidden_layers"); ++i) { + layers_.push_back(register_module("layers." + std::to_string(i), c, i, d)); + } + INFINICORE_NN_MODULE_INIT(norm, h, c->get("rms_norm_eps"), dt, d); +} +infinicore::Tensor GlmModel::forward(const infinilm::InfinilmModel::Input &i) const { + auto x = embed_tokens_->forward(i.input_ids.value()); + infinicore::Tensor r; + for (auto &l : layers_) { + l->forward(i.position_ids.value(), x, r); + } + norm_->forward_inplace(x, r); + return x; +} +GlmForCausalLM::GlmForCausalLM(std::shared_ptr c, const infinicore::Device &d) { + model_config_ = c; + INFINICORE_NN_MODULE_INIT(model, c, d); + INFINICORE_NN_MODULE_INIT(lm_head, c->get("hidden_size"), c->get("vocab_size"), false, c->get_dtype(), d); +} +infinilm::InfinilmModel::Output GlmForCausalLM::forward(const Input &i) const { + auto x = model_->forward(i); + return {lm_head_->forward(x)}; +} +std::shared_ptr create_glm_config(std::shared_ptr c) { + auto j = c->get_config_json(); + auto qh = j.at("qk_nope_head_dim").get() + j.at("qk_rope_head_dim").get(); + j["head_dim"] = qh; + if (!j.contains("rope_theta") && j.contains("rope_parameters")) { + j["rope_theta"] = j["rope_parameters"].value("rope_theta", 10000.0); + } + j["partial_rotary_factor"] = double(j.at("qk_rope_head_dim").get()) / double(qh); + j["num_experts"] = j.at("n_routed_experts"); + j["mlp_bias"] = false; + j["quantization_config"] = {{"quant_method", "glm_w8a8"}}; + auto n = std::make_shared(j); + n->set_rope_algo(infinicore::nn::RoPE::Algo::GPT_J); + return n; +} +} // namespace infinilm::models::glm_moe_dsa +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL(glm_moe_dsa, infinilm::models::glm_moe_dsa::GlmForCausalLM, infinilm::models::glm_moe_dsa::create_glm_config); +} diff --git a/csrc/models/glm_moe_dsa/glm_model.hpp b/csrc/models/glm_moe_dsa/glm_model.hpp new file mode 100644 index 000000000..f052f140f --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_model.hpp @@ -0,0 +1,44 @@ +#pragma once +#include "../../layers/linear/linear.hpp" +#include "../../layers/mlp/mlp.hpp" +#include "../infinilm_model.hpp" +#include "glm_attention.hpp" +#include "glm_moe.hpp" +#include "glm_vocab_parallel.hpp" +#include "infinicore/nn/embedding.hpp" +#include "infinicore/nn/rmsnorm.hpp" +namespace infinilm::models::glm_moe_dsa { +class GlmDecoder final : public infinicore::nn::Module { +public: + GlmDecoder(std::shared_ptr, size_t, const infinicore::Device &); + void forward(const infinicore::Tensor &, infinicore::Tensor &, infinicore::Tensor &) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); + INFINICORE_NN_MODULE(GlmAttention, self_attn); + INFINICORE_NN_MODULE(GlmDenseMLP, dense_mlp); + INFINICORE_NN_MODULE(GlmMoE, moe_mlp); + bool moe_{false}; +}; +class GlmModel final : public infinicore::nn::Module { +public: + GlmModel(std::shared_ptr, const infinicore::Device &); + infinicore::Tensor forward(const infinilm::InfinilmModel::Input &) const; + +private: + INFINICORE_NN_MODULE(GlmVocabEmbedding, embed_tokens); + INFINICORE_NN_MODULE_VEC(GlmDecoder, layers); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); +}; +class GlmForCausalLM final : public infinilm::InfinilmModel { +public: + GlmForCausalLM(std::shared_ptr, const infinicore::Device &); + Output forward(const Input &) const override; + +private: + INFINICORE_NN_MODULE(GlmModel, model); + INFINICORE_NN_MODULE(GlmVocabLMHead, lm_head); +}; +std::shared_ptr create_glm_config(std::shared_ptr); +} // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_moe.cpp b/csrc/models/glm_moe_dsa/glm_moe.cpp new file mode 100644 index 000000000..63d1ad6f8 --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_moe.cpp @@ -0,0 +1,117 @@ +#include "glm_moe.hpp" +#include "../../global_state/global_state.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/ops/cast.hpp" +#include "infinicore/ops/distributed/allreduce.hpp" +#include "infinicore/ops/grouped_topk_vllm.hpp" +#include "infinicore/ops/moe_argsort_bincount.hpp" +#include "infinicore/ops/moe_expand_input.hpp" +#include "infinicore/ops/moe_silu_and_mul_quant.hpp" +#include "infinicore/ops/moe_sum_vllm.hpp" +#include "infinicore/ops/w4a8_group_gemm.hpp" +#include +namespace infinilm::models::glm_moe_dsa { +GlmTopKRouter::GlmTopKRouter(std::shared_ptr c, const infinicore::Device &d) { + auto h = c->get("hidden_size"); + num_experts_ = c->get("num_experts"); + top_k_ = c->get("num_experts_per_tok"); + num_expert_group_ = c->get_or_alias("num_expert_group", "n_group", 1); + topk_group_ = c->get_or("topk_group", 1); + renormalize_ = c->get_or("norm_topk_prob", true); + routed_scaling_factor_ = c->get_or("routed_scaling_factor", 1); + if (!num_experts_ || !top_k_ || top_k_ > num_experts_) { + throw std::runtime_error("GlmTopKRouter: invalid config"); + } + INFINICORE_NN_PARAMETER_INIT(weight, ({num_experts_, h}, c->get_dtype(), d)); + INFINICORE_NN_PARAMETER_INIT(e_score_correction_bias, ({num_experts_}, infinicore::DataType::F32, d)); +} +void GlmTopKRouter::process_weights_after_loading() { + runtime_bias_ = infinicore::Tensor::empty({num_experts_}, weight_->dtype(), weight_->device()); + infinicore::op::cast_(runtime_bias_, e_score_correction_bias_); +} +std::tuple GlmTopKRouter::forward(const infinicore::Tensor &x) const { + auto logits = infinicore::op::linear(x, weight_, std::nullopt, 1); + auto w = infinicore::Tensor::empty({x->size(0), top_k_}, infinicore::DataType::F32, x->device()); + auto ids = infinicore::Tensor::empty({x->size(0), top_k_}, infinicore::DataType::I32, x->device()); + infinicore::op::grouped_topk_vllm_(w, ids, logits, num_expert_group_, topk_group_, renormalize_, routed_scaling_factor_, runtime_bias_, "sigmoid"); + + return {w, ids}; +} +GlmW4A8Experts::GlmW4A8Experts(std::shared_ptr c, const infinicore::Device &d) { + hidden_ = c->get("hidden_size"); + nexpert_ = c->get("num_experts"); + topk_ = c->get("num_experts_per_tok"); + auto &r = infinilm::global_state::get_tensor_model_parallel_rank_info(); + tp_ = r.tp_size; + comm_ = r.comm; + auto full_i = c->get("moe_intermediate_size"); + inter_ = full_i / tp_; + size_t pi = inter_ / 2; + w1_ = infinicore::Tensor::empty({nexpert_, inter_ * 2, hidden_ / 2}, infinicore::DataType::I8, d); + s1_ = infinicore::Tensor::empty({nexpert_, inter_ * 2, 1}, infinicore::DataType::F32, d); + w2_ = infinicore::Tensor::empty({nexpert_, hidden_, pi}, infinicore::DataType::I8, d); + s2_ = infinicore::Tensor::empty({nexpert_, hidden_, 1}, infinicore::DataType::F32, d); + for (size_t i = 0; i < nexpert_; ++i) { + auto p = std::to_string(i); + auto gw = w1_->narrow({{0, i, 1}, {1, 0, inter_}, {2, 0, hidden_ / 2}})->view({inter_, hidden_ / 2}); + auto uw = w1_->narrow({{0, i, 1}, {1, inter_, inter_}, {2, 0, hidden_ / 2}})->view({inter_, hidden_ / 2}); + auto dw = w2_->narrow({{0, i, 1}, {1, 0, hidden_}, {2, 0, pi}})->view({hidden_, pi}); + auto gs = s1_->narrow({{0, i, 1}, {1, 0, inter_}, {2, 0, 1}})->view({1, inter_}); + auto us = s1_->narrow({{0, i, 1}, {1, inter_, inter_}, {2, 0, 1}})->view({1, inter_}); + auto ds = s2_->narrow({{0, i, 1}, {1, 0, hidden_}, {2, 0, 1}})->view({1, hidden_}); + register_parameter(p + ".gate_proj.weight", infinicore::nn::Parameter(gw, 0, r.tp_rank, r.tp_size)); + register_parameter(p + ".up_proj.weight", infinicore::nn::Parameter(uw, 0, r.tp_rank, r.tp_size)); + register_parameter(p + ".down_proj.weight", infinicore::nn::Parameter(dw, 1, r.tp_rank, r.tp_size)); + register_parameter(p + ".gate_proj.weight_scale", infinicore::nn::Parameter(gs, 1, r.tp_rank, r.tp_size)); + register_parameter(p + ".up_proj.weight_scale", infinicore::nn::Parameter(us, 1, r.tp_rank, r.tp_size)); + register_parameter(p + ".down_proj.weight_scale", infinicore::nn::Parameter(ds)); + } +} +infinicore::Tensor GlmW4A8Experts::forward(const infinicore::Tensor &x, const infinicore::Tensor &ids, const infinicore::Tensor &tw) const { + if (!w1_) { + throw std::runtime_error("GlmW4A8Experts: weights not ready"); + } + size_t m = x->size(0), total = m * topk_; + bool dec = m == 1; + int64_t fmt = dec ? 2 : 1; + auto cnt = infinicore::Tensor::empty({nexpert_}, infinicore::DataType::I32, x->device()), sorted = infinicore::Tensor::empty({total}, infinicore::DataType::I32, x->device()), inv = infinicore::Tensor::empty({total}, infinicore::DataType::I32, x->device()); + infinicore::op::moe_argsort_bincount_with_inv_pos_(cnt, sorted, inv, ids, nexpert_); + auto gc = dec ? cnt : cnt->to(infinicore::Device::Type::CPU); + auto a1 = infinicore::Tensor::empty({total, hidden_}, infinicore::DataType::I8, x->device()), a1s = infinicore::Tensor::empty({total, 1}, infinicore::DataType::F32, x->device()); + infinicore::op::moe_expand_input_with_inv_pos_(a1, a1s, x, inv, topk_, 128, fmt); + auto a2 = infinicore::Tensor::empty({total, inter_ * 2}, x->dtype(), x->device()); + infinicore::op::w4a8_group_gemm_(a2, a1, w1_, a1s, s1_, gc, std::nullopt, std::nullopt, true, dec); + auto a2q = infinicore::Tensor::empty({total, inter_}, infinicore::DataType::I8, x->device()), a2s = infinicore::Tensor::empty({total, 1}, infinicore::DataType::F32, x->device()); + infinicore::op::moe_silu_and_mul_quant_(a2q, a2s, a2, fmt); + auto a3 = infinicore::Tensor::empty({total, hidden_}, x->dtype(), x->device()); + infinicore::op::w4a8_group_gemm_(a3, a2q, w2_, a2s, s2_, gc, sorted, std::nullopt, true, dec); + auto out = infinicore::Tensor::empty({m, hidden_}, x->dtype(), x->device()); + infinicore::op::moe_sum_vllm_(out, a3->view({m, topk_, hidden_}), tw); + if (tp_ > 1 && comm_) { + infinicore::op::distributed::allreduce_(out, out, INFINICCL_SUM, comm_); + } + return out; +} +GlmMoE::GlmMoE(std::shared_ptr c, const infinicore::Device &d) { + INFINICORE_NN_MODULE_INIT(gate, c, d); + INFINICORE_NN_MODULE_INIT(experts, c, d); + auto n = c->get_or("n_shared_experts", 0); + shared_ = n > 0; + if (shared_) { + auto j = c->get_config_json(); + j["intermediate_size"] = c->get("moe_intermediate_size") * n; + auto sc = std::make_shared(j); + INFINICORE_NN_MODULE_INIT(shared_experts, sc, d); + } +} +infinicore::Tensor GlmMoE::forward(const infinicore::Tensor &x) const { + auto s = x->shape(); + auto f = x->view({s[0] * s[1], s[2]}); + auto [w, i] = gate_->forward(f); + auto r = experts_->forward(f, i, w)->view(s); + if (!shared_) { + return r; + } + return infinicore::op::add(r, shared_experts_->forward(x)); +} +} // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_moe.hpp b/csrc/models/glm_moe_dsa/glm_moe.hpp new file mode 100644 index 000000000..887a95a4e --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_moe.hpp @@ -0,0 +1,49 @@ +#pragma once +#include "../../config/model_config.hpp" +#include "../../layers/mlp/mlp.hpp" +#include "../../layers/moe/legacy/moe_mlp.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" +#include +#include +#include +#include +namespace infinilm::models::glm_moe_dsa { +using GlmDenseMLP = infinilm::layers::mlp::MLP; +using GlmExpertMLP = infinilm::layers::moe::legacy::MoeMLP; +class GlmTopKRouter final : public infinicore::nn::Module { +public: + GlmTopKRouter(std::shared_ptr, const infinicore::Device &); + std::tuple forward(const infinicore::Tensor &) const; + void process_weights_after_loading() override; + +private: + INFINICORE_NN_PARAMETER(weight); + INFINICORE_NN_PARAMETER(e_score_correction_bias); + infinicore::Tensor runtime_bias_; + size_t num_experts_{0}, top_k_{0}, num_expert_group_{0}, topk_group_{0}; + bool renormalize_{false}; + float routed_scaling_factor_{1}; +}; +class GlmW4A8Experts final : public infinicore::nn::Module { +public: + GlmW4A8Experts(std::shared_ptr, const infinicore::Device &); + infinicore::Tensor forward(const infinicore::Tensor &, const infinicore::Tensor &, const infinicore::Tensor &) const; + +private: + infinicore::Tensor w1_, s1_, w2_, s2_; + size_t hidden_{0}, inter_{0}, nexpert_{0}, topk_{0}, tp_{1}; + infinicclComm_t comm_{nullptr}; +}; +class GlmMoE final : public infinicore::nn::Module { +public: + GlmMoE(std::shared_ptr, const infinicore::Device &); + infinicore::Tensor forward(const infinicore::Tensor &) const; + +private: + INFINICORE_NN_MODULE(GlmTopKRouter, gate); + INFINICORE_NN_MODULE(GlmW4A8Experts, experts); + INFINICORE_NN_MODULE(GlmDenseMLP, shared_experts); + bool shared_{false}; +}; +} // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_vocab_parallel.cpp b/csrc/models/glm_moe_dsa/glm_vocab_parallel.cpp new file mode 100644 index 000000000..cc1fc2a01 --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_vocab_parallel.cpp @@ -0,0 +1,42 @@ +#include "glm_vocab_parallel.hpp" +#include "../../global_state/global_state.hpp" +#include "infinicore/ops.hpp" +#include "infinicore/ops/distributed/allgather.hpp" +#include "infinicore/ops/distributed/allreduce.hpp" +#include "infinicore/ops/vocab_parallel_embedding.hpp" +namespace infinilm::models::glm_moe_dsa { +GlmVocabEmbedding::GlmVocabEmbedding(size_t vocab, size_t hidden, const infinicore::DataType &dt, const infinicore::Device &d) { + auto &r = infinilm::global_state::get_tensor_model_parallel_rank_info(); + hidden_ = hidden; + start_ = vocab * r.tp_rank / r.tp_size; + end_ = vocab * (r.tp_rank + 1) / r.tp_size; + comm_ = r.comm; + INFINICORE_NN_PARAMETER_INIT(weight, ({vocab, hidden}, dt, d, 0, r.tp_rank, r.tp_size)); +} +infinicore::Tensor GlmVocabEmbedding::forward(const infinicore::Tensor &ids) const { + auto s = ids->shape(); + s.push_back(hidden_); + auto out = infinicore::Tensor::empty(s, weight_->dtype(), weight_->device()); + infinicore::op::vocab_parallel_embedding_(out, ids, weight_, start_, end_); + if (comm_) { + infinicore::op::distributed::allreduce_(out, out, INFINICCL_SUM, comm_); + } + return out; +} +GlmVocabLMHead::GlmVocabLMHead(size_t hidden, size_t vocab, bool, const infinicore::DataType &dt, const infinicore::Device &d) { + auto &r = infinilm::global_state::get_tensor_model_parallel_rank_info(); + vocab_ = vocab; + world_ = r.tp_size; + comm_ = r.comm; + INFINICORE_NN_PARAMETER_INIT(weight, ({vocab, hidden}, dt, d, 0, r.tp_rank, r.tp_size)); +} +infinicore::Tensor GlmVocabLMHead::forward(const infinicore::Tensor &x) const { + auto local = infinicore::op::linear(x, weight_, std::nullopt, 1); + if (world_ == 1) { + return local; + } + auto t = local->permute({2, 0, 1})->contiguous(); + auto g = infinicore::op::distributed::allgather(t, world_, comm_); + return g->permute({1, 2, 0})->contiguous(); +} +} // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_vocab_parallel.hpp b/csrc/models/glm_moe_dsa/glm_vocab_parallel.hpp new file mode 100644 index 000000000..70e2491f7 --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_vocab_parallel.hpp @@ -0,0 +1,26 @@ +#pragma once +#include "infinicore/nn/module.hpp" +#include "infinicore/tensor.hpp" +#include +namespace infinilm::models::glm_moe_dsa { +class GlmVocabEmbedding final : public infinicore::nn::Module { +public: + GlmVocabEmbedding(size_t, size_t, const infinicore::DataType &, const infinicore::Device &); + infinicore::Tensor forward(const infinicore::Tensor &) const; + +private: + INFINICORE_NN_PARAMETER(weight); + size_t hidden_{0}, start_{0}, end_{0}; + infinicclComm_t comm_{nullptr}; +}; +class GlmVocabLMHead final : public infinicore::nn::Module { +public: + GlmVocabLMHead(size_t, size_t, bool, const infinicore::DataType &, const infinicore::Device &); + infinicore::Tensor forward(const infinicore::Tensor &) const; + +private: + INFINICORE_NN_PARAMETER(weight); + size_t vocab_{0}, world_{1}; + infinicclComm_t comm_{nullptr}; +}; +} // namespace infinilm::models::glm_moe_dsa diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 8f84c9b97..488235df0 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -64,6 +64,27 @@ def _is_internal_moe_packed_weight(key: str) -> bool: ) +def _is_glm_base_inference_unused_weight(key: str, config: dict) -> bool: + """Weights intentionally outside the base 78-layer forward. + + Layers at num_hidden_layers and above belong to the optional MTP + predictor. The DSA indexer is not consulted by the current full-attention + path; for sequences no longer than index_topk, full attention is exact. + Keep this allowlist GLM-specific so other model loaders remain strict. + """ + if config.get("model_type") != "glm_moe_dsa": + return False + if ".self_attn.indexer." in key: + return True + prefix = "model.layers." + if key.startswith(prefix): + rest = key[len(prefix) :] + layer = rest.split(".", 1)[0] + if layer.isdigit() and int(layer) >= int(config.get("num_hidden_layers", 0)): + return True + return False + + def check_parameters(model_keys: list, already_loaded_keys: list): model_keys = set(model_keys) already_loaded_keys = set(already_loaded_keys) @@ -230,6 +251,13 @@ def load_model_state_dict_by_file( if remapper is not None: model_param = remapper(model_param, config=model.hf_config) + if model_type == "glm_moe_dsa": + model_param = { + key: value + for key, value in model_param.items() + if not _is_glm_base_inference_unused_weight(key, model.hf_config) + } + already_loaded_keys.extend(model_param.keys()) # --------------------------------------------------------- # diff --git a/xmake.lua b/xmake.lua index aab1a0c70..68b1e9eaf 100644 --- a/xmake.lua +++ b/xmake.lua @@ -35,6 +35,7 @@ target_end() target("_infinilm") add_packages("pybind11") + add_defines("_GLIBCXX_USE_CXX11_ABI=0") set_default(false) add_rules("python.module", {soabi = true}) set_languages("cxx17") From 374a23e5de3bde7eb86b52ee8d31dd82fa2868ac Mon Sep 17 00:00:00 2001 From: wooway777 Date: Wed, 15 Jul 2026 11:25:44 +0000 Subject: [PATCH 2/5] pepe: optimize deepseek v2 for iluvatar --- csrc/engine/compiler/paged_compiler.cpp | 2 + csrc/engine/infer_engine.cpp | 3 +- csrc/engine/rank_worker.hpp | 2 + csrc/global_state/forward_context.hpp | 16 +- csrc/layers/mlp/mlp.cpp | 4 +- csrc/layers/moe/router/topk_router.cpp | 85 ++++- csrc/layers/moe/router/topk_router.hpp | 6 + .../deepseek_v2_allocate_kv_cache_tensors.cpp | 82 ++--- .../deepseek_v2/deepseek_v2_attention.cpp | 160 ---------- .../deepseek_v2/deepseek_v2_attention.hpp | 54 ---- .../deepseek_v2/deepseek_v2_decoder_layer.cpp | 11 +- .../deepseek_v2/deepseek_v2_decoder_layer.hpp | 6 +- .../deepseek_v2/deepseek_v2_for_causal_lm.cpp | 84 +++-- .../deepseek_v2/deepseek_v2_mla_attention.cpp | 302 ++++++++++-------- .../deepseek_v2/deepseek_v2_mla_attention.hpp | 12 +- csrc/models/deepseek_v2/deepseek_v2_moe.cpp | 213 ++++++------ csrc/models/deepseek_v2/deepseek_v2_moe.hpp | 44 +-- csrc/pybind11/engine/engine.hpp | 6 +- examples/bench.py | 25 +- python/infinilm/base_config.py | 2 +- python/infinilm/config/engine_config.py | 23 +- python/infinilm/infer_engine.py | 7 + .../processors/basic_llm_processor.py | 2 + 23 files changed, 520 insertions(+), 631 deletions(-) delete mode 100644 csrc/models/deepseek_v2/deepseek_v2_attention.cpp delete mode 100644 csrc/models/deepseek_v2/deepseek_v2_attention.hpp diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index f267794e2..9f534b087 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -23,6 +23,7 @@ PagedCompiler::PagedCompiler(const std::shared_ptr &model, RankBa void PagedCompiler::compile() { if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { size_t nblocks = dynamic_cast(model_->get_cache_config())->num_blocks(); + size_t block_size = dynamic_cast(model_->get_cache_config())->block_size(); size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); compiled_map_decode_.clear(); block_tables_holder_ = infinicore::Tensor::empty( @@ -60,6 +61,7 @@ void PagedCompiler::compile() { input.cu_seqlens, input.block_tables, input.slot_mapping, + static_cast(nblocks * block_size), }; return input; }; diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index a5221eb37..c47266d88 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -189,7 +189,8 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { input.input_offsets, input.cu_seqlens, input.block_tables, - input.slot_mapping}; + input.slot_mapping, + max_context_len}; infinilm::global_state::get_forward_context().mamba_metadata = { input.input_offsets, diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index d396ef6f1..162f8bbac 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -72,6 +72,8 @@ class RankWorker { std::optional target_hidden_states; /// Sample logits at every packed input position instead of one token per request. bool sample_all_positions{false}; + /// Maximum total sequence length in the current request batch. + std::optional max_context_len; float temperature{1}; diff --git a/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index f395b531a..3dc982f5d 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -17,6 +17,8 @@ struct AttentionMetadata { std::optional block_tables; /// Slot ids for each token `[seq]`. Used for paged cache. std::optional slot_mapping; + /// Maximum total sequence length in the current request batch. + std::optional max_context_len; AttentionMetadata() = default; @@ -25,12 +27,14 @@ struct AttentionMetadata { std::optional input_offsets, std::optional cu_seqlens, std::optional block_tables, - std::optional slot_mapping) : past_sequence_lengths(past_sequence_lengths), - total_sequence_lengths(total_sequence_lengths), - input_offsets(input_offsets), - cu_seqlens(cu_seqlens), - block_tables(block_tables), - slot_mapping(slot_mapping) {} + std::optional slot_mapping, + std::optional max_context_len = std::nullopt) : past_sequence_lengths(past_sequence_lengths), + total_sequence_lengths(total_sequence_lengths), + input_offsets(input_offsets), + cu_seqlens(cu_seqlens), + block_tables(block_tables), + slot_mapping(slot_mapping), + max_context_len(max_context_len) {} AttentionMetadata(const infinilm::InfinilmModel::Input &input) : AttentionMetadata(input.past_sequence_lengths, input.total_sequence_lengths, diff --git a/csrc/layers/mlp/mlp.cpp b/csrc/layers/mlp/mlp.cpp index f7604c505..fd4ccc61d 100644 --- a/csrc/layers/mlp/mlp.cpp +++ b/csrc/layers/mlp/mlp.cpp @@ -15,6 +15,8 @@ MLP::MLP(std::shared_ptr model_config, const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); int tp_rank = rank_info.tp_rank; int tp_size = rank_info.tp_size; + const bool reduce_results = model_config->get_or("reduce_results", true); + auto communicator = reduce_results ? rank_info.comm : nullptr; auto quantization_method = model_config->get_quantization_method(); auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; @@ -23,7 +25,7 @@ MLP::MLP(std::shared_ptr model_config, quantization_method, use_bias_, dtype, device, rank_info); down_proj_ = this->register_module( "down_proj", intermediate_size_, hidden_size_, quantization_method, - use_bias_, dtype, device, tp_rank, tp_size, rank_info.comm); + use_bias_, dtype, device, tp_rank, tp_size, communicator); } infinicore::Tensor MLP::forward(const infinicore::Tensor &hidden_states) const { diff --git a/csrc/layers/moe/router/topk_router.cpp b/csrc/layers/moe/router/topk_router.cpp index e70e55cf8..65cc58e31 100644 --- a/csrc/layers/moe/router/topk_router.cpp +++ b/csrc/layers/moe/router/topk_router.cpp @@ -1,6 +1,9 @@ #include "topk_router.hpp" #include "infinicore/ops.hpp" +#include "infinicore/ops/cast.hpp" +#include "infinicore/ops/grouped_topk_vllm.hpp" +#include "infinicore/ops/moe_topk_vllm.hpp" #include #include @@ -18,6 +21,9 @@ TopKRouterBackend parse_router_backend(const std::string &backend) { if (backend == "fused_gate" || backend == "noaux_tc") { return TopKRouterBackend::FusedGate; } + if (backend == "vllm_topk") { + return TopKRouterBackend::VllmTopK; + } throw std::runtime_error("Unsupported MoE router backend: " + backend); } @@ -49,8 +55,11 @@ TopKRouter::TopKRouter(std::shared_ptr model_conf topk_group_ = model_config->get_or("topk_group", 0); num_fused_shared_experts_ = model_config->get_or("num_fused_shared_experts", 0); apply_routed_scaling_factor_on_output_ = model_config->get_or("apply_routed_scaling_factor_on_output", false); + topk_method_ = model_config->get_or("topk_method", "greedy"); + scoring_func_ = model_config->get_or("scoring_func", "softmax"); router_backend_ = parse_router_backend(router_backend_name(model_config)); - use_correction_bias_ = model_config->get_or({"e_score_correction_bias", "moe_router_use_correction_bias"}, false) || router_backend_ == TopKRouterBackend::FusedGate; + use_correction_bias_ = model_config->get_or({"e_score_correction_bias", "moe_router_use_correction_bias"}, false) + || router_backend_ == TopKRouterBackend::FusedGate || topk_method_ == "noaux_tc"; ASSERT((num_experts_ > 0) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_experts_)); INFINICORE_NN_PARAMETER_INIT( @@ -75,6 +84,44 @@ TopKRouter::TopKRouter(std::shared_ptr model_conf throw std::runtime_error("fused_gate MoE router requires num_experts_per_tok > num_fused_shared_experts"); } } + + if (router_backend_ == TopKRouterBackend::VllmTopK) { + if (topk_method_ != "greedy" && topk_method_ != "noaux_tc") { + throw std::runtime_error("vllm_topk MoE router supports greedy or noaux_tc"); + } + if (scoring_func_ != "softmax" && scoring_func_ != "sigmoid") { + throw std::runtime_error("vllm_topk MoE router supports softmax or sigmoid scoring"); + } + if (moe_softcapping_ != 0.0f) { + throw std::runtime_error("vllm_topk MoE router does not support moe_softcapping"); + } + if (topk_method_ == "noaux_tc") { + if (!use_correction_bias_) { + throw std::runtime_error("vllm_topk noaux_tc requires correction bias"); + } + if (num_expert_group_ == 0 || topk_group_ == 0) { + throw std::runtime_error("vllm_topk noaux_tc requires num_expert_group/n_group and topk_group"); + } + if (num_experts_ % num_expert_group_ != 0) { + throw std::runtime_error("vllm_topk noaux_tc requires num_experts divisible by num_expert_group"); + } + } else { + if ((num_expert_group_ != 0 && num_expert_group_ != 1) + || (topk_group_ != 0 && topk_group_ != 1)) { + throw std::runtime_error("vllm_topk greedy supports only a single expert group"); + } + if (routed_scaling_factor_ != 1.0f) { + throw std::runtime_error("vllm_topk greedy currently requires routed_scaling_factor=1"); + } + } + } +} + +void TopKRouter::process_weights_after_loading() { + if (router_backend_ == TopKRouterBackend::VllmTopK && use_correction_bias_) { + runtime_correction_bias_ = infinicore::Tensor::empty({num_experts_}, weight_->dtype(), weight_->device()); + infinicore::op::cast_(runtime_correction_bias_, e_score_correction_bias_); + } } std::tuple TopKRouter::forward(const infinicore::Tensor &hidden_states) const { @@ -118,6 +165,42 @@ std::tuple TopKRouter::forward(const inf routed_scaling_factor_, apply_routed_scaling_factor_on_output_); break; + case TopKRouterBackend::VllmTopK: { + if (topk_method_ == "noaux_tc") { + if (!runtime_correction_bias_) { + throw std::runtime_error("vllm_topk correction bias was not prepared after weight loading"); + } + infinicore::op::grouped_topk_vllm_(router_scores, + router_indices, + router_logits, + num_expert_group_, + topk_group_, + norm_topk_prob_, + routed_scaling_factor_, + runtime_correction_bias_, + scoring_func_); + } else { + auto token_expert_indices = infinicore::Tensor::empty( + {ntoken, num_experts_per_tok_}, infinicore::DataType::I32, hidden_states->device()); + const auto &bias = use_correction_bias_ ? runtime_correction_bias_ : infinicore::Tensor(); + if (scoring_func_ == "softmax") { + infinicore::op::moe_topk_softmax_vllm_(router_scores, + router_indices, + token_expert_indices, + router_logits, + norm_topk_prob_, + bias); + } else { + infinicore::op::moe_topk_sigmoid_vllm_(router_scores, + router_indices, + token_expert_indices, + router_logits, + norm_topk_prob_, + bias); + } + } + break; + } } return std::make_tuple(router_scores, router_indices); diff --git a/csrc/layers/moe/router/topk_router.hpp b/csrc/layers/moe/router/topk_router.hpp index ce1d468b7..7c8ed1610 100644 --- a/csrc/layers/moe/router/topk_router.hpp +++ b/csrc/layers/moe/router/topk_router.hpp @@ -5,6 +5,7 @@ #include #include +#include #include namespace infinilm::layers::moe { @@ -13,12 +14,14 @@ enum class TopKRouterBackend { Softmax, Sigmoid, FusedGate, + VllmTopK, }; class TopKRouter : public infinicore::nn::Module { public: TopKRouter(std::shared_ptr model_config, const infinicore::Device &device); + void process_weights_after_loading() override; std::tuple forward(const infinicore::Tensor &hidden_states) const; @@ -26,6 +29,7 @@ class TopKRouter : public infinicore::nn::Module { INFINICORE_NN_PARAMETER(weight); INFINICORE_NN_PARAMETER(e_score_correction_bias); + infinicore::Tensor runtime_correction_bias_; size_t num_experts_{0}; size_t num_experts_per_tok_{0}; size_t num_expert_group_{0}; @@ -37,6 +41,8 @@ class TopKRouter : public infinicore::nn::Module { bool apply_routed_scaling_factor_on_output_{false}; bool use_correction_bias_{false}; TopKRouterBackend router_backend_{TopKRouterBackend::Softmax}; + std::string topk_method_{"greedy"}; + std::string scoring_func_{"softmax"}; }; } // namespace infinilm::layers::moe diff --git a/csrc/models/deepseek_v2/deepseek_v2_allocate_kv_cache_tensors.cpp b/csrc/models/deepseek_v2/deepseek_v2_allocate_kv_cache_tensors.cpp index b4dbc77ad..05c87bf94 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_allocate_kv_cache_tensors.cpp +++ b/csrc/models/deepseek_v2/deepseek_v2_allocate_kv_cache_tensors.cpp @@ -6,7 +6,6 @@ #include "../../utils.hpp" #include -#include #include namespace infinilm::models::deepseek_v2 { @@ -15,68 +14,39 @@ std::vector deepseek_v2_allocate_kv_cache_tensors( const cache::CacheConfig *cache_config, const std::shared_ptr &text_config, const backends::AttentionBackend &attention_backend) { - if (nullptr == cache_config) { + if (cache_config == nullptr) { return {}; } - if (nullptr == text_config) { - throw std::runtime_error("infinilm::models::deepseek_v2::deepseek_v2_allocate_kv_cache_tensors: text_config is null"); + if (text_config == nullptr) { + throw std::runtime_error("deepseek_v2_allocate_kv_cache_tensors: text_config is null"); } - - const size_t num_hidden_layers = text_config->get("num_hidden_layers"); - const size_t kv_lora_rank = text_config->get("kv_lora_rank"); - const size_t qk_rope_head_dim = text_config->get("qk_rope_head_dim"); - const size_t mla_head_dim = kv_lora_rank + qk_rope_head_dim; - constexpr size_t num_mla_kv_heads = 1; - const auto &dtype = text_config->get_kv_cache_dtype(); - - std::vector kv_cache_vec; - switch (attention_backend) { - case backends::AttentionBackend::STATIC_ATTN: { - auto static_kv_cache_config = dynamic_cast(cache_config); - if (nullptr == static_kv_cache_config) { - throw std::runtime_error("infinilm::models::deepseek_v2::deepseek_v2_allocate_kv_cache_tensors: invalid static kv cache config type"); - } - const size_t max_position_embeddings = text_config->get("max_position_embeddings"); - kv_cache_vec.reserve(num_hidden_layers); - for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { - auto kv_cache = cache::StaticKVCache::create_layer_kv_cache( - mla_head_dim, - mla_head_dim, - num_mla_kv_heads, - num_mla_kv_heads, - max_position_embeddings, - dtype, - *static_kv_cache_config); - kv_cache_vec.push_back(kv_cache); - } - break; + if (attention_backend == backends::AttentionBackend::STATIC_ATTN) { + throw std::runtime_error("DeepSeek V2 requires the vLLM-style paged MLA cache"); } - case backends::AttentionBackend::FLASH_ATTN: { - ; + auto paged_config = dynamic_cast(cache_config); + if (paged_config == nullptr) { + throw std::runtime_error("deepseek_v2_allocate_kv_cache_tensors: expected paged KV cache config"); } - case backends::AttentionBackend::PAGED_ATTN: { - auto paged_kv_cache_config = dynamic_cast(cache_config); - if (nullptr == paged_kv_cache_config) { - throw std::runtime_error("infinilm::models::deepseek_v2::deepseek_v2_allocate_kv_cache_tensors: invalid paged kv cache config type"); - } - const size_t mla_cache_dim = mla_head_dim + kv_lora_rank; - const auto &device = global_state::get_tensor_model_parallel_rank_info().device; - kv_cache_vec.reserve(num_hidden_layers); - for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { - auto kv_cache = infinicore::Tensor::empty( - {paged_kv_cache_config->num_blocks(), num_mla_kv_heads, paged_kv_cache_config->block_size(), mla_cache_dim}, - dtype, - device); - set_zeros(kv_cache); - infinicore::context::syncStream(); - kv_cache_vec.push_back(kv_cache); - } - break; + if (paged_config->block_size() != 16) { + throw std::runtime_error( + "deepseek_v2_allocate_kv_cache_tensors: the current Iluvatar MLA SO requires block_size=16"); } - default: - throw std::runtime_error("infinilm::models::deepseek_v2::deepseek_v2_allocate_kv_cache_tensors: Unsupported attention backend: " + std::to_string(static_cast(attention_backend))); + + const size_t num_layers = text_config->get("num_hidden_layers"); + const size_t cache_dim = text_config->get("kv_lora_rank") + + text_config->get("qk_rope_head_dim"); + const auto &dtype = text_config->get_kv_cache_dtype(); + const auto &device = global_state::get_tensor_model_parallel_rank_info().device; + std::vector caches; + caches.reserve(num_layers); + for (size_t layer = 0; layer < num_layers; ++layer) { + auto cache = infinicore::Tensor::empty( + {paged_config->num_blocks(), paged_config->block_size(), cache_dim}, dtype, device); + set_zeros(cache); + caches.push_back(cache); } - return kv_cache_vec; + infinicore::context::syncStream(); + return caches; } } // namespace infinilm::models::deepseek_v2 diff --git a/csrc/models/deepseek_v2/deepseek_v2_attention.cpp b/csrc/models/deepseek_v2/deepseek_v2_attention.cpp deleted file mode 100644 index 156172793..000000000 --- a/csrc/models/deepseek_v2/deepseek_v2_attention.cpp +++ /dev/null @@ -1,160 +0,0 @@ -#include "deepseek_v2_attention.hpp" - -#include "../../global_state/global_state.hpp" -#include "../../layers/attention/attention.hpp" -#include "../../layers/rotary_embedding/rotary_embedding.hpp" -#include "../../utils.hpp" -#include "deepseek_v2_utils.hpp" -#include "infinicore/ops.hpp" -#include "infinicore/ops/broadcast_to.hpp" -#include "infinicore/ops/cat.hpp" -#include "infinicore/ops/pad.hpp" - -#include - -namespace infinilm::models::deepseek_v2 { - -DeepseekV2Attention::DeepseekV2Attention(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device) { - layer_idx_ = layer_idx; - hidden_size_ = model_config->get("hidden_size"); - qk_nope_head_dim_ = model_config->get("qk_nope_head_dim"); - qk_rope_head_dim_ = model_config->get("qk_rope_head_dim"); - q_head_dim_ = qk_nope_head_dim_ + qk_rope_head_dim_; - v_head_dim_ = model_config->get("v_head_dim"); - - const auto &dtype{model_config->get_dtype()}; - const size_t total_num_heads = model_config->get("num_attention_heads"); - const size_t kv_lora_rank = model_config->get("kv_lora_rank"); - const bool attention_bias = model_config->get_or("attention_bias", false); - const double rms_norm_eps = model_config->get("rms_norm_eps"); - - const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); - const int tp_rank = rank_info.tp_rank; - const int tp_size = rank_info.tp_size; - if ((total_num_heads < static_cast(tp_size)) || (total_num_heads % static_cast(tp_size) != 0)) { - throw std::runtime_error("DeepseekV2Attention: num_attention_heads must be divisible by tp_size"); - } - num_attention_heads_ = total_num_heads / static_cast(tp_size); - attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend; - - auto quantization_method = model_config->get_quantization_method(); - INFINICORE_NN_MODULE_INIT(q_proj, hidden_size_, total_num_heads * q_head_dim_, quantization_method, false, dtype, device, tp_rank, tp_size); - INFINICORE_NN_MODULE_INIT(kv_a_proj_with_mqa, hidden_size_, kv_lora_rank + qk_rope_head_dim_, attention_bias, dtype, device); - INFINICORE_NN_MODULE_INIT(kv_a_layernorm, kv_lora_rank, rms_norm_eps, dtype, device); - INFINICORE_NN_MODULE_INIT(kv_b_proj, kv_lora_rank, total_num_heads * (qk_nope_head_dim_ + v_head_dim_), quantization_method, false, dtype, device, tp_rank, tp_size); - INFINICORE_NN_MODULE_INIT(o_proj, total_num_heads * v_head_dim_, hidden_size_, quantization_method, attention_bias, dtype, device, tp_rank, tp_size, rank_info.comm); - - rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); - softmax_scale_ = deepseek_v2_attention_softmax_scale(model_config, static_cast(q_head_dim_)); - - attn_ = std::make_shared( - num_attention_heads_, q_head_dim_, softmax_scale_, num_attention_heads_, layer_idx_, - kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_); - infinilm::layers::attention::init_kv_cache_quant_params( - [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }, - device, kv_cache_k_scale_, kv_cache_v_scale_); -} - -infinicore::Tensor DeepseekV2Attention::position_ids_for_rope_(const infinicore::Tensor &position_ids) const { - auto pos_shape = position_ids->shape(); - if (pos_shape.size() == 2) { - return position_ids->narrow({{0, 0, 1}})->contiguous()->view({pos_shape[1]}); - } - if (pos_shape.size() == 1) { - return position_ids->contiguous(); - } - throw std::runtime_error("DeepseekV2Attention: unexpected position_ids shape"); -} - -infinicore::Tensor DeepseekV2Attention::trim_value_padding_(const infinicore::Tensor &attn_output) const { - const auto shape = attn_output->shape(); - const size_t batch_size = shape[0]; - const size_t seq_len = shape[1]; - return attn_output->view({batch_size, seq_len, num_attention_heads_, q_head_dim_}) - ->narrow({{3, 0, v_head_dim_}}) - ->contiguous() - ->view({batch_size, seq_len, num_attention_heads_ * v_head_dim_}); -} - -infinicore::Tensor DeepseekV2Attention::forward(const infinicore::Tensor &positions, - const infinicore::Tensor &hidden_states) const { - if (::infinilm::backends::AttentionBackend::STATIC_ATTN == attention_backend_) { - return forward_static_(positions, hidden_states); - } - return forward_paged_(positions, hidden_states); -} - -infinicore::Tensor DeepseekV2Attention::forward_static_(const infinicore::Tensor &position_ids, - const infinicore::Tensor &hidden_states) const { - auto shape = hidden_states->shape(); - const size_t batch_size = shape[0]; - const size_t seq_len = shape[1]; - auto hidden_states_mutable = hidden_states; - - auto q = q_proj_->forward(hidden_states_mutable)->view({batch_size, seq_len, num_attention_heads_, q_head_dim_}); - auto q_nope = q->narrow({{3, 0, qk_nope_head_dim_}}); - auto q_pe = q->narrow({{3, qk_nope_head_dim_, qk_rope_head_dim_}})->contiguous(); - - auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable); - auto compressed_kv = compressed->narrow({{2, 0, kv_a_layernorm_->normalized_shape()}})->contiguous(); - auto k_pe = compressed->narrow({{2, kv_a_layernorm_->normalized_shape(), qk_rope_head_dim_}})->contiguous(); - - auto kv_norm = kv_a_layernorm_->forward(compressed_kv); - auto kv = kv_b_proj_->forward(kv_norm)->view({batch_size, seq_len, num_attention_heads_, qk_nope_head_dim_ + v_head_dim_}); - auto k_nope = kv->narrow({{3, 0, qk_nope_head_dim_}}); - auto value_states = kv->narrow({{3, qk_nope_head_dim_, v_head_dim_}})->contiguous(); - - auto pos_ids = position_ids_for_rope_(position_ids); - q_pe = rotary_emb_->forward(q_pe, pos_ids, true); - auto k_pe_broadcast = infinicore::op::broadcast_to(k_pe->view({batch_size, seq_len, 1, qk_rope_head_dim_}), - {static_cast(batch_size), static_cast(seq_len), static_cast(num_attention_heads_), static_cast(qk_rope_head_dim_)}); - k_pe_broadcast = rotary_emb_->forward(k_pe_broadcast, pos_ids, true); - - auto query_states = infinicore::op::cat({q_nope, q_pe}, 3); - auto key_states = infinicore::op::cat({k_nope, k_pe_broadcast}, 3); - auto value_padded = infinicore::op::pad(value_states, {0, static_cast(q_head_dim_ - v_head_dim_)}, "constant", 0.0); - - auto attn_output = attn_->forward(query_states, key_states, value_padded); - auto trimmed_output = trim_value_padding_(attn_output); - return o_proj_->forward(trimmed_output); -} - -infinicore::Tensor DeepseekV2Attention::forward_paged_(const infinicore::Tensor &position_ids, - const infinicore::Tensor &hidden_states) const { - auto shape = hidden_states->shape(); - const size_t batch_size = shape[0]; - const size_t seq_len = shape[1]; - ASSERT_EQ(batch_size, 1); - auto hidden_states_mutable = hidden_states; - - auto q = q_proj_->forward(hidden_states_mutable)->view({seq_len, num_attention_heads_, q_head_dim_}); - auto q_nope = q->narrow({{2, 0, qk_nope_head_dim_}}); - auto q_pe = q->narrow({{2, qk_nope_head_dim_, qk_rope_head_dim_}})->contiguous(); - - auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable)->view({seq_len, kv_a_layernorm_->normalized_shape() + qk_rope_head_dim_}); - auto compressed_kv = compressed->narrow({{1, 0, kv_a_layernorm_->normalized_shape()}})->contiguous(); - auto k_pe = compressed->narrow({{1, kv_a_layernorm_->normalized_shape(), qk_rope_head_dim_}})->contiguous(); - - auto kv_norm = kv_a_layernorm_->forward(compressed_kv); - auto kv = kv_b_proj_->forward(kv_norm)->view({seq_len, num_attention_heads_, qk_nope_head_dim_ + v_head_dim_}); - auto k_nope = kv->narrow({{2, 0, qk_nope_head_dim_}}); - auto value_states = kv->narrow({{2, qk_nope_head_dim_, v_head_dim_}})->contiguous(); - - auto pos_ids = position_ids_for_rope_(position_ids); - q_pe = rotary_emb_->forward(q_pe, pos_ids, true); - auto k_pe_broadcast = infinicore::op::broadcast_to(k_pe->view({seq_len, 1, qk_rope_head_dim_}), - {static_cast(seq_len), static_cast(num_attention_heads_), static_cast(qk_rope_head_dim_)}); - k_pe_broadcast = rotary_emb_->forward(k_pe_broadcast, pos_ids, true); - - auto query_states = infinicore::op::cat({q_nope, q_pe}, 2); - auto key_states = infinicore::op::cat({k_nope, k_pe_broadcast}, 2); - auto value_padded = infinicore::op::pad(value_states, {0, static_cast(q_head_dim_ - v_head_dim_)}, "constant", 0.0); - - auto attn_output = attn_->forward(query_states, key_states, value_padded); - auto trimmed_output = trim_value_padding_(attn_output); - return o_proj_->forward(trimmed_output); -} - -} // namespace infinilm::models::deepseek_v2 diff --git a/csrc/models/deepseek_v2/deepseek_v2_attention.hpp b/csrc/models/deepseek_v2/deepseek_v2_attention.hpp deleted file mode 100644 index afedb071f..000000000 --- a/csrc/models/deepseek_v2/deepseek_v2_attention.hpp +++ /dev/null @@ -1,54 +0,0 @@ -#pragma once - -#include "../../config/model_config.hpp" -#include "../../layers/attention/attention.hpp" -#include "../../layers/linear/linear.hpp" -#include "infinicore/nn/module.hpp" -#include "infinicore/nn/rmsnorm.hpp" -#include "infinicore/nn/rope.hpp" -#include "infinicore/tensor.hpp" - -#include - -namespace infinilm::models::deepseek_v2 { - -class DeepseekV2Attention : public infinicore::nn::Module { -public: - DeepseekV2Attention(std::shared_ptr model_config, - size_t layer_idx, - const infinicore::Device &device); - - infinicore::Tensor forward(const infinicore::Tensor &positions, - const infinicore::Tensor &hidden_states) const; - -private: - infinicore::Tensor forward_static_(const infinicore::Tensor &positions, - const infinicore::Tensor &hidden_states) const; - infinicore::Tensor forward_paged_(const infinicore::Tensor &positions, - const infinicore::Tensor &hidden_states) const; - infinicore::Tensor trim_value_padding_(const infinicore::Tensor &attn_output) const; - infinicore::Tensor position_ids_for_rope_(const infinicore::Tensor &position_ids) const; - - size_t layer_idx_{0}; - size_t hidden_size_{0}; - size_t num_attention_heads_{0}; - size_t qk_nope_head_dim_{0}; - size_t qk_rope_head_dim_{0}; - size_t q_head_dim_{0}; - size_t v_head_dim_{0}; - float softmax_scale_{1.0f}; - infinilm::backends::AttentionBackend attention_backend_; - - INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_proj); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, kv_a_proj_with_mqa); - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, kv_a_layernorm); - INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, kv_b_proj); - INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, o_proj); - - std::shared_ptr rotary_emb_; - std::shared_ptr attn_; - infinicore::nn::Parameter kv_cache_k_scale_; - infinicore::nn::Parameter kv_cache_v_scale_; -}; - -} // namespace infinilm::models::deepseek_v2 diff --git a/csrc/models/deepseek_v2/deepseek_v2_decoder_layer.cpp b/csrc/models/deepseek_v2/deepseek_v2_decoder_layer.cpp index 9a1e0c0b1..2843bc86b 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_decoder_layer.cpp +++ b/csrc/models/deepseek_v2/deepseek_v2_decoder_layer.cpp @@ -1,7 +1,5 @@ #include "deepseek_v2_decoder_layer.hpp" -#include "../../global_state/global_state.hpp" - namespace infinilm::models::deepseek_v2 { DeepseekV2DecoderLayer::DeepseekV2DecoderLayer(std::shared_ptr model_config, @@ -12,11 +10,7 @@ DeepseekV2DecoderLayer::DeepseekV2DecoderLayer(std::shared_ptrget("rms_norm_eps"); INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device); INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device); - if (infinilm::global_state::get_infinilm_config().use_mla) { - self_attn_ = std::make_shared(this->register_module("self_attn", model_config, layer_idx, device)); - } else { - self_attn_ = std::make_shared(this->register_module("self_attn", model_config, layer_idx, device)); - } + INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); const size_t first_k_dense_replace = model_config->get_or("first_k_dense_replace", 0); const size_t moe_layer_freq = model_config->get_or("moe_layer_freq", 1); @@ -35,8 +29,7 @@ DeepseekV2DecoderLayer::forward(const infinicore::Tensor &positions, infinicore::Tensor &hidden_states, infinicore::Tensor &residual) const { input_layernorm_->forward_inplace(hidden_states, residual); - hidden_states = std::visit( - [&](auto &attn_ptr) { return attn_ptr->forward(positions, hidden_states); }, *self_attn_); + hidden_states = self_attn_->forward(positions, hidden_states); post_attention_layernorm_->forward_inplace(hidden_states, residual); hidden_states = use_moe_ ? moe_mlp_->forward(hidden_states) : dense_mlp_->forward(hidden_states); return {hidden_states, residual}; diff --git a/csrc/models/deepseek_v2/deepseek_v2_decoder_layer.hpp b/csrc/models/deepseek_v2/deepseek_v2_decoder_layer.hpp index 34541fa5e..7cf785279 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_decoder_layer.hpp +++ b/csrc/models/deepseek_v2/deepseek_v2_decoder_layer.hpp @@ -1,7 +1,6 @@ #pragma once #include "../../config/model_config.hpp" -#include "deepseek_v2_attention.hpp" #include "deepseek_v2_mla_attention.hpp" #include "deepseek_v2_moe.hpp" #include "infinicore/device.hpp" @@ -11,12 +10,9 @@ #include #include -#include namespace infinilm::models::deepseek_v2 { -using DeepseekV2SelfAttention = std::variant, std::shared_ptr>; - class DeepseekV2DecoderLayer : public infinicore::nn::Module { public: DeepseekV2DecoderLayer(std::shared_ptr model_config, @@ -30,7 +26,7 @@ class DeepseekV2DecoderLayer : public infinicore::nn::Module { private: INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); - INFINICORE_NN_MODULE(DeepseekV2SelfAttention, self_attn); + INFINICORE_NN_MODULE(DeepseekV2MLAAttention, self_attn); INFINICORE_NN_MODULE(DeepseekV2MLP, dense_mlp); INFINICORE_NN_MODULE(DeepseekV2MoE, moe_mlp); bool use_moe_{false}; diff --git a/csrc/models/deepseek_v2/deepseek_v2_for_causal_lm.cpp b/csrc/models/deepseek_v2/deepseek_v2_for_causal_lm.cpp index 3ccc17d3d..d72b90470 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_for_causal_lm.cpp +++ b/csrc/models/deepseek_v2/deepseek_v2_for_causal_lm.cpp @@ -11,27 +11,24 @@ namespace infinilm::models::deepseek_v2 { DeepseekV2Model::DeepseekV2Model(std::shared_ptr model_config, const infinicore::Device &device) { const auto &dtype{model_config->get_dtype()}; - const size_t vocab_size = model_config->get("vocab_size"); const size_t hidden_size = model_config->get("hidden_size"); - const size_t num_hidden_layers = model_config->get("num_hidden_layers"); - const double rms_norm_eps = model_config->get("rms_norm_eps"); - - INFINICORE_NN_MODULE_INIT(embed_tokens, vocab_size, hidden_size, std::nullopt, dtype, device); - layers_.reserve(num_hidden_layers); - for (size_t i = 0; i < num_hidden_layers; ++i) { - layers_.push_back(this->register_module("layers." + std::to_string(i), model_config, i, device)); + INFINICORE_NN_MODULE_INIT( + embed_tokens, model_config->get("vocab_size"), hidden_size, std::nullopt, dtype, device); + const size_t num_layers = model_config->get("num_hidden_layers"); + layers_.reserve(num_layers); + for (size_t layer = 0; layer < num_layers; ++layer) { + layers_.push_back(this->register_module( + "layers." + std::to_string(layer), model_config, layer, device)); } - INFINICORE_NN_MODULE_INIT(norm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT( + norm, hidden_size, model_config->get("rms_norm_eps"), dtype, device); } infinicore::Tensor DeepseekV2Model::forward(const infinilm::InfinilmModel::Input &input) const { - auto input_ids = input.input_ids.value(); - auto positions = input.position_ids.value(); - auto hidden_states = embed_tokens_->forward(input_ids); - + auto hidden_states = embed_tokens_->forward(input.input_ids.value()); infinicore::Tensor residual; for (const auto &layer : layers_) { - layer->forward(positions, hidden_states, residual); + layer->forward(input.position_ids.value(), hidden_states, residual); } norm_->forward_inplace(hidden_states, residual); return hidden_states; @@ -40,59 +37,50 @@ infinicore::Tensor DeepseekV2Model::forward(const infinilm::InfinilmModel::Input DeepseekV2ForCausalLM::DeepseekV2ForCausalLM(std::shared_ptr model_config, const infinicore::Device &device) { model_config_ = model_config; - const auto &dtype{model_config->get_dtype()}; - const size_t hidden_size = model_config->get("hidden_size"); - const size_t vocab_size = model_config->get("vocab_size"); INFINICORE_NN_MODULE_INIT(model, model_config, device); - INFINICORE_NN_MODULE_INIT(lm_head, hidden_size, vocab_size, false, dtype, device); + INFINICORE_NN_MODULE_INIT(lm_head, + model_config->get("hidden_size"), + model_config->get("vocab_size"), + false, + model_config->get_dtype(), + device); } infinilm::InfinilmModel::Output DeepseekV2ForCausalLM::forward(const infinilm::InfinilmModel::Input &input) const { auto hidden_states = model_->forward(input); - auto logits = lm_head_->forward(hidden_states); - return {logits}; + return {lm_head_->forward(hidden_states)}; } void DeepseekV2ForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { - const auto &infinilm_config = infinilm::global_state::get_infinilm_config(); - if (!infinilm_config.use_mla || cache_config == nullptr) { + if (cache_config == nullptr) { InfinilmModel::reset_cache(cache_config); return; } - cache_config_ = cache_config->unique_copy(); - - auto &kv_cache_vec = infinilm::global_state::get_forward_context().kv_cache_vec; - kv_cache_vec.clear(); - auto new_kv_cache_vec = deepseek_v2_allocate_kv_cache_tensors(cache_config, model_config_, infinilm_config.attention_backend); - kv_cache_vec = std::move(new_kv_cache_vec); + global_state::get_forward_context().kv_cache_vec = deepseek_v2_allocate_kv_cache_tensors( + cache_config, model_config_, global_state::get_infinilm_config().attention_backend); } -std::shared_ptr create_deepseek_v2_model_config(std::shared_ptr model_config) { - const std::string model_type = model_config->get("model_type"); - if ("deepseek_v2" != model_type) { +std::shared_ptr +create_deepseek_v2_model_config(std::shared_ptr model_config) { + if (model_config->get("model_type") != "deepseek_v2") { throw std::runtime_error("create_deepseek_v2_model_config: model_type is not deepseek_v2"); } - - auto &config_json = model_config->get_config_json(); - const size_t q_head_dim = config_json.at("qk_nope_head_dim").get() + config_json.at("qk_rope_head_dim").get(); - config_json["head_dim"] = q_head_dim; - - const size_t qk_rope_head_dim = config_json.at("qk_rope_head_dim").get(); - config_json["partial_rotary_factor"] = static_cast(qk_rope_head_dim) / static_cast(q_head_dim); - - config_json["num_experts"] = config_json.value("n_routed_experts", 0); - config_json["mlp_bias"] = false; - if (!config_json.contains("attention_output_bias")) { - config_json["attention_output_bias"] = config_json.value("attention_bias", false); + auto &json = model_config->get_config_json(); + const size_t nope_dim = json.at("qk_nope_head_dim").get(); + const size_t rope_dim = json.at("qk_rope_head_dim").get(); + json["head_dim"] = nope_dim + rope_dim; + json["partial_rotary_factor"] = static_cast(rope_dim) / static_cast(nope_dim + rope_dim); + json["num_experts"] = json.value("n_routed_experts", 0); + json["mlp_bias"] = false; + json["moe_router_backend"] = "vllm_topk"; + if (!json.contains("attention_output_bias")) { + json["attention_output_bias"] = json.value("attention_bias", false); } - if (!config_json.contains("dtype") && config_json.contains("torch_dtype")) { - config_json["dtype"] = config_json["torch_dtype"]; + if (!json.contains("dtype") && json.contains("torch_dtype")) { + json["dtype"] = json["torch_dtype"]; } - - // Use GPT-J style for DeepseekV2 model_config->set_rope_algo(infinicore::nn::RoPE::Algo::GPT_J); - return model_config; } diff --git a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp index 3c8a3f8b6..7dfcbf857 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp +++ b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp @@ -5,10 +5,13 @@ #include "../../layers/rotary_embedding/rotary_embedding.hpp" #include "../../utils.hpp" #include "deepseek_v2_utils.hpp" -#include "infinicore/ops.hpp" +#include "infinicore/ops/broadcast_to.hpp" #include "infinicore/ops/cat.hpp" +#include "infinicore/ops/concat_and_cache_mla.hpp" +#include "infinicore/ops/concat_mla_q.hpp" +#include "infinicore/ops/matmul.hpp" #include "infinicore/ops/mha_varlen.hpp" -#include "infinicore/ops/pad.hpp" +#include "infinicore/ops/paged_attention_mla.hpp" #include @@ -26,9 +29,10 @@ DeepseekV2MLAAttention::DeepseekV2MLAAttention(std::shared_ptrget("kv_lora_rank"); mla_head_dim_ = kv_lora_rank_ + qk_rope_head_dim_; - if (model_config->get_or("q_lora_rank", 0) != 0) { - throw std::runtime_error("DeepseekV2MLAAttention: q_lora_rank is not supported yet"); - } + const auto &config_json = model_config->get_config_json(); + q_lora_rank_ = config_json.contains("q_lora_rank") && !config_json["q_lora_rank"].is_null() + ? config_json["q_lora_rank"].get() + : 0; const auto &dtype{model_config->get_dtype()}; const size_t total_num_heads = model_config->get("num_attention_heads"); @@ -38,32 +42,91 @@ DeepseekV2MLAAttention::DeepseekV2MLAAttention(std::shared_ptr(tp_size)) || (total_num_heads % static_cast(tp_size) != 0)) { + if (total_num_heads < static_cast(tp_size) + || total_num_heads % static_cast(tp_size) != 0) { throw std::runtime_error("DeepseekV2MLAAttention: num_attention_heads must be divisible by tp_size"); } num_attention_heads_ = total_num_heads / static_cast(tp_size); attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend; + if (attention_backend_ == infinilm::backends::AttentionBackend::STATIC_ATTN) { + throw std::runtime_error("DeepseekV2MLAAttention requires paged or flash attention; the dense MHA path was removed"); + } auto quantization_method = model_config->get_quantization_method(); - INFINICORE_NN_MODULE_INIT(q_proj, hidden_size_, total_num_heads * q_head_dim_, quantization_method, false, dtype, device, tp_rank, tp_size); - INFINICORE_NN_MODULE_INIT(kv_a_proj_with_mqa, hidden_size_, kv_lora_rank_ + qk_rope_head_dim_, attention_bias, dtype, device); + if (q_lora_rank_ == 0) { + INFINICORE_NN_MODULE_INIT(q_proj, + hidden_size_, + total_num_heads * q_head_dim_, + quantization_method, + false, + dtype, + device, + tp_rank, + tp_size); + } else { + INFINICORE_NN_MODULE_INIT(q_a_proj, + hidden_size_, + q_lora_rank_, + quantization_method, + false, + dtype, + device); + INFINICORE_NN_MODULE_INIT(q_a_layernorm, q_lora_rank_, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(q_b_proj, + q_lora_rank_, + total_num_heads * q_head_dim_, + quantization_method, + false, + dtype, + device, + tp_rank, + tp_size); + } + INFINICORE_NN_MODULE_INIT(kv_a_proj_with_mqa, + hidden_size_, + kv_lora_rank_ + qk_rope_head_dim_, + quantization_method, + attention_bias, + dtype, + device); INFINICORE_NN_MODULE_INIT(kv_a_layernorm, kv_lora_rank_, rms_norm_eps, dtype, device); - INFINICORE_NN_MODULE_INIT(kv_b_proj, kv_lora_rank_, total_num_heads * (qk_nope_head_dim_ + v_head_dim_), quantization_method, false, dtype, device, tp_rank, tp_size); - INFINICORE_NN_MODULE_INIT(o_proj, total_num_heads * v_head_dim_, hidden_size_, quantization_method, attention_bias, dtype, device, tp_rank, tp_size, rank_info.comm); + INFINICORE_NN_MODULE_INIT(kv_b_proj, + kv_lora_rank_, + total_num_heads * (qk_nope_head_dim_ + v_head_dim_), + quantization_method, + false, + dtype, + device, + tp_rank, + tp_size); + INFINICORE_NN_MODULE_INIT(o_proj, + total_num_heads * v_head_dim_, + hidden_size_, + quantization_method, + attention_bias, + dtype, + device, + tp_rank, + tp_size, + rank_info.comm); rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); softmax_scale_ = deepseek_v2_attention_softmax_scale(model_config, static_cast(q_head_dim_)); - - latent_attn_ = std::make_shared( - num_attention_heads_, mla_head_dim_, softmax_scale_, 1, layer_idx_, - kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_); infinilm::layers::attention::init_kv_cache_quant_params( - [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }, - device, kv_cache_k_scale_, kv_cache_v_scale_); + [this](const std::string &name, infinicore::nn::Parameter parameter) { + this->register_parameter(name, std::move(parameter)); + }, + device, + kv_cache_k_scale_, + kv_cache_v_scale_); + if (!kv_cache_k_scale_) { + kv_cache_k_scale_ = infinicore::nn::Parameter( + infinicore::Tensor::ones({1}, infinicore::DataType::F32, device)); + } } infinicore::Tensor DeepseekV2MLAAttention::position_ids_for_rope_(const infinicore::Tensor &position_ids) const { - auto pos_shape = position_ids->shape(); + const auto pos_shape = position_ids->shape(); if (pos_shape.size() == 2) { return position_ids->narrow({{0, 0, 1}})->contiguous()->view({pos_shape[1]}); } @@ -74,32 +137,27 @@ infinicore::Tensor DeepseekV2MLAAttention::position_ids_for_rope_(const infinico } infinicore::Tensor DeepseekV2MLAAttention::kv_b_weight_3d_() const { - return kv_b_proj_->weight()->view({num_attention_heads_, qk_nope_head_dim_ + v_head_dim_, kv_lora_rank_}); + return kv_b_proj_->weight()->view( + {num_attention_heads_, qk_nope_head_dim_ + v_head_dim_, kv_lora_rank_}); } infinicore::Tensor DeepseekV2MLAAttention::project_q_nope_to_latent_(const infinicore::Tensor &q_nope) const { - const size_t ntokens = q_nope->shape()[0]; + const size_t num_tokens = q_nope->shape()[0]; auto q_nope_by_head = q_nope->permute({1, 0, 2})->contiguous(); auto w_uk_t = kv_b_weight_3d_()->narrow({{1, 0, qk_nope_head_dim_}})->contiguous(); auto q_latent = infinicore::op::matmul(q_nope_by_head, w_uk_t); - return q_latent->permute({1, 0, 2})->contiguous()->view({ntokens, num_attention_heads_, kv_lora_rank_}); + return q_latent->permute({1, 0, 2}) + ->contiguous() + ->view({num_tokens, num_attention_heads_, kv_lora_rank_}); } infinicore::Tensor DeepseekV2MLAAttention::project_latent_to_value_(const infinicore::Tensor &attn_output, size_t batch_size, size_t seq_len) const { - const size_t ntokens = batch_size * seq_len; - const auto out_shape = attn_output->shape(); - const size_t out_head_dim = out_shape.back(); - infinicore::Tensor latent; - if (out_head_dim == kv_lora_rank_) { - latent = attn_output->view({ntokens, num_attention_heads_, kv_lora_rank_}); - } else { - latent = attn_output->view({ntokens, num_attention_heads_, mla_head_dim_}) - ->narrow({{2, 0, kv_lora_rank_}}) - ->contiguous(); - } - auto latent_by_head = latent->permute({1, 0, 2})->contiguous(); + const size_t num_tokens = batch_size * seq_len; + auto latent_by_head = attn_output->view({num_tokens, num_attention_heads_, kv_lora_rank_}) + ->permute({1, 0, 2}) + ->contiguous(); auto w_uv = kv_b_weight_3d_() ->narrow({{1, qk_nope_head_dim_, v_head_dim_}}) ->permute({0, 2, 1}) @@ -111,77 +169,42 @@ infinicore::Tensor DeepseekV2MLAAttention::project_latent_to_value_(const infini return o_proj_->forward(value); } -infinicore::Tensor DeepseekV2MLAAttention::forward(const infinicore::Tensor &positions, +infinicore::Tensor DeepseekV2MLAAttention::forward(const infinicore::Tensor &position_ids, const infinicore::Tensor &hidden_states) const { - if (::infinilm::backends::AttentionBackend::STATIC_ATTN == attention_backend_) { - return forward_static_(positions, hidden_states); - } - return forward_paged_(positions, hidden_states); -} - -infinicore::Tensor DeepseekV2MLAAttention::forward_static_(const infinicore::Tensor &position_ids, - const infinicore::Tensor &hidden_states) const { - auto shape = hidden_states->shape(); + const auto shape = hidden_states->shape(); const size_t batch_size = shape[0]; const size_t seq_len = shape[1]; - const size_t ntokens = batch_size * seq_len; - auto hidden_states_mutable = hidden_states; - - auto q = q_proj_->forward(hidden_states_mutable)->view({ntokens, num_attention_heads_, q_head_dim_}); - auto q_nope = q->narrow({{2, 0, qk_nope_head_dim_}})->contiguous(); - auto q_pe = q->narrow({{2, qk_nope_head_dim_, qk_rope_head_dim_}})->contiguous(); - - auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable)->view({ntokens, kv_lora_rank_ + qk_rope_head_dim_}); - auto compressed_kv = compressed->narrow({{1, 0, kv_lora_rank_}})->contiguous(); - auto k_pe = compressed->narrow({{1, kv_lora_rank_, qk_rope_head_dim_}})->contiguous(); - - auto kv_norm = kv_a_layernorm_->forward(compressed_kv); - auto pos_shape = position_ids->shape(); - auto pos_ids = pos_shape.size() == 2 ? position_ids->contiguous()->view({ntokens}) : position_ids_for_rope_(position_ids); - q_pe = rotary_emb_->forward(q_pe, pos_ids, true); - auto k_pe_rope = rotary_emb_->forward(k_pe->view({ntokens, 1, qk_rope_head_dim_}), pos_ids, true); - - auto q_latent = project_q_nope_to_latent_(q_nope); - auto query_states = infinicore::op::cat({q_latent, q_pe}, 2)->view({batch_size, seq_len, num_attention_heads_, mla_head_dim_}); - auto key_states = infinicore::op::cat({kv_norm->view({ntokens, 1, kv_lora_rank_}), k_pe_rope}, 2) - ->view({batch_size, seq_len, 1, mla_head_dim_}); - auto value_states = infinicore::op::pad(kv_norm->view({batch_size, seq_len, 1, kv_lora_rank_}), - {0, static_cast(qk_rope_head_dim_)}, "constant", 0.0); - - auto attn_output = latent_attn_->forward(query_states, key_states, value_states); - return project_latent_to_value_(attn_output, batch_size, seq_len); -} + if (batch_size != 1) { + throw std::runtime_error("DeepseekV2MLAAttention currently expects batch_size=1 in the paged engine"); + } -infinicore::Tensor DeepseekV2MLAAttention::forward_paged_(const infinicore::Tensor &position_ids, - const infinicore::Tensor &hidden_states) const { - auto shape = hidden_states->shape(); - const size_t batch_size = shape[0]; - const size_t seq_len = shape[1]; - ASSERT_EQ(batch_size, 1); auto hidden_states_mutable = hidden_states; - - auto q = q_proj_->forward(hidden_states_mutable)->view({seq_len, num_attention_heads_, q_head_dim_}); + infinicore::Tensor q_linear; + if (q_lora_rank_ == 0) { + q_linear = q_proj_->forward(hidden_states_mutable); + } else { + auto q_a = q_a_proj_->forward(hidden_states_mutable); + auto q_a_norm = q_a_layernorm_->forward(q_a); + q_linear = q_b_proj_->forward(q_a_norm); + } + auto q = q_linear->view({seq_len, num_attention_heads_, q_head_dim_}); auto q_nope = q->narrow({{2, 0, qk_nope_head_dim_}})->contiguous(); auto q_pe = q->narrow({{2, qk_nope_head_dim_, qk_rope_head_dim_}})->contiguous(); - auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable)->view({seq_len, kv_lora_rank_ + qk_rope_head_dim_}); + auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable) + ->view({seq_len, kv_lora_rank_ + qk_rope_head_dim_}); auto compressed_kv = compressed->narrow({{1, 0, kv_lora_rank_}})->contiguous(); auto k_pe = compressed->narrow({{1, kv_lora_rank_, qk_rope_head_dim_}})->contiguous(); - auto kv_norm = kv_a_layernorm_->forward(compressed_kv); + auto pos_ids = position_ids_for_rope_(position_ids); q_pe = rotary_emb_->forward(q_pe, pos_ids, true); - auto k_pe_rope = rotary_emb_->forward(k_pe->view({seq_len, 1, qk_rope_head_dim_}), pos_ids, true); - - auto q_latent = project_q_nope_to_latent_(q_nope); - auto query_states = infinicore::op::cat({q_latent, q_pe}, 2); - auto key_states = infinicore::op::cat({kv_norm->view({seq_len, 1, kv_lora_rank_}), k_pe_rope}, 2); - auto value_states = kv_norm->view({seq_len, 1, kv_lora_rank_}); + auto k_pe_rope = rotary_emb_->forward( + k_pe->view({seq_len, 1, qk_rope_head_dim_}), pos_ids, true); auto &forward_context = infinilm::global_state::get_forward_context(); auto &attn_metadata = forward_context.attn_metadata; auto &kv_cache = forward_context.kv_cache_vec[layer_idx_]; - auto block_tables = attn_metadata.block_tables; auto slot_mapping = attn_metadata.slot_mapping; auto total_sequence_lengths = attn_metadata.total_sequence_lengths; @@ -190,59 +213,72 @@ infinicore::Tensor DeepseekV2MLAAttention::forward_paged_(const infinicore::Tens ASSERT(block_tables.has_value()); ASSERT(slot_mapping.has_value()); ASSERT(total_sequence_lengths.has_value()); + ASSERT(attn_metadata.max_context_len.has_value()); - auto k_cache_layer = kv_cache->narrow({{3, 0, mla_head_dim_}}); - auto v_cache_layer = kv_cache->narrow({{3, mla_head_dim_, kv_lora_rank_}}); - infinicore::op::paged_caching_(k_cache_layer, v_cache_layer, key_states, value_states, slot_mapping.value()); + if (hidden_states->device().getType() != infinicore::Device::Type::ILUVATAR) { + throw std::runtime_error("DeepseekV2MLAAttention: the vLLM-style MLA cache path currently requires Iluvatar"); + } + infinicore::op::concat_and_cache_mla_(kv_norm, + k_pe_rope->view({seq_len, qk_rope_head_dim_}), + kv_cache, + slot_mapping.value(), + "auto", + kv_cache_k_scale_); - auto attn_output = infinicore::Tensor::empty({seq_len, num_attention_heads_, kv_lora_rank_}, query_states->dtype(), query_states->device()); - const bool is_prefill = (seq_len != total_sequence_lengths.value()->shape()[0]); + const bool is_prefill = seq_len != total_sequence_lengths.value()->shape()[0]; if (is_prefill) { + if (attention_backend_ != infinilm::backends::AttentionBackend::FLASH_ATTN) { + throw std::runtime_error("DeepseekV2MLAAttention: prefill requires --attn=flash-attn"); + } ASSERT(input_offsets.has_value()); ASSERT(cu_seqlens.has_value()); - if (attention_backend_ == ::infinilm::backends::AttentionBackend::FLASH_ATTN) { - const size_t num_reqs = total_sequence_lengths.value()->shape()[0]; - ASSERT(num_reqs > 0); - ASSERT_EQ(seq_len % num_reqs, 0); - const int max_seqlen = static_cast(seq_len / num_reqs); - // K/V are dense prompt tensors. Passing block_tables here makes the - // varlen backend interpret them as a paged KV cache. - infinicore::op::mha_varlen_( - attn_output, - query_states, - key_states, - value_states, - input_offsets.value(), - cu_seqlens.value(), - std::nullopt, - max_seqlen, - max_seqlen, - std::nullopt, - softmax_scale_); - } else { - ASSERT_EQ(attention_backend_, ::infinilm::backends::AttentionBackend::PAGED_ATTN); - infinicore::op::paged_attention_prefill_( - attn_output, - query_states, - k_cache_layer, - v_cache_layer, - block_tables.value(), - total_sequence_lengths.value(), - input_offsets.value(), - std::nullopt, - softmax_scale_); - } - } else { - infinicore::op::paged_attention_( - attn_output, - query_states, - k_cache_layer, - v_cache_layer, - block_tables.value(), - total_sequence_lengths.value(), - std::nullopt, - softmax_scale_); + const size_t num_requests = total_sequence_lengths.value()->shape()[0]; + ASSERT(num_requests > 0); + ASSERT_EQ(seq_len % num_requests, 0); + const int max_seqlen = static_cast(seq_len / num_requests); + + auto kv_b = kv_b_proj_->forward(kv_norm)->view( + {seq_len, num_attention_heads_, qk_nope_head_dim_ + v_head_dim_}); + auto key_nope = kv_b->narrow({{2, 0, qk_nope_head_dim_}})->contiguous(); + auto value_states = kv_b->narrow({{2, qk_nope_head_dim_, v_head_dim_}})->contiguous(); + auto key_pe = infinicore::op::broadcast_to( + k_pe_rope, + {static_cast(seq_len), + static_cast(num_attention_heads_), + static_cast(qk_rope_head_dim_)}) + ->contiguous(); + auto query_states = infinicore::op::cat({q_nope, q_pe}, 2); + auto key_states = infinicore::op::cat({key_nope, key_pe}, 2); + auto attn_output = infinicore::Tensor::empty( + {seq_len, num_attention_heads_, v_head_dim_}, query_states->dtype(), query_states->device()); + infinicore::op::mha_varlen_(attn_output, + query_states, + key_states, + value_states, + input_offsets.value(), + cu_seqlens.value(), + std::nullopt, + max_seqlen, + max_seqlen, + std::nullopt, + softmax_scale_); + auto projected = attn_output->view( + {batch_size, seq_len, num_attention_heads_ * v_head_dim_}); + return o_proj_->forward(projected); } + + auto q_latent = project_q_nope_to_latent_(q_nope); + auto query_states = infinicore::op::concat_mla_q(q_latent, q_pe); + auto attn_output = infinicore::Tensor::empty( + {seq_len, num_attention_heads_, kv_lora_rank_}, query_states->dtype(), query_states->device()); + const int64_t max_context_len = attn_metadata.max_context_len.value(); + infinicore::op::paged_attention_mla_(attn_output, + query_states, + kv_cache, + softmax_scale_, + block_tables.value(), + total_sequence_lengths.value(), + max_context_len); return project_latent_to_value_(attn_output, batch_size, seq_len); } diff --git a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.hpp b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.hpp index 26b42d97e..6044fca0e 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.hpp +++ b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.hpp @@ -2,7 +2,6 @@ #include "../../backends/attention_backends.hpp" #include "../../config/model_config.hpp" -#include "../../layers/attention/attention.hpp" #include "../../layers/linear/linear.hpp" #include "infinicore/nn/module.hpp" #include "infinicore/nn/rmsnorm.hpp" @@ -13,7 +12,7 @@ namespace infinilm::models::deepseek_v2 { -class DeepseekV2MLAAttention : public infinicore::nn::Module { +class DeepseekV2MLAAttention final : public infinicore::nn::Module { public: DeepseekV2MLAAttention(std::shared_ptr model_config, size_t layer_idx, @@ -23,10 +22,6 @@ class DeepseekV2MLAAttention : public infinicore::nn::Module { const infinicore::Tensor &hidden_states) const; private: - infinicore::Tensor forward_static_(const infinicore::Tensor &positions, - const infinicore::Tensor &hidden_states) const; - infinicore::Tensor forward_paged_(const infinicore::Tensor &positions, - const infinicore::Tensor &hidden_states) const; infinicore::Tensor position_ids_for_rope_(const infinicore::Tensor &position_ids) const; infinicore::Tensor kv_b_weight_3d_() const; infinicore::Tensor project_q_nope_to_latent_(const infinicore::Tensor &q_nope) const; @@ -41,19 +36,22 @@ class DeepseekV2MLAAttention : public infinicore::nn::Module { size_t qk_rope_head_dim_{0}; size_t q_head_dim_{0}; size_t v_head_dim_{0}; + size_t q_lora_rank_{0}; size_t kv_lora_rank_{0}; size_t mla_head_dim_{0}; float softmax_scale_{1.0f}; infinilm::backends::AttentionBackend attention_backend_; INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, q_a_proj); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, q_a_layernorm); + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_b_proj); INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, kv_a_proj_with_mqa); INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, kv_a_layernorm); INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, kv_b_proj); INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, o_proj); std::shared_ptr rotary_emb_; - std::shared_ptr latent_attn_; infinicore::nn::Parameter kv_cache_k_scale_; infinicore::nn::Parameter kv_cache_v_scale_; }; diff --git a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp b/csrc/models/deepseek_v2/deepseek_v2_moe.cpp index a18351ac8..f1e6c1f38 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_moe.cpp +++ b/csrc/models/deepseek_v2/deepseek_v2_moe.cpp @@ -2,132 +2,125 @@ #include "../../global_state/global_state.hpp" #include "../../utils.hpp" -#include "infinicore/ops.hpp" #include "infinicore/ops/distributed/allreduce.hpp" +#include "infinicore/ops/moe_argsort_bincount.hpp" +#include "infinicore/ops/moe_expand_input.hpp" +#include "infinicore/ops/moe_silu_and_mul_quant.hpp" +#include "infinicore/ops/moe_sum_vllm.hpp" +#include "infinicore/ops/w16a16_group_gemm.hpp" -#include "spdlog/spdlog.h" #include namespace infinilm::models::deepseek_v2 { -namespace { - -bool supports_fused_deepseek_moe(infinicore::Device::Type device_type) { - switch (device_type) { - case infinicore::Device::Type::NVIDIA: - case infinicore::Device::Type::ALI: - case infinicore::Device::Type::HYGON: - case infinicore::Device::Type::ILUVATAR: - case infinicore::Device::Type::METAX: - case infinicore::Device::Type::MOORE: - return true; - default: - return false; - } -} - -} // namespace - -DeepseekV2TopKRouter::DeepseekV2TopKRouter(std::shared_ptr model_config, - const infinicore::Device &device) { - const auto &dtype{model_config->get_dtype()}; - const size_t hidden_size = model_config->get("hidden_size"); - num_experts_ = model_config->get("num_experts"); - num_experts_per_tok_ = model_config->get("num_experts_per_tok"); - norm_topk_prob_ = model_config->get("norm_topk_prob"); - - ASSERT((num_experts_ > 0) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_experts_)); - INFINICORE_NN_PARAMETER_INIT(weight, ({num_experts_, hidden_size}, dtype, device)); -} - -std::tuple -DeepseekV2TopKRouter::forward(const infinicore::Tensor &hidden_states) const { - ASSERT(hidden_states->ndim() == 2); - const size_t ntoken = hidden_states->shape()[0]; - auto router_logits = infinicore::op::linear(hidden_states, weight_, std::nullopt, 1.0f); - auto router_scores = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::F32, hidden_states->device()); - auto router_indices = infinicore::Tensor::empty({ntoken, num_experts_per_tok_}, infinicore::DataType::I32, hidden_states->device()); - infinicore::op::topksoftmax(router_scores, router_indices, router_logits, num_experts_per_tok_, norm_topk_prob_); - return {router_scores, router_indices}; -} DeepseekV2Experts::DeepseekV2Experts(std::shared_ptr model_config, const infinicore::Device &device) { + const auto &dtype = model_config->get_dtype(); + if (dtype != infinicore::DataType::F16 && dtype != infinicore::DataType::BF16) { + throw std::runtime_error("DeepseekV2Experts requires fp16 or bfloat16 weights"); + } + if (model_config->get_or("hidden_act", "silu") != "silu") { + throw std::runtime_error("DeepseekV2Experts supports only SiLU activation"); + } + hidden_size_ = model_config->get("hidden_size"); - moe_intermediate_size_ = model_config->get("moe_intermediate_size"); num_experts_ = model_config->get("num_experts"); num_experts_per_tok_ = model_config->get("num_experts_per_tok"); - ASSERT((num_experts_ > 0) && (num_experts_per_tok_ > 0) && (num_experts_per_tok_ <= num_experts_)); + const size_t moe_intermediate_size = model_config->get("moe_intermediate_size"); const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); tp_size_ = static_cast(rank_info.tp_size); communicator_ = rank_info.comm; - - experts_.reserve(num_experts_); - gate_weights_.reserve(num_experts_); - up_weights_.reserve(num_experts_); - down_weights_.reserve(num_experts_); - for (size_t i = 0; i < num_experts_; ++i) { - auto expert = this->register_module(std::to_string(i), model_config, device); - gate_weights_.push_back(expert->gate_weight()); - up_weights_.push_back(expert->up_weight()); - down_weights_.push_back(expert->down_weight()); - experts_.push_back(std::move(expert)); + if (moe_intermediate_size % tp_size_ != 0) { + throw std::runtime_error("DeepseekV2Experts: moe_intermediate_size must be divisible by tp_size"); } - local_moe_intermediate_size_ = gate_weights_.empty() ? moe_intermediate_size_ : gate_weights_.front()->shape()[0]; -} - -infinicore::Tensor DeepseekV2Experts::forward_cpu_routed_(const infinicore::Tensor &hidden_states, - const infinicore::Tensor &top_k_index, - const infinicore::Tensor &top_k_weights) const { - auto top_k_weights_cpu = top_k_weights->to(infinicore::Device::Type::CPU); - auto top_k_index_cpu = top_k_index->to(infinicore::Device::Type::CPU); - auto *top_k_index_ptr = reinterpret_cast(top_k_index_cpu->data()); - auto *top_k_weights_ptr = reinterpret_cast(top_k_weights_cpu->data()); - - const size_t ntoken = hidden_states->shape()[0]; - auto final_hidden_states = infinicore::Tensor::empty(hidden_states->shape(), hidden_states->dtype(), hidden_states->device()); - for (size_t itok = 0; itok < ntoken; ++itok) { - auto hidden_states_i = hidden_states->narrow({{0, itok, 1}}); - const size_t route_row = itok * num_experts_per_tok_; - - infinicore::Tensor final_hidden_states_i; - for (size_t k = 0; k < num_experts_per_tok_; ++k) { - const int index = top_k_index_ptr[route_row + k]; - const float score = top_k_weights_ptr[route_row + k]; - ASSERT(index >= 0 && static_cast(index) < num_experts_); - experts_[index]->set_alpha(score); - auto expert_out = experts_[index]->forward(hidden_states_i); - if (k == 0) { - final_hidden_states_i = expert_out; - } else { - infinicore::op::add_(final_hidden_states_i, final_hidden_states_i, expert_out); - } - } - final_hidden_states->narrow({{0, itok, 1}})->copy_from(final_hidden_states_i); + local_moe_intermediate_size_ = moe_intermediate_size / tp_size_; + + w1_ = infinicore::Tensor::empty( + {num_experts_, local_moe_intermediate_size_ * 2, hidden_size_}, dtype, device); + w2_ = infinicore::Tensor::empty( + {num_experts_, hidden_size_, local_moe_intermediate_size_}, dtype, device); + + for (size_t expert_id = 0; expert_id < num_experts_; ++expert_id) { + const auto prefix = std::to_string(expert_id); + auto gate_weight = w1_->narrow({{0, expert_id, 1}, {1, 0, local_moe_intermediate_size_}}) + ->view({local_moe_intermediate_size_, hidden_size_}); + auto up_weight = w1_->narrow({{0, expert_id, 1}, {1, local_moe_intermediate_size_, local_moe_intermediate_size_}}) + ->view({local_moe_intermediate_size_, hidden_size_}); + auto down_weight = w2_->narrow({{0, expert_id, 1}}) + ->view({hidden_size_, local_moe_intermediate_size_}); + register_parameter(prefix + ".gate_proj.weight", + infinicore::nn::Parameter(gate_weight, 0, rank_info.tp_rank, rank_info.tp_size)); + register_parameter(prefix + ".up_proj.weight", + infinicore::nn::Parameter(up_weight, 0, rank_info.tp_rank, rank_info.tp_size)); + register_parameter(prefix + ".down_proj.weight", + infinicore::nn::Parameter(down_weight, 1, rank_info.tp_rank, rank_info.tp_size)); } - return final_hidden_states; } infinicore::Tensor DeepseekV2Experts::forward(const infinicore::Tensor &hidden_states, const infinicore::Tensor &top_k_index, - const infinicore::Tensor &top_k_weights) const { + const infinicore::Tensor &top_k_weights, + std::optional shared_output) const { ASSERT(hidden_states->ndim() == 2); - if (supports_fused_deepseek_moe(hidden_states->device().getType()) - && (hidden_states->dtype() == infinicore::DataType::BF16 - || hidden_states->dtype() == infinicore::DataType::F16)) { - try { - auto output = infinicore::op::deepseek_moe(hidden_states, top_k_index, top_k_weights, - gate_weights_, up_weights_, down_weights_, - local_moe_intermediate_size_, num_experts_); - if (tp_size_ > 1 && communicator_ != nullptr) { - infinicore::op::distributed::allreduce_(output, output, INFINICCL_SUM, communicator_); - } - return output; - } catch (const std::exception &e) { - spdlog::warn("DeepseekV2Experts: deepseek_moe unavailable on {}, falling back to CPU-routed experts: {}", - static_cast(hidden_states->device().getType()), e.what()); - } + const size_t num_tokens = hidden_states->size(0); + const size_t expanded_tokens = num_tokens * num_experts_per_tok_; + const auto &attn_metadata = infinilm::global_state::get_forward_context().attn_metadata; + const bool is_decode = attn_metadata.total_sequence_lengths.has_value() + && num_tokens == attn_metadata.total_sequence_lengths.value()->size(0); + + auto tokens_per_experts_gpu = infinicore::Tensor::empty( + {num_experts_}, infinicore::DataType::I32, hidden_states->device()); + auto sorted_indices = infinicore::Tensor::empty( + {expanded_tokens}, infinicore::DataType::I32, hidden_states->device()); + auto inv_pos = infinicore::Tensor::empty( + {expanded_tokens}, infinicore::DataType::I32, hidden_states->device()); + infinicore::op::moe_argsort_bincount_with_inv_pos_( + tokens_per_experts_gpu, sorted_indices, inv_pos, top_k_index, num_experts_); + auto tokens_per_experts = is_decode + ? tokens_per_experts_gpu + : tokens_per_experts_gpu->to(infinicore::Device::Type::CPU); + + auto expanded_input = infinicore::Tensor::empty( + {expanded_tokens, hidden_size_}, hidden_states->dtype(), hidden_states->device()); + infinicore::op::moe_expand_input_with_inv_pos_( + expanded_input, std::nullopt, hidden_states, inv_pos, num_experts_per_tok_, 128, 0); + + auto gate_up = infinicore::Tensor::empty( + {expanded_tokens, local_moe_intermediate_size_ * 2}, hidden_states->dtype(), hidden_states->device()); + infinicore::op::w16a16_group_gemm_(gate_up, + expanded_input, + w1_, + tokens_per_experts, + std::nullopt, + std::nullopt, + true, + is_decode); + + auto activated = infinicore::Tensor::empty( + {expanded_tokens, local_moe_intermediate_size_}, hidden_states->dtype(), hidden_states->device()); + infinicore::op::moe_silu_and_mul_quant_(activated, std::nullopt, gate_up, 0); + + auto expert_output = infinicore::Tensor::empty( + {expanded_tokens, hidden_size_}, hidden_states->dtype(), hidden_states->device()); + infinicore::op::w16a16_group_gemm_(expert_output, + activated, + w2_, + tokens_per_experts, + sorted_indices, + std::nullopt, + true, + is_decode); + + auto output = infinicore::Tensor::empty( + {num_tokens, hidden_size_}, hidden_states->dtype(), hidden_states->device()); + infinicore::op::moe_sum_vllm_(output, + expert_output->view({num_tokens, num_experts_per_tok_, hidden_size_}), + top_k_weights, + shared_output); + if (tp_size_ > 1 && communicator_ != nullptr) { + infinicore::op::distributed::allreduce_(output, output, INFINICCL_SUM, communicator_); } - return forward_cpu_routed_(hidden_states, top_k_index, top_k_weights); + return output; } DeepseekV2MoE::DeepseekV2MoE(std::shared_ptr model_config, @@ -140,6 +133,7 @@ DeepseekV2MoE::DeepseekV2MoE(std::shared_ptr mode if (has_shared_experts_) { auto shared_config_json = model_config->get_config_json(); shared_config_json["intermediate_size"] = model_config->get("moe_intermediate_size") * n_shared_experts; + shared_config_json["reduce_results"] = false; auto shared_config = std::make_shared(shared_config_json); INFINICORE_NN_MODULE_INIT(shared_experts, shared_config, device); } @@ -148,15 +142,14 @@ DeepseekV2MoE::DeepseekV2MoE(std::shared_ptr mode infinicore::Tensor DeepseekV2MoE::forward(const infinicore::Tensor &hidden_states) const { ASSERT(hidden_states->ndim() == 3); const auto shape = hidden_states->shape(); - auto hidden_states_reshaped = hidden_states->view({shape[0] * shape[1], shape[2]}); + auto flat_hidden_states = hidden_states->view({shape[0] * shape[1], shape[2]}); + auto [routing_weights, selected_experts] = gate_->forward(flat_hidden_states); - auto [routing_weights, selected_experts] = gate_->forward(hidden_states_reshaped); - auto final_hidden_states = experts_->forward(hidden_states_reshaped, selected_experts, routing_weights)->view(shape); + std::optional shared_output; if (has_shared_experts_) { - auto shared_out = shared_experts_->forward(hidden_states); - final_hidden_states = infinicore::op::add(final_hidden_states, shared_out); + shared_output = shared_experts_->forward(hidden_states)->view({shape[0] * shape[1], shape[2]}); } - return final_hidden_states; + return experts_->forward(flat_hidden_states, selected_experts, routing_weights, shared_output)->view(shape); } } // namespace infinilm::models::deepseek_v2 diff --git a/csrc/models/deepseek_v2/deepseek_v2_moe.hpp b/csrc/models/deepseek_v2/deepseek_v2_moe.hpp index 11af3669c..e838172da 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_moe.hpp +++ b/csrc/models/deepseek_v2/deepseek_v2_moe.hpp @@ -1,58 +1,38 @@ #pragma once #include "../../config/model_config.hpp" -#include "../../layers/common_modules.hpp" -#include "../../layers/linear/linear.hpp" #include "../../layers/mlp/mlp.hpp" -#include "../../layers/moe/legacy/moe_mlp.hpp" +#include "../../layers/moe/router/topk_router.hpp" #include "infinicore/device.hpp" #include "infinicore/nn/module.hpp" #include "infinicore/tensor.hpp" #include #include +#include +#include #include -#include namespace infinilm::models::deepseek_v2 { using DeepseekV2MLP = infinilm::layers::mlp::MLP; -using DeepseekV2ExpertMLP = infinilm::layers::moe::legacy::MoeMLP; -class DeepseekV2TopKRouter : public infinicore::nn::Module { -public: - DeepseekV2TopKRouter(std::shared_ptr model_config, - const infinicore::Device &device); - - std::tuple forward(const infinicore::Tensor &hidden_states) const; +using DeepseekV2TopKRouter = infinilm::layers::moe::TopKRouter; -protected: - INFINICORE_NN_PARAMETER(weight); - size_t num_experts_per_tok_{0}; - size_t num_experts_{0}; - bool norm_topk_prob_{false}; -}; - -class DeepseekV2Experts : public infinicore::nn::Module { +class DeepseekV2Experts final : public infinicore::nn::Module { public: DeepseekV2Experts(std::shared_ptr model_config, const infinicore::Device &device); infinicore::Tensor forward(const infinicore::Tensor &hidden_states, const infinicore::Tensor &top_k_index, - const infinicore::Tensor &top_k_weights) const; - -protected: - infinicore::Tensor forward_cpu_routed_(const infinicore::Tensor &hidden_states, - const infinicore::Tensor &top_k_index, - const infinicore::Tensor &top_k_weights) const; + const infinicore::Tensor &top_k_weights, + std::optional shared_output = std::nullopt) const; - INFINICORE_NN_MODULE_VEC(DeepseekV2ExpertMLP, experts); - std::vector gate_weights_; - std::vector up_weights_; - std::vector down_weights_; +private: + infinicore::Tensor w1_; + infinicore::Tensor w2_; size_t hidden_size_{0}; - size_t moe_intermediate_size_{0}; size_t local_moe_intermediate_size_{0}; size_t num_experts_per_tok_{0}; size_t num_experts_{0}; @@ -60,14 +40,14 @@ class DeepseekV2Experts : public infinicore::nn::Module { infinicclComm_t communicator_{nullptr}; }; -class DeepseekV2MoE : public infinicore::nn::Module { +class DeepseekV2MoE final : public infinicore::nn::Module { public: DeepseekV2MoE(std::shared_ptr model_config, const infinicore::Device &device); infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; -protected: +private: INFINICORE_NN_MODULE(DeepseekV2TopKRouter, gate); INFINICORE_NN_MODULE(DeepseekV2Experts, experts); INFINICORE_NN_MODULE(DeepseekV2MLP, shared_experts); diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index b73ce06ae..a8c72fa37 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -150,6 +150,7 @@ inline void bind_infer_engine(py::module &m) { std::optional> visual_token_ranges, std::optional target_hidden_states, bool sample_all_positions, + std::optional max_context_len, py::kwargs kwargs) { InferEngine::Input input{ std::move(input_ids), @@ -170,6 +171,7 @@ inline void bind_infer_engine(py::module &m) { std::move(visual_token_ranges), std::move(target_hidden_states), sample_all_positions, + max_context_len, }; // Explicit defaults @@ -220,7 +222,8 @@ inline void bind_infer_engine(py::module &m) { py::arg("image_req_ids") = std::nullopt, py::arg("visual_token_ranges") = std::nullopt, py::arg("target_hidden_states") = std::nullopt, - py::arg("sample_all_positions") = false) + py::arg("sample_all_positions") = false, + py::arg("max_context_len") = std::nullopt) .def_readwrite("input_ids", &InferEngine::Input::input_ids) .def_readwrite("position_ids", &InferEngine::Input::position_ids) .def_readwrite("past_sequence_lengths", &InferEngine::Input::past_sequence_lengths) @@ -238,6 +241,7 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("image_req_ids", &InferEngine::Input::image_req_ids) .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) + .def_readwrite("max_context_len", &InferEngine::Input::max_context_len) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) diff --git a/examples/bench.py b/examples/bench.py index bd424e836..a23fa8e9b 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -8,6 +8,7 @@ import numpy as np from infinilm.base_config import BaseConfig from infinilm.cache import PagedKVCacheConfig, StaticKVCacheConfig +from infinilm.config.engine_config import EngineConfig from infinilm.distributed import DistConfig from infinilm.infer_engine import GenerationConfig, InferEngine from infinilm.llm.llm import LLM @@ -370,7 +371,19 @@ def run( device_str = cfg.get_device_str(cfg.device) - _PAGED_KV_BLOCK_SIZE = cfg.block_size + runtime_config = EngineConfig( + model_path=cfg.model, + device=device_str, + tensor_parallel_size=cfg.tp, + cache_type="paged" if cfg.enable_paged_attn else "static", + block_size=cfg.block_size, + enable_graph=cfg.enable_graph, + attn_backend=cfg.attn, + use_mla=cfg.use_mla, + weight_load_mode=cfg.weight_load_mode, + ) + + _PAGED_KV_BLOCK_SIZE = runtime_config.block_size # -------------------------------------------------------- # # 解析参数 # -------------------------------------------------------- # @@ -404,7 +417,7 @@ def run( output_len = [output_len] cases_dict = get_test_cases( - model_path, batch_size, input_len, output_len, use_mla=cfg.use_mla + model_path, batch_size, input_len, output_len, use_mla=runtime_config.use_mla ) # -------------------------------------------------------- # # 测试 @@ -425,8 +438,10 @@ def run( else: cache_config = None - if enable_paged_attn and attn_backend == "default": - attn_backend = "paged-attn" + if enable_paged_attn: + attn_backend = runtime_config.attn_backend + if attn_backend == "default": + attn_backend = "paged-attn" test = TestModel( model_path, @@ -438,7 +453,7 @@ def run( cache_config=cache_config, enable_graph=enable_graph, attn_backend=attn_backend, - use_mla=cfg.use_mla, + use_mla=runtime_config.use_mla, weight_load_mode=cfg.weight_load_mode, moe_ep_backend=moe_ep_backend, moe_ep_size=ep, diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 3b8f6533d..1e97c05a8 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -226,7 +226,7 @@ def _add_common_args(self): self.parser.add_argument( "--use-mla", action="store_true", - help="use DeepSeek V2 MLA attention when supported", + help="deprecated compatibility flag; MLA models are detected automatically", ) self.parser.add_argument( "--enable-paged-attn", diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index 1bb1f733e..8ccee1b73 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -1,4 +1,6 @@ +import json from dataclasses import dataclass +from pathlib import Path from typing import Optional from infinilm.config.kv_transfer import KVTransferConfig @@ -62,10 +64,29 @@ class EngineConfig: def __post_init__(self) -> None: if self.num_draft_tokens < 1: raise ValueError("num_draft_tokens must be >= 1") - if self.weight_load_mode not in {"async", "sync"}: raise ValueError("weight_load_mode must be either 'async' or 'sync'") + config_path = Path(self.model_path) / "config.json" + try: + with config_path.open("r", encoding="utf-8") as config_file: + model_config = json.load(config_file) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Unable to read model config: {config_path}") from exc + text_config = model_config.get("text_config", model_config) + model_type = text_config.get("model_type", model_config.get("model_type")) + + # DeepSeek V2 exposes only the verified MLA path. Keep use_mla as a + # compatibility input, but derive the effective capability from the model. + self.use_mla = self.use_mla or model_type == "deepseek_v2" + if model_type == "deepseek_v2": + if self.cache_type != "paged": + raise ValueError( + "DeepSeek V2 MLA requires paged cache; pass --enable-paged-attn" + ) + self.block_size = 16 + self.attn_backend = "flash-attn" + if ( self.kv_transfer_config is not None and self.kv_transfer_config.kv_connector diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 3e874046e..1fb62245b 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -230,6 +230,7 @@ def _build_input( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=False, + max_context_len=None, temperature=None, top_k=None, top_p=None, @@ -288,6 +289,7 @@ def convert_tensor_list(tensor_list_): visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, + max_context_len=max_context_len, temperature=temperature, top_k=top_k, top_p=top_p, @@ -313,6 +315,7 @@ def forward( image_req_ids=None, visual_token_ranges=None, target_hidden_states=None, + max_context_len=None, temperature=None, top_k=None, top_p=None, @@ -385,6 +388,7 @@ def convert_tensor_list(tensor_list_): image_req_ids=image_req_ids, visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, + max_context_len=max_context_len, temperature=temperature, top_k=top_k, top_p=top_p, @@ -414,6 +418,7 @@ def forward_raw( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=True, + max_context_len=None, temperature=None, top_k=None, top_p=None, @@ -436,6 +441,7 @@ def forward_raw( visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, + max_context_len=max_context_len, temperature=temperature, top_k=top_k, top_p=top_p, @@ -593,6 +599,7 @@ def generate( cu_seqlens=cu_seqlens, block_tables=block_tables, slot_mapping=slot_mapping, + max_context_len=past_seq_len + seq_len, mamba_init_state_indices=mamba_init_state_indices, mamba_final_state_indices=mamba_final_state_indices, image_bound=image_bound if iter == 0 else None, diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index a6fbc33ac..b3bef06d5 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -153,6 +153,7 @@ def _build_model_input_from_static_scheduler_output( ), "block_tables": None, "slot_mapping": None, + "max_context_len": total_kv_len, "temperature": temperature, "top_k": top_k, "top_p": top_p, @@ -257,6 +258,7 @@ def _build_model_input_from_batch_scheduler_output( "cu_seqlens": infinicore.from_list(cu_seqlens, dtype=infinicore.int32), "block_tables": infinicore.from_list(block_tables, dtype=infinicore.int32), "slot_mapping": infinicore.from_list(slot_mapping, dtype=infinicore.int64), + "max_context_len": max(seq_lens), "temperature": temperature, "top_k": top_k, "top_p": top_p, From 2cbe88dc01e4feb61acf1a9793f7df289b9f6bef Mon Sep 17 00:00:00 2001 From: wooway777 Date: Thu, 23 Jul 2026 08:36:03 +0000 Subject: [PATCH 3/5] pepe: glm 5.2 service --- csrc/engine/compiler/paged_compiler.cpp | 472 +++++++++++++----- csrc/engine/compiler/paged_compiler.hpp | 14 +- .../compiler/static_batching_compiler.cpp | 1 + .../distributed/communication_group.cpp | 87 +++- .../distributed/communication_group.hpp | 29 +- csrc/engine/distributed/dist_config.cpp | 25 +- csrc/engine/distributed/dist_config.hpp | 8 +- csrc/engine/infer_engine.cpp | 10 +- csrc/engine/rank_barrier.cpp | 9 +- csrc/engine/rank_barrier.hpp | 5 +- csrc/engine/rank_worker.cpp | 85 +++- csrc/engine/rank_worker.hpp | 11 + csrc/global_state/forward_context.hpp | 26 +- csrc/layers/linear/fused_linear.cpp | 72 +++ csrc/layers/linear/fused_linear.hpp | 23 + csrc/layers/quantization/glm_w8a8.cpp | 31 +- csrc/layers/quantization/glm_w8a8.hpp | 17 +- .../rotary_embedding/rotary_embedding.cpp | 16 + .../rotary_embedding/rotary_embedding.hpp | 3 + .../deepseek_v2/deepseek_v2_indexer.cpp | 260 ++++++++++ .../deepseek_v2/deepseek_v2_indexer.hpp | 53 ++ .../deepseek_v2/deepseek_v2_mla_attention.cpp | 324 ++++++++++-- .../deepseek_v2/deepseek_v2_mla_attention.hpp | 15 +- .../glm_dsa_allocate_cache_tensors.cpp | 132 +++++ .../glm_dsa_allocate_cache_tensors.hpp | 26 + csrc/models/glm_moe_dsa/glm_model.cpp | 422 +++++++++++++++- csrc/models/glm_moe_dsa/glm_model.hpp | 44 +- csrc/models/glm_moe_dsa/glm_moe.cpp | 23 +- csrc/models/glm_moe_dsa/glm_moe.hpp | 6 +- .../models/glm_moe_dsa/glm_vocab_parallel.cpp | 14 +- csrc/models/infinilm_model.hpp | 10 + csrc/pybind11/engine/engine.hpp | 28 +- csrc/utils.hpp | 38 +- python/infinilm/base_config.py | 40 ++ python/infinilm/config/engine_config.py | 77 ++- python/infinilm/distributed/dist_config.py | 9 +- python/infinilm/infer_engine.py | 57 ++- python/infinilm/llm/llm.py | 87 +++- .../infinilm/llm/model_runner/model_runner.py | 22 +- python/infinilm/llm/request.py | 3 + python/infinilm/llm/scheduler.py | 290 ++++++++--- python/infinilm/modeling_utils.py | 8 +- .../processors/basic_llm_processor.py | 33 +- python/infinilm/processors/processor.py | 24 + .../infinilm/processors/qwen3_5_processor.py | 2 + python/infinilm/server/inference_server.py | 280 ++++++++--- 46 files changed, 2807 insertions(+), 464 deletions(-) create mode 100644 csrc/models/deepseek_v2/deepseek_v2_indexer.cpp create mode 100644 csrc/models/deepseek_v2/deepseek_v2_indexer.hpp create mode 100644 csrc/models/glm_moe_dsa/glm_dsa_allocate_cache_tensors.cpp create mode 100644 csrc/models/glm_moe_dsa/glm_dsa_allocate_cache_tensors.hpp diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index 9f534b087..af6d428b4 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -2,156 +2,360 @@ #include "../../global_state/global_state.hpp" #include "../../utils.hpp" +#include + namespace infinilm::engine { -PagedCompiler::PagedCompiler(const std::shared_ptr &model, RankBarrier *barrier) +PagedCompiler::PagedCompiler( + const std::shared_ptr &model, + RankBarrier *barrier) : GraphCompiler(model, barrier) { - for (size_t b = 1; b < 64; ++b) { - decode_batch_sizes_.push_back(b); - } - for (size_t b = 64; b < 128; b += 16) { - decode_batch_sizes_.push_back(b); - } - for (size_t b = 128; b < 256; b += 32) { - decode_batch_sizes_.push_back(b); +} + +void PagedCompiler::compile() { + compiled_map_decode_.clear(); + graph_disabled_batches_.clear(); + initialized_ = false; + num_blocks_ = 0; + block_size_ = 0; + + const auto *paged_config = dynamic_cast( + model_->get_cache_config()); + if (paged_config == nullptr) { + return; } - for (size_t b = 256; b <= 512; b += 64) { - decode_batch_sizes_.push_back(b); + num_blocks_ = paged_config->num_blocks(); + block_size_ = paged_config->block_size(); + initialized_ = num_blocks_ > 0 && block_size_ > 0; +} + +InfinilmModel::Input PagedCompiler::make_decode_input( + size_t b, + size_t block_per_req) const { + InfinilmModel::Input input; + const auto device = infinicore::context::getDevice(); + input.input_ids = infinicore::Tensor::empty( + {1, b}, infinicore::DataType::I64, device); + input.position_ids = infinicore::Tensor::empty( + {b}, infinicore::DataType::I64, device); + input.total_sequence_lengths = infinicore::Tensor::empty( + {b}, infinicore::DataType::I32, device); + set_zeros(input.input_ids.value()); + set_zeros(input.position_ids.value()); + std::vector total_sequence_lengths_vec(b, 1); + infinicore::context::memcpyH2D( + input.total_sequence_lengths.value()->data(), + total_sequence_lengths_vec.data(), + b * sizeof(int32_t), + false); + + input.input_offsets = infinicore::Tensor::empty( + {b + 1}, infinicore::DataType::I32, device); + std::vector input_offsets_vec(b + 1, 0); + for (size_t i = 0; i <= b; ++i) { + input_offsets_vec[i] = static_cast(i); } + infinicore::context::memcpyH2D( + input.input_offsets.value()->data(), + input_offsets_vec.data(), + (b + 1) * sizeof(int32_t), + false); + input.request_ids = infinicore::Tensor::empty( + {b}, infinicore::DataType::I32, device); + infinicore::context::memcpyH2D( + input.request_ids.value()->data(), + input_offsets_vec.data(), + b * sizeof(int32_t), + false); + input.cu_seqlens = infinicore::Tensor::empty( + {b + 1}, infinicore::DataType::I32, device); + infinicore::context::memcpyH2D( + input.cu_seqlens.value()->data(), + input_offsets_vec.data(), + (b + 1) * sizeof(int32_t), + false); + input.block_tables = infinicore::Tensor::empty( + {b, block_per_req}, infinicore::DataType::I32, device); + set_zeros(input.block_tables.value()); + input.slot_mapping = infinicore::Tensor::empty( + {b}, infinicore::DataType::I64, device); + // Graph::instantiate runs several warmups. Padding slots prevent those + // dummy forwards from modifying a live request's MLA/indexer KV caches. + set_minus_one_device_async(input.slot_mapping.value()); + + infinilm::global_state::get_forward_context().attn_metadata = { + input.past_sequence_lengths, + input.total_sequence_lengths, + input.input_offsets, + input.request_ids, + input.cu_seqlens, + input.block_tables, + input.slot_mapping, + static_cast(block_per_req * block_size_), + }; + return input; } -void PagedCompiler::compile() { - if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { - size_t nblocks = dynamic_cast(model_->get_cache_config())->num_blocks(); - size_t block_size = dynamic_cast(model_->get_cache_config())->block_size(); - size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); - compiled_map_decode_.clear(); - block_tables_holder_ = infinicore::Tensor::empty( - {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice()); - set_zeros(block_tables_holder_); - - auto make_decode_input = [&](size_t b) { - InfinilmModel::Input input; - input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); - input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice()); - set_zeros(input.input_ids.value()); - set_zeros(input.position_ids.value()); - set_zeros(input.total_sequence_lengths.value()); - std::vector total_sequence_lengths_vec(b, 1); - infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int32_t), false); - input.input_offsets = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); - std::vector input_offsets_vec(b + 1, 0); - for (size_t i = 0; i <= b; i++) { - input_offsets_vec[i] = i; - } - infinicore::context::memcpyH2D(input.input_offsets.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); - input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); - infinicore::context::memcpyH2D(input.cu_seqlens.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); - const size_t block_per_req = nblocks; - input.block_tables = block_tables_holder_->as_strided({b, block_per_req}, {(ptrdiff_t)block_per_req, 1}); - input.slot_mapping = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); - set_zeros(input.slot_mapping.value()); - - // Attention reads attn_metadata from thread-local forward context. - infinilm::global_state::get_forward_context().attn_metadata = { - input.past_sequence_lengths, - input.total_sequence_lengths, - input.input_offsets, - input.cu_seqlens, - input.block_tables, - input.slot_mapping, - static_cast(nblocks * block_size), - }; - return input; - }; - - { - const size_t warmup_batch_size = std::min(max_batch_size, static_cast(64)); - auto input = make_decode_input(warmup_batch_size); - model_->forward(input); - infinicore::context::syncStream(); - // Warmup runs the eager Marlin path and may leave per-layer lock - // workspaces dirty. Reset before CUDA graph capture so capture - // starts from the same all-zero lock state as normal execution. - model_->reset_runtime_state(); - infinicore::context::syncStream(); +void PagedCompiler::compile_decode( + size_t batch_size, + size_t block_per_req) { + if (!initialized_ || batch_size == 0 + || batch_size > model_->max_decode_graph_batch_size() + || block_per_req == 0 || block_per_req > num_blocks_ + || compiled_map_decode_.find(batch_size) != compiled_map_decode_.end() + || graph_disabled_batches_.find(batch_size) + != graph_disabled_batches_.end()) { + return; + } + + auto input = make_decode_input(batch_size, block_per_req); + + // Warm the exact static decode shape, then capture it once. All ranks + // enter warmup/capture in lock-step so TP collectives are instantiated in + // the same order. + barrier_->wait(); + auto fail_collectively = [&](bool local_ok, + const char *stage, + const std::string &local_error) { + const bool all_ok = barrier_->wait(local_ok); + if (all_ok) { + return false; + } + if (infinicore::context::isGraphRecording()) { + infinicore::context::cancelGraphRecording(); } + graph_disabled_batches_.insert(batch_size); + const auto &rank = infinilm::global_state::get_tensor_model_parallel_rank_info(); + if (!local_ok) { + spdlog::warn( + "[{}] disabling decode graph batch {} after {} failure: {}", + rank.tp_rank, + batch_size, + stage, + local_error); + } else if (rank.tp_rank == 0) { + spdlog::warn( + "disabling decode graph batch {} after a peer failed during {}", + batch_size, + stage); + } + return true; + }; + + bool local_ok = true; + std::string local_error; + try { + (void)model_->forward(input); + infinicore::context::syncStream(); + model_->reset_runtime_state(); + infinicore::context::syncStream(); + } catch (const std::exception &e) { + local_ok = false; + local_error = e.what(); + } catch (...) { + local_ok = false; + local_error = "unknown exception"; + } + if (fail_collectively(local_ok, "warmup", local_error)) { + return; + } + + InfinilmModel::Output output; + local_ok = true; + local_error.clear(); + try { + infinicore::context::startGraphRecording(); + // Graph-aware memsets put Marlin lock resets inside every replay instead + // of paying separate eager launches from get_compiled(). + model_->reset_runtime_state(); + output = model_->forward(input); + } catch (const std::exception &e) { + local_ok = false; + local_error = e.what(); + } catch (...) { + local_ok = false; + local_error = "unknown exception"; + } + if (fail_collectively(local_ok, "recording", local_error)) { + return; + } - for (size_t b : decode_batch_sizes_) { - auto input = make_decode_input(b); - - barrier_->wait(); - (void)model_->forward(input); - infinicore::context::syncStream(); - // Capture must not start with stale Marlin locks from previous - // warmup/capture attempts. This reset is intentionally outside - // graph capture; the current implementation still pays a memset - // before every graph replay in get_compiled(). - model_->reset_runtime_state(); - infinicore::context::syncStream(); - infinicore::context::startGraphRecording(); - auto output = model_->forward(input); - auto graph = infinicore::context::stopGraphRecording(); - barrier_->wait(); - - auto shared_output = std::shared_ptr( - new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)}); - - compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)}; + std::shared_ptr graph; + local_ok = true; + local_error.clear(); + try { + graph = infinicore::context::stopGraphRecording(); + if (graph == nullptr) { + throw std::runtime_error("graph recording returned no graph"); } + } catch (const std::exception &e) { + local_ok = false; + local_error = e.what(); + } catch (...) { + local_ok = false; + local_error = "unknown exception"; + } + if (fail_collectively(local_ok, "instantiation", local_error)) { + return; + } + + auto shared_output = std::make_shared( + InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)}); + auto padding_total_sequence_lengths = infinicore::Tensor::empty( + {batch_size}, infinicore::DataType::I32, + infinicore::context::getDevice()); + padding_total_sequence_lengths->copy_from( + input.total_sequence_lengths.value()); + auto padding_request_ids = infinicore::Tensor::empty( + {batch_size}, infinicore::DataType::I32, infinicore::context::getDevice()); + padding_request_ids->copy_from(input.request_ids.value()); + compiled_map_decode_[batch_size] = CompiledResult{ + std::move(input), + std::make_tuple(std::move(graph), std::move(shared_output)), + std::move(padding_total_sequence_lengths), + std::move(padding_request_ids), + {}, + }; + if (global_state::get_tensor_model_parallel_rank_info().tp_rank == 0) { + spdlog::info("compiled paged decode graph for batch {}", batch_size); } } -PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &input) { - if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { - size_t batch_size = input.block_tables.value()->size(0); - size_t block_per_req = input.block_tables.value()->size(1); - - // only support decode only batch - if (batch_size != input.input_ids.value()->size(1)) { - return {nullptr, nullptr}; - } else { - auto result = compiled_map_decode_.find(batch_size); - if (result == compiled_map_decode_.end()) { - return {nullptr, nullptr}; - } - auto &graph_input = result->second.input; - - graph_input.input_ids.value()->copy_from(input.input_ids.value()); - graph_input.position_ids.value()->copy_from(input.position_ids.value()); - graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value()); - graph_input.input_offsets.value()->copy_from(input.input_offsets.value()); - graph_input.cu_seqlens.value()->copy_from(input.cu_seqlens.value()); - - const size_t compiled_block_per_req = graph_input.block_tables.value()->size(1); - if (block_per_req > compiled_block_per_req) { - // Runtime width exceeds compiled graph slot; fall back to eager path. - return {nullptr, nullptr}; - } - - // Initialize only the active graph rows to -1, then overwrite the - // runtime logical region. Avoid clearing the full preallocated - // holder on every decode token. - auto &graph_block_tables = graph_input.block_tables.value(); - set_minus_one_device_async(graph_block_tables); - graph_block_tables->narrow({{1, 0, block_per_req}})->copy_from(input.block_tables.value()); - graph_input.slot_mapping.value()->copy_from(input.slot_mapping.value()); - // CUDA graph replay reuses the same per-layer Marlin workspaces. - // The graph itself does not contain a workspace reset, so enqueue - // one on the same stream before launch. This is correct but costs - // decode latency; the intended follow-up is a reusable global - // zero workspace/lock buffer shared by all Marlin layers. - model_->reset_runtime_state(); - - auto graph = std::get<0>(result->second.compiled); - auto shared_output = std::shared_ptr(new InfinilmModel::Output{std::get<1>(result->second.compiled)->logits->resume_from_blob_()}); - - return std::make_tuple(graph, shared_output); +PagedCompiler::Compiled PagedCompiler::get_compiled( + const InfinilmModel::Input &input) { + if (!initialized_ || !input.block_tables.has_value() + || !input.input_ids.has_value()) { + return {nullptr, nullptr}; + } + const size_t batch_size = input.block_tables.value()->size(0); + const size_t graph_batch_size = model_->decode_graph_batch_size(batch_size); + const size_t block_per_req = input.block_tables.value()->size(1); + + // One input token per active request is the decode-only graph contract. + // Prefill, mixed batches, oversized batches, and dynamic widths stay eager. + if (batch_size == 0 || batch_size > model_->max_decode_graph_batch_size() + || batch_size != input.input_ids.value()->size(1)) { + return {nullptr, nullptr}; + } + if (graph_batch_size < batch_size + || graph_batch_size > model_->max_decode_graph_batch_size()) { + return {nullptr, nullptr}; + } + if (block_per_req == 0 || block_per_req > num_blocks_) { + return {nullptr, nullptr}; + } + if (graph_disabled_batches_.find(graph_batch_size) + != graph_disabled_batches_.end()) { + return {nullptr, nullptr}; + } + if (compiled_map_decode_.find(graph_batch_size) + == compiled_map_decode_.end()) { + if (batch_size != graph_batch_size + && global_state::get_tensor_model_parallel_rank_info().tp_rank == 0) { + spdlog::info( + "padding paged decode graph batch {} to bucket {}", + batch_size, + graph_batch_size); } - } else { + // Match vLLM's persistent block table: one fixed-width graph input is + // reused as requests grow across cache blocks. Capturing the first + // observed width would force all later, longer decode steps eager. + compile_decode(graph_batch_size, num_blocks_); + } + auto result = compiled_map_decode_.find(graph_batch_size); + if (result == compiled_map_decode_.end()) { return {nullptr, nullptr}; } + auto &graph_input = result->second.input; + + const size_t compiled_block_per_req = graph_input.block_tables.value()->size(1); + if (block_per_req > compiled_block_per_req) { + return {nullptr, nullptr}; + } + + const size_t padding_size = graph_batch_size - batch_size; + auto copy_prefix = [batch_size]( + infinicore::Tensor &dst, + const infinicore::Tensor &src, + size_t dim) { + dst->narrow({{dim, 0, batch_size}})->copy_from(src); + }; + copy_prefix(graph_input.input_ids.value(), input.input_ids.value(), 1); + copy_prefix(graph_input.position_ids.value(), input.position_ids.value(), 0); + copy_prefix( + graph_input.total_sequence_lengths.value(), + input.total_sequence_lengths.value(), + 0); + copy_prefix(graph_input.request_ids.value(), input.request_ids.value(), 0); + + // Decode-only offsets are canonical [0, 1, ..., batch]. Keep the graph + // bucket's preinitialized dummy suffix and update only the active prefix. + graph_input.input_offsets.value() + ->narrow({{0, 0, batch_size + 1}}) + ->copy_from(input.input_offsets.value()); + graph_input.cu_seqlens.value() + ->narrow({{0, 0, batch_size + 1}}) + ->copy_from(input.cu_seqlens.value()); + + if (padding_size > 0) { + auto input_ids_tail = graph_input.input_ids.value()->narrow( + {{1, batch_size, padding_size}}); + auto position_ids_tail = graph_input.position_ids.value()->narrow( + {{0, batch_size, padding_size}}); + set_zeros_device_async(input_ids_tail); + set_zeros_device_async(position_ids_tail); + graph_input.total_sequence_lengths.value() + ->narrow({{0, batch_size, padding_size}}) + ->copy_from(result->second.padding_total_sequence_lengths->narrow( + {{0, batch_size, padding_size}})); + graph_input.request_ids.value() + ->narrow({{0, batch_size, padding_size}}) + ->copy_from(result->second.padding_request_ids->narrow( + {{0, batch_size, padding_size}})); + } + + auto &graph_block_tables = graph_input.block_tables.value(); + const size_t staging_key = batch_size * (compiled_block_per_req + 1) + block_per_req; + auto staging_it = result->second.block_tables_staging.find(staging_key); + if (staging_it == result->second.block_tables_staging.end()) { + staging_it = result->second.block_tables_staging.emplace( + staging_key, + infinicore::Tensor::empty( + {batch_size, block_per_req}, + infinicore::DataType::I32, + infinicore::context::getDevice())) + .first; + } + auto &block_tables_staging = staging_it->second; + block_tables_staging->copy_from(input.block_tables.value()); + + // The destination prefix is strided when the runtime width is smaller + // than the graph width. Keep a persistent device source so the queued + // rearrange never observes a freed temporary allocation. + set_minus_one_device_async(graph_block_tables); + graph_block_tables + ->narrow( + {{0, 0, batch_size}, + {1, 0, block_per_req}}) + ->copy_from(block_tables_staging); + graph_input.slot_mapping.value() + ->narrow({{0, 0, batch_size}}) + ->copy_from(input.slot_mapping.value()); + if (padding_size > 0) { + // Dummy attention may read block 0, but must never write any cache. + auto block_tables_tail = graph_block_tables->narrow( + {{0, batch_size, padding_size}}); + auto slot_mapping_tail = graph_input.slot_mapping.value()->narrow( + {{0, batch_size, padding_size}}); + set_zeros_device_async(block_tables_tail); + set_minus_one_device_async(slot_mapping_tail); + } + + auto graph = std::get<0>(result->second.compiled); + auto shared_output = std::make_shared( + InfinilmModel::Output{ + std::get<1>(result->second.compiled)->logits->resume_from_blob_()}); + return std::make_tuple(std::move(graph), std::move(shared_output)); } } // namespace infinilm::engine diff --git a/csrc/engine/compiler/paged_compiler.hpp b/csrc/engine/compiler/paged_compiler.hpp index a1125864d..c9a303182 100644 --- a/csrc/engine/compiler/paged_compiler.hpp +++ b/csrc/engine/compiler/paged_compiler.hpp @@ -3,6 +3,7 @@ #include "graph_compiler.hpp" #include +#include namespace infinilm::engine { class PagedCompiler : public GraphCompiler { @@ -14,18 +15,25 @@ class PagedCompiler : public GraphCompiler { Compiled get_compiled(const InfinilmModel::Input &input) override; private: - std::vector decode_batch_sizes_; + InfinilmModel::Input make_decode_input(size_t batch_size, size_t block_per_req) const; + void compile_decode(size_t batch_size, size_t block_per_req); - infinicore::Tensor block_tables_holder_; + bool initialized_ = false; + size_t num_blocks_ = 0; + size_t block_size_ = 0; struct CompiledResult { InfinilmModel::Input input; Compiled compiled; + infinicore::Tensor padding_total_sequence_lengths; + infinicore::Tensor padding_request_ids; + std::unordered_map block_tables_staging; }; std::unordered_map< - size_t, // num_requests + size_t, // static graph batch bucket CompiledResult> compiled_map_decode_; + std::unordered_set graph_disabled_batches_; }; } // namespace infinilm::engine diff --git a/csrc/engine/compiler/static_batching_compiler.cpp b/csrc/engine/compiler/static_batching_compiler.cpp index af2f4799f..fdefdbfcf 100644 --- a/csrc/engine/compiler/static_batching_compiler.cpp +++ b/csrc/engine/compiler/static_batching_compiler.cpp @@ -23,6 +23,7 @@ void StaticBatchingCompiler::compile() { input.past_sequence_lengths, input.total_sequence_lengths, input.input_offsets, + input.request_ids, input.cu_seqlens, input.block_tables, input.slot_mapping, diff --git a/csrc/engine/distributed/communication_group.cpp b/csrc/engine/distributed/communication_group.cpp index 782faa9ec..ffb69a339 100644 --- a/csrc/engine/distributed/communication_group.cpp +++ b/csrc/engine/distributed/communication_group.cpp @@ -5,23 +5,68 @@ namespace infinilm::engine::distributed { CommunicationGroup::CommunicationGroup(const DistConfig &dist_config, infinicore::Device::Type device_type) : dist_config_(dist_config), device_type_(device_type), - communicators_(std::vector(dist_config.tp_device_ids.size(), nullptr)) { + tp_communicators_(dist_config.world_size(), nullptr), + pp_communicators_(dist_config.world_size(), nullptr), + world_communicators_(dist_config.pipeline_parallel_size > 1 + ? dist_config.world_size() + : 0, + nullptr) { - size_t world_size = dist_config_.tp_device_ids.size(); + const int tp_size = dist_config_.tensor_parallel_size; + const int pp_size = dist_config_.pipeline_parallel_size; + const int world_size = dist_config_.world_size(); + if (tp_size < 1 || pp_size < 1 + || static_cast(dist_config_.tp_device_ids.size()) != world_size) { + throw std::runtime_error( + "DistConfig device count must equal tensor_parallel_size * pipeline_parallel_size"); + } size_t device_count = infinicore::context::getDeviceCount(device_type); - if (device_count < world_size) { + if (device_count < static_cast(world_size)) { throw std::runtime_error("infinilm::engine::distributed::CommunicationGroup error, world size is larger than the number of available GPUs. world size: " + std::to_string(world_size) + ", device count: " + std::to_string(device_count)); } if (infinicore::context::getDevice().getType() != device_type_) { infinicore::context::setDevice(infinicore::Device(device_type_, 0)); } - if (world_size > 1) { + if (pp_size == 1 && world_size > 1) { + RUN_INFINI(infinicclCommInitAll( + (infiniDevice_t)infinicore::context::getDevice().getType(), + tp_communicators_.data(), world_size, + dist_config.tp_device_ids.data())); + return; + } + + if (pp_size > 1) { RUN_INFINI(infinicclCommInitAll( (infiniDevice_t)infinicore::context::getDevice().getType(), - communicators_.data(), - dist_config.tp_device_ids.size(), + world_communicators_.data(), world_size, dist_config.tp_device_ids.data())); + + for (int stage = 0; stage < pp_size; ++stage) { + std::vector stage_comms(tp_size, nullptr); + const int offset = stage * tp_size; + RUN_INFINI(infinicclCommInitAll( + (infiniDevice_t)infinicore::context::getDevice().getType(), + stage_comms.data(), tp_size, + dist_config.tp_device_ids.data() + offset)); + for (int lane = 0; lane < tp_size; ++lane) { + tp_communicators_[offset + lane] = stage_comms[lane]; + } + } + + for (int lane = 0; lane < tp_size; ++lane) { + std::vector lane_devices(pp_size); + std::vector lane_comms(pp_size, nullptr); + for (int stage = 0; stage < pp_size; ++stage) { + lane_devices[stage] = dist_config.tp_device_ids[stage * tp_size + lane]; + } + RUN_INFINI(infinicclCommInitAll( + (infiniDevice_t)infinicore::context::getDevice().getType(), + lane_comms.data(), pp_size, lane_devices.data())); + for (int stage = 0; stage < pp_size; ++stage) { + pp_communicators_[stage * tp_size + lane] = lane_comms[stage]; + } + } } } @@ -31,21 +76,37 @@ const DistConfig &CommunicationGroup::get_dist_config() const { RankInfo CommunicationGroup::get_rank_info(int rank) const { RankInfo info; - info.tp_size = dist_config_.tp_device_ids.size(); - info.tp_rank = rank; + info.tp_size = dist_config_.tensor_parallel_size; + info.tp_rank = rank % info.tp_size; + info.pp_size = dist_config_.pipeline_parallel_size; + info.pp_rank = rank / info.tp_size; + info.world_size = dist_config_.world_size(); + info.global_rank = rank; info.device = infinicore::Device(device_type_, dist_config_.tp_device_ids[rank]); - info.comm = communicators_[rank]; + info.comm = tp_communicators_[rank]; + info.pp_comm = pp_communicators_[rank]; + info.world_comm = world_communicators_.empty() + ? tp_communicators_[rank] + : world_communicators_[rank]; return info; } int CommunicationGroup::get_world_size() const { - return dist_config_.tp_device_ids.size(); + return dist_config_.world_size(); +} + +int CommunicationGroup::get_output_rank() const { + return (dist_config_.pipeline_parallel_size - 1) + * dist_config_.tensor_parallel_size; } CommunicationGroup::~CommunicationGroup() { - if (communicators_.size() > 1) { - for (auto &comm : communicators_) { - infinicclCommDestroy(comm); + for (auto *group : {&pp_communicators_, &tp_communicators_, + &world_communicators_}) { + for (auto &comm : *group) { + if (comm != nullptr) { + infinicclCommDestroy(comm); + } } } } diff --git a/csrc/engine/distributed/communication_group.hpp b/csrc/engine/distributed/communication_group.hpp index e4f3c81a8..d18c18212 100644 --- a/csrc/engine/distributed/communication_group.hpp +++ b/csrc/engine/distributed/communication_group.hpp @@ -18,15 +18,32 @@ struct RankInfo { int tp_size; // Tensor parallelism rank number of this rank int tp_rank; - // Communicator handle + int pp_size; + int pp_rank; + int world_size; + int global_rank; + // Tensor-, pipeline-, and world-parallel communicator handles. infinicclComm_t comm; + infinicclComm_t pp_comm; + infinicclComm_t world_comm; RankInfo(infinicore::Device _device = infinicore::context::getDevice()) - : tp_size(1), tp_rank(0), device(_device), comm(nullptr){}; + : tp_size(1), tp_rank(0), pp_size(1), pp_rank(0), world_size(1), + global_rank(0), device(_device), comm(nullptr), pp_comm(nullptr), + world_comm(nullptr){}; + + bool is_pipeline_first_stage() const { return pp_rank == 0; } + bool is_pipeline_last_stage() const { return pp_rank + 1 == pp_size; } + bool is_output_rank() const { + return is_pipeline_last_stage() && tp_rank == 0; + } std::string to_string() const { std::stringstream ss; - ss << "RankInfo: device=" << device.toString() << ", tp_size=" << tp_size << ", tp_rank=" << tp_rank; + ss << "RankInfo: device=" << device.toString() + << ", global_rank=" << global_rank << "/" << world_size + << ", tp_rank=" << tp_rank << "/" << tp_size + << ", pp_rank=" << pp_rank << "/" << pp_size; return ss.str(); } }; @@ -42,12 +59,16 @@ class CommunicationGroup { int get_world_size() const; + int get_output_rank() const; + ~CommunicationGroup(); protected: DistConfig dist_config_; infinicore::Device::Type device_type_; - std::vector communicators_; + std::vector tp_communicators_; + std::vector pp_communicators_; + std::vector world_communicators_; }; } // namespace infinilm::engine::distributed diff --git a/csrc/engine/distributed/dist_config.cpp b/csrc/engine/distributed/dist_config.cpp index 44a73daf1..e35add7c9 100644 --- a/csrc/engine/distributed/dist_config.cpp +++ b/csrc/engine/distributed/dist_config.cpp @@ -1,18 +1,30 @@ #include "dist_config.hpp" +#include + namespace infinilm::engine::distributed { DistConfig::DistConfig() : tp_device_ids{0} {} -DistConfig::DistConfig(int tp_size) - : tp_device_ids(tp_size, 0) { - for (int i = 0; i < tp_size; ++i) { +DistConfig::DistConfig(int tp_size, int pp_size) + : tp_device_ids(tp_size * pp_size, 0), + tensor_parallel_size(tp_size), + pipeline_parallel_size(pp_size) { + if (tp_size < 1 || pp_size < 1) { + throw std::invalid_argument("TP and PP sizes must be positive"); + } + for (int i = 0; i < tp_size * pp_size; ++i) { tp_device_ids[i] = i; } } DistConfig::DistConfig(const std::vector &tp_device_ids_) - : tp_device_ids(tp_device_ids_) {} + : tp_device_ids(tp_device_ids_), + tensor_parallel_size(static_cast(tp_device_ids_.size())) {} + +int DistConfig::world_size() const { + return tensor_parallel_size * pipeline_parallel_size; +} DistConfig::operator std::string() const { std::string repr = "DistConfig(tp_device_ids=["; @@ -22,7 +34,10 @@ DistConfig::operator std::string() const { repr += ", "; } } - repr += "], moe_ep_backend=" + moe_ep_backend + ", moe_ep_size=" + std::to_string(moe_ep_size) + ")"; + repr += "], tp_size=" + std::to_string(tensor_parallel_size) + + ", pp_size=" + std::to_string(pipeline_parallel_size) + + ", moe_ep_backend=" + moe_ep_backend + + ", moe_ep_size=" + std::to_string(moe_ep_size) + ")"; return repr; } diff --git a/csrc/engine/distributed/dist_config.hpp b/csrc/engine/distributed/dist_config.hpp index 4affcc9be..08af13160 100644 --- a/csrc/engine/distributed/dist_config.hpp +++ b/csrc/engine/distributed/dist_config.hpp @@ -7,15 +7,19 @@ namespace infinilm::engine::distributed { struct DistConfig { - // Device IDs for each rank in tensor parallelism + // Device IDs for all ranks, ordered as contiguous TP groups per PP stage. std::vector tp_device_ids; + int tensor_parallel_size{1}; + int pipeline_parallel_size{1}; std::string moe_ep_backend{"disabled"}; size_t moe_ep_size{1}; DistConfig(); - explicit DistConfig(int tp_size); + explicit DistConfig(int tp_size, int pp_size = 1); explicit DistConfig(const std::vector &tp_device_ids_); + int world_size() const; + explicit operator std::string() const; }; diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index c47266d88..560ab1206 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -170,6 +170,7 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { to_device(past_sequence_lengths), // @todo: on device in the future to_device(total_sequence_lengths), to_device(input_offsets), + to_device(request_ids), to_device(cu_seqlens), to_device(block_tables), to_device(slot_mapping), @@ -181,16 +182,19 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { to_device_vec(image_grid_thw), image_req_ids, visual_token_ranges, - to_device(target_hidden_states)}; + to_device(target_hidden_states), + sample_all_positions}; infinilm::global_state::get_forward_context().attn_metadata = { input.past_sequence_lengths, input.total_sequence_lengths, input.input_offsets, + input.request_ids, input.cu_seqlens, input.block_tables, input.slot_mapping, - max_context_len}; + max_context_len, + is_mixed_batch}; infinilm::global_state::get_forward_context().mamba_metadata = { input.input_offsets, @@ -214,7 +218,7 @@ InferEngine::Output InferEngine::forward(const InferEngine::Input &input) { worker->wait(); } - return workers_[0]->get_output(); + return workers_[communication_group_.get_output_rank()]->get_output(); } void InferEngine::compile() { diff --git a/csrc/engine/rank_barrier.cpp b/csrc/engine/rank_barrier.cpp index 5e852ac61..4bbdb0804 100644 --- a/csrc/engine/rank_barrier.cpp +++ b/csrc/engine/rank_barrier.cpp @@ -3,17 +3,20 @@ namespace infinilm::engine { RankBarrier::RankBarrier(size_t num_ranks) : thread_count_(num_ranks), generation_(0), arrived_(0) {} -void RankBarrier::wait() { +bool RankBarrier::wait(bool success) { std::unique_lock lock(mutex_); - int gen = generation_; + const size_t gen = generation_; + generation_success_ = generation_success_ && success; if (++arrived_ == thread_count_) { - // last thread + completed_results_.push_back(generation_success_); generation_++; arrived_ = 0; + generation_success_ = true; cv_.notify_all(); } else { cv_.wait(lock, [&] { return gen != generation_; }); } + return completed_results_[gen]; } } // namespace infinilm::engine diff --git a/csrc/engine/rank_barrier.hpp b/csrc/engine/rank_barrier.hpp index dd068e994..142e16514 100644 --- a/csrc/engine/rank_barrier.hpp +++ b/csrc/engine/rank_barrier.hpp @@ -2,18 +2,21 @@ #include #include +#include namespace infinilm::engine { class RankBarrier { public: explicit RankBarrier(size_t nranks); - void wait(); + bool wait(bool success = true); private: const size_t thread_count_; size_t arrived_; size_t generation_; + bool generation_success_{true}; + std::vector completed_results_; std::mutex mutex_; std::condition_variable cv_; }; diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 11696e1eb..9d31fb873 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -1,6 +1,7 @@ #include "rank_worker.hpp" #include "../models/model_factory.hpp" #include "infinicore/ops.hpp" +#include "infinicore/ops/distributed/allreduce.hpp" #include #include @@ -415,11 +416,25 @@ void RankWorker::thread_loop() { { std::lock_guard lk(mutex_); + if (local_args.reuse_last_output) { + if (!last_output_ids_ || !local_args.input_ids.has_value()) { + throw std::runtime_error( + "reuse_last_output requires a previous device output and an input shape"); + } + if (last_output_ids_->numel() != local_args.input_ids.value()->numel()) { + throw std::runtime_error("reuse_last_output input shape mismatch"); + } + local_args.input_ids = last_output_ids_->view( + local_args.input_ids.value()->shape()); + } + infinicore::Tensor logits; infinicore::Tensor hidden_states; // All-position speculative/MTP runs need eager mode because // hidden states are not part of compiled graph outputs. - if (!local_args.sample_all_positions && compiler_ != nullptr) { + if (local_args.allow_graph_replay + && !local_args.sample_all_positions + && compiler_ != nullptr) { auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device::cpu())); if (graph != nullptr && output != nullptr) { graph->run(); @@ -434,43 +449,62 @@ void RankWorker::thread_loop() { hidden_states = model_output.hidden_states; } - // Random sampling (rank 0 only) - if (rank_info_.tp_rank == 0) { + // Normal forward keeps the historical CPU/I64 output contract. + // Device-resident generation samples on the final PP stage's + // TP rank 0, then broadcasts compact I32 IDs to every rank. + const bool output_rank = rank_info_.is_output_rank(); + if (output_rank || local_args.keep_output_device) { auto temperature{local_args.temperature}; auto top_p{local_args.top_p}; auto top_k{local_args.top_k}; - const auto &logits_shape{logits->shape()}; - const auto &vocab_size{logits_shape[2]}; - const auto &total_len{logits_shape[1]}; - const auto &batch_size{logits_shape[0]}; - auto n_req = local_args.input_offsets.value()->size(0) - 1; int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data(); const bool sample_all_positions = local_args.sample_all_positions; const size_t n_out = sample_all_positions ? static_cast(input_offsets[n_req]) : n_req; - auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)}; - - for (size_t i{0}; i < n_out; ++i) { - size_t score_idx = i; - if (!sample_all_positions) { - score_idx = static_cast(input_offsets[i + 1] - 1); + const auto output_dtype = local_args.keep_output_device + ? infinicore::DataType::I32 + : infinicore::DataType::I64; + auto output_ids = output_rank + ? infinicore::Tensor::empty({n_out}, output_dtype, rank_info_.device) + : infinicore::Tensor::zeros({n_out}, output_dtype, rank_info_.device); + + if (output_rank) { + const auto &logits_shape{logits->shape()}; + const auto &vocab_size{logits_shape[2]}; + const auto &total_len{logits_shape[1]}; + const auto &batch_size{logits_shape[0]}; + const bool logits_are_selected = !sample_all_positions + && batch_size * total_len == n_req; + for (size_t i{0}; i < n_out; ++i) { + size_t score_idx = i; + if (!sample_all_positions && !logits_are_selected) { + score_idx = static_cast(input_offsets[i + 1] - 1); + } + auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; + auto out{output_ids->narrow({{0, i, 1}})->view({})}; + float random_val = std::uniform_real_distribution(0, 1)(rng_); + infinicore::op::random_sample_( + out, score, random_val, top_p, top_k, temperature); } - auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; - auto out{output_ids->narrow({{0, i, 1}})->view({})}; - float random_val = std::uniform_real_distribution(0, 1)(rng_); - infinicore::op::random_sample_( - out, score, random_val, top_p, top_k, temperature); } - output_ids = output_ids->to(infinicore::Device::cpu()); - - infinicore::context::syncStream(); - - auto out{Output{output_ids, logits, hidden_states}}; + if (local_args.keep_output_device) { + if (rank_info_.world_size > 1) { + infinicore::op::distributed::allreduce_( + output_ids, output_ids, INFINICCL_SUM, + rank_info_.world_comm); + } + last_output_ids_ = output_ids; + } else { + output_ids = output_ids->to(infinicore::Device::cpu()); + infinicore::context::syncStream(); + } - output_ = std::move(out); + if (output_rank) { + output_ = Output{output_ids, logits, hidden_states}; + } } job_done_ = true; @@ -490,6 +524,7 @@ void RankWorker::thread_loop() { } else if (local_cmd == Command::RESET_CACHE) { try { model_->reset_cache(local_cache_config != nullptr ? local_cache_config.get() : nullptr); + last_output_ids_.reset(); { std::lock_guard lk(mutex_); job_done_ = true; diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index 162f8bbac..2ab286a81 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -46,6 +46,8 @@ class RankWorker { std::optional total_sequence_lengths; /// Offsets of each request in a continous-batched sequence, of shape `[num_requests + 1]`. std::optional input_offsets; + /// Request id for each flattened input token, of shape [num_tokens]. + std::optional request_ids; /// Cumulative total sequence lengths for each request, of shape `[num_requests + 1]`. std::optional cu_seqlens; /// Block ids for each request `[batch, max_block_table_length]`. Used for paged cache. @@ -74,6 +76,14 @@ class RankWorker { bool sample_all_positions{false}; /// Maximum total sequence length in the current request batch. std::optional max_context_len; + /// Keep sampled token IDs on each rank's device for generation. + bool keep_output_device{false}; + /// Reuse the previous sampled token IDs as this rank's model input. + bool reuse_last_output{false}; + /// Permit the graph compiler to replay this input; false forces eager mode. + bool allow_graph_replay{true}; + /// Whether this forward combines decode and prefill requests. + bool is_mixed_batch{false}; float temperature{1}; @@ -174,6 +184,7 @@ class RankWorker { // Output (protected by mutex) Output output_; + infinicore::Tensor last_output_ids_; // Thread sync std::thread thread_; diff --git a/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index 3dc982f5d..0b0020a11 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -11,6 +11,8 @@ struct AttentionMetadata { std::optional total_sequence_lengths; /// Offsets of each request in a continous-batched sequence, of shape `[num_requests + 1]`. std::optional input_offsets; + /// Request id for each flattened input token, of shape [num_tokens]. + std::optional request_ids; /// Cumulative total sequence lengths for each request, of shape `[num_requests + 1]`. std::optional cu_seqlens; /// Block ids for each request `[batch, max_block_table_length]`. Used for paged cache. @@ -19,26 +21,33 @@ struct AttentionMetadata { std::optional slot_mapping; /// Maximum total sequence length in the current request batch. std::optional max_context_len; + /// True when the packed input combines decode and prefill requests. + bool is_mixed_batch{false}; AttentionMetadata() = default; AttentionMetadata(std::optional past_sequence_lengths, std::optional total_sequence_lengths, std::optional input_offsets, + std::optional request_ids, std::optional cu_seqlens, std::optional block_tables, std::optional slot_mapping, - std::optional max_context_len = std::nullopt) : past_sequence_lengths(past_sequence_lengths), - total_sequence_lengths(total_sequence_lengths), - input_offsets(input_offsets), - cu_seqlens(cu_seqlens), - block_tables(block_tables), - slot_mapping(slot_mapping), - max_context_len(max_context_len) {} + std::optional max_context_len = std::nullopt, + bool is_mixed_batch = false) : past_sequence_lengths(past_sequence_lengths), + total_sequence_lengths(total_sequence_lengths), + input_offsets(input_offsets), + request_ids(request_ids), + cu_seqlens(cu_seqlens), + block_tables(block_tables), + slot_mapping(slot_mapping), + max_context_len(max_context_len), + is_mixed_batch(is_mixed_batch) {} AttentionMetadata(const infinilm::InfinilmModel::Input &input) : AttentionMetadata(input.past_sequence_lengths, input.total_sequence_lengths, input.input_offsets, + input.request_ids, input.cu_seqlens, input.block_tables, input.slot_mapping) {} @@ -64,6 +73,9 @@ struct ForwardContext { MambaMetadata mamba_metadata; MultiModalMetadata mm_metadata; std::vector kv_cache_vec; + std::vector mla_vendor_cache_vec; + std::vector indexer_cache_vec; + std::optional dsa_topk_indices; std::vector conv_state_vec; std::vector ssm_state_vec; }; diff --git a/csrc/layers/linear/fused_linear.cpp b/csrc/layers/linear/fused_linear.cpp index 1bcb96c94..c4f67a75b 100644 --- a/csrc/layers/linear/fused_linear.cpp +++ b/csrc/layers/linear/fused_linear.cpp @@ -3,6 +3,78 @@ #include namespace infinilm::layers::linear { +namespace { + +size_t merged_output_size(const std::vector &output_sizes) { + size_t total = 0; + for (const auto size : output_sizes) { + total += size; + } + return total; +} + +} // namespace + +// --------------------------------------------------------- +// Merged Replicated Linear +// --------------------------------------------------------- +MergedReplicatedLinear::MergedReplicatedLinear( + size_t input_size, + const std::vector &output_sizes, + const std::vector ¶m_names, + RegisterParamFn register_fn, + std::shared_ptr quantization, + bool bias, + const infinicore::DataType &dtype, + const infinicore::Device &device) + : infinilm::nn::Linear( + input_size, + merged_output_size(output_sizes), + quantization == nullptr + ? std::make_shared() + : quantization, + bias, + dtype, + device), + output_sizes_(output_sizes), + register_fn_(std::move(register_fn)) { + if (output_sizes_.empty() || output_sizes_.size() != param_names.size()) { + throw std::runtime_error( + "MergedReplicatedLinear expects non-empty, equally sized output_sizes and param_names"); + } + size_t offset = 0; + for (size_t i = 0; i < output_sizes_.size(); ++i) { + split_infos_.push_back({param_names[i], offset, output_sizes_[i]}); + offset += output_sizes_[i]; + } + auto params = this->split_params(split_infos_, 0, 1, -1); + for (auto &sp : params) { + register_fn_(sp.full_name, std::move(sp.param)); + } +} + +std::vector +MergedReplicatedLinear::forward_split(infinicore::Tensor &input) const { + auto output = this->forward(input); + const size_t dim = output->ndim() - 1; + std::vector result; + result.reserve(output_sizes_.size()); + size_t offset = 0; + for (const auto size : output_sizes_) { + result.push_back(output->narrow({{dim, offset, size}})); + offset += size; + } + return result; +} + +void MergedReplicatedLinear::process_weights_after_loading() { + BaseLinear::process_weights_after_loading(); + auto params = this->split_params(split_infos_, 0, 1, -1); + for (auto &sp : params) { + register_fn_(sp.full_name, std::move(sp.param)); + } +} + // --------------------------------------------------------- // QKV Parallel Linear // --------------------------------------------------------- diff --git a/csrc/layers/linear/fused_linear.hpp b/csrc/layers/linear/fused_linear.hpp index 8773a081c..0ea46ec0d 100644 --- a/csrc/layers/linear/fused_linear.hpp +++ b/csrc/layers/linear/fused_linear.hpp @@ -7,6 +7,29 @@ namespace infinilm::layers::linear { using RegisterParamFn = std::function; +// A replicated equivalent of vLLM's MergedColumnParallelLinear with +// disable_tp=True. Checkpoint shards keep their original names while sharing +// one runtime GEMM and one fused output buffer. +class MergedReplicatedLinear : public infinilm::nn::Linear { +public: + MergedReplicatedLinear(size_t input_size, + const std::vector &output_sizes, + const std::vector ¶m_names, + RegisterParamFn register_fn, + std::shared_ptr quantization = nullptr, + bool bias = false, + const infinicore::DataType &dtype = infinicore::DataType::F32, + const infinicore::Device &device = infinicore::Device()); + + std::vector forward_split(infinicore::Tensor &input) const; + void process_weights_after_loading() override; + +private: + std::vector output_sizes_; + RegisterParamFn register_fn_; + std::vector split_infos_; +}; + class QKVParallelLinear : public infinilm::nn::ColumnParallelLinear { public: explicit QKVParallelLinear(size_t hidden_size, diff --git a/csrc/layers/quantization/glm_w8a8.cpp b/csrc/layers/quantization/glm_w8a8.cpp index 6ccc75b6f..3f2992581 100644 --- a/csrc/layers/quantization/glm_w8a8.cpp +++ b/csrc/layers/quantization/glm_w8a8.cpp @@ -16,7 +16,12 @@ std::vector GlmW8A8::get_param_layout( int /*tp_num_heads*/, const infinicore::DataType &dtype, bool bias) const { std::vector result; - result.push_back({"weight", {out_features, in_features}, infinicore::DataType::I8, split_dim, tp_rank, tp_size}); + if (runtime_layout_) { + const int runtime_split_dim = split_dim >= 0 ? 1 - split_dim : -1; + result.push_back({"weight", {in_features, out_features}, infinicore::DataType::I8, runtime_split_dim, tp_rank, tp_size}); + } else { + result.push_back({"weight", {out_features, in_features}, infinicore::DataType::I8, split_dim, tp_rank, tp_size}); + } const int scale_split_dim = split_dim == 0 ? 0 : -1; result.push_back({"weight_scale", {out_features, 1}, infinicore::DataType::F32, scale_split_dim, scale_split_dim == 0 ? tp_rank : 0, scale_split_dim == 0 ? tp_size : 1}); if (bias) { @@ -29,11 +34,14 @@ infinicore::Tensor GlmW8A8::forward( const ParamsMap ¶ms, const infinicore::Tensor &input, bool has_bias, float alpha) const { auto weight = params.at("weight"); + if (!runtime_layout_) { + throw std::runtime_error("GlmW8A8 weights must be converted to runtime layout after loading"); + } auto weight_scale = params.at("weight_scale"); if (weight->ndim() != 2 || weight->dtype() != infinicore::DataType::I8) { - throw std::runtime_error("GlmW8A8 expects int8 weight [out,in]"); + throw std::runtime_error("GlmW8A8 expects int8 runtime weight [in,out]"); } - const size_t out_features = weight->size(0), in_features = weight->size(1); + const size_t in_features = weight->size(0), out_features = weight->size(1); if (weight_scale->ndim() != 2 || weight_scale->size(0) != out_features || weight_scale->size(1) != 1 || weight_scale->dtype() != infinicore::DataType::F32) { throw std::runtime_error("GlmW8A8 expects float32 weight_scale [out,1]"); } @@ -61,7 +69,7 @@ infinicore::Tensor GlmW8A8::forward( bias = params.at("bias"); } auto out = infinicore::Tensor::empty({m, out_features}, x->dtype(), x->device()); - infinicore::op::scaled_mm_w8a8_(out, x_i8, weight, x_scale, effective_scale, bias, true); + infinicore::op::scaled_mm_w8a8_(out, x_i8, weight, x_scale, effective_scale, bias, false); if (shape.size() == 2) { return out; } @@ -78,8 +86,9 @@ std::vector GlmW8A8::split_params( const auto &scale = params.at("weight_scale"); auto bias = params.find("bias"); for (const auto &s : splits) { + const int weight_dim = runtime_layout_ ? 1 : 0; result.push_back({s.prefix + ".weight", infinicore::nn::Parameter( - weight->narrow({{0, s.start, s.size}}), 0, tp_rank, tp_size, s.num_shards)}); + weight->narrow({{static_cast(weight_dim), s.start, s.size}}), weight_dim, tp_rank, tp_size, s.num_shards)}); result.push_back({s.prefix + ".weight_scale", infinicore::nn::Parameter( scale->narrow({{0, s.start, s.size}}), 0, tp_rank, tp_size, s.num_shards)}); if (bias != params.end()) { @@ -89,4 +98,16 @@ std::vector GlmW8A8::split_params( } return result; } + +std::shared_ptr GlmW8A8::process_weights_after_loading( + ParamsMap ¶ms, const infinicore::Device &device, + int /*split_dim*/) const { + (void)device; + if (runtime_layout_) { + return nullptr; + } + auto weight = params.at("weight"); + params["weight"] = weight->permute({1, 0})->contiguous(); + return std::make_shared(get_config(), true); +} } // namespace infinilm::quantization diff --git a/csrc/layers/quantization/glm_w8a8.hpp b/csrc/layers/quantization/glm_w8a8.hpp index d9164a616..4bbe32f98 100644 --- a/csrc/layers/quantization/glm_w8a8.hpp +++ b/csrc/layers/quantization/glm_w8a8.hpp @@ -3,15 +3,15 @@ namespace infinilm::quantization { -// Dense GLM W8A8 checkpoint/runtime layout. We retain [out, in] so CUINFER -// scaled_mm can consume checkpoint shards directly with trans_weight=true. +// Dense GLM W8A8 checkpoint layout is [out, in]. After loading, weights are +// transposed once to the vLLM/CUINFER runtime layout [in, out]. class GlmW8A8 final : public BaseQuantization { public: - explicit GlmW8A8(const nlohmann::json &quant_config) - : BaseQuantization(quant_config) {} + explicit GlmW8A8(const nlohmann::json &quant_config, bool runtime_layout = false) + : BaseQuantization(quant_config), runtime_layout_(runtime_layout) {} QuantScheme get_quant_scheme() const override { return QuantScheme::GLM_W8A8; } - int get_fused_split_dim() const override { return 0; } + int get_fused_split_dim() const override { return runtime_layout_ ? 1 : 0; } std::vector get_param_layout( size_t in_features, size_t out_features, @@ -27,5 +27,12 @@ class GlmW8A8 final : public BaseQuantization { const std::unordered_map ¶ms, const std::vector &splits, int narrow_dim, int tp_rank, int tp_size, int tp_num_heads) const override; + + std::shared_ptr process_weights_after_loading( + ParamsMap ¶ms, const infinicore::Device &device, + int split_dim = -1) const override; + +private: + bool runtime_layout_; }; } // namespace infinilm::quantization diff --git a/csrc/layers/rotary_embedding/rotary_embedding.cpp b/csrc/layers/rotary_embedding/rotary_embedding.cpp index 13f924dac..7f4333f56 100644 --- a/csrc/layers/rotary_embedding/rotary_embedding.cpp +++ b/csrc/layers/rotary_embedding/rotary_embedding.cpp @@ -1,5 +1,6 @@ #include "rotary_embedding.hpp" #include "../../config/model_config.hpp" +#include "infinicore/ops/cat.hpp" #include "rotary_embedding_factory.hpp" #include #include @@ -14,6 +15,7 @@ namespace { // Cache dictionary to avoid redundant allocations of RoPE instances. // thread_local ensures it is only visible within this compilation unit. thread_local std::unordered_map> _ROPE_DICT; +thread_local std::unordered_map _ROPE_COS_SIN_DICT; std::string make_cache_key(size_t head_dim, size_t rotary_dim, @@ -112,4 +114,18 @@ get_rope(const std::shared_ptr &model_config, false); } +infinicore::Tensor +get_rope_cos_sin_cache(const std::shared_ptr &rope) { + const auto *key = rope.get(); + auto it = _ROPE_COS_SIN_DICT.find(key); + if (it != _ROPE_COS_SIN_DICT.end()) { + return it->second; + } + + auto cache = infinicore::op::cat( + {rope->cos_cache(), rope->sin_cache()}, 1); + auto [inserted, _] = _ROPE_COS_SIN_DICT.emplace(key, std::move(cache)); + return inserted->second; +} + } // namespace infinilm::layers::rotary_embedding diff --git a/csrc/layers/rotary_embedding/rotary_embedding.hpp b/csrc/layers/rotary_embedding/rotary_embedding.hpp index 7bd05f055..525eee944 100644 --- a/csrc/layers/rotary_embedding/rotary_embedding.hpp +++ b/csrc/layers/rotary_embedding/rotary_embedding.hpp @@ -36,4 +36,7 @@ std::shared_ptr get_rope(const std::shared_ptr &model_config, const infinicore::Device &device); +infinicore::Tensor +get_rope_cos_sin_cache(const std::shared_ptr &rope); + } // namespace infinilm::layers::rotary_embedding diff --git a/csrc/models/deepseek_v2/deepseek_v2_indexer.cpp b/csrc/models/deepseek_v2/deepseek_v2_indexer.cpp new file mode 100644 index 000000000..3e5dca2f7 --- /dev/null +++ b/csrc/models/deepseek_v2/deepseek_v2_indexer.cpp @@ -0,0 +1,260 @@ +#include "deepseek_v2_indexer.hpp" + +#include "../../global_state/global_state.hpp" +#include "../../layers/rotary_embedding/rotary_embedding.hpp" +#include "../../layers/rotary_embedding/rotary_embedding_factory.hpp" +#include "../../utils.hpp" +#include "infinicore/ops/add.hpp" +#include "infinicore/ops/broadcast_to.hpp" +#include "infinicore/ops/cast.hpp" +#include "infinicore/ops/cat.hpp" +#include "infinicore/ops/distributed/allreduce.hpp" +#include "infinicore/ops/dsa.hpp" +#include "infinicore/ops/fp8_indexer_logits.hpp" +#include "infinicore/ops/fp8_indexer_quant.hpp" +#include "infinicore/ops/mul.hpp" + +#include +#include +#include +#include + +namespace infinilm::models::deepseek_v2 { +namespace { +void debug_dump_indexer(const infinicore::Tensor &tensor, + const std::string &name, + size_t layer_idx) { + if (layer_idx != 0 || std::getenv("INFINILM_GLM_DEBUG_DUMP") == nullptr) { + return; + } + const auto &rank = infinilm::global_state::get_tensor_model_parallel_rank_info(); + tensor->debug("/tmp/glmdbg_indexer_" + name + "_rank_" + + std::to_string(rank.tp_rank) + ".bin"); +} +} // namespace + +DeepseekV32Indexer::DeepseekV32Indexer( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + layer_idx_ = layer_idx; + num_heads_ = model_config->get("index_n_heads"); + head_dim_ = model_config->get("index_head_dim"); + rope_dim_ = model_config->get("qk_rope_head_dim"); + topk_tokens_ = model_config->get("index_topk"); + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + tp_rank_ = rank_info.tp_rank; + tp_size_ = rank_info.tp_size; + communicator_ = rank_info.comm; + const size_t hidden_size = model_config->get("hidden_size"); + const size_t q_lora_rank = model_config->get("q_lora_rank"); + const auto &dtype = model_config->get_dtype(); + auto quantization = model_config->get_quantization_method(); + + INFINICORE_NN_MODULE_INIT( + wq_b, q_lora_rank, num_heads_ * head_dim_, quantization, false, dtype, device); + INFINICORE_NN_MODULE_INIT( + wk, hidden_size, head_dim_, quantization, false, dtype, device); + INFINICORE_NN_MODULE_INIT( + weights_proj, hidden_size, num_heads_, false, dtype, device); + fused_wk_weights_proj_ = std::make_shared( + hidden_size, head_dim_ + num_heads_, false, dtype, device); + INFINICORE_NN_MODULE_INIT(k_norm, head_dim_, 1e-6, dtype, device); + + auto scaling = infinilm::layers::rotary_embedding::make_scaling_config(model_config); + rotary_emb_ = infinilm::layers::rotary_embedding::get_rope( + head_dim_, + rope_dim_, + model_config->get("max_position_embeddings"), + model_config->get("rope_theta"), + infinicore::nn::RoPE::Algo::GPT_J, + dtype, + device, + std::move(scaling)); + cos_sin_cache_ = infinilm::layers::rotary_embedding::get_rope_cos_sin_cache(rotary_emb_); + weights_scale_ = 1.0f / std::sqrt(static_cast(head_dim_ * num_heads_)); + one_i32_ = infinicore::Tensor::ones({1}, infinicore::DataType::I32, device); +} + +void DeepseekV32Indexer::process_weights_after_loading() { + const auto wk_weight = wk_->weight(); + const auto wk_scale = wk_->weight_scale(); + const auto weights_weight = weights_proj_->weight(); + if (!wk_weight || wk_weight->dtype() != infinicore::DataType::I8 + || wk_weight->ndim() != 2) { + throw std::runtime_error( + "DeepseekV32Indexer expects wk int8 weight runtime [hidden,head_dim]"); + } + if (!wk_scale || wk_scale->dtype() != infinicore::DataType::F32 + || wk_scale->ndim() != 2 + || wk_scale->size(0) != wk_weight->size(1) + || wk_scale->size(1) != 1) { + throw std::runtime_error( + "DeepseekV32Indexer expects wk float32 weight_scale [head_dim,1]"); + } + if (!weights_weight || weights_weight->dtype() != k_norm_->dtype() + || weights_weight->ndim() != 2 + || weights_weight->size(1) != wk_weight->size(0)) { + throw std::runtime_error( + "DeepseekV32Indexer expects BF16 weights_proj [num_heads,hidden]"); + } + + auto wk_out_in = wk_weight->permute({1, 0})->contiguous(); + auto wk_f32 = infinicore::Tensor::empty( + wk_out_in->shape(), infinicore::DataType::F32, wk_out_in->device()); + infinicore::op::cast_(wk_f32, wk_out_in); + auto expanded_scale = infinicore::op::broadcast_to( + wk_scale, + {static_cast(wk_out_in->size(0)), + static_cast(wk_out_in->size(1))}); + auto dequantized_f32 = infinicore::op::mul(wk_f32, expanded_scale); + auto dequantized = infinicore::Tensor::empty( + wk_out_in->shape(), weights_weight->dtype(), wk_out_in->device()); + infinicore::op::cast_(dequantized, dequantized_f32); + + auto fused_weight = infinicore::op::cat({dequantized, weights_weight}, 0); + fused_wk_weights_proj_->weight()->copy_from(fused_weight); + wk_->release_parameters(); + weights_proj_->release_parameters(); +} + +void DeepseekV32Indexer::sync_topk_indices_(infinicore::Tensor topk_indices) const { + if (tp_size_ == 1) { + return; + } + if (communicator_ == nullptr) { + throw std::runtime_error("DeepseekV32Indexer requires a TP communicator"); + } + + if (tp_rank_ != 0) { + set_zeros_device_async(topk_indices); + } + infinicore::op::distributed::allreduce_( + topk_indices, topk_indices, INFINICCL_SUM, communicator_); +} + +void DeepseekV32Indexer::forward( + const infinicore::Tensor &hidden_states, + const infinicore::Tensor &q_lora, + const infinicore::Tensor &positions, + infinicore::Tensor topk_indices) const { + const size_t num_tokens = positions->numel(); + if (topk_indices->shape() + != std::vector{num_tokens, topk_tokens_}) { + throw std::runtime_error("DeepseekV32Indexer: invalid topk buffer shape"); + } + + auto &forward_context = infinilm::global_state::get_forward_context(); + auto &metadata = forward_context.attn_metadata; + if (layer_idx_ >= forward_context.indexer_cache_vec.size()) { + throw std::runtime_error("DeepseekV32Indexer: indexer cache is not allocated"); + } + if (!metadata.slot_mapping.has_value() + || !metadata.input_offsets.has_value() + || !metadata.cu_seqlens.has_value() + || !metadata.total_sequence_lengths.has_value() + || !metadata.request_ids.has_value() + || !metadata.block_tables.has_value() + || !metadata.max_context_len.has_value()) { + throw std::runtime_error("DeepseekV32Indexer: incomplete paged attention metadata"); + } + const size_t num_requests = metadata.total_sequence_lengths.value()->numel(); + + // Rank-local indexer compute wins once larger batches amortize its fixed + // GEMM cost. Batch two with long contexts is faster on rank 0 plus a + // small top-k collective than repeating the indexer on all four TP ranks. + const bool use_replicated_indexer = num_requests > 2; + if (!use_replicated_indexer && tp_rank_ != 0) { + sync_topk_indices_(topk_indices); + debug_dump_indexer(topk_indices, "selected", layer_idx_); + return; + } + + auto hidden_mutable = hidden_states; + auto q_lora_mutable = q_lora; + auto q_raw = wq_b_->forward(q_lora_mutable) + ->view({num_tokens, num_heads_, head_dim_}); + auto k_weights = fused_wk_weights_proj_->forward(hidden_mutable) + ->view({num_tokens, head_dim_ + num_heads_}); + auto k_raw = k_weights->narrow({{1, 0, head_dim_}}); + auto weights_raw = k_weights->narrow({{1, head_dim_, num_heads_}}); + debug_dump_indexer(q_raw, "q_raw", layer_idx_); + debug_dump_indexer(k_raw, "k_raw", layer_idx_); + debug_dump_indexer(weights_raw, "weights_raw", layer_idx_); + debug_dump_indexer(k_weights, "k_weights", layer_idx_); + + auto &k_cache = forward_context.indexer_cache_vec[layer_idx_]; + auto q_fp8 = infinicore::Tensor::empty( + q_raw->shape(), infinicore::DataType::F8, q_raw->device()); + auto weights_fp32 = infinicore::Tensor::empty( + {num_tokens, num_heads_}, infinicore::DataType::F32, q_raw->device()); + const bool use_fused_fp8 = std::getenv("INFINILM_GLM_DISABLE_FUSED_FP8_INDEXER") == nullptr; + if (use_fused_fp8) { + infinicore::op::fused_fp8_indexer_( + q_fp8, weights_fp32, k_cache, q_raw, k_weights, + k_norm_->weight(), k_norm_->bias(), positions, cos_sin_cache_, + metadata.slot_mapping.value(), rope_dim_, k_norm_->eps(), + weights_scale_); + } else { + auto q = infinicore::Tensor::empty( + q_raw->shape(), q_raw->dtype(), q_raw->device()); + auto k = infinicore::Tensor::empty( + {num_tokens, head_dim_}, q_raw->dtype(), q_raw->device()); + auto weights = infinicore::Tensor::empty( + {num_tokens, num_heads_}, q_raw->dtype(), q_raw->device()); + auto empty_cache = infinicore::Tensor::empty( + {0}, q_raw->dtype(), q_raw->device()); + auto empty_slots = infinicore::Tensor::empty( + {0}, infinicore::DataType::I64, q_raw->device()); + infinicore::op::fused_deepseek_v2_indexer_postprocess_( + q, k, weights, empty_cache, empty_slots, q_raw, k_weights, + k_norm_->weight(), k_norm_->bias(), positions, cos_sin_cache_, + 0, false, k_norm_->eps(), weights_scale_); + debug_dump_indexer(q, "q", layer_idx_); + debug_dump_indexer(k, "k", layer_idx_); + infinicore::op::indexer_k_quant_and_cache_( + k, k_cache, metadata.slot_mapping.value(), + static_cast(head_dim_), "ue8m0"); + debug_dump_indexer(weights, "weights", layer_idx_); + infinicore::op::fp8_indexer_quant_( + q_fp8, weights_fp32, q, weights); + } + debug_dump_indexer(k_cache, "k_cache", layer_idx_); + debug_dump_indexer(q_fp8, "q_fp8", layer_idx_); + debug_dump_indexer(weights_fp32, "weights_fp32", layer_idx_); + + const int64_t max_context_len = metadata.max_context_len.value(); + auto logits = infinicore::Tensor::empty( + {num_tokens, static_cast(max_context_len)}, + infinicore::DataType::F32, + q_raw->device()); + const bool is_prefill = num_tokens != num_requests; + infinicore::op::fp8_indexer_logits_( + logits, q_fp8, k_cache, metadata.block_tables.value(), weights_fp32, + positions, metadata.request_ids.value()); + debug_dump_indexer(logits, "logits", layer_idx_); + + if (is_prefill) { + auto cu_seqlen_ks = infinicore::Tensor::empty( + {num_tokens}, infinicore::DataType::I32, q_raw->device()); + set_zeros_device_async(cu_seqlen_ks); + auto positions_i32 = infinicore::Tensor::empty( + {num_tokens}, infinicore::DataType::I32, q_raw->device()); + infinicore::op::cast_(positions_i32, positions); + auto ones = infinicore::op::broadcast_to( + one_i32_, {static_cast(num_tokens)}) + ->contiguous(); + auto cu_seqlen_ke = infinicore::op::add(positions_i32, ones); + infinicore::op::select_prefill_topk_block_indices_( + topk_indices, logits, cu_seqlen_ks, cu_seqlen_ke); + } else { + infinicore::op::select_decode_topk_block_indices_( + topk_indices, logits, metadata.total_sequence_lengths.value()); + } + if (!use_replicated_indexer) { + sync_topk_indices_(topk_indices); + } + debug_dump_indexer(topk_indices, "selected", layer_idx_); +} + +} // namespace infinilm::models::deepseek_v2 diff --git a/csrc/models/deepseek_v2/deepseek_v2_indexer.hpp b/csrc/models/deepseek_v2/deepseek_v2_indexer.hpp new file mode 100644 index 000000000..0965d047c --- /dev/null +++ b/csrc/models/deepseek_v2/deepseek_v2_indexer.hpp @@ -0,0 +1,53 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/linear/linear.hpp" +#include "infinicore/nn/layer_norm.hpp" +#include "infinicore/nn/module.hpp" +#include "infinicore/nn/rope.hpp" +#include "infinicore/tensor.hpp" +#include + +#include + +namespace infinilm::models::deepseek_v2 { + +class DeepseekV32Indexer final : public infinicore::nn::Module { +public: + DeepseekV32Indexer(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + void forward(const infinicore::Tensor &hidden_states, + const infinicore::Tensor &q_lora, + const infinicore::Tensor &positions, + infinicore::Tensor topk_indices) const; + + void process_weights_after_loading() override; + +private: + void sync_topk_indices_(infinicore::Tensor topk_indices) const; + + size_t layer_idx_{0}; + size_t num_heads_{0}; + size_t head_dim_{0}; + size_t rope_dim_{0}; + size_t topk_tokens_{0}; + float weights_scale_{1.0f}; + int tp_rank_{0}; + int tp_size_{1}; + infinicclComm_t communicator_{nullptr}; + + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, wq_b); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, wk); + INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, weights_proj); + std::shared_ptr + fused_wk_weights_proj_; + INFINICORE_NN_MODULE(infinicore::nn::LayerNorm, k_norm); + + std::shared_ptr rotary_emb_; + infinicore::Tensor cos_sin_cache_; + infinicore::Tensor one_i32_; +}; + +} // namespace infinilm::models::deepseek_v2 diff --git a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp index 7dfcbf857..afc875651 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp +++ b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp @@ -5,17 +5,37 @@ #include "../../layers/rotary_embedding/rotary_embedding.hpp" #include "../../utils.hpp" #include "deepseek_v2_utils.hpp" +#include "infinicore/context/context.hpp" +#include "infinicore/ops/bmm_strided.hpp" #include "infinicore/ops/broadcast_to.hpp" +#include "infinicore/ops/cast.hpp" #include "infinicore/ops/cat.hpp" #include "infinicore/ops/concat_and_cache_mla.hpp" #include "infinicore/ops/concat_mla_q.hpp" +#include "infinicore/ops/dsa.hpp" +#include "infinicore/ops/fp8_mla_rmsnorm_cache.hpp" +#include "infinicore/ops/fused_rotary_embedding.hpp" #include "infinicore/ops/matmul.hpp" #include "infinicore/ops/mha_varlen.hpp" +#include "infinicore/ops/mul.hpp" #include "infinicore/ops/paged_attention_mla.hpp" +#include #include +#include namespace infinilm::models::deepseek_v2 { +namespace { +void debug_dump_dsa(const infinicore::Tensor &tensor, const std::string &name, size_t layer_idx) { + if (layer_idx != 0 || std::getenv("INFINILM_GLM_DEBUG_DUMP") == nullptr) { + return; + } + const auto &rank = infinilm::global_state::get_tensor_model_parallel_rank_info(); + if (rank.tp_rank == 0) { + tensor->debug("/tmp/glmdbg_dsa_" + name + ".bin"); + } +} +} // namespace DeepseekV2MLAAttention::DeepseekV2MLAAttention(std::shared_ptr model_config, size_t layer_idx, @@ -48,6 +68,19 @@ DeepseekV2MLAAttention::DeepseekV2MLAAttention(std::shared_ptr(tp_size); attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend; + use_sparse_ = config_json.contains("index_topk"); + if (use_sparse_) { + if (q_lora_rank_ == 0) { + throw std::runtime_error("Sparse MLA requires q_lora_rank"); + } + const auto indexer_types = config_json.value("indexer_types", nlohmann::json::array()); + if (layer_idx_ < indexer_types.size()) { + skip_topk_ = indexer_types[layer_idx_].get() == "shared"; + } + if (!skip_topk_) { + INFINICORE_NN_MODULE_INIT(indexer, model_config, layer_idx_, device); + } + } if (attention_backend_ == infinilm::backends::AttentionBackend::STATIC_ATTN) { throw std::runtime_error("DeepseekV2MLAAttention requires paged or flash attention; the dense MHA path was removed"); } @@ -64,13 +97,17 @@ DeepseekV2MLAAttention::DeepseekV2MLAAttention(std::shared_ptr( + hidden_size_, + std::vector{q_lora_rank_, kv_lora_rank_ + qk_rope_head_dim_}, + std::vector{"q_a_proj", "kv_a_proj_with_mqa"}, + [this](const std::string &name, infinicore::nn::Parameter parameter) { + this->register_parameter(name, std::move(parameter)); + }, + quantization_method, + false, + dtype, + device); INFINICORE_NN_MODULE_INIT(q_a_layernorm, q_lora_rank_, rms_norm_eps, dtype, device); INFINICORE_NN_MODULE_INIT(q_b_proj, q_lora_rank_, @@ -82,13 +119,15 @@ DeepseekV2MLAAttention::DeepseekV2MLAAttention(std::shared_ptr(q_head_dim_)); infinilm::layers::attention::init_kv_cache_quant_params( [this](const std::string &name, infinicore::nn::Parameter parameter) { @@ -136,37 +176,96 @@ infinicore::Tensor DeepseekV2MLAAttention::position_ids_for_rope_(const infinico throw std::runtime_error("DeepseekV2MLAAttention: unexpected position_ids shape"); } -infinicore::Tensor DeepseekV2MLAAttention::kv_b_weight_3d_() const { - return kv_b_proj_->weight()->view( +void DeepseekV2MLAAttention::process_weights_after_loading() { + if (fused_qkv_a_proj_) { + fused_qkv_a_proj_->process_weights_after_loading(); + } + const auto weight = kv_b_proj_->weight(); + const auto weight_scale = kv_b_proj_->weight_scale(); + if (!weight || weight->dtype() != infinicore::DataType::I8 + || weight->ndim() != 2) { + throw std::runtime_error( + "DeepseekV2MLAAttention expects kv_b_proj int8 weight runtime [in,out]"); + } + if (!weight_scale || weight_scale->dtype() != infinicore::DataType::F32 + || weight_scale->ndim() != 2 + || weight_scale->size(0) != weight->size(1) + || weight_scale->size(1) != 1) { + throw std::runtime_error( + "DeepseekV2MLAAttention expects kv_b_proj float32 weight_scale [out,1]"); + } + + auto out_in_weight = weight->permute({1, 0})->contiguous(); + auto weight_f32 = infinicore::Tensor::empty( + out_in_weight->shape(), infinicore::DataType::F32, out_in_weight->device()); + infinicore::op::cast_(weight_f32, out_in_weight); + auto expanded_scale = infinicore::op::broadcast_to( + weight_scale, {static_cast(out_in_weight->size(0)), + static_cast(out_in_weight->size(1))}); + auto dequantized_f32 = infinicore::op::mul(weight_f32, expanded_scale); + auto dequantized = infinicore::Tensor::empty( + out_in_weight->shape(), kv_a_layernorm_->dtype(), out_in_weight->device()); + infinicore::op::cast_(dequantized, dequantized_f32); + + auto weight_3d = dequantized->view( {num_attention_heads_, qk_nope_head_dim_ + v_head_dim_, kv_lora_rank_}); + w_uk_ = weight_3d->narrow({{1, 0, qk_nope_head_dim_}})->contiguous(); + w_uv_ = weight_3d->narrow({{1, qk_nope_head_dim_, v_head_dim_}}) + ->permute({0, 2, 1}) + ->contiguous(); + kv_b_proj_->release_parameters(); +} + +void DeepseekV2MLAAttention::reset_runtime_state() const { + if (fused_qkv_a_proj_) { + fused_qkv_a_proj_->reset_runtime_state(); + } } infinicore::Tensor DeepseekV2MLAAttention::project_q_nope_to_latent_(const infinicore::Tensor &q_nope) const { + if (!w_uk_) { + throw std::runtime_error("DeepseekV2MLAAttention absorbed query weight is not prepared"); + } const size_t num_tokens = q_nope->shape()[0]; - auto q_nope_by_head = q_nope->permute({1, 0, 2})->contiguous(); - auto w_uk_t = kv_b_weight_3d_()->narrow({{1, 0, qk_nope_head_dim_}})->contiguous(); - auto q_latent = infinicore::op::matmul(q_nope_by_head, w_uk_t); - return q_latent->permute({1, 0, 2}) - ->contiguous() - ->view({num_tokens, num_attention_heads_, kv_lora_rank_}); + auto q_latent = infinicore::Tensor::empty( + {num_tokens, num_attention_heads_, kv_lora_rank_}, + q_nope->dtype(), + q_nope->device()); + infinicore::op::bmm_strided_( + q_latent->permute({1, 0, 2}), + q_nope->permute({1, 0, 2}), + w_uk_); + return q_latent; } infinicore::Tensor DeepseekV2MLAAttention::project_latent_to_value_(const infinicore::Tensor &attn_output, size_t batch_size, size_t seq_len) const { + if (!w_uv_) { + throw std::runtime_error("DeepseekV2MLAAttention absorbed value weight is not prepared"); + } const size_t num_tokens = batch_size * seq_len; auto latent_by_head = attn_output->view({num_tokens, num_attention_heads_, kv_lora_rank_}) - ->permute({1, 0, 2}) - ->contiguous(); - auto w_uv = kv_b_weight_3d_() - ->narrow({{1, qk_nope_head_dim_, v_head_dim_}}) - ->permute({0, 2, 1}) - ->contiguous(); - auto value = infinicore::op::matmul(latent_by_head, w_uv) - ->permute({1, 0, 2}) - ->contiguous() - ->view({batch_size, seq_len, num_attention_heads_ * v_head_dim_}); - return o_proj_->forward(value); + ->permute({1, 0, 2}); + auto value = infinicore::Tensor::empty( + {num_tokens, num_attention_heads_, v_head_dim_}, + attn_output->dtype(), + attn_output->device()); + infinicore::op::bmm_strided_( + value->permute({1, 0, 2}), + latent_by_head, + w_uv_); + auto value_flat = value->view( + {batch_size, seq_len, num_attention_heads_ * v_head_dim_}); + + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + auto output = o_proj_->forward(value_flat); + if (use_sparse_ && rank_info.tp_size > 1) { + // Finish sparse attention and the TP output projection before the + // decoder advances to the next layer. + infinicore::context::syncStream(); + } + return output; } infinicore::Tensor DeepseekV2MLAAttention::forward(const infinicore::Tensor &position_ids, @@ -180,27 +279,42 @@ infinicore::Tensor DeepseekV2MLAAttention::forward(const infinicore::Tensor &pos auto hidden_states_mutable = hidden_states; infinicore::Tensor q_linear; + infinicore::Tensor compressed; + infinicore::Tensor q_lora_norm; if (q_lora_rank_ == 0) { q_linear = q_proj_->forward(hidden_states_mutable); + compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable); } else { - auto q_a = q_a_proj_->forward(hidden_states_mutable); - auto q_a_norm = q_a_layernorm_->forward(q_a); - q_linear = q_b_proj_->forward(q_a_norm); + auto qkv_a = fused_qkv_a_proj_->forward_split(hidden_states_mutable); + if (qkv_a.size() != 2) { + throw std::runtime_error( + "DeepseekV2MLAAttention fused_qkv_a_proj must return two output slices"); + } + auto q_a = qkv_a[0]; + compressed = qkv_a[1]; + q_lora_norm = q_a_layernorm_->forward(q_a); + q_linear = q_b_proj_->forward(q_lora_norm); } auto q = q_linear->view({seq_len, num_attention_heads_, q_head_dim_}); - auto q_nope = q->narrow({{2, 0, qk_nope_head_dim_}})->contiguous(); + auto q_nope = q->narrow({{2, 0, qk_nope_head_dim_}}); auto q_pe = q->narrow({{2, qk_nope_head_dim_, qk_rope_head_dim_}})->contiguous(); - auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable) - ->view({seq_len, kv_lora_rank_ + qk_rope_head_dim_}); + compressed = compressed->view({seq_len, kv_lora_rank_ + qk_rope_head_dim_}); auto compressed_kv = compressed->narrow({{1, 0, kv_lora_rank_}})->contiguous(); auto k_pe = compressed->narrow({{1, kv_lora_rank_, qk_rope_head_dim_}})->contiguous(); - auto kv_norm = kv_a_layernorm_->forward(compressed_kv); auto pos_ids = position_ids_for_rope_(position_ids); - q_pe = rotary_emb_->forward(q_pe, pos_ids, true); - auto k_pe_rope = rotary_emb_->forward( - k_pe->view({seq_len, 1, qk_rope_head_dim_}), pos_ids, true); + auto k_pe_rope = k_pe->view({seq_len, 1, qk_rope_head_dim_}); + infinicore::op::fused_rotary_embedding_( + q_pe, + k_pe_rope, + pos_ids, + static_cast(qk_rope_head_dim_), + rope_cos_sin_cache_, + rotary_emb_->algo() == infinicore::nn::RoPE::Algo::GPT_NEOX); + debug_dump_dsa(q_nope, "q_nope", layer_idx_); + debug_dump_dsa(q_pe, "q_pe", layer_idx_); + debug_dump_dsa(k_pe_rope, "k_pe_rope", layer_idx_); auto &forward_context = infinilm::global_state::get_forward_context(); auto &attn_metadata = forward_context.attn_metadata; @@ -209,23 +323,135 @@ infinicore::Tensor DeepseekV2MLAAttention::forward(const infinicore::Tensor &pos auto slot_mapping = attn_metadata.slot_mapping; auto total_sequence_lengths = attn_metadata.total_sequence_lengths; auto input_offsets = attn_metadata.input_offsets; + auto request_ids = attn_metadata.request_ids; auto cu_seqlens = attn_metadata.cu_seqlens; ASSERT(block_tables.has_value()); ASSERT(slot_mapping.has_value()); ASSERT(total_sequence_lengths.has_value()); + ASSERT(request_ids.has_value()); ASSERT(attn_metadata.max_context_len.has_value()); if (hidden_states->device().getType() != infinicore::Device::Type::ILUVATAR) { throw std::runtime_error("DeepseekV2MLAAttention: the vLLM-style MLA cache path currently requires Iluvatar"); } - infinicore::op::concat_and_cache_mla_(kv_norm, - k_pe_rope->view({seq_len, qk_rope_head_dim_}), - kv_cache, - slot_mapping.value(), - "auto", - kv_cache_k_scale_); + const bool use_fp8_ds_mla = kv_cache->dtype() == infinicore::DataType::U8; + const bool use_fused_fp8_cache = use_sparse_ && use_fp8_ds_mla + && std::getenv("INFINILM_GLM_DISABLE_FUSED_FP8_MLA_CACHE") == nullptr; + const bool use_vendor_sparse = use_sparse_ && use_fp8_ds_mla + && !forward_context.mla_vendor_cache_vec.empty(); + if (use_vendor_sparse + && forward_context.mla_vendor_cache_vec.size() + != forward_context.kv_cache_vec.size()) { + throw std::runtime_error( + "GLM FP8 Sparse MLA vendor shadow layer count mismatch"); + } + if (use_vendor_sparse && !use_fused_fp8_cache) { + throw std::runtime_error( + "GLM FP8 Sparse MLA vendor shadow requires fused cache producer"); + } + infinicore::Tensor kv_norm; + if (use_fused_fp8_cache) { + auto rope_cache = k_pe_rope->view({seq_len, qk_rope_head_dim_}); + if (use_vendor_sparse) { + infinicore::op::fp8_mla_rmsnorm_dual_cache_( + kv_cache, + forward_context.mla_vendor_cache_vec[layer_idx_], + compressed_kv, + kv_a_layernorm_->weight(), + rope_cache, + slot_mapping.value(), + kv_a_layernorm_->eps()); + } else { + infinicore::op::fp8_mla_rmsnorm_cache_( + kv_cache, + compressed_kv, + kv_a_layernorm_->weight(), + rope_cache, + slot_mapping.value(), + kv_a_layernorm_->eps()); + } + } else { + kv_norm = kv_a_layernorm_->forward(compressed_kv); + debug_dump_dsa(kv_norm, "kv_norm", layer_idx_); + infinicore::op::concat_and_cache_mla_( + kv_norm, + k_pe_rope->view({seq_len, qk_rope_head_dim_}), + kv_cache, + slot_mapping.value(), + use_fp8_ds_mla ? "fp8_ds_mla" : "auto", + kv_cache_k_scale_); + } + debug_dump_dsa(kv_cache, "mla_cache", layer_idx_); const bool is_prefill = seq_len != total_sequence_lengths.value()->shape()[0]; + if (use_sparse_) { + auto &topk_indices_opt = forward_context.dsa_topk_indices; + if (!topk_indices_opt.has_value()) { + throw std::runtime_error("Sparse MLA top-k buffer is not allocated"); + } + auto topk_indices = topk_indices_opt.value(); + if (!skip_topk_) { + indexer_->forward(hidden_states, q_lora_norm, pos_ids, topk_indices); + debug_dump_dsa(topk_indices, "topk_local", layer_idx_); + } + + const size_t num_tokens = seq_len; + if (request_ids.value()->numel() < num_tokens) { + throw std::runtime_error( + "Sparse MLA request_ids is shorter than the flattened token batch"); + } + auto token_request_ids = request_ids.value(); + if (token_request_ids->numel() != num_tokens) { + token_request_ids = token_request_ids->narrow({{0, 0, num_tokens}}); + } + auto global_indices = infinicore::Tensor::empty( + topk_indices->shape(), infinicore::DataType::I32, hidden_states->device()); + const int64_t block_size = static_cast(kv_cache->size(1)); + // Match the Iluvatar vLLM production path: without prefill + // workspace, request-local indices use the decode mapping kernel. + infinicore::op::map_decode_request_block_indices_( + global_indices, + token_request_ids, + block_tables.value(), + topk_indices, + block_size); + + auto topk_lens = infinicore::Tensor::empty( + {num_tokens}, infinicore::DataType::I32, hidden_states->device()); + auto sparse_indices = global_indices->view( + {num_tokens, 1, global_indices->size(1)}); + infinicore::op::topk_indices_context_lens_(topk_lens, sparse_indices); + debug_dump_dsa(global_indices, "topk_global", layer_idx_); + debug_dump_dsa(topk_lens, "topk_lens", layer_idx_); + + auto q_latent = project_q_nope_to_latent_(q_nope); + auto query_states = infinicore::op::concat_mla_q(q_latent, q_pe); + debug_dump_dsa(q_latent, "q_latent", layer_idx_); + debug_dump_dsa(query_states, "query_states", layer_idx_); + auto attn_output = infinicore::Tensor::empty( + {num_tokens, num_attention_heads_, kv_lora_rank_}, + query_states->dtype(), + query_states->device()); + auto attention_kv_cache = use_vendor_sparse + ? forward_context.mla_vendor_cache_vec[layer_idx_] + : kv_cache; + auto sparse_kv_cache = attention_kv_cache->view( + {static_cast( + attention_kv_cache->size(0) * attention_kv_cache->size(1)), + 1, + static_cast(attention_kv_cache->size(2))}); + infinicore::op::sparse_flash_mla_( + attn_output, + query_states, + sparse_kv_cache, + sparse_indices, + topk_lens, + softmax_scale_); + debug_dump_dsa(attn_output, "attn_output", layer_idx_); + auto sparse_output = project_latent_to_value_(attn_output, batch_size, seq_len); + debug_dump_dsa(sparse_output, "projected", layer_idx_); + return sparse_output; + } if (is_prefill) { if (attention_backend_ != infinilm::backends::AttentionBackend::FLASH_ATTN) { throw std::runtime_error("DeepseekV2MLAAttention: prefill requires --attn=flash-attn"); diff --git a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.hpp b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.hpp index 6044fca0e..b64653c81 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.hpp +++ b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.hpp @@ -2,7 +2,8 @@ #include "../../backends/attention_backends.hpp" #include "../../config/model_config.hpp" -#include "../../layers/linear/linear.hpp" +#include "../../layers/linear/fused_linear.hpp" +#include "deepseek_v2_indexer.hpp" #include "infinicore/nn/module.hpp" #include "infinicore/nn/rmsnorm.hpp" #include "infinicore/nn/rope.hpp" @@ -21,9 +22,11 @@ class DeepseekV2MLAAttention final : public infinicore::nn::Module { infinicore::Tensor forward(const infinicore::Tensor &positions, const infinicore::Tensor &hidden_states) const; + void process_weights_after_loading() override; + void reset_runtime_state() const override; + private: infinicore::Tensor position_ids_for_rope_(const infinicore::Tensor &position_ids) const; - infinicore::Tensor kv_b_weight_3d_() const; infinicore::Tensor project_q_nope_to_latent_(const infinicore::Tensor &q_nope) const; infinicore::Tensor project_latent_to_value_(const infinicore::Tensor &attn_output, size_t batch_size, @@ -41,17 +44,23 @@ class DeepseekV2MLAAttention final : public infinicore::nn::Module { size_t mla_head_dim_{0}; float softmax_scale_{1.0f}; infinilm::backends::AttentionBackend attention_backend_; + bool use_sparse_{false}; + bool skip_topk_{false}; + infinicore::Tensor w_uk_; + infinicore::Tensor w_uv_; INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_proj); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, q_a_proj); + INFINICORE_NN_MODULE(infinilm::layers::linear::MergedReplicatedLinear, fused_qkv_a_proj); INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, q_a_layernorm); INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, q_b_proj); INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, kv_a_proj_with_mqa); INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, kv_a_layernorm); INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, kv_b_proj); INFINICORE_NN_MODULE(infinilm::layers::linear::RowParallelLinear, o_proj); + INFINICORE_NN_MODULE(DeepseekV32Indexer, indexer); std::shared_ptr rotary_emb_; + infinicore::Tensor rope_cos_sin_cache_; infinicore::nn::Parameter kv_cache_k_scale_; infinicore::nn::Parameter kv_cache_v_scale_; }; diff --git a/csrc/models/glm_moe_dsa/glm_dsa_allocate_cache_tensors.cpp b/csrc/models/glm_moe_dsa/glm_dsa_allocate_cache_tensors.cpp new file mode 100644 index 000000000..4a0c21f18 --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_dsa_allocate_cache_tensors.cpp @@ -0,0 +1,132 @@ +#include "glm_dsa_allocate_cache_tensors.hpp" + +#include "../../cache/kv_cache.hpp" +#include "../../global_state/global_state.hpp" +#include "../../utils.hpp" + +#include +#include +#include + +namespace infinilm::models::glm_moe_dsa { + +namespace { +bool layer_uses_indexer_cache( + const std::shared_ptr &model_config, + size_t layer_idx) { + const auto &config_json = model_config->get_config_json(); + const auto indexer_types = config_json.value( + "indexer_types", nlohmann::json::array()); + if (layer_idx < indexer_types.size()) { + return indexer_types[layer_idx].get() != "shared"; + } + return true; +} +} // namespace + +GlmDsaCacheTensors glm_dsa_allocate_cache_tensors( + const cache::CacheConfig *cache_config, + const std::shared_ptr &model_config, + const backends::AttentionBackend &attention_backend, + size_t layer_start, + size_t layer_end) { + if (cache_config == nullptr) { + return {}; + } + if (model_config == nullptr) { + throw std::runtime_error("glm_dsa_allocate_cache_tensors: model config is null"); + } + if (attention_backend == backends::AttentionBackend::STATIC_ATTN) { + throw std::runtime_error("GLM-5.2 DSA requires paged cache"); + } + const auto *paged_config = dynamic_cast(cache_config); + if (paged_config == nullptr) { + throw std::runtime_error("glm_dsa_allocate_cache_tensors: expected paged cache config"); + } + if (paged_config->block_size() != 64) { + throw std::runtime_error("GLM-5.2 DSA requires block_size=64"); + } + + const size_t num_layers = model_config->get("num_hidden_layers"); + if (layer_start > layer_end || layer_end > num_layers) { + throw std::runtime_error( + "glm_dsa_allocate_cache_tensors: invalid pipeline layer range"); + } + const size_t kv_lora_rank = model_config->get("kv_lora_rank"); + const size_t rope_dim = model_config->get("qk_rope_head_dim"); + const bool use_vendor_shadow = std::getenv("INFINILM_GLM_FP8_SPARSE_VENDOR") != nullptr; + const size_t vendor_cache_stride = kv_lora_rank + rope_dim; + const size_t mla_cache_stride = kv_lora_rank + 4 * sizeof(float) + + rope_dim * sizeof(uint16_t); + const size_t index_dim = model_config->get("index_head_dim"); + const auto &rank_info = global_state::get_tensor_model_parallel_rank_info(); + const auto &device = rank_info.device; + + GlmDsaCacheTensors caches; + caches.mla.reserve(num_layers); + if (use_vendor_shadow) { + caches.mla_vendor.reserve(num_layers); + } + caches.indexer.reserve(num_layers); + for (size_t layer = 0; layer < num_layers; ++layer) { + if (layer < layer_start || layer >= layer_end) { + caches.mla.emplace_back(); + if (use_vendor_shadow) { + caches.mla_vendor.emplace_back(); + } + caches.indexer.emplace_back(); + continue; + } + auto mla = infinicore::Tensor::empty( + {paged_config->num_blocks(), paged_config->block_size(), mla_cache_stride}, + infinicore::DataType::U8, + device); + set_zeros(mla); + caches.mla.push_back(std::move(mla)); + + if (use_vendor_shadow) { + auto vendor_cache = infinicore::Tensor::empty( + {paged_config->num_blocks(), paged_config->block_size(), + vendor_cache_stride}, + infinicore::DataType::BF16, + device); + set_zeros(vendor_cache); + caches.mla_vendor.push_back(std::move(vendor_cache)); + } + + infinicore::Tensor indexer; + if (layer_uses_indexer_cache(model_config, layer)) { + indexer = infinicore::Tensor::empty( + {paged_config->num_blocks(), paged_config->block_size(), index_dim + sizeof(float)}, + infinicore::DataType::U8, device); + set_zeros(indexer); + } + caches.indexer.push_back(std::move(indexer)); + } + if (rank_info.tp_rank == 0) { + size_t indexer_layers = 0; + for (size_t layer = layer_start; layer < layer_end; ++layer) { + indexer_layers += layer_uses_indexer_cache(model_config, layer) ? 1 : 0; + } + spdlog::info( + "GLM DSA paged cache: physical_blocks={}, kernel_block_size={}, " + "mla_layout=fp8_ds_mla({}+4xfp32+{}xbf16={}B), " + "index_dim={}+fp32_scale, layers=[{},{}), indexer_layers={}, " + "indexer_cache_tp_ranks={}, vendor_shadow={}, " + "vendor_token_bytes={}", + paged_config->num_blocks(), + paged_config->block_size(), + kv_lora_rank, rope_dim, mla_cache_stride, + index_dim, + layer_start, + layer_end, + indexer_layers, + rank_info.tp_size, + use_vendor_shadow ? "enabled" : "disabled", + vendor_cache_stride * sizeof(uint16_t)); + } + infinicore::context::syncStream(); + return caches; +} + +} // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_dsa_allocate_cache_tensors.hpp b/csrc/models/glm_moe_dsa/glm_dsa_allocate_cache_tensors.hpp new file mode 100644 index 000000000..398b6e151 --- /dev/null +++ b/csrc/models/glm_moe_dsa/glm_dsa_allocate_cache_tensors.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "../../backends/attention_backends.hpp" +#include "../../cache/cache.hpp" +#include "../../config/model_config.hpp" +#include "infinicore/tensor.hpp" + +#include +#include + +namespace infinilm::models::glm_moe_dsa { + +struct GlmDsaCacheTensors { + std::vector mla; + std::vector mla_vendor; + std::vector indexer; +}; + +GlmDsaCacheTensors glm_dsa_allocate_cache_tensors( + const cache::CacheConfig *cache_config, + const std::shared_ptr &model_config, + const backends::AttentionBackend &attention_backend, + size_t layer_start, + size_t layer_end); + +} // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_model.cpp b/csrc/models/glm_moe_dsa/glm_model.cpp index 331abae70..21a7e9220 100644 --- a/csrc/models/glm_moe_dsa/glm_model.cpp +++ b/csrc/models/glm_moe_dsa/glm_model.cpp @@ -1,12 +1,45 @@ #include "glm_model.hpp" #include "../../global_state/global_state.hpp" +#include "../../utils.hpp" #include "../models_registry.hpp" +#include "glm_dsa_allocate_cache_tensors.hpp" +#include "infinicore/context/context.hpp" #include "infinicore/ops.hpp" +#include "infinicore/ops/cat.hpp" +#include "infinicore/ops/distributed/p2p.hpp" +#include +#include #include +#include #include #include +#include namespace infinilm::models::glm_moe_dsa { +namespace { +void debug_dump(const infinicore::Tensor &tensor, const std::string &name) { + if (std::getenv("INFINILM_GLM_DEBUG_DUMP") == nullptr) { + return; + } + const auto &rank = infinilm::global_state::get_tensor_model_parallel_rank_info(); + if (rank.global_rank != 0) { + return; + } + tensor->debug("/tmp/glmdbg_" + name + ".bin"); +} + +infinicore::Tensor i32_tensor_on_device( + const std::vector &values, + const infinicore::Device &device) { + auto tensor = infinicore::Tensor::empty( + {values.size()}, infinicore::DataType::I32, device); + infinicore::context::memcpyH2D( + tensor->data(), values.data(), values.size() * sizeof(int32_t), false); + return tensor; +} +} // namespace + GlmDecoder::GlmDecoder(std::shared_ptr c, size_t i, const infinicore::Device &d) { + layer_idx_ = i; auto h = c->get("hidden_size"); auto e = c->get("rms_norm_eps"); auto dt = c->get_dtype(); @@ -22,36 +55,401 @@ GlmDecoder::GlmDecoder(std::shared_ptr c, size_t } void GlmDecoder::forward(const infinicore::Tensor &p, infinicore::Tensor &x, infinicore::Tensor &r) const { input_layernorm_->forward_inplace(x, r); + debug_dump(x, "layer_" + std::to_string(layer_idx_) + "_input_norm"); x = self_attn_->forward(p, x); + debug_dump(x, "layer_" + std::to_string(layer_idx_) + "_attn"); post_attention_layernorm_->forward_inplace(x, r); + debug_dump(x, "layer_" + std::to_string(layer_idx_) + "_post_attn_norm"); x = moe_ ? moe_mlp_->forward(x) : dense_mlp_->forward(x); + debug_dump(x, "layer_" + std::to_string(layer_idx_) + "_mlp"); } + +namespace { +std::vector make_pipeline_boundaries( + const std::shared_ptr &config, + int pp_size) { + const size_t num_layers = config->get("num_hidden_layers"); + if (pp_size < 1 || static_cast(pp_size) > num_layers) { + throw std::runtime_error( + "GLM pipeline parallel size must be within the model layer count"); + } + const auto indexer_types = config->get_config_json() + .at("indexer_types") + .get>(); + std::vector boundaries(static_cast(pp_size) + 1, 0); + boundaries.back() = num_layers; + for (int stage = 1; stage < pp_size; ++stage) { + const size_t ideal = num_layers * static_cast(stage) + / static_cast(pp_size); + const size_t min_layer = boundaries[static_cast(stage) - 1] + 1; + const size_t max_layer = num_layers + - static_cast(pp_size - stage); + size_t best = ideal; + size_t best_distance = std::numeric_limits::max(); + for (size_t layer = min_layer; layer <= max_layer; ++layer) { + const bool fresh_indexer = layer >= indexer_types.size() + || indexer_types[layer] != "shared"; + const size_t distance = layer > ideal ? layer - ideal : ideal - layer; + if (fresh_indexer && distance < best_distance) { + best = layer; + best_distance = distance; + } + } + if (best_distance == std::numeric_limits::max()) { + best = std::max(min_layer, std::min(ideal, max_layer)); + } + boundaries[static_cast(stage)] = best; + } + return boundaries; +} +} // namespace + GlmModel::GlmModel(std::shared_ptr c, const infinicore::Device &d) { - auto dt = c->get_dtype(); - auto h = c->get("hidden_size"); - INFINICORE_NN_MODULE_INIT(embed_tokens, c->get("vocab_size"), h, dt, d); - for (size_t i = 0; i < c->get("num_hidden_layers"); ++i) { + dtype_ = c->get_dtype(); + hidden_size_ = c->get("hidden_size"); + const size_t num_layers = c->get("num_hidden_layers"); + const auto &rank = infinilm::global_state::get_tensor_model_parallel_rank_info(); + pp_size_ = rank.pp_size; + pp_rank_ = rank.pp_rank; + pp_comm_ = rank.pp_comm; + const auto boundaries = make_pipeline_boundaries(c, pp_size_); + layer_start_ = boundaries[static_cast(pp_rank_)]; + layer_end_ = boundaries[static_cast(pp_rank_) + 1]; + const auto indexer_types = c->get_config_json() + .at("indexer_types") + .get>(); + stage_boundary_needs_topk_.reserve( + static_cast(std::max(0, pp_size_ - 1))); + for (int stage = 0; stage + 1 < pp_size_; ++stage) { + const size_t next_layer = boundaries[static_cast(stage) + 1]; + stage_boundary_needs_topk_.push_back( + next_layer < indexer_types.size() + && indexer_types[next_layer] == "shared"); + } + + if (rank.is_pipeline_first_stage()) { + INFINICORE_NN_MODULE_INIT( + embed_tokens, c->get("vocab_size"), hidden_size_, dtype_, d); + } + for (size_t i = layer_start_; i < layer_end_; ++i) { layers_.push_back(register_module("layers." + std::to_string(i), c, i, d)); } - INFINICORE_NN_MODULE_INIT(norm, h, c->get("rms_norm_eps"), dt, d); + if (rank.is_pipeline_last_stage()) { + INFINICORE_NN_MODULE_INIT( + norm, hidden_size_, c->get("rms_norm_eps"), dtype_, d); + } + index_topk_ = c->get("index_topk"); +} + +void GlmModel::begin_pipeline_batch() const { + pipeline_send_lifetimes_.clear(); +} + +void GlmModel::transfer_pipeline_state_( + size_t source_stage, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + bool transfer_indexer_state) const { + if (pp_size_ == 1) { + return; + } + if (pp_comm_ == nullptr) { + throw std::runtime_error("GLM pipeline stage requires a PP communicator"); + } + if (!hidden_states || hidden_states->ndim() != 3 + || hidden_states->size(0) != 1 + || hidden_states->size(2) != hidden_size_) { + throw std::runtime_error("GLM pipeline hidden state has an invalid shape"); + } + + const bool is_source = static_cast(pp_rank_) == source_stage; + const bool is_destination = static_cast(pp_rank_) == source_stage + 1; + if (!is_source && !is_destination) { + return; + } + if (is_source) { + if (!residual || residual->shape() != hidden_states->shape()) { + throw std::runtime_error("GLM pipeline residual state is unavailable"); + } + } else { + hidden_states = infinicore::Tensor::empty( + {1, hidden_states->size(1), hidden_size_}, dtype_, + hidden_states->device()); + residual = infinicore::Tensor::empty( + hidden_states->shape(), dtype_, hidden_states->device()); + } + + std::vector pipeline_state{hidden_states, residual}; + if (transfer_indexer_state) { + auto &forward_context = infinilm::global_state::get_forward_context(); + if (!forward_context.dsa_topk_indices.has_value()) { + throw std::runtime_error("GLM pipeline indexer state is unavailable"); + } + auto &topk = forward_context.dsa_topk_indices.value(); + pipeline_state.push_back(topk); + } + + if (is_source) { + infinicore::op::distributed::send_grouped( + pipeline_state, static_cast(source_stage + 1), pp_comm_); + pipeline_send_lifetimes_.insert( + pipeline_send_lifetimes_.end(), + pipeline_state.begin(), pipeline_state.end()); + } else { + infinicore::op::distributed::recv_grouped_( + pipeline_state, static_cast(source_stage), pp_comm_); + if (!infinicore::context::isGraphRecording()) { + infinicore::context::syncStream(); + } + } } + infinicore::Tensor GlmModel::forward(const infinilm::InfinilmModel::Input &i) const { - auto x = embed_tokens_->forward(i.input_ids.value()); + if (!i.input_ids.has_value() || !i.position_ids.has_value()) { + throw std::runtime_error("GLM pipeline requires input and position IDs"); + } + const size_t num_tokens = i.position_ids.value()->numel(); + infinicore::Tensor x; + if (pp_rank_ == 0) { + x = embed_tokens_->forward(i.input_ids.value()); + debug_dump(x, "embed"); + } else { + x = infinicore::Tensor::empty( + {1, num_tokens, hidden_size_}, dtype_, i.position_ids.value()->device()); + set_zeros_device_async(x); + } + auto &forward_context = infinilm::global_state::get_forward_context(); + forward_context.dsa_topk_indices = infinicore::Tensor::empty( + {num_tokens, index_topk_}, + infinicore::DataType::I32, + x->device()); + set_minus_one_device_async(forward_context.dsa_topk_indices.value()); infinicore::Tensor r; - for (auto &l : layers_) { - l->forward(i.position_ids.value(), x, r); + for (int stage = 0; stage < pp_size_; ++stage) { + if (stage == pp_rank_) { + for (auto &l : layers_) { + l->forward(i.position_ids.value(), x, r); + } + } + if (stage + 1 < pp_size_) { + transfer_pipeline_state_( + static_cast(stage), x, r, + stage_boundary_needs_topk_[static_cast(stage)]); + } + } + if (pp_rank_ == pp_size_ - 1) { + norm_->forward_inplace(x, r); + debug_dump(x, "final"); } - norm_->forward_inplace(x, r); return x; } GlmForCausalLM::GlmForCausalLM(std::shared_ptr c, const infinicore::Device &d) { model_config_ = c; INFINICORE_NN_MODULE_INIT(model, c, d); - INFINICORE_NN_MODULE_INIT(lm_head, c->get("hidden_size"), c->get("vocab_size"), false, c->get_dtype(), d); + const auto &rank = infinilm::global_state::get_tensor_model_parallel_rank_info(); + is_output_stage_ = rank.is_pipeline_last_stage(); + if (is_output_stage_) { + INFINICORE_NN_MODULE_INIT(lm_head, c->get("hidden_size"), c->get("vocab_size"), false, c->get_dtype(), d); + } } infinilm::InfinilmModel::Output GlmForCausalLM::forward(const Input &i) const { - auto x = model_->forward(i); - return {lm_head_->forward(x)}; + auto select_logits_input = [](infinicore::Tensor x, const Input &input) { + if (!input.sample_all_positions && input.input_offsets.has_value()) { + const auto &input_offsets = input.input_offsets.value(); + if (input_offsets->numel() < 2) { + throw std::runtime_error( + "GLM logits selection requires at least one request"); + } + const size_t num_requests = input_offsets->numel() - 1; + const size_t num_tokens = x->size(0) * x->size(1); + if (num_tokens != num_requests) { + if (num_requests == 1) { + x = x->narrow({{1, x->size(1) - 1, 1}}); + } else { + auto selected = infinicore::Tensor::empty( + {1, num_requests, x->size(2)}, x->dtype(), x->device()); + infinicore::op::select_last_token_hidden_( + selected, x, input_offsets); + x = selected; + } + } + } + return x; + }; + + const auto &rank = infinilm::global_state::get_tensor_model_parallel_rank_info(); + const size_t num_requests = i.total_sequence_lengths.has_value() + ? i.total_sequence_lengths.value()->numel() + : 1; + const size_t num_tokens = i.position_ids.has_value() + ? i.position_ids.value()->numel() + : 0; + const bool decode_batch = num_tokens == num_requests; + model_->begin_pipeline_batch(); + if (rank.pp_size == 1 || num_requests <= 1 || decode_batch) { + auto x = model_->forward(i); + if (!is_output_stage_) { + return {x}; + } + return {lm_head_->forward(select_logits_input(x, i))}; + } + if (!i.input_ids.has_value() || !i.position_ids.has_value() + || !i.input_offsets.has_value() || !i.request_ids.has_value() + || !i.cu_seqlens.has_value() || !i.block_tables.has_value() + || !i.slot_mapping.has_value()) { + throw std::runtime_error( + "GLM PP microbatching requires complete paged-attention metadata"); + } + + std::vector input_offsets(num_requests + 1); + std::vector cu_seqlens(num_requests + 1); + if (decode_batch) { + for (size_t request = 0; request <= num_requests; ++request) { + input_offsets[request] = static_cast(request); + cu_seqlens[request] = static_cast(request); + } + } else { + infinicore::context::syncStream(); + infinicore::context::memcpyD2H( + input_offsets.data(), i.input_offsets.value()->data(), + input_offsets.size() * sizeof(int32_t)); + infinicore::context::memcpyD2H( + cu_seqlens.data(), i.cu_seqlens.value()->data(), + cu_seqlens.size() * sizeof(int32_t)); + } + + auto &forward_context = infinilm::global_state::get_forward_context(); + const auto full_attn_metadata = forward_context.attn_metadata; + infinicore::Tensor output_hidden; + infinicore::Tensor local_output; + constexpr size_t long_prefill_threshold = 512; + constexpr size_t long_prefill_microbatch_tokens = 2048; + const bool group_requests = num_tokens >= long_prefill_threshold; + size_t request_begin = 0; + while (request_begin < num_requests) { + size_t request_end = request_begin; + while (request_end < num_requests) { + const size_t candidate_tokens = static_cast( + input_offsets[request_end + 1] - input_offsets[request_begin]); + if (request_end > request_begin + && (!group_requests + || candidate_tokens > long_prefill_microbatch_tokens)) { + break; + } + ++request_end; + if (group_requests + && candidate_tokens >= long_prefill_microbatch_tokens) { + break; + } + } + + const size_t micro_requests = request_end - request_begin; + const size_t token_start = static_cast(input_offsets[request_begin]); + const size_t token_end = static_cast(input_offsets[request_end]); + if (token_end <= token_start || token_end > num_tokens) { + throw std::runtime_error("GLM PP microbatch has invalid token offsets"); + } + const size_t token_count = token_end - token_start; + std::vector micro_input_offsets(micro_requests + 1); + std::vector micro_cu_seqlens(micro_requests + 1); + std::vector micro_request_ids(token_count); + for (size_t local_request = 0; local_request < micro_requests; + ++local_request) { + const size_t global_request = request_begin + local_request; + micro_input_offsets[local_request] = input_offsets[global_request] - input_offsets[request_begin]; + micro_cu_seqlens[local_request] = cu_seqlens[global_request] - cu_seqlens[request_begin]; + const size_t local_token_begin = static_cast( + input_offsets[global_request] - input_offsets[request_begin]); + const size_t local_token_end = static_cast( + input_offsets[global_request + 1] - input_offsets[request_begin]); + std::fill( + micro_request_ids.begin() + local_token_begin, + micro_request_ids.begin() + local_token_end, + static_cast(local_request)); + } + micro_input_offsets[micro_requests] = static_cast(token_count); + micro_cu_seqlens[micro_requests] = cu_seqlens[request_end] - cu_seqlens[request_begin]; + Input micro = i; + micro.input_ids = i.input_ids.value()->narrow({{1, token_start, token_count}}); + micro.position_ids = i.position_ids.value()->ndim() == 1 + ? i.position_ids.value()->narrow({{0, token_start, token_count}}) + : i.position_ids.value()->narrow({{1, token_start, token_count}}); + if (i.past_sequence_lengths.has_value()) { + micro.past_sequence_lengths = i.past_sequence_lengths.value()->narrow( + {{0, request_begin, micro_requests}}); + } + micro.total_sequence_lengths = i.total_sequence_lengths.value()->narrow( + {{0, request_begin, micro_requests}}); + micro.block_tables = i.block_tables.value()->narrow( + {{0, request_begin, micro_requests}}); + micro.slot_mapping = i.slot_mapping.value()->narrow({{0, token_start, token_count}}); + micro.input_offsets = i32_tensor_on_device( + micro_input_offsets, micro.position_ids.value()->device()); + micro.request_ids = i32_tensor_on_device( + micro_request_ids, micro.position_ids.value()->device()); + micro.cu_seqlens = i32_tensor_on_device( + micro_cu_seqlens, micro.position_ids.value()->device()); + if (infinicore::context::isGraphRecording()) { + graph_microbatch_constants_.push_back(micro.input_offsets.value()); + graph_microbatch_constants_.push_back(micro.request_ids.value()); + graph_microbatch_constants_.push_back(micro.cu_seqlens.value()); + } + + forward_context.attn_metadata = { + micro.past_sequence_lengths, + micro.total_sequence_lengths, + micro.input_offsets, + micro.request_ids, + micro.cu_seqlens, + micro.block_tables, + micro.slot_mapping, + full_attn_metadata.max_context_len, + false}; + auto x = model_->forward(micro); + local_output = x; + if (is_output_stage_) { + auto selected = select_logits_input(x, micro); + if (!output_hidden) { + output_hidden = infinicore::Tensor::empty( + {1, num_requests, selected->size(2)}, + selected->dtype(), selected->device()); + } + output_hidden + ->narrow({{1, request_begin, micro_requests}}) + ->copy_from(selected); + } + request_begin = request_end; + } + forward_context.attn_metadata = full_attn_metadata; + if (!is_output_stage_) { + return {local_output}; + } + return {lm_head_->forward(output_hidden)}; +} + +void GlmForCausalLM::reset_cache(const cache::CacheConfig *cache_config) { + auto &forward_context = infinilm::global_state::get_forward_context(); + forward_context.dsa_topk_indices.reset(); + forward_context.kv_cache_vec.clear(); + forward_context.mla_vendor_cache_vec.clear(); + forward_context.indexer_cache_vec.clear(); + infinicore::context::syncStream(); + infinicore::context::trimMemory(); + if (cache_config == nullptr) { + cache_config_.reset(); + return; + } + cache_config_ = cache_config->unique_copy(); + auto caches = glm_dsa_allocate_cache_tensors( + cache_config, + model_config_, + global_state::get_infinilm_config().attention_backend, + model_->layer_start(), + model_->layer_end()); + forward_context.kv_cache_vec = std::move(caches.mla); + forward_context.mla_vendor_cache_vec = std::move(caches.mla_vendor); + forward_context.indexer_cache_vec = std::move(caches.indexer); + forward_context.dsa_topk_indices.reset(); } std::shared_ptr create_glm_config(std::shared_ptr c) { auto j = c->get_config_json(); diff --git a/csrc/models/glm_moe_dsa/glm_model.hpp b/csrc/models/glm_moe_dsa/glm_model.hpp index f052f140f..551149118 100644 --- a/csrc/models/glm_moe_dsa/glm_model.hpp +++ b/csrc/models/glm_moe_dsa/glm_model.hpp @@ -1,12 +1,14 @@ #pragma once #include "../../layers/linear/linear.hpp" #include "../../layers/mlp/mlp.hpp" +#include "../deepseek_v2/deepseek_v2_mla_attention.hpp" #include "../infinilm_model.hpp" -#include "glm_attention.hpp" #include "glm_moe.hpp" #include "glm_vocab_parallel.hpp" #include "infinicore/nn/embedding.hpp" #include "infinicore/nn/rmsnorm.hpp" +#include +#include namespace infinilm::models::glm_moe_dsa { class GlmDecoder final : public infinicore::nn::Module { public: @@ -16,29 +18,67 @@ class GlmDecoder final : public infinicore::nn::Module { private: INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); - INFINICORE_NN_MODULE(GlmAttention, self_attn); + INFINICORE_NN_MODULE(infinilm::models::deepseek_v2::DeepseekV2MLAAttention, self_attn); INFINICORE_NN_MODULE(GlmDenseMLP, dense_mlp); INFINICORE_NN_MODULE(GlmMoE, moe_mlp); bool moe_{false}; + size_t layer_idx_{0}; }; class GlmModel final : public infinicore::nn::Module { public: GlmModel(std::shared_ptr, const infinicore::Device &); infinicore::Tensor forward(const infinilm::InfinilmModel::Input &) const; + void begin_pipeline_batch() const; + + size_t layer_start() const { + return layer_start_; + } + size_t layer_end() const { + return layer_end_; + } private: + void transfer_pipeline_state_(size_t source_stage, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + bool transfer_indexer_state) const; + INFINICORE_NN_MODULE(GlmVocabEmbedding, embed_tokens); INFINICORE_NN_MODULE_VEC(GlmDecoder, layers); INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); + size_t index_topk_{0}; + size_t hidden_size_{0}; + size_t layer_start_{0}; + size_t layer_end_{0}; + int pp_size_{1}; + int pp_rank_{0}; + infinicclComm_t pp_comm_{nullptr}; + std::vector stage_boundary_needs_topk_; + mutable std::vector pipeline_send_lifetimes_; + infinicore::DataType dtype_; }; class GlmForCausalLM final : public infinilm::InfinilmModel { public: GlmForCausalLM(std::shared_ptr, const infinicore::Device &); Output forward(const Input &) const override; + void reset_cache(const cache::CacheConfig *) override; + size_t max_decode_graph_batch_size() const override { + return 16; + } + size_t decode_graph_batch_size(size_t batch_size) const override { + for (const size_t bucket : {1UL, 2UL, 4UL, 8UL, 16UL}) { + if (batch_size <= bucket) { + return bucket; + } + } + return batch_size; + } private: INFINICORE_NN_MODULE(GlmModel, model); INFINICORE_NN_MODULE(GlmVocabLMHead, lm_head); + bool is_output_stage_{true}; + mutable std::vector graph_microbatch_constants_; }; std::shared_ptr create_glm_config(std::shared_ptr); } // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_moe.cpp b/csrc/models/glm_moe_dsa/glm_moe.cpp index 63d1ad6f8..6d02987c7 100644 --- a/csrc/models/glm_moe_dsa/glm_moe.cpp +++ b/csrc/models/glm_moe_dsa/glm_moe.cpp @@ -67,12 +67,19 @@ GlmW4A8Experts::GlmW4A8Experts(std::shared_ptr c, register_parameter(p + ".down_proj.weight_scale", infinicore::nn::Parameter(ds)); } } -infinicore::Tensor GlmW4A8Experts::forward(const infinicore::Tensor &x, const infinicore::Tensor &ids, const infinicore::Tensor &tw) const { +infinicore::Tensor GlmW4A8Experts::forward(const infinicore::Tensor &x, + const infinicore::Tensor &ids, + const infinicore::Tensor &tw, + std::optional shared_output) const { if (!w1_) { throw std::runtime_error("GlmW4A8Experts: weights not ready"); } size_t m = x->size(0), total = m * topk_; - bool dec = m == 1; + auto &metadata = infinilm::global_state::get_forward_context().attn_metadata; + if (!metadata.total_sequence_lengths.has_value()) { + throw std::runtime_error("GlmW4A8Experts: missing sequence metadata"); + } + const bool dec = m == metadata.total_sequence_lengths.value()->numel(); int64_t fmt = dec ? 2 : 1; auto cnt = infinicore::Tensor::empty({nexpert_}, infinicore::DataType::I32, x->device()), sorted = infinicore::Tensor::empty({total}, infinicore::DataType::I32, x->device()), inv = infinicore::Tensor::empty({total}, infinicore::DataType::I32, x->device()); infinicore::op::moe_argsort_bincount_with_inv_pos_(cnt, sorted, inv, ids, nexpert_); @@ -86,7 +93,8 @@ infinicore::Tensor GlmW4A8Experts::forward(const infinicore::Tensor &x, const in auto a3 = infinicore::Tensor::empty({total, hidden_}, x->dtype(), x->device()); infinicore::op::w4a8_group_gemm_(a3, a2q, w2_, a2s, s2_, gc, sorted, std::nullopt, true, dec); auto out = infinicore::Tensor::empty({m, hidden_}, x->dtype(), x->device()); - infinicore::op::moe_sum_vllm_(out, a3->view({m, topk_, hidden_}), tw); + infinicore::op::moe_sum_vllm_( + out, a3->view({m, topk_, hidden_}), tw, shared_output); if (tp_ > 1 && comm_) { infinicore::op::distributed::allreduce_(out, out, INFINICCL_SUM, comm_); } @@ -100,6 +108,7 @@ GlmMoE::GlmMoE(std::shared_ptr c, const infinicor if (shared_) { auto j = c->get_config_json(); j["intermediate_size"] = c->get("moe_intermediate_size") * n; + j["reduce_results"] = false; auto sc = std::make_shared(j); INFINICORE_NN_MODULE_INIT(shared_experts, sc, d); } @@ -108,10 +117,10 @@ infinicore::Tensor GlmMoE::forward(const infinicore::Tensor &x) const { auto s = x->shape(); auto f = x->view({s[0] * s[1], s[2]}); auto [w, i] = gate_->forward(f); - auto r = experts_->forward(f, i, w)->view(s); - if (!shared_) { - return r; + std::optional shared_output; + if (shared_) { + shared_output = shared_experts_->forward(x)->view({s[0] * s[1], s[2]}); } - return infinicore::op::add(r, shared_experts_->forward(x)); + return experts_->forward(f, i, w, shared_output)->view(s); } } // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/glm_moe_dsa/glm_moe.hpp b/csrc/models/glm_moe_dsa/glm_moe.hpp index 887a95a4e..76386fc5a 100644 --- a/csrc/models/glm_moe_dsa/glm_moe.hpp +++ b/csrc/models/glm_moe_dsa/glm_moe.hpp @@ -6,6 +6,7 @@ #include "infinicore/tensor.hpp" #include #include +#include #include #include namespace infinilm::models::glm_moe_dsa { @@ -28,7 +29,10 @@ class GlmTopKRouter final : public infinicore::nn::Module { class GlmW4A8Experts final : public infinicore::nn::Module { public: GlmW4A8Experts(std::shared_ptr, const infinicore::Device &); - infinicore::Tensor forward(const infinicore::Tensor &, const infinicore::Tensor &, const infinicore::Tensor &) const; + infinicore::Tensor forward(const infinicore::Tensor &, + const infinicore::Tensor &, + const infinicore::Tensor &, + std::optional = std::nullopt) const; private: infinicore::Tensor w1_, s1_, w2_, s2_; diff --git a/csrc/models/glm_moe_dsa/glm_vocab_parallel.cpp b/csrc/models/glm_moe_dsa/glm_vocab_parallel.cpp index cc1fc2a01..20bbea987 100644 --- a/csrc/models/glm_moe_dsa/glm_vocab_parallel.cpp +++ b/csrc/models/glm_moe_dsa/glm_vocab_parallel.cpp @@ -35,8 +35,16 @@ infinicore::Tensor GlmVocabLMHead::forward(const infinicore::Tensor &x) const { if (world_ == 1) { return local; } - auto t = local->permute({2, 0, 1})->contiguous(); - auto g = infinicore::op::distributed::allgather(t, world_, comm_); - return g->permute({1, 2, 0})->contiguous(); + + // Match vLLM GroupCoordinator::all_gather(dim=-1): gather the contiguous + // logits directly, expose the rank-major leading dimension, then move the + // rank dimension next to local_vocab and flatten them together. This + // removes the old pre-gather permute().contiguous() copy. + auto gathered = infinicore::op::distributed::allgather(local, world_, comm_); + auto rank_major = gathered->view( + {world_, local->size(0), local->size(1), local->size(2)}); + auto vocab_major = rank_major->permute({1, 2, 0, 3})->contiguous(); + return vocab_major->view( + {local->size(0), local->size(1), world_ * local->size(2)}); } } // namespace infinilm::models::glm_moe_dsa diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index ac994fd6d..7d69345f5 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -28,6 +28,8 @@ class InfinilmModel : public infinicore::nn::Module { std::optional total_sequence_lengths; /// Offsets of each request in a continous-batched sequence, of shape `[num_requests + 1]`. std::optional input_offsets; + /// Request id for each flattened input token, of shape [num_tokens]. + std::optional request_ids; /// Cumulative total sequence lengths for each request, of shape `[num_requests + 1]`. std::optional cu_seqlens; /// Block ids for each request `[batch, max_block_table_length]`. Used for paged cache. @@ -55,6 +57,8 @@ class InfinilmModel : public infinicore::nn::Module { std::optional> visual_token_ranges; /// Target model hidden states consumed by draft/MTP models. std::optional target_hidden_states; + /// Keep logits for every packed input position (e.g. raw/MTP calls). + bool sample_all_positions{false}; }; struct Output { @@ -70,6 +74,12 @@ class InfinilmModel : public infinicore::nn::Module { virtual const cache::CacheConfig *get_cache_config() const { return cache_config_.get(); } + virtual size_t max_decode_graph_batch_size() const { + return 512; + } + virtual size_t decode_graph_batch_size(size_t batch_size) const { + return batch_size; + } void process_weights_after_loading(); void reset_runtime_state() const; diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index a8c72fa37..0f03fc704 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -11,12 +11,14 @@ namespace infinilm::engine::distributed { inline void bind_dist_config(py::module &m) { py::class_(m, "DistConfig") .def(py::init<>(), "Default constructor, empty device list") - .def(py::init(), py::arg("tp_size"), - "Constructor with tensor parallel size, auto-assigns device IDs 0..tp_size-1") + .def(py::init(), py::arg("tp_size"), py::arg("pp_size") = 1, + "Constructor with tensor and pipeline parallel sizes") .def(py::init &>(), py::arg("tp_device_ids"), "Constructor with explicit device IDs") .def_readwrite("tp_device_ids", &DistConfig::tp_device_ids, - "List of device IDs used in tensor parallelism") + "List of device IDs used by all distributed ranks") + .def_readwrite("tensor_parallel_size", &DistConfig::tensor_parallel_size) + .def_readwrite("pipeline_parallel_size", &DistConfig::pipeline_parallel_size) .def_readwrite("moe_ep_backend", &DistConfig::moe_ep_backend, "MoE expert-parallel backend") .def_readwrite("moe_ep_size", &DistConfig::moe_ep_size, @@ -137,6 +139,7 @@ inline void bind_infer_engine(py::module &m) { std::optional past_sequence_lengths, std::optional total_sequence_lengths, std::optional input_offsets, + std::optional request_ids, std::optional cu_seqlens, std::optional block_tables, std::optional slot_mapping, @@ -158,6 +161,7 @@ inline void bind_infer_engine(py::module &m) { std::move(past_sequence_lengths), std::move(total_sequence_lengths), std::move(input_offsets), + std::move(request_ids), std::move(cu_seqlens), std::move(block_tables), std::move(slot_mapping), @@ -184,6 +188,10 @@ inline void bind_infer_engine(py::module &m) { "temperature", "top_p", "top_k", + "keep_output_device", + "reuse_last_output", + "allow_graph_replay", + "is_mixed_batch", }; for (auto &item : kwargs) { @@ -200,6 +208,14 @@ inline void bind_infer_engine(py::module &m) { input.top_p = py::cast(item.second); } else if (key == "top_k") { input.top_k = py::cast(item.second); + } else if (key == "keep_output_device") { + input.keep_output_device = py::cast(item.second); + } else if (key == "reuse_last_output") { + input.reuse_last_output = py::cast(item.second); + } else if (key == "allow_graph_replay") { + input.allow_graph_replay = py::cast(item.second); + } else if (key == "is_mixed_batch") { + input.is_mixed_batch = py::cast(item.second); } } @@ -210,6 +226,7 @@ inline void bind_infer_engine(py::module &m) { py::arg("past_sequence_lengths") = std::nullopt, py::arg("total_sequence_lengths") = std::nullopt, py::arg("input_offsets") = std::nullopt, + py::arg("request_ids") = std::nullopt, py::arg("cu_seqlens") = std::nullopt, py::arg("block_tables") = std::nullopt, py::arg("slot_mapping") = std::nullopt, @@ -229,6 +246,7 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("past_sequence_lengths", &InferEngine::Input::past_sequence_lengths) .def_readwrite("total_sequence_lengths", &InferEngine::Input::total_sequence_lengths) .def_readwrite("input_offsets", &InferEngine::Input::input_offsets) + .def_readwrite("request_ids", &InferEngine::Input::request_ids) .def_readwrite("cu_seqlens", &InferEngine::Input::cu_seqlens) .def_readwrite("block_tables", &InferEngine::Input::block_tables) .def_readwrite("slot_mapping", &InferEngine::Input::slot_mapping) @@ -243,6 +261,10 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) .def_readwrite("max_context_len", &InferEngine::Input::max_context_len) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) + .def_readwrite("keep_output_device", &InferEngine::Input::keep_output_device) + .def_readwrite("reuse_last_output", &InferEngine::Input::reuse_last_output) + .def_readwrite("allow_graph_replay", &InferEngine::Input::allow_graph_replay) + .def_readwrite("is_mixed_batch", &InferEngine::Input::is_mixed_batch) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) .def_readwrite("top_p", &InferEngine::Input::top_p); diff --git a/csrc/utils.hpp b/csrc/utils.hpp index de1bf263f..9b93a455a 100644 --- a/csrc/utils.hpp +++ b/csrc/utils.hpp @@ -1,10 +1,12 @@ #pragma once -#include #include +#include +#include #include #include #include +#include #include #include #include @@ -127,9 +129,36 @@ inline void set_zeros(infinicore::Tensor &tensor) { infinicore::context::memcpyH2D(tensor->data(), zeros.data(), tensor->nbytes(), false); } +namespace infinilm::utils_detail { + +class DeviceMemsetGraphOperator final : public infinicore::graph::GraphOperator { +public: + DeviceMemsetGraphOperator(const infinicore::Tensor &tensor, int value) + : tensor_(tensor), value_(value) {} + + void run() const override { + infinicore::context::setDeviceMemoryAsync( + tensor_->data(), value_, tensor_->nbytes(), infinicore::context::getStream()); + } + +private: + mutable infinicore::Tensor tensor_; + int value_; +}; + +inline void set_device_memory_async_graph_aware(infinicore::Tensor &tensor, int value) { + auto op = std::make_shared(tensor, value); + if (infinicore::context::isGraphRecording()) { + infinicore::context::addGraphOperator(op); + } else { + op->run(); + } +} + +} // namespace infinilm::utils_detail + inline void set_zeros_device_async(infinicore::Tensor &tensor) { - infinicore::context::setDeviceMemoryAsync( - tensor->data(), 0, tensor->nbytes(), infinicore::context::getStream()); + infinilm::utils_detail::set_device_memory_async_graph_aware(tensor, 0); } inline void set_minus_one(infinicore::Tensor &tensor) { @@ -138,8 +167,7 @@ inline void set_minus_one(infinicore::Tensor &tensor) { } inline void set_minus_one_device_async(infinicore::Tensor &tensor) { - infinicore::context::setDeviceMemoryAsync( - tensor->data(), 0xFF, tensor->nbytes(), infinicore::context::getStream()); + infinilm::utils_detail::set_device_memory_async_graph_aware(tensor, 0xFF); } // Hash combine utility (similar to boost::hash_combine) diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 1e97c05a8..d879b3ce9 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -63,6 +63,7 @@ def __init__(self): self.num_draft_tokens = self.args.num_draft_tokens self.device = self.args.device self.tp = self.args.tp + self.pp = self.args.pp self.dp = self.args.dp self.ep = self.args.ep self.moe_ep_backend = self.args.moe_ep_backend @@ -81,6 +82,8 @@ def __init__(self): self.batch_size = self.args.batch_size self.max_batch_size = self.args.max_batch_size + self.max_num_batched_tokens = self.args.max_num_batched_tokens + self.max_num_mixed_prefill_tokens = self.args.max_num_mixed_prefill_tokens self.input_len = self.args.input_len self.output_len = self.args.output_len self.max_new_tokens = self.args.max_new_tokens @@ -112,6 +115,8 @@ def __init__(self): self.port = self.args.port self.endpoint = self.args.endpoint self.ignore_eos = self.args.ignore_eos + self.admission_workers = self.args.admission_workers + self.prefill_coalesce_ms = self.args.prefill_coalesce_ms # PD separation (KV transfer) self.kv_transfer_config = self.args.kv_transfer_config @@ -199,6 +204,9 @@ def _add_common_args(self): ), ) self.parser.add_argument("--tp", "--tensor-parallel-size", type=int, default=1) + self.parser.add_argument( + "--pp", "--pipeline-parallel-size", type=int, default=1 + ) self.parser.add_argument("--dp", "--data-parallel-size", type=int, default=1) self.parser.add_argument( "--ep", "--expert-parallel-size", type=int, default=None @@ -269,6 +277,24 @@ def _add_common_args(self): default=8, help="maximum batch size for server", ) + self.parser.add_argument( + "--max-num-batched-tokens", + type=int, + default=None, + help=( + "maximum tokens scheduled in one model step; " + "defaults are selected by model type" + ), + ) + self.parser.add_argument( + "--max-num-mixed-prefill-tokens", + type=int, + default=None, + help=( + "maximum prefill tokens allowed beside active decodes; " + "defaults are selected by model type" + ), + ) self.parser.add_argument( "--input-len", type=parse_list, default=10, help="input sequence length" ) @@ -367,6 +393,20 @@ def _add_common_args(self): self.parser.add_argument( "--endpoint", type=str, default="/completions", help="API endpoint" ) + self.parser.add_argument( + "--admission-workers", + type=int, + default=4, + help="number of request preprocessing workers used by the API server", + ) + self.parser.add_argument( + "--prefill-coalesce-ms", + type=float, + default=2.0, + help=( + "idle-to-prefill admission window in milliseconds; set to 0 to disable" + ), + ) self.parser.add_argument( "--ignore-eos", diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index 8ccee1b73..4c9c9158a 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -1,5 +1,5 @@ import json -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Optional @@ -21,6 +21,8 @@ class EngineConfig: moe_ep_size: MoE expert-parallel size. cache_type: Cache type ('paged' or 'static'). max_batch_size: Maximum batch size for inference (only for paged cache). + max_num_batched_tokens: Maximum tokens scheduled in one model step. + max_num_mixed_prefill_tokens: Maximum prefill tokens beside active decodes. max_tokens: Default maximum tokens to generate. num_blocks: Number of KV cache blocks (only for paged cache). block_size: Size of each KV cache block (only for paged cache). @@ -42,6 +44,7 @@ class EngineConfig: device: str = "cuda" dtype: str = "float16" tensor_parallel_size: int = 1 + pipeline_parallel_size: int = 1 moe_ep_backend: str = "disabled" moe_ep_size: int = 1 cache_type: str = "paged" # "paged" or "static" @@ -60,12 +63,27 @@ class EngineConfig: skip_load: bool = False skip_legacy_moe: bool = False kv_transfer_config: Optional[KVTransferConfig] = None + max_num_batched_tokens: Optional[int] = None + max_num_mixed_prefill_tokens: Optional[int] = None + kernel_block_size: int = field(init=False, default=0) def __post_init__(self) -> None: + if self.tensor_parallel_size < 1: + raise ValueError("tensor_parallel_size must be positive") + if self.pipeline_parallel_size < 1: + raise ValueError("pipeline_parallel_size must be positive") if self.num_draft_tokens < 1: raise ValueError("num_draft_tokens must be >= 1") if self.weight_load_mode not in {"async", "sync"}: raise ValueError("weight_load_mode must be either 'async' or 'sync'") + if self.num_blocks < 1: + raise ValueError("num_blocks must be positive") + if self.block_size < 1: + raise ValueError("block_size must be positive") + + # The scheduler block can be a multiple of the attention-kernel block. + # This mirrors vLLM's virtual block splitting for kernels such as DSA. + self.kernel_block_size = self.block_size config_path = Path(self.model_path) / "config.json" try: @@ -75,17 +93,43 @@ def __post_init__(self) -> None: raise ValueError(f"Unable to read model config: {config_path}") from exc text_config = model_config.get("text_config", model_config) model_type = text_config.get("model_type", model_config.get("model_type")) + if self.pipeline_parallel_size > 1 and model_type != "glm_moe_dsa": + raise ValueError( + "pipeline parallelism is currently supported only for glm_moe_dsa" + ) + + if self.max_num_batched_tokens is None: + self.max_num_batched_tokens = 4096 if model_type == "glm_moe_dsa" else 1024 + if self.max_num_batched_tokens < 1: + raise ValueError("max_num_batched_tokens must be positive") + if self.max_num_mixed_prefill_tokens is None: + self.max_num_mixed_prefill_tokens = ( + 256 if model_type == "glm_moe_dsa" else self.max_num_batched_tokens + ) + if self.max_num_mixed_prefill_tokens < 1: + raise ValueError("max_num_mixed_prefill_tokens must be positive") # DeepSeek V2 exposes only the verified MLA path. Keep use_mla as a # compatibility input, but derive the effective capability from the model. - self.use_mla = self.use_mla or model_type == "deepseek_v2" + self.use_mla = self.use_mla or model_type in {"deepseek_v2", "glm_moe_dsa"} if model_type == "deepseek_v2": if self.cache_type != "paged": raise ValueError( "DeepSeek V2 MLA requires paged cache; pass --enable-paged-attn" ) self.block_size = 16 + self.kernel_block_size = 16 self.attn_backend = "flash-attn" + elif model_type == "glm_moe_dsa": + if self.cache_type != "paged": + raise ValueError( + "GLM-5.2 DSA requires paged cache; pass --enable-paged-attn" + ) + self.kernel_block_size = 64 + if self.block_size % self.kernel_block_size != 0: + raise ValueError( + "GLM-5.2 scheduler block_size must be a multiple of 64" + ) if ( self.kv_transfer_config is not None @@ -93,3 +137,32 @@ def __post_init__(self) -> None: and self.cache_type != "paged" ): raise ValueError("kv_transfer_config requires cache_type='paged'") + if ( + self.kv_transfer_config is not None + and self.kv_transfer_config.kv_connector + and self.cache_block_size_factor != 1 + ): + raise ValueError( + "KV transfer does not yet support scheduler/kernel block splitting" + ) + + @property + def cache_block_size_factor(self) -> int: + """Number of kernel blocks contained in one scheduler block.""" + if self.cache_type != "paged": + return 1 + return self.block_size // self.kernel_block_size + + @property + def num_kernel_blocks(self) -> int: + """Physical kernel pages allocated by the model worker.""" + if self.cache_type != "paged": + return 0 + return self.num_blocks * self.cache_block_size_factor + + @property + def max_cache_tokens(self) -> int: + """Total token slots shared by all active paged-cache requests.""" + if self.cache_type != "paged": + return self.max_cache_len + return self.num_blocks * self.block_size diff --git a/python/infinilm/distributed/dist_config.py b/python/infinilm/distributed/dist_config.py index 1d245d364..91645a384 100644 --- a/python/infinilm/distributed/dist_config.py +++ b/python/infinilm/distributed/dist_config.py @@ -6,6 +6,7 @@ class DistConfig: def __init__( self, tp_size=None, + pp_size=1, tp_device_ids=None, moe_ep_backend="disabled", moe_ep_size=1, @@ -16,8 +17,10 @@ def __init__( raise ValueError("Provide either tp_size OR tp_device_ids, not both") if tp_size is not None: - self._underlying = _infinilm.DistConfig(tp_size) + self._underlying = _infinilm.DistConfig(tp_size, pp_size) elif tp_device_ids is not None: + if pp_size != 1: + raise ValueError("pp_size with explicit device IDs is not supported") self._underlying = _infinilm.DistConfig(tp_device_ids) else: self._underlying = _infinilm.DistConfig() @@ -48,6 +51,10 @@ def moe_ep_size(self): def moe_ep_size(self, value): self._underlying.moe_ep_size = int(value) + @property + def pipeline_parallel_size(self): + return self._underlying.pipeline_parallel_size + def __repr__(self): return repr(self._underlying) diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 1fb62245b..8c84e67ed 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -217,6 +217,7 @@ def _build_input( past_kv_lengths=None, total_kv_lengths=None, input_offsets=None, + request_ids=None, cu_seqlens=None, block_tables=None, slot_mapping=None, @@ -231,6 +232,10 @@ def _build_input( target_hidden_states=None, sample_all_positions=False, max_context_len=None, + keep_output_device=False, + reuse_last_output=False, + allow_graph_replay=True, + is_mixed_batch=False, temperature=None, top_k=None, top_p=None, @@ -245,6 +250,7 @@ def unwrap_tensor(tensor): past_kv_lengths = unwrap_tensor(past_kv_lengths) total_kv_lengths = unwrap_tensor(total_kv_lengths) input_offsets = unwrap_tensor(input_offsets) + request_ids = unwrap_tensor(request_ids) block_tables = unwrap_tensor(block_tables) cu_seqlens = unwrap_tensor(cu_seqlens) slot_mapping = unwrap_tensor(slot_mapping) @@ -276,6 +282,7 @@ def convert_tensor_list(tensor_list_): past_sequence_lengths=past_kv_lengths, total_sequence_lengths=total_kv_lengths, input_offsets=input_offsets, + request_ids=request_ids, cu_seqlens=cu_seqlens, block_tables=block_tables, slot_mapping=slot_mapping, @@ -290,6 +297,10 @@ def convert_tensor_list(tensor_list_): target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, max_context_len=max_context_len, + keep_output_device=keep_output_device, + reuse_last_output=reuse_last_output, + allow_graph_replay=allow_graph_replay, + is_mixed_batch=is_mixed_batch, temperature=temperature, top_k=top_k, top_p=top_p, @@ -303,6 +314,7 @@ def forward( past_kv_lengths=None, total_kv_lengths=None, input_offsets=None, + request_ids=None, cu_seqlens=None, block_tables=None, slot_mapping=None, @@ -315,6 +327,10 @@ def forward( image_req_ids=None, visual_token_ranges=None, target_hidden_states=None, + keep_output_device=False, + reuse_last_output=False, + allow_graph_replay=True, + is_mixed_batch=False, max_context_len=None, temperature=None, top_k=None, @@ -335,6 +351,7 @@ def forward( input_offsets = ( input_offsets._underlying if input_offsets is not None else None ) + request_ids = request_ids._underlying if request_ids is not None else None block_tables = ( block_tables._underlying if block_tables is not None else None ) @@ -376,6 +393,7 @@ def convert_tensor_list(tensor_list_): past_kv_lengths=past_kv_lengths, total_kv_lengths=total_kv_lengths, input_offsets=input_offsets, + request_ids=request_ids, cu_seqlens=cu_seqlens, block_tables=block_tables, slot_mapping=slot_mapping, @@ -388,6 +406,10 @@ def convert_tensor_list(tensor_list_): image_req_ids=image_req_ids, visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, + keep_output_device=keep_output_device, + reuse_last_output=reuse_last_output, + allow_graph_replay=allow_graph_replay, + is_mixed_batch=is_mixed_batch, max_context_len=max_context_len, temperature=temperature, top_k=top_k, @@ -408,6 +430,7 @@ def forward_raw( past_kv_lengths=None, total_kv_lengths=None, input_offsets=None, + request_ids=None, cu_seqlens=None, block_tables=None, slot_mapping=None, @@ -431,6 +454,7 @@ def forward_raw( past_kv_lengths=past_kv_lengths, total_kv_lengths=total_kv_lengths, input_offsets=input_offsets, + request_ids=request_ids, cu_seqlens=cu_seqlens, block_tables=block_tables, slot_mapping=slot_mapping, @@ -474,6 +498,8 @@ def generate( seq_len = initial_seqlen batch_size = initial_batch_size + keep_output_device = initial_batch_size > 1 or not generation_config.stop_on_eos + if batch_size != 1 and generation_config.max_new_tokens is None: raise ValueError( "When `batch_size > 1`, `max_new_tokens` must be specified." @@ -481,6 +507,7 @@ def generate( if _measure_and_log_time: time_measurements = [] + decode_start_time = None block_tables = None max_blocks_per_batch = 0 @@ -576,6 +603,14 @@ def generate( input_offsets = infinicore.from_list( [seq_len * i for i in range(batch_size + 1)], dtype=infinicore.int32 ) + request_ids = infinicore.from_list( + [ + request_id + for request_id in range(batch_size) + for _ in range(seq_len) + ], + dtype=infinicore.int32, + ) mamba_init_state_indices = None mamba_final_state_indices = None @@ -596,6 +631,7 @@ def generate( past_kv_lengths=past_kv_lengths, total_kv_lengths=total_kv_lengths, input_offsets=input_offsets, + request_ids=request_ids, cu_seqlens=cu_seqlens, block_tables=block_tables, slot_mapping=slot_mapping, @@ -605,6 +641,8 @@ def generate( image_bound=image_bound if iter == 0 else None, tgt_sizes=tgt_sizes if iter == 0 else None, temperature=generation_config.temperature, + keep_output_device=keep_output_device, + reuse_last_output=keep_output_device and iter > 0, top_k=generation_config.top_k, top_p=generation_config.top_p, ) @@ -625,9 +663,24 @@ def generate( past_seq_len = past_seq_len + seq_len if _measure_and_log_time: - end_time = time.perf_counter() + if keep_output_device: + if iter == 0: + infinicore.sync_device() + end_time = time.perf_counter() + time_measurements.append(end_time - start_time) + decode_start_time = end_time + elif iter == generation_config.max_new_tokens - 1: + infinicore.sync_device() + end_time = time.perf_counter() + decode_steps = generation_config.max_new_tokens - 1 + avg_decode_time = (end_time - decode_start_time) / decode_steps + time_measurements.extend([avg_decode_time] * decode_steps) + else: + end_time = time.perf_counter() + time_measurements.append(end_time - start_time) - time_measurements.append((end_time - start_time)) + if keep_output_device and not _measure_and_log_time: + infinicore.sync_device() if _measure_and_log_time: print( diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59e9e94ba..21f0c05d0 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -83,9 +83,21 @@ def __init__(self, config: EngineConfig): num_mamba_cache_blocks = max(2, config.num_blocks // 4) max_num_batched_tokens = int( - os.getenv("INFINILM_MAX_NUM_BATCHED_TOKENS", max_position_embeddings) + os.getenv( + "INFINILM_MAX_NUM_BATCHED_TOKENS", + str(config.max_num_batched_tokens), + ) + ) + if max_num_batched_tokens < 1: + raise ValueError("max_num_batched_tokens must be positive") + max_num_mixed_prefill_tokens = int( + os.getenv( + "INFINILM_MAX_NUM_MIXED_PREFILL_TOKENS", + str(config.max_num_mixed_prefill_tokens), + ) ) - assert 1024 <= max_num_batched_tokens <= max_position_embeddings + if max_num_mixed_prefill_tokens < 1: + raise ValueError("max_num_mixed_prefill_tokens must be positive") self.scheduler = Scheduler( max_batch_size=config.max_batch_size, @@ -93,10 +105,25 @@ def __init__(self, config: EngineConfig): block_size=config.block_size, max_num_batched_tokens=max_num_batched_tokens, connector=connector, + max_num_mixed_prefill_tokens=max_num_mixed_prefill_tokens, has_mamba_cache=has_mamba_cache, num_mamba_cache_blocks=num_mamba_cache_blocks, + max_model_len=max_position_embeddings, + allow_mixed_batch=( + getattr(self.processor, "supports_mixed_batch", False) + and self.model_runner.speculative_runner is None + ), + ) + logger.info( + "Scheduler paged cache: num_blocks=%s, block_size=%s, " + "effective_max_model_len=%s, max_num_batched_tokens=%s, " + "max_num_mixed_prefill_tokens=%s", + config.num_blocks, + config.block_size, + self.scheduler.max_model_len, + max_num_batched_tokens, + max_num_mixed_prefill_tokens, ) - logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}") if has_mamba_cache: logger.info( "Using Mamba cache with num_blocks=%s, zero_state_index=0", @@ -122,6 +149,11 @@ def add_request(self, request: InferenceRequest): """Add a request to the scheduler.""" self.scheduler.add_request(request) + def should_coalesce_prefill(self) -> bool: + """Return whether the next step starts a new prefill wave.""" + should_coalesce = getattr(self.scheduler, "should_coalesce_prefill", None) + return bool(should_coalesce and should_coalesce()) + def step(self) -> tuple[bool, list[tuple]]: """Run one inference step. @@ -177,6 +209,7 @@ def _update_requests( case _: raise ValueError(f"Unsupported cache_type: {self.cache_type}") pending = [] + completed_step_requests = [] for req, token_ids in zip(requests, sampled_tokens): if req.is_aborted(): logger.info( @@ -188,6 +221,23 @@ def _update_requests( req.mark_canceled() continue + chunk_end = req.prefill_chunk_end + if chunk_end is not None and chunk_end < req.get_prompt_length(): + self.scheduler.requeue_prefill_chunk(req, chunk_end) + continue + + if chunk_end is not None: + # Allocation for an intermediate chunk deliberately deferred + # hash registration until every prompt KV was materialized. + if self.cache_type == "paged": + self.scheduler.cache_manager.commit_blocks_hash( + req.block_table, + req.get_input_tokens(), + req.get_prompt_length(), + ) + req.prefill_chunk_end = None + + completed_step_requests.append(req) if not isinstance(token_ids, list): token_ids = [token_ids] @@ -246,7 +296,7 @@ def _update_requests( continue pending.append((req.output_queue.async_q, output)) - self.scheduler.complete_requests(requests) + self.scheduler.complete_requests(completed_step_requests) return pending def _check_request_finished(self, req: InferenceRequest, token_id: int) -> bool: @@ -316,6 +366,7 @@ def __init__( device: str = "cuda", dtype: str = "float16", tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, moe_ep_backend: str = "disabled", moe_ep_size: int = 1, cache_type: str = "paged", @@ -333,6 +384,8 @@ def __init__( weight_load_mode: str = "async", skip_load: bool = False, skip_legacy_moe: bool = False, + max_num_batched_tokens: Optional[int] = None, + max_num_mixed_prefill_tokens: Optional[int] = None, ): """Initialize LLM. @@ -362,10 +415,13 @@ def __init__( device=device, dtype=dtype, tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, moe_ep_backend=moe_ep_backend, moe_ep_size=moe_ep_size, cache_type=cache_type, max_batch_size=max_batch_size, + max_num_batched_tokens=max_num_batched_tokens, + max_num_mixed_prefill_tokens=max_num_mixed_prefill_tokens, max_tokens=max_tokens, num_blocks=num_blocks, block_size=block_size, @@ -523,6 +579,7 @@ def __init__( device: str = "cuda", dtype: str = "float16", tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, moe_ep_backend: str = "disabled", moe_ep_size: int = 1, cache_type: str = "paged", @@ -540,6 +597,9 @@ def __init__( use_mla: bool = False, weight_load_mode: str = "async", skip_legacy_moe: bool = False, + prefill_coalesce_ms: float = 0.0, + max_num_batched_tokens: Optional[int] = None, + max_num_mixed_prefill_tokens: Optional[int] = None, ): """Initialize AsyncLLMEngine. @@ -564,6 +624,7 @@ def __init__( kv_connector_extra_config: Extra config dict for KV connector. use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. + prefill_coalesce_ms: Idle-to-prefill admission window in milliseconds. """ config = EngineConfig( model_path=model_path, @@ -572,10 +633,13 @@ def __init__( device=device, dtype=dtype, tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, moe_ep_backend=moe_ep_backend, moe_ep_size=moe_ep_size, cache_type=cache_type, max_batch_size=max_batch_size, + max_num_batched_tokens=max_num_batched_tokens, + max_num_mixed_prefill_tokens=max_num_mixed_prefill_tokens, max_tokens=max_tokens, num_blocks=num_blocks, block_size=block_size, @@ -598,6 +662,9 @@ def __init__( self._loop: Optional[asyncio.AbstractEventLoop] = None self._healthy = True self._abort_queue: Optional[janus.Queue] = None + if prefill_coalesce_ms < 0: + raise ValueError("prefill_coalesce_ms must be non-negative") + self._prefill_coalesce_seconds = prefill_coalesce_ms / 1000.0 def is_healthy(self) -> bool: return bool(self._healthy) @@ -688,6 +755,11 @@ def _step_loop(self): while self._running: try: self._drain_abort_queue() + if ( + self._prefill_coalesce_seconds > 0 + and self.engine.should_coalesce_prefill() + ): + time.sleep(self._prefill_coalesce_seconds) did_work, pending = self.engine.step() if not did_work: time.sleep(0.003) @@ -718,6 +790,7 @@ def add_request( prompt: Optional[str] = None, prompt_token_ids: Optional[List[int]] = None, sampling_params: Optional[SamplingParams] = None, + chat_template_kwargs: Optional[dict] = None, request_id: Optional[str] = None, # For server use request_data: Optional[dict] = None, @@ -774,7 +847,9 @@ def add_request( ) prompt = self.engine.apply_chat_template( - messages, add_generation_prompt=add_generation_prompt + messages, + add_generation_prompt=add_generation_prompt, + chat_template_kwargs=chat_template_kwargs, ) mm_inputs = resolve_multimodal_inputs(messages) @@ -829,6 +904,7 @@ def add_chat_request( request_id: Optional[str] = None, request_data: Optional[dict] = None, add_generation_prompt: bool = True, + chat_template_kwargs: Optional[dict] = None, **kwargs, ) -> InferenceRequest: """Add a chat request to the engine. @@ -847,6 +923,7 @@ def add_chat_request( messages=messages, apply_chat_template=True, add_generation_prompt=add_generation_prompt, + chat_template_kwargs=chat_template_kwargs, sampling_params=sampling_params, request_id=request_id, request_data=request_data, diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index 45f1eaee1..2128e7adb 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -59,9 +59,18 @@ def __init__(self, config: EngineConfig): ) elif config.cache_type == "paged": cache_config = PagedKVCacheConfig( - num_blocks=config.num_blocks, block_size=config.block_size + num_blocks=config.num_kernel_blocks, + block_size=config.kernel_block_size, + ) + logger.info( + "Using Paged KV Cache: manager_blocks=%s, manager_block_size=%s, " + "kernel_blocks=%s, kernel_block_size=%s, token_capacity=%s", + config.num_blocks, + config.block_size, + config.num_kernel_blocks, + config.kernel_block_size, + config.max_cache_tokens, ) - logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}") else: raise ValueError(f"Unsupported cache_type: {config.cache_type}") @@ -71,6 +80,7 @@ def __init__(self, config: EngineConfig): device=self.device, distributed_config=DistConfig( config.tensor_parallel_size, + pp_size=config.pipeline_parallel_size, moe_ep_backend=config.moe_ep_backend, moe_ep_size=config.moe_ep_size, ), @@ -103,6 +113,10 @@ def __init__(self, config: EngineConfig): # Initialize processor self.processor = AutoInfinilmProcessor.from_pretrained(config.model_path) + if config.cache_type == "paged": + self.processor.configure_paged_cache( + config.block_size, config.kernel_block_size + ) # Initialize KV connector self.kv_connector = None @@ -115,7 +129,9 @@ def __init__(self, config: EngineConfig): ) kv_cache_list = self.model_engine.get_kv_cache() - assert len(kv_cache_list) == self.config.tensor_parallel_size + assert len(kv_cache_list) == ( + self.config.tensor_parallel_size * self.config.pipeline_parallel_size + ) kv_caches = {} for rank_idx, kv_cache_vec in enumerate(kv_cache_list): diff --git a/python/infinilm/llm/request.py b/python/infinilm/llm/request.py index abccfe539..80511fe79 100644 --- a/python/infinilm/llm/request.py +++ b/python/infinilm/llm/request.py @@ -162,6 +162,9 @@ def __init__( self.num_computed_tokens: int = 0 # Total tokens computed (local + remote) self.num_blocks: int = 0 + # Exclusive prompt offset for the prefill chunk scheduled this step. + self.prefill_chunk_end: Optional[int] = None + # Mamba cache management. None means no mamba cache row is currently owned. self.mamba_cache_index: Optional[int] = None diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index a1153e916..1b7a0f860 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -3,6 +3,7 @@ """ import logging +import os import queue from typing import List, Optional @@ -14,6 +15,10 @@ logger = logging.getLogger(__name__) +class RequestCapacityError(ValueError): + """Raised when one request can never fit in the configured KV cache.""" + + class SpeculativeCacheOps: """Limited cache operations needed by speculative verification.""" @@ -52,20 +57,41 @@ def __init__( scheduled_requests: List[InferenceRequest], is_prefill: bool = False, speculative_cache_ops: Optional[SpeculativeCacheOps] = None, + prefill_request_ids: Optional[set[str]] = None, ): self.scheduled_requests = scheduled_requests self.num_requests = len(scheduled_requests) - self.is_prefill = is_prefill + if prefill_request_ids is None: + prefill_request_ids = ( + {req.request_id for req in scheduled_requests} if is_prefill else set() + ) + scheduled_request_ids = {req.request_id for req in scheduled_requests} + unknown_ids = set(prefill_request_ids) - scheduled_request_ids + if unknown_ids: + raise ValueError( + f"Prefill request IDs are not part of the scheduled batch: {unknown_ids}" + ) + + self.prefill_request_ids = frozenset(prefill_request_ids) + self.is_prefill = bool(scheduled_requests) and ( + len(self.prefill_request_ids) == len(scheduled_requests) + ) + self.is_mixed = bool(self.prefill_request_ids) and not self.is_prefill + self.is_decode_only = bool(scheduled_requests) and not self.prefill_request_ids self.speculative_cache_ops = speculative_cache_ops self.kv_connector_metadata = None + def is_prefill_request(self, request: InferenceRequest) -> bool: + """Return whether a request contributes prompt tokens in this batch.""" + return request.request_id in self.prefill_request_ids + class Scheduler: """Request scheduler with integrated BlockManager for KV cache management. Scheduling logic: - 1. Running queue: Check for new blocks needed, update slot_mapping - 2. Waiting queue: Try block reuse (prefix caching), allocate new blocks + 1. Running queue: Schedule one decode token per active request + 2. Waiting queue: Fill remaining batch/token budget with prefills 3. Reference counting: Free blocks when requests complete """ @@ -78,6 +104,9 @@ def __init__( connector=None, has_mamba_cache: bool = False, num_mamba_cache_blocks: int | None = None, + allow_mixed_batch: bool = False, + max_model_len: int | None = None, + max_num_mixed_prefill_tokens: int | None = None, ): self.waiting_queue = janus.Queue() self.running_queue = janus.Queue() @@ -98,14 +127,57 @@ def __init__( ) self.speculative_cache_ops = SpeculativeCacheOps(self.cache_manager) self.block_size = block_size + cache_token_capacity = num_blocks * block_size + self.max_model_len = min( + max_model_len if max_model_len is not None else cache_token_capacity, + cache_token_capacity, + ) + if max_num_batched_tokens < 1: + raise ValueError("max_num_batched_tokens must be positive") self.max_num_batched_tokens = max_num_batched_tokens + # Keep one long request from entering the vendor sparse-prefill kernel + # as a single pathological launch while preserving the full batch budget. + max_prefill_chunk_tokens = int( + os.environ.get("INFINILM_MAX_PREFILL_CHUNK_TOKENS", "2048") + ) + if max_prefill_chunk_tokens < 1: + raise ValueError("INFINILM_MAX_PREFILL_CHUNK_TOKENS must be positive") + self.max_prefill_chunk_tokens = min( + max_num_batched_tokens, max_prefill_chunk_tokens + ) + if max_num_mixed_prefill_tokens is None: + max_num_mixed_prefill_tokens = max_num_batched_tokens + if max_num_mixed_prefill_tokens < 1: + raise ValueError("max_num_mixed_prefill_tokens must be positive") + self.max_num_mixed_prefill_tokens = max_num_mixed_prefill_tokens self.connector = connector + self.allow_mixed_batch = allow_mixed_batch + self._schedule_batch_id = 0 + + def validate_request(self, request: InferenceRequest) -> None: + max_tokens = request.sampling_params.max_tokens or 0 + requested_length = request.get_prompt_length() + max_tokens + if requested_length > self.max_model_len: + raise RequestCapacityError( + f"This request requires up to {requested_length} tokens " + f"({request.get_prompt_length()} prompt + {max_tokens} output), " + f"but the server's effective max model length is " + f"{self.max_model_len} tokens" + ) def add_request(self, request: InferenceRequest): if request is not None: + self.validate_request(request) request.status = RequestStatus.WAITING self.waiting_queue.sync_q.put(request) + def should_coalesce_prefill(self) -> bool: + """Return True only at an idle-to-prefill scheduling boundary.""" + return ( + self.running_queue.sync_q.qsize() == 0 + and self.waiting_queue.sync_q.qsize() > 0 + ) + def _exceeds_token_budget( self, current_num_batched_tokens: int, @@ -124,18 +196,92 @@ def _exceeds_token_budget( > self.max_num_batched_tokens ) + def _prefill_chunk_size( + self, + remaining_tokens: int, + current_num_batched_tokens: int, + current_num_prefill_tokens: int, + num_decode_requests: int, + ) -> int: + available = self.max_num_batched_tokens - current_num_batched_tokens + if num_decode_requests: + available = min( + available, + self.max_num_mixed_prefill_tokens - current_num_prefill_tokens, + ) + chunk_size = min(remaining_tokens, self.max_prefill_chunk_tokens, available) + if chunk_size < remaining_tokens: + # Intermediate boundaries must be scheduler-block aligned so that + # completed prefix blocks can be committed without partial hashes. + chunk_size = (chunk_size // self.block_size) * self.block_size + return max(chunk_size, 0) + + def requeue_prefill_chunk(self, request: InferenceRequest, chunk_end: int): + request.num_local_cached_tokens = chunk_end + request.num_computed_tokens = chunk_end + request.slot_mapping = [] + request.prefill_chunk_end = None + request.status = RequestStatus.WAITING + self.waiting_queue.sync_q.put(request) + + def _exceeds_mixed_prefill_budget( + self, + current_num_prefill_tokens: int, + num_tokens_this_step: int, + num_decode_requests: int, + ) -> bool: + if num_decode_requests == 0: + return False + return ( + current_num_prefill_tokens + num_tokens_this_step + > self.max_num_mixed_prefill_tokens + ) + def schedule(self) -> Optional[SchedulerOutput]: """Schedule and return batch of requests to execute.""" deferred_requests = [] scheduled_requests = [] - is_prefill = False + prefill_request_ids = set() current_num_batched_tokens = 0 - current_prefill_extra_blocks = 0 + current_num_prefill_tokens = 0 + current_reserved_extra_blocks = 0 - # Process Waiting queue (prefill phase) + # Protect inter-token latency by scheduling active decodes first. Waiting + # prefills can still join the same forward when batch/token budget remains. while ( len(scheduled_requests) < self.max_batch_size and current_num_batched_tokens < self.max_num_batched_tokens + ): + try: + req = self.running_queue.sync_q.get_nowait() + except queue.Empty: + break + if req.is_finished(): + self.complete_requests([req]) + continue + + try: + req.block_table, new_slot = self.cache_manager.append_slot( + req.block_table, req.get_total_length(), req.get_all_token_ids() + ) + except RuntimeError as e: + raise RuntimeError("No available cache blocks for new token") from e + + req.slot_mapping = [new_slot] + req.num_blocks = len(req.block_table) + req.num_local_cached_tokens = req.get_total_length() - 1 + scheduled_requests.append(req) + current_num_batched_tokens += 1 + current_reserved_extra_blocks += self._get_prefill_extra_blocks(req) + + num_decode_requests = len(scheduled_requests) + + # Fill the rest of the batch with waiting prefills. This forms a mixed + # prefill/decode batch whenever running requests leave capacity. + while ( + (self.allow_mixed_batch or not scheduled_requests) + and len(scheduled_requests) < self.max_batch_size + and current_num_batched_tokens < self.max_num_batched_tokens ): try: req = self.waiting_queue.sync_q.get_nowait() @@ -180,27 +326,23 @@ def schedule(self) -> Optional[SchedulerOutput]: if load_kv_async: num_computed_tokens -= 1 num_new_tokens = req.get_prompt_length() - num_computed_tokens + num_tokens_this_step = self._prefill_chunk_size( + num_new_tokens, + current_num_batched_tokens, + current_num_prefill_tokens, + num_decode_requests, + ) - # Early token budget check: skip can_accept_request and allocate_slots - # for requests that would exceed the per-schedule token budget. - if not load_kv_async: - num_tokens_this_step = ( - req.get_prompt_length() - num_local_computed_tokens - ) - if self._exceeds_token_budget( - current_num_batched_tokens, - num_tokens_this_step, - len(scheduled_requests), - ): - if num_local_computed_tokens > 0: - self.cache_manager.free_blocks(cached_block_table) - deferred_requests.append(req) - break + if not load_kv_async and num_tokens_this_step == 0: + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break if not self.can_accept_request( req, num_local_computed_tokens, - current_prefill_extra_blocks, + current_reserved_extra_blocks, ): logger.warning( "Insufficient KV cache blocks for request %s, deferring.", @@ -218,7 +360,9 @@ def schedule(self) -> Optional[SchedulerOutput]: num_computed_tokens=num_computed_tokens, cached_block_table=cached_block_table, blocks_blueprint=blocks_blueprint, - delay_cache_blocks=load_kv_async, + delay_cache_blocks=( + load_kv_async or num_tokens_this_step < num_new_tokens + ), ) if req_blocks is None: @@ -243,7 +387,11 @@ def schedule(self) -> Optional[SchedulerOutput]: break req.block_table = req_blocks - req.slot_mapping = slot_mapping + chunk_end = num_computed_tokens + num_tokens_this_step + req.slot_mapping = self.cache_manager.update_blocks_slot( + req_blocks, num_computed_tokens, chunk_end + ) + req.prefill_chunk_end = chunk_end req.num_blocks = len(req_blocks) req.num_local_cached_tokens = num_local_computed_tokens req.num_computed_tokens = num_computed_tokens @@ -257,19 +405,21 @@ def schedule(self) -> Optional[SchedulerOutput]: ) else: load_kv_async = False - num_tokens_this_step = ( - req.get_prompt_length() - req.num_local_cached_tokens - ) - if self._exceeds_token_budget( + remaining_tokens = req.get_prompt_length() - req.num_local_cached_tokens + num_tokens_this_step = self._prefill_chunk_size( + remaining_tokens, current_num_batched_tokens, - num_tokens_this_step, - len(scheduled_requests), - ): + current_num_prefill_tokens, + num_decode_requests, + ) + if num_tokens_this_step == 0: deferred_requests.append(req) break - self.cache_manager.update_blocks_hash( - req.block_table, req.num_local_cached_tokens + chunk_end = req.num_local_cached_tokens + num_tokens_this_step + req.slot_mapping = self.cache_manager.update_blocks_slot( + req.block_table, req.num_local_cached_tokens, chunk_end ) + req.prefill_chunk_end = chunk_end if load_kv_async: req.status = RequestStatus.WAITING_FOR_REMOTE_KVS @@ -279,11 +429,12 @@ def schedule(self) -> Optional[SchedulerOutput]: ) // self.block_size continue - current_prefill_extra_blocks += self._get_prefill_extra_blocks(req) + current_reserved_extra_blocks += self._get_prefill_extra_blocks(req) scheduled_requests.append(req) + prefill_request_ids.add(req.request_id) - num_tokens_this_step = req.get_prompt_length() - req.num_local_cached_tokens current_num_batched_tokens += num_tokens_this_step + current_num_prefill_tokens += num_tokens_this_step req.status = RequestStatus.RUNNING @@ -291,43 +442,6 @@ def schedule(self) -> Optional[SchedulerOutput]: for req in deferred_requests: self.waiting_queue.sync_q.put(req) - # Return prefill batch if any waiting requests were scheduled - if scheduled_requests: - is_prefill = True - scheduler_output = SchedulerOutput( - scheduled_requests=scheduled_requests, - is_prefill=is_prefill, - speculative_cache_ops=self.speculative_cache_ops, - ) - if self.connector is not None: - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output - - # Process Running queue (decode phase) - while len(scheduled_requests) < self.max_batch_size: - try: - req = self.running_queue.sync_q.get_nowait() - except queue.Empty: - break - # Skip requests that were already finished (e.g., timed out/canceled while running) - if req.is_finished(): - self.complete_requests([req]) - continue - - # Decode phase: allocate slot for newly generated token - try: - req.block_table, new_slot = self.cache_manager.append_slot( - req.block_table, req.get_total_length(), req.get_all_token_ids() - ) - req.slot_mapping = [new_slot] - req.num_blocks = len(req.block_table) - req.num_local_cached_tokens = req.get_total_length() - 1 - scheduled_requests.append(req) - - except RuntimeError as e: - raise RuntimeError("No available cache blocks for new token") from e - # Promote completed remote KV transfers (lower priority than running queue). # Cleanup (is_finished, failed re-queue) runs unconditionally; batch append only if slots remain. if self.connector is not None and self.remote_kv_requests: @@ -350,17 +464,37 @@ def schedule(self) -> Optional[SchedulerOutput]: ) self.update_waiting_for_remote_kv(req) req.status = RequestStatus.RUNNING + current_num_batched_tokens += 1 + current_reserved_extra_blocks += self._get_prefill_extra_blocks( + req + ) scheduled_requests.append(req) else: break # Defer promotion to next schedule() if batch is full # Return decode batch if any running requests were scheduled if scheduled_requests: - is_prefill = False + self._schedule_batch_id += 1 + if prefill_request_ids: + prefill_tokens = sum( + req.prefill_chunk_end - req.num_local_cached_tokens + for req in scheduled_requests + if req.request_id in prefill_request_ids and req.prefill_chunk_end + ) + logger.info( + "Scheduler prefill batch: id=%s requests=%s prefills=%s " + "prefill_tokens=%s waiting=%s mixed=%s", + self._schedule_batch_id, + len(scheduled_requests), + len(prefill_request_ids), + prefill_tokens, + self.waiting_queue.sync_q.qsize(), + bool(len(prefill_request_ids) != len(scheduled_requests)), + ) scheduler_output = SchedulerOutput( scheduled_requests=scheduled_requests, - is_prefill=is_prefill, speculative_cache_ops=self.speculative_cache_ops, + prefill_request_ids=prefill_request_ids, ) if self.connector is not None: @@ -478,13 +612,7 @@ def can_accept_request( running_queue_size = self.running_queue.sync_q.qsize() for _ in range(running_queue_size): req = self.running_queue.sync_q.get() - remaining_tokens = ( - req.sampling_params.max_tokens - req.get_num_generated_tokens() - ) - num_blocks_needed = ( - remaining_tokens + self.block_size - 1 - ) // self.block_size - total_required_blocks += num_blocks_needed + total_required_blocks += self._get_prefill_extra_blocks(req) self.running_queue.sync_q.put(req) # Calculate blocks needed for the new request diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 488235df0..328bc7ba5 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -68,14 +68,12 @@ def _is_glm_base_inference_unused_weight(key: str, config: dict) -> bool: """Weights intentionally outside the base 78-layer forward. Layers at num_hidden_layers and above belong to the optional MTP - predictor. The DSA indexer is not consulted by the current full-attention - path; for sequences no longer than index_topk, full attention is exact. - Keep this allowlist GLM-specific so other model loaders remain strict. + predictor. Base-model DSA indexer weights are loaded normally; only MTP + predictor layers remain outside this model. Keep this allowlist GLM-specific + so other model loaders remain strict. """ if config.get("model_type") != "glm_moe_dsa": return False - if ".self_attn.indexer." in key: - return True prefix = "model.layers." if key.startswith(prefix): rest = key[len(prefix) :] diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index b3bef06d5..de80ad49e 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -8,6 +8,8 @@ @register_processor("default") class BasicLLMProcessor(InfinilmProcessor): + supports_mixed_batch = True + def __init__(self, model_dir_path: str): self.tokenizer = AutoTokenizer.from_pretrained( model_dir_path, trust_remote_code=True @@ -148,6 +150,9 @@ def _build_model_input_from_static_scheduler_output( "input_offsets": infinicore.from_list( input_offsets, dtype=infinicore.int32 ), + "request_ids": infinicore.from_list( + [0] * len(input_ids[0]), dtype=infinicore.int32 + ), "cu_seqlens": infinicore.from_list( [0, total_kv_len], dtype=infinicore.int32 ), @@ -198,21 +203,25 @@ def _build_model_input_from_batch_scheduler_output( position_ids = [] cu_seqlens = [0] - max_block_table_len = max( - len(req.block_table) for req in scheduler_output.scheduled_requests + cache_block_size_factor = getattr(self, "cache_block_size_factor", 1) + max_block_table_len = ( + max(len(req.block_table) for req in scheduler_output.scheduled_requests) + * cache_block_size_factor ) current_offset = 0 for req in scheduler_output.scheduled_requests: + request_is_prefill = scheduler_output.is_prefill_request(req) num_cached = req.num_local_cached_tokens - if scheduler_output.is_prefill: + if request_is_prefill: # Prefill phase req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] + chunk_end = req.prefill_chunk_end or len(req_tokens) + tokens_to_compute = req_tokens[num_cached:chunk_end] tokens.extend(tokens_to_compute) compute_len = len(tokens_to_compute) - seq_len = len(req_tokens) + seq_len = chunk_end seq_lens.append(seq_len) current_offset += compute_len @@ -241,12 +250,19 @@ def _build_model_input_from_batch_scheduler_output( position_ids.append(seq_len - 1) # Pad block_table to same length - padded_block_table = req.block_table + [-1] * ( - max_block_table_len - len(req.block_table) + kernel_block_table = self.expand_block_table_for_kernel(req.block_table) + padded_block_table = kernel_block_table + [-1] * ( + max_block_table_len - len(kernel_block_table) ) block_tables.append(padded_block_table) cu_seqlens.append(cu_seqlens[-1] + seq_len) + request_ids = [ + request_idx + for request_idx in range(len(seq_offsets) - 1) + for _ in range(seq_offsets[request_idx + 1] - seq_offsets[request_idx]) + ] + return { "input_ids": infinicore.from_list([tokens], dtype=infinicore.int64), "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int64), @@ -255,10 +271,13 @@ def _build_model_input_from_batch_scheduler_output( ), "total_kv_lengths": infinicore.from_list(seq_lens, dtype=infinicore.int32), "input_offsets": infinicore.from_list(seq_offsets, dtype=infinicore.int32), + "request_ids": infinicore.from_list(request_ids, dtype=infinicore.int32), "cu_seqlens": infinicore.from_list(cu_seqlens, dtype=infinicore.int32), "block_tables": infinicore.from_list(block_tables, dtype=infinicore.int32), "slot_mapping": infinicore.from_list(slot_mapping, dtype=infinicore.int64), "max_context_len": max(seq_lens), + "allow_graph_replay": scheduler_output.is_decode_only, + "is_mixed_batch": scheduler_output.is_mixed, "temperature": temperature, "top_k": top_k, "top_p": top_p, diff --git a/python/infinilm/processors/processor.py b/python/infinilm/processors/processor.py index a2952bc1e..9fb22214b 100644 --- a/python/infinilm/processors/processor.py +++ b/python/infinilm/processors/processor.py @@ -29,6 +29,30 @@ def build_model_inputs(self, scheduler_output, **kwargs) -> dict: """Build batched infinilm model inputs from the scheduler output.""" raise NotImplementedError("build_model_inputs is not implemented yet") + def configure_paged_cache( + self, scheduler_block_size: int, kernel_block_size: int + ) -> None: + """Configure worker-side virtual block splitting.""" + if scheduler_block_size % kernel_block_size != 0: + raise ValueError( + "scheduler_block_size must be divisible by kernel_block_size" + ) + self.cache_block_size_factor = scheduler_block_size // kernel_block_size + + def expand_block_table_for_kernel(self, block_table: list[int]) -> list[int]: + """Map scheduler block IDs to consecutive physical kernel block IDs.""" + factor = getattr(self, "cache_block_size_factor", 1) + if factor == 1: + return list(block_table) + expanded = [] + for block_id in block_table: + if block_id < 0: + expanded.extend([-1] * factor) + else: + first_kernel_block = block_id * factor + expanded.extend(range(first_kernel_block, first_kernel_block + factor)) + return expanded + def get_tokenizer(self): """Return the text tokenizer associated with this processor.""" raise NotImplementedError("get_tokenizer is not implemented yet") diff --git a/python/infinilm/processors/qwen3_5_processor.py b/python/infinilm/processors/qwen3_5_processor.py index e550de5fd..ae60612b3 100644 --- a/python/infinilm/processors/qwen3_5_processor.py +++ b/python/infinilm/processors/qwen3_5_processor.py @@ -12,6 +12,8 @@ @register_processor("qwen3_5") class Qwen35Processor(BasicLLMProcessor): + supports_mixed_batch = False + def __init__(self, model_dir_path: str): self.pixel_values_dtype = None config_path = os.path.join(model_dir_path, "config.json") diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 242593e25..5fa99556c 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -9,7 +9,9 @@ import sys import time import uuid +from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager +from functools import partial from typing import Optional import uvicorn @@ -18,6 +20,7 @@ from infinilm.base_config import BaseConfig from infinilm.config import KVTransferConfig from infinilm.llm import AsyncLLMEngine, FinishReason, SamplingParams +from infinilm.llm.scheduler import RequestCapacityError from infinilm.moe_config import configure_moe_ep_backend logger = logging.getLogger(__name__) @@ -98,6 +101,7 @@ def __init__( device: str = "cuda", dtype: str = "float16", tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, moe_ep_backend: str = "disabled", moe_ep_size: int = 1, cache_type: str = "paged", @@ -117,6 +121,10 @@ def __init__( weight_load_mode: str = "async", ignore_eos: bool = False, kv_transfer_config: Optional[KVTransferConfig] = None, + admission_workers: int = 4, + prefill_coalesce_ms: float = 2.0, + max_num_batched_tokens: Optional[int] = None, + max_num_mixed_prefill_tokens: Optional[int] = None, ): """Initialize inference server. @@ -130,6 +138,8 @@ def __init__( cache_type: Cache type ('paged' or 'static'). max_tokens: Default maximum tokens to generate. max_batch_size: Maximum batch size for inference (only for paged cache). + max_num_batched_tokens: Maximum tokens scheduled in one model step. + max_num_mixed_prefill_tokens: Maximum prefill tokens beside decodes. num_blocks: Number of KV cache blocks (only for paged cache). block_size: Size of each KV cache block (only for paged cache). max_cache_len: Maximum sequence length (only for static cache). @@ -144,6 +154,8 @@ def __init__( weight_load_mode: Weight loading mode across tensor-parallel workers. ignore_eos: Whether to ignore EOS tokens during generation. kv_transfer_config: Optional configuration for the KV transfer mechanism. + admission_workers: Number of API request preprocessing workers. + prefill_coalesce_ms: Idle-to-prefill admission window in milliseconds. """ self.model_path = model_path # vLLM-like served model id: directory name of model_path @@ -151,11 +163,14 @@ def __init__( self.device = device self.dtype = dtype self.tensor_parallel_size = tensor_parallel_size + self.pipeline_parallel_size = pipeline_parallel_size self.moe_ep_backend = moe_ep_backend self.moe_ep_size = moe_ep_size self.cache_type = cache_type self.max_tokens = max_tokens self.max_batch_size = max_batch_size + self.max_num_batched_tokens = max_num_batched_tokens + self.max_num_mixed_prefill_tokens = max_num_mixed_prefill_tokens self.num_blocks = num_blocks self.block_size = block_size self.max_cache_len = max_cache_len @@ -170,8 +185,17 @@ def __init__( self.weight_load_mode = weight_load_mode self.ignore_eos = ignore_eos self.kv_transfer_config = kv_transfer_config + if admission_workers < 1: + raise ValueError("admission_workers must be positive") + if prefill_coalesce_ms < 0: + raise ValueError("prefill_coalesce_ms must be non-negative") + self.admission_workers = admission_workers + self.prefill_coalesce_ms = prefill_coalesce_ms self.engine: AsyncLLMEngine = None + self._admission_executor: Optional[ThreadPoolExecutor] = None + self._profiler = None + self._profile_dir = None def start(self): """Start the HTTP server.""" @@ -185,33 +209,55 @@ def _create_app(self): @asynccontextmanager async def lifespan(app: FastAPI): - self.engine = AsyncLLMEngine( - model_path=self.model_path, - device=self.device, - dtype=self.dtype, - tensor_parallel_size=self.tensor_parallel_size, - moe_ep_backend=self.moe_ep_backend, - moe_ep_size=self.moe_ep_size, - cache_type=self.cache_type, - max_batch_size=self.max_batch_size, - max_tokens=self.max_tokens, - num_blocks=self.num_blocks, - block_size=self.block_size, - max_cache_len=self.max_cache_len, - temperature=self.temperature, - top_p=self.top_p, - top_k=self.top_k, - enable_graph=self.enable_graph, - attn_backend=self.attn_backend, - use_mla=self.use_mla, - weight_load_mode=self.weight_load_mode, - kv_transfer_config=self.kv_transfer_config, + self._admission_executor = ThreadPoolExecutor( + max_workers=self.admission_workers, + thread_name_prefix="InfiniLMAdmission", ) - self.engine.start() - logger.info(f"Engine initialized with model at {self.model_path}") - logger.info(f" enable_graph: {self.enable_graph}") - yield - self.engine.stop() + try: + self.engine = AsyncLLMEngine( + model_path=self.model_path, + device=self.device, + dtype=self.dtype, + tensor_parallel_size=self.tensor_parallel_size, + pipeline_parallel_size=self.pipeline_parallel_size, + moe_ep_backend=self.moe_ep_backend, + moe_ep_size=self.moe_ep_size, + cache_type=self.cache_type, + max_batch_size=self.max_batch_size, + max_num_batched_tokens=self.max_num_batched_tokens, + max_num_mixed_prefill_tokens=(self.max_num_mixed_prefill_tokens), + max_tokens=self.max_tokens, + num_blocks=self.num_blocks, + block_size=self.block_size, + max_cache_len=self.max_cache_len, + temperature=self.temperature, + top_p=self.top_p, + top_k=self.top_k, + enable_graph=self.enable_graph, + attn_backend=self.attn_backend, + use_mla=self.use_mla, + weight_load_mode=self.weight_load_mode, + kv_transfer_config=self.kv_transfer_config, + prefill_coalesce_ms=self.prefill_coalesce_ms, + ) + self.engine.start() + logger.info(f"Engine initialized with model at {self.model_path}") + logger.info(f" enable_graph: {self.enable_graph}") + logger.info( + " admission_workers: %s, prefill_coalesce_ms: %.3f, " + "max_num_batched_tokens: %s, " + "max_num_mixed_prefill_tokens: %s", + self.admission_workers, + self.prefill_coalesce_ms, + self.engine.config.max_num_batched_tokens, + self.engine.config.max_num_mixed_prefill_tokens, + ) + yield + finally: + if self.engine is not None: + self.engine.stop() + self._admission_executor.shutdown(wait=True, cancel_futures=True) + self._admission_executor = None app = FastAPI(lifespan=lifespan) self._register_routes(app) @@ -246,13 +292,37 @@ async def chat_completions(request: Request): stream = data.get("stream", False) request_id = f"cmpl-{uuid.uuid4().hex}" + try: + req, sampling_params = await self._prepare_chat_request( + request_id, data + ) + except RequestCapacityError as e: + return JSONResponse( + content={ + "error": { + "message": str(e), + "type": "invalid_request_error", + "param": None, + "code": None, + } + }, + status_code=400, + ) + except Exception as e: + logger.error( + f"Failed to prepare request {request_id}: {e}", exc_info=True + ) + return JSONResponse(content={"error": str(e)}, status_code=500) + if stream: return StreamingResponse( - self._stream_chat(request_id, data, request), + self._stream_chat(request_id, data, request, req, sampling_params), media_type="text/event-stream", ) else: - response = await self._chat(request_id, data, request) + response = await self._chat( + request_id, data, request, req, sampling_params + ) if isinstance(response, JSONResponse): return response return JSONResponse(content=response) @@ -268,6 +338,58 @@ async def health(): return JSONResponse(content={"status": "unhealthy"}, status_code=503) return {"status": "healthy"} + @app.post("/start_profile") + async def start_profile(): + profile_root = os.getenv("INFINILM_TORCH_PROFILE_DIR") + if not profile_root: + return JSONResponse( + content={"error": "INFINILM_TORCH_PROFILE_DIR is not set"}, + status_code=404, + ) + if self._profiler is not None: + return JSONResponse( + content={"error": "Profiler is already running"}, status_code=409 + ) + + import torch + + profile_dir = os.path.join( + os.path.abspath(profile_root), + f"infinilm_{int(time.time())}_{os.getpid()}", + ) + os.makedirs(profile_dir, exist_ok=False) + self._profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + on_trace_ready=torch.profiler.tensorboard_trace_handler(profile_dir), + record_shapes=False, + profile_memory=False, + with_stack=False, + with_flops=False, + with_modules=False, + ) + self._profile_dir = profile_dir + self._profiler.start() + logger.info("Profiler started: %s", profile_dir) + return {"status": "started", "profile_dir": profile_dir} + + @app.post("/stop_profile") + async def stop_profile(): + if self._profiler is None: + return JSONResponse( + content={"error": "Profiler is not running"}, status_code=409 + ) + + profiler = self._profiler + profile_dir = self._profile_dir + self._profiler = None + self._profile_dir = None + profiler.stop() + logger.info("Profiler stopped: %s", profile_dir) + return {"status": "stopped", "profile_dir": profile_dir} + def _models_payload(): return { "object": "list", @@ -341,10 +463,20 @@ def pick(key: str, default): return default # Accept common alias - max_tokens = pick("max_tokens", self.max_tokens) - if max_tokens is None: - # Some clients use max_new_tokens - max_tokens = pick("max_new_tokens", self.max_tokens) + if data.get("max_tokens") is not None: + max_tokens = data["max_tokens"] + elif data.get("max_completion_tokens") is not None: + max_tokens = data["max_completion_tokens"] + elif sp.get("max_tokens") is not None: + max_tokens = sp["max_tokens"] + elif sp.get("max_completion_tokens") is not None: + max_tokens = sp["max_completion_tokens"] + elif data.get("max_new_tokens") is not None: + max_tokens = data["max_new_tokens"] + elif sp.get("max_new_tokens") is not None: + max_tokens = sp["max_new_tokens"] + else: + max_tokens = self.max_tokens stop = pick("stop", None) if isinstance(stop, str): @@ -356,27 +488,50 @@ def pick(key: str, default): top_k=int(pick("top_k", self.top_k)), max_tokens=int(max_tokens) if max_tokens is not None else None, stop=stop, - ignore_eos=self.ignore_eos, + ignore_eos=bool(pick("ignore_eos", self.ignore_eos)), + ) + + async def _prepare_chat_request(self, request_id: str, data: dict): + if self._admission_executor is None: + raise RuntimeError("Admission executor is not running") + loop = asyncio.get_running_loop() + started = time.perf_counter() + result = await loop.run_in_executor( + self._admission_executor, + partial(self._prepare_chat_request_sync, request_id, data), + ) + logger.debug( + "Prepared request %s in %.3f ms", + request_id, + (time.perf_counter() - started) * 1000.0, + ) + return result + + def _prepare_chat_request_sync(self, request_id: str, data: dict): + messages = data.get("messages", []) + sampling_params = self._build_sampling_params(data) + req = self.engine.add_chat_request( + messages=messages, + sampling_params=sampling_params, + request_id=request_id, + request_data=data, + add_generation_prompt=bool(data.get("add_generation_prompt", True)), + chat_template_kwargs=data.get("chat_template_kwargs") or {}, ) + return req, sampling_params - async def _stream_chat(self, request_id: str, data: dict, http_request: Request): + async def _stream_chat( + self, + request_id: str, + data: dict, + http_request: Request, + req, + sampling_params: SamplingParams, + ): """Handle streaming chat request.""" - req = None _abort_reason = FinishReason.CANCELED try: - messages = data.get("messages", []) - sampling_params = self._build_sampling_params(data) - - req = self.engine.add_chat_request( - messages=messages, - sampling_params=sampling_params, - request_id=request_id, - request_data=data, - add_generation_prompt=bool(data.get("add_generation_prompt", True)), - chat_template_kwargs=data.get("chat_template_kwargs") or {}, - ) - async for token_output in self.engine.stream_request( req, timeout=DEFAULT_STREAM_TIMEOUT, @@ -464,24 +619,18 @@ async def _stream_chat(self, request_id: str, data: dict, http_request: Request) self.engine.add_aborted_req(req, _abort_reason) yield "data: [DONE]\n\n" - async def _chat(self, request_id: str, data: dict, http_request: Request): + async def _chat( + self, + request_id: str, + data: dict, + http_request: Request, + req, + sampling_params: SamplingParams, + ): """Handle non-streaming chat request.""" - req = None _abort_reason = FinishReason.CANCELED try: - messages = data.get("messages", []) - sampling_params = self._build_sampling_params(data) - - req = self.engine.add_chat_request( - messages=messages, - sampling_params=sampling_params, - request_id=request_id, - request_data=data, - add_generation_prompt=bool(data.get("add_generation_prompt", True)), - chat_template_kwargs=data.get("chat_template_kwargs") or {}, - ) - # Collect all generated tokens output_text = "" async for token_output in self.engine.stream_request( @@ -510,7 +659,6 @@ async def _chat(self, request_id: str, data: dict, http_request: Request): if token_output.finished: break - output_text = output_text.strip() finish_reason = self._convert_finish_reason(req.finish_reason) response = completion_json( @@ -593,9 +741,10 @@ def main(): cfg.tp, cfg.dp, cfg.ep, cfg.moe_ep_backend, cfg.model ) logger.info( - "MoE EP backend: %s TP=%s DP=%s EP=%s", + "MoE EP backend: %s TP=%s PP=%s DP=%s EP=%s", moe_ep_backend, cfg.tp, + cfg.pp, cfg.dp, ep, ) @@ -605,11 +754,14 @@ def main(): device=device, dtype=cfg.dtype, tensor_parallel_size=cfg.tp, + pipeline_parallel_size=cfg.pp, moe_ep_backend=moe_ep_backend, moe_ep_size=ep, cache_type="paged" if cfg.enable_paged_attn else "static", max_tokens=cfg.max_new_tokens, max_batch_size=cfg.max_batch_size, + max_num_batched_tokens=cfg.max_num_batched_tokens, + max_num_mixed_prefill_tokens=cfg.max_num_mixed_prefill_tokens, num_blocks=cfg.num_blocks, block_size=cfg.block_size, max_cache_len=cfg.max_cache_len, @@ -624,6 +776,8 @@ def main(): weight_load_mode=cfg.weight_load_mode, ignore_eos=cfg.ignore_eos, kv_transfer_config=kv_transfer_config, + admission_workers=cfg.admission_workers, + prefill_coalesce_ms=cfg.prefill_coalesce_ms, ) server.start() From 299b8cfa19599ced34169affb631fc43ad362ffd Mon Sep 17 00:00:00 2001 From: wooway777 Date: Mon, 27 Jul 2026 07:38:47 +0000 Subject: [PATCH 4/5] pepe: adjust decode graph batch size --- csrc/models/glm_moe_dsa/glm_model.hpp | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/csrc/models/glm_moe_dsa/glm_model.hpp b/csrc/models/glm_moe_dsa/glm_model.hpp index 551149118..b4b20e69d 100644 --- a/csrc/models/glm_moe_dsa/glm_model.hpp +++ b/csrc/models/glm_moe_dsa/glm_model.hpp @@ -66,12 +66,13 @@ class GlmForCausalLM final : public infinilm::InfinilmModel { return 16; } size_t decode_graph_batch_size(size_t batch_size) const override { - for (const size_t bucket : {1UL, 2UL, 4UL, 8UL, 16UL}) { - if (batch_size <= bucket) { - return bucket; - } + // Exact graphs avoid dummy-request padding in the latency-sensitive + // 1-8 concurrency range. Larger decode batches may still share the + // batch-16 graph to cap capture count and memory use. + if (batch_size <= 8) { + return batch_size; } - return batch_size; + return batch_size <= 16 ? 16 : batch_size; } private: From 62043089e28418cb4403c4a1091f84639fba8959 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Mon, 27 Jul 2026 07:58:33 +0000 Subject: [PATCH 5/5] pepe: support pp in bench.py --- examples/bench.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/examples/bench.py b/examples/bench.py index a23fa8e9b..f52bc06cc 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -190,6 +190,7 @@ def __init__( num_draft_tokens=4, infini_device=infinicore.device("cpu", 0), tp=1, + pp=1, skip_load=False, cache_config=None, enable_graph=False, @@ -205,6 +206,7 @@ def __init__( self.model_path = model_path self.device_str = infini_device.type self.tp = tp + self.pp = pp self.cache_config = cache_config self.enable_graph = enable_graph self.attn_backend = attn_backend @@ -232,6 +234,7 @@ def __init__( device=infini_device, distributed_config=DistConfig( tp, + pp_size=pp, moe_ep_backend=moe_ep_backend, moe_ep_size=moe_ep_size, ), @@ -276,6 +279,7 @@ def __init__( self.model_path = model_path self.device_str = infini_device.type self.tp = tp + self.pp = pp self.cache_config = cache_config self.enable_graph = enable_graph self.attn_backend = attn_backend @@ -306,6 +310,7 @@ def run( num_draft_tokens=self.num_draft_tokens, device=self.device_str, tensor_parallel_size=self.tp, + pipeline_parallel_size=self.pp, cache_type="paged" if self.cache_config is not None else "static", max_batch_size=batch_size, max_tokens=output_len, @@ -392,11 +397,12 @@ def run( infini_device = infinicore.device(device_str, 0) tp = cfg.tp + pp = cfg.pp dp = cfg.dp moe_ep_backend, ep = configure_moe_ep_backend( tp, dp, cfg.ep, cfg.moe_ep_backend, model_path ) - print(f"MoE EP backend: {moe_ep_backend} TP={tp} DP={dp} EP={ep}") + print(f"MoE EP backend: {moe_ep_backend} TP={tp} PP={pp} DP={dp} EP={ep}") skip_load = cfg.skip_load @@ -449,6 +455,7 @@ def run( num_draft_tokens=cfg.num_draft_tokens, infini_device=infini_device, tp=tp, + pp=pp, skip_load=skip_load, cache_config=cache_config, enable_graph=enable_graph,