Add session option for a BNHS GroupQueryAttention Value cache layout

[Javier Martinez (javier-intel)]

Description

Situation

com.microsoft.GroupQueryAttention requires the Value KV-cache in BNSH layout - (batch_size, num_heads, sequence_length, head_size) - for both the past_value input and the present_value output. Applications allocate those buffers themselves and bind them across decode steps.

Obstacle

Some execution providers execute GQA faster when the Value cache is BNHS - (batch_size, num_heads, head_size, sequence_length) - because the second attention matmul (attn_weights @ V) becomes an NT gemm. The operator schema cannot simply change: it is a stable contrib op and most EPs are BNSH-only. There was also no way for an application to discover an EP's preference, nor to tell a session which layout its buffers use.

Resolution

The GQA node stays BNSH and the conversion moves into the graph, where an EP compiler can absorb it:

past_value (BNHS, graph input) -> Transpose[0,1,3,2] -> GQA -> Transpose[0,1,3,2] -> present_value (BNHS, graph output)

Three pieces:

1. EP advertises its preference. New well-known OrtEpDevice metadata key gqa_preferred_value_layout (kOrtEpDevice_EpMetadataKey_GqaPreferredValueLayout), values "BNSH" (assumed when absent) or "BNHS". No new C API - applications read it through the existing OrtApi::EpDevice_EpMetadata.
2. Application selects the layout. New session option session.gqa_value_layout (kOrtSessionOptionsGqaValueLayout), "BNSH" (default) or "BNHS". Any other value fails session initialization.
3. ORT core inserts the conversion. New GqaValueLayoutTransformer inserts the two Transpose(perm=[0,1,3,2]) nodes and swaps the last two dimensions of the past_value graph input and present_value graph output declared shapes, so InferenceSession::ValidateInputsOutputs accepts the application's buffers.

An EP that reports "BNHS" fuses Transpose -> GQA -> Transpose into a single operation that reads BNHS directly and aliases past_value/present_value to one buffer, so the transposes never materialize. An EP that does not fuse them executes them: still correct, but a full copy of the Value cache in each direction per step and no past/present buffer sharing. A post-partitioning check logs a warning naming any GQA node whose flanking transposes survived, so that cost is diagnosable rather than silent.

Implementation notes

• The transformer is invoked directly from InferenceSession::TransformGraph rather than registered as a Level 1 optimizer. It must run at every optimization level including ORT_DISABLE_ALL (registered transformers at Level 1 and above are skipped there), and it must run after the Level 1 TransposeOptimizer, whose job is moving, merging and cancelling Transpose nodes, so the pattern reaches GetCapability intact.
• Applied to the main graph only. Subgraphs (a BeamSearch decoder body, a Loop carried value) are not the application's boundary. Nodes whose past_value is not a graph input, or whose present_value is not a graph output, are skipped with a warning.
• Idempotent. A model saved via session.optimized_model_filepath already carries the transform and may be reloaded with the option still set; the transformer detects the existing pattern and no-ops.
• A 4-bit quantized Value cache is rejected with an error. Two 4-bit values are packed per byte along head_size, so a byte-wise Transpose cannot express the layout change and the declared-shape swap would be wrong.
• k_scale/v_scale need no change. The GQA node is BNSH on both sides after the transform, so a PER_CHANNEL scale still has to broadcast against a BNSH tensor. Applications supply v_scale in the model-declared [1, num_heads_k, 1, head_size] shape regardless of the cache layout chosen.
• The Key cache is unaffected.

Files

File	Change
include/onnxruntime/core/session/onnxruntime_ep_device_ep_metadata_keys.h	New metadata key
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h	New session option
onnxruntime/core/optimizer/gqa_value_layout_transformer.{h,cc}	New transformer and the unfused-transpose diagnostic
onnxruntime/core/session/inference_session.cc	Option validation, transformer invocation, post-partition diagnostic
onnxruntime/test/optimizer/gqa_value_layout_transformer_test.cc	New tests
onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc, onnxruntime/test/autoep/test_registration.cc	Example EP advertises the key, with a round-trip assertion
docs/design/GQA_Value_Tensor_Layout.md	Design document

Testing

Ten new cases are added, covering transpose insertion and boundary shape swapping, idempotency, the past-only and present-only variants, both skip conditions, the 4-bit rejection, and three session-level tests for the option plumbing — including that the transform applies at ORT_DISABLE_ALL, which pins down the placement decision. Behavior is unchanged unless the new session option is set to "BNHS".

Motivation and Context

Applications that manage a KV cache across decode steps (onnxruntime-genai and similar) must allocate the Value cache in whichever layout their target EP executes best, but ORT offered no mechanism to negotiate that. Without one, an EP whose GQA implementation prefers BNHS either gives up the gain or the application guesses, with no way to stay correct when a layer falls back to a BNSH-only provider.

This change adds the negotiation - EP advertises, application selects, ORT core adapts the graph - while keeping the GQA schema and every existing kernel untouched. Expressing the layout change as ordinary Transpose nodes means correctness does not depend on the EP fusing them: a provider that cannot simply runs them.

Follow-ups not included here, tracked in the design document:

• CPU-fallback numerical parity tests.
• The compiling EP's fusion support.
• The ORT-format load path. PartitionOrtFormatModel does not go through TransformGraph, so for now the transform must be applied at conversion time, which the idempotency guard makes safe.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds end-to-end support for opting into a BNHS (vs BNSH) Value KV-cache layout for com.microsoft.GroupQueryAttention, including EP discoverability, session option plumbing, a graph transformer to adapt boundaries, logging for unfused transposes, and coverage via new unit tests plus a design doc.

Changes:

• Introduces GqaValueLayoutTransformer to insert Transpose(perm=[0,1,3,2]) around GQA Value cache and swap boundary shapes to BNHS.
• Adds a public session option (session.gqa_value_layout) and EP metadata key (gqa_preferred_value_layout) and wires both into the runtime plus example plugin/tests.
• Adds a design document and unit tests validating transformer behavior and session-option behavior (incl. ORT_DISABLE_ALL).

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/core/optimizer/gqa_value_layout_transformer.h	Declares transformer + layout constants and unfused-transpose logging helper.
onnxruntime/core/optimizer/gqa_value_layout_transformer.cc	Implements BNHS boundary adaptation, validation, idempotency detection, and post-partition diagnostics.
onnxruntime/core/session/inference_session.cc	Wires the session option to run the transformer and logs unfused transposes post-partition.
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h	Documents and exposes session.gqa_value_layout.
include/onnxruntime/core/session/onnxruntime_ep_device_ep_metadata_keys.h	Adds gqa_preferred_value_layout metadata key for EPs.
onnxruntime/test/optimizer/gqa_value_layout_transformer_test.cc	New unit tests for transformer behavior and session-option plumbing.
onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc	Example plugin now advertises BNHS preference via EP metadata.
onnxruntime/test/autoep/test_registration.cc	Validates the new EP metadata key round-trips.
docs/design/GQA_Value_Tensor_Layout.md	Design spec for BNHS Value layout support and rationale/contract.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Chi Lo (chilo-ms)]

Some other comments:

1. Minimal builds will likely fail to link. inference_session.cc unconditionally references GqaValueLayoutTransformer and  LogUnfusedGqaValueLayoutTransposes , but gqa_value_layout_transformer.cc is not included by the minimal/extended-minimal source lists in cmake/onnxruntime_optimizer.cmake .

2. ORT-format sessions silently ignore BNHS. .ort loading bypasses TransformGraph, so setting BNHS on an untransformed ORT model does nothing. With dynamic or coincident dimensions this can pass validation and produce incorrect results. Support PartitionOrtFormatModel or explicitly reject BNHS there.

3. Numerical fallback coverage should not be deferred. Correct execution on a non-fusing EP is a core promise of this design, yet the PR only tests graph structure. Add CPU end-to-end coverage, including aliased past/present buffers, before merging.

[Javier Martinez (javier-intel)]

Some other comments:

1. Minimal builds will likely fail to link. inference_session.cc unconditionally references GqaValueLayoutTransformer and  LogUnfusedGqaValueLayoutTransposes , but gqa_value_layout_transformer.cc is not included by the minimal/extended-minimal source lists in cmake/onnxruntime_optimizer.cmake .

2. ORT-format sessions silently ignore BNHS. .ort loading bypasses TransformGraph, so setting BNHS on an untransformed ORT model does nothing. With dynamic or coincident dimensions this can pass validation and produce incorrect results. Support PartitionOrtFormatModel or explicitly reject BNHS there.

3. Numerical fallback coverage should not be deferred. Correct execution on a non-fusing EP is a core promise of this design, yet the PR only tests graph structure. Add CPU end-to-end coverage, including aliased past/present buffers, before merging.

1. Minimal build breakage: no change needed. All three references (inference_session.cc:1640,1646,1741) are inside the #if !defined(ORT_MINIMAL_BUILD) block opened at line 1334, and cmake/adjust_global_compile_flags.cmake:70-76 nests ORT_EXTENDED_MINIMAL_BUILD inside onnxruntime_MINIMAL_BUILD, so extended-minimal defines ORT_MINIMAL_BUILD too. Added a comment at the call site and §4.3 recording this, so in the future it doesn't have to re-derived.

2. ORT-format silently ignores the option: now rejected. PartitionOrtFormatModel errors out when the option is not "BNSH", with a message pointing at converting the model with the transform applied. Tests: RejectsOrtFormatModel and AllowsOrtFormatModelWithTheDefaultLayout.

3. Fallback numerics: added. BnhsMatchesBnshOnCpu (bit-exact output and present_value) and BnhsWithAliasedCacheBufferMatchesSeparateBuffersOnCpu (one buffer bound to both sides via IOBinding).

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

Suppressed comments (3)

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:182

• These per-node INFO logs can become noisy for models with many GQA nodes and can clutter session initialization logs. Consider lowering them to VERBOSE/DEBUG, or aggregating into a single INFO summary (e.g., counts of transformed/skipped nodes) while keeping WARNING logs for actionable skip conditions and unfused-transpose diagnostics.

    if (AlreadyTransformed(graph, node, logger)) {
      LOGS(logger, INFO) << "GroupQueryAttention node '" << DescribeNode(node)
                         << "' already uses the BNHS Value layout. Skipping.";
      continue;
    }

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:279

• These per-node INFO logs can become noisy for models with many GQA nodes and can clutter session initialization logs. Consider lowering them to VERBOSE/DEBUG, or aggregating into a single INFO summary (e.g., counts of transformed/skipped nodes) while keeping WARNING logs for actionable skip conditions and unfused-transpose diagnostics.

    LOGS(logger, INFO) << "Applied the BNHS Value layout to GroupQueryAttention node '" << DescribeNode(node) << "'.";

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:221

• The message formats (consumers.size() - 1) directly. While this path is expected to have at least one consumer (this node), using size_t subtraction can underflow and print a huge number if the graph is malformed or consumer discovery returns empty unexpectedly. Safer: compute other_consumers = consumers.size() > 0 ? consumers.size() - 1 : 0 for logging, or log consumers.size() and clarify whether it includes the current node.

      const auto consumers = graph.GetConsumerNodes(boundary_arg->Name());
      if (consumers.size() != 1 || consumers[0] != &node) {
        LOGS(logger, WARNING) << "GroupQueryAttention node '" << DescribeNode(node) << "' shares its past_value graph "
                              << "input ('" << boundary_arg->Name() << "') with " << (consumers.size() - 1)
                              << " other node(s). The BNHS Value layout will not be applied to this node.";
        continue;
      }

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Suppressed comments (3)

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:186

• Validation runs before checking whether this GQA's cache tensors are application boundaries. Consequently, a main-graph GQA with an internal 4-bit cache fails session initialization even though the documented behavior is to warn and skip nodes whose past_value is not a graph input or whose present_value is not a graph output. Move ValidateNode after the boundary and sole-ownership eligibility checks so only nodes that will actually be transformed are rejected.

    ORT_RETURN_IF_ERROR(ValidateNode(node));

onnxruntime/core/session/inference_session.cc:1755

• This diagnostic also runs while saving an ORT-format model, where partitioning uses kAssignOnly and deliberately does not call Compile or remove fused nodes (graph_partitioner.h:34). The transposes therefore survive by design and trigger a false warning even if the runtime EP will fuse them when the ORT model is loaded. Suppress this check during the assign-only conversion path.

    if (gqa_value_layout != kGqaValueLayoutBNSH) {
      LogUnfusedGqaValueLayoutTransposes(graph, *session_logger_);

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:315

• The claim that boundary buffer sharing is lost is contradicted by BnhsWithAliasedCacheBufferMatchesSeparateBuffersOnCpu, which successfully binds one buffer to both boundary values. What is lost is the GQA kernel's in-place/shared-buffer path, so the diagnostic should describe the extra internal BNSH buffers rather than tell users past/present sharing is impossible.

                          << "nodes. The BNHS Value cache will be transposed at runtime: expect a full copy of the "
                          << "cache per step and no past/present buffer sharing. Either select the '"
                          << kGqaValueLayoutBNSH << "' layout for '" << kOrtSessionOptionsGqaValueLayout

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Suppressed comments (6)

onnxruntime/core/session/inference_session.cc:55

• This guard does not match the implementation's build availability. cmake/onnxruntime_optimizer.cmake:15-64 omits gqa_value_layout_transformer.cc from extended-minimal sources, so extended-minimal compiles these calls but cannot link them. Plain minimal has the inverse problem: this header is excluded while PartitionOrtFormatModel unconditionally references kGqaValueLayoutBNSH. Align the source/feature guards and place the layout-value constants in an always-available header.

#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
#include "core/optimizer/gqa_value_layout_transformer.h"
#endif

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:201

• Skipping the whole node here violates the session-wide layout contract when present_value is still a graph output (the new SkipsWhenPastValueIsNotAGraphInput test constructs exactly that case). Initialization succeeds but the application receives BNSH despite selecting BNHS. Transform the eligible output boundary independently, or fail initialization instead of continuing.

    if (has_past_value && !graph.IsInputsIncludingInitializers(node.InputDefs()[kPastValueInputIndex])) {
      LOGS(logger, WARNING) << "GroupQueryAttention node '" << DescribeNode(node) << "' has a past_value input ('"
                            << node.InputDefs()[kPastValueInputIndex]->Name() << "') that is not a graph input. "
                            << "The BNHS Value layout will not be applied to this node.";
      continue;

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:208

• Skipping the whole node here leaves a graph-input past_value declared BNSH even though the application selected BNHS (as exercised structurally by SkipsWhenPresentValueIsNotAGraphOutput). A later run can therefore reject the BNHS buffer or misinterpret it when dimensions are dynamic. Transform the eligible input side independently, or return an initialization error.

    if (has_present_value && !graph.IsOutput(node.OutputDefs()[kPresentValueOutputIndex])) {
      LOGS(logger, WARNING) << "GroupQueryAttention node '" << DescribeNode(node) << "' has a present_value output ('"
                            << node.OutputDefs()[kPresentValueOutputIndex]->Name() << "') that is not a graph output. "
                            << "The BNHS Value layout will not be applied to this node.";
      continue;

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:234

• Continuing after this warning leaves present_value at the graph boundary in BNSH despite the BNHS session option, so callers observe a layout that contradicts the session contract. Preserve a BNSH value for internal consumers while transposing the graph output, or fail initialization for this unsupported topology.

      if (!consumers.empty()) {
        LOGS(logger, WARNING) << "GroupQueryAttention node '" << DescribeNode(node) << "' has a present_value graph "
                              << "output ('" << boundary_arg->Name() << "') that is also consumed by " << consumers.size()
                              << " node(s) inside the graph. The BNHS Value layout will not be applied to this node.";
        continue;

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:288

• This diagnostic misses the main compiling-EP failure mode. GraphPartitioner::PlaceNode replaces any capability with a MetaDef by a fused node (core/framework/graph_partitioner.cc:597-615), so an EP that compiles the GQA alone leaves both inserted Transposes in the graph but removes the GroupQueryAttention node; this loop then emits no warning. Detect surviving layout Transposes independently of whether the original GQA still exists (or retain explicit markers that can be checked after partitioning).

void LogUnfusedGqaValueLayoutTransposes(const Graph& graph, const logging::Logger& logger) {
  for (const auto& node : graph.Nodes()) {
    if (node.OpType() != "GroupQueryAttention" || node.Domain() != kMSDomain) {
      continue;

docs/design/GQA_Value_Tensor_Layout.md:109

• This design snippet contradicts the implemented contract in gqa_value_layout_transformer.h:44-47, where ShouldOnlyApplyOnce() is deliberately not overridden so repeated application exercises structural idempotency. Remove this override from the documented class definition.

  bool ShouldOnlyApplyOnce() const override { return true; }

[Copilot]

Pull request overview

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

Suppressed comments (10)

onnxruntime/core/session/inference_session.cc:1646

• This condition also enables the transformer in extended-minimal builds, but cmake/onnxruntime_optimizer.cmake:15-64 does not add gqa_value_layout_transformer.cc to that build. The calls here and after partitioning will therefore have no linked implementation (and minimal GraphTransformer::Apply also skips Resolve). Restrict this feature to !defined(ORT_MINIMAL_BUILD) and reject the option for both minimal variants, or add complete extended-minimal support.

#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)

onnxruntime/core/session/inference_session.cc:2349

• Pure minimal builds exclude gqa_value_layout_transformer.h, so kGqaValueLayoutBNSH is undefined here and inference_session.cc will not compile. Move the accepted-value constants to an always-available session-options header (or otherwise avoid referencing the transformer header from the ORT-format path).

  ORT_RETURN_IF(sess_options.config_options.GetConfigOrDefault(kOrtSessionOptionsGqaValueLayout,
                                                               kGqaValueLayoutBNSH) != kGqaValueLayoutBNSH,

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:288

• After partitioning, a compiling EP can replace a capability containing only the GQA node with a fused node while leaving both transposes in the graph. This loop then finds no GroupQueryAttention node and emits no warning, even though the expensive fallback remains. Track or inspect the inserted transpose nodes directly so the diagnostic also covers standalone GQA compilation.

void LogUnfusedGqaValueLayoutTransposes(const Graph& graph, const logging::Logger& logger) {
  for (const auto& node : graph.Nodes()) {
    if (node.OpType() != "GroupQueryAttention" || node.Domain() != kMSDomain) {
      continue;

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:223

• Logging and continuing leaves a graph-bound past_value declared and consumed as BNSH even though the session successfully accepted the BNHS option. With static unequal dimensions this causes a run-time validation failure; with dynamic or square dimensions it can silently compute on the wrong layout. Either transform all GQA consumers of the boundary or fail session initialization instead of returning success.

      if (consumers.size() != 1 || consumers[0] != &node) {
        LOGS(logger, ERROR) << "GroupQueryAttention node '" << DescribeNode(node) << "' shares its past_value graph "
                              << "input ('" << boundary_arg->Name() << "') with " << (consumers.size() - 1)
                              << " other node(s). The BNHS Value layout will not be applied to this node.";
        continue;

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:234

• This continue likewise lets initialization succeed while leaving the application-visible present_value in BNSH despite the BNHS session option. An internal consumer can remain on the original BNSH value while only the graph output is routed through a transpose; if that rewrite is not implemented, initialization must fail rather than expose the wrong layout.

      if (!consumers.empty()) {
        LOGS(logger, ERROR) << "GroupQueryAttention node '" << DescribeNode(node) << "' has a present_value graph "
                              << "output ('" << boundary_arg->Name() << "') that is also consumed by " << consumers.size()
                              << " node(s) inside the graph. The BNHS Value layout will not be applied to this node.";
        continue;

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:186

• The idempotency check runs before the mandatory 4-bit validation. A pre-transformed or manually authored 4-bit graph matches AlreadyTransformed and skips the rejection, allowing byte-wise transposes that this code identifies as incorrect. Validate every GQA before taking the already-transformed fast path.

This issue also appears on line 285 of the same file.

    if (AlreadyTransformed(graph, node, logger)) {
      LOGS(logger, INFO) << "GroupQueryAttention node '" << DescribeNode(node)
                         << "' already uses the BNHS Value layout. Skipping.";
      continue;
    }

    ORT_RETURN_IF_ERROR(ValidateNode(node));

onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:315

• The warning's claim that past/present buffer sharing is unavailable is contradicted by BnhsWithAliasedCacheBufferMatchesSeparateBuffersOnCpu: callers can still bind one boundary buffer, although the GQA kernel cannot alias its internal BNSH operands and extra intermediates are required. Describe the lost kernel-level in-place optimization rather than saying buffer sharing is impossible.

                          << "nodes. The BNHS Value cache will be transposed at runtime: expect a full copy of the "
                          << "cache per step and no past/present buffer sharing. Either select the '"
                          << kGqaValueLayoutBNSH << "' layout for '" << kOrtSessionOptionsGqaValueLayout

onnxruntime/test/optimizer/gqa_value_layout_transformer_test.cc:725

• This assertion always succeeds for the current non-square cache because ExpectTensorsEqual first rejects the BNSH and BNHS shapes ([B,N,S,H] versus [B,N,H,S]). It therefore does not establish that the data is non-transpose-invariant as claimed. Compare the raw element sequences while intentionally ignoring shape, or compare against an independently constructed non-invariant pattern.

  ASSERT_FALSE(ExpectTensorsEqual(bnsh_fetches[present_value_index], bnhs_fetches[present_value_index],
                                  "present_value")
                   .IsOK())
      << "BNSH and BNHS present_value are byte-identical, so this test cannot detect a layout bug.";

onnxruntime/core/session/inference_session.cc:1651

• ORT_RETURN_IF_NOT constructs a FAIL status, but the documented contract says an unsupported option value returns INVALID_ARGUMENT. Callers inspecting the status code will receive the wrong API result; construct an ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, ...) instead.

    ORT_RETURN_IF_NOT(gqa_value_layout == kGqaValueLayoutBNHS,
                      "Invalid value for session option '", kOrtSessionOptionsGqaValueLayout, "': '",
                      gqa_value_layout, "'. Expected '", kGqaValueLayoutBNSH, "' or '", kGqaValueLayoutBNHS, "'.");

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h:639

• The option documentation incorrectly says non-fusion prevents callers from sharing the boundary buffer; the new aliased-buffer CPU test demonstrates that binding one buffer still works. What is lost is the GQA kernel's direct aliasing/in-place path because the transpose intermediates separate its operands.

// Transpose -> GroupQueryAttention -> Transpose sequence into a single operation; an EP that does
// not will execute the transposes, which is correct but costs a full copy of the Value cache in
// each direction per step and prevents past/present buffer sharing.

[Copilot]

🔵 Needs a closer look

Correctness-sensitive graph rewriting across ONNX, ORT-format, provider-partitioning, and minimal-build paths warrants final human review.

Review details

Suppressed comments (1)

docs/design/GQA_Value_Tensor_Layout.md:656

• This test summary is inaccurate: the ORT-format test rejects a "BNHS" request, while the new AllowsOrtFormatModelWithTheDefaultLayout test explicitly verifies that setting "BNSH" succeeds. Please state which option value is rejected.

- An ORT format model fails session initialization with `ORT_INVALID_ARGUMENT` when the option is set, and loads normally when it
  is not.

• Files reviewed: 11/12 changed files
• Comments generated: 2
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Default ORT-format loading incurs an unnecessary quadratic-style graph scan in minimal builds.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 11/12 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Boundary detection has correctness gaps around device copies and missing ONNX imports, and subgraph behavior contradicts the stated contract.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (4)

Previously missed (2) — in code that hasn't changed since the last review.

onnxruntime/core/optimizer/gqa_value_layout_boundaries.cc:88

• This treats any custom-domain node named MemcpyFromHost/MemcpyToHost as an ORT device copy. The repository's canonical predicate requires kOnnxDomain (onnxruntime/core/framework/utils.cc:100-103); without that check, boundary tracing can cross an unrelated custom op and misclassify the cache layout.
  onnxruntime/core/optimizer/gqa_value_layout_transformer.cc:254
• A model containing only com.microsoft nodes can validly omit the default ONNX opset import. Returning success here then inserts ONNX-domain Transpose nodes, but Graph::SetOpSchemaFromRegistryForNode cannot resolve a node whose domain is absent from DomainToVersionMap, so Graph::Resolve() fails after the graph was mutated. Reject this case before transformation (or explicitly add a supported import) and add a no-default-opset regression test.

onnxruntime/core/optimizer/gqa_value_layout_boundaries.cc:260

• This has the symmetric reload gap for GQA -> MemcpyToHost -> Transpose -> boundary, which occurs when GQA is on a device EP and the Transpose falls back to CPU. Because only direct Transpose consumers are examined, the converted output is classified as out of scope and explicit BNSH is not rejected. Walk through device-copy consumers before testing for the Transpose and add a cross-EP saved-model test.

  for (const Node* consumer : ConsumersOf(graph, arg->Name())) {
    if (consumer == nullptr || !IsGqaValueLayoutTranspose(*consumer) || consumer->OutputDefs().empty()) {
      continue;

onnxruntime/core/session/inference_session.cc:1694

• The PR description says subgraph GQA caches are outside the application boundary and are skipped with a warning, but this rejects every BNHS session containing any subgraph GQA. That also rejects a Loop whose GQA cache is entirely internal and unrelated to any application-bound cache. Either trace whether the subgraph cache reaches a main-graph boundary and reject only that case, or update the stated API contract and PR description to document the broader restriction.

    const GqaNodeCounts gqa_nodes = CountGqaNodes(graph);
    if (gqa_nodes.in_subgraphs != 0) {
      ORT_RETURN_IF_ERROR_SESSIONID_(ORT_MAKE_STATUS(

• Files reviewed: 11/12 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

Boundary tracing can miss unfused transposes, and the supported quantized path lacks mixed-alias coverage.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

onnxruntime/core/session/inference_session.cc:1701

• The PR description says subgraph GQA nodes are outside the application boundary and the transform is applied only to the main graph, but this rejects every BNHS session containing any subgraph GQA—even when that cache is wholly internal, as the design document's open item acknowledges. Either trace whether the subgraph cache reaches a main-graph boundary before failing, or update the PR contract to explicitly document this broader rejection.

    const GqaNodeCounts gqa_nodes = CountGqaNodes(graph);
    if (gqa_nodes.in_subgraphs != 0) {
      ORT_RETURN_IF_ERROR_SESSIONID_(ORT_MAKE_STATUS(

• Files reviewed: 14/15 changed files
• Comments generated: 2
• Review effort level: Balanced

[Hariharan Seshadri (hariharans29)]

Review — PR #32139: Add session option for a BNHS GroupQueryAttention Value cache layout

Substantial infrastructure PR (+5,011 / −27 across 17 files, 29 commits, 163 conversation items, 5 participants). Well past the point where a line-by-line review is useful; my read is at the architecture and correctness-contract level, plus a sanity check on the CUDA kernel-level change that fell out of the design.

Summary of what the PR does

Introduces a session option session.gqa_value_layout (values "BNSH" default, "BNHS") that lets an application declare the layout of the past_value/present_value KV-cache buffers it will bind. The GQA operator schema stays BNSH; when the application chooses BNHS, ORT inserts Transpose(perm=[0,1,3,2]) on both sides of every application-visible GQA node and swaps the boundary shapes to BNHS. An EP that prefers BNHS is expected to fuse Transpose → GQA → Transpose into its native BNHS-consuming kernel; one that doesn't executes the transposes at correct-but-slow cost.

Also adds a companion EP metadata key gqa_preferred_value_layout on OrtEpDevice so applications can discover EP preference without a new C API surface.

Architecture — the load-bearing design decisions

1. Transform invoked directly from InferenceSession::TransformGraph, not registered as a Level-1 optimizer. Two reasons in the comment:

   • Must run at ORT_DISABLE_ALL because it changes the session's input/output contract, not just the graph.
   • Must run after the Level-1 TransposeOptimizer, so the inserted Transpose → GQA → Transpose pattern reaches GetCapability intact for a fusing EP.

   Both correct. The location is unusual for a graph transform, but the reasoning is exactly right — anything registered at Level 1+ would be skipped at DISABLE_ALL, and running before TransposeOptimizer would risk the pattern being merged away before an EP could see it. ✓

2. CMake gating via onnxruntime_ENABLE_GQA_VALUE_LAYOUT (default ON, cmake_dependent_option auto-disables in MINIMAL_BUILD, EXTENDED_MINIMAL_BUILD, and DISABLE_CONTRIB_OPS). Compiles out the transformer and helpers entirely in disabled builds, and Initialize() rejects the option with INVALID_ARGUMENT if the app sets it in such a build. Chi Lo's opening concern about minimal-build linkage was closed correctly this way. ✓

3. Failure-mode enumeration is exhaustive. From onnxruntime_session_options_config_keys.h:

   • past_value graph input read by multiple nodes, or present_value graph output also consumed inside the graph → fail (can't reshape a shared cache for one reader);
   • node already carries the transform on only one side → fail (asymmetric = broken contract);
   • 4-bit quantized Value cache → fail (byte-packed along head_size, byte-wise transpose can't express the layout change);
   • Value cache tensor is not rank 4 → fail;
   • reached through a device-copy node → fail (transform can't be inserted across a copy);
   • GQA node is inside a Loop body / BeamSearch decoder → fail (operator and boundary in different graphs);
   • ORT-format model → fail (the ORT-format load path bypasses TransformGraph).

   Every one of these has a rationale about why silently leaving BNSH would let the app bind BNHS buffers to a BNSH boundary and silently misread — that framing (silent misread vs loud failure) is the correct correctness axis for this kind of contract change. ✓

4. Symmetric enforcement in both directions. Setting "BNSH" explicitly on a model that already carries BNHS conversion (e.g., one saved from a BNHS session via session.optimized_model_filepath and reloaded) also fails — otherwise the app would bind BNSH buffers to a BNHS boundary. Setting nothing on a converted model loads with a warning. This is the exactly-correct three-state matrix: explicit-BNSH + BNHS-model = fail, explicit-BNHS + BNHS-model = load, unset + BNHS-model = warn.

5. Post-partition diagnostic ReportUnfusedGqaValueLayoutTransposes anchored on boundary names (not GQA node identity), for exactly the right reason spelled out in the header:

   "A compiling EP may claim the GQA node and replace it with a fused node while leaving the flanking Transposes in the graph; both full-cache copies still execute, but there is no GQA node left to search from."

   Load-bearing: without this, a silently-non-fusing EP turns into a large per-step cost with nothing in the logs to explain it. ✓

6. Device-copy tracing. The TraceGqaBoundary{Back,Forward}ThroughDeviceCopies helpers walk up to kMaxDeviceCopyHops = 4 hops through MemcpyFromHost/MemcpyToHost nodes to reach the boundary. Rationale in the comment: MemcpyTransformer runs inside TransformGraph before serialization, so a model saved from a non-CPU session can have a copy spliced between a boundary and the provider-side nodes. Handling this was Tianlei's 22717ef fix; without it, boundary detection would miss BNHS-converted graphs that had been round-tripped through a GPU session. 4 hops is generous — MemcpyTransformer inserts at most one copy per boundary — but bounds the recursion safely.

The CUDA kernel-level change worth calling out

group_query_attention.cc had:

ORT_ENFORCE(past_key_shared == past_value_shared,
            "past_key/present_key and past_value/present_value must be both shared or both separate.");
parameters.past_present_share_buffer = past_key_shared;

Now:

parameters.past_present_share_buffer = past_key_shared && past_value_shared;
...
if (past_key_shared != past_value_shared) {
  // Nonshared preprocessing overwrites present KV, so preserve the aliased past cache first.
  ...
  cudaMemcpyAsync(separate_past_buffer.get(), shared_past->DataRaw(), past_bytes, cudaMemcpyDeviceToDevice, Stream(context));
  ...
}

This is a substantive change to the CUDA kernel's aliasing contract, not just cosmetic. It's necessary because after BNHS conversion, the Value cache flows through ORT-allocated Transpose intermediates on both sides of GQA, so past_value == present_value alias no longer holds, while past_key == present_key can still hold (Key cache untouched). The kernel now supports mixed aliasing by staging the aliased cache aside first.

Constraint that comes with this: sliding-window cache still requires both pairs aliased (eviction rewrites in place), so sliding_window_cache=1 + BNHS is rejected at runtime. Test CudaCacheAliasingRejectsMixedSlidingWindow pins this. Reasonable and documented in the option doc.

Verification story: the three new CUDA tests (CudaCacheAliasingUnfused, CudaCacheAliasingFlash, CudaCacheAliasingRejectsMixedSlidingWindow) walk all four (share_key × share_value) combinations, compare output against a single-shared reference, and verify the kernel-log-emitted SDPA kernel name. Rigorous.

Given the CUDA-owned nature of this change: Tianlei Wu (CUDA EP maintainer) authored the fix commit b6ba85a "fix(gqa): preserve cache aliases with BNHS layout", so effectively he's a co-author of the CUDA half. That's the right expertise applied. ✓

Test coverage

• gqa_value_layout_transformer_test.cc — 2,563 LOC — covers transpose insertion, boundary shape swap, idempotency, past-only and present-only variants, both skip conditions, 4-bit rejection, session-option plumbing including ORT_DISABLE_ALL, subgraph rejection, ORT-format handling, device-copy paths, initializer-backed boundary handling, and mixed-alias paths. Ten scenarios documented by the author, plus everything the reviewer commits added.
• CPU numerical parity: BnhsMatchesBnshOnCpu (bit-exact output and present_value) and BnhsWithAliasedCacheBufferMatchesSeparateBuffersOnCpu (one buffer bound to both sides via IOBinding). Chi Lo's third opening concern about deferred fallback numerics is closed here.
• CUDA parity: three tests above.
• Minimal-build rejection: RejectsGqaValueLayoutOptionWhenDisabled in ort_model_only_test.cc.
• Example EP integration: round-trip metadata-key test in test_registration.cc.

Test-to-code ratio is well above 1:1 for the reviewable core. Proportionate to the correctness surface.

Design doc

docs/design/GQA_Value_Tensor_Layout.md — 785 LOC. Not rendered inline in the PR diff, but the scope + PR-description structure (Situation / Obstacle / Resolution / Implementation notes / Files / Testing) strongly suggests a real design document rather than filler. Author has kept it updated across the PR's iterations (see commit 8771c0a "Documentation update to match the code"). Someone auditing this in 6 months should be pointed at that doc.

Review activity — how the design got hardened

• Chi Lo's three opening concerns (minimal build, ORT-format silent-ignore, CPU fallback numerics) are all explicitly closed with commits + tests.
• 6+ Copilot review rounds — each closed either with code changes or with the reviewer's rationale. Copilot correctly flagged real issues (converted-layout detectability after EP fusion, subgraph GQA rejection, boundary-tracing gaps around device copies), all addressed.
• Tianlei's 4 fix commits landed the device-copy tracing, the mixed-alias CUDA path, and the build-configuration gating.
• github-advanced-security flagged 2 weeks ago — resolved.

This is a change that has been through real technical review, not a rubber-stamp. The final commit 5e8e595 "fix(gqa): gate Value layout support by build configuration" is a build-config tightening; not a correctness fix.

Concerns to raise

1. CI status on the tip commit. 5e8e595 shows 30/81 checks OK — that's much lower than the prior f6a38a7 at 77/91. Likely mid-run given the timing (1 hour before I looked), but worth watching. Given the number of code paths touched (bundled EPs, plugin EP, contrib ops, minimal builds, ORT format), it's not implausible that a build-configuration commit temporarily broke a leg.

   Need to confirm CI settles green before merge — the check-count-drop combined with a commit specifically about build-configuration gating is exactly the pattern where a #if slip would show up.

2. One cpplint nit flagged in gqa_value_layout_boundaries.cc:315 — Add #include <string> for std::string. Mechanical.

3. Sliding-window + BNHS rejection. The current behavior: BNHS conversion breaks the both-shared aliasing requirement for sliding-window, so those sessions must not enable BNHS. This is documented, tested (CudaCacheAliasingRejectsMixedSlidingWindow), and correct — but it's a real limitation. If a downstream user hits it, the error path is at runtime (in group_query_attention.cc's sliding_window_cache=1 guard), not at session init. Consider a session-init-time check that rejects the combination up front instead. Non-blocking but would improve the error UX.

4. The session.gqa_value_layout name is fine, but consider whether "value_layout" is precise enough — a future BNHS-for-Key extension would need a different option. Naming as session.gqa_value_cache_layout would be more specific. Bikeshed; non-blocking.

5. Post-partition diagnostic can't distinguish an EP that intentionally didn't fuse (because the shape is out-of-envelope) from one that silently declined despite advertising "BNHS" preference. The warning message currently just names the boundary. A future refinement could match against gqa_preferred_value_layout and escalate the warning when the EP claimed BNHS support but left transposes. Not this PR's job.

Scope note

The design doc lists three follow-ups not in this PR:

• CPU-fallback numerical parity tests → now included, closed
• Compiling EP fusion support → out of scope for this PR (each EP does its own fusion)
• ORT-format load path → now handled by rejection at load time, not by extending the transform

So the design doc's "follow-ups not included here" list is stale — items 1 and 3 are now addressed. Worth a quick pass on the doc to remove or update those. Non-blocking.

Recommendation

Approve — pending:

1. CI settles green on the tip commit 5e8e595 (or subsequent). The 30/81 snapshot needs to resolve. This is the one item I'd hold on.
2. #include <string> fix in gqa_value_layout_boundaries.cc.
3. Design-doc follow-up list update to reflect that items 1 and 3 are now in-PR. Small documentation-only item.

Non-blocking:

4. Consider a session-init-time reject for sliding_window_cache=1 + BNHS instead of the runtime guard, for a cleaner error UX.
5. Consider session.gqa_value_cache_layout for naming precision (future-proof against a Key-layout companion option).

Overall assessment

This is one of the more carefully-designed infrastructure PRs I've seen recently. The failure-mode enumeration is exhaustive and framed on the correct axis (silent misread vs loud failure). The transform placement (TransformGraph not Level-1 registry) is unusual but the reasoning is exactly right. The idempotency contract is spelled out and tested. The post-partition diagnostic is anchored on boundaries rather than on GQA-node identity, which correctly handles the "compiling EP replaces the GQA node" case. The CUDA mixed-aliasing change is co-authored by the CUDA EP maintainer with rigorous test coverage.

The scope is large (5K LOC) but proportional to the correctness surface — a subtle boundary-layout mismatch would produce silently wrong outputs on dynamic-or-square-dimension caches, which is exactly the kind of bug that would ship undetected without the exhaustive fail-fast validation this PR implements.

Once CI is green and the two mechanical items are addressed, this is ready to land.