From 559b762f33b9e5ae546d8bec8265d1878ba21fba Mon Sep 17 00:00:00 2001 From: Enguikong Date: Wed, 5 Aug 2026 20:49:25 +0800 Subject: [PATCH 1/7] feat(npu): use FIA for eager decode attention --- xllm/core/layers/npu_torch/attention.cpp | 56 ++++++++++++++++++++---- 1 file changed, 48 insertions(+), 8 deletions(-) diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index f881966661..c0dae92e9f 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -173,14 +173,54 @@ void AttentionImpl::decoder_forward(torch::Tensor& query, tiling_data, output); } else { - // Standard PagedAttention path - xllm::kernel::npu::batch_decode(query, - k_cache, - v_cache.value_or(torch::Tensor()), - scale_, - block_table, - kv_seq_lens, - output); + // Eager decode uses non-causal FIA with one query token per sequence. + CHECK(v_cache.has_value() && v_cache->defined()) + << "FIA decode requires a value cache"; + torch::Tensor query_tnd = query.view({-1, num_heads_, head_size_}); + torch::Tensor output_tnd = output.view({-1, num_heads_, head_size_}); + + std::vector expanded_kv_seq_lens; + const std::vector* kv_seq_lens_vec = + &attn_metadata.kv_seq_lens_host_vec; + if (attn_metadata.expanded_decode.enabled) { + expanded_kv_seq_lens.reserve( + attn_metadata.expanded_decode.kv_seq_lens_host_vec.size()); + for (int32_t kv_seq_len : + attn_metadata.expanded_decode.kv_seq_lens_host_vec) { + expanded_kv_seq_lens.emplace_back(kv_seq_len); + } + kv_seq_lens_vec = &expanded_kv_seq_lens; + } + CHECK_EQ(static_cast(kv_seq_lens_vec->size()), + query_tnd.size(0)) + << "FIA decode KV lengths must match query tokens"; + + torch::Tensor key_view = + k_cache.view({k_cache.size(0), k_cache.size(1), -1}); + torch::Tensor value_view = + v_cache->view({v_cache->size(0), v_cache->size(1), -1}); + + std::vector actual_q_lens; + actual_q_lens.reserve(static_cast(query_tnd.size(0))); + for (int64_t token_idx = 0; token_idx < query_tnd.size(0); ++token_idx) { + actual_q_lens.emplace_back(token_idx + 1); + } + + auto fia_result = xllm::kernel::npu::npu_fused_infer_attention( + query_tnd, + key_view, + value_view, + /*atten_mask=*/std::nullopt, + std::make_optional(block_table), + actual_q_lens, + *kv_seq_lens_vec, + num_heads_, + num_kv_heads_, + scale_, + /*block_size=*/k_cache.size(1), + /*sparse_mode=*/0, + /*input_layout=*/"TND"); + output_tnd.copy_(std::get<0>(fia_result).view_as(output_tnd)); } } From 29723edc018982e116d58d87ceb695c042b86e87 Mon Sep 17 00:00:00 2001 From: Enguikong Date: Thu, 13 Aug 2026 19:45:12 +0800 Subject: [PATCH 2/7] feat(npu): migrate Qwen3.5 decode attention from PA to FIA Replace batch_decode (ATB PagedAttention) with npu_fused_infer_attention (aclnnFusedInferAttentionScoreV3) for all Qwen3.5 decode paths: eager, regular graph (task-group capture/replay + bucket workspace sharing), and expanded spec/MTP graph (per-token expanded kv_seq_lens). This is a prerequisite for DCP which needs per-rank softmax LSE that only FIA can emit via softmaxLseFlag. DCP itself is NOT included. Key changes: - 4-branch decode routing in attention.cpp gated by is_qwen3_5_model_type() - FIA .out wrapper + _get_max_workspace in npu_fused_infer_attention.cpp - FusedInferAttentionGraphTask with workspace signature for graph capture - GDN/FIA capture-order coexistence in acl_graph_executor_impl - Expanded kv_seq_lens for spec/MTP verify in mtp_worker_impl - --disable_fia_decode runtime switch to fall back to PA without rebuild - Comprehensive tests for FIA ops, graph capture/replay, routing isolation Verified: 4-model family (4B/9B/27B/35B) correctness, GSM8K precision (4-model untruncated layer answer diff=0), Qwen3-Next isolation (FIA=0), MTP 0.86 memory, graph perf matrix 32/36 PASS + 4 bs1/8192 waived. Co-Authored-By: Claude --- tests/core/kernels/npu/npu_xllm_ops_test.cpp | 98 +++++ .../core/runtime/acl_graph_executor_test.cpp | 27 ++ .../runtime/acl_graph_task_update_test.cpp | 403 +++++++++++++++--- .../framework/config/execution_config.cpp | 10 + xllm/core/framework/config/execution_config.h | 5 +- .../kernels/npu/npu_fused_infer_attention.cpp | 192 +++++++++ xllm/core/kernels/npu/npu_ops_api.h | 27 ++ xllm/core/layers/common/attention_metadata.h | 7 + .../common/attention_metadata_builder.cpp | 2 + xllm/core/layers/npu_torch/attention.cpp | 244 +++++++++-- xllm/core/layers/npu_torch/attention.h | 4 +- .../npu_torch/qwen3_gated_delta_net_base.cpp | 1 + .../layers/npu_torch/qwen3_next_attention.cpp | 24 +- .../layers/npu_torch/qwen3_next_attention.h | 3 + .../npu/acl_graph_task_update_context.h | 55 +++ xllm/core/runtime/acl_graph_executor_impl.cpp | 250 +++++++++-- xllm/core/runtime/acl_graph_executor_impl.h | 8 +- .../runtime/acl_graph_persistent_param.cpp | 5 +- .../core/runtime/acl_graph_persistent_param.h | 22 +- xllm/core/runtime/executor_impl.h | 2 + xllm/core/runtime/mtp_worker_impl.cpp | 17 +- 21 files changed, 1247 insertions(+), 159 deletions(-) diff --git a/tests/core/kernels/npu/npu_xllm_ops_test.cpp b/tests/core/kernels/npu/npu_xllm_ops_test.cpp index 9bf7acf78b..497fc3da32 100644 --- a/tests/core/kernels/npu/npu_xllm_ops_test.cpp +++ b/tests/core/kernels/npu/npu_xllm_ops_test.cpp @@ -244,6 +244,104 @@ TEST_F(NpuXllmOpsTest, EmbeddedInterpreterSeesOps) { .item(); } +TEST_F(NpuXllmOpsTest, + FusedInferAttentionDecodeOutMatchesEagerAcrossBlockBoundary) { + py::gil_scoped_acquire gil; + constexpr int64_t kBlockSize = 128; + constexpr int64_t kQueryHeads = 16; + constexpr int64_t kKvHeads = 4; + constexpr int64_t kHeadDim = 256; + constexpr int64_t kNumPhysicalBlocks = 4; + constexpr double kScale = 1.0 / 16.0; + const std::vector actual_seq_lengths = {1, 2, 3}; + const std::vector actual_seq_lengths_kv = {127, 128, 129}; + + torch::manual_seed(20260811); + const torch::TensorOptions cpu_float = + torch::TensorOptions().dtype(torch::kFloat32); + torch::Tensor query = torch::randn({3, kQueryHeads, kHeadDim}, cpu_float) + .to(torch::kBFloat16) + .to(torch::kPrivateUse1) + .contiguous(); + torch::Tensor key = + torch::randn({kNumPhysicalBlocks, kBlockSize, kKvHeads * kHeadDim}, + cpu_float) + .to(torch::kBFloat16) + .to(torch::kPrivateUse1) + .contiguous(); + torch::Tensor value = + torch::randn({kNumPhysicalBlocks, kBlockSize, kKvHeads * kHeadDim}, + cpu_float) + .to(torch::kBFloat16) + .to(torch::kPrivateUse1) + .contiguous(); + torch::Tensor block_table = + torch::tensor({{0, 0}, {1, 0}, {2, 3}}, + torch::TensorOptions().dtype(torch::kInt32)) + .to(torch::kPrivateUse1); + + auto eager_result = xllm::kernel::npu::npu_fused_infer_attention( + query, + key, + value, + /*atten_mask=*/std::nullopt, + std::make_optional(block_table), + actual_seq_lengths, + actual_seq_lengths_kv, + kQueryHeads, + kKvHeads, + kScale, + kBlockSize, + /*sparse_mode=*/0, + /*input_layout=*/"TND"); + torch::Tensor eager_output = std::get<0>(eager_result); + + torch::Tensor workspace = + xllm::kernel::npu::npu_fused_infer_attention_decode_get_max_workspace( + query, + key, + value, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + kQueryHeads, + kKvHeads, + kScale, + kBlockSize); + ASSERT_TRUE(workspace.defined()); + EXPECT_GT(workspace.numel(), 0); + EXPECT_EQ(workspace.device(), query.device()); + + torch::Tensor out = torch::zeros_like(eager_output); + torch::Tensor softmax_lse = torch::empty({0}, query.options()); + const void* out_data = out.const_data_ptr(); + xllm::kernel::npu::npu_fused_infer_attention_decode_out(query, + key, + value, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + kQueryHeads, + kKvHeads, + kScale, + kBlockSize, + workspace, + out, + softmax_lse); + + EXPECT_EQ(out.const_data_ptr(), out_data); + EXPECT_EQ(out.sizes(), eager_output.sizes()); + EXPECT_EQ(out.scalar_type(), torch::kBFloat16); + EXPECT_EQ(softmax_lse.numel(), 0); + const torch::Tensor actual = out.cpu().to(torch::kFloat32); + const torch::Tensor expected = eager_output.cpu().to(torch::kFloat32); + EXPECT_TRUE(torch::allclose(actual, + expected, + /*rtol=*/1e-3, + /*atol=*/2e-3)) + << "max abs diff = " << (actual - expected).abs().max().item(); +} + TEST_F(NpuXllmOpsTest, Qwen35_27B_TP4_FullAttentionMatchesReference) { py::gil_scoped_acquire gil; if (!is_ascend950_device()) { diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index ff7042cda7..60fbe250ad 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -22,6 +22,7 @@ limitations under the License. #include #include #include +#include #include #include "common/metrics.h" @@ -45,6 +46,7 @@ limitations under the License. #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/qwen3_next_attention.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" @@ -138,6 +140,7 @@ TEST(AclGraphStaticGraphTaskSignatureTest, .spec_width = 5, .block_table_width = 64, .max_kv_seq_len = 256, + .expanded_kv_seq_lens = {252, 253, 254, 255, 256}, }; const auto captured = npu::make_static_graph_task_signature(params); @@ -148,6 +151,30 @@ TEST(AclGraphStaticGraphTaskSignatureTest, EXPECT_FALSE(npu::make_static_graph_task_signature(params).has_value()); } +TEST(Qwen35FiaRoutingTest, UsesExactModelTypeWhitelist) { + const std::vector supported_model_types = { + "qwen3_5", + "qwen3_5_text", + "qwen3_5_moe", + "qwen3_5_moe_text", + "qwen3_5_mtp", + "qwen3_5_moe_mtp", + }; + for (const std::string& model_type : supported_model_types) { + EXPECT_TRUE(layer::is_qwen3_5_model_type(model_type)) << model_type; + } + + const std::vector unsupported_model_types = { + "", + "qwen3_next", + "qwen3_5_future", + "qwen3_50", + }; + for (const std::string& model_type : unsupported_model_types) { + EXPECT_FALSE(layer::is_qwen3_5_model_type(model_type)) << model_type; + } +} + namespace { const KVCache& first_full_attention_cache( const std::vector& kv_caches) { diff --git a/tests/core/runtime/acl_graph_task_update_test.cpp b/tests/core/runtime/acl_graph_task_update_test.cpp index 14bbade942..b63f11fb3f 100644 --- a/tests/core/runtime/acl_graph_task_update_test.cpp +++ b/tests/core/runtime/acl_graph_task_update_test.cpp @@ -21,11 +21,13 @@ limitations under the License. #include #include +#include #include #include "core/framework/batch/batch.h" #include "core/framework/block/block.h" #include "core/framework/block/block_manager_impl.h" +#include "core/framework/config/execution_config.h" #include "core/framework/kv_cache/kv_cache.h" #include "core/framework/model/model_args.h" #include "core/framework/model/model_output.h" @@ -34,8 +36,10 @@ limitations under the License. #include "core/framework/request/stopping_checker.h" #include "core/framework/sampling/sampling_params.h" #include "core/kernels/ops_api.h" +#include "core/layers/common/attention_metadata_builder.h" #include "core/layers/npu/npu_lm_head_impl.h" #include "core/layers/npu/npu_word_embedding_impl.h" +#include "core/layers/npu_torch/attention.h" #include "core/platform/npu/acl_graph_task_update_context.h" #include "core/runtime/acl_graph_executor_impl.h" #include "core/runtime/base_executor_impl.h" @@ -91,7 +95,11 @@ constexpr int64_t kHiddenSize = 2048; constexpr int64_t kMaxSeqLen = 256; constexpr int64_t kVocabSize = 1000; constexpr int64_t kNumBlocks = 100; -constexpr int64_t kBlockSize = 4; +constexpr int64_t kBlockSize = 128; +constexpr int64_t kAttentionNumHeads = 8; +constexpr int64_t kAttentionNumKvHeads = 1; +constexpr int64_t kAttentionHeadDim = 256; +constexpr double kAttentionScale = 1.0 / 16.0; constexpr torch::ScalarType kDtype = torch::kFloat16; @@ -99,8 +107,14 @@ constexpr torch::ScalarType kDtype = torch::kFloat16; class HybridConv1dMockLM final : public CausalLM { public: - HybridConv1dMockLM(const ModelArgs& args, const torch::Device& device) - : args_(args), device_(device) { + HybridConv1dMockLM(const ModelArgs& args, + const torch::Device& device, + bool enable_fia_decode = true, + int32_t attention_repetitions = 1) + : args_(args), + device_(device), + attention_repetitions_(attention_repetitions) { + CHECK_GT(attention_repetitions_, 0); linear_ = register_module( "linear", torch::nn::Linear(torch::nn::LinearOptions(kHiddenSize, kHiddenSize))); @@ -120,6 +134,24 @@ class HybridConv1dMockLM final : public CausalLM { torch::randn({kMaxSeqLen, kHiddenSize}, torch::dtype(kDtype).device(device))); + if (enable_fia_decode) { + attention_ = + register_module("attention", + layer::Attention(kAttentionNumHeads, + kAttentionHeadDim, + kAttentionScale, + kAttentionNumKvHeads, + /*sliding_window=*/-1, + /*enable_fia_decode=*/true)); + } else { + attention_ = register_module("attention", + layer::Attention(kAttentionNumHeads, + kAttentionHeadDim, + kAttentionScale, + kAttentionNumKvHeads, + /*sliding_window=*/-1)); + } + this->to(device); } @@ -190,6 +222,7 @@ class HybridConv1dMockLM final : public CausalLM { task.pad_slot_id = npu::kCausalConv1dGraphPadSlotId; task.run_mode = npu::kCausalConv1dRunModeUpdate; task.branch = branch; + task.capture_order = graph_context->next_capture_order++; task.handle = handle; task.event = std::move(event); graph_context->causal_conv1d_tasks.emplace_back(std::move(task)); @@ -222,6 +255,51 @@ class HybridConv1dMockLM final : public CausalLM { break; } + for (auto& kv_cache : kv_caches) { + if (kv_cache.empty() || !kv_cache.get_k_cache().defined()) { + continue; + } + + for (int32_t attention_index = 0; + attention_index < attention_repetitions_; + ++attention_index) { + layer::AttentionMetadata attn_metadata = + layer::AttentionMetadataBuilder::build(params, + /*enable_mla=*/false, + /*attn_mask=*/std::nullopt, + device_); + torch::Tensor query = hidden.to(torch::kBFloat16).contiguous(); + torch::Tensor key = query + .slice(/*dim=*/1, + /*start=*/0, + kAttentionNumKvHeads * kAttentionHeadDim) + .contiguous(); + torch::Tensor value = key.clone(); + torch::Tensor attention_output = std::get<0>( + attention_->forward(attn_metadata, query, key, value, kv_cache)); + hidden = hidden + attention_output.to(hidden.scalar_type()); + } + + if (register_graph_task) { + saw_causal_conv_graph_task_ |= + !graph_context->causal_conv1d_tasks.empty(); + saw_fia_graph_task_ |= + !graph_context->fused_infer_attention_tasks.empty(); + fia_graph_task_count_ = + graph_context->fused_infer_attention_tasks.size(); + all_fia_graph_tasks_share_workspace_ = fia_graph_task_count_ > 1; + for (size_t task_index = 1; task_index < fia_graph_task_count_; + ++task_index) { + all_fia_graph_tasks_share_workspace_ &= + graph_context->fused_infer_attention_tasks[task_index] + .workspace.data_ptr() == + graph_context->fused_infer_attention_tasks.front() + .workspace.data_ptr(); + } + } + break; + } + hidden = linear_->forward(hidden); return ModelOutput(hidden); @@ -242,6 +320,17 @@ class HybridConv1dMockLM final : public CausalLM { void load_model(std::unique_ptr loader) override {} torch::Device device() const override { return device_; } + bool saw_causal_conv_and_fia_graph_tasks() const { + return saw_causal_conv_graph_task_ && saw_fia_graph_task_; + } + bool saw_causal_conv_graph_task() const { + return saw_causal_conv_graph_task_; + } + bool saw_fia_graph_task() const { return saw_fia_graph_task_; } + size_t fia_graph_task_count() const { return fia_graph_task_count_; } + bool all_fia_graph_tasks_share_workspace() const { + return all_fia_graph_tasks_share_workspace_; + } void prepare_expert_weight(int32_t, const std::vector&) override {} void update_expert_weight(int32_t) override {} layer::NpuLmHead get_npu_lm_head() override { @@ -257,9 +346,15 @@ class HybridConv1dMockLM final : public CausalLM { ModelArgs args_; torch::Device device_; torch::nn::Linear linear_{nullptr}; + layer::Attention attention_{nullptr}; torch::Tensor conv_weight_; torch::Tensor token_embedding_table_; torch::Tensor pos_embedding_table_; + int32_t attention_repetitions_ = 1; + bool saw_causal_conv_graph_task_ = false; + bool saw_fia_graph_task_ = false; + size_t fia_graph_task_count_ = 0; + bool all_fia_graph_tasks_share_workspace_ = false; }; class AclGraphTaskUpdateTest : public ::testing::Test { @@ -267,18 +362,34 @@ class AclGraphTaskUpdateTest : public ::testing::Test { void SetUp() override { sequences_.reserve(100); + auto& execution_config = ExecutionConfig::get_instance(); + original_enable_graph_ = execution_config.enable_graph(); + original_enable_graph_double_buffer_ = + execution_config.enable_graph_double_buffer(); + original_enable_graph_mode_decode_no_padding_ = + execution_config.enable_graph_mode_decode_no_padding(); + original_acl_graph_decode_batch_size_limit_ = + execution_config.acl_graph_decode_batch_size_limit(); + execution_config.enable_graph(true); + execution_config.enable_graph_double_buffer(false); + execution_config.enable_graph_mode_decode_no_padding(false); + execution_config.acl_graph_decode_batch_size_limit(32); + model_args_.model_type("test_hybrid_model"); model_args_.dtype("float16"); model_args_.hidden_size(kHiddenSize); model_args_.max_position_embeddings(kMaxSeqLen); model_args_.vocab_size(kVocabSize); model_args_.n_layers(2); + model_args_.n_heads(kAttentionNumHeads); + model_args_.n_kv_heads(kAttentionNumKvHeads); + model_args_.head_dim(kAttentionHeadDim); model_args_.layer_types({"linear_attention", "full_attention"}); device_ = std::make_unique("npu:0"); options_.num_decoding_tokens(1); options_.block_size(kBlockSize); - options_.max_seqs_per_batch(16); + options_.max_seqs_per_batch(32); model_ = std::make_unique(model_args_, *device_); @@ -289,7 +400,7 @@ class AclGraphTaskUpdateTest : public ::testing::Test { sampling_param_.frequency_penalty = 0.0f; stopping_checker_.set_max_generated_tokens(20); - seq_params_.seq_capacity = 100; + seq_params_.seq_capacity = kMaxSeqLen; seq_params_.stopping_checker = &stopping_checker_; seq_params_.sampling_param = &sampling_param_; seq_params_.skip_special_tokens = true; @@ -302,7 +413,17 @@ class AclGraphTaskUpdateTest : public ::testing::Test { mm_data_ = MMData(); } - void TearDown() override {} + void TearDown() override { + reset_sequences(); + auto& execution_config = ExecutionConfig::get_instance(); + execution_config.enable_graph(original_enable_graph_); + execution_config.enable_graph_double_buffer( + original_enable_graph_double_buffer_); + execution_config.enable_graph_mode_decode_no_padding( + original_enable_graph_mode_decode_no_padding_); + execution_config.acl_graph_decode_batch_size_limit( + original_acl_graph_decode_batch_size_limit_); + } void reset_sequences() { for (auto& sequence : sequences_) { @@ -328,9 +449,10 @@ class AclGraphTaskUpdateTest : public ::testing::Test { kv_caches.emplace_back( LinearAttentionKVCacheTensors{conv_cache, ssm_cache}); - auto k_cache = torch::randn({kNumBlocks, kBlockSize * kHiddenSize}, - torch::dtype(kDtype).device(*device_)); - auto v_cache = k_cache.clone(); + auto k_cache = torch::zeros( + {kNumBlocks, kBlockSize, kAttentionNumKvHeads, kAttentionHeadDim}, + torch::dtype(torch::kBFloat16).device(*device_)); + auto v_cache = torch::zeros_like(k_cache); kv_caches.emplace_back(KVCacheTensors{k_cache, v_cache}); return kv_caches; } @@ -404,6 +526,87 @@ class AclGraphTaskUpdateTest : public ::testing::Test { return batch; } + std::vector> create_mixed_boundary_prompts( + uint32_t batch_size, + int32_t token_seed) { + const std::vector prompt_lengths = {126, 127, 128}; + std::vector> prompts; + prompts.reserve(batch_size); + for (uint32_t batch_index = 0; batch_index < batch_size; ++batch_index) { + const int64_t prompt_length = + prompt_lengths[batch_index % prompt_lengths.size()]; + std::vector prompt; + prompt.reserve(static_cast(prompt_length)); + for (int64_t token_index = 0; token_index < prompt_length; + ++token_index) { + prompt.emplace_back((token_seed + static_cast(batch_index) + + static_cast(token_index)) % + kVocabSize); + } + prompts.emplace_back(std::move(prompt)); + } + return prompts; + } + + void expect_fia_padding_replay_matches_eager(uint32_t capture_batch_size, + uint32_t replay_batch_size, + uint32_t expected_bucket) { + auto capture_prompts = + create_mixed_boundary_prompts(capture_batch_size, /*token_seed=*/10); + auto capture_batch = + create_decode_batch_with_prompts(capture_prompts, /*token_seed=*/100); + auto capture_fi = capture_batch->prepare_forward_input( + options_.num_decoding_tokens(), 0, model_args_); + capture_fi = capture_fi.to(*device_, kDtype); + populate_query_start_loc(capture_fi.input_params); + + auto kv_graph = create_hybrid_kv_caches(); + auto graph_exec = std::make_unique( + model_.get(), model_args_, *device_, options_); + EXPECT_EQ(graph_exec->bucket_num_tokens_for_test(capture_batch_size), + expected_bucket); + graph_exec->run({capture_fi.token_ids}, + {capture_fi.positions}, + kv_graph, + {capture_fi.input_params}); + ASSERT_TRUE(model_->saw_causal_conv_and_fia_graph_tasks()); + + reset_sequences(); + auto replay_prompts = + create_mixed_boundary_prompts(replay_batch_size, /*token_seed=*/200); + auto replay_batch = + create_decode_batch_with_prompts(replay_prompts, /*token_seed=*/300); + auto replay_fi = replay_batch->prepare_forward_input( + options_.num_decoding_tokens(), 0, model_args_); + replay_fi = replay_fi.to(*device_, kDtype); + populate_query_start_loc(replay_fi.input_params); + + auto kv_eager = clone_kv_caches(kv_graph); + auto graph_out = graph_exec->run({replay_fi.token_ids}, + {replay_fi.positions}, + kv_graph, + {replay_fi.input_params}); + auto eager_out = model_->forward({replay_fi.token_ids}, + {replay_fi.positions}, + kv_eager, + {replay_fi.input_params}); + + const int64_t real_tokens = static_cast(replay_batch_size); + ASSERT_EQ(graph_out.hidden_states.size(0), real_tokens); + torch::Tensor graph_real = + graph_out.hidden_states.slice(0, 0, real_tokens).to(torch::kFloat32); + torch::Tensor eager_real = + eager_out.hidden_states.slice(0, 0, real_tokens).to(torch::kFloat32); + EXPECT_TRUE(torch::allclose(eager_real, + graph_real, + /*rtol=*/1e-2, + /*atol=*/1e-2)) + << "FIA padding replay mismatch for capture_bs=" << capture_batch_size + << ", replay_bs=" << replay_batch_size << ", bucket=" << expected_bucket + << ", max_abs_diff=" + << (eager_real - graph_real).abs().max().item(); + } + void setup_spec_verify_input(ForwardInput& fi, int32_t num_sequences, int32_t num_spec_tokens) { @@ -452,6 +655,30 @@ class AclGraphTaskUpdateTest : public ::testing::Test { fi.input_params.graph.expanded_kv_seq_lens = torch::tensor(expanded_kv_vec, torch::kInt32).to(*device_); + torch::Tensor host_block_tables = + fi.input_params.attention.host.block_tables.contiguous(); + auto host_block_table_accessor = host_block_tables.accessor(); + std::vector expanded_cache_slots; + expanded_cache_slots.reserve(static_cast(total_tokens)); + for (int32_t sequence_index = 0; sequence_index < num_sequences; + ++sequence_index) { + const int32_t kv_len = + fi.input_params.attention.host + .kv_seq_lens[static_cast(sequence_index)]; + for (int32_t token_index = 0; token_index < num_spec_tokens; + ++token_index) { + const int32_t position = kv_len + token_index; + const int32_t logical_block = position / kBlockSize; + const int32_t block_offset = position % kBlockSize; + const int32_t physical_block = + host_block_table_accessor[sequence_index][logical_block]; + expanded_cache_slots.emplace_back(physical_block * kBlockSize + + block_offset); + } + } + fi.input_params.attention.host.new_cache_slots = + std::move(expanded_cache_slots); + auto block_tables = fi.input_params.attention.device.block_tables; int64_t block_table_stride = block_tables.size(1); auto expanded_bt = @@ -489,6 +716,10 @@ class AclGraphTaskUpdateTest : public ::testing::Test { MMData mm_data_; std::vector sequences_; IncrementalDecoder fake_decoder_ = IncrementalDecoder("", 1, false, false); + bool original_enable_graph_ = false; + bool original_enable_graph_double_buffer_ = true; + bool original_enable_graph_mode_decode_no_padding_ = false; + int32_t original_acl_graph_decode_batch_size_limit_ = 16; }; TEST_F(AclGraphTaskUpdateTest, CaptureReplayVsEagerDecodeBranch) { @@ -522,6 +753,95 @@ TEST_F(AclGraphTaskUpdateTest, CaptureReplayVsEagerDecodeBranch) { << "Decode branch: eager vs graph mismatch"; } +TEST_F(AclGraphTaskUpdateTest, + MultipleFiaInvocationsShareWorkspaceWithinBucket) { + auto shared_workspace_model = + std::make_unique(model_args_, + *device_, + /*enable_fia_decode=*/true, + /*attention_repetitions=*/2); + auto batch = create_decode_batch(/*batch_size=*/2); + ASSERT_FALSE(batch->empty()); + + auto forward_input = batch->prepare_forward_input( + options_.num_decoding_tokens(), 0, model_args_); + forward_input = forward_input.to(*device_, kDtype); + populate_query_start_loc(forward_input.input_params); + + auto kv_eager = create_hybrid_kv_caches(); + auto eager_out = + shared_workspace_model->forward({forward_input.token_ids}, + {forward_input.positions}, + kv_eager, + {forward_input.input_params}); + + auto kv_graph = create_hybrid_kv_caches(); + auto graph_exec = std::make_unique( + shared_workspace_model.get(), model_args_, *device_, options_); + auto graph_out = graph_exec->run({forward_input.token_ids}, + {forward_input.positions}, + kv_graph, + {forward_input.input_params}); + + EXPECT_EQ(shared_workspace_model->fia_graph_task_count(), 2); + EXPECT_TRUE(shared_workspace_model->all_fia_graph_tasks_share_workspace()); + EXPECT_TRUE(torch::allclose(eager_out.hidden_states.to(torch::kFloat32), + graph_out.hidden_states.to(torch::kFloat32), + /*rtol=*/1e-2, + /*atol=*/1e-2)) + << "Shared-workspace FIA capture/replay must match eager, max_abs_diff=" + << (eager_out.hidden_states.to(torch::kFloat32) - + graph_out.hidden_states.to(torch::kFloat32)) + .abs() + .max() + .item(); +} + +TEST_F(AclGraphTaskUpdateTest, + NonQwenDefaultAttentionKeepsPagedAttentionInEagerAndGraph) { + ModelArgs non_qwen_args = model_args_; + non_qwen_args.model_type("minimax_m2"); + non_qwen_args.dtype("bfloat16"); + auto non_qwen_model = std::make_unique( + non_qwen_args, *device_, /*enable_fia_decode=*/false); + + auto batch = create_decode_batch(/*batch_size=*/2); + ASSERT_FALSE(batch->empty()); + auto forward_input = batch->prepare_forward_input( + options_.num_decoding_tokens(), 0, non_qwen_args); + forward_input = forward_input.to(*device_, kDtype); + populate_query_start_loc(forward_input.input_params); + + auto kv_eager = create_hybrid_kv_caches(); + auto eager_out = non_qwen_model->forward({forward_input.token_ids}, + {forward_input.positions}, + kv_eager, + {forward_input.input_params}); + + auto kv_graph = create_hybrid_kv_caches(); + auto graph_exec = std::make_unique( + non_qwen_model.get(), non_qwen_args, *device_, options_); + auto graph_out = graph_exec->run({forward_input.token_ids}, + {forward_input.positions}, + kv_graph, + {forward_input.input_params}); + + EXPECT_TRUE(non_qwen_model->saw_causal_conv_graph_task()); + EXPECT_FALSE(non_qwen_model->saw_fia_graph_task()); + EXPECT_EQ(eager_out.hidden_states.sizes(), graph_out.hidden_states.sizes()); + EXPECT_TRUE(torch::allclose(eager_out.hidden_states.to(torch::kFloat32), + graph_out.hidden_states.to(torch::kFloat32), + /*rtol=*/1e-2, + /*atol=*/1e-2)) + << "Non-Qwen default Attention must keep PA eager/graph behavior, " + << "max_abs_diff=" + << (eager_out.hidden_states.to(torch::kFloat32) - + graph_out.hidden_states.to(torch::kFloat32)) + .abs() + .max() + .item(); +} + TEST_F(AclGraphTaskUpdateTest, ReplayWithDifferentParamsProducesDifferentOutputs) { std::vector> prompts_run1 = {{1, 3, 5, 7}, {2, 4, 6, 8}}; @@ -574,62 +894,19 @@ TEST_F(AclGraphTaskUpdateTest, } TEST_F(AclGraphTaskUpdateTest, PaddingBatchDoesNotPolluteRealSequences) { - constexpr uint32_t kCaptureBatchSize = 4; - constexpr uint32_t kReplayBatchSize = 3; - - auto capture_batch = create_decode_batch(kCaptureBatchSize); - auto capture_fi = capture_batch->prepare_forward_input( - options_.num_decoding_tokens(), 0, model_args_); - capture_fi = capture_fi.to(*device_, kDtype); - populate_query_start_loc(capture_fi.input_params); - - auto kv_graph = create_hybrid_kv_caches(); - auto graph_exec = std::make_unique( - model_.get(), model_args_, *device_, options_); - graph_exec->run({capture_fi.token_ids}, - {capture_fi.positions}, - kv_graph, - {capture_fi.input_params}); - - reset_sequences(); - - auto replay_batch = create_decode_batch(kReplayBatchSize); - auto replay_fi = replay_batch->prepare_forward_input( - options_.num_decoding_tokens(), 0, model_args_); - replay_fi = replay_fi.to(*device_, kDtype); - populate_query_start_loc(replay_fi.input_params); - - auto kv_eager = clone_kv_caches(kv_graph); - - auto graph_out = graph_exec->run({replay_fi.token_ids}, - {replay_fi.positions}, - kv_graph, - {replay_fi.input_params}); - auto eager_out = model_->forward({replay_fi.token_ids}, - {replay_fi.positions}, - kv_eager, - {replay_fi.input_params}); - - const int64_t real_tokens = - static_cast(kReplayBatchSize) * options_.num_decoding_tokens(); - EXPECT_EQ(graph_out.hidden_states.size(0), real_tokens); - - auto eager_real = - eager_out.hidden_states.slice(0, 0, real_tokens).to(torch::kFloat32); - auto graph_real = - graph_out.hidden_states.slice(0, 0, real_tokens).to(torch::kFloat32); + expect_fia_padding_replay_matches_eager( + /*capture_batch_size=*/4, /*replay_batch_size=*/3, /*expected_bucket=*/4); +} - auto diff = (eager_real - graph_real).abs(); - float max_diff = diff.max().item(); - float mean_diff = diff.mean().item(); - LOG(INFO) << "Padding test max_diff=" << max_diff - << " mean_diff=" << mean_diff; +TEST_F(AclGraphTaskUpdateTest, FiaPaddingReplaysAcrossBucketEight) { + expect_fia_padding_replay_matches_eager( + /*capture_batch_size=*/5, /*replay_batch_size=*/7, /*expected_bucket=*/8); +} - EXPECT_TRUE(torch::allclose(eager_real, - graph_real, - /*rtol=*/1e-2, - /*atol=*/1e-2)) - << "Padding batch pollutes real sequence output, max_diff=" << max_diff; +TEST_F(AclGraphTaskUpdateTest, FiaPaddingReplaysAcrossBucketThirtyTwo) { + expect_fia_padding_replay_matches_eager(/*capture_batch_size=*/17, + /*replay_batch_size=*/17, + /*expected_bucket=*/32); } TEST_F(AclGraphTaskUpdateTest, CaptureReplayVsEagerSpecVerifyBranch) { diff --git a/xllm/core/framework/config/execution_config.cpp b/xllm/core/framework/config/execution_config.cpp index 6d2d657011..33489c6309 100644 --- a/xllm/core/framework/config/execution_config.cpp +++ b/xllm/core/framework/config/execution_config.cpp @@ -95,6 +95,12 @@ DEFINE_string( "aclgraph (NPU decode graph with eager prefill), " "or any torch.compile backend name."); +DEFINE_bool( + disable_fia_decode, + false, + "When true, Qwen3.5 decode attention uses PagedAttention instead of FIA. " + "Useful as a runtime rollback switch without rebuilding."); + namespace xllm { void ExecutionConfig::from_flags() { @@ -112,6 +118,7 @@ void ExecutionConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(output_shm_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(random_seed); XLLM_CONFIG_ASSIGN_FROM_FLAG(python_graph_backend); + XLLM_CONFIG_ASSIGN_FROM_FLAG(disable_fia_decode); } void ExecutionConfig::from_json(const JsonReader& json) { @@ -129,6 +136,7 @@ void ExecutionConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(output_shm_size); XLLM_CONFIG_ASSIGN_FROM_JSON(random_seed); XLLM_CONFIG_ASSIGN_FROM_JSON(python_graph_backend); + XLLM_CONFIG_ASSIGN_FROM_JSON(disable_fia_decode); } void ExecutionConfig::append_config_json( @@ -162,6 +170,8 @@ void ExecutionConfig::append_config_json( config_json, default_config, random_seed); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, python_graph_backend); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, disable_fia_decode); } ExecutionConfig& ExecutionConfig::get_instance() { diff --git a/xllm/core/framework/config/execution_config.h b/xllm/core/framework/config/execution_config.h index e9351083fc..fe403a9584 100644 --- a/xllm/core/framework/config/execution_config.h +++ b/xllm/core/framework/config/execution_config.h @@ -54,7 +54,8 @@ class ExecutionConfig final { "input_shm_size", "output_shm_size", "random_seed", - "python_graph_backend"}}; + "python_graph_backend", + "disable_fia_decode"}}; return kOptionCategory; } @@ -85,6 +86,8 @@ class ExecutionConfig final { PROPERTY(int32_t, random_seed) = -1; PROPERTY(std::string, python_graph_backend) = "off"; + + PROPERTY(bool, disable_fia_decode) = false; }; } // namespace xllm diff --git a/xllm/core/kernels/npu/npu_fused_infer_attention.cpp b/xllm/core/kernels/npu/npu_fused_infer_attention.cpp index 097c029721..dfffef423c 100644 --- a/xllm/core/kernels/npu/npu_fused_infer_attention.cpp +++ b/xllm/core/kernels/npu/npu_fused_infer_attention.cpp @@ -13,8 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include +#include + +#include #include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp" #include "core/kernels/npu/npu_ops_api.h" @@ -24,6 +28,61 @@ namespace { constexpr int64_t kSwaIntMax = 2147483647; +using OptionalTensorRef = const std::optional&; +using OptionalSymIntArrayRef = c10::OptionalArrayRef; +using FusedInferAttentionOutSignature = + std::tuple(const at::Tensor&, + const at::Tensor&, + const at::Tensor&, + OptionalTensorRef, + OptionalTensorRef, + OptionalSymIntArrayRef, + OptionalSymIntArrayRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalSymIntArrayRef, + OptionalTensorRef, + OptionalTensorRef, + OptionalTensorRef, + int64_t, + double, + int64_t, + int64_t, + c10::string_view, + int64_t, + int64_t, + int64_t, + int64_t, + int64_t, + int64_t, + int64_t, + bool, + OptionalTensorRef, + c10::ArrayRef); + +std::vector to_sym_ints(const std::vector& values) { + std::vector sym_ints; + sym_ints.reserve(values.size()); + for (int64_t value : values) { + sym_ints.emplace_back(value); + } + return sym_ints; +} + torch::Tensor ascend950_packed_causal_attention( const torch::Tensor& query, const torch::Tensor& key, @@ -151,6 +210,139 @@ std::optional to_optional_tensor( namespace xllm::kernel::npu { +torch::Tensor npu_fused_infer_attention_decode_get_max_workspace( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const torch::Tensor& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size) { + std::vector actual_seq_lengths_sym = + to_sym_ints(actual_seq_lengths); + std::vector actual_seq_lengths_kv_sym = + to_sym_ints(actual_seq_lengths_kv); + const std::optional none_tensor = std::nullopt; + const at::OptionalSymIntArrayRef none_int_array = std::nullopt; + + return at_npu::native::custom_ops:: + _npu_fused_infer_attention_score_get_max_workspace( + query, + key, + value, + none_tensor, + none_tensor, + actual_seq_lengths_sym, + actual_seq_lengths_kv_sym, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + block_table, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_int_array, + none_tensor, + none_tensor, + none_tensor, + num_heads, + scale, + kSwaIntMax, + /*next_tokens=*/0, + "TND", + num_key_value_heads, + /*sparse_mode=*/0, + /*inner_precise=*/0, + block_size, + /*antiquant_mode=*/0, + /*key_antiquant_mode=*/0, + /*value_antiquant_mode=*/0, + /*softmax_lse_flag=*/false); +} + +void npu_fused_infer_attention_decode_out( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const torch::Tensor& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size, + const torch::Tensor& workspace, + torch::Tensor& output, + torch::Tensor& softmax_lse) { + std::vector actual_seq_lengths_sym = + to_sym_ints(actual_seq_lengths); + std::vector actual_seq_lengths_kv_sym = + to_sym_ints(actual_seq_lengths_kv); + const std::optional none_tensor = std::nullopt; + const at::OptionalSymIntArrayRef none_int_array = std::nullopt; + const std::optional workspace_tensor = workspace; + const std::array outputs = {output, softmax_lse}; + static const auto op = + c10::Dispatcher::singleton() + .findSchemaOrThrow("npu::npu_fused_infer_attention_score", "out") + .typed(); + + op.call(query, + key, + value, + none_tensor, + none_tensor, + actual_seq_lengths_sym, + actual_seq_lengths_kv_sym, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + block_table, + none_tensor, + none_tensor, + none_tensor, + none_tensor, + none_int_array, + none_tensor, + none_tensor, + none_tensor, + num_heads, + scale, + kSwaIntMax, + /*next_tokens=*/0, + "TND", + num_key_value_heads, + /*sparse_mode=*/0, + /*inner_precise=*/0, + block_size, + /*antiquant_mode=*/0, + /*key_antiquant_mode=*/0, + /*value_antiquant_mode=*/0, + /*softmax_lse_flag=*/false, + workspace_tensor, + outputs); +} + std::tuple npu_fused_infer_attention( const torch::Tensor& query, const torch::Tensor& key, diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index e757e41f3a..4589727cf3 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -74,6 +74,33 @@ std::tuple npu_fused_infer_attention( bool softmax_lse_flag = false, bool is_causal = true); +torch::Tensor npu_fused_infer_attention_decode_get_max_workspace( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const torch::Tensor& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size); + +void npu_fused_infer_attention_decode_out( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const torch::Tensor& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size, + const torch::Tensor& workspace, + torch::Tensor& output, + torch::Tensor& softmax_lse); + void batch_chunked_paged_prefill(const torch::Tensor& query, const torch::Tensor& k_cache, const torch::Tensor& v_cache, diff --git a/xllm/core/layers/common/attention_metadata.h b/xllm/core/layers/common/attention_metadata.h index 2b02963dcb..ffe44ee205 100644 --- a/xllm/core/layers/common/attention_metadata.h +++ b/xllm/core/layers/common/attention_metadata.h @@ -31,6 +31,12 @@ namespace ffi = tvm::ffi; #include "dsa_metadata.h" #include "layers/common/kv_shard_batch_metadata.h" +#if defined(USE_NPU) +namespace xllm::npu { +class AclGraphTaskUpdateContext; +} +#endif + namespace xllm::layer { struct ExpandedDecodeMetadata { @@ -207,6 +213,7 @@ struct AttentionMetadata { #if defined(USE_NPU) // for npu + std::shared_ptr acl_graph_task_update_context; torch::Tensor q_seq_lens_host; torch::Tensor kv_seq_lens_host; // For ACL graph execution - fixed-address device tiling data for diff --git a/xllm/core/layers/common/attention_metadata_builder.cpp b/xllm/core/layers/common/attention_metadata_builder.cpp index acf653afce..e98eeb573d 100644 --- a/xllm/core/layers/common/attention_metadata_builder.cpp +++ b/xllm/core/layers/common/attention_metadata_builder.cpp @@ -355,6 +355,8 @@ AttentionMetadata build_attention_metadata( #endif #if defined(USE_NPU) + attn_metadata.acl_graph_task_update_context = + params.graph.acl_graph_task_update_context; // Determine if we should use ACL graph mode: // - --enable_graph=true // - Must be decode phase or spec-verify chunked prefill diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index c0dae92e9f..341efddefd 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -15,9 +15,150 @@ limitations under the License. #include "attention.h" +#include +#include +#include + +#include "core/platform/npu/acl_graph_task_update_context.h" #include "kernels/npu/npu_ops_api.h" #include "kernels/ops_api.h" +namespace { +std::vector make_decode_actual_seq_lengths(int64_t num_tokens) { + std::vector actual_seq_lengths; + actual_seq_lengths.reserve(static_cast(num_tokens)); + for (int64_t token_idx = 0; token_idx < num_tokens; ++token_idx) { + actual_seq_lengths.emplace_back(token_idx + 1); + } + return actual_seq_lengths; +} + +xllm::npu::FusedInferAttentionWorkspaceSignature +make_fused_infer_attention_workspace_signature( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const torch::Tensor& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size) { + return xllm::npu::FusedInferAttentionWorkspaceSignature{ + .query_dtype = query.scalar_type(), + .key_dtype = key.scalar_type(), + .value_dtype = value.scalar_type(), + .block_table_dtype = block_table.scalar_type(), + .device_index = query.device().index(), + .query_shape = query.sizes().vec(), + .key_shape = key.sizes().vec(), + .value_shape = value.sizes().vec(), + .block_table_shape = block_table.sizes().vec(), + .actual_seq_lengths = actual_seq_lengths, + .actual_seq_lengths_kv = actual_seq_lengths_kv, + .num_heads = num_heads, + .num_key_value_heads = num_key_value_heads, + .block_size = block_size, + .scale = scale, + }; +} + +void run_fused_infer_attention_graph( + const std::shared_ptr& graph_context, + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + const torch::Tensor& block_table, + const std::vector& actual_seq_lengths, + const std::vector& actual_seq_lengths_kv, + int64_t num_heads, + int64_t num_key_value_heads, + double scale, + int64_t block_size, + xllm::npu::FusedInferAttentionGraphBranch branch, + torch::Tensor& output) { + CHECK(graph_context != nullptr && graph_context->capturing) + << "FIA graph update can only be registered during capture"; + + const xllm::npu::FusedInferAttentionWorkspaceSignature workspace_signature = + make_fused_infer_attention_workspace_signature(query, + key, + value, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + num_heads, + num_key_value_heads, + scale, + block_size); + torch::Tensor workspace = graph_context->fused_infer_attention_workspace; + if (workspace.defined()) { + CHECK(graph_context->fused_infer_attention_workspace_signature.has_value()); + CHECK(graph_context->fused_infer_attention_workspace_signature.value() == + workspace_signature) + << "FIA graph layers in one bucket require different workspaces"; + } else { + workspace = + xllm::kernel::npu::npu_fused_infer_attention_decode_get_max_workspace( + query, + key, + value, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + num_heads, + num_key_value_heads, + scale, + block_size); + CHECK(workspace.defined()) << "FIA graph workspace must be defined"; + graph_context->fused_infer_attention_workspace_signature = + workspace_signature; + graph_context->fused_infer_attention_workspace = workspace; + } + torch::Tensor softmax_lse = torch::empty({0}, query.options()); + c10_npu::NPUStream stream = c10_npu::getCurrentNPUStream(); + auto event = std::make_shared(ACL_EVENT_EXTERNAL); + event->block(stream); + event->reset(stream); + + c10_npu::graph_task_group_begin(stream); + xllm::kernel::npu::npu_fused_infer_attention_decode_out(query, + key, + value, + block_table, + actual_seq_lengths, + actual_seq_lengths_kv, + num_heads, + num_key_value_heads, + scale, + block_size, + workspace, + output, + softmax_lse); + c10_npu::NPUTaskGroupHandle handle = c10_npu::graph_task_group_end(stream); + + xllm::npu::FusedInferAttentionGraphTask task; + task.output = output; + task.softmax_lse = std::move(softmax_lse); + task.query = query; + task.key = key; + task.value = value; + task.block_table = block_table; + task.workspace = std::move(workspace); + task.actual_seq_lengths = actual_seq_lengths; + task.num_heads = num_heads; + task.num_key_value_heads = num_key_value_heads; + task.scale = scale; + task.block_size = block_size; + task.branch = branch; + task.capture_order = graph_context->next_capture_order++; + task.handle = handle; + task.event = std::move(event); + graph_context->fused_infer_attention_tasks.emplace_back(std::move(task)); +} +} // namespace + namespace xllm { namespace layer { @@ -25,12 +166,14 @@ AttentionImpl::AttentionImpl(int64_t num_heads, int64_t head_size, float scale, int64_t num_kv_heads, - int64_t sliding_window) + int64_t sliding_window, + bool enable_fia_decode) : num_heads_(num_heads), head_size_(head_size), num_kv_heads_(num_kv_heads), sliding_window_(sliding_window), - scale_(scale) { + scale_(scale), + enable_fia_decode_(enable_fia_decode) { if (sliding_window_ > -1) { sliding_window_ = sliding_window_ - 1; } @@ -161,9 +304,23 @@ void AttentionImpl::decoder_forward(torch::Tensor& query, kv_seq_lens = attn_metadata.kv_seq_lens; } - if (tiling_data.defined()) { - // Use CustomPagedAttention for ACL graph mode to avoid .to(kCPU) operations + std::vector expanded_kv_seq_lens; + const std::vector* kv_seq_lens_vec = + &attn_metadata.kv_seq_lens_host_vec; + if (attn_metadata.expanded_decode.enabled) { + expanded_kv_seq_lens.reserve( + attn_metadata.expanded_decode.kv_seq_lens_host_vec.size()); + for (int32_t kv_seq_len : + attn_metadata.expanded_decode.kv_seq_lens_host_vec) { + expanded_kv_seq_lens.emplace_back(kv_seq_len); + } + kv_seq_lens_vec = &expanded_kv_seq_lens; + } + const bool use_fia_graph_decode = + enable_fia_decode_ && + (!attn_metadata.is_spec_verify || attn_metadata.expanded_decode.enabled); + if (tiling_data.defined() && !use_fia_graph_decode) { xllm::kernel::npu::batch_decode_acl_graph(query, k_cache, v_cache.value_or(torch::Tensor()), @@ -172,40 +329,53 @@ void AttentionImpl::decoder_forward(torch::Tensor& query, kv_seq_lens, tiling_data, output); - } else { - // Eager decode uses non-causal FIA with one query token per sequence. - CHECK(v_cache.has_value() && v_cache->defined()) - << "FIA decode requires a value cache"; - torch::Tensor query_tnd = query.view({-1, num_heads_, head_size_}); - torch::Tensor output_tnd = output.view({-1, num_heads_, head_size_}); - - std::vector expanded_kv_seq_lens; - const std::vector* kv_seq_lens_vec = - &attn_metadata.kv_seq_lens_host_vec; - if (attn_metadata.expanded_decode.enabled) { - expanded_kv_seq_lens.reserve( - attn_metadata.expanded_decode.kv_seq_lens_host_vec.size()); - for (int32_t kv_seq_len : - attn_metadata.expanded_decode.kv_seq_lens_host_vec) { - expanded_kv_seq_lens.emplace_back(kv_seq_len); - } - kv_seq_lens_vec = &expanded_kv_seq_lens; - } - CHECK_EQ(static_cast(kv_seq_lens_vec->size()), - query_tnd.size(0)) - << "FIA decode KV lengths must match query tokens"; - - torch::Tensor key_view = - k_cache.view({k_cache.size(0), k_cache.size(1), -1}); - torch::Tensor value_view = - v_cache->view({v_cache->size(0), v_cache->size(1), -1}); - - std::vector actual_q_lens; - actual_q_lens.reserve(static_cast(query_tnd.size(0))); - for (int64_t token_idx = 0; token_idx < query_tnd.size(0); ++token_idx) { - actual_q_lens.emplace_back(token_idx + 1); - } + return; + } + if (!tiling_data.defined() && !enable_fia_decode_) { + xllm::kernel::npu::batch_decode(query, + k_cache, + v_cache.value_or(torch::Tensor()), + scale_, + block_table, + kv_seq_lens, + output); + return; + } + + CHECK(v_cache.has_value() && v_cache->defined()) + << "FIA decode requires a value cache"; + CHECK(block_table.defined()) << "FIA decode requires a block table"; + torch::Tensor query_tnd = query.view({-1, num_heads_, head_size_}); + torch::Tensor output_tnd = output.view({-1, num_heads_, head_size_}); + CHECK_EQ(static_cast(kv_seq_lens_vec->size()), query_tnd.size(0)) + << "FIA decode KV lengths must match query tokens"; + + torch::Tensor key_view = k_cache.view({k_cache.size(0), k_cache.size(1), -1}); + torch::Tensor value_view = + v_cache->view({v_cache->size(0), v_cache->size(1), -1}); + std::vector actual_q_lens = + make_decode_actual_seq_lengths(query_tnd.size(0)); + + if (tiling_data.defined()) { + const xllm::npu::FusedInferAttentionGraphBranch graph_branch = + attn_metadata.expanded_decode.enabled + ? xllm::npu::FusedInferAttentionGraphBranch::kSpecVerify + : xllm::npu::FusedInferAttentionGraphBranch::kDecode; + run_fused_infer_attention_graph(attn_metadata.acl_graph_task_update_context, + query_tnd, + key_view, + value_view, + block_table, + actual_q_lens, + *kv_seq_lens_vec, + num_heads_, + num_kv_heads_, + scale_, + k_cache.size(1), + graph_branch, + output_tnd); + } else { auto fia_result = xllm::kernel::npu::npu_fused_infer_attention( query_tnd, key_view, diff --git a/xllm/core/layers/npu_torch/attention.h b/xllm/core/layers/npu_torch/attention.h index 6fe09e9c99..1d758523bf 100644 --- a/xllm/core/layers/npu_torch/attention.h +++ b/xllm/core/layers/npu_torch/attention.h @@ -34,7 +34,8 @@ class AttentionImpl : public torch::nn::Module { int64_t head_size, float scale, int64_t num_kv_heads, - int64_t sliding_window); + int64_t sliding_window, + bool enable_fia_decode = false); std::tuple> forward( const AttentionMetadata& attn_metadata, @@ -63,6 +64,7 @@ class AttentionImpl : public torch::nn::Module { float scale_; int64_t num_kv_heads_; int64_t sliding_window_; + bool enable_fia_decode_ = false; }; TORCH_MODULE(Attention); diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp index e6695654e9..28078e936e 100644 --- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp +++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp @@ -376,6 +376,7 @@ torch::Tensor run_causal_conv1d_graph_update( task.pad_slot_id = xllm::npu::kCausalConv1dGraphPadSlotId; task.run_mode = xllm::npu::kCausalConv1dRunModeUpdate; task.branch = branch; + task.capture_order = graph_context->next_capture_order++; task.handle = handle; task.event = std::move(event); graph_context->causal_conv1d_tasks.emplace_back(std::move(task)); diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index b2f572de4b..dc272ed076 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp @@ -17,14 +17,22 @@ limitations under the License. #include +#include #include #include #include "common/flash_comm1_context.h" +#include "core/framework/config/execution_config.h" namespace xllm { namespace layer { +bool is_qwen3_5_model_type(const std::string& model_type) { + return model_type == "qwen3_5" || model_type == "qwen3_5_text" || + model_type == "qwen3_5_moe" || model_type == "qwen3_5_moe_text" || + model_type == "qwen3_5_mtp" || model_type == "qwen3_5_moe_mtp"; +} + Qwen3NextAttentionImpl::Qwen3NextAttentionImpl( const ModelArgs& args, const QuantArgs& quant_args, @@ -101,12 +109,16 @@ Qwen3NextAttentionImpl::Qwen3NextAttentionImpl( options)); // 6. Attention - attn_ = register_module("attn", - Attention(num_heads_, - head_dim_, - scaling_, - num_kv_heads_, - args.sliding_window())); + attn_ = register_module( + "attn", + Attention(num_heads_, + head_dim_, + scaling_, + num_kv_heads_, + args.sliding_window(), + /*enable_fia_decode=*/ + is_qwen3_5_model_type(args.model_type()) && + !ExecutionConfig::get_instance().disable_fia_decode())); // 7. Fused split_qkv_rmsnorm_mrope kernel setup rotary_dim_ = static_cast(head_dim_ * args.partial_rotary_factor()); diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.h b/xllm/core/layers/npu_torch/qwen3_next_attention.h index f5b456a550..2f712cfe74 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.h +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.h @@ -17,6 +17,7 @@ limitations under the License. #include +#include #include #include "attention.h" @@ -33,6 +34,8 @@ limitations under the License. namespace xllm { namespace layer { +bool is_qwen3_5_model_type(const std::string& model_type); + class Qwen3NextAttentionImpl : public torch::nn::Module { public: Qwen3NextAttentionImpl() = default; diff --git a/xllm/core/platform/npu/acl_graph_task_update_context.h b/xllm/core/platform/npu/acl_graph_task_update_context.h index eeae5f86b0..f2d8fdb67b 100644 --- a/xllm/core/platform/npu/acl_graph_task_update_context.h +++ b/xllm/core/platform/npu/acl_graph_task_update_context.h @@ -46,6 +46,11 @@ enum class CausalConv1dGraphBranch { kSpecVerify, }; +enum class FusedInferAttentionGraphBranch { + kDecode, + kSpecVerify, +}; + struct CausalConv1dGraphTask { torch::Tensor output; torch::Tensor x; @@ -56,6 +61,47 @@ struct CausalConv1dGraphTask { int64_t pad_slot_id = kCausalConv1dGraphPadSlotId; int64_t run_mode = kCausalConv1dRunModeUpdate; CausalConv1dGraphBranch branch = CausalConv1dGraphBranch::kDecode; + uint64_t capture_order = 0; + c10_npu::NPUTaskGroupHandle handle{}; + std::shared_ptr event; +}; + +struct FusedInferAttentionWorkspaceSignature { + torch::ScalarType query_dtype; + torch::ScalarType key_dtype; + torch::ScalarType value_dtype; + torch::ScalarType block_table_dtype; + c10::DeviceIndex device_index; + std::vector query_shape; + std::vector key_shape; + std::vector value_shape; + std::vector block_table_shape; + std::vector actual_seq_lengths; + std::vector actual_seq_lengths_kv; + int64_t num_heads; + int64_t num_key_value_heads; + int64_t block_size; + double scale; + + bool operator==(const FusedInferAttentionWorkspaceSignature&) const = default; +}; + +struct FusedInferAttentionGraphTask { + torch::Tensor output; + torch::Tensor softmax_lse; + torch::Tensor query; + torch::Tensor key; + torch::Tensor value; + torch::Tensor block_table; + torch::Tensor workspace; + std::vector actual_seq_lengths; + int64_t num_heads = 0; + int64_t num_key_value_heads = 0; + double scale = 0.0; + int64_t block_size = 0; + FusedInferAttentionGraphBranch branch = + FusedInferAttentionGraphBranch::kDecode; + uint64_t capture_order = 0; c10_npu::NPUTaskGroupHandle handle{}; std::shared_ptr event; }; @@ -64,13 +110,22 @@ class AclGraphTaskUpdateContext final { public: void begin_capture() { capturing = true; + next_capture_order = 0; causal_conv1d_tasks.clear(); + fused_infer_attention_tasks.clear(); + fused_infer_attention_workspace_signature.reset(); + fused_infer_attention_workspace = torch::Tensor(); } void end_capture() { capturing = false; } bool capturing = false; + uint64_t next_capture_order = 0; std::vector causal_conv1d_tasks; + std::vector fused_infer_attention_tasks; + std::optional + fused_infer_attention_workspace_signature; + torch::Tensor fused_infer_attention_workspace; }; } // namespace xllm::npu diff --git a/xllm/core/runtime/acl_graph_executor_impl.cpp b/xllm/core/runtime/acl_graph_executor_impl.cpp index 9ee655b8e9..852ea4b68f 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.cpp +++ b/xllm/core/runtime/acl_graph_executor_impl.cpp @@ -35,6 +35,7 @@ limitations under the License. #endif #include "core/common/metrics.h" #include "core/framework/speculative/mtp_async_state.h" +#include "core/kernels/npu/npu_ops_api.h" #include "core/kernels/npu/tilelang/tilelang_ops_api.h" #include "core/kernels/ops_api.h" #include "core/platform/device.h" @@ -457,24 +458,64 @@ bool AclGraph::capture(CausalLM* model, bool AclGraph::update_graph_tasks(const ModelInputParams& params) { if (graph_task_context_ == nullptr || - graph_task_context_->causal_conv1d_tasks.empty()) { + (graph_task_context_->causal_conv1d_tasks.empty() && + graph_task_context_->fused_infer_attention_tasks.empty())) { return false; } - const std::vector empty_host_args; - CHECK(!params.parallel.query_start_loc.empty()) - << "causal_conv1d graph update requires padded query_start_loc"; - CHECK(!params.embedding.linear_state_ids.empty()) - << "causal_conv1d graph update requires padded cache indices"; - - std::vector linear_state_indices_host( - params.embedding.linear_state_ids.begin(), - params.embedding.linear_state_ids.end()); - c10_npu::NPUStream update_stream = update_stream_.value(); c10_npu::NPUStreamGuard stream_guard(update_stream); + auto& causal_conv1d_tasks = graph_task_context_->causal_conv1d_tasks; + auto& fused_infer_attention_tasks = + graph_task_context_->fused_infer_attention_tasks; + + const bool has_causal_conv1d_tasks = !causal_conv1d_tasks.empty(); + std::vector empty_host_args; + std::vector linear_state_indices_host; + if (has_causal_conv1d_tasks) { + CHECK(!params.parallel.query_start_loc.empty()) + << "causal_conv1d graph update requires padded query_start_loc"; + CHECK(!params.embedding.linear_state_ids.empty()) + << "causal_conv1d graph update requires padded cache indices"; + linear_state_indices_host.assign(params.embedding.linear_state_ids.begin(), + params.embedding.linear_state_ids.end()); + } + + const bool has_fused_infer_attention_tasks = + !fused_infer_attention_tasks.empty(); + size_t graph_batch_size = 0; + std::vector actual_seq_lengths_kv; + if (has_fused_infer_attention_tasks) { + const auto& first_task = fused_infer_attention_tasks.front(); + graph_batch_size = static_cast(first_task.query.size(0)); + const FusedInferAttentionGraphBranch graph_branch = first_task.branch; + for (const FusedInferAttentionGraphTask& task : + fused_infer_attention_tasks) { + CHECK(task.branch == graph_branch) + << "FIA graph tasks must share one decode branch"; + CHECK_EQ(static_cast(task.query.size(0)), graph_batch_size) + << "FIA graph tasks must share the captured batch size"; + } + const std::vector& source_kv_seq_lens = + graph_branch == FusedInferAttentionGraphBranch::kSpecVerify + ? params.graph.expanded_kv_seq_lens_vec + : params.attention.host.kv_seq_lens; + if (graph_branch == FusedInferAttentionGraphBranch::kSpecVerify) { + CHECK(params.is_spec_verify) + << "spec FIA graph task requires spec-verify params"; + CHECK(params.graph.use_expanded_decode_for_spec_verify_attention) + << "spec FIA graph task requires expanded decode metadata"; + } + CHECK_LE(source_kv_seq_lens.size(), graph_batch_size) + << "FIA graph update KV lengths exceed captured query tokens"; + + actual_seq_lengths_kv.assign(graph_batch_size, 1); + for (size_t index = 0; index < source_kv_seq_lens.size(); ++index) { + actual_seq_lengths_kv[index] = source_kv_seq_lens[index]; + } + } - for (auto& task : graph_task_context_->causal_conv1d_tasks) { + auto update_causal_conv1d_task = [&](CausalConv1dGraphTask& task) { CHECK_EQ(params.parallel.query_start_loc.back(), task.x.size(0)) << "causal_conv1d graph update host args must be padded to the " "capture x.shape[0]"; @@ -509,17 +550,134 @@ bool AclGraph::update_graph_tasks(const ModelInputParams& params) { if (task.event != nullptr) { task.event->record(update_stream); } + }; + + auto update_fused_infer_attention_task = + [&](FusedInferAttentionGraphTask& task) { + c10_npu::graph_task_update_begin(update_stream, task.handle); + kernel::npu::npu_fused_infer_attention_decode_out( + task.query, + task.key, + task.value, + task.block_table, + task.actual_seq_lengths, + actual_seq_lengths_kv, + task.num_heads, + task.num_key_value_heads, + task.scale, + task.block_size, + task.workspace, + task.output, + task.softmax_lse); + c10_npu::graph_task_update_end(update_stream); + if (task.event != nullptr) { + task.event->record(update_stream); + } + }; + + size_t causal_conv1d_index = 0; + size_t fused_infer_attention_index = 0; + while (causal_conv1d_index < causal_conv1d_tasks.size() || + fused_infer_attention_index < fused_infer_attention_tasks.size()) { + const bool causal_conv1d_available = + causal_conv1d_index < causal_conv1d_tasks.size(); + const bool fused_infer_attention_available = + fused_infer_attention_index < fused_infer_attention_tasks.size(); + if (causal_conv1d_available && fused_infer_attention_available) { + CHECK_NE(causal_conv1d_tasks[causal_conv1d_index].capture_order, + fused_infer_attention_tasks[fused_infer_attention_index] + .capture_order) + << "ACL graph task capture order must be unique"; + } + const bool update_causal_conv1d = + causal_conv1d_available && + (!fused_infer_attention_available || + causal_conv1d_tasks[causal_conv1d_index].capture_order < + fused_infer_attention_tasks[fused_infer_attention_index] + .capture_order); + if (update_causal_conv1d) { + update_causal_conv1d_task(causal_conv1d_tasks[causal_conv1d_index]); + ++causal_conv1d_index; + } else { + update_fused_infer_attention_task( + fused_infer_attention_tasks[fused_infer_attention_index]); + ++fused_infer_attention_index; + } } return true; } -void AclGraph::signal_static_graph_tasks( +void AclGraph::prepare_static_graph_tasks( + const SpecVerifyGraphTaskSignal& signal, const c10_npu::NPUStream& signal_stream) { CHECK(graph_task_context_ != nullptr); - for (auto& task : graph_task_context_->causal_conv1d_tasks) { + c10_npu::NPUStreamGuard stream_guard(signal_stream); + auto& causal_conv1d_tasks = graph_task_context_->causal_conv1d_tasks; + auto& fused_infer_attention_tasks = + graph_task_context_->fused_infer_attention_tasks; + + auto signal_causal_conv1d_task = [&](CausalConv1dGraphTask& task) { CHECK(task.event != nullptr) << "static graph-task replay requires a captured ready event"; task.event->record(signal_stream); + }; + + auto update_fused_infer_attention_task = + [&](FusedInferAttentionGraphTask& task) { + CHECK(task.branch == FusedInferAttentionGraphBranch::kSpecVerify) + << "static MTP FIA task must use the spec-verify branch"; + CHECK_EQ(static_cast(task.query.size(0)), + signal.expanded_kv_seq_lens.size()) + << "static MTP FIA KV lengths must match captured query rows"; + c10_npu::graph_task_update_begin(signal_stream, task.handle); + kernel::npu::npu_fused_infer_attention_decode_out( + task.query, + task.key, + task.value, + task.block_table, + task.actual_seq_lengths, + signal.expanded_kv_seq_lens, + task.num_heads, + task.num_key_value_heads, + task.scale, + task.block_size, + task.workspace, + task.output, + task.softmax_lse); + c10_npu::graph_task_update_end(signal_stream); + CHECK(task.event != nullptr) + << "static graph-task replay requires a captured ready event"; + task.event->record(signal_stream); + }; + + size_t causal_conv1d_index = 0; + size_t fused_infer_attention_index = 0; + while (causal_conv1d_index < causal_conv1d_tasks.size() || + fused_infer_attention_index < fused_infer_attention_tasks.size()) { + const bool causal_conv1d_available = + causal_conv1d_index < causal_conv1d_tasks.size(); + const bool fused_infer_attention_available = + fused_infer_attention_index < fused_infer_attention_tasks.size(); + if (causal_conv1d_available && fused_infer_attention_available) { + CHECK_NE(causal_conv1d_tasks[causal_conv1d_index].capture_order, + fused_infer_attention_tasks[fused_infer_attention_index] + .capture_order) + << "ACL graph task capture order must be unique"; + } + const bool signal_causal_conv1d = + causal_conv1d_available && + (!fused_infer_attention_available || + causal_conv1d_tasks[causal_conv1d_index].capture_order < + fused_infer_attention_tasks[fused_infer_attention_index] + .capture_order); + if (signal_causal_conv1d) { + signal_causal_conv1d_task(causal_conv1d_tasks[causal_conv1d_index]); + ++causal_conv1d_index; + } else { + update_fused_infer_attention_task( + fused_infer_attention_tasks[fused_infer_attention_index]); + ++fused_infer_attention_index; + } } } @@ -633,6 +791,11 @@ void AclGraph::update_spec_verify_attention_tiling( spec_verify_kv_split_core_count_); } +bool AclGraph::has_fused_infer_attention_graph_tasks() const { + return graph_task_context_ != nullptr && + !graph_task_context_->fused_infer_attention_tasks.empty(); +} + ModelOutput AclGraph::replay(CausalLM* model, const torch::Tensor& tokens, const torch::Tensor& positions, @@ -685,13 +848,18 @@ ModelOutput AclGraph::replay(CausalLM* model, tokens, params, actual_num_tokens, num_tokens_); } else { auto [k_cache, v_cache] = find_attention_plan_kv_cache(kv_cache); - graph_params = persistent_param_.update(tokens, - k_cache, - v_cache, - positions, - params, - num_tokens_, - needs_graph_metadata); + graph_params = + persistent_param_.update(tokens, + k_cache, + v_cache, + positions, + params, + num_tokens_, + needs_graph_metadata, + /*skip_token_update=*/false, + /*for_capture=*/false, + /*update_paged_attention_plan=*/ + !has_fused_infer_attention_graph_tasks()); if (needs_graph_metadata) { CHECK(graph_params.has_value()) << "ACL graph replay requires persistent params for graph metadata"; @@ -714,20 +882,14 @@ ModelOutput AclGraph::replay(CausalLM* model, params.graph.spec_verify_static_graph_tasks_prepared; CHECK(!static_graph_tasks_prepared || use_static_graph_tasks) << "prepared static graph tasks do not match the replay signature"; - if (use_static_graph_tasks && !static_graph_tasks_prepared) { - // Cold/fallback path: the final-draft pre-submit could not find this graph - // variant. Signal its task-ready events immediately before replay; steady - // supported-width cycles use the compute-stream pre-submit path instead. - CHECK(update_stream_.has_value()); - signal_static_graph_tasks(update_stream_.value()); - } graph_.replay(); if (model->is_hybrid_linear_attention()) { CHECK(graph_params.has_value()) << "update() should return ModelInputParams for graph task update"; - if (use_static_graph_tasks) { - // This graph variant's task-ready event was recorded before replay. - } else { + if (!use_static_graph_tasks || !static_graph_tasks_prepared) { + // Cold static variants update after replay starts and unblock through the + // captured task-ready events. Steady variants were pre-submitted before + // the final draft and must not be updated twice. update_graph_tasks(graph_params.value()); } } @@ -759,7 +921,10 @@ void AclGraph::prepare_replay_inputs(const torch::Tensor& tokens, params, num_tokens_, /*return_capture_params=*/false, - /*skip_token_update=*/true); + /*skip_token_update=*/true, + /*for_capture=*/false, + /*update_paged_attention_plan=*/ + !has_fused_infer_attention_graph_tasks()); replay_inputs_prepared_.store(true, std::memory_order_release); } @@ -770,7 +935,26 @@ bool AclGraph::prepare_static_mtp_graph_tasks( make_static_graph_task_signature(signal)) { return false; } - signal_static_graph_tasks(signal_stream); + if (graph_task_context_ == nullptr) { + return false; + } + const auto& fused_infer_attention_tasks = + graph_task_context_->fused_infer_attention_tasks; + if (!fused_infer_attention_tasks.empty()) { + const size_t graph_batch_size = + static_cast(fused_infer_attention_tasks.front().query.size(0)); + if (signal.expanded_kv_seq_lens.size() != graph_batch_size) { + return false; + } + for (const FusedInferAttentionGraphTask& task : + fused_infer_attention_tasks) { + CHECK(task.branch == FusedInferAttentionGraphBranch::kSpecVerify) + << "static MTP graph cannot contain regular FIA decode tasks"; + CHECK_EQ(static_cast(task.query.size(0)), graph_batch_size) + << "static MTP FIA tasks must share captured query rows"; + } + } + prepare_static_graph_tasks(signal, signal_stream); return true; } diff --git a/xllm/core/runtime/acl_graph_executor_impl.h b/xllm/core/runtime/acl_graph_executor_impl.h index 50e9a4367c..f1b94350c6 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.h +++ b/xllm/core/runtime/acl_graph_executor_impl.h @@ -145,9 +145,11 @@ class AclGraph { const torch::Tensor& positions, ModelInputParams& params); void update_spec_verify_attention_tiling(const ModelInputParams& params); + bool has_fused_infer_attention_graph_tasks() const; bool update_graph_tasks(const ModelInputParams& params); - void signal_static_graph_tasks(const c10_npu::NPUStream& signal_stream); + void prepare_static_graph_tasks(const SpecVerifyGraphTaskSignal& signal, + const c10_npu::NPUStream& signal_stream); bool static_graph_task_signature_matches( const ModelInputParams& params) const; void capture_static_graph_task_signature(const ModelInputParams& params); @@ -214,6 +216,10 @@ class AclGraphExecutorImpl : public ExecutorImpl { size_t get_graph_memory_pool_count(); size_t get_graph_capture_stream_count() const; + [[nodiscard]] uint32_t bucket_num_tokens_for_test(uint32_t num_tokens) const { + return get_bucket_num_tokens(num_tokens); + } + private: // not own CausalLM* model_; diff --git a/xllm/core/runtime/acl_graph_persistent_param.cpp b/xllm/core/runtime/acl_graph_persistent_param.cpp index 420d8600ec..f00fbaf3c3 100644 --- a/xllm/core/runtime/acl_graph_persistent_param.cpp +++ b/xllm/core/runtime/acl_graph_persistent_param.cpp @@ -844,7 +844,8 @@ std::optional GraphPersistentParam::update( uint32_t padded_num_tokens, bool return_capture_params, bool skip_token_update, - bool for_capture) { + bool for_capture, + bool update_paged_attention_plan) { CHECK_GT(padded_num_tokens, 0) << "padded_num_tokens must be > 0"; const uint32_t actual_num_tokens = tokens.size(0); const bool is_decode = params.meta.batch_forward_type.is_decode(); @@ -1210,7 +1211,7 @@ std::optional GraphPersistentParam::update( persistent_host_q_seq_lens_.begin()); } - if (uses_paged_attention_tiling()) { + if (uses_paged_attention_tiling() && update_paged_attention_plan) { aclrtStream stream = c10_npu::getCurrentNPUStream().stream(); if (k_cache.defined() && v_cache.defined() && k_cache.numel() > 0 && diff --git a/xllm/core/runtime/acl_graph_persistent_param.h b/xllm/core/runtime/acl_graph_persistent_param.h index 75bd12e89b..d0465e8caf 100644 --- a/xllm/core/runtime/acl_graph_persistent_param.h +++ b/xllm/core/runtime/acl_graph_persistent_param.h @@ -78,15 +78,19 @@ class GraphPersistentParam final { // host parameters can be bucketed for graph tiling/workspace. During replay, // return_capture_params may still be true for metadata refresh, but // for_capture must stay false so dynamic host metadata uses actual lengths. - std::optional update(const torch::Tensor& tokens, - const torch::Tensor& k_cache, - const torch::Tensor& v_cache, - const torch::Tensor& positions, - const ModelInputParams& params, - uint32_t padded_num_tokens, - bool return_capture_params = false, - bool skip_token_update = false, - bool for_capture = false); + // update_paged_attention_plan can be disabled when FIA graph tasks replace + // the paged-attention task for the current graph. + std::optional update( + const torch::Tensor& tokens, + const torch::Tensor& k_cache, + const torch::Tensor& v_cache, + const torch::Tensor& positions, + const ModelInputParams& params, + uint32_t padded_num_tokens, + bool return_capture_params = false, + bool skip_token_update = false, + bool for_capture = false, + bool update_paged_attention_plan = true); void update_tokens(const torch::Tensor& tokens, const ModelInputParams& params, diff --git a/xllm/core/runtime/executor_impl.h b/xllm/core/runtime/executor_impl.h index a4c0bbd312..5866910d40 100644 --- a/xllm/core/runtime/executor_impl.h +++ b/xllm/core/runtime/executor_impl.h @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include #include "common/macros.h" #include "framework/batch/batch.h" @@ -38,6 +39,7 @@ struct SpecVerifyGraphTaskSignal { int64_t spec_width = 0; int64_t block_table_width = 0; int64_t max_kv_seq_len = 0; + std::vector expanded_kv_seq_lens; }; class ExecutorImpl { diff --git a/xllm/core/runtime/mtp_worker_impl.cpp b/xllm/core/runtime/mtp_worker_impl.cpp index 720e44bc7b..ece945f7b7 100644 --- a/xllm/core/runtime/mtp_worker_impl.cpp +++ b/xllm/core/runtime/mtp_worker_impl.cpp @@ -3340,19 +3340,24 @@ bool MTPWorkerImpl::prepare_static_mtp_graph_tasks_before_final_draft( const int64_t verify_block_table_width = spec_verify_block_table_width(block_tables); const auto& kv_seq_lens = input.input_params.attention.host.kv_seq_lens; - if (kv_seq_lens.empty()) { + if (kv_seq_lens.size() != 1) { return false; } - const int64_t spec_verify_max_kv_seq_len = - static_cast( - *std::max_element(kv_seq_lens.begin(), kv_seq_lens.end())) + - options_.num_speculative_tokens(); + const int64_t spec_width = options_.num_speculative_tokens() + 1; + const int64_t base_kv_seq_len = kv_seq_lens.front(); + std::vector expanded_kv_seq_lens; + expanded_kv_seq_lens.reserve(static_cast(spec_width)); + for (int64_t token_idx = 0; token_idx < spec_width; ++token_idx) { + expanded_kv_seq_lens.emplace_back(base_kv_seq_len + token_idx); + } + const int64_t spec_verify_max_kv_seq_len = expanded_kv_seq_lens.back(); const SpecVerifyGraphTaskSignal signal{ .linear_state_id = input.input_params.embedding.linear_state_ids.front(), .num_accepted_tokens = accepted_prefix_lengths.front(), - .spec_width = options_.num_speculative_tokens() + 1, + .spec_width = spec_width, .block_table_width = verify_block_table_width, .max_kv_seq_len = spec_verify_max_kv_seq_len, + .expanded_kv_seq_lens = std::move(expanded_kv_seq_lens), }; return impl_->prepare_static_mtp_graph_tasks(signal, *compute_stream_); #else From 103a25a941596241d46844574718a1e5caff2aaa Mon Sep 17 00:00:00 2001 From: Enguikong Date: Fri, 14 Aug 2026 13:30:51 +0800 Subject: [PATCH 3/7] test(npu): align FIA graph fixtures with runtime metadata --- tests/core/kernels/npu/npu_xllm_ops_test.cpp | 1 - .../core/runtime/acl_graph_executor_test.cpp | 2 +- .../runtime/acl_graph_task_update_test.cpp | 30 ++++++++++++++----- xllm/core/runtime/acl_graph_executor_impl.h | 4 --- 4 files changed, 24 insertions(+), 13 deletions(-) diff --git a/tests/core/kernels/npu/npu_xllm_ops_test.cpp b/tests/core/kernels/npu/npu_xllm_ops_test.cpp index 497fc3da32..c0ff8291f4 100644 --- a/tests/core/kernels/npu/npu_xllm_ops_test.cpp +++ b/tests/core/kernels/npu/npu_xllm_ops_test.cpp @@ -309,7 +309,6 @@ TEST_F(NpuXllmOpsTest, kScale, kBlockSize); ASSERT_TRUE(workspace.defined()); - EXPECT_GT(workspace.numel(), 0); EXPECT_EQ(workspace.device(), query.device()); torch::Tensor out = torch::zeros_like(eager_output); diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index 60fbe250ad..c929bef76f 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -1245,7 +1245,7 @@ TEST(AclGraphPersistentParamTest, SpecVerifyMetadataUsesTokenCapacity) { TEST(AclGraphPersistentParamTest, GenericSpecVerifyCaptureKeepsPersistentBlockTableWidth) { constexpr int32_t kSpecWidth = 6; - constexpr int64_t kActiveBlockTableWidth = 2; + constexpr int64_t kActiveBlockTableWidth = 5; ModelArgs args; args.model_type("deepseek_v4"); args.dtype("float32"); diff --git a/tests/core/runtime/acl_graph_task_update_test.cpp b/tests/core/runtime/acl_graph_task_update_test.cpp index b63f11fb3f..1ff1fea485 100644 --- a/tests/core/runtime/acl_graph_task_update_test.cpp +++ b/tests/core/runtime/acl_graph_task_update_test.cpp @@ -29,6 +29,7 @@ limitations under the License. #include "core/framework/block/block_manager_impl.h" #include "core/framework/config/execution_config.h" #include "core/framework/kv_cache/kv_cache.h" +#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_loader.h" @@ -37,6 +38,7 @@ limitations under the License. #include "core/framework/sampling/sampling_params.h" #include "core/kernels/ops_api.h" #include "core/layers/common/attention_metadata_builder.h" +#include "core/layers/common/expanded_decode_metadata_builder.h" #include "core/layers/npu/npu_lm_head_impl.h" #include "core/layers/npu/npu_word_embedding_impl.h" #include "core/layers/npu_torch/attention.h" @@ -168,6 +170,9 @@ class HybridConv1dMockLM final : public CausalLM { const bool register_graph_task = graph_context != nullptr && graph_context->capturing; + layer::AttentionMetadataBuildOptions metadata_build_options; + metadata_build_options.materialize_linear_state_validity = + !params.enable_graph; for (auto& kv_cache : kv_caches) { if (kv_cache.empty() || !kv_cache.get_conv_cache().defined()) { continue; @@ -267,7 +272,8 @@ class HybridConv1dMockLM final : public CausalLM { layer::AttentionMetadataBuilder::build(params, /*enable_mla=*/false, /*attn_mask=*/std::nullopt, - device_); + device_, + metadata_build_options); torch::Tensor query = hidden.to(torch::kBFloat16).contiguous(); torch::Tensor key = query .slice(/*dim=*/1, @@ -287,6 +293,10 @@ class HybridConv1dMockLM final : public CausalLM { !graph_context->fused_infer_attention_tasks.empty(); fia_graph_task_count_ = graph_context->fused_infer_attention_tasks.size(); + if (!graph_context->fused_infer_attention_tasks.empty()) { + captured_fia_batch_size_ = static_cast( + graph_context->fused_infer_attention_tasks.front().query.size(0)); + } all_fia_graph_tasks_share_workspace_ = fia_graph_task_count_ > 1; for (size_t task_index = 1; task_index < fia_graph_task_count_; ++task_index) { @@ -327,6 +337,7 @@ class HybridConv1dMockLM final : public CausalLM { return saw_causal_conv_graph_task_; } bool saw_fia_graph_task() const { return saw_fia_graph_task_; } + uint32_t captured_fia_batch_size() const { return captured_fia_batch_size_; } size_t fia_graph_task_count() const { return fia_graph_task_count_; } bool all_fia_graph_tasks_share_workspace() const { return all_fia_graph_tasks_share_workspace_; @@ -353,6 +364,7 @@ class HybridConv1dMockLM final : public CausalLM { int32_t attention_repetitions_ = 1; bool saw_causal_conv_graph_task_ = false; bool saw_fia_graph_task_ = false; + uint32_t captured_fia_batch_size_ = 0; size_t fia_graph_task_count_ = 0; bool all_fia_graph_tasks_share_workspace_ = false; }; @@ -563,13 +575,12 @@ class AclGraphTaskUpdateTest : public ::testing::Test { auto kv_graph = create_hybrid_kv_caches(); auto graph_exec = std::make_unique( model_.get(), model_args_, *device_, options_); - EXPECT_EQ(graph_exec->bucket_num_tokens_for_test(capture_batch_size), - expected_bucket); graph_exec->run({capture_fi.token_ids}, {capture_fi.positions}, kv_graph, {capture_fi.input_params}); ASSERT_TRUE(model_->saw_causal_conv_and_fia_graph_tasks()); + EXPECT_EQ(model_->captured_fia_batch_size(), expected_bucket); reset_sequences(); auto replay_prompts = @@ -620,6 +631,8 @@ class AclGraphTaskUpdateTest : public ::testing::Test { fi.input_params.attention.host.q_seq_lens.assign( static_cast(num_sequences), num_spec_tokens); + fi.input_params.linear_state_validity_mask = build_linear_state_mask( + fi.input_params.attention.host.kv_cache_tokens_nums, num_sequences); fi.input_params.num_accepted_tokens_host.assign( static_cast(num_sequences), 1); @@ -641,7 +654,6 @@ class AclGraphTaskUpdateTest : public ::testing::Test { fi.token_ids = torch::tensor(token_ids_vec, torch::kInt32).to(*device_); fi.positions = torch::tensor(positions_vec, torch::kInt32).to(*device_); - fi.input_params.graph.use_expanded_decode_for_spec_verify_attention = true; std::vector expanded_kv_vec; expanded_kv_vec.reserve(static_cast(total_tokens)); for (int32_t s = 0; s < num_sequences; ++s) { @@ -651,8 +663,7 @@ class AclGraphTaskUpdateTest : public ::testing::Test { expanded_kv_vec.push_back(kv_len + t + 1); } } - fi.input_params.graph.expanded_kv_seq_lens_vec = expanded_kv_vec; - fi.input_params.graph.expanded_kv_seq_lens = + auto expanded_kv_seq_lens = torch::tensor(expanded_kv_vec, torch::kInt32).to(*device_); torch::Tensor host_block_tables = @@ -689,7 +700,12 @@ class AclGraphTaskUpdateTest : public ::testing::Test { expanded_bt[s * num_spec_tokens + t] = block_tables[s]; } } - fi.input_params.graph.expanded_block_tables = expanded_bt; + layer::ExpandedDecodeMetadataBuilder::populate_expanded_layout( + fi.input_params, + expanded_kv_seq_lens, + expanded_bt, + expanded_kv_vec, + kBlockSize); std::vector q_cu_vec; q_cu_vec.reserve(static_cast(num_sequences + 1)); diff --git a/xllm/core/runtime/acl_graph_executor_impl.h b/xllm/core/runtime/acl_graph_executor_impl.h index f1b94350c6..2d25bff93c 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.h +++ b/xllm/core/runtime/acl_graph_executor_impl.h @@ -216,10 +216,6 @@ class AclGraphExecutorImpl : public ExecutorImpl { size_t get_graph_memory_pool_count(); size_t get_graph_capture_stream_count() const; - [[nodiscard]] uint32_t bucket_num_tokens_for_test(uint32_t num_tokens) const { - return get_bucket_num_tokens(num_tokens); - } - private: // not own CausalLM* model_; From 3a2adb324c21609b5fcf83fc80231f70294e39a8 Mon Sep 17 00:00:00 2001 From: Enguikong Date: Fri, 14 Aug 2026 16:17:21 +0800 Subject: [PATCH 4/7] test(npu): cover Qwen3.5 FIA rollback routing --- tests/core/runtime/acl_graph_executor_test.cpp | 14 ++++++++++++++ .../core/layers/npu_torch/qwen3_next_attention.cpp | 8 ++++++-- xllm/core/layers/npu_torch/qwen3_next_attention.h | 1 + 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index c929bef76f..caa2be8d43 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -175,6 +175,20 @@ TEST(Qwen35FiaRoutingTest, UsesExactModelTypeWhitelist) { } } +TEST(Qwen35FiaRoutingTest, DisableFlagOverridesWhitelist) { + ExecutionConfig& execution_config = ExecutionConfig::get_instance(); + const bool original_disable_fia_decode = + execution_config.disable_fia_decode(); + + execution_config.disable_fia_decode(false); + EXPECT_TRUE(layer::should_enable_qwen3_5_fia_decode("qwen3_5")); + + execution_config.disable_fia_decode(true); + EXPECT_FALSE(layer::should_enable_qwen3_5_fia_decode("qwen3_5")); + + execution_config.disable_fia_decode(original_disable_fia_decode); +} + namespace { const KVCache& first_full_attention_cache( const std::vector& kv_caches) { diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index dc272ed076..d07674dfd1 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp @@ -33,6 +33,11 @@ bool is_qwen3_5_model_type(const std::string& model_type) { model_type == "qwen3_5_mtp" || model_type == "qwen3_5_moe_mtp"; } +bool should_enable_qwen3_5_fia_decode(const std::string& model_type) { + return is_qwen3_5_model_type(model_type) && + !ExecutionConfig::get_instance().disable_fia_decode(); +} + Qwen3NextAttentionImpl::Qwen3NextAttentionImpl( const ModelArgs& args, const QuantArgs& quant_args, @@ -117,8 +122,7 @@ Qwen3NextAttentionImpl::Qwen3NextAttentionImpl( num_kv_heads_, args.sliding_window(), /*enable_fia_decode=*/ - is_qwen3_5_model_type(args.model_type()) && - !ExecutionConfig::get_instance().disable_fia_decode())); + should_enable_qwen3_5_fia_decode(args.model_type()))); // 7. Fused split_qkv_rmsnorm_mrope kernel setup rotary_dim_ = static_cast(head_dim_ * args.partial_rotary_factor()); diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.h b/xllm/core/layers/npu_torch/qwen3_next_attention.h index 2f712cfe74..e641c9be0e 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.h +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.h @@ -35,6 +35,7 @@ namespace xllm { namespace layer { bool is_qwen3_5_model_type(const std::string& model_type); +bool should_enable_qwen3_5_fia_decode(const std::string& model_type); class Qwen3NextAttentionImpl : public torch::nn::Module { public: From 5b48803137d58e0622c1dbb51b70b13953d7da16 Mon Sep 17 00:00:00 2001 From: Enguikong Date: Fri, 14 Aug 2026 16:45:54 +0800 Subject: [PATCH 5/7] style(npu): use class for FIA workspace signature --- xllm/core/platform/npu/acl_graph_task_update_context.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xllm/core/platform/npu/acl_graph_task_update_context.h b/xllm/core/platform/npu/acl_graph_task_update_context.h index f2d8fdb67b..f3b3a76569 100644 --- a/xllm/core/platform/npu/acl_graph_task_update_context.h +++ b/xllm/core/platform/npu/acl_graph_task_update_context.h @@ -66,7 +66,8 @@ struct CausalConv1dGraphTask { std::shared_ptr event; }; -struct FusedInferAttentionWorkspaceSignature { +class FusedInferAttentionWorkspaceSignature { + public: torch::ScalarType query_dtype; torch::ScalarType key_dtype; torch::ScalarType value_dtype; From a29981c1ce363ad84cd5c4ca9ef3c2b082d12f63 Mon Sep 17 00:00:00 2001 From: Enguikong Date: Thu, 20 Aug 2026 12:36:07 +0800 Subject: [PATCH 6/7] feat(npu): make Qwen3.5 FIA decode opt-in --- .../framework/config/config_json_test.cpp | 31 ++++++++++++++++--- .../core/runtime/acl_graph_executor_test.cpp | 28 ++++++++++------- xllm/core/common/global_flags.h | 2 ++ .../framework/config/execution_config.cpp | 12 +++---- xllm/core/framework/config/execution_config.h | 4 +-- .../layers/npu_torch/qwen3_next_attention.cpp | 4 +-- 6 files changed, 56 insertions(+), 25 deletions(-) diff --git a/tests/core/framework/config/config_json_test.cpp b/tests/core/framework/config/config_json_test.cpp index 3aba034951..418835c00b 100644 --- a/tests/core/framework/config/config_json_test.cpp +++ b/tests/core/framework/config/config_json_test.cpp @@ -41,7 +41,8 @@ inline constexpr std::string_view kInlineConfig = R"json({ "max_seqs_per_batch": 64, "model_impl": "py", "disable_graph_warmup": true, - "python_graph_backend": "cudagraphs" + "python_graph_backend": "cudagraphs", + "enable_fia_decode": true })json"; inline constexpr std::string_view kUpdatedConfig = R"json({ @@ -114,7 +115,8 @@ class ConfigFlagGuard final { old_model_impl_(FLAGS_model_impl), old_python_model_path_(FLAGS_python_model_path), old_disable_graph_warmup_(FLAGS_disable_graph_warmup), - old_python_graph_backend_(FLAGS_python_graph_backend) {} + old_python_graph_backend_(FLAGS_python_graph_backend), + old_enable_fia_decode_(FLAGS_enable_fia_decode) {} ~ConfigFlagGuard() { FLAGS_block_size = old_block_size_; @@ -126,6 +128,7 @@ class ConfigFlagGuard final { FLAGS_python_model_path = old_python_model_path_; FLAGS_disable_graph_warmup = old_disable_graph_warmup_; FLAGS_python_graph_backend = old_python_graph_backend_; + FLAGS_enable_fia_decode = old_enable_fia_decode_; } private: @@ -138,6 +141,7 @@ class ConfigFlagGuard final { std::string old_python_model_path_; bool old_disable_graph_warmup_; std::string old_python_graph_backend_; + bool old_enable_fia_decode_; }; class StartupConfigGuard final { @@ -150,6 +154,7 @@ class StartupConfigGuard final { old_model_impl_(model_config_.model_impl()), old_python_model_path_(model_config_.python_model_path()), old_python_graph_backend_(execution_config_.python_graph_backend()), + old_enable_fia_decode_(execution_config_.enable_fia_decode()), old_block_size_(kv_cache_config_.block_size()), old_enable_prefix_cache_(kv_cache_config_.enable_prefix_cache()), old_max_tokens_per_batch_(scheduler_config_.max_tokens_per_batch()), @@ -160,7 +165,8 @@ class StartupConfigGuard final { ~StartupConfigGuard() { model_config_.model_impl(old_model_impl_) .python_model_path(old_python_model_path_); - execution_config_.python_graph_backend(old_python_graph_backend_); + execution_config_.python_graph_backend(old_python_graph_backend_) + .enable_fia_decode(old_enable_fia_decode_); kv_cache_config_.block_size(old_block_size_) .enable_prefix_cache(old_enable_prefix_cache_); scheduler_config_.max_tokens_per_batch(old_max_tokens_per_batch_) @@ -176,6 +182,7 @@ class StartupConfigGuard final { std::string old_model_impl_; std::string old_python_model_path_; std::string old_python_graph_backend_; + bool old_enable_fia_decode_; int32_t old_block_size_; bool old_enable_prefix_cache_; int32_t old_max_tokens_per_batch_; @@ -269,11 +276,13 @@ TEST(ConfigJsonTest, FromJsonUsesParsedOverrides) { EXPECT_EQ(model_config.python_model_path(), ""); EXPECT_TRUE(execution_config.disable_graph_warmup()); EXPECT_EQ(execution_config.python_graph_backend(), "cudagraphs"); + EXPECT_TRUE(execution_config.enable_fia_decode()); EXPECT_EQ(FLAGS_model_impl, "py"); EXPECT_EQ(FLAGS_python_model_path, old_python_model_path); EXPECT_TRUE(FLAGS_disable_graph_warmup); EXPECT_EQ(FLAGS_python_graph_backend, "cudagraphs"); + EXPECT_TRUE(FLAGS_enable_fia_decode); EXPECT_EQ(kv_cache_config.kv_cache_dtype(), "auto"); EXPECT_EQ(kv_cache_config.indexer_cache_dtype(), "auto"); @@ -340,6 +349,17 @@ TEST(ConfigJsonTest, RegistersOnlyContextParallelCommandLineOption) { google::GetCommandLineFlagInfo(removed_flag.c_str(), &flag_info)); } +TEST(ConfigJsonTest, RegistersExplicitFiaDecodeCommandLineOption) { + google::CommandLineFlagInfo flag_info; + EXPECT_TRUE(google::GetCommandLineFlagInfo("enable_fia_decode", &flag_info)); + EXPECT_EQ(flag_info.default_value, "false"); + EXPECT_FALSE( + google::GetCommandLineFlagInfo("disable_fia_decode", &flag_info)); + + const ExecutionConfig execution_config; + EXPECT_FALSE(execution_config.enable_fia_decode()); +} + TEST(ConfigJsonTest, LoadJsonFileReadsConfigFixture) { // The fixture sets more keys than ConfigFlagGuard restores, and from_json // writes every resolved value back into its FLAGS_ global. FlagSaver reverts @@ -614,7 +634,9 @@ TEST(ConfigJsonTest, DumpStartupConfigWritesNonDefaultValuesOnly) { ModelConfig::get_instance().model_impl("python").python_model_path( "/tmp/xllm-python-model"); - ExecutionConfig::get_instance().python_graph_backend("cudagraphs"); + ExecutionConfig::get_instance() + .python_graph_backend("cudagraphs") + .enable_fia_decode(true); KVCacheConfig::get_instance().block_size(256).enable_prefix_cache(false); SchedulerConfig::get_instance() .max_tokens_per_batch(2048) @@ -629,6 +651,7 @@ TEST(ConfigJsonTest, DumpStartupConfigWritesNonDefaultValuesOnly) { EXPECT_EQ(config_json.at("model_impl").get(), "python"); EXPECT_EQ(config_json.at("python_graph_backend").get(), "cudagraphs"); + EXPECT_TRUE(config_json.at("enable_fia_decode").get()); EXPECT_EQ(config_json.at("block_size").get(), 256); EXPECT_FALSE(config_json.at("enable_prefix_cache").get()); EXPECT_EQ(config_json.at("max_tokens_per_batch").get(), 2048); diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index caa2be8d43..6257c68570 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -175,18 +175,24 @@ TEST(Qwen35FiaRoutingTest, UsesExactModelTypeWhitelist) { } } -TEST(Qwen35FiaRoutingTest, DisableFlagOverridesWhitelist) { +TEST(Qwen35FiaRoutingTest, RequiresExplicitEnableFlag) { ExecutionConfig& execution_config = ExecutionConfig::get_instance(); - const bool original_disable_fia_decode = - execution_config.disable_fia_decode(); - - execution_config.disable_fia_decode(false); - EXPECT_TRUE(layer::should_enable_qwen3_5_fia_decode("qwen3_5")); - - execution_config.disable_fia_decode(true); - EXPECT_FALSE(layer::should_enable_qwen3_5_fia_decode("qwen3_5")); - - execution_config.disable_fia_decode(original_disable_fia_decode); + const bool original_enable_fia_decode = execution_config.enable_fia_decode(); + + execution_config.enable_fia_decode(false); + EXPECT_FALSE(layer::should_enable_qwen3_5_fia_decode("qwen3_5_text")); + EXPECT_FALSE(layer::should_enable_qwen3_5_fia_decode("qwen3_5_moe_text")); + EXPECT_FALSE(layer::should_enable_qwen3_5_fia_decode("qwen3_5_mtp")); + EXPECT_FALSE(layer::should_enable_qwen3_5_fia_decode("qwen3_5_moe_mtp")); + + execution_config.enable_fia_decode(true); + EXPECT_TRUE(layer::should_enable_qwen3_5_fia_decode("qwen3_5_text")); + EXPECT_TRUE(layer::should_enable_qwen3_5_fia_decode("qwen3_5_moe_text")); + EXPECT_TRUE(layer::should_enable_qwen3_5_fia_decode("qwen3_5_mtp")); + EXPECT_TRUE(layer::should_enable_qwen3_5_fia_decode("qwen3_5_moe_mtp")); + EXPECT_FALSE(layer::should_enable_qwen3_5_fia_decode("qwen3_next")); + + execution_config.enable_fia_decode(original_enable_fia_decode); } namespace { diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index 5ca9c15ab7..d55722498a 100644 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -167,6 +167,8 @@ DECLARE_int32(acl_graph_decode_batch_size_limit); DECLARE_string(python_graph_backend); +DECLARE_bool(enable_fia_decode); + 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 33489c6309..823932cf4e 100644 --- a/xllm/core/framework/config/execution_config.cpp +++ b/xllm/core/framework/config/execution_config.cpp @@ -96,10 +96,10 @@ DEFINE_string( "or any torch.compile backend name."); DEFINE_bool( - disable_fia_decode, + enable_fia_decode, false, - "When true, Qwen3.5 decode attention uses PagedAttention instead of FIA. " - "Useful as a runtime rollback switch without rebuilding."); + "Enable FIA for Qwen3.5 decode attention. Applies to both target and MTP " + "draft models. Prefill attention is unaffected."); namespace xllm { @@ -118,7 +118,7 @@ void ExecutionConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(output_shm_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(random_seed); XLLM_CONFIG_ASSIGN_FROM_FLAG(python_graph_backend); - XLLM_CONFIG_ASSIGN_FROM_FLAG(disable_fia_decode); + XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_fia_decode); } void ExecutionConfig::from_json(const JsonReader& json) { @@ -136,7 +136,7 @@ void ExecutionConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(output_shm_size); XLLM_CONFIG_ASSIGN_FROM_JSON(random_seed); XLLM_CONFIG_ASSIGN_FROM_JSON(python_graph_backend); - XLLM_CONFIG_ASSIGN_FROM_JSON(disable_fia_decode); + XLLM_CONFIG_ASSIGN_FROM_JSON(enable_fia_decode); } void ExecutionConfig::append_config_json( @@ -171,7 +171,7 @@ void ExecutionConfig::append_config_json( APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, python_graph_backend); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( - config_json, default_config, disable_fia_decode); + config_json, default_config, enable_fia_decode); } ExecutionConfig& ExecutionConfig::get_instance() { diff --git a/xllm/core/framework/config/execution_config.h b/xllm/core/framework/config/execution_config.h index fe403a9584..3244ffab46 100644 --- a/xllm/core/framework/config/execution_config.h +++ b/xllm/core/framework/config/execution_config.h @@ -55,7 +55,7 @@ class ExecutionConfig final { "output_shm_size", "random_seed", "python_graph_backend", - "disable_fia_decode"}}; + "enable_fia_decode"}}; return kOptionCategory; } @@ -87,7 +87,7 @@ class ExecutionConfig final { PROPERTY(std::string, python_graph_backend) = "off"; - PROPERTY(bool, disable_fia_decode) = false; + PROPERTY(bool, enable_fia_decode) = false; }; } // namespace xllm diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index d07674dfd1..24432e0953 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp @@ -34,8 +34,8 @@ bool is_qwen3_5_model_type(const std::string& model_type) { } bool should_enable_qwen3_5_fia_decode(const std::string& model_type) { - return is_qwen3_5_model_type(model_type) && - !ExecutionConfig::get_instance().disable_fia_decode(); + return ExecutionConfig::get_instance().enable_fia_decode() && + is_qwen3_5_model_type(model_type); } Qwen3NextAttentionImpl::Qwen3NextAttentionImpl( From 68357655b06ff804318749e012398c8730231720 Mon Sep 17 00:00:00 2001 From: Enguikong Date: Mon, 24 Aug 2026 12:11:19 +0800 Subject: [PATCH 7/7] perf(npu): avoid FIA metadata work on PA path --- .../core/runtime/acl_graph_executor_test.cpp | 2 +- .../runtime/acl_graph_task_update_test.cpp | 35 +++++++++++++++++++ xllm/core/layers/npu_torch/attention.cpp | 26 +++++++------- xllm/core/runtime/acl_graph_executor_impl.cpp | 16 +++++++-- xllm/core/runtime/executor_impl.h | 2 +- xllm/core/runtime/mtp_worker_impl.cpp | 9 ++--- 6 files changed, 65 insertions(+), 25 deletions(-) diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index 6257c68570..4a6ca5663e 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -139,8 +139,8 @@ TEST(AclGraphStaticGraphTaskSignatureTest, .num_accepted_tokens = 4, .spec_width = 5, .block_table_width = 64, + .base_kv_seq_len = 252, .max_kv_seq_len = 256, - .expanded_kv_seq_lens = {252, 253, 254, 255, 256}, }; const auto captured = npu::make_static_graph_task_signature(params); diff --git a/tests/core/runtime/acl_graph_task_update_test.cpp b/tests/core/runtime/acl_graph_task_update_test.cpp index 1ff1fea485..db109bb0ed 100644 --- a/tests/core/runtime/acl_graph_task_update_test.cpp +++ b/tests/core/runtime/acl_graph_task_update_test.cpp @@ -858,6 +858,41 @@ TEST_F(AclGraphTaskUpdateTest, .item(); } +TEST_F(AclGraphTaskUpdateTest, + Qwen35OptOutKeepsPagedAttentionForExpandedSpecVerify) { + constexpr int32_t kNumSequences = 2; + constexpr int32_t kNumSpecTokens = 4; + auto pa_model = std::make_unique( + model_args_, *device_, /*enable_fia_decode=*/false); + + auto batch = create_decode_batch(/*batch_size=*/kNumSequences); + ASSERT_FALSE(batch->empty()); + auto forward_input = batch->prepare_forward_input( + options_.num_decoding_tokens(), 0, model_args_); + forward_input = forward_input.to(*device_, kDtype); + setup_spec_verify_input(forward_input, kNumSequences, kNumSpecTokens); + + auto kv_eager = create_hybrid_kv_caches(); + auto eager_out = pa_model->forward(forward_input.token_ids, + forward_input.positions, + kv_eager, + forward_input.input_params); + + auto kv_graph = create_hybrid_kv_caches(); + auto graph_exec = std::make_unique( + pa_model.get(), model_args_, *device_, options_); + auto graph_out = graph_exec->run(forward_input.token_ids, + forward_input.positions, + kv_graph, + forward_input.input_params); + + EXPECT_TRUE(pa_model->saw_causal_conv_graph_task()); + EXPECT_FALSE(pa_model->saw_fia_graph_task()); + EXPECT_EQ(eager_out.hidden_states.sizes(), graph_out.hidden_states.sizes()); + EXPECT_TRUE(torch::isfinite(eager_out.hidden_states).all().item()); + EXPECT_TRUE(torch::isfinite(graph_out.hidden_states).all().item()); +} + TEST_F(AclGraphTaskUpdateTest, ReplayWithDifferentParamsProducesDifferentOutputs) { std::vector> prompts_run1 = {{1, 3, 5, 7}, {2, 4, 6, 8}}; diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index 341efddefd..f7663e29db 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -304,19 +304,6 @@ void AttentionImpl::decoder_forward(torch::Tensor& query, kv_seq_lens = attn_metadata.kv_seq_lens; } - std::vector expanded_kv_seq_lens; - const std::vector* kv_seq_lens_vec = - &attn_metadata.kv_seq_lens_host_vec; - if (attn_metadata.expanded_decode.enabled) { - expanded_kv_seq_lens.reserve( - attn_metadata.expanded_decode.kv_seq_lens_host_vec.size()); - for (int32_t kv_seq_len : - attn_metadata.expanded_decode.kv_seq_lens_host_vec) { - expanded_kv_seq_lens.emplace_back(kv_seq_len); - } - kv_seq_lens_vec = &expanded_kv_seq_lens; - } - const bool use_fia_graph_decode = enable_fia_decode_ && (!attn_metadata.is_spec_verify || attn_metadata.expanded_decode.enabled); @@ -343,6 +330,19 @@ void AttentionImpl::decoder_forward(torch::Tensor& query, return; } + std::vector expanded_kv_seq_lens; + const std::vector* kv_seq_lens_vec = + &attn_metadata.kv_seq_lens_host_vec; + if (attn_metadata.expanded_decode.enabled) { + expanded_kv_seq_lens.reserve( + attn_metadata.expanded_decode.kv_seq_lens_host_vec.size()); + for (int32_t kv_seq_len : + attn_metadata.expanded_decode.kv_seq_lens_host_vec) { + expanded_kv_seq_lens.emplace_back(kv_seq_len); + } + kv_seq_lens_vec = &expanded_kv_seq_lens; + } + CHECK(v_cache.has_value() && v_cache->defined()) << "FIA decode requires a value cache"; CHECK(block_table.defined()) << "FIA decode requires a block table"; diff --git a/xllm/core/runtime/acl_graph_executor_impl.cpp b/xllm/core/runtime/acl_graph_executor_impl.cpp index 852ea4b68f..4bab761d61 100644 --- a/xllm/core/runtime/acl_graph_executor_impl.cpp +++ b/xllm/core/runtime/acl_graph_executor_impl.cpp @@ -615,6 +615,14 @@ void AclGraph::prepare_static_graph_tasks( auto& causal_conv1d_tasks = graph_task_context_->causal_conv1d_tasks; auto& fused_infer_attention_tasks = graph_task_context_->fused_infer_attention_tasks; + std::vector expanded_kv_seq_lens; + if (!fused_infer_attention_tasks.empty()) { + CHECK_GT(signal.spec_width, 0); + expanded_kv_seq_lens.reserve(static_cast(signal.spec_width)); + for (int64_t token_idx = 0; token_idx < signal.spec_width; ++token_idx) { + expanded_kv_seq_lens.emplace_back(signal.base_kv_seq_len + token_idx); + } + } auto signal_causal_conv1d_task = [&](CausalConv1dGraphTask& task) { CHECK(task.event != nullptr) @@ -627,7 +635,7 @@ void AclGraph::prepare_static_graph_tasks( CHECK(task.branch == FusedInferAttentionGraphBranch::kSpecVerify) << "static MTP FIA task must use the spec-verify branch"; CHECK_EQ(static_cast(task.query.size(0)), - signal.expanded_kv_seq_lens.size()) + expanded_kv_seq_lens.size()) << "static MTP FIA KV lengths must match captured query rows"; c10_npu::graph_task_update_begin(signal_stream, task.handle); kernel::npu::npu_fused_infer_attention_decode_out( @@ -636,7 +644,7 @@ void AclGraph::prepare_static_graph_tasks( task.value, task.block_table, task.actual_seq_lengths, - signal.expanded_kv_seq_lens, + expanded_kv_seq_lens, task.num_heads, task.num_key_value_heads, task.scale, @@ -943,7 +951,9 @@ bool AclGraph::prepare_static_mtp_graph_tasks( if (!fused_infer_attention_tasks.empty()) { const size_t graph_batch_size = static_cast(fused_infer_attention_tasks.front().query.size(0)); - if (signal.expanded_kv_seq_lens.size() != graph_batch_size) { + if (signal.spec_width != static_cast(graph_batch_size) || + signal.max_kv_seq_len != + signal.base_kv_seq_len + signal.spec_width - 1) { return false; } for (const FusedInferAttentionGraphTask& task : diff --git a/xllm/core/runtime/executor_impl.h b/xllm/core/runtime/executor_impl.h index 5866910d40..a2fa9d3cad 100644 --- a/xllm/core/runtime/executor_impl.h +++ b/xllm/core/runtime/executor_impl.h @@ -38,8 +38,8 @@ struct SpecVerifyGraphTaskSignal { int64_t num_accepted_tokens = 0; int64_t spec_width = 0; int64_t block_table_width = 0; + int64_t base_kv_seq_len = 0; int64_t max_kv_seq_len = 0; - std::vector expanded_kv_seq_lens; }; class ExecutorImpl { diff --git a/xllm/core/runtime/mtp_worker_impl.cpp b/xllm/core/runtime/mtp_worker_impl.cpp index ece945f7b7..c775869303 100644 --- a/xllm/core/runtime/mtp_worker_impl.cpp +++ b/xllm/core/runtime/mtp_worker_impl.cpp @@ -3345,19 +3345,14 @@ bool MTPWorkerImpl::prepare_static_mtp_graph_tasks_before_final_draft( } const int64_t spec_width = options_.num_speculative_tokens() + 1; const int64_t base_kv_seq_len = kv_seq_lens.front(); - std::vector expanded_kv_seq_lens; - expanded_kv_seq_lens.reserve(static_cast(spec_width)); - for (int64_t token_idx = 0; token_idx < spec_width; ++token_idx) { - expanded_kv_seq_lens.emplace_back(base_kv_seq_len + token_idx); - } - const int64_t spec_verify_max_kv_seq_len = expanded_kv_seq_lens.back(); + const int64_t spec_verify_max_kv_seq_len = base_kv_seq_len + spec_width - 1; const SpecVerifyGraphTaskSignal signal{ .linear_state_id = input.input_params.embedding.linear_state_ids.front(), .num_accepted_tokens = accepted_prefix_lengths.front(), .spec_width = spec_width, .block_table_width = verify_block_table_width, + .base_kv_seq_len = base_kv_seq_len, .max_kv_seq_len = spec_verify_max_kv_seq_len, - .expanded_kv_seq_lens = std::move(expanded_kv_seq_lens), }; return impl_->prepare_static_mtp_graph_tasks(signal, *compute_stream_); #else