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/4] 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/4] 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/4] 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/4] 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)