Pr 532

csrc/config/model_config.hpp

@@ -88,6 +88,18 @@ class ModelConfig {
         return quant_config.get_quantization_method();
     }
 
+    std::string get_moe_weight_method(const infinicore::Device &device) const {
+        return quant_config.get_moe_weight_method(device);
+    }
+
+    bool is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const {
+        return quant_config.is_moe_w16a16_marlin_enabled(device);
+    }
+
+    bool is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const {
+        return quant_config.is_moe_w8a8_marlin_enabled(device);
+    }
+
     infinicore::DataType get_dtype() const;
     infinilm::quantization::QuantScheme get_quant_scheme() const;
 

csrc/config/quant_config.cpp

@@ -1,6 +1,65 @@
 #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 +79,39 @@ 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 infinicore::Device &device) const {
+    auto configured_method = explicit_moe_weight_method(quantization_config);
+    if (!configured_method.empty()) {
+        return configured_method;
+    }
+    if (quantization_method != nullptr) {
+        auto method = quantization_method->get_moe_weight_method(device);
+        if (method != "dense") {
+            return method;
+        }
+    }
+    if (device.getType() == infinicore::Device::Type::HYGON && is_unquantized_config(quantization_config)) {
+        return "hygon_w16a16_marlin";
+    }
+    return "dense";
+}
+
+bool QuantConfig::is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const {
+    return is_w16a16_marlin_method(get_moe_weight_method(device));
+}
+
+bool QuantConfig::is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const {
+    return is_w8a8_marlin_method(get_moe_weight_method(device));
+}
+
 } // namespace infinilm::config

csrc/config/quant_config.hpp

@@ -1,9 +1,11 @@
 #pragma once
-#include "../utils.hpp"
 #include "../layers/quantization/quantization.hpp"
+#include "../utils.hpp"
+#include "infinicore/device.hpp"
 #include "nlohmann/json.hpp"
 #include <optional>
 #include <spdlog/spdlog.h>
+#include <string>
 
 namespace infinilm::config {
 
@@ -15,6 +17,9 @@ class QuantConfig {
     QuantConfig(const nlohmann::json &json);
 
     std::shared_ptr<infinilm::quantization::BaseQuantization> get_quantization_method() const;
+    std::string get_moe_weight_method(const infinicore::Device &device) const;
+    bool is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const;
+    bool is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const;
 
     infinilm::quantization::QuantScheme get_quant_scheme() const {
         if (quantization_method != nullptr) {

csrc/engine/compiler/paged_compiler.cpp

@@ -63,13 +63,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());
@@ -86,7 +90,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());
@@ -138,8 +144,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_) {
@@ -152,8 +160,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();
@@ -180,20 +190,29 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &
             if (result == compiled_map_decode_.end()) {
                 return {nullptr, nullptr};
             }
-            auto &graph_input = result->second.input;
 
-            graph_input.input_ids.value()->copy_from(input.input_ids.value());
-            graph_input.position_ids.value()->copy_from(input.position_ids.value());
-            graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value());
-            graph_input.input_offsets.value()->copy_from(input.input_offsets.value());
-            graph_input.cu_seqlens.value()->copy_from(input.cu_seqlens.value());
+            // Decode graphs are captured with one token per request, so their
+            // input offsets are the fixed sequence [0, 1, ..., batch_size].
+            // Reuse the captured tensor only after validating that the runtime
+            // input has the same layout; otherwise fall back to eager mode.
+            const auto &runtime_input_offsets = input.input_offsets.value();
+            if (!runtime_input_offsets->is_contiguous() || runtime_input_offsets->size(0) != batch_size + 1) {
+                return {nullptr, nullptr};
+            }
+            auto &graph_input = result->second.input;
 
             const size_t compiled_block_per_req = graph_input.block_tables.value()->size(1);
             if (block_per_req > compiled_block_per_req) {
-                // Runtime width exceeds compiled graph slot; fall back to eager path.
+                // Runtime width exceeds compiled graph slot; fall back before
+                // enqueueing copies that the eager path cannot consume.
                 return {nullptr, nullptr};
             }
 
+            graph_input.input_ids.value()->copy_from(input.input_ids.value());
+            graph_input.position_ids.value()->copy_from(input.position_ids.value());
+            graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value());
+            graph_input.cu_seqlens.value()->copy_from(input.cu_seqlens.value());
+
             // Initialize only the active graph rows to -1, then overwrite the
             // runtime logical region. Avoid clearing the full preallocated
             // holder on every decode token.
@@ -218,7 +237,9 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &
             // one on the same stream before launch. This is correct but costs
             // decode latency; the intended follow-up is a reusable global
             // zero workspace/lock buffer shared by all Marlin layers.
-            model_->reset_runtime_state();
+            if (decode_graph_needs_runtime_state_reset_) {
+                model_->reset_runtime_state();
+            }
 
             auto graph = std::get<0>(result->second.compiled);
             if (graph != nullptr) {

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

@@ -184,6 +184,7 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         image_req_ids,
         visual_token_ranges,
         to_device(target_hidden_states)};
+    input.last_token_only = !sample_all_positions;
 
     infinilm::global_state::get_forward_context().attn_metadata = {
         input.past_sequence_lengths,

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<size_t>(input_offsets[i + 1] - 1);
+                                    score_idx = total_len == n_req
+                                                  ? i
+                                                  : static_cast<size_t>(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({})};

csrc/layers/attention/backends/flash_attn.cpp

@@ -41,6 +41,23 @@ infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer,
     ASSERT(block_tables.has_value());
     ASSERT(slot_mapping.has_value());
 
+    if (query->device().getType() == infinicore::Device::Type::HYGON) {
+        return infinicore::op::paged_flash_attention(
+            query,
+            key,
+            value,
+            kv_cache,
+            total_sequence_lengths.value(),
+            input_offsets,
+            cu_seqlens,
+            block_tables.value(),
+            slot_mapping.value(),
+            num_heads_,
+            num_kv_heads_,
+            head_dim_,
+            scale_);
+    }
+
     // 1. update paged kv cache
     auto [k_total, v_total] = do_kv_cache_update(layer, key, value, kv_cache, slot_mapping.value());