From de044c2b5d085ee9609e10d2196a639c0913ea32 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Thu, 9 Jul 2026 09:26:09 +0800 Subject: [PATCH 01/14] feat(hygon): wire Marlin MoE inference path --- csrc/config/model_config.hpp | 8 + csrc/config/quant_config.cpp | 62 ++++++ csrc/config/quant_config.hpp | 3 + csrc/layers/moe/common/moe_types.hpp | 14 ++ csrc/layers/moe/experts/fused_moe_experts.cpp | 51 +++++ csrc/layers/moe/experts/fused_moe_experts.hpp | 4 + .../moe/runner/cuda_fused_moe_runner.cpp | 202 +++++++++++++++++- .../moe/runner/cuda_fused_moe_runner.hpp | 11 +- 8 files changed, 350 insertions(+), 5 deletions(-) diff --git a/csrc/config/model_config.hpp b/csrc/config/model_config.hpp index dc0e89287..8c239ca1e 100644 --- a/csrc/config/model_config.hpp +++ b/csrc/config/model_config.hpp @@ -88,6 +88,14 @@ class ModelConfig { return quant_config.get_quantization_method(); } + std::string get_moe_weight_method() const { + return quant_config.get_moe_weight_method(); + } + + bool is_moe_w16a16_marlin_enabled() const { + return quant_config.is_moe_w16a16_marlin_enabled(); + } + 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 da261ce09..801d0fdc9 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -1,6 +1,33 @@ #include "quant_config.hpp" +#include +#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; +} + +std::string env_string(const char *name) { + const char *value = std::getenv(name); + if (value == nullptr || value[0] == '\0') { + return {}; + } + return lower_string(value); +} + +bool truthy_env(const char *name) { + auto value = env_string(name); + return value == "1" || value == "true" || value == "on" || value == "yes"; +} + +} // namespace QuantConfig::QuantConfig(const nlohmann::json &json) : quantization_config(json) { this->quantization_method = get_quantization_method(); } @@ -20,6 +47,9 @@ 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 +57,36 @@ 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 { + auto env_method = env_string("INFINILM_MOE_WEIGHT_METHOD"); + if (!env_method.empty()) { + return env_method; + } + if (truthy_env("INFINILM_HYGON_MOE_W16A16_MARLIN")) { + return "w16a16_marlin"; + } + if (quantization_config.is_object()) { + 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 (method == "w16a16_marlin" || method == "hygon_w16a16_marlin") { + return "w16a16_marlin"; + } + } + } + return "dense"; +} + +bool QuantConfig::is_moe_w16a16_marlin_enabled() const { + auto method = get_moe_weight_method(); + return method == "w16a16_marlin" || method == "hygon_w16a16_marlin"; +} + } // namespace infinilm::config diff --git a/csrc/config/quant_config.hpp b/csrc/config/quant_config.hpp index fb0b8abf3..fa1d08513 100644 --- a/csrc/config/quant_config.hpp +++ b/csrc/config/quant_config.hpp @@ -3,6 +3,7 @@ #include "../layers/quantization/quantization.hpp" #include "nlohmann/json.hpp" #include +#include #include namespace infinilm::config { @@ -15,6 +16,8 @@ class QuantConfig { QuantConfig(const nlohmann::json &json); std::shared_ptr get_quantization_method() const; + std::string get_moe_weight_method() const; + bool is_moe_w16a16_marlin_enabled() const; infinilm::quantization::QuantScheme get_quant_scheme() const { if (quantization_method != nullptr) { diff --git a/csrc/layers/moe/common/moe_types.hpp b/csrc/layers/moe/common/moe_types.hpp index 8b784cf22..13b4a53a7 100644 --- a/csrc/layers/moe/common/moe_types.hpp +++ b/csrc/layers/moe/common/moe_types.hpp @@ -20,6 +20,11 @@ enum class CombineInputFormat { DeepEPLL, }; +enum class MoeWeightBackend { + Dense, + HygonW16A16Marlin, +}; + struct DispatchOutput { DispatchOutputFormat format = DispatchOutputFormat::Standard; infinicore::Tensor hidden_states; @@ -53,6 +58,7 @@ struct CombineInput { struct MoeWeights { infinicore::Tensor packed_w13; infinicore::Tensor packed_w2; + MoeWeightBackend backend = MoeWeightBackend::Dense; bool empty() const { return !packed_w13 && !packed_w2; @@ -61,6 +67,10 @@ struct MoeWeights { bool has_packed_dense_weights() const { return packed_w13 && packed_w2; } + + bool is_hygon_w16a16_marlin() const { + return backend == MoeWeightBackend::HygonW16A16Marlin; + } }; struct MoeWorkspace { @@ -69,6 +79,8 @@ struct MoeWorkspace { infinicore::Tensor ep_gathered_topk_ids; infinicore::Tensor ep_reduced_hidden_states; infinicore::Tensor fused_moe_output; + infinicore::Tensor marlin_cache13; + infinicore::Tensor marlin_cache2; infinicore::Tensor sorted_token_ids; infinicore::Tensor expert_ids; @@ -85,6 +97,8 @@ struct MoeWorkspace { size_t ep_gathered_tokens_capacity = 0; size_t ep_reduced_tokens_capacity = 0; size_t fused_moe_output_tokens_capacity = 0; + size_t marlin_cache13_capacity = 0; + size_t marlin_cache2_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/experts/fused_moe_experts.cpp b/csrc/layers/moe/experts/fused_moe_experts.cpp index b456395d6..a94d09b23 100644 --- a/csrc/layers/moe/experts/fused_moe_experts.cpp +++ b/csrc/layers/moe/experts/fused_moe_experts.cpp @@ -3,16 +3,23 @@ #include "../../../global_state/global_state.hpp" #include "../ep/ep_config.hpp" +#include "infinicore/ops/moe_w16a16_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(); + enable_hygon_w16a16_marlin_ = model_config->is_moe_w16a16_marlin_enabled(); ASSERT(num_experts_ > 0); const auto ep_config = make_ep_config(); @@ -69,6 +76,50 @@ FusedMoeExperts::FusedMoeExperts(std::shared_ptr 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_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 3f231ec64..cb68a5a26 100644 --- a/csrc/layers/moe/experts/fused_moe_experts.hpp +++ b/csrc/layers/moe/experts/fused_moe_experts.hpp @@ -17,6 +17,8 @@ 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); @@ -24,6 +26,8 @@ class FusedMoeExperts : public infinicore::nn::Module { size_t num_experts_{0}; size_t hidden_size_{0}; size_t intermediate_size_per_partition_{0}; + bool enable_hygon_w16a16_marlin_{false}; + bool w16a16_marlin_packed_{false}; MoeWeights moe_weights_; }; diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index c37abcc88..720c13270 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -3,7 +3,14 @@ #include "infinicore/context/context.hpp" #include "infinicore/ops/moe_align.hpp" #include "infinicore/ops/moe_fused_dense.hpp" +#include "infinicore/ops/moe_w16a16_marlin.hpp" +#include "nlohmann/json.hpp" + +#include +#include +#include +#include #include #include #include @@ -11,6 +18,19 @@ namespace infinilm::layers::moe { +struct HygonW16A16MarlinGemmConfig { + int mode = 103; + int delta = 1; + size_t block_size_m = 16; + bool found = false; +}; + +struct HygonW16A16MarlinRuntimeConfig { + HygonW16A16MarlinGemmConfig gemm1; + HygonW16A16MarlinGemmConfig gemm2; + bool supported = false; +}; + CudaFusedMoeRunner::CudaFusedMoeRunner(size_t num_local_experts, size_t hidden_size, size_t intermediate_size_per_partition, @@ -22,6 +42,99 @@ CudaFusedMoeRunner::CudaFusedMoeRunner(size_t num_local_experts, namespace { +std::string env_or_default(const char *name, const char *default_value) { + const char *value = std::getenv(name); + return (value != nullptr && value[0] != '\0') ? std::string(value) : std::string(default_value); +} + +HygonW16A16MarlinGemmConfig load_lightop_marlin_config(size_t n, + size_t k, + size_t m) { + HygonW16A16MarlinGemmConfig result; + const std::string config_dir = env_or_default( + "INFINILM_LIGHTOP_CONFIG_DIR", + "/usr/local/lib/python3.10/dist-packages/lightop/configs"); + const std::string device_name = env_or_default("INFINILM_HYGON_LIGHTOP_DEVICE_NAME", "gfx936"); + const std::string num_cus = env_or_default("INFINILM_HYGON_LIGHTOP_NUM_CUS", "80"); + const std::string allow_asm_env = env_or_default("INFINILM_HYGON_LIGHTOP_ALLOW_ASM", "0"); + const bool allow_asm = allow_asm_env == "1" || allow_asm_env == "true" || + allow_asm_env == "TRUE" || allow_asm_env == "on" || + allow_asm_env == "ON" || allow_asm_env == "yes" || + allow_asm_env == "YES"; + const std::string file_name = config_dir + "/MOE_W16A16_CUDA_MARLIN_" + + std::to_string(n) + "_" + std::to_string(k) + "_" + + device_name + "_" + num_cus + ".json"; + std::ifstream file(file_name); + if (!file.is_open()) { + return result; + } + + nlohmann::json config_json; + file >> config_json; + const std::string shape_key = std::to_string(n) + "_" + std::to_string(k); + if (!config_json.contains(shape_key) || !config_json.at(shape_key).is_object()) { + return result; + } + const auto &configs = config_json.at(shape_key); + + auto usable = [&](size_t token) -> bool { + const auto key = std::to_string(token); + if (!configs.contains(key) || !configs.at(key).is_object()) { + return false; + } + const int mode = configs.at(key).value("MODE", result.mode); + return allow_asm || mode < 1000; + }; + + size_t chosen = 0; + bool has_choice = false; + size_t chosen_ge = std::numeric_limits::max(); + size_t closest_diff = std::numeric_limits::max(); + for (auto it = configs.begin(); it != configs.end(); ++it) { + size_t token = 0; + try { + token = static_cast(std::stoull(it.key())); + } catch (const std::exception &) { + continue; + } + if (!usable(token)) { + continue; + } + if (token >= m && token < chosen_ge) { + chosen_ge = token; + chosen = token; + has_choice = true; + } + const size_t diff = token > m ? token - m : m - token; + if (diff < closest_diff) { + closest_diff = diff; + if (chosen_ge == std::numeric_limits::max()) { + chosen = token; + has_choice = true; + } + } + } + if (!has_choice) { + return result; + } + + const auto &cfg = configs.at(std::to_string(chosen)); + result.mode = cfg.value("MODE", result.mode); + result.delta = cfg.value("DELTA", result.delta); + result.block_size_m = cfg.value("BLOCK_SIZE_M", result.block_size_m); + result.found = cfg.contains("MODE"); + return result; +} +HygonW16A16MarlinRuntimeConfig select_hygon_w16a16_marlin_config(size_t m, + size_t hidden_size, + size_t intermediate_size_per_partition) { + HygonW16A16MarlinRuntimeConfig config; + config.gemm1 = load_lightop_marlin_config(intermediate_size_per_partition * 2, hidden_size, m); + config.gemm2 = load_lightop_marlin_config(hidden_size, intermediate_size_per_partition, m); + config.supported = config.gemm1.found && config.gemm2.found; + return config; +} + bool same_device(const infinicore::Tensor &tensor, const infinicore::Device &device) { return tensor && tensor->device().getType() == device.getType() && tensor->device().getIndex() == device.getIndex(); } @@ -76,11 +189,29 @@ void check_packed_weight_tensor(const infinicore::Tensor &tensor, CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, const MoeWeights &weights, MoeWorkspace &workspace) const { + size_t block_size = align_block_size_; + HygonW16A16MarlinRuntimeConfig marlin_config; + if (weights.is_hygon_w16a16_marlin()) { + const auto &hidden_shape = dispatch_output.hidden_states->shape(); + if (hidden_shape.size() != 2) { + throw std::runtime_error("Hygon W16A16 Marlin MoE runner requires hidden states [M, K]"); + } + marlin_config = select_hygon_w16a16_marlin_config( + hidden_shape[0], hidden_size_, intermediate_size_per_partition_); + if (!marlin_config.supported) { + throw std::runtime_error("No lightop W16A16 Marlin MoE config found for this Hygon shape"); + } + block_size = marlin_config.gemm1.block_size_m; + } + auto runner_input = prepare_runner_input( dispatch_output, - workspace); + workspace, + block_size); - auto runner_output = run_fused_core(runner_input, weights, workspace); + auto runner_output = weights.is_hygon_w16a16_marlin() + ? run_hygon_w16a16_marlin_core(runner_input, weights, workspace, marlin_config) + : run_fused_core(runner_input, weights, workspace); return CombineInput{ CombineInputFormat::Standard, @@ -91,14 +222,14 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, } CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchOutput &dispatch_output, - MoeWorkspace &workspace) const { + 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"); } 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 @@ -201,4 +332,67 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_fused_core(const CudaFusedMoeRu }; } +CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w16a16_marlin_core( + const CudaFusedMoeRunnerInput &runner_input, + const MoeWeights &weights, + MoeWorkspace &workspace, + const HygonW16A16MarlinRuntimeConfig &config) const { + if (!weights.has_packed_dense_weights() || !weights.is_hygon_w16a16_marlin()) { + throw std::runtime_error("Hygon W16A16 Marlin MoE runner requires packed Marlin weights"); + } + const size_t top_k = runner_input.topk_output.topk_ids->shape()[1]; + const size_t num_tokens = runner_input.hidden_states->shape()[0]; + const size_t cache13_required = num_tokens * top_k * std::max(intermediate_size_per_partition_ * 2, hidden_size_); + const size_t cache2_required = num_tokens * top_k * 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 = num_tokens; + + if (!same_device(workspace.marlin_cache13, runner_input.hidden_states->device()) || + workspace.marlin_cache13->dtype() != runner_input.hidden_states->dtype() || + workspace.marlin_cache13_capacity < cache13_required) { + if (infinicore::context::isGraphRecording()) { + throw std::runtime_error("MoE Marlin cache13 workspace was not initialized before graph capture"); + } + workspace.marlin_cache13 = infinicore::Tensor::empty( + {cache13_required}, runner_input.hidden_states->dtype(), runner_input.hidden_states->device()); + workspace.marlin_cache13_capacity = cache13_required; + } + if (!same_device(workspace.marlin_cache2, runner_input.hidden_states->device()) || + workspace.marlin_cache2->dtype() != runner_input.hidden_states->dtype() || + workspace.marlin_cache2_capacity < cache2_required) { + if (infinicore::context::isGraphRecording()) { + throw std::runtime_error("MoE Marlin cache2 workspace was not initialized before graph capture"); + } + workspace.marlin_cache2 = infinicore::Tensor::empty( + {cache2_required}, runner_input.hidden_states->dtype(), runner_input.hidden_states->device()); + workspace.marlin_cache2_capacity = cache2_required; + } + + infinicore::op::moe_w16a16_marlin_fused_dense_( + workspace.fused_moe_output, + workspace.marlin_cache13, + workspace.marlin_cache2, + runner_input.hidden_states, + weights.packed_w13, + weights.packed_w2, + runner_input.topk_output.topk_weights, + runner_input.routing_metadata.sorted_token_ids, + runner_input.routing_metadata.expert_ids, + runner_input.routing_metadata.num_tokens_post_padded, + top_k, + config.gemm1.mode, + config.gemm1.delta, + config.gemm2.mode, + config.gemm2.delta); + + return CudaFusedMoeRunnerOutput{ + workspace.fused_moe_output, + }; +} + } // namespace infinilm::layers::moe diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp index 2a1a1f94e..84b6628f0 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp @@ -4,6 +4,8 @@ namespace infinilm::layers::moe { +struct HygonW16A16MarlinRuntimeConfig; + struct CudaFusedMoeRunnerInput { infinicore::Tensor hidden_states; TopKOutput topk_output; @@ -27,12 +29,19 @@ class CudaFusedMoeRunner final : public MoeRunnerCore { private: CudaFusedMoeRunnerInput prepare_runner_input(const DispatchOutput &dispatch_output, - MoeWorkspace &workspace) const; + MoeWorkspace &workspace, + size_t block_size) const; CudaFusedMoeRunnerOutput run_fused_core(const CudaFusedMoeRunnerInput &runner_input, const MoeWeights &weights, MoeWorkspace &workspace) const; + CudaFusedMoeRunnerOutput run_hygon_w16a16_marlin_core( + const CudaFusedMoeRunnerInput &runner_input, + const MoeWeights &weights, + MoeWorkspace &workspace, + const HygonW16A16MarlinRuntimeConfig &config) const; + size_t num_local_experts_ = 0; size_t hidden_size_ = 0; size_t intermediate_size_per_partition_ = 0; From 96814d18c174ff7d655213298a57b95dc01ea6bc Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 10 Jul 2026 11:47:08 +0800 Subject: [PATCH 02/14] docs-qwen3-moe-w8a8-dispatch-findings --- .../qwen3moe_w8a8_status_and_vllm_dispatch.md | 151 ++++++++++++++++++ 1 file changed, 151 insertions(+) create mode 100644 docs/qwen3moe_w8a8_status_and_vllm_dispatch.md diff --git a/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md b/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md new file mode 100644 index 000000000..579eb6681 --- /dev/null +++ b/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md @@ -0,0 +1,151 @@ +# Qwen3-MoE W8A8 InfiniLM/vLLM Status + +Date: 2026-07-10 + +## Remote Environment + +- Host: `qinyiqun@10.211.3.28` +- SSH key: `C:\Users\qinyi\.ssh\bw1000` +- Container: `qinyiqun` +- InfiniCore: `/home/qinyiqun/InfiniCore` +- InfiniLM: `/home/qinyiqun/InfiniLM` +- FP model: `/home_aclsylqidf/shared/Qwen3-30B-A3B` +- W8A8 model: `/home_aclsylqidf/shared/Qwen3-30B-A3B-Channel-INT8-w8a8` + +Runtime setup inside the container: + +```bash +unset http_proxy https_proxy all_proxy HTTP_PROXY HTTPS_PROXY ALL_PROXY +export PATH=/root/.local/bin:/opt/dtk/cuda/cuda/bin:$PATH +export XMAKE_ROOT=y +export LD_LIBRARY_PATH=/usr/local/lib/python3.10/dist-packages/torch/lib:/root/.infini/lib:${LD_LIBRARY_PATH:-} +export PYTHONPATH=/usr/local/:${PYTHONPATH:-} +``` + +InfiniCore configure must include `--graph=y`: + +```bash +xmake f --hygon-dcu=true --aten=true --flash-attn=/usr/local/lib/python3.10/dist-packages/ --cuda=/opt/dtk/cuda/cuda --ccl=true --graph=y -cv -y +xmake build +xmake install +xmake build _infinicore +xmake install _infinicore +pip install -e . +``` + +## Current Benchmark Contract + +- Model family: `Qwen3-30B-A3B` +- Target path: W8A8 quantized model +- Parallelism: `TP=2`, `DP=1`, `EP=1` +- MoE communication: TP only, no DeepEP/allgather EP path +- Benchmark length going forward: `input_len=4096`, `output_len=1280` +- Important guardrail: pass only one `input_len` value. A comma-separated input length list can hang the current benchmark path. +- Device/profiling tools: `hy-smi` and Hygon trace. + +## Current InfiniLM Findings + +The long-run stall was isolated to the W8A8 MoE path with long prefill. FP graph runs and short W8A8 decode runs can complete, so the issue is not simply long output length. + +Observed behavior before the temporary workaround: + +- W8A8 `4096/128` graph timed out. +- W8A8 `4096/128` no-graph segfaulted. +- Backtraces showed one rank waiting in `RankWorker::wait`, while the other rank was inside the W8A8 Marlin MoE path and teardown/exit handling. + +A temporary internal slice loop was added around the Hygon W8A8 Marlin MoE path: + +- Files: `csrc/layers/moe/runner/cuda_fused_moe_runner.cpp`, `.hpp` +- Env: `INFINILM_HYGON_W8A8_MOE_SLICE_TOKENS` +- Debug env: `INFINILM_DEBUG_W8A8_MOE_LOOP` + +With a small slice cap, graph runs can complete, but this is not the final target because it still uses the Marlin-packed path and does not match vLLM's W8A8 channel layout/kernel flow. + +Representative temporary result: + +- W8A8 `4096/1280`, graph, slice cap `512`: prefill about `5644 tok/s`, decode about `88.8 tok/s` + +## vLLM W8A8 MoE Path + +vLLM package path: + +- `/usr/local/lib/python3.10/dist-packages/vllm` +- Runtime version in logs: `v0.15.1` + +Important vLLM env: + +- `VLLM_FUSED_MOE_CHUNK_SIZE=16384` +- `VLLM_W8A8_BACKEND=3` + +Main call chain: + +1. `CompressedTensorsW8A8Int8MoEMethod.apply()` +2. `fused_experts(...)` +3. `lmslim.layers.fused_moe.fuse_moe_int8.fused_experts_impl_int8` + +vLLM does not repack Qwen3 MoE weights into the InfiniLM Marlin layout. It keeps ordinary channel-wise int8 tensors: + +- `w1`: `[E, 768, 2048]` +- `w2`: `[E, 2048, 384]` +- `w1_scale`: `[E, 768, 1]` +- `w2_scale`: `[E, 2048, 1]` +- `E=128`, `top_k=8` + +Operator sequence per chunk: + +1. Per-token quantize hidden states. +2. Align/count/sort tokens by expert. +3. GEMM1: `lightop.moe_gemm_w8a8(...)` +4. Activation and quantize: `fuse_silu_mul_quant(...)` +5. GEMM2: `lightop.moe_gemm_w8a8(...)` +6. Reduce top-k outputs: `moe_sum` / `moe_reduce_dispatch` + +## Representative vLLM Size Dispatch + +These are the useful anchor cases for InfiniLM implementation. We do not need to reproduce every tiny graph-capture size immediately. + +| Effective M | GEMM1 shape | GEMM1 config/kernel | GEMM2 shape | GEMM2 config/kernel | Notes | +| --- | --- | --- | --- | --- | --- | +| `1..32` | `N=768,K=2048` | small-M `lightop.moe_gemm_w8a8`, often `BLOCK_M=16` | `N=2048,K=384` | small-M `lightop.moe_gemm_w8a8` | decode/graph capture sizes | +| `896` | `N=768,K=2048` | `BLOCK_M=64, MODE=517, DELTA=1`, HIP NT prefill up | `N=2048,K=384` | `BLOCK_M=32, MODE=568, DELTA=2`, HIP NT prefill down | tail chunk | +| `4096` | `N=768,K=2048` | `BLOCK_M=128, MODE=1000, DELTA=1`, `MOE_W8A8_I8_PERCHANNEL_ASM_TN_MT128x256x128_WGM1_UP` | `N=2048,K=384` | `BLOCK_M=64, MODE=517, DELTA=2`, HIP NT prefill down | target single request prefill | +| `10240` | `N=768,K=2048` | `BLOCK_M=128, MODE=1000, DELTA=1`, same ASM UP kernel | `N=2048,K=384` | `BLOCK_M=64, MODE=523, DELTA=2`, HIP NT prefill down | vLLM chunked prefill example | + +vLLM with 16 concurrent 8K prompts enabled chunked prefill with `max_num_batched_tokens=10240`. The observed MoE effective sizes were `10240`, `8256`, `896`, plus small graph-capture sizes. This means scheduler chunking, not only `VLLM_FUSED_MOE_CHUNK_SIZE`, controls the actual large-M MoE calls. + +## vLLM W8A8 Dense Linear Path + +Main call chain: + +1. `CompressedTensorsW8A8Int8.apply_weights()` +2. `apply_int8_linear(..., w8a8_strategy=3)` +3. `per_token_quant_int8(...)` +4. `ops.blaslt_scaled_mm(...)` +5. backend 3: `hipblaslt_w8a8_channelwise_gemm` + +Representative kernels: + +- `M=1,N=4096,K=2048`: small `Cijk_Alik_Bljk_I8BS_MT64x16x256...` +- `M=4096,N=4096,K=2048`: large `Cijk_Alik_Bljk_I8BS_MT256x256x128...` + +## Implementation Direction + +The next code change should move InfiniLM W8A8 MoE toward vLLM's ordinary channel-wise path: + +1. Add a new W8A8 channel MoE backend in InfiniLM, keeping `[E,N,K]` weights and `[E,N,1]` scales instead of calling `moe_w8a8_marlin_pack`. +2. Add/route an InfiniCore wrapper around ordinary `lightop.moe_gemm_w8a8`, not the current `moe_gemm_marlin_w8a8` adaptor. +3. Reuse the existing MoE workspace pattern where possible: int8 hidden cache, int8 intermediate cache, per-token scales, BF16 intermediate/output buffers. +4. Select configs by effective `M` and GEMM shape, matching the vLLM anchors above first: small decode, `896`, `4096`, `10240`. +5. Default the MoE chunk cap to `16384` for the ordinary channel path, matching vLLM's fused MoE chunk cap. Scheduler-level chunking is still needed later for 16 concurrency and 8K-10K contexts. + +## Useful Remote Artifacts + +- vLLM probe log: `/tmp/vllm_w8a8_kernel_probe_i8192_c16_o16_20260710_110631.server.log` +- MoE micro traces: + - `/tmp/hygon_trace_lmslim_w8a8_moe_m10240_20260710_111507` + - `/tmp/hygon_trace_lmslim_w8a8_moe_m896_20260710_111603` + - `/tmp/hygon_trace_lmslim_w8a8_moe_m4096_20260710_111646` +- Dense linear micro traces: + - `/tmp/hygon_trace_vllm_w8a8_linear_m1_n4096_k2048_20260710_113331` + - `/tmp/hygon_trace_vllm_w8a8_linear_m4096_n4096_k2048_20260710_113413` + From b4a70eab15c3a92f159583ce3adf5f74ed00b6c2 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 10 Jul 2026 18:17:43 +0800 Subject: [PATCH 03/14] Share Hygon Marlin MoE workspace across layers --- csrc/layers/moe/fused_moe.cpp | 37 +++++++++++++++++++++++++++++++---- csrc/layers/moe/fused_moe.hpp | 2 +- 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/csrc/layers/moe/fused_moe.cpp b/csrc/layers/moe/fused_moe.cpp index fe1301a8a..1e6072076 100644 --- a/csrc/layers/moe/fused_moe.cpp +++ b/csrc/layers/moe/fused_moe.cpp @@ -10,6 +10,34 @@ 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 +57,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 81b3cad77..7301a5d66 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 From b687546fe005f424432e210bb77045aff71afa9a Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Sun, 12 Jul 2026 14:20:45 +0800 Subject: [PATCH 04/14] feat(hygon): enable optimized Qwen3 MoE W8A8 graph inference Requires InfiniCore b65badcc for the Hygon TP8 graph kernels and operators. --- csrc/config/model_config.hpp | 12 + csrc/config/quant_config.cpp | 86 +++++- csrc/config/quant_config.hpp | 4 + csrc/engine/compiler/paged_compiler.cpp | 39 ++- csrc/engine/compiler/paged_compiler.hpp | 2 + csrc/engine/infer_engine.cpp | 1 + csrc/engine/rank_worker.cpp | 4 +- .../causal_lm_templates/text_causal_lm.hpp | 12 + csrc/layers/linear/base_linear.cpp | 4 + csrc/layers/linear/base_linear.hpp | 1 + csrc/layers/linear/linear.cpp | 3 +- csrc/layers/moe/common/moe_types.hpp | 16 +- .../moe/dispatcher/standard_dispatcher.cpp | 7 +- csrc/layers/moe/experts/fused_moe_experts.cpp | 98 +++++- csrc/layers/moe/experts/fused_moe_experts.hpp | 4 + .../moe/runner/cuda_fused_moe_runner.cpp | 286 ++++++++++++++++-- .../moe/runner/cuda_fused_moe_runner.hpp | 12 + csrc/layers/quantization/awq_marlin.hpp | 1 + .../layers/quantization/base_quantization.hpp | 10 + .../quantization/compressed_tensors.cpp | 88 ++++++ .../quantization/compressed_tensors.hpp | 2 + csrc/layers/quantization/gptq_marlin.hpp | 1 + csrc/models/infinilm_model.cpp | 24 ++ csrc/models/infinilm_model.hpp | 6 + csrc/models/qwen3/qwen3_attention.cpp | 32 +- .../qwen3moe_w8a8_status_and_vllm_dispatch.md | 16 +- python/infinilm/modeling_utils.py | 23 +- 27 files changed, 702 insertions(+), 92 deletions(-) diff --git a/csrc/config/model_config.hpp b/csrc/config/model_config.hpp index 8c239ca1e..5a41a1ac8 100644 --- a/csrc/config/model_config.hpp +++ b/csrc/config/model_config.hpp @@ -92,10 +92,22 @@ class ModelConfig { return quant_config.get_moe_weight_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 { return quant_config.is_moe_w16a16_marlin_enabled(); } + 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 801d0fdc9..d7de7e612 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -27,6 +27,53 @@ bool truthy_env(const char *name) { return value == "1" || value == "true" || value == "on" || value == "yes"; } +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(); @@ -59,6 +106,10 @@ QuantConfig::get_quantization_method() const { } std::string QuantConfig::get_moe_weight_method() const { + return get_moe_weight_method(infinicore::Device(infinicore::Device::Type::CPU, 0)); +} + +std::string QuantConfig::get_moe_weight_method(const infinicore::Device &device) const { auto env_method = env_string("INFINILM_MOE_WEIGHT_METHOD"); if (!env_method.empty()) { return env_method; @@ -66,27 +117,32 @@ std::string QuantConfig::get_moe_weight_method() const { if (truthy_env("INFINILM_HYGON_MOE_W16A16_MARLIN")) { return "w16a16_marlin"; } - if (quantization_config.is_object()) { - 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 (method == "w16a16_marlin" || method == "hygon_w16a16_marlin") { - return "w16a16_marlin"; - } + 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 { - auto method = get_moe_weight_method(); - return method == "w16a16_marlin" || method == "hygon_w16a16_marlin"; + return is_w16a16_marlin_method(get_moe_weight_method()); +} + +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 fa1d08513..a562ea155 100644 --- a/csrc/config/quant_config.hpp +++ b/csrc/config/quant_config.hpp @@ -1,6 +1,7 @@ #pragma once #include "../utils.hpp" #include "../layers/quantization/quantization.hpp" +#include "infinicore/device.hpp" #include "nlohmann/json.hpp" #include #include @@ -17,7 +18,10 @@ class QuantConfig { std::shared_ptr get_quantization_method() const; std::string get_moe_weight_method() const; + std::string get_moe_weight_method(const infinicore::Device &device) const; bool is_moe_w16a16_marlin_enabled() 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 df3fd1cb4..6b9dace38 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -47,13 +47,17 @@ void PagedCompiler::compile() { const bool has_mamba_state = has_mamba_cache(forward_context); 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({b}, infinicore::DataType::I64, infinicore::context::getDevice()); input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice()); @@ -70,7 +74,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()); @@ -122,8 +128,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_) { @@ -136,8 +144,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(); @@ -166,18 +176,19 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & } 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 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.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. @@ -202,7 +213,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); auto shared_output = std::shared_ptr(new InfinilmModel::Output{std::get<1>(result->second.compiled)->logits->resume_from_blob_()}); diff --git a/csrc/engine/compiler/paged_compiler.hpp b/csrc/engine/compiler/paged_compiler.hpp index a1125864d..807f99a53 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 a5221eb37..0568cec4d 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -182,6 +182,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 871b48d7b..acb2d24c2 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/causal_lm_templates/text_causal_lm.hpp b/csrc/layers/causal_lm_templates/text_causal_lm.hpp index d359a7f85..fdc1e0774 100644 --- a/csrc/layers/causal_lm_templates/text_causal_lm.hpp +++ b/csrc/layers/causal_lm_templates/text_causal_lm.hpp @@ -4,8 +4,11 @@ #include "../../models/infinilm_model.hpp" #include "../linear/linear.hpp" #include "infinicore/device.hpp" +#include "infinicore/ops/select_last_token_hidden_states.hpp" + #include + namespace infinilm::layers::causal_lm_templates { /** @@ -54,6 +57,15 @@ class TextCausalLM : public InfinilmModel { if (!is_last_pp_stage()) { return {infinicore::Tensor(), 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 = lm_head_->forward(hidden_states); return {logits, hidden_states}; } diff --git a/csrc/layers/linear/base_linear.cpp b/csrc/layers/linear/base_linear.cpp index f7d1c80de..b20b667aa 100644 --- a/csrc/layers/linear/base_linear.cpp +++ b/csrc/layers/linear/base_linear.cpp @@ -75,6 +75,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 a36954836..cf367d60c 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/linear/linear.cpp b/csrc/layers/linear/linear.cpp index 84982409f..2825f6018 100644 --- a/csrc/layers/linear/linear.cpp +++ b/csrc/layers/linear/linear.cpp @@ -88,7 +88,8 @@ infinicore::Tensor RowParallelLinear::forward(infinicore::Tensor &input) const { auto output = BaseLinear::forward(input); if ((tp_size_ > 1) && (communicator_ != nullptr)) { - infinicore::op::distributed::allreduce_(output, output, INFINICCL_SUM, communicator_); + return infinicore::op::distributed::allreduce( + output, INFINICCL_SUM, communicator_); } return output; } diff --git a/csrc/layers/moe/common/moe_types.hpp b/csrc/layers/moe/common/moe_types.hpp index 13b4a53a7..1b16ed710 100644 --- a/csrc/layers/moe/common/moe_types.hpp +++ b/csrc/layers/moe/common/moe_types.hpp @@ -23,6 +23,7 @@ enum class CombineInputFormat { enum class MoeWeightBackend { Dense, HygonW16A16Marlin, + HygonW8A8Marlin, }; struct DispatchOutput { @@ -58,6 +59,8 @@ 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 { @@ -68,9 +71,17 @@ struct MoeWeights { 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 { @@ -81,6 +92,10 @@ struct MoeWorkspace { infinicore::Tensor fused_moe_output; infinicore::Tensor marlin_cache13; infinicore::Tensor marlin_cache2; + infinicore::Tensor marlin_input_i8; + infinicore::Tensor marlin_input_scale; + infinicore::Tensor marlin_cache2_i8; + infinicore::Tensor marlin_cache2_scale; infinicore::Tensor sorted_token_ids; infinicore::Tensor expert_ids; @@ -96,7 +111,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 marlin_cache13_capacity = 0; size_t marlin_cache2_capacity = 0; size_t blockscale_offsets_capacity = 0; diff --git a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp index 27e5c0b23..96bc6205f 100644 --- a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp +++ b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp @@ -30,11 +30,8 @@ infinicore::Tensor StandardDispatcher::combine(const CombineInput &combine_input MoeWorkspace &workspace) const { (void)workspace; if (tp_size_ > 1 && communicator_ != nullptr) { - infinicore::op::distributed::allreduce_( - combine_input.hidden_states, - combine_input.hidden_states, - INFINICCL_SUM, - communicator_); + return infinicore::op::distributed::allreduce( + combine_input.hidden_states, INFINICCL_SUM, communicator_); } return 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 a94d09b23..40fc8ec2d 100644 --- a/csrc/layers/moe/experts/fused_moe_experts.cpp +++ b/csrc/layers/moe/experts/fused_moe_experts.cpp @@ -4,6 +4,7 @@ #include "../ep/ep_config.hpp" #include "infinicore/ops/moe_w16a16_marlin.hpp" +#include "infinicore/ops/moe_w8a8_marlin.hpp" #include @@ -19,7 +20,16 @@ FusedMoeExperts::FusedMoeExperts(std::shared_ptr 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(); - enable_hygon_w16a16_marlin_ = model_config->is_moe_w16a16_marlin_enabled(); + 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(); @@ -39,17 +49,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_ @@ -72,6 +96,27 @@ 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_; @@ -80,6 +125,53 @@ FusedMoeExperts::FusedMoeExperts(std::shared_ptr } 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; } diff --git a/csrc/layers/moe/experts/fused_moe_experts.hpp b/csrc/layers/moe/experts/fused_moe_experts.hpp index cb68a5a26..a8f2a53b5 100644 --- a/csrc/layers/moe/experts/fused_moe_experts.hpp +++ b/csrc/layers/moe/experts/fused_moe_experts.hpp @@ -22,12 +22,16 @@ class FusedMoeExperts : public infinicore::nn::Module { 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/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index 720c13270..630723a08 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -4,6 +4,7 @@ #include "infinicore/ops/moe_align.hpp" #include "infinicore/ops/moe_fused_dense.hpp" #include "infinicore/ops/moe_w16a16_marlin.hpp" +#include "infinicore/ops/moe_w8a8_marlin.hpp" #include "nlohmann/json.hpp" @@ -14,11 +15,10 @@ #include #include #include -#include namespace infinilm::layers::moe { -struct HygonW16A16MarlinGemmConfig { +struct HygonMarlinGemmConfig { int mode = 103; int delta = 1; size_t block_size_m = 16; @@ -26,8 +26,14 @@ struct HygonW16A16MarlinGemmConfig { }; struct HygonW16A16MarlinRuntimeConfig { - HygonW16A16MarlinGemmConfig gemm1; - HygonW16A16MarlinGemmConfig gemm2; + HygonMarlinGemmConfig gemm1; + HygonMarlinGemmConfig gemm2; + bool supported = false; +}; + +struct HygonW8A8MarlinRuntimeConfig { + HygonMarlinGemmConfig gemm1; + HygonMarlinGemmConfig gemm2; bool supported = false; }; @@ -47,23 +53,34 @@ std::string env_or_default(const char *name, const char *default_value) { return (value != nullptr && value[0] != '\0') ? std::string(value) : std::string(default_value); } -HygonW16A16MarlinGemmConfig load_lightop_marlin_config(size_t n, - size_t k, - size_t m) { - HygonW16A16MarlinGemmConfig result; +bool env_flag_enabled(const char *name, bool default_value) { + const char *value = std::getenv(name); + if (value == nullptr || value[0] == '\0') { + return default_value; + } + std::string s(value); + return s == "1" || s == "true" || s == "TRUE" || s == "on" || s == "ON" || s == "yes" || s == "YES"; +} + +constexpr size_t kHygonW8A8MoeSliceTokens = 16384; + +HygonMarlinGemmConfig load_lightop_marlin_config(size_t n, + size_t k, + size_t m, + const std::string &file_prefix, + const char *device_name_default, + bool allow_asm, + bool num_cus_with_cu_prefix) { + HygonMarlinGemmConfig result; const std::string config_dir = env_or_default( "INFINILM_LIGHTOP_CONFIG_DIR", "/usr/local/lib/python3.10/dist-packages/lightop/configs"); - const std::string device_name = env_or_default("INFINILM_HYGON_LIGHTOP_DEVICE_NAME", "gfx936"); + const std::string device_name = env_or_default("INFINILM_HYGON_LIGHTOP_DEVICE_NAME", device_name_default); const std::string num_cus = env_or_default("INFINILM_HYGON_LIGHTOP_NUM_CUS", "80"); - const std::string allow_asm_env = env_or_default("INFINILM_HYGON_LIGHTOP_ALLOW_ASM", "0"); - const bool allow_asm = allow_asm_env == "1" || allow_asm_env == "true" || - allow_asm_env == "TRUE" || allow_asm_env == "on" || - allow_asm_env == "ON" || allow_asm_env == "yes" || - allow_asm_env == "YES"; - const std::string file_name = config_dir + "/MOE_W16A16_CUDA_MARLIN_" + + const std::string num_cus_suffix = num_cus_with_cu_prefix ? ("_CU" + num_cus) : ("_" + num_cus); + const std::string file_name = config_dir + "/" + file_prefix + "_" + std::to_string(n) + "_" + std::to_string(k) + "_" + - device_name + "_" + num_cus + ".json"; + device_name + num_cus_suffix + ".json"; std::ifstream file(file_name); if (!file.is_open()) { return result; @@ -129,8 +146,27 @@ HygonW16A16MarlinRuntimeConfig select_hygon_w16a16_marlin_config(size_t m, size_t hidden_size, size_t intermediate_size_per_partition) { HygonW16A16MarlinRuntimeConfig config; - config.gemm1 = load_lightop_marlin_config(intermediate_size_per_partition * 2, hidden_size, m); - config.gemm2 = load_lightop_marlin_config(hidden_size, intermediate_size_per_partition, m); + const bool allow_asm = env_flag_enabled("INFINILM_HYGON_LIGHTOP_ALLOW_ASM", false); + config.gemm1 = load_lightop_marlin_config( + intermediate_size_per_partition * 2, hidden_size, m, + "MOE_W16A16_CUDA_MARLIN", "gfx936", allow_asm, false); + config.gemm2 = load_lightop_marlin_config( + hidden_size, intermediate_size_per_partition, m, + "MOE_W16A16_CUDA_MARLIN", "gfx936", allow_asm, false); + config.supported = config.gemm1.found && config.gemm2.found; + return config; +} + +HygonW8A8MarlinRuntimeConfig select_hygon_w8a8_marlin_config(size_t m, + size_t hidden_size, + size_t intermediate_size_per_partition) { + HygonW8A8MarlinRuntimeConfig config; + config.gemm1 = load_lightop_marlin_config( + intermediate_size_per_partition * 2, hidden_size, m, + "MOE_BLOCKINT8_CUDA_MARLIN", "GFX936", true, true); + config.gemm2 = load_lightop_marlin_config( + hidden_size, intermediate_size_per_partition, m, + "MOE_BLOCKINT8_CUDA_MARLIN", "GFX936", true, true); config.supported = config.gemm1.found && config.gemm2.found; return config; } @@ -176,7 +212,7 @@ void check_packed_weight_tensor(const infinicore::Tensor &tensor, 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( @@ -191,17 +227,37 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, MoeWorkspace &workspace) const { size_t block_size = align_block_size_; HygonW16A16MarlinRuntimeConfig marlin_config; - if (weights.is_hygon_w16a16_marlin()) { + HygonW8A8MarlinRuntimeConfig w8a8_marlin_config; + if (weights.is_hygon_w16a16_marlin() || weights.is_hygon_w8a8_marlin()) { const auto &hidden_shape = dispatch_output.hidden_states->shape(); if (hidden_shape.size() != 2) { - throw std::runtime_error("Hygon W16A16 Marlin MoE runner requires hidden states [M, K]"); + throw std::runtime_error("Hygon Marlin MoE runner requires hidden states [M, K]"); } - marlin_config = select_hygon_w16a16_marlin_config( - hidden_shape[0], hidden_size_, intermediate_size_per_partition_); - if (!marlin_config.supported) { - throw std::runtime_error("No lightop W16A16 Marlin MoE config found for this Hygon shape"); + if (weights.is_hygon_w16a16_marlin()) { + marlin_config = select_hygon_w16a16_marlin_config( + hidden_shape[0], hidden_size_, intermediate_size_per_partition_); + if (!marlin_config.supported) { + throw std::runtime_error("No lightop W16A16 Marlin MoE config found for this Hygon shape"); + } + block_size = marlin_config.gemm1.block_size_m; + } else { + if (hidden_shape[0] > kHygonW8A8MoeSliceTokens) { + auto runner_output = run_hygon_w8a8_marlin_core_sliced( + dispatch_output, weights, workspace); + return CombineInput{ + CombineInputFormat::Standard, + runner_output.hidden_states, + dispatch_output.topk_output, + MoeRoutingMetadata{}, + }; + } + w8a8_marlin_config = select_hygon_w8a8_marlin_config( + hidden_shape[0], hidden_size_, intermediate_size_per_partition_); + if (!w8a8_marlin_config.supported) { + throw std::runtime_error("No lightop W8A8 Marlin MoE config found for this Hygon shape"); + } + block_size = w8a8_marlin_config.gemm1.block_size_m; } - block_size = marlin_config.gemm1.block_size_m; } auto runner_input = prepare_runner_input( @@ -211,7 +267,10 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, auto runner_output = weights.is_hygon_w16a16_marlin() ? run_hygon_w16a16_marlin_core(runner_input, weights, workspace, marlin_config) - : run_fused_core(runner_input, weights, workspace); + : (weights.is_hygon_w8a8_marlin() + ? run_hygon_w8a8_marlin_core( + runner_input, weights, workspace, w8a8_marlin_config) + : run_fused_core(runner_input, weights, workspace)); return CombineInput{ CombineInputFormat::Standard, @@ -316,7 +375,6 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_fused_core(const CudaFusedMoeRu 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, @@ -350,8 +408,6 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w16a16_marlin_core( runner_input.hidden_states->shape(), runner_input.hidden_states->dtype(), runner_input.hidden_states->device()); - workspace.fused_moe_output_tokens_capacity = num_tokens; - if (!same_device(workspace.marlin_cache13, runner_input.hidden_states->device()) || workspace.marlin_cache13->dtype() != runner_input.hidden_states->dtype() || workspace.marlin_cache13_capacity < cache13_required) { @@ -395,4 +451,174 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w16a16_marlin_core( }; } +CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w8a8_marlin_core( + const CudaFusedMoeRunnerInput &runner_input, + const MoeWeights &weights, + MoeWorkspace &workspace, + const HygonW8A8MarlinRuntimeConfig &config) const { + if (!weights.has_packed_w8a8_marlin_weights() || !weights.is_hygon_w8a8_marlin()) { + throw std::runtime_error("Hygon W8A8 Marlin MoE runner requires packed Marlin weights and scales"); + } + const size_t top_k = runner_input.topk_output.topk_ids->shape()[1]; + const size_t num_tokens = runner_input.hidden_states->shape()[0]; + const size_t cache13_required = num_tokens * top_k * std::max(intermediate_size_per_partition_ * 2, hidden_size_); + + check_packed_weight_tensor( + weights.packed_w13, + "w13", + runner_input.hidden_states->device(), + infinicore::DataType::I8, + {num_local_experts_, hidden_size_ / 64, intermediate_size_per_partition_ * 2 * 64}); + check_packed_weight_tensor( + weights.packed_w2, + "w2", + runner_input.hidden_states->device(), + infinicore::DataType::I8, + {num_local_experts_, intermediate_size_per_partition_ / 64, hidden_size_ * 64}); + check_packed_weight_tensor( + weights.packed_w13_scale, + "w13_scale", + runner_input.hidden_states->device(), + infinicore::DataType::F32, + {num_local_experts_, intermediate_size_per_partition_ * 2, 1}); + check_packed_weight_tensor( + weights.packed_w2_scale, + "w2_scale", + runner_input.hidden_states->device(), + infinicore::DataType::F32, + {num_local_experts_, hidden_size_, 1}); + + ensure_tensor( + workspace.fused_moe_output, + runner_input.hidden_states->shape(), + runner_input.hidden_states->dtype(), + runner_input.hidden_states->device()); + if (!same_device(workspace.marlin_cache13, runner_input.hidden_states->device()) || + workspace.marlin_cache13->dtype() != runner_input.hidden_states->dtype() || + workspace.marlin_cache13_capacity < cache13_required) { + if (infinicore::context::isGraphRecording()) { + throw std::runtime_error("MoE W8A8 Marlin cache13 workspace was not initialized before graph capture"); + } + workspace.marlin_cache13 = infinicore::Tensor::empty( + {cache13_required}, runner_input.hidden_states->dtype(), runner_input.hidden_states->device()); + workspace.marlin_cache13_capacity = cache13_required; + } + ensure_tensor( + workspace.marlin_input_i8, + {num_tokens, hidden_size_}, + infinicore::DataType::I8, + runner_input.hidden_states->device()); + ensure_tensor( + workspace.marlin_input_scale, + {num_tokens, 1}, + infinicore::DataType::F32, + runner_input.hidden_states->device()); + ensure_tensor( + workspace.marlin_cache2_i8, + {num_tokens * top_k, intermediate_size_per_partition_}, + infinicore::DataType::I8, + runner_input.hidden_states->device()); + ensure_tensor( + workspace.marlin_cache2_scale, + {num_tokens * top_k, 1}, + infinicore::DataType::F32, + runner_input.hidden_states->device()); + + infinicore::op::moe_w8a8_marlin_fused_dense_( + workspace.fused_moe_output, + workspace.marlin_cache13, + workspace.marlin_cache2_i8, + workspace.marlin_input_i8, + workspace.marlin_input_scale, + workspace.marlin_cache2_scale, + runner_input.hidden_states, + weights.packed_w13, + weights.packed_w2, + weights.packed_w13_scale, + weights.packed_w2_scale, + runner_input.topk_output.topk_weights, + runner_input.routing_metadata.sorted_token_ids, + runner_input.routing_metadata.expert_ids, + runner_input.routing_metadata.num_tokens_post_padded, + top_k, + config.gemm1.mode, + config.gemm1.block_size_m, + config.gemm1.delta, + config.gemm2.mode, + config.gemm2.delta); + + return CudaFusedMoeRunnerOutput{ + workspace.fused_moe_output, + }; +} + +CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w8a8_marlin_core_sliced( + const DispatchOutput &dispatch_output, + const MoeWeights &weights, + MoeWorkspace &workspace) const { + const auto &hidden_shape = dispatch_output.hidden_states->shape(); + if (hidden_shape.size() != 2) { + throw std::runtime_error("Hygon W8A8 sliced MoE runner requires hidden states [M, K]"); + } + const size_t num_tokens = hidden_shape[0]; + if (infinicore::context::isGraphRecording()) { + throw std::runtime_error("Hygon W8A8 sliced MoE runner cannot allocate/copy slice outputs during graph capture"); + } + + ensure_tensor( + workspace.fused_moe_output, + dispatch_output.hidden_states->shape(), + dispatch_output.hidden_states->dtype(), + dispatch_output.hidden_states->device()); + MoeWorkspace slice_workspace; + const auto full_slice_config = select_hygon_w8a8_marlin_config( + kHygonW8A8MoeSliceTokens, hidden_size_, intermediate_size_per_partition_); + if (!full_slice_config.supported) { + throw std::runtime_error("No lightop W8A8 Marlin MoE config found for full Hygon slice"); + } + size_t offset = 0; + while (offset < num_tokens) { + const size_t slice_tokens = std::min(kHygonW8A8MoeSliceTokens, num_tokens - offset); + const auto slice_config = slice_tokens == kHygonW8A8MoeSliceTokens + ? full_slice_config + : select_hygon_w8a8_marlin_config( + slice_tokens, hidden_size_, intermediate_size_per_partition_); + if (!slice_config.supported) { + throw std::runtime_error("No lightop W8A8 Marlin MoE config found for sliced Hygon shape"); + } + + const auto hidden_slice = dispatch_output.hidden_states->narrow({{0, offset, slice_tokens}}); + const TopKOutput topk_slice{ + dispatch_output.topk_output.topk_weights->narrow({{0, offset, slice_tokens}}), + dispatch_output.topk_output.topk_ids->narrow({{0, offset, slice_tokens}}), + dispatch_output.topk_output.router_logits + ? dispatch_output.topk_output.router_logits->narrow({{0, offset, slice_tokens}}) + : infinicore::Tensor(), + }; + const DispatchOutput dispatch_slice{ + DispatchOutputFormat::Standard, + hidden_slice, + infinicore::Tensor(), + topk_slice, + infinicore::Tensor(), + }; + auto slice_input = prepare_runner_input( + dispatch_slice, + slice_workspace, + slice_config.gemm1.block_size_m); + auto slice_output = run_hygon_w8a8_marlin_core( + slice_input, + weights, + slice_workspace, + slice_config); + workspace.fused_moe_output->narrow({{0, offset, slice_tokens}})->copy_from(slice_output.hidden_states); + + offset += slice_tokens; + } + + return CudaFusedMoeRunnerOutput{ + workspace.fused_moe_output, + }; +} + } // namespace infinilm::layers::moe diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp index 84b6628f0..a6fc19845 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp @@ -5,6 +5,7 @@ namespace infinilm::layers::moe { struct HygonW16A16MarlinRuntimeConfig; +struct HygonW8A8MarlinRuntimeConfig; struct CudaFusedMoeRunnerInput { infinicore::Tensor hidden_states; @@ -42,6 +43,17 @@ class CudaFusedMoeRunner final : public MoeRunnerCore { MoeWorkspace &workspace, const HygonW16A16MarlinRuntimeConfig &config) const; + CudaFusedMoeRunnerOutput run_hygon_w8a8_marlin_core( + const CudaFusedMoeRunnerInput &runner_input, + const MoeWeights &weights, + MoeWorkspace &workspace, + const HygonW8A8MarlinRuntimeConfig &config) const; + + CudaFusedMoeRunnerOutput run_hygon_w8a8_marlin_core_sliced( + const DispatchOutput &dispatch_output, + const MoeWeights &weights, + MoeWorkspace &workspace) const; + size_t num_local_experts_ = 0; size_t hidden_size_ = 0; size_t intermediate_size_per_partition_ = 0; diff --git a/csrc/layers/quantization/awq_marlin.hpp b/csrc/layers/quantization/awq_marlin.hpp index a22fb43b3..5b6bc9428 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 6c37f4121..56350a616 100644 --- a/csrc/layers/quantization/base_quantization.hpp +++ b/csrc/layers/quantization/base_quantization.hpp @@ -69,6 +69,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 @@ -100,6 +108,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 ff5617a1e..726595549 100644 --- a/csrc/layers/quantization/compressed_tensors.cpp +++ b/csrc/layers/quantization/compressed_tensors.cpp @@ -2,10 +2,78 @@ #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 +96,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 dcf65c2e0..2a088728e 100644 --- a/csrc/layers/quantization/compressed_tensors.hpp +++ b/csrc/layers/quantization/compressed_tensors.hpp @@ -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 c16c79438..1fc345b92 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 5d284a316..2c6a652da 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 ac994fd6d..b850a960a 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 { @@ -74,6 +76,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, @@ -86,5 +90,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 7d9beb043..57b840560 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,32 @@ 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/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md b/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md index 579eb6681..e8256838d 100644 --- a/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md +++ b/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md @@ -47,23 +47,20 @@ pip install -e . The long-run stall was isolated to the W8A8 MoE path with long prefill. FP graph runs and short W8A8 decode runs can complete, so the issue is not simply long output length. -Observed behavior before the temporary workaround: +Observed behavior before long-prefill slicing: - W8A8 `4096/128` graph timed out. - W8A8 `4096/128` no-graph segfaulted. - Backtraces showed one rank waiting in `RankWorker::wait`, while the other rank was inside the W8A8 Marlin MoE path and teardown/exit handling. -A temporary internal slice loop was added around the Hygon W8A8 Marlin MoE path: +The Hygon W8A8 Marlin MoE path now chunks long prefill internally: - Files: `csrc/layers/moe/runner/cuda_fused_moe_runner.cpp`, `.hpp` -- Env: `INFINILM_HYGON_W8A8_MOE_SLICE_TOKENS` -- Debug env: `INFINILM_DEBUG_W8A8_MOE_LOOP` +- Fixed chunk size: `16384` tokens, matching vLLM's production chunk size +- The sliced path is selected before full-input routing metadata is prepared, so each token is aligned only once +- No W8A8 slice or debug environment switches are required -With a small slice cap, graph runs can complete, but this is not the final target because it still uses the Marlin-packed path and does not match vLLM's W8A8 channel layout/kernel flow. - -Representative temporary result: - -- W8A8 `4096/1280`, graph, slice cap `512`: prefill about `5644 tok/s`, decode about `88.8 tok/s` +This keeps long-prefill workspace bounded while decode continues to use the graph-captured Marlin path directly. ## vLLM W8A8 MoE Path @@ -148,4 +145,3 @@ The next code change should move InfiniLM W8A8 MoE toward vLLM's ordinary channe - Dense linear micro traces: - `/tmp/hygon_trace_vllm_w8a8_linear_m1_n4096_k2048_20260710_113331` - `/tmp/hygon_trace_vllm_w8a8_linear_m4096_n4096_k2048_20260710_113413` - diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 9bcefa7f8..9e0f9de8c 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 From 777795472e7aba597dce969b29538f1fa8d6acf0 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 13 Jul 2026 15:18:16 +0800 Subject: [PATCH 05/14] fix-hygon-select-moe-path-from-model --- csrc/config/quant_config.cpp | 21 ------- .../moe/runner/cuda_fused_moe_runner.cpp | 56 +++++++++++++------ 2 files changed, 38 insertions(+), 39 deletions(-) diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index d7de7e612..14cce4afa 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -2,7 +2,6 @@ #include #include -#include namespace infinilm::config { namespace { @@ -14,19 +13,6 @@ std::string lower_string(std::string value) { return value; } -std::string env_string(const char *name) { - const char *value = std::getenv(name); - if (value == nullptr || value[0] == '\0') { - return {}; - } - return lower_string(value); -} - -bool truthy_env(const char *name) { - auto value = env_string(name); - return value == "1" || value == "true" || value == "on" || value == "yes"; -} - bool is_w16a16_marlin_method(const std::string &method) { return method == "w16a16_marlin" || method == "hygon_w16a16_marlin"; } @@ -110,13 +96,6 @@ std::string QuantConfig::get_moe_weight_method() const { } std::string QuantConfig::get_moe_weight_method(const infinicore::Device &device) const { - auto env_method = env_string("INFINILM_MOE_WEIGHT_METHOD"); - if (!env_method.empty()) { - return env_method; - } - if (truthy_env("INFINILM_HYGON_MOE_W16A16_MARLIN")) { - return "w16a16_marlin"; - } auto configured_method = explicit_moe_weight_method(quantization_config); if (!configured_method.empty()) { return configured_method; diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index 630723a08..a8796a4aa 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -53,23 +53,20 @@ std::string env_or_default(const char *name, const char *default_value) { return (value != nullptr && value[0] != '\0') ? std::string(value) : std::string(default_value); } -bool env_flag_enabled(const char *name, bool default_value) { - const char *value = std::getenv(name); - if (value == nullptr || value[0] == '\0') { - return default_value; - } - std::string s(value); - return s == "1" || s == "true" || s == "TRUE" || s == "on" || s == "ON" || s == "yes" || s == "YES"; -} - constexpr size_t kHygonW8A8MoeSliceTokens = 16384; +enum class HygonMarlinModePolicy { + LegacyOnly, + LegacyAndBf16Mode1000, + All, +}; + HygonMarlinGemmConfig load_lightop_marlin_config(size_t n, size_t k, size_t m, const std::string &file_prefix, const char *device_name_default, - bool allow_asm, + HygonMarlinModePolicy mode_policy, bool num_cus_with_cu_prefix) { HygonMarlinGemmConfig result; const std::string config_dir = env_or_default( @@ -100,7 +97,9 @@ HygonMarlinGemmConfig load_lightop_marlin_config(size_t n, return false; } const int mode = configs.at(key).value("MODE", result.mode); - return allow_asm || mode < 1000; + return mode < 1000 || + mode_policy == HygonMarlinModePolicy::All || + (mode_policy == HygonMarlinModePolicy::LegacyAndBf16Mode1000 && mode == 1000); }; size_t chosen = 0; @@ -144,15 +143,18 @@ HygonMarlinGemmConfig load_lightop_marlin_config(size_t n, } HygonW16A16MarlinRuntimeConfig select_hygon_w16a16_marlin_config(size_t m, size_t hidden_size, - size_t intermediate_size_per_partition) { + size_t intermediate_size_per_partition, + infinicore::DataType hidden_dtype) { HygonW16A16MarlinRuntimeConfig config; - const bool allow_asm = env_flag_enabled("INFINILM_HYGON_LIGHTOP_ALLOW_ASM", false); + const auto mode_policy = hidden_dtype == infinicore::DataType::BF16 + ? HygonMarlinModePolicy::LegacyAndBf16Mode1000 + : HygonMarlinModePolicy::LegacyOnly; config.gemm1 = load_lightop_marlin_config( intermediate_size_per_partition * 2, hidden_size, m, - "MOE_W16A16_CUDA_MARLIN", "gfx936", allow_asm, false); + "MOE_W16A16_CUDA_MARLIN", "gfx936", mode_policy, false); config.gemm2 = load_lightop_marlin_config( hidden_size, intermediate_size_per_partition, m, - "MOE_W16A16_CUDA_MARLIN", "gfx936", allow_asm, false); + "MOE_W16A16_CUDA_MARLIN", "gfx936", mode_policy, false); config.supported = config.gemm1.found && config.gemm2.found; return config; } @@ -163,10 +165,10 @@ HygonW8A8MarlinRuntimeConfig select_hygon_w8a8_marlin_config(size_t m, HygonW8A8MarlinRuntimeConfig config; config.gemm1 = load_lightop_marlin_config( intermediate_size_per_partition * 2, hidden_size, m, - "MOE_BLOCKINT8_CUDA_MARLIN", "GFX936", true, true); + "MOE_BLOCKINT8_CUDA_MARLIN", "GFX936", HygonMarlinModePolicy::All, true); config.gemm2 = load_lightop_marlin_config( hidden_size, intermediate_size_per_partition, m, - "MOE_BLOCKINT8_CUDA_MARLIN", "GFX936", true, true); + "MOE_BLOCKINT8_CUDA_MARLIN", "GFX936", HygonMarlinModePolicy::All, true); config.supported = config.gemm1.found && config.gemm2.found; return config; } @@ -235,7 +237,8 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, } if (weights.is_hygon_w16a16_marlin()) { marlin_config = select_hygon_w16a16_marlin_config( - hidden_shape[0], hidden_size_, intermediate_size_per_partition_); + hidden_shape[0], hidden_size_, intermediate_size_per_partition_, + dispatch_output.hidden_states->dtype()); if (!marlin_config.supported) { throw std::runtime_error("No lightop W16A16 Marlin MoE config found for this Hygon shape"); } @@ -398,6 +401,23 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w16a16_marlin_core( if (!weights.has_packed_dense_weights() || !weights.is_hygon_w16a16_marlin()) { throw std::runtime_error("Hygon W16A16 Marlin MoE runner requires packed Marlin weights"); } + const auto activation_dtype = runner_input.hidden_states->dtype(); + if (activation_dtype != infinicore::DataType::BF16 && + activation_dtype != infinicore::DataType::F16) { + throw std::runtime_error("Hygon W16A16 Marlin MoE runner requires BF16 or FP16 activations"); + } + check_packed_weight_tensor( + weights.packed_w13, + "w13", + runner_input.hidden_states->device(), + activation_dtype, + {num_local_experts_, hidden_size_ / 16, intermediate_size_per_partition_ * 2 * 16}); + check_packed_weight_tensor( + weights.packed_w2, + "w2", + runner_input.hidden_states->device(), + activation_dtype, + {num_local_experts_, intermediate_size_per_partition_ / 16, hidden_size_ * 16}); const size_t top_k = runner_input.topk_output.topk_ids->shape()[1]; const size_t num_tokens = runner_input.hidden_states->shape()[0]; const size_t cache13_required = num_tokens * top_k * std::max(intermediate_size_per_partition_ * 2, hidden_size_); From c32739e7ef74b2862c8af5543367f648eb19ddeb Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 13 Jul 2026 13:01:52 +0000 Subject: [PATCH 06/14] feat(hygon): select LightOP configs from device --- .../moe/runner/cuda_fused_moe_runner.cpp | 65 +++++++++++++++---- 1 file changed, 52 insertions(+), 13 deletions(-) diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index a8796a4aa..4b80e3524 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -5,10 +5,12 @@ #include "infinicore/ops/moe_fused_dense.hpp" #include "infinicore/ops/moe_w16a16_marlin.hpp" #include "infinicore/ops/moe_w8a8_marlin.hpp" +#include "infinicore/adaptor/lightop_adaptor.hpp" #include "nlohmann/json.hpp" #include +#include #include #include #include @@ -53,6 +55,28 @@ std::string env_or_default(const char *name, const char *default_value) { return (value != nullptr && value[0] != '\0') ? std::string(value) : std::string(default_value); } +std::string normalize_hygon_gpu_target(std::string target, bool uppercase) { + const auto feature_pos = target.find(':'); + if (feature_pos != std::string::npos) { + target.resize(feature_pos); + } + std::transform(target.begin(), target.end(), target.begin(), [uppercase](unsigned char ch) { + return static_cast(uppercase ? std::toupper(ch) : std::tolower(ch)); + }); + + std::string lowercase = target; + std::transform(lowercase.begin(), lowercase.end(), lowercase.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + if (lowercase.size() <= 3 || lowercase.compare(0, 3, "gfx") != 0 || + !std::all_of(lowercase.begin() + 3, lowercase.end(), [](unsigned char ch) { + return std::isalnum(ch) != 0; + })) { + throw std::runtime_error("Invalid Hygon GPU target for lightop config: " + target); + } + return target; +} + constexpr size_t kHygonW8A8MoeSliceTokens = 16384; enum class HygonMarlinModePolicy { @@ -65,15 +89,21 @@ HygonMarlinGemmConfig load_lightop_marlin_config(size_t n, size_t k, size_t m, const std::string &file_prefix, - const char *device_name_default, + const infinicore::adaptor::lightop::DeviceInfo &device_info, HygonMarlinModePolicy mode_policy, + bool uppercase_device_name, bool num_cus_with_cu_prefix) { HygonMarlinGemmConfig result; const std::string config_dir = env_or_default( "INFINILM_LIGHTOP_CONFIG_DIR", "/usr/local/lib/python3.10/dist-packages/lightop/configs"); - const std::string device_name = env_or_default("INFINILM_HYGON_LIGHTOP_DEVICE_NAME", device_name_default); - const std::string num_cus = env_or_default("INFINILM_HYGON_LIGHTOP_NUM_CUS", "80"); + if (device_info.gpu_target.empty() || device_info.compute_units <= 0) { + throw std::runtime_error("Unable to query Hygon device properties for lightop config"); + } + const std::string device_name = normalize_hygon_gpu_target( + device_info.gpu_target, + uppercase_device_name); + const std::string num_cus = std::to_string(device_info.compute_units); const std::string num_cus_suffix = num_cus_with_cu_prefix ? ("_CU" + num_cus) : ("_" + num_cus); const std::string file_name = config_dir + "/" + file_prefix + "_" + std::to_string(n) + "_" + std::to_string(k) + "_" + @@ -144,31 +174,35 @@ HygonMarlinGemmConfig load_lightop_marlin_config(size_t n, HygonW16A16MarlinRuntimeConfig select_hygon_w16a16_marlin_config(size_t m, size_t hidden_size, size_t intermediate_size_per_partition, - infinicore::DataType hidden_dtype) { + infinicore::DataType hidden_dtype, + size_t device_index) { HygonW16A16MarlinRuntimeConfig config; + const auto device_info = infinicore::adaptor::lightop::device_info(device_index); const auto mode_policy = hidden_dtype == infinicore::DataType::BF16 ? HygonMarlinModePolicy::LegacyAndBf16Mode1000 : HygonMarlinModePolicy::LegacyOnly; config.gemm1 = load_lightop_marlin_config( intermediate_size_per_partition * 2, hidden_size, m, - "MOE_W16A16_CUDA_MARLIN", "gfx936", mode_policy, false); + "MOE_W16A16_CUDA_MARLIN", device_info, mode_policy, false, false); config.gemm2 = load_lightop_marlin_config( hidden_size, intermediate_size_per_partition, m, - "MOE_W16A16_CUDA_MARLIN", "gfx936", mode_policy, false); + "MOE_W16A16_CUDA_MARLIN", device_info, mode_policy, false, false); config.supported = config.gemm1.found && config.gemm2.found; return config; } HygonW8A8MarlinRuntimeConfig select_hygon_w8a8_marlin_config(size_t m, size_t hidden_size, - size_t intermediate_size_per_partition) { + size_t intermediate_size_per_partition, + size_t device_index) { HygonW8A8MarlinRuntimeConfig config; + const auto device_info = infinicore::adaptor::lightop::device_info(device_index); config.gemm1 = load_lightop_marlin_config( intermediate_size_per_partition * 2, hidden_size, m, - "MOE_BLOCKINT8_CUDA_MARLIN", "GFX936", HygonMarlinModePolicy::All, true); + "MOE_BLOCKINT8_CUDA_MARLIN", device_info, HygonMarlinModePolicy::All, true, true); config.gemm2 = load_lightop_marlin_config( hidden_size, intermediate_size_per_partition, m, - "MOE_BLOCKINT8_CUDA_MARLIN", "GFX936", HygonMarlinModePolicy::All, true); + "MOE_BLOCKINT8_CUDA_MARLIN", device_info, HygonMarlinModePolicy::All, true, true); config.supported = config.gemm1.found && config.gemm2.found; return config; } @@ -238,7 +272,8 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, if (weights.is_hygon_w16a16_marlin()) { marlin_config = select_hygon_w16a16_marlin_config( hidden_shape[0], hidden_size_, intermediate_size_per_partition_, - dispatch_output.hidden_states->dtype()); + dispatch_output.hidden_states->dtype(), + dispatch_output.hidden_states->device().getIndex()); if (!marlin_config.supported) { throw std::runtime_error("No lightop W16A16 Marlin MoE config found for this Hygon shape"); } @@ -255,7 +290,8 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, }; } w8a8_marlin_config = select_hygon_w8a8_marlin_config( - hidden_shape[0], hidden_size_, intermediate_size_per_partition_); + hidden_shape[0], hidden_size_, intermediate_size_per_partition_, + dispatch_output.hidden_states->device().getIndex()); if (!w8a8_marlin_config.supported) { throw std::runtime_error("No lightop W8A8 Marlin MoE config found for this Hygon shape"); } @@ -591,8 +627,10 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w8a8_marlin_core_sliced( dispatch_output.hidden_states->dtype(), dispatch_output.hidden_states->device()); MoeWorkspace slice_workspace; + const auto device_index = dispatch_output.hidden_states->device().getIndex(); const auto full_slice_config = select_hygon_w8a8_marlin_config( - kHygonW8A8MoeSliceTokens, hidden_size_, intermediate_size_per_partition_); + kHygonW8A8MoeSliceTokens, hidden_size_, intermediate_size_per_partition_, + device_index); if (!full_slice_config.supported) { throw std::runtime_error("No lightop W8A8 Marlin MoE config found for full Hygon slice"); } @@ -602,7 +640,8 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w8a8_marlin_core_sliced( const auto slice_config = slice_tokens == kHygonW8A8MoeSliceTokens ? full_slice_config : select_hygon_w8a8_marlin_config( - slice_tokens, hidden_size_, intermediate_size_per_partition_); + slice_tokens, hidden_size_, intermediate_size_per_partition_, + device_index); if (!slice_config.supported) { throw std::runtime_error("No lightop W8A8 Marlin MoE config found for sliced Hygon shape"); } From eb7122c980a3125fb844898ed87f4fe0fc14f86d Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Tue, 14 Jul 2026 06:11:39 +0800 Subject: [PATCH 07/14] fix(hygon): slice large W16A16 MoE prefills --- .../moe/runner/cuda_fused_moe_runner.cpp | 97 +++++++++++++++++++ .../moe/runner/cuda_fused_moe_runner.hpp | 5 + 2 files changed, 102 insertions(+) diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index 4b80e3524..ccec6154e 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -77,6 +77,7 @@ std::string normalize_hygon_gpu_target(std::string target, bool uppercase) { return target; } +constexpr size_t kHygonW16A16MoeSliceTokens = 16384; constexpr size_t kHygonW8A8MoeSliceTokens = 16384; enum class HygonMarlinModePolicy { @@ -270,6 +271,16 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, throw std::runtime_error("Hygon Marlin MoE runner requires hidden states [M, K]"); } if (weights.is_hygon_w16a16_marlin()) { + if (hidden_shape[0] > kHygonW16A16MoeSliceTokens) { + auto runner_output = run_hygon_w16a16_marlin_core_sliced( + dispatch_output, weights, workspace); + return CombineInput{ + CombineInputFormat::Standard, + runner_output.hidden_states, + dispatch_output.topk_output, + MoeRoutingMetadata{}, + }; + } marlin_config = select_hygon_w16a16_marlin_config( hidden_shape[0], hidden_size_, intermediate_size_per_partition_, dispatch_output.hidden_states->dtype(), @@ -456,6 +467,10 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w16a16_marlin_core( {num_local_experts_, intermediate_size_per_partition_ / 16, hidden_size_ * 16}); const size_t top_k = runner_input.topk_output.topk_ids->shape()[1]; const size_t num_tokens = runner_input.hidden_states->shape()[0]; + if (num_tokens > kHygonW16A16MoeSliceTokens) { + throw std::runtime_error( + "Hygon W16A16 Marlin MoE core requires inputs above 16384 tokens to be sliced"); + } const size_t cache13_required = num_tokens * top_k * std::max(intermediate_size_per_partition_ * 2, hidden_size_); const size_t cache2_required = num_tokens * top_k * intermediate_size_per_partition_; @@ -507,6 +522,88 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w16a16_marlin_core( }; } +CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w16a16_marlin_core_sliced( + const DispatchOutput &dispatch_output, + const MoeWeights &weights, + MoeWorkspace &workspace) const { + const auto &hidden_shape = dispatch_output.hidden_states->shape(); + if (hidden_shape.size() != 2) { + throw std::runtime_error("Hygon W16A16 sliced MoE runner requires hidden states [M, K]"); + } + const size_t num_tokens = hidden_shape[0]; + if (infinicore::context::isGraphRecording()) { + throw std::runtime_error("Hygon W16A16 sliced MoE runner cannot allocate/copy slice outputs during graph capture"); + } + + ensure_tensor( + workspace.fused_moe_output, + dispatch_output.hidden_states->shape(), + dispatch_output.hidden_states->dtype(), + dispatch_output.hidden_states->device()); + MoeWorkspace slice_workspace; + const auto activation_dtype = dispatch_output.hidden_states->dtype(); + const auto device_index = dispatch_output.hidden_states->device().getIndex(); + const auto full_slice_config = select_hygon_w16a16_marlin_config( + kHygonW16A16MoeSliceTokens, + hidden_size_, + intermediate_size_per_partition_, + activation_dtype, + device_index); + if (!full_slice_config.supported) { + throw std::runtime_error("No lightop W16A16 Marlin MoE config found for full Hygon slice"); + } + + size_t offset = 0; + while (offset < num_tokens) { + const size_t slice_tokens = std::min(kHygonW16A16MoeSliceTokens, num_tokens - offset); + const auto slice_config = slice_tokens == kHygonW16A16MoeSliceTokens + ? full_slice_config + : select_hygon_w16a16_marlin_config( + slice_tokens, + hidden_size_, + intermediate_size_per_partition_, + activation_dtype, + device_index); + if (!slice_config.supported) { + throw std::runtime_error("No lightop W16A16 Marlin MoE config found for sliced Hygon shape"); + } + + const auto hidden_slice = dispatch_output.hidden_states->narrow({{0, offset, slice_tokens}}); + const TopKOutput topk_slice{ + dispatch_output.topk_output.topk_weights->narrow({{0, offset, slice_tokens}}), + dispatch_output.topk_output.topk_ids->narrow({{0, offset, slice_tokens}}), + dispatch_output.topk_output.router_logits + ? dispatch_output.topk_output.router_logits->narrow({{0, offset, slice_tokens}}) + : infinicore::Tensor(), + }; + const DispatchOutput dispatch_slice{ + DispatchOutputFormat::Standard, + hidden_slice, + infinicore::Tensor(), + topk_slice, + infinicore::Tensor(), + }; + auto slice_input = prepare_runner_input( + dispatch_slice, + slice_workspace, + slice_config.gemm1.block_size_m); + auto slice_output = run_hygon_w16a16_marlin_core( + slice_input, + weights, + slice_workspace, + slice_config); + workspace.fused_moe_output + ->narrow({{0, offset, slice_tokens}}) + ->copy_from(slice_output.hidden_states); + + offset += slice_tokens; + } + + return CudaFusedMoeRunnerOutput{ + workspace.fused_moe_output, + }; +} + CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w8a8_marlin_core( const CudaFusedMoeRunnerInput &runner_input, const MoeWeights &weights, diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp index a6fc19845..b42be70d2 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp @@ -43,6 +43,11 @@ class CudaFusedMoeRunner final : public MoeRunnerCore { MoeWorkspace &workspace, const HygonW16A16MarlinRuntimeConfig &config) const; + CudaFusedMoeRunnerOutput run_hygon_w16a16_marlin_core_sliced( + const DispatchOutput &dispatch_output, + const MoeWeights &weights, + MoeWorkspace &workspace) const; + CudaFusedMoeRunnerOutput run_hygon_w8a8_marlin_core( const CudaFusedMoeRunnerInput &runner_input, const MoeWeights &weights, From 0a722288c42fa7884344a10c525398ba3ed2382e Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Thu, 23 Jul 2026 18:29:50 +0800 Subject: [PATCH 08/14] feat(eval): add token-level PPL scoring CLI --- csrc/engine/infer_engine.cpp | 45 +- csrc/engine/rank_worker.cpp | 125 +- csrc/engine/rank_worker.hpp | 15 + csrc/pybind11/engine/engine.hpp | 32 +- python/infinilm/infer_engine.py | 132 ++ test/engine/test_nll_validation.py | 145 +++ test/ppl/qwen3_235b/README.md | 200 +++ test/ppl/qwen3_235b/scripts/_gpu_guard.py | 100 ++ test/ppl/qwen3_235b/scripts/_ppl_common.py | 338 +++++ .../calculate_infinilm_precision_ppl.py | 84 ++ .../qwen3_235b/scripts/calculate_true_ppl.py | 305 +++++ .../infinilm/infinilm_ppl_Qwen3_235B.py | 334 +++++ .../scripts/prepare_ppl_corpus_Qwen3_235B.py | 328 +++++ .../scripts/transformers/_pytorch_runner.py | 1144 +++++++++++++++++ .../transformers/pytorch_ppl_Qwen3_235B.py | 491 +++++++ 15 files changed, 3777 insertions(+), 41 deletions(-) create mode 100644 test/engine/test_nll_validation.py create mode 100644 test/ppl/qwen3_235b/README.md create mode 100755 test/ppl/qwen3_235b/scripts/_gpu_guard.py create mode 100755 test/ppl/qwen3_235b/scripts/_ppl_common.py create mode 100644 test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py create mode 100755 test/ppl/qwen3_235b/scripts/calculate_true_ppl.py create mode 100755 test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py create mode 100755 test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py create mode 100755 test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py create mode 100755 test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 0568cec4d..6dcd07ccf 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -144,6 +144,44 @@ std::vector InferEngine::state_dict_keys() { //------------------------------------------------------ // forward //------------------------------------------------------ +void InferEngine::Input::validate() const { + if (!return_nll) { + if (labels.has_value()) { + throw std::invalid_argument("labels require return_nll=true"); + } + if (score_start != 0) { + throw std::invalid_argument("score_start requires return_nll=true"); + } + return; + } + + if (!input_ids.has_value() || !input_ids.value()) { + throw std::invalid_argument("NLL scoring requires input_ids"); + } + if (!labels.has_value() || !labels.value()) { + throw std::invalid_argument("NLL scoring requires labels"); + } + + const auto &ids = input_ids.value(); + const auto &target = labels.value(); + if (ids->dtype() != infinicore::DataType::I64 + || target->dtype() != infinicore::DataType::I64) { + throw std::invalid_argument("NLL input_ids and labels must use I64 dtype"); + } + if (ids->ndim() != 2 || target->ndim() != 2) { + throw std::invalid_argument("NLL input_ids and labels must be rank-2 tensors"); + } + if (ids->shape() != target->shape()) { + throw std::invalid_argument("NLL input_ids and labels must have identical shapes"); + } + if (ids->size(0) != 1) { + throw std::invalid_argument("NLL scoring currently requires batch_size=1"); + } + if (score_start >= ids->size(1)) { + throw std::invalid_argument("NLL score_start must select at least one token"); + } +} + infinilm::InfinilmModel::Input InferEngine::Input::to_model_input(infinicore::Device device) const { @@ -182,8 +220,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; - + input.last_token_only = !sample_all_positions && !return_nll; infinilm::global_state::get_forward_context().attn_metadata = { input.past_sequence_lengths, input.total_sequence_lengths, @@ -205,6 +242,10 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { } InferEngine::Output InferEngine::forward(const InferEngine::Input &input) { + // Validate before dispatch so malformed NLL requests cannot fail only one + // rank and leave the remaining workers waiting at a collective. + input.validate(); + // Trigger each worker to run inference for (auto &worker : workers_) { worker->run(input); diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index acb2d24c2..276fc6d09 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -418,9 +418,13 @@ void RankWorker::thread_loop() { 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 && rank_info_.pp_size == 1) { + // Full-position and NLL runs need eager mode because generation + // graphs return last-token logits and omit hidden states. PP graph + // compilation is not supported yet. + if (!local_args.sample_all_positions + && !local_args.return_nll + && compiler_ != nullptr + && rank_info_.pp_size == 1) { auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device::cpu())); if (graph != nullptr && output != nullptr) { graph->run(); @@ -435,6 +439,11 @@ void RankWorker::thread_loop() { hidden_states = model_output.hidden_states; } + if (local_args.return_nll && rank_info_.pp_size > 1) { + throw std::runtime_error( + "NLL scoring with pipeline parallelism is not supported"); + } + if (rank_info_.pp_size > 1 && rank_info_.pp_stage + 1 != rank_info_.pp_size) { infinicore::Tensor output_ids; if (rank_info_.pp_stage == 0 && rank_info_.tp_rank == 0) { @@ -464,54 +473,96 @@ void RankWorker::thread_loop() { continue; } - // Random sampling (rank 0 only) + // Sampling and scoring both consume replicated full-vocabulary + // logits, so only rank 0 needs to materialize the result. if (rank_info_.tp_rank == 0) { - 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()}; + if (logits_shape.size() != 3) { + throw std::runtime_error("InferEngine expected rank-3 logits"); + } 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(); + if (local_args.return_nll) { + auto labels = local_args.labels.value()->to(rank_info_.device); + if (labels->dtype() != infinicore::DataType::I64 + || labels->ndim() != 2 + || labels->size(0) != batch_size + || labels->size(1) != total_len) { + throw std::runtime_error( + "NLL labels must be I64 with shape [batch, sequence]"); + } - 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)}; + const auto score_len = total_len - local_args.score_start; + auto score_logits = logits->narrow( + {{1, local_args.score_start, score_len}}); + auto score_labels = labels->narrow( + {{1, local_args.score_start, score_len}}); + + auto token_nll = infinicore::Tensor::empty( + score_labels->shape(), + infinicore::DataType::F32, + rank_info_.device); + infinicore::op::cross_entropy_( + token_nll, score_logits, score_labels); + token_nll = token_nll->to(infinicore::Device::cpu()); + infinicore::context::syncStream(); + output_ = Output{ + infinicore::Tensor{}, + infinicore::Tensor{}, + infinicore::Tensor{}, + token_nll, + score_len, + }; + } else { + auto temperature{local_args.temperature}; + auto top_p{local_args.top_p}; + auto top_k{local_args.top_k}; + 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 = 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({})}; + float random_val = std::uniform_real_distribution(0, 1)(rng_); + infinicore::op::random_sample_( + out, score, random_val, top_p, top_k, temperature); + } - for (size_t i{0}; i < n_out; ++i) { - size_t score_idx = i; - if (!sample_all_positions) { - score_idx = total_len == n_req - ? i - : static_cast(input_offsets[i + 1] - 1); + if (rank_info_.pp_size > 1) { + infinicore::op::distributed::send( + output_ids, + 0, + rank_info_.world_comm); } - 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); - } - if (rank_info_.pp_size > 1) { - infinicore::op::distributed::send( + // Tensor::to(CPU) uses the synchronous D2H contract. + output_ids = output_ids->to(infinicore::Device::cpu()); + output_ = Output{ output_ids, + logits, + hidden_states, + infinicore::Tensor{}, 0, - rank_info_.world_comm); + }; } - - output_ids = output_ids->to(infinicore::Device::cpu()); - - infinicore::context::syncStream(); - - auto out{Output{output_ids, logits, hidden_states}}; - - output_ = std::move(out); } - job_done_ = true; } cv_.notify_all(); diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index d396ef6f1..19f2f1f28 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -73,12 +73,25 @@ class RankWorker { /// Sample logits at every packed input position instead of one token per request. bool sample_all_positions{false}; + /// Shifted causal-LM labels. Present only for explicit NLL scoring. + std::optional labels; + + /// First logits/label position included in NLL scoring. + size_t score_start{0}; + + /// Compute token NLL instead of sampling output IDs. + bool return_nll{false}; + float temperature{1}; int top_k{50}; float top_p{1}; + /// Validate invariants shared by Python and native callers before a + /// request is dispatched to any rank worker. + void validate() const; + infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const; }; @@ -86,6 +99,8 @@ class RankWorker { infinicore::Tensor output_ids; infinicore::Tensor logits; infinicore::Tensor hidden_states; + infinicore::Tensor nll; + size_t scored_tokens{0}; }; RankWorker(std::shared_ptr infinilm_config, diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index 44e412619..1b7268c6e 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -158,6 +158,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 labels, py::kwargs kwargs) { InferEngine::Input input{ std::move(input_ids), @@ -178,6 +179,7 @@ inline void bind_infer_engine(py::module &m) { std::move(visual_token_ranges), std::move(target_hidden_states), sample_all_positions, + std::move(labels), }; // Explicit defaults @@ -190,6 +192,8 @@ inline void bind_infer_engine(py::module &m) { "temperature", "top_p", "top_k", + "score_start", + "return_nll", }; for (auto &item : kwargs) { @@ -206,9 +210,24 @@ 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 == "score_start") { + if (py::isinstance(item.second)) { + throw py::type_error("score_start must be an integer, not bool"); + } + const auto score_start = py::cast(item.second); + if (score_start < 0) { + throw py::value_error("score_start must be non-negative"); + } + input.score_start = static_cast(score_start); + } else if (key == "return_nll") { + if (!py::isinstance(item.second)) { + throw py::type_error("return_nll must be a bool"); + } + input.return_nll = py::cast(item.second); } } + input.validate(); return input; }), py::arg("input_ids") = std::nullopt, @@ -228,7 +247,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("labels") = 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) @@ -247,6 +267,9 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) + .def_readwrite("labels", &InferEngine::Input::labels) + .def_readwrite("score_start", &InferEngine::Input::score_start) + .def_readwrite("return_nll", &InferEngine::Input::return_nll) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) .def_readwrite("top_p", &InferEngine::Input::top_p); @@ -254,7 +277,12 @@ inline void bind_infer_engine(py::module &m) { py::class_(infer_engine, "Output") .def_readwrite("output_ids", &InferEngine::Output::output_ids, "Sampled token IDs") .def_readwrite("logits", &InferEngine::Output::logits, "Raw logits tensor") - .def_readwrite("hidden_states", &InferEngine::Output::hidden_states, "Raw hidden states tensor"); + .def_readwrite("hidden_states", &InferEngine::Output::hidden_states, "Raw hidden states tensor") + .def_readwrite("nll", &InferEngine::Output::nll, "Per-token NLL tensor") + .def_readwrite( + "scored_tokens", + &InferEngine::Output::scored_tokens, + "Number of scored tokens"); } } // namespace infinilm::engine diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 9d78dd809..8f083a0ee 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -1,4 +1,5 @@ import json +import operator import os import time from dataclasses import dataclass @@ -18,6 +19,45 @@ } +def _validate_nll_score_inputs(input_ids, labels, score_start): + """Validate an explicit batch-1, shifted-token NLL request.""" + for name, tensor in (("input_ids", input_ids), ("labels", labels)): + if tensor is None: + raise TypeError(f"{name} must be an infinicore tensor") + missing = [ + attr + for attr in ("ndim", "shape", "dtype", "_underlying") + if not hasattr(tensor, attr) + ] + if missing: + raise TypeError( + f"{name} must be an infinicore tensor; missing {', '.join(missing)}" + ) + if tensor.dtype != infinicore.int64: + raise ValueError(f"{name} must use infinicore.int64 dtype") + if tensor.ndim != 2: + raise ValueError(f"{name} must be a rank-2 tensor") + + input_shape = tuple(input_ids.shape) + label_shape = tuple(labels.shape) + if input_shape != label_shape: + raise ValueError("input_ids and labels must have identical shapes") + if input_shape[0] != 1: + raise ValueError("score_nll currently requires batch_size=1") + + if isinstance(score_start, bool): + raise TypeError("score_start must be an integer, not bool") + try: + score_start = operator.index(score_start) + except TypeError as error: + raise TypeError("score_start must be an integer") from error + + seq_len = input_shape[1] + if score_start < 0 or score_start >= seq_len: + raise ValueError("score_start must select at least one token") + return seq_len, score_start + + def _apply_torch_dtype_defaults(config: dict) -> dict: if config.get("torch_dtype") is None: config["torch_dtype"] = config.get("dtype") or _MODEL_DEFAULTS.get( @@ -652,6 +692,98 @@ def generate( return output_ids + def score_nll(self, input_ids, labels, *, score_start=0): + """Return summed shifted-token NLL and token count for a batch-1 window. + + This explicit evaluation path bypasses graph replay and never changes the + behavior of ``forward``/``generate``. ``input_ids`` and ``labels`` must + have the same ``[1, sequence]`` shape; callers perform the causal shift. + """ + try: + seq_len, score_start = _validate_nll_score_inputs( + input_ids, labels, score_start + ) + + block_tables = None + slot_mapping = None + if self.enable_paged_attn: + cache_config = self.get_cache_config() + if cache_config is None: + raise RuntimeError("paged attention requires a cache configuration") + paged_block_size = cache_config.block_size() + max_blocks_per_batch = ( + seq_len + paged_block_size - 1 + ) // paged_block_size + if max_blocks_per_batch > cache_config.num_blocks(): + raise ValueError( + "NLL sequence requires more paged KV-cache blocks than " + "the current cache configuration provides" + ) + block_tables = infinicore.from_list( + [list(range(max_blocks_per_batch))], + dtype=infinicore.int32, + ) + slot_mapping = infinicore.from_list( + list(range(seq_len)), dtype=infinicore.int64 + ) + position_ids = infinicore.from_list( + list(range(seq_len)), dtype=infinicore.int64 + ) + else: + position_ids = infinicore.from_list( + [list(range(seq_len))], dtype=infinicore.int64 + ) + past_kv_lengths = infinicore.from_list([0], dtype=infinicore.int32) + total_kv_lengths = infinicore.from_list( + [seq_len], dtype=infinicore.int32 + ) + cu_seqlens = infinicore.from_list( + [0, seq_len], dtype=infinicore.int32 + ) + input_offsets = infinicore.from_list( + [0, seq_len], dtype=infinicore.int32 + ) + + output = super().forward( + super().Input( + input_ids._underlying, + position_ids=position_ids._underlying, + past_sequence_lengths=past_kv_lengths._underlying, + total_sequence_lengths=total_kv_lengths._underlying, + input_offsets=input_offsets._underlying, + cu_seqlens=cu_seqlens._underlying, + block_tables=( + block_tables._underlying + if block_tables is not None + else None + ), + slot_mapping=( + slot_mapping._underlying + if slot_mapping is not None + else None + ), + labels=labels._underlying, + score_start=score_start, + return_nll=True, + ) + ) + token_nll = infinicore.Tensor(output.nll).to_numpy() + scored_tokens = int(output.scored_tokens) + expected_scored_tokens = seq_len - score_start + if scored_tokens != expected_scored_tokens: + raise RuntimeError( + "score_nll returned an invalid scored-token count: " + f"expected {expected_scored_tokens}, got {scored_tokens}" + ) + if token_nll.size != expected_scored_tokens: + raise RuntimeError( + "score_nll returned a token-loss vector with an invalid size" + ) + return float(token_nll.astype("float64").sum()), scored_tokens + except BaseException as e: + handle_oom_and_exit(e) + raise + def reset_cache(self, cache_config): infinicore.sync_device() self.enable_paged_attn = isinstance(cache_config, PagedKVCacheConfig) diff --git a/test/engine/test_nll_validation.py b/test/engine/test_nll_validation.py new file mode 100644 index 000000000..cc2f64ab8 --- /dev/null +++ b/test/engine/test_nll_validation.py @@ -0,0 +1,145 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + + +class FakeTensor: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.ndim = len(self.shape) + self.dtype = dtype + self._underlying = object() + + +@pytest.fixture +def validator(monkeypatch): + """Load the pure validator without importing the hardware runtime.""" + package_root = Path(__file__).resolve().parents[2] / "python" / "infinilm" + int64_dtype = object() + + infinilm_package = types.ModuleType("infinilm") + infinilm_package.__path__ = [str(package_root)] + monkeypatch.setitem(sys.modules, "infinilm", infinilm_package) + + fake_infinicore = types.ModuleType("infinicore") + fake_infinicore.int64 = int64_dtype + fake_infinicore.Tensor = type("Tensor", (), {}) + monkeypatch.setitem(sys.modules, "infinicore", fake_infinicore) + + cache_module = types.ModuleType("infinilm.cache") + cache_module.PagedKVCacheConfig = type("PagedKVCacheConfig", (), {}) + monkeypatch.setitem(sys.modules, "infinilm.cache", cache_module) + + distributed_module = types.ModuleType("infinilm.distributed") + distributed_module.DistConfig = type("DistConfig", (), {}) + monkeypatch.setitem(sys.modules, "infinilm.distributed", distributed_module) + + native_engine = type("InferEngine", (), {}) + lib_module = types.ModuleType("infinilm.lib") + lib_module._infinilm = types.SimpleNamespace(InferEngine=native_engine) + monkeypatch.setitem(sys.modules, "infinilm.lib", lib_module) + + exception_module = types.ModuleType("infinilm.exception_utils") + exception_module.handle_oom_and_exit = lambda error: None + monkeypatch.setitem(sys.modules, "infinilm.exception_utils", exception_module) + + modeling_module = types.ModuleType("infinilm.modeling_utils") + modeling_module.parse_dtype = lambda dtype: dtype + monkeypatch.setitem(sys.modules, "infinilm.modeling_utils", modeling_module) + + module_name = "infinilm.infer_engine" + module_path = package_root / "infer_engine.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + assert spec.loader is not None + spec.loader.exec_module(module) + + return module._validate_nll_score_inputs, int64_dtype + + +def make_tensor(shape, int64_dtype, dtype=None): + return FakeTensor(shape, int64_dtype if dtype is None else dtype) + + +def test_validate_nll_score_inputs_accepts_valid_window(validator): + validate, int64_dtype = validator + input_ids = make_tensor((1, 8), int64_dtype) + labels = make_tensor((1, 8), int64_dtype) + + assert validate(input_ids, labels, 3) == (8, 3) + + +@pytest.mark.parametrize("name", ["input_ids", "labels"]) +def test_validate_nll_score_inputs_requires_tensor_protocol(validator, name): + validate, int64_dtype = validator + tensors = { + "input_ids": make_tensor((1, 8), int64_dtype), + "labels": make_tensor((1, 8), int64_dtype), + } + tensors[name] = object() + + with pytest.raises(TypeError, match=name): + validate(tensors["input_ids"], tensors["labels"], 0) + + +@pytest.mark.parametrize("name", ["input_ids", "labels"]) +def test_validate_nll_score_inputs_requires_int64(validator, name): + validate, int64_dtype = validator + tensors = { + "input_ids": make_tensor((1, 8), int64_dtype), + "labels": make_tensor((1, 8), int64_dtype), + } + tensors[name] = make_tensor((1, 8), int64_dtype, dtype=object()) + + with pytest.raises(ValueError, match=f"{name} must use infinicore.int64"): + validate(tensors["input_ids"], tensors["labels"], 0) + + +@pytest.mark.parametrize( + ("input_shape", "label_shape", "message"), + [ + ((8,), (8,), "rank-2"), + ((1, 8), (1, 7), "identical shapes"), + ((2, 8), (2, 8), "batch_size=1"), + ], +) +def test_validate_nll_score_inputs_rejects_invalid_shapes( + validator, input_shape, label_shape, message +): + validate, int64_dtype = validator + with pytest.raises(ValueError, match=message): + validate( + make_tensor(input_shape, int64_dtype), + make_tensor(label_shape, int64_dtype), + 0, + ) + + +@pytest.mark.parametrize("score_start", [-1, 8]) +def test_validate_nll_score_inputs_rejects_empty_score_range( + validator, score_start +): + validate, int64_dtype = validator + with pytest.raises(ValueError, match="select at least one token"): + validate( + make_tensor((1, 8), int64_dtype), + make_tensor((1, 8), int64_dtype), + score_start, + ) + + +@pytest.mark.parametrize("score_start", [True, 1.5, "1"]) +def test_validate_nll_score_inputs_requires_integer_score_start( + validator, score_start +): + validate, int64_dtype = validator + with pytest.raises(TypeError, match="score_start must be an integer"): + validate( + make_tensor((1, 8), int64_dtype), + make_tensor((1, 8), int64_dtype), + score_start, + ) diff --git a/test/ppl/qwen3_235b/README.md b/test/ppl/qwen3_235b/README.md new file mode 100644 index 000000000..559fd7c1c --- /dev/null +++ b/test/ppl/qwen3_235b/README.md @@ -0,0 +1,200 @@ +# Qwen3-235B true PPL CLI + +This directory contains reproducible token-level perplexity tools for: + +- Transformers BF16 on TP8 +- InfiniLM BF16 on TP8 +- InfiniLM W8A8 on TP8 + +The runners consume the same frozen token manifest and calculate causal, +shifted-token cross entropy: + +```text +mean_nll = sum(-log p(x_t | x_&1 | tee "$LOG_DIR/infinilm_w8a8_smoke.log" +rc=${PIPESTATUS[0]} +echo "INFINILM_W8A8_SMOKE_EXIT_CODE=$rc" +hy-smi --showpids +``` + +## Full WikiText-2 runs + +`--max-scored-tokens 0` scores every target token in the manifest. Use the same +`window`, `stride` and `max-scored-tokens` values for every backend. + +Transformers BF16: + +```bash +set -o pipefail +timeout --signal=TERM --kill-after=60s 21600s \ + python -u "$PPL_ROOT/scripts/transformers/pytorch_ppl_Qwen3_235B.py" \ + --model "$MODEL_BF16" \ + --token-manifest "$TOKEN_MANIFEST" \ + --window 256 \ + --stride 128 \ + --max-scored-tokens 0 \ + --tp-size 8 \ + --attention eager \ + --json-output "$LOG_DIR/transformers_bf16_full.json" \ + 2>&1 | tee "$LOG_DIR/transformers_bf16_full.log" +``` + +The Transformers entry point launches `torchrun` itself. Do not wrap it in a +second `torchrun` command. Eager attention is the validated Hygon path. + +InfiniLM BF16: + +```bash +set -o pipefail +timeout --signal=TERM --kill-after=60s 21600s \ + python -u "$PPL_ROOT/scripts/infinilm/infinilm_ppl_Qwen3_235B.py" \ + --model "$MODEL_BF16" \ + --token-manifest "$TOKEN_MANIFEST" \ + --window 256 \ + --stride 128 \ + --max-scored-tokens 0 \ + --tp-size 8 \ + --attention flash-attn \ + --json-output "$LOG_DIR/infinilm_bf16_full.json" \ + 2>&1 | tee "$LOG_DIR/infinilm_bf16_full.log" +``` + +InfiniLM W8A8: + +```bash +set -o pipefail +timeout --signal=TERM --kill-after=60s 21600s \ + python -u "$PPL_ROOT/scripts/infinilm/infinilm_ppl_Qwen3_235B.py" \ + --model "$MODEL_W8A8" \ + --token-manifest "$TOKEN_MANIFEST" \ + --window 256 \ + --stride 128 \ + --max-scored-tokens 0 \ + --tp-size 8 \ + --attention flash-attn \ + --json-output "$LOG_DIR/infinilm_w8a8_full.json" \ + 2>&1 | tee "$LOG_DIR/infinilm_w8a8_full.log" +``` + +For a bounded formal run, replace `0` with the same positive token count in all +three commands, for example `10240`. + +## Compare results + +Transformers BF16 versus InfiniLM W8A8: + +```bash +python -u "$PPL_ROOT/scripts/calculate_true_ppl.py" \ + --inputs \ + "$LOG_DIR/transformers_bf16_full.json" \ + "$LOG_DIR/infinilm_w8a8_full.json" \ + --max-ppl-increase-percent 20 \ + --json-out "$LOG_DIR/ppl_transformers_vs_w8a8.json" +``` + +InfiniLM BF16 versus InfiniLM W8A8: + +```bash +python -u "$PPL_ROOT/scripts/calculate_infinilm_precision_ppl.py" \ + --inputs \ + "$LOG_DIR/infinilm_bf16_full.json" \ + "$LOG_DIR/infinilm_w8a8_full.json" \ + --max-ppl-increase-percent 20 \ + --json-out "$LOG_DIR/ppl_bf16_vs_w8a8.json" +``` + +Exit code `0` means the configured PPL increase threshold passed, `1` means it +failed, and `2` means the input files are invalid or describe different +workloads. + +## Scope + +PPL is a quality test. InfiniLM intentionally disables graph only for the +explicit `score_nll` path because it must retain full token logits/losses. +Normal generation and formal performance tests keep their existing graph path. +Do not report PPL scoring throughput as inference performance. diff --git a/test/ppl/qwen3_235b/scripts/_gpu_guard.py b/test/ppl/qwen3_235b/scripts/_gpu_guard.py new file mode 100755 index 000000000..1ceaedddb --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/_gpu_guard.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +"""Fail closed unless all eight Hygon devices are idle.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + + +EXPECTED_DEVICES = set(range(8)) +SMI_TIMEOUT_SECONDS = 60 + + +def _local_gpu_processes() -> list[str]: + users: list[str] = [] + own_pid = os.getpid() + for process_dir in Path("/proc").glob("[0-9]*"): + try: + pid = int(process_dir.name) + except ValueError: + continue + if pid == own_pid: + continue + try: + targets = [entry.resolve() for entry in (process_dir / "fd").iterdir()] + except OSError: + continue + if not any( + str(target) == "/dev/kfd" or str(target).startswith("/dev/dri/renderD") + for target in targets + ): + continue + try: + command = (process_dir / "cmdline").read_bytes().replace(b"\0", b" ").decode( + "utf-8", errors="replace" + ).strip() + except OSError: + command = "" + users.append(f"pid={pid} command={command or '[unknown]'}") + return sorted(users) + + +def require_idle_gpu() -> None: + try: + result = subprocess.run( + ["hy-smi"], + check=False, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + timeout=SMI_TIMEOUT_SECONDS, + ) + except (OSError, subprocess.TimeoutExpired) as error: + print(f"拒绝启动：hy-smi 空闲检查失败：{error}", file=sys.stderr) + raise SystemExit(90) from error + if result.returncode != 0: + print( + f"拒绝启动：hy-smi 退出码为 {result.returncode}。\n{result.stdout}", + file=sys.stderr, + ) + raise SystemExit(90) + + utilization: dict[int, tuple[float, float]] = {} + for line in result.stdout.splitlines(): + fields = line.split() + if ( + len(fields) >= 7 + and fields[0].isdigit() + and fields[5].endswith("%") + and fields[6].endswith("%") + ): + try: + utilization[int(fields[0])] = ( + float(fields[5][:-1]), + float(fields[6][:-1]), + ) + except ValueError: + continue + if set(utilization) != EXPECTED_DEVICES: + print( + "拒绝启动：hy-smi 未完整报告 0-7 号设备。\n" + result.stdout, + file=sys.stderr, + ) + raise SystemExit(90) + + busy_devices = { + device: values + for device, values in utilization.items() + if values[0] > 0.0 or values[1] > 0.0 + } + local_users = _local_gpu_processes() + if busy_devices or local_users: + print( + "拒绝启动：GPU 未完全空闲；" + f"设备占用={busy_devices}，容器内进程={local_users}。\n{result.stdout}", + file=sys.stderr, + ) + raise SystemExit(90) diff --git a/test/ppl/qwen3_235b/scripts/_ppl_common.py b/test/ppl/qwen3_235b/scripts/_ppl_common.py new file mode 100755 index 000000000..bc720547f --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/_ppl_common.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 +"""Shared, deterministic corpus and sliding-window helpers for true PPL tests.""" + +from __future__ import annotations + +import ast +import array +import hashlib +import json +import operator +import re +import struct +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Iterator, Sequence + + +CORPUS_SCHEMA = "qw235_ppl_token_ids_v1" +SCORING_METHOD = ( + "sliding_window_shifted_cross_entropy_fp32_compute_fp64_accumulation" +) +SHA256_RE = re.compile(r"^[0-9a-f]{64}$") + + +def canonical_json_bytes(value: object) -> bytes: + return json.dumps( + value, ensure_ascii=True, sort_keys=True, separators=(",", ":") + ).encode("ascii") + + +def _canonical_int_sequence_sha256(values: Iterable[int], label: str) -> str: + """Hash an integer sequence exactly like compact JSON ``[1,2,3]``.""" + digest = hashlib.sha256() + digest.update(b"[") + for index, value in enumerate(values): + if isinstance(value, bool): + raise ValueError(f"{label}[{index}] 不是非负整数") + try: + parsed = operator.index(value) + except TypeError as error: + raise ValueError(f"{label}[{index}] 不是非负整数") from error + if parsed < 0: + raise ValueError(f"{label}[{index}] 不是非负整数：{value!r}") + if index: + digest.update(b",") + digest.update(str(parsed).encode("ascii")) + digest.update(b"]") + return digest.hexdigest() + + +def canonical_token_ids_sha256(token_ids: Iterable[int]) -> str: + return _canonical_int_sequence_sha256(token_ids, "token_ids") + + +def canonical_indices_sha256(indices: Iterable[int]) -> str: + return _canonical_int_sequence_sha256(indices, "indices") + + +@dataclass(frozen=True) +class PplCorpusManifest: + path: Path + payload: dict[str, Any] + token_ids: tuple[int, ...] + manifest_sha256: str + token_ids_sha256: str + + @property + def token_count(self) -> int: + return len(self.token_ids) + + +@dataclass(frozen=True) +class SlidingWindow: + """One causal-LM window using half-open global token index ranges. + + ``token_start:token_end`` is model input. Targets in + ``score_start:score_end`` are scored. ``prediction_*`` select the matching + logits before the causal shift, while ``target_*`` select labels locally. + """ + + index: int + token_start: int + token_end: int + score_start: int + score_end: int + token_ids: tuple[int, ...] + + @property + def scored_token_count(self) -> int: + return self.score_end - self.score_start + + @property + def prediction_start(self) -> int: + return self.score_start - self.token_start - 1 + + @property + def prediction_end(self) -> int: + return self.score_end - self.token_start - 1 + + @property + def target_start(self) -> int: + return self.score_start - self.token_start + + @property + def target_end(self) -> int: + return self.score_end - self.token_start + + +def _required_positive_int(payload: dict[str, Any], key: str, path: Path) -> int: + value = payload.get(key) + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{path} 的 {key} 必须是正整数") + parsed = value + if parsed <= 0: + raise ValueError(f"{path} 的 {key} 必须是正整数") + return parsed + + +def _required_sha(payload: dict[str, Any], key: str, path: Path) -> str: + value = str(payload.get(key, "")).lower() + if not SHA256_RE.fullmatch(value): + raise ValueError(f"{path} 的 {key} 不是有效 SHA256") + return value + + +def _load_npy(manifest_path: Path, relative_name: object) -> list[int]: + relative = Path(str(relative_name)) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"{manifest_path} 的 token_ids_file 必须是安全相对路径") + base = manifest_path.parent.resolve() + token_path = (base / relative).resolve() + try: + token_path.relative_to(base) + except ValueError as error: + raise ValueError(f"token_ids_file 越出 manifest 目录：{relative}") from error + try: + import numpy as np + except ImportError: + return _load_int64_npy_without_numpy(token_path) + try: + array = np.load(token_path, allow_pickle=False) + except FileNotFoundError: + raise ValueError(f"token_ids_file 不存在：{token_path}") from None + if array.ndim != 1 or array.dtype.kind not in "iu": + raise ValueError(f"{token_path} 必须是一维整数 .npy 数组") + return [int(value) for value in array.tolist()] + + +def _load_int64_npy_without_numpy(path: Path) -> list[int]: + try: + with path.open("rb") as handle: + if handle.read(6) != b"\x93NUMPY": + raise ValueError(f"{path} 不是有效 .npy 文件") + version = handle.read(2) + if version == b"\x01\x00": + header_length = struct.unpack(" None: + """Write a portable NumPy v1.0, one-dimensional little-endian int64 file.""" + output = Path(path) + values = list(token_ids) + # Validate before creating a partial file. + canonical_token_ids_sha256(values) + header_text = repr( + {"descr": " 65535: + raise ValueError(".npy header 超过 v1.0 长度限制") + output.parent.mkdir(parents=True, exist_ok=True) + packed = array.array("q", (int(value) for value in values)) + if packed.itemsize != 8: + raise RuntimeError("当前 Python 平台的 signed long long 不是 64 bit") + if sys.byteorder != "little": + packed.byteswap() + with output.open("wb") as handle: + handle.write(b"\x93NUMPY") + handle.write(b"\x01\x00") + handle.write(struct.pack(" PplCorpusManifest: + """Load and fully verify an inline or relative-``.npy`` corpus manifest.""" + manifest_path = Path(path) + try: + payload = json.loads(manifest_path.read_text(encoding="utf-8")) + except FileNotFoundError: + raise ValueError(f"PPL manifest 不存在：{manifest_path}") from None + except json.JSONDecodeError as error: + raise ValueError(f"PPL manifest JSON 无效：{manifest_path}: {error.msg}") from error + if not isinstance(payload, dict): + raise ValueError(f"PPL manifest 必须是 JSON 对象：{manifest_path}") + if payload.get("schema") != CORPUS_SCHEMA: + raise ValueError( + f"{manifest_path} schema 必须为 {CORPUS_SCHEMA!r}，" + f"实际为 {payload.get('schema')!r}" + ) + manifest_hash = _required_sha(payload, "manifest_sha256", manifest_path) + semantic_payload = dict(payload) + semantic_payload.pop("manifest_sha256", None) + calculated_manifest_hash = hashlib.sha256( + canonical_json_bytes(semantic_payload) + ).hexdigest() + if manifest_hash != calculated_manifest_hash: + raise ValueError(f"{manifest_path} 的 manifest_sha256 校验失败") + for key in ("source_sha256", "tokenizer_sha256"): + _required_sha(payload, key, manifest_path) + + has_inline = "token_ids" in payload + has_file = "token_ids_file" in payload + if has_inline == has_file: + raise ValueError( + f"{manifest_path} 必须且只能包含 token_ids 或 token_ids_file 之一" + ) + if has_inline: + raw_ids = payload["token_ids"] + if not isinstance(raw_ids, list): + raise ValueError(f"{manifest_path} 的 token_ids 必须是数组") + token_ids = list(raw_ids) + else: + token_ids = _load_npy(manifest_path, payload["token_ids_file"]) + + # The canonical hash validates type, integrality, sign, order and contents. + calculated_token_hash = canonical_token_ids_sha256(token_ids) + expected_token_hash = _required_sha( + payload, "token_ids_sha256", manifest_path + ) + if calculated_token_hash != expected_token_hash: + raise ValueError(f"{manifest_path} 的 token_ids_sha256 校验失败") + token_count = _required_positive_int(payload, "token_count", manifest_path) + if token_count != len(token_ids) or token_count < 2: + raise ValueError( + f"{manifest_path} token_count={token_count}，实际 token 数={len(token_ids)}" + ) + return PplCorpusManifest( + path=manifest_path, + payload=payload, + token_ids=tuple(int(value) for value in token_ids), + manifest_sha256=manifest_hash, + token_ids_sha256=expected_token_hash, + ) + + +def iter_sliding_windows( + token_ids: Sequence[int], + window_size: int, + stride: int, + max_scored_tokens: int | None = None, +) -> Iterator[SlidingWindow]: + """Yield windows that score global token indices ``1..N-1`` exactly once. + + The first window scores ``1:end``. Every later window scores only + ``previous_end:end``; overlapped prefix tokens provide context but are not + counted again. ``stride`` must be smaller than ``window_size`` so the first + new target in every later window retains its immediately preceding token. + """ + if isinstance(window_size, bool) or not isinstance(window_size, int): + raise ValueError("window_size 必须是整数") + if isinstance(stride, bool) or not isinstance(stride, int): + raise ValueError("stride 必须是整数") + if window_size < 2: + raise ValueError("window_size 必须至少为 2") + if stride < 1 or stride >= window_size: + raise ValueError("stride 必须满足 1 <= stride < window_size") + if len(token_ids) < 2: + raise ValueError("至少需要 2 个 token 才能计算 PPL") + if max_scored_tokens is not None: + if ( + isinstance(max_scored_tokens, bool) + or not isinstance(max_scored_tokens, int) + or max_scored_tokens < 1 + ): + raise ValueError("max_scored_tokens 必须是正整数") + + score_limit = len(token_ids) + if max_scored_tokens is not None: + score_limit = min(score_limit, 1 + max_scored_tokens) + previous_end = 1 + index = 0 + while previous_end < score_limit: + if index == 0: + token_end = min(score_limit, window_size) + token_start = 0 + score_start = 1 + else: + token_end = min(score_limit, previous_end + stride) + token_start = max(0, token_end - window_size) + score_start = previous_end + window = SlidingWindow( + index=index, + token_start=token_start, + token_end=token_end, + score_start=score_start, + score_end=token_end, + token_ids=tuple(int(value) for value in token_ids[token_start:token_end]), + ) + if window.prediction_start < 0 or window.prediction_end > len(window.token_ids): + raise AssertionError("内部错误：滑窗缺少 causal predecessor") + yield window + previous_end = token_end + index += 1 diff --git a/test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py b/test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py new file mode 100644 index 000000000..5bd49e86c --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Compare InfiniLM BF16 and W8A8 true-PPL result JSON files.""" + +from __future__ import annotations + +import argparse +import json +import os +from dataclasses import asdict +from pathlib import Path +from typing import Any, Sequence + +from calculate_true_ppl import _validate_same_workload, load_result + + +SCHEMA = "qwen3_235b_infinilm_precision_ppl_comparison/v1" + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--inputs", nargs=2, type=Path, required=True) + parser.add_argument("--max-ppl-increase-percent", type=float, default=20.0) + parser.add_argument("--json-out", type=Path, required=True) + args = parser.parse_args(argv) + if args.max_ppl_increase_percent < 0: + parser.error("--max-ppl-increase-percent must be non-negative") + return args + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") + temporary.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + try: + results = [load_result(path) for path in args.inputs] + if any(result.backend != "infinilm" for result in results): + raise ValueError("both inputs must be InfiniLM results") + precisions: list[str] = [] + for path in args.inputs: + payload = json.loads(path.read_text(encoding="utf-8")) + precision = str(payload.get("precision", "")).strip().upper() + if precision not in {"BF16", "W8A8"}: + raise ValueError(f"{path} has invalid precision: {precision!r}") + precisions.append(precision) + by_precision = dict(zip(precisions, results, strict=True)) + if set(by_precision) != {"BF16", "W8A8"}: + raise ValueError("inputs must contain one BF16 and one W8A8 result") + baseline = by_precision["BF16"] + candidate = by_precision["W8A8"] + _validate_same_workload(baseline, candidate) + increase = (candidate.ppl / baseline.ppl - 1.0) * 100.0 + threshold = float(args.max_ppl_increase_percent) + passed = increase <= threshold + payload = { + "schema": SCHEMA, + "status": "PASS" if passed else "FAIL", + "baseline": asdict(baseline), + "candidate": asdict(candidate), + "ppl_increase_percent": increase, + "max_ppl_increase_percent": threshold, + } + _atomic_json(args.json_out, payload) + except (OSError, ValueError, KeyError, TypeError, json.JSONDecodeError) as error: + print(f"PPL comparison error: {error}") + return 2 + + print(f"InfiniLM BF16 PPL: {baseline.ppl:.6f}") + print(f"InfiniLM W8A8 PPL: {candidate.ppl:.6f}") + print(f"W8A8 PPL increase: {increase:.2f}%") + print(f"Quality threshold: <= {threshold:.2f}%") + print(f"Result: {'PASS' if passed else 'FAIL'}") + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/calculate_true_ppl.py b/test/ppl/qwen3_235b/scripts/calculate_true_ppl.py new file mode 100755 index 000000000..cfefc2324 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/calculate_true_ppl.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +"""严格比较 Transformers 与 InfiniLM 的真实 token-level PPL 结果。""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import re +import sys +import tempfile +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Sequence + +from _ppl_common import SCORING_METHOD, canonical_indices_sha256 + +RESULT_SCHEMA = "qwen3_235b_true_ppl_result/v1" +COMPARISON_SCHEMA = "qwen3_235b_true_ppl_comparison/v1" +SHA256_RE = re.compile(r"^[0-9a-f]{64}$") + + +@dataclass(frozen=True) +class PplResult: + path: str + backend: str + model: str + corpus_manifest_sha256: str + corpus_token_ids_sha256: str + window_size: int + stride: int + scoring_method: str + first_scored_token_index: int + last_scored_token_index_exclusive: int + scored_token_count: int + scored_token_indices_sha256: str + total_nll: float + mean_nll: float + ppl: float + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("paths", nargs="*", type=Path) + parser.add_argument( + "--inputs", + nargs="+", + type=Path, + default=[], + help="Transformers 与 InfiniLM 结果 JSON，顺序可以互换", + ) + parser.add_argument( + "--max-ppl-increase-percent", + type=float, + default=20.0, + help="InfiniLM 相对 Transformers 的最大 PPL 增幅，默认 20%%", + ) + parser.add_argument("--json-out", type=Path) + parser.add_argument("--verbose", action="store_true", help="打印完整 JSON") + args = parser.parse_args(argv) + args.inputs = [*args.paths, *args.inputs] + if len(args.inputs) != 2: + parser.error("必须提供两个结果 JSON：Transformers 与 InfiniLM") + if ( + not math.isfinite(args.max_ppl_increase_percent) + or args.max_ppl_increase_percent < 0 + ): + parser.error("--max-ppl-increase-percent 必须是有限非负数") + return args + + +def _required(payload: dict[str, Any], key: str, path: Path) -> Any: + if key not in payload: + raise ValueError(f"{path} 缺少字段 {key}") + return payload[key] + + +def _sha256(payload: dict[str, Any], key: str, path: Path) -> str: + value = str(_required(payload, key, path)).lower() + if not SHA256_RE.fullmatch(value): + raise ValueError(f"{path} 的 {key} 不是有效 SHA256") + return value + + +def _positive_int(payload: dict[str, Any], key: str, path: Path) -> int: + value = _required(payload, key, path) + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{path} 的 {key} 必须是正整数") + parsed = value + if parsed <= 0: + raise ValueError(f"{path} 的 {key} 必须是正整数") + return parsed + + +def _nonnegative_int(payload: dict[str, Any], key: str, path: Path) -> int: + value = _required(payload, key, path) + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{path} 的 {key} 必须是非负整数") + parsed = value + if parsed < 0: + raise ValueError(f"{path} 的 {key} 必须是非负整数") + return parsed + + +def _finite(payload: dict[str, Any], key: str, path: Path) -> float: + try: + value = float(_required(payload, key, path)) + except (TypeError, ValueError) as error: + raise ValueError(f"{path} 的 {key} 必须是有限数") from error + if not math.isfinite(value): + raise ValueError(f"{path} 的 {key} 必须是有限数") + return value + + +def load_result(path: Path) -> PplResult: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError: + raise ValueError(f"结果文件不存在：{path}") from None + except json.JSONDecodeError as error: + raise ValueError(f"结果 JSON 无效：{path}: {error.msg}") from error + if not isinstance(payload, dict): + raise ValueError(f"结果 JSON 必须是对象：{path}") + if payload.get("status") != "PASS": + raise ValueError( + f"{path} 不是成功的 PPL 结果：status={payload.get('status')!r}" + ) + if payload.get("schema") != RESULT_SCHEMA: + raise ValueError( + f"{path} schema 必须为 {RESULT_SCHEMA!r}，实际为 {payload.get('schema')!r}" + ) + + backend = str(_required(payload, "backend", path)).strip().lower() + if backend not in {"transformers", "infinilm"}: + raise ValueError(f"{path} backend 必须是 transformers 或 infinilm") + model = str(_required(payload, "model", path)).strip() + if not model: + raise ValueError(f"{path} model 不能为空") + window_size = _positive_int(payload, "window_size", path) + stride = _positive_int(payload, "stride", path) + if stride >= window_size: + raise ValueError(f"{path} stride 必须小于 window_size") + scoring_method = str(_required(payload, "scoring_method", path)).strip() + if scoring_method != SCORING_METHOD: + raise ValueError( + f"{path} scoring_method 必须为 {SCORING_METHOD!r}" + ) + first_index = _nonnegative_int(payload, "first_scored_token_index", path) + last_index = _positive_int( + payload, "last_scored_token_index_exclusive", path + ) + scored_count = _positive_int(payload, "scored_token_count", path) + if last_index <= first_index or last_index - first_index != scored_count: + raise ValueError( + f"{path} 的计分范围 [{first_index}, {last_index}) 与 " + f"scored_token_count={scored_count} 不一致" + ) + if first_index != 1: + raise ValueError(f"{path} first_scored_token_index 必须为 1") + scored_indices_hash = _sha256( + payload, "scored_token_indices_sha256", path + ) + expected_indices_hash = canonical_indices_sha256( + range(first_index, last_index) + ) + if scored_indices_hash != expected_indices_hash: + raise ValueError(f"{path} 的 scored_token_indices_sha256 校验失败") + + total_nll = _finite(payload, "total_nll", path) + reported_mean = _finite(payload, "mean_nll", path) + reported_ppl = _finite(payload, "ppl", path) + if total_nll < 0 or reported_mean < 0 or reported_ppl < 1: + raise ValueError(f"{path} 的 NLL/PPL 超出有效范围") + calculated_mean = total_nll / scored_count + if calculated_mean > math.log(sys.float_info.max): + raise ValueError(f"{path} 的 mean NLL 过大，PPL 溢出") + calculated_ppl = math.exp(calculated_mean) + if not math.isclose(reported_mean, calculated_mean, rel_tol=1e-6, abs_tol=1e-8): + raise ValueError( + f"{path} 的 mean_nll 与 total_nll/scored_token_count 不一致" + ) + if not math.isclose(reported_ppl, calculated_ppl, rel_tol=1e-6, abs_tol=1e-8): + raise ValueError(f"{path} 的 ppl 与 exp(mean_nll) 不一致") + + return PplResult( + path=str(path), + backend=backend, + model=model, + corpus_manifest_sha256=_sha256( + payload, "corpus_manifest_sha256", path + ), + corpus_token_ids_sha256=_sha256( + payload, "corpus_token_ids_sha256", path + ), + window_size=window_size, + stride=stride, + scoring_method=scoring_method, + first_scored_token_index=first_index, + last_scored_token_index_exclusive=last_index, + scored_token_count=scored_count, + scored_token_indices_sha256=scored_indices_hash, + total_nll=total_nll, + mean_nll=calculated_mean, + ppl=calculated_ppl, + ) + + +def _ordered(results: Sequence[PplResult]) -> tuple[PplResult, PplResult]: + by_backend = {result.backend: result for result in results} + if len(by_backend) != 2 or set(by_backend) != {"transformers", "infinilm"}: + raise ValueError("必须且只能包含一份 Transformers 和一份 InfiniLM 结果") + return by_backend["transformers"], by_backend["infinilm"] + + +def _validate_same_workload(baseline: PplResult, candidate: PplResult) -> None: + fields = ( + "corpus_manifest_sha256", + "corpus_token_ids_sha256", + "window_size", + "stride", + "scoring_method", + "first_scored_token_index", + "last_scored_token_index_exclusive", + "scored_token_count", + "scored_token_indices_sha256", + ) + mismatches = [ + f"{field}: {getattr(baseline, field)!r} != {getattr(candidate, field)!r}" + for field in fields + if getattr(baseline, field) != getattr(candidate, field) + ] + if mismatches: + raise ValueError("两侧 PPL 工作负载不一致：" + "; ".join(mismatches)) + + +def compare( + baseline: PplResult, candidate: PplResult, threshold_percent: float +) -> dict[str, object]: + _validate_same_workload(baseline, candidate) + increase_percent = (candidate.ppl / baseline.ppl - 1.0) * 100.0 + passed = increase_percent <= threshold_percent + return { + "schema": COMPARISON_SCHEMA, + "status": "PASS" if passed else "FAIL", + "baseline": asdict(baseline), + "candidate": asdict(candidate), + "ppl_increase_percent": increase_percent, + "max_ppl_increase_percent": threshold_percent, + "pass": passed, + } + + +def _atomic_json(path: Path, payload: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, ensure_ascii=False, sort_keys=True, indent=2) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + except BaseException: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + raise + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + try: + baseline, candidate = _ordered([load_result(path) for path in args.inputs]) + report = compare( + baseline, candidate, float(args.max_ppl_increase_percent) + ) + if args.json_out is not None: + _atomic_json(args.json_out, report) + except (OSError, RuntimeError, ValueError) as error: + print(f"错误：{error}", file=sys.stderr) + return 2 + + print("真实 PPL 对比") + print( + f"Transformers：PPL={baseline.ppl:.6f} " + f"NLL={baseline.total_nll:.6f} Token={baseline.scored_token_count}" + ) + print( + f"InfiniLM： PPL={candidate.ppl:.6f} " + f"NLL={candidate.total_nll:.6f} Token={candidate.scored_token_count}" + ) + print(f"PPL 增幅：{report['ppl_increase_percent']:.2f}%") + print( + f"验收要求：增幅 <= {args.max_ppl_increase_percent:.2f}% " + f"结果={report['status']}" + ) + if args.verbose: + print(json.dumps(report, ensure_ascii=False, sort_keys=True, indent=2)) + return 0 if report["pass"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py b/test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py new file mode 100755 index 000000000..9a2f35629 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python3 +"""Calculate true token-level PPL with the current InfiniLM C++ TP engine.""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import os +import sys +import time +from pathlib import Path +from typing import Any, Sequence + + +SCRIPT_DIR = Path(__file__).resolve().parent +SCRIPTS_DIR = SCRIPT_DIR.parent +if str(SCRIPTS_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPTS_DIR)) + +from _gpu_guard import require_idle_gpu +from _ppl_common import ( + SCORING_METHOD, + canonical_indices_sha256, + iter_sliding_windows, + load_manifest, +) + + +RESULT_SCHEMA = "qwen3_235b_true_ppl_result/v1" +DEFAULT_MODEL = "/data1/Qwen3_235B" +EXPECTED_MODEL_TYPE = "qwen3_moe" +EXPECTED_VOCAB_SIZE = 151936 +PAGED_KV_BLOCK_SIZE = 256 + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="True shifted-token PPL for Qwen3_235B with InfiniLM TP8" + ) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--token-manifest", required=True) + parser.add_argument("--window", type=int, default=256) + parser.add_argument("--stride", type=int, default=128) + parser.add_argument( + "--max-scored-tokens", + type=int, + default=10240, + help="maximum target tokens to score; 0 scores the full manifest", + ) + parser.add_argument("--tp-size", type=int, default=8) + parser.add_argument("--attention", default="flash-attn") + parser.add_argument("--json-output") + args = parser.parse_args(argv) + + if not Path(args.model).is_dir(): + parser.error(f"model directory does not exist: {args.model}") + if not Path(args.token_manifest).is_file(): + parser.error(f"token manifest does not exist: {args.token_manifest}") + if args.window < 2: + parser.error("--window must be at least 2") + if args.stride < 1 or args.stride >= args.window: + parser.error("--stride must satisfy 1 <= stride < window") + if args.max_scored_tokens < 0: + parser.error("--max-scored-tokens must be non-negative") + if args.tp_size < 1: + parser.error("--tp-size must be positive") + + args.model = str(Path(args.model).resolve()) + args.token_manifest = str(Path(args.token_manifest).resolve()) + if args.json_output: + args.json_output = str(Path(args.json_output).resolve()) + return args + + +def _atomic_json(path_value: str, payload: dict[str, Any]) -> None: + path = Path(path_value) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") + temporary.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + +def _read_model_config(model_path: str) -> dict[str, Any]: + config_path = Path(model_path) / "config.json" + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError(f"cannot read model config {config_path}: {error}") from error + if not isinstance(config, dict): + raise RuntimeError(f"model config must be an object: {config_path}") + if config.get("model_type") != EXPECTED_MODEL_TYPE: + raise RuntimeError( + f"expected model_type={EXPECTED_MODEL_TYPE!r}, " + f"got {config.get('model_type')!r}" + ) + if int(config.get("vocab_size", 0)) != EXPECTED_VOCAB_SIZE: + raise RuntimeError( + f"expected vocab_size={EXPECTED_VOCAB_SIZE}, " + f"got {config.get('vocab_size')!r}" + ) + return config + + +def _is_quantized(config: dict[str, Any]) -> bool: + quantization = config.get("quantization_config") + return isinstance(quantization, dict) and bool(quantization) + + +def _run(args: argparse.Namespace) -> dict[str, Any]: + import infinicore + from infinilm.cache import PagedKVCacheConfig + from infinilm.distributed import DistConfig + from infinilm.infer_engine import InferEngine + from infinilm.modeling_utils import load_model_state_dict_by_file + + corpus = load_manifest(args.token_manifest) + model_config = _read_model_config(args.model) + if any(token >= EXPECTED_VOCAB_SIZE for token in corpus.token_ids): + raise RuntimeError("token manifest contains an ID outside the model vocabulary") + + available_targets = corpus.token_count - 1 + scored_token_count = ( + available_targets + if args.max_scored_tokens == 0 + else min(args.max_scored_tokens, available_targets) + ) + max_targets = None if args.max_scored_tokens == 0 else scored_token_count + windows = list( + iter_sliding_windows( + corpus.token_ids, + args.window, + args.stride, + max_targets, + ) + ) + if sum(window.scored_token_count for window in windows) != scored_token_count: + raise RuntimeError("sliding-window plan does not match scored token count") + + first_scored_token_index = 1 + last_scored_token_index_exclusive = 1 + scored_token_count + indices_sha256 = canonical_indices_sha256( + range(first_scored_token_index, last_scored_token_index_exclusive) + ) + precision = "W8A8" if _is_quantized(model_config) else "BF16" + config_payload = { + "backend": "infinilm", + "model": args.model, + "precision": precision, + "tp_size": args.tp_size, + "attention": args.attention, + "graph_enabled": False, + "window_size": args.window, + "stride": args.stride, + "scored_token_count": scored_token_count, + "scoring_method": SCORING_METHOD, + "corpus_manifest_sha256": corpus.manifest_sha256, + "corpus_token_ids_sha256": corpus.token_ids_sha256, + } + print( + "INFINILM_QWEN3_235B_PPL_CONFIG " + + json.dumps(config_payload, ensure_ascii=False, sort_keys=True), + flush=True, + ) + + device = infinicore.device("cuda", 0) + load_start = time.perf_counter() + model = InferEngine( + args.model, + device=device, + distributed_config=DistConfig(args.tp_size), + # This InfiniLM branch's flash-attn backend consumes the paged KV-cache + # layout while retaining flash-attn as the attention implementation. + cache_config=PagedKVCacheConfig( + num_blocks=(args.window + PAGED_KV_BLOCK_SIZE - 1) + // PAGED_KV_BLOCK_SIZE, + block_size=PAGED_KV_BLOCK_SIZE, + ), + enable_graph_compiling=False, + attention_backend=args.attention, + ) + if not hasattr(model, "score_nll"): + raise RuntimeError( + "installed InfiniLM lacks InferEngine.score_nll; rebuild the PPL scoring patch" + ) + load_model_state_dict_by_file(model, args.model, dtype=model.dtype) + model_load_seconds = time.perf_counter() - load_start + + window_nll_values: list[float] = [] + window_results: list[dict[str, Any]] = [] + infinicore.sync_device() + scoring_start = time.perf_counter() + for window in windows: + input_tokens = list(window.token_ids[:-1]) + label_tokens = list(window.token_ids[1:]) + if not input_tokens or len(input_tokens) != len(label_tokens): + raise RuntimeError(f"invalid causal shift in window {window.index}") + input_ids = infinicore.from_list( + [input_tokens], dtype=infinicore.int64 + ) + labels = infinicore.from_list( + [label_tokens], dtype=infinicore.int64 + ) + nll, returned_tokens = model.score_nll( + input_ids, + labels, + score_start=window.prediction_start, + ) + if returned_tokens != window.scored_token_count: + raise RuntimeError( + f"window {window.index} scored {returned_tokens} tokens, " + f"expected {window.scored_token_count}" + ) + if not math.isfinite(nll) or nll < 0: + raise RuntimeError(f"window {window.index} returned invalid NLL {nll}") + window_nll_values.append(nll) + window_results.append( + { + "index": window.index, + "context_start": window.token_start, + "target_start": window.score_start, + "target_end": window.score_end, + "input_token_count": len(window.token_ids), + "scored_token_count": returned_tokens, + "nll": nll, + } + ) + print( + f"PPL window {window.index + 1}/{len(windows)} " + f"tokens={returned_tokens} nll={nll:.6f}", + flush=True, + ) + + infinicore.sync_device() + scoring_seconds = time.perf_counter() - scoring_start + total_nll = math.fsum(window_nll_values) + mean_nll = total_nll / scored_token_count + try: + ppl = math.exp(mean_nll) + except OverflowError as error: + raise RuntimeError(f"PPL overflow at mean NLL={mean_nll}") from error + if not math.isfinite(ppl): + raise RuntimeError(f"PPL is not finite: {ppl}") + + result = { + "schema": RESULT_SCHEMA, + "status": "PASS", + "backend": "infinilm", + "model": args.model, + "precision": precision, + "tp_size": args.tp_size, + "attention": args.attention, + "graph_enabled": False, + "corpus_manifest": args.token_manifest, + "corpus_manifest_sha256": corpus.manifest_sha256, + "corpus_token_ids_sha256": corpus.token_ids_sha256, + "corpus_token_count": corpus.token_count, + "window_size": args.window, + "stride": args.stride, + "scoring_method": SCORING_METHOD, + "first_scored_token_index": first_scored_token_index, + "last_scored_token_index_exclusive": last_scored_token_index_exclusive, + "scored_token_indices_sha256": indices_sha256, + "scored_token_count": scored_token_count, + "total_nll": total_nll, + "mean_nll": mean_nll, + "ppl": ppl, + "windows": window_results, + "window_count": len(window_results), + "scoring_seconds": scoring_seconds, + "scored_tokens_per_second": scored_token_count / scoring_seconds, + "model_load_seconds": model_load_seconds, + "vocab_size": EXPECTED_VOCAB_SIZE, + } + if args.json_output: + _atomic_json(args.json_output, result) + print( + "INFINILM_QWEN3_235B_PPL_RESULT " + + json.dumps(result, ensure_ascii=False, sort_keys=True), + flush=True, + ) + print( + f"InfiniLM {precision} true PPL: {ppl:.6f} " + f"(mean NLL={mean_nll:.6f}, tokens={scored_token_count})", + flush=True, + ) + + del model + gc.collect() + infinicore.sync_device() + return result + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + try: + # Validate the workload before checking or reserving GPUs. + load_manifest(args.token_manifest) + require_idle_gpu() + _run(args) + except BaseException as error: + completion = { + "schema": RESULT_SCHEMA, + "status": "ERROR", + "exit_code": 1, + "error": { + "type": type(error).__name__, + "message": str(error), + }, + } + print( + "INFINILM_QWEN3_235B_PPL_COMPLETE " + + json.dumps(completion, ensure_ascii=False, sort_keys=True), + flush=True, + ) + raise + print( + "INFINILM_QWEN3_235B_PPL_COMPLETE " + + json.dumps( + {"schema": RESULT_SCHEMA, "status": "PASS", "exit_code": 0}, + ensure_ascii=False, + sort_keys=True, + ), + flush=True, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py b/test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py new file mode 100755 index 000000000..b501e453c --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py @@ -0,0 +1,328 @@ +#!/usr/bin/env python3 +"""将本地纯文本固化为 Qwen3_235B PPL 测试使用的 token manifest。""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import operator +import os +import sys +import tempfile +from pathlib import Path +from typing import Any, Sequence + +from _ppl_common import ( + CORPUS_SCHEMA, + canonical_json_bytes, + canonical_token_ids_sha256, + write_token_ids_npy, +) + + +def _jsonable(value: Any) -> Any: + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + return str(value) + + +def _tokenizer_fingerprint(tokenizer: Any) -> tuple[str, str]: + backend = getattr(tokenizer, "backend_tokenizer", None) + if backend is not None and hasattr(backend, "to_str"): + try: + backend_payload: object = json.loads(backend.to_str()) + except (TypeError, ValueError, json.JSONDecodeError): + backend_payload = backend.to_str() + method = "backend_tokenizer+special_tokens/v1" + semantics = { + "backend_tokenizer": backend_payload, + "special_tokens_map": _jsonable( + getattr(tokenizer, "special_tokens_map", {}) + ), + } + else: + if not hasattr(tokenizer, "get_vocab"): + raise RuntimeError("tokenizer 既没有 backend_tokenizer，也没有 get_vocab()") + method = "vocab+init_kwargs+special_tokens/v1" + semantics = { + "vocab": tokenizer.get_vocab(), + "init_kwargs": _jsonable(getattr(tokenizer, "init_kwargs", {})), + "special_tokens_map": _jsonable( + getattr(tokenizer, "special_tokens_map", {}) + ), + } + return hashlib.sha256(canonical_json_bytes(semantics)).hexdigest(), method + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--input", + nargs="+", + type=Path, + required=True, + help="本地 UTF-8 WikiText/raw text 文件，可按顺序提供多个文件", + ) + parser.add_argument( + "--tokenizer", + required=True, + help="本地 Qwen3_235B 模型或 tokenizer 目录", + ) + parser.add_argument("--output", required=True, type=Path, help="输出 JSON manifest") + parser.add_argument( + "--max-tokens", + type=int, + help="仅保留开头 N 个 token；省略时保留全部 token", + ) + parser.add_argument( + "--document-separator", + default="\n\n", + help=r"多个输入文件之间的分隔符，默认 '\n\n'", + ) + parser.add_argument( + "--storage", + choices=("inline", "npy"), + default="inline", + help="token IDs 内联到 JSON（默认），或写入相对路径 .npy 文件", + ) + parser.add_argument( + "--token-ids-file", + type=Path, + help="--storage=npy 时的相对路径；默认 .tokens.npy", + ) + parser.add_argument( + "--allow-download", + action="store_true", + help="允许 Transformers 访问网络；默认只读取本地文件/缓存", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="传给 AutoTokenizer；Qwen3 官方 tokenizer 通常不需要", + ) + parser.add_argument("--overwrite", action="store_true", help="覆盖已有输出") + args = parser.parse_args(argv) + + if args.max_tokens is not None and args.max_tokens < 2: + parser.error("--max-tokens 必须至少为 2") + if args.token_ids_file is not None and args.storage != "npy": + parser.error("--token-ids-file 只能与 --storage=npy 一起使用") + if len({path.name for path in args.input}) != len(args.input): + parser.error("输入文件名不能重复，否则 manifest 无法稳定区分来源") + return args + + +def _read_sources( + paths: Sequence[Path], separator: str +) -> tuple[str, list[dict[str, object]]]: + documents: list[str] = [] + files: list[dict[str, object]] = [] + for path in paths: + if not path.is_file(): + raise FileNotFoundError(f"输入文件不存在：{path}") + raw = path.read_bytes() + try: + text = raw.decode("utf-8") + except UnicodeDecodeError as error: + raise ValueError(f"输入文件不是有效 UTF-8：{path}: {error}") from error + documents.append(text) + files.append( + { + "name": path.name, + "byte_count": len(raw), + "sha256": hashlib.sha256(raw).hexdigest(), + } + ) + return separator.join(documents), files + + +def _load_tokenizer(identifier: str, allow_download: bool, trust_remote_code: bool) -> Any: + try: + from transformers import AutoTokenizer + except ImportError as error: + raise RuntimeError("缺少 transformers，无法加载 tokenizer") from error + return AutoTokenizer.from_pretrained( + identifier, + local_files_only=not allow_download, + trust_remote_code=trust_remote_code, + ) + + +def _encode(tokenizer: Any, text: str) -> list[int]: + encoded = tokenizer( + text, + add_special_tokens=False, + truncation=False, + return_attention_mask=False, + return_token_type_ids=False, + ) + raw_ids = encoded["input_ids"] + if not isinstance(raw_ids, (list, tuple)) or ( + raw_ids and isinstance(raw_ids[0], (list, tuple)) + ): + raise RuntimeError("tokenizer 必须为单条文本返回一维 input_ids") + token_ids: list[int] = [] + for index, value in enumerate(raw_ids): + if isinstance(value, bool): + raise RuntimeError(f"input_ids[{index}] 不是有效整数") + try: + token = operator.index(value) + except TypeError as error: + raise RuntimeError(f"input_ids[{index}] 不是有效整数") from error + if token < 0: + raise RuntimeError(f"input_ids[{index}] 不是非负整数：{value!r}") + token_ids.append(token) + if len(token_ids) < 2: + raise RuntimeError("语料 token 数不足 2，无法计算 causal LM PPL") + return token_ids + + +def _atomic_json(path: Path, payload: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, ensure_ascii=False, sort_keys=True, indent=2) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + except BaseException: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + raise + + +def _npy_relative_path(output: Path, requested: Path | None) -> Path: + relative = requested or Path(f"{output.stem}.tokens.npy") + if relative.is_absolute() or ".." in relative.parts or relative.name in {"", "."}: + raise ValueError("--token-ids-file 必须是 manifest 目录内的安全相对路径") + if relative.suffix != ".npy": + raise ValueError("--token-ids-file 必须以 .npy 结尾") + return relative + + +def _write_npy(path: Path, token_ids: Sequence[int], overwrite: bool) -> None: + if path.exists() and not overwrite: + raise FileExistsError(f"token 文件已存在（可加 --overwrite）：{path}") + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp.npy") + try: + write_token_ids_npy(temporary, token_ids) + os.replace(temporary, path) + finally: + try: + temporary.unlink() + except FileNotFoundError: + pass + + +def build_manifest( + *, + source_text: str, + source_files: list[dict[str, object]], + separator: str, + tokenizer: Any, + tokenizer_label: str, + token_ids: list[int], + original_token_count: int, + max_tokens: int | None, + storage: str, + token_ids_file: str | None, +) -> dict[str, object]: + tokenizer_hash, fingerprint_method = _tokenizer_fingerprint(tokenizer) + token_hash = canonical_token_ids_sha256(token_ids) + payload: dict[str, object] = { + "schema": CORPUS_SCHEMA, + "source_sha256": hashlib.sha256(source_text.encode("utf-8")).hexdigest(), + "tokenizer_sha256": tokenizer_hash, + "token_count": len(token_ids), + "token_ids_sha256": token_hash, + "source": { + "encoding": "utf-8", + "document_separator": separator, + "file_count": len(source_files), + "files": source_files, + }, + "tokenizer": { + "name": Path(tokenizer_label.rstrip("/")).name or tokenizer_label, + "class": tokenizer.__class__.__name__, + "vocab_size": int(getattr(tokenizer, "vocab_size", 0)), + "fingerprint_method": fingerprint_method, + }, + "tokenization": { + "add_special_tokens": False, + "original_token_count": original_token_count, + "max_tokens": max_tokens, + "truncated": len(token_ids) != original_token_count, + }, + } + if storage == "inline": + payload["token_ids"] = token_ids + else: + if token_ids_file is None: + raise ValueError("npy storage 缺少 token_ids_file") + payload["token_ids_file"] = token_ids_file + payload["token_ids_dtype"] = "int64" + payload["manifest_sha256"] = hashlib.sha256( + canonical_json_bytes(payload) + ).hexdigest() + return payload + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + try: + if args.output.exists() and not args.overwrite: + raise FileExistsError(f"输出已存在（可加 --overwrite）：{args.output}") + source_text, source_files = _read_sources(args.input, args.document_separator) + tokenizer = _load_tokenizer( + args.tokenizer, args.allow_download, args.trust_remote_code + ) + all_token_ids = _encode(tokenizer, source_text) + original_token_count = len(all_token_ids) + token_ids = ( + all_token_ids[: args.max_tokens] + if args.max_tokens is not None + else all_token_ids + ) + + relative_npy: Path | None = None + if args.storage == "npy": + relative_npy = _npy_relative_path(args.output, args.token_ids_file) + + manifest = build_manifest( + source_text=source_text, + source_files=source_files, + separator=args.document_separator, + tokenizer=tokenizer, + tokenizer_label=args.tokenizer, + token_ids=token_ids, + original_token_count=original_token_count, + max_tokens=args.max_tokens, + storage=args.storage, + token_ids_file=str(relative_npy) if relative_npy is not None else None, + ) + if relative_npy is not None: + _write_npy(args.output.parent / relative_npy, token_ids, args.overwrite) + _atomic_json(args.output, manifest) + except (FileNotFoundError, FileExistsError, RuntimeError, ValueError) as error: + print(f"错误：{error}", file=sys.stderr) + return 2 + + print(f"PPL 语料已固化：{args.output}") + print(f"Token 数：{manifest['token_count']}") + print(f"Token SHA256：{manifest['token_ids_sha256']}") + print(f"Manifest SHA256：{manifest['manifest_sha256']}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py b/test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py new file mode 100755 index 000000000..027b5c878 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py @@ -0,0 +1,1144 @@ +#!/usr/bin/env python3 +"""Minimal Transformers TP benchmark runner for Qwen3_235B-A22B. + +The scenario wrappers are intentionally directly executable. When started as +``python wrapper.py`` this module replaces the process with a torchrun agent; +the original timeout therefore remains responsible for the agent, which in +turn terminates every worker on SIGTERM. +""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import importlib.metadata +import json +import os +import statistics +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + + +SCRIPT_ROOT = str(Path(__file__).resolve().parents[1]) +if SCRIPT_ROOT not in sys.path: + sys.path.insert(0, SCRIPT_ROOT) + +from _gpu_guard import require_idle_gpu as _require_idle_gpu + + +MODEL_NAME = "Qwen3_235B" +DEFAULT_MODEL = "/data1/Qwen3_235B" +DEFAULT_PROMPT_FILE = "examples/bench_prompt.md" +FALLBACK_PROMPT = """High-performance language-model inference processes a prompt in a prefill phase +and then produces one token per request during decode. Tensor-parallel ranks must +exchange identical partial results, while the key/value cache keeps each request +lane isolated. The benchmark uses deterministic prompt tokens and greedy decoding +so that every reported run has an exact, auditable token count.""" +MEASURED_INPUT_LENGTHS = 1 +REPEATS_PER_INPUT_LENGTH = 3 +MEASURED_ITERATIONS = REPEATS_PER_INPUT_LENGTH +MEASUREMENT_SEMANTICS = "one_fixed_shape_x_three_measurements" +SMOKE_OUTPUT_TOKENS = 64 +HYGON_TP_PLAN = { + "lm_head": "colwise_gather_output", + "model.layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "model.layers.*.mlp.experts.down_proj": "rowwise", + "model.layers.*.mlp.experts": "moe_tp_experts", +} +EXPECTED_QWEN3_235B_ARCHITECTURE = { + "hidden_size": 4096, + "intermediate_size": 12288, + "head_dim": 128, + "num_attention_heads": 64, + "num_key_value_heads": 4, + "num_hidden_layers": 94, + "num_experts": 128, + "num_experts_per_tok": 8, + "moe_intermediate_size": 1536, + "vocab_size": 151936, +} + + +@dataclass(frozen=True) +class Scenario: + name: str + batch_size: int + input_tokens: int + output_tokens: int + + @property + def input_lengths(self) -> tuple[int]: + return (self.input_tokens,) + + @property + def total_context_tokens(self) -> int: + return self.input_tokens + self.output_tokens + + +def _parse_args(scenario: Scenario) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Transformers Qwen3_235B TP8 benchmark: " + f"batch={scenario.batch_size}, input={scenario.input_tokens}, " + f"output={scenario.output_tokens}, " + f"total={scenario.total_context_tokens} tokens" + ) + ) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--prompt-file", default=DEFAULT_PROMPT_FILE) + parser.add_argument("--output-tokens", type=int, default=scenario.output_tokens) + parser.add_argument("--tp-size", type=int, default=8) + parser.add_argument( + "--smoke", + action="store_true", + help=( + "load the full model but run only batch=1, input=16 and output=64 " + "to validate the TP/attention/cache path" + ), + ) + parser.add_argument( + "--attention", + choices=("eager",), + default="eager", + help="BW1100 correctness path; SDPA is unsupported for this model stack", + ) + args = parser.parse_args() + + if not Path(args.model).is_dir(): + parser.error(f"model directory does not exist: {args.model}") + if args.smoke: + args.output_tokens = SMOKE_OUTPUT_TOKENS + if args.output_tokens < 2: + parser.error("--output-tokens must be at least 2 to measure decode speed") + if args.tp_size < 1: + parser.error("--tp-size must be positive") + if len(set(scenario.input_lengths)) != MEASURED_INPUT_LENGTHS: + parser.error( + f"scenario must define exactly {MEASURED_INPUT_LENGTHS} lengths" + ) + return args + + +def _effective_scenario(scenario: Scenario, smoke: bool) -> Scenario: + if not smoke: + return scenario + return Scenario(f"{scenario.name}_smoke", 1, 16, SMOKE_OUTPUT_TOKENS) + + +def _validate_qwen3_235b_architecture(model_config: Any) -> dict[str, int]: + if getattr(model_config, "model_type", None) != "qwen3_moe": + raise RuntimeError( + "this benchmark requires model_type='qwen3_moe', got " + f"{getattr(model_config, 'model_type', None)!r}" + ) + actual: dict[str, int] = {} + mismatches: list[str] = [] + for field, expected in EXPECTED_QWEN3_235B_ARCHITECTURE.items(): + value = getattr(model_config, field, None) + try: + parsed = int(value) + except (TypeError, ValueError): + mismatches.append(f"{field}={value!r} (expected {expected})") + continue + actual[field] = parsed + if parsed != expected: + mismatches.append(f"{field}={parsed} (expected {expected})") + architectures = tuple(getattr(model_config, "architectures", None) or ()) + if "Qwen3MoeForCausalLM" not in architectures: + mismatches.append( + "architectures does not contain 'Qwen3MoeForCausalLM': " + f"{architectures!r}" + ) + if mismatches: + raise RuntimeError( + "checkpoint is not the expected Qwen3_235B-A22B architecture: " + + "; ".join(mismatches) + ) + return actual + + +def _build_qwen3_moe_tp_plan( + model_config: Any, + tp_size: int, + scenario: Scenario, + output_tokens: int, +) -> tuple[str | dict[str, str], dict[str, Any]]: + """Use the correctness-first TP8 layout validated on BW1100. + + Attention stays replicated. Qwen3_235B has four KV heads, so attempting to + tensor-parallelize attention over eight ranks produces an invalid local GQA + layout in this Transformers/DTK stack. Only MoE experts and the LM head are + sharded, matching the working Hygon container example. + """ + if getattr(model_config, "model_type", None) != "qwen3_moe": + raise RuntimeError( + "this benchmark only supports model_type='qwen3_moe', got " + f"{getattr(model_config, 'model_type', None)!r}" + ) + + global_query_heads = int(model_config.num_attention_heads) + global_kv_heads = int(model_config.num_key_value_heads) + head_dim = int(model_config.head_dim) + num_hidden_layers = int(model_config.num_hidden_layers) + if min( + global_query_heads, + global_kv_heads, + head_dim, + num_hidden_layers, + tp_size, + ) < 1: + raise RuntimeError("TP and attention dimensions must all be positive") + if global_query_heads % global_kv_heads: + raise RuntimeError( + f"global Q heads ({global_query_heads}) must be divisible by global " + f"KV heads ({global_kv_heads})" + ) + tp_plan = dict(HYGON_TP_PLAN) + maximum_sequence_tokens = scenario.input_tokens + output_tokens + dtype_bytes = 2 # BF16 K and V elements. + kv_cache_bytes_per_rank = ( + scenario.batch_size + * maximum_sequence_tokens + * num_hidden_layers + * 2 + * global_kv_heads + * head_dim + * dtype_bytes + ) + plan_payload = tp_plan if isinstance(tp_plan, dict) else {"mode": tp_plan} + metadata = { + "tp_plan_mode": f"qwen3_moe_tp{tp_size}_experts_lm_head_only", + "tp_plan_sha256": _stable_hash(plan_payload), + "attention_strategy": "replicated_eager", + "kv_projection_strategy": "replicated", + "kv_cache_replication_factor_across_tp_ranks": tp_size, + "global_query_heads": global_query_heads, + "global_kv_heads": global_kv_heads, + "head_dim": head_dim, + "num_hidden_layers": num_hidden_layers, + "local_query_heads": global_query_heads, + "local_kv_heads": global_kv_heads, + "local_gqa_groups": global_query_heads // global_kv_heads, + "maximum_sequence_tokens": maximum_sequence_tokens, + "estimated_dense_bf16_kv_cache_gib_per_rank": ( + kv_cache_bytes_per_rank / (1024**3) + ), + "kv_cache_estimate_excludes_allocator_and_cache_metadata": True, + } + return tp_plan, metadata + + +def _validate_and_set_local_gqa( + model: Any, tp_metadata: dict[str, Any] +) -> dict[str, Any]: + """Verify that attention stayed fully replicated on every TP rank.""" + base_model_prefix = getattr(model, "base_model_prefix", None) + base_model = getattr(model, base_model_prefix, None) + layers = getattr(base_model, "layers", None) + if layers is None: + raise RuntimeError( + f"cannot locate {base_model_prefix!r}.layers on loaded model" + ) + + head_dim = int(tp_metadata["head_dim"]) + expected_query_heads = int(tp_metadata["local_query_heads"]) + expected_kv_heads = int(tp_metadata["local_kv_heads"]) + expected_groups = int(tp_metadata["local_gqa_groups"]) + expected_query_width = expected_query_heads * head_dim + expected_kv_width = expected_kv_heads * head_dim + observed_groups: set[int] = set() + + for layer_index, layer in enumerate(layers): + attention = getattr(layer, "self_attn", None) + if attention is None: + raise RuntimeError(f"layer {layer_index} has no self_attn module") + query_width = int(attention.q_proj.out_features) + key_width = int(attention.k_proj.out_features) + value_width = int(attention.v_proj.out_features) + if query_width != expected_query_width: + raise RuntimeError( + f"layer {layer_index} local Q width={query_width}, expected " + f"{expected_query_width} ({expected_query_heads} heads)" + ) + if key_width != expected_kv_width or value_width != expected_kv_width: + raise RuntimeError( + f"layer {layer_index} local K/V widths={key_width}/{value_width}, " + f"expected {expected_kv_width} ({expected_kv_heads} heads)" + ) + observed_groups.add(int(attention.num_key_value_groups)) + + if observed_groups != {expected_groups}: + raise RuntimeError( + "attention GQA metadata changed despite replicated attention: " + f"observed={sorted(observed_groups)}, expected={expected_groups}" + ) + + expected_layers = int(tp_metadata["num_hidden_layers"]) + if len(layers) != expected_layers: + raise RuntimeError( + f"loaded model has {len(layers)} transformer layers, expected " + f"{expected_layers}" + ) + return { + "validated_attention_layers": len(layers), + "local_query_projection_width": expected_query_width, + "local_kv_projection_width": expected_kv_width, + "attention_replication_validated": True, + "local_gqa_groups": expected_groups, + } + + +def _launch_torchrun(args: argparse.Namespace) -> None: + env = os.environ.copy() + target_path = ( + "/root/.local/bin:/opt/dtk/cuda/cuda/bin:/opt/dtk/bin:/opt/dtk/hip/bin" + ) + target_library_path = ":".join( + ( + "/usr/local/lib/python3.10/dist-packages/torch/lib", + "/opt/dtk/dcc/gcvm/lib", + "/opt/dtk/hip/lib", + "/opt/dtk/llvm/lib", + "/opt/dtk/lib", + "/opt/dtk/lib64", + "/opt/hyhal/lib", + "/opt/hyhal/lib64", + "/opt/dtk/dushmem/lib", + "/opt/dtk/opencl/lib", + "/opt/ucx/lib", + "/opt/mpi/lib", + "/opt/hwloc/lib", + ) + ) + env["PATH"] = f"{target_path}:{env.get('PATH', '')}" + inherited_library_path = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = ( + f"{target_library_path}:{inherited_library_path}" + if inherited_library_path + else target_library_path + ) + inherited_python_path = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"/usr/local:{inherited_python_path}" + if inherited_python_path + else "/usr/local" + ) + visible_devices = ",".join(str(index) for index in range(args.tp_size)) + env.setdefault("HIP_VISIBLE_DEVICES", visible_devices) + env.setdefault("CUDA_VISIBLE_DEVICES", visible_devices) + env.setdefault("OMP_NUM_THREADS", "1") + env.setdefault("TOKENIZERS_PARALLELISM", "false") + env.setdefault("PYTHONUNBUFFERED", "1") + env.setdefault("HSA_FORCE_FINE_GRAIN_PCIE", "1") + env.setdefault("NCCL_DEBUG", "WARN") + + script = str(Path(sys.argv[0]).resolve()) + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc-per-node={args.tp_size}", + "--max-restarts=0", + "--monitor-interval=1", + script, + *sys.argv[1:], + ] + os.execvpe(sys.executable, command, env) + + +def _package_version(name: str) -> str | None: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return None + + +def _install_hygon_grouped_mm_guard(torch: Any) -> bool: + """Install the grouped-MM fallback validated in the BW1100 image.""" + if getattr(torch.version, "hip", None) is None: + return False + + from transformers.integrations import moe as transformers_moe + + if getattr(transformers_moe, "_hygon_grouped_mm_guard_installed", False): + return True + + def grouped_mm(input_tensor: Any, weight: Any, offs: Any) -> Any: + ends = [int(value) for value in offs.detach().cpu().tolist()] + output = input_tensor.new_empty( + (input_tensor.shape[0], weight.shape[-1]), dtype=weight.dtype + ) + start = 0 + for expert_index, end in enumerate(ends): + if end > start: + output[start:end] = input_tensor[start:end].to(weight.dtype).matmul( + weight[expert_index] + ) + start = end + return output + + transformers_moe._grouped_mm = grouped_mm + transformers_moe._hygon_grouped_mm_guard_installed = True + return True + + +def _stable_hash(value: Any) -> str: + payload = json.dumps(value, ensure_ascii=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _emit(rank: int, tag: str, payload: dict[str, Any]) -> None: + if rank != 0: + return + if tag == "PYTORCH_QWEN3_235B_CONFIG": + print( + "[Transformers] " + f"model={payload['model']} batch={payload['batch_size']} " + f"input={payload['input_lengths'][0]} " + f"output={payload['output_tokens_per_request']} " + f"tp={payload['tp_size']} attention={payload['attention_implementation']}", + flush=True, + ) + print( + f" load weights over! {payload['model_load_seconds'] * 1000.0:.2f} ms ", + flush=True, + ) + elif tag == "PYTORCH_QWEN3_235B_COMPLETE": + print(f"Transformers benchmark status: {payload['status']}", flush=True) + + +def _print_infinilm_style_metrics( + rank: int, measurement: dict[str, Any], decoded_output: str +) -> None: + if rank != 0: + return + print( + f"\n Generation completed in {measurement['generation_seconds'] * 1000.0:.2f} ms", + flush=True, + ) + print( + f" Batchsize={measurement['batch_size']} " + f"Per_Batch_Input_Len={measurement['input_tokens_per_request']} " + f"Per_Batch_New_Tokens={measurement['output_tokens_per_request']}", + flush=True, + ) + print( + f"\n Prefill TTFT: {measurement['ttft_seconds'] * 1000.0:.2f} ms " + f"Throughput: {measurement['prefill_tokens_per_second']:.2f} tok/s", + flush=True, + ) + print( + f"\n Decode Avg ITL: {measurement['inter_token_latency_ms']:.2f} ms " + f"Throughput: {measurement['decode_tokens_per_second']:.2f} tok/s\n", + flush=True, + ) + print(decoded_output or "（未生成可显示文本）", flush=True) + + +def _validate_decoded_output(decoded_output: str) -> dict[str, Any]: + text = decoded_output.strip() + if not text: + raise RuntimeError("generated output decoded to an empty string") + if "\ufffd" in text: + raise RuntimeError("generated output contains Unicode replacement characters") + printable_characters = sum( + character.isprintable() or character in "\n\t" for character in text + ) + printable_ratio = printable_characters / len(text) + cjk_characters = sum("\u3400" <= character <= "\u9fff" for character in text) + url_fragments = text.lower().count("http") + if printable_ratio < 0.95: + raise RuntimeError( + f"generated output printable ratio is too low: {printable_ratio:.3f}" + ) + if cjk_characters < 8: + raise RuntimeError( + f"generated output is not a substantive Chinese response: CJK={cjk_characters}" + ) + if url_fragments > 2: + raise RuntimeError( + f"generated output contains suspicious URL fragments: {url_fragments}" + ) + return { + "nonempty_decoded_text": True, + "no_replacement_characters": True, + "printable_ratio": printable_ratio, + "cjk_character_count": cjk_characters, + "url_fragment_count": url_fragments, + } + + +def _max_across_ranks(value: float, torch: Any, dist: Any, device: Any) -> float: + tensor = torch.tensor(value, dtype=torch.float64, device=device) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return float(tensor.item()) + + +def _make_prompt_base( + tokenizer: Any, prompt_file: str +) -> tuple[list[int], dict[str, Any]]: + prompt_path = Path(prompt_file).resolve() + prompt_file_exists = prompt_path.is_file() + prompt_text = ( + prompt_path.read_text(encoding="utf-8") + if prompt_file_exists + else FALLBACK_PROMPT + ).strip() + if not prompt_text: + raise RuntimeError(f"benchmark prompt file is empty: {prompt_path}") + if not getattr(tokenizer, "chat_template", None): + raise RuntimeError("model tokenizer does not define a chat template") + rendered_prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt_text}], + tokenize=False, + add_generation_prompt=True, + ) + # Match InfiniLM's benchmark path: use tokenizer.encode defaults after the + # model's chat template, then repeat this exact base sequence to each length. + token_ids = list(tokenizer.encode(rendered_prompt)) + if not token_ids: + raise RuntimeError("the fixed benchmark prompt tokenized to an empty list") + return token_ids, { + "prompt_source": ( + "file_chat_template" if prompt_file_exists else "embedded_fallback" + ), + "prompt_file": str(prompt_path), + "prompt_file_exists": prompt_file_exists, + "prompt_file_sha256": hashlib.sha256( + prompt_text.encode("utf-8") + ).hexdigest(), + "rendered_prompt_sha256": hashlib.sha256( + rendered_prompt.encode("utf-8") + ).hexdigest(), + } + + +def _repeat_prompt(token_ids: Sequence[int], target_length: int) -> list[int]: + if target_length < 1: + raise ValueError("target_length must be positive") + base = list(token_ids) + if target_length <= len(base): + # Preserve the assistant-generation suffix instead of cutting it off. + result = base[-target_length:] + else: + prefix_length = target_length - len(base) + repeats = (prefix_length + len(base) - 1) // len(base) + result = (base * repeats)[:prefix_length] + base + if len(result) != target_length: + raise RuntimeError(f"expected {target_length} prompt tokens, got {len(result)}") + return result + + +def _materialize_logits(logits: Any) -> Any: + # A replicated TP lm_head returns Tensor. Keep this guard for TP plans that + # leave the vocabulary output as a DTensor. + if type(logits).__name__ == "DTensor" and hasattr(logits, "full_tensor"): + return logits.full_tensor() + return logits + + +def _forward( + model: Any, + logits_limit_argument: str, + input_ids: Any, + past_key_values: Any | None = None, +) -> tuple[Any, Any]: + kwargs: dict[str, Any] = { + "input_ids": input_ids, + "past_key_values": past_key_values, + "use_cache": True, + "return_dict": True, + logits_limit_argument: 1, + } + outputs = model(**kwargs) + if outputs.past_key_values is None: + raise RuntimeError("model did not return past_key_values with use_cache=True") + logits = _materialize_logits(outputs.logits) + if logits.ndim != 3 or logits.shape[0] != input_ids.shape[0]: + raise RuntimeError(f"unexpected logits shape: {tuple(logits.shape)}") + if logits.shape[1] != 1: + raise RuntimeError( + f"expected one retained logits position, got shape {tuple(logits.shape)}" + ) + return logits[:, -1, :], outputs.past_key_values + + +def _validate_output( + generated: Any, + last_logits: Any, + batch_size: int, + output_tokens: int, + vocab_size: int, + torch: Any, + dist: Any, +) -> tuple[Any, dict[str, Any]]: + expected_shape = (batch_size, output_tokens) + if tuple(generated.shape) != expected_shape: + raise RuntimeError( + f"expected generated shape {expected_shape}, got {tuple(generated.shape)}" + ) + + rank_minimum = generated.clone() + rank_maximum = generated.clone() + dist.all_reduce(rank_minimum, op=dist.ReduceOp.MIN) + dist.all_reduce(rank_maximum, op=dist.ReduceOp.MAX) + rank_consensus = bool(torch.equal(rank_minimum, rank_maximum)) + + valid_ids = bool( + torch.logical_and(generated >= 0, generated < vocab_size).all().item() + ) + finite_logits = bool(torch.isfinite(last_logits).all().item()) + checks = torch.tensor( + [int(rank_consensus), int(valid_ids), int(finite_logits)], + dtype=torch.int32, + device=generated.device, + ) + dist.all_reduce(checks, op=dist.ReduceOp.MIN) + rank_consensus, valid_ids, finite_logits = [bool(value) for value in checks.tolist()] + if not (rank_consensus and valid_ids and finite_logits): + raise RuntimeError( + "correctness validation failed: " + f"rank_consensus={rank_consensus}, valid_ids={valid_ids}, " + f"finite_logits={finite_logits}" + ) + + generated_cpu = generated.cpu() + matrix = generated_cpu.tolist() + return generated_cpu, { + "exact_output_shape": True, + "rank_consensus": rank_consensus, + "valid_token_ids": valid_ids, + "finite_last_logits": finite_logits, + "output_token_ids_sha256": _stable_hash(matrix), + "first_request_first_16_tokens": matrix[0][:16], + } + + +def _run_iteration( + model: Any, + prompt_base: Sequence[int], + batch_size: int, + input_tokens: int, + output_tokens: int, + vocab_size: int, + logits_limit_argument: str, + torch: Any, + dist: Any, + device: Any, +) -> dict[str, Any]: + prompt = _repeat_prompt(prompt_base, input_tokens) + input_ids = ( + torch.tensor(prompt, dtype=torch.long, device=device) + .unsqueeze(0) + .expand(batch_size, -1) + .contiguous() + ) + + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + prefill_start = time.perf_counter() + logits, past_key_values = _forward( + model, logits_limit_argument, input_ids, past_key_values=None + ) + next_token = torch.argmax(logits, dim=-1) + torch.cuda.synchronize(device) + prefill_local_seconds = time.perf_counter() - prefill_start + prefill_seconds = _max_across_ranks( + prefill_local_seconds, torch, dist, device + ) + + generated_tokens = [next_token] + decode_start = time.perf_counter() + for _ in range(output_tokens - 1): + logits, past_key_values = _forward( + model, + logits_limit_argument, + next_token.unsqueeze(1), + past_key_values=past_key_values, + ) + next_token = torch.argmax(logits, dim=-1) + generated_tokens.append(next_token) + torch.cuda.synchronize(device) + decode_local_seconds = time.perf_counter() - decode_start + decode_seconds = _max_across_ranks(decode_local_seconds, torch, dist, device) + + peak_allocated_gib = _max_across_ranks( + torch.cuda.max_memory_allocated(device) / (1024**3), torch, dist, device + ) + peak_reserved_gib = _max_across_ranks( + torch.cuda.max_memory_reserved(device) / (1024**3), torch, dist, device + ) + generated = torch.stack(generated_tokens, dim=1) + generated_cpu, correctness = _validate_output( + generated, + logits, + batch_size, + output_tokens, + vocab_size, + torch, + dist, + ) + + total_prompt_tokens = batch_size * input_tokens + decode_token_count = batch_size * (output_tokens - 1) + generated_token_count = batch_size * output_tokens + total_seconds = prefill_seconds + decode_seconds + result = { + "batch_size": batch_size, + "input_tokens_per_request": input_tokens, + "prompt_token_ids_sha256": _stable_hash(prompt), + "output_tokens_per_request": output_tokens, + "total_context_tokens_per_request": input_tokens + output_tokens, + "total_prompt_tokens": total_prompt_tokens, + "total_generated_tokens": generated_token_count, + "ttft_seconds": prefill_seconds, + "prefill_tokens_per_second": total_prompt_tokens / prefill_seconds, + "decode_seconds": decode_seconds, + "decode_tokens_per_second": decode_token_count / decode_seconds, + "decode_tokens_per_second_per_request": ( + (output_tokens - 1) / decode_seconds + ), + "inter_token_latency_ms": decode_seconds * 1000.0 / (output_tokens - 1), + "generation_seconds": total_seconds, + "generated_tokens_per_second": generated_token_count / total_seconds, + "peak_memory_allocated_gib_max_rank": peak_allocated_gib, + "peak_memory_reserved_gib_max_rank": peak_reserved_gib, + "correctness": correctness, + "first_request_output_token_ids": generated_cpu[0].tolist(), + } + + del generated_cpu, generated, generated_tokens, logits, next_token + del past_key_values, input_ids + return result + + +def _median_summary(measurements: Sequence[dict[str, Any]]) -> dict[str, Any]: + fields = ( + "ttft_seconds", + "prefill_tokens_per_second", + "decode_seconds", + "decode_tokens_per_second", + "decode_tokens_per_second_per_request", + "inter_token_latency_ms", + "generation_seconds", + "generated_tokens_per_second", + "peak_memory_allocated_gib_max_rank", + "peak_memory_reserved_gib_max_rank", + ) + return { + f"median_{field}": statistics.median( + float(measurement[field]) for measurement in measurements + ) + for field in fields + } + + +def _per_length_medians( + measurements: Sequence[dict[str, Any]], input_lengths: Sequence[int] +) -> list[dict[str, Any]]: + summaries: list[dict[str, Any]] = [] + for input_tokens in input_lengths: + records = [ + measurement + for measurement in measurements + if int(measurement["input_tokens_per_request"]) == input_tokens + ] + if len(records) != REPEATS_PER_INPUT_LENGTH: + raise RuntimeError( + f"input length {input_tokens}: recorded {len(records)} repeats; " + f"expected {REPEATS_PER_INPUT_LENGTH}" + ) + summaries.append( + { + "input_tokens_per_request": input_tokens, + "measured_repeats": len(records), + **_median_summary(records), + } + ) + return summaries + + +def _overall_median_of_per_length_medians( + per_length_medians: Sequence[dict[str, Any]], +) -> dict[str, float]: + metric_names = [ + name for name in per_length_medians[0] if name.startswith("median_") + ] + return { + name: statistics.median( + float(length_summary[name]) for length_summary in per_length_medians + ) + for name in metric_names + } + + +def _run_worker_impl(args: argparse.Namespace, scenario: Scenario) -> int: + import inspect + + import torch + import torch.distributed as dist + import transformers + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + + grouped_mm_fallback = _install_hygon_grouped_mm_guard(torch) + + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", world_size)) + if rank != 0: + transformers.utils.logging.disable_progress_bar() + if world_size != args.tp_size: + raise RuntimeError( + f"expected WORLD_SIZE={args.tp_size}, got WORLD_SIZE={world_size}" + ) + if local_world_size != args.tp_size: + raise RuntimeError( + "this benchmark requires all TP ranks on one host: " + f"LOCAL_WORLD_SIZE={local_world_size}, TP={args.tp_size}" + ) + if not torch.cuda.is_available(): + raise RuntimeError("torch.cuda is unavailable") + if torch.cuda.device_count() < local_world_size: + raise RuntimeError( + f"need {local_world_size} visible GPUs, found {torch.cuda.device_count()}" + ) + + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + # Each torchrun worker owns one device. Avoid manual_seed_all(), which can + # make every worker initialize a context on all eight visible GPUs. + torch.random.default_generator.manual_seed(0) + torch.cuda.manual_seed(0) + + model_config = AutoConfig.from_pretrained( + args.model, + local_files_only=True, + trust_remote_code=False, + ) + architecture_signature = _validate_qwen3_235b_architecture(model_config) + quantization_config = getattr(model_config, "quantization_config", None) + if quantization_config: + raise RuntimeError( + "the Transformers benchmark is BF16-only; refusing quantized " + f"checkpoint {args.model!r} with quantization_config=" + f"{quantization_config!r}" + ) + tp_plan, tp_metadata = _build_qwen3_moe_tp_plan( + model_config, + args.tp_size, + scenario, + args.output_tokens, + ) + + load_start = time.perf_counter() + model = AutoModelForCausalLM.from_pretrained( + args.model, + config=model_config, + dtype=torch.bfloat16, + attn_implementation=args.attention, + tp_plan=tp_plan, + local_files_only=True, + low_cpu_mem_usage=True, + trust_remote_code=False, + ) + tp_validation = _validate_and_set_local_gqa(model, tp_metadata) + model.eval() + torch.cuda.synchronize(device) + if not dist.is_initialized(): + raise RuntimeError( + "Transformers tp_plan='auto' did not initialize torch.distributed" + ) + load_seconds = _max_across_ranks( + time.perf_counter() - load_start, torch, dist, device + ) + + tp_plan = getattr(model, "_tp_plan", None) + if not tp_plan: + raise RuntimeError("model loaded without a non-empty Transformers TP plan") + resolved_attention = getattr(model.config, "_attn_implementation", None) + if resolved_attention != args.attention: + raise RuntimeError( + f"requested attention={args.attention!r}, loaded model resolved " + f"attention={resolved_attention!r}" + ) + forward_parameters = inspect.signature(model.forward).parameters + if "logits_to_keep" in forward_parameters: + logits_limit_argument = "logits_to_keep" + elif "num_logits_to_keep" in forward_parameters: + logits_limit_argument = "num_logits_to_keep" + else: + raise RuntimeError( + "model.forward has no logits_to_keep argument; refusing to materialize " + "full [batch, context, vocab] logits for this benchmark" + ) + + tokenizer = AutoTokenizer.from_pretrained( + args.model, + local_files_only=True, + trust_remote_code=False, + use_fast=True, + ) + prompt_base, prompt_metadata = _make_prompt_base(tokenizer, args.prompt_file) + vocab_size = int(model.config.vocab_size) + if any(token < 0 or token >= vocab_size for token in prompt_base): + raise RuntimeError("fixed prompt contains a token outside model vocabulary") + + maximum_position_embeddings = int( + getattr(model.config, "max_position_embeddings", 0) or 0 + ) + maximum_requested = scenario.input_tokens + args.output_tokens + if maximum_position_embeddings and maximum_requested > maximum_position_embeddings: + raise RuntimeError( + f"requested sequence length {maximum_requested} exceeds " + f"max_position_embeddings={maximum_position_embeddings}" + ) + + config = { + "framework": "transformers", + "scenario": scenario.name, + "model_name": MODEL_NAME, + "model": str(Path(args.model).absolute()), + "model_realpath": str(Path(args.model).resolve()), + "model_class": type(model).__name__, + "dtype": "bfloat16", + "checkpoint_quantized": False, + "validated_qwen3_235b_architecture": architecture_signature, + "attention_implementation": resolved_attention, + "tp_plan": tp_metadata["tp_plan_mode"], + "tp_plan_rules": tp_plan, + "tp_plan_rule_count": len(tp_plan) if isinstance(tp_plan, dict) else None, + "tp_size": args.tp_size, + "smoke": args.smoke, + "batch_size": scenario.batch_size, + "input_lengths": list(scenario.input_lengths), + "output_tokens_per_request": args.output_tokens, + "total_context_tokens_per_request": maximum_requested, + "measured_iterations": MEASURED_ITERATIONS, + "measured_input_lengths": MEASURED_INPUT_LENGTHS, + "repeats_per_input_length": REPEATS_PER_INPUT_LENGTH, + "measurement_semantics": MEASUREMENT_SEMANTICS, + "model_load_seconds": load_seconds, + "fixed_prompt_base_tokens": len(prompt_base), + "fixed_prompt_base_sha256": _stable_hash(prompt_base), + **prompt_metadata, + "torch_version": torch.__version__, + "transformers_version": transformers.__version__, + "flash_attn_version": _package_version("flash-attn"), + "hygon_transformers_grouped_mm_fallback": grouped_mm_fallback, + "gpu_name": torch.cuda.get_device_name(device), + **tp_metadata, + **tp_validation, + } + _emit(rank, "PYTORCH_QWEN3_235B_CONFIG", config) + + measurements: list[dict[str, Any]] = [] + iteration = 0 + for length_index, input_tokens in enumerate(scenario.input_lengths, start=1): + with torch.inference_mode(): + shape_warmup = _run_iteration( + model, + prompt_base, + scenario.batch_size, + input_tokens, + args.output_tokens, + vocab_size, + logits_limit_argument, + torch, + dist, + device, + ) + shape_warmup_hash = shape_warmup["correctness"][ + "output_token_ids_sha256" + ] + shape_warmup_prompt_hash = shape_warmup["prompt_token_ids_sha256"] + _emit( + rank, + "PYTORCH_QWEN3_235B_SHAPE_WARMUP", + { + "scenario": scenario.name, + "length_index": length_index, + "batch_size": scenario.batch_size, + "input_tokens_per_request": input_tokens, + "prompt_token_ids_sha256": shape_warmup_prompt_hash, + "output_tokens_per_request": args.output_tokens, + "output_token_ids_sha256": shape_warmup_hash, + "correctness": shape_warmup["correctness"], + }, + ) + del shape_warmup + gc.collect() + + for repeat in range(1, REPEATS_PER_INPUT_LENGTH + 1): + iteration += 1 + with torch.inference_mode(): + measurement = _run_iteration( + model, + prompt_base, + scenario.batch_size, + input_tokens, + args.output_tokens, + vocab_size, + logits_limit_argument, + torch, + dist, + device, + ) + measured_hash = measurement["correctness"][ + "output_token_ids_sha256" + ] + measured_prompt_hash = measurement["prompt_token_ids_sha256"] + if measured_prompt_hash != shape_warmup_prompt_hash: + raise RuntimeError( + f"input length {input_tokens} repeat {repeat}: measured prompt " + f"hash {measured_prompt_hash} does not match exact-shape " + f"warmup prompt hash {shape_warmup_prompt_hash}" + ) + # Hygon BF16 kernels can make numerically valid MoE routing choices + # differ across independent runs. Keep the replay hash observable, + # while treating per-run shape/range/finite/rank checks as correctness. + measurement["correctness"]["output_matches_exact_shape_warmup"] = ( + measured_hash == shape_warmup_hash + ) + measurement["exact_shape_warmup_output_sha256"] = shape_warmup_hash + measurement = { + "scenario": scenario.name, + "iteration": iteration, + "length_index": length_index, + "repeat": repeat, + **measurement, + } + output_token_ids = measurement.pop("first_request_output_token_ids") + decoded_output = tokenizer.decode( + output_token_ids, skip_special_tokens=True + ).strip() + measurement["correctness"]["decoded_output"] = ( + _validate_decoded_output(decoded_output) + ) + measurements.append(measurement) + _print_infinilm_style_metrics(rank, measurement, decoded_output) + _emit(rank, "PYTORCH_QWEN3_235B_ITERATION", measurement) + gc.collect() + + if len(measurements) != MEASURED_ITERATIONS: + raise RuntimeError( + f"expected {MEASURED_ITERATIONS} measurements, got {len(measurements)}" + ) + per_length_medians = _per_length_medians( + measurements, scenario.input_lengths + ) + overall_medians = _overall_median_of_per_length_medians(per_length_medians) + summary = { + "scenario": scenario.name, + "status": "PASS", + "measured_iterations": len(measurements), + "measured_input_lengths": MEASURED_INPUT_LENGTHS, + "repeats_per_input_length": REPEATS_PER_INPUT_LENGTH, + "measurement_semantics": MEASUREMENT_SEMANTICS, + "input_lengths": list(scenario.input_lengths), + "batch_size": scenario.batch_size, + "output_tokens_per_request": args.output_tokens, + "total_context_tokens_per_request": maximum_requested, + "per_length_medians": per_length_medians, + "overall_aggregate": { + "aggregation_method": "median_of_three_fixed_shape_measurements", + "measurement_count": len(measurements), + "mixed_input_lengths": False, + **overall_medians, + }, + # Compatibility aliases for existing table consumers. Their scope is the + # explicitly labeled overall aggregate above, not a single input length. + **overall_medians, + "output_token_ids_sha256": [ + item["correctness"]["output_token_ids_sha256"] for item in measurements + ], + } + _emit(rank, "PYTORCH_QWEN3_235B_SUMMARY", summary) + return len(measurements) + + +def _run_worker(args: argparse.Namespace, scenario: Scenario) -> None: + import torch.distributed as dist + + rank = int(os.environ["RANK"]) + measured_iterations = 0 + caught: BaseException | None = None + caught_traceback: Any = None + teardown_errors: list[str] = [] + process_group_was_initialized = False + try: + measured_iterations = _run_worker_impl(args, scenario) + except BaseException as error: + caught = error + caught_traceback = error.__traceback__ + finally: + process_group_was_initialized = dist.is_initialized() + if process_group_was_initialized: + if caught is None: + try: + dist.barrier() + except BaseException as error: + caught = error + caught_traceback = error.__traceback__ + teardown_errors.append( + f"barrier: {type(error).__name__}: {error}" + ) + try: + dist.destroy_process_group() + except BaseException as error: + if caught is None: + caught = error + caught_traceback = error.__traceback__ + teardown_errors.append( + f"destroy_process_group: {type(error).__name__}: {error}" + ) + + teardown_complete = not dist.is_initialized() and not teardown_errors + status = ( + "PASS" + if caught is None + and measured_iterations == MEASURED_ITERATIONS + and teardown_complete + else "ERROR" + ) + completion: dict[str, Any] = { + "scenario": scenario.name, + "status": status, + "exit_code": 0 if status == "PASS" else 1, + "measured_iterations": measured_iterations, + "measured_input_lengths": MEASURED_INPUT_LENGTHS, + "repeats_per_input_length": REPEATS_PER_INPUT_LENGTH, + "measurement_semantics": MEASUREMENT_SEMANTICS, + "process_group_was_initialized": process_group_was_initialized, + "distributed_teardown_complete": teardown_complete, + } + if caught is not None: + completion["error"] = { + "type": type(caught).__name__, + "message": str(caught), + } + if teardown_errors: + completion["teardown_errors"] = teardown_errors + _emit(rank, "PYTORCH_QWEN3_235B_COMPLETE", completion) + + if caught is not None: + raise caught.with_traceback(caught_traceback) + + +def main(scenario: Scenario) -> None: + args = _parse_args(scenario) + scenario = _effective_scenario(scenario, args.smoke) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if "LOCAL_RANK" not in os.environ and world_size == 1: + _require_idle_gpu() + _launch_torchrun(args) + raise AssertionError("os.execvpe returned unexpectedly") + _run_worker(args, scenario) diff --git a/test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py b/test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py new file mode 100755 index 000000000..5474ca990 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py @@ -0,0 +1,491 @@ +#!/usr/bin/env python3 +"""Calculate true token-level PPL for Qwen3_235B with Transformers TP8. + +The input is a framework-neutral token manifest. Both the Transformers and +InfiniLM runners must consume the same manifest so their PPL values score the +same target tokens instead of independently tokenizing the source corpus. +""" + +from __future__ import annotations + +import argparse +import gc +import inspect +import json +import math +import os +import sys +import time +from pathlib import Path +from typing import Any + + +SCRIPT_DIR = Path(__file__).resolve().parent +SCRIPTS_DIR = SCRIPT_DIR.parent +for import_path in (SCRIPT_DIR, SCRIPTS_DIR): + if str(import_path) not in sys.path: + sys.path.insert(0, str(import_path)) + +import _pytorch_runner as benchmark_runner +from _ppl_common import ( + SCORING_METHOD, + canonical_indices_sha256, + iter_sliding_windows, + load_manifest, +) + + +RESULT_SCHEMA = "qwen3_235b_true_ppl_result/v1" +DEFAULT_MODEL = "/data1/Qwen3_235B" +DEFAULT_WINDOW_SIZE = 256 +DEFAULT_STRIDE = 128 +DEFAULT_MAX_SCORED_TOKENS = 10240 +EXPECTED_VOCAB_SIZE = 151936 + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="True shifted-token PPL for Qwen3_235B BF16 on Hygon TP8" + ) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--token-manifest", required=True) + parser.add_argument("--window", type=int, default=DEFAULT_WINDOW_SIZE) + parser.add_argument("--stride", type=int, default=DEFAULT_STRIDE) + parser.add_argument( + "--max-scored-tokens", + type=int, + default=DEFAULT_MAX_SCORED_TOKENS, + help="maximum shifted target tokens to score; 0 scores the full manifest", + ) + parser.add_argument("--tp-size", type=int, default=8) + parser.add_argument( + "--attention", + choices=("eager",), + default="eager", + help="BW1100 correctness path; SDPA is unsupported for this model stack", + ) + parser.add_argument( + "--json-output", + help="optional rank-0 result path; the result is always printed as JSON", + ) + args = parser.parse_args() + + model_path = Path(args.model) + manifest_path = Path(args.token_manifest) + if not model_path.is_dir(): + parser.error(f"model directory does not exist: {model_path}") + if not manifest_path.is_file(): + parser.error(f"token manifest does not exist: {manifest_path}") + if args.window < 2: + parser.error("--window must be at least 2") + if args.stride < 1 or args.stride >= args.window: + parser.error("--stride must satisfy 1 <= stride < window") + if args.max_scored_tokens < 0: + parser.error("--max-scored-tokens must be non-negative") + if args.tp_size < 1: + parser.error("--tp-size must be positive") + + args.model = str(model_path.resolve()) + args.token_manifest = str(manifest_path.resolve()) + if args.json_output: + args.json_output = str(Path(args.json_output).resolve()) + return args + +def _write_json_atomic(path_value: str, payload: dict[str, Any]) -> None: + path = Path(path_value) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") + temporary.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + +def _run_worker_impl(args: argparse.Namespace) -> dict[str, Any]: + import torch + import torch.distributed as dist + import torch.nn.functional as functional + import transformers + from transformers import AutoConfig, AutoModelForCausalLM + + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", world_size)) + if world_size != args.tp_size or local_world_size != args.tp_size: + raise RuntimeError( + "PPL runner requires one-host TP with WORLD_SIZE=" + f"LOCAL_WORLD_SIZE={args.tp_size}; got {world_size}/{local_world_size}" + ) + if not torch.cuda.is_available() or torch.cuda.device_count() < local_world_size: + raise RuntimeError( + f"need {local_world_size} visible GPUs, found {torch.cuda.device_count()}" + ) + + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + # RCCL initializes lazily. Reserve communicator memory before the 235B + # checkpoint consumes nearly all device memory. + communicator_probe = torch.ones(1, dtype=torch.int32, device=device) + dist.all_reduce(communicator_probe) + if int(communicator_probe.item()) != args.tp_size: + raise RuntimeError("Transformers TP8 RCCL communicator probe failed") + dist.barrier() + torch.cuda.synchronize(device) + del communicator_probe + torch.random.default_generator.manual_seed(0) + torch.cuda.manual_seed(0) + if rank != 0: + transformers.utils.logging.disable_progress_bar() + + corpus_manifest = load_manifest(args.token_manifest) + token_ids = corpus_manifest.token_ids + corpus = { + "manifest_path": str(corpus_manifest.path.resolve()), + "manifest_sha256": corpus_manifest.manifest_sha256, + "token_ids_sha256": corpus_manifest.token_ids_sha256, + "token_count": corpus_manifest.token_count, + "source_sha256": corpus_manifest.payload["source_sha256"], + "tokenizer_sha256": corpus_manifest.payload["tokenizer_sha256"], + "source": corpus_manifest.payload.get("source"), + "tokenizer": corpus_manifest.payload.get("tokenizer"), + } + model_config = AutoConfig.from_pretrained( + args.model, local_files_only=True, trust_remote_code=False + ) + architecture = benchmark_runner._validate_qwen3_235b_architecture(model_config) + if getattr(model_config, "quantization_config", None): + raise RuntimeError("Transformers PPL baseline requires the BF16 checkpoint") + vocab_size = int(model_config.vocab_size) + if vocab_size != EXPECTED_VOCAB_SIZE: + raise RuntimeError( + f"expected complete Qwen3_235B vocabulary {EXPECTED_VOCAB_SIZE}, " + f"got {vocab_size}" + ) + if any(token >= vocab_size for token in token_ids): + raise RuntimeError("token manifest contains an ID outside the model vocabulary") + maximum_positions = int( + getattr(model_config, "max_position_embeddings", 0) or 0 + ) + if maximum_positions and args.window > maximum_positions: + raise RuntimeError( + f"window={args.window} exceeds max_position_embeddings={maximum_positions}" + ) + + scoring_scenario = benchmark_runner.Scenario("true_ppl", 1, args.window, 1) + tp_plan, tp_metadata = benchmark_runner._build_qwen3_moe_tp_plan( + model_config, args.tp_size, scoring_scenario, 1 + ) + grouped_mm_fallback = benchmark_runner._install_hygon_grouped_mm_guard(torch) + load_start = time.perf_counter() + model = AutoModelForCausalLM.from_pretrained( + args.model, + config=model_config, + dtype=torch.bfloat16, + attn_implementation=args.attention, + tp_plan=tp_plan, + local_files_only=True, + low_cpu_mem_usage=True, + trust_remote_code=False, + ) + tp_validation = benchmark_runner._validate_and_set_local_gqa(model, tp_metadata) + model.eval() + torch.cuda.synchronize(device) + load_seconds = benchmark_runner._max_across_ranks( + time.perf_counter() - load_start, torch, dist, device + ) + + loaded_tp_plan = getattr(model, "_tp_plan", None) + if not loaded_tp_plan: + raise RuntimeError("model loaded without a non-empty Transformers TP plan") + resolved_attention = getattr(model.config, "_attn_implementation", None) + if resolved_attention != args.attention: + raise RuntimeError( + f"requested attention={args.attention!r}, resolved={resolved_attention!r}" + ) + forward_parameters = inspect.signature(model.forward).parameters + if "logits_to_keep" in forward_parameters: + logits_limit_argument = "logits_to_keep" + elif "num_logits_to_keep" in forward_parameters: + logits_limit_argument = "num_logits_to_keep" + else: + raise RuntimeError( + "model.forward has no logits_to_keep argument; refusing full-context " + "vocabulary materialization" + ) + + available_targets = len(token_ids) - 1 + scored_token_count = ( + available_targets + if args.max_scored_tokens == 0 + else min(args.max_scored_tokens, available_targets) + ) + if scored_token_count < 1: + raise RuntimeError("the selected corpus range contains no shifted target token") + first_scored_token_index = 1 + last_scored_token_index_exclusive = 1 + scored_token_count + scored_token_indices_sha256 = canonical_indices_sha256( + range(first_scored_token_index, last_scored_token_index_exclusive) + ) + scoring_method = SCORING_METHOD + + config = { + "backend": "transformers", + "model": args.model, + "dtype": "bfloat16", + "tp_size": args.tp_size, + "attention": resolved_attention, + "window_size": args.window, + "stride": args.stride, + "requested_max_scored_tokens": args.max_scored_tokens, + "scored_token_count": scored_token_count, + "scoring_method": scoring_method, + "first_scored_token_index": first_scored_token_index, + "last_scored_token_index_exclusive": last_scored_token_index_exclusive, + "scored_token_indices_sha256": scored_token_indices_sha256, + "corpus_manifest": corpus, + "vocab_size": vocab_size, + "architecture": architecture, + "model_load_seconds": load_seconds, + "tp_plan": tp_metadata, + "tp_validation": tp_validation, + "hygon_transformers_grouped_mm_fallback": grouped_mm_fallback, + "torch_version": torch.__version__, + "transformers_version": transformers.__version__, + } + if rank == 0: + print( + "PYTORCH_QWEN3_235B_PPL_CONFIG " + + json.dumps(config, ensure_ascii=False, sort_keys=True), + flush=True, + ) + + total_nll = 0.0 + windows: list[dict[str, Any]] = [] + scored_by_windows = 0 + torch.cuda.synchronize(device) + scoring_start = time.perf_counter() + with torch.inference_mode(): + for window in iter_sliding_windows( + token_ids, + args.window, + args.stride, + None if args.max_scored_tokens == 0 else args.max_scored_tokens, + ): + display_index = window.index + 1 + target_count = window.scored_token_count + input_slice = window.token_ids + expected_prediction_start = len(input_slice) - target_count - 1 + if ( + window.prediction_start != expected_prediction_start + or window.prediction_end != len(input_slice) - 1 + ): + raise RuntimeError( + f"window {display_index} retained-logits alignment is invalid" + ) + input_ids = torch.tensor( + input_slice, dtype=torch.long, device=device + ).unsqueeze(0) + outputs = model( + input_ids=input_ids, + use_cache=False, + return_dict=True, + **{logits_limit_argument: target_count + 1}, + ) + logits = benchmark_runner._materialize_logits(outputs.logits) + expected_shape = (1, target_count + 1, vocab_size) + if tuple(logits.shape) != expected_shape: + raise RuntimeError( + "incomplete or unexpected logits: " + f"got {tuple(logits.shape)}, expected {expected_shape}" + ) + score_logits = logits[:, :-1, :] + labels = torch.tensor( + input_slice[window.target_start : window.target_end], + dtype=torch.long, + device=device, + ).unsqueeze(0) + finite = torch.isfinite(score_logits).all().to(dtype=torch.int32) + dist.all_reduce(finite, op=dist.ReduceOp.MIN) + if not bool(finite.item()): + raise RuntimeError( + f"window {display_index} contains non-finite logits" + ) + window_nll_tensor = functional.cross_entropy( + score_logits.float().reshape(-1, vocab_size), + labels.reshape(-1), + reduction="sum", + ) + window_nll = float(window_nll_tensor.double().item()) + nll_min = torch.tensor(window_nll, dtype=torch.float64, device=device) + nll_max = nll_min.clone() + dist.all_reduce(nll_min, op=dist.ReduceOp.MIN) + dist.all_reduce(nll_max, op=dist.ReduceOp.MAX) + rank_delta_per_token = float((nll_max - nll_min).item()) / target_count + if rank_delta_per_token > 1e-4: + raise RuntimeError( + f"window {display_index} rank NLL mismatch: " + f"delta/token={rank_delta_per_token:.6g}" + ) + total_nll += window_nll + scored_by_windows += target_count + windows.append( + { + "index": window.index, + "token_start": window.token_start, + "token_end": window.token_end, + "score_start": window.score_start, + "score_end": window.score_end, + "input_token_count": len(input_slice), + "scored_token_count": target_count, + "nll": window_nll, + } + ) + del outputs, logits, score_logits, labels, window_nll_tensor, input_ids + + torch.cuda.synchronize(device) + scoring_seconds = benchmark_runner._max_across_ranks( + time.perf_counter() - scoring_start, torch, dist, device + ) + if scored_by_windows != scored_token_count: + raise RuntimeError( + f"scored {scored_by_windows} tokens, expected {scored_token_count}" + ) + + mean_nll = total_nll / scored_token_count + if not math.isfinite(mean_nll): + raise RuntimeError(f"mean NLL is not finite: {mean_nll}") + try: + ppl = math.exp(mean_nll) + except OverflowError as error: + raise RuntimeError(f"PPL overflows float64 at mean NLL={mean_nll}") from error + if not math.isfinite(ppl): + raise RuntimeError(f"PPL is not finite: {ppl}") + + result: dict[str, Any] = {} + if rank == 0: + result = { + "schema": RESULT_SCHEMA, + "status": "PASS", + "backend": "transformers", + "model": args.model, + "dtype": "bfloat16", + "tp_size": args.tp_size, + "attention": resolved_attention, + "corpus_manifest": args.token_manifest, + "corpus_manifest_sha256": corpus["manifest_sha256"], + "corpus_token_ids_sha256": corpus["token_ids_sha256"], + "corpus_token_count": corpus["token_count"], + "window_size": args.window, + "stride": args.stride, + "scoring_method": scoring_method, + "first_scored_token_index": first_scored_token_index, + "last_scored_token_index_exclusive": ( + last_scored_token_index_exclusive + ), + "scored_token_indices_sha256": scored_token_indices_sha256, + "scored_token_count": scored_token_count, + "total_nll": total_nll, + "mean_nll": mean_nll, + "ppl": ppl, + "windows": windows, + "window_count": len(windows), + "scoring_seconds": scoring_seconds, + "scored_tokens_per_second": scored_token_count / scoring_seconds, + "model_load_seconds": load_seconds, + "vocab_size": vocab_size, + "full_vocab_logits_validated_every_window": True, + } + if args.json_output: + _write_json_atomic(args.json_output, result) + print( + "PYTORCH_QWEN3_235B_PPL_RESULT " + + json.dumps(result, ensure_ascii=False, sort_keys=True), + flush=True, + ) + print( + f"Transformers true PPL: {ppl:.6f} " + f"(mean NLL={mean_nll:.6f}, tokens={scored_token_count})", + flush=True, + ) + + del model + gc.collect() + torch.cuda.empty_cache() + return result + + +def _run_worker(args: argparse.Namespace) -> int: + import torch.distributed as dist + + rank = int(os.environ["RANK"]) + caught: BaseException | None = None + caught_traceback: Any = None + teardown_errors: list[str] = [] + try: + _run_worker_impl(args) + except BaseException as error: + caught = error + caught_traceback = error.__traceback__ + finally: + initialized = dist.is_initialized() + if initialized: + if caught is None: + try: + dist.barrier() + except BaseException as error: + caught = error + caught_traceback = error.__traceback__ + teardown_errors.append( + f"barrier: {type(error).__name__}: {error}" + ) + try: + dist.destroy_process_group() + except BaseException as error: + if caught is None: + caught = error + caught_traceback = error.__traceback__ + teardown_errors.append( + f"destroy_process_group: {type(error).__name__}: {error}" + ) + status = "PASS" if caught is None and not teardown_errors else "ERROR" + if rank == 0: + completion: dict[str, Any] = { + "schema": RESULT_SCHEMA, + "status": status, + "exit_code": 0 if status == "PASS" else 1, + "distributed_teardown_complete": not dist.is_initialized(), + } + if caught is not None: + completion["error"] = { + "type": type(caught).__name__, + "message": str(caught), + } + if teardown_errors: + completion["teardown_errors"] = teardown_errors + print( + "PYTORCH_QWEN3_235B_PPL_COMPLETE " + + json.dumps(completion, ensure_ascii=False, sort_keys=True), + flush=True, + ) + if caught is not None: + raise caught.with_traceback(caught_traceback) + return 0 + + +def main() -> int: + args = _parse_args() + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if "LOCAL_RANK" not in os.environ and world_size == 1: + # Fail on corpus/schema errors before reserving all eight devices. + load_manifest(args.token_manifest) + benchmark_runner._require_idle_gpu() + benchmark_runner._launch_torchrun(args) + raise AssertionError("os.execvpe returned unexpectedly") + return _run_worker(args) + + +if __name__ == "__main__": + raise SystemExit(main()) From 55abe99b11b13568b09e940dd964c5cac4f52424 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Thu, 23 Jul 2026 18:31:07 +0800 Subject: [PATCH 09/14] feat(bench): support paired multi-length cases --- README.md | 9 ++ examples/bench.py | 225 +++++++++++++++++++++------------ python/infinilm/base_config.py | 16 ++- test/bench/test_bench_cases.py | 161 +++++++++++++++++++++++ 4 files changed, 326 insertions(+), 85 deletions(-) create mode 100644 test/bench/test_bench_cases.py diff --git a/README.md b/README.md index fa8af49c2..8c993f19b 100644 --- a/README.md +++ b/README.md @@ -159,6 +159,15 @@ > 注意：`--cache-dir` 应指向包含 `ceval___ceval-exam` 和 `cais___mmlu` 等数据集子目录的父目录，而不是直接指向这些子目录 - 试验中功能 + - 单次加载模型测试多组长度 + ```bash + python examples/bench.py --device nvidia --model= --batch-size=4 --input-len=2048,4096 --output-len=512,128 --warmup + ``` + `--input-len` 和 `--output-len` 按位置组成 `(2048, 512)`、 + `(4096, 128)` 两个 case，不生成笛卡尔积。任意一侧只有一个值时， + 该值会广播到另一侧的所有长度；两个参数都只有一个值时，行为与原单 + case 命令一致。模型只加载一次，每个不同的 `(batch_size, input_len)` + prefill shape 各 warmup 一次。 - Warm Up ```bash python examples/bench.py --device nvidia --model= --warmup diff --git a/examples/bench.py b/examples/bench.py index b991078cd..755c8c29d 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -89,6 +89,65 @@ def _normalize_config(config, model_type): # OUTPUT_LENS = [256, 1024, 4096] +def pair_sequence_lengths( + input_len_list: list[int], output_len_list: list[int] +) -> list[tuple[int, int]]: + """Return positional input/output length pairs. + + Lists with equal lengths are paired by position, not expanded as a + Cartesian product. If either side contains one value, that value is + broadcast across the other side. + """ + if not input_len_list or not output_len_list: + raise ValueError("input_len and output_len must not be empty") + if any(length <= 0 for length in input_len_list): + raise ValueError(f"input_len values must be positive: {input_len_list}") + if any(length <= 0 for length in output_len_list): + raise ValueError(f"output_len values must be positive: {output_len_list}") + + if len(input_len_list) == len(output_len_list): + return list(zip(input_len_list, output_len_list)) + if len(input_len_list) == 1: + return [(input_len_list[0], output_len) for output_len in output_len_list] + if len(output_len_list) == 1: + return [(input_len, output_len_list[0]) for input_len in input_len_list] + raise ValueError( + "input_len and output_len must have the same number of values, " + "or one side must contain a single value for broadcasting: " + f"input_len={input_len_list}, output_len={output_len_list}" + ) + + +def get_paged_kv_cache_num_blocks(cases, block_size: int) -> int: + """Return the shared paged-cache capacity required by sequential cases.""" + if block_size <= 0: + raise ValueError(f"block_size must be positive: {block_size}") + + case_list = list(cases) + if not case_list: + raise ValueError("at least one benchmark case is required") + + return max( + ( + (case["input_len"] + case["output_len"] + block_size - 1) + // block_size + ) + * case["batch_size"] + for case in case_list + ) + + +def get_warmup_shapes(cases) -> OrderedDict: + """Map each prefill shape to the largest output capacity it needs.""" + warmup_shapes = OrderedDict() + for case in cases: + shape = (case["batch_size"], case["input_len"]) + warmup_shapes[shape] = max( + warmup_shapes.get(shape, 0), case["output_len"] + ) + return warmup_shapes + + def read_json_file(file_path): """Load and return JSON content from file_path.""" with open(file_path, "r") as file: @@ -102,9 +161,19 @@ def get_test_cases( output_len_list: list[int], use_mla: bool = False, ): + """Generate cases from batch sizes and positional length pairs. + + Batch sizes are combined with each input/output pair. The two length lists + themselves are paired by position (or single-value broadcast), never as a + Cartesian product. Returned cases are ordered by ascending KV-cache usage. + """ model_path = os.path.expanduser(model_path) - """Generate cases ordered by ascending KV cache memory usage.""" + if not batch_size_list or any( + batch_size <= 0 for batch_size in batch_size_list + ): + raise ValueError(f"batch_size values must be positive: {batch_size_list}") + # Load model config to derive attention dimensions config = read_json_file(os.path.join(model_path, "config.json")) model_type = config.get("model_type", "") @@ -125,32 +194,34 @@ def get_test_cases( num_key_value_heads = config.get("num_key_value_heads") num_hidden_layers = config.get("num_hidden_layers") - # Enumerate all batch/input/output combinations and compute KV cache size + length_pairs = pair_sequence_lengths(input_len_list, output_len_list) + + # Each input/output list position is one case. A one-element list is + # broadcast so one input length can still be tested with many output lengths. case_list = [] for batch_size in batch_size_list: - for input_len in input_len_list: - for output_len in output_len_list: - for data_type in ["bfloat16"]: - data_type_bytes = DATA_TYPE_BYTES[data_type] - - total_seq_len = input_len + output_len - kvcache_memory_bytes = ( - data_type_bytes - * (batch_size * total_seq_len * num_key_value_heads * head_dim) - * num_hidden_layers - ) - kvcache_memory_gb = kvcache_memory_bytes / (1024 * 1024 * 1024) - - case_list.append( - { - "idx": len(case_list), - "batch_size": batch_size, - "input_len": input_len, - "output_len": output_len, - "data_type": data_type, - "kvcache_memory": round(kvcache_memory_gb, 3), - } - ) + for input_len, output_len in length_pairs: + for data_type in ["bfloat16"]: + data_type_bytes = DATA_TYPE_BYTES[data_type] + + total_seq_len = input_len + output_len + kvcache_memory_bytes = ( + data_type_bytes + * (batch_size * total_seq_len * num_key_value_heads * head_dim) + * num_hidden_layers + ) + kvcache_memory_gb = kvcache_memory_bytes / (1024 * 1024 * 1024) + + case_list.append( + { + "idx": len(case_list), + "batch_size": batch_size, + "input_len": input_len, + "output_len": output_len, + "data_type": data_type, + "kvcache_memory": round(kvcache_memory_gb, 3), + } + ) # Sort by KV cache size and wrap in OrderedDict with index keys case_dict = OrderedDict( @@ -414,15 +485,11 @@ def run( # -------------------------------------------------------- # if enable_paged_attn: paged_kv_block_size = _PAGED_KV_BLOCK_SIZE - max_num_blocks = max( - [ - ( - (c_["input_len"] + c_["output_len"] + (paged_kv_block_size - 1)) - // paged_kv_block_size - ) - * c_["batch_size"] - for _, c_ in cases_dict.items() - ] + # Cases run sequentially and each generate call rebuilds block tables + # from block zero, so the shared cache needs the largest case capacity, + # not the sum of all case capacities. + max_num_blocks = get_paged_kv_cache_num_blocks( + cases_dict.values(), paged_kv_block_size ) cache_config = PagedKVCacheConfig(max_num_blocks, paged_kv_block_size) else: @@ -454,57 +521,49 @@ def run( if cfg.warmup: warmup_steps = 1 - # warmup cache capacity - warmup_case = next(iter(cases_dict.values())) - warmup_batch = warmup_case["batch_size"] - warmup_input_len = warmup_case["input_len"] - warmup_decode_len = 5 - - if enable_paged_attn: - warmup_num_blocks = ( - (warmup_input_len + warmup_decode_len + paged_kv_block_size - 1) - // paged_kv_block_size - ) * warmup_batch - warmup_cache_config = PagedKVCacheConfig( - warmup_num_blocks, paged_kv_block_size - ) - else: - warmup_cache_config = StaticKVCacheConfig( - max_batch_size=warmup_batch, - max_cache_len=warmup_input_len + warmup_decode_len, - ) - - test.model.reset_cache(warmup_cache_config) - - warmup_prompt_ids = repeat_prompt(test.input_ids_list[0], warmup_input_len) - warmup_ids = [warmup_prompt_ids] * warmup_batch - - input_ids_infini = infinicore.from_list(warmup_ids, dtype=infinicore.int64) + # Warm every distinct prefill shape once. Repeated benchmark cases keep + # a single warmup, while mixed input lengths do not include first-use + # graph/kernel setup in their measured run. + warmup_shapes = get_warmup_shapes(cases_dict.values()) + + for warmup_idx, ((warmup_batch, warmup_input_len), max_output_len) in enumerate( + warmup_shapes.items(), start=1 + ): + warmup_decode_len = min(5, max_output_len) + if not enable_paged_attn: + # Reserve the largest complete case for this prefill shape, + # even though warmup itself only runs a few decode steps. + warmup_cache_config = StaticKVCacheConfig( + max_batch_size=warmup_batch, + max_cache_len=warmup_input_len + max_output_len, + ) + test.model.reset_cache(warmup_cache_config) - print( - f"\033[93m[warmup] batch={warmup_batch}, input_len={warmup_input_len}, " - f"will prefill + {warmup_decode_len} decode steps\033[0m" - ) - print("=================== warmup start ===================") - - for _ in range(warmup_steps): - _ = test.model.generate( - input_ids_infini, - GenerationConfig( - max_new_tokens=warmup_decode_len, # decode kernel warmup - temperature=cfg.temperature, - top_k=cfg.top_k, - top_p=cfg.top_p, - stop_on_eos=False, - ), - _measure_and_log_time=False, + warmup_prompt_ids = repeat_prompt(test.input_ids_list[0], warmup_input_len) + warmup_ids = [warmup_prompt_ids] * warmup_batch + input_ids_infini = infinicore.from_list( + warmup_ids, dtype=infinicore.int64 ) - print("=================== warmup done ====================") - - # reset cache back to benchmark config - if cache_config is not None: - test.model.reset_cache(cache_config) + print( + f"\033[93m[warmup {warmup_idx}/{len(warmup_shapes)}] " + f"batch={warmup_batch}, input_len={warmup_input_len}, " + f"will prefill + {warmup_decode_len} decode steps\033[0m" + ) + print("=================== warmup start ===================") + for _ in range(warmup_steps): + _ = test.model.generate( + input_ids_infini, + GenerationConfig( + max_new_tokens=warmup_decode_len, + temperature=cfg.temperature, + top_k=cfg.top_k, + top_p=cfg.top_p, + stop_on_eos=False, + ), + _measure_and_log_time=False, + ) + print("=================== warmup done ====================") # ---------------------------------------------------------------------------- # # Warmup done @@ -518,7 +577,7 @@ def run( output_len = case["output_len"] if not enable_paged_attn: - # reset cache if static kvcache is used + # Each static-cache case gets its exact full generation capacity. initial_capacity = input_len + output_len test.model.reset_cache( StaticKVCacheConfig( diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index aa7d11890..c92c21efb 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -305,10 +305,22 @@ def _add_common_args(self): help="maximum batch size for server", ) self.parser.add_argument( - "--input-len", type=parse_list, default=10, help="input sequence length" + "--input-len", + type=parse_list, + default=10, + help=( + "input sequence length; examples/bench.py pairs comma-separated " + "input/output values by position and broadcasts a single value" + ), ) self.parser.add_argument( - "--output-len", type=parse_list, default=20, help="output sequence length" + "--output-len", + type=parse_list, + default=20, + help=( + "output sequence length; examples/bench.py pairs comma-separated " + "input/output values by position and broadcasts a single value" + ), ) self.parser.add_argument( "--max-new-tokens", diff --git a/test/bench/test_bench_cases.py b/test/bench/test_bench_cases.py new file mode 100644 index 000000000..bd5410ef1 --- /dev/null +++ b/test/bench/test_bench_cases.py @@ -0,0 +1,161 @@ +import importlib.util +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import mock_open, patch + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _module(name, **attributes): + module = types.ModuleType(name) + for key, value in attributes.items(): + setattr(module, key, value) + return module + + +def _load_bench_module(): + infinicore = _module("infinicore") + infinicore.nn = types.SimpleNamespace(Module=object) + infinicore.device = lambda *_args, **_kwargs: object() + + infinilm = _module("infinilm") + infinilm.__path__ = [] + stub_modules = { + "infinicore": infinicore, + "infinilm": infinilm, + "infinilm.modeling_utils": _module( + "infinilm.modeling_utils", load_model_state_dict_by_file=lambda *_args: None + ), + "infinilm.distributed": _module("infinilm.distributed", DistConfig=object), + "infinilm.infer_engine": _module( + "infinilm.infer_engine", GenerationConfig=object, InferEngine=object + ), + "infinilm.base_config": _module("infinilm.base_config", BaseConfig=object), + "infinilm.cache": _module( + "infinilm.cache", + StaticKVCacheConfig=object, + PagedKVCacheConfig=object, + ), + "infinilm.moe_config": _module( + "infinilm.moe_config", configure_moe_ep_backend=lambda *_args: (None, None) + ), + "infinilm.processors": _module( + "infinilm.processors", AutoInfinilmProcessor=object + ), + "numpy": _module("numpy"), + "tqdm": _module("tqdm", tqdm=lambda values, **_kwargs: values), + } + + spec = importlib.util.spec_from_file_location( + "infinilm_bench_case_test", REPO_ROOT / "examples" / "bench.py" + ) + module = importlib.util.module_from_spec(spec) + old_cwd = os.getcwd() + try: + os.chdir(REPO_ROOT) + with patch.dict(sys.modules, stub_modules), patch( + "builtins.open", mock_open(read_data="test prompt") + ): + spec.loader.exec_module(module) + finally: + os.chdir(old_cwd) + return module + + +bench = _load_bench_module() + + +class PairSequenceLengthsTest(unittest.TestCase): + def test_pairs_equal_length_lists_by_position(self): + self.assertEqual( + bench.pair_sequence_lengths([1024, 4096], [128, 256]), + [(1024, 128), (4096, 256)], + ) + + def test_broadcasts_a_single_value_on_either_side(self): + self.assertEqual( + bench.pair_sequence_lengths([1024], [128, 256]), + [(1024, 128), (1024, 256)], + ) + self.assertEqual( + bench.pair_sequence_lengths([1024, 4096], [128]), + [(1024, 128), (4096, 128)], + ) + + def test_rejects_unpairable_or_nonpositive_lengths(self): + with self.assertRaises(ValueError): + bench.pair_sequence_lengths([1024, 2048], [64, 128, 256]) + with self.assertRaises(ValueError): + bench.pair_sequence_lengths([0], [128]) + + +class GetTestCasesTest(unittest.TestCase): + @staticmethod + def _model_dir(): + model_dir = tempfile.TemporaryDirectory() + config_path = Path(model_dir.name) / "config.json" + config_path.write_text( + """{ + "model_type": "qwen3", + "hidden_size": 128, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "num_hidden_layers": 4 + }""", + encoding="utf-8", + ) + return model_dir + + def test_preserves_single_case_behavior(self): + with self._model_dir() as model_dir: + cases = bench.get_test_cases(model_dir, [4], [2048], [512]) + + self.assertEqual(len(cases), 1) + case = next(iter(cases.values())) + self.assertEqual( + (case["batch_size"], case["input_len"], case["output_len"]), + (4, 2048, 512), + ) + + def test_combines_batches_with_pairs_without_cartesian_lengths(self): + with self._model_dir() as model_dir: + cases = bench.get_test_cases( + model_dir, [1, 2], [1024, 4096], [128, 256] + ) + + actual = { + (case["batch_size"], case["input_len"], case["output_len"]) + for case in cases.values() + } + self.assertEqual( + actual, + { + (1, 1024, 128), + (1, 4096, 256), + (2, 1024, 128), + (2, 4096, 256), + }, + ) + self.assertEqual(len(cases), 4) + + def test_cache_helpers_cover_each_shape_at_full_capacity(self): + cases = [ + {"batch_size": 4, "input_len": 1024, "output_len": 128}, + {"batch_size": 1, "input_len": 4096, "output_len": 256}, + {"batch_size": 4, "input_len": 1024, "output_len": 512}, + ] + + self.assertEqual(bench.get_paged_kv_cache_num_blocks(cases, 256), 24) + self.assertEqual( + bench.get_warmup_shapes(cases), + {(4, 1024): 512, (1, 4096): 256}, + ) + + +if __name__ == "__main__": + unittest.main() From 4118f3b220243035c92b8c30054f7e2fd5a4188a Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 24 Jul 2026 11:15:11 +0800 Subject: [PATCH 10/14] perf(hygon): narrow reusable MoE routing workspace --- .../moe/runner/cuda_fused_moe_runner.cpp | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index ccec6154e..bee6ebc87 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -371,10 +371,16 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchO {1}, infinicore::DataType::I32, device); } + // Prefill grows the reusable workspace. Limit decode views to the current + // shape so fused routing kernels do not initialize the oversized buffer. + 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, @@ -383,8 +389,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_, @@ -395,8 +401,8 @@ 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, }, }; From 2d233c91594bae81594644a02845e2d2b4ebd093 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 24 Jul 2026 11:33:04 +0800 Subject: [PATCH 11/14] perf(hygon): use paged KV cache in FlashAttention graphs --- csrc/engine/compiler/paged_compiler.cpp | 17 ++++++++- csrc/layers/attention/backends/flash_attn.cpp | 35 +++++++++++++++++-- 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index 6b9dace38..e3de84332 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -174,6 +174,22 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & if (result == compiled_map_decode_.end()) { return {nullptr, nullptr}; } + + // 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}; + } + const auto *offsets = reinterpret_cast(runtime_input_offsets->data()); + for (size_t i = 0; i <= batch_size; ++i) { + if (offsets[i] != static_cast(i)) { + return {nullptr, nullptr}; + } + } auto &graph_input = result->second.input; const size_t compiled_block_per_req = graph_input.block_tables.value()->size(1); @@ -186,7 +202,6 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::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()); // Initialize only the active graph rows to -1, then overwrite the diff --git a/csrc/layers/attention/backends/flash_attn.cpp b/csrc/layers/attention/backends/flash_attn.cpp index ec7e37722..e8612af0a 100644 --- a/csrc/layers/attention/backends/flash_attn.cpp +++ b/csrc/layers/attention/backends/flash_attn.cpp @@ -5,6 +5,8 @@ #include "infinicore/ops/mha_kvcache.hpp" #include "infinicore/ops/mha_varlen.hpp" +#include + namespace infinilm::layers::attention::backends { FlashAttentionImpl::FlashAttentionImpl(size_t num_heads, @@ -50,6 +52,12 @@ infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, // 2. Compute attention infinicore::Tensor attn_output = infinicore::Tensor::empty({seq_len, num_heads_, head_dim_}, query->dtype(), query->device()); if (is_prefill) { + const auto cache_block_size = kv_cache->shape()[2]; + const auto max_cache_seqlen = block_tables.value()->shape()[1] * cache_block_size; + if (seq_len > static_cast(std::numeric_limits::max()) + || max_cache_seqlen > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("FlashAttention sequence length exceeds int range"); + } infinicore::op::mha_varlen_( attn_output, query, @@ -58,8 +66,8 @@ infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, input_offsets.value(), cu_seqlens.value(), block_tables.value(), - max_position_embeddings_, - max_position_embeddings_, + static_cast(seq_len), + static_cast(max_cache_seqlen), std::nullopt, scale_); } else { @@ -89,6 +97,29 @@ std::tuple FlashAttentionImpl::do_kv_cac const infinicore::Tensor slot_mapping) const { auto k_cache_layer = kv_cache->narrow({{0, 0, 1}})->squeeze(0); auto v_cache_layer = kv_cache->narrow({{0, 1, 1}})->squeeze(0); + const auto &cache_shape = k_cache_layer->shape(); + const bool use_hygon_paged_attention = + key->device().getType() == infinicore::Device::Type::HYGON + && cache_shape.size() == 4 + && cache_shape[1] == 64; + if (use_hygon_paged_attention) { + const auto num_blocks = cache_shape[0]; + const auto block_size = cache_shape[1]; + const auto num_kv_heads = cache_shape[2]; + const auto head_dim = cache_shape[3]; + auto k_cache_vllm = k_cache_layer->view( + {num_blocks, num_kv_heads, block_size, head_dim}); + auto v_cache_vllm = v_cache_layer->view( + {num_blocks, num_kv_heads, head_dim, block_size}); + infinicore::op::paged_caching_( + k_cache_vllm, + v_cache_vllm, + key, + value, + slot_mapping); + return {k_cache_vllm, v_cache_vllm}; + } + infinicore::op::paged_caching_( k_cache_layer->permute({0, 2, 1, 3}), // permute to BHSD for paged_caching_ v_cache_layer->permute({0, 2, 1, 3}), From 0ea02094a8afe9136f34b6bdcea224e0839e1a37 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 27 Jul 2026 14:39:40 +0800 Subject: [PATCH 12/14] fix(hygon): select safe paged attention layouts --- csrc/layers/attention/backends/flash_attn.cpp | 49 ++++++----- examples/bench.py | 86 ++++++++++++------- python/infinilm/base_config.py | 15 +++- 3 files changed, 98 insertions(+), 52 deletions(-) diff --git a/csrc/layers/attention/backends/flash_attn.cpp b/csrc/layers/attention/backends/flash_attn.cpp index e8612af0a..6efd991cf 100644 --- a/csrc/layers/attention/backends/flash_attn.cpp +++ b/csrc/layers/attention/backends/flash_attn.cpp @@ -6,6 +6,7 @@ #include "infinicore/ops/mha_varlen.hpp" #include +#include namespace infinilm::layers::attention::backends { @@ -34,6 +35,14 @@ infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, const infinicore::Tensor &value, infinicore::Tensor &kv_cache, const infinilm::global_state::AttentionMetadata &attn_metadata) const { + // The Hygon flash-attn extension uses process-global launch state while + // capturing graphs. InfiniLM TP ranks are threads in the same process. + static std::mutex hygon_flash_attention_mutex; + std::unique_lock hygon_lock(hygon_flash_attention_mutex, std::defer_lock); + if (query->device().getType() == infinicore::Device::Type::HYGON) { + hygon_lock.lock(); + } + auto total_sequence_lengths = attn_metadata.total_sequence_lengths; auto input_offsets = attn_metadata.input_offsets; auto block_tables = attn_metadata.block_tables; @@ -71,17 +80,14 @@ infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, std::nullopt, scale_); } else { - // FA2 decode path: flash::mha_fwd_kvcache - // In paged-attn mode, seq_len = actual batch_size (one query token per sequence). - // q_reshaped: [seq_len, num_heads, head_dim] → [seq_len, 1, num_heads, head_dim] - // k/v cache: [num_blocks, block_size, num_kv_heads, head_dim] + // In paged-attn mode, seq_len is the batch size (one query token per sequence). auto q_for_fa = query->view({seq_len, 1, num_heads_, head_dim_}); auto attn_out_4d = infinicore::op::mha_kvcache( q_for_fa, - k_total, // [num_blocks, block_size, num_kv_heads, head_dim] + k_total, v_total, - total_sequence_lengths.value(), // [seq_len] int32 (one entry per sequence) - block_tables.value(), // [seq_len, max_num_blocks_per_seq] int32 + total_sequence_lengths.value(), + block_tables.value(), std::nullopt, scale_); attn_output = attn_out_4d->view({seq_len, num_heads_, head_dim_}); @@ -98,30 +104,33 @@ std::tuple FlashAttentionImpl::do_kv_cac auto k_cache_layer = kv_cache->narrow({{0, 0, 1}})->squeeze(0); auto v_cache_layer = kv_cache->narrow({{0, 1, 1}})->squeeze(0); const auto &cache_shape = k_cache_layer->shape(); - const bool use_hygon_paged_attention = + const bool use_hygon_lightop_paged_attention = key->device().getType() == infinicore::Device::Type::HYGON && cache_shape.size() == 4 - && cache_shape[1] == 64; - if (use_hygon_paged_attention) { + && cache_shape[1] == 64 + && cache_shape[2] == num_kv_heads_ + && cache_shape[3] == head_dim_ + && num_heads_ == 8 + && num_kv_heads_ == 1 + && head_dim_ == 128; + if (use_hygon_lightop_paged_attention) { const auto num_blocks = cache_shape[0]; const auto block_size = cache_shape[1]; - const auto num_kv_heads = cache_shape[2]; - const auto head_dim = cache_shape[3]; - auto k_cache_vllm = k_cache_layer->view( - {num_blocks, num_kv_heads, block_size, head_dim}); - auto v_cache_vllm = v_cache_layer->view( - {num_blocks, num_kv_heads, head_dim, block_size}); + auto k_cache_lightop = k_cache_layer->view( + {num_blocks, num_kv_heads_, block_size, head_dim_}); + auto v_cache_lightop = v_cache_layer->view( + {num_blocks, num_kv_heads_, head_dim_, block_size}); infinicore::op::paged_caching_( - k_cache_vllm, - v_cache_vllm, + k_cache_lightop, + v_cache_lightop, key, value, slot_mapping); - return {k_cache_vllm, v_cache_vllm}; + return {k_cache_lightop, v_cache_lightop}; } infinicore::op::paged_caching_( - k_cache_layer->permute({0, 2, 1, 3}), // permute to BHSD for paged_caching_ + k_cache_layer->permute({0, 2, 1, 3}), v_cache_layer->permute({0, 2, 1, 3}), key, value, diff --git a/examples/bench.py b/examples/bench.py index 755c8c29d..b6f79b085 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -234,21 +234,33 @@ def get_test_cases( return case_dict -prompt_path = ( - "examples/bench_prompt.md" - if os.path.isfile("examples/bench_prompt.md") - else "InfiniLM/examples/bench_prompt.md" -) -with open(prompt_path, "r") as f: - prompt = f.read() - - -def repeat_prompt(input_ids: list[int], target_length: int): +def repeat_tokens(input_ids: list[int], target_length: int): num = len(input_ids) repeat_times = (target_length + num - 1) // num return (input_ids * repeat_times)[:target_length] +def split_chat_prompt_tokens(tokenizer, rendered_prompt: str, user_prompt: str): + """Split one rendered chat prompt around its user-content token span.""" + full_ids = tokenizer.encode(rendered_prompt) + content_ids = tokenizer.encode(user_prompt, add_special_tokens=False) + if not content_ids: + raise ValueError("bench prompt must contain at least one token") + + last_start = len(full_ids) - len(content_ids) + for start in range(last_start + 1): + if full_ids[start : start + len(content_ids)] == content_ids: + return ( + full_ids[:start], + content_ids, + full_ids[start + len(content_ids) :], + ) + + raise ValueError( + "Could not locate the user prompt inside the rendered chat template" + ) + + class TestModel: model: infinicore.nn.Module input_ids_list: list[int] @@ -269,6 +281,7 @@ def __init__( moe_ep_backend="disabled", moe_ep_size=1, enable_prefix_caching=False, + prompt="How are you", ) -> None: model_path = os.path.expanduser(model_path) self.draft_model_path = draft_model_path @@ -292,7 +305,13 @@ def __init__( add_generation_prompt=True, tokenize=False, ) - self.input_ids_list = [self.tokenizer.encode(input_content)] + prefix_ids, content_ids, suffix_ids = split_chat_prompt_tokens( + self.tokenizer, input_content, prompt + ) + self.prompt_prefix_ids = prefix_ids + self.prompt_content_ids = content_ids + self.prompt_suffix_ids = suffix_ids + self.input_ids_list = [prefix_ids + content_ids + suffix_ids] self.model = None return @@ -336,24 +355,30 @@ def __init__( tokenize=False, ) - input_ids_list = [ - self.tokenizer.encode( - input_content, - ) - ] - + prefix_ids, content_ids, suffix_ids = split_chat_prompt_tokens( + self.tokenizer, input_content, prompt + ) + self.prompt_prefix_ids = prefix_ids + self.prompt_content_ids = content_ids + self.prompt_suffix_ids = suffix_ids + self.input_ids_list = [prefix_ids + content_ids + suffix_ids] self.model = model - self.input_ids_list = input_ids_list - self.draft_model_path = draft_model_path - self.model_path = model_path - self.device_str = infini_device.type - self.tp = tp - self.cache_config = cache_config - self.enable_graph = enable_graph - self.attn_backend = attn_backend - self.use_mla = use_mla - self.weight_load_mode = weight_load_mode - self.skip_load = skip_load + + def build_input_ids(self, target_length: int) -> list[int]: + template_tokens = len(self.prompt_prefix_ids) + len(self.prompt_suffix_ids) + if target_length < template_tokens: + raise ValueError( + f"input_len={target_length} is shorter than the chat template " + f"overhead ({template_tokens} tokens)" + ) + content_length = target_length - template_tokens + input_ids = ( + self.prompt_prefix_ids + + repeat_tokens(self.prompt_content_ids, content_length) + + self.prompt_suffix_ids + ) + assert len(input_ids) == target_length + return input_ids def run( self, @@ -364,7 +389,7 @@ def run( top_p=1.0, temperature=1.0, ): - input_ids = repeat_prompt(self.input_ids_list[0], target_length=input_len) + input_ids = self.build_input_ids(input_len) input_ids_list = [input_ids] * batch_size # ---------------------------------------------------------------------------- # @@ -513,6 +538,7 @@ def run( moe_ep_backend=moe_ep_backend, moe_ep_size=ep, enable_prefix_caching=False, + prompt=cfg.prompt, ) # ---------------------------------------------------------------------------- # @@ -539,7 +565,7 @@ def run( ) test.model.reset_cache(warmup_cache_config) - warmup_prompt_ids = repeat_prompt(test.input_ids_list[0], warmup_input_len) + warmup_prompt_ids = test.build_input_ids(warmup_input_len) warmup_ids = [warmup_prompt_ids] * warmup_batch input_ids_infini = infinicore.from_list( warmup_ids, dtype=infinicore.int64 diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index c92c21efb..ad72da296 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -78,7 +78,15 @@ def __init__(self): self.enable_prefix_caching = self.args.enable_prefix_caching self.use_mla = self.args.use_mla 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 @@ -272,7 +280,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" From 91a45c807bdba6107f5a5dd152fe42d6d5af9601 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Tue, 28 Jul 2026 09:47:39 +0800 Subject: [PATCH 13/14] refactor(hygon): route optimized inference through InfiniCore --- csrc/layers/attention/backends/flash_attn.cpp | 145 +--- csrc/layers/attention/backends/flash_attn.hpp | 26 +- csrc/layers/moe/common/moe_types.hpp | 13 +- .../moe/runner/cuda_fused_moe_runner.cpp | 786 ++++-------------- .../moe/runner/cuda_fused_moe_runner.hpp | 47 +- 5 files changed, 199 insertions(+), 818 deletions(-) diff --git a/csrc/layers/attention/backends/flash_attn.cpp b/csrc/layers/attention/backends/flash_attn.cpp index 6efd991cf..eb2e37975 100644 --- a/csrc/layers/attention/backends/flash_attn.cpp +++ b/csrc/layers/attention/backends/flash_attn.cpp @@ -2,11 +2,6 @@ #include "../../../utils.hpp" #include "infinicore/ops.hpp" -#include "infinicore/ops/mha_kvcache.hpp" -#include "infinicore/ops/mha_varlen.hpp" - -#include -#include namespace infinilm::layers::attention::backends { @@ -16,127 +11,39 @@ FlashAttentionImpl::FlashAttentionImpl(size_t num_heads, size_t num_kv_heads, size_t layer_idx) : num_heads_(num_heads), - head_size_(head_size), scale_(scale), num_kv_heads_(num_kv_heads), - layer_idx_(layer_idx), head_dim_(head_size) { - - const infinilm::global_state::InfinilmConfig &infinilm_config = infinilm::global_state::get_infinilm_config(); - if (!infinilm_config.model_config) { - throw std::runtime_error("infinilm::layers::attention::backends::FlashAttentionImpl: model_config is null"); - } - max_position_embeddings_ = infinilm_config.model_config->get("max_position_embeddings"); + (void)layer_idx; } -infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, - const infinicore::Tensor &query, - const infinicore::Tensor &key, - const infinicore::Tensor &value, - infinicore::Tensor &kv_cache, - const infinilm::global_state::AttentionMetadata &attn_metadata) const { - // The Hygon flash-attn extension uses process-global launch state while - // capturing graphs. InfiniLM TP ranks are threads in the same process. - static std::mutex hygon_flash_attention_mutex; - std::unique_lock hygon_lock(hygon_flash_attention_mutex, std::defer_lock); - if (query->device().getType() == infinicore::Device::Type::HYGON) { - hygon_lock.lock(); - } - - auto total_sequence_lengths = attn_metadata.total_sequence_lengths; - auto input_offsets = attn_metadata.input_offsets; - auto block_tables = attn_metadata.block_tables; - auto slot_mapping = attn_metadata.slot_mapping; - auto cu_seqlens = attn_metadata.cu_seqlens; - - ASSERT(block_tables.has_value()); - ASSERT(slot_mapping.has_value()); - - // 1. update paged kv cache - auto [k_total, v_total] = do_kv_cache_update(layer, key, value, kv_cache, slot_mapping.value()); - - size_t seq_len = query->shape()[0]; - bool is_prefill = (seq_len != total_sequence_lengths.value()->shape()[0]); - - // 2. Compute attention - infinicore::Tensor attn_output = infinicore::Tensor::empty({seq_len, num_heads_, head_dim_}, query->dtype(), query->device()); - if (is_prefill) { - const auto cache_block_size = kv_cache->shape()[2]; - const auto max_cache_seqlen = block_tables.value()->shape()[1] * cache_block_size; - if (seq_len > static_cast(std::numeric_limits::max()) - || max_cache_seqlen > static_cast(std::numeric_limits::max())) { - throw std::runtime_error("FlashAttention sequence length exceeds int range"); - } - infinicore::op::mha_varlen_( - attn_output, - query, - k_total, - v_total, - input_offsets.value(), - cu_seqlens.value(), - block_tables.value(), - static_cast(seq_len), - static_cast(max_cache_seqlen), - std::nullopt, - scale_); - } else { - // In paged-attn mode, seq_len is the batch size (one query token per sequence). - auto q_for_fa = query->view({seq_len, 1, num_heads_, head_dim_}); - auto attn_out_4d = infinicore::op::mha_kvcache( - q_for_fa, - k_total, - v_total, - total_sequence_lengths.value(), - block_tables.value(), - std::nullopt, - scale_); - attn_output = attn_out_4d->view({seq_len, num_heads_, head_dim_}); - } - attn_output = attn_output->view({1, seq_len, num_heads_ * head_dim_}); - return attn_output; -} - -std::tuple FlashAttentionImpl::do_kv_cache_update(const AttentionLayer &layer, - const infinicore::Tensor key, - const infinicore::Tensor value, - infinicore::Tensor &kv_cache, - const infinicore::Tensor slot_mapping) const { - auto k_cache_layer = kv_cache->narrow({{0, 0, 1}})->squeeze(0); - auto v_cache_layer = kv_cache->narrow({{0, 1, 1}})->squeeze(0); - const auto &cache_shape = k_cache_layer->shape(); - const bool use_hygon_lightop_paged_attention = - key->device().getType() == infinicore::Device::Type::HYGON - && cache_shape.size() == 4 - && cache_shape[1] == 64 - && cache_shape[2] == num_kv_heads_ - && cache_shape[3] == head_dim_ - && num_heads_ == 8 - && num_kv_heads_ == 1 - && head_dim_ == 128; - if (use_hygon_lightop_paged_attention) { - const auto num_blocks = cache_shape[0]; - const auto block_size = cache_shape[1]; - auto k_cache_lightop = k_cache_layer->view( - {num_blocks, num_kv_heads_, block_size, head_dim_}); - auto v_cache_lightop = v_cache_layer->view( - {num_blocks, num_kv_heads_, head_dim_, block_size}); - infinicore::op::paged_caching_( - k_cache_lightop, - v_cache_lightop, - key, - value, - slot_mapping); - return {k_cache_lightop, v_cache_lightop}; - } - - infinicore::op::paged_caching_( - k_cache_layer->permute({0, 2, 1, 3}), - v_cache_layer->permute({0, 2, 1, 3}), +infinicore::Tensor FlashAttentionImpl::forward( + const AttentionLayer &layer, + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const { + (void)layer; + + ASSERT(attn_metadata.total_sequence_lengths.has_value()); + ASSERT(attn_metadata.block_tables.has_value()); + ASSERT(attn_metadata.slot_mapping.has_value()); + + return infinicore::op::paged_flash_attention( + query, key, value, - slot_mapping); - - return {k_cache_layer, v_cache_layer}; + kv_cache, + attn_metadata.total_sequence_lengths.value(), + attn_metadata.input_offsets, + attn_metadata.cu_seqlens, + attn_metadata.block_tables.value(), + attn_metadata.slot_mapping.value(), + num_heads_, + num_kv_heads_, + head_dim_, + scale_); } } // namespace infinilm::layers::attention::backends diff --git a/csrc/layers/attention/backends/flash_attn.hpp b/csrc/layers/attention/backends/flash_attn.hpp index 93f61e8ba..a592dddb5 100644 --- a/csrc/layers/attention/backends/flash_attn.hpp +++ b/csrc/layers/attention/backends/flash_attn.hpp @@ -2,7 +2,6 @@ #include "../../../global_state/global_state.hpp" #include "infinicore/tensor.hpp" -#include namespace infinilm::layers::attention { class AttentionLayer; @@ -29,26 +28,19 @@ class FlashAttentionImpl { * @param attn_metadata: Attention metadata. * @return Attention output, shape `[1, num_tokens, num_heads * head_dim]`. */ - infinicore::Tensor forward(const AttentionLayer &layer, - const infinicore::Tensor &query, - const infinicore::Tensor &key, - const infinicore::Tensor &value, - infinicore::Tensor &kv_cache, - const infinilm::global_state::AttentionMetadata &attn_metadata) const; - - std::tuple do_kv_cache_update(const AttentionLayer &layer, - const infinicore::Tensor key, - const infinicore::Tensor value, - infinicore::Tensor &kv_cache, - const infinicore::Tensor slot_mapping) const; + infinicore::Tensor forward( + const AttentionLayer &layer, + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const; private: size_t num_heads_; - size_t head_size_; float scale_; size_t num_kv_heads_; - size_t layer_idx_; - size_t head_dim_; // Note: head_dim equals to head_size - size_t max_position_embeddings_; + size_t head_dim_; }; + } // namespace infinilm::layers::attention::backends diff --git a/csrc/layers/moe/common/moe_types.hpp b/csrc/layers/moe/common/moe_types.hpp index 1b16ed710..84c9dcbd8 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 @@ -72,7 +73,8 @@ struct MoeWeights { } bool has_packed_w8a8_marlin_weights() const { - return packed_w13 && packed_w2 && packed_w13_scale && packed_w2_scale; + return packed_w13 && packed_w2 + && packed_w13_scale && packed_w2_scale; } bool is_hygon_w16a16_marlin() const { @@ -90,12 +92,7 @@ struct MoeWorkspace { infinicore::Tensor ep_gathered_topk_ids; infinicore::Tensor ep_reduced_hidden_states; infinicore::Tensor fused_moe_output; - infinicore::Tensor marlin_cache13; - infinicore::Tensor marlin_cache2; - infinicore::Tensor marlin_input_i8; - infinicore::Tensor marlin_input_scale; - infinicore::Tensor marlin_cache2_i8; - infinicore::Tensor marlin_cache2_scale; + infinicore::op::HygonMoeMarlinWorkspace hygon_marlin; infinicore::Tensor sorted_token_ids; infinicore::Tensor expert_ids; @@ -111,8 +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 marlin_cache13_capacity = 0; - size_t marlin_cache2_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/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index bee6ebc87..629f93459 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -1,224 +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 "infinicore/ops/moe_w16a16_marlin.hpp" -#include "infinicore/ops/moe_w8a8_marlin.hpp" -#include "infinicore/adaptor/lightop_adaptor.hpp" -#include "nlohmann/json.hpp" - -#include -#include -#include -#include -#include #include #include #include namespace infinilm::layers::moe { -struct HygonMarlinGemmConfig { - int mode = 103; - int delta = 1; - size_t block_size_m = 16; - bool found = false; -}; - -struct HygonW16A16MarlinRuntimeConfig { - HygonMarlinGemmConfig gemm1; - HygonMarlinGemmConfig gemm2; - bool supported = false; -}; - -struct HygonW8A8MarlinRuntimeConfig { - HygonMarlinGemmConfig gemm1; - HygonMarlinGemmConfig gemm2; - bool supported = false; -}; - -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 { -std::string env_or_default(const char *name, const char *default_value) { - const char *value = std::getenv(name); - return (value != nullptr && value[0] != '\0') ? std::string(value) : std::string(default_value); -} - -std::string normalize_hygon_gpu_target(std::string target, bool uppercase) { - const auto feature_pos = target.find(':'); - if (feature_pos != std::string::npos) { - target.resize(feature_pos); - } - std::transform(target.begin(), target.end(), target.begin(), [uppercase](unsigned char ch) { - return static_cast(uppercase ? std::toupper(ch) : std::tolower(ch)); - }); - - std::string lowercase = target; - std::transform(lowercase.begin(), lowercase.end(), lowercase.begin(), [](unsigned char ch) { - return static_cast(std::tolower(ch)); - }); - if (lowercase.size() <= 3 || lowercase.compare(0, 3, "gfx") != 0 || - !std::all_of(lowercase.begin() + 3, lowercase.end(), [](unsigned char ch) { - return std::isalnum(ch) != 0; - })) { - throw std::runtime_error("Invalid Hygon GPU target for lightop config: " + target); - } - return target; -} - -constexpr size_t kHygonW16A16MoeSliceTokens = 16384; -constexpr size_t kHygonW8A8MoeSliceTokens = 16384; - -enum class HygonMarlinModePolicy { - LegacyOnly, - LegacyAndBf16Mode1000, - All, -}; - -HygonMarlinGemmConfig load_lightop_marlin_config(size_t n, - size_t k, - size_t m, - const std::string &file_prefix, - const infinicore::adaptor::lightop::DeviceInfo &device_info, - HygonMarlinModePolicy mode_policy, - bool uppercase_device_name, - bool num_cus_with_cu_prefix) { - HygonMarlinGemmConfig result; - const std::string config_dir = env_or_default( - "INFINILM_LIGHTOP_CONFIG_DIR", - "/usr/local/lib/python3.10/dist-packages/lightop/configs"); - if (device_info.gpu_target.empty() || device_info.compute_units <= 0) { - throw std::runtime_error("Unable to query Hygon device properties for lightop config"); - } - const std::string device_name = normalize_hygon_gpu_target( - device_info.gpu_target, - uppercase_device_name); - const std::string num_cus = std::to_string(device_info.compute_units); - const std::string num_cus_suffix = num_cus_with_cu_prefix ? ("_CU" + num_cus) : ("_" + num_cus); - const std::string file_name = config_dir + "/" + file_prefix + "_" + - std::to_string(n) + "_" + std::to_string(k) + "_" + - device_name + num_cus_suffix + ".json"; - std::ifstream file(file_name); - if (!file.is_open()) { - return result; - } - - nlohmann::json config_json; - file >> config_json; - const std::string shape_key = std::to_string(n) + "_" + std::to_string(k); - if (!config_json.contains(shape_key) || !config_json.at(shape_key).is_object()) { - return result; - } - const auto &configs = config_json.at(shape_key); - - auto usable = [&](size_t token) -> bool { - const auto key = std::to_string(token); - if (!configs.contains(key) || !configs.at(key).is_object()) { - return false; - } - const int mode = configs.at(key).value("MODE", result.mode); - return mode < 1000 || - mode_policy == HygonMarlinModePolicy::All || - (mode_policy == HygonMarlinModePolicy::LegacyAndBf16Mode1000 && mode == 1000); - }; - - size_t chosen = 0; - bool has_choice = false; - size_t chosen_ge = std::numeric_limits::max(); - size_t closest_diff = std::numeric_limits::max(); - for (auto it = configs.begin(); it != configs.end(); ++it) { - size_t token = 0; - try { - token = static_cast(std::stoull(it.key())); - } catch (const std::exception &) { - continue; - } - if (!usable(token)) { - continue; - } - if (token >= m && token < chosen_ge) { - chosen_ge = token; - chosen = token; - has_choice = true; - } - const size_t diff = token > m ? token - m : m - token; - if (diff < closest_diff) { - closest_diff = diff; - if (chosen_ge == std::numeric_limits::max()) { - chosen = token; - has_choice = true; - } - } - } - if (!has_choice) { - return result; - } - - const auto &cfg = configs.at(std::to_string(chosen)); - result.mode = cfg.value("MODE", result.mode); - result.delta = cfg.value("DELTA", result.delta); - result.block_size_m = cfg.value("BLOCK_SIZE_M", result.block_size_m); - result.found = cfg.contains("MODE"); - return result; -} -HygonW16A16MarlinRuntimeConfig select_hygon_w16a16_marlin_config(size_t m, - size_t hidden_size, - size_t intermediate_size_per_partition, - infinicore::DataType hidden_dtype, - size_t device_index) { - HygonW16A16MarlinRuntimeConfig config; - const auto device_info = infinicore::adaptor::lightop::device_info(device_index); - const auto mode_policy = hidden_dtype == infinicore::DataType::BF16 - ? HygonMarlinModePolicy::LegacyAndBf16Mode1000 - : HygonMarlinModePolicy::LegacyOnly; - config.gemm1 = load_lightop_marlin_config( - intermediate_size_per_partition * 2, hidden_size, m, - "MOE_W16A16_CUDA_MARLIN", device_info, mode_policy, false, false); - config.gemm2 = load_lightop_marlin_config( - hidden_size, intermediate_size_per_partition, m, - "MOE_W16A16_CUDA_MARLIN", device_info, mode_policy, false, false); - config.supported = config.gemm1.found && config.gemm2.found; - return config; -} - -HygonW8A8MarlinRuntimeConfig select_hygon_w8a8_marlin_config(size_t m, - size_t hidden_size, - size_t intermediate_size_per_partition, - size_t device_index) { - HygonW8A8MarlinRuntimeConfig config; - const auto device_info = infinicore::adaptor::lightop::device_info(device_index); - config.gemm1 = load_lightop_marlin_config( - intermediate_size_per_partition * 2, hidden_size, m, - "MOE_BLOCKINT8_CUDA_MARLIN", device_info, HygonMarlinModePolicy::All, true, true); - config.gemm2 = load_lightop_marlin_config( - hidden_size, intermediate_size_per_partition, m, - "MOE_BLOCKINT8_CUDA_MARLIN", device_info, HygonMarlinModePolicy::All, true, true); - config.supported = config.gemm1.found && config.gemm2.found; - return config; -} - -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); } @@ -237,91 +61,86 @@ 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 packed tensor dtype mismatch for " + name); + 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 { - size_t block_size = align_block_size_; - HygonW16A16MarlinRuntimeConfig marlin_config; - HygonW8A8MarlinRuntimeConfig w8a8_marlin_config; - if (weights.is_hygon_w16a16_marlin() || weights.is_hygon_w8a8_marlin()) { - const auto &hidden_shape = dispatch_output.hidden_states->shape(); - if (hidden_shape.size() != 2) { - throw std::runtime_error("Hygon Marlin MoE runner requires hidden states [M, K]"); - } - if (weights.is_hygon_w16a16_marlin()) { - if (hidden_shape[0] > kHygonW16A16MoeSliceTokens) { - auto runner_output = run_hygon_w16a16_marlin_core_sliced( - dispatch_output, weights, workspace); - return CombineInput{ - CombineInputFormat::Standard, - runner_output.hidden_states, - dispatch_output.topk_output, - MoeRoutingMetadata{}, - }; - } - marlin_config = select_hygon_w16a16_marlin_config( - hidden_shape[0], hidden_size_, intermediate_size_per_partition_, - dispatch_output.hidden_states->dtype(), - dispatch_output.hidden_states->device().getIndex()); - if (!marlin_config.supported) { - throw std::runtime_error("No lightop W16A16 Marlin MoE config found for this Hygon shape"); - } - block_size = marlin_config.gemm1.block_size_m; - } else { - if (hidden_shape[0] > kHygonW8A8MoeSliceTokens) { - auto runner_output = run_hygon_w8a8_marlin_core_sliced( - dispatch_output, weights, workspace); - return CombineInput{ - CombineInputFormat::Standard, - runner_output.hidden_states, - dispatch_output.topk_output, - MoeRoutingMetadata{}, - }; - } - w8a8_marlin_config = select_hygon_w8a8_marlin_config( - hidden_shape[0], hidden_size_, intermediate_size_per_partition_, - dispatch_output.hidden_states->device().getIndex()); - if (!w8a8_marlin_config.supported) { - throw std::runtime_error("No lightop W8A8 Marlin MoE config found for this Hygon shape"); - } - block_size = w8a8_marlin_config.gemm1.block_size_m; +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, - block_size); - - auto runner_output = weights.is_hygon_w16a16_marlin() - ? run_hygon_w16a16_marlin_core(runner_input, weights, workspace, marlin_config) - : (weights.is_hygon_w8a8_marlin() - ? run_hygon_w8a8_marlin_core( - runner_input, weights, workspace, w8a8_marlin_config) - : run_fused_core(runner_input, weights, workspace)); - + align_block_size_); + auto runner_output = + run_fused_core(runner_input, weights, workspace); return CombineInput{ CombineInputFormat::Standard, runner_output.hidden_states, @@ -330,53 +149,73 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, }; } -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; +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 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); - 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 size_t max_num_tokens_padded = + num_pairs < align_num_experts + ? num_pairs * block_size + : 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); - workspace.sorted_token_ids_capacity = sorted_token_ids_capacity; - } - if (!workspace.expert_ids || workspace.expert_ids_capacity < max_num_blocks) { + {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 (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); - } - - // Prefill grows the reusable workspace. Limit decode views to the current - // shape so fused routing kernels do not initialize the oversized buffer. - 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}}); - + workspace.num_tokens_post_padded = + infinicore::Tensor::empty( + {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_( sorted_token_ids, @@ -408,24 +247,30 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchO }; } -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(), @@ -446,341 +291,4 @@ CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_fused_core(const CudaFusedMoeRu }; } -CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w16a16_marlin_core( - const CudaFusedMoeRunnerInput &runner_input, - const MoeWeights &weights, - MoeWorkspace &workspace, - const HygonW16A16MarlinRuntimeConfig &config) const { - if (!weights.has_packed_dense_weights() || !weights.is_hygon_w16a16_marlin()) { - throw std::runtime_error("Hygon W16A16 Marlin MoE runner requires packed Marlin weights"); - } - const auto activation_dtype = runner_input.hidden_states->dtype(); - if (activation_dtype != infinicore::DataType::BF16 && - activation_dtype != infinicore::DataType::F16) { - throw std::runtime_error("Hygon W16A16 Marlin MoE runner requires BF16 or FP16 activations"); - } - check_packed_weight_tensor( - weights.packed_w13, - "w13", - runner_input.hidden_states->device(), - activation_dtype, - {num_local_experts_, hidden_size_ / 16, intermediate_size_per_partition_ * 2 * 16}); - check_packed_weight_tensor( - weights.packed_w2, - "w2", - runner_input.hidden_states->device(), - activation_dtype, - {num_local_experts_, intermediate_size_per_partition_ / 16, hidden_size_ * 16}); - const size_t top_k = runner_input.topk_output.topk_ids->shape()[1]; - const size_t num_tokens = runner_input.hidden_states->shape()[0]; - if (num_tokens > kHygonW16A16MoeSliceTokens) { - throw std::runtime_error( - "Hygon W16A16 Marlin MoE core requires inputs above 16384 tokens to be sliced"); - } - const size_t cache13_required = num_tokens * top_k * std::max(intermediate_size_per_partition_ * 2, hidden_size_); - const size_t cache2_required = num_tokens * top_k * 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()); - if (!same_device(workspace.marlin_cache13, runner_input.hidden_states->device()) || - workspace.marlin_cache13->dtype() != runner_input.hidden_states->dtype() || - workspace.marlin_cache13_capacity < cache13_required) { - if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE Marlin cache13 workspace was not initialized before graph capture"); - } - workspace.marlin_cache13 = infinicore::Tensor::empty( - {cache13_required}, runner_input.hidden_states->dtype(), runner_input.hidden_states->device()); - workspace.marlin_cache13_capacity = cache13_required; - } - if (!same_device(workspace.marlin_cache2, runner_input.hidden_states->device()) || - workspace.marlin_cache2->dtype() != runner_input.hidden_states->dtype() || - workspace.marlin_cache2_capacity < cache2_required) { - if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE Marlin cache2 workspace was not initialized before graph capture"); - } - workspace.marlin_cache2 = infinicore::Tensor::empty( - {cache2_required}, runner_input.hidden_states->dtype(), runner_input.hidden_states->device()); - workspace.marlin_cache2_capacity = cache2_required; - } - - infinicore::op::moe_w16a16_marlin_fused_dense_( - workspace.fused_moe_output, - workspace.marlin_cache13, - workspace.marlin_cache2, - runner_input.hidden_states, - weights.packed_w13, - weights.packed_w2, - runner_input.topk_output.topk_weights, - runner_input.routing_metadata.sorted_token_ids, - runner_input.routing_metadata.expert_ids, - runner_input.routing_metadata.num_tokens_post_padded, - top_k, - config.gemm1.mode, - config.gemm1.delta, - config.gemm2.mode, - config.gemm2.delta); - - return CudaFusedMoeRunnerOutput{ - workspace.fused_moe_output, - }; -} - -CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w16a16_marlin_core_sliced( - const DispatchOutput &dispatch_output, - const MoeWeights &weights, - MoeWorkspace &workspace) const { - const auto &hidden_shape = dispatch_output.hidden_states->shape(); - if (hidden_shape.size() != 2) { - throw std::runtime_error("Hygon W16A16 sliced MoE runner requires hidden states [M, K]"); - } - const size_t num_tokens = hidden_shape[0]; - if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("Hygon W16A16 sliced MoE runner cannot allocate/copy slice outputs during graph capture"); - } - - ensure_tensor( - workspace.fused_moe_output, - dispatch_output.hidden_states->shape(), - dispatch_output.hidden_states->dtype(), - dispatch_output.hidden_states->device()); - MoeWorkspace slice_workspace; - const auto activation_dtype = dispatch_output.hidden_states->dtype(); - const auto device_index = dispatch_output.hidden_states->device().getIndex(); - const auto full_slice_config = select_hygon_w16a16_marlin_config( - kHygonW16A16MoeSliceTokens, - hidden_size_, - intermediate_size_per_partition_, - activation_dtype, - device_index); - if (!full_slice_config.supported) { - throw std::runtime_error("No lightop W16A16 Marlin MoE config found for full Hygon slice"); - } - - size_t offset = 0; - while (offset < num_tokens) { - const size_t slice_tokens = std::min(kHygonW16A16MoeSliceTokens, num_tokens - offset); - const auto slice_config = slice_tokens == kHygonW16A16MoeSliceTokens - ? full_slice_config - : select_hygon_w16a16_marlin_config( - slice_tokens, - hidden_size_, - intermediate_size_per_partition_, - activation_dtype, - device_index); - if (!slice_config.supported) { - throw std::runtime_error("No lightop W16A16 Marlin MoE config found for sliced Hygon shape"); - } - - const auto hidden_slice = dispatch_output.hidden_states->narrow({{0, offset, slice_tokens}}); - const TopKOutput topk_slice{ - dispatch_output.topk_output.topk_weights->narrow({{0, offset, slice_tokens}}), - dispatch_output.topk_output.topk_ids->narrow({{0, offset, slice_tokens}}), - dispatch_output.topk_output.router_logits - ? dispatch_output.topk_output.router_logits->narrow({{0, offset, slice_tokens}}) - : infinicore::Tensor(), - }; - const DispatchOutput dispatch_slice{ - DispatchOutputFormat::Standard, - hidden_slice, - infinicore::Tensor(), - topk_slice, - infinicore::Tensor(), - }; - auto slice_input = prepare_runner_input( - dispatch_slice, - slice_workspace, - slice_config.gemm1.block_size_m); - auto slice_output = run_hygon_w16a16_marlin_core( - slice_input, - weights, - slice_workspace, - slice_config); - workspace.fused_moe_output - ->narrow({{0, offset, slice_tokens}}) - ->copy_from(slice_output.hidden_states); - - offset += slice_tokens; - } - - return CudaFusedMoeRunnerOutput{ - workspace.fused_moe_output, - }; -} - -CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w8a8_marlin_core( - const CudaFusedMoeRunnerInput &runner_input, - const MoeWeights &weights, - MoeWorkspace &workspace, - const HygonW8A8MarlinRuntimeConfig &config) const { - if (!weights.has_packed_w8a8_marlin_weights() || !weights.is_hygon_w8a8_marlin()) { - throw std::runtime_error("Hygon W8A8 Marlin MoE runner requires packed Marlin weights and scales"); - } - const size_t top_k = runner_input.topk_output.topk_ids->shape()[1]; - const size_t num_tokens = runner_input.hidden_states->shape()[0]; - const size_t cache13_required = num_tokens * top_k * std::max(intermediate_size_per_partition_ * 2, hidden_size_); - - check_packed_weight_tensor( - weights.packed_w13, - "w13", - runner_input.hidden_states->device(), - infinicore::DataType::I8, - {num_local_experts_, hidden_size_ / 64, intermediate_size_per_partition_ * 2 * 64}); - check_packed_weight_tensor( - weights.packed_w2, - "w2", - runner_input.hidden_states->device(), - infinicore::DataType::I8, - {num_local_experts_, intermediate_size_per_partition_ / 64, hidden_size_ * 64}); - check_packed_weight_tensor( - weights.packed_w13_scale, - "w13_scale", - runner_input.hidden_states->device(), - infinicore::DataType::F32, - {num_local_experts_, intermediate_size_per_partition_ * 2, 1}); - check_packed_weight_tensor( - weights.packed_w2_scale, - "w2_scale", - runner_input.hidden_states->device(), - infinicore::DataType::F32, - {num_local_experts_, hidden_size_, 1}); - - ensure_tensor( - workspace.fused_moe_output, - runner_input.hidden_states->shape(), - runner_input.hidden_states->dtype(), - runner_input.hidden_states->device()); - if (!same_device(workspace.marlin_cache13, runner_input.hidden_states->device()) || - workspace.marlin_cache13->dtype() != runner_input.hidden_states->dtype() || - workspace.marlin_cache13_capacity < cache13_required) { - if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE W8A8 Marlin cache13 workspace was not initialized before graph capture"); - } - workspace.marlin_cache13 = infinicore::Tensor::empty( - {cache13_required}, runner_input.hidden_states->dtype(), runner_input.hidden_states->device()); - workspace.marlin_cache13_capacity = cache13_required; - } - ensure_tensor( - workspace.marlin_input_i8, - {num_tokens, hidden_size_}, - infinicore::DataType::I8, - runner_input.hidden_states->device()); - ensure_tensor( - workspace.marlin_input_scale, - {num_tokens, 1}, - infinicore::DataType::F32, - runner_input.hidden_states->device()); - ensure_tensor( - workspace.marlin_cache2_i8, - {num_tokens * top_k, intermediate_size_per_partition_}, - infinicore::DataType::I8, - runner_input.hidden_states->device()); - ensure_tensor( - workspace.marlin_cache2_scale, - {num_tokens * top_k, 1}, - infinicore::DataType::F32, - runner_input.hidden_states->device()); - - infinicore::op::moe_w8a8_marlin_fused_dense_( - workspace.fused_moe_output, - workspace.marlin_cache13, - workspace.marlin_cache2_i8, - workspace.marlin_input_i8, - workspace.marlin_input_scale, - workspace.marlin_cache2_scale, - runner_input.hidden_states, - weights.packed_w13, - weights.packed_w2, - weights.packed_w13_scale, - weights.packed_w2_scale, - runner_input.topk_output.topk_weights, - runner_input.routing_metadata.sorted_token_ids, - runner_input.routing_metadata.expert_ids, - runner_input.routing_metadata.num_tokens_post_padded, - top_k, - config.gemm1.mode, - config.gemm1.block_size_m, - config.gemm1.delta, - config.gemm2.mode, - config.gemm2.delta); - - return CudaFusedMoeRunnerOutput{ - workspace.fused_moe_output, - }; -} - -CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_hygon_w8a8_marlin_core_sliced( - const DispatchOutput &dispatch_output, - const MoeWeights &weights, - MoeWorkspace &workspace) const { - const auto &hidden_shape = dispatch_output.hidden_states->shape(); - if (hidden_shape.size() != 2) { - throw std::runtime_error("Hygon W8A8 sliced MoE runner requires hidden states [M, K]"); - } - const size_t num_tokens = hidden_shape[0]; - if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("Hygon W8A8 sliced MoE runner cannot allocate/copy slice outputs during graph capture"); - } - - ensure_tensor( - workspace.fused_moe_output, - dispatch_output.hidden_states->shape(), - dispatch_output.hidden_states->dtype(), - dispatch_output.hidden_states->device()); - MoeWorkspace slice_workspace; - const auto device_index = dispatch_output.hidden_states->device().getIndex(); - const auto full_slice_config = select_hygon_w8a8_marlin_config( - kHygonW8A8MoeSliceTokens, hidden_size_, intermediate_size_per_partition_, - device_index); - if (!full_slice_config.supported) { - throw std::runtime_error("No lightop W8A8 Marlin MoE config found for full Hygon slice"); - } - size_t offset = 0; - while (offset < num_tokens) { - const size_t slice_tokens = std::min(kHygonW8A8MoeSliceTokens, num_tokens - offset); - const auto slice_config = slice_tokens == kHygonW8A8MoeSliceTokens - ? full_slice_config - : select_hygon_w8a8_marlin_config( - slice_tokens, hidden_size_, intermediate_size_per_partition_, - device_index); - if (!slice_config.supported) { - throw std::runtime_error("No lightop W8A8 Marlin MoE config found for sliced Hygon shape"); - } - - const auto hidden_slice = dispatch_output.hidden_states->narrow({{0, offset, slice_tokens}}); - const TopKOutput topk_slice{ - dispatch_output.topk_output.topk_weights->narrow({{0, offset, slice_tokens}}), - dispatch_output.topk_output.topk_ids->narrow({{0, offset, slice_tokens}}), - dispatch_output.topk_output.router_logits - ? dispatch_output.topk_output.router_logits->narrow({{0, offset, slice_tokens}}) - : infinicore::Tensor(), - }; - const DispatchOutput dispatch_slice{ - DispatchOutputFormat::Standard, - hidden_slice, - infinicore::Tensor(), - topk_slice, - infinicore::Tensor(), - }; - auto slice_input = prepare_runner_input( - dispatch_slice, - slice_workspace, - slice_config.gemm1.block_size_m); - auto slice_output = run_hygon_w8a8_marlin_core( - slice_input, - weights, - slice_workspace, - slice_config); - workspace.fused_moe_output->narrow({{0, offset, slice_tokens}})->copy_from(slice_output.hidden_states); - - offset += slice_tokens; - } - - return CudaFusedMoeRunnerOutput{ - workspace.fused_moe_output, - }; -} - } // namespace infinilm::layers::moe diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp index b42be70d2..be557fa03 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp @@ -4,9 +4,6 @@ namespace infinilm::layers::moe { -struct HygonW16A16MarlinRuntimeConfig; -struct HygonW8A8MarlinRuntimeConfig; - struct CudaFusedMoeRunnerInput { infinicore::Tensor hidden_states; TopKOutput topk_output; @@ -19,43 +16,25 @@ 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); - - CombineInput run(const DispatchOutput &dispatch_output, - const MoeWeights &weights, - MoeWorkspace &workspace) const override; - -private: - 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; - - CudaFusedMoeRunnerOutput run_hygon_w16a16_marlin_core( - const CudaFusedMoeRunnerInput &runner_input, - const MoeWeights &weights, - MoeWorkspace &workspace, - const HygonW16A16MarlinRuntimeConfig &config) const; + CudaFusedMoeRunner( + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size_per_partition, + size_t align_block_size); - CudaFusedMoeRunnerOutput run_hygon_w16a16_marlin_core_sliced( + CombineInput run( const DispatchOutput &dispatch_output, const MoeWeights &weights, - MoeWorkspace &workspace) const; + MoeWorkspace &workspace) const override; - CudaFusedMoeRunnerOutput run_hygon_w8a8_marlin_core( - const CudaFusedMoeRunnerInput &runner_input, - const MoeWeights &weights, +private: + CudaFusedMoeRunnerInput prepare_runner_input( + const DispatchOutput &dispatch_output, MoeWorkspace &workspace, - const HygonW8A8MarlinRuntimeConfig &config) const; + size_t block_size) const; - CudaFusedMoeRunnerOutput run_hygon_w8a8_marlin_core_sliced( - const DispatchOutput &dispatch_output, + CudaFusedMoeRunnerOutput run_fused_core( + const CudaFusedMoeRunnerInput &runner_input, const MoeWeights &weights, MoeWorkspace &workspace) const; From 6a486f249c9d926b4a1e9c50a1f2770d582b910c Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Wed, 5 Aug 2026 15:32:36 +0800 Subject: [PATCH 14/14] perf-hygon-shard-lm-head-across-tp-ranks --- .../causal_lm_templates/text_causal_lm.hpp | 51 +++++++++++++++++-- 1 file changed, 47 insertions(+), 4 deletions(-) diff --git a/csrc/layers/causal_lm_templates/text_causal_lm.hpp b/csrc/layers/causal_lm_templates/text_causal_lm.hpp index fdc1e0774..135a165ad 100644 --- a/csrc/layers/causal_lm_templates/text_causal_lm.hpp +++ b/csrc/layers/causal_lm_templates/text_causal_lm.hpp @@ -4,6 +4,7 @@ #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 @@ -42,10 +43,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); } } @@ -66,7 +80,7 @@ class TextCausalLM : public InfinilmModel { hidden_states, input.input_offsets.value()); } - auto logits = lm_head_->forward(hidden_states); + auto logits = gather_logits(lm_head_->forward(hidden_states)); return {logits, hidden_states}; } @@ -74,20 +88,49 @@ 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