Fix flash attention state accumulation

[Zijie-Tian]

Summary

• extend flash attention state tensor to store accumulated VKQ values
• adjust ops to read/store full state, keeping numerators across segments
• update API documentation and validation for new state layout
• enhance test to use expanded state tensor and verify exact match

Testing

• cmake --build build-x86_64 --config Release -j12
• ./build-x86_64/bin/test-flash-attn-state

---

https://chatgpt.com/codex/tasks/task_e_6854917486ac83329eb2492474af1ecc

[Copilot]

Pull Request Overview

Extends the flash-attention state tensor to include accumulated VKQ values, updates the relevant operators and API validation for the new layout, and enhances the test to work with the expanded state tensor.

• Extend state tensor shape from [2, n_heads * seq_len] to [2 + head_dim, n_heads * seq_len]
• Update ggml_flash_attn_ext_with_state and CPU op implementations to read/write VKQ slots
• Update API docs and bolster tests for the new layout

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
tests/test-flash-attn-state.cpp	Adjusted reset/print routines for expanded state; updated creation of state tensor
ggml/src/ggml.c	Updated state-dimension assertions in ggml_flash_attn_ext_with_state
ggml/src/ggml-cpu/ops.cpp	Expanded validation and read/write paths to handle VKQ values
ggml/include/ggml.h	Updated function comment to document new [M, S, VKQ...] layout

Comments suppressed due to low confidence (2)

tests/test-flash-attn-state.cpp:376

• [nitpick] After printing the final state, the test does not verify the VKQ portion of the state tensor. Add assertions to confirm that the VKQ entries match expected values across segments.

    int row = state->ne[0];

tests/test-flash-attn-state.cpp:154

• The variable head_dim is not defined in main() before it’s used here. Declare and initialize head_dim (e.g. from v->ne[0] or state->ne[0] - 2) prior to creating the state tensor.

    ggml_tensor * state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2 + head_dim, n_heads * seq_len);