diff --git a/csrc/config/model_config.hpp b/csrc/config/model_config.hpp index dc0e8928..72339160 100644 --- a/csrc/config/model_config.hpp +++ b/csrc/config/model_config.hpp @@ -88,6 +88,18 @@ class ModelConfig { return quant_config.get_quantization_method(); } + std::string get_moe_weight_method(const infinicore::Device &device) const { + return quant_config.get_moe_weight_method(device); + } + + bool is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const { + return quant_config.is_moe_w16a16_marlin_enabled(device); + } + + bool is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const { + return quant_config.is_moe_w8a8_marlin_enabled(device); + } + infinicore::DataType get_dtype() const; infinilm::quantization::QuantScheme get_quant_scheme() const; diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index da261ce0..57ef20dd 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -1,6 +1,65 @@ #include "quant_config.hpp" +#include +#include + namespace infinilm::config { +namespace { + +std::string lower_string(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + return value; +} + +bool is_w16a16_marlin_method(const std::string &method) { + return method == "w16a16_marlin" || method == "hygon_w16a16_marlin"; +} + +bool is_w8a8_marlin_method(const std::string &method) { + return method == "slimquant_marlin" || method == "slimquant_compressed_tensors_marlin" || method == "w8a8_marlin" || method == "hygon_w8a8_marlin"; +} + +bool is_unquantized_config(const nlohmann::json &quantization_config) { + if (quantization_config.is_null()) { + return true; + } + if (!quantization_config.is_object()) { + return false; + } + auto it = quantization_config.find("quant_method"); + if (it == quantization_config.end() || it->is_null()) { + return true; + } + if (!it->is_string()) { + return false; + } + auto method = lower_string(it->get()); + return method.empty() || method == "none" || method == "dense"; +} + +std::string explicit_moe_weight_method(const nlohmann::json &quantization_config) { + if (!quantization_config.is_object()) { + return {}; + } + for (const char *key : {"moe_weight_method", "weight_method", "moe_kernel_method"}) { + auto it = quantization_config.find(key); + if (it != quantization_config.end() && it->is_string()) { + return lower_string(it->get()); + } + } + auto it = quantization_config.find("quant_method"); + if (it != quantization_config.end() && it->is_string()) { + auto method = lower_string(it->get()); + if (is_w16a16_marlin_method(method) || is_w8a8_marlin_method(method)) { + return method; + } + } + return {}; +} + +} // namespace QuantConfig::QuantConfig(const nlohmann::json &json) : quantization_config(json) { this->quantization_method = get_quantization_method(); } @@ -20,6 +79,8 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); } else if (quant_method == "gptq") { return std::make_shared(quantization_config); + } else if (quantization_config["quant_method"] == "w16a16_marlin" || quantization_config["quant_method"] == "hygon_w16a16_marlin") { + return std::make_shared(quantization_config); } else { return std::make_shared(quantization_config); } @@ -27,4 +88,30 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); // Default case if no matching scheme } + +std::string QuantConfig::get_moe_weight_method(const infinicore::Device &device) const { + auto configured_method = explicit_moe_weight_method(quantization_config); + if (!configured_method.empty()) { + return configured_method; + } + if (quantization_method != nullptr) { + auto method = quantization_method->get_moe_weight_method(device); + if (method != "dense") { + return method; + } + } + if (device.getType() == infinicore::Device::Type::HYGON && is_unquantized_config(quantization_config)) { + return "hygon_w16a16_marlin"; + } + return "dense"; +} + +bool QuantConfig::is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const { + return is_w16a16_marlin_method(get_moe_weight_method(device)); +} + +bool QuantConfig::is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const { + return is_w8a8_marlin_method(get_moe_weight_method(device)); +} + } // namespace infinilm::config diff --git a/csrc/config/quant_config.hpp b/csrc/config/quant_config.hpp index fb0b8abf..2b1d7e5f 100644 --- a/csrc/config/quant_config.hpp +++ b/csrc/config/quant_config.hpp @@ -1,9 +1,11 @@ #pragma once -#include "../utils.hpp" #include "../layers/quantization/quantization.hpp" +#include "../utils.hpp" +#include "infinicore/device.hpp" #include "nlohmann/json.hpp" #include #include +#include namespace infinilm::config { @@ -15,6 +17,9 @@ class QuantConfig { QuantConfig(const nlohmann::json &json); std::shared_ptr get_quantization_method() const; + std::string get_moe_weight_method(const infinicore::Device &device) const; + bool is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const; + bool is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const; infinilm::quantization::QuantScheme get_quant_scheme() const { if (quantization_method != nullptr) { diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index dee3123c..06e7d42a 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -71,13 +71,17 @@ void PagedCompiler::compile() { } size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); + decode_graph_needs_runtime_state_reset_ = model_->needs_runtime_state_reset(); compiled_map_decode_.clear(); + // b * ceil(nblocks / b) is at most nblocks + b - 1. All decode + // graphs share this holder and only the selected graph runs at once. block_tables_holder_ = infinicore::Tensor::empty( - {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice()); + {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.last_token_only = true; input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice()); input.position_ids = infinicore::Tensor::empty( position_id_axes > 1 @@ -98,7 +102,9 @@ void PagedCompiler::compile() { 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; + // Give each request its fair share of the global cache capacity. + // Wider runtime tables safely fall back to eager in get_compiled(). + const size_t block_per_req = (nblocks + b - 1) / b; 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()); @@ -150,8 +156,10 @@ void PagedCompiler::compile() { // 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(); + if (decode_graph_needs_runtime_state_reset_) { + model_->reset_runtime_state(); + infinicore::context::syncStream(); + } } for (size_t b : decode_batch_sizes_) { @@ -164,8 +172,10 @@ void PagedCompiler::compile() { // 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(); + if (decode_graph_needs_runtime_state_reset_) { + model_->reset_runtime_state(); + infinicore::context::syncStream(); + } infinicore::context::startGraphRecording(); auto output = model_->forward(input); auto graph = infinicore::context::stopGraphRecording(); @@ -192,20 +202,29 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & 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()); + // Decode graphs are captured with one token per request, so their + // input offsets are the fixed sequence [0, 1, ..., batch_size]. + // Reuse the captured tensor only after validating that the runtime + // input has the same layout; otherwise fall back to eager mode. + const auto &runtime_input_offsets = input.input_offsets.value(); + if (!runtime_input_offsets->is_contiguous() || runtime_input_offsets->size(0) != batch_size + 1) { + 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) { - // Runtime width exceeds compiled graph slot; fall back to eager path. + // Runtime width exceeds compiled graph slot; fall back before + // enqueueing copies that the eager path cannot consume. return {nullptr, nullptr}; } + 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.cu_seqlens.value()->copy_from(input.cu_seqlens.value()); + // Initialize only the active graph rows to -1, then overwrite the // runtime logical region. Avoid clearing the full preallocated // holder on every decode token. @@ -230,7 +249,9 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & // 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(); + if (decode_graph_needs_runtime_state_reset_) { + model_->reset_runtime_state(); + } auto graph = std::get<0>(result->second.compiled); if (graph != nullptr) { diff --git a/csrc/engine/compiler/paged_compiler.hpp b/csrc/engine/compiler/paged_compiler.hpp index a1125864..807f99a5 100644 --- a/csrc/engine/compiler/paged_compiler.hpp +++ b/csrc/engine/compiler/paged_compiler.hpp @@ -18,6 +18,8 @@ class PagedCompiler : public GraphCompiler { infinicore::Tensor block_tables_holder_; + bool decode_graph_needs_runtime_state_reset_ = true; + struct CompiledResult { InfinilmModel::Input input; Compiled compiled; diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index ba01e035..e43aa36b 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -184,6 +184,7 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { image_req_ids, visual_token_ranges, to_device(target_hidden_states)}; + input.last_token_only = !sample_all_positions; infinilm::global_state::get_forward_context().attn_metadata = { input.past_sequence_lengths, diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 871b48d7..acb2d24c 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -485,7 +485,9 @@ void RankWorker::thread_loop() { 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); + score_idx = total_len == n_req + ? i + : 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({})}; diff --git a/csrc/layers/attention/backends/flash_attn.cpp b/csrc/layers/attention/backends/flash_attn.cpp index ec7e3772..d593128e 100644 --- a/csrc/layers/attention/backends/flash_attn.cpp +++ b/csrc/layers/attention/backends/flash_attn.cpp @@ -41,6 +41,23 @@ infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, ASSERT(block_tables.has_value()); ASSERT(slot_mapping.has_value()); + if (query->device().getType() == infinicore::Device::Type::HYGON) { + return infinicore::op::paged_flash_attention( + query, + key, + value, + kv_cache, + total_sequence_lengths.value(), + input_offsets, + cu_seqlens, + block_tables.value(), + slot_mapping.value(), + num_heads_, + num_kv_heads_, + head_dim_, + scale_); + } + // 1. update paged kv cache auto [k_total, v_total] = do_kv_cache_update(layer, key, value, kv_cache, slot_mapping.value()); diff --git a/csrc/layers/causal_lm_templates/text_causal_lm.hpp b/csrc/layers/causal_lm_templates/text_causal_lm.hpp index d359a7f8..5dfbfa4a 100644 --- a/csrc/layers/causal_lm_templates/text_causal_lm.hpp +++ b/csrc/layers/causal_lm_templates/text_causal_lm.hpp @@ -4,6 +4,9 @@ #include "../../models/infinilm_model.hpp" #include "../linear/linear.hpp" #include "infinicore/device.hpp" +#include "infinicore/ops/distributed/allgather.hpp" +#include "infinicore/ops/select_last_token_hidden_states.hpp" + #include namespace infinilm::layers::causal_lm_templates { @@ -39,10 +42,23 @@ class TextCausalLM : public InfinilmModel { const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); pp_size_ = static_cast(rank_info.pp_size); pp_stage_ = static_cast(rank_info.pp_stage); + tp_size_ = static_cast(rank_info.tp_size); + tp_rank_ = static_cast(rank_info.tp_rank); + vocab_parallel_ = device.getType() == infinicore::Device::Type::HYGON + && tp_size_ > 1 + && vocab_size % tp_size_ == 0; model_ = this->register_module("model", model_config, device); if (is_last_pp_stage()) { - lm_head_ = this->register_module("lm_head", hidden_size, vocab_size, false, dtype, device); + lm_head_ = this->register_module( + "lm_head", + hidden_size, + vocab_size, + false, + dtype, + device, + vocab_parallel_ ? tp_rank_ : 0, + vocab_parallel_ ? tp_size_ : 1); } } @@ -54,7 +70,16 @@ class TextCausalLM : public InfinilmModel { if (!is_last_pp_stage()) { return {infinicore::Tensor(), hidden_states}; } - auto logits = lm_head_->forward(hidden_states); + + if (input.last_token_only) { + if (!input.input_offsets.has_value()) { + throw std::runtime_error("TextCausalLM: last_token_only requires input_offsets"); + } + hidden_states = infinicore::op::select_last_token_hidden_states( + hidden_states, input.input_offsets.value()); + } + + auto logits = gather_logits(lm_head_->forward(hidden_states)); return {logits, hidden_states}; } @@ -62,20 +87,48 @@ class TextCausalLM : public InfinilmModel { if (!lm_head_) { throw std::runtime_error("TextCausalLM::logits_from_hidden called on a non-last pipeline stage"); } - return lm_head_->forward(const_cast(hidden_states)); + return gather_logits( + lm_head_->forward(const_cast(hidden_states))); } Model &model() { return *model_; } protected: INFINICORE_NN_MODULE(Model, model); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, lm_head); private: bool is_last_pp_stage() const { return pp_stage_ + 1 == pp_size_; } + infinicore::Tensor gather_logits(const infinicore::Tensor &local_logits) const { + if (!vocab_parallel_) { + return local_logits; + } + + const auto &local_shape = local_logits->shape(); + if (local_shape.empty() || local_shape.back() == 0) { + throw std::runtime_error("TextCausalLM: invalid local logits shape"); + } + const size_t local_vocab_size = local_shape.back(); + const size_t num_rows = local_logits->numel() / local_vocab_size; + auto local_flat = local_logits->view({num_rows, local_vocab_size}); + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + auto gathered = infinicore::op::distributed::allgather( + local_flat, tp_size_, rank_info.comm); + + auto output_shape = local_shape; + output_shape.back() *= tp_size_; + return gathered->view({tp_size_, num_rows, local_vocab_size}) + ->permute({1, 0, 2}) + ->contiguous() + ->view(output_shape); + } + size_t pp_size_{1}; size_t pp_stage_{0}; + size_t tp_size_{1}; + size_t tp_rank_{0}; + bool vocab_parallel_{false}; }; } // namespace infinilm::layers::causal_lm_templates diff --git a/csrc/layers/linear/base_linear.cpp b/csrc/layers/linear/base_linear.cpp index dc4c77f6..79b10ca3 100644 --- a/csrc/layers/linear/base_linear.cpp +++ b/csrc/layers/linear/base_linear.cpp @@ -88,6 +88,10 @@ void BaseLinear::reset_runtime_state() const { quantization_->reset_runtime_state(); } +bool BaseLinear::needs_runtime_state_reset() const { + return quantization_->needs_runtime_state_reset(); +} + // Backward compatible accessors infinicore::Tensor BaseLinear::weight() const { diff --git a/csrc/layers/linear/base_linear.hpp b/csrc/layers/linear/base_linear.hpp index 8b452544..0566adb8 100644 --- a/csrc/layers/linear/base_linear.hpp +++ b/csrc/layers/linear/base_linear.hpp @@ -47,6 +47,7 @@ class BaseLinear : public infinicore::nn::Module { std::shared_ptr get_quantization() const { return quantization_; } void process_weights_after_loading() override; void reset_runtime_state() const override; + bool needs_runtime_state_reset() const; // Split fused linear parameters into named sub-parameters std::vector split_params( diff --git a/csrc/layers/moe/common/moe_types.hpp b/csrc/layers/moe/common/moe_types.hpp index 8b784cf2..84c9dcbd 100644 --- a/csrc/layers/moe/common/moe_types.hpp +++ b/csrc/layers/moe/common/moe_types.hpp @@ -2,6 +2,7 @@ #include "topk_output.hpp" +#include "infinicore/ops/hygon_moe_marlin.hpp" #include "infinicore/tensor.hpp" #include @@ -20,6 +21,12 @@ enum class CombineInputFormat { DeepEPLL, }; +enum class MoeWeightBackend { + Dense, + HygonW16A16Marlin, + HygonW8A8Marlin, +}; + struct DispatchOutput { DispatchOutputFormat format = DispatchOutputFormat::Standard; infinicore::Tensor hidden_states; @@ -53,6 +60,9 @@ struct CombineInput { struct MoeWeights { infinicore::Tensor packed_w13; infinicore::Tensor packed_w2; + infinicore::Tensor packed_w13_scale; + infinicore::Tensor packed_w2_scale; + MoeWeightBackend backend = MoeWeightBackend::Dense; bool empty() const { return !packed_w13 && !packed_w2; @@ -61,6 +71,19 @@ struct MoeWeights { bool has_packed_dense_weights() const { return packed_w13 && packed_w2; } + + bool has_packed_w8a8_marlin_weights() const { + return packed_w13 && packed_w2 + && packed_w13_scale && packed_w2_scale; + } + + bool is_hygon_w16a16_marlin() const { + return backend == MoeWeightBackend::HygonW16A16Marlin; + } + + bool is_hygon_w8a8_marlin() const { + return backend == MoeWeightBackend::HygonW8A8Marlin; + } }; struct MoeWorkspace { @@ -69,6 +92,7 @@ struct MoeWorkspace { infinicore::Tensor ep_gathered_topk_ids; infinicore::Tensor ep_reduced_hidden_states; infinicore::Tensor fused_moe_output; + infinicore::op::HygonMoeMarlinWorkspace hygon_marlin; infinicore::Tensor sorted_token_ids; infinicore::Tensor expert_ids; @@ -84,7 +108,6 @@ struct MoeWorkspace { size_t expert_ids_capacity = 0; size_t ep_gathered_tokens_capacity = 0; size_t ep_reduced_tokens_capacity = 0; - size_t fused_moe_output_tokens_capacity = 0; size_t blockscale_offsets_capacity = 0; size_t permutation_capacity = 0; size_t prepared_num_experts = 0; diff --git a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp index 27e5c0b2..b0a4d015 100644 --- a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp +++ b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp @@ -30,6 +30,10 @@ infinicore::Tensor StandardDispatcher::combine(const CombineInput &combine_input MoeWorkspace &workspace) const { (void)workspace; if (tp_size_ > 1 && communicator_ != nullptr) { + if (combine_input.hidden_states->device().getType() == infinicore::Device::Type::HYGON) { + return infinicore::op::distributed::allreduce( + combine_input.hidden_states, INFINICCL_SUM, communicator_); + } infinicore::op::distributed::allreduce_( combine_input.hidden_states, combine_input.hidden_states, diff --git a/csrc/layers/moe/experts/fused_moe_experts.cpp b/csrc/layers/moe/experts/fused_moe_experts.cpp index b456395d..3a679386 100644 --- a/csrc/layers/moe/experts/fused_moe_experts.cpp +++ b/csrc/layers/moe/experts/fused_moe_experts.cpp @@ -3,16 +3,32 @@ #include "../../../global_state/global_state.hpp" #include "../ep/ep_config.hpp" +#include "infinicore/ops/moe_w16a16_marlin.hpp" +#include "infinicore/ops/moe_w8a8_marlin.hpp" + +#include + +#include #include namespace infinilm::layers::moe { FusedMoeExperts::FusedMoeExperts(std::shared_ptr model_config, const infinicore::Device &device) { + device_ = device; num_experts_ = model_config->get("num_experts"); hidden_size_ = model_config->get("hidden_size"); const size_t intermediate_size = model_config->get("moe_intermediate_size"); const auto dtype = model_config->get_dtype(); + const auto moe_weight_method = model_config->get_moe_weight_method(device); + enable_hygon_w16a16_marlin_ = model_config->is_moe_w16a16_marlin_enabled(device); + enable_hygon_w8a8_marlin_ = model_config->is_moe_w8a8_marlin_enabled(device); + if (enable_hygon_w16a16_marlin_ && enable_hygon_w8a8_marlin_) { + throw std::runtime_error("Only one Hygon MoE Marlin weight method can be enabled"); + } + if (moe_weight_method != "dense" && !enable_hygon_w16a16_marlin_ && !enable_hygon_w8a8_marlin_) { + throw std::runtime_error("Unsupported MoE weight method: " + moe_weight_method); + } ASSERT(num_experts_ > 0); const auto ep_config = make_ep_config(); @@ -32,17 +48,31 @@ FusedMoeExperts::FusedMoeExperts(std::shared_ptr const size_t expert_tp_rank = ep_enabled ? 0 : tp_rank; const size_t expert_tp_size = ep_enabled ? 1 : tp_size; + const auto expert_weight_dtype = enable_hygon_w8a8_marlin_ ? infinicore::DataType::I8 : dtype; w13_weight_ = infinicore::nn::Parameter( {num_local_experts, intermediate_size_per_partition_ * 2, hidden_size_}, - dtype, + expert_weight_dtype, device); w2_weight_ = infinicore::nn::Parameter( {num_local_experts, hidden_size_, intermediate_size_per_partition_}, - dtype, + expert_weight_dtype, device); this->register_parameter("w13_weight", w13_weight_); this->register_parameter("w2_weight", w2_weight_); + if (enable_hygon_w8a8_marlin_) { + w13_weight_scale_ = infinicore::nn::Parameter( + {num_local_experts, intermediate_size_per_partition_ * 2, 1}, + infinicore::DataType::F32, + device); + w2_weight_scale_ = infinicore::nn::Parameter( + {num_local_experts, hidden_size_, 1}, + infinicore::DataType::F32, + device); + this->register_parameter("w13_weight_scale", w13_weight_scale_); + this->register_parameter("w2_weight_scale", w2_weight_scale_); + } + for (size_t local_expert = 0; local_expert < num_local_experts; ++local_expert) { const size_t global_expert = expert_placement.local_expert_start + local_expert; auto gate_weight = w13_weight_ @@ -65,10 +95,117 @@ FusedMoeExperts::FusedMoeExperts(std::shared_ptr this->register_parameter( prefix + "down_proj.weight", infinicore::nn::Parameter(down_weight, 1, expert_tp_rank, expert_tp_size)); + + if (enable_hygon_w8a8_marlin_) { + auto gate_scale = w13_weight_scale_ + ->narrow({{0, local_expert, 1}, {1, 0, intermediate_size_per_partition_}}) + ->squeeze(0); + auto up_scale = w13_weight_scale_ + ->narrow({{0, local_expert, 1}, {1, intermediate_size_per_partition_, intermediate_size_per_partition_}}) + ->squeeze(0); + auto down_scale = w2_weight_scale_ + ->narrow({{0, local_expert, 1}}) + ->squeeze(0); + this->register_parameter( + prefix + "gate_proj.weight_scale", + infinicore::nn::Parameter(gate_scale, 0, expert_tp_rank, expert_tp_size)); + this->register_parameter( + prefix + "up_proj.weight_scale", + infinicore::nn::Parameter(up_scale, 0, expert_tp_rank, expert_tp_size)); + this->register_parameter( + prefix + "down_proj.weight_scale", + infinicore::nn::Parameter(down_scale)); + } } moe_weights_.packed_w13 = w13_weight_; moe_weights_.packed_w2 = w2_weight_; + moe_weights_.backend = MoeWeightBackend::Dense; +} + +void FusedMoeExperts::process_weights_after_loading() { + if (enable_hygon_w8a8_marlin_ && !w8a8_marlin_packed_) { + if (device_.getType() != infinicore::Device::Type::HYGON) { + throw std::runtime_error("slimquant_marlin MoE weight method is only supported on HYGON"); + } + const auto ep_config = make_ep_config(); + if (ep_config.backend != EPBackend::Disabled) { + throw std::runtime_error("slimquant_marlin MoE weight method currently supports TP-split experts only; disable MoE EP"); + } + if (!w13_weight_ || !w2_weight_ || !w13_weight_scale_ || !w2_weight_scale_) { + throw std::runtime_error("slimquant_marlin MoE weight method requires loaded int8 w13/w2 weights and scales"); + } + if (w13_weight_->dtype() != infinicore::DataType::I8 || w2_weight_->dtype() != infinicore::DataType::I8 || w13_weight_scale_->dtype() != infinicore::DataType::F32 || w2_weight_scale_->dtype() != infinicore::DataType::F32) { + throw std::runtime_error("slimquant_marlin MoE weight method requires int8 weights and fp32 weight scales"); + } + if (hidden_size_ % 64 != 0 || intermediate_size_per_partition_ % 64 != 0) { + throw std::runtime_error("slimquant_marlin MoE weight method requires hidden/intermediate sizes divisible by 64"); + } + + spdlog::debug( + "Packing MoE weights with Hygon W8A8 slimquant Marlin layout: experts={}, hidden={}, intermediate_per_partition={}", + w13_weight_->size(0), hidden_size_, intermediate_size_per_partition_); + + auto packed_w13 = infinicore::op::moe_w8a8_marlin_pack(w13_weight_); + auto packed_w2 = infinicore::op::moe_w8a8_marlin_pack(w2_weight_); + + parameters_.clear(); + w13_weight_ = infinicore::nn::Parameter(packed_w13); + w2_weight_ = infinicore::nn::Parameter(packed_w2); + w13_weight_scale_ = infinicore::nn::Parameter(w13_weight_scale_); + w2_weight_scale_ = infinicore::nn::Parameter(w2_weight_scale_); + this->register_parameter("w13_weight", w13_weight_); + this->register_parameter("w2_weight", w2_weight_); + this->register_parameter("w13_weight_scale", w13_weight_scale_); + this->register_parameter("w2_weight_scale", w2_weight_scale_); + + moe_weights_.packed_w13 = w13_weight_; + moe_weights_.packed_w2 = w2_weight_; + moe_weights_.packed_w13_scale = w13_weight_scale_; + moe_weights_.packed_w2_scale = w2_weight_scale_; + moe_weights_.backend = MoeWeightBackend::HygonW8A8Marlin; + w8a8_marlin_packed_ = true; + return; + } + + if (!enable_hygon_w16a16_marlin_ || w16a16_marlin_packed_) { + return; + } + if (device_.getType() != infinicore::Device::Type::HYGON) { + throw std::runtime_error("w16a16_marlin MoE weight method is only supported on HYGON"); + } + + const auto ep_config = make_ep_config(); + if (ep_config.backend != EPBackend::Disabled) { + throw std::runtime_error("w16a16_marlin MoE weight method currently supports TP-split experts only; disable MoE EP"); + } + if (!w13_weight_ || !w2_weight_) { + throw std::runtime_error("w16a16_marlin MoE weight method requires loaded dense w13/w2 weights"); + } + if (w13_weight_->dtype() != infinicore::DataType::F16 && w13_weight_->dtype() != infinicore::DataType::BF16) { + throw std::runtime_error("w16a16_marlin MoE weight method requires FP16 or BF16 weights"); + } + if (hidden_size_ % 32 != 0 || intermediate_size_per_partition_ % 16 != 0 || (intermediate_size_per_partition_ * 2) % 32 != 0) { + throw std::runtime_error("w16a16_marlin MoE weight method requires aligned hidden/intermediate sizes"); + } + + spdlog::debug( + "Packing MoE weights with Hygon W16A16 Marlin layout: experts={}, hidden={}, intermediate_per_partition={}", + w13_weight_->size(0), hidden_size_, intermediate_size_per_partition_); + + auto packed_w13 = infinicore::op::moe_w16a16_marlin_pack(w13_weight_); + auto packed_w2 = infinicore::op::moe_w16a16_marlin_pack(w2_weight_); + + parameters_.clear(); + w13_weight_ = infinicore::nn::Parameter(packed_w13); + w2_weight_ = infinicore::nn::Parameter(packed_w2); + this->register_parameter("w13_weight", w13_weight_); + this->register_parameter("w2_weight", w2_weight_); + + moe_weights_.packed_w13 = w13_weight_; + moe_weights_.packed_w2 = w2_weight_; + moe_weights_.backend = MoeWeightBackend::HygonW16A16Marlin; + w16a16_marlin_packed_ = true; } const MoeWeights &FusedMoeExperts::moe_weights() const { diff --git a/csrc/layers/moe/experts/fused_moe_experts.hpp b/csrc/layers/moe/experts/fused_moe_experts.hpp index 3f231ec6..a8f2a53b 100644 --- a/csrc/layers/moe/experts/fused_moe_experts.hpp +++ b/csrc/layers/moe/experts/fused_moe_experts.hpp @@ -17,13 +17,21 @@ class FusedMoeExperts : public infinicore::nn::Module { const MoeWeights &moe_weights() const; + void process_weights_after_loading() override; + protected: INFINICORE_NN_PARAMETER(w13_weight); INFINICORE_NN_PARAMETER(w2_weight); + INFINICORE_NN_PARAMETER(w13_weight_scale); + INFINICORE_NN_PARAMETER(w2_weight_scale); size_t num_experts_{0}; size_t hidden_size_{0}; size_t intermediate_size_per_partition_{0}; + bool enable_hygon_w16a16_marlin_{false}; + bool enable_hygon_w8a8_marlin_{false}; + bool w16a16_marlin_packed_{false}; + bool w8a8_marlin_packed_{false}; MoeWeights moe_weights_; }; diff --git a/csrc/layers/moe/fused_moe.cpp b/csrc/layers/moe/fused_moe.cpp index fe1301a8..8defcbc8 100644 --- a/csrc/layers/moe/fused_moe.cpp +++ b/csrc/layers/moe/fused_moe.cpp @@ -10,6 +10,30 @@ namespace infinilm::layers::moe { +namespace { + +std::shared_ptr make_workspace( + const EPConfig &ep_config, + const std::shared_ptr &model_config, + const infinicore::Device &device) { + const bool use_hygon_marlin = device.getType() == infinicore::Device::Type::HYGON && ep_config.backend == EPBackend::Disabled && (model_config->is_moe_w8a8_marlin_enabled(device) || model_config->is_moe_w16a16_marlin_enabled(device)); + if (!use_hygon_marlin) { + return std::make_shared(); + } + + // Decoder layers execute sequentially on each rank, including graph replay. + // Reuse their large Marlin scratch buffers instead of retaining one copy per layer. + static thread_local std::weak_ptr shared_workspace; + auto workspace = shared_workspace.lock(); + if (!workspace) { + workspace = std::make_shared(); + shared_workspace = workspace; + } + return workspace; +} + +} // namespace + FusedMoE::FusedMoE(std::shared_ptr model_config, const infinicore::Device &device, size_t layer_id) { @@ -29,21 +53,22 @@ FusedMoE::FusedMoE(std::shared_ptr model_config, intermediate_size_per_partition = intermediate_size / tp_size; } + workspace_ = make_workspace(ep_config, model_config, device); dispatcher_ = make_dispatcher(ep_config, num_experts); runner_ = std::make_shared( expert_placement.local_num_experts, hidden_size, intermediate_size_per_partition, model_config->get_or("moe_align_block_size", 16)); - dispatcher_->initialize(device, workspace_); + dispatcher_->initialize(device, *workspace_); } infinicore::Tensor FusedMoE::forward(const infinicore::Tensor &hidden_states, const TopKOutput &topk_output, const MoeWeights &weights) const { - auto dispatch_output = dispatcher_->dispatch(hidden_states, topk_output, workspace_); - auto combine_input = runner_->run(dispatch_output, weights, workspace_); - return dispatcher_->combine(combine_input, workspace_); + auto dispatch_output = dispatcher_->dispatch(hidden_states, topk_output, *workspace_); + auto combine_input = runner_->run(dispatch_output, weights, *workspace_); + return dispatcher_->combine(combine_input, *workspace_); } } // namespace infinilm::layers::moe diff --git a/csrc/layers/moe/fused_moe.hpp b/csrc/layers/moe/fused_moe.hpp index 81b3cad7..7301a5d6 100644 --- a/csrc/layers/moe/fused_moe.hpp +++ b/csrc/layers/moe/fused_moe.hpp @@ -24,7 +24,7 @@ class FusedMoE final : public infinicore::nn::Module { private: std::shared_ptr dispatcher_; std::shared_ptr runner_; - mutable MoeWorkspace workspace_; + std::shared_ptr workspace_; }; } // namespace infinilm::layers::moe diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index c37abcc8..e0561b11 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -1,38 +1,48 @@ #include "cuda_fused_moe_runner.hpp" #include "infinicore/context/context.hpp" +#include "infinicore/ops/hygon_moe_marlin.hpp" #include "infinicore/ops/moe_align.hpp" #include "infinicore/ops/moe_fused_dense.hpp" #include #include #include -#include namespace infinilm::layers::moe { -CudaFusedMoeRunner::CudaFusedMoeRunner(size_t num_local_experts, - size_t hidden_size, - size_t intermediate_size_per_partition, - size_t align_block_size) +CudaFusedMoeRunner::CudaFusedMoeRunner( + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size_per_partition, + size_t align_block_size) : num_local_experts_(num_local_experts), hidden_size_(hidden_size), - intermediate_size_per_partition_(intermediate_size_per_partition), + intermediate_size_per_partition_( + intermediate_size_per_partition), align_block_size_(align_block_size) {} namespace { -bool same_device(const infinicore::Tensor &tensor, const infinicore::Device &device) { - return tensor && tensor->device().getType() == device.getType() && tensor->device().getIndex() == device.getIndex(); +bool same_device( + const infinicore::Tensor &tensor, + const infinicore::Device &device) { + return tensor + && tensor->device().getType() == device.getType() + && tensor->device().getIndex() == device.getIndex(); } -void ensure_tensor(infinicore::Tensor &tensor, - const infinicore::Shape &shape, - infinicore::DataType dtype, - const infinicore::Device &device) { - if (!same_device(tensor, device) || tensor->dtype() != dtype || tensor->shape() != shape) { +void ensure_tensor( + infinicore::Tensor &tensor, + const infinicore::Shape &shape, + infinicore::DataType dtype, + const infinicore::Device &device) { + if (!same_device(tensor, device) + || tensor->dtype() != dtype + || tensor->shape() != shape) { if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE runner workspace tensor was not initialized before graph capture"); + throw std::runtime_error( + "MoE runner workspace tensor was not initialized before graph capture"); } tensor = infinicore::Tensor::empty(shape, dtype, device); } @@ -51,37 +61,83 @@ std::string shape_to_string(const infinicore::Shape &shape) { return oss.str(); } -void check_packed_weight_tensor(const infinicore::Tensor &tensor, - const std::string &name, - const infinicore::Device &device, - const infinicore::DataType dtype, - const infinicore::Shape &shape) { +void check_packed_weight_tensor( + const infinicore::Tensor &tensor, + const std::string &name, + const infinicore::Device &device, + infinicore::DataType dtype, + const infinicore::Shape &shape) { if (!tensor) { - throw std::runtime_error("MoE fused dense core requires " + name); + throw std::runtime_error( + "MoE fused dense core requires " + name); } - if (tensor->device().getType() != device.getType() || tensor->device().getIndex() != device.getIndex()) { - throw std::runtime_error("MoE fused dense core requires packed weights on the hidden_states device"); + if (tensor->device().getType() != device.getType() + || tensor->device().getIndex() != device.getIndex()) { + throw std::runtime_error( + "MoE fused dense core requires packed weights on the hidden_states device"); } if (tensor->dtype() != dtype) { - throw std::runtime_error("MoE fused dense core requires packed weights to have the same dtype as hidden_states"); + throw std::runtime_error( + "MoE fused dense core packed tensor dtype mismatch for " + + name); } if (tensor->shape() != shape) { throw std::runtime_error( - "MoE fused dense core packed weight shape mismatch for " + name + ": expected " + shape_to_string(shape) + ", got " + shape_to_string(tensor->shape())); + "MoE fused dense core packed weight shape mismatch for " + + name + ": expected " + shape_to_string(shape) + + ", got " + shape_to_string(tensor->shape())); } } } // namespace -CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, - const MoeWeights &weights, - MoeWorkspace &workspace) const { +CombineInput CudaFusedMoeRunner::run( + const DispatchOutput &dispatch_output, + const MoeWeights &weights, + MoeWorkspace &workspace) const { + if (weights.is_hygon_w16a16_marlin() + || weights.is_hygon_w8a8_marlin()) { + const auto format = weights.is_hygon_w16a16_marlin() + ? infinicore::op::HygonMoeMarlinWeightFormat::W16A16 + : infinicore::op::HygonMoeMarlinWeightFormat::W8A8; + const infinicore::op::HygonMoeMarlinWeights marlin_weights{ + weights.packed_w13, + weights.packed_w2, + weights.packed_w13_scale, + weights.packed_w2_scale, + format, + }; + const auto output = infinicore::op::hygon_moe_marlin_fused( + dispatch_output.hidden_states, + dispatch_output.topk_output.topk_weights, + dispatch_output.topk_output.topk_ids, + dispatch_output.expert_map, + marlin_weights, + workspace.hygon_marlin, + num_local_experts_, + hidden_size_, + intermediate_size_per_partition_, + align_block_size_); + + MoeRoutingMetadata routing_metadata; + if (output.has_routing_metadata) { + routing_metadata.sorted_token_ids = output.sorted_token_ids; + routing_metadata.expert_ids = output.expert_ids; + routing_metadata.num_tokens_post_padded = output.num_tokens_post_padded; + } + return CombineInput{ + CombineInputFormat::Standard, + output.hidden_states, + dispatch_output.topk_output, + routing_metadata, + }; + } + auto runner_input = prepare_runner_input( dispatch_output, - workspace); - + workspace, + align_block_size_); auto runner_output = run_fused_core(runner_input, weights, workspace); - return CombineInput{ CombineInputFormat::Standard, runner_output.hidden_states, @@ -90,51 +146,70 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, }; } -CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchOutput &dispatch_output, - MoeWorkspace &workspace) const { +CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input( + const DispatchOutput &dispatch_output, + MoeWorkspace &workspace, + size_t block_size) const { const auto &topk_ids = dispatch_output.topk_output.topk_ids; const auto &topk_shape = topk_ids->shape(); if (topk_shape.size() != 2) { - throw std::runtime_error("MoE runner requires topk_ids to be a 2D tensor"); + throw std::runtime_error( + "MoE runner requires topk_ids to be a 2D tensor"); } const size_t num_pairs = topk_shape[0] * topk_shape[1]; - const size_t block_size = align_block_size_; const size_t align_num_experts = num_local_experts_ + 1; const size_t max_num_tokens_padded = num_pairs < align_num_experts ? num_pairs * block_size - : num_pairs + align_num_experts * (block_size - 1); + : num_pairs + + align_num_experts * (block_size - 1); const size_t sorted_token_ids_capacity = ((max_num_tokens_padded + 3) / 4) * 4; const size_t max_num_blocks = (max_num_tokens_padded + block_size - 1) / block_size; const auto device = topk_ids->device(); - if (!workspace.sorted_token_ids || workspace.sorted_token_ids_capacity < sorted_token_ids_capacity) { + if (!workspace.sorted_token_ids + || workspace.sorted_token_ids_capacity + < sorted_token_ids_capacity) { if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE sorted_token_ids workspace was not initialized before graph capture"); + throw std::runtime_error( + "MoE sorted_token_ids workspace was not initialized before graph capture"); } workspace.sorted_token_ids = infinicore::Tensor::empty( - {sorted_token_ids_capacity}, infinicore::DataType::I32, device); + {sorted_token_ids_capacity}, + infinicore::DataType::I32, + device); workspace.sorted_token_ids_capacity = sorted_token_ids_capacity; } - if (!workspace.expert_ids || workspace.expert_ids_capacity < max_num_blocks) { + if (!workspace.expert_ids + || workspace.expert_ids_capacity < max_num_blocks) { if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE expert_ids workspace was not initialized before graph capture"); + throw std::runtime_error( + "MoE expert_ids workspace was not initialized before graph capture"); } workspace.expert_ids = infinicore::Tensor::empty( - {max_num_blocks}, infinicore::DataType::I32, device); + {max_num_blocks}, + infinicore::DataType::I32, + device); workspace.expert_ids_capacity = max_num_blocks; } if (!workspace.num_tokens_post_padded) { if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE num_tokens_post_padded workspace was not initialized before graph capture"); + throw std::runtime_error( + "MoE num_tokens_post_padded workspace was not initialized before graph capture"); } workspace.num_tokens_post_padded = infinicore::Tensor::empty( - {1}, infinicore::DataType::I32, device); + {1}, + infinicore::DataType::I32, + device); } + auto sorted_token_ids = workspace.sorted_token_ids->narrow( + {{0, 0, sorted_token_ids_capacity}}); + auto expert_ids = workspace.expert_ids->narrow( + {{0, 0, max_num_blocks}}); if (dispatch_output.expert_map) { infinicore::op::moe_align_with_expert_map_( - workspace.sorted_token_ids, - workspace.expert_ids, + sorted_token_ids, + expert_ids, workspace.num_tokens_post_padded, topk_ids, dispatch_output.expert_map, @@ -143,8 +218,8 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchO true); } else { infinicore::op::moe_align_( - workspace.sorted_token_ids, - workspace.expert_ids, + sorted_token_ids, + expert_ids, workspace.num_tokens_post_padded, topk_ids, num_local_experts_, @@ -155,37 +230,42 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchO dispatch_output.hidden_states, dispatch_output.topk_output, MoeRoutingMetadata{ - workspace.sorted_token_ids, - workspace.expert_ids, + sorted_token_ids, + expert_ids, workspace.num_tokens_post_padded, }, }; } -CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_fused_core(const CudaFusedMoeRunnerInput &runner_input, - const MoeWeights &weights, - MoeWorkspace &workspace) const { +CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_fused_core( + const CudaFusedMoeRunnerInput &runner_input, + const MoeWeights &weights, + MoeWorkspace &workspace) const { if (!weights.has_packed_dense_weights()) { - throw std::runtime_error("MoE fused dense runner requires load-time packed w13/w2 weights"); + throw std::runtime_error( + "MoE fused dense runner requires load-time packed w13/w2 weights"); } check_packed_weight_tensor( weights.packed_w13, "w13", runner_input.hidden_states->device(), runner_input.hidden_states->dtype(), - {num_local_experts_, intermediate_size_per_partition_ * 2, hidden_size_}); + {num_local_experts_, + intermediate_size_per_partition_ * 2, + hidden_size_}); check_packed_weight_tensor( weights.packed_w2, "w2", runner_input.hidden_states->device(), runner_input.hidden_states->dtype(), - {num_local_experts_, hidden_size_, intermediate_size_per_partition_}); + {num_local_experts_, + hidden_size_, + intermediate_size_per_partition_}); ensure_tensor( workspace.fused_moe_output, runner_input.hidden_states->shape(), runner_input.hidden_states->dtype(), runner_input.hidden_states->device()); - workspace.fused_moe_output_tokens_capacity = runner_input.hidden_states->shape()[0]; infinicore::op::moe_fused_dense_( workspace.fused_moe_output, runner_input.hidden_states, diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp index 2a1a1f94..be557fa0 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp @@ -16,22 +16,27 @@ struct CudaFusedMoeRunnerOutput { class CudaFusedMoeRunner final : public MoeRunnerCore { public: - CudaFusedMoeRunner(size_t num_local_experts, - size_t hidden_size, - size_t intermediate_size_per_partition, - size_t align_block_size); + CudaFusedMoeRunner( + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size_per_partition, + size_t align_block_size); - CombineInput run(const DispatchOutput &dispatch_output, - const MoeWeights &weights, - MoeWorkspace &workspace) const override; + CombineInput run( + const DispatchOutput &dispatch_output, + const MoeWeights &weights, + MoeWorkspace &workspace) const override; private: - CudaFusedMoeRunnerInput prepare_runner_input(const DispatchOutput &dispatch_output, - MoeWorkspace &workspace) const; - - CudaFusedMoeRunnerOutput run_fused_core(const CudaFusedMoeRunnerInput &runner_input, - const MoeWeights &weights, - MoeWorkspace &workspace) const; + CudaFusedMoeRunnerInput prepare_runner_input( + const DispatchOutput &dispatch_output, + MoeWorkspace &workspace, + size_t block_size) const; + + CudaFusedMoeRunnerOutput run_fused_core( + const CudaFusedMoeRunnerInput &runner_input, + const MoeWeights &weights, + MoeWorkspace &workspace) const; size_t num_local_experts_ = 0; size_t hidden_size_ = 0; diff --git a/csrc/layers/quantization/awq_marlin.hpp b/csrc/layers/quantization/awq_marlin.hpp index a22fb43b..5b6bc942 100644 --- a/csrc/layers/quantization/awq_marlin.hpp +++ b/csrc/layers/quantization/awq_marlin.hpp @@ -29,6 +29,7 @@ class AWQMarlin : public BaseQuantization { int tp_rank, int tp_size, int tp_num_heads) const override; void reset_runtime_state() const override; + bool needs_runtime_state_reset() const override { return true; } private: infinicore::Tensor get_workspace( diff --git a/csrc/layers/quantization/base_quantization.hpp b/csrc/layers/quantization/base_quantization.hpp index 4b17cc94..6142a4a3 100644 --- a/csrc/layers/quantization/base_quantization.hpp +++ b/csrc/layers/quantization/base_quantization.hpp @@ -79,6 +79,14 @@ class BaseQuantization : public std::enable_shared_from_this { // Default: raw size is already logical size. virtual size_t get_logical_dim_size(size_t raw_size) const { return raw_size; } + // Optional MoE weight backend selected from the same quantization config as Linear. + // Backends can specialize by device while keeping model code independent from + // vendor-specific kernel names. + virtual std::string get_moe_weight_method(const infinicore::Device &device) const { + (void)device; + return "dense"; + } + // Split fused linear parameters into named sub-parameters (for QKV/GateUp) // params: the fused linear's registered parameters (by name) // splits: description of each shard @@ -111,6 +119,8 @@ class BaseQuantization : public std::enable_shared_from_this { // runtime state can keep the default no-op implementation. virtual void reset_runtime_state() const {} + virtual bool needs_runtime_state_reset() const { return false; } + template T get(const std::string &key) const { if (!quant_config_.contains(key)) { diff --git a/csrc/layers/quantization/compressed_tensors.cpp b/csrc/layers/quantization/compressed_tensors.cpp index ff5617a1..9428875e 100644 --- a/csrc/layers/quantization/compressed_tensors.cpp +++ b/csrc/layers/quantization/compressed_tensors.cpp @@ -2,10 +2,68 @@ #include "infinicore/ops/linear_w8a8i8.hpp" #include "infinicore/ops/mul_scalar.hpp" +#include +#include #include #include +#include namespace infinilm::quantization { +namespace { + +std::string lower_string(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + return value; +} + +bool string_field_equals(const nlohmann::json &json, const char *key, const char *expected) { + auto it = json.find(key); + return it != json.end() && it->is_string() && lower_string(it->get()) == expected; +} + +bool bool_field_equals(const nlohmann::json &json, const char *key, bool expected) { + auto it = json.find(key); + return it != json.end() && it->is_boolean() && it->get() == expected; +} + +bool integer_field_equals(const nlohmann::json &json, const char *key, int expected) { + auto it = json.find(key); + return it != json.end() && it->is_number_integer() && it->get() == expected; +} + +bool has_linear_or_moe_target(const nlohmann::json &group) { + auto targets = group.find("targets"); + if (targets == group.end() || !targets->is_array()) { + return false; + } + for (const auto &target : *targets) { + if (!target.is_string()) { + continue; + } + const auto value = lower_string(target.get()); + if (value == "linear" || value == "fusedmoe") { + return true; + } + } + return false; +} + +bool is_dynamic_token_w8a8_group(const nlohmann::json &group) { + auto weights_it = group.find("weights"); + auto input_it = group.find("input_activations"); + if (weights_it == group.end() || input_it == group.end() || !weights_it->is_object() || !input_it->is_object()) { + return false; + } + const auto &weights = *weights_it; + const auto &input = *input_it; + const bool weight_ok = string_field_equals(weights, "type", "int") && string_field_equals(weights, "strategy", "channel") && integer_field_equals(weights, "num_bits", 8) && bool_field_equals(weights, "symmetric", true); + const bool input_ok = string_field_equals(input, "type", "int") && string_field_equals(input, "strategy", "token") && integer_field_equals(input, "num_bits", 8) && bool_field_equals(input, "dynamic", true) && bool_field_equals(input, "symmetric", true); + return weight_ok && input_ok; +} + +} // namespace std::vector CompressedTensors::get_param_layout( size_t in_features, size_t out_features, @@ -28,6 +86,26 @@ std::vector CompressedTensors::get_param_layout( return descs; } +std::string CompressedTensors::get_moe_weight_method(const infinicore::Device &device) const { + if (device.getType() != infinicore::Device::Type::HYGON || !quant_config_.is_object()) { + return "dense"; + } + if (!string_field_equals(quant_config_, "quant_method", "compressed-tensors")) { + return "dense"; + } + auto groups = quant_config_.find("config_groups"); + if (groups == quant_config_.end() || !groups->is_object()) { + return "dense"; + } + for (const auto &item : groups->items()) { + const auto &group = item.value(); + if (group.is_object() && has_linear_or_moe_target(group) && is_dynamic_token_w8a8_group(group)) { + return "slimquant_marlin"; + } + } + return "dense"; +} + infinicore::Tensor CompressedTensors::forward( const ParamsMap ¶ms, const infinicore::Tensor &input, diff --git a/csrc/layers/quantization/compressed_tensors.hpp b/csrc/layers/quantization/compressed_tensors.hpp index dcf65c2e..15438da9 100644 --- a/csrc/layers/quantization/compressed_tensors.hpp +++ b/csrc/layers/quantization/compressed_tensors.hpp @@ -6,7 +6,7 @@ namespace infinilm::quantization { class CompressedTensors : public BaseQuantization { public: explicit CompressedTensors(const nlohmann::json &quant_config) - : BaseQuantization(quant_config) {}; + : BaseQuantization(quant_config){}; QuantScheme get_quant_scheme() const override { return QuantScheme::COMPRESSED_TENSOR_W8A8I8; @@ -25,6 +25,8 @@ class CompressedTensors : public BaseQuantization { bool has_bias, float alpha = 1.0f) const override; + std::string get_moe_weight_method(const infinicore::Device &device) const override; + std::vector split_params( const std::unordered_map ¶ms, const std::vector &splits, diff --git a/csrc/layers/quantization/gptq_marlin.hpp b/csrc/layers/quantization/gptq_marlin.hpp index c16c7943..1fc345b9 100644 --- a/csrc/layers/quantization/gptq_marlin.hpp +++ b/csrc/layers/quantization/gptq_marlin.hpp @@ -31,6 +31,7 @@ class GPTQMarlin : public BaseQuantization { int tp_rank, int tp_size, int tp_num_heads) const override; void reset_runtime_state() const override; + bool needs_runtime_state_reset() const override { return true; } private: infinicore::Tensor get_workspace( diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp index 5d284a31..2c6a652d 100644 --- a/csrc/models/infinilm_model.cpp +++ b/csrc/models/infinilm_model.cpp @@ -1,6 +1,7 @@ #include "infinilm_model.hpp" #include "../cache/kv_cache.hpp" #include "../global_state/global_state.hpp" +#include "../layers/linear/base_linear.hpp" #include "../utils.hpp" #include @@ -104,6 +105,15 @@ void InfinilmModel::reset_runtime_state() const { } } +bool InfinilmModel::needs_runtime_state_reset() const { + for (const auto &[_, sub] : children()) { + if (needs_runtime_state_reset_recursive_(sub.get())) { + return true; + } + } + return false; +} + void InfinilmModel::process_weights_recursive_(infinicore::nn::Module *module) { for (const auto &[_, sub] : module->children()) { process_weights_recursive_(sub.get()); @@ -118,4 +128,18 @@ void InfinilmModel::reset_runtime_state_recursive_(const infinicore::nn::Module module->reset_runtime_state(); } +bool InfinilmModel::needs_runtime_state_reset_recursive_(const infinicore::nn::Module *module) { + if (const auto *linear = dynamic_cast(module)) { + if (linear->needs_runtime_state_reset()) { + return true; + } + } + for (const auto &[_, sub] : module->children()) { + if (needs_runtime_state_reset_recursive_(sub.get())) { + return true; + } + } + return false; +} + } // namespace infinilm diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index 27275ee2..6b7d6f34 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -55,6 +55,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; + /// Return one logit row per request instead of logits for every input token. + bool last_token_only{false}; }; struct Output { @@ -77,6 +79,8 @@ class InfinilmModel : public infinicore::nn::Module { void process_weights_after_loading(); void reset_runtime_state() const; + bool needs_runtime_state_reset() const; + protected: std::vector default_allocate_kv_cache_tensors( const cache::CacheConfig *cache_config, @@ -89,5 +93,7 @@ class InfinilmModel : public infinicore::nn::Module { private: static void process_weights_recursive_(infinicore::nn::Module *module); static void reset_runtime_state_recursive_(const infinicore::nn::Module *module); + + static bool needs_runtime_state_reset_recursive_(const infinicore::nn::Module *module); }; } // namespace infinilm diff --git a/csrc/models/qwen3/qwen3_attention.cpp b/csrc/models/qwen3/qwen3_attention.cpp index 7d9beb04..c7b794ae 100644 --- a/csrc/models/qwen3/qwen3_attention.cpp +++ b/csrc/models/qwen3/qwen3_attention.cpp @@ -2,6 +2,7 @@ #include "../../global_state/global_state.hpp" #include "../../layers/attention/attention.hpp" #include "../../utils.hpp" +#include "infinicore/ops/rms_rotary_embedding.hpp" namespace infinilm::models::qwen3 { @@ -121,8 +122,6 @@ infinicore::Tensor Qwen3Attention::forward_paged_(const infinicore::Tensor &posi auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_}); auto k_reshaped = k->view({seq_len, num_key_value_heads_, head_dim_}); auto v_reshaped = v->view({seq_len, num_key_value_heads_, head_dim_}); - q_reshaped = q_norm_->forward(q_reshaped); - k_reshaped = k_norm_->forward(k_reshaped); // 3. Prepare position_ids for RoPE auto pos_shape = position_ids->shape(); @@ -136,9 +135,31 @@ infinicore::Tensor Qwen3Attention::forward_paged_(const infinicore::Tensor &posi throw std::runtime_error("Unexpected position_ids shape"); } - // 4. Apply RoPE to QK - rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); - rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + // 4. Apply Q/K RMSNorm and RoPE. + const bool can_use_hygon_fused_rms_rope = qkv_proj_->get_quantization()->get_quant_scheme() == infinilm::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8 + && q_reshaped->device().getType() == infinicore::Device::Type::HYGON + && rotary_emb_->rotary_dim() == head_dim_ + && !rotary_emb_->mrope_section().has_value() + && infinicore::op::rms_rotary_embedding_fuse_available(q_reshaped->device()); + if (can_use_hygon_fused_rms_rope) { + q_reshaped = q_reshaped->contiguous(); + k_reshaped = k_reshaped->contiguous(); + auto pos_ids_fused = pos_ids_for_rope->is_contiguous() ? pos_ids_for_rope : pos_ids_for_rope->contiguous(); + infinicore::op::rms_rotary_embedding_fuse_(q_reshaped, + k_reshaped, + pos_ids_fused, + static_cast(head_dim_), + rotary_emb_->cos_sin_cache(), + rotary_emb_->algo() == infinicore::nn::RoPE::Algo::GPT_NEOX, + q_norm_->weight(), + k_norm_->weight(), + static_cast(q_norm_->eps())); + } else { + q_reshaped = q_norm_->forward(q_reshaped); + k_reshaped = k_norm_->forward(k_reshaped); + rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + } // 5. Attn Backend calculate auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 4ae7665c..bbcd0233 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -79,7 +79,15 @@ def __init__(self): self.use_mla = self.args.use_mla self.pre_transpose = self.args.pre_transpose self.num_blocks = self.args.num_blocks - self.block_size = self.args.block_size + if self.args.block_size is None: + platform = ( + self.detect_device() + if self.device.lower() == "auto" + else self.device.lower() + ) + self.block_size = 64 if platform == "hygon" else 256 + else: + self.block_size = self.args.block_size self.max_cache_len = self.args.max_cache_len self.kv_cache_dtype = self.args.kv_cache_dtype self.skip_load = self.args.skip_load @@ -278,7 +286,10 @@ def _add_common_args(self): "--num-blocks", type=int, default=512, help="number of KV cache blocks" ) self.parser.add_argument( - "--block-size", type=int, default=256, help="size of each KV cache block" + "--block-size", + type=int, + default=None, + help="size of each KV cache block (default: 64 on Hygon, 256 otherwise)", ) self.parser.add_argument( "--max-cache-len", type=int, default=4096, help="maximum cache length" diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 486df03b..70a9b21f 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -58,13 +58,17 @@ def parse_dtype(dtype_str: str): def _is_internal_moe_packed_weight(key: str) -> bool: # InfiniLM registers packed MoE parameters internally. HF checkpoints - # provide per-expert gate/up/down weights instead, so these packed tensors - # are expected missing keys during non-strict checkpoint loading. - return ( - key.endswith(".mlp.experts.w13_weight") - or key.endswith(".mlp.experts.w2_weight") - or key.endswith(".mlp.experts.w1") - or key.endswith(".mlp.experts.w2") + # provide per-expert gate/up/down weights and scales, so these internal + # packed tensors are expected missing keys during non-strict loading. + return key.endswith( + ( + ".mlp.experts.w13_weight", + ".mlp.experts.w2_weight", + ".mlp.experts.w1", + ".mlp.experts.w2", + ".mlp.experts.w13_weight_scale", + ".mlp.experts.w2_weight_scale", + ) ) @@ -268,7 +272,10 @@ def load_model_state_dict_by_file( # --------------------------------------------------------- # model_param_infini = {} for key in model_param.keys(): - model_param_infini[key] = infinicore.from_torch(model_param[key]) + tensor = model_param[key] + if key.endswith(".weight_scale") and tensor.dtype != torch.float32: + tensor = tensor.to(torch.float32) + model_param_infini[key] = infinicore.from_torch(tensor) model.load_state_dict(model_param_infini, strict=False) infinicore.sync_device() del model_param_infini