diff --git a/csrc/engine/compiled_prefill_flags.hpp b/csrc/engine/compiled_prefill_flags.hpp new file mode 100644 index 000000000..1c5f80476 --- /dev/null +++ b/csrc/engine/compiled_prefill_flags.hpp @@ -0,0 +1,111 @@ +#pragma once + +#include +#include +#include + +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 diff --git a/csrc/engine/compiler/attn_metadata_utils.hpp b/csrc/engine/compiler/attn_metadata_utils.hpp new file mode 100644 index 000000000..c1520ddbe --- /dev/null +++ b/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 +#include +#include + +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(cpu->data()); + for (size_t i = 0; i + 1 < n; ++i) { + meta.max_query_length = std::max(meta.max_query_length, static_cast(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(cpu->data()); + for (size_t i = 0; i < n; ++i) { + meta.max_sequence_length = std::max(meta.max_sequence_length, static_cast(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(cpu->data()); + for (size_t i = 0; i + 1 < n; ++i) { + meta.max_sequence_length = std::max(meta.max_sequence_length, static_cast(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( + 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 diff --git a/csrc/engine/compiler/general_compiler.cpp b/csrc/engine/compiler/general_compiler.cpp index 84ee670d4..595121c0d 100644 --- a/csrc/engine/compiler/general_compiler.cpp +++ b/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 &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 &model, RankBarrier *barrier) + : GraphCompiler(model, barrier) { static_batching_compiler_ = std::make_unique(model_, barrier); paged_compiler_ = std::make_unique(model_, barrier); + if (native_piecewise_prefill_enabled()) { + piecewise_prefill_compiler_ = std::make_unique(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 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 &GeneralCompiler::native_capture_buckets() const { + static const std::vector 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 diff --git a/csrc/engine/compiler/general_compiler.hpp b/csrc/engine/compiler/general_compiler.hpp index e8b84b5d9..b4cd0a779 100644 --- a/csrc/engine/compiler/general_compiler.hpp +++ b/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 + 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 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 &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 static_batching_compiler_; std::unique_ptr paged_compiler_; + std::unique_ptr piecewise_prefill_compiler_; }; } // namespace infinilm::engine diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index 55a05da4b..fc4c73e51 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -1,6 +1,8 @@ #include "paged_compiler.hpp" +#include "../compiled_prefill_flags.hpp" #include "../../global_state/global_state.hpp" #include "../../utils.hpp" +#include "infinicore/context/context.hpp" #include #include @@ -209,6 +211,9 @@ void PagedCompiler::compile() { // Capture must not start with stale state from previous attempts. model_->reset_runtime_state(); infinicore::context::syncStream(); + // Decode FULL capture: FA in-graph when INFINI_FA_INGRAPH=decode|both. + infinicore::context::InferencePhaseGuard phase_guard( + infinicore::context::InferencePhase::Decode); GraphRecordingGuard recording; auto output = model_->forward(input); auto graph = recording.finish(); @@ -334,4 +339,29 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & return std::make_tuple(graph, shared_output); } +void PagedCompiler::record_graph_hit(bool is_prefill) { + if (is_prefill) { + ++prefill_graph_hits_; + } else { + ++decode_graph_hits_; + } +} + +void PagedCompiler::record_graph_miss(bool is_prefill) { + if (is_prefill) { + ++prefill_graph_misses_; + } else { + ++decode_graph_misses_; + } +} + +PagedCompiler::GraphStats PagedCompiler::graph_stats() const { + return GraphStats{ + prefill_graph_hits_, + prefill_graph_misses_, + decode_graph_hits_, + decode_graph_misses_, + }; +} + } // namespace infinilm::engine diff --git a/csrc/engine/compiler/paged_compiler.hpp b/csrc/engine/compiler/paged_compiler.hpp index a1125864d..d397701f0 100644 --- a/csrc/engine/compiler/paged_compiler.hpp +++ b/csrc/engine/compiler/paged_compiler.hpp @@ -7,12 +7,29 @@ namespace infinilm::engine { class PagedCompiler : public GraphCompiler { public: + struct GraphStats { + size_t prefill_graph_hits{0}; + size_t prefill_graph_misses{0}; + size_t decode_graph_hits{0}; + size_t decode_graph_misses{0}; + size_t piecewise_segment_replays{0}; + size_t piecewise_prefill_hits{0}; + size_t piecewise_prefill_misses{0}; + size_t piecewise_decode_hits{0}; + size_t piecewise_decode_misses{0}; + size_t piecewise_decode_device_segments{0}; + }; + PagedCompiler(const std::shared_ptr &model, RankBarrier *barrier); void compile() override; Compiled get_compiled(const InfinilmModel::Input &input) override; + void record_graph_hit(bool is_prefill); + void record_graph_miss(bool is_prefill); + GraphStats graph_stats() const; + private: std::vector decode_batch_sizes_; @@ -27,5 +44,10 @@ class PagedCompiler : public GraphCompiler { size_t, // num_requests CompiledResult> compiled_map_decode_; + + size_t prefill_graph_hits_{0}; + size_t prefill_graph_misses_{0}; + size_t decode_graph_hits_{0}; + size_t decode_graph_misses_{0}; }; } // namespace infinilm::engine diff --git a/csrc/engine/compiler/piecewise_bucket_policy.hpp b/csrc/engine/compiler/piecewise_bucket_policy.hpp new file mode 100644 index 000000000..d30243d21 --- /dev/null +++ b/csrc/engine/compiler/piecewise_bucket_policy.hpp @@ -0,0 +1,145 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace infinilm::engine { + +constexpr size_t kPiecewisePowerLadderFloor = 512; +constexpr size_t kPiecewisePowerLadderCap = 8192; +constexpr size_t kPiecewiseOverflowTailBucket = 8448; + +inline size_t compile_max_seq_from_env(size_t default_max = kPiecewiseOverflowTailBucket) { + const char *raw = std::getenv("INFINI_COMPILE_MAX_SEQ"); + if (raw == nullptr || raw[0] == '\0') { + return default_max; + } + return static_cast(std::stoul(raw)); +} + +inline bool vllm_capture_ladder_enabled() { + const char *raw = std::getenv("INFINI_VLLM_CAPTURE_LADDER"); + if (raw == nullptr || raw[0] == '\0') { + return false; + } + const std::string v(raw); + return v != "0" && v != "false" && v != "no" && v != "off"; +} + +/// Shared step budget / chunk size. Prefer ``INFINI_MAX_NUM_BATCHED_TOKENS`` +/// (Python canonical); fall back to deprecated ``INFINI_PREFILL_CHUNK_SIZE``. +inline size_t prefill_chunk_size_from_env(size_t default_size = 512) { + if (const char *raw = std::getenv("INFINI_MAX_NUM_BATCHED_TOKENS")) { + if (raw[0] != '\0') { + return std::max(1, static_cast(std::stoul(raw))); + } + } + if (const char *raw = std::getenv("INFINI_PREFILL_CHUNK_SIZE")) { + if (raw[0] != '\0') { + return std::max(1, static_cast(std::stoul(raw))); + } + } + return default_size; +} + +inline std::vector vllm_piecewise_capture_sizes(size_t chunk_cap) { + std::vector buckets; + const size_t cap = std::min(chunk_cap, kPiecewisePowerLadderFloor); + for (size_t s = 1; s <= cap; s *= 2) { + buckets.push_back(s); + if (s > SIZE_MAX / 2) { + break; + } + } + return buckets; +} + +inline std::vector merge_bucket_ladders(const std::vector &a, + const std::vector &b) { + std::vector merged = a; + merged.insert(merged.end(), b.begin(), b.end()); + std::sort(merged.begin(), merged.end()); + merged.erase(std::unique(merged.begin(), merged.end()), merged.end()); + return merged; +} + +inline std::vector piecewise_compile_buckets(size_t max_seq_len) { + std::vector buckets; + for (size_t s = kPiecewisePowerLadderFloor; s <= std::min(max_seq_len, kPiecewisePowerLadderCap); s *= 2) { + buckets.push_back(s); + } + if (max_seq_len > kPiecewisePowerLadderCap) { + buckets.push_back(max_seq_len); + } else if (max_seq_len > 0 && (buckets.empty() || buckets.back() != max_seq_len)) { + buckets.push_back(max_seq_len); + } + if (max_seq_len >= kPiecewisePowerLadderCap && max_seq_len < kPiecewiseOverflowTailBucket) { + buckets.push_back(kPiecewiseOverflowTailBucket); + } + std::sort(buckets.begin(), buckets.end()); + buckets.erase(std::unique(buckets.begin(), buckets.end()), buckets.end()); + return buckets; +} + +inline std::vector piecewise_compile_buckets_vllm(size_t max_seq_len, size_t chunk_cap) { + return merge_bucket_ladders( + vllm_piecewise_capture_sizes(std::min(chunk_cap, kPiecewisePowerLadderFloor)), + piecewise_compile_buckets(max_seq_len)); +} + +inline std::vector piecewise_capture_buckets_vllm(size_t max_seq_len, size_t chunk_cap) { + auto buckets = piecewise_compile_buckets_vllm(max_seq_len, chunk_cap); + buckets.erase( + std::remove_if(buckets.begin(), buckets.end(), + [](size_t b) { return b > kPiecewisePowerLadderCap; }), + buckets.end()); + return buckets; +} + +inline std::vector piecewise_capture_buckets(size_t max_seq_len) { + auto buckets = piecewise_compile_buckets(max_seq_len); + // Pad ladder may include 8448 overflow tail; graphs capture through 8192 only. + buckets.erase( + std::remove_if(buckets.begin(), buckets.end(), + [](size_t b) { return b > kPiecewisePowerLadderCap; }), + buckets.end()); + return buckets; +} + +inline std::vector build_bs_to_padded_bucket(const std::vector &capture_sizes) { + if (capture_sizes.empty()) { + return {0}; + } + std::vector sizes = capture_sizes; + std::sort(sizes.begin(), sizes.end(), std::greater()); + const size_t max_capture = sizes.front(); + std::vector table(max_capture + 1, 0); + for (size_t i = 0; i < sizes.size(); ++i) { + const size_t end = sizes[i]; + const size_t start = (i + 1 < sizes.size()) ? sizes[i + 1] : 0; + for (size_t bs = start; bs < end; ++bs) { + table[bs] = (bs == start) ? start : end; + } + } + table[max_capture] = max_capture; + return table; +} + +inline size_t padded_bucket_for_seq_len(size_t seq_len, + const std::vector &bs_to_padded, + size_t fallback) { + if (seq_len < bs_to_padded.size() && bs_to_padded[seq_len] > 0) { + return bs_to_padded[seq_len]; + } + return fallback; +} + +inline size_t graph_replay_bucket_for_padded(size_t padded_bucket) { + return padded_bucket; +} + +} // namespace infinilm::engine diff --git a/csrc/engine/compiler/piecewise_prefill_compiler.cpp b/csrc/engine/compiler/piecewise_prefill_compiler.cpp new file mode 100644 index 000000000..1dd330b97 --- /dev/null +++ b/csrc/engine/compiler/piecewise_prefill_compiler.cpp @@ -0,0 +1,1706 @@ +#include "piecewise_prefill_compiler.hpp" + +#include "../../global_state/ar_profile.hpp" +#include "../../global_state/global_state.hpp" +#include "../../global_state/hang_trace.hpp" +#include "../compiled_prefill_flags.hpp" +#include "../../global_state/piecewise_inductor_flags.hpp" +#include "../../utils.hpp" +#include "../dispatch_hist.hpp" +#include "piecewise_bucket_policy.hpp" +#include "attn_metadata_utils.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "infinicore/context/context.hpp" +#include "infinicore/graph/graph.hpp" +#include "infinicore/ops/inductor_segment.hpp" + +namespace infinilm::engine { + +namespace { + +bool graph_capture_audit_enabled_() { + const char *v = std::getenv("INFINI_GRAPH_CAPTURE_AUDIT"); + return v != nullptr && v[0] != '\0' && std::string(v) != "0"; +} + +/// When MoE is device-capturable, post_attn must be a single hcGraph segment +/// (true one-piece: o_proj+AR+LN+MoE). Fail fast under capture audit. +/// FORCE_OP_LIST skips device capture on purpose → device_segments=0 is OK. +void assert_post_one_device_seg_( + const char *where, + size_t layer, + const std::shared_ptr &g) { + if (!graph_capture_audit_enabled_() + || !infinicore::context::moeTritonCaptureAllowed() + || !g) { + return; + } + if (const char *fol = std::getenv("INFINI_GRAPH_FORCE_OP_LIST"); + fol != nullptr && fol[0] != '\0' && std::string(fol) != "0") { + return; + } + const size_t n = g->device_segment_count(); + if (n == 1) { + return; + } + spdlog::error( + "[capture_audit] {} layer={} post device_segments={} expected=1 " + "(MoE capturable — true one-piece post requires a single device seg)", + where, + layer, + n); + throw std::runtime_error( + std::string("[capture_audit] post device_segments!=1 when MoE capturable (") + + where + ")"); +} + +bool rank_worker_profile_enabled() { + static int cached = -1; + if (cached < 0) { + const char *raw = std::getenv("INFINI_RANK_WORKER_PROFILE"); + cached = (raw != nullptr && raw[0] == '1' && raw[1] == '\0') ? 1 : 0; + } + return cached == 1; +} + +/// When set with RANK_WORKER_PROFILE, syncDevice after each Graph::run so +/// *_graph_run_ms is exclusive GPU+host (mirrors decode_phase_profile exclusive_sync). +bool rank_worker_profile_sync() { + static int cached = -1; + if (cached < 0) { + const char *raw = std::getenv("INFINI_RANK_WORKER_PROFILE_SYNC"); + cached = (raw != nullptr && raw[0] == '1' && raw[1] == '\0') ? 1 : 0; + } + return cached == 1; +} + +/// Prefill TP=1: Graph::run may skip its end syncStream (INFINI_GRAPH_SKIP_TAIL_SYNC). +/// Callers must drain once before logits / sample D2H (Gate C). +bool graph_skip_tail_sync_enabled() { + static int cached = -1; + if (cached < 0) { + const char *raw = std::getenv("INFINI_GRAPH_SKIP_TAIL_SYNC"); + cached = (raw != nullptr && raw[0] != '\0' && std::string(raw) != "0") ? 1 : 0; + } + return cached == 1; +} + +void append_graph_run_split( + std::string &out, + const char *prefix, + const infinicore::graph::Graph::LastRunProfile &p) { + char buf[256]; + std::snprintf( + buf, + sizeof(buf), + " %s_launch_ms=%.3f %s_tail_sync_ms=%.3f %s_hostop_ms=%.3f %s_op_list_ms=%.3f", + prefix, + p.launch_ms, + prefix, + p.tail_sync_ms, + prefix, + p.hostop_ms, + prefix, + p.op_list_ms); + out.append(buf); +} + +/// Trim redundant post-allreduce barrier (next layer's pre-graph barrier subsumes it). +/// Default on; opt out for bisect: INFINI_PIECEWISE_KEEP_BARRIERS=1 +bool piecewise_trim_barriers() { + static int cached = -1; + if (cached < 0) { + const char *raw = std::getenv("INFINI_PIECEWISE_KEEP_BARRIERS"); + cached = (raw != nullptr && raw[0] == '1' && raw[1] == '\0') ? 0 : 1; + } + return cached == 1; +} + +bool repro_skip_midchunk_eager() { + return infinilm::global_state::repro_skip_midchunk_eager(); +} + +bool repro_allow_midchunk_cg() { + return infinilm::global_state::repro_allow_midchunk_cg(); +} + +bool repro_mixed_mid_cg() { + return infinilm::global_state::repro_mixed_mid_cg(); +} + +bool enable_mixed_mid_cg() { + return infinilm::global_state::enable_mixed_mid_cg(); +} + +bool keep_mixed_mid_capture() { + return infinilm::global_state::keep_mixed_mid_capture(); +} + +/// Skip homo mid-chunk dual capture (saves ~2x device-graph execs). Default off; +/// QY KS38 sets INFINI_NATIVE_CG_SKIP_MID_CAPTURE=1 from run-wrap. +bool skip_homo_mid_capture() { + const char *raw = std::getenv("INFINI_NATIVE_CG_SKIP_MID_CAPTURE"); + return raw != nullptr && raw[0] == '1' && raw[1] == '\0'; +} + +bool scoped_inductor_pre_attn() { + return infinilm::global_state::scoped_inductor_pre_attn_enabled(); +} + +bool repro_skip_final_inductor() { + return infinilm::global_state::repro_skip_final_inductor(); +} + +bool pre_attn_qkv_only() { + return infinilm::global_state::pre_attn_qkv_only(); +} + +bool piecewise_pad_up_cg_enabled() { + return infinilm::global_state::piecewise_pad_up_cg_enabled(); +} + +bool bucket_inductor_capture_enabled(size_t bucket) { + return infinilm::global_state::piecewise_inductor_segment_enabled() + && infinilm::global_state::scoped_inductor_pre_attn_enabled() + && infinilm::global_state::bucket_is_inductor_eligible(bucket); +} + +bool layer_capture_inductor_pre_attn(size_t layer, size_t bucket) { + return bucket_inductor_capture_enabled(bucket) + && infinicore::op::inductor_segment_impl::has_package( + infinicore::op::PiecewiseInductorSegmentId::PreAttn, layer, bucket); +} + +void verify_inductor_packages_(const std::shared_ptr &model, + const std::vector &capture_buckets) { + if (!infinilm::global_state::piecewise_inductor_segment_enabled()) { + return; + } + if (!infinilm::global_state::scoped_inductor_pre_attn_enabled()) { + return; + } + const size_t num_layers = model->native_piecewise_num_layers(); + const int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank(); + for (size_t bucket : capture_buckets) { + if (!infinilm::global_state::bucket_is_inductor_eligible(bucket)) { + continue; + } + for (size_t layer = 0; layer < num_layers; ++layer) { + if (!infinicore::op::inductor_segment_impl::has_package( + infinicore::op::PiecewiseInductorSegmentId::PreAttn, layer, bucket)) { + throw std::runtime_error( + "piecewise inductor: missing AOT pre_attn package layer=" + + std::to_string(layer) + " bucket=" + std::to_string(bucket) + + " tp_rank=" + std::to_string(tp_rank) + + " (compile with aot_compile_piecewise_segments.py --buckets " + + std::to_string(bucket) + " --layers all)"); + } + } + } +} + +double monotonic_ms() { + using clock = std::chrono::steady_clock; + return std::chrono::duration(clock::now().time_since_epoch()).count(); +} + +size_t device_free_bytes_() { + // Post-refactor InfiniRT does not expose the legacy infinirtGetMemInfo + // path used for MetaX capture audit; return 0 (logging only). + (void)infinicore::context::getDevice(); + return 0; +} + +size_t tensor_nbytes_(const infinicore::Tensor &t) { + return t ? t->nbytes() : 0; +} + +size_t graph_arena_nbytes_(const std::shared_ptr &g) { + return g ? g->capture_arena_bytes() : 0; +} + +size_t bucket_arena_nbytes_(const PiecewisePrefillCompiler::BucketGraphs &g) { + size_t n = graph_arena_nbytes_(g.lm_head); + for (const auto &seg : g.pre_attn) { + n += graph_arena_nbytes_(seg); + } + for (const auto &seg : g.post_attn) { + n += graph_arena_nbytes_(seg); + } + return n; +} + +double bytes_to_gib_(size_t n) { + return static_cast(n) / (1024.0 * 1024.0 * 1024.0); +} + +size_t compute_prefill_len(const InfinilmModel::Input &input) { + if (input.input_offsets.has_value()) { + const auto &offsets = input.input_offsets.value(); + const size_t n = offsets->size(0); + if (n >= 2) { + auto cpu_offsets = offsets->to(infinicore::Device{infinicore::Device::Type::kCpu}); + const auto *data = reinterpret_cast(cpu_offsets->data()); + return static_cast(data[n - 1] - data[0]); + } + } + if (input.input_ids.has_value()) { + return input.input_ids.value()->size(1); + } + return 0; +} + +void set_attn_metadata(const InfinilmModel::Input &input) { + attn_metadata_utils::set_attn_metadata(input); +} + +void set_attn_metadata_for_varlen_batch(const InfinilmModel::Input &compiled, + const InfinilmModel::Input &runtime) { + attn_metadata_utils::set_attn_metadata_for_varlen_batch(compiled, runtime); +} + +void zero_tensor_tail_seq_(infinicore::Tensor &tensor, size_t valid_len, size_t bucket) { + if (valid_len >= bucket) { + return; + } + auto tail = tensor->narrow({{1, valid_len, bucket - valid_len}}); + set_zeros(tail); +} + +void clear_stale_bucket_tails_(global_state::PiecewisePrefillState &piecewise, + infinicore::Tensor &logits_holder, + size_t valid_seq_len, + size_t bucket) { + if (valid_seq_len >= bucket) { + return; + } + zero_tensor_tail_seq_(piecewise.hidden_states, valid_seq_len, bucket); + zero_tensor_tail_seq_(piecewise.residual, valid_seq_len, bucket); + for (auto &staging : piecewise.layer_staging) { + zero_tensor_tail_seq_(staging.q_rope, valid_seq_len, bucket); + zero_tensor_tail_seq_(staging.k_rope, valid_seq_len, bucket); + zero_tensor_tail_seq_(staging.v_rope, valid_seq_len, bucket); + zero_tensor_tail_seq_(staging.attn_output, valid_seq_len, bucket); + } + if (piecewise.ar_staging) { + zero_tensor_tail_seq_(piecewise.ar_staging, valid_seq_len, bucket); + } + if (logits_holder) { + zero_tensor_tail_seq_(logits_holder, valid_seq_len, bucket); + } +} + +} // namespace + +PiecewisePrefillCompiler::PiecewisePrefillCompiler(const std::shared_ptr &model, + RankBarrier *barrier) + : model_(model), barrier_(barrier) { + enabled_ = native_piecewise_prefill_enabled() && model_->supports_native_piecewise_prefill(); + if (!enabled_) { + return; + } + max_seq_len_ = compile_max_seq_from_env(); + prefill_chunk_size_ = prefill_chunk_size_from_env(); + const size_t chunk_cap = prefill_chunk_size_; + const bool vllm_ladder = vllm_capture_ladder_enabled(); + if (const char *raw = std::getenv("INFINI_NATIVE_CG_CAPTURE_BUCKETS")) { + capture_buckets_.clear(); + std::string spec(raw); + size_t start = 0; + while (start < spec.size()) { + const size_t comma = spec.find(',', start); + const std::string token = spec.substr(start, comma == std::string::npos ? std::string::npos : comma - start); + if (!token.empty()) { + capture_buckets_.push_back(static_cast(std::stoul(token))); + } + if (comma == std::string::npos) { + break; + } + start = comma + 1; + } + // Runtime pad ladder: keep full vLLM ladder when enabled; capture list may be a subset + // (e.g. Qwen3-32B captures B8192 only to avoid multi-bucket CG OOM under TP=4). + if (vllm_ladder) { + bs_to_padded_ = build_bs_to_padded_bucket( + piecewise_compile_buckets_vllm(max_seq_len_, chunk_cap)); + } else { + bs_to_padded_ = build_bs_to_padded_bucket(capture_buckets_); + } + } else if (vllm_ladder) { + capture_buckets_ = piecewise_capture_buckets_vllm(max_seq_len_, chunk_cap); + bs_to_padded_ = build_bs_to_padded_bucket(capture_buckets_); + } else { + auto pad_ladder = piecewise_compile_buckets(max_seq_len_); + bs_to_padded_ = build_bs_to_padded_bucket(pad_ladder); + capture_buckets_ = piecewise_capture_buckets(max_seq_len_); + } + std::sort(capture_buckets_.begin(), capture_buckets_.end(), std::greater()); +} + +void PiecewisePrefillCompiler::allocate_shared_banks_(size_t max_bucket, + size_t num_layers, + size_t n_req) { + shared_banks_ = SharedPrefillBanks{}; + shared_banks_.max_bucket = max_bucket; + shared_banks_.num_layers = num_layers; + const auto device = infinicore::context::getDevice(); + const auto &model_config = infinilm::global_state::get_infinilm_config().model_config; + const auto dtype = model_config->get_dtype(); + const size_t hidden = model_config->get("hidden_size"); + const size_t vocab_size = model_config->get("vocab_size"); + const size_t tp_size = std::max( + 1, static_cast(infinilm::global_state::get_tensor_model_parallel_world_size())); + const size_t num_heads = model_config->get("num_attention_heads") / static_cast(tp_size); + const size_t total_kv = model_config->get("num_key_value_heads"); + const size_t num_kv_heads = total_kv < tp_size ? 1 : total_kv / tp_size; + const size_t head_dim = model_config->get_head_dim(); + + shared_banks_.layer_staging.resize(num_layers); + for (size_t i = 0; i < num_layers; ++i) { + auto &st = shared_banks_.layer_staging[i]; + st.q_rope = infinicore::Tensor::empty({1, max_bucket, num_heads, head_dim}, dtype, device); + st.k_rope = infinicore::Tensor::empty({1, max_bucket, num_kv_heads, head_dim}, dtype, device); + st.v_rope = infinicore::Tensor::empty({1, max_bucket, num_kv_heads, head_dim}, dtype, device); + st.attn_output = infinicore::Tensor::empty({1, max_bucket, num_heads * head_dim}, dtype, device); + } + shared_banks_.hidden_states = infinicore::Tensor::empty({1, max_bucket, hidden}, dtype, device); + shared_banks_.residual = infinicore::Tensor::empty({1, max_bucket, hidden}, dtype, device); + shared_banks_.ar_staging = infinicore::Tensor::empty({1, max_bucket, hidden}, dtype, device); + + shared_banks_.input_ids = + infinicore::Tensor::empty({1, max_bucket}, infinicore::DataType::kInt64, device); + shared_banks_.position_ids = + infinicore::Tensor::empty({max_bucket}, infinicore::DataType::kInt64, device); + shared_banks_.slot_mapping = + infinicore::Tensor::empty({max_bucket}, infinicore::DataType::kInt64, device); + shared_banks_.past_sequence_lengths = + infinicore::Tensor::empty({n_req}, infinicore::DataType::kInt32, device); + shared_banks_.total_sequence_lengths = + infinicore::Tensor::empty({n_req}, infinicore::DataType::kInt32, device); + shared_banks_.input_offsets = + infinicore::Tensor::empty({n_req + 1}, infinicore::DataType::kInt32, device); + shared_banks_.cu_seqlens = + infinicore::Tensor::empty({n_req + 1}, infinicore::DataType::kInt32, device); + shared_banks_.logits_holder = + infinicore::Tensor::empty({1, max_bucket, vocab_size}, dtype, device); + + size_t staging = tensor_nbytes_(shared_banks_.hidden_states) + + tensor_nbytes_(shared_banks_.residual) + + tensor_nbytes_(shared_banks_.ar_staging); + for (const auto &st : shared_banks_.layer_staging) { + staging += tensor_nbytes_(st.q_rope) + tensor_nbytes_(st.k_rope) + + tensor_nbytes_(st.v_rope) + tensor_nbytes_(st.attn_output); + } + const size_t io = tensor_nbytes_(shared_banks_.input_ids) + + tensor_nbytes_(shared_banks_.position_ids) + + tensor_nbytes_(shared_banks_.slot_mapping) + + tensor_nbytes_(shared_banks_.past_sequence_lengths) + + tensor_nbytes_(shared_banks_.total_sequence_lengths) + + tensor_nbytes_(shared_banks_.input_offsets) + + tensor_nbytes_(shared_banks_.cu_seqlens) + + tensor_nbytes_(shared_banks_.logits_holder); + shared_banks_.staging_physical_bytes = staging; + shared_banks_.io_physical_bytes = io; + spdlog::info( + "native piecewise CG: shared max-bucket banks max_bucket={} layers={} " + "staging_GiB={:.3f} io_GiB={:.3f}", + max_bucket, + num_layers, + bytes_to_gib_(staging), + bytes_to_gib_(io)); +} + +void PiecewisePrefillCompiler::bind_bucket_staging_(size_t bucket, size_t num_layers) { + if (!shared_banks_.hidden_states || bucket > shared_banks_.max_bucket) { + throw std::runtime_error( + "bind_bucket_staging_: shared banks missing or bucket exceeds max"); + } + if (num_layers > shared_banks_.num_layers) { + throw std::runtime_error("bind_bucket_staging_: num_layers exceeds shared bank"); + } + auto &piecewise = infinilm::global_state::get_forward_context().piecewise; + piecewise.bucket_seq_len = bucket; + piecewise.layer_staging.clear(); + piecewise.layer_staging.resize(num_layers); + for (size_t i = 0; i < num_layers; ++i) { + const auto &src = shared_banks_.layer_staging[i]; + auto &st = piecewise.layer_staging[i]; + st.q_rope = src.q_rope->narrow({{1, 0, bucket}}); + st.k_rope = src.k_rope->narrow({{1, 0, bucket}}); + st.v_rope = src.v_rope->narrow({{1, 0, bucket}}); + st.attn_output = src.attn_output->narrow({{1, 0, bucket}}); + } + piecewise.hidden_states = shared_banks_.hidden_states->narrow({{1, 0, bucket}}); + piecewise.residual = shared_banks_.residual->narrow({{1, 0, bucket}}); + piecewise.ar_staging = shared_banks_.ar_staging->narrow({{1, 0, bucket}}); +} + +InfinilmModel::Input PiecewisePrefillCompiler::make_bucket_input_(size_t bucket, + size_t nblocks, + size_t n_req, + bool mid_chunk_capture) const { + if (!shared_banks_.input_ids || bucket > shared_banks_.max_bucket) { + throw std::runtime_error("make_bucket_input_: shared I/O banks not allocated"); + } + InfinilmModel::Input input; + input.input_ids = shared_banks_.input_ids->narrow({{1, 0, bucket}}); + input.position_ids = shared_banks_.position_ids->narrow({{0, 0, bucket}}); + input.past_sequence_lengths = shared_banks_.past_sequence_lengths; + input.total_sequence_lengths = shared_banks_.total_sequence_lengths; + input.input_offsets = shared_banks_.input_offsets; + input.cu_seqlens = shared_banks_.cu_seqlens; + input.slot_mapping = shared_banks_.slot_mapping->narrow({{0, 0, bucket}}); + set_zeros(input.input_ids.value()); + set_zeros(input.past_sequence_lengths.value()); + set_zeros(input.total_sequence_lengths.value()); + + const size_t chunk_size = + prefill_chunk_size_ > 0 ? prefill_chunk_size_ : prefill_chunk_size_from_env(); + // Mid dual-capture: representative continuing chunk after one full prior chunk. + const size_t past = + mid_chunk_capture ? (bucket == chunk_size ? bucket : chunk_size) : 0; + const int64_t pos_start = mid_chunk_capture + ? static_cast(past) + : ((bucket < chunk_size) ? static_cast(chunk_size) : int64_t{0}); + std::vector position_ids_vec(bucket); + std::iota(position_ids_vec.begin(), position_ids_vec.end(), pos_start); + infinicore::context::memcpyH2D( + input.position_ids.value()->data(), position_ids_vec.data(), bucket * sizeof(int64_t), false); + + const int32_t per_req = static_cast(bucket / std::max(1, n_req)); + std::vector past_lengths_vec(n_req, static_cast(past)); + std::vector total_lengths_vec( + n_req, static_cast(past) + per_req); + infinicore::context::memcpyH2D( + input.past_sequence_lengths.value()->data(), + past_lengths_vec.data(), + n_req * sizeof(int32_t), + false); + infinicore::context::memcpyH2D( + input.total_sequence_lengths.value()->data(), + total_lengths_vec.data(), + n_req * sizeof(int32_t), + false); + + std::vector input_offsets_vec(n_req + 1, 0); + for (size_t i = 0; i <= n_req; ++i) { + input_offsets_vec[i] = static_cast(std::min(bucket, i * per_req)); + } + input_offsets_vec[n_req] = static_cast(bucket); + infinicore::context::memcpyH2D( + input.input_offsets.value()->data(), + input_offsets_vec.data(), + (n_req + 1) * sizeof(int32_t), + false); + + std::vector cu_seqlens_vec(n_req + 1, 0); + for (size_t i = 0; i <= n_req; ++i) { + cu_seqlens_vec[i] = static_cast(std::min(bucket, i * per_req)); + } + cu_seqlens_vec[n_req] = static_cast(bucket); + infinicore::context::memcpyH2D( + input.cu_seqlens.value()->data(), + cu_seqlens_vec.data(), + (n_req + 1) * sizeof(int32_t), + false); + + const size_t block_per_req = nblocks; + input.block_tables = block_tables_holder_->as_strided({n_req, block_per_req}, {(ptrdiff_t)block_per_req, 1}); + const auto *paged_config = dynamic_cast(model_->get_cache_config()); + const size_t block_size = paged_config != nullptr ? paged_config->block_size() : 256; + const size_t span_tokens = past + bucket; + const size_t blocks_needed = (span_tokens + block_size - 1) / block_size; + for (size_t row = 0; row < n_req; ++row) { + std::vector block_row(block_per_req, -1); + const size_t row_offset = row * blocks_needed; + for (size_t b = 0; b < blocks_needed && b < block_per_req; ++b) { + block_row[b] = static_cast(row_offset + b); + } + auto row_tensor = input.block_tables.value()->narrow({{0, row, 1}}); + infinicore::context::memcpyH2D( + row_tensor->data(), block_row.data(), block_per_req * sizeof(int32_t), false); + } + + std::vector slot_mapping_vec(bucket); + std::iota(slot_mapping_vec.begin(), + slot_mapping_vec.end(), + static_cast(past)); + infinicore::context::memcpyH2D( + input.slot_mapping.value()->data(), slot_mapping_vec.data(), bucket * sizeof(int64_t), false); + // Mid capture is a continuing chunk (not final); leave flag empty → treated as final for + // lm_head dry-run only when we explicitly run lm_head below. Runtime mid uses is_final=false. + return input; +} + +InfinilmModel::Input PiecewisePrefillCompiler::make_mixed_mid_bucket_input_(size_t bucket, + size_t nblocks) const { + if (!shared_banks_.input_ids || bucket > shared_banks_.max_bucket) { + throw std::runtime_error("make_mixed_mid_bucket_input_: shared I/O banks not allocated"); + } + if (bucket < 2) { + throw std::runtime_error("make_mixed_mid_bucket_input_: bucket must be >= 2"); + } + const size_t n_req = mixed_mid_capture_req_; + // Match v1 scheduler RUNNING order: decode rows often precede mid-prefills. + // Dominant LongBench MIXED@2048: 1 decode token + continuing mid fills bucket-1. + const size_t decode_q = 1; + const size_t mid_q = bucket - 1; + const size_t decode_past = bucket; + const size_t mid_past = bucket; + const size_t decode_total = decode_past + decode_q; + const size_t mid_total = mid_past + mid_q; + + InfinilmModel::Input input; + input.input_ids = shared_banks_.input_ids->narrow({{1, 0, bucket}}); + input.position_ids = shared_banks_.position_ids->narrow({{0, 0, bucket}}); + // Narrow metadata to capture n_req so FA max_seqlens do not see zero-pad tails. + input.past_sequence_lengths = shared_banks_.past_sequence_lengths->narrow({{0, 0, n_req}}); + input.total_sequence_lengths = shared_banks_.total_sequence_lengths->narrow({{0, 0, n_req}}); + input.input_offsets = shared_banks_.input_offsets->narrow({{0, 0, n_req + 1}}); + input.cu_seqlens = shared_banks_.cu_seqlens->narrow({{0, 0, n_req + 1}}); + input.slot_mapping = shared_banks_.slot_mapping->narrow({{0, 0, bucket}}); + set_zeros(input.input_ids.value()); + set_zeros(input.past_sequence_lengths.value()); + set_zeros(input.total_sequence_lengths.value()); + + std::vector position_ids_vec(bucket); + position_ids_vec[0] = static_cast(decode_past); + for (size_t i = 0; i < mid_q; ++i) { + position_ids_vec[decode_q + i] = static_cast(mid_past + i); + } + infinicore::context::memcpyH2D( + input.position_ids.value()->data(), position_ids_vec.data(), bucket * sizeof(int64_t), false); + + const std::vector past_lengths_vec{ + static_cast(decode_past), static_cast(mid_past)}; + const std::vector total_lengths_vec{ + static_cast(decode_total), static_cast(mid_total)}; + infinicore::context::memcpyH2D( + input.past_sequence_lengths.value()->data(), + past_lengths_vec.data(), + n_req * sizeof(int32_t), + false); + infinicore::context::memcpyH2D( + input.total_sequence_lengths.value()->data(), + total_lengths_vec.data(), + n_req * sizeof(int32_t), + false); + + const std::vector input_offsets_vec{ + 0, static_cast(decode_q), static_cast(bucket)}; + // cu_seqlens = cumulative KV lengths (matches processor MIXED packing). + const std::vector cu_seqlens_vec{ + 0, static_cast(decode_total), static_cast(decode_total + mid_total)}; + infinicore::context::memcpyH2D( + input.input_offsets.value()->data(), + input_offsets_vec.data(), + (n_req + 1) * sizeof(int32_t), + false); + infinicore::context::memcpyH2D( + input.cu_seqlens.value()->data(), + cu_seqlens_vec.data(), + (n_req + 1) * sizeof(int32_t), + false); + + const size_t block_per_req = nblocks; + input.block_tables = block_tables_holder_->as_strided({n_req, block_per_req}, {(ptrdiff_t)block_per_req, 1}); + const auto *paged_config = dynamic_cast(model_->get_cache_config()); + const size_t block_size = paged_config != nullptr ? paged_config->block_size() : 256; + const size_t spans[2] = {decode_total, mid_total}; + size_t row_block_cursor = 0; + for (size_t row = 0; row < n_req; ++row) { + std::vector block_row(block_per_req, -1); + const size_t blocks_needed = (spans[row] + block_size - 1) / block_size; + for (size_t b = 0; b < blocks_needed && b < block_per_req; ++b) { + block_row[b] = static_cast(row_block_cursor + b); + } + row_block_cursor += blocks_needed; + auto row_tensor = input.block_tables.value()->narrow({{0, row, 1}}); + infinicore::context::memcpyH2D( + row_tensor->data(), block_row.data(), block_per_req * sizeof(int32_t), false); + } + + std::vector slot_mapping_vec(bucket); + slot_mapping_vec[0] = static_cast(decode_past); + for (size_t i = 0; i < mid_q; ++i) { + slot_mapping_vec[decode_q + i] = static_cast(mid_past + i); + } + infinicore::context::memcpyH2D( + input.slot_mapping.value()->data(), slot_mapping_vec.data(), bucket * sizeof(int64_t), false); + + // Row0 decode (lm_head + sample); row1 continuing mid (no sample). + input.is_final_prefill_chunk = {true, false}; + return input; +} + +void PiecewisePrefillCompiler::capture_bucket_(size_t bucket, bool mid_chunk_capture) { + const auto rank_device = infinilm::global_state::get_tensor_model_parallel_rank_info().device; + infinicore::context::setDevice(rank_device); + // Prefill: FA host-break + AOT MoE under full_and_piecewise. + infinicore::context::InferencePhaseGuard phase_guard( + infinicore::context::InferencePhase::Prefill); + + auto &piecewise_flag = infinilm::global_state::get_forward_context().piecewise; + struct CaptureGuard { + infinilm::global_state::PiecewisePrefillState &pw; + explicit CaptureGuard(infinilm::global_state::PiecewisePrefillState &p) : pw(p) { pw.compile_capture_active = true; } + ~CaptureGuard() { pw.compile_capture_active = false; } + } capture_guard(piecewise_flag); + + const size_t nblocks = dynamic_cast(model_->get_cache_config())->num_blocks(); + const size_t num_layers = model_->native_piecewise_num_layers(); + bind_bucket_staging_(bucket, num_layers); + auto bucket_input = make_bucket_input_(bucket, nblocks, max_capture_req_, mid_chunk_capture); + set_attn_metadata(bucket_input); + + BucketGraphs graphs; + graphs.input = std::move(bucket_input); + + auto &piecewise = infinilm::global_state::get_forward_context().piecewise; + piecewise.valid_seq_len = bucket; + // Dry-run + QKV CG must use the C++ pre path (not inductor). Inductor AOT does not + // populate MiniCPM gate_score_cache_; allocating that buffer inside startGraphRecording + // corrupts the QKV-only graph on MetaX. + piecewise.allow_inductor_pre_attn = false; + piecewise.phase = global_state::PiecewiseCapturePhase::None; + + auto &hidden = piecewise.hidden_states; + auto &residual = piecewise.residual; + + size_t capture_layers = num_layers; + if (const char *raw = std::getenv("INFINI_NATIVE_CG_MAX_LAYERS")) { + capture_layers = std::min(num_layers, static_cast(std::stoul(raw))); + } + + graphs.logits_holder = shared_banks_.logits_holder->narrow({{1, 0, bucket}}); + + // Eager warmup dry-run before capture (NONE-equivalent). + // Dry-run / capture use valid_seq_len==bucket (exact). Runtime mid-chunk + // with seq_len==graph_bucket shares this staging; pad-up CG replays the + // same banks with valid_seq_len=L + zeroed tails. + // Dry-run must allocate MiniCPM gate buffers before CG (outside recording). + const bool qkv_only = pre_attn_qkv_only(); + model_->native_piecewise_embed(graphs.input, hidden); + for (size_t layer = 0; layer < capture_layers; ++layer) { + if (qkv_only) { + model_->native_piecewise_pre_attn_qkv_layer(layer, graphs.input, hidden, residual); + model_->native_piecewise_pre_attn_rope_layer(layer, graphs.input); + } else { + model_->native_piecewise_pre_attn_layer(layer, graphs.input, hidden, residual); + } + if (infinilm::global_state::piecewise_inductor_segment_enabled()) { + infinicore::context::syncDevice(); + barrier_->wait(); + } + model_->native_piecewise_eager_attn_layer(layer, graphs.input); + if (infinilm::global_state::piecewise_inductor_segment_enabled()) { + infinicore::context::syncDevice(); + barrier_->wait(); + } + model_->native_piecewise_post_attn_cg_layer(layer, graphs.input, hidden, residual); + if (infinilm::global_state::piecewise_inductor_segment_enabled()) { + infinicore::context::syncDevice(); + barrier_->wait(); + } + if (infinilm::global_state::piecewise_inductor_segment_enabled()) { + infinicore::context::syncDevice(); + } + } + // Mid chunks do not sample; skip lm_head dry-run/capture for mid key. + if (!mid_chunk_capture) { + model_->native_piecewise_lm_head(graphs.input, hidden, residual, graphs.logits_holder); + } + // Reset KV so capture FA sees the same empty-cache shape as dry-run start. + // Stale KV + many cudaGraphExecs on QY previously LOOR'd mid-ladder FA. + { + 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::syncDevice(); + } + graphs.pre_attn.resize(capture_layers); + graphs.post_attn.resize(capture_layers); + + set_zeros(piecewise.residual); + model_->native_piecewise_embed(graphs.input, hidden); + + for (size_t layer = 0; layer < capture_layers; ++layer) { + piecewise.active_layer = layer; + piecewise.phase = global_state::PiecewiseCapturePhase::PreAttn; + + barrier_->wait(); + // Product: one pre piece [CG: LN+QKV+RoPE]. Kill-switch QKV_ONLY restores + // Phase-2c QKV CG + eager RoPE. Inductor AOT packages bake RoPE — keep + // allow_inductor_pre_attn=false so C++ path enters the shape-keyed CG. + piecewise.allow_inductor_pre_attn = false; + if (qkv_only) { + infinicore::context::startGraphRecording(); + model_->native_piecewise_pre_attn_qkv_layer(layer, graphs.input, hidden, residual); + graphs.pre_attn[layer] = infinicore::context::stopGraphRecording(); + infinicore::context::syncStream(); + model_->native_piecewise_pre_attn_rope_layer(layer, graphs.input); + infinicore::context::syncStream(); + } else { + infinicore::context::startGraphRecording(); + model_->native_piecewise_pre_attn_layer(layer, graphs.input, hidden, residual); + graphs.pre_attn[layer] = infinicore::context::stopGraphRecording(); + infinicore::context::syncStream(); + } + + piecewise.phase = global_state::PiecewiseCapturePhase::EagerAttn; + model_->native_piecewise_eager_attn_layer(layer, graphs.input); + + piecewise.phase = global_state::PiecewiseCapturePhase::PostAttn; + barrier_->wait(); + infinicore::context::startGraphRecording(); + model_->native_piecewise_post_attn_cg_layer(layer, graphs.input, hidden, residual); + graphs.post_attn[layer] = infinicore::context::stopGraphRecording(); + assert_post_one_device_seg_("piecewise_prefill_post", layer, graphs.post_attn[layer]); + // Drain before next layer: QY GraphInstantiate failures otherwise sticky-LOOR FA. + infinicore::context::syncDevice(); + barrier_->wait(); + } + + if (!mid_chunk_capture) { + piecewise.phase = global_state::PiecewiseCapturePhase::LmHead; + barrier_->wait(); + infinicore::context::startGraphRecording(); + model_->native_piecewise_lm_head(graphs.input, hidden, residual, graphs.logits_holder); + graphs.lm_head = infinicore::context::stopGraphRecording(); + barrier_->wait(); + } + + piecewise.phase = global_state::PiecewiseCapturePhase::None; + graphs.hidden_states = piecewise.hidden_states; + graphs.residual = piecewise.residual; + graphs.ar_staging = piecewise.ar_staging; + graphs.layer_staging = piecewise.layer_staging; + if (mid_chunk_capture) { + compiled_mid_[bucket] = std::move(graphs); + } else { + compiled_[bucket] = std::move(graphs); + } + const size_t captured_segments = + capture_layers * 2 + (mid_chunk_capture ? 0 : 1); + spdlog::info( + "native piecewise CG: captured bucket={} mid={} layers={} segments={} pre={}", + bucket, + mid_chunk_capture, + capture_layers, + captured_segments, + qkv_only ? "qkv_only" : "qkv_rope"); +} + +void PiecewisePrefillCompiler::capture_mixed_mid_bucket_(size_t bucket) { + const auto rank_device = infinilm::global_state::get_tensor_model_parallel_rank_info().device; + infinicore::context::setDevice(rank_device); + infinicore::context::InferencePhaseGuard phase_guard( + infinicore::context::InferencePhase::Prefill); + + auto &piecewise_flag = infinilm::global_state::get_forward_context().piecewise; + struct CaptureGuard { + infinilm::global_state::PiecewisePrefillState &pw; + explicit CaptureGuard(infinilm::global_state::PiecewisePrefillState &p) : pw(p) { + pw.compile_capture_active = true; + } + ~CaptureGuard() { pw.compile_capture_active = false; } + } capture_guard(piecewise_flag); + + const size_t nblocks = + dynamic_cast(model_->get_cache_config())->num_blocks(); + const size_t num_layers = model_->native_piecewise_num_layers(); + bind_bucket_staging_(bucket, num_layers); + auto bucket_input = make_mixed_mid_bucket_input_(bucket, nblocks); + set_attn_metadata(bucket_input); + + BucketGraphs graphs; + graphs.input = std::move(bucket_input); + + auto &piecewise = infinilm::global_state::get_forward_context().piecewise; + piecewise.valid_seq_len = bucket; + // MIXED mid dual-capture: never bake inductor-in-CG (same as homo mid). + piecewise.allow_inductor_pre_attn = false; + piecewise.phase = global_state::PiecewiseCapturePhase::None; + + auto &hidden = piecewise.hidden_states; + auto &residual = piecewise.residual; + + size_t capture_layers = num_layers; + if (const char *raw = std::getenv("INFINI_NATIVE_CG_MAX_LAYERS")) { + capture_layers = std::min(num_layers, static_cast(std::stoul(raw))); + } + + graphs.logits_holder = shared_banks_.logits_holder->narrow({{1, 0, bucket}}); + + // Eager warmup dry-run (includes lm_head — decode row needs logits). + model_->native_piecewise_embed(graphs.input, hidden); + for (size_t layer = 0; layer < capture_layers; ++layer) { + model_->native_piecewise_pre_attn_layer(layer, graphs.input, hidden, residual); + if (infinilm::global_state::piecewise_inductor_segment_enabled()) { + infinicore::context::syncDevice(); + barrier_->wait(); + } + model_->native_piecewise_eager_attn_layer(layer, graphs.input); + if (infinilm::global_state::piecewise_inductor_segment_enabled()) { + infinicore::context::syncDevice(); + barrier_->wait(); + } + model_->native_piecewise_post_attn_cg_layer(layer, graphs.input, hidden, residual); + if (infinilm::global_state::piecewise_inductor_segment_enabled()) { + infinicore::context::syncDevice(); + barrier_->wait(); + } + if (infinilm::global_state::piecewise_inductor_segment_enabled()) { + infinicore::context::syncDevice(); + } + } + model_->native_piecewise_lm_head(graphs.input, hidden, residual, graphs.logits_holder); + // Shape-keyed QKV comes from compiled_[bucket]; do not store MIXED pre_attn. + graphs.pre_attn.resize(capture_layers); + graphs.post_attn.resize(capture_layers); + + set_zeros(piecewise.residual); + model_->native_piecewise_embed(graphs.input, hidden); + + for (size_t layer = 0; layer < capture_layers; ++layer) { + piecewise.active_layer = layer; + piecewise.phase = global_state::PiecewiseCapturePhase::PreAttn; + + barrier_->wait(); + piecewise.allow_inductor_pre_attn = false; + // Eager QKV+RoPE for capture continuity only — prefer path uses final-bucket QKV CG. + model_->native_piecewise_pre_attn_layer(layer, graphs.input, hidden, residual); + infinicore::context::syncStream(); + + piecewise.phase = global_state::PiecewiseCapturePhase::EagerAttn; + model_->native_piecewise_eager_attn_layer(layer, graphs.input); + + piecewise.phase = global_state::PiecewiseCapturePhase::PostAttn; + barrier_->wait(); + infinicore::context::startGraphRecording(); + model_->native_piecewise_post_attn_cg_layer(layer, graphs.input, hidden, residual); + graphs.post_attn[layer] = infinicore::context::stopGraphRecording(); + assert_post_one_device_seg_("piecewise_prefill_post", layer, graphs.post_attn[layer]); + barrier_->wait(); + } + + piecewise.phase = global_state::PiecewiseCapturePhase::LmHead; + barrier_->wait(); + infinicore::context::startGraphRecording(); + model_->native_piecewise_lm_head(graphs.input, hidden, residual, graphs.logits_holder); + graphs.lm_head = infinicore::context::stopGraphRecording(); + barrier_->wait(); + + piecewise.phase = global_state::PiecewiseCapturePhase::None; + graphs.hidden_states = piecewise.hidden_states; + graphs.residual = piecewise.residual; + graphs.ar_staging = piecewise.ar_staging; + graphs.layer_staging = piecewise.layer_staging; + compiled_mixed_mid_[bucket] = std::move(graphs); + // post per layer + lm_head (no MIXED pre CG). + const size_t captured_segments = capture_layers + 1; + spdlog::info( + "native piecewise CG: captured bucket={} mixed_mid=1 n_req={} layers={} segments={} " + "pre=skip_use_final_qkv", + bucket, + mixed_mid_capture_req_, + capture_layers, + captured_segments); +} + +void PiecewisePrefillCompiler::warmup_inductor_segments_(size_t nblocks, size_t n_req) { + if (!infinilm::global_state::piecewise_inductor_segment_enabled()) { + return; + } + const size_t num_layers = model_->native_piecewise_num_layers(); + if (num_layers == 0 || capture_buckets_.empty()) { + return; + } + size_t warmed_buckets = 0; + for (size_t bucket : capture_buckets_) { + if (!infinilm::global_state::bucket_is_inductor_eligible(bucket)) { + continue; + } + bind_bucket_staging_(bucket, num_layers); + auto bucket_input = make_bucket_input_(bucket, nblocks, n_req); + const auto &positions = bucket_input.position_ids.value(); + auto positions_padded = infinicore::Tensor::zeros( + {1, bucket}, infinicore::DataType::kInt64, infinicore::context::getDevice()); + auto &piecewise = infinilm::global_state::get_forward_context().piecewise; + for (size_t layer = 0; layer < num_layers; ++layer) { + if (!infinicore::op::inductor_segment_impl::has_package( + infinicore::op::PiecewiseInductorSegmentId::PreAttn, layer, bucket)) { + continue; + } + /* inductor warmup stubbed */; + infinicore::context::syncDevice(); + } + ++warmed_buckets; + } + if (warmed_buckets > 0) { + auto &piecewise = infinilm::global_state::get_forward_context().piecewise; + piecewise.layer_staging.clear(); + piecewise.hidden_states = infinicore::Tensor(); + piecewise.residual = infinicore::Tensor(); + piecewise.ar_staging = infinicore::Tensor(); + infinicore::context::setDevice( + infinilm::global_state::get_tensor_model_parallel_rank_info().device); + infinicore::context::syncDevice(); + spdlog::info( + "piecewise inductor: eager AOT warmup buckets={} (per-rank, before CG capture)", + warmed_buckets); + } +} + +void PiecewisePrefillCompiler::compile() { + if (!native_piecewise_prefill_enabled() || !model_->supports_native_piecewise_prefill()) { + enabled_ = false; + return; + } + if (model_->get_cache_config() == nullptr + || dynamic_cast(model_->get_cache_config()) == nullptr) { + spdlog::debug("PiecewisePrefillCompiler: defer capture until paged cache is configured"); + return; + } + enabled_ = true; + + const auto *paged_config = + dynamic_cast(model_->get_cache_config()); + const size_t nblocks = paged_config->num_blocks(); + // Capture width ≥ MAX_BATCH_SIZE so mixed/multi-req PIECEWISE can replay with + // runtime_n_req ≤ compiled_n_req (vLLM mixed_mode=PIECEWISE). Opt-in override: + // INFINI_MAX_PREFILL_BATCH. Default 8 (Band C). Set to 1 only for Gate-D bisect. + max_capture_req_ = 8; + if (const char *raw = std::getenv("INFINI_MAX_PREFILL_BATCH")) { + max_capture_req_ = std::max(1, std::stoul(raw)); + } else if (const char *raw = std::getenv("MAX_BATCH_SIZE")) { + max_capture_req_ = std::max(1, std::stoul(raw)); + } + size_t max_bucket = capture_buckets_.empty() ? 0 : capture_buckets_.front(); + block_tables_holder_ = infinicore::Tensor::empty( + {max_capture_req_ * nblocks}, infinicore::DataType::kInt32, infinicore::context::getDevice()); + set_zeros(block_tables_holder_); + const size_t num_layers = model_->native_piecewise_num_layers(); + const size_t free_before = device_free_bytes_(); + // M2: allocate shared banks before inductor warmup so warmup also uses prefix views. + allocate_shared_banks_(max_bucket, num_layers, max_capture_req_); + spdlog::info( + "native piecewise CG: capture warmup n_req={} (metadata only, hidden [1,bucket] shared max={})", + max_capture_req_, + max_bucket); + + const int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank(); + const int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size(); + if (tp_size > 1) { + barrier_->wait(); + } + verify_inductor_packages_(model_, capture_buckets_); + warmup_inductor_segments_(nblocks, max_capture_req_); + if (tp_size > 1) { + barrier_->wait(); + } + + compiled_.clear(); + compiled_mid_.clear(); + compiled_mixed_mid_.clear(); + mixed_mid_capture_req_ = std::min(2, max_capture_req_); + for (size_t bucket : capture_buckets_) { + capture_bucket_(bucket, /*mid_chunk_capture=*/false); + infinicore::context::syncDevice(); + // Dual capture: exact chunk-sized bucket also gets homo mid keys. + // MIXED dual-capture (compiled_mixed_mid_) is opt-in via KEEP_MIXED_MID_CAPTURE; + // product MIXED prefer replays final compiled_ banks (vLLM-like). + // QY: SKIP_MID_CAPTURE avoids a second full layer ladder of cudaGraphExec. + if (!skip_homo_mid_capture() && bucket == prefill_chunk_size_ && prefill_chunk_size_ > 0) { + capture_bucket_(bucket, /*mid_chunk_capture=*/true); + infinicore::context::syncDevice(); + if (keep_mixed_mid_capture() && mixed_mid_capture_req_ >= 2 && bucket >= 2) { + capture_mixed_mid_bucket_(bucket); + infinicore::context::syncDevice(); + } + } + } + infinicore::context::syncDevice(); + const size_t free_after = device_free_bytes_(); + // Physical staging/I-O = shared max bank (O(max)); arenas = MoE scratch via + // shared Python workspace after M3 (still summed per Graph CaptureArena). + const size_t staging_bytes = shared_banks_.staging_physical_bytes; + const size_t io_bytes = + shared_banks_.io_physical_bytes + tensor_nbytes_(block_tables_holder_); + size_t arena_bytes = 0; + for (const auto &kv : compiled_) { + arena_bytes += bucket_arena_nbytes_(kv.second); + } + for (const auto &kv : compiled_mid_) { + arena_bytes += bucket_arena_nbytes_(kv.second); + } + for (const auto &kv : compiled_mixed_mid_) { + arena_bytes += bucket_arena_nbytes_(kv.second); + } + const size_t staging_io_bytes = staging_bytes + io_bytes; + const size_t free_delta = + free_before >= free_after ? free_before - free_after : 0; + std::ostringstream oss; + for (size_t i = 0; i < capture_buckets_.size(); ++i) { + if (i > 0) { + oss << ','; + } + oss << capture_buckets_[i]; + } + spdlog::info( + "native piecewise CG: capture_buckets=[{}] max_seq={} chunk_size={} mid_keys={} " + "mixed_mid_keys={}", + oss.str(), + max_seq_len_, + prefill_chunk_size_, + compiled_mid_.size(), + compiled_mixed_mid_.size()); + // M1: greppable one-line CG VRAM split (before /health). Staging is physical O(max) after M2. + spdlog::info( + "cg_mem_budget tag=piecewise_prefill buckets={} mid_buckets={} mixed_mid_buckets={} " + "staging_bytes={} io_bytes={} " + "staging_io_bytes={} arena_bytes={} free_delta_bytes={} " + "staging_GiB={:.3f} io_GiB={:.3f} staging_io_GiB={:.3f} arena_GiB={:.3f} " + "free_delta_GiB={:.3f}", + compiled_.size(), + compiled_mid_.size(), + compiled_mixed_mid_.size(), + staging_bytes, + io_bytes, + staging_io_bytes, + arena_bytes, + free_delta, + bytes_to_gib_(staging_bytes), + bytes_to_gib_(io_bytes), + bytes_to_gib_(staging_io_bytes), + bytes_to_gib_(arena_bytes), + bytes_to_gib_(free_delta)); +} + +size_t PiecewisePrefillCompiler::padded_bucket_for(size_t seq_len) const { + return padded_bucket_for_seq_len(seq_len, bs_to_padded_, max_seq_len_); +} + +void PiecewisePrefillCompiler::copy_runtime_into_bucket_(BucketGraphs &bucket_graphs, + const InfinilmModel::Input &runtime, + size_t valid_seq_len) const { + const size_t bucket = bucket_graphs.input.input_ids.value()->size(1); + auto &graph_input = bucket_graphs.input; + + graph_input.input_ids.value() + ->narrow({{1, 0, valid_seq_len}}) + ->copy_from(runtime.input_ids.value()); + graph_input.position_ids.value() + ->narrow({{0, 0, valid_seq_len}}) + ->copy_from(runtime.position_ids.value()); + + + const size_t runtime_n_req = runtime.block_tables.value()->size(0); + const size_t compiled_n_req = graph_input.block_tables.value()->size(0); + if (runtime_n_req > compiled_n_req) { + throw std::runtime_error("block_tables batch exceeds compiled capture warmup width"); + } + + if (graph_input.past_sequence_lengths.has_value() && runtime.past_sequence_lengths.has_value()) { + graph_input.past_sequence_lengths.value() + ->narrow({{0, 0, runtime_n_req}}) + ->copy_from(runtime.past_sequence_lengths.value()); + if (runtime_n_req < compiled_n_req) { + auto past_tail = graph_input.past_sequence_lengths.value()->narrow( + {{0, runtime_n_req, compiled_n_req - runtime_n_req}}); + set_zeros(past_tail); + } + } + graph_input.total_sequence_lengths.value() + ->narrow({{0, 0, runtime_n_req}}) + ->copy_from(runtime.total_sequence_lengths.value()); + if (runtime_n_req < compiled_n_req) { + auto total_tail = graph_input.total_sequence_lengths.value()->narrow( + {{0, runtime_n_req, compiled_n_req - runtime_n_req}}); + set_zeros(total_tail); + } + graph_input.input_offsets.value() + ->narrow({{0, 0, runtime_n_req + 1}}) + ->copy_from(runtime.input_offsets.value()); + graph_input.cu_seqlens.value() + ->narrow({{0, 0, runtime_n_req + 1}}) + ->copy_from(runtime.cu_seqlens.value()); + // Pad unused capture-width rows with the last cumulative value (not zeros). + // Zero-fill would make cu_seqlens/input_offsets non-monotonic + // (…, total_tokens, 0, …) → negative FA varlen spans if anything reads + // compiled_n_req width. + if (runtime_n_req < compiled_n_req) { + const size_t off_tail = compiled_n_req - runtime_n_req; + auto last_off = graph_input.input_offsets.value() + ->narrow({{0, runtime_n_req, 1}}) + ->to(infinicore::Device{infinicore::Device::Type::kCpu}); + const int32_t last_val = + *reinterpret_cast(last_off->data()); + std::vector pad(off_tail, last_val); + auto offsets_tail = graph_input.input_offsets.value()->narrow( + {{0, runtime_n_req + 1, off_tail}}); + infinicore::context::memcpyH2D( + offsets_tail->data(), pad.data(), pad.size() * sizeof(int32_t), false); + auto cu_tail = graph_input.cu_seqlens.value()->narrow( + {{0, runtime_n_req + 1, off_tail}}); + infinicore::context::memcpyH2D( + cu_tail->data(), pad.data(), pad.size() * sizeof(int32_t), false); + } + + const size_t block_per_req = runtime.block_tables.value()->size(1); + const size_t compiled_block_per_req = graph_input.block_tables.value()->size(1); + if (block_per_req > compiled_block_per_req) { + throw std::runtime_error("block_tables width exceeds compiled bucket"); + } + + auto &graph_block_tables = graph_input.block_tables.value(); + set_minus_one(graph_block_tables); + graph_block_tables + ->narrow({{0, 0, runtime_n_req}, {1, 0, block_per_req}}) + ->copy_from(runtime.block_tables.value()); + + if (runtime_n_req < compiled_n_req) { + auto stale_rows = graph_block_tables->narrow( + {{0, runtime_n_req, compiled_n_req - runtime_n_req}, {1, 0, compiled_block_per_req}}); + set_minus_one(stale_rows); + } + + graph_input.slot_mapping.value() + ->narrow({{0, 0, valid_seq_len}}) + ->copy_from(runtime.slot_mapping.value()); + if (valid_seq_len < bucket) { + // Prefer decode-style pad hygiene: zero slot tails (not -1). Prefill FA is + // eager, but post CG may still touch padded width; -1 slots ATU under + // rearrange when pad-up CG is on. + auto slot_tail = graph_input.slot_mapping.value()->narrow({{0, valid_seq_len, bucket - valid_seq_len}}); + set_zeros(slot_tail); + auto ids_tail = graph_input.input_ids.value()->narrow({{1, valid_seq_len, bucket - valid_seq_len}}); + set_zeros(ids_tail); + auto pos_tail = graph_input.position_ids.value()->narrow({{0, valid_seq_len, bucket - valid_seq_len}}); + set_zeros(pos_tail); + } +} + +std::optional PiecewisePrefillCompiler::run_prefill(const InfinilmModel::Input &input) { + if (!enabled_) { + return std::nullopt; + } + // Phase-scoped MoE: Prefill must not take Triton-in-graph path. + // Under full_and_piecewise MoE is host-break in Prefill (Decode-only adaptive). + infinicore::context::InferencePhaseGuard phase_guard( + infinicore::context::InferencePhase::Prefill); + last_prefill_executed_ = false; + const bool profile = rank_worker_profile_enabled(); + const double t_total0 = profile ? monotonic_ms() : 0.0; + const size_t seq_len = compute_prefill_len(input); + const size_t padded = padded_bucket_for(seq_len); + const size_t graph_bucket = graph_replay_bucket_for_padded(padded); + const size_t runtime_n_req = input.block_tables.has_value() ? input.block_tables.value()->size(0) : 1; + // Decode rows set is_final=True for sampling; that must not hide mid-prefill + // rows (MIXED). mid_chunk ← any False flag; need_lm_head ← any True / empty. + const bool mid_chunk = + InfinilmModel::has_nonfinal_prefill_chunk(input.is_final_prefill_chunk); + const bool need_lm_head = + InfinilmModel::any_final_prefill_chunk(input.is_final_prefill_chunk); + if (profile) { + spdlog::info( + "rank_worker_profile: piecewise run_prefill begin seq_len={} padded={} graph_bucket={} " + "n_req={} mid_chunk={} need_lm_head={}", + seq_len, + padded, + graph_bucket, + runtime_n_req, + mid_chunk, + need_lm_head); + } + if (compiled_.find(graph_bucket) == compiled_.end()) { + ++prefill_misses_; + return std::nullopt; + } + if (runtime_n_req > max_capture_req_) { + ++prefill_misses_; + spdlog::debug( + "piecewise run_prefill: runtime_n_req={} > max_capture_req={} → eager", + runtime_n_req, + max_capture_req_); + return std::nullopt; + } + const bool inductor_mode = infinilm::global_state::piecewise_inductor_segment_enabled(); + // Pad-up (seq_len < graph_bucket): product default replays token-bucket CG with + // valid_seq_len=L + zeroed tails (INFINI_PIECEWISE_PAD_UP_CG, default ON). + // Kill-switch =0 restores legacy eager-only pad-up (PlannedMeta / SIGSEGV history). + const bool pad_up = seq_len != graph_bucket; + const bool pad_up_cg_on = piecewise_pad_up_cg_enabled(); + const bool pad_up_blocks_cg = pad_up && !pad_up_cg_on; + const bool qkv_only = pre_attn_qkv_only(); + + const bool have_mid_graphs = [&]() { + auto it = compiled_mid_.find(graph_bucket); + if (it == compiled_mid_.end()) { + return false; + } + return !it->second.pre_attn.empty() && it->second.pre_attn[0] + && !it->second.post_attn.empty() && it->second.post_attn[0]; + }(); + const bool have_mixed_mid_graphs = [&]() { + auto it = compiled_mixed_mid_.find(graph_bucket); + if (it == compiled_mixed_mid_.end()) { + return false; + } + // Legacy KEEP_MIXED_MID_CAPTURE banks: post + lm_head only. + return !it->second.post_attn.empty() && it->second.post_attn[0] + && static_cast(it->second.lm_head); + }(); + (void)have_mixed_mid_graphs; // telemetry / legacy banks only; prefer uses final compiled_. + // Runtime past of request 0 (representative). Continuing homo mids (past>0) use + // compiled_mid_; first mid (past==0) uses shape-keyed final compiled_ pre+post + // (fused LN+QKV+RoPE CG by default; QKV_ONLY → eager RoPE). + int32_t runtime_past0 = 0; + if (input.past_sequence_lengths.has_value() + && input.past_sequence_lengths.value()->size(0) > 0) { + auto cpu = input.past_sequence_lengths.value()->to(infinicore::Device{infinicore::Device::Type::kCpu}); + runtime_past0 = reinterpret_cast(cpu->data())[0]; + } + const bool past_matches_mid_capture = runtime_past0 > 0; + // Shape-keyed QKV + post live on compiled_[bucket] (final capture). + const bool have_shape_qkv = + compiled_.count(graph_bucket) > 0 + && !compiled_.at(graph_bucket).pre_attn.empty() + && compiled_.at(graph_bucket).pre_attn[0]; + const bool have_final_post = + compiled_.count(graph_bucket) > 0 + && !compiled_.at(graph_bucket).post_attn.empty() + && compiled_.at(graph_bucket).post_attn[0]; + // Telemetry only: former layout gate 1+(bucket-1); prefer no longer requires it. + bool mixed_layout_matches_capture = false; + int32_t mixed_q0 = -1; + int32_t mixed_q1 = -1; + int32_t runtime_mid_past = runtime_past0; + if (runtime_n_req >= 2 && input.input_offsets.has_value() + && input.input_offsets.value()->size(0) >= runtime_n_req + 1) { + auto off_cpu = input.input_offsets.value()->to(infinicore::Device{infinicore::Device::Type::kCpu}); + const auto *offs = reinterpret_cast(off_cpu->data()); + mixed_q0 = offs[1] - offs[0]; + mixed_q1 = offs[2] - offs[1]; + const int32_t expect_mid_q = static_cast(graph_bucket) - 1; + mixed_layout_matches_capture = (mixed_q0 == 1 && mixed_q1 == expect_mid_q); + if (input.past_sequence_lengths.has_value() + && input.past_sequence_lengths.value()->size(0) >= 2) { + auto past_cpu = input.past_sequence_lengths.value()->to(infinicore::Device{infinicore::Device::Type::kCpu}); + // Row1 past under decode-first packs (telemetry). + runtime_mid_past = reinterpret_cast(past_cpu->data())[1]; + } + } + // vLLM-like: MIXED exact-width (or pad-up CG) PIECEWISE uses final token-bucket + // graphs (compiled_[bucket]). ENABLE/REPRO is the product kill-switch. + const bool prefer_mixed_piecewise = + mid_chunk && need_lm_head && !pad_up_blocks_cg && have_shape_qkv && have_final_post + && (enable_mixed_mid_cg() || repro_mixed_mid_cg()); + // Alias for greppability with prior probe strings. + const bool prefer_mixed_mid_graphs = prefer_mixed_piecewise; + // Product: exact homogeneous mid + past>0 → mid dual-capture graphs. + // Part A bisect: ALLOW_MIDCHUNK_CG forces final graphs on mid. + const bool prefer_mid_graphs = + !prefer_mixed_piecewise + && mid_chunk && !need_lm_head && !pad_up_blocks_cg && have_mid_graphs + && past_matches_mid_capture && !repro_allow_midchunk_cg(); + // Product: homogeneous first mid (past==0) → final shape-keyed pre+post (not + // compiled_mid_, which was captured with past=chunk_size). Do not route onto mid banks. + const bool prefer_past0_mid_final_graphs = + !prefer_mixed_piecewise + && mid_chunk && !need_lm_head && !pad_up_blocks_cg && have_shape_qkv + && runtime_past0 == 0; + // bucket_graphs: prefer_mid → compiled_mid_; else compiled_ (mixed + past0 + final). + auto &bucket_graphs = prefer_mid_graphs + ? compiled_mid_.at(graph_bucket) + : compiled_.at(graph_bucket); + + // Mid CG ok: past>0 mid / MIXED final-bank / past0→final, or repro lifts onto final graphs. + const bool mid_cg_ok = + prefer_mid_graphs + || prefer_mixed_piecewise + || prefer_past0_mid_final_graphs + || (mid_chunk && !pad_up_blocks_cg + && (repro_allow_midchunk_cg() || repro_skip_midchunk_eager())); + const bool force_mid_eager = mid_chunk && !mid_cg_ok; + // Step-level eager hist: pad_up_eager > mid_chunk > missing > exact_cg. + // Successful pad-up CG counts as exact_cg (not pw_eager_pad_up). + { + const bool missing_graph = + (!have_shape_qkv + && (bucket_graphs.pre_attn.empty() || !bucket_graphs.pre_attn[0])) + || bucket_graphs.post_attn.empty() || !bucket_graphs.post_attn[0]; + const bool use_pad_up_cg = + pad_up && pad_up_cg_on && !missing_graph && !force_mid_eager; + const bool pad_up_eager = pad_up && !use_pad_up_cg; + const bool mid_chunk_eager = force_mid_eager && !pad_up_eager; + dispatch_hist::record_piecewise_eager(pad_up_eager, mid_chunk_eager, missing_graph); + } + const double t_copy0 = profile ? monotonic_ms() : 0.0; + copy_runtime_into_bucket_(bucket_graphs, input, seq_len); + // Keep sample/lm_head row flags aligned with the live MIXED pack (capture baked + // decode-first {true,false}; runtime must match for eager lm_head gather). + if (!input.is_final_prefill_chunk.empty()) { + bucket_graphs.input.is_final_prefill_chunk = input.is_final_prefill_chunk; + } + set_attn_metadata_for_varlen_batch(bucket_graphs.input, input); + if (profile) { + spdlog::info( + "rank_worker_profile: piecewise copy_runtime+metadata_ms={:.3f}", + monotonic_ms() - t_copy0); + } + + auto &piecewise = infinilm::global_state::get_forward_context().piecewise; + piecewise.valid_seq_len = seq_len; + piecewise.bucket_seq_len = padded; + piecewise.hidden_states = bucket_graphs.hidden_states; + piecewise.residual = bucket_graphs.residual; + piecewise.ar_staging = bucket_graphs.ar_staging; + piecewise.layer_staging = bucket_graphs.layer_staging; + if (!mid_chunk) { + // Exact-width and pad-up final chunks both use inductor when enabled *and* + // not taking native pad-up CG (kill-switch / miss falls back to eager AOT). + // Pad-up still sets piecewise.valid_seq_len=L; AOT pads positions/hidden to bucket. + piecewise.allow_inductor_pre_attn = + inductor_mode && !repro_skip_final_inductor() + && (pad_up_blocks_cg || !pad_up); + } else if (prefer_past0_mid_final_graphs || prefer_mixed_piecewise || prefer_mid_graphs) { + // Product mid CG prefers: fused pre CG (or QKV_ONLY+eager RoPE) — never + // fused PreAttn AOTI (packages bake RoPE; incompatible with host-break RoPE). + piecewise.allow_inductor_pre_attn = false; + } else if (repro_skip_midchunk_eager()) { + // A2: stress inductor-in-CG mid (out of product default). + piecewise.allow_inductor_pre_attn = inductor_mode; + } else { + // Remaining mid paths: keep inductor out. + piecewise.allow_inductor_pre_attn = false; + } + + // Fresh residual each replay (matches capture warmup); hidden prefix comes from embed. + set_zeros(piecewise.residual); + // Pad-up: clear stale CG tails on staging/hidden/logits before embed so pad + // positions do not pollute layernorm / MoE (RC-4). + if (seq_len < graph_bucket) { + clear_stale_bucket_tails_(piecewise, bucket_graphs.logits_holder, seq_len, graph_bucket); + } + + model_->native_piecewise_embed(bucket_graphs.input, piecewise.hidden_states); + + + // Always walk the full model; layers without a captured graph run eager + // (QY may cap capture via INFINI_NATIVE_CG_MAX_LAYERS to fit cudaGraphExec budget). + const size_t num_layers = model_->native_piecewise_num_layers(); + const size_t captured_post = bucket_graphs.post_attn.size(); + const size_t captured_pre = bucket_graphs.pre_attn.size(); + (void)captured_pre; + // Mid with dual mid graphs: CG replay; inductor-in-CG layers stay eager unless A2. + // Final-chunk inductor pre-attn needs the same eager post replay (B4 tail). + // Pad-up: CG when PAD_UP_CG on; legacy kill-switch forces eager post/lm_head. + // MIXED prefer: pre CG + eager attn + post CG from compiled_; lm_head always eager + // when mid+need_lm_head (live is_final_prefill_chunk; do not replay capture-baked gather). + const bool use_eager_post = + pad_up_blocks_cg + || force_mid_eager + || (inductor_mode && piecewise.allow_inductor_pre_attn); + const bool use_eager_lm_head = + pad_up_blocks_cg + || force_mid_eager + || (need_lm_head && mid_chunk) + || prefer_mid_graphs + || prefer_past0_mid_final_graphs + || !bucket_graphs.lm_head + || (inductor_mode && piecewise.allow_inductor_pre_attn); + { + // Telemetry: first past=0 mid, first past>0 homo mid, first MIXED mid prefer, + // and first short final pad-up CG. + static bool logged_mid_past0 = false; + static bool logged_mid_past_gt0 = false; + static bool logged_mixed_mid = false; + static bool logged_pad_up_cg = false; + if (pad_up && pad_up_cg_on && !force_mid_eager && !logged_pad_up_cg) { + logged_pad_up_cg = true; + spdlog::info( + "pad_up_cg=true valid_seq_len={} graph_bucket={} mid_chunk={} " + "pre_rope_in_cg={}", + seq_len, + graph_bucket, + mid_chunk, + !qkv_only); + } + if (mid_chunk && !pad_up_blocks_cg && graph_bucket >= 2048) { + const bool do_log = (runtime_past0 <= 0 && !logged_mid_past0) + || (runtime_past0 > 0 && !need_lm_head && !logged_mid_past_gt0) + || (prefer_mixed_piecewise && !logged_mixed_mid) + || (need_lm_head && !logged_mixed_mid); + if (do_log) { + if (runtime_past0 <= 0 && !need_lm_head) { + logged_mid_past0 = true; + } else if (!need_lm_head) { + logged_mid_past_gt0 = true; + } + if (prefer_mixed_piecewise || need_lm_head) { + logged_mixed_mid = true; + } + const bool pre_cg = + have_shape_qkv && !pad_up_blocks_cg && !force_mid_eager; + spdlog::info( + "mid_chunk_cg_probe: mid_chunk={} past_len={} mid_past={} seq_len={} " + "graph_bucket={} force_mid_eager={} prefer_mid_graphs={} " + "prefer_past0_mid_final_graphs={} prefer_mixed_piecewise={} " + "prefer_mixed_mid_graphs={} " + "mid_cg_ok={} past_matches_mid_capture={} " + "mixed_layout_matches_capture={} mixed_q0={} mixed_q1={} n_req={} " + "allow_inductor_pre_attn={} use_eager_post={} use_eager_lm_head={} " + "pre_qkv_cg={} pre_rope_in_cg={} pre_rope_eager={} " + "have_shape_qkv={} have_final_post={} pad_up={} pad_up_cg_on={}", + mid_chunk, + runtime_past0, + runtime_mid_past, + seq_len, + graph_bucket, + force_mid_eager, + prefer_mid_graphs, + prefer_past0_mid_final_graphs, + prefer_mixed_piecewise, + prefer_mixed_mid_graphs, + mid_cg_ok, + past_matches_mid_capture, + mixed_layout_matches_capture, + mixed_q0, + mixed_q1, + runtime_n_req, + piecewise.allow_inductor_pre_attn, + use_eager_post, + use_eager_lm_head, + pre_cg, + pre_cg && !qkv_only, + pre_cg && qkv_only, + have_shape_qkv, + have_final_post, + pad_up, + pad_up_cg_on); + } + } + } + double layers_ms = 0.0; + const double t_layers0 = profile ? monotonic_ms() : 0.0; + for (size_t layer = 0; layer < num_layers; ++layer) { + const double t_layer0 = profile ? monotonic_ms() : 0.0; + barrier_->wait(); + const bool inductor_layer_mid = + mid_chunk && layer_capture_inductor_pre_attn(layer, graph_bucket) + && !repro_skip_midchunk_eager(); + // Product: fused pre CG (LN+QKV+RoPE). QKV_ONLY kill-switch: QKV CG + eager RoPE. + // Prefer shape-keyed CG over inductor-in-pre (AOT packages bake RoPE). + // Legacy pad-up kill-switch / force_mid_eager / mid-inductor stress → eager full pre. + const bool force_eager_pre = + pad_up_blocks_cg + || force_mid_eager + || inductor_layer_mid; + + // Prefer path pre: bucket_graphs is compiled_ for MIXED/past0/final, or + // compiled_mid_ for homo continuing mid. + std::shared_ptr pre_graph; + if (!force_eager_pre) { + if (layer < bucket_graphs.pre_attn.size() + && bucket_graphs.pre_attn[layer]) { + pre_graph = bucket_graphs.pre_attn[layer]; + } else if (have_shape_qkv && layer < compiled_.at(graph_bucket).pre_attn.size() + && compiled_.at(graph_bucket).pre_attn[layer]) { + pre_graph = compiled_.at(graph_bucket).pre_attn[layer]; + } + } + + double pre_graph_run_ms = 0.0; + double pre_eager_ms = 0.0; + infinicore::graph::Graph::LastRunProfile pre_split{}; + if (pre_graph) { + const double t_pre_g0 = profile ? monotonic_ms() : 0.0; + pre_graph->run(); + pre_split = pre_graph->last_run_profile(); + // When Graph::run skips its tail sync, PROFILE_SYNC still needs a drain + // for exclusive graph_run_ms; otherwise FA must see pre outputs. + if (profile && rank_worker_profile_sync()) { + if (pre_split.skipped_tail_sync || !pre_split.used_device) { + infinicore::context::syncDevice(); + } + } else if (pre_split.skipped_tail_sync) { + // FA is eager and reads QKV from pre CG — one stream drain. + infinicore::context::syncStream(); + } + if (profile) { + pre_graph_run_ms = monotonic_ms() - t_pre_g0; + } + ++segment_replays_; + if (qkv_only) { + // Phase-2c: RoPE must see QKV staging writes; graph replay may be async. + // Opt-in skip via INFINI_PIECEWISE_SKIP_QKV_ROPE_SYNC=1. + if (!global_state::skip_qkv_rope_sync()) { + infinicore::context::syncDevice(); + } + const double t_rope0 = profile ? monotonic_ms() : 0.0; + model_->native_piecewise_pre_attn_rope_layer(layer, bucket_graphs.input); + if (profile) { + pre_eager_ms = monotonic_ms() - t_rope0; + } + } + // else: fused pre already includes RoPE — no host break / sync. + } else { + // Eager full pre (QKV+RoPE / inductor package when allowed). + const double t_pre_e0 = profile ? monotonic_ms() : 0.0; + model_->native_piecewise_pre_attn_layer( + layer, bucket_graphs.input, piecewise.hidden_states, piecewise.residual); + if (profile) { + pre_eager_ms = monotonic_ms() - t_pre_e0; + } + } + const double t_pre_attn = profile ? monotonic_ms() : 0.0; + piecewise.phase = global_state::PiecewiseCapturePhase::EagerAttn; + model_->native_piecewise_eager_attn_layer(layer, bucket_graphs.input); + const double t_eager_attn = profile ? monotonic_ms() : 0.0; + barrier_->wait(); + double post_graph_run_ms = 0.0; + double post_eager_ms = 0.0; + infinicore::graph::Graph::LastRunProfile post_split{}; + const bool have_post_graph = + layer < captured_post && static_cast(bucket_graphs.post_attn[layer]); + if (use_eager_post || inductor_layer_mid || !have_post_graph) { + const double t_post_e0 = profile ? monotonic_ms() : 0.0; + model_->native_piecewise_post_attn_cg_layer( + layer, bucket_graphs.input, piecewise.hidden_states, piecewise.residual); + barrier_->wait(); + if (profile) { + post_eager_ms = monotonic_ms() - t_post_e0; + } + } else { + const double t_post_g0 = profile ? monotonic_ms() : 0.0; + bucket_graphs.post_attn[layer]->run(); + post_split = bucket_graphs.post_attn[layer]->last_run_profile(); + // Double-sync cut: Graph::run already synced unless SKIP_TAIL_SYNC. + // PROFILE_SYNC exclusive timing: skip redundant syncDevice when Graph + // already drained the stream (LIGHT_SYNC spirit / TP=1). + if (profile && rank_worker_profile_sync()) { + if (post_split.skipped_tail_sync || !post_split.used_device) { + infinicore::context::syncDevice(); + } + } + // Mid-layer: next consumer is another Graph::run / FA on same stream — + // SKIP_TAIL_SYNC leaves work ordered; drain once before lm_head below. + if (profile) { + post_graph_run_ms = monotonic_ms() - t_post_g0; + } + ++segment_replays_; + barrier_->wait(); + } + if (profile) { + // Barrier after Graph::run folds into post_attn_ms only; graph_run stays + // Graph::run(+optional sync) exclusive. + std::string extra; + if (pre_graph) { + append_graph_run_split(extra, "pre_graph", pre_split); + } + if (post_graph_run_ms > 0.0 || post_split.used_device ) { + append_graph_run_split(extra, "post_graph", post_split); + } + spdlog::info( + "rank_worker_profile: piecewise layer={} pre_attn_ms={:.3f} " + "pre_graph_run_ms={:.3f} pre_eager_ms={:.3f} eager_attn_ms={:.3f} " + "post_attn_ms={:.3f} post_graph_run_ms={:.3f} post_eager_ms={:.3f} " + "layer_total_ms={:.3f}{}", + layer, + t_pre_attn - t_layer0, + pre_graph_run_ms, + pre_eager_ms, + t_eager_attn - t_pre_attn, + monotonic_ms() - t_eager_attn, + post_graph_run_ms, + post_eager_ms, + monotonic_ms() - t_layer0, + extra); + } + } + if (profile) { + layers_ms = monotonic_ms() - t_layers0; + } + if (need_lm_head) { + const double t_lm0 = profile ? monotonic_ms() : 0.0; + double lm_graph_run_ms = 0.0; + double lm_eager_ms = 0.0; + barrier_->wait(); + // Gate C: one stream drain before logits when post GraphExec stayed async. + if (graph_skip_tail_sync_enabled()) { + infinicore::context::syncStream(); + } + if (use_eager_lm_head) { + const double t_e0 = profile ? monotonic_ms() : 0.0; + model_->native_piecewise_lm_head( + bucket_graphs.input, piecewise.hidden_states, piecewise.residual, bucket_graphs.logits_holder); + if (profile) { + lm_eager_ms = monotonic_ms() - t_e0; + } + } else { + const double t_g0 = profile ? monotonic_ms() : 0.0; + bucket_graphs.lm_head->run(); + if (profile && rank_worker_profile_sync()) { + const auto &lm_split = bucket_graphs.lm_head->last_run_profile(); + if (lm_split.skipped_tail_sync || !lm_split.used_device) { + infinicore::context::syncDevice(); + } + } + if (profile) { + lm_graph_run_ms = monotonic_ms() - t_g0; + } + ++segment_replays_; + } + barrier_->wait(); + if (profile) { + spdlog::info( + "rank_worker_profile: piecewise lm_head_ms={:.3f} " + "lm_head_graph_run_ms={:.3f} lm_head_eager_ms={:.3f}", + monotonic_ms() - t_lm0, + lm_graph_run_ms, + lm_eager_ms); + } + } else if (graph_skip_tail_sync_enabled()) { + // Mid-chunk without lm_head: still drain before returning (KV consumers). + infinicore::context::syncStream(); + } + + piecewise.phase = global_state::PiecewiseCapturePhase::None; + ++prefill_hits_; + last_prefill_executed_ = true; + if (profile) { + spdlog::info( + "rank_worker_profile: piecewise run_prefill end layers_total_ms={:.3f} total_ms={:.3f} " + "mid_chunk={} need_lm_head={} pad_up={}", + layers_ms, + monotonic_ms() - t_total0, + mid_chunk, + need_lm_head, + pad_up); + } + if (global_state::ar_profile::enabled() + && global_state::get_tensor_model_parallel_rank() == 0) { + global_state::ar_profile::log_barrier_chunk_summary( + "piecewise_replay", seq_len, runtime_n_req); + } + if (!need_lm_head) { + return std::nullopt; + } + return bucket_graphs.logits_holder->narrow({{1, 0, seq_len}}); +} + +} // namespace infinilm::engine diff --git a/csrc/engine/compiler/piecewise_prefill_compiler.hpp b/csrc/engine/compiler/piecewise_prefill_compiler.hpp new file mode 100644 index 000000000..d24e084d1 --- /dev/null +++ b/csrc/engine/compiler/piecewise_prefill_compiler.hpp @@ -0,0 +1,107 @@ +#pragma once + +#include "../../global_state/piecewise_prefill_state.hpp" +#include "../../models/infinilm_model.hpp" +#include "../rank_barrier.hpp" +#include "graph_compiler.hpp" +#include "paged_compiler.hpp" + +#include +#include +#include +#include + +namespace infinilm::engine { + +class PiecewisePrefillCompiler { +public: + struct BucketGraphs { + InfinilmModel::Input input; + std::vector> pre_attn; + std::vector> post_attn; + std::shared_ptr lm_head; + infinicore::Tensor logits_holder; + infinicore::Tensor hidden_states; + infinicore::Tensor residual; + infinicore::Tensor ar_staging; + std::vector layer_staging; + }; + + /// M2: one max-bucket staging/I-O bank; smaller buckets use prefix views. + struct SharedPrefillBanks { + size_t max_bucket{0}; + size_t num_layers{0}; + size_t staging_physical_bytes{0}; + size_t io_physical_bytes{0}; + infinicore::Tensor hidden_states; + infinicore::Tensor residual; + infinicore::Tensor ar_staging; + std::vector layer_staging; + infinicore::Tensor input_ids; + infinicore::Tensor position_ids; + infinicore::Tensor past_sequence_lengths; + infinicore::Tensor total_sequence_lengths; + infinicore::Tensor input_offsets; + infinicore::Tensor cu_seqlens; + infinicore::Tensor slot_mapping; + infinicore::Tensor logits_holder; + }; + + PiecewisePrefillCompiler(const std::shared_ptr &model, RankBarrier *barrier); + + void compile(); + bool enabled() const { return enabled_; } + + /// Replay piecewise graphs for a prefill request. Returns logits on hit. + std::optional run_prefill(const InfinilmModel::Input &input); + + size_t padded_bucket_for(size_t seq_len) const; + size_t max_capture_req() const { return max_capture_req_; } + const std::vector &capture_buckets() const { return capture_buckets_; } + size_t segment_replays() const { return segment_replays_; } + size_t prefill_hits() const { return prefill_hits_; } + size_t prefill_misses() const { return prefill_misses_; } + bool last_prefill_executed() const { return last_prefill_executed_; } + +private: + void allocate_shared_banks_(size_t max_bucket, size_t num_layers, size_t n_req); + /// Bind prefix views of the shared max bank into piecewise TLS for ``bucket``. + void bind_bucket_staging_(size_t bucket, size_t num_layers); + /// ``mid_chunk_capture``: past_len=chunk_size, positions continue from past (exact mid key). + InfinilmModel::Input make_bucket_input_(size_t bucket, + size_t nblocks, + size_t n_req, + bool mid_chunk_capture = false) const; + /// MIXED dual-capture: n_req=2 continuing mid (~bucket-1 toks) + 1 decode tok @ chunk bucket. + InfinilmModel::Input make_mixed_mid_bucket_input_(size_t bucket, size_t nblocks) const; + void capture_bucket_(size_t bucket, bool mid_chunk_capture = false); + void capture_mixed_mid_bucket_(size_t bucket); + void warmup_inductor_segments_(size_t nblocks, size_t n_req); + void copy_runtime_into_bucket_(BucketGraphs &bucket_graphs, + const InfinilmModel::Input &runtime, + size_t valid_seq_len) const; + + std::shared_ptr model_; + RankBarrier *barrier_; + bool enabled_{false}; + size_t max_seq_len_{0}; + size_t prefill_chunk_size_{0}; + std::vector capture_buckets_; + std::vector bs_to_padded_; + size_t max_capture_req_{1}; + infinicore::Tensor block_tables_holder_; + SharedPrefillBanks shared_banks_; + std::unordered_map compiled_; + /// Exact-width mid-chunk graphs for ``bucket == prefill_chunk_size`` (dual capture). + std::unordered_map compiled_mid_; + /// MIXED mid+decode graphs for ``bucket == prefill_chunk_size`` (n_req=2, includes lm_head). + std::unordered_map compiled_mixed_mid_; + /// Capture width used for ``compiled_mixed_mid_`` (normally 2). + size_t mixed_mid_capture_req_{2}; + size_t segment_replays_{0}; + size_t prefill_hits_{0}; + size_t prefill_misses_{0}; + bool last_prefill_executed_{false}; +}; + +} // namespace infinilm::engine diff --git a/csrc/engine/cudagraph_dispatcher.hpp b/csrc/engine/cudagraph_dispatcher.hpp new file mode 100644 index 000000000..c699dd44f --- /dev/null +++ b/csrc/engine/cudagraph_dispatcher.hpp @@ -0,0 +1,246 @@ +#pragma once + +/// vLLM-shaped CUDA-graph runtime dispatch (thin InfiniLM mirror). +/// +/// Selects FULL / PIECEWISE / NONE from a BatchDescriptor derived from input +/// shape (not scheduler PREFILL/DECODE/MIXED). Under ``full_and_piecewise``: +/// - uniform decode batches in ``INFINI_DECODE_CG_BATCHES`` → FULL +/// (exact hit, or pad-up to next FULL key unless ``INFINI_DECODE_CG_PAD_UP=0``) +/// - ragged / mixed / multi-req prefill (``!uniform_decode``): pad-up +/// ``num_tokens`` to the next ``INFINI_NATIVE_CG_CAPTURE_BUCKETS`` entry +/// (vLLM ``bs_to_padded_graph_size`` / ``mixed_mode=PIECEWISE``) → PIECEWISE +/// with ``key.num_tokens = padded``; eager only when +/// ``num_tokens > max_capture_size`` +/// - else → NONE (eager) +/// ``eager`` policy → always NONE. + +#include "compiled_prefill_flags.hpp" +#include "compiler/piecewise_bucket_policy.hpp" + +#include +#include +#include +#include +#include +#include +#include + +namespace infinilm::engine { + +enum class CudaGraphRuntimeMode { + None = 0, + Piecewise = 1, + Full = 2, +}; + +inline const char *cudagraph_runtime_mode_cstr(CudaGraphRuntimeMode mode) { + switch (mode) { + case CudaGraphRuntimeMode::Full: + return "FULL"; + case CudaGraphRuntimeMode::Piecewise: + return "PIECEWISE"; + case CudaGraphRuntimeMode::None: + default: + return "NONE"; + } +} + +struct BatchDescriptor { + size_t num_tokens{0}; + size_t num_reqs{0}; + /// True when block_tables.batch == input_ids.width and every row schedules + /// exactly one new token (decode-shaped). + bool uniform_decode{false}; +}; + +inline bool operator==(const BatchDescriptor &a, const BatchDescriptor &b) { + return a.num_tokens == b.num_tokens && a.num_reqs == b.num_reqs + && a.uniform_decode == b.uniform_decode; +} + +namespace detail { + +inline std::vector parse_csv_sizes_(const char *raw) { + std::vector out; + if (raw == nullptr || raw[0] == '\0') { + return out; + } + std::string spec(raw); + size_t start = 0; + while (start < spec.size()) { + const size_t comma = spec.find(',', start); + const std::string token = + spec.substr(start, comma == std::string::npos ? std::string::npos : comma - start); + if (!token.empty()) { + out.push_back(static_cast(std::stoul(token))); + } + if (comma == std::string::npos) { + break; + } + start = comma + 1; + } + std::sort(out.begin(), out.end()); + out.erase(std::unique(out.begin(), out.end()), out.end()); + return out; +} + +} // namespace detail + +class CudagraphDispatcher { +public: + CudagraphDispatcher() = default; + + /// Load FULL / PIECEWISE key sets from env for the active cudagraph policy. + void initialize_from_env() { + full_keys_.clear(); + piecewise_keys_.clear(); + bs_to_padded_.clear(); + bs_to_padded_decode_.clear(); + max_capture_ = 0; + max_decode_bs_ = 0; + const char *policy = cudagraph_policy(); + if (std::strcmp(policy, "eager") == 0) { + // No capture keys — dispatch always NONE. + return; + } + if (std::strcmp(policy, "full_and_piecewise") != 0 && policy[0] != '\0') { + // Unknown policy string already normalized to "" by cudagraph_policy(). + return; + } + // full_and_piecewise or legacy (empty policy with companion envs). + if (std::strcmp(policy, "full_and_piecewise") == 0 + || std::strcmp(policy, "") == 0) { + for (size_t b : detail::parse_csv_sizes_(std::getenv("INFINI_DECODE_CG_BATCHES"))) { + full_keys_.insert(b); + } + for (size_t b : + detail::parse_csv_sizes_(std::getenv("INFINI_NATIVE_CG_CAPTURE_BUCKETS"))) { + piecewise_keys_.insert(b); + } + rebuild_pad_table_(); + rebuild_decode_pad_table_(); + } + } + + void add_full_key(size_t num_tokens) { + full_keys_.insert(num_tokens); + rebuild_decode_pad_table_(); + } + void add_piecewise_key(size_t num_tokens) { + piecewise_keys_.insert(num_tokens); + rebuild_pad_table_(); + } + + const std::set &full_keys() const { return full_keys_; } + const std::set &piecewise_keys() const { return piecewise_keys_; } + + /// Priority FULL > PIECEWISE > NONE. Returns (mode, padded key descriptor). + /// Prefill pad-up matches vLLM ``bs_to_padded_graph_size``: in-range + /// ``num_tokens`` maps to the next capture bucket; key.num_tokens is padded. + /// Decode pad-up (default on under full_and_piecewise): exact miss maps to + /// the next ``INFINI_DECODE_CG_BATCHES`` entry; kill-switch + /// ``INFINI_DECODE_CG_PAD_UP=0``. + std::pair dispatch(const BatchDescriptor &desc) const { + if (std::strcmp(cudagraph_policy(), "eager") == 0) { + return {CudaGraphRuntimeMode::None, desc}; + } + if (desc.uniform_decode && !full_keys_.empty()) { + if (full_keys_.count(desc.num_tokens) > 0) { + BatchDescriptor key = desc; + key.num_reqs = desc.num_tokens; // uniform decode: 1 token / req + return {CudaGraphRuntimeMode::Full, key}; + } + if (decode_cg_pad_up_enabled_() && desc.num_tokens <= max_decode_bs_) { + const size_t padded = padded_bucket_for_seq_len( + desc.num_tokens, bs_to_padded_decode_, /*fallback=*/0); + if (padded > 0 && full_keys_.count(padded) > 0) { + BatchDescriptor key = desc; + key.num_tokens = padded; + key.num_reqs = padded; + return {CudaGraphRuntimeMode::Full, key}; + } + } + } + // Ragged / mixed / homogeneous prefill: vLLM-style pad-up of num_tokens + // (no num_reqs==1 gate; capture must allow runtime_n_req ≤ max_capture_req). + if (!desc.uniform_decode && !piecewise_keys_.empty()) { + if (desc.num_tokens > max_capture_) { + return {CudaGraphRuntimeMode::None, desc}; + } + const size_t padded = + padded_bucket_for_seq_len(desc.num_tokens, bs_to_padded_, /*fallback=*/0); + if (padded > 0 && piecewise_keys_.count(padded) > 0) { + BatchDescriptor key = desc; + key.num_tokens = padded; + return {CudaGraphRuntimeMode::Piecewise, key}; + } + } + return {CudaGraphRuntimeMode::None, desc}; + } + + /// Classify why ``dispatch`` returned NONE (for profile / hang_trace histograms). + /// ``is_mixed`` is the scheduler MIXED / ragged flag from RankWorker. + const char *none_reason(const BatchDescriptor &desc, bool is_mixed) const { + if (std::strcmp(cudagraph_policy(), "eager") == 0) { + return "eager_policy"; + } + if (desc.uniform_decode) { + if (!full_keys_.empty() && desc.num_tokens > max_decode_bs_) { + return "decode_bs_over_max"; + } + return "decode_bs_miss"; + } + if (!piecewise_keys_.empty() && desc.num_tokens > max_capture_) { + return "over_max"; + } + // Rare: no keys / pad table miss. Keep legacy labels for hist continuity. + if (is_mixed) { + return "mixed"; + } + if (desc.num_reqs > 1) { + return "multi_req_prefill"; + } + return "bucket_miss"; + } + +private: + /// Decode batch pad-up default on; ``INFINI_DECODE_CG_PAD_UP=0`` disables. + static bool decode_cg_pad_up_enabled_() { + const char *v = std::getenv("INFINI_DECODE_CG_PAD_UP"); + if (v != nullptr && v[0] == '0' && v[1] == '\0') { + return false; + } + return true; + } + + void rebuild_pad_table_() { + bs_to_padded_.clear(); + max_capture_ = 0; + if (piecewise_keys_.empty()) { + return; + } + std::vector caps(piecewise_keys_.begin(), piecewise_keys_.end()); + bs_to_padded_ = build_bs_to_padded_bucket(caps); + max_capture_ = *piecewise_keys_.rbegin(); + } + + void rebuild_decode_pad_table_() { + bs_to_padded_decode_.clear(); + max_decode_bs_ = 0; + if (full_keys_.empty()) { + return; + } + std::vector caps(full_keys_.begin(), full_keys_.end()); + bs_to_padded_decode_ = build_bs_to_padded_bucket(caps); + max_decode_bs_ = *full_keys_.rbegin(); + } + + std::set full_keys_; + std::set piecewise_keys_; + std::vector bs_to_padded_; + std::vector bs_to_padded_decode_; + size_t max_capture_{0}; + size_t max_decode_bs_{0}; +}; + +} // namespace infinilm::engine diff --git a/csrc/engine/dispatch_hist.cpp b/csrc/engine/dispatch_hist.cpp new file mode 100644 index 000000000..68b910d26 --- /dev/null +++ b/csrc/engine/dispatch_hist.cpp @@ -0,0 +1,216 @@ +#include "dispatch_hist.hpp" + +#include +#include +#include +#include + +namespace infinilm::engine::dispatch_hist { +namespace { + +struct Counters { + std::atomic full{0}; + std::atomic piecewise{0}; + std::atomic none{0}; + std::atomic none_eager_policy{0}; + std::atomic none_mixed{0}; + std::atomic none_multi_req_prefill{0}; + std::atomic none_bucket_miss{0}; + std::atomic none_decode_bs_miss{0}; + std::atomic none_decode_bs_over_max{0}; + std::atomic none_over_max{0}; + std::atomic none_other{0}; + // Piecewise eager vs exact-CG (step-level; not per-layer). + std::atomic pw_exact_cg{0}; + std::atomic pw_eager_pad_up{0}; + std::atomic pw_eager_mid_chunk{0}; + std::atomic pw_eager_missing_graph{0}; + std::atomic steps_since_dump{0}; +}; + +Counters &counters() { + static Counters c; + return c; +} + +bool rank_worker_profile_enabled() { + static int cached = -1; + if (cached < 0) { + const char *raw = std::getenv("INFINI_RANK_WORKER_PROFILE"); + cached = (raw != nullptr && raw[0] == '1' && raw[1] == '\0') ? 1 : 0; + } + return cached == 1; +} + +bool hang_trace_enabled() { + static int cached = -1; + if (cached < 0) { + const char *raw = std::getenv("INFINI_HANG_TRACE"); + cached = (raw != nullptr && raw[0] == '1' && raw[1] == '\0') ? 1 : 0; + } + return cached == 1; +} + +bool cg_eager_hist_enabled() { + static int cached = -1; + if (cached < 0) { + const char *raw = std::getenv("INFINI_CG_EAGER_HIST"); + // Default off (quiet). LongBench / restart_serve set =1. + cached = (raw != nullptr && raw[0] == '1' && raw[1] == '\0') ? 1 : 0; + } + return cached == 1; +} + +uint64_t dump_every_n() { + static uint64_t cached = 0; + if (cached == 0) { + const char *raw = std::getenv("INFINI_CG_EAGER_HIST_EVERY"); + if (raw != nullptr && raw[0] != '\0') { + const long v = std::strtol(raw, nullptr, 10); + cached = v > 0 ? static_cast(v) : 200; + } else { + cached = 200; + } + } + return cached; +} + +void dump_line(const char *tag) { + const auto &h = counters(); + spdlog::info( + "cg_eager_hist tag={} FULL={} PIECEWISE={} NONE={} " + "pw_exact_cg={} pw_eager_pad_up={} pw_eager_mid_chunk={} pw_eager_missing={} " + "none_over_max={} none_mixed={} none_bucket_miss={} none_decode_bs_miss={} " + "none_decode_bs_over_max={} none_other={}", + tag != nullptr ? tag : "-", + h.full.load(std::memory_order_relaxed), + h.piecewise.load(std::memory_order_relaxed), + h.none.load(std::memory_order_relaxed), + h.pw_exact_cg.load(std::memory_order_relaxed), + h.pw_eager_pad_up.load(std::memory_order_relaxed), + h.pw_eager_mid_chunk.load(std::memory_order_relaxed), + h.pw_eager_missing_graph.load(std::memory_order_relaxed), + h.none_over_max.load(std::memory_order_relaxed), + h.none_mixed.load(std::memory_order_relaxed), + h.none_bucket_miss.load(std::memory_order_relaxed), + h.none_decode_bs_miss.load(std::memory_order_relaxed), + h.none_decode_bs_over_max.load(std::memory_order_relaxed), + h.none_other.load(std::memory_order_relaxed)); +} + +void atexit_dump() { + if (cg_eager_hist_enabled()) { + dump_line("atexit"); + } +} + +} // namespace + +void record_mode(CudaGraphRuntimeMode mode, const char *none_reason) { + auto &h = counters(); + switch (mode) { + case CudaGraphRuntimeMode::Full: + h.full.fetch_add(1, std::memory_order_relaxed); + break; + case CudaGraphRuntimeMode::Piecewise: + h.piecewise.fetch_add(1, std::memory_order_relaxed); + break; + case CudaGraphRuntimeMode::None: + default: + h.none.fetch_add(1, std::memory_order_relaxed); + if (none_reason == nullptr) { + h.none_other.fetch_add(1, std::memory_order_relaxed); + } else if (std::strcmp(none_reason, "eager_policy") == 0) { + h.none_eager_policy.fetch_add(1, std::memory_order_relaxed); + } else if (std::strcmp(none_reason, "mixed") == 0) { + h.none_mixed.fetch_add(1, std::memory_order_relaxed); + } else if (std::strcmp(none_reason, "multi_req_prefill") == 0) { + h.none_multi_req_prefill.fetch_add(1, std::memory_order_relaxed); + } else if (std::strcmp(none_reason, "bucket_miss") == 0) { + h.none_bucket_miss.fetch_add(1, std::memory_order_relaxed); + } else if (std::strcmp(none_reason, "decode_bs_miss") == 0) { + h.none_decode_bs_miss.fetch_add(1, std::memory_order_relaxed); + } else if (std::strcmp(none_reason, "decode_bs_over_max") == 0) { + h.none_decode_bs_over_max.fetch_add(1, std::memory_order_relaxed); + } else if (std::strcmp(none_reason, "over_max") == 0) { + h.none_over_max.fetch_add(1, std::memory_order_relaxed); + } else { + h.none_other.fetch_add(1, std::memory_order_relaxed); + } + break; + } + maybe_dump_eager("periodic"); +} + +void record_piecewise_eager(bool pad_up, bool mid_chunk_eager, bool missing_graph) { + auto &h = counters(); + if (pad_up) { + h.pw_eager_pad_up.fetch_add(1, std::memory_order_relaxed); + } else if (mid_chunk_eager) { + h.pw_eager_mid_chunk.fetch_add(1, std::memory_order_relaxed); + } else if (missing_graph) { + h.pw_eager_missing_graph.fetch_add(1, std::memory_order_relaxed); + } else { + // Final exact CG and exact-width mid CG (dual mid graphs) both land here. + h.pw_exact_cg.fetch_add(1, std::memory_order_relaxed); + } +} + +void log_profile_hist(const char *tag) { + if (!rank_worker_profile_enabled() && !hang_trace_enabled()) { + return; + } + const auto &h = counters(); + spdlog::info( + "{}: dispatch_hist FULL={} PIECEWISE={} NONE={} " + "none_reason[eager_policy={} mixed={} multi_req_prefill={} " + "bucket_miss={} decode_bs_miss={} decode_bs_over_max={} over_max={} other={}] " + "pw[exact_cg={} pad_up={} mid_chunk={} missing={}]", + tag, + h.full.load(std::memory_order_relaxed), + h.piecewise.load(std::memory_order_relaxed), + h.none.load(std::memory_order_relaxed), + h.none_eager_policy.load(std::memory_order_relaxed), + h.none_mixed.load(std::memory_order_relaxed), + h.none_multi_req_prefill.load(std::memory_order_relaxed), + h.none_bucket_miss.load(std::memory_order_relaxed), + h.none_decode_bs_miss.load(std::memory_order_relaxed), + h.none_decode_bs_over_max.load(std::memory_order_relaxed), + h.none_over_max.load(std::memory_order_relaxed), + h.none_other.load(std::memory_order_relaxed), + h.pw_exact_cg.load(std::memory_order_relaxed), + h.pw_eager_pad_up.load(std::memory_order_relaxed), + h.pw_eager_mid_chunk.load(std::memory_order_relaxed), + h.pw_eager_missing_graph.load(std::memory_order_relaxed)); +} + +void maybe_dump_eager(const char *tag) { + if (!cg_eager_hist_enabled()) { + return; + } + ensure_atexit_dump(); + auto &h = counters(); + const uint64_t n = h.steps_since_dump.fetch_add(1, std::memory_order_relaxed) + 1; + if (n >= dump_every_n()) { + h.steps_since_dump.store(0, std::memory_order_relaxed); + dump_line(tag != nullptr ? tag : "periodic"); + } +} + +void dump_eager_final(const char *tag) { + if (!cg_eager_hist_enabled()) { + return; + } + ensure_atexit_dump(); + dump_line(tag != nullptr ? tag : "teardown"); +} + +void ensure_atexit_dump() { + static std::atomic registered{0}; + int expected = 0; + if (registered.compare_exchange_strong(expected, 1, std::memory_order_relaxed)) { + std::atexit(atexit_dump); + } +} + +} // namespace infinilm::engine::dispatch_hist diff --git a/csrc/engine/dispatch_hist.hpp b/csrc/engine/dispatch_hist.hpp new file mode 100644 index 000000000..c39c22795 --- /dev/null +++ b/csrc/engine/dispatch_hist.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include "cudagraph_dispatcher.hpp" + +namespace infinilm::engine::dispatch_hist { + +/// Record dispatcher mode (FULL / PIECEWISE / NONE + NONE reason). Rank-0 only. +void record_mode(CudaGraphRuntimeMode mode, const char *none_reason); + +/// One step-level count for a PIECEWISE run_prefill entry after path is known. +/// Priority: pad_up_eager > mid_chunk_eager > missing_graph > exact_cg. +/// ``pad_up`` here means *eager* pad-up (kill-switch / miss); successful pad-up +/// CG counts as exact_cg. ``mid_chunk_eager`` is true only when the mid path +/// still ran eager; exact-width / pad-up CG mid replay counts as exact_cg. +void record_piecewise_eager(bool pad_up, bool mid_chunk_eager, bool missing_graph); + +/// Legacy noisy dump gated by INFINI_RANK_WORKER_PROFILE / hang_trace. +void log_profile_hist(const char *tag); + +/// Periodic / teardown one-line dump gated by INFINI_CG_EAGER_HIST=1. +/// Call after each mode record (period) and on process teardown / RankWorker::close. +void maybe_dump_eager(const char *tag); +void dump_eager_final(const char *tag = "teardown"); + +/// Ensure atexit final dump is registered (idempotent). +void ensure_atexit_dump(); + +} // namespace infinilm::engine::dispatch_hist diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index b8eeecc3d..40e60fa0c 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -104,9 +104,10 @@ InferEngine::InferEngine( && device_type != infinicore::Device::Type::kCambricon && device_type != infinicore::Device::Type::kIluvatar && device_type != infinicore::Device::Type::kHygon - && device_type != infinicore::Device::Type::kAscend) { + && device_type != infinicore::Device::Type::kAscend + && device_type != infinicore::Device::Type::kQy) { throw std::invalid_argument( - "flash-attn is only available on NVIDIA, MetaX, Moore, Cambricon, Iluvatar, Hygon, and Ascend devices"); + "flash-attn is only available on NVIDIA, MetaX, Moore, Cambricon, Iluvatar, Hygon, Ascend, and QY devices"); } const auto *paged_cache_config = dynamic_cast(cache_config); if (paged_cache_config == nullptr) { diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index ddfe1ee4e..ca0369de0 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -1,5 +1,7 @@ #include "rank_worker.hpp" +#include "compiled_prefill_flags.hpp" #include "../models/model_factory.hpp" +#include "infinicore/context/context.hpp" #include "infinicore/ops.hpp" #include "infinicore/ops/distributed/send_recv.hpp" #include @@ -418,22 +420,70 @@ void RankWorker::thread_loop() { infinicore::Tensor logits; infinicore::Tensor hidden_states; + bool piecewise_ran = false; // 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}, - false, - true)); - if (graph != nullptr && output != nullptr) { - graph->run(); - logits = output->logits; - hidden_states = output->hidden_states; + auto *general_compiler = dynamic_cast(compiler_.get()); + const bool has_bt = local_args.block_tables.has_value() + && local_args.input_ids.has_value(); + const bool uniform_decode = has_bt + && local_args.block_tables.value()->size(0) + == local_args.input_ids.value()->size(1) + && local_args.input_ids.value()->size(1) > 0; + auto model_input_cpu = local_args.to_model_input( + infinicore::Device{infinicore::Device::Type::kCpu}, + false, + true); + + // F&P native piecewise prefill (FA host-break segments). + if (!uniform_decode && general_compiler != nullptr + && general_compiler->native_piecewise_enabled()) { + infinicore::context::InferencePhaseGuard phase_guard( + infinicore::context::InferencePhase::Prefill); + if (auto pw_logits = + general_compiler->run_native_piecewise_prefill(model_input_cpu)) { + logits = *pw_logits; + piecewise_ran = true; + general_compiler->record_graph_hit(/*is_prefill=*/true); + } else if (general_compiler->native_piecewise_last_prefill_executed()) { + piecewise_ran = true; + } else { + general_compiler->record_graph_miss(/*is_prefill=*/true); + } + } + + // FULL decode: FA in-graph under FA_INGRAPH=decode. + if (!logits && !piecewise_ran && uniform_decode) { + auto [graph, output] = compiler_->get_compiled(model_input_cpu); + if (graph != nullptr && output != nullptr) { + infinicore::context::InferencePhaseGuard phase_guard( + infinicore::context::InferencePhase::Decode); + graph->run(); + logits = output->logits; + hidden_states = output->hidden_states; + if (general_compiler != nullptr) { + general_compiler->record_graph_hit(/*is_prefill=*/false); + } + } else if (general_compiler != nullptr) { + general_compiler->record_graph_miss(/*is_prefill=*/false); + } } } // Fall back to eager mode - if (!logits) { + if (!logits && !piecewise_ran) { + auto &piecewise_eager = + global_state::get_forward_context().piecewise; + piecewise_eager.valid_seq_len = 0; + piecewise_eager.phase = + global_state::PiecewiseCapturePhase::None; + const bool looks_decode = local_args.block_tables.has_value() + && local_args.input_ids.has_value() + && local_args.block_tables.value()->size(0) + == local_args.input_ids.value()->size(1); + infinicore::context::InferencePhaseGuard phase_guard( + looks_decode ? infinicore::context::InferencePhase::Decode + : infinicore::context::InferencePhase::Prefill); auto model_args = local_args.to_model_input( rank_info_.device, attention_backend_ == backends::AttentionBackend::STATIC_ATTN); @@ -493,7 +543,11 @@ void RankWorker::thread_loop() { 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 + // QY batch_greedy (one random_sample_ on a flattened logits + // view) returns garbage ids when n_out>1 under V1 multi-request. + // Keep NVIDIA batch path; force QY through the per-index loop. + const bool batch_greedy = (rank_info_.device.type() == infinicore::Device::Type::kNvidia + || rank_info_.device.type() == infinicore::Device::Type::kQy) && parameter_greedy && n_out > 0 && logits_positions == n_out @@ -501,7 +555,9 @@ void RankWorker::thread_loop() { && (logits_dtype == infinicore::DataType::kFloat16 || logits_dtype == infinicore::DataType::kBFloat16 || logits_dtype == infinicore::DataType::kFloat32) - && (sample_all_positions || logits_are_last_token_only); + && (sample_all_positions || logits_are_last_token_only) + && !(rank_info_.device.type() == infinicore::Device::Type::kQy + && n_out > 1); if (batch_greedy) { float random_val = 0.0f; for (size_t i{0}; i < n_out; ++i) { diff --git a/csrc/global_state/ar_profile.hpp b/csrc/global_state/ar_profile.hpp new file mode 100644 index 000000000..53f4fea5a --- /dev/null +++ b/csrc/global_state/ar_profile.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "infinicore/tensor.hpp" + +#include +#include +#include +#include + +namespace infinilm::global_state::ar_profile { + +inline bool enabled() { return false; } + +struct Counters { + double decode_pre_barrier_ms{0}; + double decode_graph_run_ms{0}; + double decode_post_sync_ms{0}; + size_t decode_graph_steps{0}; +}; + +inline Counters &counters() { + static Counters c; + return c; +} + +inline void log_barrier_chunk_summary(const char *, + std::optional = std::nullopt, + size_t = 0) {} + +template +inline void allreduce_hidden_valid_contiguous(infinicore::Tensor &hidden, + size_t valid_len, + infinicore::Tensor & /*ar_staging*/, + Fn &&fn) { + (void)valid_len; + fn(hidden); +} + +} // namespace infinilm::global_state::ar_profile diff --git a/csrc/global_state/decode_phase_profile.hpp b/csrc/global_state/decode_phase_profile.hpp new file mode 100644 index 000000000..933023b1e --- /dev/null +++ b/csrc/global_state/decode_phase_profile.hpp @@ -0,0 +1,28 @@ +#pragma once + +namespace infinilm::global_state::decode_phase_profile { + +inline bool enabled() { return false; } +inline bool recording() { return false; } +inline double monotonic_ms() { return 0.0; } + +struct Counters { + double attn_ms{0}; + double moe_ms{0}; + double dense_mlp_ms{0}; + double other_layer_ms{0}; + double eager_forward_ms{0}; +}; + +inline Counters &counters() { + static Counters c; + return c; +} + +inline void reset() {} +inline bool &decode_step_active() { + static bool active = false; + return active; +} + +} // namespace infinilm::global_state::decode_phase_profile diff --git a/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index 1bc608424..c727fee1e 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -1,6 +1,7 @@ #pragma once #include "../models/infinilm_model.hpp" +#include "piecewise_prefill_state.hpp" namespace infinilm::global_state { @@ -74,6 +75,7 @@ struct ForwardContext { std::vector kv_cache_vec; std::vector conv_state_vec; std::vector ssm_state_vec; + PiecewisePrefillState piecewise; }; void initialize_forward_context(ForwardContext &forward_context); diff --git a/csrc/global_state/hang_trace.hpp b/csrc/global_state/hang_trace.hpp new file mode 100644 index 000000000..deba45098 --- /dev/null +++ b/csrc/global_state/hang_trace.hpp @@ -0,0 +1,13 @@ +#pragma once + +#include + +namespace infinilm::global_state::hang_trace { + +inline bool enabled() { return false; } + +struct ScopedBracket { + ScopedBracket(const char *, int = 0) {} +}; + +} // namespace infinilm::global_state::hang_trace diff --git a/csrc/global_state/piecewise_inductor_flags.hpp b/csrc/global_state/piecewise_inductor_flags.hpp new file mode 100644 index 000000000..2194f66cd --- /dev/null +++ b/csrc/global_state/piecewise_inductor_flags.hpp @@ -0,0 +1,96 @@ +#pragma once + +#include +#include + +namespace infinilm::global_state { + +/// When true, piecewise pre/post segments use AOTInductor kernels (M4 path). +inline bool piecewise_inductor_segment_enabled() { + const char *v = std::getenv("INFINI_PIECEWISE_INDUCTOR_SEGMENT"); + return v != nullptr && v[0] != '\0' && std::string(v) != "0"; +} + +/// Production default ON when inductor segment is enabled; set to 0 to bisect. +inline bool scoped_inductor_pre_attn_enabled() { + const char *v = std::getenv("INFINI_PIECEWISE_SCOPED_INDUCTOR"); + if (v == nullptr || v[0] == '\0') { + return piecewise_inductor_segment_enabled(); + } + return v[0] != '0' && std::string(v) != "false"; +} + +inline bool repro_skip_midchunk_eager() { + const char *v = std::getenv("INFINI_PIECEWISE_REPRO_SKIP_MIDCHUNK_EAGER"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; +} + +/// Part A / Gate D bisect: allow exact-width mid-chunk CG replay (final graphs). +/// Keeps allow_inductor_pre_attn=false for mid unless SKIP_MIDCHUNK_EAGER also set. +/// Do not enable in product defaults. +inline bool repro_allow_midchunk_cg() { + const char *v = std::getenv("INFINI_PIECEWISE_REPRO_ALLOW_MIDCHUNK_CG"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; +} + +/// MIXED mid CG bisect: force prefer final token-bucket PIECEWISE banks for MIXED +/// exact-width mid (same as ENABLE_MIXED_MID_CG). Microbench only — not a product default. +/// Legacy name kept; no longer means layout-matched ``compiled_mixed_mid_`` prefer. +inline bool repro_mixed_mid_cg() { + const char *v = std::getenv("INFINI_PIECEWISE_REPRO_MIXED_MID_CG"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; +} + +/// Product opt-in: MIXED exact-width mid may replay final PIECEWISE banks +/// (``compiled_[bucket]`` QKV+post, eager RoPE/attn/lm_head). Default OFF (kill-switch: +/// leave unset → MIXED mid stays eager). Does **not** mean layout-matched +/// ``compiled_mixed_mid_`` prefer — dual-capture is skipped unless KEEP_MIXED_MID_CAPTURE. +inline bool enable_mixed_mid_cg() { + const char *v = std::getenv("INFINI_PIECEWISE_ENABLE_MIXED_MID_CG"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; +} + +/// Legacy/bisect: still capture decode-first ``compiled_mixed_mid_`` banks (VRAM/time). +/// Product default skips this dual-capture; prefer uses final ``compiled_`` instead. +inline bool keep_mixed_mid_capture() { + const char *v = std::getenv("INFINI_PIECEWISE_KEEP_MIXED_MID_CAPTURE"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; +} + +/// Opt-in: skip host ``syncDevice()`` between QKV CG replay and eager RoPE. +/// Only relevant when ``PRE_ATTN_QKV_ONLY=1`` (Phase-2c split). Product fused +/// pre (LN+QKV+RoPE in one CG) does not need this fence. Default: sync ON. +inline bool skip_qkv_rope_sync() { + const char *v = std::getenv("INFINI_PIECEWISE_SKIP_QKV_ROPE_SYNC"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; +} + +/// Kill-switch: capture/replay Phase-2c QKV-only pre CG + host-eager RoPE. +/// Default OFF → product fused pre ``[CG: LN+QKV+RoPE]``. +inline bool pre_attn_qkv_only() { + const char *v = std::getenv("INFINI_PIECEWISE_PRE_ATTN_QKV_ONLY"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; +} + +/// Prefill pad-up CG (seq_len < graph_bucket). Default ON when unset; +/// ``INFINI_PIECEWISE_PAD_UP_CG=0`` restores legacy eager-only pad-up. +inline bool piecewise_pad_up_cg_enabled() { + const char *v = std::getenv("INFINI_PIECEWISE_PAD_UP_CG"); + if (v != nullptr && v[0] == '0' && v[1] == '\0') { + return false; + } + return true; +} + +inline bool repro_skip_final_inductor() { + const char *v = std::getenv("INFINI_PIECEWISE_REPRO_SKIP_FINAL_INDUCTOR"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; +} + +/// Tail bucket eligible for inductor pre_attn inside hcGraph (B4 only). +inline bool bucket_is_inductor_eligible(size_t bucket) { + constexpr size_t kInductorTailBucket = 4; + return bucket == kInductorTailBucket; +} + +} // namespace infinilm::global_state diff --git a/csrc/global_state/piecewise_prefill_state.hpp b/csrc/global_state/piecewise_prefill_state.hpp new file mode 100644 index 000000000..b1caa6cb6 --- /dev/null +++ b/csrc/global_state/piecewise_prefill_state.hpp @@ -0,0 +1,45 @@ +#pragma once + +#include "infinicore/tensor.hpp" + +#include +#include + +namespace infinilm::global_state { + +enum class PiecewiseCapturePhase { + None, + PreAttn, + EagerAttn, + PostAttn, + LmHead, +}; + +/// Per-layer staging tensors shared between captured pre-attn graphs and eager attn. +struct PiecewiseLayerStaging { + infinicore::Tensor q_rope; + infinicore::Tensor k_rope; + infinicore::Tensor v_rope; + infinicore::Tensor attn_output; +}; + +/// Thread-local piecewise prefill state (capture + replay). +struct PiecewisePrefillState { + PiecewiseCapturePhase phase{PiecewiseCapturePhase::None}; + size_t active_layer{0}; + /// Actual token count for this request (may be < bucket size). + size_t valid_seq_len{0}; + /// Padded bucket length for the active replay. + size_t bucket_seq_len{0}; + infinicore::Tensor hidden_states; + infinicore::Tensor residual; + /// Shared contiguous buffer for sequential per-layer allreduce (one buffer per bucket). + infinicore::Tensor ar_staging; + std::vector layer_staging; + /// True while native piecewise CG bucket capture is running (compile path). + bool compile_capture_active{false}; + /// Per-replay gate: when false, piecewise_pre_attn uses infiniop/CG even if inductor env is on. + bool allow_inductor_pre_attn{true}; +}; + +} // namespace infinilm::global_state diff --git a/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp b/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp index f48f640cd..107d38e26 100644 --- a/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp +++ b/csrc/infinicore/include/infinicore/adaptor/aten_adaptor.hpp @@ -8,7 +8,7 @@ #if defined(ENABLE_HYGON_API) #include #include -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) #include #include #include @@ -41,7 +41,9 @@ inline at::ScalarType to_at_dtype(DataType dtype) { inline at::Device to_at_device(const Device &device) { // PyTorch ATen only exposes standard device types (e.g. kCPU/kCUDA). // Treat CUDA-compatible devices as CUDA devices for ATen interoperability. - if (device.type() == Device::Type::kNvidia || device.type() == Device::Type::kMetax || device.type() == Device::Type::kHygon) { + if (device.type() == Device::Type::kNvidia || device.type() == Device::Type::kMetax + || device.type() == Device::Type::kHygon || device.type() == Device::Type::kQy + || device.type() == Device::Type::kThead || device.type() == Device::Type::kIluvatar) { return at::Device(at::kCUDA, device.index()); } else if (device.type() == Device::Type::kCpu) { return at::Device(at::kCPU); @@ -60,7 +62,7 @@ at::Tensor to_aten_tensor(const infinicore::Tensor &t); #if defined(ENABLE_HYGON_API) c10::hip::HIPStream get_hip_stream(); -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) +#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) c10::cuda::CUDAStream get_cuda_stream(); #endif @@ -68,6 +70,9 @@ c10::cuda::CUDAStream get_cuda_stream(); c10::musa::MUSAStream get_musa_stream(); #endif +/** Align ATen/CUDA current stream with InfiniCore's active stream. */ +void set_aten_stream_to_infinicore(); + } // namespace infinicore::adaptor #endif // ENABLE_ATEN diff --git a/csrc/infinicore/include/infinicore/adaptor/denglin_fused_moe_adaptor.hpp b/csrc/infinicore/include/infinicore/adaptor/denglin_fused_moe_adaptor.hpp new file mode 100644 index 000000000..a6cb99da8 --- /dev/null +++ b/csrc/infinicore/include/infinicore/adaptor/denglin_fused_moe_adaptor.hpp @@ -0,0 +1,43 @@ +#pragma once + +#if defined(ENABLE_ATEN) && defined(ENABLE_QY_API) +#include +#include +#include +#include + +namespace infinicore::adaptor::denglin_fused_moe { + +/** True when vllm._dl_C.invoke_fused_moe_opt is loadable (same backend as dl-vllm). */ +bool available(); + +/** + * Thin wrapper around Denglin/vLLM ``invoke_fused_moe_opt`` (torch.ops._dl_C). + * Matches vLLM ``dl_fused_moe.dl_invoke_fused_moe`` calling convention. + */ +void invoke_fused_moe_opt( + at::Tensor &a, + at::Tensor &b, + at::Tensor &c, + const std::optional &b_bias, + const std::optional &b_scale, + const std::optional &b_zp, + at::Tensor &topk_weights, + at::Tensor &topk_ids, + at::Tensor &sorted_token_ids, + at::Tensor &expert_ids, + at::Tensor &num_tokens_post_padded, + bool mul_routed_weight, + int64_t top_k, + int64_t block_size_m, + int64_t block_size_n, + int64_t block_size_k, + bool use_fp8_w8a8, + bool use_int8_w8a16, + bool use_int4_w4a16, + bool use_mxfp4_w4a16, + const std::vector &block_shape, + int64_t m); + +} // namespace infinicore::adaptor::denglin_fused_moe +#endif diff --git a/csrc/infinicore/include/infinicore/context/context.hpp b/csrc/infinicore/include/infinicore/context/context.hpp index 33e92e4db..250dd6be6 100644 --- a/csrc/infinicore/include/infinicore/context/context.hpp +++ b/csrc/infinicore/include/infinicore/context/context.hpp @@ -7,6 +7,7 @@ #include +#include #include namespace infinicore { @@ -51,6 +52,55 @@ void addGraphOperator(std::shared_ptr op); std::shared_ptr stopGraphRecording(); void cancelGraphRecording() noexcept; +/// ``INFINI_CUDAGRAPH_POLICY``: empty (legacy), ``eager``, or ``full_and_piecewise``. +/// Unknown values are treated as empty / legacy. +const char *cudagraphPolicy(); + +/// True when device-graph segment capture/replay should run. +/// ``eager`` → false; ``full_and_piecewise`` → true unless +/// ``INFINICORE_DISABLE_DEVICE_GRAPH_SEGMENTS`` is set; empty/legacy honors +/// that kill-switch when present, otherwise enables segments. +bool deviceGraphSegmentsEnabled(); + +/// Engine TLS inference phase for phase-scoped MoE / FA capture policy. +/// FA uses ``INFINI_FA_INGRAPH``; MoE uses ``INFINI_MOE_INGRAPH``. +enum class InferencePhase : uint8_t { + Unknown = 0, + Prefill = 1, + Decode = 2, +}; + +void setInferencePhase(InferencePhase phase); +InferencePhase getInferencePhase(); + +/// FA may enter a device segment — phase allow-list via ``INFINI_FA_INGRAPH`` +/// (``off`` | ``decode`` | ``both``; truthy ``1`` → ``both``). +/// Product ``full_and_piecewise`` defaults to ``decode`` when unset +/// (FULL decode segs≈1; prefill FA stays host-break). +/// Legacy ``INFINI_FA_FORCE_CAPTURE`` → ``both``. Kill-switch: ``0`` / ``off``. +bool faInGraphAllowed(); + +/// Triton MoE under stream capture — phase allow-list via ``INFINI_MOE_INGRAPH`` +/// (``off`` | ``decode`` | ``prefill`` | ``both``; default ``off``). +/// ``INFINI_MOE_FORCE_HOST_BREAK=1`` → always host-break. +/// ``eager`` policy → host-break. +bool moeTritonCaptureAllowed(); + +/// RAII restore of ``InferencePhase`` for graph capture / forward scopes. +class InferencePhaseGuard { +public: + explicit InferencePhaseGuard(InferencePhase phase) + : prev_(getInferencePhase()) { + setInferencePhase(phase); + } + ~InferencePhaseGuard() { setInferencePhase(prev_); } + InferencePhaseGuard(const InferencePhaseGuard &) = delete; + InferencePhaseGuard &operator=(const InferencePhaseGuard &) = delete; + +private: + InferencePhase prev_; +}; + } // namespace context } // namespace infinicore diff --git a/csrc/infinicore/include/infinicore/graph/graph.hpp b/csrc/infinicore/include/infinicore/graph/graph.hpp index b06249ed2..cb458fa91 100644 --- a/csrc/infinicore/include/infinicore/graph/graph.hpp +++ b/csrc/infinicore/include/infinicore/graph/graph.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -57,11 +58,29 @@ class Graph { std::shared_ptr allocation_lease_; public: + /// Split of the most recent ``run()`` wall (ms). Zeroed when profile off. + struct LastRunProfile { + double launch_ms{0.0}; ///< sum of device GraphLaunch walls + double tail_sync_ms{0.0}; ///< end syncStream after device launch + double hostop_ms{0.0}; ///< capture-unsafe / host segments + double op_list_ms{0.0}; ///< eager op-list / fallback segments + bool used_device{false}; + bool skipped_tail_sync{false}; + }; + Graph(); ~Graph() noexcept; void run() const; + const LastRunProfile &last_run_profile() const { return last_run_profile_; } + + /// Number of capture-safe segments that own a device graph exec. + size_t device_segment_count() const; + + /// Bytes retained by capture arena (0 when arena accounting unavailable). + size_t capture_arena_bytes() const { return 0; } + protected: void add_operator(std::shared_ptr op); void instantiate(); @@ -79,6 +98,7 @@ class Graph { ~Segment() noexcept; void run() const; + bool has_device_exec() const; bool capture_safe; std::vector> ops; @@ -86,6 +106,8 @@ class Graph { }; std::vector> segments_; + mutable LastRunProfile last_run_profile_{}; + friend class ::infinicore::Runtime; }; } // namespace infinicore::graph @@ -102,6 +124,23 @@ class Graph { static void execute(__VA_ARGS__); \ }; +// Same as INFINICORE_GRAPH_OP_CLASS, but keeps the op in eager segments during +// CUDA-graph capture (host D2H/H2D, ATen temps, or non-capturable kernels). +#define INFINICORE_GRAPH_OP_CLASS_CAPTURE_UNSAFE(__OP_NAME__, ...) \ + class __OP_NAME__ : public graph::DispatchableGraphOperator { \ + public: \ + using schema = void (*)(__VA_ARGS__); \ + using plan_schema = void *(*)(__VA_ARGS__); \ + static common::OpDispatcher &plan_dispatcher(); \ + static common::OpDispatcher &run_dispatcher(); \ + static common::OpDispatcher &cleanup_dispatcher(); \ + __OP_NAME__(__VA_ARGS__); \ + static void execute(__VA_ARGS__); \ + bool is_device_graph_capture_safe() const override { \ + return false; \ + } \ + }; + #define INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(__OP_NAME__) \ common::OpDispatcher<__OP_NAME__::plan_schema> &__OP_NAME__::plan_dispatcher() { \ static common::OpDispatcher<__OP_NAME__::plan_schema> dispatcher_; \ @@ -116,15 +155,21 @@ class Graph { return dispatcher_; \ } -#define INFINICORE_GRAPH_OP_DISPATCH(__DEVICE_TYPE__, ...) \ - do { \ - const auto device_type = (__DEVICE_TYPE__); \ - const auto plan = plan_dispatcher().lookup(device_type); \ - const auto runner = run_dispatcher().lookup(device_type); \ - const auto deleter = cleanup_dispatcher().lookup(device_type); \ - planned_meta_ = plan(__VA_ARGS__); \ - runner_ = runner; \ - deleter_ = deleter; \ +#define INFINICORE_GRAPH_OP_DISPATCH(__DEVICE_TYPE__, ...) \ + do { \ + const auto device_type = (__DEVICE_TYPE__); \ + try { \ + const auto plan = plan_dispatcher().lookup(device_type); \ + const auto runner = run_dispatcher().lookup(device_type); \ + const auto deleter = cleanup_dispatcher().lookup(device_type); \ + planned_meta_ = plan(__VA_ARGS__); \ + runner_ = runner; \ + deleter_ = deleter; \ + } catch (const std::exception &e) { \ + throw std::runtime_error(std::string(e.what()) + " [" \ + + std::string(__PRETTY_FUNCTION__) \ + + "]"); \ + } \ } while (0) #define INFINICORE_DETAIL_FIRST_ARG(__FIRST__, ...) __FIRST__ diff --git a/csrc/infinicore/include/infinicore/ops/inductor_segment.hpp b/csrc/infinicore/include/infinicore/ops/inductor_segment.hpp new file mode 100644 index 000000000..c06d6ce66 --- /dev/null +++ b/csrc/infinicore/include/infinicore/ops/inductor_segment.hpp @@ -0,0 +1,51 @@ +#pragma once + +#include "infinicore/tensor.hpp" + +#include +#include + +namespace infinicore::op { + +enum class PiecewiseInductorSegmentId : int { + PreAttn = 0, + PostAttn = 1, + Moe = 2, +}; + +namespace inductor_segment_impl { + +struct PreAttnExternalWeightTensors { + Tensor ln_weight; + Tensor qkv_weight; + Tensor qkv_bias; +}; + +struct MoeExternalWeightTensors { + Tensor gate_weight; + Tensor expert_weight; +}; + +inline bool has_package(PiecewiseInductorSegmentId, size_t, size_t) { + return false; +} + +inline void run_pre_attn(size_t, + size_t, + const Tensor &, + const Tensor &, + const Tensor &, + const PreAttnExternalWeightTensors &) { + throw std::runtime_error("inductor_segment: AOT packages not available on this build"); +} + +} // namespace inductor_segment_impl + +inline void inductor_warmup_pre_attn_bucket(size_t, + size_t, + size_t, + const inductor_segment_impl::PreAttnExternalWeightTensors &) { + // no-op stub (MetaX AOT trimmed for QY) +} + +} // namespace infinicore::op diff --git a/csrc/infinicore/include/infinicore/ops/moe_align.hpp b/csrc/infinicore/include/infinicore/ops/moe_align.hpp index 2ce7ca8f4..e50c89754 100644 --- a/csrc/infinicore/include/infinicore/ops/moe_align.hpp +++ b/csrc/infinicore/include/infinicore/ops/moe_align.hpp @@ -9,6 +9,7 @@ namespace infinicore::op { +// Capture-safe for QY full_and_piecewise; ATen path skips sync under capture. INFINICORE_GRAPH_OP_CLASS(MoeAlign, Tensor, Tensor, Tensor, const Tensor &, const size_t, const size_t, const bool); INFINICORE_GRAPH_OP_CLASS(MoeAlignWithExpertMap, Tensor, Tensor, Tensor, const Tensor &, const Tensor &, const size_t, const size_t, const bool); diff --git a/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp b/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp index 0a34c734c..3346ad262 100644 --- a/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp +++ b/csrc/infinicore/include/infinicore/ops/moe_fused_dense.hpp @@ -7,6 +7,8 @@ namespace infinicore::op { +// Capture-safe for QY full_and_piecewise (vLLM / pre-refactor InfiniLM proven). +// Impls skip host sync while the stream is capturing (see sync_unless_capturing). INFINICORE_GRAPH_OP_CLASS(MoeFusedDense, Tensor, const Tensor &, diff --git a/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp b/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp index e796c1dc8..ff9ad290b 100644 --- a/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp +++ b/csrc/infinicore/include/infinicore/ops/moe_topk_sigmoid.hpp @@ -9,6 +9,7 @@ namespace infinicore::op { +// Capture-safe for QY full_and_piecewise; ATen path skips sync under capture. INFINICORE_GRAPH_OP_CLASS(MoeTopkSigmoid, Tensor, Tensor, diff --git a/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp b/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp index 5ee5e0d0b..7a312291a 100644 --- a/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp +++ b/csrc/infinicore/include/infinicore/ops/mul_scalar.hpp @@ -6,6 +6,7 @@ namespace infinicore::op { +// MiniCPM5 MoE routing-weight scale; capture-safe under full_and_piecewise. INFINICORE_GRAPH_OP_CLASS(MulScalar, Tensor, const Tensor &, double); Tensor mul_scalar(const Tensor &a, double alpha); diff --git a/csrc/infinicore/include/infinicore/ops/rope.hpp b/csrc/infinicore/include/infinicore/ops/rope.hpp index 8fd630ce1..d5411edfb 100644 --- a/csrc/infinicore/include/infinicore/ops/rope.hpp +++ b/csrc/infinicore/include/infinicore/ops/rope.hpp @@ -8,6 +8,7 @@ namespace infinicore::op { +// Capture-safe for QY full_and_piecewise; pack path skips sync under capture. INFINICORE_GRAPH_OP_CLASS(RoPE, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, infinicore::nn::RoPE::Algo); // Internal diff --git a/csrc/infinicore/src/adaptor/aten_adaptor.cc b/csrc/infinicore/src/adaptor/aten_adaptor.cc index 87eb033f4..55b3a52d2 100644 --- a/csrc/infinicore/src/adaptor/aten_adaptor.cc +++ b/csrc/infinicore/src/adaptor/aten_adaptor.cc @@ -1,6 +1,8 @@ #ifdef ENABLE_ATEN #include "infinicore/adaptor/aten_adaptor.hpp" +#include + namespace infinicore::adaptor { at::Tensor to_aten_tensor(const infinicore::Tensor &t) { @@ -51,6 +53,14 @@ c10::musa::MUSAStream get_musa_stream() { } #endif +void set_aten_stream_to_infinicore() { +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) + c10::cuda::setCurrentCUDAStream(get_cuda_stream()); +#else + throw std::runtime_error("ATen stream synchronization is unsupported on this device"); +#endif +} + } // namespace infinicore::adaptor #endif // ENABLE_ATEN diff --git a/csrc/infinicore/src/adaptor/denglin_fused_moe_adaptor.cc b/csrc/infinicore/src/adaptor/denglin_fused_moe_adaptor.cc new file mode 100644 index 000000000..d7798072f --- /dev/null +++ b/csrc/infinicore/src/adaptor/denglin_fused_moe_adaptor.cc @@ -0,0 +1,132 @@ +#if defined(ENABLE_ATEN) && defined(ENABLE_QY_API) +#include "infinicore/adaptor/denglin_fused_moe_adaptor.hpp" +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include +#include +#include +#include +#include +#include + +namespace infinicore::adaptor::denglin_fused_moe { +namespace { + +// Mangled name of free function in vllm._dl_C: +// invoke_fused_moe_opt(at::Tensor&, at::Tensor&, at::Tensor&, +// optional const& x3, at::Tensor& x5, bool, long x4, bool x4, +// vector const&, long) +using invoke_fused_moe_opt_fn = void (*)( + at::Tensor &, + at::Tensor &, + at::Tensor &, + const std::optional &, + const std::optional &, + const std::optional &, + at::Tensor &, + at::Tensor &, + at::Tensor &, + at::Tensor &, + at::Tensor &, + bool, + int64_t, + int64_t, + int64_t, + int64_t, + bool, + bool, + bool, + bool, + const std::vector &, + int64_t); + +struct Symbols { + void *handle = nullptr; + invoke_fused_moe_opt_fn invoke = nullptr; + std::string error; +}; + +Symbols &symbols() { + static Symbols syms; + static std::once_flag once; + std::call_once(once, []() { + const char *env_path = std::getenv("INFINICORE_DENGLIN_DL_C_SO"); + // Prefer exact override, then common dl-vLLM image layouts (3.12/3.10). + const char *paths[] = { + env_path, + "/usr/local/lib/python3.12/dist-packages/vllm/_dl_C.cpython-312-x86_64-linux-gnu.so", + "/usr/local/lib/python3.10/dist-packages/vllm/_dl_C.cpython-310-x86_64-linux-gnu.so", + "/usr/lib/python3/dist-packages/vllm/_dl_C.cpython-312-x86_64-linux-gnu.so", + }; + constexpr const char *kSym = "_Z20invoke_fused_moe_optRN2at6TensorES1_S1_RKSt8optionalIS0_ES5_S5_" + "S1_S1_S1_S1_S1_bllllbbbbRKSt6vectorIlSaIlEEl"; + + for (const char *path : paths) { + if (path == nullptr || path[0] == '\0') { + continue; + } + void *handle = dlopen(path, RTLD_LAZY | RTLD_LOCAL); + if (!handle) { + syms.error = dlerror() ? dlerror() : "dlopen failed"; + continue; + } + auto fn = reinterpret_cast(dlsym(handle, kSym)); + if (!fn) { + syms.error = dlerror() ? dlerror() : "dlsym invoke_fused_moe_opt failed"; + dlclose(handle); + continue; + } + syms.handle = handle; + syms.invoke = fn; + syms.error.clear(); + return; + } + }); + return syms; +} + +} // namespace + +bool available() { + return symbols().invoke != nullptr; +} + +void invoke_fused_moe_opt( + at::Tensor &a, + at::Tensor &b, + at::Tensor &c, + const std::optional &b_bias, + const std::optional &b_scale, + const std::optional &b_zp, + at::Tensor &topk_weights, + at::Tensor &topk_ids, + at::Tensor &sorted_token_ids, + at::Tensor &expert_ids, + at::Tensor &num_tokens_post_padded, + bool mul_routed_weight, + int64_t top_k, + int64_t block_size_m, + int64_t block_size_n, + int64_t block_size_k, + bool use_fp8_w8a8, + bool use_int8_w8a16, + bool use_int4_w4a16, + bool use_mxfp4_w4a16, + const std::vector &block_shape, + int64_t m) { + auto &syms = symbols(); + if (!syms.invoke) { + throw std::runtime_error( + "Denglin fused MoE unavailable (vllm._dl_C.invoke_fused_moe_opt): " + syms.error); + } + infinicore::adaptor::set_aten_stream_to_infinicore(); + syms.invoke( + a, b, c, b_bias, b_scale, b_zp, + topk_weights, topk_ids, sorted_token_ids, expert_ids, num_tokens_post_padded, + mul_routed_weight, top_k, block_size_m, block_size_n, block_size_k, + use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16, use_mxfp4_w4a16, + block_shape, m); +} + +} // namespace infinicore::adaptor::denglin_fused_moe +#endif diff --git a/csrc/infinicore/src/context/context_impl.cc b/csrc/infinicore/src/context/context_impl.cc index a243eb58f..b9e404fca 100644 --- a/csrc/infinicore/src/context/context_impl.cc +++ b/csrc/infinicore/src/context/context_impl.cc @@ -4,8 +4,11 @@ #include "../utils.hpp" #include +#include +#include #include #include +#include #include namespace infinicore { @@ -20,6 +23,7 @@ constexpr std::array(Device::Type::kCount)> kD Device::Type::kIluvatar, Device::Type::kHygon, Device::Type::kThead, + Device::Type::kQy, Device::Type::kCpu, }; @@ -237,6 +241,7 @@ ContextImpl::ContextImpl() { initializeDeviceType(); initializeDeviceType(); initializeDeviceType(); + initializeDeviceType(); const bool has_runtime = std::any_of( runtime_table_.begin(), runtime_table_.end(), @@ -365,6 +370,180 @@ void cancelGraphRecording() noexcept { ContextImpl::singleton().cancelGraphRecording(); } +const char *cudagraphPolicy() { + const char *v = std::getenv("INFINI_CUDAGRAPH_POLICY"); + if (v == nullptr || v[0] == '\0') { + return ""; + } + const std::string s(v); + if (s == "eager") { + return "eager"; + } + if (s == "full_and_piecewise") { + return "full_and_piecewise"; + } + // Reject unknown values — fall back to legacy (empty). + return ""; +} + +bool deviceGraphSegmentsEnabled() { + // Kill-switch always wins (diagnostic escape hatch). + if (std::getenv("INFINICORE_DISABLE_DEVICE_GRAPH_SEGMENTS") != nullptr) { + return false; + } + const char *policy = cudagraphPolicy(); + if (std::strcmp(policy, "eager") == 0) { + return false; + } + // full_and_piecewise and empty/legacy → enable segmented device graphs. + return true; +} + +namespace { +thread_local InferencePhase g_inference_phase = InferencePhase::Unknown; + +bool env_truthy_(const char *name) { + const char *v = std::getenv(name); + return v != nullptr && v[0] != '\0' && std::string(v) != "0" + && std::strcmp(v, "false") != 0 && std::strcmp(v, "no") != 0 + && std::strcmp(v, "off") != 0; +} + +void warn_once_(bool &flag, const char *msg) { + if (flag) { + return; + } + flag = true; + spdlog::warn("{}", msg); +} + +enum class MoeIngraphMode : uint8_t { Off = 0, Decode, Prefill, Both }; + +MoeIngraphMode moeIngraphMode_() { + const char *raw = std::getenv("INFINI_MOE_INGRAPH"); + if (raw != nullptr && raw[0] != '\0') { + std::string s(raw); + for (char &c : s) { + if (c >= 'A' && c <= 'Z') { + c = static_cast(c - 'A' + 'a'); + } + } + if (s == "off" || s == "0" || s == "false" || s == "no") { + return MoeIngraphMode::Off; + } + if (s == "decode") { + return MoeIngraphMode::Decode; + } + if (s == "prefill") { + return MoeIngraphMode::Prefill; + } + if (s == "both") { + return MoeIngraphMode::Both; + } + static bool unknown_warned = false; + warn_once_(unknown_warned, + "INFINI_MOE_INGRAPH has unknown value (expected " + "off|decode|prefill|both); treating as off"); + return MoeIngraphMode::Off; + } + if (env_truthy_("INFINI_MOE_FORCE_CAPTURE") + || env_truthy_("INFINI_MOE_TRITON_CAPTURE")) { + static bool legacy_warned = false; + warn_once_( + legacy_warned, + "INFINI_MOE_FORCE_CAPTURE / INFINI_MOE_TRITON_CAPTURE are deprecated; " + "shim to INFINI_MOE_INGRAPH=decode"); + return MoeIngraphMode::Decode; + } + return MoeIngraphMode::Off; +} + +enum class FaIngraphMode : uint8_t { Off = 0, Decode, Both }; + +FaIngraphMode faIngraphMode_() { + const char *raw = std::getenv("INFINI_FA_INGRAPH"); + if (raw != nullptr && raw[0] != '\0') { + std::string s(raw); + for (char &c : s) { + if (c >= 'A' && c <= 'Z') { + c = static_cast(c - 'A' + 'a'); + } + } + if (s == "off" || s == "0" || s == "false" || s == "no") { + return FaIngraphMode::Off; + } + if (s == "decode") { + return FaIngraphMode::Decode; + } + if (s == "both" || s == "1" || s == "true" || s == "yes" || s == "on") { + return FaIngraphMode::Both; + } + static bool unknown_warned = false; + warn_once_(unknown_warned, + "INFINI_FA_INGRAPH has unknown value (expected " + "off|decode|both); treating as off"); + return FaIngraphMode::Off; + } + if (env_truthy_("INFINI_FA_FORCE_CAPTURE")) { + static bool warned = false; + warn_once_( + warned, + "INFINI_FA_FORCE_CAPTURE is deprecated; use " + "INFINI_FA_INGRAPH=decode (product) or both/1 (Prefill+Decode)"); + return FaIngraphMode::Both; + } + // Product full_and_piecewise defaults to decode (FULL segs≈1; prefill FA HB). + if (std::strcmp(cudagraphPolicy(), "full_and_piecewise") == 0) { + return FaIngraphMode::Decode; + } + return FaIngraphMode::Off; +} +} // namespace + +void setInferencePhase(InferencePhase phase) { + g_inference_phase = phase; +} + +InferencePhase getInferencePhase() { + return g_inference_phase; +} + +bool faInGraphAllowed() { + const FaIngraphMode mode = faIngraphMode_(); + const InferencePhase phase = getInferencePhase(); + switch (mode) { + case FaIngraphMode::Off: + return false; + case FaIngraphMode::Decode: + return phase == InferencePhase::Decode; + case FaIngraphMode::Both: + return phase == InferencePhase::Prefill || phase == InferencePhase::Decode; + } + return false; +} + +bool moeTritonCaptureAllowed() { + if (env_truthy_("INFINI_MOE_FORCE_HOST_BREAK")) { + return false; + } + if (std::strcmp(cudagraphPolicy(), "eager") == 0) { + return false; + } + const MoeIngraphMode mode = moeIngraphMode_(); + const InferencePhase phase = getInferencePhase(); + switch (mode) { + case MoeIngraphMode::Off: + return false; + case MoeIngraphMode::Decode: + return phase == InferencePhase::Decode; + case MoeIngraphMode::Prefill: + return phase == InferencePhase::Prefill; + case MoeIngraphMode::Both: + return phase == InferencePhase::Prefill || phase == InferencePhase::Decode; + } + return false; +} + std::shared_ptr reinstantiateBlob(std::shared_ptr blob) { setDevice(blob->device()); return ContextImpl::singleton().getCurrentRuntime()->reinstantiateBlob(blob); diff --git a/csrc/infinicore/src/graph/graph.cc b/csrc/infinicore/src/graph/graph.cc index 423cdf8c1..855c0b454 100644 --- a/csrc/infinicore/src/graph/graph.cc +++ b/csrc/infinicore/src/graph/graph.cc @@ -8,11 +8,36 @@ #include #endif +#include #include +#include #include +#include #include namespace infinicore::graph { +namespace { + +bool graph_run_profile_enabled() { + static int cached = -1; + if (cached < 0) { + const char *g = std::getenv("INFINI_GRAPH_RUN_PROFILE"); + cached = (g != nullptr && g[0] != '\0' && std::string(g) != "0") ? 1 : 0; + } + return cached == 1; +} + +bool graph_skip_tail_sync_enabled() { + const char *v = std::getenv("INFINI_GRAPH_SKIP_TAIL_SYNC"); + return v != nullptr && v[0] != '\0' && std::string(v) != "0"; +} + +double graph_monotonic_ms() { + using clock = std::chrono::steady_clock; + return std::chrono::duration(clock::now().time_since_epoch()).count(); +} + +} // namespace #ifdef USE_INFINIRT_GRAPH namespace rt_runtime = ::infini::rt::runtime; @@ -206,9 +231,17 @@ Graph::Segment::Segment(bool capture_safe_) : capture_safe(capture_safe_) { Graph::Segment::~Segment() noexcept = default; +bool Graph::Segment::has_device_exec() const { +#ifdef USE_INFINIRT_GRAPH + return device_graph_ != nullptr && device_graph_->exec != nullptr; +#else + return false; +#endif +} + void Graph::Segment::run() const { #ifdef USE_INFINIRT_GRAPH - if (device_graph_ != nullptr && device_graph_->exec != nullptr) { + if (has_device_exec()) { device_graph_->launch(); return; } @@ -221,6 +254,16 @@ void Graph::Segment::run() const { Graph::Graph() { } +size_t Graph::device_segment_count() const { + size_t n = 0; + for (const auto &segment : segments_) { + if (segment && segment->capture_safe && segment->has_device_exec()) { + ++n; + } + } + return n; +} + void Graph::retain_runtime(std::shared_ptr<::infinicore::Runtime> runtime, std::shared_ptr allocation_lease) { runtime_lease_ = std::move(runtime); @@ -229,14 +272,84 @@ void Graph::retain_runtime(std::shared_ptr<::infinicore::Runtime> runtime, void Graph::run() const { (void)context::isGraphRecording(); + last_run_profile_ = LastRunProfile{}; + const bool profile = graph_run_profile_enabled(); + if (segments_.empty()) { + const double t0 = profile ? graph_monotonic_ms() : 0.0; for (const auto &op : op_list_) { op->run(); } + if (profile) { + last_run_profile_.op_list_ms = graph_monotonic_ms() - t0; + } + if (profile) { + spdlog::info( + "graph_run_profile: launch_ms={:.3f} tail_sync_ms={:.3f} hostop_ms={:.3f} " + "op_list_ms={:.3f} used_device={} skip_tail={}", + last_run_profile_.launch_ms, + last_run_profile_.tail_sync_ms, + last_run_profile_.hostop_ms, + last_run_profile_.op_list_ms, + 0, + 0); + } return; } + + bool used_device = false; + bool device_launched_since_sync = false; for (const auto &segment : segments_) { + if (segment->has_device_exec()) { + const double t_l0 = profile ? graph_monotonic_ms() : 0.0; + segment->run(); + if (profile) { + last_run_profile_.launch_ms += graph_monotonic_ms() - t_l0; + } + used_device = true; + device_launched_since_sync = true; + continue; + } + + // Host / eager segments must observe prior device-segment outputs. + if (device_launched_since_sync) { + infinicore::context::syncStream(); + device_launched_since_sync = false; + } + const double t0 = profile ? graph_monotonic_ms() : 0.0; segment->run(); + if (profile) { + if (segment->capture_safe) { + last_run_profile_.op_list_ms += graph_monotonic_ms() - t0; + } else { + last_run_profile_.hostop_ms += graph_monotonic_ms() - t0; + } + } + } + + if (used_device && device_launched_since_sync) { + if (graph_skip_tail_sync_enabled()) { + last_run_profile_.skipped_tail_sync = true; + } else { + const double t_s0 = profile ? graph_monotonic_ms() : 0.0; + infinicore::context::syncStream(); + if (profile) { + last_run_profile_.tail_sync_ms += graph_monotonic_ms() - t_s0; + } + } + } + + last_run_profile_.used_device = used_device; + if (profile) { + spdlog::info( + "graph_run_profile: launch_ms={:.3f} tail_sync_ms={:.3f} hostop_ms={:.3f} " + "op_list_ms={:.3f} used_device={} skip_tail={}", + last_run_profile_.launch_ms, + last_run_profile_.tail_sync_ms, + last_run_profile_.hostop_ms, + last_run_profile_.op_list_ms, + used_device ? 1 : 0, + last_run_profile_.skipped_tail_sync ? 1 : 0); } } @@ -254,12 +367,27 @@ void Graph::instantiate() { } infinicore::context::syncStream(); - // Keep recorded-operator replay available as a diagnostic escape hatch. - if (std::getenv("INFINICORE_DISABLE_DEVICE_GRAPH_SEGMENTS") != nullptr) { - spdlog::info("device graph segments disabled; replaying recorded operators"); + // Policy: eager / DISABLE kill-switch → recorded-operator replay only. + if (!context::deviceGraphSegmentsEnabled()) { + const char *policy = context::cudagraphPolicy(); + if (std::strcmp(policy, "eager") == 0) { + spdlog::info( + "INFINI_CUDAGRAPH_POLICY=eager; replaying recorded operators without device segments"); + } else { + spdlog::info("device graph segments disabled; replaying recorded operators"); + } return; } + static bool policy_logged = false; + if (!policy_logged) { + policy_logged = true; + const char *policy = context::cudagraphPolicy(); + spdlog::info( + "INFINI_CUDAGRAPH_POLICY={} (device graph segments enabled)", + (policy[0] == '\0') ? "legacy" : policy); + } + for (const auto &op : op_list_) { const bool capture_safe = op->is_device_graph_capture_safe(); if (segments_.empty() diff --git a/csrc/infinicore/src/nn/rope.cc b/csrc/infinicore/src/nn/rope.cc index aced91856..9b5785a91 100644 --- a/csrc/infinicore/src/nn/rope.cc +++ b/csrc/infinicore/src/nn/rope.cc @@ -78,6 +78,8 @@ void RoPE::initialize_cache() { INFINICORE_NN_BUFFER_INIT(cos_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); #ifdef ENABLE_INFINIOPS_API + // QY: prefer sin/cos + ATen RoPE. InfiniOps RotaryEmbedding no-ops on 4D + // [B,S,H,D] (MiniCPM5 static) and 3D reshape still mismatches HF GPT-NeoX. if ((device_.type() == Device::Type::kNvidia || device_.type() == Device::Type::kMetax || device_.type() == Device::Type::kIluvatar diff --git a/csrc/infinicore/src/ops/cat/cat.cc b/csrc/infinicore/src/ops/cat/cat.cc index c19e0896b..dfe54d859 100644 --- a/csrc/infinicore/src/ops/cat/cat.cc +++ b/csrc/infinicore/src/ops/cat/cat.cc @@ -14,7 +14,8 @@ bool use_slice_copy_cat(Device::Type device_type, int dim, int ndim) { && (device_type == Device::Type::kNvidia || device_type == Device::Type::kHygon || device_type == Device::Type::kIluvatar - || device_type == Device::Type::kThead); + || device_type == Device::Type::kThead + || device_type == Device::Type::kQy); } class CatInfo { diff --git a/csrc/infinicore/src/ops/causal_softmax/causal_softmax_aten.cc b/csrc/infinicore/src/ops/causal_softmax/causal_softmax_aten.cc new file mode 100644 index 000000000..6bdc3af03 --- /dev/null +++ b/csrc/infinicore/src/ops/causal_softmax/causal_softmax_aten.cc @@ -0,0 +1,77 @@ +#include "infinicore/ops/causal_softmax.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include + +#include +#include + +namespace infinicore::op::causal_softmax_impl::aten { +namespace { + +struct PlannedMeta { + Tensor output; + Tensor input; +}; + +void *plan(Tensor output, const Tensor &input) { + INFINICORE_ASSERT(output->device().type() == Device::Type::kQy); + return new PlannedMeta{output, input}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + using namespace ::infinicore::adaptor; + set_aten_stream_to_infinicore(); + + auto input = to_aten_tensor(p->input); + auto output = to_aten_tensor(p->output); + INFINICORE_ASSERT(input.dim() == 2 || input.dim() == 3); + + const auto seq_len = input.size(-2); + const auto total_seq_len = input.size(-1); + INFINICORE_ASSERT(seq_len <= total_seq_len); + + const bool reuses_input = input.is_same(output); + const bool needs_mask = seq_len != 1 || !reuses_input; + + at::Tensor scores; + if (needs_mask) { + const int64_t score_diagonal = total_seq_len - seq_len; + const int64_t mask_diagonal = total_seq_len - seq_len + 1; + // Match InfiniOps causal_softmax: tril(scores) + triu(-inf mask). + auto tril_scores = at::tril(input, score_diagonal); + auto mask = at::full({seq_len, total_seq_len}, + -std::numeric_limits::infinity(), + input.options()); + mask = at::triu(mask, mask_diagonal); + // Broadcast mask over leading dims if present. + scores = tril_scores + mask; + } else { + scores = input; + } + + auto probs = at::softmax(scores, /*dim=*/-1); + output.copy_(probs); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + CausalSoftmax::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + CausalSoftmax::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + CausalSoftmax::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::causal_softmax_impl::aten +#endif diff --git a/csrc/infinicore/src/ops/infiniops_impl.hpp b/csrc/infinicore/src/ops/infiniops_impl.hpp index a9051f1dc..3b8483689 100644 --- a/csrc/infinicore/src/ops/infiniops_impl.hpp +++ b/csrc/infinicore/src/ops/infiniops_impl.hpp @@ -65,6 +65,8 @@ inline infini::ops::Device toInfiniOpsDevice(const Device &device) { return infini::ops::Device{infini::ops::Device::Type::kHygon, static_cast(device.index())}; case Device::Type::kThead: return infini::ops::Device{infini::ops::Device::Type::kThead, static_cast(device.index())}; + case Device::Type::kQy: + return infini::ops::Device{infini::ops::Device::Type::kQy, static_cast(device.index())}; default: throw std::runtime_error("InfiniOps backend does not support this device type."); } @@ -80,6 +82,7 @@ inline bool isSupportedDevice(Device::Type device_type) { case Device::Type::kAscend: case Device::Type::kHygon: case Device::Type::kThead: + case Device::Type::kQy: return true; default: return false; @@ -148,6 +151,7 @@ void registerSupportedDevices(Dispatcher &dispatcher, Function function) { dispatcher.registerDevice(Device::Type::kAscend, function); dispatcher.registerDevice(Device::Type::kHygon, function); dispatcher.registerDevice(Device::Type::kThead, function); + dispatcher.registerDevice(Device::Type::kQy, function); } struct TensorMeta { diff --git a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc index 12fdf2b74..21ee9af5a 100644 --- a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc +++ b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache.cc @@ -1,5 +1,8 @@ #include "infinicore/ops/mha_kvcache.hpp" #include "../../utils.hpp" +#include "infinicore/context/context.hpp" + +#include namespace infinicore::op { @@ -70,6 +73,26 @@ bool MhaKVCache::supports_device_graph_capture( && k_cache->size(3) == 128 && seqlens_k->shape() == Shape({16}) && block_table->shape() == Shape({16, 512}); + // MiniCPM5-2B TP1 decode: fused QKV Q view, paged KV [NUM_BLOCKS,256,2,128]. + const bool minicpm5_2b_shape = q->shape() == Shape({1, 1, 16, 128}) + && out->shape() == Shape({1, 1, 16, 128}) + && k_cache->size(0) > 0 + && k_cache->size(1) == 256 + && k_cache->size(2) == 2 + && k_cache->size(3) == 128 + && seqlens_k->shape() == Shape({1}) + && block_table->size(0) == 1 + && block_table->size(1) > 0; + // MiniCPM5-16A3 TP1 decode: gated attention Q|gate interleaved head stride. + const bool minicpm5_16a3_shape = q->shape() == Shape({1, 1, 32, 128}) + && out->shape() == Shape({1, 1, 32, 128}) + && k_cache->size(0) > 0 + && k_cache->size(1) == 256 + && k_cache->size(2) == 2 + && k_cache->size(3) == 128 + && seqlens_k->shape() == Shape({1}) + && block_table->size(0) == 1 + && block_table->size(1) > 0; const bool p13_layout = q->is_contiguous() && out->is_contiguous(); // P09 reads Q directly from the fused QKV projection while output is dense. const bool p09_layout = q->strides() == Strides({5120, 5120, 128, 1}) @@ -83,33 +106,75 @@ bool MhaKVCache::supports_device_graph_capture( // P12 has the same fused QKV projection layout at batch size one. const bool p12_layout = q->strides() == Strides({4608, 4608, 128, 1}) && out->strides() == Strides({4096, 4096, 128, 1}); + // MiniCPM5-2B: QKV width (16+2+2)*128=2560; dense out width 16*128=2048. + const bool minicpm5_2b_layout + = (q->strides() == Strides({2560, 2560, 128, 1}) + && out->strides() == Strides({2048, 2048, 128, 1})) + || (q->is_contiguous() && out->is_contiguous()); + // MiniCPM5-16A3 gated: QKV width (64+2+2)*128=8704 with per-head Q|gate + // stride 256; dense out width 32*128=4096. + const bool minicpm5_16a3_layout + = (q->strides() == Strides({8704, 8704, 256, 1}) + && out->strides() == Strides({4096, 4096, 128, 1})) + || (q->is_contiguous() && out->is_contiguous()); const bool reviewed_shape_and_layout = (p12_shape && p12_layout) || (p13_shape && p13_layout) || (p09_shape && p09_layout) || (p11_shape && p11_layout) - || (p14_shape && p14_layout); + || (p14_shape && p14_layout) + || (minicpm5_2b_shape && minicpm5_2b_layout) + || (minicpm5_16a3_shape && minicpm5_16a3_layout); const auto dtype = q->dtype(); - return device.type() == Device::Type::kNvidia - && same_device(q) - && same_device(k_cache) - && same_device(v_cache) - && same_device(seqlens_k) - && same_device(block_table) - && reviewed_shape_and_layout - && out->shape() == q->shape() - && k_cache->size(0) > 0 - && v_cache->shape() == k_cache->shape() - && (dtype == DataType::kFloat16 - || dtype == DataType::kBFloat16) - && out->dtype() == dtype - && k_cache->dtype() == dtype - && v_cache->dtype() == dtype - && seqlens_k->dtype() == DataType::kInt32 - && block_table->dtype() == DataType::kInt32 - && k_cache->is_contiguous() - && v_cache->is_contiguous() - && seqlens_k->is_contiguous() - && block_table->is_contiguous(); + const bool supported_device = device.type() == Device::Type::kNvidia + || device.type() == Device::Type::kQy; + const bool ok = supported_device + && same_device(q) + && same_device(k_cache) + && same_device(v_cache) + && same_device(seqlens_k) + && same_device(block_table) + && reviewed_shape_and_layout + && out->shape() == q->shape() + && k_cache->size(0) > 0 + && v_cache->shape() == k_cache->shape() + && (dtype == DataType::kFloat16 + || dtype == DataType::kBFloat16) + && out->dtype() == dtype + && k_cache->dtype() == dtype + && v_cache->dtype() == dtype + && seqlens_k->dtype() == DataType::kInt32 + && block_table->dtype() == DataType::kInt32 + && k_cache->is_contiguous() + && v_cache->is_contiguous() + && seqlens_k->is_contiguous() + && block_table->is_contiguous(); + if (!ok + && device.type() == Device::Type::kQy + && (minicpm5_2b_shape || minicpm5_16a3_shape)) { + static bool logged = false; + if (!logged) { + logged = true; + const auto &qs = q->shape(); + const auto &qstr = q->strides(); + const auto &ostr = out->strides(); + const auto &ks = k_cache->shape(); + const auto &bts = block_table->shape(); + std::fprintf( + stderr, + "[mha_kvcache] QY graph-capture allowlist miss: " + "q=[%ld,%ld,%ld,%ld] q_strides=[%ld,%ld,%ld,%ld] " + "out_strides=[%ld,%ld,%ld,%ld] k_cache=[%ld,%ld,%ld,%ld] " + "block_table=[%ld,%ld] contiguous(q/out)=%d/%d\n", + (long)qs[0], (long)qs[1], (long)qs[2], (long)qs[3], + (long)qstr[0], (long)qstr[1], (long)qstr[2], (long)qstr[3], + (long)ostr[0], (long)ostr[1], (long)ostr[2], (long)ostr[3], + (long)ks[0], (long)ks[1], (long)ks[2], (long)ks[3], + (long)bts[0], (long)bts[1], + q->is_contiguous() ? 1 : 0, + out->is_contiguous() ? 1 : 0); + } + } + return ok; } MhaKVCache::MhaKVCache(Tensor out, @@ -122,6 +187,7 @@ MhaKVCache::MhaKVCache(Tensor out, float scale) : device_graph_capture_safe_( context::isGraphRecording() + && context::faInGraphAllowed() && supports_device_graph_capture( out, q, diff --git a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc index 7a03a736d..15268fada 100644 --- a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc +++ b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc @@ -45,7 +45,8 @@ bool is_supported(const Tensor &out, && device_type != Device::Type::kCambricon && device_type != Device::Type::kIluvatar && device_type != Device::Type::kHygon - && device_type != Device::Type::kAscend) + && device_type != Device::Type::kAscend + && device_type != Device::Type::kQy) || q->ndim() != 4 || out->ndim() != 4 || k_cache->ndim() != 4 @@ -119,9 +120,11 @@ struct PlannedMeta { std::unique_ptr make_graph_safe_provider( PlannedMeta &planned) { + const auto device_type = planned.q.device.type(); + const auto implementation_index = implementation_index_for_device(device_type); const auto config = ::infinicore::op::infiniops::configForImplementation< infini::ops::FlashAttnWithKvcache>( - infini::ops::Device::Type::kNvidia, 16); + device_type, implementation_index); const auto q = planned.q.tensor(planned.q_tensor); const auto k_cache = planned.k_cache.tensor(planned.k_cache_tensor); @@ -323,6 +326,9 @@ static bool registered = []() { MhaKVCache::plan_dispatcher().registerDevice(Device::Type::kAscend, &plan); MhaKVCache::run_dispatcher().registerDevice(Device::Type::kAscend, &run); MhaKVCache::cleanup_dispatcher().registerDevice(Device::Type::kAscend, &cleanup); + MhaKVCache::plan_dispatcher().registerDevice(Device::Type::kQy, &plan); + MhaKVCache::run_dispatcher().registerDevice(Device::Type::kQy, &run); + MhaKVCache::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup); return true; }(); diff --git a/csrc/infinicore/src/ops/moe_align/moe_align_aten.cc b/csrc/infinicore/src/ops/moe_align/moe_align_aten.cc new file mode 100644 index 000000000..e95149139 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_align/moe_align_aten.cc @@ -0,0 +1,211 @@ +#include "infinicore/ops/moe_align.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include + +#include +#include +#include + +#include +#include +#include +#include + +namespace infinicore::op::moe_align_impl::aten { +namespace { + +at::Tensor ensure_i32_contig(const at::Tensor &t) { + INFINICORE_ASSERT(t.is_contiguous()); + if (t.scalar_type() == at::kInt) { + return t; + } + INFINICORE_ASSERT(t.scalar_type() == at::kLong); + // Small [M*topk] cast — avoid ATen .to() OOR on Denglin. + auto out = at::empty(t.sizes(), t.options().dtype(at::kInt)); + const int64_t n = t.numel(); + std::vector host(static_cast(n)); + cudaMemcpy(host.data(), t.data_ptr(), + static_cast(n * static_cast(sizeof(int64_t))), + cudaMemcpyDeviceToHost); + std::vector host32(static_cast(n)); + for (int64_t i = 0; i < n; ++i) { + host32[static_cast(i)] = static_cast(host[static_cast(i)]); + } + cudaMemcpy(out.data_ptr(), host32.data(), + static_cast(n * static_cast(sizeof(int32_t))), + cudaMemcpyHostToDevice); + return out; +} + +struct AlignWorkspacePool { + c10::DeviceIndex device_index = -1; + int64_t num_experts = -1; + at::Tensor workspace; + + at::Tensor &ensure(const at::Tensor &like, int64_t num_experts_v, int64_t timer_m) { + const int64_t n_ints = static_cast( + ::infini::ops::qy::moe_align_workspace_bytes( + static_cast(num_experts_v)) / + sizeof(int32_t)); + if (workspace.defined() && device_index == like.device().index() && + num_experts == num_experts_v && workspace.numel() == n_ints) { + return workspace; + } + { + workspace = at::empty({n_ints}, like.options().dtype(at::kInt)); + } + device_index = like.device().index(); + num_experts = num_experts_v; + return workspace; + } +}; + +AlignWorkspacePool &align_workspace_pool() { + thread_local AlignWorkspacePool pool; + return pool; +} + +void run_moe_align_gpu(at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded, + at::Tensor topk_ids, + const at::Tensor *expert_map, + int64_t num_experts, + int64_t block_size, + bool pad_sorted_token_ids) { + INFINICORE_ASSERT(sorted_token_ids.is_contiguous()); + INFINICORE_ASSERT(expert_ids.is_contiguous()); + INFINICORE_ASSERT(num_tokens_post_padded.is_contiguous()); + INFINICORE_ASSERT(topk_ids.is_contiguous()); + + auto topk_i32 = ensure_i32_contig(topk_ids); + at::Tensor map_i32; + const void *map_ptr = nullptr; + if (expert_map != nullptr) { + INFINICORE_ASSERT(expert_map->is_contiguous()); + map_i32 = ensure_i32_contig(*expert_map); + map_ptr = map_i32.data_ptr(); + } + + auto &workspace = align_workspace_pool().ensure( + sorted_token_ids, num_experts, static_cast(topk_ids.numel())); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + ::infini::ops::qy::moe_align_launch( + sorted_token_ids.data_ptr(), expert_ids.data_ptr(), + num_tokens_post_padded.data_ptr(), topk_i32.data_ptr(), map_ptr, + static_cast(topk_i32.numel()), + static_cast(num_experts), static_cast(block_size), + static_cast(sorted_token_ids.numel()), + static_cast(expert_ids.numel()), pad_sorted_token_ids, + workspace.data_ptr(), reinterpret_cast(stream)); + // Sync only when profiling; otherwise rely on stream order into fuse. + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("QY moe_align launch: ") + + cudaGetErrorString(err)); + } +} + +struct PlannedMeta { + Tensor sorted_token_ids; + Tensor expert_ids; + Tensor num_tokens_post_padded; + Tensor topk_ids; + Tensor expert_map; + bool has_expert_map; + size_t num_experts; + size_t block_size; + bool pad_sorted_token_ids; +}; + +void *plan_no_map(Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + const size_t num_experts, + const size_t block_size, + const bool pad_sorted_token_ids) { + INFINICORE_ASSERT(sorted_token_ids->device().type() == Device::Type::kQy); + return new PlannedMeta{ + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + Tensor{}, + false, + num_experts, + block_size, + pad_sorted_token_ids}; +} + +void *plan_with_map(Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded, + const Tensor &topk_ids, + const Tensor &expert_map, + const size_t num_experts, + const size_t block_size, + const bool pad_sorted_token_ids) { + INFINICORE_ASSERT(sorted_token_ids->device().type() == Device::Type::kQy); + return new PlannedMeta{ + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + expert_map, + true, + num_experts, + block_size, + pad_sorted_token_ids}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + using namespace ::infinicore::adaptor; + set_aten_stream_to_infinicore(); + + + at::Tensor expert_map_aten; + const at::Tensor *expert_map_ptr = nullptr; + if (p->has_expert_map) { + expert_map_aten = to_aten_tensor(p->expert_map); + expert_map_ptr = &expert_map_aten; + } + + run_moe_align_gpu( + to_aten_tensor(p->sorted_token_ids), + to_aten_tensor(p->expert_ids), + to_aten_tensor(p->num_tokens_post_padded), + to_aten_tensor(p->topk_ids), + expert_map_ptr, + static_cast(p->num_experts), + static_cast(p->block_size), + p->pad_sorted_token_ids); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + MoeAlign::plan_dispatcher().registerDevice(Device::Type::kQy, &plan_no_map, true); + MoeAlign::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MoeAlign::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + + MoeAlignWithExpertMap::plan_dispatcher().registerDevice(Device::Type::kQy, &plan_with_map, true); + MoeAlignWithExpertMap::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MoeAlignWithExpertMap::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::moe_align_impl::aten +#endif diff --git a/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc index dea77265e..b0f4f9111 100644 --- a/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc +++ b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense.cc @@ -18,9 +18,9 @@ MoeFusedDense::MoeFusedDense(Tensor output, INFINICORE_ASSERT_TENSORS_SAME_DEVICE( output, hidden_states, w13, w2, topk_weights, topk_ids, sorted_token_ids, expert_ids, num_tokens_post_padded); - throw std::runtime_error( - "`moe_fused_dense` requires an InfiniOps dense expert provider; " - "quantized WNA16/Marlin providers do not accept floating-point expert weights"); + INFINICORE_GRAPH_OP_DISPATCH( + output->device().type(), output, hidden_states, w13, w2, topk_weights, topk_ids, + sorted_token_ids, expert_ids, num_tokens_post_padded); } void MoeFusedDense::execute(Tensor output, diff --git a/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense_qy_dl.cc b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense_qy_dl.cc new file mode 100644 index 000000000..25149e8d1 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_fused_dense/moe_fused_dense_qy_dl.cc @@ -0,0 +1,258 @@ +#if defined(ENABLE_ATEN) && defined(ENABLE_QY_API) + +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/adaptor/denglin_fused_moe_adaptor.hpp" +#include "infinicore/graph/graph.hpp" +#include "infinicore/ops/moe_fused_dense.hpp" + +#include +#include + +#include +#include + +#include +#include +#include + +namespace infinicore::op::moe_fused_dense_impl::qy_dl { +namespace { + +// Shape-keyed scratch pool: carve once per (device, dtype, topk, h, inter), +// grow capacity_m as needed; callers narrow to the live m. +struct FuseScratchPool { + c10::DeviceIndex device_index = -1; + at::ScalarType dtype = at::ScalarType::Undefined; + int64_t capacity_m = 0; + int64_t topk = 0; + int64_t h = 0; + int64_t inter = 0; + at::Tensor cache1_storage; + at::Tensor activated_storage; + at::Tensor cache3_storage; + + bool same_shape_key(const at::Tensor &like, int64_t topk_v, int64_t h_v, + int64_t inter_v) const { + return device_index == like.device().index() && dtype == like.scalar_type() && + topk == topk_v && h == h_v && inter == inter_v; + } + + void ensure(const at::Tensor &like, int64_t m_v, int64_t topk_v, int64_t h_v, + int64_t inter_v) { + const int64_t inter2 = inter_v * 2; + if (cache1_storage.defined() && activated_storage.defined() && + cache3_storage.defined() && same_shape_key(like, topk_v, h_v, inter_v) && + capacity_m >= m_v) { + return; + } + const int64_t alloc_m = m_v; + { + cache1_storage = at::empty({alloc_m, topk_v, inter2}, like.options()); + } + { + activated_storage = + at::empty({alloc_m * topk_v, inter_v}, like.options()); + } + { + cache3_storage = at::empty({alloc_m, topk_v, h_v}, like.options()); + } + device_index = like.device().index(); + dtype = like.scalar_type(); + capacity_m = alloc_m; + topk = topk_v; + h = h_v; + inter = inter_v; + } + + at::Tensor cache1(int64_t m_v) const { + return cache1_storage.narrow(/*dim=*/0, /*start=*/0, /*length=*/m_v); + } + at::Tensor activated(int64_t m_v) const { + return activated_storage.narrow(/*dim=*/0, /*start=*/0, + /*length=*/m_v * topk); + } + at::Tensor cache3(int64_t m_v) const { + return cache3_storage.narrow(/*dim=*/0, /*start=*/0, /*length=*/m_v); + } +}; + +FuseScratchPool &fuse_scratch_pool() { + thread_local FuseScratchPool pool; + return pool; +} + +} // namespace + +struct PlannedMeta { + graph::GraphTensor output; + graph::GraphTensor hidden_states; + graph::GraphTensor w13; + graph::GraphTensor w2; + graph::GraphTensor topk_weights; + graph::GraphTensor topk_ids; + graph::GraphTensor sorted_token_ids; + graph::GraphTensor expert_ids; + graph::GraphTensor num_tokens_post_padded; +}; + +void *plan(Tensor output, + const Tensor &hidden_states, + const Tensor &w13, + const Tensor &w2, + const Tensor &topk_weights, + const Tensor &topk_ids, + const Tensor &sorted_token_ids, + const Tensor &expert_ids, + const Tensor &num_tokens_post_padded) { + if (!adaptor::denglin_fused_moe::available()) { + throw std::runtime_error( + "QY moe_fused_dense requires vllm._dl_C (set INFINICORE_DENGLIN_DL_C_SO if needed)"); + } + return new PlannedMeta{ + graph::GraphTensor(output), + graph::GraphTensor(hidden_states), + graph::GraphTensor(w13), + graph::GraphTensor(w2), + graph::GraphTensor(topk_weights), + graph::GraphTensor(topk_ids), + graph::GraphTensor(sorted_token_ids), + graph::GraphTensor(expert_ids), + graph::GraphTensor(num_tokens_post_padded)}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + adaptor::set_aten_stream_to_infinicore(); + + auto hidden = adaptor::to_aten_tensor(p->hidden_states); + auto w13 = adaptor::to_aten_tensor(p->w13); + auto w2 = adaptor::to_aten_tensor(p->w2); + auto topk_weights = adaptor::to_aten_tensor(p->topk_weights); + auto topk_ids = adaptor::to_aten_tensor(p->topk_ids); + auto sorted_token_ids = adaptor::to_aten_tensor(p->sorted_token_ids); + auto expert_ids = adaptor::to_aten_tensor(p->expert_ids); + auto num_tokens_post_padded = adaptor::to_aten_tensor(p->num_tokens_post_padded); + auto output = adaptor::to_aten_tensor(p->output); + + const int64_t m = hidden.size(0); + const int64_t h = hidden.size(1); + const int64_t topk = topk_ids.size(1); + const int64_t inter2 = w13.size(1); + const int64_t inter = inter2 / 2; + + // Workspace capacity rounding makes sorted_cap/expert_cap recovery inexact + // (e.g. 17 instead of 16 for typical prefill). MiniCPM5 / CudaFusedMoeRunner + // default moe_align_block_size is 16 — use that for _dl_C. + constexpr int64_t block_size_m = 16; + constexpr int64_t block_size_n = 64; + constexpr int64_t block_size_k = 64; + const std::vector block_shape; // empty == vLLM default + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + // Stage 1: hidden @ w13 -> [M, topk, 2*inter] (vLLM fused_experts_impl) + auto &pool = fuse_scratch_pool(); + pool.ensure(hidden, m, topk, h, inter); + auto cache1 = pool.cache1(m); + auto activated = pool.activated(m); + auto cache3 = pool.cache3(m); + auto a1 = hidden.view({-1, h}); + auto c1 = cache1.view({-1, topk, inter2}); + { + adaptor::denglin_fused_moe::invoke_fused_moe_opt( + a1, w13, c1, + std::nullopt, std::nullopt, std::nullopt, + topk_weights, topk_ids, sorted_token_ids, expert_ids, num_tokens_post_padded, + /*mul_routed_weight=*/false, + topk, block_size_m, block_size_n, block_size_k, + false, false, false, false, + block_shape, m); + // Rely on stream order to silu (no host sync on score path). + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_fused_dense stage1 (w13) launch: ") + + cudaGetErrorString(err)); + } + } + + // SiLU/GLU activation via thin QY kernel (avoids D2H/H2D and ATen OOR). + at::Tensor flat_gate_up; + { + auto reshaped = cache1.reshape({m * topk, inter2}); + flat_gate_up = reshaped.is_contiguous() ? reshaped : reshaped.contiguous(); + } + { + ::infini::ops::qy::silu_and_mul_launch( + activated.data_ptr(), flat_gate_up.data_ptr(), + static_cast(m * topk), static_cast(inter), + static_cast(p->hidden_states->dtype()), + reinterpret_cast(stream)); + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_fused_dense silu launch: ") + + cudaGetErrorString(err)); + } + } + + // Stage 2: activated @ w2 -> [M, topk, H], top_k=1, mul routed weights + at::Tensor a2; + { + a2 = activated.is_contiguous() ? activated : activated.contiguous(); + } + auto c2 = cache3.view({-1, /*top_k=*/1, h}); + { + adaptor::denglin_fused_moe::invoke_fused_moe_opt( + a2, w2, c2, + std::nullopt, std::nullopt, std::nullopt, + topk_weights, topk_ids, sorted_token_ids, expert_ids, num_tokens_post_padded, + /*mul_routed_weight=*/true, + /*top_k=*/1, block_size_m, block_size_n, block_size_k, + false, false, false, false, + block_shape, m); + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_fused_dense stage2 (w2) launch: ") + + cudaGetErrorString(err)); + } + } + + // Expert combine via thin QY moe_sum (arbitrary topk). + auto cache3_contig = cache3.is_contiguous() ? cache3 : cache3.contiguous(); + if (!output.is_contiguous()) { + throw std::runtime_error( + "QY moe_fused_dense combine requires contiguous output"); + } + { + ::infini::ops::qy::moe_sum_launch( + output.data_ptr(), cache3_contig.data_ptr(), + static_cast(m), static_cast(topk), + static_cast(h), static_cast(p->hidden_states->dtype()), + reinterpret_cast(stream)); + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("QY moe_fused_dense combine launch: ") + + cudaGetErrorString(err)); + } + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +static bool registered = []() { + // Override the ALLDEVICE infiniop registration for Denglin QY. + MoeFusedDense::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + MoeFusedDense::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MoeFusedDense::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::moe_fused_dense_impl::qy_dl + +#endif diff --git a/csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid_aten.cc b/csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid_aten.cc new file mode 100644 index 000000000..3bdf28b40 --- /dev/null +++ b/csrc/infinicore/src/ops/moe_topk_sigmoid/moe_topk_sigmoid_aten.cc @@ -0,0 +1,116 @@ +#include "infinicore/ops/moe_topk_sigmoid.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" + +#include + +#include +#include + +#include +#include +#include + +namespace infinicore::op::moe_topk_sigmoid_impl::aten { +namespace { + +struct PlannedMeta { + Tensor topk_weights; + Tensor topk_indices; + Tensor gating_output; + Tensor correction_bias; + bool renormalize; +}; + +void *plan(Tensor topk_weights, + Tensor topk_indices, + const Tensor &gating_output, + const Tensor &correction_bias, + const bool renormalize) { + INFINICORE_ASSERT(topk_weights->device().type() == Device::Type::kQy); + return new PlannedMeta{ + topk_weights, + topk_indices, + gating_output, + correction_bias, + renormalize}; +} + +void run(void *planned_meta) { + auto *planned = reinterpret_cast(planned_meta); + using namespace ::infinicore::adaptor; + set_aten_stream_to_infinicore(); + + auto gating_dev = to_aten_tensor(planned->gating_output); + auto weights_dev = to_aten_tensor(planned->topk_weights); + auto indices_dev = to_aten_tensor(planned->topk_indices); + + INFINICORE_ASSERT(gating_dev.is_contiguous()); + INFINICORE_ASSERT(weights_dev.is_contiguous()); + INFINICORE_ASSERT(indices_dev.is_contiguous()); + INFINICORE_ASSERT(gating_dev.dim() == 2); + INFINICORE_ASSERT(weights_dev.dim() == 2); + + const auto tokens = gating_dev.size(0); + const auto experts = gating_dev.size(1); + const auto topk = weights_dev.size(1); + INFINICORE_ASSERT(indices_dev.size(0) == tokens); + INFINICORE_ASSERT(indices_dev.size(1) == topk); + + + const void *bias_ptr = nullptr; + int bias_dtype_tag = 0; + if (planned->correction_bias) { + auto bias = to_aten_tensor(planned->correction_bias); + INFINICORE_ASSERT(bias.is_contiguous()); + INFINICORE_ASSERT(bias.numel() == experts); + // f32/bf16/fp16 bias consumed on device (kernel converts to float). + if (bias.scalar_type() != at::kFloat && + bias.scalar_type() != at::kHalf && + bias.scalar_type() != at::kBFloat16) { + throw std::runtime_error( + "QY moe_topk_sigmoid: unsupported correction_bias dtype"); + } + bias_ptr = bias.data_ptr(); + bias_dtype_tag = static_cast(planned->correction_bias->dtype()); + } + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + // Kernel writes the actual weight dtype (f16/bf16/f32) and index dtype + // (i32/i64) — no D2H/H2D post-cast. + ::infini::ops::qy::moe_topk_sigmoid_launch( + weights_dev.data_ptr(), indices_dev.data_ptr(), gating_dev.data_ptr(), + bias_ptr, static_cast(tokens), static_cast(experts), + static_cast(topk), planned->renormalize, + static_cast(planned->gating_output->dtype()), + static_cast(planned->topk_weights->dtype()), + static_cast(planned->topk_indices->dtype()), bias_dtype_tag, + reinterpret_cast(stream)); + // Launch-error check only. Forced stream sync was inflating host timers and + // serializing MoE; consumers (align / fused dense) sync as needed. + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("QY moe_topk_sigmoid launch: ") + + cudaGetErrorString(err)); + } +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + MoeTopkSigmoid::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + MoeTopkSigmoid::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MoeTopkSigmoid::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::moe_topk_sigmoid_impl::aten +#endif diff --git a/csrc/infinicore/src/ops/mul_scalar/mul_scalar_aten.cc b/csrc/infinicore/src/ops/mul_scalar/mul_scalar_aten.cc new file mode 100644 index 000000000..4f407c471 --- /dev/null +++ b/csrc/infinicore/src/ops/mul_scalar/mul_scalar_aten.cc @@ -0,0 +1,58 @@ +#include "infinicore/ops/mul_scalar.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/context/context.hpp" + +#include + +namespace infinicore::op::mul_scalar_impl::aten { +namespace { + +struct PlannedMeta { + Tensor c; + Tensor a; + double alpha; +}; + +void *plan(Tensor c, const Tensor &a, double alpha) { + INFINICORE_ASSERT(c->device().type() == Device::Type::kQy); + return new PlannedMeta{c, a, alpha}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + + INFINICORE_ASSERT(p->c->device() == p->a->device()); + INFINICORE_ASSERT(p->c->dtype() == p->a->dtype()); + INFINICORE_ASSERT(p->c->numel() == p->a->numel()); + INFINICORE_ASSERT(p->c->is_contiguous()); + + Tensor a = p->a->is_contiguous() ? p->a : p->a->contiguous(); + INFINICORE_ASSERT(a->is_contiguous()); + INFINICORE_ASSERT(a->numel() == p->c->numel()); + INFINICORE_ASSERT(a->dtype() == p->c->dtype()); + + ::infini::ops::qy::mul_scalar_launch( + p->c->data(), a->data(), static_cast(a->numel()), + static_cast(a->dtype()), p->alpha, + reinterpret_cast(context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + MulScalar::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + MulScalar::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + MulScalar::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::mul_scalar_impl::aten +#endif diff --git a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen.cc b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen.cc index 872eb5d98..7ee20da38 100644 --- a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen.cc +++ b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen.cc @@ -1,11 +1,13 @@ #include "infinicore/ops/mha_varlen.hpp" #include "../../utils.hpp" +#include "infinicore/context/context.hpp" namespace infinicore::op { INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MultiheadAttentionVarlen); // The Ascend provider synchronizes cumulative sequence lengths on the host. +// Prefill FA (varlen) is capture-safe only when phase policy allows FA in-graph. MultiheadAttentionVarlen::MultiheadAttentionVarlen(Tensor out, const Tensor &q, const Tensor &k, @@ -17,9 +19,10 @@ MultiheadAttentionVarlen::MultiheadAttentionVarlen(Tensor out, int max_seqlen_k, std::optional alibi_slopes, float scale) - : device_graph_capture_safe_(out->device().type() != Device::Type::kNvidia - && out->device().type() - != Device::Type::kAscend) { + : device_graph_capture_safe_( + context::faInGraphAllowed() + && out->device().type() != Device::Type::kNvidia + && out->device().type() != Device::Type::kAscend) { INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k, v, cum_seqlens_q, cum_seqlens_kv); if (block_table.has_value()) { INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, block_table.value()); diff --git a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc index 043ce070a..f4c6b55d4 100644 --- a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc +++ b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc @@ -7,8 +7,11 @@ #include #include +#include #include +#include + namespace infinicore::op::mha_varlen_impl::infiniops { namespace { @@ -47,7 +50,8 @@ bool is_supported(const Tensor &out, && device_type != Device::Type::kCambricon && device_type != Device::Type::kIluvatar && device_type != Device::Type::kHygon - && device_type != Device::Type::kAscend) + && device_type != Device::Type::kAscend + && device_type != Device::Type::kQy) || q->ndim() != 3 || out->ndim() != 3 || ((paged && (k->ndim() != 4 || v->ndim() != 4)) @@ -144,6 +148,40 @@ void *plan(Tensor out, int max_seqlen_k, std::optional alibi_slopes, float scale) { + // Piecewise / paged paths often pass a non-contiguous block_table view + // (e.g. shape [1,1] prefix of a wider bank). InfiniOps FA requires contig. + if (block_table.has_value() && block_table.value() + && !block_table.value()->is_contiguous()) { + block_table = block_table.value()->contiguous(); + } + if (!is_supported( + out, q, k, v, cum_seqlens_q, cum_seqlens_k, block_table, + max_seqlen_q, max_seqlen_k, alibi_slopes)) { + auto shape_str = [](const Tensor &t) { + std::ostringstream oss; + oss << "ndim=" << t->ndim() << " shape=["; + for (size_t i = 0; i < t->ndim(); ++i) { + if (i) oss << ','; + oss << t->size(i); + } + oss << "] contig=" << t->is_contiguous() + << " dtype=" << static_cast(t->dtype()); + return oss.str(); + }; + spdlog::error( + "mha_varlen is_supported failed: max_q={} max_k={} q({}) k({}) v({}) " + "out({}) cu_q({}) cu_k({}) bt({}) device={}", + max_seqlen_q, + max_seqlen_k, + shape_str(q), + shape_str(k), + shape_str(v), + shape_str(out), + shape_str(cum_seqlens_q), + shape_str(cum_seqlens_k), + block_table ? shape_str(*block_table) : "null", + static_cast(out->device().type())); + } INFINICORE_ASSERT(is_supported( out, q, k, v, cum_seqlens_q, cum_seqlens_k, block_table, max_seqlen_q, max_seqlen_k, alibi_slopes)); @@ -266,6 +304,12 @@ static bool registered = []() { Device::Type::kAscend, &run); MultiheadAttentionVarlen::cleanup_dispatcher().registerDevice( Device::Type::kAscend, &cleanup); + MultiheadAttentionVarlen::plan_dispatcher().registerDevice( + Device::Type::kQy, &plan); + MultiheadAttentionVarlen::run_dispatcher().registerDevice( + Device::Type::kQy, &run); + MultiheadAttentionVarlen::cleanup_dispatcher().registerDevice( + Device::Type::kQy, &cleanup); return true; }(); diff --git a/csrc/infinicore/src/ops/random_sample/random_sample.cc b/csrc/infinicore/src/ops/random_sample/random_sample.cc index 6d93af3c1..55bbb654f 100644 --- a/csrc/infinicore/src/ops/random_sample/random_sample.cc +++ b/csrc/infinicore/src/ops/random_sample/random_sample.cc @@ -25,13 +25,15 @@ bool tryGreedyWithInfiniOps( && device_type != Device::Type::kMoore && device_type != Device::Type::kCambricon && device_type != Device::Type::kAscend - && device_type != Device::Type::kHygon) + && device_type != Device::Type::kHygon + && device_type != Device::Type::kQy) || (random_value != 0.0f && top_p != 0.0f && top_k != 1 && temperature != 0.0f) || (logits->ndim() != 1 - && !(device_type == Device::Type::kNvidia && batched)) + && !(device_type == Device::Type::kNvidia && batched) + && !(device_type == Device::Type::kQy && batched)) || logits->numel() == 0 || !logits->is_contiguous() || (dtype != DataType::kFloat16 && dtype != DataType::kBFloat16 && dtype != DataType::kFloat32) diff --git a/csrc/infinicore/src/ops/rope/rope_aten.cc b/csrc/infinicore/src/ops/rope/rope_aten.cc new file mode 100644 index 000000000..fd57fa8ac --- /dev/null +++ b/csrc/infinicore/src/ops/rope/rope_aten.cc @@ -0,0 +1,332 @@ +#include "infinicore/ops/rope.hpp" + +#include "../../utils.hpp" + +#ifdef ENABLE_ATEN +#include "infinicore/adaptor/aten_adaptor.hpp" +#include "infinicore/nn/rope.hpp" + +#include +#include + +#include +#include + +#include +#include + +namespace infinicore::op::rope_impl::aten { +namespace { + +// QY FA+paged long-prefill RoPE: Denglin ATen index_select/gather and strided +// copy_/mul on gated as_strided Q/K hit LaunchOutOfResources. Pack/unpack with +// cudaMemcpyAsync (contiguous head vectors), rotate on a compact workspace, +// gather cos/sin via cache-table slice. +constexpr int64_t kQyRopeChunkTokens = 128; + + + +struct PlannedMeta { + Tensor x_out; + Tensor x; + Tensor pos; + Tensor sin; + Tensor cos; + infinicore::nn::RoPE::Algo algo; +}; + +void *plan(Tensor x_out, + const Tensor &x, + const Tensor &pos, + const Tensor &sin, + const Tensor &cos, + infinicore::nn::RoPE::Algo algo) { + INFINICORE_ASSERT(x_out->device().type() == Device::Type::kQy); + return new PlannedMeta{x_out, x, pos, sin, cos, algo}; +} + +int64_t seq_dim_of(const at::Tensor &x) { + if (x.dim() == 2) { + return 0; + } + INFINICORE_ASSERT(x.dim() >= 3); + return x.dim() - 3; +} + +at::Tensor d2h_contig(const at::Tensor &dev) { + INFINICORE_ASSERT(dev.is_contiguous()); + auto host = at::empty(dev.sizes(), dev.options().device(at::kCPU)); + cudaMemcpy(host.data_ptr(), + dev.data_ptr(), + static_cast(dev.numel() * static_cast(dev.element_size())), + cudaMemcpyDeviceToHost); + return host; +} + +at::Tensor gather_cache_rows(const at::Tensor &table, const at::Tensor &pos_chunk) { + const int64_t rows = pos_chunk.numel(); + INFINICORE_ASSERT(pos_chunk.dim() == 1); + INFINICORE_ASSERT(table.dim() == 2); + if (rows == 0) { + return at::empty({0, table.size(1)}, table.options()); + } + + // Host indices once — avoid ATen index_select / .contiguous() on device. + INFINICORE_ASSERT(pos_chunk.is_contiguous()); + auto pos_host = d2h_contig(pos_chunk); + const auto *pos_ptr = pos_host.data_ptr(); + + auto out = at::empty({rows, table.size(1)}, table.options()); + const auto elem = static_cast(table.element_size()); + auto *dst_base = static_cast(out.data_ptr()); + const auto *src_base = static_cast(table.data_ptr()); + const int64_t row_bytes = table.size(1) * elem; + const int64_t table_rows = table.size(0); + + const int64_t p0 = pos_ptr[0]; + bool consecutive = p0 >= 0 && (p0 + rows) <= table_rows; + for (int64_t i = 0; consecutive && i < rows; ++i) { + consecutive = (pos_ptr[i] == p0 + i); + } + if (consecutive && table.stride(0) == table.size(1)) { + cudaMemcpy(dst_base, + src_base + p0 * table.stride(0) * elem, + static_cast(rows * row_bytes), + cudaMemcpyDeviceToDevice); + return out; + } + + for (int64_t i = 0; i < rows; ++i) { + const int64_t p = pos_ptr[i]; + INFINICORE_ASSERT(p >= 0 && p < table_rows); + cudaMemcpy(dst_base + i * row_bytes, + src_base + p * table.stride(0) * elem, + static_cast(row_bytes), + cudaMemcpyDeviceToDevice); + } + return out; +} + +// Pack strided [seq, heads, dim] (or [seq, dim]) into contiguous via QY kernel. +void pack_head_vectors(const at::Tensor &src, at::Tensor &dst, cudaStream_t stream) { + INFINICORE_ASSERT(src.sizes() == dst.sizes()); + INFINICORE_ASSERT(dst.is_contiguous()); + if (src.is_contiguous()) { + cudaMemcpyAsync(dst.data_ptr(), + src.data_ptr(), + static_cast(src.numel() * static_cast(src.element_size())), + cudaMemcpyDeviceToDevice, + stream); + return; + } + const int elem = static_cast(src.element_size()); + const int64_t dim = src.size(-1); + INFINICORE_ASSERT(src.stride(-1) == 1); + + if (src.dim() == 2) { + ::infini::ops::qy::rope_pack_launch( + dst.data_ptr(), src.data_ptr(), static_cast(src.size(0)), + /*heads=*/1, static_cast(dim), src.stride(0), /*s1=*/0, + /*s2=*/1, elem, reinterpret_cast(stream)); + return; + } + INFINICORE_ASSERT(src.dim() == 3); + ::infini::ops::qy::rope_pack_launch( + dst.data_ptr(), src.data_ptr(), static_cast(src.size(0)), + static_cast(src.size(1)), static_cast(dim), + src.stride(0), src.stride(1), src.stride(2), elem, + reinterpret_cast(stream)); +} + +void unpack_head_vectors(const at::Tensor &src, at::Tensor &dst, cudaStream_t stream) { + INFINICORE_ASSERT(src.sizes() == dst.sizes()); + INFINICORE_ASSERT(src.is_contiguous()); + if (dst.is_contiguous()) { + cudaMemcpyAsync(dst.data_ptr(), + src.data_ptr(), + static_cast(src.numel() * static_cast(src.element_size())), + cudaMemcpyDeviceToDevice, + stream); + return; + } + const int elem = static_cast(src.element_size()); + const int64_t dim = src.size(-1); + INFINICORE_ASSERT(dst.stride(-1) == 1); + + if (dst.dim() == 2) { + ::infini::ops::qy::rope_unpack_launch( + dst.data_ptr(), src.data_ptr(), static_cast(dst.size(0)), + /*heads=*/1, static_cast(dim), dst.stride(0), /*s1=*/0, + /*s2=*/1, elem, reinterpret_cast(stream)); + return; + } + INFINICORE_ASSERT(dst.dim() == 3); + ::infini::ops::qy::rope_unpack_launch( + dst.data_ptr(), src.data_ptr(), static_cast(dst.size(0)), + static_cast(dst.size(1)), static_cast(dim), + dst.stride(0), dst.stride(1), dst.stride(2), elem, + reinterpret_cast(stream)); +} + +at::Tensor rotate_contig_gpu(const at::Tensor &x_contig, + const at::Tensor &pos_chunk, + const at::Tensor &sin_table, + const at::Tensor &cos_table, + infinicore::nn::RoPE::Algo algo, + int dtype_tag, + cudaStream_t stream) { + INFINICORE_ASSERT(x_contig.is_contiguous()); + INFINICORE_ASSERT(pos_chunk.is_contiguous()); + INFINICORE_ASSERT(sin_table.is_contiguous()); + INFINICORE_ASSERT(cos_table.is_contiguous()); + + const int64_t seqlen = x_contig.size(0); + const int64_t nhead = (x_contig.dim() == 2) ? 1 : x_contig.size(1); + const int64_t head_dim = x_contig.size(-1); + INFINICORE_ASSERT(pos_chunk.numel() == seqlen); + INFINICORE_ASSERT(sin_table.size(1) == head_dim / 2); + INFINICORE_ASSERT(cos_table.size(1) == head_dim / 2); + + const int algo_tag = (algo == infinicore::nn::RoPE::Algo::GPT_J) ? 0 : 1; + if (algo != infinicore::nn::RoPE::Algo::GPT_NEOX && algo != infinicore::nn::RoPE::Algo::GPT_J) { + throw std::runtime_error("Unsupported RoPE algorithm for QY ATen path"); + } + + auto out = at::empty(x_contig.sizes(), x_contig.options()); + ::infini::ops::qy::rope_launch( + out.data_ptr(), x_contig.data_ptr(), + pos_chunk.data_ptr(), sin_table.data_ptr(), + cos_table.data_ptr(), static_cast(seqlen), + static_cast(nhead), static_cast(head_dim), algo_tag, + dtype_tag, reinterpret_cast(stream)); + return out; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + using namespace ::infinicore::adaptor; + set_aten_stream_to_infinicore(); + + + auto x = to_aten_tensor(p->x); + auto pos = to_aten_tensor(p->pos).to(at::kLong).reshape({-1}).contiguous(); + auto sin_table = to_aten_tensor(p->sin); + auto cos_table = to_aten_tensor(p->cos); + auto out = to_aten_tensor(p->x_out); + + const int64_t seq_len = pos.numel(); + const int64_t seq_dim = seq_dim_of(x); + INFINICORE_ASSERT(seq_dim == 0); + INFINICORE_ASSERT(x.size(0) == seq_len); + INFINICORE_ASSERT(out.sizes() == x.sizes()); + INFINICORE_ASSERT(x.stride(-1) == 1); + INFINICORE_ASSERT(out.stride(-1) == 1); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + + // Prefer full-table + pos lookup in the kernel to avoid per-chunk gather. + // Fall back to gathered rows only if tables are not contiguous. + at::Tensor sin_use = sin_table; + at::Tensor cos_use = cos_table; + at::Tensor pos_use = pos; + if (!sin_table.is_contiguous() || !cos_table.is_contiguous()) { + sin_use = gather_cache_rows(sin_table, pos); + cos_use = gather_cache_rows(cos_table, pos); + // Remap positions to [0..seq) against gathered rows. + auto pos_remap = at::empty({seq_len}, pos.options().device(at::kCPU)); + auto *ptr = pos_remap.data_ptr(); + for (int64_t i = 0; i < seq_len; ++i) { + ptr[i] = i; + } + pos_use = at::empty({seq_len}, pos.options()); + cudaMemcpy(pos_use.data_ptr(), ptr, + static_cast(seq_len * static_cast(sizeof(int64_t))), + cudaMemcpyHostToDevice); + } + + const int algo_tag = + (p->algo == infinicore::nn::RoPE::Algo::GPT_J) ? 0 : 1; + if (p->algo != infinicore::nn::RoPE::Algo::GPT_NEOX && + p->algo != infinicore::nn::RoPE::Algo::GPT_J) { + throw std::runtime_error("Unsupported RoPE algorithm for QY ATen path"); + } + const int dtype_tag = static_cast(p->x->dtype()); + INFINICORE_ASSERT(sin_use.is_contiguous()); + INFINICORE_ASSERT(cos_use.is_contiguous()); + INFINICORE_ASSERT(pos_use.is_contiguous()); + + // Fast path: strided in-place/out RoPE (no pack/unpack). Covers MiniCPM5 + // gated Q [seq,heads,dim] strides like [8704,256,1] and contiguous K. + // Also [1,seq,heads,dim] static views with batch=1. + const bool same_layout = out.strides() == x.strides(); + int64_t nhead = 1; + int64_t head_dim = x.size(-1); + int64_t s0 = 0; + int64_t s1 = 0; + bool use_strided = false; + if (same_layout && x.stride(-1) == 1) { + if (x.dim() == 2) { + nhead = 1; + s0 = x.stride(0); + s1 = 0; + use_strided = true; + } else if (x.dim() == 3) { + nhead = x.size(1); + s0 = x.stride(0); + s1 = x.stride(1); + use_strided = true; + } else if (x.dim() == 4 && x.size(0) == 1) { + nhead = x.size(2); + s0 = x.stride(1); + s1 = x.stride(2); + use_strided = true; + } + } + if (use_strided) { + ::infini::ops::qy::rope_strided_launch( + out.data_ptr(), x.data_ptr(), pos_use.data_ptr(), + sin_use.data_ptr(), cos_use.data_ptr(), + static_cast(seq_len), static_cast(nhead), + static_cast(head_dim), s0, s1, /*stride2=*/1, algo_tag, + dtype_tag, reinterpret_cast(stream)); + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("QY rope_strided launch: ") + + cudaGetErrorString(err)); + } + } else { + // Legacy pack → contiguous rotate → unpack. + at::Tensor x_contig = + x.is_contiguous() ? x : at::empty(x.sizes(), x.options()); + if (!x.is_contiguous()) { + pack_head_vectors(x, x_contig, stream); + } + auto rotated = rotate_contig_gpu(x_contig, pos_use, sin_use, cos_use, + p->algo, dtype_tag, stream); + unpack_head_vectors(rotated, out, stream); + auto err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("QY rope_unpack launch: ") + + cudaGetErrorString(err)); + } + } + + (void)kQyRopeChunkTokens; +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +} // namespace + +static bool registered = []() { + RoPE::plan_dispatcher().registerDevice(Device::Type::kQy, &plan, true); + RoPE::run_dispatcher().registerDevice(Device::Type::kQy, &run, true); + RoPE::cleanup_dispatcher().registerDevice(Device::Type::kQy, &cleanup, true); + return true; +}(); + +} // namespace infinicore::op::rope_impl::aten +#endif diff --git a/csrc/infinicore/src/pybind11/context.hpp b/csrc/infinicore/src/pybind11/context.hpp index a59b8fde3..ec2d16efa 100644 --- a/csrc/infinicore/src/pybind11/context.hpp +++ b/csrc/infinicore/src/pybind11/context.hpp @@ -3,6 +3,7 @@ #include #include +#include #include "infinicore.hpp" @@ -34,6 +35,38 @@ inline void bind(py::module &m) { m.def("start_graph_recording", &startGraphRecording, "Start graph recording"); m.def("stop_graph_recording", &stopGraphRecording, "Stop graph recording and return the graph"); m.def("cancel_graph_recording", &cancelGraphRecording, "Cancel the active graph recording"); + m.def("fa_in_graph_allowed", + &faInGraphAllowed, + "Phase-adaptive FA under stream capture (INFINI_FA_INGRAPH)"); + m.def("moe_triton_capture_allowed", + &moeTritonCaptureAllowed, + "Phase-adaptive MoE under stream capture (INFINI_MOE_INGRAPH)"); + m.def( + "set_inference_phase", + [](const std::string &phase) { + if (phase == "decode" || phase == "Decode") { + setInferencePhase(InferencePhase::Decode); + } else if (phase == "prefill" || phase == "Prefill") { + setInferencePhase(InferencePhase::Prefill); + } else { + setInferencePhase(InferencePhase::Unknown); + } + }, + py::arg("phase"), + "Set TLS InferencePhase: decode | prefill | unknown"); + m.def( + "get_inference_phase", + []() -> std::string { + switch (getInferencePhase()) { + case InferencePhase::Decode: + return "decode"; + case InferencePhase::Prefill: + return "prefill"; + default: + return "unknown"; + } + }, + "Current TLS InferencePhase string"); } } // namespace infinicore::context diff --git a/csrc/infinicore/src/pybind11/device.hpp b/csrc/infinicore/src/pybind11/device.hpp index b8c8e76a6..482968c84 100644 --- a/csrc/infinicore/src/pybind11/device.hpp +++ b/csrc/infinicore/src/pybind11/device.hpp @@ -20,7 +20,8 @@ inline void bind(py::module &m) { .value("MOORE", Device::Type::kMoore) .value("ILUVATAR", Device::Type::kIluvatar) .value("HYGON", Device::Type::kHygon) - .value("THEAD", Device::Type::kThead); + .value("THEAD", Device::Type::kThead) + .value("QY", Device::Type::kQy); device .def(py::init(), diff --git a/csrc/layers/attention/attention.cpp b/csrc/layers/attention/attention.cpp index 9392bd94f..921738850 100644 --- a/csrc/layers/attention/attention.cpp +++ b/csrc/layers/attention/attention.cpp @@ -139,6 +139,170 @@ infinicore::Tensor Attention::forward_paged_(const infinicore::Tensor &position_ return output; } +void Attention::forward_pre_attn_qkv_piecewise(const infinicore::Tensor &, + const infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + auto &piecewise = global_state::get_forward_context().piecewise; + auto hidden_states_mutable = hidden_states; + auto shape = hidden_states->shape(); + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + ASSERT_EQ(batch_size, 1); + + auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + const size_t valid_len = piecewise.valid_seq_len > 0 ? piecewise.valid_seq_len : seq_len; + + auto q_heads = q->view({1, seq_len, num_attention_heads_, head_dim_}); + auto k_heads = k->view({1, seq_len, num_key_value_heads_, head_dim_}); + auto v_heads = v->view({1, seq_len, num_key_value_heads_, head_dim_}); + + if (valid_len < seq_len) { + staging.q_rope->narrow({{1, 0, valid_len}})->copy_from(q_heads->narrow({{1, 0, valid_len}})); + staging.k_rope->narrow({{1, 0, valid_len}})->copy_from(k_heads->narrow({{1, 0, valid_len}})); + staging.v_rope->narrow({{1, 0, valid_len}})->copy_from(v_heads->narrow({{1, 0, valid_len}})); + auto q_tail = staging.q_rope->narrow({{1, valid_len, seq_len - valid_len}}); + auto k_tail = staging.k_rope->narrow({{1, valid_len, seq_len - valid_len}}); + auto v_tail = staging.v_rope->narrow({{1, valid_len, seq_len - valid_len}}); + set_zeros(q_tail); + set_zeros(k_tail); + set_zeros(v_tail); + } else { + staging.q_rope->copy_from(q_heads); + staging.k_rope->copy_from(k_heads); + staging.v_rope->copy_from(v_heads); + } +} + +void Attention::forward_pre_attn_rope_piecewise(const infinicore::Tensor &position_ids, + global_state::PiecewiseLayerStaging &staging) const { + auto &piecewise = global_state::get_forward_context().piecewise; + const size_t seq_len = staging.q_rope->size(1); + const size_t valid_len = piecewise.valid_seq_len > 0 ? piecewise.valid_seq_len : seq_len; + const bool addr_stable = infinicore::context::isGraphRecording() + || piecewise.compile_capture_active; + + auto pos_shape = position_ids->shape(); + infinicore::Tensor pos_ids_for_rope = position_ids; + if (pos_shape.size() == 2) { + auto pos_narrowed = position_ids->narrow({{0, 0, 1}}); + if (addr_stable) { + pos_ids_for_rope = pos_narrowed->view({pos_shape[1]}); + } else { + pos_ids_for_rope = pos_narrowed->contiguous()->view({pos_shape[1]}); + } + } else if (pos_shape.size() == 1) { + if (!addr_stable && !position_ids->is_contiguous()) { + pos_ids_for_rope = position_ids->contiguous(); + } + } else { + throw std::runtime_error("Attention: Unexpected position_ids shape"); + } + if (pos_ids_for_rope->size(0) > valid_len) { + auto narrowed = pos_ids_for_rope->narrow({{0, 0, valid_len}}); + if (addr_stable) { + if (!narrowed->is_contiguous()) { + throw std::runtime_error( + "Attention::forward_pre_attn_rope_piecewise: " + "position_ids narrow not contiguous under graph recording"); + } + pos_ids_for_rope = narrowed; + } else { + pos_ids_for_rope = narrowed->contiguous(); + } + } + + auto q_rope = staging.q_rope->view({seq_len, num_attention_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + auto k_rope = staging.k_rope->view({seq_len, num_key_value_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + if (!q_rope->is_contiguous()) { + if (addr_stable) { + throw std::runtime_error( + "Attention::forward_pre_attn_rope_piecewise: " + "q_rope not contiguous under graph recording"); + } + q_rope = q_rope->contiguous(); + } + if (!k_rope->is_contiguous()) { + if (addr_stable) { + throw std::runtime_error( + "Attention::forward_pre_attn_rope_piecewise: " + "k_rope not contiguous under graph recording"); + } + k_rope = k_rope->contiguous(); + } + rotary_emb_->forward(q_rope, pos_ids_for_rope, true); + rotary_emb_->forward(k_rope, pos_ids_for_rope, true); +} + +void Attention::forward_pre_attn_piecewise(const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + forward_pre_attn_qkv_piecewise(position_ids, hidden_states, staging); + forward_pre_attn_rope_piecewise(position_ids, staging); +} + +void Attention::forward_eager_attn_piecewise(const infinicore::Tensor &, + global_state::PiecewiseLayerStaging &staging) const { + auto &piecewise = global_state::get_forward_context().piecewise; + const size_t seq_len = staging.q_rope->size(1); + const size_t valid_len = piecewise.valid_seq_len > 0 ? piecewise.valid_seq_len : seq_len; + + auto q = staging.q_rope->view({seq_len, num_attention_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + auto k = staging.k_rope->view({seq_len, num_key_value_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + auto v = staging.v_rope->view({seq_len, num_key_value_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + if (!q->is_contiguous()) { + q = q->contiguous(); + } + if (!k->is_contiguous()) { + k = k->contiguous(); + } + if (!v->is_contiguous()) { + v = v->contiguous(); + } + + auto attn_output = attn_->forward(q, k, v); + if (valid_len < seq_len) { + staging.attn_output->narrow({{1, 0, valid_len}})->copy_from(attn_output->narrow({{1, 0, valid_len}})); + auto attn_tail = staging.attn_output->narrow({{1, valid_len, seq_len - valid_len}}); + set_zeros(attn_tail); + } else { + staging.attn_output->copy_from(attn_output); + } +} + +void Attention::forward_post_attn_piecewise_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + forward_post_attn_piecewise_graph_into(hidden_states, staging); + forward_post_attn_piecewise_allreduce_into(hidden_states, staging); +} + +void Attention::forward_post_attn_piecewise_graph_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + // Dest RowParallelLinear has no matmul-only/allreduce split; reuse fused path. + forward_post_attn_piecewise_cg_into(hidden_states, staging); +} + +void Attention::forward_post_attn_piecewise_cg_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + auto &piecewise = global_state::get_forward_context().piecewise; + const size_t seq_len = staging.attn_output->size(1); + const size_t valid_len = piecewise.valid_seq_len > 0 ? piecewise.valid_seq_len : seq_len; + auto hidden_narrow = hidden_states->narrow({{1, 0, valid_len}}); + auto attn_for_proj = staging.attn_output->narrow({{1, 0, valid_len}}); + auto projected = o_proj_->forward(attn_for_proj); + hidden_narrow->copy_from(projected); + if (valid_len < seq_len) { + auto hidden_tail = hidden_states->narrow({{1, valid_len, seq_len - valid_len}}); + set_zeros(hidden_tail); + } +} + +void Attention::forward_post_attn_piecewise_allreduce_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + (void)hidden_states; + (void)staging; + // TP>1 deferred allreduce not ported on post-refactor RowParallelLinear. +} + void init_kv_cache_quant_params(std::function register_fn, const infinicore::Device &device, infinicore::nn::Parameter &kv_cache_k_scale, diff --git a/csrc/layers/attention/attention.hpp b/csrc/layers/attention/attention.hpp index 671b90f5c..912791f12 100644 --- a/csrc/layers/attention/attention.hpp +++ b/csrc/layers/attention/attention.hpp @@ -3,6 +3,7 @@ #include "../../backends/attention_backends.hpp" #include "../../config/model_config.hpp" #include "../../global_state/global_state.hpp" +#include "../../global_state/piecewise_prefill_state.hpp" #include "../linear/linear.hpp" #include "backends/attention_layer.hpp" #include "infinicore/nn/module.hpp" @@ -20,6 +21,26 @@ class Attention : public infinicore::nn::Module { infinicore::Tensor forward(const infinicore::Tensor &positions, const infinicore::Tensor &hidden_states) const; + /// Piecewise prefill: QKV + RoPE into staging. + void forward_pre_attn_piecewise(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + void forward_pre_attn_qkv_piecewise(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + void forward_pre_attn_rope_piecewise(const infinicore::Tensor &positions, + global_state::PiecewiseLayerStaging &staging) const; + void forward_eager_attn_piecewise(const infinicore::Tensor &positions, + global_state::PiecewiseLayerStaging &staging) const; + void forward_post_attn_piecewise_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + void forward_post_attn_piecewise_graph_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + void forward_post_attn_piecewise_cg_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + void forward_post_attn_piecewise_allreduce_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + void process_weights_after_loading() override { qkv_proj_->process_weights_after_loading(); } diff --git a/csrc/layers/attention/backends/flash_attn.cpp b/csrc/layers/attention/backends/flash_attn.cpp index e37c513f6..c1e9f59e2 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, @@ -51,12 +53,19 @@ infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, auto attn_output = infinicore::Tensor::empty( {seq_len, num_heads_, head_dim_}, query->dtype(), query->device()); if (is_prefill) { - const size_t max_query_length = attn_metadata.max_query_length > 0 - ? attn_metadata.max_query_length - : max_position_embeddings_; - const size_t max_sequence_length = attn_metadata.max_sequence_length > 0 - ? attn_metadata.max_sequence_length - : max_position_embeddings_; + // Prefer batch max from metadata; never fall back to max_position_embeddings + // (MiniCPM5=131072) which exceeds block_table capacity and breaks FA. + int max_seqlen_q = attn_metadata.max_query_length > 0 + ? static_cast(attn_metadata.max_query_length) + : static_cast(seq_len); + int max_seqlen_k = attn_metadata.max_sequence_length > 0 + ? static_cast(attn_metadata.max_sequence_length) + : max_seqlen_q; + if (block_tables.value()->ndim() >= 2) { + const int table_cap = static_cast(block_tables.value()->size(1) * 256); + max_seqlen_k = std::min(max_seqlen_k, table_cap); + max_seqlen_q = std::min(max_seqlen_q, table_cap); + } infinicore::op::mha_varlen_( attn_output, query, @@ -65,8 +74,8 @@ infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, input_offsets.value(), cu_seqlens.value(), block_tables.value(), - max_query_length, - max_sequence_length, + max_seqlen_q, + max_seqlen_k, std::nullopt, scale_); } else { diff --git a/csrc/layers/attention/backends/static_attn.cpp b/csrc/layers/attention/backends/static_attn.cpp index 69af850f8..992ddad29 100644 --- a/csrc/layers/attention/backends/static_attn.cpp +++ b/csrc/layers/attention/backends/static_attn.cpp @@ -5,6 +5,8 @@ #include "infinicore/ops/per_tensor_dequant_i8.hpp" #include "infinicore/ops/per_tensor_quant_i8.hpp" +#include + namespace infinilm::layers::attention::backends { namespace { @@ -100,14 +102,42 @@ infinicore::Tensor StaticAttentionImpl::forward(const AttentionLayer &layer, auto K = k_total->view({batch_size * num_kv_heads_, total_seq_len, head_dim_}); auto V = v_total->view({batch_size * num_kv_heads_, total_seq_len, value_head_dim}); - auto K_transposed = K->permute({0, 2, 1}); // [bs * n_kv_head, head_dim, total_seq_len] - - auto attn_weight = infinicore::op::matmul(Q, K_transposed, scale_); // [bs * n_kv_head, ng * seq_len, total_seq_len] - - auto attn_weight_softmax = attn_weight->view({batch_size * num_heads_, seq_len, total_seq_len}); - infinicore::op::causal_softmax_(attn_weight_softmax, attn_weight_softmax); - - auto out = infinicore::op::matmul(attn_weight, V); // [bs * n_kv_head, ng * seq_len, value_head_dim] + // dl static-attn GEMM/softmax launch fails around ~600+ query tokens + // ("too many resources requested for launch"). Chunk Q along the sequence + // so each causal matmul stays small; FA/paged remains M2. + constexpr size_t kQyAttnChunkTokens = 128; + const bool chunk_qy = q_reshaped->device().type() == infinicore::Device::Type::kQy + && (seq_len > kQyAttnChunkTokens || total_seq_len > 512); + + auto run_chunk = [&](size_t q_begin, size_t q_rows) -> infinicore::Tensor { + // Queries at absolute positions [cache_pos + q_begin, ...); keys needed + // are only [0, cache_pos + q_begin + q_rows) under causal masking. + const size_t abs_begin = cache_pos + q_begin; + const size_t kv_len = abs_begin + q_rows; + auto Q_chunk = Q->narrow({{1, ngroup * q_begin, ngroup * q_rows}}); + auto K_chunk = K->narrow({{1, 0, kv_len}}); + auto V_chunk = V->narrow({{1, 0, kv_len}}); + auto K_T = K_chunk->permute({0, 2, 1}); + auto attn_weight = infinicore::op::matmul(Q_chunk, K_T, scale_); + auto attn_weight_softmax = attn_weight->view({batch_size * num_heads_, q_rows, kv_len}); + infinicore::op::causal_softmax_(attn_weight_softmax, attn_weight_softmax); + return infinicore::op::matmul(attn_weight, V_chunk); + }; + + infinicore::Tensor out; + if (!chunk_qy) { + out = run_chunk(/*q_begin=*/0, seq_len); + } else { + out = infinicore::Tensor::empty( + {batch_size * num_kv_heads_, ngroup * seq_len, value_head_dim}, + Q->dtype(), + Q->device()); + for (size_t begin = 0; begin < seq_len; begin += kQyAttnChunkTokens) { + const size_t rows = std::min(kQyAttnChunkTokens, seq_len - begin); + auto chunk_out = run_chunk(begin, rows); + out->narrow({{1, ngroup * begin, ngroup * rows}})->copy_from(chunk_out); + } + } return out->view({batch_size, num_heads_, seq_len, value_head_dim}) ->permute({0, 2, 1, 3}) diff --git a/csrc/layers/causal_lm_templates/piecewise_text_causal_lm.hpp b/csrc/layers/causal_lm_templates/piecewise_text_causal_lm.hpp new file mode 100644 index 000000000..6f2dff7e9 --- /dev/null +++ b/csrc/layers/causal_lm_templates/piecewise_text_causal_lm.hpp @@ -0,0 +1,100 @@ +#pragma once + +#include "text_causal_lm.hpp" + +namespace infinilm::layers::causal_lm_templates { + +/// TextCausalLM with native piecewise CUDAGraph hooks (TextDecoderLayer-based models). +template +class PiecewiseTextCausalLM : public TextCausalLM { +public: + using TextCausalLM::TextCausalLM; + + bool supports_native_piecewise_prefill() const override { return true; } + + size_t native_piecewise_num_layers() const override { + return this->model_->num_layers(); + } + + void native_piecewise_embed(const InfinilmModel::Input &input, + infinicore::Tensor &hidden_states) const override { + this->model_->piecewise_embed(input, hidden_states); + } + + void native_piecewise_pre_attn_layer(size_t layer_idx, + const InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const override { + this->model_->piecewise_pre_attn_layer(layer_idx, input, hidden_states, residual); + } + + void native_piecewise_pre_attn_qkv_layer(size_t layer_idx, + const InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const override { + this->model_->piecewise_pre_attn_qkv_layer(layer_idx, input, hidden_states, residual); + } + + void native_piecewise_pre_attn_rope_layer(size_t layer_idx, + const InfinilmModel::Input &input) const override { + this->model_->piecewise_pre_attn_rope_layer(layer_idx, input); + } + + void native_piecewise_eager_attn_layer(size_t layer_idx, + const InfinilmModel::Input &input) const override { + this->model_->piecewise_eager_attn_layer(layer_idx, input); + } + + void native_piecewise_post_attn_layer(size_t layer_idx, + const InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const override { + this->model_->piecewise_post_attn_layer(layer_idx, input, hidden_states, residual); + } + + void native_piecewise_post_attn_cg_layer(size_t layer_idx, + const InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const override { + this->model_->piecewise_post_attn_cg_layer(layer_idx, input, hidden_states, residual); + } + + void native_piecewise_post_attn_decode_cg_layer(size_t layer_idx, + const InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const override { + this->model_->piecewise_post_attn_decode_cg_layer(layer_idx, input, hidden_states, residual); + } + + void native_piecewise_eager_moe_layer(size_t layer_idx, + const InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const override { + this->model_->piecewise_eager_moe_layer(layer_idx, input, hidden_states, residual); + } + + void native_piecewise_post_attn_graph_layer(size_t layer_idx, + const InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const override { + this->model_->piecewise_post_attn_graph_layer(layer_idx, input, hidden_states, residual); + } + + void native_piecewise_post_attn_allreduce_layer(size_t layer_idx, + const InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const override { + this->model_->piecewise_post_attn_allreduce_layer(layer_idx, input, hidden_states, residual); + } + + void native_piecewise_lm_head(const InfinilmModel::Input &, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + infinicore::Tensor &logits_out) const override { + this->model_->piecewise_lm_head(hidden_states, residual); + auto logits = this->lm_head_->forward(hidden_states); + logits_out->copy_from(logits); + } +}; + +} // namespace infinilm::layers::causal_lm_templates diff --git a/csrc/layers/causal_lm_templates/text_decoder_layer.hpp b/csrc/layers/causal_lm_templates/text_decoder_layer.hpp index c01afcf10..6dd86868c 100644 --- a/csrc/layers/causal_lm_templates/text_decoder_layer.hpp +++ b/csrc/layers/causal_lm_templates/text_decoder_layer.hpp @@ -1,6 +1,9 @@ #pragma once #include "../../config/model_config.hpp" +#include "../../global_state/global_state.hpp" +#include "../../global_state/piecewise_inductor_flags.hpp" +#include "../../global_state/piecewise_prefill_state.hpp" #include "infinicore/device.hpp" #include "infinicore/nn/module.hpp" #include "infinicore/nn/rmsnorm.hpp" @@ -62,6 +65,73 @@ class TextDecoderLayer : public infinicore::nn::Module { size_t layer_idx() const { return layer_idx_; } + void piecewise_pre_attn(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + (void)global_state::piecewise_inductor_segment_enabled(); + piecewise_pre_attn_qkv(positions, hidden_states, residual, staging); + piecewise_pre_attn_rope(positions, staging); + } + + void piecewise_pre_attn_qkv(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + input_layernorm_->forward_inplace(hidden_states, residual); + self_attn_->forward_pre_attn_qkv_piecewise(positions, hidden_states, staging); + } + + void piecewise_pre_attn_rope(const infinicore::Tensor &positions, + global_state::PiecewiseLayerStaging &staging) const { + self_attn_->forward_pre_attn_rope_piecewise(positions, staging); + } + + void piecewise_eager_attn(const infinicore::Tensor &positions, + global_state::PiecewiseLayerStaging &staging) const { + self_attn_->forward_eager_attn_piecewise(positions, staging); + } + + void piecewise_post_attn(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + piecewise_post_attn_cg(hidden_states, residual, staging); + } + + void piecewise_post_attn_cg(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + self_attn_->forward_post_attn_piecewise_cg_into(hidden_states, staging); + post_attention_layernorm_->forward_inplace(hidden_states, residual); + auto mlp_out = mlp_->forward(hidden_states); + hidden_states->copy_from(mlp_out); + } + + /// Dense MLP models capture the full post segment (no MoE break). + void piecewise_post_attn_decode_cg(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + piecewise_post_attn_cg(hidden_states, residual, staging); + } + + void piecewise_eager_moe(infinicore::Tensor &, + infinicore::Tensor &, + global_state::PiecewiseLayerStaging &) const {} + + void piecewise_post_attn_graph(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + // Without deferred allreduce split, reuse the fused CG post path. + piecewise_post_attn_cg(hidden_states, residual, staging); + } + + void piecewise_post_attn_allreduce(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + (void)residual; + self_attn_->forward_post_attn_piecewise_allreduce_into(hidden_states, staging); + } + protected: INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); diff --git a/csrc/layers/causal_lm_templates/text_model.hpp b/csrc/layers/causal_lm_templates/text_model.hpp index bafe98d14..eee57de50 100644 --- a/csrc/layers/causal_lm_templates/text_model.hpp +++ b/csrc/layers/causal_lm_templates/text_model.hpp @@ -2,7 +2,9 @@ #include "../../config/model_config.hpp" #include "../../global_state/global_state.hpp" +#include "../../global_state/piecewise_prefill_state.hpp" #include "../../models/infinilm_model.hpp" +#include "../../utils.hpp" #include "infinicore/nn/embedding.hpp" #include "infinicore/nn/rmsnorm.hpp" #include "infinicore/ops.hpp" @@ -130,6 +132,106 @@ class TextModel : public infinicore::nn::Module { return embed_tokens_->forward(input_ids); } + size_t num_layers() const { return layers_.size(); } + + void piecewise_embed(const infinilm::InfinilmModel::Input &input, + infinicore::Tensor &hidden_states) const { + auto &piecewise = infinilm::global_state::get_forward_context().piecewise; + const size_t bucket = hidden_states->size(1); + const size_t valid_len = piecewise.valid_seq_len > 0 ? piecewise.valid_seq_len : bucket; + auto embedded = embed_tokens_->forward(input.input_ids.value()); + hidden_states->narrow({{1, 0, valid_len}})->copy_from(embedded->narrow({{1, 0, valid_len}})); + if (valid_len < bucket) { + auto tail = hidden_states->narrow({{1, valid_len, bucket - valid_len}}); + set_zeros(tail); + } + } + + void piecewise_pre_attn_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_pre_attn( + input.position_ids.value(), hidden_states, residual, staging); + } + + void piecewise_pre_attn_qkv_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_pre_attn_qkv( + input.position_ids.value(), hidden_states, residual, staging); + } + + void piecewise_pre_attn_rope_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &input) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_pre_attn_rope(input.position_ids.value(), staging); + } + + void piecewise_eager_attn_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &input) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_eager_attn(input.position_ids.value(), staging); + } + + void piecewise_post_attn_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_post_attn(hidden_states, residual, staging); + } + + void piecewise_post_attn_cg_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_post_attn_cg(hidden_states, residual, staging); + } + + void piecewise_post_attn_decode_cg_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_post_attn_decode_cg(hidden_states, residual, staging); + } + + void piecewise_eager_moe_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_eager_moe(hidden_states, residual, staging); + } + + void piecewise_post_attn_graph_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_post_attn_graph(hidden_states, residual, staging); + } + + void piecewise_post_attn_allreduce_layer(size_t layer_idx, + const infinilm::InfinilmModel::Input &, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + auto &staging = infinilm::global_state::get_forward_context().piecewise.layer_staging.at(layer_idx); + layers_.at(layer_idx)->piecewise_post_attn_allreduce(hidden_states, residual, staging); + } + + infinicore::Tensor piecewise_lm_head(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + norm_->forward_inplace(hidden_states, residual); + return hidden_states; + } + + protected: INFINICORE_NN_MODULE(infinicore::nn::Embedding, embed_tokens); INFINICORE_NN_MODULE_VEC(DecoderLayer, layers); diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index e33e16495..ba47ee22b 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -4,6 +4,9 @@ #include "infinicore/ops/moe_align.hpp" #include "infinicore/ops/moe_fused_dense.hpp" +#include +#include +#include #include #include #include @@ -76,17 +79,53 @@ void check_packed_weight_tensor(const infinicore::Tensor &tensor, CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, const MoeWeights &weights, MoeWorkspace &workspace) const { - auto runner_input = prepare_runner_input( - dispatch_output, - workspace); + // dl fused MoE materializes [M, topk, H] intermediates. A 7k-token + // prefill OOMs a 32GB KS38; vLLM chunks, but its 16k chunk is still too + // big here. Run ≤1024-token chunks on QY (graph capture stays one shot). + constexpr size_t kQyChunkTokens = 1024; + const auto &hidden = dispatch_output.hidden_states; + const size_t num_tokens = hidden ? hidden->shape()[0] : 0; + const bool chunk_qy = hidden + && hidden->device().type() == infinicore::Device::Type::kQy + && !infinicore::context::isGraphRecording() + && num_tokens > kQyChunkTokens + && hidden->shape().size() == 2; + if (!chunk_qy) { + auto runner_input = prepare_runner_input(dispatch_output, workspace); + auto runner_output = run_fused_core(runner_input, weights, workspace); + return CombineInput{ + CombineInputFormat::Standard, + runner_output.hidden_states, + runner_input.topk_output, + runner_input.routing_metadata, + }; + } - auto runner_output = run_fused_core(runner_input, weights, workspace); + ensure_tensor( + workspace.fused_moe_output, + hidden->shape(), + hidden->dtype(), + hidden->device()); + workspace.fused_moe_output_tokens_capacity = num_tokens; + const auto &topk = dispatch_output.topk_output; + for (size_t begin = 0; begin < num_tokens; begin += kQyChunkTokens) { + const size_t rows = std::min(kQyChunkTokens, num_tokens - begin); + DispatchOutput slice = dispatch_output; + slice.hidden_states = hidden->narrow({{0, begin, rows}})->contiguous(); + slice.topk_output.topk_weights = topk.topk_weights->narrow({{0, begin, rows}})->contiguous(); + slice.topk_output.topk_ids = topk.topk_ids->narrow({{0, begin, rows}})->contiguous(); + // Fresh align buffers so block_size recovered from numel matches this chunk. + MoeWorkspace chunk_workspace; + auto runner_input = prepare_runner_input(slice, chunk_workspace); + auto runner_output = run_fused_core(runner_input, weights, chunk_workspace); + workspace.fused_moe_output->narrow({{0, begin, rows}})->copy_from(runner_output.hidden_states); + } return CombineInput{ CombineInputFormat::Standard, - runner_output.hidden_states, - runner_input.topk_output, - runner_input.routing_metadata, + workspace.fused_moe_output, + dispatch_output.topk_output, + MoeRoutingMetadata{}, }; } diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp index 5d284a316..5fc88d294 100644 --- a/csrc/models/infinilm_model.cpp +++ b/csrc/models/infinilm_model.cpp @@ -118,4 +118,84 @@ void InfinilmModel::reset_runtime_state_recursive_(const infinicore::nn::Module module->reset_runtime_state(); } + +void InfinilmModel::native_piecewise_embed(const Input &, + infinicore::Tensor &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_embed not supported"); +} + +void InfinilmModel::native_piecewise_pre_attn_layer(size_t, + const Input &, + infinicore::Tensor &, + infinicore::Tensor &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_pre_attn_layer not supported"); +} + +void InfinilmModel::native_piecewise_pre_attn_qkv_layer(size_t, + const Input &, + infinicore::Tensor &, + infinicore::Tensor &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_pre_attn_qkv_layer not supported"); +} + +void InfinilmModel::native_piecewise_pre_attn_rope_layer(size_t, + const Input &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_pre_attn_rope_layer not supported"); +} + +void InfinilmModel::native_piecewise_eager_attn_layer(size_t, + const Input &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_eager_attn_layer not supported"); +} + +void InfinilmModel::native_piecewise_post_attn_layer(size_t, + const Input &, + infinicore::Tensor &, + infinicore::Tensor &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_post_attn_layer not supported"); +} + +void InfinilmModel::native_piecewise_post_attn_cg_layer(size_t, + const Input &, + infinicore::Tensor &, + infinicore::Tensor &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_post_attn_cg_layer not supported"); +} + +void InfinilmModel::native_piecewise_post_attn_decode_cg_layer(size_t, + const Input &, + infinicore::Tensor &, + infinicore::Tensor &) const { + throw std::runtime_error( + "InfinilmModel: native_piecewise_post_attn_decode_cg_layer not supported"); +} + +void InfinilmModel::native_piecewise_eager_moe_layer(size_t, + const Input &, + infinicore::Tensor &, + infinicore::Tensor &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_eager_moe_layer not supported"); +} + +void InfinilmModel::native_piecewise_post_attn_graph_layer(size_t, + const Input &, + infinicore::Tensor &, + infinicore::Tensor &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_post_attn_graph_layer not supported"); +} + +void InfinilmModel::native_piecewise_post_attn_allreduce_layer(size_t, + const Input &, + infinicore::Tensor &, + infinicore::Tensor &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_post_attn_allreduce_layer not supported"); +} + +void InfinilmModel::native_piecewise_lm_head(const Input &, + infinicore::Tensor &, + infinicore::Tensor &, + infinicore::Tensor &) const { + throw std::runtime_error("InfinilmModel: native_piecewise_lm_head not supported"); +} + } // namespace infinilm diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index 02677318d..2ee8aa286 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -6,6 +6,7 @@ #include "infinicore/nn/module.hpp" #include "infinicore/tensor.hpp" +#include #include #include @@ -57,8 +58,20 @@ class InfinilmModel : public infinicore::nn::Module { std::optional target_hidden_states; /// Preserve logits for every packed position for speculative/MTP callers. bool sample_all_positions{false}; + /// Per-request final-chunk flags. Empty means all rows are final (lm_head + sampling). + std::vector is_final_prefill_chunk; }; + static bool any_final_prefill_chunk(const std::vector &flags) { + return flags.empty() + || std::any_of(flags.begin(), flags.end(), [](bool v) { return v; }); + } + + static bool has_nonfinal_prefill_chunk(const std::vector &flags) { + return !flags.empty() + && std::any_of(flags.begin(), flags.end(), [](bool v) { return !v; }); + } + struct Output { /// Logits. infinicore::Tensor logits; @@ -76,6 +89,52 @@ class InfinilmModel : public infinicore::nn::Module { return model_config_; } + /// Native piecewise CUDAGraph prefill. Default: unsupported. + virtual bool supports_native_piecewise_prefill() const { return false; } + virtual size_t native_piecewise_num_layers() const { return 0; } + virtual void native_piecewise_embed(const Input &input, + infinicore::Tensor &hidden_states) const; + virtual void native_piecewise_pre_attn_layer(size_t layer_idx, + const Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + virtual void native_piecewise_pre_attn_qkv_layer(size_t layer_idx, + const Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + virtual void native_piecewise_pre_attn_rope_layer(size_t layer_idx, + const Input &input) const; + virtual void native_piecewise_eager_attn_layer(size_t layer_idx, + const Input &input) const; + virtual void native_piecewise_post_attn_layer(size_t layer_idx, + const Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + virtual void native_piecewise_post_attn_cg_layer(size_t layer_idx, + const Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + virtual void native_piecewise_post_attn_decode_cg_layer(size_t layer_idx, + const Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + virtual void native_piecewise_eager_moe_layer(size_t layer_idx, + const Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + virtual void native_piecewise_post_attn_graph_layer(size_t layer_idx, + const Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + virtual void native_piecewise_post_attn_allreduce_layer(size_t layer_idx, + const Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + virtual void native_piecewise_lm_head(const Input &input, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + infinicore::Tensor &logits_out) const; + void process_weights_after_loading(); void reset_runtime_state() const; diff --git a/csrc/models/llama/llama_for_causal_lm.hpp b/csrc/models/llama/llama_for_causal_lm.hpp index e58e062ed..c7fc1f754 100644 --- a/csrc/models/llama/llama_for_causal_lm.hpp +++ b/csrc/models/llama/llama_for_causal_lm.hpp @@ -1,5 +1,6 @@ #pragma once +#include "../../layers/causal_lm_templates/piecewise_text_causal_lm.hpp" #include "../../layers/common_modules.hpp" #include @@ -13,7 +14,7 @@ using LlamaDecoderLayer = infinilm::layers::causal_lm_templates::TextDecoderLaye using LlamaModel = infinilm::layers::causal_lm_templates::TextModel; -using LlamaForCausalLM = infinilm::layers::causal_lm_templates::TextCausalLM; +using LlamaForCausalLM = infinilm::layers::causal_lm_templates::PiecewiseTextCausalLM; } // namespace infinilm::models::llama diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_attention.cpp b/csrc/models/minicpm5_moe/minicpm5_moe_attention.cpp new file mode 100644 index 000000000..778f96749 --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_attention.cpp @@ -0,0 +1,468 @@ +#include "minicpm5_moe_attention.hpp" + +#include "../../global_state/piecewise_prefill_state.hpp" + +#include "../../global_state/global_state.hpp" +#include "../../layers/attention/attention.hpp" +#include "../../utils.hpp" +#include "infinicore/context/context.hpp" + +#include +#include +#include +#include + +namespace infinilm::models::minicpm5_moe { + +MiniCPM5MoEAttention::MiniCPM5MoEAttention( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + layer_idx_ = layer_idx; + const auto &dtype{model_config->get_dtype()}; + hidden_size_ = model_config->get("hidden_size"); + head_dim_ = model_config->get("head_dim"); + use_gated_attention_ = model_config->get_or("use_gated_attention", true); + + size_t total_num_heads = model_config->get("num_attention_heads"); + size_t total_num_kv_heads = model_config->get("num_key_value_heads"); + bool use_bias = model_config->get_or("attention_bias", false); + bool use_output_bias = model_config->get_or("attention_output_bias", false); + + attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend; + const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + int tp_rank = infinilm::global_state::get_tensor_model_parallel_rank(); + int tp_size = infinilm::global_state::get_tensor_model_parallel_world_size(); + + num_attention_heads_ = total_num_heads / tp_size; + num_key_value_heads_ = total_num_kv_heads < static_cast(tp_size) + ? 1 + : total_num_kv_heads / tp_size; + + // Gated path doubles q_proj (per-head [q | gate]), same packing as qwen3_next. + const size_t q_heads_for_proj = total_num_heads * (use_gated_attention_ ? 2 : 1); + + auto quantization_method = model_config->get_quantization_method(); + auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { + this->register_parameter(n, std::move(p)); + }; + qkv_proj_ = std::make_shared( + hidden_size_, head_dim_, q_heads_for_proj, total_num_kv_heads, + "q_proj", "k_proj", "v_proj", register_fn, + quantization_method, use_bias, dtype, device, rank_info); + o_proj_ = this->register_module( + "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method, + use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm); + + rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); + + float scaling = 1.0f / std::sqrt(static_cast(head_dim_)); + attn_ = std::make_shared( + num_attention_heads_, head_dim_, scaling, num_key_value_heads_, layer_idx_, + kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_); + + infinilm::layers::attention::init_kv_cache_quant_params( + register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_); +} + +infinicore::Tensor MiniCPM5MoEAttention::forward( + const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const { + if (::infinilm::backends::AttentionBackend::STATIC_ATTN == attention_backend_) { + return forward_static_(positions, hidden_states); + } + return forward_paged_(positions, hidden_states); +} + +infinicore::Tensor MiniCPM5MoEAttention::forward_static_( + const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + auto shape = hidden_states->shape(); + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + + auto [q_proj_out, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + + infinicore::Tensor q_reshaped; + infinicore::Tensor gate; + if (use_gated_attention_) { + auto q_proj_heads = q_proj_out->as_strided( + {batch_size * seq_len, num_attention_heads_, head_dim_ * 2}, + {q_proj_out->stride(1), static_cast(head_dim_ * 2), 1}); + auto q = q_proj_heads->narrow({{2, 0, head_dim_}}); + gate = q_proj_heads->narrow({{2, head_dim_, head_dim_}}) + ->contiguous() + ->view({batch_size, seq_len, num_attention_heads_ * head_dim_}); + // Keep real gated strides for in-place QY RoPE; view as 4-D for static attn. + q_reshaped = q->as_strided( + {batch_size, seq_len, num_attention_heads_, head_dim_}, + {q_proj_out->stride(0), q_proj_out->stride(1), + static_cast(head_dim_ * 2), 1}); + } else { + q_reshaped = q_proj_out->view({batch_size, seq_len, num_attention_heads_, head_dim_}); + } + + auto k_reshaped = k->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + auto v_reshaped = v->view({batch_size, seq_len, num_key_value_heads_, head_dim_}); + + auto pos_shape = position_ids->shape(); + infinicore::Tensor pos_ids_for_rope = position_ids; + if (pos_shape.size() == 2) { + auto pos_narrowed = position_ids->narrow({{0, 0, 1}}); + pos_ids_for_rope = pos_narrowed->contiguous()->view({pos_shape[1]}); + } else if (pos_shape.size() == 1) { + pos_ids_for_rope = position_ids->contiguous(); + } else { + throw std::runtime_error( + "MiniCPM5MoEAttention: Unexpected position_ids shape"); + } + + rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + + auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); + if (use_gated_attention_) { + attn_output = infinicore::op::mul(attn_output, infinicore::op::sigmoid(gate)); + } + return o_proj_->forward(attn_output); +} + +infinicore::Tensor MiniCPM5MoEAttention::forward_paged_( + const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const { + auto hidden_states_mutable = hidden_states; + auto shape = hidden_states->shape(); + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + + ASSERT_EQ(batch_size, 1); + + auto [q_proj_out, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + + infinicore::Tensor q_reshaped; + infinicore::Tensor gate; + if (use_gated_attention_) { + // Keep strided gated Q [seq,heads,dim] strides [8704,256,1]; QY RoPE + // rotates in-place via rope_strided_launch (no pack/contiguous tax). + auto q_proj_heads = q_proj_out->as_strided( + {seq_len, num_attention_heads_, head_dim_ * 2}, + {q_proj_out->stride(1), static_cast(head_dim_ * 2), 1}); + auto q = q_proj_heads->narrow({{2, 0, head_dim_}}); + gate = q_proj_heads->narrow({{2, head_dim_, head_dim_}}) + ->contiguous() + ->view({seq_len, num_attention_heads_ * head_dim_}); + q_reshaped = q; + } else { + q_reshaped = q_proj_out->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_}); + + auto pos_shape = position_ids->shape(); + infinicore::Tensor pos_ids_for_rope = position_ids; + if (pos_shape.size() == 2) { + auto pos_narrowed = position_ids->narrow({{0, 0, 1}}); + pos_ids_for_rope = pos_narrowed->view({pos_shape[1]}); + } else if (pos_shape.size() == 1) { + pos_ids_for_rope = position_ids; + } else { + throw std::runtime_error("MiniCPM5MoEAttention: Unexpected position_ids shape"); + } + + rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + + auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); + if (use_gated_attention_) { + attn_output = infinicore::op::mul( + attn_output, infinicore::op::sigmoid(gate)->view(attn_output->shape())); + } + return o_proj_->forward(attn_output); +} + +namespace { + +bool same_trailing_dims_(const infinicore::Shape &a, const infinicore::Shape &b) { + if (a.size() != b.size() || a.empty()) { + return false; + } + for (size_t i = 1; i < a.size(); ++i) { + if (a[i] != b[i]) { + return false; + } + } + return true; +} + +void ensure_same_shape_buf(infinicore::Tensor &buf, const infinicore::Tensor &proto) { + if (buf && buf->shape() == proto->shape() && buf->dtype() == proto->dtype() + && buf->device() == proto->device()) { + return; + } + buf = infinicore::Tensor::empty(proto->shape(), proto->dtype(), proto->device()); +} + +void store_gate_score(infinicore::Tensor &gate_buf, + infinicore::Tensor &gate_write_view, + const infinicore::Tensor &gate) { + // Grow-only base buffer: multi-bucket piecewise CG records copy destinations by + // device pointer. Shrinking (exact-shape realloc) after bucket=2048 → 16 made + // earlier graphs write into freed/undersized memory → EM collapse / garbage. + // During CG: only copy_from into a write view prepared outside recording + // (narrow/empty HostOps inside startGraphRecording kill MetaX decode CG). + const bool recording = infinicore::context::isGraphRecording(); + infinicore::Tensor g = gate->is_contiguous() ? gate : gate->contiguous(); + const auto &gshape = g->shape(); + + const bool need_alloc = !gate_buf || gate_buf->dtype() != g->dtype() + || gate_buf->device() != g->device() + || !same_trailing_dims_(gate_buf->shape(), gshape) + || gate_buf->shape()[0] < gshape[0]; + + if (need_alloc) { + if (recording) { + throw std::runtime_error( + "MiniCPM5MoEAttention: gate_score_cache_ grow required during CG " + "(dry-run must grow-only allocate via C++ QKV path)"); + } + auto alloc_shape = gshape; + if (gate_buf && same_trailing_dims_(gate_buf->shape(), gshape) + && gate_buf->dtype() == g->dtype() && gate_buf->device() == g->device()) { + alloc_shape[0] = std::max(gate_buf->shape()[0], gshape[0]); + } + gate_buf = infinicore::Tensor::empty(alloc_shape, g->dtype(), g->device()); + gate_write_view = infinicore::Tensor{}; + } + + const bool view_ok = gate_write_view && gate_write_view->shape() == gshape + && gate_write_view->dtype() == g->dtype() + && gate_write_view->device() == g->device(); + if (!view_ok) { + if (recording) { + throw std::runtime_error( + "MiniCPM5MoEAttention: gate_score_write_view_ unset/mismatch during CG " + "(dry-run must narrow write view outside recording)"); + } + if (gate_buf->shape() == gshape) { + gate_write_view = gate_buf; + } else { + gate_write_view = gate_buf->narrow({{0, 0, gshape[0]}}); + } + } + gate_write_view->copy_from(g); +} + +infinicore::Tensor apply_gate_sigmoid_mul(const infinicore::Tensor &attn_output, + const infinicore::Tensor &gate_score, + infinicore::Tensor &gate_sigmoid_buf) { + auto gate_view = gate_score; + if (attn_output->shape().size() == 3 && gate_score->shape().size() == 2 + && attn_output->shape()[0] == 1) { + gate_view = gate_score->view({1, gate_score->shape()[0], gate_score->shape()[1]}); + } + auto gate_in = gate_view->contiguous(); + ensure_same_shape_buf(gate_sigmoid_buf, gate_in); + gate_sigmoid_buf = infinicore::op::sigmoid(gate_in); + return infinicore::op::mul(attn_output->contiguous(), gate_sigmoid_buf); +} + +} // namespace + +void MiniCPM5MoEAttention::forward_pre_attn_qkv_piecewise( + const infinicore::Tensor &, + const infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + auto &piecewise = global_state::get_forward_context().piecewise; + auto hidden_states_mutable = hidden_states; + auto shape = hidden_states->shape(); + size_t batch_size = shape[0]; + size_t seq_len = shape[1]; + ASSERT_EQ(batch_size, 1); + + auto [qg, k, v] = qkv_proj_->forward_split(hidden_states_mutable); + const size_t valid_len = piecewise.valid_seq_len > 0 ? piecewise.valid_seq_len : seq_len; + + infinicore::Tensor q_heads; + if (use_gated_attention_) { + // Gate extract into stable per-layer cache (device copy; dry-run allocates). + // Gate *apply* stays with eager attn (host-break). + auto qg_view = qg->view({1, seq_len, num_attention_heads_, 2 * head_dim_}); + auto q = qg_view->narrow({{3, 0, head_dim_}})->contiguous(); + auto gate = qg_view->narrow({{3, head_dim_, head_dim_}})->contiguous(); + store_gate_score(gate_score_cache_, + gate_score_write_view_, + gate->view({seq_len, num_attention_heads_ * head_dim_})); + q_heads = q->view({1, seq_len, num_attention_heads_, head_dim_}); + } else { + q_heads = qg->view({1, seq_len, num_attention_heads_, head_dim_}); + } + auto k_heads = k->view({1, seq_len, num_key_value_heads_, head_dim_}); + auto v_heads = v->view({1, seq_len, num_key_value_heads_, head_dim_}); + + if (valid_len < seq_len) { + staging.q_rope->narrow({{1, 0, valid_len}})->copy_from(q_heads->narrow({{1, 0, valid_len}})); + staging.k_rope->narrow({{1, 0, valid_len}})->copy_from(k_heads->narrow({{1, 0, valid_len}})); + staging.v_rope->narrow({{1, 0, valid_len}})->copy_from(v_heads->narrow({{1, 0, valid_len}})); + auto q_tail = staging.q_rope->narrow({{1, valid_len, seq_len - valid_len}}); + auto k_tail = staging.k_rope->narrow({{1, valid_len, seq_len - valid_len}}); + auto v_tail = staging.v_rope->narrow({{1, valid_len, seq_len - valid_len}}); + set_zeros(q_tail); + set_zeros(k_tail); + set_zeros(v_tail); + } else { + staging.q_rope->copy_from(q_heads); + staging.k_rope->copy_from(k_heads); + staging.v_rope->copy_from(v_heads); + } +} + +void MiniCPM5MoEAttention::forward_pre_attn_rope_piecewise( + const infinicore::Tensor &position_ids, + global_state::PiecewiseLayerStaging &staging) const { + auto &piecewise = global_state::get_forward_context().piecewise; + const size_t seq_len = staging.q_rope->size(1); + const size_t valid_len = piecewise.valid_seq_len > 0 ? piecewise.valid_seq_len : seq_len; + // Same address-stable contract as Attention::forward_pre_attn_rope_piecewise: + // no alloc/rebind of position_ids / QK under recording; past/pad via contents. + const bool addr_stable = infinicore::context::isGraphRecording() + || piecewise.compile_capture_active; + + auto pos_shape = position_ids->shape(); + infinicore::Tensor pos_ids_for_rope = position_ids; + if (pos_shape.size() == 2) { + auto pos_narrowed = position_ids->narrow({{0, 0, 1}}); + if (addr_stable) { + pos_ids_for_rope = pos_narrowed->view({pos_shape[1]}); + } else { + pos_ids_for_rope = pos_narrowed->contiguous()->view({pos_shape[1]}); + } + } else if (pos_shape.size() == 1) { + if (!addr_stable && !position_ids->is_contiguous()) { + pos_ids_for_rope = position_ids->contiguous(); + } + } else { + throw std::runtime_error("MiniCPM5MoEAttention: Unexpected position_ids shape"); + } + if (pos_ids_for_rope->size(0) > valid_len) { + auto narrowed = pos_ids_for_rope->narrow({{0, 0, valid_len}}); + if (addr_stable) { + if (!narrowed->is_contiguous()) { + throw std::runtime_error( + "MiniCPM5MoEAttention::forward_pre_attn_rope_piecewise: " + "position_ids narrow not contiguous under graph recording"); + } + pos_ids_for_rope = narrowed; + } else { + pos_ids_for_rope = narrowed->contiguous(); + } + } + + auto q_rope = staging.q_rope->view({seq_len, num_attention_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + auto k_rope = staging.k_rope->view({seq_len, num_key_value_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + if (!q_rope->is_contiguous()) { + if (addr_stable) { + throw std::runtime_error( + "MiniCPM5MoEAttention::forward_pre_attn_rope_piecewise: " + "q_rope not contiguous under graph recording"); + } + q_rope = q_rope->contiguous(); + } + if (!k_rope->is_contiguous()) { + if (addr_stable) { + throw std::runtime_error( + "MiniCPM5MoEAttention::forward_pre_attn_rope_piecewise: " + "k_rope not contiguous under graph recording"); + } + k_rope = k_rope->contiguous(); + } + rotary_emb_->forward(q_rope, pos_ids_for_rope, true); + rotary_emb_->forward(k_rope, pos_ids_for_rope, true); +} + +void MiniCPM5MoEAttention::forward_pre_attn_piecewise( + const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + forward_pre_attn_qkv_piecewise(position_ids, hidden_states, staging); + forward_pre_attn_rope_piecewise(position_ids, staging); +} + +void MiniCPM5MoEAttention::forward_eager_attn_piecewise( + const infinicore::Tensor &, + global_state::PiecewiseLayerStaging &staging) const { + auto &piecewise = global_state::get_forward_context().piecewise; + const size_t seq_len = staging.q_rope->size(1); + const size_t valid_len = piecewise.valid_seq_len > 0 ? piecewise.valid_seq_len : seq_len; + + auto q = staging.q_rope->view({seq_len, num_attention_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + auto k = staging.k_rope->view({seq_len, num_key_value_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + auto v = staging.v_rope->view({seq_len, num_key_value_heads_, head_dim_})->narrow({{0, 0, valid_len}}); + if (!q->is_contiguous()) q = q->contiguous(); + if (!k->is_contiguous()) k = k->contiguous(); + if (!v->is_contiguous()) v = v->contiguous(); + + auto attn_output = attn_->forward(q, k, v); + if (use_gated_attention_ && gate_score_cache_) { + // Prefer grow-only cache + valid_len prefix: write_view_ may still be the + // last capture/eager seq (e.g. 16) after a larger bucket CG replay wrote + // into the shared base buffer. + auto gate = gate_score_cache_; + if (gate->shape().size() == 2 && gate->shape()[0] >= valid_len) { + gate = gate->narrow({{0, 0, valid_len}}); + } else if (gate_score_write_view_ && gate_score_write_view_->shape().size() == 2 + && gate_score_write_view_->shape()[0] >= valid_len) { + gate = gate_score_write_view_->narrow({{0, 0, valid_len}}); + } + attn_output = apply_gate_sigmoid_mul(attn_output, gate, gate_sigmoid_buf_); + } + + if (valid_len < seq_len) { + staging.attn_output->narrow({{1, 0, valid_len}})->copy_from(attn_output->narrow({{1, 0, valid_len}})); + auto attn_tail = staging.attn_output->narrow({{1, valid_len, seq_len - valid_len}}); + set_zeros(attn_tail); + } else { + staging.attn_output->copy_from(attn_output); + } +} + +void MiniCPM5MoEAttention::forward_post_attn_piecewise_into( + infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + forward_post_attn_piecewise_graph_into(hidden_states, staging); + forward_post_attn_piecewise_allreduce_into(hidden_states, staging); +} + +void MiniCPM5MoEAttention::forward_post_attn_piecewise_graph_into( + infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + // Dest RowParallelLinear has no matmul-only/allreduce split; reuse fused path. + forward_post_attn_piecewise_cg_into(hidden_states, staging); +} + +void MiniCPM5MoEAttention::forward_post_attn_piecewise_cg_into( + infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + auto &piecewise = global_state::get_forward_context().piecewise; + const size_t seq_len = staging.attn_output->size(1); + const size_t valid_len = piecewise.valid_seq_len > 0 ? piecewise.valid_seq_len : seq_len; + auto hidden_narrow = hidden_states->narrow({{1, 0, valid_len}}); + auto attn_for_proj = staging.attn_output->narrow({{1, 0, valid_len}}); + auto projected = o_proj_->forward(attn_for_proj); + hidden_narrow->copy_from(projected); + if (valid_len < seq_len) { + auto hidden_tail = hidden_states->narrow({{1, valid_len, seq_len - valid_len}}); + set_zeros(hidden_tail); + } +} + +void MiniCPM5MoEAttention::forward_post_attn_piecewise_allreduce_into( + infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const { + (void)hidden_states; + (void)staging; + // TP>1 deferred allreduce not ported on post-refactor RowParallelLinear. +} + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_attention.hpp b/csrc/models/minicpm5_moe/minicpm5_moe_attention.hpp new file mode 100644 index 000000000..4324f6440 --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_attention.hpp @@ -0,0 +1,90 @@ +#pragma once + +#include "../../layers/common_modules.hpp" +#include "../../global_state/piecewise_prefill_state.hpp" + +namespace infinilm::models::minicpm5_moe { + +/** + * GQA attention with optional gated output (doubled q_proj, per-head [q|gate]). + * Matches HF MiniCPM5MoEAttention / vllm_minicpm5 gated path (no q/k RMSNorm). + */ +class MiniCPM5MoEAttention : public infinicore::nn::Module { +public: + MiniCPM5MoEAttention(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states) const; + + + void forward_pre_attn_piecewise(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + + void forward_pre_attn_qkv_piecewise(const infinicore::Tensor &positions, + const infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + + void forward_pre_attn_rope_piecewise(const infinicore::Tensor &positions, + global_state::PiecewiseLayerStaging &staging) const; + + void forward_eager_attn_piecewise(const infinicore::Tensor &positions, + global_state::PiecewiseLayerStaging &staging) const; + + void forward_post_attn_piecewise_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + + void forward_post_attn_piecewise_graph_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + + void forward_post_attn_piecewise_cg_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + + void forward_post_attn_piecewise_allreduce_into(infinicore::Tensor &hidden_states, + global_state::PiecewiseLayerStaging &staging) const; + + void process_weights_after_loading() override { + qkv_proj_->process_weights_after_loading(); + } + + void reset_runtime_state() const override { + qkv_proj_->reset_runtime_state(); + } + + size_t layer_idx() const { return layer_idx_; } + size_t num_heads() const { return num_attention_heads_; } + size_t num_kv_heads() const { return num_key_value_heads_; } + size_t head_dim() const { return head_dim_; } + size_t hidden_size() const { return hidden_size_; } + +protected: + infinicore::Tensor forward_static_(const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const; + infinicore::Tensor forward_paged_(const infinicore::Tensor &position_ids, + const infinicore::Tensor &hidden_states) const; + + std::shared_ptr qkv_proj_; + std::shared_ptr o_proj_; + std::shared_ptr rotary_emb_; + + std::shared_ptr attn_; + ::infinilm::backends::AttentionBackend attention_backend_; + size_t layer_idx_; + size_t num_attention_heads_; + size_t num_key_value_heads_; + size_t hidden_size_; + size_t head_dim_; + bool use_gated_attention_{true}; + + INFINICORE_NN_PARAMETER(kv_cache_k_scale); + INFINICORE_NN_PARAMETER(kv_cache_v_scale); + + mutable infinicore::Tensor gate_score_cache_; + mutable infinicore::Tensor gate_score_write_view_; + mutable infinicore::Tensor gate_sigmoid_buf_; +}; + + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.cpp b/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.cpp new file mode 100644 index 000000000..38f51d917 --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.cpp @@ -0,0 +1,135 @@ +#include "minicpm5_moe_decoder_layer.hpp" + +#include "../../global_state/ar_profile.hpp" +#include "../../global_state/decode_phase_profile.hpp" +#include "../../global_state/global_state.hpp" +#include "../../global_state/piecewise_inductor_flags.hpp" +#include "infinicore/context/context.hpp" + +namespace infinilm::models::minicpm5_moe { + +MiniCPM5MoEDecoderLayer::MiniCPM5MoEDecoderLayer( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + layer_idx_ = layer_idx; + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const double rms_norm_eps = model_config->get("rms_norm_eps"); + INFINICORE_NN_MODULE_INIT(self_attn, model_config, layer_idx, device); + INFINICORE_NN_MODULE_INIT(input_layernorm, hidden_size, rms_norm_eps, dtype, device); + INFINICORE_NN_MODULE_INIT(post_attention_layernorm, hidden_size, rms_norm_eps, dtype, device); + + const size_t first_dense_layers = model_config->get_or("first_k_dense_replace", 0); + use_moe_ = model_config->get_or("n_routed_experts", 0) > 0 + && layer_idx >= first_dense_layers; + if (use_moe_) { + moe_mlp_ = this->register_module( + "mlp", model_config, layer_idx, device); + } else { + auto dense_config = make_minicpm5_mlp_config( + model_config, model_config->get("intermediate_size")); + dense_mlp_ = this->register_module( + "mlp", dense_config, device); + } +} + +infinicore::Tensor MiniCPM5MoEDecoderLayer::mlp_forward( + const infinicore::Tensor &hidden_states) const { + return use_moe_ ? moe_mlp_->forward(hidden_states) : dense_mlp_->forward(hidden_states); +} + +std::tuple +MiniCPM5MoEDecoderLayer::forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const { + input_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = self_attn_->forward(positions, hidden_states); + post_attention_layernorm_->forward_inplace(hidden_states, residual); + hidden_states = mlp_forward(hidden_states); + return {hidden_states, residual}; +} + +void MiniCPM5MoEDecoderLayer::piecewise_pre_attn( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + (void)global_state::piecewise_inductor_segment_enabled(); + piecewise_pre_attn_qkv(positions, hidden_states, residual, staging); + piecewise_pre_attn_rope(positions, staging); +} + +void MiniCPM5MoEDecoderLayer::piecewise_pre_attn_qkv( + const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + input_layernorm_->forward_inplace(hidden_states, residual); + self_attn_->forward_pre_attn_qkv_piecewise(positions, hidden_states, staging); +} + +void MiniCPM5MoEDecoderLayer::piecewise_pre_attn_rope( + const infinicore::Tensor &positions, + global_state::PiecewiseLayerStaging &staging) const { + self_attn_->forward_pre_attn_rope_piecewise(positions, staging); +} + +void MiniCPM5MoEDecoderLayer::piecewise_post_attn_cg( + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + self_attn_->forward_post_attn_piecewise_cg_into(hidden_states, staging); + post_attention_layernorm_->forward_inplace(hidden_states, residual); + auto mlp_out = mlp_forward(hidden_states); + hidden_states->copy_from(mlp_out); +} + +void MiniCPM5MoEDecoderLayer::piecewise_post_attn_decode_cg( + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + self_attn_->forward_post_attn_piecewise_cg_into(hidden_states, staging); + post_attention_layernorm_->forward_inplace(hidden_states, residual); + if (!use_moe_) { + auto mlp_out = dense_mlp_->forward(hidden_states); + hidden_states->copy_from(mlp_out); + } else if (infinicore::context::moeTritonCaptureAllowed()) { + auto mlp_out = moe_mlp_->forward(hidden_states); + hidden_states->copy_from(mlp_out); + } +} + +void MiniCPM5MoEDecoderLayer::piecewise_eager_moe( + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + (void)residual; + (void)staging; + if (!use_moe_) { + return; + } + if (infinicore::context::moeTritonCaptureAllowed()) { + return; + } + auto mlp_out = moe_mlp_->forward(hidden_states); + hidden_states->copy_from(mlp_out); +} + +void MiniCPM5MoEDecoderLayer::piecewise_post_attn_graph( + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + // Without deferred allreduce split, reuse the fused CG post path. + piecewise_post_attn_cg(hidden_states, residual, staging); +} + +void MiniCPM5MoEDecoderLayer::piecewise_post_attn_allreduce( + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + (void)residual; + self_attn_->forward_post_attn_piecewise_allreduce_into(hidden_states, staging); +} + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.hpp b/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.hpp new file mode 100644 index 000000000..83b834920 --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_decoder_layer.hpp @@ -0,0 +1,92 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../global_state/piecewise_prefill_state.hpp" +#include "../../layers/mlp/mlp.hpp" +#include "minicpm5_moe_attention.hpp" +#include "minicpm5_moe_sparse_moe_block.hpp" + +#include +#include +#include + +#include +#include +#include + +namespace infinilm::models::minicpm5_moe { + +class MiniCPM5MoEDecoderLayer : public infinicore::nn::Module { +public: + MiniCPM5MoEDecoderLayer(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + std::tuple + forward(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual) const; + + void piecewise_pre_attn(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const; + + /// Layernorm + past-independent QKV (CG-capturable). No RoPE. + void piecewise_pre_attn_qkv(const infinicore::Tensor &positions, + infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const; + + /// Eager RoPE on staging Q/K (host-break). + void piecewise_pre_attn_rope(const infinicore::Tensor &positions, + global_state::PiecewiseLayerStaging &staging) const; + + void piecewise_eager_attn(const infinicore::Tensor &positions, + global_state::PiecewiseLayerStaging &staging) const { + self_attn_->forward_eager_attn_piecewise(positions, staging); + } + + void piecewise_post_attn(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const { + piecewise_post_attn_cg(hidden_states, residual, staging); + } + + void piecewise_post_attn_cg(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const; + + /// Decode-piecewise CG segment: o_proj + post-norm (+ dense MLP). MoE stays eager. + void piecewise_post_attn_decode_cg(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const; + + /// Eager MoE between Graph::run segments (no-op on dense MLP layers). + void piecewise_eager_moe(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const; + + void piecewise_post_attn_graph(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const; + + void piecewise_post_attn_allreduce(infinicore::Tensor &hidden_states, + infinicore::Tensor &residual, + global_state::PiecewiseLayerStaging &staging) const; + + +protected: + INFINICORE_NN_MODULE(MiniCPM5MoEAttention, self_attn); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); + INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, dense_mlp); + INFINICORE_NN_MODULE(MiniCPM5MoESparseMoeBlock, moe_mlp); + bool use_moe_{false}; + size_t layer_idx_{0}; + +private: + infinicore::Tensor mlp_forward(const infinicore::Tensor &hidden_states) const; +}; + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.cpp b/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.cpp new file mode 100644 index 000000000..588bfa058 --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.cpp @@ -0,0 +1,71 @@ +#include "minicpm5_moe_for_causal_lm.hpp" + +#include "../models_registry.hpp" + +#include +#include + +namespace infinilm::models::minicpm5_moe { + +std::shared_ptr +create_minicpm5_moe_model_config(std::shared_ptr model_config) { + const std::string model_type = model_config->get("model_type"); + if (model_type != "minicpm5_moe") { + throw std::runtime_error( + "create_minicpm5_moe_model_config: model_type is not minicpm5_moe"); + } + + auto &config_json = model_config->get_config_json(); + + if (!config_json.contains("head_dim")) { + config_json["head_dim"] = model_config->get("hidden_size") + / model_config->get("num_attention_heads"); + } + if (!config_json.contains("attention_bias")) { + config_json["attention_bias"] = false; + } + if (!config_json.contains("attention_output_bias")) { + config_json["attention_output_bias"] = config_json.value("attention_bias", false); + } + if (!config_json.contains("mlp_bias")) { + config_json["mlp_bias"] = false; + } + if (!config_json.contains("rope_theta")) { + config_json["rope_theta"] = 10000.0; + } + if (!config_json.contains("n_group")) { + config_json["n_group"] = 1; + } + if (!config_json.contains("topk_group")) { + config_json["topk_group"] = 1; + } + if (!config_json.contains("norm_topk_prob")) { + config_json["norm_topk_prob"] = true; + } + if (!config_json.contains("dtype") && config_json.contains("torch_dtype")) { + config_json["dtype"] = config_json["torch_dtype"]; + } + + config_json["num_experts"] = config_json.value("n_routed_experts", 0); + config_json["e_score_correction_bias"] = true; + config_json["apply_routed_scaling_factor_on_output"] = false; + // n_group/topk_group default to 1 → grouped selection is a no-op; use sigmoid + // (same escape hatch as kimi_k25 until fused_gate is graph-safe on dl). + config_json["moe_router_backend"] = "sigmoid"; + if (!config_json.contains("use_gated_attention")) { + config_json["use_gated_attention"] = true; + } + + // Llama / MiniCPM rotate_half = GPT_NEOX. + model_config->set_rope_algo(infinicore::nn::RoPE::Algo::GPT_NEOX); + return model_config; +} + +} // namespace infinilm::models::minicpm5_moe + +namespace { +INFINILM_REGISTER_CAUSAL_LM_MODEL( + minicpm5_moe, + infinilm::models::minicpm5_moe::MiniCPM5MoEForCausalLM, + infinilm::models::minicpm5_moe::create_minicpm5_moe_model_config); +} // namespace diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.hpp b/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.hpp new file mode 100644 index 000000000..c2667cc4e --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_for_causal_lm.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../../layers/causal_lm_templates/piecewise_text_causal_lm.hpp" +#include "../../layers/causal_lm_templates/text_model.hpp" +#include "minicpm5_moe_decoder_layer.hpp" + +#include + +namespace infinilm::models::minicpm5_moe { + +using MiniCPM5MoEModel = infinilm::layers::causal_lm_templates::TextModel; +using MiniCPM5MoEForCausalLM = + infinilm::layers::causal_lm_templates::PiecewiseTextCausalLM; + +std::shared_ptr +create_minicpm5_moe_model_config(std::shared_ptr model_config); + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.cpp b/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.cpp new file mode 100644 index 000000000..7d5f9d38b --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.cpp @@ -0,0 +1,58 @@ +#include "minicpm5_moe_sparse_moe_block.hpp" + +#include +#include + +namespace infinilm::models::minicpm5_moe { + +std::shared_ptr +make_minicpm5_mlp_config(const std::shared_ptr &model_config, + size_t intermediate_size) { + auto config_json = model_config->get_config_json(); + config_json["intermediate_size"] = intermediate_size; + return std::make_shared(std::move(config_json)); +} + +MiniCPM5MoESparseMoeBlock::MiniCPM5MoESparseMoeBlock( + std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + INFINICORE_NN_MODULE_INIT(gate, model_config, device); + INFINICORE_NN_MODULE_INIT(experts, model_config, device); + INFINICORE_NN_MODULE_INIT(fused_moe, model_config, device, layer_idx); + + routed_scaling_factor_ = model_config->get_or("routed_scaling_factor", 1.0f); + const size_t n_shared_experts = model_config->get_or("n_shared_experts", 0); + has_shared_experts_ = n_shared_experts > 0; + if (has_shared_experts_) { + const size_t shared_intermediate_size = model_config->get("moe_intermediate_size") * n_shared_experts; + auto shared_config = make_minicpm5_mlp_config(model_config, shared_intermediate_size); + INFINICORE_NN_MODULE_INIT(shared_experts, shared_config, device); + } +} + +infinicore::Tensor MiniCPM5MoESparseMoeBlock::forward( + const infinicore::Tensor &hidden_states) const { + ASSERT(hidden_states->ndim() == 3); + const auto shape = hidden_states->shape(); + auto flattened = hidden_states->view({shape[0] * shape[1], shape[2]}); + + auto [routing_weights, selected_experts] = gate_->forward(flattened); + // Fold routed_scaling into topk weights (M×topk) instead of mul_scalar on + // the full hidden [B,S,H] — dl ATen mul_Scalar OORs on long prefill. + if (routed_scaling_factor_ != 1.0f) { + routing_weights = infinicore::op::mul_scalar(routing_weights, routed_scaling_factor_); + } + infinilm::layers::moe::TopKOutput topk_output{ + routing_weights, + selected_experts, + infinicore::Tensor(), + }; + auto routed = fused_moe_->forward(flattened, topk_output, experts_->moe_weights())->view(shape); + if (has_shared_experts_) { + return infinicore::op::add(routed, shared_experts_->forward(hidden_states)); + } + return routed; +} + +} // namespace infinilm::models::minicpm5_moe diff --git a/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.hpp b/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.hpp new file mode 100644 index 000000000..38a117f40 --- /dev/null +++ b/csrc/models/minicpm5_moe/minicpm5_moe_sparse_moe_block.hpp @@ -0,0 +1,39 @@ +#pragma once + +#include "../../config/model_config.hpp" +#include "../../layers/mlp/mlp.hpp" +#include "../../layers/moe/experts/fused_moe_experts.hpp" +#include "../../layers/moe/fused_moe.hpp" +#include "../../layers/moe/router/topk_router.hpp" + +#include +#include + +namespace infinilm::models::minicpm5_moe { + +std::shared_ptr +make_minicpm5_mlp_config(const std::shared_ptr &model_config, + size_t intermediate_size); + +/** + * Sparse MoE + shared expert (DeepSeek-V3 / MiniCPM5 style). + * Router: sigmoid + e_score_correction_bias; routed_scaling applied on routed output. + */ +class MiniCPM5MoESparseMoeBlock final : public infinicore::nn::Module { +public: + MiniCPM5MoESparseMoeBlock(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + +protected: + INFINICORE_NN_MODULE(infinilm::layers::moe::TopKRouter, gate); + INFINICORE_NN_MODULE(infinilm::layers::moe::FusedMoeExperts, experts); + INFINICORE_NN_MODULE(infinilm::layers::moe::FusedMoE, fused_moe); + INFINICORE_NN_MODULE(infinilm::layers::mlp::MLP, shared_experts); + bool has_shared_experts_{false}; + float routed_scaling_factor_{1.0f}; +}; + +} // namespace infinilm::models::minicpm5_moe diff --git a/python/infinicore/device.py b/python/infinicore/device.py index 3e9461569..775e9999e 100644 --- a/python/infinicore/device.py +++ b/python/infinicore/device.py @@ -16,6 +16,7 @@ "iluvatar": (_infinicore.Device.Type.ILUVATAR, "iluvatar"), "hygon": (_infinicore.Device.Type.HYGON, "hygon"), "thead": (_infinicore.Device.Type.THEAD, "thead"), + "qy": (_infinicore.Device.Type.QY, "qy"), } _NATIVE_TO_NAME = {native: name for native, name in set(_TYPE_ALIASES.values())} diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py index 6459dda6b..249c35dea 100644 --- a/python/infinicore/tensor.py +++ b/python/infinicore/tensor.py @@ -181,16 +181,20 @@ def from_torch(torch_tensor, *, device=None): ) index = 0 if owner.device.index is None else owner.device.index + _cuda_compat = ("cuda", "qy", "metax", "iluvatar", "hygon", "thead") if device is None: if owner.device.type == "cuda": current = get_device() - if current.type not in ("cpu", "cuda"): + if current.type in ("qy", "thead", "hygon", "iluvatar", "metax"): + target = Device(current.type, index) + elif current.type in ("cpu", "cuda"): + target = Device("nvidia", index) + else: raise ValueError( "Torch reports a CUDA-compatible device without its vendor; " "pass device= explicitly" ) - target = Device("nvidia", index) else: target = Device(owner.device.type, index) elif isinstance(device, str) and ":" not in device: @@ -201,7 +205,7 @@ def from_torch(torch_tensor, *, device=None): if target.index != index: raise ValueError("device index must match the source Torch tensor") if owner.device.type == "cuda": - if target.type not in ("cuda", "metax", "iluvatar", "hygon", "thead"): + if target.type not in _cuda_compat: raise ValueError("device is not compatible with a Torch CUDA tensor") elif target.type != owner.device.type: raise ValueError("device type must match the source Torch tensor") @@ -218,11 +222,11 @@ def from_torch(torch_tensor, *, device=None): ) result = empty(list(owner.shape), dtype=dtype, device=target) - if owner.device.type == "cuda": + if owner.device.type == "cuda" or target.type == "qy": torch.cuda.synchronize(owner.device) result.copy_(borrowed) - if owner.device.type == "cuda": + if owner.device.type == "cuda" or target.type == "qy": sync_stream() return result diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 857b71c01..ab1dd6c6d 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -3,6 +3,7 @@ import json import os import shutil +import subprocess import warnings from infinilm.config.attention import resolve_attention_backend @@ -201,7 +202,7 @@ def _add_common_args(self): type=str, default="auto", help=( - "device platform: auto, cpu, nvidia, metax, moore, iluvatar, " + "device platform: auto, cpu, nvidia, qy, metax, moore, iluvatar, " "cambricon, ascend, hygon, thead, or backend name " "(cuda/mlu/musa/npu)" ), @@ -524,11 +525,28 @@ def detect_device(self): env_checks = [ ("metax", ["MACA_PATH", "MACA_HOME", "MACA_ROOT"]), ("hygon", ["DTK_HOME", "DTK_PATH"]), + ("qy", ["QY_ROOT", "QY_HOME", "QY_PATH"]), ] for device_name, env_names in env_checks: if any(os.getenv(env_name) for env_name in env_names): return device_name + # QY exposes inventory via dlsmi; require a successful probe. + dlsmi = shutil.which("dlsmi") + if dlsmi: + try: + subprocess.run( + [dlsmi], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=True, + timeout=2, + ) + except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired): + pass + else: + return "qy" + command_checks = [ ("cambricon", ["cnmon"]), ("ascend", ["npu-smi"]), @@ -557,6 +575,8 @@ def get_device_str(self, device): "musa": "musa", "npu": "npu", "nvidia": "cuda", + # Keep InfiniCore device as `qy` (kQy); torch tensors still use CUDA. + "qy": "qy", "cambricon": "mlu", "ascend": "npu", "metax": "metax", diff --git a/python/infinilm/compile/__init__.py b/python/infinilm/compile/__init__.py new file mode 100644 index 000000000..6ee8ce73d --- /dev/null +++ b/python/infinilm/compile/__init__.py @@ -0,0 +1,39 @@ +"""Compile / CUDA-graph policy helpers.""" + +from .env import ( + apply_cudagraph_policy_env, + cudagraph_policy, + long_prefill_threshold, + max_num_batched_tokens, + moe_aot_step_max_tokens, + moe_aot_step_token_ladder, + moe_aot_step_total_allowed, + prefill_chunk_size, + prefill_chunked_enabled, + prefill_native_cg_enabled, + schedule_homogeneous_enabled, + schedule_no_mixed_enabled, + schedule_preempt_enabled, + v1_scheduler_enabled, + CUDAGRAPH_POLICY_EAGER, + CUDAGRAPH_POLICY_FULL_AND_PIECEWISE, +) + +__all__ = [ + "CUDAGRAPH_POLICY_EAGER", + "CUDAGRAPH_POLICY_FULL_AND_PIECEWISE", + "apply_cudagraph_policy_env", + "cudagraph_policy", + "long_prefill_threshold", + "max_num_batched_tokens", + "moe_aot_step_max_tokens", + "moe_aot_step_token_ladder", + "moe_aot_step_total_allowed", + "prefill_chunk_size", + "prefill_chunked_enabled", + "prefill_native_cg_enabled", + "schedule_homogeneous_enabled", + "schedule_no_mixed_enabled", + "schedule_preempt_enabled", + "v1_scheduler_enabled", +] diff --git a/python/infinilm/compile/env.py b/python/infinilm/compile/env.py new file mode 100644 index 000000000..7e9d89513 --- /dev/null +++ b/python/infinilm/compile/env.py @@ -0,0 +1,192 @@ +"""CUDA-graph policy + continuous-batching / chunked-prefill env knobs.""" + +from __future__ import annotations + +import logging +import os +from typing import Optional, Tuple + +logger = logging.getLogger(__name__) + +CUDAGRAPH_POLICY_EAGER = "eager" +CUDAGRAPH_POLICY_FULL_AND_PIECEWISE = "full_and_piecewise" +_CUDAGRAPH_POLICIES = frozenset( + {CUDAGRAPH_POLICY_EAGER, CUDAGRAPH_POLICY_FULL_AND_PIECEWISE} +) + +_SCHEDULE_HOMOGENEOUS_WARNED = False +_PREFILL_CHUNK_SIZE_WARNED = False + +# MoE/piecewise AOT step ladder (powers of two through 4096). +_MOE_AOT_STEP_LADDER_MAX = 4096 +_MOE_AOT_STEP_TOKEN_LADDER: Tuple[int, ...] = tuple( + 1 << i for i in range(0, 13) if (1 << i) <= _MOE_AOT_STEP_LADDER_MAX +) + + +def _truthy(name: str, default: str = "0") -> bool: + return os.environ.get(name, default).strip().lower() not in ( + "", + "0", + "false", + "no", + "off", + ) + + +def _setdefault_env(name: str, value: str) -> None: + """Set env only when unset or empty.""" + cur = os.environ.get(name) + if cur is None or not str(cur).strip(): + os.environ[name] = value + + +def cudagraph_policy() -> str: + raw = os.environ.get("INFINI_CUDAGRAPH_POLICY", "").strip().lower() + if raw in ("", CUDAGRAPH_POLICY_EAGER): + return CUDAGRAPH_POLICY_EAGER if raw == CUDAGRAPH_POLICY_EAGER else "" + if raw == CUDAGRAPH_POLICY_FULL_AND_PIECEWISE: + return CUDAGRAPH_POLICY_FULL_AND_PIECEWISE + return "" + + +def apply_cudagraph_policy_env(policy: Optional[str] = None) -> str: + """Expand companion knobs when F&P is selected. + + Keeps existing FA_INGRAPH=decode setdefault, and setdefaults + ``INFINI_MAX_NUM_BATCHED_TOKENS=2048`` under F&P when unset. + """ + pol = (policy or os.environ.get("INFINI_CUDAGRAPH_POLICY", "")).strip().lower() + if pol == CUDAGRAPH_POLICY_FULL_AND_PIECEWISE: + os.environ.setdefault( + "INFINI_CUDAGRAPH_POLICY", CUDAGRAPH_POLICY_FULL_AND_PIECEWISE + ) + if not os.environ.get("INFINI_FA_INGRAPH", "").strip(): + os.environ["INFINI_FA_INGRAPH"] = "decode" + logger.info( + "cudagraph_policy=full_and_piecewise → INFINI_FA_INGRAPH=decode" + ) + _setdefault_env("INFINI_MAX_NUM_BATCHED_TOKENS", "2048") + return CUDAGRAPH_POLICY_FULL_AND_PIECEWISE + return cudagraph_policy() + + +def prefill_native_cg_enabled() -> bool: + """Native C++ piecewise CUDAGraph prefill (follows F&P policy when set).""" + raw = os.environ.get("INFINI_CUDAGRAPH_POLICY", "").strip().lower() + if raw in _CUDAGRAPH_POLICIES: + return raw == CUDAGRAPH_POLICY_FULL_AND_PIECEWISE + return _truthy("INFINI_PREFILL_NATIVE_CG", "0") + + +def prefill_chunked_enabled() -> bool: + """Multi-step chunked prefill for long prompts.""" + return _truthy("INFINI_PREFILL_CHUNKED", "0") + + +def max_num_batched_tokens(default: int = 8192) -> int: + """Token budget per v1 scheduler step (shared chunker / step budget). + + Canonical: ``INFINI_MAX_NUM_BATCHED_TOKENS``. Legacy + ``INFINILM_MAX_NUM_BATCHED_TOKENS`` and deprecated + ``INFINI_PREFILL_CHUNK_SIZE`` are one-release shims when MAX is unset. + """ + global _PREFILL_CHUNK_SIZE_WARNED + raw_max = os.environ.get("INFINI_MAX_NUM_BATCHED_TOKENS") + raw_legacy = os.environ.get("INFINILM_MAX_NUM_BATCHED_TOKENS") + raw_chunk = os.environ.get("INFINI_PREFILL_CHUNK_SIZE") + + if raw_chunk and not _PREFILL_CHUNK_SIZE_WARNED: + logger.warning( + "INFINI_PREFILL_CHUNK_SIZE is deprecated; use " + "INFINI_MAX_NUM_BATCHED_TOKENS (shared step budget). " + "CHUNK is ignored when MAX_BATCHED is set; if only CHUNK is set it " + "is treated as MAX_BATCHED for one release." + ) + _PREFILL_CHUNK_SIZE_WARNED = True + + if raw_max: + max_v = int(raw_max) + if raw_chunk and int(raw_chunk) != max_v: + logger.warning( + "INFINI_PREFILL_CHUNK_SIZE=%s ignored; " + "INFINI_MAX_NUM_BATCHED_TOKENS=%s wins", + raw_chunk, + raw_max, + ) + return max(1, max_v) + + if raw_legacy: + return max(1, int(raw_legacy)) + + if raw_chunk: + return max(1, int(raw_chunk)) + + if prefill_chunked_enabled() and prefill_native_cg_enabled(): + return 2048 + return default + + +def prefill_chunk_size(default: int = 8192) -> int: + """Deprecated alias of ``max_num_batched_tokens`` (one shared step budget).""" + return max_num_batched_tokens(default=default) + + +def v1_scheduler_enabled() -> bool: + """Master switch: token-budget continuous mixed scheduler (default off).""" + return _truthy("INFINI_V1_SCHEDULER", "0") + + +def schedule_homogeneous_enabled() -> bool: + """Deprecated: homogeneous PREFILL XOR DECODE is removed.""" + global _SCHEDULE_HOMOGENEOUS_WARNED + if _truthy("INFINI_SCHEDULE_HOMOGENEOUS", "0"): + if not _SCHEDULE_HOMOGENEOUS_WARNED: + logger.warning( + "INFINI_SCHEDULE_HOMOGENEOUS is deprecated and ignored; " + "v1 scheduler always uses continuous mixed batching" + ) + _SCHEDULE_HOMOGENEOUS_WARNED = True + return False + + +def schedule_no_mixed_enabled() -> bool: + """When true, v1 never mixes decode rows with prefill rows in one step.""" + return _truthy("INFINI_SCHEDULE_NO_MIXED", "0") + + +def schedule_preempt_enabled() -> bool: + """When true, RECOMPUTE-preempt a RUNNING victim on decode append_slot OOM. + + Default on. Set ``INFINI_SCHEDULE_PREEMPT=0`` for defer-only behavior. + """ + return _truthy("INFINI_SCHEDULE_PREEMPT", "1") + + +def long_prefill_threshold(default: int = 0) -> int: + """Per-request prefill cap for one v1 step. + + Default 0 (inactive) when unset. Explicit ``INFINI_LONG_PREFILL_THRESHOLD`` + wins; otherwise chunked prefill uses ``max_num_batched_tokens()``. + """ + raw = os.environ.get("INFINI_LONG_PREFILL_THRESHOLD") + if raw: + return int(raw) + if prefill_chunked_enabled(): + return max_num_batched_tokens() + return default + + +def moe_aot_step_max_tokens() -> int: + """Largest MoE/piecewise AOT bucket (shape-gate ceiling).""" + return _MOE_AOT_STEP_LADDER_MAX + + +def moe_aot_step_token_ladder() -> Tuple[int, ...]: + """MoE/piecewise AOT bucket sizes (powers of two through 4096).""" + return _MOE_AOT_STEP_TOKEN_LADDER + + +def moe_aot_step_total_allowed(total: int) -> bool: + """True if step total can be served by MoE AOT (runtime pads up to a bucket).""" + return 0 < total <= _MOE_AOT_STEP_LADDER_MAX diff --git a/python/infinilm/llm/cache_manager.py b/python/infinilm/llm/cache_manager.py index 6c045e4ee..cf1a380df 100644 --- a/python/infinilm/llm/cache_manager.py +++ b/python/infinilm/llm/cache_manager.py @@ -155,6 +155,27 @@ def get_computed_blocks( cached_block_table.append(block_id) return cached_block_table, len(cached_block_table) * self.block_size + def allocate_blocks( + self, + tokens: Sequence[int], + block_table: List[int] | None = None, + ) -> tuple[List[int], List[int], int]: + """Compat helper for tests / CB admit: allocate a fresh full sequence. + + Returns ``(block_table, slot_mapping, num_cached_tokens)`` with + ``num_cached_tokens=0`` (no prefix hit in this helper). + """ + del block_table # fresh allocation; callers pass [] / None + allocation = self.allocate_slots(len(tokens), num_computed_tokens=0) + if allocation is None: + raise RuntimeError("No available cache blocks") + req_blocks, slot_mapping = allocation + return req_blocks, slot_mapping, 0 + + def reset_req_blocks(self) -> None: + """No-op on POST BlockManager (no Aug-style req_block_ids split).""" + return + def allocate_slots( self, num_new_tokens: int, diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..970f2ea23 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -96,10 +96,15 @@ def __init__(self, config: EngineConfig): ) num_mamba_cache_blocks = max(2, config.num_blocks // 4) + # Prefer INFINI_MAX_NUM_BATCHED_TOKENS; fall back to legacy INFINILM_*. + batched_env = os.getenv("INFINI_MAX_NUM_BATCHED_TOKENS") or os.getenv( + "INFINILM_MAX_NUM_BATCHED_TOKENS" + ) max_num_batched_tokens = int( - os.getenv("INFINILM_MAX_NUM_BATCHED_TOKENS", max_position_embeddings) + batched_env if batched_env is not None else max_position_embeddings ) - assert 1024 <= max_num_batched_tokens <= max_position_embeddings + # Floor 256: QY F&P often sets step budget = capture bucket (e.g. 512). + assert 256 <= max_num_batched_tokens <= max_position_embeddings self.scheduler = Scheduler( max_batch_size=config.max_batch_size, @@ -122,6 +127,25 @@ def __init__(self, config: EngineConfig): self.cache_type = config.cache_type + # Chunked prefill step budget (shared with v1 scheduler / F&P capture max). + self.chunk_size = 0 + try: + from infinilm.compile.env import ( + prefill_chunk_size, + prefill_chunked_enabled, + ) + + if prefill_chunked_enabled(): + self.chunk_size = prefill_chunk_size() + logger.info( + "chunked prefill enabled chunk_size=%s (INFINI_PREFILL_CHUNKED=1)", + self.chunk_size, + ) + except ImportError: + pass + if self.cache_type == "paged" and hasattr(self.scheduler, "chunk_size"): + self.scheduler.chunk_size = self.chunk_size + # Get EOS token IDs from model config self.eos_token_ids = self.model_runner.eos_token_id or [] if isinstance(self.eos_token_ids, int): @@ -135,6 +159,10 @@ def __init__(self, config: EngineConfig): def add_request(self, request: InferenceRequest): """Add a request to the scheduler.""" + if self.cache_type == "paged" and self.chunk_size > 0: + request.chunk_size = self.chunk_size + if hasattr(self.scheduler, "chunk_size"): + self.scheduler.chunk_size = self.chunk_size self.scheduler.add_request(request) def close(self): @@ -158,7 +186,7 @@ def step(self) -> tuple[bool, list[tuple]]: sampled_token_ids = runner_output.sampled_token_ids self.scheduler.update_from_output(runner_output) pending = self._update_requests( - scheduler_output.scheduled_requests, + scheduler_output, sampled_token_ids, ) @@ -176,11 +204,131 @@ def step(self) -> tuple[bool, list[tuple]]: return True, pending def _update_requests( + self, + scheduler_output, + sampled_token_ids: list[int | list[int]], + ) -> List[tuple]: + """Apply sampled tokens (or mid-chunk progress) and publish KV boundaries.""" + rows = getattr(scheduler_output, "rows", None) or [] + if rows: + return self._update_requests_from_rows(scheduler_output, sampled_token_ids) + return self._update_requests_legacy( + scheduler_output.scheduled_requests, sampled_token_ids + ) + + def _advance_chunk_progress(self, req: InferenceRequest, num_tokens: int) -> None: + req.chunk_prefill_offset += num_tokens + req.num_computed_tokens = req.chunk_prefill_offset + req.num_local_cached_tokens = req.chunk_prefill_offset + self.scheduler.commit_computed_tokens(req, req.num_computed_tokens) + + def _update_requests_from_rows( + self, + scheduler_output, + sampled_token_ids: list[int | list[int]], + ) -> List[tuple]: + """V1 row-aware update: mid-chunk advances offset without sampling.""" + rows = scheduler_output.rows + pending: List[tuple] = [] + + all_mid_chunk = bool(rows) and all( + r.is_prefill_row and not r.is_final_prefill_chunk for r in rows + ) + if all_mid_chunk: + for row in rows: + req = row.request + self._advance_chunk_progress(req, row.num_scheduled_tokens) + if req.is_aborted(): + logger.info( + f"Request {req.request_id} aborted by client during chunked-prefill" + ) + continue + self.scheduler.complete_requests(scheduler_output.scheduled_requests) + return [] + + rows_needing_sample = [ + row + for row in rows + if not (row.is_prefill_row and not row.is_final_prefill_chunk) + ] + + # Model may return one sample per scheduled request (POST) or only for + # rows that need a sample (Aug). Support both. + if len(sampled_token_ids) == len(rows): + sample_by_row = list(zip(rows, sampled_token_ids)) + elif len(sampled_token_ids) == len(rows_needing_sample): + sample_iter = iter(sampled_token_ids) + sample_by_row = [] + for row in rows: + if row.is_prefill_row and not row.is_final_prefill_chunk: + sample_by_row.append((row, None)) + else: + sample_by_row.append((row, next(sample_iter))) + else: + raise RuntimeError( + "model output count does not match scheduled rows: " + f"rows={len(rows)}, needing_sample={len(rows_needing_sample)}, " + f"outputs={len(sampled_token_ids)}" + ) + + for row, token_ids in sample_by_row: + req = row.request + if row.is_prefill_row and not row.is_final_prefill_chunk: + self._advance_chunk_progress(req, row.num_scheduled_tokens) + if req.is_aborted(): + logger.info( + f"Request {req.request_id} aborted by client during chunked-prefill" + ) + continue + + if token_ids is None: + raise RuntimeError( + f"Missing sampled token for request_id={req.request_id}" + ) + + if row.is_prefill_row: + # Final prefill chunk: mark compute complete before sampling apply. + req.num_computed_tokens = req.prefill_seq_len() + req.num_local_cached_tokens = req.prefill_seq_len() + req.chunk_prefill_offset = 0 + req.chunk_size = 0 + req.is_prefill = False + self.scheduler.commit_computed_tokens(req, req.num_computed_tokens) + + pending.extend(self._apply_sampled_token_ids(req, token_ids)) + + self.scheduler.complete_requests(scheduler_output.scheduled_requests) + return pending + + def _update_requests_legacy( self, requests: List[InferenceRequest], sampled_token_ids: list[int | list[int]], ) -> List[tuple]: - """Apply sampled tokens and publish their target-model KV boundary.""" + """POST / non-row update path (optional legacy mid-chunk via request flags).""" + chunk_mid_step = ( + len(requests) > 0 + and all( + r.is_prefill and r.is_chunking() and not r.chunk_is_last() + for r in requests + ) + ) + if chunk_mid_step: + for req in requests: + chunk_len = min( + req.chunk_size, + req.prefill_seq_len() - req.chunk_prefill_offset, + ) + self._advance_chunk_progress(req, chunk_len) + if req.is_aborted(): + logger.info( + f"Request {req.request_id} aborted by client during chunked-prefill" + ) + continue + if hasattr(self.scheduler, "requeue_chunking"): + self.scheduler.requeue_chunking(req) + return [] + if len(requests) != len(sampled_token_ids): raise RuntimeError( "model output count does not match the scheduled request count: " @@ -188,92 +336,145 @@ def _update_requests( ) pending = [] for req, token_ids in zip(requests, sampled_token_ids): - # The model successfully consumed the request's current logical tokens. - # Commit this boundary before observing a concurrent client abort. - pre_output_computed_tokens = req.get_total_length() - self.scheduler.commit_computed_tokens(req, pre_output_computed_tokens) + if req.is_prefill: + req.chunk_prefill_offset = 0 + req.chunk_size = 0 + req.is_prefill = False + pending.extend(self._apply_sampled_token_ids(req, token_ids)) - if req.is_aborted(): - logger.info( - f"Request {req.request_id} aborted by client, skipping update" - ) - # close() may have set _aborted=True without setting a terminal status - # (status still RUNNING). - if not req.is_finished(): - req.mark_canceled() - continue + self.scheduler.complete_requests(requests) + return pending - if isinstance(token_ids, list): - num_computed_output_tokens = max(len(token_ids) - 1, 0) - else: - num_computed_output_tokens = 0 - token_ids = [token_ids] + def _apply_sampled_token_ids( + self, + req: InferenceRequest, + token_ids: int | list[int], + ) -> List[tuple]: + """Apply sampled tokens and publish their target-model KV boundary.""" + pending: List[tuple] = [] + # The model successfully consumed the request's current logical tokens. + # Commit this boundary before observing a concurrent client abort. + pre_output_computed_tokens = req.get_total_length() + self.scheduler.commit_computed_tokens(req, pre_output_computed_tokens) - num_appended_tokens = 0 - for token_id in token_ids: - if req.is_finished(): - break - req.append_generated_token_id(token_id) - num_appended_tokens += 1 - pending_tokens = req.generated_token_ids[req._token_decode_offset :] + if req.is_aborted(): + logger.info( + f"Request {req.request_id} aborted by client, skipping update" + ) + # close() may have set _aborted=True without setting a terminal status + # (status still RUNNING). + if not req.is_finished(): + req.mark_canceled() + return pending + + # Flatten accidental nested sample layouts (e.g. (B,1,1) → [[id]]). + while isinstance(token_ids, list) and len(token_ids) == 1 and isinstance( + token_ids[0], list + ): + token_ids = token_ids[0] + if isinstance(token_ids, list): + flat: list[int] = [] + for item in token_ids: + while isinstance(item, list): + if not item: + break + item = item[0] + if isinstance(item, list): + continue + flat.append(int(item)) + token_ids = flat + num_computed_output_tokens = max(len(token_ids) - 1, 0) + else: + num_computed_output_tokens = 0 + token_ids = [int(token_ids)] + + # Prefer len(tokenizer): HF vocab_size often excludes added tokens + # (MiniCPM5: vocab_size=130072 vs len=130560 == model config). + vocab_size = len(self.tokenizer) + + num_appended_tokens = 0 + for token_id in token_ids: + if req.is_finished(): + break + if token_id < 0 or token_id >= vocab_size: + # Hard error: OOR samples must not be soft-mapped (% vocab); + # that kept the step loop alive but corrupted generation under + # wave packing. QY batch_greedy is disabled for n_out>1 so ids + # should stay in-range after the C++ fix. + raise ValueError( + f"Out-of-range sampled token_id={token_id} for request " + f"{req.request_id} (vocab={vocab_size})" + ) + req.append_generated_token_id(token_id) + num_appended_tokens += 1 + pending_tokens = req.generated_token_ids[req._token_decode_offset :] + try: delta = self.tokenizer.decode(pending_tokens) - holds_back = bool(delta) and delta.endswith("\ufffd") + except OverflowError: + logger.error( + "tokenizer.decode OverflowError for request %s tokens=%s; " + "dropping delta", + req.request_id, + pending_tokens[-8:], + ) + delta = "" + req._token_decode_offset = len(req.generated_token_ids) + holds_back = bool(delta) and delta.endswith("\ufffd") - last_committed_text = req.generated_text + last_committed_text = req.generated_text - if not holds_back: - req.generated_text = last_committed_text + delta - req._token_decode_offset = len(req.generated_token_ids) + if not holds_back: + req.generated_text = last_committed_text + delta + req._token_decode_offset = len(req.generated_token_ids) - is_finished = self._check_request_finished(req, token_id) + is_finished = self._check_request_finished(req, token_id) - # vLLM-style replacement character handling is primarily relevant for streaming. - # For offline generation (no output queue), keep the fast incremental path. - if req._output_queue is None: - if is_finished: - req.mark_finished(req.finish_reason) + # vLLM-style replacement character handling is primarily relevant for streaming. + # For offline generation (no output queue), keep the fast incremental path. + if req._output_queue is None: + if is_finished: + req.mark_finished(req.finish_reason) + else: + if holds_back and not is_finished: + token_text = "" else: - if holds_back and not is_finished: + if is_finished and req.finish_reason in ( + FinishReason.EOS_TOKEN, + FinishReason.LENGTH, + FinishReason.STOP_STRING, + ): token_text = "" else: - if is_finished and req.finish_reason in ( - FinishReason.EOS_TOKEN, - FinishReason.LENGTH, - FinishReason.STOP_STRING, - ): - token_text = "" - else: - token_text = req.generated_text[req._text_output_offset :] - if token_text: - req._text_output_offset = len(req.generated_text) - - if is_finished: - req.mark_finished(req.finish_reason) - - output = TokenOutput( - request_id=req.request_id, - token_id=token_id, - token_text=token_text, - finished=is_finished, - finish_reason=req.finish_reason if is_finished else None, - generated_text=req.generated_text, + token_text = req.generated_text[req._text_output_offset :] + if token_text: + req._text_output_offset = len(req.generated_text) + + if is_finished: + req.mark_finished(req.finish_reason) + + output = TokenOutput( + request_id=req.request_id, + token_id=token_id, + token_text=token_text, + finished=is_finished, + finish_reason=req.finish_reason if is_finished else None, + generated_text=req.generated_text, + ) + if req.is_aborted(): + logger.info( + f"Request {req.request_id} aborted before putting token" ) - if req.is_aborted(): - logger.info( - f"Request {req.request_id} aborted before putting token" - ) - continue - pending.append((req.output_queue.async_q, output)) + continue + pending.append((req.output_queue.async_q, output)) - post_output_computed_tokens = pre_output_computed_tokens + min( - num_appended_tokens, - num_computed_output_tokens, - ) - if post_output_computed_tokens > pre_output_computed_tokens: - self.scheduler.commit_computed_tokens(req, post_output_computed_tokens) + post_output_computed_tokens = pre_output_computed_tokens + min( + num_appended_tokens, + num_computed_output_tokens, + ) + if post_output_computed_tokens > pre_output_computed_tokens: + self.scheduler.commit_computed_tokens(req, post_output_computed_tokens) - self.scheduler.complete_requests(requests) return pending def _check_request_finished(self, req: InferenceRequest, token_id: int) -> bool: @@ -756,6 +957,14 @@ def _step_loop(self): time.sleep(0.003) elif pending: self._loop.call_soon_threadsafe(self._batch_put, pending) + except OverflowError as e: + # Bad sample ids must not permanently kill the step loop under CB. + logger.error( + "Recoverable step-loop OverflowError (continuing): %s", + e, + exc_info=True, + ) + time.sleep(0.01) except Exception as e: logger.error(f"Error in step loop: {e}", exc_info=True) self._healthy = False @@ -784,6 +993,7 @@ def add_request( request_id: Optional[str] = None, # For server use request_data: Optional[dict] = None, + chat_template_kwargs: Optional[dict] = None, ) -> InferenceRequest: """Add a request to the engine. @@ -838,7 +1048,9 @@ def add_request( ) prompt = self.engine.apply_chat_template( - messages, add_generation_prompt=add_generation_prompt + messages, + add_generation_prompt=add_generation_prompt, + chat_template_kwargs=chat_template_kwargs or {}, ) mm_inputs = resolve_multimodal_inputs(messages) @@ -918,6 +1130,7 @@ def add_chat_request( sampling_params=sampling_params, request_id=request_id, request_data=request_data, + chat_template_kwargs=kwargs.get("chat_template_kwargs") or {}, ) async def stream_request( diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index 5a178a08b..93d20de4a 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -21,6 +21,40 @@ logger = logging.getLogger(__name__) +def _normalize_sampled_token_ids(raw) -> list[int | list[int]]: + """Flatten model sample tensors to per-request int or list[int]. + + Batched eager/F&P forwards may return shapes like (B,), (B,1), or (B,1,1). + Nested lists must not reach tokenizer.decode (OverflowError on list ids). + """ + + def _as_int(value) -> int: + if isinstance(value, (list, tuple)): + if not value: + raise ValueError("empty sampled token nested list") + return _as_int(value[0]) + return int(value) + + if raw is None: + return [] + if not isinstance(raw, list): + return [_as_int(raw)] + if not raw: + return [] + if all(isinstance(item, (list, tuple)) for item in raw): + # (B, K) speculative / multi-token → list[list[int]] + if any( + isinstance(sub, (list, tuple)) and sub and isinstance(sub[0], (list, tuple)) + for sub in raw + ): + return [_as_int(item) for item in raw] + # (B, 1) common sample layout → list[int] + if all(len(item) == 1 for item in raw): + return [_as_int(item) for item in raw] + return [[_as_int(tok) for tok in item] for item in raw] + return [_as_int(item) for item in raw] + + @dataclass class KVConnectorOutput: finished_sending: set[str] | None = None @@ -166,6 +200,7 @@ def _init_device(self): supported_devices = [ "cpu", "cuda", + "qy", "hygon", "thead", "iluvatar", @@ -232,6 +267,11 @@ def _model_forward(self, scheduler_output): if self.speculative_runner is not None: return self._model_forward_with_speculative(scheduler_output, model_input) + # Scheduler/processor may attach Python-only metadata (chunk finality / + # scheduling mode). Drop keys InferEngine.forward does not accept yet. + for _meta_key in ("is_final_prefill_chunk", "scheduling_mode"): + model_input.pop(_meta_key, None) + # Wake every stage before stage 0 enters forward. Each worker receives # the same metadata and then blocks in its model on the activation from # the preceding stage. Stage 0 waits for all acknowledgements afterward. @@ -249,8 +289,7 @@ def _model_forward(self, scheduler_output): if self.pipeline_control is not None: self.pipeline_control.wait_forward() sampled_tokens_list = sampled_tokens.to_numpy().tolist() - - return sampled_tokens_list + return _normalize_sampled_token_ids(sampled_tokens_list) def _model_forward_with_speculative(self, scheduler_output, model_input): return self.speculative_runner.forward(scheduler_output, model_input) diff --git a/python/infinilm/llm/request.py b/python/infinilm/llm/request.py index f4612b816..6013b3703 100644 --- a/python/infinilm/llm/request.py +++ b/python/infinilm/llm/request.py @@ -186,6 +186,7 @@ def __init__( self._generated_token_ids: List[int] = [] self._generated_token_ids_view = _SequenceView(self._generated_token_ids) self.generated_text: str = "" # generated_text == tokenizer.decode(generated_token_ids[:_token_decode_offset]) + self.is_prefill: bool = True self.status: RequestStatus = RequestStatus.WAITING self.finish_reason: Optional[FinishReason] = None @@ -206,6 +207,13 @@ def __init__( # Number of leading full blocks published to the prefix-cache index. self.num_cache_indexed_blocks: int = 0 + # Chunked-prefill state (0 = disabled, otherwise tokens per chunk / step budget) + self.chunk_size: int = 0 + # Prompt(+recompute) tokens already fed through forward as chunked-prefill + self.chunk_prefill_offset: int = 0 + # RECOMPUTE preemption count; stream/generated state is kept. + self.num_preemptions: int = 0 + # Mamba cache management. None means no mamba cache row is currently owned. self.mamba_cache_index: Optional[int] = None @@ -241,16 +249,95 @@ def prompt_token_ids(self) -> tuple[int, ...]: def generated_token_ids(self) -> Sequence[int]: return self._generated_token_ids_view + @generated_token_ids.setter + def generated_token_ids(self, token_ids: Sequence[int]) -> None: + """Replace generated tokens (tests / helpers). Prefer append_generated_token_id.""" + self._generated_token_ids.clear() + self._generated_token_ids.extend(int(t) for t in token_ids) + @property def block_hashes(self) -> Sequence[BlockHash]: return self._block_hashes_view + @property + def num_cached_tokens(self) -> int: + """Alias of ``num_local_cached_tokens`` (Aug CB / pack helpers).""" + return self.num_local_cached_tokens + + @num_cached_tokens.setter + def num_cached_tokens(self, value: int) -> None: + self.num_local_cached_tokens = int(value) + def get_prompt_length(self) -> int: return self.prompt_length def get_input_tokens(self) -> Sequence[int]: return self._prompt_token_ids + def tokens_for_kv_recompute(self) -> List[int]: + """Full token sequence needed to rebuild KV after RECOMPUTE preemption.""" + return list(self._prompt_token_ids) + list(self._generated_token_ids) + + def get_prefill_tokens(self) -> List[int]: + """Tokens to feed during the current prefill (prompt+generated after RECOMPUTE).""" + if self.is_prefill and self.num_preemptions > 0: + return self.tokens_for_kv_recompute() + return list(self._prompt_token_ids) + + def prefill_seq_len(self) -> int: + """Sequence length that must be computed during the current prefill.""" + if self.is_prefill and self.num_preemptions > 0: + return self.prompt_length + len(self._generated_token_ids) + return self.prompt_length + + def prepare_for_recompute(self, chunk_size: int) -> None: + """Reset compute/KV state for RECOMPUTE while keeping streamed output. + + Does **not** call ``mark_canceled`` (that would close client streams). + Keeps generated tokens/text and ``_text_output_offset`` so resume does + not re-emit prefixes. Preserves WAITING (not WAITING_FOR_REMOTE_KVS). + """ + self.is_prefill = True + self.status = RequestStatus.WAITING + self.finish_reason = None + self.finished_time = None + + self.block_table = [] + self.slot_mapping = [] + self.num_local_cached_tokens = 0 + self.num_computed_tokens = 0 + self.num_blocks = 0 + self.num_cache_indexed_blocks = 0 + + self.chunk_prefill_offset = 0 + self.chunk_size = chunk_size + self.num_preemptions += 1 + + def prefill_debt(self) -> int: + """Prefill tokens still to compute (0 when decoding). + + Uses ``num_computed_tokens`` / ``chunk_prefill_offset`` as the progress + cursor (POST field), not Aug's property overlay. + """ + if not self.is_prefill: + return 0 + seq_len = self.prefill_seq_len() + if self.chunk_size > 0 and self.chunk_prefill_offset > 0: + return max(0, seq_len - self.chunk_prefill_offset) + return max(0, seq_len - self.num_computed_tokens) + + def is_chunking(self) -> bool: + """Return True if this request uses multi-step chunked prefill.""" + if self.chunk_size <= 0 or not self.is_prefill: + return False + remaining = self.prefill_seq_len() - self.num_computed_tokens + return remaining > self.chunk_size + + def chunk_is_last(self) -> bool: + """Return True if the next chunk would finish the prefill sequence.""" + remaining = self.prefill_seq_len() - self.chunk_prefill_offset + return remaining <= self.chunk_size + def get_num_generated_tokens(self) -> int: return len(self._generated_token_ids) diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index c10b55f2f..da5a83c19 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -3,16 +3,109 @@ """ import logging +import os import queue -from typing import List, Optional +from collections import deque +from dataclasses import dataclass +from typing import Deque, List, Optional, Set import janus +from infinilm.compile.env import ( + long_prefill_threshold, + max_num_batched_tokens, + moe_aot_step_max_tokens, + moe_aot_step_total_allowed, + schedule_no_mixed_enabled, + schedule_preempt_enabled, + v1_scheduler_enabled, +) from infinilm.llm.cache_manager import BlockManager, MambaCacheManager from infinilm.llm.request import InferenceRequest, RequestStatus logger = logging.getLogger(__name__) +# Piecewise CG power-ladder cap. +_PACK_BUCKET_CAP = 8192 +_PREEMPT_WARN_EVERY = 10 + + +def _padded_bucket(total_q: int) -> int: + """Pad total Q token count to the next power-of-two bucket (512..8192).""" + if total_q <= 0: + return 0 + if total_q > _PACK_BUCKET_CAP: + return total_q + bucket = 512 + while bucket < total_q: + bucket *= 2 + return bucket + + +class _WaitingQueueCompat: + """Thin janus-like facade over Scheduler waiting deque for legacy callers.""" + + def __init__(self, scheduler: "Scheduler"): + self._scheduler = scheduler + + @property + def sync_q(self) -> "_WaitingQueueCompat": + return self + + def qsize(self) -> int: + return self._scheduler.waiting_size() + + def put(self, request: InferenceRequest) -> None: + self._scheduler.waiting_append(request) + + def get_nowait(self) -> InferenceRequest: + if self._scheduler.waiting_size() == 0: + raise queue.Empty + return self._scheduler.waiting_popleft() + + def empty(self) -> bool: + return self._scheduler.waiting_size() == 0 + + +@dataclass +class ScheduledRow: + """Per-request scheduling metadata for one engine step.""" + + request: InferenceRequest + num_scheduled_tokens: int + is_prefill_row: bool + is_final_prefill_chunk: bool + + +class SchedulingBudget: + """Token and sequence budget with per-request deduplication.""" + + def __init__(self, token_budget: int, max_num_seqs: int): + self._token_budget = token_budget + self._max_num_seqs = max_num_seqs + self._num_scheduled_tokens = 0 + self._scheduled_req_ids: set[str] = set() + + def remaining_tokens(self) -> int: + return self._token_budget - self._num_scheduled_tokens + + def remaining_seqs(self) -> int: + return self._max_num_seqs - len(self._scheduled_req_ids) + + def can_add(self, request_id: str, num_tokens: int) -> bool: + if num_tokens <= 0: + return False + if ( + request_id not in self._scheduled_req_ids + and self.remaining_seqs() <= 0 + ): + return False + return num_tokens <= self.remaining_tokens() + + def add(self, request_id: str, num_tokens: int) -> None: + self._num_scheduled_tokens += num_tokens + self._scheduled_req_ids.add(request_id) + class SpeculativeCacheOps: """Limited cache operations needed by speculative verification.""" @@ -44,21 +137,46 @@ def __init__( scheduled_requests: List[InferenceRequest], is_prefill: bool = False, speculative_cache_ops: Optional[SpeculativeCacheOps] = None, + rows: Optional[List[ScheduledRow]] = None, ): self.scheduled_requests = scheduled_requests self.num_requests = len(scheduled_requests) self.is_prefill = is_prefill self.speculative_cache_ops = speculative_cache_ops self.kv_connector_metadata = None + self.rows: List[ScheduledRow] = rows or [] + + @property + def total_scheduled_tokens(self) -> int: + if self.rows: + return sum(r.num_scheduled_tokens for r in self.rows) + return 0 + + def is_homogeneous_prefill(self) -> bool: + return bool(self.rows) and all(r.is_prefill_row for r in self.rows) + + def is_homogeneous_decode(self) -> bool: + return bool(self.rows) and all(not r.is_prefill_row for r in self.rows) + + @property + def scheduling_mode(self) -> str: + if not self.rows: + return "PREFILL" if self.is_prefill else "DECODE" + if self.is_homogeneous_prefill(): + return "PREFILL" + if self.is_homogeneous_decode(): + return "DECODE" + return "MIXED" class Scheduler: """Request scheduler with integrated BlockManager for KV cache management. - Scheduling logic: - 1. Running queue: Check for new blocks needed, update slot_mapping - 2. Waiting queue: Try block reuse (prefix caching), allocate new blocks - 3. Reference counting: Free blocks when requests complete + Default (``INFINI_V1_SCHEDULER=0``): POST schedule — waiting prefill first, + then running decode, with prefix-cache / mamba / remote-KV / speculative hooks. + + V1 (``INFINI_V1_SCHEDULER=1``): token-budget continuous mixed batching with + mid-chunk prefill rows, pack gates, and optional RECOMPUTE preemption. """ def __init__( @@ -71,11 +189,23 @@ def __init__( has_mamba_cache: bool = False, num_mamba_cache_blocks: int | None = None, enable_prefix_caching: bool = True, + enable_prefix_cache: bool | None = None, + max_prefill_batch_size: int | None = None, ): - self.waiting_queue = janus.Queue() + if enable_prefix_cache is not None: + enable_prefix_caching = enable_prefix_cache + + self._waiting: Deque[InferenceRequest] = deque() self.running_queue = janus.Queue() + self.chunking_queue = janus.Queue() self.max_batch_size = max_batch_size + env_prefill_batch = os.environ.get("INFINI_MAX_PREFILL_BATCH") + if max_prefill_batch_size is None and env_prefill_batch: + max_prefill_batch_size = int(env_prefill_batch) + self.max_prefill_batch_size = max_prefill_batch_size + self.max_capture_req = max_prefill_batch_size or max_batch_size + self.finished_receiving_kv_req_ids: set[str] = set() self.failed_receiving_kv_req_ids: set[str] = set() self.pending_free_blocks: dict[str, list[int]] = {} @@ -94,6 +224,29 @@ def __init__( self.max_num_batched_tokens = max_num_batched_tokens self.connector = connector self.enable_prefix_caching = enable_prefix_caching + self.chunk_size: int = 0 + self.num_preemptions_total: int = 0 + + # --- waiting deque helpers --- + + def waiting_append(self, request: InferenceRequest) -> None: + self._waiting.append(request) + + def waiting_prepend(self, request: InferenceRequest) -> None: + self._waiting.appendleft(request) + + def waiting_popleft(self) -> InferenceRequest: + if not self._waiting: + raise queue.Empty + return self._waiting.popleft() + + def waiting_size(self) -> int: + return len(self._waiting) + + @property + def waiting_queue(self): + """Compat shim: ``.sync_q.qsize()`` / ``.put`` / ``.get_nowait``.""" + return _WaitingQueueCompat(self) def add_request(self, request: InferenceRequest): if request is not None: @@ -106,7 +259,7 @@ def add_request(self, request: InferenceRequest): and not request.has_multimodal_inputs, ) request.status = RequestStatus.WAITING - self.waiting_queue.sync_q.put(request) + self.waiting_append(request) def _exceeds_token_budget( self, @@ -128,6 +281,12 @@ def _exceeds_token_budget( def schedule(self) -> Optional[SchedulerOutput]: """Schedule and return batch of requests to execute.""" + if v1_scheduler_enabled(): + return self._schedule_v1() + return self._schedule_post() + + def _schedule_post(self) -> Optional[SchedulerOutput]: + """POST schedule path (prefill-priority, non-mixed).""" deferred_requests = [] scheduled_requests = [] is_prefill = False @@ -140,7 +299,7 @@ def schedule(self) -> Optional[SchedulerOutput]: and current_num_batched_tokens < self.max_num_batched_tokens ): try: - req = self.waiting_queue.sync_q.get_nowait() + req = self.waiting_popleft() except queue.Empty: break # Skip requests that were already finished (e.g., timed out/canceled while waiting) @@ -203,6 +362,7 @@ def schedule(self) -> Optional[SchedulerOutput]: req, num_local_computed_tokens, current_prefill_extra_blocks, + _post=True, ): logger.warning( "Insufficient KV cache blocks for request %s, deferring.", @@ -287,10 +447,11 @@ def schedule(self) -> Optional[SchedulerOutput]: current_num_batched_tokens += num_tokens_this_step req.status = RequestStatus.RUNNING + req.is_prefill = True if deferred_requests: for req in deferred_requests: - self.waiting_queue.sync_q.put(req) + self.waiting_append(req) # Return prefill batch if any waiting requests were scheduled if scheduled_requests: @@ -323,6 +484,7 @@ def schedule(self) -> Optional[SchedulerOutput]: req.slot_mapping = [new_slot] req.num_blocks = len(req.block_table) req.num_local_cached_tokens = req.get_total_length() - 1 + req.is_prefill = False scheduled_requests.append(req) # Promote completed remote KV transfers (lower priority than running queue). @@ -339,7 +501,7 @@ def schedule(self) -> Optional[SchedulerOutput]: ) self.update_waiting_for_remote_kv(req) req.status = RequestStatus.WAITING - self.waiting_queue.sync_q.put(req) + self.waiting_append(req) elif req_id in self.finished_receiving_kv_req_ids: if len(scheduled_requests) < self.max_batch_size: logger.info( @@ -376,6 +538,642 @@ def schedule(self) -> Optional[SchedulerOutput]: return None + + def _schedule_v1(self) -> Optional[SchedulerOutput]: + return self._schedule_v1_mixed() + + def _schedule_v1_mixed(self) -> Optional[SchedulerOutput]: + """vLLM-like continuous batch: RUNNING (decode + mid-prefill) then WAITING.""" + budget = SchedulingBudget( + max_num_batched_tokens(default=self.max_num_batched_tokens), + self.max_batch_size, + ) + rows: List[ScheduledRow] = [] + scheduled_requests: List[InferenceRequest] = [] + chunk_size_hint = 0 + deferred_running: List[InferenceRequest] = [] + no_mixed = schedule_no_mixed_enabled() + + running_snapshot = self._drain_running_queue() + if no_mixed: + has_decode = any( + (not req.is_finished()) and (not req.is_prefill) + for req in running_snapshot + ) + filtered: List[InferenceRequest] = [] + for req in running_snapshot: + if req.is_finished(): + filtered.append(req) + continue + if has_decode: + if req.is_prefill: + deferred_running.append(req) + else: + filtered.append(req) + else: + if req.is_prefill and req.prefill_debt() > 0: + filtered.append(req) + elif req.is_prefill: + deferred_running.append(req) + else: + filtered.append(req) + running_snapshot = filtered + + for req in running_snapshot: + if req.is_finished(): + self.complete_requests([req]) + continue + if req.status != RequestStatus.RUNNING: + continue + if req.is_prefill and req.prefill_debt() > 0: + chunk_size_hint_ref = [chunk_size_hint] + if self._try_schedule_v1_prefill_row( + req, + rows, + scheduled_requests, + budget, + chunk_size_hint_ref, + defer_fn=deferred_running.append, + from_waiting=False, + ): + chunk_size_hint = chunk_size_hint_ref[0] + elif not req.is_prefill: + self._try_schedule_v1_decode_row( + req, + rows, + scheduled_requests, + budget, + deferred_running, + running_pool=running_snapshot, + ) + else: + deferred_running.append(req) + + has_decode_row = any(not r.is_prefill_row for r in rows) + skip_waiting_prefills = no_mixed and has_decode_row + prefill_batch_cap = min( + self.max_batch_size, + self.max_prefill_batch_size or self.max_batch_size, + ) + n_prefill = sum(1 for r in rows if r.is_prefill_row) + while ( + not skip_waiting_prefills + and n_prefill < prefill_batch_cap + and len(scheduled_requests) < self.max_batch_size + ): + if self.waiting_size() == 0: + break + req = self.waiting_popleft() + + if req.is_finished(): + self.complete_requests([req]) + continue + + if not self._v1_can_accept_request(req): + self.waiting_prepend(req) + break + + if not self._v1_allocate_waiting_prefill(req): + self.waiting_prepend(req) + break + + remaining = req.prefill_seq_len() - req.num_computed_tokens + if req.chunk_size > 0 and remaining > req.chunk_size: + req.chunk_prefill_offset = req.num_computed_tokens + + chunk_size_hint_ref = [chunk_size_hint] + ok = self._try_schedule_v1_prefill_row( + req, + rows, + scheduled_requests, + budget, + chunk_size_hint_ref, + defer_fn=None, + from_waiting=True, + ) + chunk_size_hint = chunk_size_hint_ref[0] + if not ok: + break + n_prefill += 1 + + for req in deferred_running: + self.running_queue.sync_q.put(req) + + if not scheduled_requests: + # Still allow remote-KV connector meta / empty step like POST. + if self.connector is not None: + out = SchedulerOutput( + scheduled_requests=[], + speculative_cache_ops=self.speculative_cache_ops, + ) + out.kv_connector_metadata = self.connector.build_connector_meta() + return out + return None + + prefill_rows = [r.request for r in rows if r.is_prefill_row] + self._log_prefill_pack(prefill_rows) + + out = SchedulerOutput( + scheduled_requests=scheduled_requests, + is_prefill=all(r.is_prefill_row for r in rows), + speculative_cache_ops=self.speculative_cache_ops, + rows=rows, + ) + if self.connector is not None: + out.kv_connector_metadata = self.connector.build_connector_meta() + return out + + def _v1_allocate_waiting_prefill(self, req: InferenceRequest) -> bool: + """Allocate KV (and optional mamba) for a WAITING prefill using POST APIs.""" + if req.block_table: + req.status = RequestStatus.RUNNING + return True + + if self.has_mamba_cache: + cached_block_table: List[int] = [] + num_local_computed_tokens = 0 + load_kv_async = False + num_external_computed_tokens = 0 + else: + if self.enable_prefix_caching: + cached_block_table, num_local_computed_tokens = ( + self.cache_manager.get_computed_blocks( + req.block_hashes, req.get_prompt_length() - 1 + ) + ) + else: + cached_block_table = [] + num_local_computed_tokens = 0 + if self.connector is not None: + ext_tokens, load_kv_async = self.connector.get_num_new_matched_tokens( + req, num_local_computed_tokens + ) + num_external_computed_tokens = ext_tokens + else: + load_kv_async = False + num_external_computed_tokens = 0 + + available_cached_tokens = ( + num_local_computed_tokens + num_external_computed_tokens + ) + num_computed_tokens = min( + available_cached_tokens, + max(req.prefill_seq_len() - 1, 0), + ) + num_new_tokens = req.prefill_seq_len() - num_computed_tokens + + if not self.can_accept_request( + req, + num_local_computed_tokens, + current_prefill_extra_blocks=0, + _post=True, + ): + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + return False + + allocation = self.cache_manager.allocate_slots( + num_new_tokens, + num_computed_tokens=num_computed_tokens, + cached_block_table=cached_block_table, + ) + if allocation is None: + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + return False + req_blocks, slot_mapping = allocation + + if self.has_mamba_cache and req.mamba_cache_index is None: + req.mamba_cache_index = self.mamba_cache_manager.allocate() + if req.mamba_cache_index is None: + self.cache_manager.free_blocks(req_blocks) + return False + + req.block_table = req_blocks + req.slot_mapping = slot_mapping + req.num_blocks = len(req_blocks) + req.num_local_cached_tokens = ( + num_local_computed_tokens if load_kv_async else num_computed_tokens + ) + req.num_cache_indexed_blocks = len(cached_block_table) + req.num_computed_tokens = num_computed_tokens + req.is_prefill = True + + if self.connector is not None: + self.connector.update_state_after_alloc( + req, req.block_table, num_external_computed_tokens, self.block_size + ) + + if load_kv_async: + req.status = RequestStatus.WAITING_FOR_REMOTE_KVS + self.remote_kv_requests[req.request_id] = req + self.pending_kv_decode_blocks += ( + req.sampling_params.max_tokens + self.block_size - 1 + ) // self.block_size + return False + + req.status = RequestStatus.RUNNING + return True + + def _step_scheduled_tokens(self, rows: List[ScheduledRow]) -> int: + return sum(r.num_scheduled_tokens for r in rows) + + def _shape_gate_allows(self, rows: List[ScheduledRow], add_tokens: int) -> bool: + return moe_aot_step_total_allowed( + self._step_scheduled_tokens(rows) + add_tokens + ) + + def _shape_headroom(self, rows: List[ScheduledRow]) -> int: + return max(0, moe_aot_step_max_tokens() - self._step_scheduled_tokens(rows)) + + def _preempt_request(self, req: InferenceRequest) -> int: + blocks_freed = len(req.block_table) if req.block_table else 0 + if req.block_table: + # POST free_blocks only drops refcounts; reclaim into free_block_ids. + self.cache_manager.free_blocks(req.block_table) + self.cache_manager.try_free_blocks(max(1, blocks_freed)) + if self.mamba_cache_manager is not None: + self.mamba_cache_manager.free(req.mamba_cache_index) + req.mamba_cache_index = None + chunk = self.chunk_size if self.chunk_size > 0 else req.chunk_size + req.prepare_for_recompute(chunk) + self.waiting_prepend(req) + self.num_preemptions_total += 1 + if self.num_preemptions_total == 1 or ( + self.num_preemptions_total % _PREEMPT_WARN_EVERY == 0 + ): + logger.warning( + "RECOMPUTE preempted request %s (num_preemptions_total=%d). " + "Chronic preemption usually means undersized NUM_BLOCKS / too-high " + "max concurrency — raise NUM_BLOCKS or lower MC.", + req.request_id[:16], + self.num_preemptions_total, + ) + return blocks_freed + + def _preempt_until_free( + self, + need_blocks: int, + *, + protected: Set[int], + candidate_pool: Optional[List[InferenceRequest]] = None, + deferred: Optional[List[InferenceRequest]] = None, + ) -> bool: + if not schedule_preempt_enabled(): + return False + if self.cache_manager.can_allocate(need_blocks): + return True + + deferred = deferred or [] + candidate_pool = candidate_pool or [] + protected_ids = set(protected) + + def _candidates() -> List[InferenceRequest]: + seen: Set[int] = set() + pool: List[InferenceRequest] = [] + for r in list(deferred) + list(candidate_pool): + rid = id(r) + if rid in protected_ids or rid in seen: + continue + if r.is_finished() or not r.block_table: + continue + seen.add(rid) + pool.append(r) + decode = [r for r in pool if not r.is_prefill] + mid_prefill = [r for r in pool if r.is_prefill] + decode.sort(key=lambda r: r.arrival_time, reverse=True) + mid_prefill.sort(key=lambda r: r.arrival_time, reverse=True) + return decode + mid_prefill + + while not self.cache_manager.can_allocate(need_blocks): + cands = _candidates() + if not cands: + return False + victim = cands[0] + try: + deferred.remove(victim) + except ValueError: + pass + try: + candidate_pool.remove(victim) + except ValueError: + pass + self._preempt_request(victim) + return True + + def _try_schedule_v1_decode_row( + self, + req: InferenceRequest, + rows: List[ScheduledRow], + scheduled_requests: List[InferenceRequest], + budget: SchedulingBudget, + deferred: List[InferenceRequest], + running_pool: Optional[List[InferenceRequest]] = None, + ) -> bool: + if len(scheduled_requests) >= self.max_batch_size: + deferred.append(req) + return False + if not budget.can_add(req.request_id, 1): + deferred.append(req) + return False + if not self._shape_gate_allows(rows, 1): + deferred.append(req) + return False + + def _append_once() -> None: + req.block_table, new_slot = self.cache_manager.append_slot( + req.block_table, + req.get_total_length(), + ) + req.slot_mapping = [new_slot] + req.num_blocks = len(req.block_table) + req.num_local_cached_tokens = req.get_total_length() - 1 + req.num_computed_tokens = req.get_total_length() - 1 + + try: + _append_once() + except RuntimeError: + protected = {id(req)} | {id(r) for r in scheduled_requests} + freed = self._preempt_until_free( + 1, + protected=protected, + candidate_pool=running_pool, + deferred=deferred, + ) + if not freed: + deferred.append(req) + return False + try: + _append_once() + except RuntimeError: + deferred.append(req) + return False + + rows.append(ScheduledRow(req, 1, False, False)) + scheduled_requests.append(req) + budget.add(req.request_id, 1) + return True + + def _try_schedule_v1_prefill_row( + self, + req: InferenceRequest, + rows: List[ScheduledRow], + scheduled_requests: List[InferenceRequest], + budget: SchedulingBudget, + chunk_size_hint_ref: List[int], + *, + defer_fn, + from_waiting: bool, + ) -> bool: + q = self._next_prefill_chunk_len(req, budget) + headroom = self._shape_headroom(rows) + if headroom <= 0: + if from_waiting: + req.status = RequestStatus.WAITING + self.waiting_append(req) + elif defer_fn is not None: + defer_fn(req) + return False + q = min(q, headroom) + if q <= 0 or not budget.can_add(req.request_id, q): + if from_waiting: + req.status = RequestStatus.WAITING + self.waiting_append(req) + elif defer_fn is not None: + defer_fn(req) + return False + if not self._shape_gate_allows(rows, q): + if from_waiting: + req.status = RequestStatus.WAITING + self.waiting_append(req) + elif defer_fn is not None: + defer_fn(req) + return False + + chunk_size_hint = chunk_size_hint_ref[0] + prefill_rows = [r for r in rows if r.is_prefill_row] + if prefill_rows: + if chunk_size_hint == 0: + chunk_size_hint = req.chunk_size + if not self._can_add_v1_prefill_row(rows, req, chunk_size_hint): + if from_waiting: + req.status = RequestStatus.WAITING + self.waiting_append(req) + elif defer_fn is not None: + defer_fn(req) + return False + elif req.chunk_size > 0 and chunk_size_hint == 0: + chunk_size_hint = req.chunk_size + + rows.append( + ScheduledRow( + req, + q, + True, + self._is_final_prefill_chunk(req, q), + ) + ) + scheduled_requests.append(req) + budget.add(req.request_id, q) + chunk_size_hint_ref[0] = chunk_size_hint + return True + + def _drain_running_queue(self) -> List[InferenceRequest]: + reqs: List[InferenceRequest] = [] + size = self.running_queue.sync_q.qsize() + for _ in range(size): + try: + reqs.append(self.running_queue.sync_q.get_nowait()) + except queue.Empty: + break + return reqs + + def _next_prefill_chunk_len( + self, req: InferenceRequest, budget: SchedulingBudget + ) -> int: + q = self._prefill_compute_len(req) + if q <= 0: + return 0 + remaining = budget.remaining_tokens() + threshold = long_prefill_threshold() + if threshold > 0: + return min(q, threshold, remaining) + return min(q, remaining) + + def _is_final_prefill_chunk(self, req: InferenceRequest, q: int) -> bool: + progress = ( + req.chunk_prefill_offset + if req.chunk_prefill_offset > 0 + else req.num_computed_tokens + ) + return progress + q >= req.prefill_seq_len() + + def _can_add_v1_prefill_row( + self, + rows: List[ScheduledRow], + candidate: InferenceRequest, + chunk_size: int, + ) -> bool: + prefill_rows = [r for r in rows if r.is_prefill_row] + if not prefill_rows: + return True + scheduled = [r.request for r in prefill_rows] + if not self._can_add_to_prefill_pack( + scheduled, candidate, chunk_size=chunk_size + ): + return False + # Homogeneous Q only (allows 2048+2048 under a 4096 chunk budget; rejects + # mixing mid-chunk 4096 with a short final remainder). + if chunk_size > 0: + pack = scheduled + [candidate] + lens = [self._prefill_compute_len(r) for r in pack] + if len(set(lens)) != 1: + return False + return True + + def _pack_token_budget(self, chunk_size: int) -> int: + if chunk_size > 0: + return min(_PACK_BUCKET_CAP, chunk_size) + return _PACK_BUCKET_CAP + + @staticmethod + def _prefill_compute_len(req: InferenceRequest) -> int: + remaining = req.prefill_seq_len() - req.num_computed_tokens + if req.is_chunking() or ( + req.chunk_size > 0 + and req.is_prefill + and req.chunk_prefill_offset > 0 + and req.prefill_seq_len() - req.chunk_prefill_offset > 0 + ): + start = ( + req.chunk_prefill_offset + if req.chunk_prefill_offset > 0 + else req.num_computed_tokens + ) + end = min(start + req.chunk_size, req.prefill_seq_len()) + return max(0, end - start) + if req.chunk_size > 0 and remaining > req.chunk_size: + return req.chunk_size + return remaining + + @staticmethod + def _prefill_is_final_chunk(req: InferenceRequest) -> bool: + if req.chunk_size > 0 and req.is_prefill and ( + req.is_chunking() or req.chunk_prefill_offset > 0 + ): + return req.chunk_is_last() + return True + + @staticmethod + def _prefix_cache_pack_compatible(reqs: List[InferenceRequest]) -> bool: + cached = [r.num_computed_tokens for r in reqs] + if any(c > 0 for c in cached): + return len(set(cached)) == 1 + return True + + def _can_add_to_prefill_pack( + self, + scheduled: List[InferenceRequest], + candidate: InferenceRequest, + *, + chunk_size: int, + ) -> bool: + if not scheduled: + return True + if candidate.chunk_size != scheduled[0].chunk_size: + return False + if scheduled[0].chunk_prefill_offset != candidate.chunk_prefill_offset: + return False + if not self._prefix_cache_pack_compatible(scheduled + [candidate]): + return False + + phase_final = self._prefill_is_final_chunk(scheduled[0]) + if self._prefill_is_final_chunk(candidate) != phase_final: + return False + + total_q = sum(self._prefill_compute_len(r) for r in scheduled) + cand_q = self._prefill_compute_len(candidate) + if chunk_size > 0 and cand_q > chunk_size: + return False + total_q += cand_q + budget = self._pack_token_budget(chunk_size) + if total_q > budget: + return False + if _padded_bucket(total_q) > budget: + return False + return True + + def _log_prefill_pack(self, scheduled: List[InferenceRequest]) -> None: + if len(scheduled) <= 1: + return + total_q = sum(self._prefill_compute_len(r) for r in scheduled) + bucket = _padded_bucket(total_q) + phase = "final" if self._prefill_is_final_chunk(scheduled[0]) else "mid" + logger.info( + "scheduled prefill pack n_req=%s total_q=%s bucket=%s phase=%s", + len(scheduled), + total_q, + bucket, + phase, + ) + + def _prefill_pack_cap(self) -> int: + return min(self.max_batch_size, self.max_capture_req) + + def requeue_chunking(self, req: InferenceRequest) -> None: + """Put a request back into the chunking queue (legacy pack / tests).""" + self.chunking_queue.sync_q.put(req) + + def requeue_running(self, req: InferenceRequest) -> None: + """Put a request back into the running queue after a forward step (v1).""" + self.running_queue.sync_q.put(req) + + def _remaining_generation_tokens(self, req: InferenceRequest) -> int: + max_tokens = req.sampling_params.max_tokens + if max_tokens is None: + return 0 + return max(0, max_tokens - req.get_num_generated_tokens()) + + def _kv_completion_tokens(self, req: InferenceRequest) -> int: + max_tokens = req.sampling_params.max_tokens + if max_tokens is None: + return req.get_total_length() + if ( + req.status == RequestStatus.WAITING + and req.num_preemptions > 0 + and req.generated_token_ids + ): + return req.get_prompt_length() + max_tokens + return req.get_prompt_length() + self._remaining_generation_tokens(req) + + def _v1_can_accept_request(self, request: InferenceRequest) -> bool: + """Aug-style admission used by V1 waiting + unit tests.""" + total_required_blocks = 0 + + running_queue_size = self.running_queue.sync_q.qsize() + for _ in range(running_queue_size): + req = self.running_queue.sync_q.get() + if req.block_table: + blocks_for_req = len(req.block_table) + else: + blocks_for_req = req.get_num_blocks_required(self.block_size) + total_at_completion = ( + self._kv_completion_tokens(req) + self.block_size - 1 + ) // self.block_size + total_required_blocks += max(blocks_for_req, total_at_completion) + self.running_queue.sync_q.put(req) + + total_length = self._kv_completion_tokens(request) + num_blocks_needed = (total_length + self.block_size - 1) // self.block_size + total_required_blocks += num_blocks_needed + total_required_blocks += self.pending_kv_decode_blocks + + return total_required_blocks <= self.cache_manager.get_total_usable_blocks() + + # Alias for tests that call can_accept_request(req) with one arg under V1. + def can_accept_request_v1(self, request: InferenceRequest) -> bool: + return self._v1_can_accept_request(request) + def update_waiting_for_remote_kv(self, request: InferenceRequest): self.remote_kv_requests.pop(request.request_id, None) self.pending_kv_decode_blocks -= ( @@ -453,9 +1251,19 @@ def complete_requests(self, requests: List[InferenceRequest]): def can_accept_request( self, request: InferenceRequest, - num_local_computed_tokens: int, + num_local_computed_tokens: int = 0, current_prefill_extra_blocks: int = 0, + *, + _post: bool = False, ) -> bool: + """Admission check. + + Under V1 (and one-arg test calls), use Aug-style ``_v1_can_accept_request`` + unless ``_post=True`` (POST schedule path). + """ + if v1_scheduler_enabled() and not _post: + return self._v1_can_accept_request(request) + if ( self.mamba_cache_manager is not None and request.mamba_cache_index is None @@ -584,6 +1392,7 @@ def get_cache_stats(self) -> dict: "num_free_blocks": self.cache_manager.get_num_free_blocks(), "usable_blocks": self.cache_manager.get_total_usable_blocks(), "num_used_blocks": len(self.cache_manager.used_block_ids), + "num_preemptions_total": self.num_preemptions_total, } if self.mamba_cache_manager is not None: stats.update( diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index f4b7c50fc..bc454413f 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -342,7 +342,10 @@ def load_model_state_dict_by_file( ) already_loaded_keys.append(key) - model.load_state_dict(model_param_infini, strict=True) + # MiniCPM5 MoE registers packed w13/w2 views absent from HF `.bin` + # checkpoints; check_parameters filters those expected missing keys. + bin_strict = model_type not in ("minicpm5_moe",) + model.load_state_dict(model_param_infini, strict=bin_strict) infinicore.sync_device() del model_param_infini del model_params diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index a6fbc33ac..425ec0539 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -1,7 +1,7 @@ from transformers import AutoTokenizer from typing_extensions import override -from ..llm.scheduler import SchedulerOutput +from ..llm.scheduler import ScheduledRow, SchedulerOutput from ..llm.static_scheduler import StaticSchedulerOutput from .processor import InfinilmProcessor, register_processor @@ -158,11 +158,154 @@ def _build_model_input_from_static_scheduler_output( "top_p": top_p, } + @staticmethod + def _append_prefill_row_metadata( + row: ScheduledRow, + *, + tokens: list, + seq_lens: list, + seq_offsets: list, + slot_mapping: list, + cached_lens: list, + position_ids: list, + cu_seqlens: list, + is_final_prefill_chunk: list, + ) -> None: + req = row.request + num_cached = req.num_local_cached_tokens + req_tokens = req.get_prefill_tokens() + q = row.num_scheduled_tokens + start = ( + req.chunk_prefill_offset + if req.chunk_prefill_offset > 0 + else num_cached + ) + end = min(start + q, len(req_tokens)) + tokens_to_compute = req_tokens[start:end] + tokens.extend(tokens_to_compute) + total_kv_len = end + seq_lens.append(total_kv_len) + current_offset = seq_offsets[-1] + len(tokens_to_compute) + seq_offsets.append(current_offset) + slot_start = start - num_cached + slot_end = end - num_cached + # slot_mapping may cover the full allocated suffix from num_cached. + if req.slot_mapping and slot_end > slot_start: + slot_mapping.extend(req.slot_mapping[slot_start:slot_end]) + else: + slot_mapping.extend(req.slot_mapping) + cached_lens.append(start) + position_ids.extend(range(start, end)) + cu_seqlens.append(cu_seqlens[-1] + total_kv_len) + is_final_prefill_chunk.append(row.is_final_prefill_chunk) + + @staticmethod + def _append_decode_row_metadata( + req, + *, + tokens: list, + seq_lens: list, + seq_offsets: list, + slot_mapping: list, + cached_lens: list, + position_ids: list, + cu_seqlens: list, + is_final_prefill_chunk: list, + ) -> None: + num_cached = req.num_local_cached_tokens + seq_len = req.get_total_length() + last_token = ( + req.generated_token_ids[-1] + if req.generated_token_ids + else req.prompt_token_ids[-1] + ) + tokens.append(last_token) + seq_lens.append(seq_len) + current_offset = seq_offsets[-1] + 1 + seq_offsets.append(current_offset) + slot_mapping.extend(req.slot_mapping) + cached_lens.append(num_cached) + position_ids.append(seq_len - 1) + cu_seqlens.append(cu_seqlens[-1] + seq_len) + is_final_prefill_chunk.append(True) + + def _build_model_input_from_rows( + self, scheduler_output: SchedulerOutput, temperature, top_p, top_k + ) -> dict: + """Per-row varlen metadata for v1 (and mixed) scheduling.""" + import infinicore + + tokens: list = [] + seq_lens: list = [] + seq_offsets: list = [0] + block_tables: list = [] + slot_mapping: list = [] + cached_lens: list = [] + position_ids: list = [] + cu_seqlens: list = [0] + is_final_prefill_chunk: list = [] + + max_block_table_len = max( + len(row.request.block_table) for row in scheduler_output.rows + ) + + for row in scheduler_output.rows: + req = row.request + if row.is_prefill_row: + self._append_prefill_row_metadata( + row, + tokens=tokens, + seq_lens=seq_lens, + seq_offsets=seq_offsets, + slot_mapping=slot_mapping, + cached_lens=cached_lens, + position_ids=position_ids, + cu_seqlens=cu_seqlens, + is_final_prefill_chunk=is_final_prefill_chunk, + ) + else: + self._append_decode_row_metadata( + req, + tokens=tokens, + seq_lens=seq_lens, + seq_offsets=seq_offsets, + slot_mapping=slot_mapping, + cached_lens=cached_lens, + position_ids=position_ids, + cu_seqlens=cu_seqlens, + is_final_prefill_chunk=is_final_prefill_chunk, + ) + padded_block_table = req.block_table + [-1] * ( + max_block_table_len - len(req.block_table) + ) + block_tables.append(padded_block_table) + + return { + "input_ids": infinicore.from_list([tokens], dtype=infinicore.int64), + "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int64), + "past_kv_lengths": infinicore.from_list( + cached_lens, dtype=infinicore.int32 + ), + "total_kv_lengths": infinicore.from_list(seq_lens, dtype=infinicore.int32), + "input_offsets": infinicore.from_list(seq_offsets, dtype=infinicore.int32), + "cu_seqlens": infinicore.from_list(cu_seqlens, dtype=infinicore.int32), + "block_tables": infinicore.from_list(block_tables, dtype=infinicore.int32), + "slot_mapping": infinicore.from_list(slot_mapping, dtype=infinicore.int64), + "temperature": temperature, + "top_k": top_k, + "top_p": top_p, + "is_final_prefill_chunk": is_final_prefill_chunk, + "scheduling_mode": scheduler_output.scheduling_mode, + } + def _build_model_input_from_batch_scheduler_output( self, scheduler_output: SchedulerOutput, temperature, top_p, top_k ) -> dict: """Construct model inputs for prefill or decode phase. + When ``scheduler_output.rows`` is present (v1), use row-aware slicing via + ``chunk_prefill_offset`` / ``ScheduledRow.num_scheduled_tokens``. + Prefill phase: - input_ids: Flattened token list (excluding cached tokens) - position_ids: Position IDs for new tokens in complete sequence @@ -188,6 +331,11 @@ def _build_model_input_from_batch_scheduler_output( "build_model_inputs called with empty scheduled_requests" ) + if getattr(scheduler_output, "rows", None): + return self._build_model_input_from_rows( + scheduler_output, temperature, top_p, top_k + ) + tokens = [] seq_lens = [] seq_offsets = [0] @@ -196,6 +344,7 @@ def _build_model_input_from_batch_scheduler_output( cached_lens = [] position_ids = [] cu_seqlens = [0] + is_final_prefill_chunk: list = [] max_block_table_len = max( len(req.block_table) for req in scheduler_output.scheduled_requests @@ -205,21 +354,47 @@ def _build_model_input_from_batch_scheduler_output( for req in scheduler_output.scheduled_requests: num_cached = req.num_local_cached_tokens if scheduler_output.is_prefill: - # Prefill phase - req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] - tokens.extend(tokens_to_compute) + req_tokens = req.get_prefill_tokens() + if ( + req.chunk_size > 0 + and req.is_prefill + and ( + req.is_chunking() + or req.chunk_prefill_offset > 0 + ) + ): + start = req.chunk_prefill_offset + end = min(start + req.chunk_size, len(req_tokens)) + tokens_to_compute = req_tokens[start:end] + compute_len = len(tokens_to_compute) + tokens.extend(tokens_to_compute) + total_kv_len = start + compute_len + seq_lens.append(total_kv_len) + current_offset += compute_len + seq_offsets.append(current_offset) + slot_start = start - num_cached + slot_end = end - num_cached + slot_mapping.extend(req.slot_mapping[slot_start:slot_end]) + cached_lens.append(start) + position_ids.extend(range(start, end)) + cu_seqlens.append(cu_seqlens[-1] + total_kv_len) + is_final_prefill_chunk.append(req.chunk_is_last()) + else: + tokens_to_compute = req_tokens[num_cached:] + tokens.extend(tokens_to_compute) - compute_len = len(tokens_to_compute) - seq_len = len(req_tokens) - seq_lens.append(seq_len) + compute_len = len(tokens_to_compute) + seq_len = len(req_tokens) + seq_lens.append(seq_len) - current_offset += compute_len - seq_offsets.append(current_offset) + current_offset += compute_len + seq_offsets.append(current_offset) - slot_mapping.extend(req.slot_mapping) - cached_lens.append(num_cached) - position_ids.extend(range(num_cached, num_cached + compute_len)) + slot_mapping.extend(req.slot_mapping) + cached_lens.append(num_cached) + position_ids.extend(range(num_cached, num_cached + compute_len)) + cu_seqlens.append(cu_seqlens[-1] + seq_len) + is_final_prefill_chunk.append(True) else: # Decode phase @@ -238,15 +413,15 @@ def _build_model_input_from_batch_scheduler_output( slot_mapping.extend(req.slot_mapping) cached_lens.append(num_cached) position_ids.append(seq_len - 1) + cu_seqlens.append(cu_seqlens[-1] + seq_len) # Pad block_table to same length padded_block_table = req.block_table + [-1] * ( max_block_table_len - len(req.block_table) ) block_tables.append(padded_block_table) - cu_seqlens.append(cu_seqlens[-1] + seq_len) - return { + model_input = { "input_ids": infinicore.from_list([tokens], dtype=infinicore.int64), "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int64), "past_kv_lengths": infinicore.from_list( @@ -261,6 +436,12 @@ def _build_model_input_from_batch_scheduler_output( "top_k": top_k, "top_p": top_p, } + if is_final_prefill_chunk: + model_input["is_final_prefill_chunk"] = is_final_prefill_chunk + model_input["scheduling_mode"] = ( + "PREFILL" if scheduler_output.is_prefill else "DECODE" + ) + return model_input @override def get_tokenizer(self): diff --git a/python/infinilm/processors/sentencepiece_processor.py b/python/infinilm/processors/sentencepiece_processor.py index c8c3c1b33..857752d55 100644 --- a/python/infinilm/processors/sentencepiece_processor.py +++ b/python/infinilm/processors/sentencepiece_processor.py @@ -16,10 +16,29 @@ def __init__(self, model_dir_path: str): super().__init__(model_dir_path) self._fix_tokenizer_decode(self.tokenizer) + @staticmethod + def _uses_sentencepiece_metaspace(tokenizer) -> bool: + """True for SentencePiece Metaspace (▁ pieces), false for ByteLevel (Ġ).""" + vocab = tokenizer.get_vocab() + return any(token.startswith("▁") for token in vocab) + @staticmethod def _fix_tokenizer_decode(tokenizer): """Fix leading space loss when tokenizer decodes incrementally. + ByteLevel tokenizers (MiniCPM5 sidecar ``tokenizer_bytelevel``) must + keep the native ``decode()`` so GPT-2 byte-unicode pieces become UTF-8. + Joining ``convert_ids_to_tokens`` leaves Chinese as mojibake (e.g. 二). + """ + if not SentencePieceProcessor._uses_sentencepiece_metaspace(tokenizer): + return + + SentencePieceProcessor._patch_metaspace_decode(tokenizer) + + @staticmethod + def _patch_metaspace_decode(tokenizer): + """Fix leading space loss when tokenizer decodes incrementally. + Problem: Fast tokenizer's Rust backend trims leading spaces derived from ▁ (U+2581) during single-token decoding, causing words to concatenate. @@ -78,6 +97,7 @@ class MistralProcessor(SentencePieceProcessor): @register_processor("minicpm") @register_processor("minicpm4") +@register_processor("minicpm5_moe") class MiniCPMProcessor(SentencePieceProcessor): pass diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 462c31084..bd23c2227 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -15,6 +15,7 @@ import uvicorn from fastapi import FastAPI, Request from fastapi.responses import JSONResponse, StreamingResponse + from infinilm.base_config import BaseConfig from infinilm.config import KVTransferConfig from infinilm.llm import AsyncLLMEngine, FinishReason, SamplingParams @@ -23,7 +24,8 @@ logger = logging.getLogger(__name__) DEFAULT_STREAM_TIMEOUT = 100.0 -DEFAULT_REQUEST_TIMEOUT = 1000.0 +# Long-prefill adapting (CPU RoPE + Q-chunked attn) can exceed 1000s at ~4–7k. +DEFAULT_REQUEST_TIMEOUT = 7200.0 def chunk_json( @@ -193,6 +195,12 @@ def __init__( def start(self): """Start the HTTP server.""" + try: + from infinilm.compile.env import apply_cudagraph_policy_env + + apply_cudagraph_policy_env() + except Exception as exc: # noqa: BLE001 — best-effort policy expand + logger.debug("apply_cudagraph_policy_env skipped: %s", exc) app = self._create_app() logger.info(f"Starting API Server at {self.host}:{self.port}...") uvicorn.run(app, host=self.host, port=self.port) diff --git a/python/tests/test_processor_v1_rows.py b/python/tests/test_processor_v1_rows.py new file mode 100644 index 000000000..14e7dfd47 --- /dev/null +++ b/python/tests/test_processor_v1_rows.py @@ -0,0 +1,112 @@ +"""Processor row-based metadata for v1 mixed scheduling.""" + +from __future__ import annotations + +import unittest +from unittest.mock import MagicMock + +from infinilm.llm.request import InferenceRequest, RequestStatus +from infinilm.llm.sampling_params import SamplingParams +from infinilm.llm.scheduler import ScheduledRow, Scheduler, SchedulerOutput +from infinilm.processors.basic_llm_processor import BasicLLMProcessor + + +def _req(rid: str, prompt_len: int, *, chunk_size: int = 4096) -> InferenceRequest: + req = InferenceRequest( + request_id=rid, + prompt_token_ids=list(range(1, prompt_len + 1)), + sampling_params=SamplingParams(max_tokens=1), + ) + req.chunk_size = chunk_size + return req + + +class _FakeTensor: + def __init__(self, values): + self._values = values + + def to_numpy(self): + import numpy as np + + return np.asarray(self._values) + + +class ProcessorV1RowsTest(unittest.TestCase): + def setUp(self): + self.processor = BasicLLMProcessor.__new__(BasicLLMProcessor) + + def test_mixed_row_metadata_decode_cu_seqlens(self): + """Decode row cu_seqlens uses total KV length, not 1.""" + sched = Scheduler( + max_batch_size=4, + num_blocks=512, + block_size=256, + enable_prefix_cache=False, + ) + decode = _req("d", 4096, chunk_size=4096) + decode.is_prefill = False + decode.generated_token_ids = [99] + decode.block_table, decode.slot_mapping, _ = ( + sched.cache_manager.allocate_blocks(decode.get_input_tokens(), []) + ) + # Extend allocation for the generated token slot used at decode. + decode.block_table, new_slot = sched.cache_manager.append_slot( + decode.block_table, decode.get_total_length() + ) + decode.slot_mapping = [new_slot] + decode.num_blocks = len(decode.block_table) + decode.status = RequestStatus.RUNNING + decode.num_local_cached_tokens = decode.get_total_length() - 1 + decode.num_computed_tokens = decode.num_local_cached_tokens + + prefill = _req("p", 4096, chunk_size=4096) + prefill.block_table, prefill.slot_mapping, cached = ( + sched.cache_manager.allocate_blocks(prefill.get_input_tokens(), []) + ) + prefill.num_blocks = len(prefill.block_table) + prefill.status = RequestStatus.RUNNING + prefill.num_local_cached_tokens = cached + prefill.num_computed_tokens = cached + prefill.chunk_prefill_offset = 0 + + rows = [ + ScheduledRow(decode, 1, False, False), + ScheduledRow(prefill, 4096, True, False), + ] + out = SchedulerOutput( + scheduled_requests=[decode, prefill], + is_prefill=False, + rows=rows, + ) + + # Avoid requiring a real infinicore extension: stub from_list. + import sys + import types + + fake_ic = types.ModuleType("infinicore") + fake_ic.int32 = "int32" + fake_ic.int64 = "int64" + + def _from_list(values, dtype=None): + return _FakeTensor(values) + + fake_ic.from_list = _from_list + sys.modules["infinicore"] = fake_ic + + model_input = self.processor._build_model_input_from_rows( + out, temperature=1.0, top_p=1.0, top_k=1 + ) + cu = model_input["cu_seqlens"].to_numpy().tolist() + self.assertEqual(model_input["scheduling_mode"], "MIXED") + self.assertEqual(len(model_input["is_final_prefill_chunk"]), 2) + self.assertTrue(model_input["is_final_prefill_chunk"][0]) + self.assertFalse(model_input["is_final_prefill_chunk"][1]) + self.assertEqual(cu[-1], 4096 + 4097) + past = model_input["past_kv_lengths"].to_numpy().tolist() + total = model_input["total_kv_lengths"].to_numpy().tolist() + self.assertEqual(past[0], 4096) + self.assertEqual(total[0], 4097) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_scheduler_pack.py b/python/tests/test_scheduler_pack.py new file mode 100644 index 000000000..0b5562f9c --- /dev/null +++ b/python/tests/test_scheduler_pack.py @@ -0,0 +1,195 @@ +"""Unit tests for continuous-batch prefill packing in Scheduler (V1 path).""" + +from __future__ import annotations + +import os +import unittest + +from infinilm.llm.request import InferenceRequest, RequestStatus +from infinilm.llm.sampling_params import SamplingParams +from infinilm.llm.scheduler import Scheduler + + +def _req(request_id: str, prompt_len: int, *, chunk_size: int = 4096) -> InferenceRequest: + req = InferenceRequest( + request_id=request_id, + prompt_token_ids=list(range(1, prompt_len + 1)), + sampling_params=SamplingParams(max_tokens=1), + ) + req.chunk_size = chunk_size + return req + + +class SchedulerPackTest(unittest.TestCase): + def setUp(self): + os.environ["INFINI_V1_SCHEDULER"] = "1" + os.environ["INFINI_MAX_NUM_BATCHED_TOKENS"] = "8192" + os.environ.pop("INFINI_MAX_PREFILL_BATCH", None) + os.environ.pop("INFINI_SCHEDULE_NO_MIXED", None) + self.scheduler = Scheduler( + max_batch_size=4, + num_blocks=512, + block_size=256, + max_prefill_batch_size=4, + enable_prefix_cache=False, + ) + + def tearDown(self): + os.environ.pop("INFINI_V1_SCHEDULER", None) + os.environ.pop("INFINI_MAX_NUM_BATCHED_TOKENS", None) + os.environ.pop("INFINI_SCHEDULE_NO_MIXED", None) + + def test_four_4096_chunks_no_two_pack_at_chunk_budget(self): + """With chunk_size=4096, V1 never packs 2×4096 mid-chunks in one step. + + RUNNING-first means a mid-chunk request continues before new WAITING + admits; assert the pack gate (one 4096 row) rather than legacy + waiting-priority four-mid drain. + """ + chunk_size = 4096 + for i in range(4): + self.scheduler.add_request(_req(f"r{i}", 8192, chunk_size=chunk_size)) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertTrue(out.is_prefill) + self.assertEqual(len(out.scheduled_requests), 1) + self.assertEqual(out.total_scheduled_tokens, 4096) + self.assertFalse(out.rows[0].is_final_prefill_chunk) + self.assertEqual(self.scheduler.waiting_size(), 3) + + # Advance mid-chunk and requeue as RUNNING; next step finishes it (not 2-pack). + req = out.scheduled_requests[0] + req.chunk_prefill_offset += chunk_size + req.num_computed_tokens = req.chunk_prefill_offset + self.scheduler.complete_requests(out.scheduled_requests) + + out2 = self.scheduler.schedule() + self.assertIsNotNone(out2) + assert out2 is not None + self.assertEqual(len(out2.scheduled_requests), 1) + self.assertEqual(out2.total_scheduled_tokens, 4096) + self.assertTrue(out2.rows[0].is_final_prefill_chunk) + self.assertEqual(out2.scheduled_requests[0].request_id, req.request_id) + + def test_smaller_chunks_still_multi_pack(self): + """Under chunk_size=4096 budget, four 2048-token prefills pack two per step.""" + chunk_size = 4096 + for i in range(4): + self.scheduler.add_request(_req(f"r{i}", 2048, chunk_size=chunk_size)) + + out1 = self.scheduler.schedule() + self.assertIsNotNone(out1) + assert out1 is not None + self.assertTrue(out1.is_prefill) + self.assertEqual(len(out1.scheduled_requests), 2) + self.assertEqual(out1.total_scheduled_tokens, 4096) + + out2 = self.scheduler.schedule() + self.assertIsNotNone(out2) + assert out2 is not None + self.assertEqual(len(out2.scheduled_requests), 2) + self.assertEqual(out2.total_scheduled_tokens, 4096) + + def test_reject_mixed_mid_and_final(self): + chunk_size = 4096 + mid = _req("mid", 8192, chunk_size=chunk_size) + mid.chunk_prefill_offset = 0 + mid.num_computed_tokens = 0 + final = _req("final", 4096, chunk_size=chunk_size) + # Prime both onto waiting; pack gate should admit only one phase per step. + self.scheduler.add_request(mid) + self.scheduler.add_request(final) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(len(out.scheduled_requests), 1) + + def test_prefix_cache_incompatible_rejects_pack(self): + chunk_size = 4096 + a = _req("a", 4096, chunk_size=chunk_size) + b = _req("b", 4096, chunk_size=chunk_size) + a.num_computed_tokens = 0 + b.num_computed_tokens = 512 + a.num_local_cached_tokens = 0 + b.num_local_cached_tokens = 512 + a.block_table = [0, 1] + b.block_table = [2, 3] + a.slot_mapping = list(range(4096)) + b.slot_mapping = list(range(4096)) + a.status = RequestStatus.WAITING + b.status = RequestStatus.WAITING + self.scheduler.waiting_append(a) + self.scheduler.waiting_append(b) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(len(out.scheduled_requests), 1) + + def test_prefill_pack_respects_capture_width(self): + """max_batch_size=8 but capture width=4 → at most 4 requests per prefill pack.""" + wide = Scheduler( + max_batch_size=8, + num_blocks=512, + block_size=256, + max_prefill_batch_size=4, + enable_prefix_cache=False, + ) + chunk_size = 512 + for i in range(6): + wide.add_request(_req(f"r{i}", 2048, chunk_size=chunk_size)) + + out = wide.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertTrue(out.is_prefill) + self.assertLessEqual(len(out.scheduled_requests), 4) + + def test_no_mixed_defers_prefill_when_decode_present(self): + """With INFINI_SCHEDULE_NO_MIXED=1, decode-only step; mid-prefill deferred.""" + os.environ["INFINI_SCHEDULE_NO_MIXED"] = "1" + chunk_size = 512 + prefill_req = _req("prefill", 2048, chunk_size=chunk_size) + prefill_req.is_prefill = True + prefill_req.status = RequestStatus.RUNNING + prefill_req.block_table = [0, 1, 2, 3] + prefill_req.slot_mapping = list(range(2048)) + prefill_req.num_computed_tokens = 0 + prefill_req.num_local_cached_tokens = 0 + prefill_req.chunk_prefill_offset = 0 + prefill_req.num_blocks = 4 + + decode_req = _req("decode", 128, chunk_size=0) + decode_req.is_prefill = False + decode_req.status = RequestStatus.RUNNING + decode_req.generated_token_ids = [99] + decode_req.block_table, decode_req.slot_mapping, _ = ( + self.scheduler.cache_manager.allocate_blocks( + decode_req.get_input_tokens(), [] + ) + ) + decode_req.num_blocks = len(decode_req.block_table) + decode_req.num_local_cached_tokens = decode_req.get_total_length() - 1 + decode_req.num_computed_tokens = decode_req.num_local_cached_tokens + + self.scheduler.running_queue.sync_q.put(prefill_req) + self.scheduler.running_queue.sync_q.put(decode_req) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertFalse(out.is_prefill) + self.assertEqual(len(out.scheduled_requests), 1) + self.assertFalse(out.scheduled_requests[0].is_prefill) + self.assertEqual(out.scheduled_requests[0].request_id, "decode") + # Prefill deferred back onto running. + self.assertEqual(self.scheduler.running_queue.sync_q.qsize(), 1) + deferred = self.scheduler.running_queue.sync_q.get_nowait() + self.assertEqual(deferred.request_id, "prefill") + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_scheduler_preempt.py b/python/tests/test_scheduler_preempt.py new file mode 100644 index 000000000..4e20f44b6 --- /dev/null +++ b/python/tests/test_scheduler_preempt.py @@ -0,0 +1,188 @@ +"""Unit tests for RECOMPUTE preemption (INFINI_SCHEDULE_PREEMPT).""" + +from __future__ import annotations + +import os +import unittest + +from infinilm.llm.request import InferenceRequest, RequestStatus +from infinilm.llm.sampling_params import SamplingParams +from infinilm.llm.scheduler import Scheduler + + +def _req( + request_id: str, + prompt_len: int, + *, + generated: list[int] | None = None, + max_tokens: int = 64, + arrival_time: float = 1.0, +) -> InferenceRequest: + req = InferenceRequest( + request_id=request_id, + prompt_token_ids=list(range(1, prompt_len + 1)), + sampling_params=SamplingParams(max_tokens=max_tokens), + arrival_time=arrival_time, + ) + if generated: + req.generated_token_ids = list(generated) + req.generated_text = "x" * len(generated) + req._text_output_offset = len(req.generated_text) + req.is_prefill = False + req.status = RequestStatus.RUNNING + return req + + +def _prime_kv_all_tokens(scheduler: Scheduler, req: InferenceRequest) -> None: + tokens = req.get_all_token_ids() + req.block_table, req.slot_mapping, cached = ( + scheduler.cache_manager.allocate_blocks(tokens, req.block_table) + ) + req.num_blocks = len(req.block_table) + req.num_local_cached_tokens = cached + req.num_computed_tokens = cached + req.status = RequestStatus.RUNNING + scheduler.cache_manager.reset_req_blocks() + + +class SchedulerPreemptTest(unittest.TestCase): + def setUp(self): + os.environ["INFINI_V1_SCHEDULER"] = "1" + os.environ["INFINI_MAX_NUM_BATCHED_TOKENS"] = "8192" + os.environ["INFINI_SCHEDULE_PREEMPT"] = "1" + os.environ.pop("INFINI_SCHEDULE_NO_MIXED", None) + os.environ.pop("INFINI_LONG_PREFILL_THRESHOLD", None) + # Tiny KV: 3 blocks × 4 = 12 tokens total. + self.scheduler = Scheduler( + max_batch_size=4, + num_blocks=3, + block_size=4, + max_prefill_batch_size=4, + enable_prefix_cache=False, + ) + self.scheduler.chunk_size = 8 + + def tearDown(self): + os.environ.pop("INFINI_V1_SCHEDULER", None) + os.environ.pop("INFINI_MAX_NUM_BATCHED_TOKENS", None) + os.environ.pop("INFINI_SCHEDULE_PREEMPT", None) + + def test_prepare_for_recompute_keeps_stream_state(self): + req = _req("r", 4, generated=[10, 11, 12], arrival_time=1.0) + req.block_table = [0, 1] + req.slot_mapping = [0, 1, 2, 3, 4] + req.num_local_cached_tokens = 6 + req.num_computed_tokens = 6 + req.num_blocks = 2 + stream_len = req._text_output_offset + gen_ids = list(req.generated_token_ids) + gen_text = req.generated_text + + req.prepare_for_recompute(chunk_size=8) + + self.assertEqual(req.status, RequestStatus.WAITING) + self.assertTrue(req.is_prefill) + self.assertEqual(list(req.generated_token_ids), gen_ids) + self.assertEqual(req.generated_text, gen_text) + self.assertEqual(req._text_output_offset, stream_len) + self.assertEqual(req.block_table, []) + self.assertEqual(req.num_cached_tokens, 0) + self.assertEqual(req.num_computed_tokens, 0) + self.assertEqual(req.chunk_size, 8) + self.assertEqual(req.num_preemptions, 1) + self.assertEqual( + req.tokens_for_kv_recompute(), + list(req.prompt_token_ids) + gen_ids, + ) + self.assertEqual(req.get_prefill_tokens(), req.tokens_for_kv_recompute()) + self.assertEqual(req.prefill_seq_len(), 4 + 3) + + def test_waiting_prepend_ahead_of_new_admits(self): + a = _req("a", 2, arrival_time=1.0) + b = _req("b", 2, arrival_time=2.0) + self.scheduler.waiting_append(a) + self.scheduler.waiting_prepend(b) + self.assertEqual(self.scheduler.waiting_popleft().request_id, "b") + self.assertEqual(self.scheduler.waiting_popleft().request_id, "a") + + def test_decode_append_oom_preempts_newest_victim(self): + """A (older) needs a new block; B (newer) is RECOMPUTE-preempted.""" + older = _req("older", 4, generated=[99], arrival_time=1.0) + newer = _req("newer", 4, arrival_time=2.0) + newer.is_prefill = False + newer.generated_token_ids = [] + newer.status = RequestStatus.RUNNING + + _prime_kv_all_tokens(self.scheduler, older) + _prime_kv_all_tokens(self.scheduler, newer) + self.assertEqual(self.scheduler.cache_manager.get_num_free_blocks(), 0) + self.assertEqual(len(older.block_table), 2) + self.assertEqual(len(newer.block_table), 1) + + self.assertEqual(older.get_total_length() % self.scheduler.block_size, 1) + + self.scheduler.running_queue.sync_q.put(older) + self.scheduler.running_queue.sync_q.put(newer) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertGreater(self.scheduler.num_preemptions_total, 0) + self.assertEqual(newer.status, RequestStatus.WAITING) + self.assertEqual(newer.num_preemptions, 1) + self.assertEqual(newer.block_table, []) + self.assertEqual(self.scheduler.waiting_size(), 1) + self.assertEqual(self.scheduler._waiting[0].request_id, "newer") + scheduled_ids = [r.request_id for r in out.scheduled_requests] + self.assertIn("older", scheduled_ids) + self.assertNotIn("newer", scheduled_ids) + self.assertEqual( + self.scheduler.get_cache_stats()["num_preemptions_total"], + self.scheduler.num_preemptions_total, + ) + + def test_preempt_off_defers_without_counter(self): + os.environ["INFINI_SCHEDULE_PREEMPT"] = "0" + older = _req("older", 4, generated=[99], arrival_time=1.0) + newer = _req("newer", 4, arrival_time=2.0) + newer.is_prefill = False + newer.status = RequestStatus.RUNNING + + _prime_kv_all_tokens(self.scheduler, older) + _prime_kv_all_tokens(self.scheduler, newer) + self.scheduler.running_queue.sync_q.put(older) + self.scheduler.running_queue.sync_q.put(newer) + + out = self.scheduler.schedule() + self.assertEqual(self.scheduler.num_preemptions_total, 0) + self.assertEqual(newer.status, RequestStatus.RUNNING) + self.assertTrue(newer.block_table) + self.assertEqual(older.status, RequestStatus.RUNNING) + self.assertTrue(older.block_table) + self.assertEqual(self.scheduler.waiting_size(), 0) + if out is not None: + scheduled_ids = [r.request_id for r in out.scheduled_requests] + self.assertNotIn("older", scheduled_ids) + + def test_readmit_allocates_prompt_plus_generated(self): + victim = _req("v", 4, generated=[7, 8], arrival_time=1.0, max_tokens=4) + victim.prepare_for_recompute(8) + self.assertEqual(victim.prefill_seq_len(), 6) + self.scheduler.waiting_prepend(victim) + + self.assertGreaterEqual( + self.scheduler.cache_manager.get_num_free_blocks(), 2 + ) + self.assertTrue(self.scheduler.can_accept_request(victim)) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(victim.status, RequestStatus.RUNNING) + self.assertEqual(len(victim.block_table), 2) + self.assertTrue(out.rows[0].is_prefill_row) + self.assertEqual(out.rows[0].num_scheduled_tokens, 6) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/tests/test_scheduler_v1.py b/python/tests/test_scheduler_v1.py new file mode 100644 index 000000000..0f4d84345 --- /dev/null +++ b/python/tests/test_scheduler_v1.py @@ -0,0 +1,361 @@ +"""Unit tests for v1 token-budget scheduler (INFINI_V1_SCHEDULER=1).""" + +from __future__ import annotations + +import os +import unittest + +from infinilm.compile.env import long_prefill_threshold, schedule_homogeneous_enabled +from infinilm.llm.request import InferenceRequest, RequestStatus +from infinilm.llm.sampling_params import SamplingParams +from infinilm.llm.scheduler import Scheduler + + +def _req( + request_id: str, + prompt_len: int, + *, + chunk_size: int = 4096, + max_tokens: int = 1, +) -> InferenceRequest: + req = InferenceRequest( + request_id=request_id, + prompt_token_ids=list(range(1, prompt_len + 1)), + sampling_params=SamplingParams(max_tokens=max_tokens), + ) + req.chunk_size = chunk_size + return req + + +def _prime_kv(scheduler: Scheduler, req: InferenceRequest) -> None: + """Allocate block table so running-queue decode can append_slot.""" + tokens = req.get_input_tokens() + req.block_table, req.slot_mapping, cached = ( + scheduler.cache_manager.allocate_blocks(tokens, req.block_table) + ) + req.num_blocks = len(req.block_table) + req.num_local_cached_tokens = cached + req.num_computed_tokens = cached + req.status = RequestStatus.RUNNING + + +class SchedulerV1Test(unittest.TestCase): + def setUp(self): + os.environ["INFINI_V1_SCHEDULER"] = "1" + os.environ["INFINI_MAX_NUM_BATCHED_TOKENS"] = "8192" + os.environ.pop("INFINI_MAX_PREFILL_BATCH", None) + os.environ.pop("INFINI_SCHEDULE_HOMOGENEOUS", None) + os.environ.pop("INFINI_PREFILL_NATIVE_CG", None) + os.environ.pop("INFINI_PREFILL_CHUNKED", None) + os.environ.pop("INFINI_PREFILL_CHUNK_SIZE", None) + os.environ.pop("INFINI_LONG_PREFILL_THRESHOLD", None) + self.scheduler = Scheduler( + max_batch_size=4, + num_blocks=512, + block_size=256, + max_prefill_batch_size=4, + enable_prefix_cache=False, + ) + + def tearDown(self): + os.environ.pop("INFINI_V1_SCHEDULER", None) + os.environ.pop("INFINI_MAX_NUM_BATCHED_TOKENS", None) + os.environ.pop("INFINI_SCHEDULE_HOMOGENEOUS", None) + os.environ.pop("INFINI_PREFILL_NATIVE_CG", None) + os.environ.pop("INFINI_PREFILL_CHUNKED", None) + os.environ.pop("INFINI_PREFILL_CHUNK_SIZE", None) + os.environ.pop("INFINI_LONG_PREFILL_THRESHOLD", None) + + def test_mixed_on_ladder_prefill_2047_and_decode_one(self): + """Running decode + waiting prefill 2047 → MIXED total 2048 (on MoE max).""" + os.environ["INFINI_MAX_NUM_BATCHED_TOKENS"] = "2048" + os.environ["INFINI_LONG_PREFILL_THRESHOLD"] = "2048" + + decode = _req("decode", 512, chunk_size=2048) + decode.is_prefill = False + decode.generated_token_ids = [99] + _prime_kv(self.scheduler, decode) + self.scheduler.running_queue.sync_q.put(decode) + + prefill = _req("prefill", 2047, chunk_size=2048) + self.scheduler.add_request(prefill) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(out.scheduling_mode, "MIXED") + self.assertEqual(out.total_scheduled_tokens, 2048) + self.assertEqual(len(out.rows), 2) + prefill_rows = [r for r in out.rows if r.is_prefill_row] + decode_rows = [r for r in out.rows if not r.is_prefill_row] + self.assertEqual(len(prefill_rows), 1) + self.assertEqual(len(decode_rows), 1) + self.assertEqual(prefill_rows[0].num_scheduled_tokens, 2047) + self.assertEqual(decode_rows[0].num_scheduled_tokens, 1) + + def test_decode_deferred_when_would_create_4097(self): + """Running mid-prefill 4096 + decode → DECODE deferred (no 4097).""" + mid = _req("mid", 8192, chunk_size=4096) + _prime_kv(self.scheduler, mid) + mid.num_computed_tokens = 0 + mid.num_local_cached_tokens = 0 + mid.chunk_prefill_offset = 0 + mid.is_prefill = True + self.scheduler.running_queue.sync_q.put(mid) + + decode = _req("decode", 512, chunk_size=4096) + decode.is_prefill = False + decode.generated_token_ids = [99] + _prime_kv(self.scheduler, decode) + self.scheduler.running_queue.sync_q.put(decode) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(out.scheduling_mode, "PREFILL") + self.assertEqual(out.total_scheduled_tokens, 4096) + self.assertEqual(len(out.rows), 1) + self.assertTrue(out.rows[0].is_prefill_row) + self.assertEqual(self.scheduler.running_queue.sync_q.qsize(), 1) + deferred = self.scheduler.running_queue.sync_q.get_nowait() + self.assertEqual(deferred.request_id, "decode") + + def test_running_decode_before_waiting_prefill_when_budget_tight(self): + """Budget 1: RUNNING decode consumes step; WAITING prefill stays waiting.""" + os.environ["INFINI_MAX_NUM_BATCHED_TOKENS"] = "1" + + decode = _req("decode", 512, chunk_size=512) + decode.is_prefill = False + decode.generated_token_ids = [1] + _prime_kv(self.scheduler, decode) + self.scheduler.running_queue.sync_q.put(decode) + + self.scheduler.add_request(_req("prefill", 4096, chunk_size=4096)) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(out.scheduling_mode, "DECODE") + self.assertEqual(out.total_scheduled_tokens, 1) + self.assertEqual(len(out.rows), 1) + self.assertFalse(out.rows[0].is_prefill_row) + self.assertEqual(self.scheduler.waiting_queue.sync_q.qsize(), 1) + + def test_homogeneous_env_ignored(self): + """INFINI_SCHEDULE_HOMOGENEOUS=1 does not change mixed RUNNING-first behavior.""" + os.environ["INFINI_SCHEDULE_HOMOGENEOUS"] = "1" + os.environ["INFINI_PREFILL_NATIVE_CG"] = "1" + self.assertFalse(schedule_homogeneous_enabled()) + + os.environ["INFINI_MAX_NUM_BATCHED_TOKENS"] = "2048" + os.environ["INFINI_LONG_PREFILL_THRESHOLD"] = "2048" + + decode = _req("decode", 512, chunk_size=2048) + decode.is_prefill = False + decode.generated_token_ids = [99] + _prime_kv(self.scheduler, decode) + self.scheduler.running_queue.sync_q.put(decode) + + self.scheduler.add_request(_req("prefill", 2047, chunk_size=2048)) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(out.scheduling_mode, "MIXED") + self.assertEqual(out.total_scheduled_tokens, 2048) + + def test_budget_exhaustion_prefill_before_decode(self): + """Budget 4096: prefill fills step; decode on next step.""" + os.environ["INFINI_MAX_NUM_BATCHED_TOKENS"] = "4096" + + self.scheduler.add_request(_req("prefill", 4096, chunk_size=4096)) + + out1 = self.scheduler.schedule() + self.assertIsNotNone(out1) + assert out1 is not None + self.assertEqual(out1.scheduling_mode, "PREFILL") + self.assertEqual(out1.total_scheduled_tokens, 4096) + self.assertEqual(len(out1.rows), 1) + self.assertTrue(out1.rows[0].is_prefill_row) + + row = out1.rows[0] + row.request.is_prefill = False + row.request.generated_token_ids = [1] + self.scheduler.requeue_running(row.request) + + out2 = self.scheduler.schedule() + self.assertIsNotNone(out2) + assert out2 is not None + self.assertEqual(out2.scheduling_mode, "DECODE") + self.assertEqual(out2.total_scheduled_tokens, 1) + + def test_seq_cap_max_batch_size_two(self): + """Only two rows scheduled when max_batch_size=2 (no chunk pack cap).""" + self.scheduler.max_batch_size = 2 + for i in range(4): + self.scheduler.add_request(_req(f"r{i}", 1024, chunk_size=0)) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(len(out.rows), 2) + + def test_single_8192_mid_chunk_no_spurious_2pack(self): + """mc=1: after first 4096 chunk, second step is 1 row (not 2-pack @8192).""" + req = _req("solo", 8192, chunk_size=4096) + self.scheduler.max_batch_size = 1 + self.scheduler.add_request(req) + + out1 = self.scheduler.schedule() + self.assertIsNotNone(out1) + assert out1 is not None + self.assertEqual(len(out1.rows), 1) + self.assertEqual(out1.total_scheduled_tokens, 4096) + self.assertFalse(out1.rows[0].is_final_prefill_chunk) + + for row in out1.rows: + row.request.chunk_prefill_offset += row.num_scheduled_tokens + row.request.num_computed_tokens = row.request.chunk_prefill_offset + self.scheduler.complete_requests(out1.scheduled_requests) + + out2 = self.scheduler.schedule() + self.assertIsNotNone(out2) + assert out2 is not None + self.assertEqual(len(out2.rows), 1) + self.assertEqual(out2.total_scheduled_tokens, 4096) + self.assertTrue(out2.rows[0].is_final_prefill_chunk) + + def test_four_8192_chunk4096_no_two_pack(self): + """4×8192 chunk4096 → pack≡chunk: one full mid-chunk per step (not 8192).""" + chunk_size = 4096 + for i in range(4): + self.scheduler.add_request(_req(f"r{i}", 8192, chunk_size=chunk_size)) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(out.scheduling_mode, "PREFILL") + self.assertEqual(len(out.rows), 1) + self.assertEqual(out.total_scheduled_tokens, 4096) + self.assertTrue(out.rows[0].is_prefill_row) + + def test_dry_run_one_decode_one_prefill_shape(self): + """RUNNING decode first; waiting prefill shrinks to MoE headroom (no 4097).""" + decode = _req("d", 512, chunk_size=512) + decode.is_prefill = False + decode.generated_token_ids = [7] + _prime_kv(self.scheduler, decode) + self.scheduler.running_queue.sync_q.put(decode) + + self.scheduler.add_request(_req("p", 4096, chunk_size=4096)) + + out = self.scheduler.schedule() + self.assertIsNotNone(out) + assert out is not None + self.assertEqual(out.scheduling_mode, "MIXED") + self.assertEqual(out.total_scheduled_tokens, 4096) + self.assertLessEqual(out.total_scheduled_tokens, 4096) + prefill_rows = [r for r in out.rows if r.is_prefill_row] + decode_rows = [r for r in out.rows if not r.is_prefill_row] + self.assertEqual(len(decode_rows), 1) + self.assertEqual(len(prefill_rows), 1) + self.assertEqual(prefill_rows[0].num_scheduled_tokens, 4095) + + def test_can_accept_rejects_third_long_seq(self): + """3×65536-token requests with num_blocks=512 → third rejected.""" + long_len = 65536 + blocks_per_req = (long_len + 1 + 255) // 256 # max_tokens=1 + self.assertGreater(blocks_per_req * 3, 512) + + for i in range(2): + req = _req(f"long{i}", long_len, max_tokens=1) + _prime_kv(self.scheduler, req) + self.scheduler.running_queue.sync_q.put(req) + + third = _req("long2", long_len, max_tokens=1) + self.assertFalse(self.scheduler.can_accept_request(third)) + + def test_can_accept_counts_allocated_blocks_not_remaining_decode(self): + """Running 65536-token decode reserves full KV; second long seq rejected.""" + long_len = 65536 + decode = _req("decode", long_len, max_tokens=512) + decode.is_prefill = False + decode.generated_token_ids = [1] + _prime_kv(self.scheduler, decode) + self.scheduler.running_queue.sync_q.put(decode) + + second_long = _req("second", long_len, max_tokens=512) + blocks_per = (long_len + 512 + 255) // 256 + self.assertGreater(blocks_per * 2, 512) + self.assertFalse(self.scheduler.can_accept_request(second_long)) + + def test_can_accept_small_max_tokens_long_prompt(self): + """16k prompt + max_tokens=512 admits multiple concurrent reqs; overflow rejected.""" + prompt_len = 16384 + max_tokens = 512 + + admitted = 0 + for i in range(16): + req = _req(f"r{i}", prompt_len, max_tokens=max_tokens) + if not self.scheduler.can_accept_request(req): + break + _prime_kv(self.scheduler, req) + self.scheduler.running_queue.sync_q.put(req) + admitted += 1 + + self.assertGreaterEqual(admitted, 2) + overflow = _req("overflow", prompt_len, max_tokens=max_tokens) + self.assertFalse(self.scheduler.can_accept_request(overflow)) + + def test_can_accept_running_decode_remaining_generation(self): + """Running decode reserves remaining generation budget, not full max_tokens.""" + prompt_len = 8192 + max_tokens = 600 + block_size = self.scheduler.block_size + + decode = _req("decode", prompt_len, max_tokens=max_tokens) + decode.is_prefill = False + decode.generated_token_ids = list(range(300)) + _prime_kv(self.scheduler, decode) + + blocks_current = (prompt_len + 300 + block_size - 1) // block_size + blocks_full_cap = (prompt_len + max_tokens + block_size - 1) // block_size + blocks_remaining_cap = ( + prompt_len + (max_tokens - 300) + block_size - 1 + ) // block_size + self.assertEqual(blocks_current, blocks_remaining_cap) + self.assertLess(blocks_remaining_cap, blocks_full_cap) + + self.assertEqual(self.scheduler._remaining_generation_tokens(decode), 300) + self.assertEqual( + self.scheduler._kv_completion_tokens(decode), + prompt_len + 300, + ) + reservation = max(blocks_current, blocks_remaining_cap) + self.assertLess(reservation, max(blocks_current, blocks_full_cap)) + + +class LongPrefillThresholdTest(unittest.TestCase): + def tearDown(self): + os.environ.pop("INFINI_LONG_PREFILL_THRESHOLD", None) + os.environ.pop("INFINI_PREFILL_CHUNKED", None) + os.environ.pop("INFINI_PREFILL_CHUNK_SIZE", None) + + def test_defaults_to_chunk_size_when_chunked_prefill_enabled(self): + os.environ["INFINI_PREFILL_CHUNKED"] = "1" + os.environ["INFINI_PREFILL_CHUNK_SIZE"] = "8192" + self.assertEqual(long_prefill_threshold(), 8192) + + def test_explicit_env_overrides_chunked_default(self): + os.environ["INFINI_PREFILL_CHUNKED"] = "1" + os.environ["INFINI_PREFILL_CHUNK_SIZE"] = "8192" + os.environ["INFINI_LONG_PREFILL_THRESHOLD"] = "4096" + self.assertEqual(long_prefill_threshold(), 4096) + + def test_defaults_to_zero_when_chunked_prefill_disabled(self): + os.environ.pop("INFINI_PREFILL_CHUNKED", None) + self.assertEqual(long_prefill_threshold(), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/build_infini_stack.py b/scripts/build_infini_stack.py index 35a4e80b7..515740935 100644 --- a/scripts/build_infini_stack.py +++ b/scripts/build_infini_stack.py @@ -9,9 +9,12 @@ from typing import Dict, List, Mapping, Optional, Sequence PROJECT_ROOT = Path(__file__).resolve().parents[1] -SUPPORTED_BACKENDS = ("nvidia", "iluvatar", "hygon", "moore") +SUPPORTED_BACKENDS = ("nvidia", "iluvatar", "hygon", "moore", "qy") DEFAULT_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops.json" MOORE_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops_moore.json" +QY_M0_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops_qy_m0.json" +QY_M1_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops_qy_m1.json" +QY_M2_OPERATOR_CONFIG = PROJECT_ROOT / "scripts/configs/infiniops_ops_qy_m2.json" SUBMODULES = { "InfiniRT": Path("submodules/InfiniRT"), "InfiniOps": Path("submodules/InfiniOps"), @@ -80,7 +83,12 @@ def git_capture(arguments: Sequence[str], cwd: Path) -> str: ) -def validate_submodule(project_root: Path, relative_path: Path) -> str: +def validate_submodule( + project_root: Path, + relative_path: Path, + *, + allow_dirty: bool = False, +) -> str: submodule_root = project_root / relative_path if not (submodule_root / "CMakeLists.txt").is_file(): raise RuntimeError( @@ -93,14 +101,20 @@ def validate_submodule(project_root: Path, relative_path: Path) -> str: relative_path, ) actual = git_capture(["rev-parse", "HEAD"], submodule_root) - if actual != expected: + if actual != expected and not allow_dirty: raise RuntimeError( f"Submodule revision mismatch for {relative_path}: " f"expected {expected}, found {actual}" ) status = git_capture(["status", "--porcelain"], submodule_root) - if status: + if status and not allow_dirty: raise RuntimeError(f"Submodule worktree is dirty: {relative_path}") + if allow_dirty and (actual != expected or status): + print( + f"warning: allowing dirty/mismatched submodule {relative_path} " + f"(expected {expected}, actual {actual})", + flush=True, + ) return actual @@ -219,6 +233,15 @@ def build_infiniops_commands( iluvatar_arch: Optional[str] = None, hygon_arch: Optional[str] = None, ) -> List[List[str]]: + # QY: enable linked FA when the ops config selects slot-16 flash providers (M2+). + ops_text = operator_config.read_text(encoding="utf-8") + qy_wants_linked = ( + backend == "qy" + and "flash_attn_varlen_func" in ops_text + and "flash_attn_with_kvcache" in ops_text + ) + linked = "ON" if (backend != "qy" or qy_wants_linked) else "OFF" + generate_bindings = "ON" if backend == "qy" else "OFF" configure = [ "cmake", "-S", @@ -227,11 +250,11 @@ def build_infiniops_commands( str(build), "-DWITH_CPU=ON", backend_cmake_option(backend), - "-DWITH_LINKED=ON", + f"-DWITH_LINKED={linked}", "-DWITH_TORCH=ON", "-DAUTO_DETECT_DEVICES=OFF", "-DAUTO_DETECT_BACKENDS=OFF", - "-DGENERATE_PYTHON_BINDINGS=OFF", + f"-DGENERATE_PYTHON_BINDINGS={generate_bindings}", f"-DPython_EXECUTABLE={sys.executable}", f"-DINFINI_RT_ROOT={prefix}", f"-DINFINI_OPS_OPS={operator_config}", @@ -239,20 +262,46 @@ def build_infiniops_commands( f"-DCMAKE_INSTALL_PREFIX={prefix}", ] append_gpu_arch(configure, backend, cuda_arch, iluvatar_arch, hygon_arch) + if backend == "qy": + qy_root = Path( + os.environ.get("QY_ROOT") + or os.environ.get("QY_HOME") + or os.environ.get("QY_PATH") + or "/usr/local/denglin" + ) + configure.append(f"-DQY_ROOT={qy_root}") + try: + import pybind11 - return [ - configure, - [ - "cmake", - "--build", - str(build), - "--target", - "infiniops", - "--parallel", - str(jobs), - ], - ["cmake", "--install", str(build)], + configure.append(f"-Dpybind11_DIR={pybind11.get_cmake_dir()}") + except Exception: + pass + + build_cmd = [ + "cmake", + "--build", + str(build), + "--target", + "infiniops", + "--parallel", + str(jobs), ] + commands = [configure, build_cmd] + # QY M0 enables GENERATE_PYTHON_BINDINGS; install expects the ops module. + if backend == "qy": + commands.append( + [ + "cmake", + "--build", + str(build), + "--target", + "ops", + "--parallel", + str(jobs), + ] + ) + commands.append(["cmake", "--install", str(build)]) + return commands def build_infiniccl_commands( @@ -267,23 +316,31 @@ def build_infiniccl_commands( iluvatar_arch: Optional[str] = None, hygon_arch: Optional[str] = None, ) -> List[List[str]]: + # QY InfiniCCL/NCCL is M3; TP1 keeps InfiniCCL CPU-only (no WITH_QY yet). + ccl_backend = backend + ccl_test = test + if backend == "qy": + ccl_backend = None + ccl_test = False + configure = [ "cmake", "-S", str(source), "-B", str(build), - backend_cmake_option(backend), - f"-DWITH_NCCL={'OFF' if backend == 'moore' else 'ON'}", + "-DWITH_NCCL=OFF" if backend in ("moore", "qy") else "-DWITH_NCCL=ON", f"-DWITH_MCCL={'ON' if backend == 'moore' else 'OFF'}", "-DWITH_OMPI=OFF", "-DWITH_MPICH=OFF", "-DAUTO_DETECT_DEVICES=OFF", "-DAUTO_DETECT_BACKENDS=OFF", - f"-DBUILD_EXAMPLES={'ON' if test else 'OFF'}", + f"-DBUILD_EXAMPLES={'ON' if ccl_test else 'OFF'}", f"-DCMAKE_BUILD_TYPE={build_type}", f"-DCMAKE_INSTALL_PREFIX={prefix}", ] + if ccl_backend is not None: + configure.insert(5, backend_cmake_option(ccl_backend)) append_gpu_arch(configure, backend, cuda_arch, iluvatar_arch, hygon_arch) if backend == "iluvatar": configure.extend( @@ -310,7 +367,7 @@ def build_infiniccl_commands( ["cmake", "--build", str(build), "--parallel", str(jobs)], ["cmake", "--install", str(build)], ] - if test: + if ccl_test: commands.append( [ str(build / "examples/ccl/all_reduce"), @@ -465,6 +522,14 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: action="store_true", help="Print commands without executing them.", ) + parser.add_argument( + "--allow-dirty-submodules", + action="store_true", + help=( + "Allow dirty or gitlink-mismatched InfiniCore submodules " + "(for local M0 feature-branch builds)." + ), + ) args = parser.parse_args(argv) if args.jobs < 1: parser.error("--jobs must be positive") @@ -479,13 +544,16 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: if args.backend == "hygon" and args.hygon_arch is None: args.hygon_arch = "gfx936" if args.operator_config is None: - if args.backend not in ("nvidia", "moore"): + if args.backend == "qy": + args.operator_config = QY_M1_OPERATOR_CONFIG + elif args.backend not in ("nvidia", "moore"): parser.error(f"--operator-config is required with --backend={args.backend}") - args.operator_config = ( - MOORE_OPERATOR_CONFIG - if args.backend == "moore" - else DEFAULT_OPERATOR_CONFIG - ) + else: + args.operator_config = ( + MOORE_OPERATOR_CONFIG + if args.backend == "moore" + else DEFAULT_OPERATOR_CONFIG + ) if args.build_root is None: args.build_root = Path("build/integration") / args.backend return args @@ -509,7 +577,11 @@ def main(argv: Optional[Sequence[str]] = None) -> int: operator_config_path = resolve_from_project(args.operator_config) operator_config = read_operator_config(operator_config_path) revisions = { - name: validate_submodule(infinicore_root, relative_path) + name: validate_submodule( + infinicore_root, + relative_path, + allow_dirty=args.allow_dirty_submodules, + ) for name, relative_path in SUBMODULES.items() } infinilm_revision = git_capture(["rev-parse", "HEAD"], PROJECT_ROOT) @@ -527,6 +599,29 @@ def main(argv: Optional[Sequence[str]] = None) -> int: ], ) ) + elif args.backend == "qy": + qy_root = Path( + build_env.setdefault( + "QY_ROOT", + build_env.get("QY_HOME") + or build_env.get("QY_PATH") + or "/usr/local/denglin", + ) + ) + build_env["CUDA_HOME"] = str(qy_root / "sdk") + build_env["CUDA_PATH"] = str(qy_root / "sdk") + build_env["LD_LIBRARY_PATH"] = os.pathsep.join( + filter( + None, + [ + str(qy_root / "sdk/lib"), + build_env.get("LD_LIBRARY_PATH", ""), + ], + ) + ) + build_env["PATH"] = os.pathsep.join( + filter(None, [str(qy_root / "sdk/bin"), build_env.get("PATH", "")]) + ) print(f"InfiniLM: {infinilm_revision}") print(f"InfiniCore: {infinicore_revision}") diff --git a/scripts/configs/infiniops_ops_qy_m0.json b/scripts/configs/infiniops_ops_qy_m0.json new file mode 100644 index 000000000..ca05294d1 --- /dev/null +++ b/scripts/configs/infiniops_ops_qy_m0.json @@ -0,0 +1,8 @@ +{ + "add": { + "implementations": [0] + }, + "gemm": { + "implementations": [0] + } +} diff --git a/scripts/configs/infiniops_ops_qy_m1.json b/scripts/configs/infiniops_ops_qy_m1.json new file mode 100644 index 000000000..83c615156 --- /dev/null +++ b/scripts/configs/infiniops_ops_qy_m1.json @@ -0,0 +1,22 @@ +{ + "add": {"implementations": [0]}, + "argmax": {"implementations": [0]}, + "convolution": {"implementations": [0]}, + "copy": {"implementations": [0]}, + "embedding": {"implementations": [0]}, + "fill": {"implementations": [0]}, + "fused_add_rms_norm": {"implementations": [0]}, + "gelu": {"implementations": [0]}, + "gemm": {"implementations": [0]}, + "mul": {"implementations": [0]}, + "relu": {"implementations": [0]}, + "rms_norm": {"implementations": [0]}, + "rotary_embedding": {"implementations": [0]}, + "sigmoid": {"implementations": [0]}, + "silu": {"implementations": [0]}, + "silu_and_mul": {"implementations": [0]}, + "softmax": {"implementations": [0]}, + "topk_softmax": {"implementations": [0]}, + "tril": {"implementations": [8]}, + "triu": {"implementations": [8]} +} diff --git a/scripts/configs/infiniops_ops_qy_m2.json b/scripts/configs/infiniops_ops_qy_m2.json new file mode 100644 index 000000000..54d0df8f4 --- /dev/null +++ b/scripts/configs/infiniops_ops_qy_m2.json @@ -0,0 +1,25 @@ +{ + "add": {"implementations": [0]}, + "argmax": {"implementations": [0]}, + "convolution": {"implementations": [0]}, + "copy": {"implementations": [0]}, + "embedding": {"implementations": [0]}, + "fill": {"implementations": [0]}, + "flash_attn_varlen_func": {"implementations": [16]}, + "flash_attn_with_kvcache": {"implementations": [16]}, + "fused_add_rms_norm": {"implementations": [0]}, + "gelu": {"implementations": [0]}, + "gemm": {"implementations": [0]}, + "mul": {"implementations": [0]}, + "relu": {"implementations": [0]}, + "reshape_and_cache_flash": {"implementations": [0]}, + "rms_norm": {"implementations": [0]}, + "rotary_embedding": {"implementations": [0]}, + "sigmoid": {"implementations": [0]}, + "silu": {"implementations": [0]}, + "silu_and_mul": {"implementations": [0]}, + "softmax": {"implementations": [0]}, + "topk_softmax": {"implementations": [0]}, + "tril": {"implementations": [8]}, + "triu": {"implementations": [8]} +} diff --git a/scripts/debug_first_token.py b/scripts/debug_first_token.py new file mode 100644 index 000000000..9ab965ac6 --- /dev/null +++ b/scripts/debug_first_token.py @@ -0,0 +1,316 @@ +#!/usr/bin/env python3 +"""Dump first-token argmax / top-5 logits for MiniCPM5-2B on QY. + +Runs one prefill forward (InferEngine) with static attn (graph off) and +optionally flash-attn + paged KV. Writes JSON under --out-dir. + +Tensor→numpy patterns follow test/models/llama/test_llama_inference.py +(via InfiniCore Tensor.to_numpy). Does not wire HookRegistry. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +import traceback +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import numpy as np + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--model", + default=os.environ.get("MODEL_DIR", "/models/minicpm5-2b"), + help="HF model directory", + ) + p.add_argument("--device", default="qy", help="InfiniCore device type") + p.add_argument("--device-index", type=int, default=0) + p.add_argument( + "--prompt", + default="你好", + help="User prompt (chat-templated like the OpenAI server)", + ) + p.add_argument( + "--modes", + default="static,fa", + help="Comma-separated modes: static and/or fa", + ) + p.add_argument( + "--out-dir", + default=os.environ.get( + "OUT_DIR", + str(Path(__file__).resolve().parents[2] / "logs"), + ), + help="Directory for m1-2b-first-token-{mode}.json", + ) + p.add_argument("--max-cache-len", type=int, default=4096) + p.add_argument("--num-blocks", type=int, default=64) + p.add_argument("--block-size", type=int, default=256) + p.add_argument("--top-k", type=int, default=5) + return p.parse_args() + + +def _logits_hash(arr: np.ndarray) -> str: + return hashlib.sha256(np.ascontiguousarray(arr).tobytes()).hexdigest()[:16] + + +def _top_k(logits_1d: np.ndarray, k: int) -> List[Dict[str, Any]]: + k = min(k, int(logits_1d.shape[0])) + idx = np.argpartition(logits_1d, -k)[-k:] + idx = idx[np.argsort(-logits_1d[idx])] + return [{"token_id": int(i), "logit": float(logits_1d[i])} for i in idx] + + +def _tokenize_chat(model_dir: str, user_prompt: str) -> Tuple[Any, List[int], str]: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) + messages = [{"role": "user", "content": user_prompt}] + rendered = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + input_ids = tokenizer.encode(rendered, add_special_tokens=False) + return tokenizer, input_ids, rendered + + +def _build_engine( + model_path: str, + device, + *, + mode: str, + max_cache_len: int, + num_blocks: int, + block_size: int, +): + from infinilm.cache.cache import PagedKVCacheConfig, StaticKVCacheConfig + from infinilm.distributed import DistConfig + from infinilm.infer_engine import InferEngine + from infinilm.modeling_utils import load_model_state_dict_by_file + + if mode == "static": + cache_config = StaticKVCacheConfig( + max_batch_size=1, max_cache_len=max_cache_len + ) + attention_backend = "default" + enable_graph = False + elif mode == "fa": + cache_config = PagedKVCacheConfig( + num_blocks=num_blocks, + block_size=block_size, + max_batch_size=1, + ) + attention_backend = "flash-attn" + enable_graph = False + else: + raise ValueError(f"unknown mode: {mode}") + + engine = InferEngine( + model_path=model_path, + device=device, + distributed_config=DistConfig(1), + cache_config=cache_config, + enable_graph_compiling=enable_graph, + attention_backend=attention_backend, + weight_load_mode="sync", + ) + load_model_state_dict_by_file(engine, model_path, dtype=engine.dtype) + return engine + + +def _tensor_to_f32_numpy(tensor) -> np.ndarray: + """InfiniCore→numpy; bf16 has no numpy dtype so go through torch.""" + import ctypes + + import infinicore + import torch + + source = tensor + if source.device.type != "cpu": + source = source.to(infinicore.device("cpu", 0)) + try: + infinicore.sync_device() + except Exception: + pass + if not source.is_contiguous(): + source = source.contiguous() + + shape = list(source.shape) + dtype_name = str(source.dtype) + if "bfloat16" in dtype_name: + torch_t = torch.empty(shape, dtype=torch.bfloat16) + if torch_t.numel(): + ctypes.memmove(torch_t.data_ptr(), source.data_ptr(), torch_t.numel() * 2) + return torch_t.float().numpy() + if "float16" in dtype_name or dtype_name.endswith("f16"): + torch_t = torch.empty(shape, dtype=torch.float16) + if torch_t.numel(): + ctypes.memmove(torch_t.data_ptr(), source.data_ptr(), torch_t.numel() * 2) + return torch_t.float().numpy() + return np.asarray(source.to_numpy(), dtype=np.float32) + + +def _prefill_logits(engine, input_ids: List[int], mode: str, block_size: int): + import infinicore + + seq_len = len(input_ids) + batch_size = 1 + past_seq_len = 0 + + ids = infinicore.from_list([input_ids], dtype=infinicore.int64) + position_ids = infinicore.from_list( + [list(range(past_seq_len, past_seq_len + seq_len))], + dtype=infinicore.int64, + ) + past_kv_lengths = infinicore.from_list( + [past_seq_len] * batch_size, dtype=infinicore.int32 + ) + total_kv_lengths = infinicore.from_list( + [past_seq_len + seq_len] * batch_size, dtype=infinicore.int32 + ) + cu_seqlens = infinicore.from_list( + [(past_seq_len + seq_len) * i for i in range(batch_size + 1)], + dtype=infinicore.int32, + ) + input_offsets = infinicore.from_list( + [seq_len * i for i in range(batch_size + 1)], + dtype=infinicore.int32, + ) + + block_tables = None + slot_mapping = None + if mode == "fa": + # Match InferEngine.generate paged prefill bookkeeping (batch=1). + max_blocks = (seq_len + block_size - 1) // block_size + max_blocks = max(max_blocks, 1) + block_tables = infinicore.from_list( + [list(range(max_blocks))], dtype=infinicore.int32 + ) + slot_mapping = infinicore.from_list( + list(range(seq_len)), dtype=infinicore.int64 + ) + ids = ids.view([1, seq_len]) + + out = engine.forward_raw( + input_ids=ids, + position_ids=position_ids, + past_kv_lengths=past_kv_lengths, + total_kv_lengths=total_kv_lengths, + input_offsets=input_offsets, + cu_seqlens=cu_seqlens, + block_tables=block_tables, + slot_mapping=slot_mapping, + sample_all_positions=True, + ) + arr = _tensor_to_f32_numpy(out["logits"]) + # Expected [B, T, V] or [B, 1, V] (last token only). + if arr.ndim == 3: + last = arr[0, -1, :] + elif arr.ndim == 2: + last = arr[-1, :] + else: + raise RuntimeError(f"unexpected logits shape: {arr.shape}") + return arr, last.astype(np.float32, copy=False) + + +def _run_mode(args: argparse.Namespace, mode: str) -> Dict[str, Any]: + import infinicore + + record: Dict[str, Any] = { + "mode": mode, + "model": args.model, + "device": args.device, + "prompt": args.prompt, + "status": "ok", + } + try: + tokenizer, input_ids, rendered = _tokenize_chat(args.model, args.prompt) + record["rendered_prompt"] = rendered + record["input_ids"] = input_ids + record["num_prompt_tokens"] = len(input_ids) + + device = infinicore.device(args.device, args.device_index) + engine = _build_engine( + args.model, + device, + mode=mode, + max_cache_len=args.max_cache_len, + num_blocks=args.num_blocks, + block_size=args.block_size, + ) + full_logits, last = _prefill_logits( + engine, input_ids, mode=mode, block_size=args.block_size + ) + argmax_id = int(np.argmax(last)) + top = _top_k(last, args.top_k) + # Attach decoded text for readability. + for item in top: + item["token"] = tokenizer.decode( + [item["token_id"]], skip_special_tokens=False + ) + record.update( + { + "argmax": argmax_id, + "argmax_token": tokenizer.decode( + [argmax_id], skip_special_tokens=False + ), + "top5": top, + "logits_shape": list(full_logits.shape), + "logits_dtype": str(full_logits.dtype), + "logits_hash": _logits_hash(last), + "logits_l2": float(np.linalg.norm(last)), + "logits_stats": { + "min": float(last.min()), + "max": float(last.max()), + "mean": float(last.mean()), + }, + } + ) + except Exception as e: + record["status"] = "error" + record["error"] = f"{type(e).__name__}: {e}" + record["traceback"] = traceback.format_exc() + return record + + +def main() -> int: + args = _parse_args() + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + modes = [m.strip() for m in args.modes.split(",") if m.strip()] + rc = 0 + for mode in modes: + print(f"=== mode={mode} ===", flush=True) + record = _run_mode(args, mode) + out_path = out_dir / f"m1-2b-first-token-{mode}.json" + with open(out_path, "w", encoding="utf-8") as f: + json.dump(record, f, ensure_ascii=False, indent=2) + f.write("\n") + print(f"wrote {out_path} status={record.get('status')}", flush=True) + if record.get("status") == "ok": + print( + f" argmax={record['argmax']} ({record['argmax_token']!r}) " + f"hash={record['logits_hash']} l2={record['logits_l2']:.4f}", + flush=True, + ) + for i, item in enumerate(record["top5"]): + print( + f" top{i + 1}: id={item['token_id']} " + f"logit={item['logit']:.6f} token={item['token']!r}", + flush=True, + ) + else: + print(f" ERROR: {record.get('error')}", flush=True) + rc = 1 + return rc + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/static/test_core_graph_sync.py b/test/static/test_core_graph_sync.py index 93b7b8e4e..15c92e9a0 100644 --- a/test/static/test_core_graph_sync.py +++ b/test/static/test_core_graph_sync.py @@ -42,10 +42,11 @@ def test_graph_replay_is_partitioned_by_capture_safety(self) -> None: self.assertIn(contract, header) warmup = instantiate.index("for (size_t iter = 0; iter < 5; ++iter)") - disable = instantiate.index("INFINICORE_DISABLE_DEVICE_GRAPH_SEGMENTS") + policy = instantiate.index("deviceGraphSegmentsEnabled") partition = instantiate.index("op->is_device_graph_capture_safe()") - self.assertLess(warmup, disable) - self.assertLess(disable, partition) + self.assertLess(warmup, policy) + self.assertLess(policy, partition) + self.assertIn("INFINI_CUDAGRAPH_POLICY", instantiate) self.assertIn("segments_.back()->ops.push_back(op)", instantiate) self.assertIn("if (!segment->capture_safe)", instantiate) self.assertIn("segment->run()", instantiate) @@ -55,7 +56,16 @@ def test_graph_replay_is_partitioned_by_capture_safety(self) -> None: self.assertIn("host_segments", instantiate) self.assertIn("if (segments_.empty())", run) - self.assertIn("segment->run()", run) + self.assertIn("last_run_profile_", run) + self.assertIn("graph_run_profile", run) + self.assertIn("has_device_exec()", run) + + context_header = read_source( + "csrc/infinicore/include/infinicore/context/context.hpp" + ) + self.assertIn("cudagraphPolicy()", context_header) + self.assertIn("deviceGraphSegmentsEnabled()", context_header) + self.assertIn("LastRunProfile", header) def test_segment_cleanup_keeps_runtime_and_allocation_leases_alive(self) -> None: header = read_source("csrc/infinicore/include/infinicore/graph/graph.hpp") diff --git a/test/static/test_infinicore_python_contracts.py b/test/static/test_infinicore_python_contracts.py index ae6e501ec..e6b4cc667 100644 --- a/test/static/test_infinicore_python_contracts.py +++ b/test/static/test_infinicore_python_contracts.py @@ -66,8 +66,7 @@ def test_ppu_detection_selects_thead(self) -> None: config._torch_device_available = lambda name: False self.assertEqual(config.get_device_str("thead"), "thead") self.assertEqual(config.detect_device(), "thead") - with self.assertRaisesRegex(ValueError, "unsupported device platform 'qy'"): - config.get_device_str("qy") + self.assertEqual(config.get_device_str("qy"), "qy") def test_rope_wrapper_preserves_arguments_and_output_identity(self) -> None: calls = [] diff --git a/test/static/test_modern_infinicore_compatibility.py b/test/static/test_modern_infinicore_compatibility.py index 29988fbb5..b49179465 100644 --- a/test/static/test_modern_infinicore_compatibility.py +++ b/test/static/test_modern_infinicore_compatibility.py @@ -267,7 +267,11 @@ def test_current_infiniops_backend_bridges_are_registered(self) -> None: engine, ) self.assertIn( - "flash-attn is only available on NVIDIA, MetaX, Moore, Cambricon, Iluvatar, Hygon, and Ascend devices", + "device_type != infinicore::Device::Type::kQy", + engine, + ) + self.assertIn( + "flash-attn is only available on NVIDIA, MetaX, Moore, Cambricon, Iluvatar, Hygon, Ascend, and QY devices", engine, ) self.assertIn( diff --git a/xmake.lua b/xmake.lua index 42e93b626..ec8a5e51b 100644 --- a/xmake.lua +++ b/xmake.lua @@ -1,6 +1,20 @@ -add_requires("pybind11") -if is_plat("linux") then - add_requireconfs("pybind11.python", { configs = { headeronly = true } }) +-- Prefer pip/system pybind11 to avoid xmake-repo network fetches in offline containers. +local PYBIND11_INCLUDE = os.getenv("PYBIND11_INCLUDE") +if not PYBIND11_INCLUDE or PYBIND11_INCLUDE == "" then + local ok, out = pcall(function() + return os.iorunv(os.getenv("PYTHON") or "python3", { + "-c", "import pybind11; print(pybind11.get_include())" + }) + end) + if ok and out and out:trim() ~= "" then + PYBIND11_INCLUDE = out:trim() + end +end +if not PYBIND11_INCLUDE or PYBIND11_INCLUDE == "" then + add_requires("pybind11") + if is_plat("linux") then + add_requireconfs("pybind11.python", { configs = { headeronly = true } }) + end end local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") @@ -16,6 +30,24 @@ local OPERATOR_CALLS_HEADER = INFINI_ROOT .. "/include/infini/operator_call_inst set_toolchains("gcc") +-- Fallback when xmake python plugin / network repos are unavailable. +rule("python.module") + on_config(function (target) + target:set("kind", "shared") + target:set("prefixname", "") + local suffix = try { + function () + return os.iorunv(os.getenv("PYTHON") or "python3", { + "-c", "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX') or '')" + }):trim() + end + } + if suffix and suffix ~= "" and suffix ~= "None" then + target:set("extension", suffix) + end + end) +rule_end() + option("cxx11-abi") set_default(nil) set_showmenu(true) @@ -50,11 +82,18 @@ if os.isdir(MUSA_ROOT .. "/include") then add_linkdirs(MUSA_ROOT .. "/lib", MUSA_ROOT .. "/lib64") end +local QY_SDK_PRESENT = os.getenv("QY_ROOT") or os.getenv("QY_HOME") or os.getenv("QY_PATH") + or os.isdir("/usr/local/denglin/sdk") or os.isdir("/usr/local/dlgpu/sdk") + target("infinicore_runtime") set_default(false) set_kind("shared") set_languages("cxx17") add_defines("ENABLE_INFINIOPS_API", "USE_INFINIRT_GRAPH") + -- Denglin QY: CUDA-compatible streams + ATen bridge for `_dl_C` MoE. + if QY_SDK_PRESENT then + add_defines("ENABLE_QY_API", "ENABLE_ATEN") + end if not is_plat("windows") then add_cxflags("-fPIC") @@ -75,6 +114,7 @@ target("infinicore_runtime") add_links("infiniops", "infiniccl", "infinirt", { public = true }) add_files("csrc/infinicore/src/*.cc") + add_files("csrc/infinicore/src/adaptor/*.cc") add_files("csrc/infinicore/src/context/*.cc") add_files("csrc/infinicore/src/context/*/*.cc") add_files("csrc/infinicore/src/graph/*.cc") @@ -91,6 +131,39 @@ target("infinicore_runtime") remove_files("csrc/infinicore/src/ops/*/*_moore.cc") on_load(function (target) + if QY_SDK_PRESENT then + local python = os.getenv("PYTHON") or "python3" + local torch_dir = try { + function () + return os.iorunv(python, {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim() + end + } + if torch_dir and torch_dir ~= "" then + target:add("includedirs", + path.join(torch_dir, "include"), + path.join(torch_dir, "include/torch/csrc/api/include"), + { public = true }) + target:add("linkdirs", path.join(torch_dir, "lib"), { public = true }) + target:add("rpathdirs", path.join(torch_dir, "lib")) + target:add("links", + "torch", "torch_cpu", "torch_cuda", "torch_python", + "c10", "c10_cuda", + { public = true }) + end + local qy_root = os.getenv("QY_ROOT") or os.getenv("QY_HOME") or os.getenv("QY_PATH") + if (not qy_root or qy_root == "") and os.isdir("/usr/local/dlgpu/sdk") then + qy_root = "/usr/local/dlgpu" + end + if (not qy_root or qy_root == "") and os.isdir("/usr/local/denglin/sdk") then + qy_root = "/usr/local/denglin" + end + if qy_root and os.isdir(qy_root .. "/sdk/include") then + target:add("includedirs", qy_root .. "/sdk/include", { public = true }) + target:add("linkdirs", qy_root .. "/sdk/lib", { public = true }) + target:add("links", "curt", { public = true }) + end + end + if not os.isfile(OPERATOR_CALLS_HEADER) then return end @@ -124,7 +197,11 @@ target("infinicore_runtime") target_end() target("_infinicore") - add_packages("pybind11") + if PYBIND11_INCLUDE and PYBIND11_INCLUDE ~= "" then + add_includedirs(PYBIND11_INCLUDE) + else + add_packages("pybind11") + end set_default(false) add_rules("python.module", { soabi = true }) set_languages("cxx17") @@ -150,7 +227,11 @@ target("_infinicore") target_end() target("_infinilm") - add_packages("pybind11") + if PYBIND11_INCLUDE and PYBIND11_INCLUDE ~= "" then + add_includedirs(PYBIND11_INCLUDE) + else + add_packages("pybind11") + end set_default(false) add_rules("python.module", { soabi = true }) set_languages("cxx17") @@ -178,6 +259,7 @@ target("_infinilm") add_files("csrc/models/internlm3/*.cpp") add_files("csrc/models/llama/*.cpp") add_files("csrc/models/minicpm4/*.cpp") + add_files("csrc/models/minicpm5_moe/*.cpp") add_files("csrc/models/minicpm_eagle/*.cpp") add_files("csrc/models/qwen2/*.cpp") add_files("csrc/models/qwen3/*.cpp")