fix(test): work around Gaudi3 sdpa_recomp_fwd GQA kernel bug in reference - #1467
Draft
yangulei wants to merge 1 commit into
Draft
fix(test): work around Gaudi3 sdpa_recomp_fwd GQA kernel bug in reference#1467yangulei wants to merge 1 commit into
yangulei wants to merge 1 commit into
Conversation
…ence Signed-off-by: Youlei Yang <youlei.yang@intel.com>
yangulei
requested review from
PatrykWo,
adobrzyn,
afierka-intel,
iboiko-habana,
jbyczkow,
kamil-kaczor,
ksmusz,
mgawarkiewicz-intel,
michalkuligowski and
xuechendi
as code owners
May 20, 2026 06:25
Contributor
There was a problem hiding this comment.
Pull request overview
Works around a Gaudi3 sdpa_recomp_fwd GQA kernel correctness issue that corrupts outputs for KV head group 1+ in certain mask/shape combinations by adjusting the test reference path to compute ground truth per KV-group.
Changes:
- Update BF16 non-sliced reference to run per-KV-group
sdpa_recomp_fwdwhen GQA is detected, then reshape/merge outputs. - Apply the same per-KV-group reference workaround for the FP8 accuracy test’s BF16 ground-truth reference.
Comment on lines
+1078
to
+1081
| """Non-sliced reference: single FusedSDPA call with full mask. | ||
|
|
||
| Uses per-KV-group processing to work around a Gaudi3 kernel bug | ||
| where sdpa_recomp_fwd corrupts output for KV head group 1+ when |
Comment on lines
+1197
to
+1200
| """BF16 ground-truth reference: single FusedSDPA call with full mask. | ||
|
|
||
| Uses per-KV-group processing to work around a Gaudi3 kernel bug | ||
| (same as TestFsdpaSlicingAccuracyBF16._run_reference). |
Comment on lines
+1202
to
+1216
| from habana_frameworks.torch.hpex.kernels.FusedSDPA import is_gqa, gqa_input_reshape_fwd, gqa_output_reshape | ||
| scale = 1.0 / (q.shape[-1]**0.5) | ||
| gqa = is_gqa(q, k) | ||
| with torch.inference_mode(): | ||
| output = FusedSDPA.apply( | ||
| q, | ||
| k, | ||
| v, | ||
| attn_mask, | ||
| 0.0, # dropout_p | ||
| False, # is_causal (mask encodes causality) | ||
| None, # scale | ||
| 'fast', # softmax_mode | ||
| True, # recompute_mode | ||
| None, # valid_sequence_lengths | ||
| 'right', # padding_side | ||
| ) | ||
| if gqa: | ||
| q_r, k_r, v_r, mask_r = gqa_input_reshape_fwd(q, k, v, attn_mask) | ||
| outputs = [] | ||
| for g in range(q_r.shape[1]): | ||
| q_g = q_r[:, g:g + 1].contiguous() | ||
| k_g = k_r[:, g:g + 1].contiguous() | ||
| v_g = v_r[:, g:g + 1].contiguous() | ||
| res = torch.ops.hpu.sdpa_recomp_fwd(q_g, k_g, v_g, mask_r, 0.0, scale, False, True, 'fast', None, | ||
| 'right') | ||
| outputs.append(res[0]) | ||
| output = gqa_output_reshape(torch.cat(outputs, dim=1)) |
Contributor
✅ CI PassedAll checks passed successfully against the following vllm commit: |
yangulei
marked this pull request as draft
May 28, 2026 01:24
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.
Work around a Gaudi3 kernel bug where
sdpa_recomp_fwdcorrupts output for KV head group 1+ when multiple GQA groups are processed together with certain mask/shape combinations.The fix uses per-KV-group processing in the test reference implementations to produce correct ground-truth outputs.