fix edge cases

[fsx950223]

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests
ci:atom	ATOM benchmark (DeepSeek-R1 + GPT-OSS)
ci:vllm	vLLM benchmark
ci:all	All of the above

Add labels via the sidebar or gh pr edit 2906 --add-label <label>

[Copilot]

Pull request overview

This PR adjusts the Triton/Gloun paged-attention decode kernel to handle a couple of edge cases around multi-token-per-program (MTP) indexing and per-token KV scale staging.

Changes:

• Fixes query_token_idx computation for QUERY_SEQ_LEN_POW2 > 1 to avoid double-counting the mtp_idx offset after max_logits_base_offsets is expanded into the flattened query-group space.
• Removes a later redundant/incorrect recomputation of query_token_idx based on the already-expanded max_logits_base_offsets.
• Ensures per-token KV scales are reshaped to a 1D [CONTEXT_PARTITION_SIZE] buffer before staging into shared memory.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.