Skip to content

feat: use FIA for qwen3.5 prefill attention on npu. #1448

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
Sinle4Cat wants to merge 10 commits into
base: preview/qwen3.5-qwen3.6
Choose a base branch
from
Open

feat: use FIA for qwen3.5 prefill attention on npu. #1448

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
ee7a139
feat(npu): use FIA for qwen3.5 prefill attention
May 11, 2026
b448a69
Support x flash attention prefill on NPU
May 13, 2026
39f5b84
bugfix: x flash attention long kv x mte excepiton
May 14, 2026
f182a6f
Merge preview/qwen3.5-qwen3.6 into x_flash
May 14, 2026
b02e42f
bugfix: route chunked prefill to x flash attention
May 15, 2026
fa77f62
bugfix: improve x flash attention long kv stability
May 16, 2026
2d41489
Merge remote-tracking branch 'origin/preview/qwen3.5-qwen3.6' into pr…
May 18, 2026
53cbf1c
Merge remote-tracking branch 'origin/preview/qwen3.5-qwen3.6' into pr…
May 19, 2026
6d7cdfb
bugfix: fix long kv xflash prefill grouping
May 20, 2026
7192530
chore: remove qwen35 gdn debug logging
May 21, 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
2 changes: 1 addition & 1 deletion third_party/xllm_ops
Open in desktop
Submodule xllm_ops updated from 79eb46 to 409d1d
9 changes: 9 additions & 0 deletions xllm/core/framework/model/model_input_params.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -420,6 +420,12 @@ struct ModelInputParams {
safe_to(graph_buffer.attn_mask, device, true);
params.graph_buffer.tiling_data =
safe_to(graph_buffer.tiling_data, device, true);
params.graph_buffer.xfa_q_cu_seq_lens =
safe_to(graph_buffer.xfa_q_cu_seq_lens, device, true);
params.graph_buffer.xfa_extra_tiling =
safe_to(graph_buffer.xfa_extra_tiling, device, true);
params.graph_buffer.xfa_attn_mask =
safe_to(graph_buffer.xfa_attn_mask, device, true);

// params for flashinfer
params.paged_kv_indptr = safe_to(paged_kv_indptr, device);
Expand Down Expand Up @@ -666,6 +672,9 @@ struct ModelInputParams {
struct GraphBuffer {
torch::Tensor attn_mask;
torch::Tensor tiling_data;
torch::Tensor xfa_q_cu_seq_lens;
torch::Tensor xfa_extra_tiling;
torch::Tensor xfa_attn_mask;
bool use_expanded_decode_for_spec_verify_attention = false;
torch::Tensor expanded_kv_seq_lens;
torch::Tensor expanded_block_tables;
Expand Down
1 change: 1 addition & 0 deletions xllm/core/kernels/npu/xllm_ops/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ cc_library(
beam_search.cpp
select_unshared_kv.cpp
beam_search_rec.cpp
x_flash_attention_infer.cpp
DEPS
atb
torch_npu
Expand Down
4 changes: 4 additions & 0 deletions xllm/core/kernels/npu/xllm_ops/beam_search_rec.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,9 @@ void beam_search_rec(const torch::Tensor& logprobs,
create_acltensor(&out_beam_count_prefix_sums_ids, out_beam_count_prefix_sums);
create_acltensor(&out_sequence_ids, out_sequence);

CHECK_GT(top_tokens.dim(), 0)
<< "beam_search_rec: top_tokens must have at least one dimension";
const int64_t top_k = top_tokens.size(-1);
uint64_t workspace_size = 0;
aclOpExecutor* executor = nullptr;
CHECK_ACL_SUCCESS(
Expand All @@ -91,6 +94,7 @@ void beam_search_rec(const torch::Tensor& logprobs,
top_logprobs_ids,
sequence_group_ids,
current_step,
top_k,
out_token_ids_ids,
out_token_index_ids,
out_log_probs_ids,
Expand Down
Loading