[Kernel][Perf] Optimize paged-attention metadata decode - #910
Open
fsx950223 wants to merge 8 commits into
Open
Conversation
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR optimizes paged-attention (PA) metadata generation and the block-1024 persistent decode path by moving metadata decode/reduction onto FlyDSL tile-programming APIs and introducing shape/device-tuned persistent grid sizing via a CSV lookup + an offline tuner.
Changes:
- Refactors PA metadata scheduler + persistent decode kernel to use FlyDSL layout/copy primitives and reduce per-token scale/reduction overhead.
- Adds a CSV-backed grid-multiplier lookup (
pa_metadata_grid_tuning.csv) with read/write utilities plus an offline CUDA-Graph-based tuner script. - Extends PA tests with metadata-specific accuracy and validation coverage (incl. block-size and dtype rejection cases).
Reviewed changes
Copilot reviewed 6 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_pa_metadata_tuning.py | Adds unit tests for tuning CSV read/write/upsert and lookup behavior. |
| tests/kernels/test_pa.py | Adds metadata accuracy + validation tests; passes per_token_kv through metadata creation. |
| scripts/tune_pa_metadata_grid.py | Adds an offline tuner that benchmarks candidate grid multipliers and writes winners into the CSV. |
| kernels/attention/pa_metadata.py | Major refactor/optimization: new tile-programming decode path (block_size=1024 only) and updated reduction kernel. |
| kernels/attention/pa_metadata_tuning.py | Implements CSV schema, normalization, cached lookup, and atomic upsert writer. |
| kernels/attention/pa_metadata_grid_tuning.csv | Provides initial validated tuning entries (gfx942/80CU shapes). |
| kernels/attention/pa_decode_fp8.py | Integrates tuning lookup into metadata generation and updates metadata API usage. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # resident per CU → more waves in flight → better latency hiding. | ||
| base_cu = props.multi_processor_count | ||
| num_sm = base_cu * _PA_METADATA_GRID_OVERSUB | ||
| num_blocks = key_cache.shape[0] |
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…tile-programming Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # python/flydsl/autotune.py
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
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.
Summary
Paged-attention metadata generation now uses the tile-programming path and shape/device-tuned persistent grids. The resulting decode-plus-reduce path is faster across representative query lengths and quantization modes while retaining a safe one-grid-per-CU default for untuned shapes.
Motivation
The previous metadata implementation mixed legacy data movement with a fixed grid oversubscription. A single grid choice performed well for some Q4 per-token workloads but regressed lower-Q and per-tensor shapes, so the scheduling choice needs to follow measured shape and device characteristics.
Changes
grid_multiplier=1.Performance (MI308X, gfx942, 80 CUs)
GPU-only p50 from 35 rocprofv3 dispatches. Times cover the metadata decode plus split reducer for block size 1024, 16 query heads, one KV head, and head dimension 128.
Geometric-mean speedup is 1.19x; 11 of 12 measured shapes improve. The sole underfill delta is 0.24 us at an approximately 12 us launch-scale runtime.
Testing
python3 -m pytest -q tests/unit/test_pa_metadata_tuning.pyFLYDSL_RUNTIME_ENABLE_CACHE=0 python3 -m pytest -q tests/kernels/test_pa.py -k metadatagit diff --checkDependencies
Breaking Changes
None. Existing callers remain valid; tuned precomputed metadata may optionally provide
per_token_kvfor scale-mode-specific lookup.