OpenVINO #37606 — Why Unsqueeze(axis=0) Kept PagedAttention’s Token Axis Stuck at One

GGUF PagedAttention, Part 1 of 2 — A shared GET_ROWS translator, two attention layouts, and an embedding axis that had to correct itself inside AdaptToGenAI

The model worked through ordinary scaled dot-product attention.

Switch it to PagedAttention, and any multi-token prefill could fail.

The first visible problems were downstream:

Q and K activations no longer matched the SDPA path

RoPE positions aligned with the wrong tensor axis

PagedAttention received incompatible Q/K/V shapes

the CPU plugin could crash during multi-token prefill

The root cause appeared much earlier.

OpenVINO’s GGUF GET_ROWS translation restored the embedding result to rank four with:

Unsqueeze(axis = 0)

That inserted a literal leading dimension of one.

For the ordinary SDPA layout, that was correct.

After SDPAToPagedAttention, it was not.

PagedAttention flattened the active tokens into the leading dimension, but the embedding path inserted another leading one in front of them.

Expected after PagedAttention:

[tokens, 1, 1, hidden]

Old embedding result:

[1, tokens, 1, hidden]

The same values were present.

They were attached to the wrong axis.

PR #37606 changed the GenAI adaptation path so the restored dimension is inserted at:

axis = 1

rather than:

axis = 0

That one-axis move preserves the original SDPA shape while allowing the PagedAttention shape to retain the real token count in front.

SDPA:
[1, tokens, hidden]
→ Unsqueeze(axis=1)
→ [1, 1, tokens, hidden]

PagedAttention:
[tokens, 1, hidden]
→ Unsqueeze(axis=1)
→ [tokens, 1, 1, hidden]

The fix was deliberately placed inside:

AdaptToGenAI

rather than in the shared translate_get_rows() implementation.

That distinction matters because the same translator is also used by llama.cpp’s raw cgraph decoder, which does not run SDPAToPagedAttention and still requires the original layout.

The translation was not universally wrong. It became wrong only after another backend changed which dimension represented a token row.

OpenVINO merged the patch into master on August 25, 2026.

The minimum shape counterexample

Let:

T = 3 tokens

H = hidden width

Before PagedAttention, GenAI supplies token IDs in the ordinary layout:

input_ids
=
[1, T]

AdaptToGenAI lifts those IDs back into the rank-four GGUF form expected by the existing graph:

[1, 1, 1, T]

The embedding GET_ROWS path removes its first two singleton dimensions:

Squeeze(axes = [0, 1])

[1, 1, 1, T]
→
[1, T]

The vocabulary-table gather then produces:

[1, T, H]

Now compare the two possible Unsqueeze positions.

Old axis

Unsqueeze(axis = 0)

[1, T, H]
→
[1, 1, T, H]

Fixed axis

Unsqueeze(axis = 1)

[1, T, H]
→
[1, 1, T, H]

They produce the same SDPA shape.

That is the first important property of the fix:

ordinary layout
→ numerically and structurally unchanged

Now let SDPAToPagedAttention rewrite the token interface.

Its PagedAttention layout represents the same three active tokens as:

input_ids
=
[T]

and inserts a trailing singleton dimension for consumers of the original two-dimensional interface:

[T, 1]

After the GGUF rank-four lift:

[1, 1, T, 1]

and the same two-axis squeeze:

[1, 1, T, 1]
→
[T, 1]

the embedding gather produces:

[T, 1, H]

Here the old and fixed paths diverge.

Old axis under PagedAttention

Unsqueeze(axis = 0)

[T, 1, H]
→
[1, T, 1, H]

Fixed axis under PagedAttention

Unsqueeze(axis = 1)

[T, 1, H]
→
[T, 1, 1, H]

The old path moves the token count away from the leading dimension.

The fixed path preserves it.

For T = 3:

old:
[1, 3, 1, H]

required:
[3, 1, 1, H]

No unusual model size is needed.

One prompt containing more than one token is enough to make the two layouts different.

Why PagedAttention wants tokens in front

Ordinary attention commonly carries a tensor in a form equivalent to:

[batch, heads, sequence, head_size]

For the single-request GGUF path discussed here:

batch = 1

and the tokens belong to the sequence dimension.

PagedAttention serves a different execution model.

Continuous batching can flatten active token rows from one or more sequences into a leading token dimension.

Its Q input is eventually expected in a form where each leading row represents one active token.

Conceptually:

[token_rows, heads × head_size]

The rest of the converted graph must therefore preserve:

token_rows

as the dimension that PagedAttention treats as its leading work unit.

If the residual stream remains:

[1, T, 1, H]

instead of:

[T, 1, 1, H]

later reshape and projection operations start from a different dimensional contract.

A reshape can preserve the total number of values.

It cannot infer that the token count belonged to another semantic axis.

The RoPE branch already followed the PagedAttention layout

The position path had been made backend-dependent.

For the ordinary layout, it produced cosines and sines shaped approximately like:

[1, T, 1, rotary_half]

After PagedAttention, it produced:

[T, 1, 1, rotary_half]

The Q/K tensors derived from the residual stream needed to follow the same token-axis transition.

Before #37606:

RoPE position state:
[T, 1, 1, rotary_half]

embedding-derived state:
[1, T, 1, hidden]

The token count lived on different dimensions in the two branches.

Broadcasting can sometimes make incompatible shapes executable.

That does not make the result correct.

A singleton dimension on one side and a T dimension on the other can expand rather than reject.

[1, T, ...]
+
[T, 1, ...]
→
[T, T, ...]

Instead of representing three token rows, the graph may construct a three-by-three cross-product of dimensions.

That is the more dangerous outcome.

shape inference succeeds

but

the graph now represents the wrong computation

Why the fix uses axis 1

Moving the new singleton to axis 1 looks unusual at first.

Why not remove the extra dimension entirely?

Because the downstream GGUF graph still expects rank-four residual tensors.

The repair must satisfy two layouts with one graph.

Before PagedAttention

Gather result:
[batch, sequence, hidden]

batch = 1

Inserting at axis 1 gives:

[batch, 1, sequence, hidden]

which is the established GGUF/SDPA residual layout.

After PagedAttention

Gather result:
[token_rows, 1, hidden]

Inserting at axis 1 gives:

[token_rows, 1, 1, hidden]

which retains the flattened token rows at the front.

The inserted dimension is fixed.

Its position is chosen so the dimensions already present in the gather result are allowed to change meaning with the backend.

fixed singleton
→ axis 1

backend-dependent leading dimension
→ allowed to remain axis 0

The resulting operation is self-correcting with respect to the incoming layout.

It does not need a runtime branch asking:

Am I executing SDPA or PagedAttention?

The input shape itself determines the final arrangement.

The shared translator was not the right ownership boundary

The problematic Unsqueeze(axis=0) lives in:

translate_get_rows()

A direct edit there would have been smaller.

It would also have changed every consumer of that translation.

The function is shared by at least two different execution paths.

GGUF → OpenVINO GenAI
→ may run SDPAToPagedAttention

GGUF raw cgraph → llama.cpp-compatible path
→ does not run SDPAToPagedAttention

For the raw cgraph path, the restored leading singleton remains part of the expected GGML-style layout.

The translator does not know that a later GenAI pass will reinterpret token rows.

Changing the translator globally would make one consumer’s later transformation dictate another consumer’s base representation.

PR #37606 instead adds:

FixEmbdAxis

inside:

AdaptToGenAI

The ownership becomes:

translate_get_rows
→ preserve the shared GGUF translation contract

AdaptToGenAI
→ adapt that translation for the dual SDPA/PagedAttention contract

This is a narrower fix.

It changes only graphs entering the GenAI adaptation path.

The patch had to find the operation structurally

A first implementation could have searched for the main embedding node by name.

Real decoder graphs made that insufficient.

The main lookup often carries a name associated with:

embd

But some architectures contain additional token-indexed embedding-table lookups.

Gemma4 includes a per-layer path identified as:

pe_tok_flat

That node does not need to end in _embd.

It is still semantically the same class of operation:

embedding table

+

indices derived from inp_tokens

→ gathered per-token values

→ restored rank-four result

If only the main embedding branch moves to axis 1:

main branch:
[T, 1, 1, feature]

while the secondary branch remains:

secondary branch:
[1, T, 1, feature]

their later addition can broadcast into:

[T, T, ..., feature]

The final patch therefore does not trust the friendly name.

FixEmbdAxis matches the graph structure.

Unsqueeze
        ↓ input
optional Convert
        ↓
Gather
        ↓ indices
optional Squeeze
        ↓
Parameter named `inp_tokens`

Only an Unsqueeze whose gather indices trace back to the token-input parameter is rewritten.

This avoids two opposite mistakes.

Name-only matching
→ can miss valid secondary embedding lookups

Rewriting every Gather/Unsqueeze pair
→ can alter unrelated row selections

The structural root provides the actual semantic boundary.

A `Convert` could sit between the gather and the unsqueeze

translate_get_rows() may insert a type conversion when the gathered table type differs from the operation’s expected output type.

The graph can therefore be:

Gather
→ Unsqueeze

or:

Gather
→ Convert
→ Unsqueeze

A matcher requiring direct adjacency would cover only the first form.

The merged matcher looks through one optional Convert before checking for the gather.

This is a small detail with a larger implication:

same semantic embedding lookup

+

different element-type path

→ must receive the same axis repair

The transformation follows the operation’s meaning rather than one exact formatting of the graph.

The original `Unsqueeze` is replaced, not mutated in place

When the matcher finds the target:

old:
Unsqueeze(gather_result, axis=0)

it creates:

new:
Unsqueeze(gather_result, axis=1)

It then preserves:

the friendly name,
runtime metadata,
and existing consumers.

Conceptually:

same gather result

same node identity for downstream diagnostics

same runtime information

different restored axis

The gather itself remains unchanged.

The vocabulary lookup still returns the same values in the same order.

Only their rank-four interpretation changes.

The tests protect both sides of the layout contract

The patch added focused tests rather than relying only on the real-model report.

Structural axis test

The test builds the same graph pattern as the GGUF translator:

Squeeze(inp_tokens)
→ Gather(embedding_table)
→ Unsqueeze(axis=0)

After AdaptToGenAI, it verifies:

the gather is unchanged

the replacement remains an Unsqueeze

the new axis is 1

Ordinary-layout functional test

For token IDs shaped:

[1, T]

the result must remain:

[1, 1, T, H]

PagedAttention-layout functional test

For the equivalent post-rewrite shape:

[T, 1]

the result must become:

[T, 1, 1, H]

The test also checks the embedding values themselves.

The axis move is allowed to change shape interpretation.

It is not allowed to reorder or alter the selected vocabulary rows.

A secondary embedding test prevents name-based regression

Another regression builds two independent embedding-table lookups rooted at:

inp_tokens

One uses the ordinary embedding-style name.

The second deliberately does not.

After the adaptation pass, both must use:

Unsqueeze(axis=1)

This test records the review-discovered Gemma4 boundary:

semantic source
→ token input dependency

not:

semantic source
→ node name suffix

A future simplification back to name matching would fail even if the main embd path continued to work.

The real transformation is also present in the test pipeline

An early test version simulated the post-PagedAttention input shape directly.

That proved the local axis transformation.

It did not prove that the complete:

AdaptToGenAI
→ SDPAToPagedAttention

pipeline remained structurally valid.

Review requested a stronger integration boundary.

The merged suite therefore constructs a small decoder containing:

real Q/K/V projections

stateful KV-cache ReadValue and Assign paths

ScaledDotProductAttention

token-derived embedding input

It then runs:

AdaptToGenAI

followed by

SDPAToPagedAttention

The test requires:

the transformation completes

a PagedAttentionExtension is created

its Q input has the expected rank

its trailing hidden width remains correct

This is structural end-to-end coverage of the actual transformation that exposed the bug.

It is not a full numerical execution of PagedAttentionExtension.

That operation also requires runtime-managed inputs such as block tables, past lengths, and subsequence boundaries that are normally supplied by the GenAI runtime.

The test boundary is therefore explicit:

real graph transformation
→ covered

complete runtime-managed PagedAttention execution
→ not performed by this frontend unit test

Real-model validation went beyond the unit graph

The PR author additionally reports end-to-end validation with:

Qwen2.5-0.5B-Instruct

SmolLM2-135M

Both models were passed through:

GGUF frontend

→ AdaptToGenAI

→ SDPAToPagedAttention

→ CPU plugin

Before the fix:

multi-token prefill crashed

After the fix:

the models compiled and inferred

Q activations matched the SDPA reference bit for bit

This is author-reported validation, not an independently reproduced Resonetic Lab test.

It nevertheless adds a runtime layer beyond the checked-in structural regressions.

The unchanged consumer was tested separately

Because the patch intentionally avoided changing translate_get_rows(), the other consumer still needed protection.

The PR author reports building llama.cpp against the modified frontend and running:

test-backend-ops:
2,146 / 2,146 passed

test-recurrent-state-rollback:
passed

test-thread-safety:
passed

The existing OpenVINO GGUF frontend suite also passed:

245 / 245

These results support the ownership choice.

GenAI-specific adaptation changed

shared raw translator unchanged

existing raw-decoder behavior retained

They do not prove every llama.cpp model and architecture is unaffected.

They establish the reported suite coverage.

What the patch directly changed in this first boundary

For the embedding-axis problem, PR #37606 directly:

adds FixEmbdAxis to AdaptToGenAI,
finds token embedding lookups structurally,
looks through an optional Convert,
traces gather indices back to inp_tokens,
replaces Unsqueeze(axis=0) with Unsqueeze(axis=1),
preserves node names and runtime metadata,
applies the rewrite before GenAI replaces the original GGUF inputs,
covers the main embedding lookup,
covers secondary token-table lookups such as the Gemma4 pattern,
tests ordinary and post-PagedAttention shapes,
and runs the real attention transformation structurally.

The patch also changes inp_out_ids row selection.

That is a separate code and indexing contract covered in Part 2.

What this article does not establish

The public evidence does not establish that:

every GGUF architecture had the same axis failure,
every PagedAttention backend crashed in the same way,
all GPU plugins reproduced the CPU failure,
full PagedAttentionExtension numerical execution is covered by the new unit test,
every RoPE implementation was affected identically,
the patch changes inference performance,
or changing translate_get_rows() globally would necessarily fail every llama.cpp model.

The supported conclusion is narrower:

When SDPAToPagedAttention moved active tokens into the leading dimension, the GGUF embedding path restored a fixed leading singleton and displaced that token axis. Moving the singleton to axis 1 inside AdaptToGenAI preserved both the original SDPA layout and the PagedAttention layout.

The complete boundary

GGUF GET_ROWS translation
        ↓
Gather produces backend-dependent leading dimensions
        ↓
old Unsqueeze(axis=0)
        ↓
leading dimension forced to 1
        ↓
SDPAToPagedAttention expects token rows in front
        ↓
embedding-derived Q/K/V and RoPE position state disagree
        ↓
wrong activations or invalid PagedAttention shapes

The repaired path is:

Gather produces:
[1, T, H] under SDPA
or
[T, 1, H] under PagedAttention
        ↓
Unsqueeze(axis=1)
        ↓
[1, 1, T, H]
or
[T, 1, 1, H]
        ↓
the existing leading dimension retains its backend meaning

The contract that remains

An inserted singleton dimension can look semantically empty.

It still has a position.

That position determines which existing dimensions are allowed to remain visible to later consumers.

The old graph treated:

axis 0

as a permanent batch placeholder.

PagedAttention changed the meaning of axis 0 from:

batch

to:

active token row

The repair did not hard-code a PagedAttention-only tensor.

It moved the placeholder out of the dimension whose meaning could change.

When one graph must support two layouts, fixed dimensions should occupy fixed roles—and dynamic backend meaning must be allowed to flow through the axes that already carry it.

That is the embedding-axis contract repaired by OpenVINO #37606.

Part 2: why `inp_out_ids` needed global row indices

The second and final article examines the output-row selection path.

The old graph used a batched gather based on this assumption:

the leading residual dimension is always 1

Once PagedAttention made that dimension equal to the token count, the same squeeze-and-gather pattern no longer represented both layouts.

Part 2 follows the repair:

activation
→ flatten to [rows, hidden]

indices
→ flatten to [rows, 1]

global row index
→ batch_index × seq_length + last_local_index

plain axis-0 Gather

It also explains why:

ordinary SDPA
→ one final prompt row survives

while:

PagedAttention
→ every already-flattened token row passes through

without maintaining two separate graph variants.

Search This Blog

Resonetic Lab | AI Systems Code Review