From 02ae0f2e2f9adb4f2238bafbedec822ac3298124 Mon Sep 17 00:00:00 2001 From: sunbaosong <13793883820@163.com> Date: Fri, 10 Jul 2026 09:34:33 +0800 Subject: [PATCH] feat: add ViT encoder ACL Graph capture for Qwen3-VL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Capture the vision encoder loop as an ACL Graph (NPUGraph) to cut operator dispatch latency and lower TTFT for VLM inference. The graph manager (core/runtime/vision_encoder_acl_graph_manager) is model-agnostic: it captures/replays purely through the VisionEncoderGraphAdapter interface in xllm::npu, so core/runtime never depends on a concrete model type. Qwen3_VisionTransformerImpl implements the adapter and registers itself via set_adapter(this). - Bucket-based lazy capture/replay with persistent input/output buffers - Stable-storage cu_seqlens: in-place copy_ into a pre-allocated device tensor (no set_data storage swap) so the captured graph reads a valid address on every replay - Segment-count guard: a request whose image count differs from the captured graph falls back to eager instead of corrupting attention - capture() wrapped in try/catch with stream restore + eager fallback on failure - head_dim sourced from args.mm_head_dim() (not derived) to support GQA - Config: --enable_encoder_graph, --encoder_graph_budgets Address PR #1866 review feedback: - Fix capture-stream sync race: when the current stream was already non-default, capture ran on it but we synced capture_stream_ (wrong/idle stream). Record the raw stream actually used and synchronize that. - Remove redundant dynamic_cast(model_) in VlmExecutorImpl (model_ is already CausalVLM*); null-check is sufficient. - Fix non-NPU backend build: the has_init_encoder_graph_manager SFINAE trait was defined inside #if defined(USE_NPU) but is referenced unconditionally from CausalVLMImpl::init_encoder_graph_manager, so any non-NPU TU compiling causal_vlm.h failed with 'not a member of xllm::detail'. Moved the trait out of the USE_NPU guard (it uses only ModelArgs/torch::Device; on non-NPU no model provides the method, so it evaluates false and the if-constexpr body is discarded — encoder graph stays NPU-only at runtime). - Style (per .agents/skills/code-review/references/custom-code-style.md): project-root-relative #include, class final, braces on single-line if, int32_t over plain int with static_cast, no auto for std::string, annotate constant args, clang-format line wrapping. Post-build-verification: - Drop the qwen3.h direct #include of rotary_embedding_util.h that this commit had added: a clean build without it succeeds, so layer::rotary::get_concat_rotary_embedding is still resolved transitively (via npu_qwen3_decoder_layer_impl.h / llm_model_base.h). The include was a build side-effect of the CMakeLists change, not required by the feature. Production tuning (benefit-driven, folded in): 4-config (1/2/4/8 TP) x 7-token sweep showed graph only helps at small token counts on high-TP (collective-launch-bound, e.g. 8-card 256 +30.3%) and the captured encoder kernel is slower than eager past ~512 tokens (-3..-5%, TP-independent), so capture is capped at 512. - Hard cap: VisionEncoderAclGraphManager drops any budget > kEncoderGraphMaxReplayTokens (512) at construction via filter_budgets() with a WARNING; select_budget then returns -1 for actual_tokens > 512 so larger images fall back to eager. No config knob (cap is fixed). - Default --encoder_graph_budgets "1024,2048,4096,8192" -> "256,512" so the feature is useful out-of-box (small exact/near-exact buckets). - copy_inputs_to_persistent: drop the redundant full-budget zero_() on the head [0, actual) (immediately overwritten by copy_); zero only the padding tail [actual, bucket) when actual < bucket, since padding rows join the last segment's attention. Exact-budget does no zeroing. - Drop uncommitted ViT-forward timing diagnostics (Timer / [ViTForward] LOG / aclrtSynchronizeDevice) and the now-unused timer.h / acl_rt.h includes from qwen3_vl.h. - Strip non-essential comments; keep only 2 (the 512-cap rationale and the padding-tail zero invariant) plus mandatory license headers. Verified: build_fia.sh incremental build succeeds; pre-commit clang-format passes. Co-Authored-By: Claude --- xllm/core/common/global_flags.h | 3 + .../framework/config/execution_config.cpp | 16 + xllm/core/framework/config/execution_config.h | 6 + xllm/core/framework/model/causal_vlm.h | 9 + xllm/core/framework/model/model_traits.h | 12 + xllm/core/framework/model_context.cpp | 3 +- xllm/core/layers/npu/npu_base_layer.cpp | 3 +- xllm/core/runtime/options.h | 4 + .../vision_encoder_acl_graph_manager.cpp | 324 ++++++++++++++++++ .../vision_encoder_acl_graph_manager.h | 126 +++++++ .../runtime/vision_encoder_graph_adapter.h | 46 +++ xllm/core/runtime/vlm_executor_impl.cpp | 13 + xllm/models/CMakeLists.txt | 4 + xllm/models/vlm/npu/qwen3_vl.h | 100 +++++- 14 files changed, 654 insertions(+), 15 deletions(-) create mode 100644 xllm/core/runtime/vision_encoder_acl_graph_manager.cpp create mode 100644 xllm/core/runtime/vision_encoder_acl_graph_manager.h create mode 100644 xllm/core/runtime/vision_encoder_graph_adapter.h diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index a560f85c1c..2ad4abb54d 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -151,6 +151,9 @@ DECLARE_int32(max_tokens_for_graph_mode); DECLARE_int32(acl_graph_decode_batch_size_limit); +DECLARE_bool(enable_encoder_graph); +DECLARE_string(encoder_graph_budgets); + DECLARE_bool(enable_chunked_prefill); DECLARE_string(master_node_addr); diff --git a/xllm/core/framework/config/execution_config.cpp b/xllm/core/framework/config/execution_config.cpp index 78aa0cc08e..a7ae8a48c3 100644 --- a/xllm/core/framework/config/execution_config.cpp +++ b/xllm/core/framework/config/execution_config.cpp @@ -62,6 +62,14 @@ DEFINE_int32(acl_graph_decode_batch_size_limit, "When actual decode batch_size > this value, ACL graph decode " "falls back to eager mode to avoid OOM."); +DEFINE_bool(enable_encoder_graph, + false, + "Whether to enable ACL graph for vision encoder"); + +DEFINE_string(encoder_graph_budgets, + "256,512", + "Comma-separated token budgets for encoder graph buckets"); + DEFINE_bool(enable_shm, false, "Whether to enable shared memory for executing model."); @@ -91,6 +99,8 @@ void ExecutionConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_graph_vmm_pool); XLLM_CONFIG_ASSIGN_FROM_FLAG(max_tokens_for_graph_mode); XLLM_CONFIG_ASSIGN_FROM_FLAG(acl_graph_decode_batch_size_limit); + XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_encoder_graph); + XLLM_CONFIG_ASSIGN_FROM_FLAG(encoder_graph_budgets); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_shm); XLLM_CONFIG_ASSIGN_FROM_FLAG(use_contiguous_input_buffer); XLLM_CONFIG_ASSIGN_FROM_FLAG(input_shm_size); @@ -106,6 +116,8 @@ void ExecutionConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(enable_graph_vmm_pool); XLLM_CONFIG_ASSIGN_FROM_JSON(max_tokens_for_graph_mode); XLLM_CONFIG_ASSIGN_FROM_JSON(acl_graph_decode_batch_size_limit); + XLLM_CONFIG_ASSIGN_FROM_JSON(enable_encoder_graph); + XLLM_CONFIG_ASSIGN_FROM_JSON(encoder_graph_budgets); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_shm); XLLM_CONFIG_ASSIGN_FROM_JSON(use_contiguous_input_buffer); XLLM_CONFIG_ASSIGN_FROM_JSON(input_shm_size); @@ -130,6 +142,10 @@ void ExecutionConfig::append_config_json( config_json, default_config, max_tokens_for_graph_mode); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, acl_graph_decode_batch_size_limit); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, enable_encoder_graph); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, encoder_graph_budgets); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, enable_shm); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/execution_config.h b/xllm/core/framework/config/execution_config.h index 6b7c335d3b..52b5e4fb62 100644 --- a/xllm/core/framework/config/execution_config.h +++ b/xllm/core/framework/config/execution_config.h @@ -47,6 +47,8 @@ class ExecutionConfig final { "enable_graph_vmm_pool", "max_tokens_for_graph_mode", "acl_graph_decode_batch_size_limit", + "enable_encoder_graph", + "encoder_graph_budgets", "enable_shm", "use_contiguous_input_buffer", "input_shm_size", @@ -69,6 +71,10 @@ class ExecutionConfig final { PROPERTY(int32_t, acl_graph_decode_batch_size_limit) = 16; + PROPERTY(bool, enable_encoder_graph) = false; + + PROPERTY(std::string, encoder_graph_budgets) = "256,512"; + PROPERTY(bool, enable_shm) = false; PROPERTY(bool, use_contiguous_input_buffer) = true; diff --git a/xllm/core/framework/model/causal_vlm.h b/xllm/core/framework/model/causal_vlm.h index 69009ae48f..092f772a05 100644 --- a/xllm/core/framework/model/causal_vlm.h +++ b/xllm/core/framework/model/causal_vlm.h @@ -36,6 +36,8 @@ class CausalVLM : public CausalLM { virtual torch::Tensor get_input_embeddings( const torch::Tensor& input_ids, const ModelInputParams& input_params) = 0; + virtual void init_encoder_graph_manager(const ModelArgs& args, + const torch::Device& device) {} }; template @@ -161,6 +163,13 @@ class CausalVLMImpl : public CausalVLM { torch::Device device() const override { return options_.device(); } + void init_encoder_graph_manager(const ModelArgs& args, + const torch::Device& device) override { + if constexpr (detail::has_init_encoder_graph_manager::value) { + model_->init_encoder_graph_manager(args, device); + } + } + const torch::TensorOptions& options() const override { return options_; } private: diff --git a/xllm/core/framework/model/model_traits.h b/xllm/core/framework/model/model_traits.h index 2c68690782..9bf1f77265 100644 --- a/xllm/core/framework/model/model_traits.h +++ b/xllm/core/framework/model/model_traits.h @@ -25,6 +25,8 @@ namespace xllm { struct ModelInputParams; struct ModelGraphMetadataState; +class ModelLoader; + namespace layer { class LmHead; class WordEmbedding; @@ -167,6 +169,16 @@ struct has_pooler()))>> : std::true_type {}; +template +struct has_init_encoder_graph_manager : std::false_type {}; + +template +struct has_init_encoder_graph_manager< + T, + std::void_t()->init_encoder_graph_manager( + std::declval(), + std::declval()))>> : std::true_type {}; + #if defined(USE_NPU) template struct has_get_npu_lm_head : std::false_type {}; diff --git a/xllm/core/framework/model_context.cpp b/xllm/core/framework/model_context.cpp index 822c44faa5..e45cbcc39c 100644 --- a/xllm/core/framework/model_context.cpp +++ b/xllm/core/framework/model_context.cpp @@ -41,7 +41,8 @@ bool should_enable_async_tiling_copy_stream() { // ATB copy-stream teardown is not reversible for the same context on the // current CANN/PTA stack, so contexts that may enter graph capture must not // pre-create the helper stream. - if (::xllm::ExecutionConfig::get_instance().enable_graph()) { + if (::xllm::ExecutionConfig::get_instance().enable_graph() || + ::xllm::ExecutionConfig::get_instance().enable_encoder_graph()) { return false; } return util::get_bool_env("ATB_USE_TILING_COPY_STREAM", false); diff --git a/xllm/core/layers/npu/npu_base_layer.cpp b/xllm/core/layers/npu/npu_base_layer.cpp index 6074f4937b..fd0e05520f 100644 --- a/xllm/core/layers/npu/npu_base_layer.cpp +++ b/xllm/core/layers/npu/npu_base_layer.cpp @@ -81,7 +81,8 @@ atb::Status BaseLayer::execute_node(atb_speed::Model::Node& node, // However, libtorch_npu current stream is set to default stream after // capture ends, causing inconsistency between ATB context and the actual // execution stream - if (::xllm::ExecutionConfig::get_instance().enable_graph()) { + if (::xllm::ExecutionConfig::get_instance().enable_graph() || + ::xllm::ExecutionConfig::get_instance().enable_encoder_graph()) { void* stream = c10_npu::getCurrentNPUStream(device_.index()).stream(); context_->SetExecuteStream(stream); } diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h index 161b9200a8..c241d0caed 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -248,6 +248,10 @@ struct Options { // maximum number of tokens for graph execution PROPERTY(int32_t, max_tokens_for_graph_mode) = 2048; + PROPERTY(bool, enable_encoder_graph) = false; + + PROPERTY(std::string, encoder_graph_budgets) = "256,512"; + // beam width for beam search PROPERTY(int32_t, beam_width) = 128; diff --git a/xllm/core/runtime/vision_encoder_acl_graph_manager.cpp b/xllm/core/runtime/vision_encoder_acl_graph_manager.cpp new file mode 100644 index 0000000000..4d5291d045 --- /dev/null +++ b/xllm/core/runtime/vision_encoder_acl_graph_manager.cpp @@ -0,0 +1,324 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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. +==============================================================================*/ + +#include "core/runtime/vision_encoder_acl_graph_manager.h" + +#include + +#include + +namespace xllm::npu { + +namespace { + +// Hard cap: graph only helps at small token counts; the captured encoder is +// slower than eager past ~512, so larger images fall back to eager. +constexpr int64_t kEncoderGraphMaxReplayTokens = 512; + +std::vector filter_budgets(const std::vector& budgets) { + std::vector valid; + valid.reserve(budgets.size()); + for (int64_t b : budgets) { + if (b > kEncoderGraphMaxReplayTokens) { + LOG(WARNING) << "[EncoderGraph] dropping budget " << b << " > cap " + << kEncoderGraphMaxReplayTokens + << " (graph only helps at small token counts; larger " + << "images fall back to eager)"; + } else { + valid.push_back(b); + } + } + if (!budgets.empty() && valid.empty()) { + LOG(WARNING) << "[EncoderGraph] no budget <= " + << kEncoderGraphMaxReplayTokens + << "; encoder graph effectively disabled"; + } + return valid; +} + +} // namespace + +VisionEncoderAclGraphManager::VisionEncoderAclGraphManager( + const ModelArgs& args, + const torch::Device& device, + const std::vector& budgets) + : budgets_(filter_budgets(budgets)), + device_(device), + hidden_size_(args.mm_hidden_size()), + head_dim_(args.mm_head_dim()), + d_model_(args.mm_projection_dim()), + spatial_merge_size_(args.mm_spatial_merge_size()), + device_index_(device.index()), + dtype_(args.dtype() == "bfloat16" ? torch::kBFloat16 : torch::kFloat16) { + max_budget_ = budgets_.empty() ? 0 : budgets_.back(); + capture_stream_.emplace( + c10_npu::getStreamFromPool(/*isHighPriority=*/true, device_index_)); + LOG(INFO) << "VisionEncoderAclGraphManager created with " << budgets_.size() + << " budgets, max_budget=" << max_budget_ + << ", hidden_size=" << hidden_size_ << ", head_dim=" << head_dim_ + << ", d_model=" << d_model_; +} + +void VisionEncoderAclGraphManager::initialize_persistent_param( + EncoderPersistentParam& param, + int64_t budget, + int32_t num_deepstack, + size_t num_segments) { + auto opts = torch::TensorOptions().device(device_).dtype(dtype_); + auto int_opts = torch::TensorOptions().device(device_).dtype(torch::kInt32); + + param.hidden_states = torch::zeros({budget, hidden_size_}, opts); + param.cos_pos = torch::zeros({budget, head_dim_}, opts); + param.sin_pos = torch::zeros({budget, head_dim_}, opts); + param.cu_seqlens = + torch::zeros({static_cast(num_segments)}, int_opts); + param.cu_seqlens_vec.clear(); + param.cu_seqlens_vec.reserve(num_segments); + + param.output_hidden = torch::zeros({budget, hidden_size_}, opts); + int64_t merge_sq = spatial_merge_size_ * spatial_merge_size_; + int64_t ds_tokens = budget / merge_sq; + param.deepstack_outs.reserve(num_deepstack); + for (int32_t i = 0; i < num_deepstack; ++i) { + param.deepstack_outs.emplace_back( + torch::zeros({ds_tokens, d_model_}, opts)); + } +} + +int64_t VisionEncoderAclGraphManager::select_budget( + int64_t actual_tokens) const { + for (auto b : budgets_) { + if (b >= actual_tokens) { + return b; + } + } + return -1; +} + +bool VisionEncoderAclGraphManager::can_replay(int64_t actual_tokens) const { + return select_budget(actual_tokens) > 0; +} + +void VisionEncoderAclGraphManager::copy_inputs_to_persistent( + EncoderPersistentParam& param, + const torch::Tensor& hidden_states, + const torch::Tensor& cos_pos, + const torch::Tensor& sin_pos, + const torch::Tensor& /*cu_seqlens*/, + const std::vector& cu_seqlens_vec, + int64_t actual_tokens, + int64_t bucket_size) { + // Zero the padding tail: the last segment is extended to bucket_size below, + // so padding rows join real attention and stale input would leak K/V. + if (actual_tokens < bucket_size) { + param.hidden_states.slice(0, actual_tokens, bucket_size).zero_(); + param.cos_pos.slice(0, actual_tokens, bucket_size).zero_(); + param.sin_pos.slice(0, actual_tokens, bucket_size).zero_(); + } + param.hidden_states.slice(0, 0, actual_tokens).copy_(hidden_states); + param.cos_pos.slice(0, 0, actual_tokens).copy_(cos_pos); + param.sin_pos.slice(0, 0, actual_tokens).copy_(sin_pos); + + param.cu_seqlens_vec = cu_seqlens_vec; + if (!param.cu_seqlens_vec.empty()) { + param.cu_seqlens_vec.back() = static_cast(bucket_size); + } + param.cu_seqlens.copy_(torch::tensor( + param.cu_seqlens_vec, torch::TensorOptions().dtype(torch::kInt32))); +} + +bool VisionEncoderAclGraphManager::capture(VisionEncoderGraphAdapter* adapter, + torch::Tensor& hidden_states, + torch::Tensor& cos_pos, + torch::Tensor& sin_pos, + torch::Tensor& cu_seqlens, + std::vector& cu_seqlens_vec, + int64_t budget) { + CHECK(adapter != nullptr) << "encoder graph adapter must be set"; + const auto& deepstack_indexes = adapter->deepstack_indexes(); + const size_t num_segments = cu_seqlens_vec.size(); + + LOG(INFO) << "Capturing encoder graph, budget=" << budget + << ", actual_tokens=" << hidden_states.size(0) + << ", num_segments=" << num_segments; + + auto bucket_graph = std::make_unique(); + bucket_graph->num_tokens = budget; + bucket_graph->num_segments = num_segments; + + initialize_persistent_param(bucket_graph->param, + budget, + static_cast(deepstack_indexes.size()), + num_segments); + + int64_t actual_tokens = hidden_states.size(0); + copy_inputs_to_persistent(bucket_graph->param, + hidden_states, + cos_pos, + sin_pos, + cu_seqlens, + cu_seqlens_vec, + actual_tokens, + budget); + + auto& param = bucket_graph->param; + + torch::npu::synchronize(); + + { + LOG(INFO) << "[EncoderGraph] Warm-up run with budget=" << budget; + torch::Tensor warm_hidden = param.hidden_states.clone(); + for (int32_t i = 0; i < adapter->num_encoder_layers(); ++i) { + warm_hidden = adapter->forward_encoder_layer( + /*layer_idx=*/i, + warm_hidden, + param.cos_pos, + param.sin_pos, + param.cu_seqlens, + param.cu_seqlens_vec); + } + torch::npu::synchronize(); + LOG(INFO) << "[EncoderGraph] Warm-up completed successfully"; + } + + param.hidden_states.zero_(); + param.hidden_states.slice(0, 0, actual_tokens).copy_(hidden_states); + + bool need_restore = false; + aclrtStream capture_raw_stream = nullptr; + if (c10_npu::getCurrentNPUStream(device_index_) == + c10_npu::getDefaultNPUStream(device_index_)) { + c10_npu::setCurrentNPUStream(*capture_stream_); + capture_raw_stream = capture_stream_->stream(); + aclrtSynchronizeStream(capture_raw_stream); + need_restore = true; + } else { + capture_raw_stream = c10_npu::getCurrentNPUStream(device_index_).stream(); + } + + bool captured = false; + try { + bucket_graph->graph.capture_begin( + {0, 0}, aclmdlRICaptureMode::ACL_MODEL_RI_CAPTURE_MODE_THREAD_LOCAL); + + torch::Tensor hidden = param.hidden_states; + for (int32_t i = 0; i < adapter->num_encoder_layers(); ++i) { + hidden = adapter->forward_encoder_layer( + /*layer_idx=*/i, + hidden, + param.cos_pos, + param.sin_pos, + param.cu_seqlens, + param.cu_seqlens_vec); + + for (size_t k = 0; k < deepstack_indexes.size(); ++k) { + if (deepstack_indexes[k] == i) { + auto ds_out = adapter->forward_deepstack_merger( + static_cast(k), hidden); + param.deepstack_outs[k].copy_(ds_out); + } + } + } + param.output_hidden.copy_(hidden); + + bucket_graph->graph.capture_end(); + captured = true; + } catch (const std::exception& e) { + LOG(ERROR) << "[EncoderGraph] capture threw for budget=" << budget << ": " + << e.what() << "; falling back to eager"; + captured = false; + } + + if (need_restore) { + c10_npu::setCurrentNPUStream(c10_npu::getDefaultNPUStream(device_index_)); + } + + if (!captured) { + return false; + } + + aclrtSynchronizeStream(capture_raw_stream); + bucket_graph->graph.replay(); + + graphs_[budget] = std::move(bucket_graph); + LOG(INFO) << "Encoder graph captured successfully for budget=" << budget; + return true; +} + +std::optional VisionEncoderAclGraphManager::replay( + torch::Tensor& hidden_states, + torch::Tensor& cos_pos, + torch::Tensor& sin_pos, + torch::Tensor& cu_seqlens, + std::vector& cu_seqlens_vec, + int64_t actual_num_tokens) { + int64_t budget = select_budget(actual_num_tokens); + if (budget <= 0) { + return std::nullopt; + } + + auto it = graphs_.find(budget); + if (it == graphs_.end()) { + CHECK(adapter_ != nullptr) << "Encoder adapter not set for lazy capture"; + if (!capture(adapter_, + hidden_states, + cos_pos, + sin_pos, + cu_seqlens, + cu_seqlens_vec, + budget)) { + return std::nullopt; + } + it = graphs_.find(budget); + } + + auto& bucket = it->second; + if (bucket->num_segments != cu_seqlens_vec.size()) { + LOG_FIRST_N(WARNING, 1) + << "[EncoderGraph] segment count mismatch (captured=" + << bucket->num_segments << ", request=" << cu_seqlens_vec.size() + << ") for budget=" << budget << "; falling back to eager. " + << "This message is logged only once."; + return std::nullopt; + } + + copy_inputs_to_persistent(bucket->param, + hidden_states, + cos_pos, + sin_pos, + cu_seqlens, + cu_seqlens_vec, + actual_num_tokens, + budget); + + bucket->graph.replay(); + + EncoderGraphOutput output; + output.hidden_states = + bucket->param.output_hidden.slice(0, 0, actual_num_tokens); + + int64_t merge_sq = spatial_merge_size_ * spatial_merge_size_; + int64_t ds_actual = actual_num_tokens / merge_sq; + const auto& deepstack_indexes = adapter_->deepstack_indexes(); + output.deepstack_features.reserve(deepstack_indexes.size()); + for (size_t k = 0; k < deepstack_indexes.size(); ++k) { + output.deepstack_features.push_back( + bucket->param.deepstack_outs[k].slice(0, 0, ds_actual)); + } + + return output; +} + +} // namespace xllm::npu diff --git a/xllm/core/runtime/vision_encoder_acl_graph_manager.h b/xllm/core/runtime/vision_encoder_acl_graph_manager.h new file mode 100644 index 0000000000..12a1ddc2a9 --- /dev/null +++ b/xllm/core/runtime/vision_encoder_acl_graph_manager.h @@ -0,0 +1,126 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 +#include + +#if defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wattributes" +#endif + +#include "torch_npu/csrc/core/npu/NPUGraph.h" + +#if defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + +#include +#include +#include + +#include "core/framework/model/model_args.h" +#include "core/runtime/vision_encoder_graph_adapter.h" + +namespace xllm::npu { + +struct EncoderGraphOutput { + torch::Tensor hidden_states; + std::vector deepstack_features; +}; + +class VisionEncoderAclGraphManager final { + public: + VisionEncoderAclGraphManager(const ModelArgs& args, + const torch::Device& device, + const std::vector& budgets); + + ~VisionEncoderAclGraphManager() = default; + + bool can_replay(int64_t actual_tokens) const; + + bool capture(VisionEncoderGraphAdapter* adapter, + torch::Tensor& hidden_states, + torch::Tensor& cos_pos, + torch::Tensor& sin_pos, + torch::Tensor& cu_seqlens, + std::vector& cu_seqlens_vec, + int64_t budget); + + std::optional replay(torch::Tensor& hidden_states, + torch::Tensor& cos_pos, + torch::Tensor& sin_pos, + torch::Tensor& cu_seqlens, + std::vector& cu_seqlens_vec, + int64_t actual_num_tokens); + + void set_adapter(VisionEncoderGraphAdapter* adapter) { adapter_ = adapter; } + + private: + struct EncoderPersistentParam { + torch::Tensor hidden_states; + torch::Tensor cos_pos; + torch::Tensor sin_pos; + torch::Tensor cu_seqlens; + std::vector cu_seqlens_vec; + + torch::Tensor output_hidden; + std::vector deepstack_outs; + }; + + struct BucketGraph { + c10_npu::NPUGraph graph; + uint32_t num_tokens = 0; + size_t num_segments = 0; + EncoderPersistentParam param; + }; + + int64_t select_budget(int64_t actual_tokens) const; + + void copy_inputs_to_persistent(EncoderPersistentParam& param, + const torch::Tensor& hidden_states, + const torch::Tensor& cos_pos, + const torch::Tensor& sin_pos, + const torch::Tensor& cu_seqlens, + const std::vector& cu_seqlens_vec, + int64_t actual_tokens, + int64_t bucket_size); + + void initialize_persistent_param(EncoderPersistentParam& param, + int64_t budget, + int32_t num_deepstack, + size_t num_segments); + + std::vector budgets_; + torch::Device device_; + int64_t max_budget_; + int64_t hidden_size_; + int64_t head_dim_; + int64_t d_model_; + int64_t spatial_merge_size_; + torch::ScalarType dtype_; + + std::unordered_map> graphs_; + + std::optional capture_stream_; + c10::DeviceIndex device_index_; + + VisionEncoderGraphAdapter* adapter_ = nullptr; +}; + +} // namespace xllm::npu diff --git a/xllm/core/runtime/vision_encoder_graph_adapter.h b/xllm/core/runtime/vision_encoder_graph_adapter.h new file mode 100644 index 0000000000..ba148440b6 --- /dev/null +++ b/xllm/core/runtime/vision_encoder_graph_adapter.h @@ -0,0 +1,46 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 +#include + +namespace xllm::npu { + +class VisionEncoderGraphAdapter { + public: + virtual ~VisionEncoderGraphAdapter() = default; + + virtual int32_t num_encoder_layers() const = 0; + + virtual const std::vector& deepstack_indexes() const = 0; + + virtual torch::Tensor forward_encoder_layer( + int32_t layer_idx, + torch::Tensor& hidden, + torch::Tensor& cos_pos, + torch::Tensor& sin_pos, + torch::Tensor& cu_seqlens, + std::vector& cu_seqlens_vec) = 0; + + virtual torch::Tensor forward_deepstack_merger( + int32_t deepstack_slot, + const torch::Tensor& hidden) = 0; +}; + +} // namespace xllm::npu diff --git a/xllm/core/runtime/vlm_executor_impl.cpp b/xllm/core/runtime/vlm_executor_impl.cpp index 7cfe4fb049..b71a8e3b95 100644 --- a/xllm/core/runtime/vlm_executor_impl.cpp +++ b/xllm/core/runtime/vlm_executor_impl.cpp @@ -43,6 +43,19 @@ VlmExecutorImpl::VlmExecutorImpl(CausalLM* model, llm_executor_ = ExecutorImplFactory::get_instance().create_executor_impl( model, args, device, options, Platform::type_str()); } + if (::xllm::ExecutionConfig::get_instance().enable_encoder_graph()) { + if (model_) { + model_->init_encoder_graph_manager(args, device); + LOG(INFO) << "[EncoderGraph] encoder graph manager initialized, " + << "enable_encoder_graph=true"; + } else { + LOG(WARNING) << "[EncoderGraph] enable_encoder_graph=true but model is " + "not CausalVLM, skip init"; + } + } else { + LOG(INFO) << "[EncoderGraph] enable_encoder_graph=false, encoder graph " + << "disabled (use --enable_encoder_graph=True to enable)"; + } } ForwardInput VlmExecutorImpl::prepare_inputs(Batch& batch) { diff --git a/xllm/models/CMakeLists.txt b/xllm/models/CMakeLists.txt index 30acbf2151..c27bb09705 100644 --- a/xllm/models/CMakeLists.txt +++ b/xllm/models/CMakeLists.txt @@ -16,11 +16,14 @@ cc_library( models.h dit/utils/lanczos_resample.h $<$:dit/utils/dit_block_weight_manager.h> + $<$:../core/runtime/vision_encoder_acl_graph_manager.h> + $<$:../core/runtime/vision_encoder_graph_adapter.h> SRCS model_registry.cpp dit/utils/lanczos_resample.cpp $<$:dit/utils/dit_block_weight_manager.cpp> vlm/utils/multimodal_utils.cpp + $<$:../core/runtime/vision_encoder_acl_graph_manager.cpp> DEPS :model :mposition @@ -30,4 +33,5 @@ cc_library( :model_loader :layers $<$:npu_layers> + $<$:torch_npu> ) diff --git a/xllm/models/vlm/npu/qwen3_vl.h b/xllm/models/vlm/npu/qwen3_vl.h index 3c19ee7c51..eac75ab002 100644 --- a/xllm/models/vlm/npu/qwen3_vl.h +++ b/xllm/models/vlm/npu/qwen3_vl.h @@ -17,9 +17,11 @@ limitations under the License. #include +#include #include #include "core/common/global_flags.h" +#include "core/framework/config/execution_config.h" #include "core/framework/kv_cache/kv_cache.h" #include "core/framework/model/model_input_params.h" #include "core/framework/model/model_output.h" @@ -27,7 +29,7 @@ limitations under the License. #include "core/layers/npu/npu_qwen3_vision_encoder_layer_impl.h" #include "core/layers/npu/npu_rms_norm_impl.h" #include "core/platform/device.h" -#include "core/util/timer.h" +#include "core/runtime/vision_encoder_acl_graph_manager.h" #include "models/llm/npu/qwen3.h" #include "models/model_registry.h" #include "models/vlm/mposition/mposition.h" @@ -309,7 +311,8 @@ class Qwen3_VisionPatchMergerImpl : public torch::nn::Module { }; TORCH_MODULE(Qwen3_VisionPatchMerger); -class Qwen3_VisionTransformerImpl : public torch::nn::Module { +class Qwen3_VisionTransformerImpl : public torch::nn::Module, + public VisionEncoderGraphAdapter { public: Qwen3_VisionTransformerImpl(const ModelContext& context) : options_(context.get_tensor_options()) { @@ -516,17 +519,36 @@ class Qwen3_VisionTransformerImpl : public torch::nn::Module { cu_seqlens_cpu.data_ptr() + cu_seqlens_cpu.numel()); std::vector deepstack_feature_lists; deepstack_feature_lists.reserve(deepstack_visual_indexes_.size()); - for (int idx = 0; idx < blocks_->size(); ++idx) { - hidden_states = layers_[idx]( - hidden_states, m_cos, m_sin, cu_seqlens, cu_seqlens_vec, idx); - auto it = std::find(deepstack_visual_indexes_.begin(), - deepstack_visual_indexes_.end(), - idx); - - if (it != deepstack_visual_indexes_.end()) { - int index = std::distance(deepstack_visual_indexes_.begin(), it); - deepstack_feature_lists.push_back( - deepstack_merger_layers_[index](hidden_states)); + + int64_t actual_tokens = hidden_states.size(0); + bool replayed = false; + if (encoder_graph_manager_ && + encoder_graph_manager_->can_replay(actual_tokens)) { + auto result = encoder_graph_manager_->replay(hidden_states, + m_cos, + m_sin, + cu_seqlens, + cu_seqlens_vec, + actual_tokens); + if (result.has_value()) { + hidden_states = result->hidden_states; + deepstack_feature_lists = std::move(result->deepstack_features); + replayed = true; + } + } + if (!replayed) { + for (int32_t idx = 0; idx < static_cast(layers_.size()); ++idx) { + hidden_states = layers_[idx]( + hidden_states, m_cos, m_sin, cu_seqlens, cu_seqlens_vec, idx); + auto it = std::find(deepstack_visual_indexes_.begin(), + deepstack_visual_indexes_.end(), + idx); + if (it != deepstack_visual_indexes_.end()) { + int32_t index = static_cast( + std::distance(deepstack_visual_indexes_.begin(), it)); + deepstack_feature_lists.push_back( + deepstack_merger_layers_[index](hidden_states)); + } } } // adapter @@ -588,6 +610,52 @@ class Qwen3_VisionTransformerImpl : public torch::nn::Module { } } + void init_encoder_graph_manager(const ModelArgs& args, + const torch::Device& device) { + std::string budgets_str = + ::xllm::ExecutionConfig::get_instance().encoder_graph_budgets(); + LOG(INFO) << "[EncoderGraph] init_encoder_graph_manager called, " + << "budgets_str='" << budgets_str << "', device=" << device; + std::vector budgets; + std::istringstream iss(budgets_str); + std::string tok; + while (std::getline(iss, tok, ',')) { + if (tok.empty()) continue; + budgets.push_back(std::stoll(tok)); + } + CHECK(!budgets.empty()) + << "[EncoderGraph] --encoder_graph_budgets has no valid values: '" + << budgets_str << "'"; + encoder_graph_manager_ = + std::make_unique(args, device, budgets); + encoder_graph_manager_->set_adapter(this); + LOG(INFO) << "[EncoderGraph] encoder_graph_manager_ created successfully, " + << "layers=" << layers_.size() + << ", deepstack_mergers=" << deepstack_merger_layers_.size() + << ", deepstack_indexes=" << deepstack_visual_indexes_.size(); + } + + int32_t num_encoder_layers() const override { + return static_cast(layers_.size()); + } + const std::vector& deepstack_indexes() const override { + return deepstack_visual_indexes_; + } + torch::Tensor forward_encoder_layer( + int32_t layer_idx, + torch::Tensor& hidden, + torch::Tensor& cos_pos, + torch::Tensor& sin_pos, + torch::Tensor& cu_seqlens, + std::vector& cu_seqlens_vec) override { + return layers_[layer_idx]( + hidden, cos_pos, sin_pos, cu_seqlens, cu_seqlens_vec, layer_idx); + } + torch::Tensor forward_deepstack_merger(int32_t deepstack_slot, + const torch::Tensor& hidden) override { + return deepstack_merger_layers_[deepstack_slot](hidden); + } + private: int hidden_size_ = 0; int num_heads_ = 0; @@ -615,6 +683,7 @@ class Qwen3_VisionTransformerImpl : public torch::nn::Module { int device_id = 0; bool is_emb_weight_loaded = false; torch::TensorOptions options_; + std::unique_ptr encoder_graph_manager_; }; TORCH_MODULE(Qwen3_VisionTransformer); @@ -827,6 +896,11 @@ class Qwen3_VLForConditionalGenerationImpl : public torch::nn::Module { language_model_->set_npu_word_embedding(npu_word_embedding); } + void init_encoder_graph_manager(const ModelArgs& args, + const torch::Device& device) { + visual_->init_encoder_graph_manager(args, device); + } + private: ModelArgs model_args_; torch::TensorOptions options_;