Skip to content

[LLVM][FlyDSL] Bump up LLVM and Adapt FlyDSL to upstream LLVM/MLIR API changes - #945

Merged
coderfeli merged 20 commits into
mainfrom
jli/update-llvm-basemain
Aug 9, 2026
Merged

coderfeli merged 20 commits into
mainfrom
jli/update-llvm-basemain

Conversation

@jli-melchior

@jli-melchior jli-melchior commented Jul 31, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

Adapt FlyDSL to LLVM API changes. This PR updates both C++ and Python layers to match upstream LLVM/ROCDL breaking changes.

Technical Details

C++ changes

  • OpaqueProperties → PropertyRef
  • getType() → getSourceElementType() for GEP
  • getSizeInBits() → getTypeSizeInBits()
  • MFMA/WMMA enum attribute constructors

Python compat wrappers (expr/rocdl/init.py)

  • MFMA blgp: int → #rocdl<mfma_perm_b ...> attribute
  • MFMA scale cbsz/blgp: int → #rocdl<matrix_format ...> attribute
  • WMMA fmtA/fmtB/modC/scaleType/fmtScale: int → ROCDL enum attributes
  • sched_barrier/sched_group_barrier: int mask → #rocdl<sched_group_mask ...> attribute
  • buffer_load/buffer_store/buffer_load_lds: aux from positional Value → keyword Attribute
  • buffer_atomic_fadd/buffer_atomic_fmax: added leading res type, aux as keyword Attribute
  • tensor_load_to_lds/tensor_store_from_lds: cache_policy as keyword Attribute
  • global_prefetch: scope from positional → cache_policy keyword

Compiler/backend

  • gpu.func kernel attribute: accept both "kernel" and "gpu.kernel"
  • waves_per_eu: use native rocdl.waves_per_eu attribute instead of passthrough string

Build infrastructure

  • llvm-hash.txt → llvm-build-info.json

Known regressions from LLVM base update (to investigate separately)

  • test_conv3d_implicit_fp8[1-96-4-8-9-96-1-1]: rel_err 0.31 (threshold 0.05)
  • test_rmsnorm_dynamicquant: quant diff 116 on M=64, N=2000, f32

Test Plan

Test Result

Submission Checklist

@jli-melchior jli-melchior changed the title [llvm] bump up llvm for upstream and internal coexec opt and adapt flydsl api [wip][llvm][flydsl] bump up llvm for upstream and internal coexec opt and adapt flydsl api Jul 31, 2026
@jli-melchior jli-melchior changed the title [wip][llvm][flydsl] bump up llvm for upstream and internal coexec opt and adapt flydsl api [WIP][LLVM][Flydsl] bump up llvm for upstream and internal coexec opt and adapt flydsl api Jul 31, 2026
@Boss2002n
Boss2002n force-pushed the jli/update-llvm-basemain branch 3 times, most recently from 0694bb6 to fc5250c Compare August 5, 2026 05:46
@coderfeli

Copy link
Copy Markdown
Collaborator

@jli-melchior still failed.

@Boss2002n
Boss2002n force-pushed the jli/update-llvm-basemain branch from 58cc8b2 to 9bb26a4 Compare August 5, 2026 07:28
@jli-melchior jli-melchior changed the title [WIP][LLVM][Flydsl] bump up llvm for upstream and internal coexec opt and adapt flydsl api [WIP][LLVM][FlyDSL] Bump up LLVM and Adapt FlyDSL to upstream LLVM/MLIR API changes Aug 5, 2026
@Boss2002n
Boss2002n force-pushed the jli/update-llvm-basemain branch from 9d6415a to 966c4b0 Compare August 6, 2026 02:44
@jli-melchior jli-melchior changed the title [WIP][LLVM][FlyDSL] Bump up LLVM and Adapt FlyDSL to upstream LLVM/MLIR API changes [LLVM][FlyDSL] Bump up LLVM and Adapt FlyDSL to upstream LLVM/MLIR API changes Aug 7, 2026
@Boss2002n
Boss2002n force-pushed the jli/update-llvm-basemain branch from 2fe5e75 to 67b0bda Compare August 7, 2026 02:19
@coderfeli
coderfeli requested a review from sjfeng1999 August 7, 2026 06:41
@Boss2002n
Boss2002n force-pushed the jli/update-llvm-basemain branch from 9e029b6 to 69523e6 Compare August 7, 2026 07:38
Comment thread python/flydsl/expr/rocdl/__init__.py Resolved
Comment thread python/flydsl/expr/rocdl/__init__.py Resolved
jli-melchior and others added 15 commits August 8, 2026 09:01
Update build scripts to read repo/commit from internal-llvm-build-info.json
and tag build/install directories with the commit short hash. Fix C++ API
breakage introduced by the internal LLVM:

- OpaqueProperties → PropertyRef in inferReturnTypes
- Buffer op aux parameter from Value → Attribute
- Atomic buffer ops now require explicit result Type
- MFMA blgp parameter from IntegerAttr → MFMAPermB enum
- MFMA scale cbsz/blgp from IntegerAttr → MatrixFormat enum
- WMMA modC from raw uint16_t → WMMACModifier enum
- WMMA scale params from raw casts → typed ROCDL enums
- TDM cachePolicy from uint32_t → Attribute
- DenseMapInfo: remove sentinel keys (tombstone-free hashing)

Co-Authored-By: Claude <noreply@anthropic.com>
Consolidate internal-llvm-build-info.json and llvm-hash.txt into a
single thirdparty/llvm-build-info.json with "upstream" and "internal"
entries. Build scripts select by key name:
- build_llvm.sh reads ['upstream']
- build_internal_llvm.sh / build_internal_flydsl.sh read ['internal']

Update CI workflows (build-whl, flydsl, build-custom-llvm-tools) to
use llvm-build-info.json for cache keys and commit lookups.

Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
The internal LLVM update changed several ROCDL ops: operands became
keyword-only attributes (aux, mask, cache_policy) and WMMA sign params
were removed. Add compatibility wrappers in expr/rocdl so downstream
callers (e.g. aiter) continue to work without source changes.

- raw_ptr_buffer_load: accept aux as positional, convert Value/int → Attr
- raw_ptr_buffer_load_lds: same aux handling
- sched_barrier / sched_group_barrier: convert int mask → SchedGroupMask enum attr
- tensor_load_to_lds / tensor_store_from_lds: accept cache_policy as positional
- tdm_ops: pass cache_policy as keyword to underlying ROCDL ops

Co-Authored-By: Claude <noreply@anthropic.com>
The internal LLVM update changed WMMA op parameters (fmtA, fmtB, modC,
scaleAType, scaleBType, fmtScaleA, fmtScaleB) from plain integers to
typed ROCDL enum attributes. Add int-to-enum conversion in the existing
wmma_scale_f32_16x16x128_f8f6f4, wmma_scale_f32_32x16x128_f4, and
wmma_f32_16x16x128_fp8_fp8 wrappers so downstream callers passing ints
continue to work.

Co-Authored-By: Claude <noreply@anthropic.com>
RawPtrBufferLoadOp/StoreOp aux changed from a positional Value operand
to a keyword-only Attribute in the internal LLVM update. Convert the
cache_modifier int to IntegerAttr and pass as keyword.

Co-Authored-By: Claude <noreply@anthropic.com>
global_prefetch's second param changed from positional `scope` to
keyword-only `cache_policy` in the LLVM binding regeneration.
l2_prefetch_tile was the only call site missed — fix it to use
cache_policy=_cache_policy_attr(scope), consistent with the
tensor_load_to_lds/tensor_store_from_lds fixes in the same file.

Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
The upstream LLVM changed blgp from I32Attr to ROCDL_MFMAPermBAttr and
cbsz/blgp in mfma_scale from I32Attr to ROCDL_MatrixFormatAttr. The
AttrBuilder is not registered for these enum types, so passing a plain
int causes std::bad_cast at op creation time.

Add _blgp_attr() to convert int blgp values to #rocdl<mfma_perm_b ...>
attributes in _split_mfma_operands, and reuse _wmma_fmt() for the
mfma_scale cbsz/blgp (ROCDL_MatrixFormatAttr).

Co-Authored-By: Claude <noreply@anthropic.com>
…tribute

The upstream LLVM added a leading `res` (result type) parameter to
raw_ptr_buffer_atomic_fadd/fmax and changed `aux` from positional Value
to keyword-only Attribute. The old wrapper passed vdata as the first
arg, which was misinterpreted as the result type causing std::bad_cast.

Infer the result type from vdata.type, convert int aux to IntegerAttr,
and pass aux as keyword.

Co-Authored-By: Claude <noreply@anthropic.com>
Upstream LLVM renamed the gpu.func kernel marker attribute from
"gpu.kernel" to "kernel". The _iter_gpu_kernel_funcs helper failed to
match any kernel functions, so lower_compile_hints silently skipped the
waves_per_eu passthrough override.

Accept both "kernel" and "gpu.kernel" for backward compatibility.

Co-Authored-By: Claude <noreply@anthropic.com>
Upstream LLVM now lowers rocdl.waves_per_eu directly to the LLVM IR
"amdgpu-waves-per-eu" function attribute. The old approach of deleting
the native attribute and adding a passthrough string no longer works.

Set rocdl.waves_per_eu directly with the override value. Update tests
to match the new single-value format ("2" instead of "2,2").

Co-Authored-By: Claude <noreply@anthropic.com>
Upstream LLVM/MLIR changed rocdl.mfma.scale to print cbsz and blgp as
enum keywords (fp8_e4m3, fp8_e5m2, fp4_e2m1) instead of raw integers.
Update the three CHECK lines in mma_atom_stateful.mlir to match.
jli-melchior and others added 4 commits August 8, 2026 09:01
- mma_scale_gfx1250.mlir: wmma.scale/scale16 ops now print named enum
  keyword attrs (fmtA, fmtB, modC, scaleAType, ...) instead of integer
  dict attrs; modC=1 prints as "neg" instead of "1 : i16".
- tdm_gfx1250.mlir: tensor.load.to.lds / tensor.store.from.lds changed
  cachepolicy from keyword syntax ("cachepolicy 0") to positional (", 0").
The upstream LLVM narrowed scaleAType/scaleBType to a 2-value enum
(WMMAMatrixScale: row0=0, row1=1). The old verifier accepted [0,3] but
the static_cast to the enum produces malformed IR for values 2 and 3.

Tighten the verifier range to [0,1] and fix the fp4_opsel FileCheck
test to use opselB=1 (row1) instead of opselB=2 (out-of-range).

No shipping kernels are affected — all gfx1250 WMMAScale kernels
default opsel to 0.

Co-Authored-By: Claude <noreply@anthropic.com>
Update test_wmma_scale_type_roundtrip to use opsel_b=1 instead of
opsel_b=2, matching the tightened verifier range [0,1].

Co-Authored-By: Claude <noreply@anthropic.com>
@Boss2002n
Boss2002n force-pushed the jli/update-llvm-basemain branch from 84f3026 to c212abb Compare August 8, 2026 09:01
mfma_scale_f32_16x16x128_f8f6f4 gained a wrapper converting cbsz/blgp to
ROCDL MatrixFormat enum attributes, but the 32x32x64 variant did not. It
is the only rocdl.mfma* op called from kernels/ without one.

kernels/attention/flash_attn_utils.py passes plain ints 0, 0 for those
two operands. Upstream emits no Python enum bindings for the rocdl
dialect (mlir/python/CMakeLists.txt declares it without
GEN_ENUM_BINDINGS, unlike the NVVM block right above it), and
ROCDL_I32Enum derives from I32Enum rather than EnumAttrInfo, so
EnumInfo::asEnumAttr() is nullopt and the int-accepting builder path in
EnumPythonBindingGen never fires either. With no builder registered,
OpPythonBindingGen passes the raw int straight through and
Operation.create rejects it:

  TypeError: Invalid attribute value for the key "cbsz" (std::bad_cast)

The path is live, not dead code: instrumenting a gfx950 fp8 run of
tests/kernels/test_flash_attn_fwd.py shows _mfma_acc_fp8_wide traced 16
times, reached via flash_attn_fp8_gfx950.py.

Factor the shared body into _mfma_scale_call so every scale variant
routes cbsz/blgp through _wmma_fmt, and move the call site to the
list-operand form its sibling already uses.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@coderfeli
coderfeli merged commit 494b2cb into main Aug 9, 2026
14 of 15 checks passed
@coderfeli
coderfeli deleted the jli/update-llvm-basemain branch August 9, 2026 04:01
coderfeli pushed a commit that referenced this pull request Aug 9, 2026
Phase 1 probe (/tmp/w4a8_probe) proved the recorded A-operand blocker is GONE
on the current toolchain (post-#945): feeding a computed nibble-unpacked value
into the MFMA A-fragment lowers cleanly and is bit-exact. So W4A8 keeps the
new pipeline's B-first layout (weight = MFMA A-operand) and unpacks in-register.

This commit adds the int4 weight views + unpack helper (layout_helpers.py) and
wires gemm1's weight load. Numerics not yet correct: the packed 64b tiled_copy
splits the K/ki dims differently from int8's 128b load; fixing the packed load
layout next.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
coderfeli pushed a commit that referenced this pull request Aug 9, 2026
Phase 1 (probe scripts in /tmp/w4a8_probe, not committed) conclusively shows the
recorded W4A8 A-operand blocker is GONE on the current toolchain (post-#945):
feeding a COMPUTED nibble-unpacked int4->int8 value into the MFMA A-fragment
(tiled_mma.make_fragment_A + .store) lowers cleanly AND is bit-exact vs an
int4xint8 reference (probe_correct.py, maxdiff 0). So no B-operand restructuring
is needed for the compiler; W4A8 can keep the B-first weight=A layout and unpack
in-register.

Phase 2 is INCOMPLETE and int4 is gated OFF. Exhaustive probing shows the
packed-int4 weight cannot be driven through the pipeline's generic
flat_divide + make_tiled_copy_A in the int8 MFMA A-fragment's ki-separated
kpack layout: a half-width packed load always addresses K contiguously, whereas
int8's tiled_copy reads ki0/ki1 from kpack groups 1024 int8 apart. Groups 0,1
of the fragment come out bit-exact but groups 2,3 read the wrong K region, and
no static reorder fixes it (the values differ, not just their order). The
correct fix is an explicit ki-correct preshuffle-address loader mirroring the
legacy load_b_pack_k32; that is the remaining work.

This commit lands only the verified/reusable parts and keeps every existing
dtype byte-identical:
- gemm1 builder body is UNCHANGED vs 2c714ba (verified by diff); only the
  compile_moe_gemm1 arg validation gains an 'int4' branch that raises
  NotImplementedError with the diagnosis, plus a docstring note.
- layout_helpers gains make_gateup_weight_view_int4 / make_weight_view_int4
  (int8's view with byte strides halved, element_num=8) and
  unpack_int4_weight_frag (the 7-op even/odd sign-extend unpack), documented as
  scaffolding for the follow-up loader.
- gemm2 and all tests untouched.

Verified: 47/47 existing tests (test_moe_gemm_2stage + test_moe_reduce) pass;
int4 fails fast; black+ruff clean.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
coderfeli pushed a commit that referenced this pull request Aug 9, 2026
Brings the branch up to date with main (10 commits), including 494b2cb
(#945, LLVM/MLIR bump) which the W4A8 A-fragment unblock depends on.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
JohnQin-llm added a commit that referenced this pull request Aug 19, 2026
#919 made `_init_dualwave_thread_mapping` re-derive (head, q_block) from the
linear workgroup id with the head as the slow axis, so a head's blocks stay on
one XCD and its K/V stays resident in that XCD's L2 across the concurrent
wavefronts. It applied to every qualifying dualwave dense shape, bf16 included,
gated only on `not SPLITK and not CAUSAL and NUM_HEADS_Q % NUM_XCD_GFX950 == 0`.

#950 put a trait in front of that condition:

    -    if const_expr(not traits.SPLITK and not traits.CAUSAL and ...):
    +    if const_expr(traits.XCD_SWIZZLE and not traits.SPLITK and ...):

`XCD_SWIZZLE` defaults to False and only the fp8 builder passes it. That PR
measured and described fp8 only, and never touched the bf16 builder -- the
string `xcd` does not appear in flash_attn_gfx950.py at all -- so the remap
silently stopped applying to bf16. No test covered the mapping, so nothing
caught it.

Bisecting upstream on MI355X, B=2 H=32 D=128 bf16 non-causal, TFLOP/s at
S = 61,380 / 180,180 / 239,580, aiter measured in the same process as control:

    cf1db2f  7/31  #919     1273.2  1287.0  1287.4
    772f070  8/6   #974     1253.7  1074.5  1053.4
    494b2cb  8/9   #945     1251.9  1073.3  1052.6
    f6d0a41  8/14  #991     1247.2  1076.2  1061.2
    47009a3  8/14  #1004    1248.2  1076.2  1061.3
    daee90e  8/17  #930     1249.0  1076.4  1060.6
    96af91b  8/17  #1014    1247.2  1076.6  1064.5
    4521c8c  8/18  #960     1247.8  1074.8  1061.2
    aiter                   ~1250   ~1253   ~1252

One step, between 7/31 and 8/6, and #950 (8/3) is the only commit in that
window that touches the mapping. Everything after it is flat to within noise.

This plumbs `xcd_swizzle` through the bf16 builder and reuses the existing
dispatch condition (non-causal, `H % NUM_XCD_GFX950 == 0`, `num_q_blocks >=
MIN_Q_BLOCKS_XCD_SWIZZLE`). Split-K is excluded structurally: it takes the
`_build_splitk` branch above this dispatch site.

The remap is bijective, so output is bit-identical. Two tests assert that over
H=8/16/32/64 and cover the `H % 8 != 0` fallback both forced and auto -- the
guard #919 shipped without.

One loose end: the restored figure is still ~1.4% below the 7/31 measurement
(1268.8 against 1287.4 at S=239,580). The candidates are the
`MIN_Q_BLOCKS_XCD_SWIZZLE` threshold #950 introduced for fp8, which #919 did
not have, and the LLVM bump in #945.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
JohnQin-llm added a commit that referenced this pull request Aug 19, 2026
#919 made `_init_dualwave_thread_mapping` re-derive (head, q_block) from the
linear workgroup id with the head as the slow axis, so a head's blocks stay on
one XCD and its K/V stays resident in that XCD's L2 across the concurrent
wavefronts. It applied to every qualifying dualwave dense shape, bf16 included,
gated only on `not SPLITK and not CAUSAL and NUM_HEADS_Q % NUM_XCD_GFX950 == 0`.

#950 put a trait in front of that condition:

    -    if const_expr(not traits.SPLITK and not traits.CAUSAL and ...):
    +    if const_expr(traits.XCD_SWIZZLE and not traits.SPLITK and ...):

`XCD_SWIZZLE` defaults to False and only the fp8 builder passes it. That PR
measured and described fp8 only, and never touched the bf16 builder -- the
string `xcd` does not appear in flash_attn_gfx950.py at all -- so the remap
silently stopped applying to bf16. No test covered the mapping, so nothing
caught it.

Bisecting upstream on MI355X, B=2 H=32 D=128 bf16 non-causal, cache cold, TFLOP/s
at S = 61,380 / 180,180 / 239,580, aiter in the same process as control:

    cf1db2f  7/31  #919     1284.6  1288.2  1286.8
    772f070  8/6   #974     1256.7  1188.1  1184.3
    494b2cb  8/9   #945     1255.6  1187.9  1187.2
    f6d0a41  8/14  #991     1256.0  1182.8  1183.6
    47009a3  8/14  #1004    1260.5  1188.6  1183.4
    daee90e  8/17  #930     1254.1  1077.0  1060.9
    96af91b  8/17  #1014    1251.8  1074.9  1060.6
    4521c8c  8/18  #960     1249.3  1073.9  1059.0
    aiter                   ~1252   ~1255   ~1255

Two independent steps. The first, between 7/31 and 8/6 (-8.0%), is #950 and is
what this commit fixes. The second, between 47009a3 and daee90e (-10.4%), is
#930 dropping the fastmath flag when it moved these kernels onto fx.maxnumf,
which reaches the raw MLIR builder without passing through resolve_fastmath;
that is a separate fix, filed separately.

This plumbs `xcd_swizzle` through the bf16 builder and reuses the existing
dispatch condition (non-causal, `H % NUM_XCD_GFX950 == 0`, `num_q_blocks >=
MIN_Q_BLOCKS_XCD_SWIZZLE`). Split-K is excluded structurally: it takes the
`_build_splitk` branch above this dispatch site.

The remap is bijective, so output is bit-identical. Two tests assert that over
H=8/16/32/64 and cover the `H % 8 != 0` fallback both forced and auto -- the
guard #919 shipped without.

One loose end: the restored figure is 1268.3 against cf1db2f's 1286.8, a 1.4%
gap. That is #930's dropped fastmath, which costs 10.4% with the remap off but
only 1.4% with it on -- the two interact. Fixing #930 separately closes it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
JohnQin-llm added a commit that referenced this pull request Aug 20, 2026
#919 made `_init_dualwave_thread_mapping` re-derive (head, q_block) from the
linear workgroup id with the head as the slow axis, so a head's blocks stay on
one XCD and its K/V stays resident in that XCD's L2 across the concurrent
wavefronts. It applied to every qualifying dualwave dense shape, bf16 included,
gated only on `not SPLITK and not CAUSAL and NUM_HEADS_Q % NUM_XCD_GFX950 == 0`.

#950 put a trait in front of that condition:

    -    if const_expr(not traits.SPLITK and not traits.CAUSAL and ...):
    +    if const_expr(traits.XCD_SWIZZLE and not traits.SPLITK and ...):

`XCD_SWIZZLE` defaults to False and only the fp8 builder passes it. That PR
measured and described fp8 only, and never touched the bf16 builder -- the
string `xcd` does not appear in flash_attn_gfx950.py at all -- so the remap
silently stopped applying to bf16. No test covered the mapping, so nothing
caught it.

Bisecting upstream on MI355X, B=2 H=32 D=128 bf16 non-causal, cache cold, TFLOP/s
at S = 61,380 / 180,180 / 239,580, aiter in the same process as control:

    cf1db2f  7/31  #919     1284.6  1288.2  1286.8
    772f070  8/6   #974     1256.7  1188.1  1184.3
    494b2cb  8/9   #945     1255.6  1187.9  1187.2
    f6d0a41  8/14  #991     1256.0  1182.8  1183.6
    47009a3  8/14  #1004    1260.5  1188.6  1183.4
    daee90e  8/17  #930     1254.1  1077.0  1060.9
    96af91b  8/17  #1014    1251.8  1074.9  1060.6
    4521c8c  8/18  #960     1249.3  1073.9  1059.0
    aiter                   ~1252   ~1255   ~1255

Two independent steps. The first, between 7/31 and 8/6 (-8.0%), is #950 and is
what this commit fixes. The second, between 47009a3 and daee90e (-10.4%), is
#930 dropping the fastmath flag when it moved these kernels onto fx.maxnumf,
which reaches the raw MLIR builder without passing through resolve_fastmath;
that is a separate fix, filed separately.

This plumbs `xcd_swizzle` through the bf16 builder and reuses the existing
dispatch condition (non-causal, `H % NUM_XCD_GFX950 == 0`, `num_q_blocks >=
MIN_Q_BLOCKS_XCD_SWIZZLE`). Split-K is excluded structurally: it takes the
`_build_splitk` branch above this dispatch site.

The remap is bijective, so output is bit-identical. Two tests assert that over
H=8/16/32/64 and cover the `H % 8 != 0` fallback both forced and auto -- the
guard #919 shipped without.

One loose end: the restored figure is 1268.3 against cf1db2f's 1286.8, a 1.4%
gap. That is #930's dropped fastmath, which costs 10.4% with the remap off but
only 1.4% with it on -- the two interact. Fixing #930 separately closes it.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
coderfeli pushed a commit that referenced this pull request Aug 20, 2026
#919 made `_init_dualwave_thread_mapping` re-derive (head, q_block) from the
linear workgroup id with the head as the slow axis, so a head's blocks stay on
one XCD and its K/V stays resident in that XCD's L2 across the concurrent
wavefronts. It applied to every qualifying dualwave dense shape, bf16 included,
gated only on `not SPLITK and not CAUSAL and NUM_HEADS_Q % NUM_XCD_GFX950 == 0`.

#950 put a trait in front of that condition:

    -    if const_expr(not traits.SPLITK and not traits.CAUSAL and ...):
    +    if const_expr(traits.XCD_SWIZZLE and not traits.SPLITK and ...):

`XCD_SWIZZLE` defaults to False and only the fp8 builder passes it. That PR
measured and described fp8 only, and never touched the bf16 builder -- the
string `xcd` does not appear in flash_attn_gfx950.py at all -- so the remap
silently stopped applying to bf16. No test covered the mapping, so nothing
caught it.

Bisecting upstream on MI355X, B=2 H=32 D=128 bf16 non-causal, cache cold, TFLOP/s
at S = 61,380 / 180,180 / 239,580, aiter in the same process as control:

    cf1db2f  7/31  #919     1284.6  1288.2  1286.8
    772f070  8/6   #974     1256.7  1188.1  1184.3
    494b2cb  8/9   #945     1255.6  1187.9  1187.2
    f6d0a41  8/14  #991     1256.0  1182.8  1183.6
    47009a3  8/14  #1004    1260.5  1188.6  1183.4
    daee90e  8/17  #930     1254.1  1077.0  1060.9
    96af91b  8/17  #1014    1251.8  1074.9  1060.6
    4521c8c  8/18  #960     1249.3  1073.9  1059.0
    aiter                   ~1252   ~1255   ~1255

Two independent steps. The first, between 7/31 and 8/6 (-8.0%), is #950 and is
what this commit fixes. The second, between 47009a3 and daee90e (-10.4%), is
#930 dropping the fastmath flag when it moved these kernels onto fx.maxnumf,
which reaches the raw MLIR builder without passing through resolve_fastmath;
that is a separate fix, filed separately.

This plumbs `xcd_swizzle` through the bf16 builder and reuses the existing
dispatch condition (non-causal, `H % NUM_XCD_GFX950 == 0`, `num_q_blocks >=
MIN_Q_BLOCKS_XCD_SWIZZLE`). Split-K is excluded structurally: it takes the
`_build_splitk` branch above this dispatch site.

The remap is bijective, so output is bit-identical. Two tests assert that over
H=8/16/32/64 and cover the `H % 8 != 0` fallback both forced and auto -- the
guard #919 shipped without.

One loose end: the restored figure is 1268.3 against cf1db2f's 1286.8, a 1.4%
gap. That is #930's dropped fastmath, which costs 10.4% with the remap off but
only 1.4% with it on -- the two interact. Fixing #930 separately closes it.

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
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.

3 participants