[Kernel][PA] Support BF16 vectorized KV with asymmetric K/V head dim

[sammysun0711]

Motivation

Based on #1064
Add BF16 vectorized-5D KV support and allow different Q/K and V/output head dimensions in pa_decode_tile kernel.

Technical Details

Extend FlyDSL PA kernel to support Models with different K/V head dim with Both FP8 and BF16 support.
E.g. MiMo-V2.5-Pro with Q/K D192 and V/O D128 does not need to physically pad V to 192.
The same tile PA path is also useful for native BF16 page-64 KV on gfx942 and gfx950.

• Add a BF16 K16-MFMA specialization for vectorized-5D K/V caches.
• Derive PV accumulation, partial buffers, output stores, and reduction width
  from the physical value dimension.
• Support BF16 Q/K192 with V128/V192 and FP8 Q/K192 with V128.
• Preserve symmetric FP8 behavior and caller-owned graph-safe workspaces.

Performance (gfx950)

Batch 32, query length 4, 16Q/1KV, Q/K D192, page 64:

Context	AITER Gluon V192	FlyDSL V128	Speedup	FlyDSL V192	Speedup
16K	0.648544 ms	0.103483 ms	6.267x	0.118604 ms	5.468x
64K	2.278582 ms	0.331712 ms	6.869x	0.389114 ms	5.856x

Test Plan

Extend test_pa_decode_tile.py to covers BF16 QL1/QL4 with V128/V192, asymmetric FP8, symmetric FP8 D192, multi-partition reduction.
python -m pytest -q tests/kernels/test_pa.py

Test Result

python3 -m pytest -q     tests/kernels/test_pa.py  
===================================================================== 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/pr2-pa-bf16-asym/tests
configfile: pytest.ini
plugins: hypothesis-6.150.2, anyio-4.14.2
collected 2327 items                                                                                                                                          

tests/kernels/test_pa.py .............................................................................................................................. [  5%]
....................................................................................................................................................... [ 11%]
....................................................................................................................................................... [ 18%]
....................................................................................................................................................... [ 24%]
....................................................................................................................................................... [ 31%]
....................................................................................................................................................... [ 37%]
....................................................................................................................................................... [ 44%]
....................................................................................................................................................... [ 50%]
....................................................................................................................................................... [ 57%]
....................................................................................................................................................... [ 63%]
....................................................................................................................................................... [ 70%]
....................................................................................................................................................... [ 76%]
....................................................................................................................................................... [ 83%]
....................................................................................................................................................... [ 89%]
....................................................................................................................................................... [ 96%]
.......................................................................................                                                                 [100%]

====================================================================== 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
======================================================== 2327 passed, 2 warnings in 198.96s (0:03:18) =========================================================

Submission Checklist

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

[Copilot]

Pull request overview

This PR extends the tile-programming paged-attention decode kernel to support (1) BF16 vectorized-5D KV caches and (2) asymmetric head dimensions where Q/K head dim can differ from V/output head dim, with accompanying correctness tests.

Changes:

• Add a BF16 KV specialization (K16 BF16 MFMA path) alongside the existing FP8 path in pa_decode_tile.
• Generalize PV accumulation/output/storage to use v_head_dim (derived from the physical value dimension) instead of assuming head_dim.
• Add focused GPU correctness tests covering BF16 vectorized-5D KV (V128/V192) and FP8 asymmetric value head dim (V128).

Reviewed changes

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

File	Description
kernels/attention/pa_decode_tile.py	Adds BF16 KV specialization and supports independent v_head_dim across the kernel and host wrapper validation/dispatch.
tests/kernels/test_pa_decode_tile.py	Adds correctness coverage for BF16 vectorized-5D KV and FP8 asymmetric value head dimensions with multi-partition reduction.

Suppressed comments (1)

kernels/attention/pa_decode_tile.py:1191

• query_group_size = num_q_heads // num_kv_heads uses integer division without validating divisibility. If num_q_heads is not divisible by num_kv_heads, the kernel will silently use the truncated group size and produce incorrect indexing.

    assert context_lengths.dtype == torch.int32, f"context_lengths must be int32, got {context_lengths.dtype}"
    query_group_size = num_q_heads // num_kv_heads
    max_blocks_per_seq = block_tables.shape[1]

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[sammysun0711]

Fixed in ad712fb, The workspace comments now use [num_seqs, num_kv_heads, num_partitions, total_rows], with total_rows = query_length * query_group_size, matching both host allocation and flattened kernel indexing.

Besides, I also added host-side validation that num_q_heads is divisible by num_kv_heads before deriving the query group size.

[fsx950223]

bf16_mfma

[sammysun0711]

Thanks for your suggestions. In ab4e847 I've updated the compile-time API in from bf16_kv: bool to kv_dtype: str, validated as fp8 or bf16. The host derives it from the physical key-cache dtype, while the kernel uses an internal is_bf16_kv boolean for specialization branches. This describes the full KV storage/layout path rather than only the selected MFMA instruction.

[coderfeli]

also change pa_decode_ps_launch ?

[sammysun0711]

also change pa_decode_ps_launch ?

Thanks for your suggestions. In 53236d0, I've updated pa_decode_ps_launch to dispatch non-sliding window requests (page sizes 16 and 64) to pa_decode_tile before the FP8-specific scale and metadata setup. This change enables BF16 vectorized KV and asymmetric V/output dimensions through the PS wrapper. It also supports automatic or preallocated V-sized workspaces, including CUDA graph capture and replay.

The scope of this PR is limited to BF16 asymmetric Paged Attention (PA) support for page sizes 16 and 64. Support for sliding-window and page-1024 paths requires modifications to common utilities (e.g., pa_metadata.py and pa_metadata_tuning.py), which can be addressed in a follow-up PR.

[Copilot]

Pull request overview

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