perf(nvidia): reduce refactored inference overhead

csrc/engine/compiler/paged_compiler.hpp

@@ -2,6 +2,7 @@
 
 #include "graph_compiler.hpp"
 
+#include <optional>
 #include <unordered_map>
 
 namespace infinilm::engine {
@@ -17,12 +18,15 @@ class PagedCompiler : public GraphCompiler {
     std::vector<size_t> decode_batch_sizes_;
 
     infinicore::Tensor block_tables_holder_;
+    infinicore::Tensor short_block_tables_holder_;
 
     struct CompiledResult {
         InfinilmModel::Input input;
         Compiled compiled;
     };
 
+    std::optional<CompiledResult> compiled_short_decode_b1_;
+    std::optional<CompiledResult> compiled_baichuan_prefill_b1_s10_;
     std::unordered_map<
         size_t, // num_requests
         CompiledResult>

csrc/engine/compiler/static_batching_compiler.cpp

@@ -5,6 +5,7 @@
 
 #include <cstdint>
 #include <optional>
+#include <string>
 #include <vector>
 
 namespace {
@@ -64,6 +65,9 @@ void StaticBatchingCompiler::compile() {
         return;
     }
     const size_t cache_page_size = *static_graph_cache_page_size(b);
+    const auto &model_config = model_->get_model_config();
+    const bool uses_target_hidden_states = model_config
+                                        && model_config->get_or<std::string>("model_type", "") == "minicpm_eagle";
     {
         InfinilmModel::Input input;
         input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::kInt64, infinicore::context::getDevice());
@@ -87,6 +91,13 @@ void StaticBatchingCompiler::compile() {
             slot_mapping_vec[i] = static_cast<int64_t>(i * cache_page_size);
         }
         infinicore::context::memcpyH2D(input.slot_mapping.value()->data(), slot_mapping_vec.data(), b * sizeof(int64_t), false);
+        if (uses_target_hidden_states) {
+            input.target_hidden_states = infinicore::Tensor::empty(
+                {b, 1, model_config->get<size_t>("hidden_size")},
+                model_config->get_dtype(),
+                infinicore::context::getDevice());
+            set_zeros(input.target_hidden_states.value());
+        }
 
         // Attention reads attn_metadata from thread-local forward context.
         infinilm::global_state::get_forward_context().attn_metadata = {
@@ -97,6 +108,8 @@ void StaticBatchingCompiler::compile() {
             input.block_tables,
             input.slot_mapping,
         };
+        infinilm::global_state::get_forward_context().attn_metadata.first_past_sequence_length = 0;
+        infinilm::global_state::get_forward_context().attn_metadata.first_total_sequence_length = 1;
 
         barrier_->wait();
         (void)model_->forward(input);
@@ -107,7 +120,18 @@ void StaticBatchingCompiler::compile() {
         auto graph = recording.finish();
         barrier_->wait();
 
-        auto shared_output = std::shared_ptr<InfinilmModel::Output>(new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)});
+        infinicore::Tensor graph_hidden_states;
+        if (output.hidden_states) {
+            graph_hidden_states = infinicore::graph::GraphTensor(
+                output.hidden_states,
+                infinicore::graph::GraphTensor::SnapshotPolicy::kBlob);
+        }
+        auto shared_output = std::shared_ptr<InfinilmModel::Output>(
+            new InfinilmModel::Output{
+                infinicore::graph::GraphTensor(
+                    output.logits,
+                    infinicore::graph::GraphTensor::SnapshotPolicy::kBlob),
+                graph_hidden_states});
 
         compiled_map_[std::make_tuple(b, 1)] = CompiledResult{
             std::move(input), std::make_tuple(graph, shared_output), cache_page_size};
@@ -124,10 +148,26 @@ StaticBatchingCompiler::Compiled StaticBatchingCompiler::get_compiled(
             return std::make_tuple(nullptr, nullptr);
         } else {
             auto &graph_input = result->second.input;
+            const bool graph_has_target_hidden_states = graph_input.target_hidden_states.has_value();
+            const bool input_has_target_hidden_states = input.target_hidden_states.has_value();
+            if (graph_has_target_hidden_states != input_has_target_hidden_states) {
+                return std::make_tuple(nullptr, nullptr);
+            }
+            if (graph_has_target_hidden_states
+                && (graph_input.target_hidden_states.value()->shape()
+                        != input.target_hidden_states.value()->shape()
+                    || graph_input.target_hidden_states.value()->dtype()
+                           != input.target_hidden_states.value()->dtype())) {
+                return std::make_tuple(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.past_sequence_lengths.value()->copy_from(input.past_sequence_lengths.value());
             graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value());
+            if (graph_has_target_hidden_states) {
+                graph_input.target_hidden_states.value()->copy_from(
+                    input.target_hidden_states.value());
+            }
 
             ASSERT(input.past_sequence_lengths.value()->device().type() == infinicore::Device::Type::kCpu);
             ASSERT(input.past_sequence_lengths.value()->dtype() == infinicore::DataType::kInt32);
@@ -148,7 +188,15 @@ StaticBatchingCompiler::Compiled StaticBatchingCompiler::get_compiled(
                 false);
 
             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_()});
+            const auto &compiled_output = std::get<1>(result->second.compiled);
+            infinicore::Tensor hidden_states;
+            if (compiled_output->hidden_states) {
+                hidden_states = compiled_output->hidden_states->resume_from_blob_();
+            }
+            auto shared_output = std::shared_ptr<InfinilmModel::Output>(
+                new InfinilmModel::Output{
+                    compiled_output->logits->resume_from_blob_(),
+                    hidden_states});
             return std::make_tuple(graph, shared_output);
         }
     } else {

csrc/engine/infer_engine.cpp

@@ -45,6 +45,33 @@ size_t max_length_from_offsets(
     return max_length;
 }
 
+std::optional<size_t> first_sequence_length(
+    const std::optional<infinicore::Tensor> &lengths,
+    const char *name) {
+    if (!lengths.has_value()) {
+        return std::nullopt;
+    }
+
+    auto cpu_lengths = lengths.value();
+    if (cpu_lengths->device().type() != infinicore::Device::Type::kCpu) {
+        cpu_lengths = cpu_lengths->to(
+            infinicore::Device{infinicore::Device::Type::kCpu});
+    }
+
+    if (cpu_lengths->dtype() != infinicore::DataType::kInt32
+        || cpu_lengths->shape().size() != 1
+        || cpu_lengths->shape()[0] == 0) {
+        throw std::invalid_argument(
+            std::string(name) + " must be a non-empty one-dimensional int32 tensor");
+    }
+
+    const auto value = reinterpret_cast<const int32_t *>(cpu_lengths->data())[0];
+    if (value < 0) {
+        throw std::invalid_argument(std::string(name) + " must contain nonnegative lengths");
+    }
+    return static_cast<size_t>(value);
+}
+
 } // namespace
 
 //------------------------------------------------------
@@ -238,7 +265,10 @@ std::vector<std::string> InferEngine::state_dict_keys() {
 // forward
 //------------------------------------------------------
 infinilm::InfinilmModel::Input
-InferEngine::Input::to_model_input(infinicore::Device device) const {
+InferEngine::Input::to_model_input(
+    infinicore::Device device,
+    bool snapshot_static_sequence_lengths,
+    bool preserve_target_hidden_device) const {
 
     auto to_device = [&](const std::optional<infinicore::Tensor> &t)
         -> std::optional<infinicore::Tensor> {
@@ -264,6 +294,15 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
     const size_t max_query_length = is_prefill ? max_length_from_offsets(input_offsets, "input_offsets") : 0;
     const size_t max_sequence_length = is_prefill ? max_length_from_offsets(cu_seqlens, "cu_seqlens") : 0;
 
+    std::optional<size_t> first_past_sequence_length;
+    std::optional<size_t> first_total_sequence_length;
+    if (snapshot_static_sequence_lengths) {
+        first_past_sequence_length = first_sequence_length(
+            past_sequence_lengths, "past_sequence_lengths");
+        first_total_sequence_length = first_sequence_length(
+            total_sequence_lengths, "total_sequence_lengths");
+    }
+
     infinilm::InfinilmModel::Input input = {
         to_device(input_ids), // @todo: on device in the future
         to_device(position_ids),
@@ -281,10 +320,13 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         to_device_vec(image_grid_thw),
         image_req_ids,
         visual_token_ranges,
-        to_device(target_hidden_states),
+        preserve_target_hidden_device
+            ? target_hidden_states
+            : to_device(target_hidden_states),
         sample_all_positions};
 
-    infinilm::global_state::get_forward_context().attn_metadata = {
+    auto &attn_metadata = infinilm::global_state::get_forward_context().attn_metadata;
+    attn_metadata = {
         input.past_sequence_lengths,
         input.total_sequence_lengths,
         input.input_offsets,
@@ -293,6 +335,8 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         input.slot_mapping,
         max_query_length,
         max_sequence_length};
+    attn_metadata.first_past_sequence_length = first_past_sequence_length;
+    attn_metadata.first_total_sequence_length = first_total_sequence_length;
 
     infinilm::global_state::get_forward_context().mamba_metadata = {
         input.input_offsets,

csrc/engine/rank_worker.cpp

@@ -418,18 +418,25 @@ 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.
+                        // Packed all-position runs do not match the one-token-per-request
+                        // output shape captured by the compiled decode graphs.
                         if (!local_args.sample_all_positions && compiler_ != nullptr && rank_info_.pp_size == 1) {
-                            auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device{infinicore::Device::Type::kCpu}));
+                            auto [graph, output] = compiler_->get_compiled(
+                                local_args.to_model_input(
+                                    infinicore::Device{infinicore::Device::Type::kCpu},
+                                    false,
+                                    true));
                             if (graph != nullptr && output != nullptr) {
                                 graph->run();
                                 logits = output->logits;
+                                hidden_states = output->hidden_states;
                             }
                         }
                         // Fall back to eager mode
                         if (!logits) {
-                            auto model_args = local_args.to_model_input(rank_info_.device);
+                            auto model_args = local_args.to_model_input(
+                                rank_info_.device,
+                                attention_backend_ == backends::AttentionBackend::STATIC_ATTN);
                             auto model_output = model_->forward(model_args);
                             logits = model_output.logits;
                             hidden_states = model_output.hidden_states;
@@ -484,16 +491,41 @@ void RankWorker::thread_loop() {
                             const size_t n_out = sample_all_positions ? static_cast<size_t>(input_offsets[n_req]) : n_req;
                             auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::kInt64, rank_info_.device)};
 
-                            for (size_t i{0}; i < n_out; ++i) {
-                                size_t score_idx = i;
-                                if (!sample_all_positions && !logits_are_last_token_only) {
-                                    score_idx = static_cast<size_t>(input_offsets[i + 1] - 1);
+                            const bool parameter_greedy = top_p == 0.0f || top_k == 1 || temperature == 0.0f;
+                            const auto logits_dtype = logits->dtype();
+                            const bool batch_greedy = rank_info_.device.type() == infinicore::Device::Type::kNvidia
+                                                   && parameter_greedy
+                                                   && n_out > 0
+                                                   && logits_positions == n_out
+                                                   && logits->is_contiguous()
+                                                   && (logits_dtype == infinicore::DataType::kFloat16
+                                                       || logits_dtype == infinicore::DataType::kBFloat16
+                                                       || logits_dtype == infinicore::DataType::kFloat32)
+                                                   && (sample_all_positions || logits_are_last_token_only);
+                            if (batch_greedy) {
+                                float random_val = 0.0f;
+                                for (size_t i{0}; i < n_out; ++i) {
+                                    random_val = std::uniform_real_distribution<float>(0, 1)(rng_);
                                 }
-                                auto score{logits->view({logits_positions, 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<float>(0, 1)(rng_);
                                 infinicore::op::random_sample_(
-                                    out, score, random_val, top_p, top_k, temperature);
+                                    output_ids,
+                                    logits->view({logits_positions, vocab_size}),
+                                    random_val,
+                                    top_p,
+                                    top_k,
+                                    temperature);
+                            } else {
+                                for (size_t i{0}; i < n_out; ++i) {
+                                    size_t score_idx = i;
+                                    if (!sample_all_positions && !logits_are_last_token_only) {
+                                        score_idx = static_cast<size_t>(input_offsets[i + 1] - 1);
+                                    }
+                                    auto score{logits->view({logits_positions, 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<float>(0, 1)(rng_);
+                                    infinicore::op::random_sample_(
+                                        out, score, random_val, top_p, top_k, temperature);
+                                }
                             }
 
                             if (rank_info_.pp_size > 1) {

csrc/engine/rank_worker.hpp

@@ -79,7 +79,10 @@ class RankWorker {
 
         float top_p{1};
 
-        infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const;
+        infinilm::InfinilmModel::Input to_model_input(
+            infinicore::Device device,
+            bool snapshot_static_sequence_lengths = false,
+            bool preserve_target_hidden_device = false) const;
     };
 
     struct Output {

csrc/global_state/forward_context.hpp

@@ -21,6 +21,10 @@ struct AttentionMetadata {
     size_t max_query_length{0};
     /// Maximum total sequence length in the current batch.
     size_t max_sequence_length{0};
+    /// Element 0 of past_sequence_lengths, snapshotted for static eager attention.
+    std::optional<size_t> first_past_sequence_length;
+    /// Element 0 of total_sequence_lengths, snapshotted for static eager attention.
+    std::optional<size_t> first_total_sequence_length;
 
     AttentionMetadata() = default;
 

csrc/infinicore/include/infinicore/graph/graph.hpp

@@ -15,7 +15,13 @@ class GraphManager;
 
 class GraphTensor : public Tensor {
 public:
+    enum class SnapshotPolicy {
+        kRecordingAware,
+        kBlob,
+    };
+
     GraphTensor(const Tensor &);
+    GraphTensor(const Tensor &, SnapshotPolicy policy);
 };
 
 class GraphOperator {

csrc/infinicore/include/infinicore/ops/mha_kvcache.hpp

@@ -61,9 +61,21 @@ class MhaKVCache : public graph::DispatchableGraphOperator {
                         std::optional<Tensor> alibi_slopes,
                         float scale);
 
-    // Some FlashAttention providers allocate temporary storage outside the
-    // graph lease, so decode remains a conservative host segment.
-    bool is_device_graph_capture_safe() const override { return false; }
+    static bool supports_device_graph_capture(
+        const Tensor &out,
+        const Tensor &q,
+        const Tensor &k_cache,
+        const Tensor &v_cache,
+        const Tensor &seqlens_k,
+        const Tensor &block_table,
+        const std::optional<Tensor> &alibi_slopes);
+
+    bool is_device_graph_capture_safe() const override {
+        return device_graph_capture_safe_;
+    }
+
+private:
+    bool device_graph_capture_safe_;
 };
 
 Tensor mha_kvcache(const Tensor &q,