Support mqa_logits blocksize is multiple of ChunkK cases - #1674
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds support for cases where the KV cache blocksize is a multiple of ChunkK (in addition to the existing support for ChunkK being a multiple of blocksize), and introduces experimental s_set_prio instructions to optimize performance on AMD CDNA3 architecture.
Key Changes
- Updated assertion logic to support bidirectional divisibility between ChunkK and KVBlockSize
- Refactored the main kernel into two code paths: one for when ChunkKPerStage divides evenly into KVBlockSize, and another for the reverse case
- Added
s_set_priocalls for instruction scheduling optimization with different priority levels (0-3)
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
op_tests/op_benchmarks/triton/bench_deepgemm_attention.py |
Updated benchmark parameters: ChunkK from 256 to 128, WavePerEU from 2 to 5, and fixed num_blocks calculation to properly account for blocksize |
aiter/ops/triton/pa_mqa_logits.py |
Relaxed assertion to allow KVBlockSize % ChunkK == 0 in addition to ChunkK % KVBlockSize == 0 |
aiter/ops/triton/gluon/pa_mqa_logits.py |
Major refactoring: added version compatibility for 2D vs 3D MFMA layouts, added s_set_prio import with fallback, introduced LoadBlockIndiceForEachStage branching for two different block loading strategies |
Comments suppressed due to low confidence (1)
aiter/ops/triton/gluon/pa_mqa_logits.py:43
- Except block directly handles BaseException.
except:
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
valarLip
previously approved these changes
Dec 18, 2025
sjfeng1999
force-pushed
the
feat-mqa_logits_large_blocksize
branch
3 times, most recently
from
December 23, 2025 06:03
8029feb to
d40fb02
Compare
sjfeng1999
force-pushed
the
feat-mqa_logits_large_blocksize
branch
from
December 24, 2025 06:18
cd2c8f7 to
c5a826a
Compare
ZhangLirong-amd
pushed a commit
that referenced
this pull request
Dec 29, 2025
* Support mqa_logits blocksize is multiple of ChunkK cases and Add s_set_prio optimization * Add var-context optimzation for pa_mqa_logits * resolve ruff suggestion
zhuyuhua-v
pushed a commit
that referenced
this pull request
Jan 14, 2026
* Support mqa_logits blocksize is multiple of ChunkK cases and Add s_set_prio optimization * Add var-context optimzation for pa_mqa_logits * resolve ruff suggestion
valarLip
pushed a commit
that referenced
this pull request
Mar 18, 2026
* Support mqa_logits blocksize is multiple of ChunkK cases and Add s_set_prio optimization * Add var-context optimzation for pa_mqa_logits * resolve ruff suggestion
1 task
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.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist