Repository navigation
[FlyDSL] [Misc] pa_decode: key the tile kernel on FlyDSL dtypes and drop the batch-first plan grid - #7388
Conversation
WAITING ON AUTHORPR Self-ReviewRevision Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision. Questions to answer in the PR description
Edit the Missing or invalid answers
Shadow mode reports what would block but keeps the status successful. Set |
…trings and flags
The schedule took kv_dtype/query_dtype as strings ("fp8", "bf16", "mxfp4", "f16")
and the traits re-derived them into parallel flags and types: is_bf16_kv, MXFP4,
IS_BF16 (which also counted e4m3 queries), Q_DTYPE (the output type), Q_IN_DTYPE,
FP8 and FP8_MAX picked by arch, KV_DTYPE (bf16 for an MXFP4 cache), and TILE_TOK
(an alias of KV_COMPUTE_BLOCK). The fp8 flavor came from the arch, not the cache.
- compile_pa_decode_tile and PaDecodeSchedule.select take the cache's and the
query's torch dtypes (kv_dtype is required) and map them through a new generic
torch_to_fx_dtype/fx_to_torch_dtype in tensor_shim; uint8 KV is MXFP4. An FP8
cache must be the arch's MFMA format (FNUZ on gfx942, OCP on gfx950).
- PaDecodeTraits carries kv_dtype and query_dtype as FlyDSL types. FP8, KV_DTYPE,
Q_IN_DTYPE, is_bf16_kv, MXFP4 and IS_BF16 are gone; Q_DTYPE is OUT_DTYPE;
FP8_MAX is torch.finfo of the format Q/P are quantized to (none for BF16 KV);
TILE_TOK is KV_COMPUTE_BLOCK. PV_BF16 keeps only its MFMA-operand meaning.
- pa_decode passes tensor dtypes; _flydsl_pointer_dtype uses the generic map.
- Kernel names are unchanged; gfx1250 keeps its string dtypes at the boundary
and refuses the dtypes it does not support (an e4m3 query became f16 there).
- test: the select refusals take torch dtypes; a FNUZ cache on gfx950 is refused.
Test (MI355X gfx950, against main e4973c4): op_tests/test_flydsl_pa_decode.py
337 passed, 37 skipped (main: 336 passed, 37 skipped; the extra case is the new
FNUZ refusal, the other 373 outcomes match); all 323 FlyDSL cache entries (272
distinct gpu.binary) byte-identical to main; graph-replay time unchanged on 132
fp8/MXFP4 shapes, outputs bitwise equal. The four select refusals pass here and
fail on main.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
fddc69d to
c4243be
Compare
There was a problem hiding this comment.
🔵 Needs a closer look
The dtype refactor affects gfx942 and gfx1250 dispatch and query loads, but native validation covers only gfx950, leaving compilation, numerical correctness, and generated-code equivalence on the other targets unconfirmed.
No open issues 🎉
What changed in this PR
Refactors FlyDSL paged-attention decode to derive compile-time formats from tensor dtypes instead of strings and duplicate flags, while preserving public APIs and existing gfx1250 merge handling.
Changes:
- Adds shared torch ↔ FlyDSL dtype conversion helpers and updates wrapper calls.
- Uses FlyDSL types in scheduling, cache keys, and kernel branches, removing redundant traits.
- Rejects unsupported dtype combinations and architecture-incompatible FP8 formats, with an added refusal test.
| File | Description |
|---|---|
op_tests/test_flydsl_pa_decode.py |
Updates dtype arguments and tests FNUZ cache rejection on gfx950. |
aiter/ops/flydsl/pa_decode.py |
Passes torch dtypes and uses shared pointer-type conversion. |
aiter/ops/flydsl/kernels/tensor_shim.py |
Adds bidirectional torch/FlyDSL dtype maps. |
aiter/ops/flydsl/kernels/pa_decode/traits.py |
Consolidates format traits and validates supported dtype combinations. |
aiter/ops/flydsl/kernels/pa_decode/pipeline.py |
Replaces format flags and tile-size aliases. |
aiter/ops/flydsl/kernels/pa_decode/op_softmax.py |
Uses dtype checks for probability packing. |
aiter/ops/flydsl/kernels/pa_decode/op_lds.py |
Derives query load sizes and cache handling from element types. |
aiter/ops/flydsl/kernels/pa_decode/op_epilog.py |
Renames the output-type trait and confines FP8 scaling to its branch. |
aiter/ops/flydsl/kernels/pa_decode/context.py |
Selects scale initialization using cache dtype membership. |
aiter/ops/flydsl/kernels/pa_decode_kernel.py |
Updates compile arguments, preserves name tags, and validates gfx1250 dtypes. |
Review effort: Balanced
zhiding512
left a comment
zhiding512
left a comment
There was a problem hiding this comment.
Reviewed at c4243be. No blocking issues from my side; 4 minor, non-blocking suggestions inline.
Logic check: went through the diff and checked all call sites and removed fields at head; no dangling references. Q load width, FP8_MAX (448 / 240), output dtype, and the scale / page-table / PV branch conditions are equivalent to main. With BF16 KV FP8_MAX becomes None, but every use is inside an FP8 or per-token branch, so it is unreachable.
gfx942 (MI308X, not covered by CI): ran op_tests/test_flydsl_pa_decode.py on this head and on main e4973c4bed:
- PR: 267 passed / 107 skipped; main: 266 passed / 107 skipped. The only difference is the new fnuz-kv-on-gfx950 test.
- The 538 FlyDSL kernel cache files produced by each run have identical contents, i.e. the generated gfx942 kernels are unchanged.
- Note: the container's torch 2.9.1 lacks
shared_memory_per_block, so both runs used the same test-only shim returning 64 KiB (without it, every test errors on both sides).
gfx1250: static check only, not run (no machine available). dtypes.fp8 is e4m3fn there, so the new FP8 check does not reject it, and the dtype tags passed down to compile_gfx1250_pa_decode are the same as before.
batch_first_plan_grid only reordered CTA dispatch for one tuned shape family: FP8 per-token KV, MTP4 with four query splits, GQA16, page 128, one KV head, sliding window 4096-8192, batch 4-24, single-tile plan. Every task slot and every output is the same with the flat grid. The special case cost a launch-grid branch, slot / partial-slot / pout-slot arithmetic in three places, and a cache-key entry. Measured on MI355X (gfx950), 50 shapes of that family (sliding window 4096/6000/8192, batch 4-24, ctx 32768), CUDA-graph replay, 2 rounds: - the 41 shapes that used it are 1.05% slower geomean with the flat grid, from -19.5% (window 8192, batch 6) to +32.9% (window 6000, batch 20); - the 9 shapes it excluded are unchanged; - outputs are bitwise equal. Test: op_tests/test_flydsl_pa_decode.py 337 passed, 37 skipped, the same outcomes as before. No test shape used the batch-first grid, and all 323 FlyDSL cache entries are byte-identical. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
fsx950223
changed the title
|
Independent gfx950 follow-up confirms the disclosed batch-first-grid regression in both eager execution and graph replay, using the same GPU-kernel-time metric. For window=6000 / batch=20, three paired rounds give:
The baseline selected Versions and method
Eager versus graph: execution mode and timing scopeEager executes the Python wrapper, validation/kernel selection and kernel submission on each call. Graph replay captures the GPU sequence once and replays it, avoiding repeated Python-wrapper execution and individual kernel submissions. Both columns here sum GPU kernel durations, including decode/reduction when dispatched, and exclude Python-wrapper time and launch gaps. They are neither full eager-call latency nor graph-event elapsed time, so their absolute numbers should not be compared directly with the PR's graph-replay/event measurements. The grid change affects GPU CTA ordering and is visible in both modes. The AITER helper already supports the two modes: from aiter.test_common import run_perftest
# launch() closes over fixed inputs, output, work_plan and preallocated workspace.
# Compile/initialize before graph capture; correctness checks stay outside timing.
launch()
kwargs = dict(num_iters=101, num_warmup=20, num_rotate_args=1)
_, eager_us = run_perftest(launch, testGraph=False, **kwargs)
_, graph_us = run_perftest(launch, testGraph=True, **kwargs)In this helper, Window-family resultsCommon configuration: FP8 per-token K/V, BF16 query, MTP4, GQA16, D128, page128, transposed V, one KV head, context32768. The plan budget is exactly window_tiles = (window + 4 - 2 + 256 - 1) // 256 + 1
max_partitions = window_tiles
workgroup_budget = batch * window_tiles
Nine of these ten points used the batch-first grid on the baseline; Rechecking the original-matrix outliersAn earlier single complete paired run of the shared FP8/BF16 correctness matrix had 186 timed points, one expected BF16-KV/FP16-query skip and a latency-ratio geomean of +0.22%. It did not select the batch-first grid. Seven points were initially slower by more than 5%, so they were included in this three-round follow-up:
Six of the seven did not reproduce their initial >5% slowdown in the medians above. In the exercised runtime caches, 187 distinct ELF payloads were identical across arms. The differing payloads were confined to window4096/6000/8192 grid variants. Binary identity supports the matrix check but does not substitute for timing or explain the eager-only outlier. Copyable reproductionUse an idle gfx950 and the same PyTorch/ROCm/FlyDSL environment for both arms. Run these commands from an AITER checkout: git fetch origin main \
refs/pull/7388/head:refs/remotes/origin/pr7388-head \
refs/heads/zhiding512/pa_benchmark:refs/remotes/origin/zhiding512/pa_benchmark
export PA_REPRO_DIR="$(mktemp -d /tmp/pa7388-repro.XXXXXX)"
git worktree add --detach "$PA_REPRO_DIR/base" e4973c4bed57a615068089fa02c88c64d644fd42
git worktree add --detach "$PA_REPRO_DIR/head" 612ed7d84b3e9c43e078dcd677218ccec58cc07f
git show 67525723f714be3af2169efa64f04c70298afdf4:op_tests/op_benchmarks/flydsl/bench_pa_decode.py \
> "$PA_REPRO_DIR/bench_pa_decode.py"Save the adapter below as Reproduction adapter: 17-point follow-up, window family or original matriximport ast
import importlib.util
import inspect
import os
import sys
from functools import partial
from pathlib import Path
import aiter
import pytest
import torch
from op_tests import test_flydsl_pa_decode as tests
benchmark_path = Path(sys.argv.pop(1)).resolve()
suite = sys.argv.pop(1)
assert suite in ("grid", "ut", "selected")
mode = os.environ.get("PA_MODE", "eager")
assert mode in ("eager", "graph")
root = Path.cwd().resolve()
assert Path(aiter.__file__).resolve().is_relative_to(root), aiter.__file__
assert Path(tests.__file__).resolve().is_relative_to(root), tests.__file__
# Expose the original test's prepared launch in memory. No kernel edits.
mark = next(m for m in tests.test_pa_decode.pytestmark
if m.name == "parametrize" and m.args[0] == "case")
tests._PA_DECODE_CASES = mark.args[1]
if not hasattr(tests, "_run_pa_decode_case"):
fn = ast.parse(inspect.getsource(tests.test_pa_decode)).body[0]
fn.name = "_run_pa_decode_case"
fn.decorator_list = []
fn.body.append(ast.parse(
"return case, launch, check, arguments, options, plan"
).body[0])
module = ast.fix_missing_locations(ast.Module(body=[fn], type_ignores=[]))
exec(compile(module, tests.__file__, "exec"), tests.__dict__)
# Keep workspace allocation the same for both execution modes.
original_prepare = tests._run_pa_decode_case
def prepare(case, device, monkeypatch):
return original_prepare(dict(case, capture=True), device, monkeypatch)
tests._run_pa_decode_case = prepare
if suite != "ut":
matrix = tests._PA_DECODE_CASES
specs = [(6000,20), (6000,12), (8192,8), (4096,5), (6000,8),
(4096,24), (4096,6), (8192,6), (8192,12), (6000,4)]
tests._PA_DECODE_CASES = []
for window, batch in specs:
tiles = (window + 4 - 2 + 256 - 1) // 256 + 1
case = dict(
kv_dtype="fp8", query_dtype=torch.bfloat16,
scale_mode="per-token", trans_v=True,
query_length=4, kv_heads=1, group_size=16,
head_dim=128, block_size=128, lengths=(32768,) * batch,
window=window, max_partitions=tiles, budget=batch * tiles,
)
tests._PA_DECODE_CASES.append(
pytest.param(case, id=f"grid-w{window}-b{batch}")
)
if suite == "selected":
selected = {
"d256-page64-per-token",
"fp8-torch.float16-per-tensor-plain-v-decode",
"fp8-torch.float16-per-tensor-transposed-v-decode",
"d256-page64-q4-np130",
"fp8-torch.bfloat16-per-tensor-transposed-v-head256",
"fp8-torch.float16-per-tensor-transposed-v-head256",
"fp8-torch.float16-per-tensor-plain-v-head256",
}
tests._PA_DECODE_CASES += [p for p in matrix if p.id in selected]
spec = importlib.util.spec_from_file_location("pa7388_benchmark", benchmark_path)
benchmark = importlib.util.module_from_spec(spec)
spec.loader.exec_module(benchmark)
benchmark.run_perftest = partial(benchmark.run_perftest, testGraph=(mode == "graph"))
original_case = benchmark.benchmark_case
def benchmark_case(*args, **kwargs):
record = original_case(*args, **kwargs)
record.update(mode=mode, timing=f"{mode}_gpu_kernel_sum")
return record
benchmark.benchmark_case = benchmark_case
benchmark.main()Run the same 17-point follow-up, with separate caches and alternating arm order: PA_GPU=6 # Replace with an idle device in your ROCR enumeration.
for round in 1 2 3; do
if [ "$round" = 2 ]; then
PA_ORDER=(head base)
else
PA_ORDER=(base head)
fi
for mode in eager graph; do
for arm in "${PA_ORDER[@]}"; do
PA_TREE="$PA_REPRO_DIR/$arm"
(
cd "$PA_TREE"
PA_MODE="$mode" ROCR_VISIBLE_DEVICES="$PA_GPU" \
FLYDSL_RUNTIME_CACHE_DIR="$PA_REPRO_DIR/cache-$arm" \
PYTHONPATH="$PA_TREE${PYTHONPATH:+:$PYTHONPATH}" \
python "$PA_REPRO_DIR/repro.py" "$PA_REPRO_DIR/bench_pa_decode.py" selected \
--rounds 1 --iters 101 --warmup 20 \
--jsonl "$PA_REPRO_DIR/$mode-$arm-round$round.jsonl" \
> "$PA_REPRO_DIR/$mode-$arm-round$round.log" 2>&1
)
done
done
done
Validation and limitsAll 204 records in the three-round, two-mode, 17-point follow-up passed. Fixed-input hashes and correctness-output hashes match across versions and modes; FP32-reference checks also pass after timing. Per-process KFD checks found zero external GPU VRAM in every accepted record. Earlier interrupted/contended attempts were excluded. The full 186-point matrix still has only its earlier single complete paired round; the later three-round coverage is the 17-point subset above. This does not establish absence of regression across every PA shape. MXFP4, CPU/Python-wrapper latency, end-to-end inference, event-timed graph-replay latency and dedicated HBM throughput were not tested here. |
Review follow-up. - compile_pa_decode_tile passes the caller's torch dtypes straight to PaDecodeSchedule.select, which converts them once with pa_fx_dtypes (replacing kv_fx_dtype), so its torch.dtype annotation holds. The gfx1250 branch, which does not call select, uses the same helper. The direct_output/MXFP4 check moves into select next to the conversion. - A value that is not a torch dtype, such as the former "fp8" strings, raises TypeError. Every unsupported dtype raises ValueError, including one FlyDSL has no type for (complex64 was a TypeError). torch_to_fx_dtype no longer passes non-dtypes through; it and fx_to_torch_dtype raise ValueError for a missing mapping. - Error messages name torch dtypes, as passed: "FP8 KV on gfx950 must be torch.float8_e4m3fn, got torch.float8_e4m3fnuz". - compile_pa_decode_tile documents its kv_dtype values, including that uint8 (dtypes.fp4x2 on torch without float4_e2m1fn_x2) selects MXFP4. Test: invalid-select gains string-kv (TypeError) and unmapped-kv (complex64, ValueError); the fnuz-kv-on-gfx950 regex matches the torch dtype names. All 6 invalid-select cases pass; string-kv, unmapped-kv and the updated fnuz case fail on the previous head. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
PaDecodeSchedule.select took OCP FP8 on gfx950 and FNUZ on every other arch, so on gfx1250, whose dtypes.fp8 is OCP e4m3fn, it refused FP8 KV. The gfx1250 branch of compile_pa_decode_tile never calls select, but the op test's _check_pipe_policy does: six d256 FP8 cases failed on gfx1250 with "FP8 KV on gfx1250 must be Float8E4M3FNUZ", and pass on main (e4973c4). FNUZ is now expected only on gfx94x, matching aiter's dtypes.fp8; gfx942 and gfx950 select as before. Test (gfx1250): op_tests/test_flydsl_pa_decode.py 217 passed, 159 skipped. The previous head had 6 failed, 209 passed, 159 skipped; the two new cases are the earlier invalid-select additions. Skip reasons are unchanged, and all gfx1250 tile kernels are byte-identical to the previous head's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
@zhiding512 thanks for the review, the gfx942 run, and the independent rerun of the batch-first family in both modes with a copyable repro. Batch-first grid. Your numbers agree with the PR's graph run. w6000/b20 is the one large loss (+37.8% here, +32.9% in the description). w6000/b12 and w6000/b8 lose a few percent, and w8192/b6 is a consistent ~19% win. I'm keeping the removal in this PR, as the description says. The grid only helped that one tuned family, and it cost a launch-grid branch, slot arithmetic in three places and a cache-key entry.
New commits
Validation of the new head (8f3c0f4) on gfx1250 only
Not rerun: gfx950 and gfx942. Both commits change only host-side dtype conversion, validation and messages; on those arches, |
Resolve PA decode conflicts with #7388, which keys the tile kernel on FlyDSL dtypes: the NHD helpers (d256_page64_pipe, nhd_max_query_length, require_nhd_pipe and the default workgroup budget) now take FlyDSL KV types, the NHD V address uses KV_COMPUTE_BLOCK instead of TILE_TOK, and pa_decode uses torch_to_fx_dtype in place of _flydsl_pointer_dtype. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Brings in #7388 (FlyDSL-dtype tile kernel keys, no batch-first plan grid). This branch's traits and kernels now use its names: kv_dtype/query_dtype FlyDSL types, OUT_DTYPE, KV_COMPUTE_BLOCK and MXFP4_KV_DTYPE in place of MXFP4/is_bf16_kv/Q_IN_DTYPE/Q_DTYPE/KV_DTYPE/TILE_TOK; the gfx950 one-partition direct-output allowance moves into the schedule's direct_output check with #7388's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
#7388 passes torch dtypes to compile_pa_decode_tile: gfx1250_tile_traits takes them as well and converts them as the gfx1250 branch does, pa_decode shares them through tile_options, and the step-width compile test passes torch dtypes. Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # aiter/ops/flydsl/kernels/pa_decode_kernel.py # aiter/ops/flydsl/pa_decode.py
Summary
Targets
main(e4973c4bed), which already contains #7056 (the MXFP4 path) and #7079.The FlyDSL pa_decode compile layer now takes the cache's and the query's torch dtypes and keys everything on FlyDSL element types from one generic torch → FlyDSL map. The compile layer is
compile_pa_decode_tile,PaDecodeSchedule.selectandPaDecodeTraits. Until now it took the strings"fp8"/"bf16"/"mxfp4"/"f16"and a set of flags re-derived from them.A second commit removes
batch_first_plan_grid, a CTA-ordering special case for one tuned shape family. Apart from that family, the generated kernels are byte-identical tomain.Motivation
On
mainthe KV and query formats arrive as strings, and the traits re-derive them into parallel names. Several of those names no longer say what they mean:IS_BF16has two definitions. It isquery_dtype == "bf16"inselect()andquery_dtype in ("bf16", "fp8")increate(). The second one means "the output is bf16", which is also whatQ_DTYPEholds. The query's own type isQ_IN_DTYPE.FP8/FP8_MAXfollow the arch, not the cache. They are chosen byis_gfx950. A cache whose fp8 flavor does not match the arch is never rejected; the MFMAs would read its bytes in the arch's format.KV_DTYPEis wrong for MXFP4. It isfx.BFloat16for an MXFP4 cache, which is packed E2M1. It is only harmless because the MXFP4 loaders never read it.is_bf16_kvandMXFP4both restatekv_dtype.TILE_TOKis an alias ofKV_COMPUTE_BLOCK.batch_first_plan_gridis a narrow special case with a small, mixed effect. It only reorders the launch grid for FP8 per-token KV, MTP4 with four query splits, GQA16, page 128, one KV head, sliding window 4096–8192, batch 4–24 and a single-tile plan; every task slot and output is the same with the flat grid. It costs a grid branch, slot arithmetic in three places and a cache-key entry. Measured on that family, it is worth about 1% on average, and some shapes are faster without it (see Performance).The cache tensor's dtype already says which format it is. Carrying it through removes the duplicates, and a wrong fp8 flavor becomes an explicit error.
Changes
kernels/tensor_shim.pygainstorch_to_fx_dtype/fx_to_torch_dtype.float8_e8m0fnu, andfloat4_e2m1fn_x2→Float4E2M1FNx2. It skips dtypes that the installed torch or FlyDSL lacks.pa_decode.py's_flydsl_pointer_dtypenow uses it.compile_pa_decode_tileandPaDecodeSchedule.selecttakekv_dtype: torch.dtype(required) andquery_dtype: torch.dtype = torch.float16.kv_fx_dtype).selectrefuses unsupported dtypes. It also refuses an FP8 cache that is not the arch's MFMA format: FNUZ on gfx942, OCP on gfx950.PaDecodeTraitscarrieskv_dtypeandquery_dtypeas FlyDSL types. They are also in the cache key.FP8,KV_DTYPE(merged intokv_dtype),Q_IN_DTYPE(nowquery_dtype),is_bf16_kv,MXFP4,IS_BF16,TILE_TOK(nowKV_COMPUTE_BLOCK).Q_DTYPEis renamedOUT_DTYPE.FP8_MAXistorch.finfoof the format Q/P are quantized to: the FP8 cache's own, E4M3 for MXFP4, andNonefor BF16 KV.PV_BF16keeps only its "bf16 MFMA operands" meaning. The "no scalar fp8 scale" checks usekv_dtype in FP8_KV_DTYPES.traits.kv_dtype/traits.query_dtypewith FlyDSL types.op_epilogcomputes1 / FP8_MAXonly on its FP8 branch.batch_first_plan_gridremoved (second commit): its condition, traits field and cache-key entry, the batch-first launch grid, and the batch-first slot / partial-slot / pout-slot arithmetic inpa_decode_kernel.py,context.pyandop_epilog.py. Every plan now launches the flat(capacity, kv_heads × query_splits, 1)grid._dtype_tag). The gfx1250 branch converts back to its own string dtypes and refuses dtypes it does not support. Before, an e4m3 query fell back to f16 there.pa_decode/pa_decode_mxfp4signatures, the generated kernels outside the batch-first family, and the gfx1250 kernels. The gfx1250 branch keeps main'sdirect_output/merge_partitions/cluster_mergehandling from [FlyDSL] Merge gfx1250 PA decode partitions in the tile kernel and speed up dense decode #7079.Performance (if applicable)
Hardware: MI355X (gfx950, 256 CUs),
rocm/atom-dev:latest(torch 2.10, ROCm 7.2.4), FlyDSL 0.4.1. This PR is compared withmaine4973c4bedon one idle GPU, with a separate FlyDSL cache per tree.Timing method:
pa_decode, decode and reduce together.Shapes:
Per-shape ranges (this PR vs main):
d128-p128-per-token-tv-g8-q1-b32-c4096) to +1.1% (d128-p128-per-tensor-tv-g16-q1-b32-c4096)These 132 shapes were timed on the first commit (
c4243beac2). None of them has a sliding window, so none ever used the batch-first grid, and the removal leaves their flat-grid code unchanged. Device time is unchanged because the kernels are byte-identical.The MXFP4 eager row uses two extra MXFP4-only rounds (4 and 5), because round 2 had a one-off outlier:
qfp8-direct-b32-c2048took 39.36 µs once, against 30.65–31.89 µs in rounds 3–5 and 30.33–30.93 µs on main. Per round, the MXFP4 eager geomean was +1.9%, +0.6%, −0.9% and +0.6%.The remaining eager cost is the dtype conversion and validation that
select()now does on every call:qbf16-direct-b1-c2048: 29.85 → 30.63 µs (+2.6%).FP8 eager calls get slightly faster: the old
_flydsl_pointer_dtypebuilt a new dict on every call.Batch-first plan grid removal
Same hardware. 50 shapes of the family the batch-first grid was tuned for:
max_partitions= window tiles, so each tile is its own task.Device time is the median of 7×300 CUDA-graph replays. Two rounds per tree, order alternated; run-to-run 0.957–1.018. Outputs are bitwise equal.
Shapes that moved more than ±3% (both rounds agree):
This is a reproducible regression on some shapes, the largest at window 6000 / batch 20, and this PR accepts it in exchange for removing the special case.
Testing
Only MI355X (gfx950) was run. MI250X and MI300X are not tested; see "Validation and unverified scope".
Documentation
Dependencies
No PR dependency: #7056 and #7079 are already in
main.Breaking Changes
Internal API only:
compile_pa_decode_tileandPaDecodeSchedule.selecttake torch dtypes forquery_dtype/kv_dtypeinstead of strings, andkv_dtypeis now required.PaDecodeTraitsloses or renamesFP8,KV_DTYPE,Q_IN_DTYPE,Q_DTYPE,TILE_TOK,is_bf16_kvandMXFP4._PA_DECODE_TILE_CACHEkeys hold FlyDSL types instead of strings.PaDecodeTraits.batch_first_plan_gridis removed, and the schedule cache key loses that entry.The in-tree callers (
pa_decode.pyand the op test) are updated. The publicpa_decode/pa_decode_mxfp4signatures are unchanged.PR Self-Review
Reviewed revision
612ed7d
Behavior and intent
The pa_decode compile layer now derives the KV and query formats from the tensors' torch dtypes, mapped once to FlyDSL types. Before, it used strings plus flags such as
IS_BF16,FP8,KV_DTYPEandQ_IN_DTYPE. Some of those disagreed with their names, andFP8was picked by arch instead of by the cache.The second commit removes
batch_first_plan_grid. The shapes that used it now launch the flat grid, which changes their speed (−19.5% to +32.9%, +1.05% geomean) but not their results.Apart from that, kernel behavior is unchanged. The only new behavior is in
select(): it raisesValueErrorfor an FP8 cache that is not the arch's MFMA format, and for dtypes it does not support. The gfx1250 branch now refuses an e4m3 query instead of treating it as f16.Validation and unverified scope
On MI355X (gfx950):
python3 -m pytest op_tests/test_flydsl_pa_decode.py: 337 passed, 37 skipped. Onmain(e4973c4bed) it is 336 passed, 37 skipped; the extra case is the new FNUZ refusal, and the other 373 outcomes match.gpu.binary) are byte-identical tomain, compared by content.pytest -k invalid-select: 4 passed here; the same test file gives4 failedonmain.psf/black@stable) and ruff 0.16.0 report clean.Unverified:
select/creategive FNUZ/240 as before).PR-specific question
The changed test is the
invalid-select-*group oftest_pa_decode.The first three cases (
mxfp4-on-gfx942,fp8-query-fp8-kv,fp8-query-bf16-kv) now pass torch dtypes toselect. Onmain(e4973c4bed),selectonly accepts strings, so it raisesUnsupported PA KV dtype torch.float4_e2m1fn_x2before reaching the refusal each case asserts; the regex does not match and the case fails. I ran this.The new
fnuz-kv-on-gfx950case fails onmainfor the same reason. With string dtypesmainhas no flavor check at all.The wrapper-level tests stop at
pa_decode's own checks. These cases are the only ones that callselectdirectly with a mocked arch, so they are what pins the schedule's own refusals.Changed op tests
op_tests/test_flydsl_pa_decode.py:_SELECT_INVALIDpasses torch dtypes, matchingselect's new signature.fnuz-kv-on-gfx950, which checks thatselectrejects afloat8_e4m3fnuzcache on gfx950, where the MFMAs read OCP e4m3.No other test covers that refusal; it is new in this PR. Nothing else in the file changes.
🤖 Generated with Claude Code