Skip to content

Repository navigation

fix(optimize-memrefs): Keep host-accessed buffers out of slot reuse - #1305

Draft
ximzhang-amd wants to merge 3 commits into
mainfrom
fix/optimize-memrefs-host-buffers
Draft

fix(optimize-memrefs): Keep host-accessed buffers out of slot reuse#1305
ximzhang-amd wants to merge 3 commits into
mainfrom
fix/optimize-memrefs-host-buffers

Conversation

@ximzhang-amd

@ximzhang-amd ximzhang-amd commented Oct 23, 2026

Copy link
Copy Markdown
Contributor

Summary

hip-optimize-memrefs no longer lets a buffer the host loads or stores (directly or through a view) take or
donate a reuse slot.

Related issue or design

Related: #1272, which avoids the same race by pass order; this PR is complementary (see "Relation to #1272"). Backport of the hip-ir-runtime downstream patch 0030 (link under Notes for reviewers).

Why

Program order is not execution order across the host/device boundary: a host memref.store takes effect at once,
while the hip op before it has only been enqueued. The pass let two host-fed scalars share a slot, so the second
store could overwrite the first before the kernel reading it ran. In a GQA decoder the total-sequence-length
staging buffer was reused for the next scalar; decode steps after a prefill nondeterministically read the step
length, repeating tokens.

What

  • Exclude host-accessed buffers (host memref.load/memref.store, also through views) from slot reuse in
    OptimizeMemRefs.cpp; they are tiny.
  • Lit tests in hip-optimize-memrefs.mlir: two host-fed scalars (the GQA case), and two 32-element host-fed
    buffers, too large for MaterializeHostScalars; note in docs/design/pool-allocs-memory-planning.md.

Relation to #1272

this PR #1272 (head 9f3d537)
Where OptimizeMemRefs: host-accessed buffers neither take nor donate a slot pipeline order: MaterializeHostScalars before OptimizeMemRefs
Covers any buffer the host loads or stores, directly or through a view static host-fed buffers of at most 16 elements (what MaterializeHostScalars redirects)
Scope 3 files, +52/-1 (+ a test commit) 11 files, +481/-94 (also other dynamic-shape memory fixes)
State ready open, in review

Measured with hip-mlir-opt (Test plan): host store, GPU read, host store, GPU read.

Pass setup 1-element buffers 32-element buffers
main slot reused (race) slot reused (race)
#1272's order kept apart slot reused (race)
this PR kept apart kept apart
this PR + #1272's order kept apart kept apart

The race actually seen (the GQA total-length scalar) is a 1-element buffer, which #1272 also fixes. No model
with a larger host-fed buffer has been seen yet; the second test keeps the pass correct by construction rather
than by pass order. If both land, the scalar case is covered twice, which is harmless.

Test plan

Build environment: Linux, no GPU; base origin/main 93c53c4, LLVM from rocmlirTriton 85f34b1a, TheRock 7.13.0, build.py --hip_arch gfx1151 (BUILD_EP=OFF, BUILD_HIP_UNIT_TESTS=ON). Baseline main: lit 624 pass / 9 unsupported / 0 fail; CTest 134/266 pass, the 132 failures are GPU execution tests (131 E2E_Execute_* + HipGpuElfContract) that cannot run without a GPU.

  • Builds; lit 624 pass / 0 fail (both new cases included); CTest identical to main.
  • The new cases fail with main's hip-optimize-memrefs (%b merged into %a) and pass with this branch's.
  • Relation to fix(compiler): correct memory handling for dynamic shapes #1272, with main's passes in fix(compiler): correct memory handling for dynamic shapes #1272's order (hip-materialize-host-scalars then
    hip-optimize-memrefs): the 1-element case moves to host-mapped scratch (offsets 0 and 64) and is fixed;
    the 32-element case stays in the pool and %b still takes %a's slot. With this branch's pass, both are
    kept apart in either order.
  • Not run here: decode-after-prefill determinism on a full GQA decoder model (needs a gfx1151 model run; the
    race is timing-dependent). The lit cases above check the slot assignment that causes it.

Notes for reviewers

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 (docs/design/pool-allocs-memory-planning.md).
  • Substantial AI assistance is disclosed, and I reviewed and understand the result.

🤖 Generated with Claude Code

@github-actions

github-actions Bot commented Oct 23, 2026

Copy link
Copy Markdown

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

github-actions Bot commented Oct 23, 2026 •
edited
Loading

Copy link
Copy Markdown

L2 Accuracy Results (EP vs CPU)

Model Combined L2 Total Elems Skipped NaN/Inf
conv_test_hybrid 1.12256E-07 64 0
GroupQueryAttention_seq256 59.2561 2621440 0
MatMulNBits_o_seq128 259.905 368640 0
QMoE_seq128 24.2583 368640 0

Threshold: 0.01 | Run: 7997 - Commit: 2199d12

@github-actions

github-actions Bot commented Oct 23, 2026 •
edited
Loading

Copy link
Copy Markdown

Binary sizes

Baseline: run 40617608668 - cf8ff09 (dll-sizes)

File This (MB) Main (MB) Delta (MB)
hipgpu.dll 59.34 59.34 0.00
custom_kernels_gfx1170.dll 5.55 5.55 0.00
custom_kernels_gfx11-generic.dll 5.79 5.79 0.00

MorphiZen EP Performance Results

Model QPS Session (s) 1st Infer (ms) CPU% Mem (MB)
full_model_seq128 8.73 7.55 266 3 1158
GroupQueryAttention_seq128 11354.50 2.55577 1 5 229
matmul_down_seq128 512.53 3.29 127 3 273

EPContext Export Performance

Model QPS Session (s) 1st Infer (ms) CPU% Mem (MB)
full_model_seq128 8.83 48.28 250 3 15505

EPContext Import Performance

Model QPS Session (s) 1st Infer (ms) CPU% Mem (MB)
full_model_seq128 8.83 9.81 245 3 15639

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 93.3 74.2 1.25 12.24
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml 1 5 128 128 120.5 47.7 1.14 5.23

OGA Wheel Smoke (Python benchmark_e2e.py)

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

VLM Benchmark Results

Model TTFT (ms) TPS Output
gemma-4-E4B-it-fp16-int4-kquant-prune-logits 4053.5 7.1 This is a high-quality, outdoor photograph featuring a beautiful, fluffy white dog, likely a West Highland White Terrier or a si
Qwen3.8-27B-fp16-ve-fp16-int4-k_quant-gs128-text-dml-onnx 4951.7 14.0 The user wants a detailed description of the image. Let me carefully observe all elements before writing. Subject: A white do

TTFT / TPS vs main

Baseline: run 40617608668 - cf8ff09 | Gate: TTFT +20% / TPS -20%

Suite Model TTFT (ms) Main Delta TPS Main Delta
oga gpt-oss-20b-webgpu-int4-rtn-block-32 93.3 92.4 +0.9% 74.2 75.0 -1.0%
oga Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml 120.5 119.4 +0.9% 47.7 47.7 -0.2%
vlm gemma-4-E4B-it-fp16-int4-kquant-prune-logits 4053.5 4096.2 -1.0% 7.1 7.3 -2.9%
vlm Qwen3.8-27B-fp16-ve-fp16-int4-k_quant-gs128-text-dml-onnx 4951.7 4928.4 +0.5% 14.0 14.0 -0.1%

Run: 7997 - Commit: 2199d12

@ximzhang-amd
ximzhang-amd marked this pull request as ready for review October 23, 2026 09:16
@amd-mingw

amd-mingw commented Oct 27, 2026

Copy link
Copy Markdown
Collaborator

/run-test models=all node=xsjnucstrhalo09

@amd-yanjunz

amd-yanjunz commented Oct 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

ROCm GenAI CI finished with result: FAILURE.

xingy16-amd and others added 3 commits October 27, 2026 20:14
@xingy16-amd @claude @ximzhang-amd
fix(optimize-memrefs): Keep host-accessed buffers out of slot reuse
4b43a82 
Program order is not execution order across the host/device boundary: a
host memref.store takes effect at once, while the hip op before it has
only been enqueued. hip-optimize-memrefs let two host-fed scalars share a
slot, so the second store could overwrite the first before the kernel
reading it ran. In a GQA decoder the total-sequence-length staging buffer
was reused for the next scalar; decode steps after a prefill
nondeterministically read the step length, repeating tokens.

Buffers the host loads or stores, directly or through a view, now neither
take nor donate a slot. They are tiny, so nothing is lost.

Co-Authored-By: Claude <noreply@anthropic.com>
Made-with: Claude Code
@ximzhang-amd @claude
test(optimize-memrefs): keep host-fed buffers too large to materializ…
cef8313 
…e out of reuse

hip-materialize-host-scalars moves only static host-fed buffers of at most
16 elements to host-mapped scratch, so a larger one stays in the pool and
reaches slot reuse whichever of the two passes runs first. Check that
hip-optimize-memrefs keeps two 32-element host-fed buffers apart.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Made-with: Claude Code
@ximzhang-amd @claude
style(optimize-memrefs): clang-format the host-access check
2199d12 
Whitespace only, as the lintrunner pre-commit hook formats it.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Made-with: Claude Code
@ximzhang-amd
ximzhang-amd force-pushed the fix/optimize-memrefs-host-buffers branch from 840273f to 2199d12 Compare October 28, 2026 02:15
@ximzhang-amd

ximzhang-amd commented Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

/run-test models=all node=xsjnucstrhalo09

@amd-yanjunz

amd-yanjunz commented Oct 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

ROCm GenAI CI finished with result: UNSTABLE.

fhanuman
fhanuman requested changes Oct 28, 2026
View reviewed changes
Comment thread lib/Dialect/Transforms/OptimizeMemRefs.cpp

Value result = allocOp.getResult();
if (result.use_empty())
if (result.use_empty() || isHostAccessed(result))

@fhanuman fhanuman Oct 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We shouldn't be doing this. This pass intended to be a generic pass, why do this pass need to know whether it's a host or device?

@ximzhang-amd ximzhang-amd Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR upstreams the fix from https://github.com/AMD-AIG-AI-SOLUTIONS/hip-ir-runtime/blob/1835dd33d2e6247c4fa0672e8d1d217421dcac75/hip-ep-ext/patches/0030-fix-optimize-memrefs-Keep-host-accessed-buffers-out-.patch. #1272 makes a similar change by running MaterializeHostScalars before OptimizeMemRefs, which keeps this pass generic.

I checked that #1272 covers it. The test model has two host-fed shape scalars, each read by a cast queued behind slow matmuls. On main, the first output comes back with the second scalar's value in 20/20 runs. With #1272's reorder alone, all 20 runs are correct, the same as with this PR alone. I also checked a Llama 3.1 8B decoder: with #1272's order, the compiled IR is identical with and without this PR.

Once #1272 is merged, this PR isn't needed, and I'll close it then.

@fhanuman fhanuman Oct 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to clarify the changes in the repo that you mentioned is POC and a lot of it is generated by AI. We need to be careful what we upstream. I have been reviewing the code to see what can be upstreamed.

@ximzhang-amd
ximzhang-amd marked this pull request as draft October 28, 2026 09:03
@fhanuman fhanuman mentioned this pull request Oct 30, 2026
Closed

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@fhanuman fhanuman fhanuman requested changes
@amd-mingw amd-mingw Awaiting requested review from amd-mingw amd-mingw is a code owner
@edelaye edelaye Awaiting requested review from edelaye edelaye is a code owner
@wcy123 wcy123 Awaiting requested review from wcy123 wcy123 is a code owner
@qianglin-amd qianglin-amd Awaiting requested review from qianglin-amd qianglin-amd is a code owner
@amd-bqiao amd-bqiao Awaiting requested review from amd-bqiao amd-bqiao is a code owner
@asparkhi asparkhi Awaiting requested review from asparkhi asparkhi is a code owner
@sambanerxilinxamd sambanerxilinxamd Awaiting requested review from sambanerxilinxamd sambanerxilinxamd is a code owner
@CaymanYang CaymanYang Awaiting requested review from CaymanYang CaymanYang is a code owner

Requested changes must be addressed to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@ximzhang-amd @amd-mingw @amd-yanjunz @fhanuman @xingy16-amd