Skip to content

[NPUW] Enable mask skipping for fused flash attention #36077

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
kkoryun wants to merge 9 commits into
base: master
Choose a base branch
from
+179 −53
Open

[NPUW] Enable mask skipping for fused flash attention #36077

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
89a9754
remove broadcast for fused hfa
kkoryun May 8, 2026
ba3fcd6
change compiler version to support GQA
kkoryun May 8, 2026
d4dbbf5
Merge branch 'master' into enabling_gqa_hfa
kkoryun May 15, 2026
68ab3ac
refactoring
kkoryun May 20, 2026
ff5cac7
added nomask subgraph
kkoryun May 26, 2026
9d133d4
revert host_flash_attention.cpp
kkoryun May 26, 2026
6a59259
Merge branch 'master' into enable_mask_skipping
kkoryun Jun 3, 2026
7b06aec
fix
kkoryun Jun 3, 2026
8bc4f47
fixed false branch for no mask
kkoryun Jun 3, 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
120 changes: 88 additions & 32 deletions src/plugins/intel_npu/src/plugin/npuw/attn/attn_subgraph.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,10 @@ struct PyramidRequestSet {

struct HFARequestSet {
enum TileIdx : std::size_t {
REGULAR_TILE = 0,
FINAL_TILE = 1,
COUNT = 2,
REGULAR_TILE_MASK = 0,
Copy link
Copy Markdown
Contributor

@esmirno esmirno Jun 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if it is regular tile with mask - may be better to have TILE_WITH_MASK and REGULAR_TILE

REGULAR_TILE_NO_MASK = 1,
FINAL_TILE = 2,
COUNT = 3,
};

std::array<ov::SoPtr<ov::IAsyncInferRequest>, COUNT> infer_requests{};
Expand Down Expand Up @@ -262,7 +263,7 @@ void ensure_pyramid_requests(ov::npuw::v1::subgraphs::InferContext& ctx, Runtime
}

void ensure_hfa_requests(ov::npuw::v1::subgraphs::InferContext& ctx, RuntimeState& state) {
if (state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE]) {
if (state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_MASK]) {
return;
}

Expand All @@ -275,25 +276,54 @@ void ensure_hfa_requests(ov::npuw::v1::subgraphs::InferContext& ctx, RuntimeStat
auto& request = get_request(ctx);
const bool is_piped = request.is_subrequest_pipelined(ctx.real_subgraph_idx);

state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE] = hfa->_compiled_tile_model->create_infer_request();
state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_MASK] =
hfa->_compiled_tile_model->create_infer_request();
if (hfa->_compiled_tile_no_mask_model) {
state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_NO_MASK] =
hfa->_compiled_tile_no_mask_model->create_infer_request();
}
state.hfa_requests.infer_requests[HFARequestSet::FINAL_TILE] = state.base_request;
if (is_piped) {
state.hfa_requests.pipeline_requests[HFARequestSet::REGULAR_TILE] =
state.hfa_requests.pipeline_requests[HFARequestSet::REGULAR_TILE_MASK] =
hfa->_compiled_tile_model->create_infer_request();
if (hfa->_compiled_tile_no_mask_model) {
Copy link
Copy Markdown
Contributor

@esmirno esmirno Jun 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is clear place to spot a problem - initial no_mask_model might refer missing a model, so better to have compiled_tile, compiled_tile_with_mask

state.hfa_requests.pipeline_requests[HFARequestSet::REGULAR_TILE_NO_MASK] =
hfa->_compiled_tile_no_mask_model->create_infer_request();
}
state.hfa_requests.pipeline_requests[HFARequestSet::FINAL_TILE] = state.base_pipeline_request;
}

const size_t num_inputs = hfa->_compiled_tile_model->inputs().size();
for (size_t input_idx = 0; input_idx < num_inputs; ++input_idx) {
const auto tile_input = hfa->_compiled_tile_model->inputs()[input_idx];
// Share final tile model tensors with the mask tile request (all inputs, including mask)
const size_t num_mask_inputs = hfa->_compiled_tile_model->inputs().size();
for (size_t input_idx = 0; input_idx < num_mask_inputs; ++input_idx) {
const auto mask_tile_input = hfa->_compiled_tile_model->inputs()[input_idx];
const auto final_tile_input = hfa->_compiled_final_tile_model->inputs()[input_idx];

auto main_tensor = state.base_request->get_tensor(final_tile_input);
state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE]->set_tensor(tile_input, main_tensor);

state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_MASK]->set_tensor(mask_tile_input, main_tensor);
if (is_piped) {
auto pipeline_tensor = state.base_pipeline_request->get_tensor(final_tile_input);
state.hfa_requests.pipeline_requests[HFARequestSet::REGULAR_TILE]->set_tensor(tile_input, pipeline_tensor);
state.hfa_requests.pipeline_requests[HFARequestSet::REGULAR_TILE_MASK]->set_tensor(mask_tile_input,
pipeline_tensor);
}
}

// Share final tile model tensors with the no-mask tile request (only the non-mask inputs)
if (hfa->_compiled_tile_no_mask_model) {
const size_t num_no_mask_inputs = hfa->_compiled_tile_no_mask_model->inputs().size();
for (size_t input_idx = 0; input_idx < num_no_mask_inputs; ++input_idx) {
const auto no_mask_tile_input = hfa->_compiled_tile_no_mask_model->inputs()[input_idx];
const auto final_tile_input = hfa->_compiled_final_tile_model->inputs()[input_idx];

auto main_tensor = state.base_request->get_tensor(final_tile_input);
state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_NO_MASK]->set_tensor(no_mask_tile_input,
main_tensor);
if (is_piped) {
auto pipeline_tensor = state.base_pipeline_request->get_tensor(final_tile_input);
state.hfa_requests.pipeline_requests[HFARequestSet::REGULAR_TILE_NO_MASK]->set_tensor(
no_mask_tile_input,
pipeline_tensor);
}
}
}

Expand All @@ -309,13 +339,26 @@ void ensure_hfa_requests(ov::npuw::v1::subgraphs::InferContext& ctx, RuntimeStat
});

const auto& tile_in = hfa->_sdpa_attention_info._tile_input_indices;
auto state_acc = state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE]->get_tensor(
auto state_acc = state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_MASK]->get_tensor(
hfa->_compiled_tile_model->inputs()[tile_in.acc]);
auto state_max = state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE]->get_tensor(
auto state_max = state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_MASK]->get_tensor(
hfa->_compiled_tile_model->inputs()[tile_in.max]);
auto state_sum = state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE]->get_tensor(
auto state_sum = state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_MASK]->get_tensor(
hfa->_compiled_tile_model->inputs()[tile_in.d]);

// Share state tensors between mask and no-mask tile requests
if (hfa->_compiled_tile_no_mask_model) {
state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_NO_MASK]->set_tensor(
hfa->_compiled_tile_no_mask_model->inputs()[tile_in.acc],
state_acc);
state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_NO_MASK]->set_tensor(
hfa->_compiled_tile_no_mask_model->inputs()[tile_in.max],
state_max);
state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_NO_MASK]->set_tensor(
hfa->_compiled_tile_no_mask_model->inputs()[tile_in.d],
state_sum);
}

runtime::host_flash_attention::HFARuntimeContext::initialize_state_tensors(state_acc, state_max, state_sum);
runtime::host_flash_attention::HFARuntimeContext::StateBuffers initial_buffers{state_acc, state_max, state_sum};
state.hfa_runtime_ctx->initialize_state_buffers(
Expand Down Expand Up @@ -772,6 +815,7 @@ ov::npuw::v1::subgraphs::RuntimeBehaviorFactory make_runtime_factory() {
OPENVINO_ASSERT(hfa_desc->is_valid(), "HFA configuration must be valid");
const int64_t tile_size = hfa_desc->_tile_size;
const int64_t total_kv_length = state.hfa_selector->context_length();
const int64_t actual_kv_length = state.hfa_selector->current_length();
const int64_t num_tiles = total_kv_length / tile_size;
OPENVINO_ASSERT(total_kv_length % tile_size == 0,
"HFA total KV length must be multiple of tile size for now");
Expand All @@ -786,8 +830,15 @@ ov::npuw::v1::subgraphs::RuntimeBehaviorFactory make_runtime_factory() {
auto present_key_tensor = hfa_inputs.at(sdpa_in.present_key);
auto attention_mask_tensor = hfa_inputs.at(sdpa_in.attention_mask);
auto present_value_tensor = hfa_inputs.at(sdpa_in.present_value);
auto& regular_tile_request = state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE];
// If the regular tile is not fully filled, need to use the mask
const bool use_mask = (actual_kv_length + 1) % tile_size != 0;
const bool use_no_mask_model = !use_mask && hfa_desc->_compiled_tile_no_mask_model;
auto& regular_tile_request =
use_no_mask_model ? state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_NO_MASK]
Comment on lines +833 to +837
: state.hfa_requests.infer_requests[HFARequestSet::REGULAR_TILE_MASK];
auto& final_tile_request = state.hfa_requests.infer_requests[HFARequestSet::FINAL_TILE];
const auto& compiled_regular_tile_model =
use_no_mask_model ? hfa_desc->_compiled_tile_no_mask_model : hfa_desc->_compiled_tile_model;
auto attention_output_tensor =
final_tile_request->get_tensor(hfa_desc->_compiled_final_tile_model->outputs()[0]);
const auto& tile_in = sdpa_info._tile_input_indices;
Expand All @@ -799,34 +850,33 @@ ov::npuw::v1::subgraphs::RuntimeBehaviorFactory make_runtime_factory() {
state_acc = current_buffer.acc;
state_max = current_buffer.max;
state_sum = current_buffer.sum;
regular_tile_request->set_tensor(hfa_desc->_compiled_tile_model->inputs()[tile_in.acc],
regular_tile_request->set_tensor(compiled_regular_tile_model->inputs()[tile_in.acc],
state_acc);
regular_tile_request->set_tensor(hfa_desc->_compiled_tile_model->inputs()[tile_in.max],
regular_tile_request->set_tensor(compiled_regular_tile_model->inputs()[tile_in.max],
state_max);
regular_tile_request->set_tensor(hfa_desc->_compiled_tile_model->inputs()[tile_in.d],
regular_tile_request->set_tensor(compiled_regular_tile_model->inputs()[tile_in.d],
state_sum);
} else {
state_acc =
regular_tile_request->get_tensor(hfa_desc->_compiled_tile_model->inputs()[tile_in.acc]);
regular_tile_request->get_tensor(compiled_regular_tile_model->inputs()[tile_in.acc]);
state_max =
regular_tile_request->get_tensor(hfa_desc->_compiled_tile_model->inputs()[tile_in.max]);
regular_tile_request->get_tensor(compiled_regular_tile_model->inputs()[tile_in.max]);
state_sum =
regular_tile_request->get_tensor(hfa_desc->_compiled_tile_model->inputs()[tile_in.d]);
regular_tile_request->get_tensor(compiled_regular_tile_model->inputs()[tile_in.d]);
runtime::host_flash_attention::HFARuntimeContext::initialize_state_tensors(state_acc,
state_max,
state_sum);
}

regular_tile_request->set_tensor(hfa_desc->_compiled_tile_model->inputs()[tile_in.q],
regular_tile_request->set_tensor(compiled_regular_tile_model->inputs()[tile_in.q],
query_tensor);
final_tile_request->set_tensor(hfa_desc->_compiled_final_tile_model->inputs()[tile_in.q],
query_tensor);
regular_tile_request->set_tensor(hfa_desc->_compiled_tile_model->outputs()[tile_out.acc],
regular_tile_request->set_tensor(compiled_regular_tile_model->outputs()[tile_out.acc],
state_acc);
regular_tile_request->set_tensor(hfa_desc->_compiled_tile_model->outputs()[tile_out.max],
regular_tile_request->set_tensor(compiled_regular_tile_model->outputs()[tile_out.max],
state_max);
regular_tile_request->set_tensor(hfa_desc->_compiled_tile_model->outputs()[tile_out.d],
state_sum);
regular_tile_request->set_tensor(compiled_regular_tile_model->outputs()[tile_out.d], state_sum);
final_tile_request->set_tensor(hfa_desc->_compiled_final_tile_model->inputs()[tile_in.acc],
state_acc);
final_tile_request->set_tensor(hfa_desc->_compiled_final_tile_model->inputs()[tile_in.max],
Expand All @@ -848,10 +898,14 @@ ov::npuw::v1::subgraphs::RuntimeBehaviorFactory make_runtime_factory() {
int64_t kv_offset,
int64_t mask_offset,
int64_t tile_length,
bool async = false) {
bool async = false,
bool process_with_mask = true) {
auto k_tile_buffer = request->get_tensor(model->inputs()[tile_in.k]);
auto v_tile_buffer = request->get_tensor(model->inputs()[tile_in.v]);
auto mask_tile_buffer = request->get_tensor(model->inputs()[tile_in.mask]);
ov::SoPtr<ov::ITensor> mask_tile_buffer;
if (process_with_mask) {
mask_tile_buffer = request->get_tensor(model->inputs()[tile_in.mask]);
}

if (can_reuse_tensor_zero_copy(k_source,
k_tile_buffer,
Expand Down Expand Up @@ -879,7 +933,7 @@ ov::npuw::v1::subgraphs::RuntimeBehaviorFactory make_runtime_factory() {
extract_and_copy_tile(v_source, v_tile_buffer, V_SEQ_DIM, kv_offset, tile_length, "V");
}

if (attention_mask_tensor) {
if (process_with_mask && attention_mask_tensor) {
if (can_reuse_tensor_zero_copy(attention_mask_tensor,
mask_tile_buffer,
MASK_KV_SEQ_DIM,
Expand Down Expand Up @@ -934,12 +988,14 @@ ov::npuw::v1::subgraphs::RuntimeBehaviorFactory make_runtime_factory() {
int64_t kv_tile_offset = 0;
for (int64_t tile_idx = 0; tile_idx < num_tiles - 1; ++tile_idx) {
process_tile(regular_tile_request,
hfa_desc->_compiled_tile_model,
compiled_regular_tile_model,
past_key_tensor,
past_value_tensor,
kv_tile_offset,
mask_tile_offset,
tile_size);
tile_size,
false,
use_mask);
Comment on lines 994 to +998
kv_tile_offset += tile_size;
mask_tile_offset += tile_size;
}
Expand Down
14 changes: 14 additions & 0 deletions src/plugins/intel_npu/src/plugin/npuw/compiled_model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1842,6 +1842,10 @@ bool ov::npuw::CompiledModel::compile_for_success(std::size_t id, const std::vec
}

hfa->set_compiled_tile_model(make_wrapped(hfa->_tile_model_to_compile, "/hfa_tile", devices));
if (hfa->_tile_no_mask_model_to_compile) {
hfa->set_compiled_tile_no_mask_model(
make_wrapped(hfa->_tile_no_mask_model_to_compile, "/hfa_tile_no_mask", devices));
}
hfa->set_compiled_final_tile_model(desc.compiled_model);
LOG_INFO("Host flash attention compilation complete for Subgraph[" << id << "]");

Expand Down Expand Up @@ -1990,6 +1994,16 @@ void ov::npuw::CompiledModel::dump_subgraph_model(std::size_t id,
std::string hfa_final_tile_model_dump_path =
ov::util::path_join({dump_dir, hfa_final_tile_model_name}).string();
ov::save_model(hfa_final_tile_model, hfa_final_tile_model_dump_path);
LOG_INFO("Wrote " << hfa_final_tile_model_dump_path);

if (hfa->_tile_no_mask_model_to_compile) {
const auto& hfa_tile_no_mask_model = hfa->_tile_no_mask_model_to_compile;
std::string hfa_tile_no_mask_model_name = format_subgraph_name(id, funcall) + "_hfa_tile_no_mask.xml";
std::string hfa_tile_no_mask_model_dump_path =
ov::util::path_join({dump_dir, hfa_tile_no_mask_model_name}).string();
ov::save_model(hfa_tile_no_mask_model, hfa_tile_no_mask_model_dump_path);
LOG_INFO("Wrote " << hfa_tile_no_mask_model_dump_path);
}
}
}

Expand Down
Loading
Loading