Repository navigation
fix(gqa): run a short first prompt on the hipBLASLt path as well - #1327
Draft
ximzhang-amd wants to merge 6 commits into
Draft
ximzhang-amd wants to merge 6 commits into
ximzhang-amd wants to merge 6 commits into
Conversation
|
Thanks for opening a PR! This project follows LLVM's incremental-development and AI-tool-use Before requesting review, please check that:
Reviewers are assigned through |
L2 Accuracy Results (EP vs CPU)
Threshold: 0.01 | Run: 7995 - Commit: |
Binary sizesBaseline: run 40617608668 -
MorphiZen EP Performance Results
EPContext Export Performance
EPContext Import Performance
OGA Benchmark Results
OGA Wheel Smoke (Python benchmark_e2e.py)
VLM Benchmark Results
TTFT / TPS vs mainBaseline: run 40617608668 -
Run: 7995 - Commit: |
Kernel Unit Test Results
Leaves: 16 (0 failed) | Cases: 289 (0 failed) Mode: lookup | Coverage: 1 | Run: 7995 - Commit: |
ximzhang-amd
force-pushed
the
fix/gqa-short-first-prompt-hipblaslt
branch
from
October 28, 2026 01:59
70e9cf1 to
186551b
Compare
Contributor
Author
|
/run-test models=all node=xsjnucstrhalo09 |
Contributor
|
ROCm GenAI CI has been triggered.
|
seqlens_k is the total length less one. The fused prefill derived the past as seqlens_k + 1 - sq and rejected a negative one, so a prompt shorter than the window -- tokens in rows 0..seqlens_k, padding after, no past -- failed in every layer. The wrapper's return code is dropped by the generated code, so the model went on to return garbage. That is the padded first prompt ORT reads with a zero past. The host now clamps the past at 0, and attends over seqlens_k + 1 keys; the prefill kernels already bound keys by that count, so padding rows attend to the prompt and nothing past it. The rope and KV-append kernels derive the past on the device the same way through gqa_past_len_from_seqlens: the append used to skip a negative past, leaving the cache unwritten, and rope would have indexed the cos/sin cache below 0. Full windows, later chunks and decode compute the same past as before. Checked on gfx1151 against the full window on the same input: a 5-token prompt gives bit-identical hidden rows 0..4 and KV positions 0..4 on phi-muffin gemm and conv. Replacing the padding rows' input changes no output bit. On orca 3.4.0 (int8 KV cache) rows 0..4 match to 2.5e-3 relative L2, with the first int8 flips of one unit nine or more layers in. Both lengths are deterministic run to run. Decode latency is unchanged. Co-Authored-By: Claude <noreply@anthropic.com> Made-with: Claude Code (cherry picked from commit 9caa18d)
test_gqa_padded_chunk_prefill_llama_shape compared only the attention output of a prompt shorter than its prefill window. The KV-cache append and RoPE kernels derive the past from seqlens_k on the device too, and an unclamped negative past made the append kernels skip the batch, leaving the cache unwritten while the output could still look right. Also compare present_key / present_value slots [0, real_count) with the CPU reference of the unpadded prompt, and require them to be non-zero. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Made-with: Claude Code
seqlens_k cannot say where a window's padding is or whether a past precedes it: 10 tokens after 40 cached ones, padded to a 128-row window, arrive exactly as a 50-token first prompt. Reading every padded window as a first prompt therefore computed such a step over the cache, writing its rows from slot 0 over the cached tokens, where the call used to be refused. A padded window is now read as a first prompt (tokens first, padding after, no past) only when total_sequence_length equals the window, which is how ORT marks a first prompt and what ORT CPU accepts. Any other padded window is refused again, before RoPE and the KV-cache write. The runtime used to ignore the scalar; it is read only for a padded prefill window. Layouts that carry the same scalars - a first prompt padded before its tokens, or a padded continuation reporting the window as its total - are not told apart and are not supported. Fixed-cache exports whose total is the cache size are refused too; their short prompts need the caller to pass the window. Tests: - the padded fixed-cache test passes total = the window, and gains a full-window case; - a padded first prompt with RoPE inside GQA (1, 22 and 128 tokens) and nonzero padding rows, compared with the same tokens unpadded on ORT CPU and with ORT CPU on the padded graph; - a padded prefill followed by a decode step on caches bound in place (past and present the same GPU buffers), compared with ORT CPU; - three padded windows that must be refused (a step after a past, total = the token count, total = the cache size), checked by the refusal diagnostic of that run and by in-place caches that come back byte for byte unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Made-with: Claude Code
The decomposed hipBLASLt pipeline derived the past the same way the fused one did and rejected a prompt shorter than the window. It is the path for head sizes the fused kernels do not template -- 96 on phi-silica -- so such a model still failed every layer on a short prompt. It now clamps the past at 0 like gqa_forward_fused. The causal mask takes past_len explicitly and the scores span total_seq keys, so padding rows attend to the prompt alone; the per-chunk key narrowing, which assumes total_seq == past_len + sq, already switches itself off here. On gfx1151, phi-silica (head_dim 96, 64-row window): a 5-token prompt gives bit-identical hidden rows and KV positions 0..4 against the full window. Co-Authored-By: Claude <noreply@anthropic.com> Made-with: Claude Code
The decomposed hipBLASLt pipeline had the same short-first-prompt bug as the fused prefill. Head size 96 no longer reaches it (the fused kernels take 64, 96, 128 and 256 now); float32 activations with an unquantized cache still do, because they are not fused_supported. Parametrize test_gqa_padded_chunk_prefill_llama_shape over float16 (fused) and float32 (decomposed hipBLASLt), adding an element-type argument to _make_gqa_fixed_cache_model. The CPU reference (onnxruntime 1.31.0) accepts both types for the unpadded model. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Made-with: Claude Code
… on hipBLASLt too The decomposed hipBLASLt pipeline clamped every negative past to 0, so a padded step after a past was computed from cache slot 0 over the cached tokens, as the fused path did before. It now takes the fused path's rule (gqa_padded_window_is_first_prompt): a padded window is a first prompt only when total_sequence_length equals the window, and is refused otherwise, before RoPE and the KV-cache write. The padded-window tests run on both paths: fp16 takes the fused kernels, fp32 the hipBLASLt pipeline. - the RoPE case against the unpadded ORT CPU result; - the in-place prefill-then-decode case; - the three refused windows, each expecting its path's diagnostic. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Made-with: Claude Code
ximzhang-amd
force-pushed
the
fix/gqa-short-first-prompt-hipblaslt
branch
from
October 28, 2026 02:51
186551b to
5683117
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The decomposed hipBLASLt GQA pipeline now runs a padded first prompt under the same rule as the fused path after
#1301: tokens first, padding after, no past, accepted only when
total_sequence_lengthequals the window. Anyother padded window is rejected.
Related issue or design
Depends on #1301 (the same fix on the fused path). Backport of the hip-ir-runtime downstream patch 0008 (link under Notes for reviewers).
Why
gqa_forward_hipblasltderived the past the same way the fused prefill did and rejected a prompt shorter than thewindow. The original motivation was head size 96 (phi-silica), which the fused kernels did not template at the
time. On current
mainhead_dim 96 runs fused (head_dim_okaccepts 64, 96, 128, 256), but the decomposedpath is still taken whenever the fused path is not supported: float32 activations with an unquantized cache,
wave64 devices (CDNA), an attention bias, non-causal attention, reused past KV, or an unsupported window/sink.
What
gqa_forward_hipblaslt: a negative past is clamped to 0 only whengqa_padded_window_is_first_prompt()(fromfix(gqa): run a first prompt shorter than its prefill window #1301) accepts the window, i.e.
total_sequence_length == sq. Otherwise the call is rejected before RoPE and theKV-cache write, with the same diagnostic as the fused path. The scalar is passed through from
wrap_group_query_attention. The causal mask takespast_lenexplicitly and the scores spantotal_seqkeys,so padding rows attend to the prompt alone.
Tests: every padded-window test from fix(gqa): run a first prompt shorter than its prefill window #1301 is parametrized over float16 (
fp16_fused) and float32(
fp32_hipblaslt), so both paths are covered the same way:test_gqa_padded_chunk_prefill_llama_shape;test_gqa_padded_first_prompt_rotary(RoPE inside GQA, nonzero padding, compared with unpadded ORT CPU);test_gqa_padded_prefill_then_decode_in_place;test_gqa_padded_window_is_rejected, which expects each path's own diagnostic._make_gqa_fixed_cache_modelgains an element-type argument.Test plan
Validation by the original author (gfx1151, from the commit message): phi-silica (head_dim 96, 64-row window),
5-token prompt bit-identical to the full window (at that time head_dim 96 used this path).
Builds (gfx1151 EP package from this branch); clang-format, ruff and
git diff --checkclean.Routing checked in
wrap_group_query_attention: float32 + unquantized cache →!fused_supported→gqa_forward_hipblaslt.HIPDNN_EP_DEBUG=1on an fp32 case showsrouting to legacy decomposed pipeline (elem=4 … d=128 sq=128 …).Ordering: the only RoPE / KV work in
gqa_forward_hipblasltbefore the check is thesq == 1decodefallback, which cannot be padded.
GPU (gfx1151, Strix Halo, TheRock 7.13, onnxruntime 1.31.0,
HIPDNN_EP_STRICT=1), padded-window tests(24 = 12 per path):
All of
test_gqa.pywith this branch's EP: 32 pass / 10 fail. The 10 aremain's pre-existinggrowing-KV-cache failures, where the EP returns
presentsized likepast; no new failure.Notes for reviewers
0008-fix-gqa-run-a-short-first-prompt-on-the-hipBLASLt-pa.patch
Co-Authored-By: Claude); the test commit and thisdescription were prepared with Claude Code and reviewed by the submitter.
Checklist
🤖 Generated with Claude Code