feat(archon): add NPU support for varlen attention - #1686
Open
262913 wants to merge 1 commit into
Open
Conversation
Enable varlen attention to run on Ascend NPU via torch_npu, complementing the existing CUDA Flash Attention backend. Key changes: - add device dispatch in _varlen_attn: CUDA uses Flash Attention, NPU uses npu_fusion_attention - implement _varlen_attn_npu and _varlen_attn_backward_npu with torch_npu APIs, preserving the same interface as CUDA - add _make_causal_mask_npu to build explicit bool causal masks for NPU (sparse_mode=1, allMask) - unify return signature across CUDA and NPU backends Signed-off-by: xixianwu <49093929+262913@users.noreply.github.com>
262913
requested review from
fishcrap,
garrett4wade and
sitabulaixizawaluduo
as code owners
September 8, 2026 14:41
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Enable archon's varlen attention to run on Ascend NPU via torch_npu, complementing the existing CUDA Flash Attention backend within the archon module.
Key changes:
Type of Change
Checklist
pre-commit run --all-files)main/review-prcommand/create-prAdditional Context
_varlen_attn uniformly returns four tensors (output, lse_or_max, softmax_sum, aux) to satisfy the fixed return count required by custom_op. The shape of output is always [T_q, H, D]. The differences are: on CUDA, lse_or_max is softmax_lse with shape [H, T_q], and both softmax_sum and aux are placeholder zeros of shape [2]; on NPU, lse_or_max is softmax_max with shape [T_q, H, S], softmax_sum also has shape [T_q, H, S], and aux is [seed, offset] as shape [2] of int64.