Create output name to id mapping.

Signed-off-by: intelgaoxiong <xiong.gao@intel.com>

Author: intelgaoxiong

src/plugins/intel_npu/src/plugin/npuw/host_flash_attention.cpp

@@ -579,6 +579,49 @@ static void build_tile_param_mapping(HostFlashAttention& hfa, const std::shared_
     std::cout << "==================================================\n" << std::endl;
 }
 
+// ============================================================================
+// Helper function: Build tile model output index mapping
+// ============================================================================
+static void build_tile_output_mapping(HostFlashAttention& hfa, const std::shared_ptr<ov::Model>& tile_model) {
+    LOG_INFO("Building HFA Tile Model output index mapping...");
+
+    // Parse tile model outputs by their tensor names
+    // Expected output order: [acc, maxx, d]
+    const auto& tile_outputs = tile_model->outputs();
+    for (std::size_t i = 0; i < tile_outputs.size(); ++i) {
+        const auto& tensor_names = tile_outputs[i].get_names();
+        if (tensor_names.empty()) {
+            LOG_WARN("Tile model output[" << i << "] has no tensor name");
+            continue;
+        }
+
+        const std::string& name = *tensor_names.begin();
+
+        // Map tensor name to enum ID
+        if (name == "acc") {
+            hfa._tile_output_index_map[HFATileOutputId::ACC] = i;
+            LOG_DEBUG("Mapped ACC to tile output[" << i << "]");
+        } else if (name == "maxx") {
+            hfa._tile_output_index_map[HFATileOutputId::MAXX] = i;
+            LOG_DEBUG("Mapped MAXX to tile output[" << i << "]");
+        } else if (name == "d") {
+            hfa._tile_output_index_map[HFATileOutputId::D] = i;
+            LOG_DEBUG("Mapped D to tile output[" << i << "]");
+        } else {
+            LOG_WARN("Unknown tile model output name: " << name);
+        }
+    }
+
+    // Print the tile output mapping
+    std::cout << "\n========== HFA Tile Model Output Mapping ==========\n";
+    std::cout << "Total entries: " << hfa._tile_output_index_map.size() << "\n";
+
+    for (const auto& [output_id, output_idx] : hfa._tile_output_index_map) {
+        std::cout << "  " << hfa_tile_output_id_to_string(output_id) << " -> output[" << output_idx << "]" << std::endl;
+    }
+    std::cout << "==================================================\n" << std::endl;
+}
+
 // ============================================================================
 // Helper function: Extract sequence dimension from Concat node
 // ============================================================================
@@ -734,6 +777,11 @@ std::optional<HostFlashAttention> HostFlashAttention::from(const std::shared_ptr
     // ========================================================================
     build_tile_param_mapping(hfa, tile_model);
 
+    // ========================================================================
+    // Step 10: Build tile model output index mapping
+    // ========================================================================
+    build_tile_output_mapping(hfa, tile_model);
+
     LOG_INFO("Successfully created HostFlashAttention with query_size=" << query_size << ", tile_size=" << query_size);
 
     return hfa;
@@ -761,6 +809,9 @@ HostFlashAttention::HostFlashAttention(const function::HostFlashAttention& func_
     // Copy HFA Tile Model input index mapping from function HFA
     _sdpa_attention_info._tile_param_index_map = func_hfa._tile_param_index_map;
 
+    // Copy HFA Tile Model output index mapping from function HFA
+    _sdpa_attention_info._tile_output_index_map = func_hfa._tile_output_index_map;
+
     // Copy query size directly from function HFA (no need to extract from model)
     _sdpa_attention_info._query_size = func_hfa._query_size;
 

src/plugins/intel_npu/src/plugin/npuw/host_flash_attention.hpp

@@ -48,6 +48,18 @@ enum class HFATileInputId : uint8_t {
     COUNT
 };
 
+// HFA Regular Tile Model output tensor identifiers
+// Represents the output layout for regular (non-final) tile models
+// Output names: [acc, maxx, d]
+enum class HFATileOutputId : uint8_t {
+    ACC = 0,   // Accumulated attention output
+    MAXX = 1,  // Maximum values for numerical stability
+    D = 2,     // Normalization denominator
+
+    // Sentinel value for enum range
+    COUNT
+};
+
 // Helper functions to convert enum values to string representations for logging/debugging
 inline const char* sdpa_input_id_to_string(SDPAInputId id) {
     switch (id) {
@@ -89,6 +101,19 @@ inline const char* hfa_tile_input_id_to_string(HFATileInputId id) {
     }
 }
 
+inline const char* hfa_tile_output_id_to_string(HFATileOutputId id) {
+    switch (id) {
+    case HFATileOutputId::ACC:
+        return "ACC";
+    case HFATileOutputId::MAXX:
+        return "MAXX";
+    case HFATileOutputId::D:
+        return "D";
+    default:
+        return "UNKNOWN";
+    }
+}
+
 namespace function {
 
 // HostFlashAttention structure definition
@@ -124,6 +149,12 @@ struct HostFlashAttention {
     // This is created after tile model generation in from() method
     std::map<HFATileInputId, std::size_t> _tile_param_index_map;
 
+    // Tile model output index mapping
+    // Maps tile output IDs (UPDATED_ACC, UPDATED_MAX, UPDATED_D) to actual output indices
+    // Only applicable to regular tile model (final tile has single output at index 0)
+    // This is created after tile model generation in from() method
+    std::map<HFATileOutputId, std::size_t> _tile_output_index_map;
+
     // Validation helpers
     bool is_valid() const {
         return _tile_model != nullptr && _final_tile_model != nullptr && _tile_size > 0;
@@ -157,6 +188,11 @@ struct HostFlashAttentionInfo {
     // This allows accessing tile model parameters by semantic name
     // Populated from function::HostFlashAttention::_tile_param_index_map
     std::map<HFATileInputId, std::size_t> _tile_param_index_map;
+
+    // Mapping from HFA Tile output identifier to actual output index in tile model
+    // This allows accessing tile model outputs by semantic name rather than hardcoded indices
+    // Populated from function::HostFlashAttention::_tile_output_index_map
+    std::map<HFATileOutputId, std::size_t> _tile_output_index_map;
 };
 
 // Compile-time host flash attention information

src/plugins/intel_npu/src/plugin/npuw/just_sync_infer_request.cpp

@@ -1377,6 +1377,21 @@ void ov::npuw::JustInferRequest::run_hfa_tiled_inference(std::size_t real_idx, s
     const std::size_t tile_idx_max = get_tile_param_idx(ov::npuw::HFATileInputId::PAST_MAX);
     const std::size_t tile_idx_d = get_tile_param_idx(ov::npuw::HFATileInputId::PAST_D);
 
+    // Pre-cache regular tile model output indices (final tile has only 1 output at index 0)
+    // Regular tile outputs: [ACC, MAXX, D]
+    const auto& tile_output_map = hfa_desc._sdpa_attention_info._tile_output_index_map;
+    auto get_tile_output_idx = [&](ov::npuw::HFATileOutputId output_id) -> std::size_t {
+        auto it = tile_output_map.find(output_id);
+        if (it == tile_output_map.end()) {
+            OPENVINO_THROW("HFA: Tile output mapping not found for output ID: ", static_cast<uint8_t>(output_id));
+        }
+        return it->second;
+    };
+
+    const std::size_t regular_tile_output_acc = get_tile_output_idx(ov::npuw::HFATileOutputId::ACC);
+    const std::size_t regular_tile_output_max = get_tile_output_idx(ov::npuw::HFATileOutputId::MAXX);
+    const std::size_t regular_tile_output_d = get_tile_output_idx(ov::npuw::HFATileOutputId::D);
+
     auto state_acc = regular_tile_request->get_tensor(hfa_desc._compiled_tile_model->inputs()[tile_idx_acc]);
     auto state_max = regular_tile_request->get_tensor(hfa_desc._compiled_tile_model->inputs()[tile_idx_max]);
     auto state_sum = regular_tile_request->get_tensor(hfa_desc._compiled_tile_model->inputs()[tile_idx_d]);
@@ -1632,10 +1647,9 @@ void ov::npuw::JustInferRequest::run_hfa_tiled_inference(std::size_t real_idx, s
             final_attention_output->copy_to(attention_output_tensor._ptr);
         } else {
             // Regular tile: Update accumulation state for next iteration
-            // Tile model outputs: [0] updated_acc, [1] updated_max, [2] updated_sum
-            auto output_acc = current_request->get_tensor(current_model->outputs()[0]);
-            auto output_max = current_request->get_tensor(current_model->outputs()[1]);
-            auto output_sum = current_request->get_tensor(current_model->outputs()[2]);
+            auto output_acc = current_request->get_tensor(current_model->outputs()[regular_tile_output_acc]);
+            auto output_max = current_request->get_tensor(current_model->outputs()[regular_tile_output_max]);
+            auto output_sum = current_request->get_tensor(current_model->outputs()[regular_tile_output_d]);
 
             // Copy updated state back to input buffers for next tile
             output_acc->copy_to(state_acc._ptr);