From e7656f5b46b8f7b04b5074a7e5ea5c67663bc364 Mon Sep 17 00:00:00 2001 From: jindonghe1 Date: Mon, 20 Jul 2026 18:14:18 +0800 Subject: [PATCH] bugfix: align DeepSeek V4 MTP hidden state flow on NPU. --- .../core/runtime/acl_graph_executor_test.cpp | 179 +++++++++++++++++- .../profile/profile_graph_warmup_test.cpp | 23 ++- .../framework/config/speculative_config.h | 7 + xllm/core/framework/model/CMakeLists.txt | 2 + xllm/core/framework/model/mtp_utils.h | 32 ++++ .../runtime/acl_graph_persistent_param.cpp | 7 +- xllm/core/runtime/mtp_worker_impl.cpp | 5 +- xllm/core/runtime/worker.cpp | 16 +- xllm/core/runtime/worker_impl.cpp | 6 +- xllm/core/scheduler/profile/graph_warmup.cpp | 8 +- xllm/core/scheduler/profile/graph_warmup.h | 6 +- .../scheduler/profile/profile_manager.cpp | 6 +- xllm/models/llm/mtp_model_base.h | 11 +- 13 files changed, 272 insertions(+), 36 deletions(-) create mode 100644 xllm/core/framework/model/mtp_utils.h diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index 7db963de46..f3e18d9523 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -24,6 +24,7 @@ limitations under the License. #include #include +#include "common/metrics.h" #include "core/framework/batch/batch.h" #include "core/framework/block/block.h" #include "core/framework/block/block_manager_impl.h" @@ -34,16 +35,20 @@ limitations under the License. #include "core/framework/kv_cache/linear_state_restore.h" #include "core/framework/model/model_args.h" #include "core/framework/model/model_output.h" +#include "core/framework/model_context.h" #include "core/framework/model_loader.h" #include "core/framework/request/sequence.h" #include "core/framework/request/stopping_checker.h" #include "core/framework/sampling/sampling_params.h" +#include "core/layers/common/attention_metadata.h" #include "core/layers/npu/npu_lm_head_impl.h" #include "core/layers/npu/npu_word_embedding_impl.h" +#include "core/layers/npu_torch/tests_utils.h" #include "core/runtime/acl_graph_executor_impl.h" #include "core/runtime/acl_graph_persistent_param.h" #include "core/runtime/base_executor_impl.h" #include "core/runtime/options.h" +#include "models/model_registry.h" // Global test environment for ACL graph executor tests class AclGraphExecutorTestEnvironment : public ::testing::Environment { @@ -231,7 +236,16 @@ class SimpleCausalLM : public CausalLM { std::vector& kv_caches, const ModelInputParams& parameters) override { auto hidden_states = forward_impl(tokens, positions, kv_caches, parameters); - return ModelOutput(hidden_states); + ModelOutput model_output(hidden_states); + if (return_aux_hidden_states_) { + model_output.aux_hidden_states = + torch::cat({hidden_states, hidden_states}, /*dim=*/-1); + } + return model_output; + } + + void set_return_aux_hidden_states(bool value) { + return_aux_hidden_states_ = value; } const torch::TensorOptions& options() const override { @@ -298,6 +312,7 @@ class SimpleCausalLM : public CausalLM { torch::Tensor block_scale_; torch::Tensor block_size_; torch::Tensor scalar_one_; + bool return_aux_hidden_states_ = false; }; class AclGraphExecutorTest : public ::testing::Test { @@ -510,6 +525,118 @@ TEST_F(AclGraphExecutorTest, GraphReplayConsistency) { << output2.hidden_states; } +TEST_F(AclGraphExecutorTest, PreservesAuxHiddenStatesAcrossGraphReplay) { + auto batch = CreateTestBatch(); + ASSERT_FALSE(batch->empty()); + + auto forward_input = + batch->prepare_forward_input(options_.num_decoding_tokens(), + /*min_decoding_batch_size=*/0, + model_args_); + forward_input = forward_input.to(*device_, /*dtype=*/torch::kFloat32); + + SimpleCausalLM* simple_model = dynamic_cast(model_.get()); + ASSERT_NE(simple_model, nullptr); + simple_model->set_return_aux_hidden_states(/*value=*/true); + options_.enable_graph_aux_hidden_states(/*value=*/true); + + ModelOutput eager_output = model_->forward({forward_input.token_ids}, + {forward_input.positions}, + kv_caches_, + {forward_input.input_params}); + auto graph_executor = std::make_unique<::xllm::npu::AclGraphExecutorImpl>( + model_.get(), model_args_, *device_, options_); + const double eager_fallbacks_before = + COUNTER_num_model_execution_total_eager.get_value(); + auto run_graph = [&]() { + return graph_executor->run({forward_input.token_ids}, + {forward_input.positions}, + kv_caches_, + {forward_input.input_params}); + }; + ModelOutput capture_output = run_graph(); + for (int32_t slot_idx = 1; + slot_idx < graph_executor->graph_slot_count_for_test(); + ++slot_idx) { + run_graph(); + } + ModelOutput replay_output = run_graph(); + EXPECT_EQ(COUNTER_num_model_execution_total_eager.get_value(), + eager_fallbacks_before) + << "ACL graph aux-hidden test unexpectedly fell back to eager"; + + ASSERT_TRUE(eager_output.aux_hidden_states.defined()); + ASSERT_TRUE(capture_output.aux_hidden_states.defined()); + ASSERT_TRUE(replay_output.aux_hidden_states.defined()); + EXPECT_EQ(capture_output.aux_hidden_states.size(/*dim=*/-1), + model_args_.hidden_size() * 2); + EXPECT_TRUE(torch::allclose(eager_output.aux_hidden_states, + capture_output.aux_hidden_states, + /*rtol=*/1e-5, + /*atol=*/1e-6)); + EXPECT_TRUE(torch::allclose(eager_output.aux_hidden_states, + replay_output.aux_hidden_states, + /*rtol=*/1e-5, + /*atol=*/1e-6)); +} + +TEST(DeepseekV4ModelTest, ReturnsPreHcHiddenStatesForMtp) { + ModelArgs model_args; + model_args.model_type("deepseek_v4"); + model_args.dtype("float32"); + model_args.vocab_size(8); + model_args.hidden_size(2); + model_args.n_layers(0); + model_args.n_heads(1); + model_args.hc_mult(2); + model_args.hc_eps(0.1f); + model_args.rms_norm_eps(1e-5f); + model_args.window_size(128); + model_args.max_position_embeddings(0); + model_args.index_head_dim(0); + model_args.qk_rope_head_dim(0); + model_args.o_lora_rank(0); + model_args.num_speculative_tokens(1); + + const torch::Device device("npu:0"); + const torch::TensorOptions tensor_options = + torch::dtype(torch::kFloat32).device(device); + layer::test::MockProcessGroup process_group( + device, /*rank=*/0, /*world_size=*/1); + ParallelArgs parallel_args( + /*rank=*/0, /*world_size=*/1, &process_group); + parallel_args.tp_group_ = &process_group; + const ModelContext context( + parallel_args, model_args, QuantArgs(), tensor_options); + const CausalLMFactory factory = + ModelRegistry::get_causallm_factory("deepseek_v4"); + ASSERT_TRUE(static_cast(factory)); + std::unique_ptr model = factory(context); + ASSERT_NE(model, nullptr); + + const torch::Tensor pre_hc_hidden = + torch::tensor({1.0f, 2.0f, 3.0f, 4.0f}, tensor_options).view({1, 2, 2}); + const torch::Tensor tokens = + torch::zeros({1}, torch::dtype(torch::kInt).device(device)); + const torch::Tensor positions = + torch::zeros({1}, torch::dtype(torch::kInt).device(device)); + ModelInputParams input_params; + input_params.meta.num_sequences = 1; + input_params.embedding.input_embedding = pre_hc_hidden; + input_params.attn_metadata = std::make_shared(); + std::vector kv_caches; + + const ModelOutput output = + model->forward(tokens, positions, kv_caches, input_params); + + ASSERT_TRUE(output.aux_hidden_states.defined()); + EXPECT_TRUE(torch::equal(output.aux_hidden_states, + pre_hc_hidden.flatten(/*start_dim=*/1))); + ASSERT_TRUE(output.hidden_states.defined()); + EXPECT_EQ(output.hidden_states.sizes(), torch::IntArrayRef({1, 2})); + EXPECT_EQ(output.aux_hidden_states.sizes(), torch::IntArrayRef({1, 4})); +} + // Test graph creation and execution with different batch sizes TEST_F(AclGraphExecutorTest, DifferentBatchSizes) { // Test with different batch sizes to ensure graph creation works @@ -929,6 +1056,56 @@ TEST(AclGraphPersistentParamTest, SpecVerifyMetadataUsesTokenCapacity) { speculative_config.enable_atb_spec_kernel(original_enable_atb_spec_kernel); } +TEST(AclGraphPersistentParamTest, AuxHiddenStatesUseGraphTokenCapacity) { + SpeculativeConfig& speculative_config = SpeculativeConfig::get_instance(); + const bool original_enable_atb_spec_kernel = + speculative_config.enable_atb_spec_kernel(); + speculative_config.enable_atb_spec_kernel(false); + + ModelArgs args; + args.model_type("deepseek_v4"); + args.dtype("float32"); + args.hidden_size(8); + args.max_position_embeddings(32); + + runtime::Options options; + options.block_size(4); + options.max_seqs_per_batch(10); + options.max_tokens_per_batch(64); + options.num_decoding_tokens(3); + options.enable_speculative_decode(true); + options.is_draft_engine(false); + + const torch::Device device("npu:0"); + const torch::TensorOptions tensor_options = + torch::dtype(torch::kFloat32).device(device); + const torch::Tensor aux_hidden_states = torch::ones({12, 16}, tensor_options); + + ::xllm::npu::GraphPersistentParam target_param(args, device, options); + target_param.set_aux_hidden_states(aux_hidden_states); + EXPECT_EQ(target_param.aux_hidden_states().size(0), 30); + + options.is_draft_engine(true); + ::xllm::npu::GraphPersistentParam draft_param(args, device, options); + draft_param.set_aux_hidden_states(aux_hidden_states.slice( + /*dim=*/0, /*start=*/0, /*end=*/options.max_seqs_per_batch())); + EXPECT_EQ(draft_param.aux_hidden_states().size(0), 10); + + speculative_config.enable_atb_spec_kernel(original_enable_atb_spec_kernel); +} + +TEST(SpeculativeConfigTest, MtpAlgorithmClassificationIsCaseInsensitive) { + EXPECT_TRUE(SpeculativeConfig::is_mtp_algorithm("MTP")); + EXPECT_TRUE(SpeculativeConfig::is_mtp_algorithm("mtp")); + EXPECT_TRUE(SpeculativeConfig::is_mtp_algorithm("Mtp")); + EXPECT_TRUE(SpeculativeConfig::is_mtp_algorithm("mTp")); + + EXPECT_FALSE(SpeculativeConfig::is_mtp_algorithm("Eagle3")); + EXPECT_FALSE(SpeculativeConfig::is_mtp_algorithm("DFlash")); + EXPECT_FALSE(SpeculativeConfig::is_mtp_algorithm("Suffix")); + EXPECT_FALSE(SpeculativeConfig::is_mtp_algorithm("unknown")); +} + TEST(AclGraphExecutorHybridTest, KvCacheSupportsLinearOnlyLayers) { auto conv_cache = torch::zeros({4, 32, 3}, torch::dtype(torch::kFloat32)); auto ssm_cache = torch::zeros({4, 8, 64, 64}, torch::dtype(torch::kFloat32)); diff --git a/tests/core/scheduler/profile/profile_graph_warmup_test.cpp b/tests/core/scheduler/profile/profile_graph_warmup_test.cpp index 2cf4247340..1a026efe62 100644 --- a/tests/core/scheduler/profile/profile_graph_warmup_test.cpp +++ b/tests/core/scheduler/profile/profile_graph_warmup_test.cpp @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "core/framework/model/mtp_utils.h" #include "core/framework/sampling/sampling_params.h" #include "framework/block/block_manager_pool.h" #include "framework/request/incremental_decoder.h" @@ -126,7 +127,7 @@ TEST(GraphWarmupTest, InjectsBootstrapEmbeddingWhenSpeculativeEnabled) { Sequence sequence = make_sequence(/*index=*/0, /*tokens=*/{1, 2, 3}); prepare_warmup_decode_sequence(&sequence, - /*hidden_size=*/128, + /*embedding_width=*/128, /*num_speculative_tokens=*/3); const torch::Tensor embedding = sequence.get_mtp_bootstrap_embedding(); @@ -136,11 +137,29 @@ TEST(GraphWarmupTest, InjectsBootstrapEmbeddingWhenSpeculativeEnabled) { EXPECT_EQ(embedding.size(1), 128); } +TEST(GraphWarmupTest, DeepseekV4MtpUsesFlattenedHyperConnectionWidth) { + ModelArgs model_args; + model_args.model_type() = "deepseek_v4_mtp"; + model_args.hidden_size() = 128; + model_args.hc_mult() = 4; + + EXPECT_EQ(mtp_hidden_state_width(model_args), 512); +} + +TEST(GraphWarmupTest, OtherMtpModelsUseDenseHiddenWidth) { + ModelArgs model_args; + model_args.model_type() = "deepseek_v32_mtp"; + model_args.hidden_size() = 128; + model_args.hc_mult() = 4; + + EXPECT_EQ(mtp_hidden_state_width(model_args), 128); +} + TEST(GraphWarmupTest, SkipsBootstrapEmbeddingWhenSpeculativeDisabled) { Sequence sequence = make_sequence(/*index=*/0, /*tokens=*/{1, 2, 3}); prepare_warmup_decode_sequence(&sequence, - /*hidden_size=*/128, + /*embedding_width=*/128, /*num_speculative_tokens=*/0); EXPECT_FALSE(sequence.get_mtp_bootstrap_embedding().defined()); diff --git a/xllm/core/framework/config/speculative_config.h b/xllm/core/framework/config/speculative_config.h index 54c98b9db9..e257fbd029 100644 --- a/xllm/core/framework/config/speculative_config.h +++ b/xllm/core/framework/config/speculative_config.h @@ -46,6 +46,13 @@ class SpeculativeConfig final { return algorithm == "Eagle3" || algorithm == "DFlash"; } + static bool is_mtp_algorithm(std::string_view algorithm) { + return algorithm.size() == 3 && + (algorithm[0] == 'M' || algorithm[0] == 'm') && + (algorithm[1] == 'T' || algorithm[1] == 't') && + (algorithm[2] == 'P' || algorithm[2] == 'p'); + } + void from_flags(); void from_json(const JsonReader& json); void append_config_json(nlohmann::ordered_json& config_json) const; diff --git a/xllm/core/framework/model/CMakeLists.txt b/xllm/core/framework/model/CMakeLists.txt index 8b6e262036..932855982a 100644 --- a/xllm/core/framework/model/CMakeLists.txt +++ b/xllm/core/framework/model/CMakeLists.txt @@ -10,6 +10,7 @@ cc_library( dit_model.h model_args.h model_input_params.h + mtp_utils.h SRCS DEPS :common @@ -18,6 +19,7 @@ cc_library( :parallel_state :processors :chat_template + :util glog::glog torch torch_python diff --git a/xllm/core/framework/model/mtp_utils.h b/xllm/core/framework/model/mtp_utils.h new file mode 100644 index 0000000000..ae5c9be1ad --- /dev/null +++ b/xllm/core/framework/model/mtp_utils.h @@ -0,0 +1,32 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include "core/framework/model/model_args.h" +#include "core/util/utils.h" + +namespace xllm { + +inline int64_t mtp_hidden_state_width(const ModelArgs& model_args) { + if (util::is_deepseek_v4_model_type(model_args.model_type())) { + return model_args.hc_mult() * model_args.hidden_size(); + } + return model_args.hidden_size(); +} + +} // namespace xllm diff --git a/xllm/core/runtime/acl_graph_persistent_param.cpp b/xllm/core/runtime/acl_graph_persistent_param.cpp index f84b3a7fe5..701359b48f 100644 --- a/xllm/core/runtime/acl_graph_persistent_param.cpp +++ b/xllm/core/runtime/acl_graph_persistent_param.cpp @@ -482,9 +482,10 @@ void GraphPersistentParam::set_aux_hidden_states(const torch::Tensor& value) { if (aux_hidden_states_.numel() == 0) { // Lazy initialization: create aux_hidden_states tensor if not already // created - const int64_t max_tokens_per_batch = options_.max_tokens_per_batch(); + const int64_t graph_token_capacity = + get_decode_graph_token_capacity(options_); auto shape = value.sizes().vec(); - shape[0] = max_tokens_per_batch; + shape[0] = graph_token_capacity; torch::Dtype dtype = util::parse_dtype(args_.dtype(), device_); if (args_.dtype() == "float" || args_.dtype() == "float32") { dtype = torch::kFloat32; @@ -492,6 +493,8 @@ void GraphPersistentParam::set_aux_hidden_states(const torch::Tensor& value) { aux_hidden_states_ = torch::zeros(shape, torch::dtype(dtype).device(device_)); } + CHECK_LE(result_tokens, aux_hidden_states_.size(0)) + << "aux hidden state output exceeds graph token capacity"; // Slice to match the actual shape auto slice = aux_hidden_states_.slice(/*dim=*/0, /*start=*/0, /*end=*/result_tokens); diff --git a/xllm/core/runtime/mtp_worker_impl.cpp b/xllm/core/runtime/mtp_worker_impl.cpp index bb59f8419a..edc9275009 100644 --- a/xllm/core/runtime/mtp_worker_impl.cpp +++ b/xllm/core/runtime/mtp_worker_impl.cpp @@ -30,6 +30,7 @@ limitations under the License. #include "core/framework/config/kv_cache_config.h" #include "core/framework/config/speculative_config.h" #include "core/framework/kv_cache/kv_cache_estimation.h" +#include "core/framework/model/mtp_utils.h" #include "core/framework/multimodal/mm_data.h" #include "spec_input_builder.h" #include "util/pretty_print.h" @@ -711,9 +712,7 @@ int64_t MTPWorkerImpl::get_embedding_placeholder_size() { // hc_mult*hidden per row. if (impl_ != nullptr) { const ModelArgs& args = impl_->context_.get_model_args(); - if (util::is_deepseek_v4_model_type(args.model_type())) { - return args.hc_mult() * args.hidden_size(); - } + return mtp_hidden_state_width(args); } return static_cast(embedding_size_); } diff --git a/xllm/core/runtime/worker.cpp b/xllm/core/runtime/worker.cpp index f5c6c41b75..78ee555e52 100644 --- a/xllm/core/runtime/worker.cpp +++ b/xllm/core/runtime/worker.cpp @@ -26,6 +26,7 @@ limitations under the License. #include #include "common/metrics.h" +#include "core/framework/config/speculative_config.h" #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_input_params.h" #include "framework/state_dict/state_dict.h" @@ -48,17 +49,18 @@ Worker::Worker(const ParallelArgs& parallel_args, const runtime::Options& options, WorkerType worker_type) { if (options.enable_speculative_decode()) { - auto algo = options.speculative_algorithm(); - LOG(INFO) << "Speculative decode is enabled, algorithm: " << algo; - if (algo == "Eagle3") { + const std::string& algorithm = options.speculative_algorithm(); + LOG(INFO) << "Speculative decode is enabled, algorithm: " << algorithm; + if (algorithm == "Eagle3") { impl_ = new Eagle3WorkerImpl(parallel_args, device, options); - } else if (algo == "DFlash") { + } else if (algorithm == "DFlash") { impl_ = new DFlashWorkerImpl(parallel_args, device, options); - } else if (algo == "Suffix") { + } else if (algorithm == "Suffix") { impl_ = new SuffixWorkerImpl(parallel_args, device, options); - } else { - // Default: MTP + } else if (SpeculativeConfig::is_mtp_algorithm(algorithm)) { impl_ = new MTPWorkerImpl(parallel_args, device, options); + } else { + LOG(FATAL) << "Unsupported speculative decoding algorithm: " << algorithm; } } else if (worker_type == WorkerType::LLM) { impl_ = new LLMWorkerImpl(parallel_args, device, options); diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 2a0418159d..bdace1b9d0 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -1517,10 +1517,10 @@ bool WorkerImpl::init_model(const std::string& model_weights_path, } #if defined(USE_NPU) + const std::string& speculative_algorithm = options_.speculative_algorithm(); if (options_.enable_speculative_decode() && - util::is_deepseek_v4_model_type(args.model_type()) && - (options_.speculative_algorithm() == "MTP" || - options_.speculative_algorithm() == "mtp")) { + SpeculativeConfig::is_mtp_algorithm(speculative_algorithm) && + util::is_deepseek_v4_model_type(args.model_type())) { args.num_speculative_tokens(options_.num_speculative_tokens()); } if (options_.speculative_algorithm() == "DFlash") { diff --git a/xllm/core/scheduler/profile/graph_warmup.cpp b/xllm/core/scheduler/profile/graph_warmup.cpp index a27c7719ed..09f1c65671 100644 --- a/xllm/core/scheduler/profile/graph_warmup.cpp +++ b/xllm/core/scheduler/profile/graph_warmup.cpp @@ -122,18 +122,18 @@ std::string next_warmup_request_id() { } void prepare_warmup_decode_sequence(Sequence* sequence, - int64_t hidden_size, + int64_t embedding_width, int32_t num_speculative_tokens) { CHECK(sequence != nullptr); if (num_speculative_tokens <= 0) { return; } - CHECK_GT(hidden_size, 0); + CHECK_GT(embedding_width, 0); // Placeholder bootstrap hidden states; the worker converts dtype/device and - // only the [1, hidden_size] shape matters for the batch input builder. + // only the [1, embedding_width] shape matters for the batch input builder. sequence->update_mtp_bootstrap_embedding( - torch::zeros({1, hidden_size}, torch::kFloat)); + torch::zeros({1, embedding_width}, /*options=*/torch::kFloat)); } } // namespace xllm diff --git a/xllm/core/scheduler/profile/graph_warmup.h b/xllm/core/scheduler/profile/graph_warmup.h index 3bdee18c20..4914fd3e94 100644 --- a/xllm/core/scheduler/profile/graph_warmup.h +++ b/xllm/core/scheduler/profile/graph_warmup.h @@ -55,11 +55,11 @@ std::string next_warmup_request_id(); // decoding is enabled (MTP), the worker's decode path requires a valid decode // state written through the MTP bootstrap channel before it validates the // per-token decode state. This injects a placeholder bootstrap embedding of -// shape [1, hidden_size] so the bootstrap path runs during graph capture; the -// embedding values are irrelevant because warmup only captures the graph. +// shape [1, embedding_width] so the bootstrap path runs during graph capture; +// the embedding values are irrelevant because warmup only captures the graph. // Does nothing when speculative decoding is disabled. void prepare_warmup_decode_sequence(Sequence* sequence, - int64_t hidden_size, + int64_t embedding_width, int32_t num_speculative_tokens); } // namespace xllm diff --git a/xllm/core/scheduler/profile/profile_manager.cpp b/xllm/core/scheduler/profile/profile_manager.cpp index d63349a615..2638182e77 100644 --- a/xllm/core/scheduler/profile/profile_manager.cpp +++ b/xllm/core/scheduler/profile/profile_manager.cpp @@ -34,6 +34,7 @@ limitations under the License. #include "core/framework/config/scheduler_config.h" #include "core/framework/config/service_config.h" #include "core/framework/config/speculative_config.h" +#include "core/framework/model/mtp_utils.h" #include "framework/batch/batch_factory.h" #include "framework/request/request_state.h" #include "scheduler/profile/graph_warmup.h" @@ -691,10 +692,7 @@ std::shared_ptr ProfileManager::generate_single_decode_request( // per-token decode state. Inject a placeholder bootstrap embedding so the // synthetic warmup/profile request takes the same bootstrap path as a real // disagg PD decode request instead of reading stale recycled decode state. - const int64_t bootstrap_width = - util::is_deepseek_v4_model_type(model_args.model_type()) - ? model_args.hc_mult() * model_args.hidden_size() - : model_args.hidden_size(); + const int64_t bootstrap_width = mtp_hidden_state_width(model_args); prepare_warmup_decode_sequence( sequence, bootstrap_width, num_speculative_tokens); diff --git a/xllm/models/llm/mtp_model_base.h b/xllm/models/llm/mtp_model_base.h index 1e0889477c..422d9a435c 100644 --- a/xllm/models/llm/mtp_model_base.h +++ b/xllm/models/llm/mtp_model_base.h @@ -23,6 +23,7 @@ limitations under the License. #include #include +#include "core/framework/model/mtp_utils.h" #include "core/framework/state_dict/utils.h" #include "core/util/utils.h" #include "llm_model_base.h" @@ -115,12 +116,7 @@ class MtpDecoderLayerImplBase : public torch::nn::Module { torch::Tensor embedding_data = input_params.embedding.input_embedding; // for dummy data parallel run, we set a empty embedding if (attn_metadata.is_dummy) { - int64_t embed_cols = model_args_.hidden_size(); - // DeepSeek-V4 targets stash the pre-hc_head 3D hidden flattened to - // [num_tokens, hc_mult*hidden] as aux_hidden_states. - if (is_deepseek_v4_mtp_model(model_args_)) { - embed_cols = model_args_.hc_mult() * model_args_.hidden_size(); - } + const int64_t embed_cols = mtp_hidden_state_width(model_args_); embedding_data = torch::zeros({embed.size(0), embed_cols}, embed.options()); } @@ -147,7 +143,8 @@ class MtpDecoderLayerImplBase : public torch::nn::Module { hidden_states = e_out.unsqueeze(1) + h_out.reshape({-1, hc_mult, hidden}); } else { - // Draft step output is the post-hc_head 2D hidden [N, hidden]. + // Preserve compatibility with callers that provide the post-hc_head + // 2D hidden [N, hidden]. previous_hidden_states = std::get<0>(hnorm_(previous_hidden_states)); hidden_states = e_proj_(enorm_out) + h_proj_(previous_hidden_states); hidden_states = hidden_states.unsqueeze(1).repeat({1, hc_mult, 1});