feat(hygon): add Qwen3-235B-A3B BF16/W8A8 inference support

README.md

@@ -159,6 +159,15 @@
         > 注意：`--cache-dir` 应指向包含 `ceval___ceval-exam` 和 `cais___mmlu` 等数据集子目录的父目录，而不是直接指向这些子目录
 
   - 试验中功能
+    - 单次加载模型测试多组长度
+      ```bash
+      python examples/bench.py --device nvidia --model=<model-path> --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=<model-path> --warmup

csrc/config/model_config.hpp

@@ -88,6 +88,26 @@ class ModelConfig {
         return quant_config.get_quantization_method();
     }
 
+    std::string get_moe_weight_method() const {
+        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;
 

csrc/config/quant_config.cpp

@@ -1,6 +1,66 @@
 #include "quant_config.hpp"
 
+#include <algorithm>
+#include <cctype>
+
 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<char>(std::tolower(ch));
+    });
+    return value;
+}
+
+bool is_w16a16_marlin_method(const std::string &method) {
+    return method == "w16a16_marlin" || method == "hygon_w16a16_marlin";
+}
+
+bool is_w8a8_marlin_method(const std::string &method) {
+    return method == "slimquant_marlin" || method == "slimquant_compressed_tensors_marlin" ||
+           method == "w8a8_marlin" || method == "hygon_w8a8_marlin";
+}
+
+bool is_unquantized_config(const nlohmann::json &quantization_config) {
+    if (quantization_config.is_null()) {
+        return true;
+    }
+    if (!quantization_config.is_object()) {
+        return false;
+    }
+    auto it = quantization_config.find("quant_method");
+    if (it == quantization_config.end() || it->is_null()) {
+        return true;
+    }
+    if (!it->is_string()) {
+        return false;
+    }
+    auto method = lower_string(it->get<std::string>());
+    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<std::string>());
+        }
+    }
+    auto it = quantization_config.find("quant_method");
+    if (it != quantization_config.end() && it->is_string()) {
+        auto method = lower_string(it->get<std::string>());
+        if (is_w16a16_marlin_method(method) || is_w8a8_marlin_method(method)) {
+            return method;
+        }
+    }
+    return {};
+}
+
+} // namespace
 QuantConfig::QuantConfig(const nlohmann::json &json) : quantization_config(json) {
     this->quantization_method = get_quantization_method();
 }
@@ -20,11 +80,48 @@ QuantConfig::get_quantization_method() const {
         return std::make_shared<infinilm::quantization::AWQ>(quantization_config);
     } else if (quant_method == "gptq") {
         return std::make_shared<infinilm::quantization::GPTQ>(quantization_config);
+    } else if (quantization_config["quant_method"] == "w16a16_marlin" ||
+               quantization_config["quant_method"] == "hygon_w16a16_marlin") {
+        return std::make_shared<infinilm::quantization::NoneQuantization>(quantization_config);
     } else {
         return std::make_shared<infinilm::quantization::NoneQuantization>(quantization_config);
     }
     // Add other schemes as needed
 
     return std::make_shared<infinilm::quantization::NoneQuantization>(quantization_config); // Default case if no matching scheme
 }
+
+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 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 {
+    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

csrc/config/quant_config.hpp

@@ -1,8 +1,10 @@
 #pragma once
 #include "../utils.hpp"
 #include "../layers/quantization/quantization.hpp"
+#include "infinicore/device.hpp"
 #include "nlohmann/json.hpp"
 #include <optional>
+#include <string>
 #include <spdlog/spdlog.h>
 
 namespace infinilm::config {
@@ -15,6 +17,11 @@ class QuantConfig {
     QuantConfig(const nlohmann::json &json);
 
     std::shared_ptr<infinilm::quantization::BaseQuantization> 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) {

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();
@@ -164,20 +174,36 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &
             if (result == compiled_map_decode_.end()) {
                 return {nullptr, nullptr};
             }
-            auto &graph_input = result->second.input;
 
-            graph_input.input_ids.value()->copy_from(input.input_ids.value());
-            graph_input.position_ids.value()->copy_from(input.position_ids.value());
-            graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value());
-            graph_input.input_offsets.value()->copy_from(input.input_offsets.value());
-            graph_input.cu_seqlens.value()->copy_from(input.cu_seqlens.value());
+            // Decode graphs are captured with one token per request, so their
+            // input offsets are the fixed sequence [0, 1, ..., batch_size].
+            // Reuse the captured tensor only after validating that the runtime
+            // input has the same layout; otherwise fall back to eager mode.
+            const auto &runtime_input_offsets = input.input_offsets.value();
+            if (!runtime_input_offsets->is_contiguous() ||
+                runtime_input_offsets->size(0) != batch_size + 1) {
+                return {nullptr, nullptr};
+            }
+            const auto *offsets = reinterpret_cast<const int32_t *>(runtime_input_offsets->data());
+            for (size_t i = 0; i <= batch_size; ++i) {
+                if (offsets[i] != static_cast<int32_t>(i)) {
+                    return {nullptr, nullptr};
+                }
+            }
+            auto &graph_input = result->second.input;
 
             const size_t compiled_block_per_req = graph_input.block_tables.value()->size(1);
             if (block_per_req > compiled_block_per_req) {
-                // Runtime width exceeds compiled graph slot; fall back to eager path.
+                // Runtime width exceeds compiled graph slot; fall back before
+                // enqueueing copies that the eager path cannot consume.
                 return {nullptr, nullptr};
             }
 
+            graph_input.input_ids.value()->copy_from(input.input_ids.value());
+            graph_input.position_ids.value()->copy_from(input.position_ids.value());
+            graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value());
+            graph_input.cu_seqlens.value()->copy_from(input.cu_seqlens.value());
+
             // Initialize only the active graph rows to -1, then overwrite the
             // runtime logical region. Avoid clearing the full preallocated
             // holder on every decode token.
@@ -202,7 +228,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<InfinilmModel::Output>(new InfinilmModel::Output{std::get<1>(result->second.compiled)->logits->resume_from_blob_()});

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;

csrc/engine/infer_engine.cpp

@@ -144,6 +144,44 @@ std::vector<std::string> 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,7 +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 && !return_nll;
     infinilm::global_state::get_forward_context().attn_metadata = {
         input.past_sequence_lengths,
         input.total_sequence_lengths,
@@ -204,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);