From 9f262aa947fc26ceae780facdcfeadddc2f36a17 Mon Sep 17 00:00:00 2001 From: tangchengxiang <2064027004@qq.com> Date: Sat, 19 Sep 2026 12:59:05 +0000 Subject: [PATCH 1/7] feat(qwen): add paged greedy MTP with FP8 block weights Execute the checkpoint's shared MTP head through the existing engine and scheduler. Verify up to four candidates, commit matching Conv/GDN states without target replay, and reclaim request resources on stop or failure. Reuse column-parallel vocabulary projection, Marlin FP8 packing and PagedCompiler draft capture. Bound recurrent rows independently from KV pages and make exact-prompt snapshots opt-in. Wire MTP through the CLI, benchmark and service entrypoints; retain three core MTP test modules. Validated NVIDIA A6000 TP1/TP2, graph recapture, forced acceptance lengths, real 27B FP8 batching/cancellation, and ordinary Qwen2 pre-transpose. Runtime support: InfiniTensor/InfiniCore#1565; graphs also require #1560. --- README.md | 74 +++ csrc/cache/kv_cache.cpp | 5 +- csrc/cache/kv_cache.hpp | 5 +- csrc/config/quant_config.cpp | 4 +- csrc/engine/compiler/paged_compiler.cpp | 207 ++++++-- csrc/engine/compiler/paged_compiler.hpp | 3 + csrc/engine/infer_engine.cpp | 164 +++++- csrc/engine/infer_engine.hpp | 1 + csrc/engine/rank_worker.cpp | 131 +++-- csrc/engine/rank_worker.hpp | 14 +- csrc/global_state/forward_context.hpp | 7 + csrc/layers/attention/backends/paged_attn.cpp | 5 + csrc/layers/linear/vocab_parallel.cpp | 68 +++ csrc/layers/linear/vocab_parallel.hpp | 23 + csrc/layers/quantization/fp8_block.cpp | 172 +++++++ csrc/layers/quantization/fp8_block.hpp | 34 ++ .../layers/quantization/none_quantization.cpp | 19 +- .../layers/quantization/none_quantization.hpp | 7 +- csrc/layers/quantization/quantization.hpp | 1 + .../quantization/quantization_scheme.hpp | 1 + csrc/models/infinilm_model.hpp | 7 + csrc/models/qwen3_5/qwen3_5_attention.hpp | 8 + csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp | 53 +- csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp | 6 +- csrc/models/qwen3_5/qwen3_5_model.hpp | 4 + csrc/models/qwen3_5/qwen3_5_mtp.cpp | 51 ++ csrc/models/qwen3_5/qwen3_5_mtp.hpp | 25 + .../qwen3_next_allocate_kv_cache_tensors.cpp | 12 +- .../qwen3_next/qwen3_next_gated_deltanet.cpp | 152 +++++- .../qwen3_next/qwen3_next_gated_deltanet.hpp | 8 + csrc/pybind11/cache/cache.hpp | 6 +- csrc/pybind11/engine/engine.hpp | 54 +- examples/bench.py | 40 +- examples/test_infer.py | 9 + python/infinilm/base_config.py | 22 +- python/infinilm/cache/cache.py | 3 +- python/infinilm/config/engine_config.py | 63 ++- python/infinilm/infer_engine.py | 51 +- python/infinilm/llm/hybrid_prefix_cache.py | 140 ++++++ python/infinilm/llm/llm.py | 104 +++- .../infinilm/llm/model_runner/model_runner.py | 11 +- .../infinilm/llm/model_runner/mtp_runner.py | 468 ++++++++++++++++++ python/infinilm/llm/request.py | 11 + python/infinilm/llm/scheduler.py | 40 +- python/infinilm/modeling_utils.py | 62 ++- python/infinilm/server/inference_server.py | 16 + test/bench/backends/infinilm.py | 19 +- test/bench/test_benchmark.py | 6 + test/layers/test_pre_transpose.py | 78 +++ test/models/qwen3_5/test_mtp_execution.py | 206 ++++++++ test/models/qwen3_5/test_mtp_model.py | 284 +++++++++++ test/models/qwen3_5/test_mtp_runner.py | 435 ++++++++++++++++ 52 files changed, 3226 insertions(+), 173 deletions(-) create mode 100644 csrc/layers/linear/vocab_parallel.cpp create mode 100644 csrc/layers/linear/vocab_parallel.hpp create mode 100644 csrc/layers/quantization/fp8_block.cpp create mode 100644 csrc/layers/quantization/fp8_block.hpp create mode 100644 csrc/models/qwen3_5/qwen3_5_mtp.cpp create mode 100644 csrc/models/qwen3_5/qwen3_5_mtp.hpp create mode 100644 python/infinilm/llm/hybrid_prefix_cache.py create mode 100644 python/infinilm/llm/model_runner/mtp_runner.py create mode 100644 test/layers/test_pre_transpose.py create mode 100644 test/models/qwen3_5/test_mtp_execution.py create mode 100644 test/models/qwen3_5/test_mtp_model.py create mode 100644 test/models/qwen3_5/test_mtp_runner.py diff --git a/README.md b/README.md index 9a8b365f0..d642d0bfc 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,80 @@ 当前版本依赖[`InfiniCore v0.2.9`](https://github.com/InfiniTensor/InfiniCore/releases/tag/v0.2.9)版本。 +### Qwen built-in MTP (experimental) + +The text-only greedy path supports one built-in MTP layer, 1–4 draft tokens, +paged KV caching and PP1. Target verification selects matching Conv/GDN +checkpoints without replay. The scheduler batches target verification while +keeping acceptance and state independent for each request. TP1 also batches +draft continuation; TP2 retains separate draft calls because its end-to-end +batching benefit has not been established. +It handles cancellation, EOS/output limits and ordinary Decode when speculative +cache capacity is unavailable. This requires the matching FP8/MTP InfiniCore build; +the release dependency listed above does not contain these additions. Track the +runtime patch in [InfiniCore #1565](https://github.com/InfiniTensor/InfiniCore/issues/1565). +Graph execution additionally requires the graph lifetime/recording fixes in +[InfiniCore #1560](https://github.com/InfiniTensor/InfiniCore/pull/1560). + +For Qwen3.8-27B-FP8 with E4M3 weights and 128×128 weight blocks, set +`quantization_config.fp8_backend` to `"marlin"` in the checkpoint configuration +for NVIDIA inference. Weights are packed once during loading using the existing +Marlin operator; a separately converted weight file is unnecessary. The default +`"compatibility"` backend dequantizes on the device at execution time and is slower. + +Example on one A6000 with the Marlin configuration: + +```bash +python -m infinilm.server.inference_server \ + --model /models/Qwen3.8-27B-FP8-marlin --device nvidia --dtype bfloat16 \ + --enable-paged-attn --enable-mtp --num-draft-tokens 2 \ + --max-batch-size 2 --num-state-rows 9 --num-blocks 40 --block-size 64 \ + --disable-prefix-caching --top-k 1 --max-new-tokens 64 +``` + +The same MTP and cache options are accepted by `examples/test_infer.py`, +`examples/bench.py` and `test/bench/test_benchmark.py`. The offline benchmark +reuses the scheduler-backed `LLM` path for MTP; it does not time model loading. + +`num_state_rows` counts the zero row, committed request states and speculative +checkpoints. Its MTP default is `1 + max_batch_size * (num_draft_tokens + 2)`, +independent of KV page count. A smaller pool can reduce concurrent admission or +use ordinary Decode when checkpoint rows are unavailable. The page budget must +also accommodate each prompt and its requested output limit. + +For exact full-prompt reuse, replace `--disable-prefix-caching` with +`--mtp-prefix-cache-mib 512`. This TP1-only LRU cache owns device copies of both +target/draft KV, recurrent state and the initial MTP outputs. Hits restore into +request-owned pages and state rows. The budget limits live snapshot tensor +storage, not model memory, allocator reservations or total process memory. +Partial-prefix matching and the Attention-only cache's SLRU policy are not +supported by this hybrid snapshot cache. Cache reset or weight loading +invalidates snapshots. + +`--enable-graph` currently requires `--num-draft-tokens 1 --max-batch-size 1`. +It captures ordinary Decode and supported short draft shapes; Prefill and target +verification remain eager. Multi-candidate and batched execution use eager. +Random sampling, multimodal requests, multi-layer MTP service execution and remote +state transfer are rejected. NVIDIA A6000 validation covers TP1 and TP2 greedy +execution, batched requests, cancellation and cache reclamation. The 27B FP8 TP2 +checks use K=2; K=1/2/4 and graph recapture are additionally checked with a tiny +checkpoint. Exact full-prompt caching remains TP1-only. Other accelerators have +not been validated for this service path. + +Control-flow and GPU integration checks: + +```bash +python -m pytest test/models/qwen3_5 -q +INFINILM_QWEN_MTP_TEST_MODEL=/models/tiny-qwen-mtp \ +INFINILM_QWEN_MTP_TEST_TP=1 python -m pytest \ + test/models/qwen3_5 -q +``` + +For the TP2 execution and batching checks, expose two GPUs and set +`INFINILM_QWEN_MTP_TEST_TP=2`; the prefix-cache check still uses TP1. +The three test modules cover CPU scheduling/lifecycle, GPU execution, and +checkpoint/model contracts. GPU checks skip when no test checkpoint is set. + ## 使用方式 #### 一、编译并安装 `InfiniCore` 编译并安装 `InfiniCore`， 详情见 InfiniCore的 [`README`](https://github.com/InfiniTensor/InfiniCore) : diff --git a/csrc/cache/kv_cache.cpp b/csrc/cache/kv_cache.cpp index 2a7780090..9c322026a 100644 --- a/csrc/cache/kv_cache.cpp +++ b/csrc/cache/kv_cache.cpp @@ -80,10 +80,11 @@ infinicore::Tensor create_layer_kv_cache( PagedKVCacheConfig::PagedKVCacheConfig( size_t num_blocks, size_t block_size, - size_t max_batch_size) + size_t max_batch_size, + size_t num_state_rows) : num_blocks_(num_blocks), block_size_(block_size), - max_batch_size_(max_batch_size) { + max_batch_size_(max_batch_size), num_state_rows_(num_state_rows) { } std::unique_ptr diff --git a/csrc/cache/kv_cache.hpp b/csrc/cache/kv_cache.hpp index 760fb6685..5f92ab66b 100644 --- a/csrc/cache/kv_cache.hpp +++ b/csrc/cache/kv_cache.hpp @@ -40,17 +40,20 @@ class PagedKVCacheConfig final : public CacheConfig { PagedKVCacheConfig( size_t num_blocks, size_t block_size = 256, - size_t max_batch_size = 1); + size_t max_batch_size = 1, + size_t num_state_rows = 0); std::unique_ptr unique_copy() const override; size_t num_blocks() const; size_t block_size() const; size_t max_batch_size() const; + size_t num_state_rows() const { return num_state_rows_; } private: size_t num_blocks_; size_t block_size_; size_t max_batch_size_; + size_t num_state_rows_; }; namespace PagedKVCache { diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index e58966d89..e1a414a42 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -14,7 +14,9 @@ QuantConfig::get_quantization_method() const { const std::string quant_method = quantization_config.value("quant_method", ""); // Determine the quantization scheme from the JSON config - if (quant_method == "compressed-tensors") { + if (quant_method == "fp8") { + return std::make_shared(quantization_config); + } else if (quant_method == "compressed-tensors") { return std::make_shared(quantization_config); } else if (quant_method == "awq") { return std::make_shared(quantization_config); diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index dee3123c9..e75c94b9f 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -23,6 +24,43 @@ bool has_mamba_cache(const infinilm::global_state::ForwardContext &forward_conte return has_state(forward_context.conv_state_vec) || has_state(forward_context.ssm_state_vec); } +class CacheStateGuard { +public: + void save(const infinicore::Tensor &state, size_t rows) { + if (!state) { + return; + } + save_region(state->narrow({{0, 1, rows}})); + } + + void save_region(const infinicore::Tensor ®ion) { + auto backup = infinicore::Tensor::empty(region->shape(), region->dtype(), region->device()); + backup->copy_from(region); + saved_.emplace_back(region, backup); + } + + void restore() { + for (auto &[region, backup] : saved_) { + region->copy_from(backup); + } + if (!saved_.empty()) { + infinicore::context::syncStream(); + saved_.clear(); + } + } + + ~CacheStateGuard() { + try { + restore(); + } catch (const std::exception &error) { + spdlog::error("Failed to restore request states after graph capture: {}", error.what()); + } + } + +private: + std::vector> saved_; +}; + } // namespace PagedCompiler::PagedCompiler(const std::shared_ptr &model, RankBarrier *barrier) @@ -70,8 +108,43 @@ void PagedCompiler::compile() { throw std::runtime_error("PagedCompiler: position_id_axes must be positive"); } - size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); + infinicore::context::syncStream(); compiled_map_decode_.clear(); + compiled_map_draft_.clear(); + const bool capture_mtp = model_->supports_token_state_checkpoints() + && infinicore::context::getDevice().getType() == infinicore::Device::Type::NVIDIA; + if (decode_batch_sizes_.empty()) { + return; + } + size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); + if (has_mamba_state) { + for (const auto *states : {&forward_context.conv_state_vec, &forward_context.ssm_state_vec}) { + for (const auto &state : *states) { + if (state) { + const size_t capacity = state->size(0) == 0 ? 0 : state->size(0) - 1; + max_batch_size = std::min(max_batch_size, capacity); + } + } + } + } + if (max_batch_size == 0) { + return; + } + CacheStateGuard state_guard; + if (has_mamba_state) { + for (const auto *states : {&forward_context.conv_state_vec, &forward_context.ssm_state_vec}) { + for (const auto &state : *states) { + state_guard.save(state, max_batch_size); + } + } + } + // Warmup and capture write into physical page zero. Preserve it so + // recapturing also remains safe while a request owns that page. + for (const auto &kv : forward_context.kv_cache_vec) { + if (capture_mtp && kv) { + state_guard.save_region(kv->narrow({{1, 0, 1}})); + } + } block_tables_holder_ = infinicore::Tensor::empty( {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice()); set_zeros(block_tables_holder_); @@ -109,7 +182,8 @@ void PagedCompiler::compile() { input.mamba_final_state_indices = infinicore::Tensor::empty( {b}, infinicore::DataType::I32, infinicore::context::getDevice()); std::vector init_state_indices_vec(b, 0); - std::vector final_state_indices_vec(b, 1); + std::vector final_state_indices_vec(b); + std::iota(final_state_indices_vec.begin(), final_state_indices_vec.end(), 1); infinicore::context::memcpyH2D( input.mamba_init_state_indices.value()->data(), init_state_indices_vec.data(), @@ -154,46 +228,110 @@ void PagedCompiler::compile() { infinicore::context::syncStream(); } - for (size_t b : decode_batch_sizes_) { - auto input = make_decode_input(b); - + auto capture = [&](InfinilmModel::Input input) { + const auto lengths = forward_context.attn_metadata.verification_sequence_lengths; + const auto tables = forward_context.attn_metadata.verification_block_tables; barrier_->wait(); (void)model_->forward(input); infinicore::context::syncStream(); - // Capture must not start with stale Marlin locks from previous - // warmup/capture attempts. This reset is intentionally outside - // graph capture; the current implementation still pays a memset - // before every graph replay in get_compiled(). model_->reset_runtime_state(); infinicore::context::syncStream(); infinicore::context::startGraphRecording(); auto output = model_->forward(input); auto graph = infinicore::context::stopGraphRecording(); barrier_->wait(); - - auto shared_output = std::shared_ptr( - new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)}); - - compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)}; + auto shared_output = std::make_shared(); + shared_output->logits = infinicore::graph::GraphTensor(output.logits); + if (output.hidden_states) { + shared_output->hidden_states = infinicore::graph::GraphTensor(output.hidden_states); + } + return CompiledResult{std::move(input), {graph, shared_output}, lengths, tables}; + }; + for (size_t b : decode_batch_sizes_) { + if (b <= max_batch_size) { + compiled_map_decode_[b] = capture(make_decode_input(b)); + } + } + if (capture_mtp) { + auto make_draft_input = [&](size_t tokens) { + auto input = make_decode_input(1); + const auto device = infinicore::context::getDevice(); + auto i32 = [&](const std::vector &values) { + auto result = infinicore::Tensor::empty({values.size()}, infinicore::DataType::I32, device); + infinicore::context::memcpyH2D(result->data(), values.data(), values.size() * sizeof(int32_t), false); + return result; + }; + input.input_ids = infinicore::Tensor::zeros({1, tokens}, infinicore::DataType::I64, device); + input.position_ids = infinicore::Tensor::zeros( + position_id_axes > 1 ? std::vector{position_id_axes, tokens} : std::vector{tokens}, + infinicore::DataType::I64, device); + input.total_sequence_lengths = i32({static_cast(tokens)}); + input.input_offsets = i32({0, static_cast(tokens)}); + input.cu_seqlens = i32({0, static_cast(tokens)}); + input.slot_mapping = infinicore::Tensor::empty({tokens}, infinicore::DataType::I64, device); + std::vector slots(tokens); + std::iota(slots.begin(), slots.end(), 0); + infinicore::context::memcpyH2D(input.slot_mapping.value()->data(), slots.data(), tokens * sizeof(int64_t), false); + input.target_hidden_states = infinicore::Tensor::zeros( + {1, tokens, model_config->get("hidden_size")}, model_config->get_dtype(), device); + forward_context.attn_metadata = global_state::AttentionMetadata(input); + forward_context.mamba_metadata = {input.input_offsets, input.mamba_init_state_indices, + input.mamba_final_state_indices, input.token_state_indices}; + if (tokens > 1) { + forward_context.attn_metadata.verification_sequence_lengths = i32({1, 2}); + forward_context.attn_metadata.verification_block_tables = infinicore::Tensor::zeros( + {tokens, nblocks}, infinicore::DataType::I32, device); + } + return input; + }; + for (size_t tokens : {1, 2}) { + compiled_map_draft_[tokens] = capture(make_draft_input(tokens)); + } } + state_guard.restore(); } } PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &input) { + // Q=2 target graphs are intentionally excluded: measured slower than + // short eager verification. They also require their own checkpoint mode. + if (input.token_state_indices) { + return {nullptr, nullptr}; + } if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { size_t batch_size = input.block_tables.value()->size(0); size_t block_per_req = input.block_tables.value()->size(1); - - // only support decode only batch - if (batch_size != input.input_ids.value()->size(1)) { + const size_t tokens = input.input_ids.value()->size(1); + const bool draft = input.target_hidden_states.has_value(); + if ((draft && batch_size != 1) + || (!draft && (batch_size != tokens || input.sample_all_positions))) { return {nullptr, nullptr}; - } else { - auto result = compiled_map_decode_.find(batch_size); - if (result == compiled_map_decode_.end()) { + } + { + auto &compiled = draft ? compiled_map_draft_ : compiled_map_decode_; + auto result = compiled.find(draft ? tokens : batch_size); + if (result == compiled.end()) { return {nullptr, nullptr}; } auto &graph_input = result->second.input; + const auto &runtime_seq_lens = input.total_sequence_lengths.value(); + if (runtime_seq_lens->device().getType() + != infinicore::Device::Type::CPU + || runtime_seq_lens->dtype() != infinicore::DataType::I32 + || runtime_seq_lens->shape().size() != 1 + || runtime_seq_lens->shape()[0] != batch_size) { + throw std::runtime_error( + "PagedCompiler expected CPU int32 " + "total_sequence_lengths for graph replay"); + } + if (draft) { + if (input.target_hidden_states.value()->shape() != graph_input.target_hidden_states.value()->shape() + || input.target_hidden_states.value()->dtype() != graph_input.target_hidden_states.value()->dtype()) { + return {nullptr, nullptr}; + } + graph_input.target_hidden_states.value()->copy_from(input.target_hidden_states.value()); + } graph_input.input_ids.value()->copy_from(input.input_ids.value()); graph_input.position_ids.value()->copy_from(input.position_ids.value()); graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value()); @@ -213,6 +351,17 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & set_minus_one_device_async(graph_block_tables); graph_block_tables->narrow({{1, 0, block_per_req}})->copy_from(input.block_tables.value()); graph_input.slot_mapping.value()->copy_from(input.slot_mapping.value()); + std::vector verification_lengths; + if (result->second.verification_lengths) { + const auto total = reinterpret_cast(input.total_sequence_lengths.value()->data())[0]; + verification_lengths.resize(tokens); + for (size_t t = 0; t < tokens; ++t) { + verification_lengths[t] = total - static_cast(tokens) + static_cast(t) + 1; + result->second.verification_tables.value()->narrow({{0, t, 1}})->copy_from(graph_block_tables); + } + infinicore::context::memcpyH2D(result->second.verification_lengths.value()->data(), + verification_lengths.data(), tokens * sizeof(int32_t), false); + } const bool graph_has_mamba_indices = graph_input.mamba_init_state_indices.has_value() && graph_input.mamba_final_state_indices.has_value(); const bool input_has_mamba_indices = input.mamba_init_state_indices.has_value() && input.mamba_final_state_indices.has_value(); @@ -234,15 +383,8 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & auto graph = std::get<0>(result->second.compiled); if (graph != nullptr) { - const auto &runtime_seq_lens = input.total_sequence_lengths.value(); - if (runtime_seq_lens->device().getType() - != infinicore::Device::Type::CPU - || runtime_seq_lens->dtype() != infinicore::DataType::I32 - || runtime_seq_lens->shape().size() != 1 - || runtime_seq_lens->shape()[0] != batch_size) { - throw std::runtime_error( - "PagedCompiler expected CPU int32 " - "total_sequence_lengths for graph replay"); + if (result->second.verification_lengths) { + graph->bind_host_int_array(result->second.verification_lengths.value(), verification_lengths.data(), tokens); } graph->bind_host_int_array( graph_input.total_sequence_lengths.value(), @@ -250,7 +392,12 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & runtime_seq_lens->data()), batch_size); } - auto shared_output = std::shared_ptr(new InfinilmModel::Output{std::get<1>(result->second.compiled)->logits->resume_from_blob_()}); + auto saved_output = std::get<1>(result->second.compiled); + auto shared_output = std::make_shared(); + shared_output->logits = saved_output->logits->resume_from_blob_(); + if (saved_output->hidden_states) { + shared_output->hidden_states = saved_output->hidden_states->resume_from_blob_(); + } return std::make_tuple(graph, shared_output); } diff --git a/csrc/engine/compiler/paged_compiler.hpp b/csrc/engine/compiler/paged_compiler.hpp index a1125864d..7a78a4315 100644 --- a/csrc/engine/compiler/paged_compiler.hpp +++ b/csrc/engine/compiler/paged_compiler.hpp @@ -21,11 +21,14 @@ class PagedCompiler : public GraphCompiler { struct CompiledResult { InfinilmModel::Input input; Compiled compiled; + std::optional verification_lengths; + std::optional verification_tables; }; std::unordered_map< size_t, // num_requests CompiledResult> compiled_map_decode_; + std::unordered_map compiled_map_draft_; }; } // namespace infinilm::engine diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 422b5df73..3e4b7f36e 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -1,5 +1,6 @@ #include "infer_engine.hpp" #include "../config/config_factory.hpp" +#include "infinicore/ops/distributed/broadcast.hpp" #include "spdlog/spdlog.h" #include #include @@ -186,7 +187,7 @@ std::vector InferEngine::state_dict_keys() { // forward //------------------------------------------------------ infinilm::InfinilmModel::Input -InferEngine::Input::to_model_input(infinicore::Device device) const { +InferEngine::Input::to_model_input(infinicore::Device device, bool for_graph) const { auto to_device = [&](const std::optional &t) -> std::optional { @@ -212,8 +213,88 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { const size_t max_query_length = is_prefill ? max_length_from_offsets(input_offsets, "input_offsets") : 0; const size_t max_sequence_length = is_prefill ? max_length_from_offsets(cu_seqlens, "cu_seqlens") : 0; + if (token_state_indices.has_value()) { + auto host_indices = [](const std::optional &value, size_t count) { + if (!value || value.value()->device().getType() != infinicore::Device::Type::CPU + || value.value()->dtype() != infinicore::DataType::I32 + || value.value()->shape() != std::vector{count} + || !value.value()->is_contiguous()) { + throw std::runtime_error("Token state checkpoints require contiguous CPU int32 indices."); + } + return reinterpret_cast(value.value()->data()); + }; + if (!input_ids || input_ids.value()->ndim() != 2 || input_ids.value()->size(0) != 1 + || !input_offsets || input_offsets.value()->numel() < 2 || target_hidden_states) { + throw std::runtime_error("Token checkpoints require packed target-model requests."); + } + const size_t tokens = input_ids.value()->size(1); + const size_t requests = input_offsets.value()->numel() - 1; + const auto *destinations = host_indices(token_state_indices, tokens); + const auto *initial = host_indices(mamba_init_state_indices, requests); + const auto *final = host_indices(mamba_final_state_indices, requests); + const auto *offsets = host_indices(input_offsets, requests + 1); + if (offsets[0] != 0 || offsets[requests] != static_cast(tokens)) { + throw std::runtime_error("Token checkpoint offsets must cover all packed tokens."); + } + std::unordered_set used; + for (size_t r = 0; r < requests; ++r) { + if (offsets[r] < 0 || offsets[r + 1] <= offsets[r] + || static_cast(offsets[r + 1]) > tokens) { + throw std::runtime_error("Token checkpoint offsets must be increasing and within the packed input."); + } + const auto length = offsets[r + 1] - offsets[r]; + if (length > 8 || destinations[offsets[r + 1] - 1] != final[r]) { + throw std::runtime_error("Each checkpoint request needs 1..8 tokens and a matching final row."); + } + if (initial[r] != 0 && !used.insert(initial[r]).second) { + throw std::runtime_error("Checkpoint requests must own distinct initial state rows."); + } + used.insert(initial[r]); + } + for (size_t t = 0; t < tokens; ++t) { + if (destinations[t] <= 0 || !used.insert(destinations[t]).second) { + throw std::runtime_error("Token checkpoints must use distinct nonzero destination rows."); + } + } + const auto &context = global_state::get_forward_context(); + for (const auto *states : {&context.conv_state_vec, &context.ssm_state_vec}) { + for (const auto &state : *states) { + if (state) { + for (auto index : used) { + if (index < 0 || static_cast(index) >= state->size(0)) { + throw std::runtime_error("Token checkpoint exceeds the allocated state pool."); + } + } + } + } + } + } + + const auto transfer_device = for_graph ? global_state::get_tensor_model_parallel_rank_info().device : device; + auto distribute = [&](const std::optional &value, int source_rank) { + if (!value || source_rank < 0 || transfer_device.getType() == infinicore::Device::Type::CPU) { + return value; + } + const auto &rank_info = global_state::get_tensor_model_parallel_rank_info(); + if (rank_info.tp_size == 1) { + // Same-device inputs also need compact storage for model kernels. + return std::optional{value.value()->contiguous()}; + } + const auto &source = value.value(); + auto local = rank_info.tp_rank == source_rank + ? source->contiguous() + : infinicore::Tensor::empty(source->shape(), source->dtype(), transfer_device); + // `Tensor::to` does not transfer between distinct GPUs. Use the existing + // TP communicator for both hidden states and device-resident candidates. + infinicore::op::distributed::broadcast_(local, local, source_rank, rank_info.comm); + return std::optional{local}; + }; + auto local_target_hidden = distribute(target_hidden_states, target_hidden_source_rank); + auto local_ids = distribute(input_ids, input_source_rank); + // MACA maps a registered user pointer to only one node. Serialize H2D // copies so TP ranks never access the same host registration concurrently. + // Collectives above must stay outside this lock so all ranks can enter. static std::mutex maca_host_copy_mutex; const bool serialize_host_copy = device.getType() == infinicore::Device::Type::METAX; @@ -223,7 +304,7 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { } infinilm::InfinilmModel::Input input = { - to_device(input_ids), // @todo: on device in the future + for_graph ? local_ids : to_device(local_ids), to_device(position_ids), to_device(past_sequence_lengths), // @todo: on device in the future to_device(total_sequence_lengths), @@ -239,8 +320,10 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { to_device_vec(image_grid_thw), image_req_ids, visual_token_ranges, - to_device(target_hidden_states), - sample_all_positions}; + for_graph ? local_target_hidden : to_device(local_target_hidden), + sample_all_positions, + to_device(token_state_indices), + top_k == 1 && !return_logits}; if (serialize_host_copy) { infinicore::context::syncStream(); @@ -256,10 +339,42 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { max_query_length, max_sequence_length}; + // The single-request fast path expands only the attention query rows. + // Each checkpoint request has at most eight tokens. GDN/Conv keep + // their original request offsets and recurrent state indices. + const bool short_draft = target_hidden_states && input_offsets + && input_offsets.value()->numel() == 2 && max_query_length <= 8; + if ((token_state_indices || short_draft) && input_offsets.value()->numel() == 2 && is_prefill && input.block_tables + && device.getType() == infinicore::Device::Type::NVIDIA + && max_sequence_length > max_query_length) { + std::vector lengths(max_query_length); + for (size_t i = 0; i < lengths.size(); ++i) { + lengths[i] = static_cast(max_sequence_length - max_query_length + i + 1); + } + auto &metadata = global_state::get_forward_context().attn_metadata; + metadata.verification_sequence_lengths = infinicore::Tensor::empty( + {lengths.size()}, infinicore::DataType::I32, device); + infinicore::context::memcpyH2D(metadata.verification_sequence_lengths.value()->data(), + lengths.data(), lengths.size() * sizeof(int32_t), false); + // Each query references the same physical KV pages. Its own length + // hides the speculative future. NVIDIA Decode kernels assume packed + // page-table rows, so materialize only these small indices, never KV. + metadata.verification_block_tables = input.block_tables.value()->as_strided( + {max_query_length, input.block_tables.value()->size(1)}, + {0, input.block_tables.value()->stride(1)}) + ->contiguous(); + } + infinilm::global_state::get_forward_context().mamba_metadata = { input.input_offsets, input.mamba_init_state_indices, - input.mamba_final_state_indices}; + input.mamba_final_state_indices, + input.token_state_indices}; + if (token_state_indices) { + const auto *offsets = reinterpret_cast(input_offsets.value()->data()); + global_state::get_forward_context().mamba_metadata.checkpoint_offsets.assign( + offsets, offsets + input_offsets.value()->numel()); + } global_state::get_forward_context().mm_metadata = { image_req_ids, @@ -269,9 +384,37 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { } InferEngine::Output InferEngine::forward(const InferEngine::Input &input) { + auto local_input = input; + auto source_rank = [&](const std::optional &value) { + if (!value || value.value()->device().getType() == infinicore::Device::Type::CPU) { + return -1; + } + for (int rank = 0; rank < communication_group_.get_world_size(); ++rank) { + if (communication_group_.get_rank_info(rank).device == value.value()->device()) { + return rank; + } + } + throw std::invalid_argument("Device inputs must belong to the target TP group."); + }; + local_input.target_hidden_source_rank = source_rank(input.target_hidden_states); + local_input.input_source_rank = get_dist_config().pp_size == 1 ? source_rank(input.input_ids) : -1; + if (input.verify_draft) { + const bool valid_shape = input.input_ids && input.input_ids.value()->ndim() == 2 + && input.input_ids.value()->size(0) == 1 + && input.input_ids.value()->size(1) >= 2 + && input.input_ids.value()->size(1) <= 5 + && input.input_offsets && input.input_offsets.value()->numel() == 2; + if (input.top_k != 1 || !input.sample_all_positions || !input.token_state_indices + || !valid_shape || get_dist_config().pp_size != 1 || input.return_device_tokens) { + throw std::invalid_argument("Device MTP acceptance requires one greedy Q=2..5 request, checkpoints, PP1 and host results."); + } + } + if (input.return_device_tokens && get_dist_config().pp_size != 1) { + throw std::invalid_argument("Device token output currently requires PP1."); + } // Trigger each worker to run inference for (auto &worker : workers_) { - worker->run(input); + worker->run(local_input); } // Wait for all workers for (auto &worker : workers_) { @@ -322,6 +465,15 @@ void InferEngine::reset_cache(const cache::CacheConfig *new_config) { this->compile(); } +std::vector>> InferEngine::get_hybrid_states() { + std::vector>> result; + for (auto &worker : workers_) { + worker->wait(); + result.push_back(worker->get_hybrid_states()); + } + return result; +} + std::vector> InferEngine::get_kv_cache() { std::vector> kv_cache_list; if (workers_.empty()) { diff --git a/csrc/engine/infer_engine.hpp b/csrc/engine/infer_engine.hpp index 31054149d..e6e2ea066 100644 --- a/csrc/engine/infer_engine.hpp +++ b/csrc/engine/infer_engine.hpp @@ -57,6 +57,7 @@ class InferEngine { void reset_cache(const cache::CacheConfig *new_config); std::vector> get_kv_cache(); + std::vector>> get_hybrid_states(); ~InferEngine(); diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 0fa0a84cf..5ba433fbf 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -2,6 +2,10 @@ #include "../models/model_factory.hpp" #include "infinicore/ops.hpp" #include "infinicore/ops/distributed/send_recv.hpp" +#include +#include +#include +#include #include #include @@ -249,6 +253,15 @@ std::vector RankWorker::get_kv_cache() { //------------------------------------------------------ // close -- request shutdown and join thread //------------------------------------------------------ +std::vector> RankWorker::get_hybrid_states() { + std::unique_lock lk(mutex_); + cv_.wait(lk, [&] { return init_done_ || should_exit_; }); + if (should_exit_ || has_job_) { + throw std::runtime_error("State access requires an idle worker."); + } + return {forward_context_.conv_state_vec, forward_context_.ssm_state_vec}; +} + void RankWorker::close() { { std::lock_guard lock(mutex_); @@ -418,21 +431,34 @@ void RankWorker::thread_loop() { infinicore::Tensor logits; infinicore::Tensor hidden_states; - // All-position speculative/MTP runs need eager mode because - // hidden states are not part of compiled graph outputs. - 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::cpu())); + infinicore::Tensor sampled_ids; + infinicore::Tensor model_input_ids; + if (local_args.token_state_indices.has_value() + && !model_->supports_token_state_checkpoints()) { + throw std::runtime_error("This model does not support per-token state checkpoints."); + } + const bool graph_candidate = !local_args.token_state_indices + && local_args.input_ids && local_args.input_offsets + && (local_args.input_ids.value()->numel() == local_args.input_offsets.value()->numel() - 1 + || (local_args.target_hidden_states && local_args.input_ids.value()->numel() <= 2)); + if (graph_candidate && compiler_ != nullptr && rank_info_.pp_size == 1) { + auto graph_input = local_args.to_model_input(infinicore::Device::cpu(), true); + auto [graph, output] = compiler_->get_compiled(graph_input); if (graph != nullptr && output != nullptr) { graph->run(); logits = output->logits; + hidden_states = output->hidden_states; + model_input_ids = graph_input.input_ids.value(); } } // Fall back to eager mode if (!logits) { auto model_args = local_args.to_model_input(rank_info_.device); + model_input_ids = model_args.input_ids.value(); auto model_output = model_->forward(model_args); logits = model_output.logits; hidden_states = model_output.hidden_states; + sampled_ids = model_output.output_ids; } if (rank_info_.pp_size > 1 && rank_info_.pp_stage + 1 != rank_info_.pp_size) { @@ -466,34 +492,37 @@ void RankWorker::thread_loop() { // Random sampling (rank 0 only) if (rank_info_.tp_rank == 0) { - auto temperature{local_args.temperature}; - auto top_p{local_args.top_p}; - auto top_k{local_args.top_k}; - - const auto &logits_shape{logits->shape()}; - const auto &vocab_size{logits_shape[2]}; - const auto &total_len{logits_shape[1]}; - const auto &batch_size{logits_shape[0]}; - - auto n_req = local_args.input_offsets.value()->size(0) - 1; - int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data(); - - const bool sample_all_positions = local_args.sample_all_positions; - const size_t logits_positions = batch_size * total_len; - const bool logits_are_last_token_only = !sample_all_positions && logits_positions == n_req; - const size_t n_out = sample_all_positions ? static_cast(input_offsets[n_req]) : n_req; - auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)}; - - for (size_t i{0}; i < n_out; ++i) { - size_t score_idx = i; - if (!sample_all_positions && !logits_are_last_token_only) { - score_idx = static_cast(input_offsets[i + 1] - 1); + auto output_ids = sampled_ids; + if (!output_ids) { + auto temperature{local_args.temperature}; + auto top_p{local_args.top_p}; + auto top_k{local_args.top_k}; + + const auto &logits_shape{logits->shape()}; + const auto &vocab_size{logits_shape[2]}; + const auto &total_len{logits_shape[1]}; + const auto &batch_size{logits_shape[0]}; + + auto n_req = local_args.input_offsets.value()->size(0) - 1; + int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data(); + + const bool sample_all_positions = local_args.sample_all_positions; + const size_t logits_positions = batch_size * total_len; + const bool logits_are_last_token_only = !sample_all_positions && logits_positions == n_req; + const size_t n_out = sample_all_positions ? static_cast(input_offsets[n_req]) : n_req; + output_ids = infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device); + + for (size_t i{0}; i < n_out; ++i) { + size_t score_idx = i; + if (!sample_all_positions && !logits_are_last_token_only) { + score_idx = static_cast(input_offsets[i + 1] - 1); + } + auto score{logits->view({logits_positions, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; + auto out{output_ids->narrow({{0, i, 1}})->view({})}; + float random_val = std::uniform_real_distribution(0, 1)(rng_); + infinicore::op::random_sample_( + out, score, random_val, top_p, top_k, temperature); } - auto score{logits->view({logits_positions, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; - auto out{output_ids->narrow({{0, i, 1}})->view({})}; - float random_val = std::uniform_real_distribution(0, 1)(rng_); - infinicore::op::random_sample_( - out, score, random_val, top_p, top_k, temperature); } if (rank_info_.pp_size > 1) { @@ -503,11 +532,41 @@ void RankWorker::thread_loop() { rank_info_.world_comm); } - output_ids = output_ids->to(infinicore::Device::cpu()); - - infinicore::context::syncStream(); + int accepted_draft_tokens = -1; + if (local_args.verify_draft) { + const size_t count = model_input_ids->size(1) - 1; + auto candidates = model_input_ids->narrow({{1, 1, count}})->view({count})->to(rank_info_.device); + auto accepted = infinicore::op::equal(output_ids->narrow({{0, 0, count}}), candidates); + auto packed = infinicore::Tensor::empty({count + 2}, infinicore::DataType::I64, rank_info_.device); + packed->narrow({{0, 0, count + 1}})->copy_from(output_ids); + auto length = packed->narrow({{0, count + 1, 1}}); + if (count == 1) { + infinicore::op::cast_(length, accepted); + } else { + // Sum consecutive prefix matches, stopping at the + // first rejection. F32 represents these 0..4 counts exactly. + auto matches = infinicore::Tensor::empty({count}, infinicore::DataType::F32, rank_info_.device); + infinicore::op::cast_(matches, accepted); + auto prefix = matches->narrow({{0, 0, 1}}); + auto total = prefix; + for (size_t i = 1; i < count; ++i) { + prefix = infinicore::op::mul(prefix, matches->narrow({{0, i, 1}})); + total = infinicore::op::add(total, prefix); + } + infinicore::op::cast_(length, total); + } + // One bounded host transfer contains the tokens and + // acceptance length. Scheduler ownership stays on CPU. + infinicore::context::syncStream(); + auto host = packed->to(infinicore::Device::cpu()); + accepted_draft_tokens = static_cast(reinterpret_cast(host->data())[count + 1]); + output_ids = host->narrow({{0, 0, static_cast(1 + accepted_draft_tokens)}}); + } else if (!local_args.return_device_tokens) { + infinicore::context::syncStream(); + output_ids = output_ids->to(infinicore::Device::cpu()); + } - auto out{Output{output_ids, logits, hidden_states}}; + auto out{Output{output_ids, logits, hidden_states, accepted_draft_tokens}}; output_ = std::move(out); } @@ -548,7 +607,9 @@ void RankWorker::thread_loop() { } else if (local_cmd == Command::COMPILE) { try { if (compiler_ != nullptr) { + spdlog::info("Graph capture begin: tp_rank={}", rank_info_.tp_rank); compiler_->compile(); + spdlog::info("Graph capture end: tp_rank={}", rank_info_.tp_rank); } { std::lock_guard lk(mutex_); diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index d396ef6f1..a9d3e5f6e 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -72,6 +72,15 @@ class RankWorker { std::optional target_hidden_states; /// Sample logits at every packed input position instead of one token per request. bool sample_all_positions{false}; + /// Per-token state-pool destinations for speculative verification. + std::optional token_state_indices; + // Resolved by `InferEngine` for device-resident draft hidden states. + // CPU inputs use ordinary per-rank H2D copies instead. + int target_hidden_source_rank{-1}; + bool return_logits{true}; + bool return_device_tokens{false}; + bool verify_draft{false}; + int input_source_rank{-1}; float temperature{1}; @@ -79,13 +88,15 @@ class RankWorker { float top_p{1}; - infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const; + infinilm::InfinilmModel::Input to_model_input(infinicore::Device device, bool for_graph = false) const; }; struct Output { infinicore::Tensor output_ids; infinicore::Tensor logits; infinicore::Tensor hidden_states; + // -1 for ordinary calls; greedy verification returns the accepted prefix length. + int accepted_draft_tokens{-1}; }; RankWorker(std::shared_ptr infinilm_config, @@ -119,6 +130,7 @@ class RankWorker { void reset_cache(const cache::CacheConfig *new_config); std::vector get_kv_cache(); + std::vector> get_hybrid_states(); // Compile the model graph if enabled. void compile(); diff --git a/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index 6d02de4c5..4eeac6302 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -21,6 +21,10 @@ struct AttentionMetadata { size_t max_query_length{0}; /// Maximum total sequence length in the current batch. size_t max_sequence_length{0}; + // Single-request checkpoint verification can reuse Decode Attention with + // one causal KV length per query. These buffers are shared by all layers. + std::optional verification_sequence_lengths; + std::optional verification_block_tables; AttentionMetadata() = default; @@ -61,6 +65,9 @@ struct MambaMetadata { std::optional init_state_indices; /// State cache indices written with the final state of each request forward. std::optional final_state_indices; + /// Optional destination after each packed token. + std::optional token_state_indices; + std::vector checkpoint_offsets; }; struct ForwardContext { diff --git a/csrc/layers/attention/backends/paged_attn.cpp b/csrc/layers/attention/backends/paged_attn.cpp index f39937ead..56a0bdaed 100644 --- a/csrc/layers/attention/backends/paged_attn.cpp +++ b/csrc/layers/attention/backends/paged_attn.cpp @@ -35,6 +35,11 @@ infinicore::Tensor PagedAttentionImpl::forward(const AttentionLayer &layer, size_t seq_len = query->shape()[0]; bool is_prefill = (seq_len != total_sequence_lengths.value()->shape()[0]); + if (attn_metadata.verification_sequence_lengths.has_value()) { + total_sequence_lengths = attn_metadata.verification_sequence_lengths; + block_tables = attn_metadata.verification_block_tables; + is_prefill = false; + } // 2. Compute attention const size_t value_head_dim = value->size(value->ndim() - 1); diff --git a/csrc/layers/linear/vocab_parallel.cpp b/csrc/layers/linear/vocab_parallel.cpp new file mode 100644 index 000000000..82aba084c --- /dev/null +++ b/csrc/layers/linear/vocab_parallel.cpp @@ -0,0 +1,68 @@ +#include "vocab_parallel.hpp" + +#include +#include +#include + +namespace infinilm::nn { + +VocabParallelLinear::VocabParallelLinear(size_t hidden_size, size_t vocab_size, + const infinicore::DataType &dtype, const infinicore::Device &device, + size_t tp_rank, size_t tp_size, infinicclComm_t communicator) + : ColumnParallelLinear(hidden_size, vocab_size, false, dtype, device, + vocab_size % tp_size == 0 ? tp_rank : 0, + vocab_size % tp_size == 0 ? tp_size : 1), + communicator_(communicator) { + auto start = infinicore::Tensor::empty({1}, infinicore::DataType::I64, infinicore::Device::cpu()); + *reinterpret_cast(start->data()) = static_cast(tp_rank_ * (vocab_size / tp_size_)); + vocab_start_ = start->to(device); +} + +infinicore::Tensor VocabParallelLinear::forward(infinicore::Tensor &input) const { + auto local = ColumnParallelLinear::forward(input); + if (tp_size_ == 1) { + return local; + } + const auto rows = local->numel() / local->size(local->ndim() - 1); + auto gathered = infinicore::op::distributed::allgather(local->view({1, rows, out_features_ / tp_size_}), tp_size_, communicator_); + auto shape = local->shape(); + shape.back() = out_features_; + return gathered->permute({1, 0, 2})->contiguous()->view(shape); +} + +infinicore::Tensor VocabParallelLinear::top_tokens(infinicore::Tensor &input) const { + auto logits = ColumnParallelLinear::forward(input); + const auto vocab = logits->size(logits->ndim() - 1); + const auto rows = logits->numel() / vocab; + logits = logits->view({rows, vocab}); + auto ids = infinicore::Tensor::empty({rows}, infinicore::DataType::I64, input->device()); + auto scores = tp_size_ > 1 + ? infinicore::Tensor::empty({rows}, logits->dtype(), input->device()) + : infinicore::Tensor{}; + for (size_t i = 0; i < rows; ++i) { + auto row = logits->narrow({{0, i, 1}})->view({vocab}); + auto id = ids->narrow({{0, i, 1}}); + // The existing greedy sampler chooses the lowest index on ties. + infinicore::op::random_sample_(id->view({}), row, 0.0f, 1.0f, 1, 1.0f); + if (tp_size_ > 1) { + infinicore::op::take_(scores->narrow({{0, i, 1}}), row, id); + } + } + if (tp_size_ == 1) { + return ids; + } + ids = infinicore::op::add(ids, vocab_start_->as_strided({rows}, {0})->contiguous()); + auto rank_ids = infinicore::op::distributed::allgather(ids->view({1, rows}), tp_size_, communicator_)->permute({1, 0})->contiguous(); + auto rank_scores = infinicore::op::distributed::allgather(scores->view({1, rows}), tp_size_, communicator_)->permute({1, 0})->contiguous(); + auto output = infinicore::Tensor::empty({rows}, infinicore::DataType::I64, input->device()); + auto winner = infinicore::Tensor::empty({1}, infinicore::DataType::I64, input->device()); + for (size_t i = 0; i < rows; ++i) { + auto row = rank_scores->narrow({{0, i, 1}})->view({tp_size_}); + infinicore::op::random_sample_(winner->view({}), row, 0.0f, 1.0f, 1, 1.0f); + // Rank order matches increasing global token ID, preserving tie rules. + infinicore::op::take_(output->narrow({{0, i, 1}}), rank_ids->narrow({{0, i, 1}}), winner); + } + return output; +} + +} // namespace infinilm::nn diff --git a/csrc/layers/linear/vocab_parallel.hpp b/csrc/layers/linear/vocab_parallel.hpp new file mode 100644 index 000000000..4ba857c3e --- /dev/null +++ b/csrc/layers/linear/vocab_parallel.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include "linear.hpp" + +namespace infinilm::nn { + +// Shares the usual `lm_head.weight` loader and column partitioning. Uneven +// vocabularies remain replicated until the parameter loader supports padding. +class VocabParallelLinear : public ColumnParallelLinear { +public: + VocabParallelLinear(size_t hidden_size, size_t vocab_size, + const infinicore::DataType &dtype, const infinicore::Device &device, + size_t tp_rank, size_t tp_size, infinicclComm_t communicator); + + infinicore::Tensor forward(infinicore::Tensor &input) const; + infinicore::Tensor top_tokens(infinicore::Tensor &input) const; + +private: + infinicclComm_t communicator_; + infinicore::Tensor vocab_start_; +}; + +} // namespace infinilm::nn diff --git a/csrc/layers/quantization/fp8_block.cpp b/csrc/layers/quantization/fp8_block.cpp new file mode 100644 index 000000000..fcfb26d84 --- /dev/null +++ b/csrc/layers/quantization/fp8_block.cpp @@ -0,0 +1,172 @@ +#include "fp8_block.hpp" +#include "marlin_support.hpp" + +#include "infinicore/context/context.hpp" +#include "infinicore/ops/add.hpp" +#include "infinicore/ops/cast.hpp" +#include "infinicore/ops/mul.hpp" +#include + +#if INFINILM_ENABLE_MARLIN && __has_include("infinicore/ops/awq_marlin_gemm.hpp") +#define INFINILM_ENABLE_FP8_MARLIN 1 +#include "infinicore/ops/awq_marlin_gemm.hpp" +#include "marlin_utils.hpp" +#else +#define INFINILM_ENABLE_FP8_MARLIN 0 +#endif + +namespace infinilm::quantization { + +FP8Block::FP8Block(const nlohmann::json &config) : NoneQuantization(config) { + if (get_or("fmt", "") != "e4m3" + || get_or>("weight_block_size", {}) != std::vector{128, 128}) { + throw std::runtime_error("FP8 compatibility path requires E4M3 with 128x128 weight blocks."); + } +} + +std::vector FP8Block::get_param_layout( + size_t in_features, size_t out_features, int split_dim, int tp_rank, + int tp_size, int tp_num_heads, const infinicore::DataType &dtype, bool bias) const { + activation_dtype_ = dtype; + if (in_features % block_size_ || out_features % block_size_ + || (split_dim == 0 && out_features % (block_size_ * tp_size)) + || (split_dim == 1 && in_features % (block_size_ * tp_size))) { + throw std::runtime_error("FP8 weight and TP partitions must align to 128-element blocks."); + } + auto layout = NoneQuantization::get_param_layout( + in_features, out_features, split_dim, tp_rank, tp_size, tp_num_heads, dtype, bias); + layout.front().dtype = infinicore::DataType::F8; + layout.push_back({"weight_scale_inv", {out_features / block_size_, in_features / block_size_}, infinicore::DataType::F32, split_dim, tp_rank, tp_size}); + return layout; +} + +std::vector FP8Block::split_params( + const std::unordered_map ¶ms, + const std::vector &splits, int narrow_dim, + int tp_rank, int tp_size, int tp_num_heads) const { + auto result = NoneQuantization::split_params(params, splits, narrow_dim, tp_rank, tp_size, tp_num_heads); + for (const auto &s : splits) { + if (s.start % block_size_ || s.size % block_size_) { + throw std::runtime_error("FP8 fused projection split must align to weight blocks."); + } + result.push_back({s.prefix + ".weight_scale_inv", + infinicore::nn::Parameter( + params.at("weight_scale_inv")->narrow({{static_cast(narrow_dim), s.start / block_size_, s.size / block_size_}}), + narrow_dim, tp_rank, tp_size, s.num_shards)}); + } + return result; +} + +infinicore::Tensor FP8Block::forward( + const ParamsMap ¶ms, const infinicore::Tensor &input, bool has_bias, float alpha) const { + auto weight = params.at("weight"); + auto shape = weight->shape(); + auto unpacked = infinicore::Tensor::empty(shape, infinicore::DataType::F32, weight->device()); + infinicore::op::cast_(unpacked, weight); + auto tiles = unpacked->view({shape[0] / block_size_, block_size_, shape[1] / block_size_, block_size_}); + auto block_scales = params.at("weight_scale_inv"); + auto scales = block_scales->as_strided(tiles->shape(), {block_scales->stride(0), 0, block_scales->stride(1), 0}); + // Despite its checkpoint name, `weight_scale_inv` multiplies the FP8 values. + infinicore::op::mul_(tiles, tiles, scales); + auto dequantized = infinicore::Tensor::empty(shape, input->dtype(), weight->device()); + infinicore::op::cast_(dequantized, unpacked); + auto dense_params = params; + dense_params["weight"] = dequantized; + return NoneQuantization::forward(dense_params, input, has_bias, alpha); +} + +infinicore::Tensor FP8Block::forward_allreduce( + const ParamsMap ¶ms, const infinicore::Tensor &input, + bool has_bias, infinicclComm_t communicator, float alpha) const { + return BaseQuantization::forward_allreduce(params, input, has_bias, communicator, alpha); +} + +#if INFINILM_ENABLE_FP8_MARLIN +namespace { +// InfiniCore's native Marlin operator supports FP8 as well as integer weights. +constexpr int64_t FP8_E4M3FN_ID = 2814749767172868LL; + +class FP8Marlin final : public FP8Block { +public: + FP8Marlin(const nlohmann::json &config, size_t n) : FP8Block(config), n_(n) {} + infinicore::Tensor forward(const ParamsMap ¶ms, const infinicore::Tensor &input, + bool bias, float alpha) const override { + if (alpha != 1.0f) { + throw std::runtime_error("FP8 Marlin currently requires linear `alpha=1`."); + } + auto contiguous = input->is_contiguous() ? input : input->contiguous(); + auto shape = input->shape(); + const size_t k = shape.back(), m = input->numel() / k; + auto output = infinicore::Tensor::empty({m, n_}, input->dtype(), input->device()); + auto weight = params.at("qweight"), scales = params.at("scales"), empty = params.at("empty"); + infinicore::op::awq_marlin_gemm_(output, contiguous->view({m, k}), weight, + empty, scales, empty, empty, empty, empty, empty, + FP8_E4M3FN_ID, true, false, true, false); + if (bias) { + infinicore::op::add_(output, output, params.at("bias")->as_strided({m, n_}, {0, 1})); + } + shape.back() = n_; + return output->view(shape); + } + std::shared_ptr process_weights_after_loading( + ParamsMap &, const infinicore::Device &, int) const override { return nullptr; } + std::vector split_params( + const std::unordered_map &, + const std::vector &, int, int, int, int) const override { return {}; } + +private: + size_t n_; +}; +} // namespace +#endif + +std::shared_ptr FP8Block::process_weights_after_loading( + ParamsMap ¶ms, const infinicore::Device &device, int) const { + const auto backend = get_or("fp8_backend", "compatibility"); + if (backend == "compatibility") { + return nullptr; + } + if (backend != "marlin" || device.getType() != infinicore::Device::Type::NVIDIA) { + throw std::runtime_error("FP8 backend must be `compatibility` or NVIDIA `marlin`."); + } +#if INFINILM_ENABLE_FP8_MARLIN + if (activation_dtype_ != infinicore::DataType::BF16 && activation_dtype_ != infinicore::DataType::F16) { + throw std::runtime_error("FP8 Marlin requires BF16 or FP16 activations."); + } + auto weight = params.at("weight"); + const size_t n = weight->size(0), k = weight->size(1); + // Pack bytes without requantization. Reuse the existing GPU Marlin repacker. + auto packed = infinicore::Tensor::from_blob(weight->data(), {n, k / 4}, infinicore::DataType::I32, device) + ->permute({1, 0}) + ->contiguous(); + auto empty = marlin::make_empty_i32(device); + auto repacked = marlin::gptq_marlin_repack(packed, empty, k, n, 8); + // Reuse original storage, also held by fused-projection checkpoint aliases. + // Keeping `weight` owns this memory; `qweight` is its packed runtime view. + auto qweight = infinicore::Tensor::from_blob(weight->data(), repacked->shape(), infinicore::DataType::I32, device); + qweight->copy_from(repacked); + + auto cpu_scales = params.at("weight_scale_inv")->to(infinicore::Device::cpu())->contiguous(); + infinicore::context::syncStream(); + const auto *src = reinterpret_cast(cpu_scales->data()); + std::vector expanded(k / block_size_ * n); + const float exponent_bias = std::ldexp(1.0f, activation_dtype_ == infinicore::DataType::BF16 ? 120 : 8); + for (size_t group = 0; group < k / block_size_; ++group) { + for (size_t row = 0; row < n; ++row) { + expanded[group * n + row] = src[(row / block_size_) * (k / block_size_) + group] * exponent_bias; + } + } + auto scale_f32 = infinicore::Tensor::from_blob(expanded.data(), {k / block_size_, n}, infinicore::DataType::F32, infinicore::Device::cpu())->to(device); + auto scales = infinicore::Tensor::empty(scale_f32->shape(), activation_dtype_, device); + infinicore::op::cast_(scales, scale_f32); + params["qweight"] = qweight; + params["scales"] = marlin::permute_scales(scales, k, n, block_size_); + params["empty"] = empty; + infinicore::context::syncStream(); + return std::make_shared(get_config(), n); +#else + throw std::runtime_error("FP8 Marlin requires an InfiniCore build with Marlin support."); +#endif +} + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/fp8_block.hpp b/csrc/layers/quantization/fp8_block.hpp new file mode 100644 index 000000000..ec17ddbe0 --- /dev/null +++ b/csrc/layers/quantization/fp8_block.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include "none_quantization.hpp" + +namespace infinilm::quantization { + +// Block-scaled E4M3 checkpoints with a portable W8A16 compatibility path and +// opt-in NVIDIA Marlin packing during the existing post-load lifecycle. +class FP8Block : public NoneQuantization { +public: + explicit FP8Block(const nlohmann::json &config); + + QuantScheme get_quant_scheme() const override { return QuantScheme::FP8_BLOCK_W8A16; } + std::vector get_param_layout( + size_t in_features, size_t out_features, int split_dim, int tp_rank, + int tp_size, int tp_num_heads, const infinicore::DataType &dtype, bool bias) const override; + std::vector split_params( + const std::unordered_map ¶ms, + const std::vector &splits, int narrow_dim, + int tp_rank, int tp_size, int tp_num_heads) const override; + infinicore::Tensor forward(const ParamsMap ¶ms, const infinicore::Tensor &input, + bool has_bias, float alpha = 1.0f) const override; + infinicore::Tensor forward_allreduce(const ParamsMap ¶ms, const infinicore::Tensor &input, + bool has_bias, infinicclComm_t communicator, + float alpha = 1.0f) const override; + std::shared_ptr process_weights_after_loading( + ParamsMap &, const infinicore::Device &, int = -1) const override; + +private: + static constexpr size_t block_size_ = 128; + mutable infinicore::DataType activation_dtype_ = infinicore::DataType::BF16; +}; + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/none_quantization.cpp b/csrc/layers/quantization/none_quantization.cpp index 3184da991..b7d5ccc81 100644 --- a/csrc/layers/quantization/none_quantization.cpp +++ b/csrc/layers/quantization/none_quantization.cpp @@ -82,12 +82,13 @@ std::vector NoneQuantization::split_params( std::vector result; auto weight_it = params.find("weight"); auto bias_it = params.find("bias"); + const int weight_dim = weight_prepacked_ ? 1 - narrow_dim : narrow_dim; for (const auto &s : splits) { result.push_back({s.prefix + ".weight", infinicore::nn::Parameter( - weight_it->second->narrow({{static_cast(narrow_dim), s.start, s.size}}), - narrow_dim, tp_rank, tp_size, s.num_shards)}); + weight_it->second->narrow({{static_cast(weight_dim), s.start, s.size}}), + weight_dim, tp_rank, tp_size, s.num_shards)}); if (bias_it != params.end()) { result.push_back({s.prefix + ".bias", infinicore::nn::Parameter( @@ -104,7 +105,7 @@ std::shared_ptr NoneQuantization::process_weights_after_loadin int /*split_dim*/) const { // Controlled by --pre-transpose CLI flag, default off. - if (!global_state::get_infinilm_config().pre_transpose) { + if (!global_state::get_infinilm_config().pre_transpose || weight_prepacked_) { return nullptr; } @@ -115,15 +116,13 @@ std::shared_ptr NoneQuantization::process_weights_after_loadin // subsequent forwards can feed it directly to GEMM. params["weight"] = weight_it->second->permute({1, 0})->contiguous(); - // Mark as pre-packed so forward() uses linear_packed. - weight_prepacked_ = true; + // A quantization object may be shared by several unprocessed linears. + auto packed = std::make_shared(get_config()); + packed->weight_prepacked_ = true; + return packed; } - // Must return non-null so that BaseLinear::process_weights_after_loading - // writes the modified params back into parameters_. - // Returning shared_from_this() triggers the "quantization changed" path - // which calls parameters_.clear() + re-insert from params. - return std::const_pointer_cast(shared_from_this()); + return nullptr; } } // namespace infinilm::quantization diff --git a/csrc/layers/quantization/none_quantization.hpp b/csrc/layers/quantization/none_quantization.hpp index 108e87d7e..69bb3a57a 100644 --- a/csrc/layers/quantization/none_quantization.hpp +++ b/csrc/layers/quantization/none_quantization.hpp @@ -6,7 +6,7 @@ namespace infinilm::quantization { class NoneQuantization : public BaseQuantization { public: explicit NoneQuantization(const nlohmann::json &quant_config) - : BaseQuantization(quant_config){}; + : BaseQuantization(quant_config) {}; NoneQuantization(); @@ -40,15 +40,14 @@ class NoneQuantization : public BaseQuantization { int narrow_dim, int tp_rank, int tp_size, int tp_num_heads) const override; - // Ascend: pre-pack weight to [IC, OC] after loading to skip runtime permute. - // Returns shared_from_this() only on Ascend; nullptr otherwise (no-op). + // Pre-pack to [IC, OC] when enabled and return a per-linear layout state. std::shared_ptr process_weights_after_loading( ParamsMap ¶ms, const infinicore::Device &device, int split_dim = -1) const override; private: - mutable bool weight_prepacked_ = false; // true when weight was pre-packed for Ascend + bool weight_prepacked_ = false; }; } // namespace infinilm::quantization diff --git a/csrc/layers/quantization/quantization.hpp b/csrc/layers/quantization/quantization.hpp index 0cc9cd7e2..74e695d59 100644 --- a/csrc/layers/quantization/quantization.hpp +++ b/csrc/layers/quantization/quantization.hpp @@ -4,6 +4,7 @@ #include "awq_marlin.hpp" #include "base_quantization.hpp" #include "compressed_tensors.hpp" +#include "fp8_block.hpp" #include "gptq.hpp" #include "gptq_marlin.hpp" #include "gptq_qy.hpp" diff --git a/csrc/layers/quantization/quantization_scheme.hpp b/csrc/layers/quantization/quantization_scheme.hpp index 455968a7a..a9b1f4969 100644 --- a/csrc/layers/quantization/quantization_scheme.hpp +++ b/csrc/layers/quantization/quantization_scheme.hpp @@ -11,6 +11,7 @@ enum class QuantScheme { GPTQ_W4A16, GPTQ_MARLIN_W4A16, MXFP4_W4A16, + FP8_BLOCK_W8A16, }; enum class KVQuantAlgo { diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index 02677318d..33f31237d 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -57,6 +57,10 @@ 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}; + /// Optional state-pool destination after each packed token. + std::optional token_state_indices; + /// A greedy caller permits a model to return token IDs without full logits. + bool greedy_output{false}; }; struct Output { @@ -64,10 +68,13 @@ class InfinilmModel : public infinicore::nn::Module { infinicore::Tensor logits; /// Optional final hidden states, used by MTP/Eagle draft models. infinicore::Tensor hidden_states; + /// Optional device token IDs, in sampling order, instead of logits. + infinicore::Tensor output_ids; }; virtual ~InfinilmModel() = default; virtual Output forward(const Input &input) const = 0; + virtual bool supports_token_state_checkpoints() const { return false; } virtual void reset_cache(const cache::CacheConfig *cache_config); virtual const cache::CacheConfig *get_cache_config() const { return cache_config_.get(); diff --git a/csrc/models/qwen3_5/qwen3_5_attention.hpp b/csrc/models/qwen3_5/qwen3_5_attention.hpp index 12baa3dd1..a1418eb0d 100644 --- a/csrc/models/qwen3_5/qwen3_5_attention.hpp +++ b/csrc/models/qwen3_5/qwen3_5_attention.hpp @@ -13,6 +13,14 @@ class Qwen35Attention : public infinicore::nn::Module { infinicore::Tensor forward(const infinicore::Tensor &positions, const infinicore::Tensor &hidden_states) 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_; } diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp index 72fe1a87f..1ff0c3b16 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.cpp @@ -1,6 +1,9 @@ #include "qwen3_5_for_causal_lm.hpp" +#include "../../cache/kv_cache.hpp" +#include "../../global_state/global_state.hpp" #include "../models_registry.hpp" +#include #include #include #include @@ -16,14 +19,45 @@ Qwen35ForCausalLM::Qwen35ForCausalLM( const auto &dtype = model_config->get_dtype(); INFINICORE_NN_MODULE_INIT(model, model_config, device); + const auto &rank = global_state::get_tensor_model_parallel_rank_info(); INFINICORE_NN_MODULE_INIT( - lm_head, hidden_size, vocab_size, false, dtype, device); + lm_head, hidden_size, vocab_size, dtype, device, rank.tp_rank, rank.tp_size, rank.comm); + if (model_config->get_or("enable_mtp", false)) { + INFINICORE_NN_MODULE_INIT(mtp, model_config, device); + } } InfinilmModel::Output Qwen35ForCausalLM::forward( const InfinilmModel::Input &input) const { - auto hidden_states = model_->forward(input); - return {lm_head_->forward(hidden_states)}; + infinicore::Tensor hidden_states; + if (input.target_hidden_states.has_value()) { + if (!mtp_) { + throw std::runtime_error("Qwen MTP weights must be enabled before draft execution."); + } + hidden_states = mtp_->forward(model_->embed_input_ids(input.input_ids.value()), + input.target_hidden_states.value(), input.position_ids.value()); + } else { + hidden_states = model_->forward(input); + } + auto head_input = hidden_states; + const bool packed_greedy = input.greedy_output && input.block_tables && input.input_offsets + && hidden_states->size(0) == 1 + && hidden_states->device().getType() == infinicore::Device::Type::NVIDIA; + if (packed_greedy && !input.sample_all_positions + && hidden_states->size(1) != input.input_offsets.value()->numel() - 1) { + head_input = infinicore::Tensor::empty( + {1, input.input_offsets.value()->numel() - 1, hidden_states->size(2)}, + hidden_states->dtype(), hidden_states->device()); + infinicore::op::select_last_token_hidden_(head_input, hidden_states, input.input_offsets.value()); + } + if (packed_greedy) { + return {{}, mtp_ ? hidden_states : infinicore::Tensor{}, lm_head_->top_tokens(head_input)}; + } + auto logits = lm_head_->forward(head_input); + if (mtp_) { + return {logits, hidden_states}; + } + return {logits}; } void Qwen35ForCausalLM::reset_cache( @@ -34,6 +68,19 @@ void Qwen35ForCausalLM::reset_cache( cache_config_ = cache_config->unique_copy(); } model_->reset_cache(cache_config); + if (mtp_ && cache_config != nullptr) { + const auto *paged = dynamic_cast(cache_config); + if (paged == nullptr) { + throw std::runtime_error("Qwen MTP requires paged KV storage."); + } + const auto head_dim = model_config_->get("head_dim"); + const auto heads = model_config_->get("num_key_value_heads"); + auto &context = global_state::get_forward_context(); + context.kv_cache_vec.push_back(cache::PagedKVCache::create_layer_kv_cache( + head_dim, head_dim, heads, heads, model_config_->get_kv_cache_dtype(), *paged)); + context.conv_state_vec.emplace_back(); + context.ssm_state_vec.emplace_back(); + } } std::shared_ptr prepare_qwen3_5_model_config(std::shared_ptr model_config) { diff --git a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp index 51211481f..5949c9bb2 100644 --- a/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp +++ b/csrc/models/qwen3_5/qwen3_5_for_causal_lm.hpp @@ -1,6 +1,8 @@ #pragma once +#include "../../layers/linear/vocab_parallel.hpp" #include "qwen3_5_model.hpp" +#include "qwen3_5_mtp.hpp" #include #include @@ -12,12 +14,14 @@ class Qwen35ForCausalLM : public InfinilmModel { const infinicore::Device &device); Output forward(const Input &input) const override; + bool supports_token_state_checkpoints() const override { return mtp_ != nullptr; } void reset_cache(const cache::CacheConfig *cache_config) override; protected: INFINICORE_NN_MODULE(Qwen35Model, model); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); + INFINICORE_NN_MODULE(Qwen35MTP, mtp); + INFINICORE_NN_MODULE(infinilm::nn::VocabParallelLinear, lm_head); }; std::shared_ptr prepare_qwen3_5_model_config( diff --git a/csrc/models/qwen3_5/qwen3_5_model.hpp b/csrc/models/qwen3_5/qwen3_5_model.hpp index e9f2080bb..7f94eed39 100644 --- a/csrc/models/qwen3_5/qwen3_5_model.hpp +++ b/csrc/models/qwen3_5/qwen3_5_model.hpp @@ -46,6 +46,10 @@ class Qwen35ModelTemplate : public Qwen35ModelBase { return language_model_->forward(input); } + infinicore::Tensor embed_input_ids(const infinicore::Tensor &input_ids) const { + return language_model_->embed_tokens(input_ids); + } + protected: INFINICORE_NN_MODULE(LanguageModel, language_model); }; diff --git a/csrc/models/qwen3_5/qwen3_5_mtp.cpp b/csrc/models/qwen3_5/qwen3_5_mtp.cpp new file mode 100644 index 000000000..ab01eded8 --- /dev/null +++ b/csrc/models/qwen3_5/qwen3_5_mtp.cpp @@ -0,0 +1,51 @@ +#include "qwen3_5_mtp.hpp" + +#include "../../global_state/global_state.hpp" +#include +#include + +namespace infinilm::models::qwen3_5 { + +Qwen35MTP::Qwen35MTP(std::shared_ptr config, const infinicore::Device &device) { + if (config->get_or("mtp_num_hidden_layers", 1) != 1 + || global_state::get_tensor_model_parallel_rank_info().pp_size != 1 + || config->get_or("mtp_use_dedicated_embeddings", false) + || (config->get_config_json().contains("quantization_config") + && !config->get_config_json()["quantization_config"].is_null() + && !config->get_config_json()["quantization_config"].empty() + && config->get_quant_scheme() != quantization::QuantScheme::FP8_BLOCK_W8A16)) { + throw std::runtime_error("Qwen MTP requires one MTP layer, PP1, shared embeddings and BF16 or block FP8 weights."); + } + const auto hidden = config->get("hidden_size"); + const auto eps = config->get("rms_norm_eps"); + const auto dtype = config->get_dtype(); + INFINICORE_NN_MODULE_INIT(pre_fc_norm_embedding, hidden, eps, dtype, device); + INFINICORE_NN_MODULE_INIT(pre_fc_norm_hidden, hidden, eps, dtype, device); + INFINICORE_NN_MODULE_INIT(fc, 2 * hidden, hidden, false, dtype, device); + INFINICORE_NN_MODULE_INIT(norm, hidden, eps, dtype, device); + + auto json = config->get_config_json(); + const size_t target_layers = config->get("num_hidden_layers"); + json["layer_types"].push_back("full_attention"); + json["num_hidden_layers"] = target_layers + 1; + auto mtp_config = std::make_shared(json); + // Keep the checkpoint's `layers.0` name and a separate KV layer after target. + layer_ = register_module("layers.0", mtp_config, target_layers, device); +} + +infinicore::Tensor Qwen35MTP::forward(const infinicore::Tensor &embeddings, + const infinicore::Tensor &target_hidden, + const infinicore::Tensor &positions) const { + if (embeddings->shape() != target_hidden->shape() + || embeddings->dtype() != target_hidden->dtype()) { + throw std::runtime_error("Qwen MTP requires matching embedding and target hidden-state shapes and dtypes."); + } + auto embed = pre_fc_norm_embedding_->forward(embeddings); + auto hidden = pre_fc_norm_hidden_->forward(target_hidden); + auto fused = infinicore::op::cat({embed, hidden}, -1); + hidden = fc_->forward(fused); + hidden = layer_->forward(positions, hidden); + return norm_->forward(hidden); +} + +} // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_5/qwen3_5_mtp.hpp b/csrc/models/qwen3_5/qwen3_5_mtp.hpp new file mode 100644 index 000000000..716c27a04 --- /dev/null +++ b/csrc/models/qwen3_5/qwen3_5_mtp.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include "infinicore/nn/rmsnorm.hpp" +#include "qwen3_5_decoderLayer.hpp" + +namespace infinilm::models::qwen3_5 { + +// The checkpoint's one-step MTP head shares the target embedding and LM head. +class Qwen35MTP : public infinicore::nn::Module { +public: + Qwen35MTP(std::shared_ptr config, const infinicore::Device &device); + + infinicore::Tensor forward(const infinicore::Tensor &embeddings, + const infinicore::Tensor &target_hidden, + const infinicore::Tensor &positions) const; + +private: + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, pre_fc_norm_embedding); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, pre_fc_norm_hidden); + INFINICORE_NN_MODULE(layers::linear::ReplicatedLinear, fc); + INFINICORE_NN_MODULE(Qwen35DecoderLayer, layer); + INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); +}; + +} // namespace infinilm::models::qwen3_5 diff --git a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp index 3aaf898cb..ea52bdfa3 100644 --- a/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp +++ b/csrc/models/qwen3_next/qwen3_next_allocate_kv_cache_tensors.cpp @@ -33,6 +33,9 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( const size_t linear_value_head_dim = text_config->get("linear_value_head_dim"); const auto &dtype{text_config->get_dtype()}; + const auto ssm_dtype = text_config->get_or("mamba_ssm_dtype", "") == "float32" + ? infinicore::DataType::F32 + : dtype; const auto &kv_cache_dtype{text_config->get_kv_cache_dtype()}; const std::vector layer_types = text_config->get>("layer_types"); @@ -57,7 +60,7 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( linear_value_head_dim, linear_num_key_heads, linear_num_value_heads, - dtype, + ssm_dtype, pool_size); kv_cache_vec.emplace_back(); @@ -121,7 +124,12 @@ AllocatedHybridCache qwen3_next_allocate_cache_tensors( if (nullptr == paged_kv_cache_config) { throw std::runtime_error("infinilm::models::qwen3_next::qwen3_next_allocate_kv_cache_tensors: invalid paged kv cache config type"); } - const size_t mamba_pool_size = std::max(2, paged_kv_cache_config->num_blocks() / 4); + const size_t mamba_pool_size = paged_kv_cache_config->num_state_rows() + ? paged_kv_cache_config->num_state_rows() + : std::max(2, paged_kv_cache_config->num_blocks() / 4); + if (mamba_pool_size < 2) { + throw std::invalid_argument("State pool needs a zero row and a writable row."); + } for (size_t layer_idx = 0; layer_idx < num_hidden_layers; ++layer_idx) { const std::string &layer_type = layer_types[layer_idx]; diff --git a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp index 022454247..2ae0bba77 100644 --- a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp +++ b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp @@ -2,6 +2,7 @@ #include "../../global_state/global_state.hpp" +#include #include #include #include @@ -15,6 +16,20 @@ #include namespace infinilm::models::qwen3_next { +namespace { +// Limit the per-token path to short speculation windows. The crossover with +// the chunked kernel depends on the device and has not been tuned generally. +constexpr size_t kMaxRecurrentVerifyTokens = 8; + +infinicore::Tensor cast_for_state(const infinicore::Tensor &input, infinicore::DataType dtype) { + if (input->dtype() == dtype) { + return input; + } + auto output = infinicore::Tensor::empty(input->shape(), dtype, input->device()); + infinicore::op::cast_(output, input); + return output; +} +} // namespace Qwen3NextCausalConv1D::Qwen3NextCausalConv1D(std::shared_ptr model_config, size_t layer_idx, @@ -85,14 +100,33 @@ infinicore::Tensor Qwen3NextCausalConv1D::forward(const infinicore::Tensor &qkv) auto &forward_context = infinilm::global_state::get_forward_context(); auto &mamba_metadata = forward_context.mamba_metadata; - auto conv_out = infinicore::op::causal_conv1d( - qkv, - forward_context.conv_state_vec[layer_idx_], - weight_->narrow({{0, 0, local_conv_dim_}}), // narrow in case load is skipped - std::nullopt, - mamba_metadata.input_offsets.value(), - mamba_metadata.init_state_indices.value(), - mamba_metadata.final_state_indices.value()); + auto weight = weight_->narrow({{0, 0, local_conv_dim_}}); // Handle skipped weight loading. + infinicore::Tensor conv_out; + if (mamba_metadata.token_state_indices.has_value()) { + // Preserve convolution history at the same token boundaries as the + // delta-rule state; restoring only one of the two is incorrect. + conv_out = infinicore::Tensor::empty(qkv->shape(), qkv->dtype(), qkv->device()); + const auto &destinations = mamba_metadata.token_state_indices.value(); + size_t request = 0; + const auto &offsets = mamba_metadata.checkpoint_offsets; + for (size_t t = 0; t < qkv->size(1); ++t) { + while (t >= static_cast(offsets[request + 1])) { ++request; } + auto initial = t == static_cast(offsets[request]) + ? mamba_metadata.init_state_indices.value()->narrow({{0, request, 1}}) + : destinations->narrow({{0, t - 1, 1}}); + infinicore::op::causal_conv1d_( + conv_out->narrow({{1, t, 1}}), + forward_context.conv_state_vec[layer_idx_], std::nullopt, + qkv->narrow({{1, t, 1}}), weight, std::nullopt, std::nullopt, + initial, destinations->narrow({{0, t, 1}})); + } + } else { + conv_out = infinicore::op::causal_conv1d( + qkv, forward_context.conv_state_vec[layer_idx_], weight, + std::nullopt, mamba_metadata.input_offsets.value(), + mamba_metadata.init_state_indices.value(), + mamba_metadata.final_state_indices.value()); + } auto conv_qkv = infinicore::op::silu(conv_out); return conv_qkv; } @@ -130,7 +164,10 @@ Qwen3NextGatedDeltaNet::Qwen3NextGatedDeltaNet(std::shared_ptrregister_module("in_proj_z", hidden_size, value_dim, false, dtype, device, tp_rank, tp_size); + auto z_quantization = model_config->get_quant_scheme() == quantization::QuantScheme::FP8_BLOCK_W8A16 + ? quantization_method + : std::make_shared(); + in_proj_z_ = this->register_module("in_proj_z", hidden_size, value_dim, z_quantization, false, dtype, device, tp_rank, tp_size); in_proj_a_ = this->register_module("in_proj_a", hidden_size, linear_num_value_heads, false, dtype, device, tp_rank, tp_size); in_proj_b_ = this->register_module("in_proj_b", hidden_size, linear_num_value_heads, false, dtype, device, tp_rank, tp_size); @@ -152,17 +189,42 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid auto qkv = in_proj_qkv_->forward(hidden_states_mutable); auto z = in_proj_z_->forward(hidden_states_mutable); - auto a = in_proj_a_->forward(hidden_states_mutable); - auto b = in_proj_b_->forward(hidden_states_mutable); + // Keep the tiny gate projections on the same GEMM shape as decode. + // In BF16, Q=1 and Q=2 can round differently and change candidate ranking + // after recurrent-state updates. Large quantized projections stay batched. + auto project_gate = [&](const auto &projection) { + const auto &metadata = infinilm::global_state::get_forward_context().mamba_metadata; + if (!metadata.token_state_indices.has_value() || seq_len == 1) { + return projection->forward(hidden_states_mutable); + } + auto output = infinicore::Tensor::empty({batch_size, seq_len, local_num_value_heads_}, + hidden_states->dtype(), hidden_states->device()); + for (size_t t = 0; t < seq_len; ++t) { + auto token = hidden_states->narrow({{1, t, 1}}); + output->narrow({{1, t, 1}})->copy_from(projection->forward(token)); + } + return output; + }; + auto a = project_gate(in_proj_a_); + auto b = project_gate(in_proj_b_); auto &forward_context = infinilm::global_state::get_forward_context(); auto &mamba_metadata = forward_context.mamba_metadata; + const bool single_request = batch_size == 1 + && mamba_metadata.input_offsets.value()->numel() == 2; + if (mamba_metadata.token_state_indices.has_value() + && (batch_size != 1 || seq_len == 0 || mamba_metadata.checkpoint_offsets.size() < 2)) { + throw std::runtime_error("GDN token checkpoints require one request with 1 to 8 tokens."); + } auto conv_qkv = this->conv1d_->forward(qkv); - auto q = conv_qkv->narrow({{2, 0, local_key_dim_}}); - auto k = conv_qkv->narrow({{2, local_key_dim_, local_key_dim_}}); - auto v = conv_qkv->narrow({{2, local_key_dim_ * 2, local_value_dim_}}); + // Existing delta-rule operators require activations and persistent state to + // share a dtype. Honor checkpoints requesting FP32 state with device casts. + auto state_qkv = cast_for_state(conv_qkv, forward_context.ssm_state_vec[layer_idx_]->dtype()); + auto q = state_qkv->narrow({{2, 0, local_key_dim_}}); + auto k = state_qkv->narrow({{2, local_key_dim_, local_key_dim_}}); + auto v = state_qkv->narrow({{2, local_key_dim_ * 2, local_value_dim_}}); bool is_decode = mamba_metadata.input_offsets.value()->shape()[0] - 1 == seq_len; infinicore::Tensor delta_out; if (is_decode) { @@ -198,6 +260,67 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid delta_out = delta_out->as_strided( {seq_len, local_num_value_heads_, value_head_dim_}, {delta_out->stride(0), delta_out->stride(2), delta_out->stride(3)}); + } else if ((single_request && seq_len <= kMaxRecurrentVerifyTokens) || mamba_metadata.token_state_indices.has_value()) { + // Short single-sequence speculation window (e.g. one draft token per + // step). The chunked kernel pads to a 128-token chunk, so a two-token + // verification pays a full chunk of work per layer. Reuse the existing + // T=1 indexed-pool operator once per token instead; this mirrors the + // recurrent path used for single-token decode and needs no kernel + // change. Optional token destinations retain intermediate states so + // callers can commit an accepted prefix without a target replay. + auto ssm_state = forward_context.ssm_state_vec[layer_idx_]; + auto q_delta = q->as_strided( + {1, seq_len, local_num_key_heads_, key_head_dim_}, + {q->stride(0), q->stride(1), static_cast(key_head_dim_), 1}); + auto k_delta = k->as_strided( + {1, seq_len, local_num_key_heads_, key_head_dim_}, + {k->stride(0), k->stride(1), static_cast(key_head_dim_), 1}); + auto v_delta = v->as_strided( + {1, seq_len, local_num_value_heads_, value_head_dim_}, + {v->stride(0), v->stride(1), static_cast(value_head_dim_), 1}); + + auto a_heads = a->as_strided( + {1, seq_len, local_num_value_heads_}, + {a->stride(0), a->stride(1), 1}); + auto b_heads = b->as_strided( + {1, seq_len, local_num_value_heads_}, + {b->stride(0), b->stride(1), 1}); + auto [g, beta] = infinicore::op::fused_gated_delta_net_gating(A_log_, a_heads, b_heads, dt_bias_); + + auto recurrent_out = infinicore::Tensor::empty( + {1, seq_len, local_num_value_heads_, value_head_dim_}, + ssm_state->dtype(), ssm_state->device()); + const auto &init_indices = mamba_metadata.init_state_indices.value(); + const auto &final_indices = mamba_metadata.final_state_indices.value(); + for (size_t t = 0; t < seq_len; ++t) { + auto step_init = (t == 0) ? init_indices : final_indices; + auto step_final = final_indices; + if (mamba_metadata.token_state_indices.has_value()) { + const auto &destinations = mamba_metadata.token_state_indices.value(); + step_final = destinations->narrow({{0, t, 1}}); + const auto &offsets = mamba_metadata.checkpoint_offsets; + size_t request = 0; + while (t >= static_cast(offsets[request + 1])) { ++request; } + step_init = t == static_cast(offsets[request]) + ? init_indices->narrow({{0, request, 1}}) + : destinations->narrow({{0, t - 1, 1}}); + } + infinicore::op::recurrent_gated_delta_rule_( + recurrent_out->narrow({{1, t, 1}}), + ssm_state, + std::nullopt, + q_delta->narrow({{1, t, 1}}), + k_delta->narrow({{1, t, 1}}), + v_delta->narrow({{1, t, 1}}), + g->narrow({{1, t, 1}}), + beta->narrow({{1, t, 1}}), + step_init, + step_final, + true); + } + delta_out = recurrent_out->as_strided( + {seq_len, local_num_value_heads_, value_head_dim_}, + {recurrent_out->stride(1), recurrent_out->stride(2), recurrent_out->stride(3)}); } else { auto ssm_state = forward_context.ssm_state_vec[layer_idx_]; auto q_delta = q->as_strided( @@ -234,6 +357,7 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid {delta_out->stride(1), delta_out->stride(2), delta_out->stride(3)}); } + delta_out = cast_for_state(delta_out, hidden_states->dtype()); auto delta_out_2d = delta_out->as_strided( {batch_size * seq_len * local_num_value_heads_, value_head_dim_}, {static_cast(value_head_dim_), 1}); diff --git a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.hpp b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.hpp index e7d64e7a9..9b80bb384 100644 --- a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.hpp +++ b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.hpp @@ -35,6 +35,14 @@ class Qwen3NextGatedDeltaNet : public infinicore::nn::Module { infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const; + void process_weights_after_loading() override { + in_proj_qkv_->process_weights_after_loading(); + } + + void reset_runtime_state() const override { + in_proj_qkv_->reset_runtime_state(); + } + private: std::shared_ptr in_proj_qkv_; std::shared_ptr in_proj_z_; diff --git a/csrc/pybind11/cache/cache.hpp b/csrc/pybind11/cache/cache.hpp index bf76ebfe0..a3f24ad3e 100644 --- a/csrc/pybind11/cache/cache.hpp +++ b/csrc/pybind11/cache/cache.hpp @@ -34,10 +34,12 @@ inline void bind_cache(py::module &m) { infinilm::cache::CacheConfig, std::shared_ptr>(m, "PagedKVCacheConfig") .def( - py::init(), + py::init(), py::arg("num_blocks"), py::arg("block_size") = 256, - py::arg("max_batch_size") = 1) + py::arg("max_batch_size") = 1, + py::arg("num_state_rows") = 0) + .def("num_state_rows", &infinilm::cache::PagedKVCacheConfig::num_state_rows) .def( "num_blocks", &infinilm::cache::PagedKVCacheConfig::num_blocks) diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index c5e85577c..d9f614f45 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -122,19 +122,39 @@ inline void bind_infer_engine(py::module &m) { return state_dict_tp_all; }) .def("process_weights_after_loading", &InferEngine::process_weights_after_loading, "Process the weights after loading on all workers (e.g., for quantization)") - .def( - "forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { + .def("forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { // IMPORTANT: Release the GIL before calling forward() to allow other Python threads // to run concurrently during inference (which may block for a long time). // Do NOT remove this — without it, the GIL is held throughout inference and will // deadlock or stall any other Python thread (e.g., request handling, scheduling). py::gil_scoped_release release; - return self.forward(input); - }, - "Run inference on all ranks with arbitrary arguments") - .def( - "reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) - .def("get_kv_cache", &InferEngine::get_kv_cache, "Get per-rank kv cache list") + return self.forward(input); }, "Run inference on all ranks with arbitrary arguments") + .def("reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) + .def("compile", &InferEngine::compile, py::call_guard()) + .def("get_kv_cache", [](InferEngine &self) { + py::list ranks; + for (const auto &rank : self.get_kv_cache()) { + py::list layers; + for (const auto &tensor : rank) { + layers.append(tensor ? py::cast(tensor) : py::none()); + } + ranks.append(layers); + } + return ranks; }) + .def("get_hybrid_states", [](InferEngine &self) { + py::list ranks; + for (const auto &rank : self.get_hybrid_states()) { + py::list kinds; + for (const auto &kind : rank) { + py::list layers; + for (const auto &tensor : kind) { + layers.append(tensor ? py::cast(tensor) : py::none()); + } + kinds.append(layers); + } + ranks.append(kinds); + } + return ranks; }) .def("get_cache_config", [](const InferEngine &self) -> std::shared_ptr { auto cfg = self.get_cache_config(); return cfg ? std::shared_ptr(cfg->unique_copy()) : nullptr; }) @@ -161,6 +181,7 @@ inline void bind_infer_engine(py::module &m) { std::optional> visual_token_ranges, std::optional target_hidden_states, bool sample_all_positions, + std::optional token_state_indices, py::kwargs kwargs) { InferEngine::Input input{ std::move(input_ids), @@ -181,6 +202,7 @@ inline void bind_infer_engine(py::module &m) { std::move(visual_token_ranges), std::move(target_hidden_states), sample_all_positions, + std::move(token_state_indices), }; // Explicit defaults @@ -193,6 +215,9 @@ inline void bind_infer_engine(py::module &m) { "temperature", "top_p", "top_k", + "return_logits", + "return_device_tokens", + "verify_draft", }; for (auto &item : kwargs) { @@ -207,6 +232,12 @@ inline void bind_infer_engine(py::module &m) { input.temperature = py::cast(item.second); } else if (key == "top_p") { input.top_p = py::cast(item.second); + } else if (key == "return_device_tokens") { + input.return_device_tokens = py::cast(item.second); + } else if (key == "verify_draft") { + input.verify_draft = py::cast(item.second); + } else if (key == "return_logits") { + input.return_logits = py::cast(item.second); } else if (key == "top_k") { input.top_k = py::cast(item.second); } @@ -231,7 +262,8 @@ inline void bind_infer_engine(py::module &m) { py::arg("image_req_ids") = std::nullopt, py::arg("visual_token_ranges") = std::nullopt, py::arg("target_hidden_states") = std::nullopt, - py::arg("sample_all_positions") = false) + py::arg("sample_all_positions") = false, + py::arg("token_state_indices") = std::nullopt) .def_readwrite("input_ids", &InferEngine::Input::input_ids) .def_readwrite("position_ids", &InferEngine::Input::position_ids) .def_readwrite("past_sequence_lengths", &InferEngine::Input::past_sequence_lengths) @@ -250,6 +282,7 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) + .def_readwrite("token_state_indices", &InferEngine::Input::token_state_indices) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) .def_readwrite("top_p", &InferEngine::Input::top_p); @@ -257,7 +290,8 @@ inline void bind_infer_engine(py::module &m) { py::class_(infer_engine, "Output") .def_readwrite("output_ids", &InferEngine::Output::output_ids, "Sampled token IDs") .def_readwrite("logits", &InferEngine::Output::logits, "Raw logits tensor") - .def_readwrite("hidden_states", &InferEngine::Output::hidden_states, "Raw hidden states tensor"); + .def_readwrite("hidden_states", &InferEngine::Output::hidden_states, "Raw hidden states tensor") + .def_readwrite("accepted_draft_tokens", &InferEngine::Output::accepted_draft_tokens); } } // namespace infinilm::engine diff --git a/examples/bench.py b/examples/bench.py index 17bfe1a6d..e8e3634e8 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -445,6 +445,9 @@ def __init__( top_k=1, use_legacy_moe=False, enable_prefix_caching=False, + enable_mtp=False, + num_state_rows=0, + mtp_prefix_cache_bytes=0, processor=None, tokenizer=None, prompt_token_segments=None, @@ -454,6 +457,7 @@ def __init__( model_path = os.path.expanduser(model_path) self.draft_model_path = draft_model_path self.num_draft_tokens = num_draft_tokens + self.enable_mtp = enable_mtp self.model_path = model_path self.device_str = infini_device.type self.tp = tp @@ -474,7 +478,9 @@ def __init__( if pp > 1 and draft_model_path is not None: raise ValueError("pipeline-parallel speculative decoding is not supported") - if draft_model_path is not None: + if enable_mtp and processed_multimodal_inputs is not None: + raise ValueError("Built-in MTP benchmarks support text prompts only.") + if draft_model_path is not None and not enable_mtp: if processed_multimodal_inputs is not None: raise ValueError("Draft-model benchmarks do not support --image") if self.processor is None: @@ -490,9 +496,14 @@ def __init__( self.model = None return - if pp > 1: + if pp > 1 or enable_mtp: self.model = LLM( model_path=model_path, + draft_model_path=draft_model_path, + num_draft_tokens=num_draft_tokens, + enable_mtp=enable_mtp, + num_state_rows=num_state_rows, + mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, device=self.device_str, dtype=cfg.dtype, tensor_parallel_size=tp, @@ -514,6 +525,7 @@ def __init__( enable_graph=enable_graph, attn_backend=attn_backend, use_mla=use_mla, + pre_transpose=pre_transpose, weight_load_mode=weight_load_mode, skip_load=skip_load, use_legacy_moe=use_legacy_moe, @@ -638,11 +650,11 @@ def get_input_tokens(self): return inputs @property - def uses_pipeline_parallel(self) -> bool: - return self.pp > 1 + def uses_llm_engine(self) -> bool: + return self.pp > 1 or self.enable_mtp def close(self) -> None: - if self.uses_pipeline_parallel: + if self.uses_llm_engine: self.model.close() def run( @@ -662,7 +674,7 @@ def run( # ---------------------------------------------------------------------------- # # 自回归生成 # ---------------------------------------------------------------------------- # - if self.draft_model_path is not None or self.uses_pipeline_parallel: + if self.draft_model_path is not None or self.uses_llm_engine: prompt_text = self.tokenizer.decode(input_ids, skip_special_tokens=False) llm = self.model if self.draft_model_path is not None: @@ -699,6 +711,9 @@ def run( use_tqdm=False, ) t2 = time.time() + actual_input_lengths = [len(output.prompt_token_ids) for output in outputs] + if any(length != input_len for length in actual_input_lengths): + print(f"[bench] tokenized input lengths: {actual_input_lengths}") if cfg.verbose and not skip_load: if output_len <= 256: for output in outputs: @@ -859,7 +874,7 @@ def run( # 测试 # -------------------------------------------------------- # if enable_paged_attn: - paged_kv_block_size = _PAGED_KV_BLOCK_SIZE + paged_kv_block_size = cfg.block_size if cfg.enable_mtp else _PAGED_KV_BLOCK_SIZE max_num_blocks = max( [ ( @@ -953,7 +968,10 @@ def run( top_p=cfg.top_p, top_k=cfg.top_k, use_legacy_moe=cfg.use_legacy_moe, - enable_prefix_caching=False, + enable_prefix_caching=cfg.enable_prefix_caching if cfg.enable_mtp else False, + enable_mtp=cfg.enable_mtp, + num_state_rows=cfg.num_state_rows, + mtp_prefix_cache_bytes=int(cfg.mtp_prefix_cache_mib * 1024**2), processor=processor, tokenizer=tokenizer, prompt_token_segments=prompt_token_segments, @@ -979,7 +997,7 @@ def run( ) print("=================== warmup start ===================") - if test.uses_pipeline_parallel: + if test.uses_llm_engine: for _ in range(warmup_steps): test.run( batch_size=warmup_batch, @@ -1030,7 +1048,7 @@ def run( print("=================== warmup done ====================") # reset cache back to benchmark config - if cache_config is not None and not test.uses_pipeline_parallel: + if cache_config is not None and not test.uses_llm_engine: test.model.reset_cache(cache_config) # ---------------------------------------------------------------------------- # @@ -1044,7 +1062,7 @@ def run( input_len = case["input_len"] output_len = case["output_len"] - if not test.uses_pipeline_parallel and not enable_paged_attn: + if not test.uses_llm_engine and not enable_paged_attn: # reset cache if static kvcache is used initial_capacity = input_len + output_len test.model.reset_cache( diff --git a/examples/test_infer.py b/examples/test_infer.py index f17a8baa0..2dde71415 100644 --- a/examples/test_infer.py +++ b/examples/test_infer.py @@ -41,6 +41,9 @@ def test( use_legacy_moe=False, enable_prefix_caching=True, pre_transpose=False, + enable_mtp=False, + num_state_rows=0, + mtp_prefix_cache_bytes=0, ): model_path = os.path.expanduser(model_path) # ---------------------------------------------------------------------------- # @@ -77,6 +80,9 @@ def test( use_legacy_moe=use_legacy_moe, enable_prefix_caching=enable_prefix_caching, pre_transpose=pre_transpose, + enable_mtp=enable_mtp, + num_state_rows=num_state_rows, + mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, ) conversations = [ @@ -185,4 +191,7 @@ def test( use_legacy_moe=cfg.use_legacy_moe, enable_prefix_caching=cfg.enable_prefix_caching, pre_transpose=cfg.pre_transpose, + enable_mtp=cfg.enable_mtp, + num_state_rows=cfg.num_state_rows, + mtp_prefix_cache_bytes=int(cfg.mtp_prefix_cache_mib * 1024**2), ) diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 4ae7665c0..add26f9e1 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -61,6 +61,9 @@ def __init__(self): self.model = self.args.model self.draft_model = self.args.draft_model self.num_draft_tokens = self.args.num_draft_tokens + self.enable_mtp = self.args.enable_mtp + self.num_state_rows = self.args.num_state_rows + self.mtp_prefix_cache_mib = self.args.mtp_prefix_cache_mib self.device = self.args.device self.tp = self.args.tp self.pp = self.args.pp @@ -192,7 +195,7 @@ def _add_common_args(self): "--num-draft-tokens", type=int, default=1, - help="number of Eagle draft tokens to verify per target step", + help="number of Eagle or built-in MTP draft tokens to verify per target step", ) self.parser.add_argument( "--device", @@ -204,6 +207,23 @@ def _add_common_args(self): "(cuda/mlu/musa/npu)" ), ) + self.parser.add_argument( + "--enable-mtp", + action="store_true", + help="use the checkpoint's built-in Qwen MTP head (greedy)", + ) + self.parser.add_argument( + "--num-state-rows", + type=int, + default=0, + help="Qwen hybrid state rows, including zero row; 0 selects automatic capacity", + ) + self.parser.add_argument( + "--mtp-prefix-cache-mib", + type=int, + default=0, + help="device storage budget for exact-prompt MTP snapshots; TP1 only", + ) self.parser.add_argument("--tp", "--tensor-parallel-size", type=int, default=1) self.parser.add_argument( "--pp", "--pipeline-parallel-size", type=int, default=1 diff --git a/python/infinilm/cache/cache.py b/python/infinilm/cache/cache.py index e52011c30..68709d701 100644 --- a/python/infinilm/cache/cache.py +++ b/python/infinilm/cache/cache.py @@ -23,7 +23,8 @@ def __init__( num_blocks: int, block_size: int = 256, max_batch_size: int = 1, + num_state_rows: int = 0, ): _infinilm.PagedKVCacheConfig.__init__( - self, num_blocks, block_size, max_batch_size + self, num_blocks, block_size, max_batch_size, num_state_rows ) diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index bccb7e758..cea158bf7 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -10,8 +10,12 @@ class EngineConfig: Attributes: model_path: Path to the model directory. - draft_model_path: Optional Eagle/MTP draft model directory. - num_draft_tokens: Number of Eagle draft tokens to verify per step. + draft_model_path: Optional external Eagle draft model directory. + num_draft_tokens: Number of draft tokens to verify per step. + enable_mtp: Use the Qwen checkpoint's built-in MTP head. + num_state_rows: Hybrid state rows, including the reserved zero row. + Zero selects an automatic capacity. + mtp_prefix_cache_bytes: Device storage budget for exact-prompt MTP snapshots. device: Device type string ('cpu', 'cuda', 'mlu', etc.). dtype: Data type string ('float16', 'bfloat16', 'float32'). tensor_parallel_size: Number of devices for tensor parallelism. @@ -69,8 +73,63 @@ class EngineConfig: use_legacy_moe: bool = False kv_transfer_config: Optional[KVTransferConfig] = None enable_prefix_caching: bool = True + enable_mtp: bool = False + num_state_rows: int = 0 + mtp_prefix_cache_bytes: int = 0 def __post_init__(self) -> None: + if self.max_batch_size < 1: + raise ValueError("`max_batch_size` must be >= 1.") + if self.num_state_rows != 0 and self.num_state_rows < 2: + raise ValueError("`num_state_rows` must be zero (automatic) or >= 2.") + if self.mtp_prefix_cache_bytes < 0: + raise ValueError("`mtp_prefix_cache_bytes` must be non-negative.") + if self.mtp_prefix_cache_bytes and not self.enable_mtp: + raise ValueError("`mtp_prefix_cache_bytes` requires `enable_mtp`.") + if self.mtp_prefix_cache_bytes and ( + not self.enable_prefix_caching or self.tensor_parallel_size != 1 + ): + raise ValueError( + "MTP prefix snapshots require prefix caching and `tensor_parallel_size=1`." + ) + if self.enable_mtp: + if self.draft_model_path is not None: + raise ValueError( + "Built-in MTP cannot be combined with `draft_model_path`." + ) + if not 1 <= self.num_draft_tokens <= 4: + raise ValueError("Qwen MTP requires `1 <= num_draft_tokens <= 4`.") + if self.enable_graph and ( + self.num_draft_tokens != 1 or self.max_batch_size != 1 + ): + raise ValueError( + "Batched or multi-candidate Qwen MTP currently requires eager mode." + ) + if self.cache_type != "paged" or self.pipeline_parallel_size != 1: + raise ValueError( + "Qwen MTP requires paged caching and `pipeline_parallel_size=1`." + ) + if self.enable_prefix_caching and not self.mtp_prefix_cache_bytes: + raise ValueError( + "Disable prefix caching or set `mtp_prefix_cache_bytes` " + "for exact-prompt MTP snapshots." + ) + if ( + self.kv_transfer_config is not None + and self.kv_transfer_config.kv_connector + ): + raise ValueError("Qwen MTP does not support remote KV/state transfer.") + if self.top_k != 1: + raise ValueError( + "Qwen MTP currently supports greedy sampling (`top_k=1`)." + ) + if self.attn_backend not in ("default", "paged-attn"): + raise ValueError("Qwen MTP requires the `paged-attn` backend.") + self.attn_backend = "paged-attn" + if not self.num_state_rows: + self.num_state_rows = 1 + self.max_batch_size * ( + self.num_draft_tokens + 2 + ) if self.num_draft_tokens < 1: raise ValueError("num_draft_tokens must be >= 1") if self.pipeline_parallel_size < 1: diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 117d82f9a..3ba833347 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -4,6 +4,7 @@ from dataclasses import dataclass import infinicore +from infinicore.lib import _infinicore from infinilm.cache import PagedKVCacheConfig from infinilm.distributed import DistConfig @@ -165,8 +166,21 @@ def __init__( moe_ep_size=1, use_legacy_moe=False, pre_transpose=False, + enable_mtp=None, ): + self.cache_generation = 0 self.hf_config = read_hf_config(model_path) + if enable_mtp is not None: + self.hf_config["enable_mtp"] = bool(enable_mtp) + if self.hf_config.get("enable_mtp", False): + text_config = self.hf_config.get("text_config", self.hf_config) + if ( + self.hf_config.get("model_type") != "qwen3_5" + or text_config.get("mtp_num_hidden_layers", 1) != 1 + ): + raise ValueError( + "Built-in MTP requires a single-layer dense Qwen model." + ) self.hf_generation_config = read_hf_generation_config(model_path) self.hf_config["use_legacy_moe"] = bool(use_legacy_moe) self.position_id_axes = _infer_position_id_axes(self.hf_config) @@ -179,6 +193,9 @@ def __init__( device = infinicore.device() if distributed_config is None: distributed_config = DistConfig(1) + self.supports_device_mtp = ( + device._underlying.type == _infinicore.Device.Type.NVIDIA + ) self.distributed_config = distributed_config if ( moe_ep_backend != "disabled" @@ -267,6 +284,7 @@ def _build_input( slot_mapping=None, mamba_init_state_indices=None, mamba_final_state_indices=None, + token_state_indices=None, pixel_values=None, image_bound=None, tgt_sizes=None, @@ -275,6 +293,9 @@ def _build_input( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=False, + return_logits=True, + return_device_tokens=False, + verify_draft=False, temperature=None, top_k=None, top_p=None, @@ -294,6 +315,7 @@ def unwrap_tensor(tensor): slot_mapping = unwrap_tensor(slot_mapping) mamba_init_state_indices = unwrap_tensor(mamba_init_state_indices) mamba_final_state_indices = unwrap_tensor(mamba_final_state_indices) + token_state_indices = unwrap_tensor(token_state_indices) target_hidden_states = unwrap_tensor(target_hidden_states) def convert_tensor_list(tensor_list_): @@ -325,6 +347,7 @@ def convert_tensor_list(tensor_list_): slot_mapping=slot_mapping, mamba_init_state_indices=mamba_init_state_indices, mamba_final_state_indices=mamba_final_state_indices, + token_state_indices=token_state_indices, pixel_values=pixel_values, image_bound=image_bound, tgt_sizes=tgt_sizes, @@ -333,6 +356,9 @@ def convert_tensor_list(tensor_list_): visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, + return_logits=return_logits, + return_device_tokens=return_device_tokens, + verify_draft=verify_draft, temperature=temperature, top_k=top_k, top_p=top_p, @@ -351,6 +377,7 @@ def forward( slot_mapping=None, mamba_init_state_indices=None, mamba_final_state_indices=None, + token_state_indices=None, pixel_values=None, image_bound=None, tgt_sizes=None, @@ -423,6 +450,7 @@ def convert_tensor_list(tensor_list_): slot_mapping=slot_mapping, mamba_init_state_indices=mamba_init_state_indices, mamba_final_state_indices=mamba_final_state_indices, + token_state_indices=token_state_indices, pixel_values=pixel_values, image_bound=image_bound, tgt_sizes=tgt_sizes, @@ -430,6 +458,7 @@ def convert_tensor_list(tensor_list_): image_req_ids=image_req_ids, visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, + return_logits=False, temperature=temperature, top_k=top_k, top_p=top_p, @@ -452,6 +481,9 @@ def forward_raw( cu_seqlens=None, block_tables=None, slot_mapping=None, + mamba_init_state_indices=None, + mamba_final_state_indices=None, + token_state_indices=None, pixel_values=None, image_bound=None, tgt_sizes=None, @@ -459,6 +491,9 @@ def forward_raw( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=True, + return_logits=True, + return_device_tokens=False, + verify_draft=False, temperature=None, top_k=None, top_p=None, @@ -474,6 +509,9 @@ def forward_raw( cu_seqlens=cu_seqlens, block_tables=block_tables, slot_mapping=slot_mapping, + mamba_init_state_indices=mamba_init_state_indices, + mamba_final_state_indices=mamba_final_state_indices, + token_state_indices=token_state_indices, pixel_values=pixel_values, image_bound=image_bound, tgt_sizes=tgt_sizes, @@ -481,6 +519,9 @@ def forward_raw( visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, + return_logits=return_logits, + return_device_tokens=return_device_tokens, + verify_draft=verify_draft, temperature=temperature, top_k=top_k, top_p=top_p, @@ -488,7 +529,8 @@ def forward_raw( ) return { "output_ids": infinicore.Tensor(output.output_ids), - "logits": infinicore.Tensor(output.logits), + "accepted_draft_tokens": output.accepted_draft_tokens, + "logits": infinicore.Tensor(output.logits) if return_logits else None, "hidden_states": infinicore.Tensor(output.hidden_states), } except BaseException as e: @@ -547,7 +589,10 @@ def generate( "Low-level generate for mamba-cache models currently requires paged attention" ) elif self.has_mamba_cache: - mamba_pool_size = max(2, self.get_cache_config().num_blocks() // 4) + cache_config = self.get_cache_config() + mamba_pool_size = max(2, cache_config.num_blocks() // 4) + if self.model_type in ("qwen3_next", "qwen3_5", "qwen3_5_moe"): + mamba_pool_size = cache_config.num_state_rows() or mamba_pool_size if batch_size > mamba_pool_size - 1: raise RuntimeError( f"Batch size {batch_size} exceeds available mamba cache rows " @@ -730,6 +775,7 @@ def generate( return output_ids def reset_cache(self, cache_config): + self.cache_generation += 1 infinicore.sync_device() self.enable_paged_attn = isinstance(cache_config, PagedKVCacheConfig) super().reset_cache(cache_config) @@ -738,6 +784,7 @@ def state_dict_keyname(self): return list(super().state_dict_keyname()) def load_state_dict(self, state_dict, strict=None): + self.cache_generation += 1 # MoE/quantized paths may register internal packed tensors that are not # present in the HF checkpoint, so callers can request non-strict loads. super().load_params( diff --git a/python/infinilm/llm/hybrid_prefix_cache.py b/python/infinilm/llm/hybrid_prefix_cache.py new file mode 100644 index 000000000..e2872e476 --- /dev/null +++ b/python/infinilm/llm/hybrid_prefix_cache.py @@ -0,0 +1,140 @@ +"""Byte-bounded, exact-prompt snapshots for single-rank greedy hybrid MTP. + +Snapshots own their device storage. Restoration writes into request-owned pages +and recurrent rows; no mutable state or partially filled KV page is shared. +""" + +import math +from collections import OrderedDict +from dataclasses import dataclass + +import infinicore + + +@dataclass +class HybridPrefixSnapshot: + kv: list + states: list + hidden: object + proposal: object + pending: int + nbytes: int + + +class HybridPrefixCache: + def __init__(self, engine, budget_bytes, block_size): + self.engine = engine + self.budget_bytes = budget_bytes + self.block_size = block_size + self.entries = OrderedDict() + self.used_bytes = 0 + self.hits = self.misses = self.evictions = self.skipped = 0 + self.generation = self._generation() + + def _generation(self): + return getattr(self.engine, "cache_generation", 0) + + def clear(self): + self.entries.clear() + self.used_bytes = 0 + self.generation = self._generation() + + def _check_generation(self): + if self.generation != self._generation(): + self.clear() + + def _storage(self): + kv = [t for t in self.engine.get_kv_cache()[0] if t._underlying is not None] + states = [ + infinicore.Tensor(t) + for group in self.engine.get_hybrid_states()[0] + for t in group + if t is not None + ] + return kv, states + + @staticmethod + def _bytes(tensor): + return ( + math.prod(tensor.shape) + * infinicore.utils.to_torch_dtype(tensor.dtype).itemsize + ) + + @staticmethod + def _clone(tensor): + # InfiniCore-owned buffers keep this copy independent of foreign allocator + # lifetimes and of the compiler's reusable graph output buffers. + result = infinicore.empty( + tensor.shape, dtype=tensor.dtype, device=tensor.device + ) + result.copy_(tensor) + return result + + def restore(self, request): + self._check_generation() + key = tuple(request.prompt_token_ids) + saved = self.entries.get(key) + if saved is None: + self.misses += 1 + return None + kv, states = self._storage() + # Metadata construction can leave the calling thread on CPU. In-place + # InfiniCore ops use that thread's handle/stream, unlike `empty()`, which + # selects the tensor's device. Select it explicitly before restoration. + infinicore.set_device(saved.hidden.device) + for dst, pages in zip(kv, saved.kv): + for block, src in zip(request.block_table, pages): + dst.narrow(1, block, 1).copy_(src) + for dst, src in zip(states, saved.states): + dst.narrow(0, request.mamba_cache_index, 1).copy_(src) + infinicore.sync_device() + self.entries.move_to_end(key) + self.hits += 1 + return saved + + def save(self, request, pending): + self._check_generation() + key = tuple(request.prompt_token_ids) + if key in self.entries: + self.entries.move_to_end(key) + return + state = request.mtp_state + if state is None: + return + kv, states = self._storage() + pages = (request.get_prompt_length() + self.block_size - 1) // self.block_size + sources = [ + [t.narrow(1, block, 1) for block in request.block_table[:pages]] for t in kv + ] + state_sources = [t.narrow(0, request.mamba_cache_index, 1) for t in states] + tensors = ( + [x for layer in sources for x in layer] + + state_sources + + [state.draft_hidden] + ) + if isinstance(state.draft_token, infinicore.Tensor): + tensors.append(state.draft_token) + nbytes = sum(self._bytes(t) for t in tensors) + # Evict before allocating the new snapshot, so the live snapshots never + # temporarily exceed the configured storage budget. + if nbytes > self.budget_bytes: + self.skipped += 1 + return + while self.used_bytes + nbytes > self.budget_bytes: + _, old = self.entries.popitem(last=False) + self.used_bytes -= old.nbytes + self.evictions += 1 + del old + snapshot = HybridPrefixSnapshot( + [[self._clone(t) for t in layer] for layer in sources], + [self._clone(t) for t in state_sources], + self._clone(state.draft_hidden), + self._clone(state.draft_token) + if isinstance(state.draft_token, infinicore.Tensor) + else state.draft_token, + pending, + nbytes, + ) + infinicore.sync_device() + self.entries[key] = snapshot + self.used_bytes += nbytes diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..0b6c583df 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -42,7 +42,19 @@ def __init__(self, config: EngineConfig): self.config = config hf_config = read_hf_config(config.model_path) has_mamba_cache = model_uses_mamba_cache(hf_config) - if has_mamba_cache and config.enable_prefix_caching: + if config.num_state_rows and hf_config["model_type"] not in ( + "qwen3_5", + "qwen3_5_moe", + "qwen3_next", + ): + raise ValueError( + "Explicit num_state_rows is currently supported for Qwen hybrid models only." + ) + if ( + has_mamba_cache + and config.enable_prefix_caching + and not (config.enable_mtp and config.mtp_prefix_cache_bytes) + ): model_type = hf_config["model_type"] raise RuntimeError( "Prefix caching is not supported for Mamba-cache model " @@ -94,7 +106,9 @@ def __init__(self, config: EngineConfig): max_position_embeddings = llm_config.get( "max_position_embeddings", config.max_cache_len ) - num_mamba_cache_blocks = max(2, config.num_blocks // 4) + num_mamba_cache_blocks = config.num_state_rows or max( + 2, config.num_blocks // 4 + ) max_num_batched_tokens = int( os.getenv("INFINILM_MAX_NUM_BATCHED_TOKENS", max_position_embeddings) @@ -135,9 +149,45 @@ def __init__(self, config: EngineConfig): def add_request(self, request: InferenceRequest): """Add a request to the scheduler.""" + if getattr(self, "_closed", False): + raise RuntimeError("LLM engine is closed") + if self.config.enable_mtp: + self.model_runner.speculative_runner.validate_request(request) + if ( + request.get_prompt_length() + request.sampling_params.max_tokens + > self.config.num_blocks * self.config.block_size + ): + raise ValueError( + "MTP prompt and output limit exceed the paged cache capacity." + ) self.scheduler.add_request(request) + @staticmethod + def _publish_terminal(request): + if request._output_queue is not None: + try: + request.output_queue.sync_q.put_nowait( + TokenOutput( + request_id=request.request_id, + token_id=-1, + token_text="", + finished=True, + finish_reason=request.finish_reason, + generated_text=request.generated_text, + ) + ) + except Exception: + # A disconnected client may already have closed its queue. + # Publishing a terminal event must not interrupt cache cleanup. + logger.debug("Terminal output queue closed for %s", request.request_id) + def close(self): + if getattr(self, "_closed", False): + return + self._closed = True + if self.config.enable_mtp: + for req in self.scheduler.cancel_all(): + self._publish_terminal(req) self.model_runner.close() def step(self) -> tuple[bool, list[tuple]]: @@ -154,7 +204,17 @@ def step(self) -> tuple[bool, list[tuple]]: if scheduler_output is None: return False, [] - runner_output = self.model_runner.execute_model(scheduler_output) + try: + runner_output = self.model_runner.execute_model(scheduler_output) + except Exception: + if self.config.enable_mtp: + for req in scheduler_output.scheduled_requests: + if not req.is_finished(): + req.mark_failed() + self.scheduler.complete_requests(scheduler_output.scheduled_requests) + for req in scheduler_output.scheduled_requests: + self._publish_terminal(req) + raise sampled_token_ids = runner_output.sampled_token_ids self.scheduler.update_from_output(runner_output) pending = self._update_requests( @@ -366,6 +426,9 @@ def __init__( skip_load: bool = False, use_legacy_moe: bool = False, enable_prefix_caching: bool = True, + enable_mtp: bool = False, + num_state_rows: int = 0, + mtp_prefix_cache_bytes: int = 0, ): """Initialize LLM. @@ -384,6 +447,14 @@ def __init__( top_p: Default top-p sampling parameter. top_k: Default top-k sampling parameter. enable_graph: Whether to enable graph compiling. + enable_mtp: Use built-in Qwen MTP (text, greedy, 1-4 candidates, + paged cache). Batched verification runs eager. With one request + and one candidate, enable_graph captures Decode and MTP drafts; + target verification remains eager. + num_state_rows: Hybrid state pool capacity including zero row. + Zero selects a concurrency/candidate-based capacity for MTP. + mtp_prefix_cache_bytes: Budget for exact-prompt MTP snapshots (TP1). + Nonzero requires enable_prefix_caching; zero disables snapshots. attn_backend: Attention backend to use ('default', 'flash-attn'). use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. @@ -418,6 +489,9 @@ def __init__( skip_load=skip_load, use_legacy_moe=use_legacy_moe, enable_prefix_caching=enable_prefix_caching, + enable_mtp=enable_mtp, + num_state_rows=num_state_rows, + mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, ) self.engine = LLMEngine(config) self.config = config @@ -594,6 +668,9 @@ def __init__( weight_load_mode: str = "async", use_legacy_moe: bool = False, enable_prefix_caching: bool = True, + enable_mtp: bool = False, + num_state_rows: int = 0, + mtp_prefix_cache_bytes: int = 0, ): """Initialize AsyncLLMEngine. @@ -612,6 +689,14 @@ def __init__( top_p: Default top-p sampling parameter. top_k: Default top-k sampling parameter. enable_graph: Whether to enable graph compiling. + enable_mtp: Use built-in Qwen MTP (text, greedy, 1-4 candidates, + paged cache). Batched verification runs eager. With one request + and one candidate, enable_graph captures Decode and MTP drafts; + target verification remains eager. + num_state_rows: Hybrid state pool capacity including zero row. + Zero selects a concurrency/candidate-based capacity for MTP. + mtp_prefix_cache_bytes: Budget for exact-prompt MTP snapshots (TP1). + Nonzero requires enable_prefix_caching; zero disables snapshots. attn_backend: Attention backend to use ('default', 'flash-attn'). kv_connector: KV connector type ('MooncakeConnector'). kv_role: Role in KV connector ('kv_producer' or 'kv_consumer'). @@ -651,6 +736,9 @@ def __init__( weight_load_mode=weight_load_mode, use_legacy_moe=use_legacy_moe, enable_prefix_caching=enable_prefix_caching, + enable_mtp=enable_mtp, + num_state_rows=num_state_rows, + mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, ) self.engine = LLMEngine(config) self.config = config @@ -681,13 +769,11 @@ def start(self): def stop(self): """Stop the background inference loop.""" - if not self._running: - logger.warning("AsyncLLMEngine is not running") - return - self._running = False if self._step_thread: - self._step_thread.join(timeout=5) + # Do not free model/state buffers while a long Prefill is still + # executing on the background thread. + self._step_thread.join() self.engine.close() logger.info("AsyncLLMEngine stopped") @@ -760,6 +846,8 @@ def _step_loop(self): logger.error(f"Error in step loop: {e}", exc_info=True) self._healthy = False self._running = False + if self.config.enable_mtp: + self.engine.close() break @staticmethod diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index a1696f848..7ff879f82 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -13,6 +13,7 @@ KVConnectorFactory, KVConnectorRole, ) +from infinilm.llm.model_runner.mtp_runner import MTPRunner from infinilm.llm.model_runner.speculative_runner import SpeculativeRunner from infinilm.modeling_utils import load_model_state_dict_by_file from infinilm.processors import AutoInfinilmProcessor @@ -64,6 +65,7 @@ def __init__(self, config: EngineConfig, initialize_processor: bool = True): num_blocks=config.num_blocks, block_size=config.block_size, max_batch_size=config.max_batch_size, + num_state_rows=config.num_state_rows, ) logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}") else: @@ -91,6 +93,7 @@ def __init__(self, config: EngineConfig, initialize_processor: bool = True): weight_load_mode=config.weight_load_mode, use_legacy_moe=config.use_legacy_moe, pre_transpose=config.pre_transpose, + enable_mtp=config.enable_mtp, ) if self.model_engine.model_type == "minicpm_eagle": @@ -107,7 +110,9 @@ def __init__(self, config: EngineConfig, initialize_processor: bool = True): ) self.speculative_runner = None - if config.draft_model_path is not None: + if config.enable_mtp: + self.speculative_runner = MTPRunner(config, self.model_engine) + elif config.draft_model_path is not None: self.speculative_runner = SpeculativeRunner( config, self.model_engine, self.device ) @@ -278,4 +283,8 @@ def close(self) -> None: self.pipeline_control.close() if self.kv_connector is not None: self.kv_connector.shutdown() + if self.speculative_runner is not None: + cache = getattr(self.speculative_runner, "prefix_cache", None) + if cache is not None: + cache.clear() self._closed = True diff --git a/python/infinilm/llm/model_runner/mtp_runner.py b/python/infinilm/llm/model_runner/mtp_runner.py new file mode 100644 index 000000000..d2d699112 --- /dev/null +++ b/python/infinilm/llm/model_runner/mtp_runner.py @@ -0,0 +1,468 @@ +"""Greedy Qwen MTP using the target engine's built-in head.""" + +import infinicore +from infinilm.llm.hybrid_prefix_cache import HybridPrefixCache +from infinilm.llm.request import MTPRequestState + + +class MTPRunner: + def __init__(self, config, target_model_engine): + hf_config = getattr(target_model_engine, "hf_config", {}) + if isinstance(hf_config, dict): + text_config = hf_config.get("text_config", hf_config) + if text_config.get("mtp_num_hidden_layers", 1) != 1: + raise ValueError("Built-in MTP requires a single MTP layer.") + self.engine = target_model_engine + self.block_size = config.block_size + self.num_draft_tokens = config.num_draft_tokens + # TP2 trials reduced draft GPU work but regressed end-to-end throughput. + # Keep the measured batching benefit limited to TP1 for now. + self.batch_draft = config.tensor_parallel_size == 1 + self.device_tokens = getattr(self.engine, "supports_device_mtp", False) + self.num_proposals = 0 + self.num_accepted = 0 + self.num_target_calls = 0 + self.num_capacity_fallbacks = 0 + budget = getattr(config, "mtp_prefix_cache_bytes", 0) + self.prefix_cache = ( + HybridPrefixCache(self.engine, budget, self.block_size) if budget else None + ) + + def _target(self, inputs, *, all_positions=False, device_verify=True): + self.num_target_calls += 1 + return self.engine.forward_raw( + **inputs, + sample_all_positions=all_positions, + return_logits=False, + verify_draft=all_positions and self.device_tokens and device_verify, + ) + + @staticmethod + def validate_request(req): + if req.has_multimodal_inputs or req.sampling_params.top_k != 1: + raise ValueError( + "Qwen MTP currently supports text-only greedy requests (`top_k=1`)." + ) + if ( + not isinstance(req.sampling_params.max_tokens, int) + or req.sampling_params.max_tokens < 1 + ): + raise ValueError("Qwen MTP requires a positive integer `max_tokens`.") + if req.get_prompt_length() < 1: + raise ValueError("Qwen MTP requires a nonempty tokenized prompt.") + + def _inputs(self, req, tokens, past, *, hidden=None, destinations=None): + def i32(values): + return infinicore.from_list(values, dtype=infinicore.int32) + + def i64(values): + return infinicore.from_list(values, dtype=infinicore.int64) + + device_ids = isinstance(tokens, infinicore.Tensor) + count = tokens.shape[-1] if device_ids else len(tokens) + # MTP consumes the next token's embedding with the current position's + # target hidden state. Its KV slots stay aligned with target slots. + position_start = past + int(hidden is not None) + positions = list(range(position_start, position_start + count)) + slots = [ + req.block_table[p // self.block_size] * self.block_size + + p % self.block_size + for p in range(past, past + count) + ] + return dict( + input_ids=tokens if device_ids else i64([tokens]), + position_ids=i64([positions] * self.engine.position_id_axes), + past_kv_lengths=i32([past]), + total_kv_lengths=i32([past + count]), + input_offsets=i32([0, count]), + cu_seqlens=i32([0, past + count]), + block_tables=i32([req.block_table]), + slot_mapping=i64(slots), + mamba_init_state_indices=i32([0 if past == 0 else req.mamba_cache_index]), + mamba_final_state_indices=i32( + [destinations[-1] if destinations else req.mamba_cache_index] + ), + token_state_indices=i32(destinations) if destinations else None, + target_hidden_states=hidden, + temperature=1.0, + top_k=1, + top_p=1.0, + ) + + @staticmethod + def _can_continue(req, output): + remaining = req.sampling_params.max_tokens + return ( + not req.is_aborted() + and req.get_num_generated_tokens() + len(output) < remaining + and ( + req.sampling_params.ignore_eos + or not any(token in req.eos_token_ids for token in output) + ) + ) + + def _draft(self, req, shifted_tokens, past, hidden, *, commit=True): + return self._draft_batch([(req, shifted_tokens, past, hidden)], commit=commit)[ + 0 + ] + + def forward(self, scheduler_output, model_input): + requests = scheduler_output.scheduled_requests + if not requests: + return [] + if scheduler_output.is_prefill: + return self._prefill( + requests, scheduler_output.speculative_cache_ops, model_input + ) + if len(requests) != 1: + return self._decode_batch(requests, scheduler_output.speculative_cache_ops) + req = requests[0] + cache_ops = scheduler_output.speculative_cache_ops + past, count, rows, inputs = self._prepare_verify(req, cache_ops) + if inputs is None: + return [self._ordinary(req, past, model_input)] + target = self._target(inputs, all_positions=True) + expected = list(map(int, target["output_ids"].to_numpy())) + if self.device_tokens: + accepted = target["accepted_draft_tokens"] + output = expected + else: + candidates = inputs["input_ids"].to_numpy()[0].tolist()[1:] + accepted = self._accepted(candidates, expected) + output = expected[: accepted + 1] + self._commit( + req, + cache_ops, + past, + count, + rows, + accepted, + output, + target["hidden_states"].narrow(1, 0, len(output)), + ) + return [output] + + @staticmethod + def _accepted(candidates, expected): + for index, (candidate, token) in enumerate(zip(candidates, expected)): + if candidate != token: + return index + return len(candidates) + + def _ordinary(self, req, past, inputs=None): + target = self._target( + inputs + if inputs is not None + else self._inputs(req, [req.generated_token_ids[-1]], past) + ) + output = [int(target["output_ids"].to_numpy()[-1])] + if req.mtp_state is not None and self._can_continue(req, output): + self._draft(req, output, past, target["hidden_states"]) + return output + + def _reserve_verify(self, req, cache_ops): + self.validate_request(req) + if cache_ops is None or req.mamba_cache_index is None: + raise RuntimeError( + "Qwen MTP requires scheduler-owned paged and recurrent caches." + ) + past, state = req.get_total_length() - 1, req.mtp_state + if state is not None and state.cached_tokens != past: + raise RuntimeError( + "Qwen MTP draft cache is not aligned with the committed target prefix." + ) + count = min( + self.num_draft_tokens, + req.sampling_params.max_tokens - req.get_num_generated_tokens() - 1, + ) + if state is None or state.draft_token is None or count < 1: + return past, 0, None + try: + table, _ = cache_ops.append_verify_slots( + list(req.block_table), req.get_total_length() + 1, count + ) + except RuntimeError: + # Allocation checks capacity before mutation. The pending token + # already has a scheduler-owned slot, so ordinary Decode can run. + self.num_capacity_fallbacks += 1 + return past, 0, None + req.block_table, req.num_blocks = table, len(table) + return past, count, state.scratch_indices[: count + 1] + + def _prepare_verify(self, req, cache_ops): + past, count, rows = self._reserve_verify(req, cache_ops) + if not count: + return past, 0, None, None + state = req.mtp_state + proposals, hidden = [state.draft_token], state.draft_hidden + for step in range(1, count): + ids = proposals[-1].view((1, 1)) if self.device_tokens else [proposals[-1]] + # Prefill filled slots `[0, past)`. The next draft writes slot `past`, + # with position `past+1`; leaving a gap would read unwritten KV. + token, hidden = self._draft(req, ids, past + step - 1, hidden, commit=False) + proposals.append(token) + return past, count, rows, self._verify_inputs(req, past, rows, proposals) + + def _verify_inputs(self, req, past, rows, proposals): + tokens = [req.generated_token_ids[-1], *proposals] + if self.device_tokens: + pending = infinicore.from_list( + [[req.generated_token_ids[-1]]], dtype=infinicore.int64 + ).to(proposals[0].device) + tokens = infinicore.cat( + [pending] + [t.view((1, 1)) for t in proposals], dim=1 + ) + return self._inputs(req, tokens, past, destinations=rows) + + def _commit( + self, req, cache_ops, past, count, rows, accepted, output, hidden, *, draft=True + ): + if not 0 <= accepted <= count or len(output) != accepted + 1: + raise RuntimeError( + "MTP verification did not return a valid acceptance length." + ) + state, old = req.mtp_state, req.mamba_cache_index + # Select the same checkpoint for Conv and GDN. The correction/bonus + # token is returned but has no target KV yet. + req.mamba_cache_index = rows[accepted] + state.scratch_indices = [ + row for row in (old, *state.scratch_indices) if row != req.mamba_cache_index + ] + req.block_table = cache_ops.rollback_to_length( + req.block_table, past + 1 + accepted + ) + req.num_blocks, req.slot_mapping = len(req.block_table), [] + continuing = self._can_continue(req, output) + if continuing and draft: + self._draft(req, output, past, hidden) + self.num_proposals += count + self.num_accepted += accepted + return continuing + + def _pack(self, items): + """Pack requests using existing paged metadata, preserving device token IDs.""" + if len(items) == 1: + return items[0] + fields = ( + "past_kv_lengths", + "total_kv_lengths", + "mamba_init_state_indices", + "mamba_final_state_indices", + ) + values = {key: [] for key in fields} + positions = [[] for _ in range(self.engine.position_id_axes)] + offsets, cu, slots, tables, destinations = [0], [0], [], [], [] + ids = [] + for item in items: + ids.append(item["input_ids"]) + offsets.append(offsets[-1] + ids[-1].shape[-1]) + for axis, row in zip(positions, item["position_ids"].to_numpy().tolist()): + axis.extend(row) + for key in fields: + values[key].extend(item[key].to_numpy().tolist()) + cu.append(cu[-1] + values["total_kv_lengths"][-1]) + slots.extend(item["slot_mapping"].to_numpy().tolist()) + tables.extend(item["block_tables"].to_numpy().tolist()) + if item["token_state_indices"] is not None: + destinations.extend(item["token_state_indices"].to_numpy().tolist()) + width = max(map(len, tables)) + device = ids[0].device + tokens = infinicore.cat([t.to(device) for t in ids], dim=1) + result = { + key: infinicore.from_list(value, dtype=infinicore.int32) + for key, value in values.items() + } + result.update( + input_ids=tokens, + position_ids=infinicore.from_list(positions, dtype=infinicore.int64), + input_offsets=infinicore.from_list(offsets, dtype=infinicore.int32), + cu_seqlens=infinicore.from_list(cu, dtype=infinicore.int32), + slot_mapping=infinicore.from_list(slots, dtype=infinicore.int64), + block_tables=infinicore.from_list( + [row + [0] * (width - len(row)) for row in tables], + dtype=infinicore.int32, + ), + token_state_indices=infinicore.from_list( + destinations, dtype=infinicore.int32 + ) + if destinations + else None, + top_k=1, + top_p=1.0, + temperature=1.0, + ) + if items[0].get("target_hidden_states") is not None: + result["target_hidden_states"] = infinicore.cat( + [item["target_hidden_states"] for item in items], dim=1 + ) + return result + + def _draft_batch(self, jobs, *, commit=True): + """Share the draft's projections across independent requests.""" + if not jobs: + return [] + if len(jobs) > 1 and not self.batch_draft: + return [self._draft(*job, commit=commit) for job in jobs] + inputs = self._pack( + [ + self._inputs(req, tokens, past, hidden=hidden) + for req, tokens, past, hidden in jobs + ] + ) + result = self.engine.forward_raw( + **inputs, + sample_all_positions=False, + return_logits=False, + return_device_tokens=self.device_tokens, + ) + sampled = result["output_ids"] + if not self.device_tokens: + sampled = list(map(int, sampled.to_numpy())) + outputs, offset = [], 0 + for index, (req, tokens, past, _) in enumerate(jobs): + count = ( + tokens.shape[-1] + if isinstance(tokens, infinicore.Tensor) + else len(tokens) + ) + offset += count + token = ( + sampled.narrow(0, index, 1) if self.device_tokens else sampled[index] + ) + hidden = result["hidden_states"].narrow(1, offset - 1, 1) + if commit: + req.mtp_state.draft_token = token + req.mtp_state.draft_hidden = hidden + req.mtp_state.cached_tokens = past + count + outputs.append((token, hidden)) + return outputs + + def _prefill(self, requests, cache_ops, model_input): + outputs, misses = {}, [] + for req in requests: + self.validate_request(req) + if cache_ops is None or req.mamba_cache_index is None: + raise RuntimeError( + "Qwen MTP requires scheduler-owned paged and recurrent caches." + ) + if req.num_local_cached_tokens or req.mtp_state is not None: + raise RuntimeError("Qwen MTP requires a fresh full prompt prefill.") + saved = self.prefix_cache.restore(req) if self.prefix_cache else None + if saved is not None: + if self._can_continue(req, [saved.pending]): + rows = cache_ops.allocate_state_rows(self.num_draft_tokens + 1) + if rows is not None: + req.mtp_state = MTPRequestState(rows) + # Requests own these small tensors too: evicting the + # snapshot must release all of its budgeted storage. + req.mtp_state.draft_hidden = self.prefix_cache._clone( + saved.hidden + ) + req.mtp_state.draft_token = ( + self.prefix_cache._clone(saved.proposal) + if isinstance(saved.proposal, infinicore.Tensor) + else saved.proposal + ) + req.mtp_state.cached_tokens = req.get_prompt_length() + else: + self.num_capacity_fallbacks += 1 + outputs[req] = [saved.pending] + else: + misses.append(req) + if misses: + # The processor already built this full-prompt batch. Only rebuild + # metadata when snapshot hits remove requests from the target batch. + inputs = ( + model_input + if len(misses) == len(requests) + else self._pack( + [self._inputs(req, list(req.prompt_token_ids), 0) for req in misses] + ) + ) + target = self._target(inputs) + sampled = list(map(int, target["output_ids"].to_numpy())) + offset = 0 + for req, pending in zip(misses, sampled): + hidden = target["hidden_states"].narrow( + 1, offset, req.get_prompt_length() + ) + offset += req.get_prompt_length() + if self._can_continue(req, [pending]): + rows = cache_ops.allocate_state_rows(self.num_draft_tokens + 1) + if rows is not None: + req.mtp_state = MTPRequestState(rows) + self._draft( + req, list(req.prompt_token_ids[1:]) + [pending], 0, hidden + ) + else: + self.num_capacity_fallbacks += 1 + if self.prefix_cache: + self.prefix_cache.save(req, pending) + outputs[req] = [pending] + return [outputs[req] for req in requests] + + def _decode_batch(self, requests, cache_ops): + outputs, prepared = {}, [] + proposals, hidden = {}, {} + for req in requests: + past, count, rows = self._reserve_verify(req, cache_ops) + if not count: + outputs[req] = self._ordinary(req, past) + else: + prepared.append((req, past, count, rows)) + proposals[req] = [req.mtp_state.draft_token] + hidden[req] = req.mtp_state.draft_hidden + for step in range(1, max((item[2] for item in prepared), default=0)): + active = [item for item in prepared if item[2] > step] + jobs = [ + ( + req, + proposals[req][-1].view((1, 1)) + if self.device_tokens + else [proposals[req][-1]], + past + step - 1, + hidden[req], + ) + for req, past, _, _ in active + ] + for (req, *_), (token, state) in zip( + active, self._draft_batch(jobs, commit=False) + ): + proposals[req].append(token) + hidden[req] = state + if prepared: + packed = self._pack( + [ + self._verify_inputs(req, past, rows, proposals[req]) + for req, past, _, rows in prepared + ] + ) + target = self._target(packed, all_positions=True, device_verify=False) + expected = list(map(int, target["output_ids"].to_numpy())) + candidates = packed["input_ids"].to_numpy()[0].tolist() + cursor, rebuild = 0, [] + for req, past, count, rows in prepared: + values = expected[cursor : cursor + count + 1] + accepted = self._accepted( + candidates[cursor + 1 : cursor + count + 1], values + ) + output = values[: accepted + 1] + state = target["hidden_states"].narrow(1, cursor, accepted + 1) + if self._commit( + req, + cache_ops, + past, + count, + rows, + accepted, + output, + state, + draft=False, + ): + rebuild.append((req, output, past, state)) + outputs[req] = output + cursor += count + 1 + # Different acceptance lengths become variable-length packed + # requests; their positions, KV tables and hidden states stay separate. + self._draft_batch(rebuild) + return [outputs[req] for req in requests] diff --git a/python/infinilm/llm/request.py b/python/infinilm/llm/request.py index f4612b816..d87f3d3f1 100644 --- a/python/infinilm/llm/request.py +++ b/python/infinilm/llm/request.py @@ -141,6 +141,16 @@ class TokenOutput: generated_text: str = "" +@dataclass +class MTPRequestState: + """Request-owned verification rows and the next built-in MTP proposal.""" + + scratch_indices: List[int] + draft_token: object = None # Device tensor when supported, otherwise an integer. + draft_hidden: object = None + cached_tokens: int = 0 + + class InferenceRequest: """Internal inference request object for managing generation state and resources.""" @@ -208,6 +218,7 @@ def __init__( # Mamba cache management. None means no mamba cache row is currently owned. self.mamba_cache_index: Optional[int] = None + self.mtp_state: Optional[MTPRequestState] = None # Qwen-style MRoPE decode offset. It is zero for pure text requests. self.mrope_position_delta: int = 0 diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index c10b55f2f..c2c3c6ef9 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -17,8 +17,17 @@ class SpeculativeCacheOps: """Limited cache operations needed by speculative verification.""" - def __init__(self, cache_manager: BlockManager): + def __init__(self, cache_manager: BlockManager, state_manager=None): self._cache_manager = cache_manager + self._state_manager = state_manager + + def allocate_state_rows(self, count: int) -> Optional[List[int]]: + if ( + self._state_manager is None + or self._state_manager.get_num_free_blocks() < count + ): + return None + return [self._state_manager.allocate() for _ in range(count)] def append_verify_slots( self, @@ -89,7 +98,9 @@ def __init__( if has_mamba_cache else None ) - self.speculative_cache_ops = SpeculativeCacheOps(self.cache_manager) + self.speculative_cache_ops = SpeculativeCacheOps( + self.cache_manager, self.mamba_cache_manager + ) self.block_size = block_size self.max_num_batched_tokens = max_num_batched_tokens self.connector = connector @@ -430,9 +441,19 @@ def complete_requests(self, requests: List[InferenceRequest]): self.cache_manager.free_blocks(req.block_table) elif req.block_table and delay_free_blocks: self.pending_free_blocks[req.request_id] = list(req.block_table) + # Ownership has been released or moved to `pending_free_blocks`. + # A repeated completion must not free pages reassigned to a + # different request in the meantime. + req.block_table = [] + req.slot_mapping = [] + req.num_blocks = 0 if self.mamba_cache_manager is not None: self.mamba_cache_manager.free(req.mamba_cache_index) req.mamba_cache_index = None + if req.mtp_state is not None: + for row in req.mtp_state.scratch_indices: + self.mamba_cache_manager.free(row) + req.mtp_state = None if req.status == RequestStatus.CANCELED: logger.info( @@ -450,6 +471,21 @@ def complete_requests(self, requests: List[InferenceRequest]): # Still running, put back in running queue self.running_queue.sync_q.put(req) + def cancel_all(self): + """Release queued requests after the engine's execution loop has stopped.""" + canceled = [] + for pending in (self.waiting_queue, self.running_queue): + while True: + try: + req = pending.sync_q.get_nowait() + except queue.Empty: + break + if not req.is_finished(): + req.mark_canceled() + self.complete_requests([req]) + canceled.append(req) + return canceled + def can_accept_request( self, request: InferenceRequest, diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index b210b0c56..81121cdb5 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -105,6 +105,7 @@ def load_state_dict( device="cpu", dtype=torch.bfloat16, preserve_fp32_suffixes: Tuple[str, ...] = (".e_score_correction_bias",), + preserve_fp8: bool = False, ) -> Dict[str, torch.Tensor]: """ Reads a `safetensor` checkpoint file. We load the checkpoint on "cpu" by default. @@ -128,8 +129,13 @@ def load_state_dict( for k in f.keys(): tensor = f.get_tensor(k) - preserve_fp32 = k.endswith(preserve_fp32_suffixes) - if tensor.is_floating_point() and not preserve_fp32: + if preserve_fp8 and k.endswith(".weight_scale_inv"): + state_dict[k] = tensor.to(device=device, dtype=torch.float32) + continue + preserve_dtype = k.endswith(preserve_fp32_suffixes) or ( + preserve_fp8 and tensor.dtype == torch.float8_e4m3fn + ) + if tensor.is_floating_point() and not preserve_dtype: tensor = tensor.to(device=device, dtype=dtype) else: tensor = tensor.to(device=device) @@ -201,6 +207,9 @@ def load_model_state_dict_by_file( t1 = time.time() model_type = model.hf_config.get("model_type", "") + preserve_fp8 = (model.hf_config.get("quantization_config") or {}).get( + "quant_method" + ) == "fp8" preserve_fp32_suffixes = (".e_score_correction_bias",) if model_type == "kimi_k3": preserve_fp32_suffixes += (".A_log", ".dt_bias") @@ -215,7 +224,6 @@ def load_model_state_dict_by_file( already_loaded_keys = [] embed_tokens_torch_unscaled = None - weights_processed = False remapper = _WEIGHT_REMAPPER.get(model_type) @@ -246,6 +254,7 @@ def load_model_state_dict_by_file( device=torch_device, dtype=torch_dtype, preserve_fp32_suffixes=preserve_fp32_suffixes, + preserve_fp8=preserve_fp8, ) # Apply model-specific weight remapping @@ -290,9 +299,6 @@ def load_model_state_dict_by_file( embed_tokens_torch_unscaled = None gc.collect() - model.process_weights_after_loading() - weights_processed = True - elif os.path.exists(os.path.join(model_path, "pytorch_model.bin")): file_path = os.path.join(model_path, "pytorch_model.bin") model_params = torch.load(file_path, weights_only=True, map_location="cpu") @@ -352,8 +358,8 @@ def load_model_state_dict_by_file( check_parameters(model_keys, already_loaded_keys) - if not weights_processed: - model.process_weights_after_loading() + # All weights, including a tied output head, must exist before packing/capture. + model.process_weights_after_loading() t2 = time.time() print(f" load weights over! {(t2 - t1) * 1000} ms \n") @@ -757,7 +763,8 @@ def _remap_videonsa(state_dict, config=None): # Model type → remap function mapping def _remap_qwen3_5(state_dict, config): """Apply Qwen3.5-specific load-time weight fixes.""" - state_dict = drop_keys(state_dict, ["mtp."]) + if not config.get("enable_mtp", False): + state_dict = drop_keys(state_dict, ["mtp."]) llm_config = config["text_config"] key_dim = llm_config["linear_key_head_dim"] * llm_config["linear_num_key_heads"] @@ -771,19 +778,32 @@ def _remap_qwen3_5(state_dict, config): to_drop = [] to_add = {} for key, tensor in state_dict.items(): - if key == "model.norm.weight" or key.endswith(norm_weight_suffixes): + if key in ( + "model.norm.weight", + "model.language_model.norm.weight", + "mtp.norm.weight", + "mtp.pre_fc_norm_embedding.weight", + "mtp.pre_fc_norm_hidden.weight", + ) or key.endswith(norm_weight_suffixes): state_dict[key] = tensor + torch.ones_like(tensor) - elif key.endswith("linear_attn.in_proj_qkv.weight"): - prefix = key[: -len("in_proj_qkv.weight")] - to_add[prefix + "in_proj_q.weight"] = state_dict[key][ - :key_dim, : - ].contiguous() - to_add[prefix + "in_proj_k.weight"] = state_dict[key][ - key_dim : key_dim * 2, : - ].contiguous() - to_add[prefix + "in_proj_v.weight"] = state_dict[key][ - key_dim * 2 :, : - ].contiguous() + elif key.endswith( + ( + "linear_attn.in_proj_qkv.weight", + "linear_attn.in_proj_qkv.weight_scale_inv", + ) + ): + suffix = key.rsplit(".", 1)[1] + prefix = key[: -len("in_proj_qkv." + suffix)] + split_dim = key_dim if suffix == "weight" else key_dim // 128 + for name, part in zip( + ("q", "k", "v"), + ( + tensor[:split_dim], + tensor[split_dim : 2 * split_dim], + tensor[2 * split_dim :], + ), + ): + to_add[prefix + "in_proj_" + name + "." + suffix] = part.contiguous() to_drop.append(key) state_dict = drop_keys(state_dict, to_drop) diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 462c31084..5af030b9f 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -125,6 +125,10 @@ def __init__( kv_transfer_config: Optional[KVTransferConfig] = None, enable_prefix_caching: bool = True, pre_transpose: bool = False, + enable_mtp: bool = False, + num_draft_tokens: int = 1, + num_state_rows: int = 0, + mtp_prefix_cache_bytes: int = 0, ): """Initialize inference server. @@ -188,6 +192,10 @@ def __init__( self.kv_transfer_config = kv_transfer_config self.enable_prefix_caching = enable_prefix_caching self.pre_transpose = pre_transpose + self.enable_mtp = enable_mtp + self.num_draft_tokens = num_draft_tokens + self.num_state_rows = num_state_rows + self.mtp_prefix_cache_bytes = mtp_prefix_cache_bytes self.engine: AsyncLLMEngine = None @@ -232,6 +240,10 @@ async def lifespan(app: FastAPI): kv_transfer_config=self.kv_transfer_config, enable_prefix_caching=self.enable_prefix_caching, pre_transpose=self.pre_transpose, + enable_mtp=self.enable_mtp, + num_draft_tokens=self.num_draft_tokens, + num_state_rows=self.num_state_rows, + mtp_prefix_cache_bytes=self.mtp_prefix_cache_bytes, ) self.engine.start() logger.info(f"Engine initialized with model at {self.model_path}") @@ -667,6 +679,10 @@ def main(): kv_transfer_config=kv_transfer_config, enable_prefix_caching=cfg.enable_prefix_caching, pre_transpose=cfg.pre_transpose, + enable_mtp=cfg.enable_mtp, + num_draft_tokens=cfg.num_draft_tokens, + num_state_rows=cfg.num_state_rows, + mtp_prefix_cache_bytes=cfg.mtp_prefix_cache_mib * 1024 * 1024, ) server.start() diff --git a/test/bench/backends/infinilm.py b/test/bench/backends/infinilm.py index fbd99379d..2dded8a27 100644 --- a/test/bench/backends/infinilm.py +++ b/test/bench/backends/infinilm.py @@ -17,8 +17,16 @@ def __init__( enable_paged_attn=False, enable_graph=False, attn_backend="default", + *, + enable_mtp=False, + num_draft_tokens=4, + num_state_rows=0, + mtp_prefix_cache_bytes=0, + num_blocks=128, + block_size=256, ): from infinilm import LLM + from infinilm.infer_engine import model_uses_mamba_cache super().__init__(benchmark) @@ -56,8 +64,14 @@ def __init__( tensor_parallel_size=tensor_parallel_size, cache_type="paged" if enable_paged_attn else "static", max_batch_size=1, - num_blocks=128, - block_size=256, + num_blocks=num_blocks, + block_size=block_size, + enable_mtp=enable_mtp, + num_draft_tokens=num_draft_tokens, + num_state_rows=num_state_rows, + mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, + enable_prefix_caching=bool(mtp_prefix_cache_bytes) + or not model_uses_mamba_cache(self.config_dict), enable_graph=enable_graph, attn_backend=attn_backend, ) @@ -92,5 +106,6 @@ def generate(self, *args, max_steps=500, topp_=1.0, topk_=1, temperature_=1.0): ) def destroy_model_instance(self): + self.model.close() del self.model print("InfiniLM model destroyed") diff --git a/test/bench/test_benchmark.py b/test/bench/test_benchmark.py index b77b11cde..cd9893c3e 100644 --- a/test/bench/test_benchmark.py +++ b/test/bench/test_benchmark.py @@ -613,6 +613,12 @@ def main(): cfg.enable_paged_attn, cfg.enable_graph, cfg.attn, + enable_mtp=cfg.enable_mtp, + num_draft_tokens=cfg.num_draft_tokens, + num_state_rows=cfg.num_state_rows, + mtp_prefix_cache_bytes=int(cfg.mtp_prefix_cache_mib * 1024**2), + num_blocks=cfg.num_blocks, + block_size=cfg.block_size, ) else: raise ValueError(f"Unsupported backend: {cfg.backend}") diff --git a/test/layers/test_pre_transpose.py b/test/layers/test_pre_transpose.py new file mode 100644 index 000000000..cd72c0fe3 --- /dev/null +++ b/test/layers/test_pre_transpose.py @@ -0,0 +1,78 @@ +"""Check fused projections sharing a quantization object after weight packing.""" + +import json + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="A CUDA device is required.") +def test_pre_transpose_and_reprocessing_preserve_logits(tmp_path): + import infinicore + from infinilm.cache import PagedKVCacheConfig + from infinilm.infer_engine import InferEngine + from infinilm.modeling_utils import load_model_state_dict_by_file + from safetensors.torch import save_file + + config = { + "model_type": "qwen2", + "hidden_size": 128, + "intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "head_dim": 64, + "vocab_size": 128, + "torch_dtype": "float16", + "rms_norm_eps": 1e-5, + "max_position_embeddings": 256, + "rope_theta": 10000, + "hidden_act": "silu", + "eos_token_id": 0, + "tie_word_embeddings": True, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + results = [] + for packed in (False, True): + engine = InferEngine( + str(tmp_path), + device=infinicore.device("cuda", 0), + cache_config=PagedKVCacheConfig(8, 256, 1), + attention_backend="paged-attn", + pre_transpose=packed, + enable_graph_compiling=True, + ) + generator = torch.Generator().manual_seed(318) + weights = {} + for name, parameter in sorted(engine.state_dict()[0].items()): + if name == "lm_head.weight": + continue + tensor = infinicore.Tensor(parameter) + values = torch.randn(tensor.shape, generator=generator) * 0.03 + if "norm" in name and name.endswith("weight"): + values.fill_(1) + values = values.to(infinicore.utils.to_torch_dtype(tensor.dtype)) + weights[name] = values + save_file(weights, tmp_path / "model.safetensors") + load_model_state_dict_by_file(engine, str(tmp_path), dtype=engine.dtype) + for _ in range(2): + engine.process_weights_after_loading() + output = engine.forward_raw( + infinicore.from_list([[17, 83, 51]], dtype=infinicore.int64), + position_ids=infinicore.from_list([0, 1, 2], dtype=infinicore.int64), + past_kv_lengths=infinicore.from_list([0], dtype=infinicore.int32), + total_kv_lengths=infinicore.from_list([3], dtype=infinicore.int32), + input_offsets=infinicore.from_list([0, 3], dtype=infinicore.int32), + cu_seqlens=infinicore.from_list([0, 3], dtype=infinicore.int32), + block_tables=infinicore.from_list([[0]], dtype=infinicore.int32), + slot_mapping=infinicore.from_list([0, 1, 2], dtype=infinicore.int64), + sample_all_positions=True, + )["logits"] + copied = torch.empty(output.shape, dtype=torch.float16) + infinicore.from_torch(copied).copy_(output) + infinicore.sync_device() + assert torch.isfinite(copied).all() + results.append(copied) + del engine + for actual in results[1:]: + torch.testing.assert_close(actual, results[0], atol=1e-3, rtol=1e-3) diff --git a/test/models/qwen3_5/test_mtp_execution.py b/test/models/qwen3_5/test_mtp_execution.py new file mode 100644 index 000000000..55b59fc51 --- /dev/null +++ b/test/models/qwen3_5/test_mtp_execution.py @@ -0,0 +1,206 @@ +"""GPU correctness: serial equivalence, checkpoint ownership, graphs and batching. + +Set INFINILM_QWEN_MTP_TEST_MODEL to a tiny checkpoint with MTP weights; +INFINILM_QWEN_MTP_TEST_TP=1 or 2 selects the parallelism (prefix caching uses TP1). +""" + +import os + +import infinicore +import pytest +from infinilm.cache import PagedKVCacheConfig +from infinilm.llm.llm import LLM +from infinilm.llm.request import InferenceRequest +from infinilm.llm.sampling_params import SamplingParams + + +def create(**options): + path = os.environ.get("INFINILM_QWEN_MTP_TEST_MODEL") + if not path: + pytest.skip("Set INFINILM_QWEN_MTP_TEST_MODEL to a tiny GPU checkpoint") + config = dict( + enable_mtp=True, + device="cuda", + dtype="bfloat16", + tensor_parallel_size=int(os.environ.get("INFINILM_QWEN_MTP_TEST_TP", "1")), + num_blocks=40, + block_size=64, + num_state_rows=13, + max_batch_size=1, + enable_prefix_caching=False, + attn_backend="paged-attn", + top_k=1, + top_p=1.0, + weight_load_mode="sync", + ) + return LLM(path, **(config | options)) + + +def request(name, prompt, count=17): + return InferenceRequest( + name, + prompt_token_ids=prompt, + sampling_params=SamplingParams(max_tokens=count, ignore_eos=True, top_k=1), + ) + + +def drain(engine, reqs, *, admit=True, recapture=False): + if admit: + for req in reqs: + engine.add_request(req) + for step in range(100): + if all(r.is_finished() for r in reqs): + break + assert engine.step()[0] + if recapture and step == 1: + engine.model_runner.model_engine.compile() + assert all(r.is_finished() for r in reqs) + assert not engine.scheduler.mamba_cache_manager.used_block_ids + cache = engine.scheduler.cache_manager + assert all(b.ref_count == 0 for b in cache.blocks) + assert cache.get_total_usable_blocks() == cache.num_blocks + return [list(r.generated_token_ids) for r in reqs] + + +@pytest.mark.parametrize("candidates,graph", [(1, True), (2, False), (4, False)]) +def test_mtp_matches_ordinary_with_reused_and_recaptured_state(candidates, graph): + llm = create(num_draft_tokens=candidates, enable_graph=graph) + engine = llm.engine + runner = engine.model_runner + mtp = runner.speculative_runner + raw = runner.model_engine + + prompt = [i % 63 + 1 for i in range(63)] + + def generate(recapture=False): + return drain(engine, [request("test", prompt, 20)], recapture=recapture)[0] + + try: + runner.speculative_runner = None + engine.config.enable_mtp = False + expected = generate() + runner.speculative_runner = mtp + engine.config.enable_mtp = True + mtp.device_tokens = False + assert generate() == expected + mtp.device_tokens = True + assert generate(recapture=graph) == expected + raw.reset_cache(PagedKVCacheConfig(40, 64, 1, llm.config.num_state_rows)) + assert generate() == expected + # Exercise every device-side acceptance length, then consume the + # selected Conv/GDN checkpoint with ordinary Q=1 continuation. + forward = raw.forward_raw + for prefix in range(candidates + 1): + verified = [] + + def force_prefix(**inputs): + if inputs.get("verify_draft"): + ids = expected[: candidates + 1].copy() + if prefix < candidates: + ids[prefix + 1] = (ids[prefix + 1] + 1) % 63 + inputs["input_ids"] = infinicore.from_list( + [ids], dtype=infinicore.int64 + ).to(inputs["input_ids"].device) + result = forward(**inputs) + if inputs.get("verify_draft"): + verified.append(result["accepted_draft_tokens"]) + return result + + req = request("test", prompt, 20) + engine.add_request(req) + engine.step() + raw.forward_raw = force_prefix + try: + engine.step() + finally: + raw.forward_raw = forward + assert verified == [prefix] + runner.speculative_runner = None + engine.config.enable_mtp = False + assert drain(engine, [req], admit=False)[0] == expected + runner.speculative_runner = mtp + engine.config.enable_mtp = True + finally: + llm.close() + + +def test_bounded_prefix_hit_eviction_and_rebuild(): + llm = create( + num_draft_tokens=1, + enable_graph=True, + tensor_parallel_size=1, + enable_prefix_caching=True, + mtp_prefix_cache_bytes=8 * 1024**2, + ) + engine = llm.engine + runner = engine.model_runner + mtp = runner.speculative_runner + prompt = [i % 59 + 1 for i in range(63)] + try: + runner.speculative_runner = None + engine.config.enable_mtp = False + expected = drain(engine, [request("ordinary", prompt)])[0] + runner.speculative_runner = mtp + engine.config.enable_mtp = True + assert drain(engine, [request("cold", prompt)])[0] == expected + cache = mtp.prefix_cache + assert cache.entries and cache.used_bytes <= cache.budget_bytes + cache.budget_bytes = cache.used_bytes + hits = cache.hits + calls = [] + forward = runner.model_engine.forward_raw + + def capture(**kw): + calls.append(kw["input_ids"].shape[-1]) + return forward(**kw) + + runner.model_engine.forward_raw = capture + assert drain(engine, [request("hit", prompt)])[0] == expected + assert cache.hits == hits + 1 and max(calls) <= 2 + runner.model_engine.forward_raw = forward + drain(engine, [request("different", [7, 2, 19, 8])]) + assert cache.evictions == 1 and cache.used_bytes <= cache.budget_bytes + assert drain(engine, [request("evicted", prompt)])[0] == expected + assert cache.evictions == 2 + generation = runner.model_engine.cache_generation + runner.model_engine.reset_cache(PagedKVCacheConfig(40, 64, 1, 13)) + assert runner.model_engine.cache_generation > generation + hits = cache.hits + assert drain(engine, [request("reset", prompt)])[0] == expected + assert cache.hits == hits + finally: + llm.close() + assert not mtp.prefix_cache.entries + + +def test_batched_mtp_matches_serial_with_cancellation(): + llm = create( + num_draft_tokens=2, + max_batch_size=3, + enable_graph=False, + ) + engine = llm.engine + runner = engine.model_runner + mtp = runner.speculative_runner + prompts = [[i % 59 + 1 for i in range(63)], [7, 2, 19, 8], [19, 9, 13, 5]] + try: + runner.speculative_runner = None + engine.config.enable_mtp = False + expected = [drain(engine, [request("serial", p)])[0] for p in prompts] + runner.speculative_runner = mtp + engine.config.enable_mtp = True + reqs = [request(str(i), p) for i, p in enumerate(prompts)] + assert drain(engine, reqs) == expected + cancel = request("cancel", prompts[0]) + keep = request("keep", prompts[1]) + engine.add_request(cancel) + engine.add_request(keep) + engine.step() + engine.step() + cancel.mark_canceled() + other = request("new", prompts[2]) + engine.add_request(other) + actual = drain(engine, [cancel, keep, other], admit=False) + assert actual[1:] == expected[1:] + finally: + llm.close() diff --git a/test/models/qwen3_5/test_mtp_model.py b/test/models/qwen3_5/test_mtp_model.py new file mode 100644 index 000000000..ff1b963af --- /dev/null +++ b/test/models/qwen3_5/test_mtp_model.py @@ -0,0 +1,284 @@ +"""Qwen checkpoint loading and GPU model contracts used by MTP. + +CPU loading tests run unconditionally. GPU checks require +INFINILM_QWEN_MTP_TEST_MODEL and optionally INFINILM_QWEN_MTP_TEST_TP=2. +""" + +import gc +import json +import os +from pathlib import Path +from types import SimpleNamespace + +import infinicore +import pytest +import torch +from infinilm.cache import PagedKVCacheConfig +from infinilm.distributed import DistConfig +from infinilm.infer_engine import InferEngine +from infinilm.llm.model_runner.mtp_runner import MTPRunner +from infinilm.modeling_utils import ( + _remap_qwen3_5, + load_model_state_dict_by_file, + load_state_dict, +) +from safetensors.torch import load_file, save_file + + +def test_mtp_is_opt_in_and_gemma_norms_are_converted(): + config = {"text_config": {"linear_key_head_dim": 2, "linear_num_key_heads": 1}} + norms = [ + "model.language_model.norm.weight", + "mtp.norm.weight", + "mtp.pre_fc_norm_embedding.weight", + "mtp.pre_fc_norm_hidden.weight", + "mtp.layers.0.input_layernorm.weight", + "mtp.layers.0.post_attention_layernorm.weight", + "mtp.layers.0.self_attn.q_norm.weight", + "mtp.layers.0.self_attn.k_norm.weight", + ] + gdn_norm = "model.language_model.layers.0.linear_attn.norm.weight" + weights = {key: torch.tensor([0.25, -0.5]) for key in norms} + weights[gdn_norm] = torch.tensor([1.0, 0.75]) + weights["mtp.fc.weight"] = torch.arange(8).reshape(2, 4).float() + + disabled = _remap_qwen3_5(dict(weights), config) + assert not any(key.startswith("mtp.") for key in disabled) + enabled = _remap_qwen3_5(dict(weights), {**config, "enable_mtp": True}) + for key in norms: + torch.testing.assert_close(enabled[key], weights[key] + 1) + for key in (gdn_norm, "mtp.fc.weight"): + torch.testing.assert_close(enabled[key], weights[key]) + + +@pytest.mark.parametrize("scale_dtype", [torch.float32, torch.bfloat16]) +def test_fp8_loading_preserves_scales_and_qkv_block_alignment(tmp_path, scale_dtype): + prefix = "model.language_model.layers.0.linear_attn." + weight = ( + torch.arange(1024 * 256) + .reshape(1024, 256) + .remainder(31) + .float() + .to(torch.float8_e4m3fn) + ) + scales = (torch.arange(16).reshape(8, 2).float() / 1000 + 0.00012345).to( + scale_dtype + ) + path = tmp_path / "model.safetensors" + save_file( + { + prefix + "in_proj_qkv.weight": weight, + prefix + "in_proj_qkv.weight_scale_inv": scales, + }, + path, + ) + loaded = load_state_dict(str(path), preserve_fp8=True) + config = {"text_config": {"linear_key_head_dim": 128, "linear_num_key_heads": 2}} + mapped = _remap_qwen3_5(loaded, config) + for name, start, end in (("q", 0, 256), ("k", 256, 512), ("v", 512, 1024)): + key = prefix + "in_proj_" + name + assert mapped[key + ".weight"].dtype == torch.float8_e4m3fn + assert mapped[key + ".weight_scale_inv"].dtype == torch.float32 + torch.testing.assert_close( + mapped[key + ".weight"].float(), weight[start:end].float(), rtol=0, atol=0 + ) + torch.testing.assert_close( + mapped[key + ".weight_scale_inv"], + scales[start // 128 : end // 128].float(), + rtol=0, + atol=0, + ) + + +@pytest.fixture +def checkpoint(): + path = os.environ.get("INFINILM_QWEN_MTP_TEST_MODEL") + if not path: + pytest.skip("Set INFINILM_QWEN_MTP_TEST_MODEL to a tiny GPU checkpoint") + return path + + +def create_model(path): + tp = int(os.environ.get("INFINILM_QWEN_MTP_TEST_TP", "1")) + model = InferEngine( + path, + device=infinicore.device("cuda", 0), + distributed_config=DistConfig(tp), + cache_config=PagedKVCacheConfig(16, 64, 1), + enable_mtp=True, + attention_backend="paged-attn", + ) + load_model_state_dict_by_file(model, path, dtype=model.dtype) + return model, tp + + +@pytest.fixture +def engine(checkpoint): + # Malformed raw inputs can stop the workers: never share this engine across tests. + return create_model(checkpoint) + + +def draft_inputs(model, hidden, ids, past=0): + # from_torch currently assumes device 0 and contiguous storage. + tensor = infinicore.strided_from_blob( + hidden.data_ptr(), + list(hidden.shape), + list(hidden.stride()), + dtype=infinicore.utils.to_infinicore_dtype(hidden.dtype), + device=infinicore.device(hidden.device.type, hidden.device.index or 0), + ) + runner = MTPRunner( + SimpleNamespace( + block_size=64, + num_draft_tokens=1, + tensor_parallel_size=len(model.distributed_config.tp_device_ids), + ), + model, + ) + req = SimpleNamespace(block_table=[0], mamba_cache_index=1) + return runner._inputs(req, ids, past, hidden=tensor) + + +def draft(model, hidden, ids, past=0): + result = model.forward_raw( + **draft_inputs(model, hidden, ids, past), sample_all_positions=True + ) + logits = torch.empty(result["logits"].shape, dtype=hidden.dtype) + infinicore.from_torch(logits).copy_(result["logits"]) + return logits + + +def test_tp_teardown_preserves_calling_device(checkpoint): + if int(os.environ.get("INFINILM_QWEN_MTP_TEST_TP", "1")) < 2: + pytest.skip("Communicator teardown requires TP2") + torch.cuda.set_device(0) + model, _ = create_model(checkpoint) + previous = torch.cuda.current_device() + del model + gc.collect() + assert torch.cuda.current_device() == previous + + +@pytest.mark.parametrize("strided", [False, True]) +def test_cpu_and_each_tp_device_produce_identical_draft_logits(engine, strided): + model, tp = engine + width = model.hf_config["text_config"]["hidden_size"] + hidden = torch.randn(1, 3, width, generator=torch.Generator().manual_seed(71)) + hidden = hidden.to(infinicore.utils.to_torch_dtype(model.dtype)) + expected = draft(model, hidden, [3, 8, 15]) + for rank in range(tp): + device_hidden = hidden.to(f"cuda:{rank}") + if strided: + storage = torch.empty( + 1, 3, 2 * width, dtype=hidden.dtype, device=device_hidden.device + ) + storage[..., ::2] = device_hidden + device_hidden = storage[..., ::2] + torch.cuda.synchronize(rank) + actual = draft(model, device_hidden, [3, 8, 15]) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + # The next call must also reuse valid MTP history on every rank. + actual_tail = draft(model, device_hidden[:, 2:], [3], past=3) + draft(model, hidden, [3, 8, 15]) + expected_tail = draft(model, hidden[:, 2:].contiguous(), [3], past=3) + torch.testing.assert_close(actual_tail, expected_tail, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "offsets,device_verify,error", + [ + ([0, 3, 2], False, "checkpoint offsets must be increasing"), + ([0, 1, 2], True, "one greedy"), + ], +) +def test_invalid_checkpoint_boundaries_are_rejected( + engine, offsets, device_verify, error, capfd +): + model, _ = engine + + def i32(values): + return infinicore.from_list(values, dtype=infinicore.int32) + + def i64(values): + return infinicore.from_list(values, dtype=infinicore.int64) + + # Two Q=1 requests bypass prefill's offset validator. The checkpoint + # validator must reject an out-of-range offset before indexing destinations. + # Worker-side validation is logged before RankWorker.wait() reports shutdown; + # acceptance-shape validation runs on the caller and raises its own message. + raised_error = error if device_verify else "RankWorker stopped during run" + with pytest.raises((ValueError, RuntimeError), match=raised_error): + model.forward_raw( + input_ids=i64([[3, 8]]), + position_ids=i64([[0, 0]] * model.position_id_axes), + past_kv_lengths=i32([0, 0]), + total_kv_lengths=i32([1, 1]), + input_offsets=i32(offsets), + cu_seqlens=i32([0, 1, 2]), + block_tables=i32([[0], [1]]), + slot_mapping=i64([0, 64]), + mamba_init_state_indices=i32([0, 0]), + mamba_final_state_indices=i32([1, 2]), + token_state_indices=i32([1, 2]), + sample_all_positions=True, + verify_draft=device_verify, + top_k=1, + ) + if not device_verify: + captured = capfd.readouterr() + assert error in captured.out + captured.err + + +@pytest.mark.parametrize("odd_vocab", [False, True]) +def test_vocab_logits_and_local_argmax(checkpoint, tmp_path, odd_vocab): + path = checkpoint + if odd_vocab: + original = Path(path) + config = json.loads((original / "config.json").read_text()) + config["text_config"]["vocab_size"] += 1 + (tmp_path / "config.json").write_text(json.dumps(config)) + weights = load_file(original / "model.safetensors") + for key in ("lm_head.weight", "model.language_model.embed_tokens.weight"): + weights[key] = torch.cat([weights[key], weights[key][-1:]], dim=0) + save_file(weights, tmp_path / "model.safetensors") + path = str(tmp_path) + model, tp = create_model(path) + width = model.hf_config["text_config"]["hidden_size"] + vocab = model.hf_config["text_config"]["vocab_size"] + dtype = infinicore.utils.to_torch_dtype(model.dtype) + hidden = torch.randn(1, 3, width, generator=torch.Generator().manual_seed(17)).to( + dtype + ) + inputs = draft_inputs(model, hidden, [3, 8, 15]) + full = model.forward_raw(**inputs, sample_all_positions=True) + fast = model.forward_raw(**inputs, sample_all_positions=True, return_logits=False) + assert ( + fast["output_ids"].to_numpy().tolist() == full["output_ids"].to_numpy().tolist() + ) + last = model.forward_raw(**inputs, sample_all_positions=False, return_logits=False) + assert ( + last["output_ids"].to_numpy().tolist() + == full["output_ids"].to_numpy().tolist()[-1:] + ) + + # Reuse the same hidden vector to force maxima at shard boundaries, then + # equal maxima on different ranks. The lowest global token ID must win. + hidden_out = torch.empty(full["hidden_states"].shape, dtype=dtype) + infinicore.from_torch(hidden_out).copy_(full["hidden_states"]) + weights = torch.zeros(vocab, width, dtype=dtype) + for candidates in ((vocab - 1,), (vocab // tp - 1,), (1, vocab - 1)): + weights.zero_() + for token in candidates: + weights[token] = hidden_out[0, -1] + model.load_state_dict( + {"lm_head.weight": infinicore.from_torch(weights)}, strict=False + ) + full = model.forward_raw(**inputs, sample_all_positions=True) + fast = model.forward_raw( + **inputs, sample_all_positions=True, return_logits=False + ) + assert ( + full["output_ids"].to_numpy().tolist() + == fast["output_ids"].to_numpy().tolist() + ) + assert int(fast["output_ids"].to_numpy()[-1]) == min(candidates) diff --git a/test/models/qwen3_5/test_mtp_runner.py b/test/models/qwen3_5/test_mtp_runner.py new file mode 100644 index 000000000..ec793f7aa --- /dev/null +++ b/test/models/qwen3_5/test_mtp_runner.py @@ -0,0 +1,435 @@ +"""CPU control-flow checks using real scheduling, pages and request updates. + +The deterministic engine below checks orchestration, not GPU numerical accuracy. +""" + +from types import SimpleNamespace + +import infinicore +import pytest +from infinilm.config.engine_config import EngineConfig +from infinilm.llm.llm import LLMEngine +from infinilm.llm.model_runner.model_runner import ModelRunner +from infinilm.llm.model_runner.mtp_runner import MTPRunner +from infinilm.llm.request import InferenceRequest, RequestStatus +from infinilm.llm.sampling_params import SamplingParams +from infinilm.llm.scheduler import Scheduler +from infinilm.processors.qwen3_5_processor import Qwen35Processor + + +def advance(state, token): + return (3 * state + token) % 97 + + +def serial_tokens(prompt, count): + state = 0 + for token in prompt: + state = advance(state, token) + output = [] + for _ in range(count): + token = (state + 13) % 97 + output.append(token) + state = advance(state, token) + return output + + +class DeterministicEngine: + position_id_axes = 3 + + def __init__(self, policy): + self.rows = {0: 0} + self.policy = policy + self.draft_calls = 0 + + def forward_raw(self, **inputs): + ids = inputs["input_ids"].to_numpy()[0].tolist() + hidden = inputs.get("target_hidden_states") + states = [] + if hidden is None: + offsets = inputs["input_offsets"].to_numpy().tolist() + initial = inputs["mamba_init_state_indices"].to_numpy().tolist() + finals = inputs["mamba_final_state_indices"].to_numpy().tolist() + destinations = inputs.get("token_state_indices") + destinations = ( + destinations.to_numpy().tolist() if destinations is not None else None + ) + sampled = [] + for r, (begin, end) in enumerate(zip(offsets, offsets[1:])): + state = self.rows[initial[r]] + for i in range(begin, end): + state = advance(state, ids[i]) + states.append(state) + if destinations is not None: + self.rows[destinations[i]] = state + self.rows[finals[r]] = state + if inputs.get("sample_all_positions", False): + sampled = [(state + 13) % 97 for state in states] + else: + sampled = [(states[end - 1] + 13) % 97 for end in offsets[1:]] + else: + self.draft_calls += 1 + previous = hidden.to_numpy()[0, :, 0].tolist() + states = [advance(int(state), token) for state, token in zip(previous, ids)] + reject = ( + (isinstance(self.policy, int) and self.draft_calls == self.policy + 1) + or self.policy == "reject" + or (self.policy == "alternate" and self.draft_calls % 2 == 0) + ) + sampled = [(state + 13 + int(reject)) % 97 for state in states] + if hidden is not None and not inputs.get("sample_all_positions", False): + offsets = inputs["input_offsets"].to_numpy().tolist() + sampled = [sampled[end - 1] for end in offsets[1:]] + return { + "output_ids": infinicore.from_list(sampled, dtype=infinicore.int64), + "hidden_states": infinicore.from_list( + [[state] for state in states], dtype=infinicore.float32 + ).view((1, len(states), 1)), + } + + +def service(policy="alternate", candidates=1, *, batch_size=1, state_rows=0): + config = EngineConfig( + "unused", + enable_mtp=True, + num_draft_tokens=candidates, + max_batch_size=batch_size, + num_state_rows=state_rows, + enable_prefix_caching=False, + num_blocks=max(16, 4 * (candidates + 3)), + block_size=4, + ) + engine = DeterministicEngine(policy) + runner = ModelRunner.__new__(ModelRunner) + runner.config = config + runner.model_engine = engine + runner.speculative_runner = MTPRunner(config, engine) + runner.processor = Qwen35Processor.__new__(Qwen35Processor) + runner.kv_connector = None + runner.pipeline_control = None + result = LLMEngine.__new__(LLMEngine) + result.config = config + result.model_runner = runner + result.scheduler = Scheduler( + max_batch_size=config.max_batch_size, + num_blocks=config.num_blocks, + block_size=config.block_size, + num_mamba_cache_blocks=config.num_state_rows, + has_mamba_cache=True, + enable_prefix_caching=False, + ) + result.tokenizer = SimpleNamespace( + decode=lambda ids: "".join(chr(0x4E00 + token) for token in ids) + ) + result.eos_token_ids = [] + return result + + +def request(name="a", prompt=None, **sampling): + return InferenceRequest( + name, + prompt_token_ids=prompt or [3, 8, 15], + sampling_params=SamplingParams(max_tokens=9, **sampling), + ) + + +def drain(engine, req): + while not req.is_finished(): + assert engine.step()[0] + return list(req.generated_token_ids) + + +def assert_released(engine): + assert not engine.scheduler.mamba_cache_manager.used_block_ids + manager = engine.scheduler.cache_manager + assert all(block.ref_count == 0 for block in manager.blocks) + assert manager.get_total_usable_blocks() == manager.num_blocks + + +@pytest.mark.parametrize("candidates,policy", [(2, "reject"), (4, "alternate")]) +def test_serial_equivalence_across_page_boundaries(candidates, policy): + engine, req = service(policy, candidates), request() + engine.add_request(req) + while not req.is_finished(): + engine.step() + if not req.is_finished(): + assert req.mtp_state.cached_tokens == req.get_total_length() - 1 + assert list(req.generated_token_ids) == serial_tokens(req.prompt_token_ids, 9) + assert_released(engine) + + +@pytest.mark.parametrize("stop_kind", ["eos", "string", "length", "prefill"]) +def test_stopping_truncates_output_and_releases_state(stop_kind): + engine, req = service("accept", candidates=4), request() + expected = serial_tokens(req.prompt_token_ids, 9) + limit = 1 if stop_kind == "prefill" else 2 + if stop_kind == "eos": + req.eos_token_ids = [expected[1]] + elif stop_kind == "string": + req.sampling_params.stop = [chr(0x4E00 + expected[1])] + else: + req.sampling_params.max_tokens = limit + output_queue = req.output_queue + engine.add_request(req) + outputs = [] + while not req.is_finished(): + _, pending = engine.step() + outputs.extend(output for _, output in pending) + assert list(req.generated_token_ids) == expected[:limit] + assert [output.token_id for output in outputs] == expected[:limit] + assert [output.finished for output in outputs] == [False] * (limit - 1) + [True] + assert_released(engine) + output_queue.close() + + +def test_kv_capacity_falls_back_then_resumes_mtp(): + engine, req = service("accept", candidates=2), request() + engine.add_request(req) + engine.step() + manager = engine.scheduler.cache_manager + occupied, _ = manager.allocate_slots(manager.get_num_free_blocks() * 4) + engine.step() # pending is in the last slot; speculative tail needs a page + assert engine.model_runner.speculative_runner.num_capacity_fallbacks == 1 + manager.free_blocks(occupied) + assert drain(engine, req) == serial_tokens(req.prompt_token_ids, 9) + assert engine.model_runner.speculative_runner.num_proposals > 0 + assert_released(engine) + + +def test_state_capacity_uses_normal_decode_without_stale_draft_cache(): + engine, req = service(candidates=2), request() + manager = engine.scheduler.mamba_cache_manager + occupied = [manager.allocate(), manager.allocate()] + engine.add_request(req) + assert drain(engine, req) == serial_tokens(req.prompt_token_ids, 9) + assert engine.model_runner.model_engine.draft_calls == 0 + for row in occupied: + manager.free(row) + assert_released(engine) + + +@pytest.mark.parametrize( + "option", + [ + {"enable_prefix_caching": True}, + {"num_draft_tokens": 5}, + {"num_draft_tokens": 2, "enable_graph": True}, + {"max_batch_size": 2, "enable_graph": True}, + {"pipeline_parallel_size": 2}, + {"cache_type": "static"}, + {"draft_model_path": "external"}, + {"top_k": 8}, + {"attn_backend": "flash-attn"}, + ], +) +def test_unsupported_engine_modes_fail_before_loading(option): + options = dict( + enable_mtp=True, + num_draft_tokens=1, + max_batch_size=1, + enable_prefix_caching=False, + ) + options.update(option) + with pytest.raises(ValueError): + EngineConfig("unused", **options) + + +@pytest.mark.parametrize( + "candidates,prefix", [(k, n) for k in (1, 2, 4) for n in range(k + 1)] +) +def test_every_acceptance_prefix_selects_checkpoint_without_target_replay( + candidates, prefix +): + engine, req = service(prefix, candidates), request() + engine.add_request(req) + engine.step() + rows = list(req.mtp_state.scratch_indices) + runner = engine.model_runner.speculative_runner + engine.step() + assert runner.num_accepted == prefix + assert runner.num_proposals == candidates + assert runner.num_target_calls == 2 + assert req.mamba_cache_index == rows[prefix] + assert drain(engine, req) == serial_tokens(req.prompt_token_ids, 9) + assert_released(engine) + + +@pytest.mark.parametrize("small_pool", [False, True]) +def test_packed_requests_arrival_cancel_and_capacity_fallback(small_pool): + engine = service( + "alternate", candidates=2, batch_size=3, state_rows=5 if small_pool else 0 + ) + requests = [ + request("a"), + request("b", prompt=[7, 2, 10]), + request("c", prompt=[19, 9, 13]), + ] + engine.add_request(requests[0]) + engine.add_request(requests[1]) + engine.step() + engine.step() + requests[0].mark_canceled() + engine.add_request(requests[2]) + for _ in range(40): + if all(r.is_finished() for r in requests): + break + assert engine.step()[0] + for req in requests[1:]: + assert list(req.generated_token_ids) == serial_tokens(req.prompt_token_ids, 9) + assert_released(engine) + if small_pool: + assert engine.model_runner.speculative_runner.num_capacity_fallbacks > 0 + + +def test_shutdown_reclaims_running_and_waiting_requests_once(): + engine = service(candidates=4) + engine.model_runner._closed = False + active, waiting = request("active"), request("waiting") + engine.add_request(active) + engine.step() + engine.step() # rotate checkpoint ownership + engine.add_request(waiting) + engine.close() + engine.close() + assert active.status == waiting.status == RequestStatus.CANCELED + assert engine.scheduler.waiting_queue.sync_q.empty() + assert engine.scheduler.running_queue.sync_q.empty() + assert_released(engine) + # Repeat completion after those pages have a different owner. + blocks, _ = engine.scheduler.cache_manager.allocate_slots(4) + engine.scheduler.complete_requests([active]) + assert all(engine.scheduler.cache_manager.blocks[b].ref_count == 1 for b in blocks) + engine.scheduler.cache_manager.free_blocks(blocks) + with pytest.raises(RuntimeError, match="closed"): + engine.add_request(request("after-close")) + + +@pytest.mark.parametrize("failure", ["prefill", "verify", "draft"]) +def test_packed_failure_releases_all_request_ownership(failure): + engine = service(candidates=2, batch_size=2) + requests = [request("a"), request("b")] + outputs = [req.output_queue for req in requests] + for req in requests: + engine.add_request(req) + if failure != "prefill": + engine.step() + forward = engine.model_runner.model_engine.forward_raw + + def fail(**kwargs): + if failure == "prefill" or ( + failure == "verify" and kwargs.get("sample_all_positions") + ): + raise RuntimeError("packed failure") + if failure == "draft" and kwargs.get("target_hidden_states") is not None: + raise RuntimeError("packed failure") + return forward(**kwargs) + + engine.model_runner.model_engine.forward_raw = fail + with pytest.raises(RuntimeError, match="packed failure"): + engine.step() + for req, queue in zip(requests, outputs): + assert req.status == RequestStatus.FAILED + output = queue.sync_q.get_nowait() + assert output.finished and output.token_id == -1 + assert output.finish_reason == req.finish_reason + req.output_queue.close() + assert_released(engine) + + +def test_mtp_config_sizes_state_pool_independently_from_pages(): + config = EngineConfig( + "unused", + enable_mtp=True, + num_draft_tokens=4, + max_batch_size=3, + num_blocks=512, + enable_prefix_caching=False, + ) + assert config.num_state_rows == 19 + config = EngineConfig( + "unused", + enable_mtp=True, + num_state_rows=5, + enable_prefix_caching=True, + mtp_prefix_cache_bytes=1024, + ) + assert config.num_state_rows == 5 + with pytest.raises(ValueError, match="tensor_parallel_size"): + EngineConfig( + "unused", + enable_mtp=True, + tensor_parallel_size=2, + mtp_prefix_cache_bytes=1024, + ) + + +def test_packed_draft_preserves_different_acceptance_lengths(): + engine = service("accept", candidates=2, batch_size=2) + raw = engine.model_runner.model_engine + original = raw.forward_raw + packed_lengths = [] + + def traced(**inputs): + result = original(**inputs) + if inputs.get("target_hidden_states") is not None: + offsets = inputs["input_offsets"].to_numpy().tolist() + if len(offsets) > 2: + packed_lengths.append([b - a for a, b in zip(offsets, offsets[1:])]) + # Reject only the first request's rolled candidate. Rebuilding + # the histories must then pack two different accepted lengths. + tokens = result["output_ids"].to_numpy().tolist() + tokens[0] = (tokens[0] + 1) % 97 + result["output_ids"] = infinicore.from_list( + tokens, dtype=infinicore.int64 + ) + return result + + raw.forward_raw = traced + # Caller-supplied IDs need not be unique; ownership follows the request. + reqs = [request("same"), request("same", prompt=[7, 2, 10])] + for req in reqs: + engine.add_request(req) + while not all(req.is_finished() for req in reqs): + engine.step() + assert [1, 1] in packed_lengths + assert any(len(set(lengths)) > 1 for lengths in packed_lengths) + for req in reqs: + assert list(req.generated_token_ids) == serial_tokens(req.prompt_token_ids, 9) + assert_released(engine) + + +@pytest.mark.parametrize("model_type,layers", [("qwen3_5_moe", 1), ("qwen3_5", 2)]) +def test_unsupported_mtp_architecture_fails_before_worker_setup( + monkeypatch, model_type, layers +): + from infinilm.infer_engine import InferEngine + + monkeypatch.setattr( + "infinilm.infer_engine.read_hf_config", + lambda _: { + "model_type": model_type, + "text_config": {"mtp_num_hidden_layers": layers}, + }, + ) + with pytest.raises(ValueError, match="single-layer dense"): + InferEngine("unused", enable_mtp=True) + + +@pytest.mark.parametrize( + "sampling,error", + [ + ({"top_k": 8}, "greedy"), + ({"max_tokens": 0}, "positive integer"), + ({"max_tokens": -1}, "positive integer"), + ({"max_tokens": None}, "positive integer"), + ({"max_tokens": 1.5}, "positive integer"), + ({"max_tokens": 10000}, "capacity"), + ], +) +def test_invalid_request_rejected_before_admission(sampling, error): + engine, req = service(), request() + for name, value in sampling.items(): + setattr(req.sampling_params, name, value) + with pytest.raises(ValueError, match=error): + engine.add_request(req) + assert engine.scheduler.waiting_queue.sync_q.empty() + assert_released(engine) From 86208ae1f200bae769fdf0a07bf2f1e4092fcc2b Mon Sep 17 00:00:00 2001 From: tangchengxiang <2064027004@qq.com> Date: Sat, 19 Sep 2026 15:23:15 +0000 Subject: [PATCH 2/7] fix(inference): isolate MTP recurrence and bound async shutdown Keep ordinary short-prefill dispatch, preserve live KV during graph recapture, and defer cleanup until a timed-out worker exits. Retain opt-in MTP acceleration features and cover ordinary projection, recapture, shutdown and remote-release contracts. --- README.md | 16 +++++ csrc/engine/compiler/paged_compiler.cpp | 2 +- .../qwen3_next/qwen3_next_gated_deltanet.cpp | 37 ++++------ python/infinilm/llm/llm.py | 52 ++++++++------ test/layers/test_pre_transpose.py | 22 +++++- test/models/qwen3_5/test_mtp_model.py | 39 ++++++++++- test/models/qwen3_5/test_mtp_runner.py | 70 ++++++++++++++++++- 7 files changed, 187 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index d642d0bfc..8275ecd41 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,8 @@ For Qwen3.8-27B-FP8 with E4M3 weights and 128×128 weight blocks, set for NVIDIA inference. Weights are packed once during loading using the existing Marlin operator; a separately converted weight file is unnecessary. The default `"compatibility"` backend dequantizes on the device at execution time and is slower. +On A6000, Marlin stores weights in FP8 and computes with BF16/FP16 activations; +it does not require native FP8 matrix multiplication. Example on one A6000 with the Marlin configuration: @@ -46,6 +48,15 @@ checkpoints. Its MTP default is `1 + max_batch_size * (num_draft_tokens + 2)`, independent of KV page count. A smaller pool can reduce concurrent admission or use ordinary Decode when checkpoint rows are unavailable. The page budget must also accommodate each prompt and its requested output limit. +GDN state precision follows the model's `mamba_ssm_dtype` even with MTP disabled. +FP32 states use twice the storage of BF16 states; `num_state_rows` controls their +capacity without changing the precision. Short multi-token GDN recurrence is +selected by speculative checkpoint metadata, not ordinary prompt length. + +Ordinary Qwen inference also uses vocabulary-parallel output projection and the +corrected norm/weight-loading path. These shared changes require ordinary-model +regression checks independently of MTP equivalence. Graph capture preserves the +KV page and recurrent rows that its warmup touches, including on recapture. For exact full-prompt reuse, replace `--disable-prefix-caching` with `--mtp-prefix-cache-mib 512`. This TP1-only LRU cache owns device copies of both @@ -80,6 +91,11 @@ For the TP2 execution and batching checks, expose two GPUs and set The three test modules cover CPU scheduling/lifecycle, GPU execution, and checkpoint/model contracts. GPU checks skip when no test checkpoint is set. +`AsyncLLMEngine.stop(timeout=5.0)` raises `TimeoutError` if a forward is still +running at the deadline. The worker retains its resources and closes the engine +after that forward returns. A caller may retry `stop()` to wait again; a stopping +or closed engine cannot be restarted. + ## 使用方式 #### 一、编译并安装 `InfiniCore` 编译并安装 `InfiniCore`， 详情见 InfiniCore的 [`README`](https://github.com/InfiniTensor/InfiniCore) : diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index e75c94b9f..8f2492240 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -141,7 +141,7 @@ void PagedCompiler::compile() { // Warmup and capture write into physical page zero. Preserve it so // recapturing also remains safe while a request owns that page. for (const auto &kv : forward_context.kv_cache_vec) { - if (capture_mtp && kv) { + if (kv) { state_guard.save_region(kv->narrow({{1, 0, 1}})); } } diff --git a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp index 2ae0bba77..efff5a465 100644 --- a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp +++ b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp @@ -17,10 +17,6 @@ namespace infinilm::models::qwen3_next { namespace { -// Limit the per-token path to short speculation windows. The crossover with -// the chunked kernel depends on the device and has not been tuned generally. -constexpr size_t kMaxRecurrentVerifyTokens = 8; - infinicore::Tensor cast_for_state(const infinicore::Tensor &input, infinicore::DataType dtype) { if (input->dtype() == dtype) { return input; @@ -260,14 +256,10 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid delta_out = delta_out->as_strided( {seq_len, local_num_value_heads_, value_head_dim_}, {delta_out->stride(0), delta_out->stride(2), delta_out->stride(3)}); - } else if ((single_request && seq_len <= kMaxRecurrentVerifyTokens) || mamba_metadata.token_state_indices.has_value()) { - // Short single-sequence speculation window (e.g. one draft token per - // step). The chunked kernel pads to a 128-token chunk, so a two-token - // verification pays a full chunk of work per layer. Reuse the existing - // T=1 indexed-pool operator once per token instead; this mirrors the - // recurrent path used for single-token decode and needs no kernel - // change. Optional token destinations retain intermediate states so - // callers can commit an accepted prefix without a target replay. + } else if (mamba_metadata.token_state_indices.has_value()) { + // Reuse the indexed Decode operator to save each speculative prefix. + // Request boundaries select independent initial states in packed batches. + // Ordinary multi-token Prefill keeps the existing chunked path. auto ssm_state = forward_context.ssm_state_vec[layer_idx_]; auto q_delta = q->as_strided( {1, seq_len, local_num_key_heads_, key_head_dim_}, @@ -291,20 +283,15 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid {1, seq_len, local_num_value_heads_, value_head_dim_}, ssm_state->dtype(), ssm_state->device()); const auto &init_indices = mamba_metadata.init_state_indices.value(); - const auto &final_indices = mamba_metadata.final_state_indices.value(); + const auto &destinations = mamba_metadata.token_state_indices.value(); + const auto &offsets = mamba_metadata.checkpoint_offsets; + size_t request = 0; for (size_t t = 0; t < seq_len; ++t) { - auto step_init = (t == 0) ? init_indices : final_indices; - auto step_final = final_indices; - if (mamba_metadata.token_state_indices.has_value()) { - const auto &destinations = mamba_metadata.token_state_indices.value(); - step_final = destinations->narrow({{0, t, 1}}); - const auto &offsets = mamba_metadata.checkpoint_offsets; - size_t request = 0; - while (t >= static_cast(offsets[request + 1])) { ++request; } - step_init = t == static_cast(offsets[request]) - ? init_indices->narrow({{0, request, 1}}) - : destinations->narrow({{0, t - 1, 1}}); - } + while (t >= static_cast(offsets[request + 1])) { ++request; } + auto step_init = t == static_cast(offsets[request]) + ? init_indices->narrow({{0, request, 1}}) + : destinations->narrow({{0, t - 1, 1}}); + auto step_final = destinations->narrow({{0, t, 1}}); infinicore::op::recurrent_gated_delta_rule_( recurrent_out->narrow({{1, t, 1}}), ssm_state, diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 0b6c583df..d7482e2dc 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -757,6 +757,8 @@ def start(self): if self._running: logger.warning("AsyncLLMEngine is already running") return + if self._step_thread is not None or getattr(self.engine, "_closed", False): + raise RuntimeError("Cannot restart a stopping or closed `AsyncLLMEngine`.") self._loop = asyncio.get_running_loop() self._abort_queue = janus.Queue() @@ -767,13 +769,21 @@ def start(self): self._step_thread.start() logger.info("AsyncLLMEngine started") - def stop(self): - """Stop the background inference loop.""" + def stop(self, timeout: float = 5.0): + """Wait for shutdown; a timeout leaves in-flight resources with the worker. + + The worker closes the engine when its current step returns. Callers may + retry `stop()` to wait again, but must not restart this engine. + """ self._running = False + self._healthy = False if self._step_thread: - # Do not free model/state buffers while a long Prefill is still - # executing on the background thread. - self._step_thread.join() + self._step_thread.join(timeout=timeout) + if self._step_thread.is_alive(): + raise TimeoutError( + "Inference is still stopping; in-flight resources are retained " + "until the worker exits." + ) self.engine.close() logger.info("AsyncLLMEngine stopped") @@ -834,21 +844,23 @@ def _drain_abort_queue(self): def _step_loop(self): """Background loop that runs inference steps.""" - while self._running: - try: - self._drain_abort_queue() - did_work, pending = self.engine.step() - if not did_work: - time.sleep(0.003) - elif pending: - self._loop.call_soon_threadsafe(self._batch_put, pending) - except Exception as e: - logger.error(f"Error in step loop: {e}", exc_info=True) - self._healthy = False - self._running = False - if self.config.enable_mtp: - self.engine.close() - break + try: + while self._running: + try: + self._drain_abort_queue() + did_work, pending = self.engine.step() + if not did_work: + time.sleep(0.003) + elif pending: + self._loop.call_soon_threadsafe(self._batch_put, pending) + except Exception as e: + logger.error(f"Error in step loop: {e}", exc_info=True) + self._healthy = False + self._running = False + finally: + # Only the execution thread can release ownership after a timed-out + # stop. `LLMEngine.close()` is idempotent for a subsequent `stop()`. + self.engine.close() @staticmethod def _batch_put(pending): diff --git a/test/layers/test_pre_transpose.py b/test/layers/test_pre_transpose.py index cd72c0fe3..45bd43959 100644 --- a/test/layers/test_pre_transpose.py +++ b/test/layers/test_pre_transpose.py @@ -32,7 +32,7 @@ def test_pre_transpose_and_reprocessing_preserve_logits(tmp_path): "tie_word_embeddings": True, } (tmp_path / "config.json").write_text(json.dumps(config)) - results = [] + results, decode_results = [], [] for packed in (False, True): engine = InferEngine( str(tmp_path), @@ -55,7 +55,7 @@ def test_pre_transpose_and_reprocessing_preserve_logits(tmp_path): weights[name] = values save_file(weights, tmp_path / "model.safetensors") load_model_state_dict_by_file(engine, str(tmp_path), dtype=engine.dtype) - for _ in range(2): + for recapture in (False, True): engine.process_weights_after_loading() output = engine.forward_raw( infinicore.from_list([[17, 83, 51]], dtype=infinicore.int64), @@ -73,6 +73,24 @@ def test_pre_transpose_and_reprocessing_preserve_logits(tmp_path): infinicore.sync_device() assert torch.isfinite(copied).all() results.append(copied) + if recapture: + engine.compile() # Must preserve page zero of an ordinary model. + decoded = engine.forward_raw( + infinicore.from_list([[23]], dtype=infinicore.int64), + position_ids=infinicore.from_list([3], dtype=infinicore.int64), + past_kv_lengths=infinicore.from_list([3], dtype=infinicore.int32), + total_kv_lengths=infinicore.from_list([4], dtype=infinicore.int32), + input_offsets=infinicore.from_list([0, 1], dtype=infinicore.int32), + cu_seqlens=infinicore.from_list([0, 4], dtype=infinicore.int32), + block_tables=infinicore.from_list([[0]], dtype=infinicore.int32), + slot_mapping=infinicore.from_list([3], dtype=infinicore.int64), + )["logits"] + copied_decode = torch.empty(decoded.shape, dtype=torch.float16) + infinicore.from_torch(copied_decode).copy_(decoded) + infinicore.sync_device() + decode_results.append(copied_decode) del engine for actual in results[1:]: torch.testing.assert_close(actual, results[0], atol=1e-3, rtol=1e-3) + for actual in decode_results[1:]: + torch.testing.assert_close(actual, decode_results[0], atol=1e-3, rtol=1e-3) diff --git a/test/models/qwen3_5/test_mtp_model.py b/test/models/qwen3_5/test_mtp_model.py index ff1b963af..22de723c0 100644 --- a/test/models/qwen3_5/test_mtp_model.py +++ b/test/models/qwen3_5/test_mtp_model.py @@ -98,14 +98,14 @@ def checkpoint(): return path -def create_model(path): +def create_model(path, enable_mtp=True): tp = int(os.environ.get("INFINILM_QWEN_MTP_TEST_TP", "1")) model = InferEngine( path, device=infinicore.device("cuda", 0), distributed_config=DistConfig(tp), cache_config=PagedKVCacheConfig(16, 64, 1), - enable_mtp=True, + enable_mtp=enable_mtp, attention_backend="paged-attn", ) load_model_state_dict_by_file(model, path, dtype=model.dtype) @@ -282,3 +282,38 @@ def test_vocab_logits_and_local_argmax(checkpoint, tmp_path, odd_vocab): == fast["output_ids"].to_numpy().tolist() ) assert int(fast["output_ids"].to_numpy()[-1]) == min(candidates) + + +def test_ordinary_qwen_head_matches_dense_projection_without_mtp(checkpoint): + weights = load_file(Path(checkpoint) / "model.safetensors")["lm_head.weight"] + reference = None + for enabled in (True, False): + model, _ = create_model(checkpoint, enable_mtp=enabled) + runner = MTPRunner( + SimpleNamespace(block_size=64, num_draft_tokens=1, tensor_parallel_size=1), + model, + ) + req = SimpleNamespace(block_table=[0], mamba_cache_index=1) + inputs = runner._inputs(req, [3, 8, 15, 6], 0) + full = model.forward_raw(**inputs, sample_all_positions=True) + logits = torch.empty(full["logits"].shape, dtype=weights.dtype) + infinicore.from_torch(logits).copy_(full["logits"]) + infinicore.sync_device() + if enabled: + hidden = torch.empty(full["hidden_states"].shape, dtype=weights.dtype) + infinicore.from_torch(hidden).copy_(full["hidden_states"]) + infinicore.sync_device() + # Independent, unsharded FP32 projection checks global vocabulary order. + dense = torch.nn.functional.linear(hidden.float(), weights.float()) + torch.testing.assert_close(logits.float(), dense, atol=0.02, rtol=0.02) + reference = logits + else: + torch.testing.assert_close(logits, reference, rtol=0, atol=0) + fast = model.forward_raw( + **inputs, return_logits=False, sample_all_positions=False + ) + assert ( + fast["output_ids"].to_numpy().tolist() == logits[0, -1:].argmax(-1).tolist() + ) + del model + gc.collect() diff --git a/test/models/qwen3_5/test_mtp_runner.py b/test/models/qwen3_5/test_mtp_runner.py index ec793f7aa..402a3621d 100644 --- a/test/models/qwen3_5/test_mtp_runner.py +++ b/test/models/qwen3_5/test_mtp_runner.py @@ -3,12 +3,13 @@ The deterministic engine below checks orchestration, not GPU numerical accuracy. """ +import threading from types import SimpleNamespace import infinicore import pytest from infinilm.config.engine_config import EngineConfig -from infinilm.llm.llm import LLMEngine +from infinilm.llm.llm import AsyncLLMEngine, LLMEngine from infinilm.llm.model_runner.model_runner import ModelRunner from infinilm.llm.model_runner.mtp_runner import MTPRunner from infinilm.llm.request import InferenceRequest, RequestStatus @@ -303,6 +304,73 @@ def test_shutdown_reclaims_running_and_waiting_requests_once(): engine.add_request(request("after-close")) +@pytest.mark.parametrize("mtp", [False, True]) +@pytest.mark.parametrize("failure", [False, True]) +def test_async_shutdown_keeps_inflight_resources_until_worker_exits(mtp, failure): + engine = service() + engine.config.enable_mtp = mtp + engine.model_runner._closed = False + entered, release = threading.Event(), threading.Event() + + def blocked_step(): + entered.set() + assert release.wait(5) + assert not getattr(engine, "_closed", False) + if failure: + raise RuntimeError("in-flight failure") + return False, [] + + engine.step = blocked_step + asynchronous = AsyncLLMEngine.__new__(AsyncLLMEngine) + asynchronous.engine = engine + asynchronous.config = engine.config + asynchronous._running = asynchronous._healthy = True + asynchronous._abort_queue = None + asynchronous._step_thread = threading.Thread(target=asynchronous._step_loop) + asynchronous._step_thread.start() + try: + assert entered.wait(5) + with pytest.raises(TimeoutError, match="resources are retained"): + asynchronous.stop(timeout=0) + assert not getattr(engine, "_closed", False) + assert not asynchronous.is_healthy() + with pytest.raises(RuntimeError, match="restart"): + asynchronous.start() + finally: + release.set() + asynchronous._step_thread.join(timeout=5) + assert not asynchronous._step_thread.is_alive() + assert engine._closed # Worker cleanup runs even without a second stop(). + asynchronous.stop(timeout=0) + asynchronous.stop(timeout=0) + + +@pytest.mark.parametrize("completion", ["finished_sending", "finished_recving"]) +def test_ordinary_remote_pages_remain_owned_until_transfer_completion(completion): + scheduler = Scheduler(num_blocks=4, block_size=4, enable_prefix_caching=False) + scheduler.connector = SimpleNamespace(request_finished=lambda *args: (True, None)) + req = request() + req.block_table, req.slot_mapping = scheduler.cache_manager.allocate_slots(4) + pages = list(req.block_table) + req.mark_canceled() + scheduler.complete_requests([req]) + scheduler.complete_requests([req]) + assert not req.block_table + assert scheduler.pending_free_blocks[req.request_id] == pages + assert scheduler.cache_manager.get_total_usable_blocks() == 3 + output = SimpleNamespace( + kv_connector_output=SimpleNamespace(**{completion: [req.request_id]}) + ) + scheduler.update_from_output(output) + assert not scheduler.pending_free_blocks + assert scheduler.cache_manager.get_total_usable_blocks() == 4 + reassigned, _ = scheduler.cache_manager.allocate_slots(16) + scheduler.update_from_output(output) + scheduler.complete_requests([req]) + assert all(scheduler.cache_manager.blocks[b].ref_count == 1 for b in reassigned) + scheduler.cache_manager.free_blocks(reassigned) + + @pytest.mark.parametrize("failure", ["prefill", "verify", "draft"]) def test_packed_failure_releases_all_request_ownership(failure): engine = service(candidates=2, batch_size=2) From 78d19f74ca5507f895e8760c01466dd62e5aa353 Mon Sep 17 00:00:00 2001 From: tangchengxiang <2064027004@qq.com> Date: Mon, 21 Sep 2026 03:08:19 +0000 Subject: [PATCH 3/7] fix(qwen): align packed verification with decode numerics Reuse Decode Attention with per-request causal lengths for packed short verification. Align NVIDIA batched gate projections with checkpointed recurrence to avoid shape-dependent BF16 state drift after cancellation and admission. Add request-isolation and causal-boundary coverage to the existing MTP execution tests. Validate the real TP2 RTX 5090 cancel/re-admit reproducer for K=1/2/4, exact Conv/GDN state comparisons, and ordinary eager/graph regressions. --- README.md | 4 ++ csrc/engine/infer_engine.cpp | 52 +++++++++++----- csrc/global_state/forward_context.hpp | 4 +- .../qwen3_next/qwen3_next_gated_deltanet.cpp | 11 ++-- test/models/qwen3_5/test_mtp_execution.py | 59 +++++++++++++++++++ 5 files changed, 109 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 8275ecd41..7dafcd4a2 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,10 @@ GDN state precision follows the model's `mamba_ssm_dtype` even with MTP disabled FP32 states use twice the storage of BF16 states; `num_state_rows` controls their capacity without changing the precision. Short multi-token GDN recurrence is selected by speculative checkpoint metadata, not ordinary prompt length. +NVIDIA batched Decode uses the same per-token GDN gate projection shape as +checkpointed verification to avoid BF16 rounding changes with batch size. +Short packed verification reuses Decode Attention with a causal KV length and +the owning request's page-table row for each query. Ordinary Qwen inference also uses vocabulary-parallel output projection and the corrected norm/weight-loading path. These shared changes require ordinary-model diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 3e4b7f36e..834e2c56c 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -339,30 +339,52 @@ InferEngine::Input::to_model_input(infinicore::Device device, bool for_graph) co max_query_length, max_sequence_length}; - // The single-request fast path expands only the attention query rows. - // Each checkpoint request has at most eight tokens. GDN/Conv keep - // their original request offsets and recurrent state indices. + // Expand only attention rows, giving every speculative query its causal + // length and its own request's page table. GDN/Conv retain packed offsets. const bool short_draft = target_hidden_states && input_offsets - && input_offsets.value()->numel() == 2 && max_query_length <= 8; - if ((token_state_indices || short_draft) && input_offsets.value()->numel() == 2 && is_prefill && input.block_tables + && max_query_length <= 8; + if ((token_state_indices || short_draft) && is_prefill && input.block_tables && device.getType() == infinicore::Device::Type::NVIDIA && max_sequence_length > max_query_length) { - std::vector lengths(max_query_length); - for (size_t i = 0; i < lengths.size(); ++i) { - lengths[i] = static_cast(max_sequence_length - max_query_length + i + 1); + auto host_offsets = input_offsets.value()->to(infinicore::Device::cpu())->contiguous(); + auto host_lengths = total_sequence_lengths.value()->to(infinicore::Device::cpu())->contiguous(); + if (input_offsets.value()->device().getType() != infinicore::Device::Type::CPU + || total_sequence_lengths.value()->device().getType() != infinicore::Device::Type::CPU) { + infinicore::context::syncStream(); } + const size_t requests = host_offsets->numel() - 1; + if (host_lengths->dtype() != infinicore::DataType::I32 + || host_lengths->shape() != std::vector{requests} + || input.block_tables.value()->ndim() != 2 + || input.block_tables.value()->size(0) != requests) { + throw std::invalid_argument("Short verification needs one KV length and page-table row per request."); + } + const auto *offsets = reinterpret_cast(host_offsets->data()); + const auto *totals = reinterpret_cast(host_lengths->data()); + const size_t tokens = input_ids.value()->numel(); + if (offsets[0] != 0 || offsets[requests] != static_cast(tokens)) { + throw std::invalid_argument("Short verification offsets must cover all query tokens."); + } + std::vector lengths(tokens); auto &metadata = global_state::get_forward_context().attn_metadata; + const auto &tables = input.block_tables.value(); + auto expanded_tables = infinicore::Tensor::empty({tokens, tables->size(1)}, tables->dtype(), device); + for (size_t r = 0; r < requests; ++r) { + const auto count = offsets[r + 1] - offsets[r]; + if (count <= 0 || totals[r] < count) { + throw std::invalid_argument("Short verification requires nonempty queries within each KV length."); + } + for (int32_t t = 0; t < count; ++t) { + lengths[offsets[r] + t] = totals[r] - count + t + 1; + } + auto repeated = tables->narrow({{0, r, 1}})->as_strided({static_cast(count), tables->size(1)}, {0, tables->stride(1)}); + expanded_tables->narrow({{0, static_cast(offsets[r]), static_cast(count)}})->copy_from(repeated); + } metadata.verification_sequence_lengths = infinicore::Tensor::empty( {lengths.size()}, infinicore::DataType::I32, device); infinicore::context::memcpyH2D(metadata.verification_sequence_lengths.value()->data(), lengths.data(), lengths.size() * sizeof(int32_t), false); - // Each query references the same physical KV pages. Its own length - // hides the speculative future. NVIDIA Decode kernels assume packed - // page-table rows, so materialize only these small indices, never KV. - metadata.verification_block_tables = input.block_tables.value()->as_strided( - {max_query_length, input.block_tables.value()->size(1)}, - {0, input.block_tables.value()->stride(1)}) - ->contiguous(); + metadata.verification_block_tables = std::move(expanded_tables); } infinilm::global_state::get_forward_context().mamba_metadata = { diff --git a/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index 4eeac6302..ff9b217b7 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -21,8 +21,8 @@ struct AttentionMetadata { size_t max_query_length{0}; /// Maximum total sequence length in the current batch. size_t max_sequence_length{0}; - // Single-request checkpoint verification can reuse Decode Attention with - // one causal KV length per query. These buffers are shared by all layers. + // Packed checkpoint verification reuses Decode Attention with one causal + // KV length and request page-table row per query, shared by all layers. std::optional verification_sequence_lengths; std::optional verification_block_tables; diff --git a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp index efff5a465..e2c52e3d4 100644 --- a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp +++ b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.cpp @@ -185,12 +185,15 @@ infinicore::Tensor Qwen3NextGatedDeltaNet::forward(const infinicore::Tensor &hid auto qkv = in_proj_qkv_->forward(hidden_states_mutable); auto z = in_proj_z_->forward(hidden_states_mutable); - // Keep the tiny gate projections on the same GEMM shape as decode. - // In BF16, Q=1 and Q=2 can round differently and change candidate ranking - // after recurrent-state updates. Large quantized projections stay batched. + // Keep tiny gate projections on a common GEMM shape for ordinary NVIDIA + // batched Decode and verification. BF16 shape-dependent rounding changes + // recurrent states even when the input prefix is identical. auto project_gate = [&](const auto &projection) { const auto &metadata = infinilm::global_state::get_forward_context().mamba_metadata; - if (!metadata.token_state_indices.has_value() || seq_len == 1) { + const bool decode = hidden_states->device().getType() == infinicore::Device::Type::NVIDIA + && metadata.input_offsets + && metadata.input_offsets.value()->numel() - 1 == seq_len; + if ((!metadata.token_state_indices && !decode) || seq_len == 1) { return projection->forward(hidden_states_mutable); } auto output = infinicore::Tensor::empty({batch_size, seq_len, local_num_value_heads_}, diff --git a/test/models/qwen3_5/test_mtp_execution.py b/test/models/qwen3_5/test_mtp_execution.py index 55b59fc51..f6ed1277e 100644 --- a/test/models/qwen3_5/test_mtp_execution.py +++ b/test/models/qwen3_5/test_mtp_execution.py @@ -8,6 +8,7 @@ import infinicore import pytest +import torch from infinilm.cache import PagedKVCacheConfig from infinilm.llm.llm import LLM from infinilm.llm.request import InferenceRequest @@ -204,3 +205,61 @@ def test_batched_mtp_matches_serial_with_cancellation(): assert actual[1:] == expected[1:] finally: llm.close() + + +def test_packed_verification_preserves_request_and_causal_boundaries(): + llm = create(num_draft_tokens=2, max_batch_size=2, enable_graph=False) + runner = llm.engine.model_runner.speculative_runner + raw = llm.engine.model_runner.model_engine + reqs = [ + request("left", [i % 59 + 1 for i in range(63)]), + request("right", [i % 53 + 1 for i in range(127)]), + ] + # Different histories and query lengths cross distinct physical page boundaries. + for row, (req, blocks) in enumerate(zip(reqs, ([0, 2], [1, 3, 4])), 1): + req.block_table = blocks + req.mamba_cache_index = row + + def logits_for(inputs): + result = raw.forward_raw(**inputs, sample_all_positions=True) + logits = result["logits"] + cpu = torch.empty(logits.shape, dtype=torch.bfloat16) + infinicore.from_torch(cpu).copy_(logits) + infinicore.sync_device() + return cpu[0] + + def verify(left, right): + return logits_for( + runner._pack( + [ + runner._inputs(reqs[0], left, 63, destinations=[3, 4, 5]), + runner._inputs(reqs[1], right, 127, destinations=[6, 7]), + ] + ) + ) + + try: + inputs = runner._pack( + [runner._inputs(req, list(req.prompt_token_ids), 0) for req in reqs] + ) + raw.forward_raw(**inputs) + expected = verify([5, 7, 9], [11, 13]) + changed_future = verify([5, 17, 19], [11, 23]) + torch.testing.assert_close( + changed_future[[0, 3]], expected[[0, 3]], rtol=0, atol=0 + ) + changed_request = verify([5, 7, 9], [29, 31]) + torch.testing.assert_close(changed_request[:3], expected[:3], rtol=0, atol=0) + # Verification writes scratch rows, so the committed initial states can + # also be consumed by ordinary Decode for an independent prefix check. + ordinary = logits_for( + runner._pack( + [ + runner._inputs(reqs[0], [5], 63), + runner._inputs(reqs[1], [11], 127), + ] + ) + ) + torch.testing.assert_close(ordinary, expected[[0, 3]], rtol=0, atol=0) + finally: + llm.close() From 70db65c29826e5fda2cccb18ee583d2f88b048c7 Mon Sep 17 00:00:00 2001 From: tangchengxiang <2064027004@qq.com> Date: Mon, 21 Sep 2026 08:01:46 +0000 Subject: [PATCH 4/7] test(mtp): share ordinary model regression setup --- README.md | 5 ++- test/layers/test_pre_transpose.py | 52 ++++++++++++++----------------- 2 files changed, 26 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index 7dafcd4a2..1c52916dc 100644 --- a/README.md +++ b/README.md @@ -17,9 +17,8 @@ batching benefit has not been established. It handles cancellation, EOS/output limits and ordinary Decode when speculative cache capacity is unavailable. This requires the matching FP8/MTP InfiniCore build; the release dependency listed above does not contain these additions. Track the -runtime patch in [InfiniCore #1565](https://github.com/InfiniTensor/InfiniCore/issues/1565). -Graph execution additionally requires the graph lifetime/recording fixes in -[InfiniCore #1560](https://github.com/InfiniTensor/InfiniCore/pull/1560). +runtime and graph lifetime/recording fixes in +[InfiniCore #1566](https://github.com/InfiniTensor/InfiniCore/pull/1566). For Qwen3.8-27B-FP8 with E4M3 weights and 128×128 weight blocks, set `quantization_config.fp8_backend` to `"marlin"` in the checkpoint configuration diff --git a/test/layers/test_pre_transpose.py b/test/layers/test_pre_transpose.py index 45bd43959..cdabda10e 100644 --- a/test/layers/test_pre_transpose.py +++ b/test/layers/test_pre_transpose.py @@ -55,40 +55,36 @@ def test_pre_transpose_and_reprocessing_preserve_logits(tmp_path): weights[name] = values save_file(weights, tmp_path / "model.safetensors") load_model_state_dict_by_file(engine, str(tmp_path), dtype=engine.dtype) - for recapture in (False, True): - engine.process_weights_after_loading() - output = engine.forward_raw( - infinicore.from_list([[17, 83, 51]], dtype=infinicore.int64), - position_ids=infinicore.from_list([0, 1, 2], dtype=infinicore.int64), - past_kv_lengths=infinicore.from_list([0], dtype=infinicore.int32), - total_kv_lengths=infinicore.from_list([3], dtype=infinicore.int32), - input_offsets=infinicore.from_list([0, 3], dtype=infinicore.int32), - cu_seqlens=infinicore.from_list([0, 3], dtype=infinicore.int32), - block_tables=infinicore.from_list([[0]], dtype=infinicore.int32), - slot_mapping=infinicore.from_list([0, 1, 2], dtype=infinicore.int64), - sample_all_positions=True, + + def forward(model, tokens, past=0): + def tensor(values, dtype=infinicore.int32): + return infinicore.from_list(values, dtype=dtype) + + end = past + len(tokens) + output = model.forward_raw( + tensor([tokens], infinicore.int64), + position_ids=tensor(list(range(past, end)), infinicore.int64), + past_kv_lengths=tensor([past]), + total_kv_lengths=tensor([end]), + input_offsets=tensor([0, len(tokens)]), + cu_seqlens=tensor([0, end]), + block_tables=tensor([[0]]), + slot_mapping=tensor(list(range(past, end)), infinicore.int64), + sample_all_positions=past == 0, )["logits"] copied = torch.empty(output.shape, dtype=torch.float16) infinicore.from_torch(copied).copy_(output) infinicore.sync_device() assert torch.isfinite(copied).all() - results.append(copied) + return copied + + for recapture in (False, True): + engine.process_weights_after_loading() + results.append(forward(engine, [17, 83, 51])) if recapture: - engine.compile() # Must preserve page zero of an ordinary model. - decoded = engine.forward_raw( - infinicore.from_list([[23]], dtype=infinicore.int64), - position_ids=infinicore.from_list([3], dtype=infinicore.int64), - past_kv_lengths=infinicore.from_list([3], dtype=infinicore.int32), - total_kv_lengths=infinicore.from_list([4], dtype=infinicore.int32), - input_offsets=infinicore.from_list([0, 1], dtype=infinicore.int32), - cu_seqlens=infinicore.from_list([0, 4], dtype=infinicore.int32), - block_tables=infinicore.from_list([[0]], dtype=infinicore.int32), - slot_mapping=infinicore.from_list([3], dtype=infinicore.int64), - )["logits"] - copied_decode = torch.empty(decoded.shape, dtype=torch.float16) - infinicore.from_torch(copied_decode).copy_(decoded) - infinicore.sync_device() - decode_results.append(copied_decode) + # Reprocessing also recaptures graphs and must preserve live KV. + engine.process_weights_after_loading() + decode_results.append(forward(engine, [23], past=3)) del engine for actual in results[1:]: torch.testing.assert_close(actual, results[0], atol=1e-3, rtol=1e-3) From dc431b9be9cd1bf1a752743b216fb2400a15830a Mon Sep 17 00:00:00 2001 From: tangchengxiang <2064027004@qq.com> Date: Wed, 23 Sep 2026 13:09:07 +0000 Subject: [PATCH 5/7] refactor(mtp): remove optional draft graph capture --- README.md | 7 +- csrc/engine/compiler/paged_compiler.cpp | 128 +++++----------------- csrc/engine/compiler/paged_compiler.hpp | 3 - csrc/engine/infer_engine.cpp | 3 + csrc/engine/rank_worker.cpp | 6 +- csrc/pybind11/engine/engine.hpp | 1 - python/infinilm/config/engine_config.py | 6 +- python/infinilm/llm/llm.py | 8 +- test/models/qwen3_5/test_mtp_execution.py | 26 +++-- test/models/qwen3_5/test_mtp_model.py | 10 ++ test/models/qwen3_5/test_mtp_runner.py | 3 +- 11 files changed, 71 insertions(+), 130 deletions(-) diff --git a/README.md b/README.md index 1c52916dc..c3f6d02bf 100644 --- a/README.md +++ b/README.md @@ -70,13 +70,12 @@ Partial-prefix matching and the Attention-only cache's SLRU policy are not supported by this hybrid snapshot cache. Cache reset or weight loading invalidates snapshots. -`--enable-graph` currently requires `--num-draft-tokens 1 --max-batch-size 1`. -It captures ordinary Decode and supported short draft shapes; Prefill and target -verification remain eager. Multi-candidate and batched execution use eager. +Built-in MTP uses eager execution for all candidate counts; do not combine +`--enable-mtp` with `--enable-graph`. Ordinary inference retains Decode graphs. Random sampling, multimodal requests, multi-layer MTP service execution and remote state transfer are rejected. NVIDIA A6000 validation covers TP1 and TP2 greedy execution, batched requests, cancellation and cache reclamation. The 27B FP8 TP2 -checks use K=2; K=1/2/4 and graph recapture are additionally checked with a tiny +checks use K=2; K=1/2/4 and ordinary graph recapture are also checked with a tiny checkpoint. Exact full-prompt caching remains TP1-only. Other accelerators have not been validated for this service path. diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index 8f2492240..251e7b444 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -110,9 +110,6 @@ void PagedCompiler::compile() { infinicore::context::syncStream(); compiled_map_decode_.clear(); - compiled_map_draft_.clear(); - const bool capture_mtp = model_->supports_token_state_checkpoints() - && infinicore::context::getDevice().getType() == infinicore::Device::Type::NVIDIA; if (decode_batch_sizes_.empty()) { return; } @@ -228,110 +225,53 @@ void PagedCompiler::compile() { infinicore::context::syncStream(); } - auto capture = [&](InfinilmModel::Input input) { - const auto lengths = forward_context.attn_metadata.verification_sequence_lengths; - const auto tables = forward_context.attn_metadata.verification_block_tables; + for (size_t b : decode_batch_sizes_) { + if (b > max_batch_size) { + continue; + } + auto input = make_decode_input(b); + barrier_->wait(); (void)model_->forward(input); infinicore::context::syncStream(); + // Capture must not start with stale Marlin locks from previous + // warmup/capture attempts. This reset is intentionally outside + // graph capture; the current implementation still pays a memset + // before every graph replay in get_compiled(). model_->reset_runtime_state(); infinicore::context::syncStream(); infinicore::context::startGraphRecording(); auto output = model_->forward(input); auto graph = infinicore::context::stopGraphRecording(); barrier_->wait(); + auto shared_output = std::make_shared(); shared_output->logits = infinicore::graph::GraphTensor(output.logits); - if (output.hidden_states) { - shared_output->hidden_states = infinicore::graph::GraphTensor(output.hidden_states); - } - return CompiledResult{std::move(input), {graph, shared_output}, lengths, tables}; - }; - for (size_t b : decode_batch_sizes_) { - if (b <= max_batch_size) { - compiled_map_decode_[b] = capture(make_decode_input(b)); - } - } - if (capture_mtp) { - auto make_draft_input = [&](size_t tokens) { - auto input = make_decode_input(1); - const auto device = infinicore::context::getDevice(); - auto i32 = [&](const std::vector &values) { - auto result = infinicore::Tensor::empty({values.size()}, infinicore::DataType::I32, device); - infinicore::context::memcpyH2D(result->data(), values.data(), values.size() * sizeof(int32_t), false); - return result; - }; - input.input_ids = infinicore::Tensor::zeros({1, tokens}, infinicore::DataType::I64, device); - input.position_ids = infinicore::Tensor::zeros( - position_id_axes > 1 ? std::vector{position_id_axes, tokens} : std::vector{tokens}, - infinicore::DataType::I64, device); - input.total_sequence_lengths = i32({static_cast(tokens)}); - input.input_offsets = i32({0, static_cast(tokens)}); - input.cu_seqlens = i32({0, static_cast(tokens)}); - input.slot_mapping = infinicore::Tensor::empty({tokens}, infinicore::DataType::I64, device); - std::vector slots(tokens); - std::iota(slots.begin(), slots.end(), 0); - infinicore::context::memcpyH2D(input.slot_mapping.value()->data(), slots.data(), tokens * sizeof(int64_t), false); - input.target_hidden_states = infinicore::Tensor::zeros( - {1, tokens, model_config->get("hidden_size")}, model_config->get_dtype(), device); - forward_context.attn_metadata = global_state::AttentionMetadata(input); - forward_context.mamba_metadata = {input.input_offsets, input.mamba_init_state_indices, - input.mamba_final_state_indices, input.token_state_indices}; - if (tokens > 1) { - forward_context.attn_metadata.verification_sequence_lengths = i32({1, 2}); - forward_context.attn_metadata.verification_block_tables = infinicore::Tensor::zeros( - {tokens, nblocks}, infinicore::DataType::I32, device); - } - return input; - }; - for (size_t tokens : {1, 2}) { - compiled_map_draft_[tokens] = capture(make_draft_input(tokens)); - } + + compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)}; } state_guard.restore(); } } PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &input) { - // Q=2 target graphs are intentionally excluded: measured slower than - // short eager verification. They also require their own checkpoint mode. - if (input.token_state_indices) { + if (input.target_hidden_states || input.token_state_indices || input.sample_all_positions) { return {nullptr, nullptr}; } if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { size_t batch_size = input.block_tables.value()->size(0); size_t block_per_req = input.block_tables.value()->size(1); - const size_t tokens = input.input_ids.value()->size(1); - const bool draft = input.target_hidden_states.has_value(); - if ((draft && batch_size != 1) - || (!draft && (batch_size != tokens || input.sample_all_positions))) { + + // Only ordinary Decode batches use captured graphs. + if (batch_size != input.input_ids.value()->size(1)) { return {nullptr, nullptr}; - } - { - auto &compiled = draft ? compiled_map_draft_ : compiled_map_decode_; - auto result = compiled.find(draft ? tokens : batch_size); - if (result == compiled.end()) { + } else { + auto result = compiled_map_decode_.find(batch_size); + if (result == compiled_map_decode_.end()) { return {nullptr, nullptr}; } auto &graph_input = result->second.input; - const auto &runtime_seq_lens = input.total_sequence_lengths.value(); - if (runtime_seq_lens->device().getType() - != infinicore::Device::Type::CPU - || runtime_seq_lens->dtype() != infinicore::DataType::I32 - || runtime_seq_lens->shape().size() != 1 - || runtime_seq_lens->shape()[0] != batch_size) { - throw std::runtime_error( - "PagedCompiler expected CPU int32 " - "total_sequence_lengths for graph replay"); - } - if (draft) { - if (input.target_hidden_states.value()->shape() != graph_input.target_hidden_states.value()->shape() - || input.target_hidden_states.value()->dtype() != graph_input.target_hidden_states.value()->dtype()) { - return {nullptr, nullptr}; - } - graph_input.target_hidden_states.value()->copy_from(input.target_hidden_states.value()); - } graph_input.input_ids.value()->copy_from(input.input_ids.value()); graph_input.position_ids.value()->copy_from(input.position_ids.value()); graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value()); @@ -351,17 +291,6 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & set_minus_one_device_async(graph_block_tables); graph_block_tables->narrow({{1, 0, block_per_req}})->copy_from(input.block_tables.value()); graph_input.slot_mapping.value()->copy_from(input.slot_mapping.value()); - std::vector verification_lengths; - if (result->second.verification_lengths) { - const auto total = reinterpret_cast(input.total_sequence_lengths.value()->data())[0]; - verification_lengths.resize(tokens); - for (size_t t = 0; t < tokens; ++t) { - verification_lengths[t] = total - static_cast(tokens) + static_cast(t) + 1; - result->second.verification_tables.value()->narrow({{0, t, 1}})->copy_from(graph_block_tables); - } - infinicore::context::memcpyH2D(result->second.verification_lengths.value()->data(), - verification_lengths.data(), tokens * sizeof(int32_t), false); - } const bool graph_has_mamba_indices = graph_input.mamba_init_state_indices.has_value() && graph_input.mamba_final_state_indices.has_value(); const bool input_has_mamba_indices = input.mamba_init_state_indices.has_value() && input.mamba_final_state_indices.has_value(); @@ -383,8 +312,15 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & auto graph = std::get<0>(result->second.compiled); if (graph != nullptr) { - if (result->second.verification_lengths) { - graph->bind_host_int_array(result->second.verification_lengths.value(), verification_lengths.data(), tokens); + const auto &runtime_seq_lens = input.total_sequence_lengths.value(); + if (runtime_seq_lens->device().getType() + != infinicore::Device::Type::CPU + || runtime_seq_lens->dtype() != infinicore::DataType::I32 + || runtime_seq_lens->shape().size() != 1 + || runtime_seq_lens->shape()[0] != batch_size) { + throw std::runtime_error( + "PagedCompiler expected CPU int32 " + "total_sequence_lengths for graph replay"); } graph->bind_host_int_array( graph_input.total_sequence_lengths.value(), @@ -392,12 +328,8 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & runtime_seq_lens->data()), batch_size); } - auto saved_output = std::get<1>(result->second.compiled); auto shared_output = std::make_shared(); - shared_output->logits = saved_output->logits->resume_from_blob_(); - if (saved_output->hidden_states) { - shared_output->hidden_states = saved_output->hidden_states->resume_from_blob_(); - } + shared_output->logits = std::get<1>(result->second.compiled)->logits->resume_from_blob_(); return std::make_tuple(graph, shared_output); } diff --git a/csrc/engine/compiler/paged_compiler.hpp b/csrc/engine/compiler/paged_compiler.hpp index 7a78a4315..a1125864d 100644 --- a/csrc/engine/compiler/paged_compiler.hpp +++ b/csrc/engine/compiler/paged_compiler.hpp @@ -21,14 +21,11 @@ class PagedCompiler : public GraphCompiler { struct CompiledResult { InfinilmModel::Input input; Compiled compiled; - std::optional verification_lengths; - std::optional verification_tables; }; std::unordered_map< size_t, // num_requests CompiledResult> compiled_map_decode_; - std::unordered_map compiled_map_draft_; }; } // namespace infinilm::engine diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 834e2c56c..31a84f58e 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -74,6 +74,9 @@ InferEngine::InferEngine( // Load model config if model_path is provided, model_path must be valid, and config.json exists this->model_config_ = infinilm::config::ConfigFactory::createConfig(config_str); + if (enable_graph_compiling && model_config_->get_or("enable_mtp", false)) { + throw std::invalid_argument("Built-in MTP requires eager execution; disable graph compiling."); + } auto infinilm_config = std::make_shared( attention_backend, this->model_config_, diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 5ba433fbf..c7dbd1fee 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -437,17 +437,15 @@ void RankWorker::thread_loop() { && !model_->supports_token_state_checkpoints()) { throw std::runtime_error("This model does not support per-token state checkpoints."); } - const bool graph_candidate = !local_args.token_state_indices + const bool graph_candidate = !local_args.token_state_indices && !local_args.target_hidden_states && local_args.input_ids && local_args.input_offsets - && (local_args.input_ids.value()->numel() == local_args.input_offsets.value()->numel() - 1 - || (local_args.target_hidden_states && local_args.input_ids.value()->numel() <= 2)); + && local_args.input_ids.value()->numel() == local_args.input_offsets.value()->numel() - 1; if (graph_candidate && compiler_ != nullptr && rank_info_.pp_size == 1) { auto graph_input = local_args.to_model_input(infinicore::Device::cpu(), true); auto [graph, output] = compiler_->get_compiled(graph_input); if (graph != nullptr && output != nullptr) { graph->run(); logits = output->logits; - hidden_states = output->hidden_states; model_input_ids = graph_input.input_ids.value(); } } diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index d9f614f45..bceb195d1 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -130,7 +130,6 @@ inline void bind_infer_engine(py::module &m) { py::gil_scoped_release release; return self.forward(input); }, "Run inference on all ranks with arbitrary arguments") .def("reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) - .def("compile", &InferEngine::compile, py::call_guard()) .def("get_kv_cache", [](InferEngine &self) { py::list ranks; for (const auto &rank : self.get_kv_cache()) { diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index cea158bf7..f830db141 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -99,11 +99,9 @@ def __post_init__(self) -> None: ) if not 1 <= self.num_draft_tokens <= 4: raise ValueError("Qwen MTP requires `1 <= num_draft_tokens <= 4`.") - if self.enable_graph and ( - self.num_draft_tokens != 1 or self.max_batch_size != 1 - ): + if self.enable_graph: raise ValueError( - "Batched or multi-candidate Qwen MTP currently requires eager mode." + "Qwen MTP requires eager execution; disable `enable_graph`." ) if self.cache_type != "paged" or self.pipeline_parallel_size != 1: raise ValueError( diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index d7482e2dc..8c2cef5dc 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -448,9 +448,7 @@ def __init__( top_k: Default top-k sampling parameter. enable_graph: Whether to enable graph compiling. enable_mtp: Use built-in Qwen MTP (text, greedy, 1-4 candidates, - paged cache). Batched verification runs eager. With one request - and one candidate, enable_graph captures Decode and MTP drafts; - target verification remains eager. + paged cache). Requires eager execution; disable enable_graph. num_state_rows: Hybrid state pool capacity including zero row. Zero selects a concurrency/candidate-based capacity for MTP. mtp_prefix_cache_bytes: Budget for exact-prompt MTP snapshots (TP1). @@ -690,9 +688,7 @@ def __init__( top_k: Default top-k sampling parameter. enable_graph: Whether to enable graph compiling. enable_mtp: Use built-in Qwen MTP (text, greedy, 1-4 candidates, - paged cache). Batched verification runs eager. With one request - and one candidate, enable_graph captures Decode and MTP drafts; - target verification remains eager. + paged cache). Requires eager execution; disable enable_graph. num_state_rows: Hybrid state pool capacity including zero row. Zero selects a concurrency/candidate-based capacity for MTP. mtp_prefix_cache_bytes: Budget for exact-prompt MTP snapshots (TP1). diff --git a/test/models/qwen3_5/test_mtp_execution.py b/test/models/qwen3_5/test_mtp_execution.py index f6ed1277e..ffbd9c9d6 100644 --- a/test/models/qwen3_5/test_mtp_execution.py +++ b/test/models/qwen3_5/test_mtp_execution.py @@ -54,7 +54,7 @@ def drain(engine, reqs, *, admit=True, recapture=False): break assert engine.step()[0] if recapture and step == 1: - engine.model_runner.model_engine.compile() + engine.model_runner.model_engine.process_weights_after_loading() assert all(r.is_finished() for r in reqs) assert not engine.scheduler.mamba_cache_manager.used_block_ids cache = engine.scheduler.cache_manager @@ -63,9 +63,9 @@ def drain(engine, reqs, *, admit=True, recapture=False): return [list(r.generated_token_ids) for r in reqs] -@pytest.mark.parametrize("candidates,graph", [(1, True), (2, False), (4, False)]) -def test_mtp_matches_ordinary_with_reused_and_recaptured_state(candidates, graph): - llm = create(num_draft_tokens=candidates, enable_graph=graph) +@pytest.mark.parametrize("candidates", [1, 2, 4]) +def test_mtp_matches_ordinary_with_reused_and_rebuilt_state(candidates): + llm = create(num_draft_tokens=candidates) engine = llm.engine runner = engine.model_runner mtp = runner.speculative_runner @@ -73,8 +73,8 @@ def test_mtp_matches_ordinary_with_reused_and_recaptured_state(candidates, graph prompt = [i % 63 + 1 for i in range(63)] - def generate(recapture=False): - return drain(engine, [request("test", prompt, 20)], recapture=recapture)[0] + def generate(): + return drain(engine, [request("test", prompt, 20)])[0] try: runner.speculative_runner = None @@ -85,7 +85,7 @@ def generate(recapture=False): mtp.device_tokens = False assert generate() == expected mtp.device_tokens = True - assert generate(recapture=graph) == expected + assert generate() == expected raw.reset_cache(PagedKVCacheConfig(40, 64, 1, llm.config.num_state_rows)) assert generate() == expected # Exercise every device-side acceptance length, then consume the @@ -125,10 +125,20 @@ def force_prefix(**inputs): llm.close() +def test_ordinary_graph_recapture_preserves_recurrent_state(): + llm = create(enable_mtp=False, enable_graph=True) + prompt = [i % 63 + 1 for i in range(63)] + try: + expected = drain(llm.engine, [request("ordinary", prompt)])[0] + actual = drain(llm.engine, [request("recapture", prompt)], recapture=True)[0] + assert actual == expected + finally: + llm.close() + + def test_bounded_prefix_hit_eviction_and_rebuild(): llm = create( num_draft_tokens=1, - enable_graph=True, tensor_parallel_size=1, enable_prefix_caching=True, mtp_prefix_cache_bytes=8 * 1024**2, diff --git a/test/models/qwen3_5/test_mtp_model.py b/test/models/qwen3_5/test_mtp_model.py index 22de723c0..98a11af4e 100644 --- a/test/models/qwen3_5/test_mtp_model.py +++ b/test/models/qwen3_5/test_mtp_model.py @@ -112,6 +112,16 @@ def create_model(path, enable_mtp=True): return model, tp +def test_native_mtp_graph_is_rejected_before_model_loading(checkpoint): + with pytest.raises(ValueError, match="MTP requires eager"): + InferEngine( + checkpoint, + device=infinicore.device("cpu", 0), + enable_mtp=True, + enable_graph_compiling=True, + ) + + @pytest.fixture def engine(checkpoint): # Malformed raw inputs can stop the workers: never share this engine across tests. diff --git a/test/models/qwen3_5/test_mtp_runner.py b/test/models/qwen3_5/test_mtp_runner.py index 402a3621d..b35bb9240 100644 --- a/test/models/qwen3_5/test_mtp_runner.py +++ b/test/models/qwen3_5/test_mtp_runner.py @@ -213,8 +213,7 @@ def test_state_capacity_uses_normal_decode_without_stale_draft_cache(): [ {"enable_prefix_caching": True}, {"num_draft_tokens": 5}, - {"num_draft_tokens": 2, "enable_graph": True}, - {"max_batch_size": 2, "enable_graph": True}, + {"enable_graph": True}, {"pipeline_parallel_size": 2}, {"cache_type": "static"}, {"draft_model_path": "external"}, From 4f6dcb347a137b0e42fde28488643d6d4e1466e7 Mon Sep 17 00:00:00 2001 From: tangchengxiang <2064027004@qq.com> Date: Wed, 23 Sep 2026 13:27:05 +0000 Subject: [PATCH 6/7] refactor(mtp): defer optional prompt snapshot caching Keep the MTP PR focused on greedy speculation, accepted-state checkpoints and scheduler integration. Remove the separate TP1 snapshot cache, its public budget option and invalidation hooks, and its dedicated tests. MTP requires prefix caching to be disabled. Reuse the existing full-prompt model input directly; retain core acceptance, lifecycle, packed execution and ordinary-graph regressions. --- README.md | 15 +- examples/bench.py | 3 - examples/test_infer.py | 3 - python/infinilm/base_config.py | 7 - python/infinilm/config/engine_config.py | 19 +-- python/infinilm/infer_engine.py | 3 - python/infinilm/llm/hybrid_prefix_cache.py | 140 ------------------ python/infinilm/llm/llm.py | 14 +- .../infinilm/llm/model_runner/model_runner.py | 4 - .../infinilm/llm/model_runner/mtp_runner.py | 75 ++-------- python/infinilm/server/inference_server.py | 4 - test/bench/backends/infinilm.py | 5 +- test/bench/test_benchmark.py | 1 - test/models/qwen3_5/test_mtp_execution.py | 50 +------ test/models/qwen3_5/test_mtp_runner.py | 10 +- 15 files changed, 26 insertions(+), 327 deletions(-) delete mode 100644 python/infinilm/llm/hybrid_prefix_cache.py diff --git a/README.md b/README.md index c3f6d02bf..1da846b3a 100644 --- a/README.md +++ b/README.md @@ -61,14 +61,8 @@ corrected norm/weight-loading path. These shared changes require ordinary-model regression checks independently of MTP equivalence. Graph capture preserves the KV page and recurrent rows that its warmup touches, including on recapture. -For exact full-prompt reuse, replace `--disable-prefix-caching` with -`--mtp-prefix-cache-mib 512`. This TP1-only LRU cache owns device copies of both -target/draft KV, recurrent state and the initial MTP outputs. Hits restore into -request-owned pages and state rows. The budget limits live snapshot tensor -storage, not model memory, allocator reservations or total process memory. -Partial-prefix matching and the Attention-only cache's SLRU policy are not -supported by this hybrid snapshot cache. Cache reset or weight loading -invalidates snapshots. +Prefix reuse is disabled for MTP: Attention KV alone cannot restore the matching +Conv/GDN state. Requests still use the existing paged KV allocator. Built-in MTP uses eager execution for all candidate counts; do not combine `--enable-mtp` with `--enable-graph`. Ordinary inference retains Decode graphs. @@ -76,8 +70,7 @@ Random sampling, multimodal requests, multi-layer MTP service execution and remo state transfer are rejected. NVIDIA A6000 validation covers TP1 and TP2 greedy execution, batched requests, cancellation and cache reclamation. The 27B FP8 TP2 checks use K=2; K=1/2/4 and ordinary graph recapture are also checked with a tiny -checkpoint. Exact full-prompt caching remains TP1-only. Other accelerators have -not been validated for this service path. +checkpoint. Other accelerators have not been validated for this service path. Control-flow and GPU integration checks: @@ -89,7 +82,7 @@ INFINILM_QWEN_MTP_TEST_TP=1 python -m pytest \ ``` For the TP2 execution and batching checks, expose two GPUs and set -`INFINILM_QWEN_MTP_TEST_TP=2`; the prefix-cache check still uses TP1. +`INFINILM_QWEN_MTP_TEST_TP=2`. The three test modules cover CPU scheduling/lifecycle, GPU execution, and checkpoint/model contracts. GPU checks skip when no test checkpoint is set. diff --git a/examples/bench.py b/examples/bench.py index e8e3634e8..4c24f61d9 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -447,7 +447,6 @@ def __init__( enable_prefix_caching=False, enable_mtp=False, num_state_rows=0, - mtp_prefix_cache_bytes=0, processor=None, tokenizer=None, prompt_token_segments=None, @@ -503,7 +502,6 @@ def __init__( num_draft_tokens=num_draft_tokens, enable_mtp=enable_mtp, num_state_rows=num_state_rows, - mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, device=self.device_str, dtype=cfg.dtype, tensor_parallel_size=tp, @@ -971,7 +969,6 @@ def run( enable_prefix_caching=cfg.enable_prefix_caching if cfg.enable_mtp else False, enable_mtp=cfg.enable_mtp, num_state_rows=cfg.num_state_rows, - mtp_prefix_cache_bytes=int(cfg.mtp_prefix_cache_mib * 1024**2), processor=processor, tokenizer=tokenizer, prompt_token_segments=prompt_token_segments, diff --git a/examples/test_infer.py b/examples/test_infer.py index 2dde71415..102703a38 100644 --- a/examples/test_infer.py +++ b/examples/test_infer.py @@ -43,7 +43,6 @@ def test( pre_transpose=False, enable_mtp=False, num_state_rows=0, - mtp_prefix_cache_bytes=0, ): model_path = os.path.expanduser(model_path) # ---------------------------------------------------------------------------- # @@ -82,7 +81,6 @@ def test( pre_transpose=pre_transpose, enable_mtp=enable_mtp, num_state_rows=num_state_rows, - mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, ) conversations = [ @@ -193,5 +191,4 @@ def test( pre_transpose=cfg.pre_transpose, enable_mtp=cfg.enable_mtp, num_state_rows=cfg.num_state_rows, - mtp_prefix_cache_bytes=int(cfg.mtp_prefix_cache_mib * 1024**2), ) diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index add26f9e1..aa008c2bc 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -63,7 +63,6 @@ def __init__(self): self.num_draft_tokens = self.args.num_draft_tokens self.enable_mtp = self.args.enable_mtp self.num_state_rows = self.args.num_state_rows - self.mtp_prefix_cache_mib = self.args.mtp_prefix_cache_mib self.device = self.args.device self.tp = self.args.tp self.pp = self.args.pp @@ -218,12 +217,6 @@ def _add_common_args(self): default=0, help="Qwen hybrid state rows, including zero row; 0 selects automatic capacity", ) - self.parser.add_argument( - "--mtp-prefix-cache-mib", - type=int, - default=0, - help="device storage budget for exact-prompt MTP snapshots; TP1 only", - ) self.parser.add_argument("--tp", "--tensor-parallel-size", type=int, default=1) self.parser.add_argument( "--pp", "--pipeline-parallel-size", type=int, default=1 diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index f830db141..7c854ba18 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -15,7 +15,6 @@ class EngineConfig: enable_mtp: Use the Qwen checkpoint's built-in MTP head. num_state_rows: Hybrid state rows, including the reserved zero row. Zero selects an automatic capacity. - mtp_prefix_cache_bytes: Device storage budget for exact-prompt MTP snapshots. device: Device type string ('cpu', 'cuda', 'mlu', etc.). dtype: Data type string ('float16', 'bfloat16', 'float32'). tensor_parallel_size: Number of devices for tensor parallelism. @@ -75,23 +74,12 @@ class EngineConfig: enable_prefix_caching: bool = True enable_mtp: bool = False num_state_rows: int = 0 - mtp_prefix_cache_bytes: int = 0 def __post_init__(self) -> None: if self.max_batch_size < 1: raise ValueError("`max_batch_size` must be >= 1.") if self.num_state_rows != 0 and self.num_state_rows < 2: raise ValueError("`num_state_rows` must be zero (automatic) or >= 2.") - if self.mtp_prefix_cache_bytes < 0: - raise ValueError("`mtp_prefix_cache_bytes` must be non-negative.") - if self.mtp_prefix_cache_bytes and not self.enable_mtp: - raise ValueError("`mtp_prefix_cache_bytes` requires `enable_mtp`.") - if self.mtp_prefix_cache_bytes and ( - not self.enable_prefix_caching or self.tensor_parallel_size != 1 - ): - raise ValueError( - "MTP prefix snapshots require prefix caching and `tensor_parallel_size=1`." - ) if self.enable_mtp: if self.draft_model_path is not None: raise ValueError( @@ -107,11 +95,8 @@ def __post_init__(self) -> None: raise ValueError( "Qwen MTP requires paged caching and `pipeline_parallel_size=1`." ) - if self.enable_prefix_caching and not self.mtp_prefix_cache_bytes: - raise ValueError( - "Disable prefix caching or set `mtp_prefix_cache_bytes` " - "for exact-prompt MTP snapshots." - ) + if self.enable_prefix_caching: + raise ValueError("Qwen MTP requires `enable_prefix_caching=False`.") if ( self.kv_transfer_config is not None and self.kv_transfer_config.kv_connector diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 3ba833347..8e43dbed7 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -168,7 +168,6 @@ def __init__( pre_transpose=False, enable_mtp=None, ): - self.cache_generation = 0 self.hf_config = read_hf_config(model_path) if enable_mtp is not None: self.hf_config["enable_mtp"] = bool(enable_mtp) @@ -775,7 +774,6 @@ def generate( return output_ids def reset_cache(self, cache_config): - self.cache_generation += 1 infinicore.sync_device() self.enable_paged_attn = isinstance(cache_config, PagedKVCacheConfig) super().reset_cache(cache_config) @@ -784,7 +782,6 @@ def state_dict_keyname(self): return list(super().state_dict_keyname()) def load_state_dict(self, state_dict, strict=None): - self.cache_generation += 1 # MoE/quantized paths may register internal packed tensors that are not # present in the HF checkpoint, so callers can request non-strict loads. super().load_params( diff --git a/python/infinilm/llm/hybrid_prefix_cache.py b/python/infinilm/llm/hybrid_prefix_cache.py deleted file mode 100644 index e2872e476..000000000 --- a/python/infinilm/llm/hybrid_prefix_cache.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Byte-bounded, exact-prompt snapshots for single-rank greedy hybrid MTP. - -Snapshots own their device storage. Restoration writes into request-owned pages -and recurrent rows; no mutable state or partially filled KV page is shared. -""" - -import math -from collections import OrderedDict -from dataclasses import dataclass - -import infinicore - - -@dataclass -class HybridPrefixSnapshot: - kv: list - states: list - hidden: object - proposal: object - pending: int - nbytes: int - - -class HybridPrefixCache: - def __init__(self, engine, budget_bytes, block_size): - self.engine = engine - self.budget_bytes = budget_bytes - self.block_size = block_size - self.entries = OrderedDict() - self.used_bytes = 0 - self.hits = self.misses = self.evictions = self.skipped = 0 - self.generation = self._generation() - - def _generation(self): - return getattr(self.engine, "cache_generation", 0) - - def clear(self): - self.entries.clear() - self.used_bytes = 0 - self.generation = self._generation() - - def _check_generation(self): - if self.generation != self._generation(): - self.clear() - - def _storage(self): - kv = [t for t in self.engine.get_kv_cache()[0] if t._underlying is not None] - states = [ - infinicore.Tensor(t) - for group in self.engine.get_hybrid_states()[0] - for t in group - if t is not None - ] - return kv, states - - @staticmethod - def _bytes(tensor): - return ( - math.prod(tensor.shape) - * infinicore.utils.to_torch_dtype(tensor.dtype).itemsize - ) - - @staticmethod - def _clone(tensor): - # InfiniCore-owned buffers keep this copy independent of foreign allocator - # lifetimes and of the compiler's reusable graph output buffers. - result = infinicore.empty( - tensor.shape, dtype=tensor.dtype, device=tensor.device - ) - result.copy_(tensor) - return result - - def restore(self, request): - self._check_generation() - key = tuple(request.prompt_token_ids) - saved = self.entries.get(key) - if saved is None: - self.misses += 1 - return None - kv, states = self._storage() - # Metadata construction can leave the calling thread on CPU. In-place - # InfiniCore ops use that thread's handle/stream, unlike `empty()`, which - # selects the tensor's device. Select it explicitly before restoration. - infinicore.set_device(saved.hidden.device) - for dst, pages in zip(kv, saved.kv): - for block, src in zip(request.block_table, pages): - dst.narrow(1, block, 1).copy_(src) - for dst, src in zip(states, saved.states): - dst.narrow(0, request.mamba_cache_index, 1).copy_(src) - infinicore.sync_device() - self.entries.move_to_end(key) - self.hits += 1 - return saved - - def save(self, request, pending): - self._check_generation() - key = tuple(request.prompt_token_ids) - if key in self.entries: - self.entries.move_to_end(key) - return - state = request.mtp_state - if state is None: - return - kv, states = self._storage() - pages = (request.get_prompt_length() + self.block_size - 1) // self.block_size - sources = [ - [t.narrow(1, block, 1) for block in request.block_table[:pages]] for t in kv - ] - state_sources = [t.narrow(0, request.mamba_cache_index, 1) for t in states] - tensors = ( - [x for layer in sources for x in layer] - + state_sources - + [state.draft_hidden] - ) - if isinstance(state.draft_token, infinicore.Tensor): - tensors.append(state.draft_token) - nbytes = sum(self._bytes(t) for t in tensors) - # Evict before allocating the new snapshot, so the live snapshots never - # temporarily exceed the configured storage budget. - if nbytes > self.budget_bytes: - self.skipped += 1 - return - while self.used_bytes + nbytes > self.budget_bytes: - _, old = self.entries.popitem(last=False) - self.used_bytes -= old.nbytes - self.evictions += 1 - del old - snapshot = HybridPrefixSnapshot( - [[self._clone(t) for t in layer] for layer in sources], - [self._clone(t) for t in state_sources], - self._clone(state.draft_hidden), - self._clone(state.draft_token) - if isinstance(state.draft_token, infinicore.Tensor) - else state.draft_token, - pending, - nbytes, - ) - infinicore.sync_device() - self.entries[key] = snapshot - self.used_bytes += nbytes diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 8c2cef5dc..bb177b8e7 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -50,11 +50,7 @@ def __init__(self, config: EngineConfig): raise ValueError( "Explicit num_state_rows is currently supported for Qwen hybrid models only." ) - if ( - has_mamba_cache - and config.enable_prefix_caching - and not (config.enable_mtp and config.mtp_prefix_cache_bytes) - ): + if has_mamba_cache and config.enable_prefix_caching: model_type = hf_config["model_type"] raise RuntimeError( "Prefix caching is not supported for Mamba-cache model " @@ -428,7 +424,6 @@ def __init__( enable_prefix_caching: bool = True, enable_mtp: bool = False, num_state_rows: int = 0, - mtp_prefix_cache_bytes: int = 0, ): """Initialize LLM. @@ -451,8 +446,6 @@ def __init__( paged cache). Requires eager execution; disable enable_graph. num_state_rows: Hybrid state pool capacity including zero row. Zero selects a concurrency/candidate-based capacity for MTP. - mtp_prefix_cache_bytes: Budget for exact-prompt MTP snapshots (TP1). - Nonzero requires enable_prefix_caching; zero disables snapshots. attn_backend: Attention backend to use ('default', 'flash-attn'). use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. @@ -489,7 +482,6 @@ def __init__( enable_prefix_caching=enable_prefix_caching, enable_mtp=enable_mtp, num_state_rows=num_state_rows, - mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, ) self.engine = LLMEngine(config) self.config = config @@ -668,7 +660,6 @@ def __init__( enable_prefix_caching: bool = True, enable_mtp: bool = False, num_state_rows: int = 0, - mtp_prefix_cache_bytes: int = 0, ): """Initialize AsyncLLMEngine. @@ -691,8 +682,6 @@ def __init__( paged cache). Requires eager execution; disable enable_graph. num_state_rows: Hybrid state pool capacity including zero row. Zero selects a concurrency/candidate-based capacity for MTP. - mtp_prefix_cache_bytes: Budget for exact-prompt MTP snapshots (TP1). - Nonzero requires enable_prefix_caching; zero disables snapshots. attn_backend: Attention backend to use ('default', 'flash-attn'). kv_connector: KV connector type ('MooncakeConnector'). kv_role: Role in KV connector ('kv_producer' or 'kv_consumer'). @@ -734,7 +723,6 @@ def __init__( enable_prefix_caching=enable_prefix_caching, enable_mtp=enable_mtp, num_state_rows=num_state_rows, - mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, ) self.engine = LLMEngine(config) self.config = config diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index 7ff879f82..f5112846a 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -283,8 +283,4 @@ def close(self) -> None: self.pipeline_control.close() if self.kv_connector is not None: self.kv_connector.shutdown() - if self.speculative_runner is not None: - cache = getattr(self.speculative_runner, "prefix_cache", None) - if cache is not None: - cache.clear() self._closed = True diff --git a/python/infinilm/llm/model_runner/mtp_runner.py b/python/infinilm/llm/model_runner/mtp_runner.py index d2d699112..40cb77682 100644 --- a/python/infinilm/llm/model_runner/mtp_runner.py +++ b/python/infinilm/llm/model_runner/mtp_runner.py @@ -1,7 +1,6 @@ """Greedy Qwen MTP using the target engine's built-in head.""" import infinicore -from infinilm.llm.hybrid_prefix_cache import HybridPrefixCache from infinilm.llm.request import MTPRequestState @@ -23,10 +22,6 @@ def __init__(self, config, target_model_engine): self.num_accepted = 0 self.num_target_calls = 0 self.num_capacity_fallbacks = 0 - budget = getattr(config, "mtp_prefix_cache_bytes", 0) - self.prefix_cache = ( - HybridPrefixCache(self.engine, budget, self.block_size) if budget else None - ) def _target(self, inputs, *, all_positions=False, device_verify=True): self.num_target_calls += 1 @@ -338,7 +333,6 @@ def _draft_batch(self, jobs, *, commit=True): return outputs def _prefill(self, requests, cache_ops, model_input): - outputs, misses = {}, [] for req in requests: self.validate_request(req) if cache_ops is None or req.mamba_cache_index is None: @@ -347,59 +341,22 @@ def _prefill(self, requests, cache_ops, model_input): ) if req.num_local_cached_tokens or req.mtp_state is not None: raise RuntimeError("Qwen MTP requires a fresh full prompt prefill.") - saved = self.prefix_cache.restore(req) if self.prefix_cache else None - if saved is not None: - if self._can_continue(req, [saved.pending]): - rows = cache_ops.allocate_state_rows(self.num_draft_tokens + 1) - if rows is not None: - req.mtp_state = MTPRequestState(rows) - # Requests own these small tensors too: evicting the - # snapshot must release all of its budgeted storage. - req.mtp_state.draft_hidden = self.prefix_cache._clone( - saved.hidden - ) - req.mtp_state.draft_token = ( - self.prefix_cache._clone(saved.proposal) - if isinstance(saved.proposal, infinicore.Tensor) - else saved.proposal - ) - req.mtp_state.cached_tokens = req.get_prompt_length() - else: - self.num_capacity_fallbacks += 1 - outputs[req] = [saved.pending] - else: - misses.append(req) - if misses: - # The processor already built this full-prompt batch. Only rebuild - # metadata when snapshot hits remove requests from the target batch. - inputs = ( - model_input - if len(misses) == len(requests) - else self._pack( - [self._inputs(req, list(req.prompt_token_ids), 0) for req in misses] - ) - ) - target = self._target(inputs) - sampled = list(map(int, target["output_ids"].to_numpy())) - offset = 0 - for req, pending in zip(misses, sampled): - hidden = target["hidden_states"].narrow( - 1, offset, req.get_prompt_length() - ) - offset += req.get_prompt_length() - if self._can_continue(req, [pending]): - rows = cache_ops.allocate_state_rows(self.num_draft_tokens + 1) - if rows is not None: - req.mtp_state = MTPRequestState(rows) - self._draft( - req, list(req.prompt_token_ids[1:]) + [pending], 0, hidden - ) - else: - self.num_capacity_fallbacks += 1 - if self.prefix_cache: - self.prefix_cache.save(req, pending) - outputs[req] = [pending] - return [outputs[req] for req in requests] + target = self._target(model_input) + sampled = list(map(int, target["output_ids"].to_numpy())) + offset = 0 + for req, pending in zip(requests, sampled): + hidden = target["hidden_states"].narrow(1, offset, req.get_prompt_length()) + offset += req.get_prompt_length() + if self._can_continue(req, [pending]): + rows = cache_ops.allocate_state_rows(self.num_draft_tokens + 1) + if rows is not None: + req.mtp_state = MTPRequestState(rows) + self._draft( + req, list(req.prompt_token_ids[1:]) + [pending], 0, hidden + ) + else: + self.num_capacity_fallbacks += 1 + return [[pending] for pending in sampled] def _decode_batch(self, requests, cache_ops): outputs, prepared = {}, [] diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 5af030b9f..bb75e62ed 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -128,7 +128,6 @@ def __init__( enable_mtp: bool = False, num_draft_tokens: int = 1, num_state_rows: int = 0, - mtp_prefix_cache_bytes: int = 0, ): """Initialize inference server. @@ -195,7 +194,6 @@ def __init__( self.enable_mtp = enable_mtp self.num_draft_tokens = num_draft_tokens self.num_state_rows = num_state_rows - self.mtp_prefix_cache_bytes = mtp_prefix_cache_bytes self.engine: AsyncLLMEngine = None @@ -243,7 +241,6 @@ async def lifespan(app: FastAPI): enable_mtp=self.enable_mtp, num_draft_tokens=self.num_draft_tokens, num_state_rows=self.num_state_rows, - mtp_prefix_cache_bytes=self.mtp_prefix_cache_bytes, ) self.engine.start() logger.info(f"Engine initialized with model at {self.model_path}") @@ -682,7 +679,6 @@ def main(): enable_mtp=cfg.enable_mtp, num_draft_tokens=cfg.num_draft_tokens, num_state_rows=cfg.num_state_rows, - mtp_prefix_cache_bytes=cfg.mtp_prefix_cache_mib * 1024 * 1024, ) server.start() diff --git a/test/bench/backends/infinilm.py b/test/bench/backends/infinilm.py index 2dded8a27..d62cc9506 100644 --- a/test/bench/backends/infinilm.py +++ b/test/bench/backends/infinilm.py @@ -21,7 +21,6 @@ def __init__( enable_mtp=False, num_draft_tokens=4, num_state_rows=0, - mtp_prefix_cache_bytes=0, num_blocks=128, block_size=256, ): @@ -69,9 +68,7 @@ def __init__( enable_mtp=enable_mtp, num_draft_tokens=num_draft_tokens, num_state_rows=num_state_rows, - mtp_prefix_cache_bytes=mtp_prefix_cache_bytes, - enable_prefix_caching=bool(mtp_prefix_cache_bytes) - or not model_uses_mamba_cache(self.config_dict), + enable_prefix_caching=not model_uses_mamba_cache(self.config_dict), enable_graph=enable_graph, attn_backend=attn_backend, ) diff --git a/test/bench/test_benchmark.py b/test/bench/test_benchmark.py index cd9893c3e..13c19a085 100644 --- a/test/bench/test_benchmark.py +++ b/test/bench/test_benchmark.py @@ -616,7 +616,6 @@ def main(): enable_mtp=cfg.enable_mtp, num_draft_tokens=cfg.num_draft_tokens, num_state_rows=cfg.num_state_rows, - mtp_prefix_cache_bytes=int(cfg.mtp_prefix_cache_mib * 1024**2), num_blocks=cfg.num_blocks, block_size=cfg.block_size, ) diff --git a/test/models/qwen3_5/test_mtp_execution.py b/test/models/qwen3_5/test_mtp_execution.py index ffbd9c9d6..645da96c9 100644 --- a/test/models/qwen3_5/test_mtp_execution.py +++ b/test/models/qwen3_5/test_mtp_execution.py @@ -1,7 +1,7 @@ """GPU correctness: serial equivalence, checkpoint ownership, graphs and batching. Set INFINILM_QWEN_MTP_TEST_MODEL to a tiny checkpoint with MTP weights; -INFINILM_QWEN_MTP_TEST_TP=1 or 2 selects the parallelism (prefix caching uses TP1). +INFINILM_QWEN_MTP_TEST_TP=1 or 2 selects the parallelism. """ import os @@ -136,54 +136,6 @@ def test_ordinary_graph_recapture_preserves_recurrent_state(): llm.close() -def test_bounded_prefix_hit_eviction_and_rebuild(): - llm = create( - num_draft_tokens=1, - tensor_parallel_size=1, - enable_prefix_caching=True, - mtp_prefix_cache_bytes=8 * 1024**2, - ) - engine = llm.engine - runner = engine.model_runner - mtp = runner.speculative_runner - prompt = [i % 59 + 1 for i in range(63)] - try: - runner.speculative_runner = None - engine.config.enable_mtp = False - expected = drain(engine, [request("ordinary", prompt)])[0] - runner.speculative_runner = mtp - engine.config.enable_mtp = True - assert drain(engine, [request("cold", prompt)])[0] == expected - cache = mtp.prefix_cache - assert cache.entries and cache.used_bytes <= cache.budget_bytes - cache.budget_bytes = cache.used_bytes - hits = cache.hits - calls = [] - forward = runner.model_engine.forward_raw - - def capture(**kw): - calls.append(kw["input_ids"].shape[-1]) - return forward(**kw) - - runner.model_engine.forward_raw = capture - assert drain(engine, [request("hit", prompt)])[0] == expected - assert cache.hits == hits + 1 and max(calls) <= 2 - runner.model_engine.forward_raw = forward - drain(engine, [request("different", [7, 2, 19, 8])]) - assert cache.evictions == 1 and cache.used_bytes <= cache.budget_bytes - assert drain(engine, [request("evicted", prompt)])[0] == expected - assert cache.evictions == 2 - generation = runner.model_engine.cache_generation - runner.model_engine.reset_cache(PagedKVCacheConfig(40, 64, 1, 13)) - assert runner.model_engine.cache_generation > generation - hits = cache.hits - assert drain(engine, [request("reset", prompt)])[0] == expected - assert cache.hits == hits - finally: - llm.close() - assert not mtp.prefix_cache.entries - - def test_batched_mtp_matches_serial_with_cancellation(): llm = create( num_draft_tokens=2, diff --git a/test/models/qwen3_5/test_mtp_runner.py b/test/models/qwen3_5/test_mtp_runner.py index b35bb9240..66ef543ab 100644 --- a/test/models/qwen3_5/test_mtp_runner.py +++ b/test/models/qwen3_5/test_mtp_runner.py @@ -416,17 +416,9 @@ def test_mtp_config_sizes_state_pool_independently_from_pages(): "unused", enable_mtp=True, num_state_rows=5, - enable_prefix_caching=True, - mtp_prefix_cache_bytes=1024, + enable_prefix_caching=False, ) assert config.num_state_rows == 5 - with pytest.raises(ValueError, match="tensor_parallel_size"): - EngineConfig( - "unused", - enable_mtp=True, - tensor_parallel_size=2, - mtp_prefix_cache_bytes=1024, - ) def test_packed_draft_preserves_different_acceptance_lengths(): From 31ef429044a114bd8b4bfb1d290edcc06342a2c6 Mon Sep 17 00:00:00 2001 From: tangchengxiang <2064027004@qq.com> Date: Wed, 23 Sep 2026 14:32:37 +0000 Subject: [PATCH 7/7] refactor(mtp): remove orphaned snapshot accessors Remove the unused worker and engine state-export methods with their Python binding after prompt snapshots were deferred. Restore the upstream KV-cache binding and drop capture tracing logs. Retain checkpoint execution, request ownership and ordinary graph capture. Validate the native rebuild, CPU contracts and all four CLI entrypoints. --- csrc/engine/infer_engine.cpp | 9 --------- csrc/engine/infer_engine.hpp | 1 - csrc/engine/rank_worker.cpp | 11 ----------- csrc/engine/rank_worker.hpp | 1 - csrc/pybind11/engine/engine.hpp | 25 +------------------------ 5 files changed, 1 insertion(+), 46 deletions(-) diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 31a84f58e..e5d721073 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -490,15 +490,6 @@ void InferEngine::reset_cache(const cache::CacheConfig *new_config) { this->compile(); } -std::vector>> InferEngine::get_hybrid_states() { - std::vector>> result; - for (auto &worker : workers_) { - worker->wait(); - result.push_back(worker->get_hybrid_states()); - } - return result; -} - std::vector> InferEngine::get_kv_cache() { std::vector> kv_cache_list; if (workers_.empty()) { diff --git a/csrc/engine/infer_engine.hpp b/csrc/engine/infer_engine.hpp index e6e2ea066..31054149d 100644 --- a/csrc/engine/infer_engine.hpp +++ b/csrc/engine/infer_engine.hpp @@ -57,7 +57,6 @@ class InferEngine { void reset_cache(const cache::CacheConfig *new_config); std::vector> get_kv_cache(); - std::vector>> get_hybrid_states(); ~InferEngine(); diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index c7dbd1fee..952dd09ba 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -253,15 +253,6 @@ std::vector RankWorker::get_kv_cache() { //------------------------------------------------------ // close -- request shutdown and join thread //------------------------------------------------------ -std::vector> RankWorker::get_hybrid_states() { - std::unique_lock lk(mutex_); - cv_.wait(lk, [&] { return init_done_ || should_exit_; }); - if (should_exit_ || has_job_) { - throw std::runtime_error("State access requires an idle worker."); - } - return {forward_context_.conv_state_vec, forward_context_.ssm_state_vec}; -} - void RankWorker::close() { { std::lock_guard lock(mutex_); @@ -605,9 +596,7 @@ void RankWorker::thread_loop() { } else if (local_cmd == Command::COMPILE) { try { if (compiler_ != nullptr) { - spdlog::info("Graph capture begin: tp_rank={}", rank_info_.tp_rank); compiler_->compile(); - spdlog::info("Graph capture end: tp_rank={}", rank_info_.tp_rank); } { std::lock_guard lk(mutex_); diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index a9d3e5f6e..abb4e9246 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -130,7 +130,6 @@ class RankWorker { void reset_cache(const cache::CacheConfig *new_config); std::vector get_kv_cache(); - std::vector> get_hybrid_states(); // Compile the model graph if enabled. void compile(); diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index bceb195d1..184641ea5 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -130,30 +130,7 @@ inline void bind_infer_engine(py::module &m) { py::gil_scoped_release release; return self.forward(input); }, "Run inference on all ranks with arbitrary arguments") .def("reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) - .def("get_kv_cache", [](InferEngine &self) { - py::list ranks; - for (const auto &rank : self.get_kv_cache()) { - py::list layers; - for (const auto &tensor : rank) { - layers.append(tensor ? py::cast(tensor) : py::none()); - } - ranks.append(layers); - } - return ranks; }) - .def("get_hybrid_states", [](InferEngine &self) { - py::list ranks; - for (const auto &rank : self.get_hybrid_states()) { - py::list kinds; - for (const auto &kind : rank) { - py::list layers; - for (const auto &tensor : kind) { - layers.append(tensor ? py::cast(tensor) : py::none()); - } - kinds.append(layers); - } - ranks.append(kinds); - } - return ranks; }) + .def("get_kv_cache", &InferEngine::get_kv_cache, "Get per-rank kv cache list") .def("get_cache_config", [](const InferEngine &self) -> std::shared_ptr { auto cfg = self.get_cache_config(); return cfg ? std::shared_ptr(cfg->unique_copy()) : nullptr; })