fix(onnx-to-hip): take multi-dynamic expand_shape extents from the shape operand (-64% TTFT on Gemma-4 26B-A4B)

[BoarQing]

Summary

buildExpandShapeOutputShape derived every dynamic output dimension of a tensor.expand_shape from the source dimension it came from. That is correct only while a reassociation group covers at most one dynamic dimension. When a Reshape splits one dynamic source dimension into several dynamic output dimensions, the source extent is their product, so each of those output dimensions was assigned the whole product.

This makes the split silently wrong for any [bs*ss, H] -> [bs, ss, H] reshape, and every consumer inherits the inflated shape. On Gemma-4 26B-A4B it cost 78% of decoder prefill and made 4K prompts unrunnable.

Related issue or design

None.

Why

The per-layer Reshape([bs*ss, 2816] -> [bs, ss, 2816]) in Gemma-4 26B-A4B lowered to

%dim_289 = tensor.dim %2525, %c0 : tensor<?x2816xf16>
%dim_290 = tensor.dim %2525, %c0 : tensor<?x2816xf16>   // same dim, twice
%expanded_291 = tensor.expand_shape %2525 [[0, 1], [2]]
    output_shape [%dim_289, %dim_290, 2816] : tensor<?x2816xf16> into tensor<?x?x2816xf16>

With bs = 1 the source dimension already equals ss, so the expansion claimed [ss, ss, 2816] rather than [1, ss, 2816]. The 30 layer input norms consuming it ran over ss^2 rows instead of ss.

It stayed silent because the norm runtime ABI passes element counts, not shapes: simplified_layer_norm recovers num_rows as input_num_elements / scale_num_elements, which divides evenly for any ss, so an ss-times-too-large buffer looks like a well-formed larger batch. At 2,049 tokens that was 5,036 ms of the 6,486 ms decoder prefill; at 4,097 tokens the N-squared tensor needed 88 GB against 63.6 GB of shared APU memory, and the host rebooted.

ReshapeShapeFold already rewrites Reshape(_, Shape(x)) into a host-visible tensor.from_elements carrying the correct per-dimension values, and its comment said this was so the conversion could consume them — but the consuming branch was never written and the conversion never read operand 1. The fix is to finish that path rather than add a new mechanism.

What

• ReshapeConversion.cpp: buildExpandShapeOutputShape now detects groups with more than one dynamic output dimension and takes those extents from the shape operand via resolveShapeOperandExtents, which walks tensor.from_elements through arith.index_cast and constants. Groups with zero or one dynamic dimension keep the existing tensor.dim derivation, so the other expand-shape sites in the model are bit-for-bit unchanged.
• The helper returns std::nullopt when a multi-dynamic group's extents are not host-readable, and all three call sites (Reshape expand, same-rank split/combine, Unsqueeze) fall back to tensor.reshape instead of emitting an expand_shape whose output_shape cannot be justified.
• NormConversion.cpp: verifyNormAxisMatchesScale rejects at conversion time any SimplifiedLayerNormalization or RMSNormalization whose axis does not span exactly the scale's element count, when the dimensions involved are static. This is the layer that can still see shapes.
• simplified_layer_norm.cpp, skip_simplified_layer_norm.cpp: add the scale_num_elements > 0 and even-divisibility preconditions that wrap_layer_normalization already had, so a mismatch fails instead of truncating.
• ReshapeShapeFold.cpp: correct the stale comment that claimed the operand was already consumed.

A same-shape guard is deliberately not added, since [ss, ss, 2816] is a legal shape and only the value is wrong.

Test plan

• test/lit/Conversion/onnx-to-hip/test_reshape.mlir: test_reshape_expand_multi_dyn pins the [?x2816, 2816] -> [?, ?, 2816] split to distinct extents from the shape operand; test_reshape_expand_multi_dyn_opaque pins the tensor.reshape fallback when the operand is opaque.
• test/lit/Conversion/hip-to-llvm/test_rms_norm.mlir: test_rms_norm_dynamic_batch_seq_3d pins input_num_elements for a [?, ?, 2816] norm to the product of three distinct dimensions.
• 362 lit tests pass. 21 numeric layer-norm tests pass on gfx1151.
• Correctness on Gemma-4 26B-A4B FP16W4-qmoe: generated text at 513 and 2,049 tokens matches the pre-fix build, and HIPDNN_EP_DEBUG confirms all 181 rank-3 norms report 2049x2816 with no remaining 4198401.

Performance

Gemma-4 26B-A4B FP16W4-qmoe, gfx1151 (Strix Halo, 63.6 GB shared), Release build, 2,049-token prompt, median of 5 after warmup.

Metric	Before	After
TTFT	7,865 ms	2,859 ms
Decoder prefill	6,486 ms	1,520 ms
Layernorm within prefill	5,036 ms	50 ms
Norm rows dispatched	4,198,401	2,049

TTFT at other lengths: 1,905 ms at 513, 2,183 ms at 1,025, and 4,983 ms at 4,097. Prefill is now near-linear in prompt length; 4,097 tokens previously exhausted shared memory and rebooted the host.

Notes for reviewers

• The vision encoder path is unaffected; its 1,446 ms is now the largest fixed component of VLM TTFT.
• Worth confirming the fallback direction is the one you want: an unreadable shape operand on a multi-dynamic group now yields tensor.reshape, which is correct but blocks fusion. I found no such case in the models tested, so the branch is exercised only by the lit test.

Checklist

• The change is focused, or links a design/series explaining its scope.
• Relevant tests were added or updated and the results are documented.
• User-facing or design documentation was updated when needed.
• Substantial AI assistance is disclosed, and I reviewed and understand the result.

AI assistance: the diagnosis (IR dumps, roofline check confirming the kernel was at 91% of achievable bandwidth and so not the problem), the patch, the tests, and the benchmark runs were produced with Cursor. All measurements above are from real runs on the machine named, and I reviewed the change.

[github-actions]

Thanks for opening a PR!

This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.

Before requesting review, please check that:

1. The change is focused. Substantial work links the relevant issue
   or design discussion.
2. The PR documents relevant test results and updates affected
   documentation.
3. If AI tools provided substantial assistance, the description
   explains what was assisted and how it was validated, and commit
   trailers identify the tool. The contributor has reviewed and
   understands the result.

Reviewers are assigned through
CODEOWNERS where ownership
is configured.

[github-actions]

L2 Accuracy Results (EP vs CPU)

Model	Combined L2	Total Elems	Skipped NaN/Inf
conv_test_hybrid	4.8668E-07	64	0
GroupQueryAttention_seq256	25.2366	2621440	0
MatMulNBits_o_seq128	259.906	368640	0
QMoE_seq128	34.957	368640	0

Threshold: 0.01 | Run: 4598 - Commit: 70b9e6b

[github-actions]

MorphiZen EP Performance Results

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	7.46	6.19	361	3	1141
GroupQueryAttention_seq128	4393.65	1.66363	10	6	209
matmul_down_seq128	516.74	2.32	74	3	255

EPContext Export Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	7.52	46.06	357	3	15485

EPContext Import Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	7.50	9.36	365	3	15659

OGA Benchmark Results

Model	Warmup	Reps	Prompt Len	Gen Tokens	TTFT (ms)	TPS	Peak Mem (GB)	GPU Mem (GB)
gpt-oss-20b-webgpu-int4-rtn-block-32	1	5	128	128	167.6	79.0	1.23	13.53
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	1	5	128	128	212.2	40.6	1.12	6.48

OGA Wheel Smoke (Python benchmark_e2e.py)

Model	TTFT (ms)	TPS
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	134	39.5

Run: 4598 - Commit: 70b9e6b

[BoarQing]

perf: http://xsjnts.xilinx.com/ROCm/compare?base_url=genaihttp%3A%2F%2Fxsjvitisaijenkins%3A8080%2Fjob%2Fdaily_hip_ep_genai%2F161%2F&base_job=daily_hip_ep_genai&base=161&curr_url=http%3A%2F%2Fxsjvitisaijenkins%3A8080%2Fjob%2Fdaily_hip_ep_genai%2F163&curr_job=daily_hip_ep_genai&curr=163(Gemma)
http://xsjvitisaijenkins:8080/job/daily_hip_ep_genai/244/ (General)

acc: http://xsjvitisaijenkins:8080/job/daily_hip_ep_genai/241/(Gemma)
http://xsjvitisaijenkins:8080/job/daily_hip_ep_genai/242/(General)