feat(qy): MiniCPM5 FA, longctx MoE/RoPE, CUDA-graph capture safety

csrc/engine/compiled_prefill_flags.hpp

@@ -0,0 +1,111 @@
+#pragma once
+
+#include <cstdlib>
+#include <cstring>
+#include <string>
+
+namespace infinilm::engine {
+namespace {
+
+inline const char *cudagraph_policy_cstr_() {
+    const char *v = std::getenv("INFINI_CUDAGRAPH_POLICY");
+    if (v == nullptr || v[0] == '\0') {
+        return "";
+    }
+    if (std::strcmp(v, "eager") == 0) {
+        return "eager";
+    }
+    if (std::strcmp(v, "full_and_piecewise") == 0) {
+        return "full_and_piecewise";
+    }
+    // Reject unknown / track_b — legacy flag behavior.
+    return "";
+}
+
+inline bool env_is_one_(const char *name) {
+    const char *v = std::getenv(name);
+    return v != nullptr && v[0] == '1' && v[1] == '\0';
+}
+
+} // namespace
+
+/// ``INFINI_CUDAGRAPH_POLICY`` when set to a known value; empty = legacy envs only.
+inline const char *cudagraph_policy() {
+    return cudagraph_policy_cstr_();
+}
+
+/// When true, C++ ``PagedCompiler`` must not capture or replay monolithic prefill CUDA graphs
+/// (torch compiled prefill owns the prefill path). Under ``full_and_piecewise``, native
+/// piecewise owns prefill (``native_piecewise_prefill_enabled``); monolithic is skipped
+/// via that check in ``PagedCompiler`` even when this returns false.
+inline bool skip_cpp_prefill_graph() {
+    if (std::strcmp(cudagraph_policy(), "eager") == 0) {
+        return true;
+    }
+    const char *v = std::getenv("INFINI_PREFILL_COMPILE");
+    if (v != nullptr && v[0] != '\0' && std::string(v) != "0") {
+        return true;
+    }
+    // Decode-only CG for MiniCPM5 MoE: skip monolithic prefill (HTC mem risk).
+    return env_is_one_("INFINI_DECODE_GRAPH_ONLY");
+}
+
+/// Skip monolithic full-forward decode CUDAGraph capture.
+/// MiniCPM5 MoE CG-1: MoE+FA are host-breaks, but remaining ops in one capture still
+/// HTC-fault on MetaX; use decode-piecewise instead (``native_piecewise_decode_enabled``).
+/// Under ``full_and_piecewise``, monolithic FULL decode is the default (unless SKIP /
+/// DECODE_GRAPH_ONLY). Explicit ``eager`` disables all decode CG.
+inline bool skip_monolithic_decode_graph() {
+    if (std::strcmp(cudagraph_policy(), "eager") == 0) {
+        return true;
+    }
+    if (env_is_one_("INFINI_SKIP_MONOLITHIC_DECODE_CG")) {
+        return true;
+    }
+    if (std::strcmp(cudagraph_policy(), "full_and_piecewise") == 0) {
+        // FULL decode unless opt-out via DECODE_GRAPH_ONLY / SKIP (checked above).
+        return env_is_one_("INFINI_DECODE_GRAPH_ONLY");
+    }
+    return env_is_one_("INFINI_DECODE_GRAPH_ONLY");
+}
+
+/// When true, use native C++ ``PiecewiseDecodeCompiler`` (per-layer pre / eager FA /
+/// post-without-MoE / eager MoE). Default on when monolithic decode CG is skipped
+/// (legacy). Under ``full_and_piecewise`` default off (monolithic FULL); under
+/// ``eager`` always off.
+inline bool native_piecewise_decode_enabled() {
+    if (std::strcmp(cudagraph_policy(), "eager") == 0) {
+        return false;
+    }
+    const char *v = std::getenv("INFINI_DECODE_PIECEWISE");
+    if (v != nullptr) {
+        return v[0] != '\0' && std::string(v) != "0";
+    }
+    if (std::strcmp(cudagraph_policy(), "full_and_piecewise") == 0) {
+        return false;
+    }
+    return skip_monolithic_decode_graph();
+}
+
+/// When true, use native C++ ``PiecewisePrefillCompiler`` for prefill (HPCC v1 path).
+/// Policy set: ``eager`` → off; ``full_and_piecewise`` → on (ignores
+/// ``INFINI_PREFILL_NATIVE_CG``). Unset policy (legacy Track-B scripts): read
+/// ``INFINI_PREFILL_NATIVE_CG`` only.
+inline bool native_piecewise_prefill_enabled() {
+    if (std::strcmp(cudagraph_policy(), "eager") == 0) {
+        return false;
+    }
+    if (std::strcmp(cudagraph_policy(), "full_and_piecewise") == 0) {
+        return true;
+    }
+    // Legacy: no known INFINI_CUDAGRAPH_POLICY → env flag until launchers migrate.
+    const char *v = std::getenv("INFINI_PREFILL_NATIVE_CG");
+    return v != nullptr && v[0] != '\0' && std::string(v) != "0";
+}
+
+/// When true, allow paged decode CUDAGraph capture/replay under tensor parallel (tp>1).
+inline bool decode_cg_tp_enabled() {
+    return env_is_one_("INFINI_DECODE_CG_TP");
+}
+
+} // namespace infinilm::engine

csrc/engine/compiler/attn_metadata_utils.hpp

@@ -0,0 +1,150 @@
+#pragma once
+
+#include "../../global_state/global_state.hpp"
+#include "../../models/infinilm_model.hpp"
+
+#include "infinicore/context/context.hpp"
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+namespace infinilm::engine::attn_metadata_utils {
+
+/// Refresh host max_query_len / max_kv_len from length tensors (once per forward).
+/// Uses Tensor::to(cpu()) — same safe path as compute_prefill_len — instead of raw
+/// memcpyD2H on possibly-narrowed / graph-backed views (hcErrorInvalidValue).
+inline void refresh_max_seqlens(infinilm::global_state::AttentionMetadata &meta) {
+    meta.max_query_length = 0;
+    meta.max_sequence_length = 0;
+
+    if (meta.input_offsets.has_value()) {
+        const auto &t = meta.input_offsets.value();
+        const size_t n = t->size(0);
+        if (n >= 2) {
+            auto cpu = t->to(infinicore::Device{infinicore::Device::Type::kCpu});
+            const auto *offs = reinterpret_cast<const int32_t *>(cpu->data());
+            for (size_t i = 0; i + 1 < n; ++i) {
+                meta.max_query_length = std::max(meta.max_query_length, static_cast<size_t>(offs[i + 1] - offs[i]));
+            }
+        }
+    }
+
+    if (meta.total_sequence_lengths.has_value()) {
+        const auto &t = meta.total_sequence_lengths.value();
+        const size_t n = t->size(0);
+        if (n > 0) {
+            auto cpu = t->to(infinicore::Device{infinicore::Device::Type::kCpu});
+            const auto *totals = reinterpret_cast<const int32_t *>(cpu->data());
+            for (size_t i = 0; i < n; ++i) {
+                meta.max_sequence_length = std::max(meta.max_sequence_length, static_cast<size_t>(totals[i]));
+            }
+        }
+    }
+
+    // Prefer cu_seqlens diffs when present (authoritative KV spans for FA varlen).
+    if (meta.cu_seqlens.has_value()) {
+        const auto &t = meta.cu_seqlens.value();
+        const size_t n = t->size(0);
+        if (n >= 2) {
+            auto cpu = t->to(infinicore::Device{infinicore::Device::Type::kCpu});
+            const auto *cu = reinterpret_cast<const int32_t *>(cpu->data());
+            for (size_t i = 0; i + 1 < n; ++i) {
+                meta.max_sequence_length = std::max(meta.max_sequence_length, static_cast<size_t>(cu[i + 1] - cu[i]));
+            }
+        }
+    }
+
+    // Clamp to block-table capacity so FA never indexes past table columns when
+    // max_position_embeddings (e.g. 131072) exceeds num_blocks * block_size.
+    if (meta.block_tables.has_value() && meta.block_tables.value()->ndim() >= 2) {
+        const size_t bt_cols = meta.block_tables.value()->size(1);
+        constexpr size_t kMaxPage = 256;
+        const size_t table_cap_hi = bt_cols * kMaxPage;
+        if (meta.max_sequence_length > table_cap_hi) {
+            meta.max_sequence_length = table_cap_hi;
+        }
+    }
+
+    if (meta.max_query_length == 0 && meta.slot_mapping.has_value()) {
+        meta.max_query_length = meta.slot_mapping.value()->shape()[0];
+    }
+    if (meta.max_sequence_length == 0) {
+        meta.max_sequence_length = meta.max_query_length;
+    }
+    if (meta.max_query_length == 0) {
+        meta.max_query_length = meta.max_sequence_length;
+    }
+}
+
+inline void set_attn_metadata(const InfinilmModel::Input &input) {
+    auto &meta = infinilm::global_state::get_forward_context().attn_metadata;
+    meta = {
+        input.past_sequence_lengths,
+        input.total_sequence_lengths,
+        input.input_offsets,
+        input.cu_seqlens,
+        input.block_tables,
+        input.slot_mapping,
+    };
+    refresh_max_seqlens(meta);
+}
+
+/// Narrow compiled graph buffers to runtime active shapes (RC-2 analog for replay).
+inline void set_attn_metadata_for_varlen_batch(const InfinilmModel::Input &compiled,
+                                               const InfinilmModel::Input &runtime) {
+    const size_t runtime_n_req = runtime.block_tables.value()->size(0);
+    const size_t block_per_req = runtime.block_tables.value()->size(1);
+    const size_t offset_len = runtime.input_offsets.value()->size(0);
+    const size_t cu_len = runtime.cu_seqlens.value()->size(0);
+    const size_t slot_len = runtime.slot_mapping.value()->shape()[0];
+
+    auto &meta = infinilm::global_state::get_forward_context().attn_metadata;
+    meta.past_sequence_lengths = compiled.past_sequence_lengths.has_value()
+                                     ? std::optional<infinicore::Tensor>(
+                                           compiled.past_sequence_lengths.value()->narrow({{0, 0, runtime_n_req}}))
+                                     : std::nullopt;
+    meta.total_sequence_lengths = compiled.total_sequence_lengths.value()->narrow({{0, 0, runtime_n_req}});
+    meta.input_offsets = compiled.input_offsets.value()->narrow({{0, 0, offset_len}});
+    meta.cu_seqlens = compiled.cu_seqlens.value()->narrow({{0, 0, cu_len}});
+    meta.block_tables = compiled.block_tables.value()->narrow({{0, 0, runtime_n_req}, {1, 0, block_per_req}});
+    // paged_caching uses slot_mapping.shape[0] as num_tokens (see paged_caching/info.h).
+    meta.slot_mapping = compiled.slot_mapping.value()->narrow({{0, 0, slot_len}});
+    // Host max seqlens must come from RUNTIME length tensors: narrowed compiled/graph
+    // views fail raw D2H (and may fail to(cpu) on graph storage). FA still uses
+    // narrowed compiled meta (CG-safe addresses); runtime lengths were already
+    // copy_from'd into those buffers in copy_runtime_into_bucket_. Do NOT swap meta
+    // to runtime tensor objects under pad-up — that breaks CG address binding (SIGSEGV).
+    //
+    // Clamp capacity from *compiled* block_tables (full num_blocks width), not the
+    // runtime table's shorter padded width. Large-past MIXED (past≫capture, e.g.
+    // ~18k) was clamping max_kv_len to runtime_bt_cols*256 and under-covering the
+    // cu_seqlens KV span → FA varlen OOB / silent death (211225 class).
+    {
+        auto io = meta.input_offsets;
+        auto tot = meta.total_sequence_lengths;
+        auto cu = meta.cu_seqlens;
+        auto bt = meta.block_tables;
+        meta.input_offsets = runtime.input_offsets;
+        meta.total_sequence_lengths = runtime.total_sequence_lengths;
+        meta.cu_seqlens = runtime.cu_seqlens;
+        meta.block_tables = bt; // compiled width for table_cap clamp
+        refresh_max_seqlens(meta);
+        const int mq = meta.max_query_length;
+        const int mk = meta.max_sequence_length;
+        meta.input_offsets = io;
+        meta.total_sequence_lengths = tot;
+        meta.cu_seqlens = cu;
+        meta.block_tables = bt;
+        meta.max_query_length = mq;
+        meta.max_sequence_length = mk;
+    }
+}
+
+/// Decode replay: same narrow helper; decode batches use fixed width == num_requests.
+inline void set_attn_metadata_for_decode_batch(const InfinilmModel::Input &compiled,
+                                               const InfinilmModel::Input &runtime) {
+    set_attn_metadata_for_varlen_batch(compiled, runtime);
+}
+
+} // namespace infinilm::engine::attn_metadata_utils

csrc/engine/compiler/general_compiler.cpp

@@ -1,20 +1,47 @@
 #include "general_compiler.hpp"
 
+#include "../compiled_prefill_flags.hpp"
+#include "../../global_state/global_state.hpp"
+#include "../../utils.hpp"
+#include "infinicore/context/context.hpp"
+
 namespace infinilm::engine {
-GeneralCompiler::GeneralCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier) : GraphCompiler(model, barrier) {
+
+namespace {
+
+void zero_kv_caches_after_compile_() {
+    auto &kv_cache_vec = infinilm::global_state::get_forward_context().kv_cache_vec;
+    for (auto &kv : kv_cache_vec) {
+        if (kv) {
+            set_zeros(kv);
+        }
+    }
+    infinicore::context::syncStream();
+}
+
+} // namespace
+
+GeneralCompiler::GeneralCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier)
+    : GraphCompiler(model, barrier) {
     static_batching_compiler_ = std::make_unique<StaticBatchingCompiler>(model_, barrier);
     paged_compiler_ = std::make_unique<PagedCompiler>(model_, barrier);
+    if (native_piecewise_prefill_enabled()) {
+        piecewise_prefill_compiler_ = std::make_unique<PiecewisePrefillCompiler>(model_, barrier);
+    }
 }
 
 void GeneralCompiler::compile() {
     static_batching_compiler_->compile();
+    if (piecewise_prefill_compiler_ != nullptr) {
+        piecewise_prefill_compiler_->compile();
+        zero_kv_caches_after_compile_();
+    }
     paged_compiler_->compile();
+    zero_kv_caches_after_compile_();
 }
 
 GeneralCompiler::Compiled GeneralCompiler::get_compiled(const InfinilmModel::Input &input) {
     GeneralCompiler::Compiled result = {nullptr, nullptr};
-
-    // try each compiler, return the first valid result
     result = static_batching_compiler_.get()->get_compiled(input);
     if (std::get<0>(result) != nullptr && std::get<1>(result) != nullptr) {
         return result;
@@ -23,4 +50,49 @@ GeneralCompiler::Compiled GeneralCompiler::get_compiled(const InfinilmModel::Inp
     return result;
 }
 
+std::optional<infinicore::Tensor> GeneralCompiler::run_native_piecewise_prefill(
+    const InfinilmModel::Input &input) {
+    if (piecewise_prefill_compiler_ == nullptr) {
+        return std::nullopt;
+    }
+    return piecewise_prefill_compiler_->run_prefill(input);
+}
+
+bool GeneralCompiler::native_piecewise_last_prefill_executed() const {
+    if (piecewise_prefill_compiler_ == nullptr) {
+        return false;
+    }
+    return piecewise_prefill_compiler_->last_prefill_executed();
+}
+
+bool GeneralCompiler::native_piecewise_enabled() const {
+    return piecewise_prefill_compiler_ != nullptr && piecewise_prefill_compiler_->enabled();
+}
+
+const std::vector<size_t> &GeneralCompiler::native_capture_buckets() const {
+    static const std::vector<size_t> empty;
+    if (piecewise_prefill_compiler_ == nullptr) {
+        return empty;
+    }
+    return piecewise_prefill_compiler_->capture_buckets();
+}
+
+void GeneralCompiler::record_graph_hit(bool is_prefill) {
+    paged_compiler_->record_graph_hit(is_prefill);
+}
+
+void GeneralCompiler::record_graph_miss(bool is_prefill) {
+    paged_compiler_->record_graph_miss(is_prefill);
+}
+
+PagedCompiler::GraphStats GeneralCompiler::graph_stats() const {
+    auto stats = paged_compiler_->graph_stats();
+    if (piecewise_prefill_compiler_ != nullptr) {
+        stats.piecewise_segment_replays = piecewise_prefill_compiler_->segment_replays();
+        stats.piecewise_prefill_hits = piecewise_prefill_compiler_->prefill_hits();
+        stats.piecewise_prefill_misses = piecewise_prefill_compiler_->prefill_misses();
+    }
+    return stats;
+}
+
 } // namespace infinilm::engine

csrc/engine/compiler/general_compiler.hpp

@@ -1,8 +1,11 @@
 #pragma once
 
 #include "paged_compiler.hpp"
+#include "piecewise_prefill_compiler.hpp"
 #include "static_batching_compiler.hpp"
 
+#include <optional>
+
 namespace infinilm::engine {
 class GeneralCompiler : public GraphCompiler {
 public:
@@ -12,8 +15,22 @@ class GeneralCompiler : public GraphCompiler {
 
     Compiled get_compiled(const InfinilmModel::Input &input) override;
 
+    std::optional<infinicore::Tensor> run_native_piecewise_prefill(const InfinilmModel::Input &input);
+
+    bool native_piecewise_last_prefill_executed() const;
+
+    bool native_piecewise_enabled() const;
+    // Piecewise decode optional; default OFF under F&P — keep flag false.
+    bool native_piecewise_decode_enabled() const { return false; }
+    const std::vector<size_t> &native_capture_buckets() const;
+
+    void record_graph_hit(bool is_prefill);
+    void record_graph_miss(bool is_prefill);
+    PagedCompiler::GraphStats graph_stats() const;
+
 private:
     std::unique_ptr<StaticBatchingCompiler> static_batching_compiler_;
     std::unique_ptr<PagedCompiler> paged_compiler_;
+    std::unique_ptr<PiecewisePrefillCompiler> piecewise_prefill_compiler_;
 };
 } // namespace infinilm::engine