Skip to content

NPU Flash Attention. #33036

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

Jump to bottom
Draft
wants to merge 33 commits into
base: master
Choose a base branch
from
Draft

NPU Flash Attention. #33036

wants to merge 33 commits into from
+2,030 −29

Conversation

@intelgaoxiong
Copy link
Contributor

@intelgaoxiong intelgaoxiong commented Nov 26, 2025

Details:

  • item1
  • ...

Tickets:

  • ticket-id

@github-actions github-actions bot added category: build OpenVINO cmake script / infra category: NPU OpenVINO NPU plugin category: NPUW NPUW plugin labels Nov 26, 2025
@intelgaoxiong intelgaoxiong force-pushed the xiong/npu_flash_attention_wip branch 2 times, most recently from 5b789e6 to f798d10 Compare November 28, 2025 09:08
@intelgaoxiong
Initialize HFA.
4eafeff 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
[Func]Generate HFA tile model.
3bfcbe8 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
[Compiled]Compile HFA tile model.
c192f89 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Set parameter and result names.
aa0a9fd 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
[Debug]Print IOs.
6cb131a 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Use fp16 for input and output (mask is fp32)
609618a 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Infer()
9ed0f38 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
wip - incorrect
f2f9d9d 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
E2e execution with correct results.
48a0d74 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Debug code clean up.
82e2e5a 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Wrap run_hfa_tiled_inference function.
0857d82 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Optimize final computing: TTFT 25s -> 5s.
ae66d88 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Optimize final computing by model infer: TTFT 5s -> 2.5s
583e100 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Optimize preparation by broadcast in model: TTFT 2.5s -> 2s
84baa29 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Decouple HFA and original model: store and infer hfa IO.
5bcd97d 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Skip compilation for the original SDPA model.
bc9e1d9 
Create HFA infer requests when creating JustSyncInferRequest.

Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Refine code for K V Mast tile extraction.
ebdb81c 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Align mask dtype to avoid convert in runtime. 2s -> 1.8s
ea6fd31 
Refactor code for HFATile model creation.

Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Fix incorrect index and clean up.
c1942ae 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Extract tile with copy_tensor_by_dim.
4dfc654 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Extract tiles only if cannot reuse source
75c38a6 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Remove perf stastics.
ee470d2 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Refactor run_hfa_tiled_inference
85ad0f7 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Access IO via created mapping.
00d60e8 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Refactor host_flash_attention.
a0ccc25 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Extract KV seq dim.
f481e7b 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Use query size for HFA chunking.
c7988fe 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Minor optimization on run_hfa_tiled_inference
b350b5d 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Create output name to id mapping.
31fe728 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong intelgaoxiong force-pushed the xiong/npu_flash_attention_wip branch 2 times, most recently from 37b9c27 to 6696123 Compare November 28, 2025 14:37
@intelgaoxiong intelgaoxiong changed the title npu flash attention NPU Flash Attention. Nov 28, 2025
@intelgaoxiong
Clean up.
1d9e498 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong intelgaoxiong force-pushed the xiong/npu_flash_attention_wip branch from 6696123 to 1d9e498 Compare November 28, 2025 15:10
@intelgaoxiong
Extract IO indices during compilation.
c9b6a7a 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong intelgaoxiong force-pushed the xiong/npu_flash_attention_wip branch from 8968944 to c9b6a7a Compare November 29, 2025 01:53
@intelgaoxiong
Refine run_hfa_tiled_inference.
ffa1d00 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
@intelgaoxiong
Fixed serialization.
5e40138 
Signed-off-by: intelgaoxiong <xiong.gao@intel.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

category: build OpenVINO cmake script / infra category: NPU OpenVINO NPU plugin category: NPUW NPUW plugin

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@intelgaoxiong