Skip to content

Add page-16 sliding-window attention support - #1055

Draft
zjing14 wants to merge 4 commits into
ROCm:mainfrom
zjing14:fp8-page16-swa
Draft

Add page-16 sliding-window attention support#1055
zjing14 wants to merge 4 commits into
ROCm:mainfrom
zjing14:fp8-page16-swa

Conversation

@zjing14

@zjing14 zjing14 commented Aug 21, 2026

Copy link
Copy Markdown

Summary

Extend FlyDSL paged attention for SBD-style short-query workloads:

  • Support page sizes 16 and 64 for vectorized paged BF16/F16 attention.
  • Add causal sliding-window masking and KV tile pruning.
  • Route page-size-16 FP8 sliding-window decode through the tile kernel.
  • Support query lengths 1–16 for FP8 paged decode.
  • Package the kernels modules in wheels and source distributions.

The existing dual-wave path remains page-size-64 only. Page-size-16 and sliding-window workloads use the short-query generic path.

Testing

  • Page-size-16 vectorized BF16 attention:
    • Query lengths 1–16
    • Global and sliding-window attention
  • Page-size-16 FP8 paged decode:
    • Query lengths 1, 15, and 16
    • Global and sliding-window attention
python -m pytest -q \
  tests/kernels/test_flash_attn_fwd.py \
  tests/kernels/test_pa.py \
  -k 'test_vectorized_paged_page16_short_query or test_small_block_accuracy'

Result: 38 passed.

- Black and Ruff checks passed.
- git diff --check passed.
- Built and inspected the source distribution; required kernels.attention and kernels.common modules are included.
- Validated downstream with FlyDSL FP8 attention and SBD on 4× MI355X.

Jing Zhang added 4 commits August 21, 2026 12:36
# Conflicts:
#	kernels/attention/flash_attn_interface.py
#	kernels/attention/flash_attn_utils.py
#	tests/kernels/test_flash_attn_fwd.py
Hoist the FP8 V-tile LDS address calculation out of repeated loads in the gfx950 dual-wave FlashAttention loop. Preserve the existing load path for non-FP8 PV configurations.
Reverse the q-block launch order for the dense gfx950 BF16 GQA causal configuration so longer-running blocks issue first. Enable the optimized path from sequence length 256 while retaining the existing guards for unsupported modes.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant