[sammysun0711]

Motivation

Fix tile paged attention for D192 query staging and serving-sized physical KV cache offsets.

Technical Details

For model w/ head dimension 192 (e.g. MiMO-V2.5-Pro), each lane stages 12 BF16 query elements.
The previous load selection staged only eight. Physical page IDs were also multiplied in 32-bit arithmetic, allowing cache offsets to wrap beyond 2 GiB.

• Use three 64-bit loads for the D192 per-lane query chunk.
• Promote physical K and V page IDs to 64-bit before offset arithmetic.
• Add a D192 PyTorch-reference regression.
• Add a selected-page regression beyond the 2 GiB offset boundary.

Correctness-only change. Existing schedules are unchanged; D192 now loads the complete query chunk.

Test Plan

pytest -q tests/kernels/test_pa.py -k 'fp8_head_dim_192_matches_torch or fp8_cache_offset_above_2gib' -s

Test Result

pytest -q tests/kernels/test_pa.py -k 'fp8_head_dim_192_matches_torch or fp8_cache_offset_above_2gib' -s
==================================================== test session starts =====================================================
platform linux -- Python 3.10.12, pytest-9.0.3, pluggy-1.6.0
rootdir: /root/workspace/mimo-opt/flydsl-upstream-prs/pr1-pa-64bit/tests
configfile: pytest.ini
plugins: hypothesis-6.150.2, anyio-4.14.2
collecting ... [aiter] import [module_aiter_core] under /root/workspace/aiter/aiter/jit/module_aiter_core.so
collected 2322 items / 2320 deselected / 2 selected                                                                          

tests/kernels/test_pa.py ..

====================================================== warnings summary ======================================================
../../../../../opt/venv/lib/python3.10/site-packages/Cython/Distutils/old_build_ext.py:15
../../../../../opt/venv/lib/python3.10/site-packages/Cython/Distutils/old_build_ext.py:15
  /opt/venv/lib/python3.10/site-packages/Cython/Distutils/old_build_ext.py:15: DeprecationWarning: dep_util is Deprecated. Use functions from setuptools instead.
    from distutils.dep_util import newer, newer_group

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======================================= 2 passed, 2320 deselected, 2 warnings in 7.04s =======================================

Submission Checklist

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

[Copilot]

Pull request overview

Fixes tile paged attention for D192 queries and 64-bit physical KV-cache offsets.

Changes:

• Completes D192 query staging with three 64-bit loads.
• Promotes K/V page-offset arithmetic to 64-bit.
• Adds D192 and large-offset regression tests.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Summary
tests/kernels/test_pa_decode_tile.py	Adds focused regressions. Two moderate findings remain (2 votes each): initialize accessed pages with finite values, and make the test K-sensitive with multiple valid tokens.
kernels/attention/pa_decode_tile.py	Updates D192 query loading and K/V address arithmetic.

Suppressed comments (1)

tests/kernels/test_pa_decode_tile.py:23

• The production kernel selects the OCP FP8 path for every gfx95* target (pa_decode_tile.py:77-79), and existing tests use the same predicate. With ARCH normalized to gfx951/gfx955, both new tests are skipped by the exact-match decorators at lines 72 and 131; if those decorators are broadened without changing line 23, the test also selects the wrong FNUZ dtype. Use a consistent gfx95* predicate for support checks and FP8 selection so this regression runs on all supported CDNA4 variants.

ARCH = str(get_rocm_arch()).split(":", 1)[0]
SUPPORTED_ARCHS = ("gfx942", "gfx950")
FP8_DTYPE = torch.float8_e4m3fn if ARCH == "gfx950" else torch.float8_e4m3fnuz

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.