Skip to content

perf(pa tile): readable tile-programming PA decode kernel, matches/beats production #825

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fsx950223 wants to merge 59 commits into
base: main
Choose a base branch
from
Open

perf(pa tile): readable tile-programming PA decode kernel, matches/beats production #825

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
59 commits
Select commit Hold shift + click to select a range
e9e2875
Add tile-programming PA-decode reference kernel
fsx950223 Jun 29, 2026
a317405
Make PA-decode tile output accumulator register-resident
fsx950223 Jun 29, 2026
260fc77
Distribute PA-decode softmax phase-2 across all 64 lanes
fsx950223 Jun 29, 2026
8ada08f
Reorient PA-decode tile to token=M (production) layout
fsx950223 Jun 30, 2026
28a3bb1
Add K prefetch at occupancy 3 via TLOOP softmax (pa_decode_ps structure)
fsx950223 Jun 30, 2026
f5cf69a
Hoist V load before QK while keeping K prefetch (hide both loads)
fsx950223 Jun 30, 2026
f44f95f
perf(pa tile): raw dwordx4 K load + raw QK MFMA
fsx950223 Jul 1, 2026
05052d8
Merge remote-tracking branch 'origin/main' into pa-decode-tile
fsx950223 Jul 1, 2026
3082ae8
refactor(pa tile): use shared kernels.utils global-load helpers
fsx950223 Jul 1, 2026
e4fcfb4
perf(pa tile): raw dwordx4 V load + raw PV MFMA
fsx950223 Jul 1, 2026
2fc39dc
perf(pa tile): carry block-table phys page across tiles (off critical…
fsx950223 Jul 1, 2026
879c2e6
perf(pa tile): raise partition count at high batch to fill occupancy
fsx950223 Jul 3, 2026
48b7b9c
perf(pa tile): use production's get_recommended_splits for partition …
fsx950223 Jul 3, 2026
e519999
perf(pa tile): call get_recommended_splits directly, drop the num_til…
fsx950223 Jul 3, 2026
887e866
perf(pa tile): chunk the Q staging load into dwordx4 pieces
fsx950223 Jul 3, 2026
6704756
perf(pa tile): skip redundant block-table reload when page is unchanged
fsx950223 Jul 3, 2026
12008e1
perf(pa tile): raw i32-word P store, drop the make_fragment_C/tiled_c…
fsx950223 Jul 3, 2026
1f75c07
perf(pa tile): use fmath.absf instead of maximumf(x,-x) for Q absmax
fsx950223 Jul 3, 2026
6153927
perf(pa tile): blocked K-cache layout for coalesced raw loads
fsx950223 Jul 3, 2026
ecbc2d6
perf(pa tile): blocked V-cache layout for coalesced raw loads
fsx950223 Jul 3, 2026
2154f0c
perf(pa tile): epilogue vectorization, HW exp2, V page-index prefetch…
fsx950223 Jul 6, 2026
cda1578
perf(pa tile): head_dim 64/128 support, match K/V cache layout to pa_…
fsx950223 Jul 8, 2026
fa7d548
perf(pa tile): bf16 query, scalar KV-scale via buffer_load, closer ma…
fsx950223 Jul 9, 2026
1f52610
perf(pa tile): close head_dim=64 gap to pa_decode_ps_kernel via sched…
fsx950223 Jul 10, 2026
692b4de
refactor(pa tile): trim comments, use torch.testing.assert_close, fix…
fsx950223 Jul 10, 2026
52faa5f
fix(pa tile): guard NaN in empty-partition epilogue, extract NP heuri…
fsx950223 Jul 10, 2026
74cfc29
Merge remote-tracking branch 'origin/main' into pa-decode-tile
fsx950223 Jul 10, 2026
ed684b8
refactor(pa tile): Trim remaining production-comparison comments
fsx950223 Jul 10, 2026
45b734b
refactor(pa tile): Trim remaining test/kernel comments
fsx950223 Jul 10, 2026
e5f41cf
refactor(common): Remove unused global_store helper
fsx950223 Jul 10, 2026
9f65de8
refactor(pa tile): Simplify redundant variables and aliases
fsx950223 Jul 10, 2026
0d6cc71
refactor(test): Trim stale aggregate-check comment in test_pa.py
fsx950223 Jul 10, 2026
ef3f408
fix(pa tile): use arch-native fp8 format (OCP e4m3fn on gfx950)
fsx950223 Jul 15, 2026
66df1a8
feat(pa tile): Add per-token KV dequant scale support
fsx950223 Jul 15, 2026
d45c0a4
Merge branch 'main' into pa-decode-tile
fsx950223 Jul 15, 2026
36eba6b
perf(pa tile): Fix per_token_kv occupancy regressions vs production
fsx950223 Jul 15, 2026
a763de4
style(pa tile): Trim verbose inline comments
fsx950223 Jul 15, 2026
13de3be
style(pa tile): Further trim inline comments
fsx950223 Jul 15, 2026
ba55cad
style(pa tile): Trim remaining inline comments
fsx950223 Jul 15, 2026
d94115f
feat(pa tile): Support MTP (query_length>1) and wide GQA (group_size>16)
fsx950223 Jul 15, 2026
0c47df4
fix(pa tile): Fix MTP register-allocator regression, unify M-tile cod…
fsx950223 Jul 16, 2026
44d09a0
style(pa tile): Trim inline comments
fsx950223 Jul 16, 2026
d70872a
refactor(pa tile): Simplify loop-carried state, remove dead aliases
fsx950223 Jul 16, 2026
b67005d
feat(pa tile): Replace small-block PS launcher with pa_decode_tile
fsx950223 Jul 16, 2026
fb74352
refactor(pa tile): Remove dead grid-partition small-block PS kernel
fsx950223 Jul 16, 2026
cdea67e
refactor(pa tile): Trim inline comments, drop dedicated tile test suite
fsx950223 Jul 16, 2026
20faaa4
perf(pa tile): narrow MTP/wide-GQA slowdown for small block_size
fsx950223 Jul 17, 2026
985c653
fix(pa tile): zero-mask fully-invalid softmax tiles, clamp reduce exp2
fsx950223 Jul 17, 2026
63e97d4
perf(pa tile): swap PV operands + KV-scale double-buffer to align LDS…
fsx950223 Jul 17, 2026
24f89c6
refactor(pa tile): unify to a single PV layout, delete old-PV path
fsx950223 Jul 17, 2026
1b3509b
perf(pa tile): extend the phase-split (barrier-merge) to block_size=64
fsx950223 Jul 17, 2026
81f14e5
perf(pa tile): extend the phase-split to head_dim=64 (PHASE1_MTILES =…
fsx950223 Jul 17, 2026
4b305cf
refactor(pa tile): simplify the M_TILES==1 (non-phase-split) branch
fsx950223 Jul 17, 2026
c320bf3
refactor(pa tile): use literal state-slot indices in the single-tile …
fsx950223 Jul 17, 2026
7b21bb6
refactor(pa tile): merge row_ok / row<TOTAL_ROWS epilogue guards
fsx950223 Jul 17, 2026
74cc4f3
refactor(pa tile): drop redundant row_ok, keep only row < TOTAL_ROWS
fsx950223 Jul 17, 2026
d5a6f31
refactor(pa tile): delete PHASE1_MTILES alias, use M_TILES directly
fsx950223 Jul 17, 2026
814e576
refactor(pa tile): inline the trivial _tile_tok0_and_page helper
fsx950223 Jul 17, 2026
d1c2516
refactor(pa tile): fm_contract alias, n_kv from gridDim.y, uniform KV…
fsx950223 Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Loading
Loading