Skip to content

feat: add opt-in FIA decode for Qwen3.5. #2240

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Enguikong wants to merge 7 commits into
base: main
Choose a base branch
from
Open

feat: add opt-in FIA decode for Qwen3.5. #2240

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
559b762
feat(npu): use FIA for eager decode attention
Enguikong Aug 5, 2026
29723ed
feat(npu): migrate Qwen3.5 decode attention from PA to FIA
Enguikong Aug 13, 2026
103a25a
test(npu): align FIA graph fixtures with runtime metadata
Enguikong Aug 14, 2026
3a2adb3
test(npu): cover Qwen3.5 FIA rollback routing
Enguikong Aug 14, 2026
5b48803
style(npu): use class for FIA workspace signature
Enguikong Aug 14, 2026
a29981c
feat(npu): make Qwen3.5 FIA decode opt-in
Enguikong Aug 20, 2026
6835765
perf(npu): avoid FIA metadata work on PA path
Enguikong Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 27 additions & 4 deletions tests/core/framework/config/config_json_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand Down Expand Up @@ -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_;
Expand All @@ -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:
Expand All @@ -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 {
Expand All @@ -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()),
Expand All @@ -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_)
Expand All @@ -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_;
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -629,6 +651,7 @@ TEST(ConfigJsonTest, DumpStartupConfigWritesNonDefaultValuesOnly) {
EXPECT_EQ(config_json.at("model_impl").get<std::string>(), "python");
EXPECT_EQ(config_json.at("python_graph_backend").get<std::string>(),
"cudagraphs");
EXPECT_TRUE(config_json.at("enable_fia_decode").get<bool>());
EXPECT_EQ(config_json.at("block_size").get<int32_t>(), 256);
EXPECT_FALSE(config_json.at("enable_prefix_cache").get<bool>());
EXPECT_EQ(config_json.at("max_tokens_per_batch").get<int32_t>(), 2048);
Expand Down
97 changes: 97 additions & 0 deletions tests/core/kernels/npu/npu_xllm_ops_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,6 +244,103 @@ TEST_F(NpuXllmOpsTest, EmbeddedInterpreterSeesOps) {
.item<float>();
}

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<int64_t> actual_seq_lengths = {1, 2, 3};
const std::vector<int64_t> 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_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<float>();
}

TEST_F(NpuXllmOpsTest, Qwen35_27B_TP4_FullAttentionMatchesReference) {
py::gil_scoped_acquire gil;
if (!is_ascend950_device()) {
Expand Down
49 changes: 48 additions & 1 deletion tests/core/runtime/acl_graph_executor_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ limitations under the License.
#include <cstdlib>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "common/metrics.h"
Expand All @@ -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"
Expand Down Expand Up @@ -137,6 +139,7 @@ TEST(AclGraphStaticGraphTaskSignatureTest,
.num_accepted_tokens = 4,
.spec_width = 5,
.block_table_width = 64,
.base_kv_seq_len = 252,
.max_kv_seq_len = 256,
};

Expand All @@ -148,6 +151,50 @@ TEST(AclGraphStaticGraphTaskSignatureTest,
EXPECT_FALSE(npu::make_static_graph_task_signature(params).has_value());
}

TEST(Qwen35FiaRoutingTest, UsesExactModelTypeWhitelist) {
const std::vector<std::string> 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<std::string> 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;
}
}

TEST(Qwen35FiaRoutingTest, RequiresExplicitEnableFlag) {
ExecutionConfig& execution_config = ExecutionConfig::get_instance();
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 {
const KVCache& first_full_attention_cache(
const std::vector<KVCache>& kv_caches) {
Expand Down Expand Up @@ -1218,7 +1265,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");
Expand Down
Loading
Loading