ONNX Runtime #22536 — Why Same-Rank MatMuls Reused the Wrong WebGPU Broadcast Shader

Rank-only cache keys, shape-dependent WGSL source, and the uniform-driven indexing rewrite that made shader reuse valid

Two MatMul calls had the same input ranks.

They had the same matrix dimensions.

They produced the same output shape.

Their batch broadcasting was different.

First call

A:

[1, 2, 2, 2]

B:

[2, 1, 2, 1]

Second call

A:

[2, 1, 2, 2]

B:

[1, 2, 2, 1]

Both inputs were rank four.

Both MatMuls used:

M = 2

K = 2

N = 1

Both produced:

[2, 2, 2, 1]

But the singleton batch dimensions had traded places.

First call:

A broadcasts batch axis 0

B broadcasts batch axis 1

Second call:

A broadcasts batch axis 1

B broadcasts batch axis 0

Before ONNX Runtime PR #22536, the WebGPU MatMul generator calculated those broadcast axes in TypeScript while creating the shader source.

It then embedded the result into WGSL indexing code.

Conceptually, one shader said:

A batch index 0 must always be zero.

B batch index 1 must always be zero.

The other shader needed to say the opposite:

A batch index 1 must always be zero.

B batch index 0 must always be zero.

The shader source was different.

The shader cache key was not.

The MatMul cache declared that its source depended on the input types and ranks, but not on the full dimensions or the resulting broadcast pattern.

The first compiled shader could therefore be reused for the second MatMul.

That cache hit reused indexing rules written for another tensor layout.

PR #22536 repaired the mismatch without expanding the cache key.

Instead, it removed the broadcast pattern from the compiled shader variant.

The shader now reads the input shapes from uniforms and decides at execution time:

input dimension is 1
→ use batch index 0

input dimension is not 1
→ use the corresponding output batch index

One rank-compatible shader can now handle both broadcasting patterns correctly. The PR was merged on November 8, 2024, as commit d9b91682.

Two calls belong to one shader-cache entry only when every difference omitted from the key has also been removed from the generated source.

---

The two tests used the same matrices in a different batch arrangement

The smallest regression pair used the same underlying values.

Matrix A contained:

0, 1, 2, 3, 4, 5, 6, 7

Those values formed two 2×2 matrices:

A0 =

[0, 1]
[2, 3]

A1 =

[4, 5]
[6, 7]

Matrix B contained:

8, 9, 10, 11

Those values formed two 2×1 matrices:

B0 =

[8]
[9]

B1 =

[10]
[11]

The four possible products were:

A0 × B0

=

[9]
[43]

A0 × B1

=

[11]
[53]

A1 × B0

=

[77]
[111]

A1 × B1

=

[95]
[137]

The arithmetic did not change between the two test cases.

Only the batch-coordinate mapping changed.

---

Case 0: A varies on the second batch axis

The first shape pair was:

A batch shape:

[1, 2]

B batch shape:

[2, 1]

A has one entry on the first batch axis, so it broadcasts there.

B has one entry on the second batch axis, so it broadcasts there.

For the output batch shape:

[2, 2]

the mapping is:

Output batch	A matrix	B matrix	Result
[0, 0]	A0	B0	[9, 43]
[0, 1]	A1	B0	[77, 111]
[1, 0]	A0	B1	[11, 53]
[1, 1]	A1	B1	[95, 137]

The expected flattened output is:

[
    9, 43,
    77, 111,
    11, 53,
    95, 137
]

---

Case 1: the singleton dimensions exchange positions

The second shape pair was:

A batch shape:

[2, 1]

B batch shape:

[1, 2]

Now A varies on the first batch axis.

B varies on the second.

The output shape is still:

[2, 2]

but its batch mapping becomes:

Output batch	A matrix	B matrix	Result
[0, 0]	A0	B0	[9, 43]
[0, 1]	A0	B1	[11, 53]
[1, 0]	A1	B0	[77, 111]
[1, 1]	A1	B1	[95, 137]

The expected flattened output is therefore:

[
    9, 43,
    11, 53,
    77, 111,
    95, 137
]

The two expected vectors contain the same four matrix products.

Their middle groups appear in a different order.

That difference is exactly what a shape-insensitive shader-cache hit could lose. The checked-in suite retains both small cases and two larger equivalents with the same reversed broadcasting pattern.

---

What the old shader generator hard-coded

MatMul applies matrix multiplication to the final two dimensions.

All earlier dimensions form the batch shape.

For an input such as:

[1, 2, 2, 2]

the shader must convert an output batch coordinate into an input batch coordinate.

The general rule is:

If an input batch dimension is 1,
that input coordinate must be 0.

Otherwise,
use the corresponding output coordinate.

The old generator did not evaluate that condition inside the shader.

It evaluated the shapes while generating the source.

Conceptually, TypeScript performed:

broadcast_dims_A =
getBroadcastDims(
    A batch shape,
    output batch shape
)

and:

broadcast_dims_B =
getBroadcastDims(
    B batch shape,
    output batch shape
)

It first emitted assignments copying output batch coordinates into the input indices.

It then emitted extra lines setting each known broadcast dimension to zero.

For the first shape pair, the resulting source was equivalent to:

a_indices[0] = 0u;
a_indices[1] = batch_indices[1];

b_indices[0] = batch_indices[0];
b_indices[1] = 0u;

For the second shape pair, the source needed to be equivalent to:

a_indices[0] = batch_indices[0];
a_indices[1] = 0u;

b_indices[0] = 0u;
b_indices[1] = batch_indices[1];

These are different programs.

They read different batch matrices for the same output coordinate.

---

The cache declared only a rank dependency

ONNX Runtime’s WebGPU program metadata distinguishes several levels of input dependency.

The relevant definitions are:

type

rank

dims

data

A rank dependency means the shader key incorporates the input data type and rank, but not every dimension value.

A dims dependency includes the actual input dimensions.

The MatMul shader declared:

inputDependencies:

['rank', 'rank']

for two-input MatMul.

Its other cache hint fields covered properties such as:

• activation type;

• vector component width;

• values processed per output;

• and layout mode.

They did not encode which batch dimension was broadcast.

For the two regression calls:

A rank:
4 in both cases

B rank:
4 in both cases

M:
2 in both cases

K:
2 in both cases

N:
1 in both cases

output shape:
[2, 2, 2, 1] in both cases

The cache-visible properties matched.

The compile-time broadcast decisions did not.

That violated the cache contract.

Cache says:

same shader variant

Source generator says:

different indexing code

Once the first program had been cached, a later call could receive the earlier shader even though its singleton axes had moved.

The PR description identifies this directly as a wrong cache hit caused by different broadcast-dependent shader code sharing an identical cache key.

---

This was not a hash collision

Two unrelated keys did not accidentally produce the same hash.

The two calls were intentionally assigned the same logical cache identity.

The key-building policy omitted dimensions because the shader metadata claimed ranks were sufficient.

The source generator then used dimensions anyway.

Accidental hash collision:

different complete keys
→ same hash by coincidence

PR #22536 failure:

incomplete equivalence definition
→ same key by design

The second problem cannot be solved by choosing a stronger hash function.

The missing dependency must either be:

added to the key

or:

removed from the compiled variant

PR #22536 chose the second approach.

---

The first proposed repair expanded the key

The PR’s earlier description stated that broadcasting information would be added to the MatMul shader cache key.

That would have made the source and key agree.

Pattern A
→ cache entry A

Pattern B
→ cache entry B

A reviewer pointed out that the complete input shapes were already uploaded to the shader as uniforms.

That created another option:

Do not specialize the source by broadcast pattern.

Read the dimensions already supplied at runtime.

The author revised the implementation accordingly.

This changed the repair from:

make the cache key more specific

to:

make the shader source less specific

Both can restore correctness.

They create different cache behavior.

---

The fixed shader derives broadcasting from uniforms

The patch introduced a shared helper:

convertOutputBatchIndicesToInputBatchIndices

For each input batch dimension, the generated WGSL now performs a runtime condition equivalent to:

if (input_shape[i] != 1u) {
    input_indices[i] =
        output_batch_indices[aligned_i];
} else {
    input_indices[i] = 0u;
}

The input rank remains part of shader structure.

The dimension values do not.

That distinction is important.

Rank affects generated types and loop structure

A rank-four input needs a rank-four index representation.

A rank-five input may require another generated index type or another set of assignments.

It can legitimately require another shader variant.

Dimension values affect runtime indexing

Whether one existing dimension contains:

1

or

2

can be read from the shape uniform.

It does not need separate WGSL source.

The current implementation still uses the runtime shape test and still declares rank-only shader dependencies.

---

Leading batch dimensions are aligned from the right

MatMul inputs do not always have identical ranks.

For example:

A batch shape:

[3, 1, 5]

B batch shape:

[5]

Broadcasting aligns dimensions from the right.

The fixed helper therefore calculates:

extending_input_rank
=
output_batch_rank
-
input_batch_rank

It ignores the extra leading output dimensions when mapping into a lower-rank input.

For each actual input batch dimension, it uses the corresponding right-aligned output coordinate—unless the input dimension is one, in which case it uses zero.

Conceptually:

Output batch coordinates:

[o0, o1, o2]

Lower-rank input batch coordinates:

[o2]

or, when that input dimension is broadcast:

[0]

This behavior is expressed using rank-specialized code and shape-driven runtime decisions.

The cache key therefore retains the part that changes code structure:

rank

while uniforms supply the part that changes only values:

dimension sizes

---

Both MatMul shader paths received the same rule

The WebGPU implementation contained more than one MatMul shader path.

The patch applied the new helper to:

• the naive MatMul shader;

• and the packed MatMul shader.

Before the repair, both paths contained logic that derived broadcast dimensions during source construction.

After the repair, both use the runtime input-shape uniforms to map output batch coordinates to input coordinates.

This matters because fixing only one generator would leave cache correctness dependent on which MatMul implementation was selected.

The shared helper makes the indexing contract explicit in one place:

output coordinate

+

input rank

+

input shape uniform

→ correct input coordinate

---

Why the runtime branch does not change MatMul arithmetic

The patch does not alter the inner multiplication:

A[row, k] × B[k, col]

It changes only which batch slice supplies A and B.

For each output batch coordinate, the shader still computes:

C_batch
=
A_selected_batch
×
B_selected_batch

The question repaired by the patch is:

Which A batch?

Which B batch?

The matrix dimensions and arithmetic loop remain unchanged.

This is why the tests can use small matrices with easily checked products.

The numerical difference comes from selecting another batch pair, not from changing floating-point accumulation.

---

Why one uniform-driven shader is a stronger cache contract

Adding the broadcast pattern to the key would have been correct.

It would also have preserved a source generator where runtime shapes affected compiled text.

The final design gives the cache a simpler equivalence class.

Same ranks

+

same source-specialized properties

→ same shader source

The dimension values can vary without invalidating the compiled pipeline because the shader reads them from uniforms.

This has two architectural benefits.

The cache key matches the source dependency

Nothing omitted from the key changes which broadcast axes are hard-coded.

There are no hard-coded broadcast axes anymore.

Shape variation does not automatically create another shader

Calls with different compatible dimensions can reuse one pipeline when the generated structure remains the same.

The PR does not publish a cache-size, compilation-time, or runtime-performance benchmark.

The supported conclusion is structural:

fewer broadcast-specific source variants are required

—not a measured speedup claim.

---

The regression pair protects cache order

A single broadcast test can prove that one generated shader works for one shape.

It cannot prove that the cache distinguishes or safely unifies two shapes.

The new tests deliberately use pairs with:

same input ranks

same matrix dimensions

same output shape

different singleton batch axes

The small pair expects:

Pattern 0:

[9, 43, 77, 111, 11, 53, 95, 137]

and:

Pattern 1:

[9, 43, 11, 53, 77, 111, 95, 137]

The larger pair repeats the same structure with:

K = 8

and larger products.

These are not generic broadcasting tests.

They target the cache equivalence boundary.

Same cache-visible rank signature

but

different correct batch mapping

The four cases remain present in ONNX Runtime’s current MatMul WebGPU test data.

---

What the merged evidence confirms

The PR, current source, and retained tests confirm that:

• WebGPU MatMul shader generation previously calculated broadcast dimensions from concrete input shapes;

• those dimensions changed the emitted indexing source;

• the shader cache declared rank-only input dependencies;

• two same-rank MatMuls with different broadcast patterns could therefore share one cache key;

• the PR identifies this as producing a wrong cache hit;

• the first proposed direction was to add broadcast information to the key;

• review redirected the implementation toward using existing shape uniforms;

• the final shader evaluates singleton dimensions at runtime;

• both naive and packed MatMul paths use the shared mapping rule;

• paired small and larger regressions cover reversed broadcast axes;

• and the uniform-driven implementation remains in current main.

---

What the evidence does not establish

The public record does not establish that:

• every WebGPU MatMul was affected;

• rank-two MatMul used the same failing batch path;

• every browser or GPU adapter reproduced an identical wrong output;

• the wrong indexing necessarily read outside a GPU allocation;

• a named production model was publicly identified;

• the issue had a security impact;

• the patch changed shader execution time measurably;

• or every WebGPU operator’s cache key is now complete.

The supported conclusion is narrower:

ONNX Runtime’s WebGPU MatMul generator emitted shape-specific broadcast indexing while declaring that rank alone determined the shader variant. The merged fix moved the shape decision into uniforms, making same-rank shader reuse consistent with the code being cached.

---

The complete boundary

MatMul request A
        ↓
broadcast axes derived from concrete shapes
        ↓
shader A hard-codes one set of zero indices
        ↓
shader cached under rank-only key

MatMul request B
        ↓
same ranks and shader hint
        ↓
same cache key
        ↓
shader A reused
        ↓
B requires different broadcast axes
        ↓
wrong input batch selected
        ↓
wrong MatMul output

The repaired path is:

MatMul request
        ↓
cache identifies rank-compatible shader
        ↓
shader reads input shapes from uniforms
        ↓
for each batch axis:
    size 1 → index 0
    otherwise → output batch index
        ↓
same shader handles either broadcast pattern

---

The contract that remains

A cache key is not merely a convenient identifier.

It is a claim of equivalence.

When two requests receive the same key, the system is asserting:

Either request can safely use
the same compiled artifact.

That assertion can be made true in two ways.

Include every source-changing property in the key.

Or:

Move the varying property out of the compiled artifact
and supply it at runtime.

ONNX Runtime chose the second for WebGPU MatMul broadcasting.

A cache may ignore a runtime difference only after the compiled program has stopped hard-coding it.

That is the shader-cache contract repaired by ONNX Runtime #22536.

---

Related material

• ONNX Runtime PR #22536 — Make WebGPU MatMul shaders independent of broadcast patterns

• Merge commit — d9b91682

• Current WebGPU MatMul shader generation

• Current packed WebGPU MatMul path

• WebGPU shader-cache dependency definitions

• Current same-rank broadcast regression cases

---

Patch status: Merged into ONNX Runtime main on November 8, 2024, as d9b91682; the uniform-driven broadcast mapping and paired regressions remain in current main.