Repository navigation
[FlyDSL] gfx950 MXFP8 bmm: grouped XCD tile order, split-K LDS flag only when split - #6440
andrei-kochin wants to merge 1 commit into
Conversation
…nly when split - xcd_swizzle (_xcd<g> in kernel names): inside each XCD's contiguous run of tiles (xcd_order), re-order tiles M-major in groups of g M tiles, as xcd_remap_bx_by does. Fewer B column blocks are live per XCD, so L2 reuse improves. The b_direct heuristic now picks _xcd4 instead of _xcd. Existing _xcd CSV rows parse and run unchanged; xcd_swizzle without xcd_order is rejected, since the digits would parse back as _sk. - LDS: reserve the 16-byte split-K flag only when K is actually split, which lets a few large tiles fit. - Tests: _xcd4 accuracy with 32x32 scales (including a partial swizzle group and a padded grid), kernel-name round trip, and the invalid xcd_swizzle-without-xcd_order name. Split out of ROCm#6060. Signed-off-by: Sergei Shliapnikov <sergei.shliapnikov@amd.com> Signed-off-by: Andrei Kochin <andrei.kochin@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
There was a problem hiding this comment.
🟡 Changes recommended
Split-K configurations accept _xcd<g> but silently bypass the grouped remapping, and the LDS boundary change lacks regression coverage.
2 open findings
What changed in this PR
Adds grouped XCD tile ordering to improve gfx950 MXFP8 BMM cache reuse while reducing non-split LDS usage.
Changes:
- Adds
_xcd<g>parsing, validation, and grouped tile remapping. - Reserves the LDS split-K flag only for split launches.
- Adds swizzle accuracy and kernel-name tests.
| File | Description |
|---|---|
aiter/ops/flydsl/kernels/bmm_a8w8_mxscale_gfx950.py |
Implements grouped XCD remapping and conditional LDS allocation. |
aiter/ops/flydsl/batched_gemm_a8w8_gfx950.py |
Extends kernel naming and heuristic selection with _xcd4. |
op_tests/test_gemm_group32_interface.py |
Tests swizzle correctness, parsing, and validation. |
🧠 Review effort: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| tile = bid_x % 8 * per_xcd + bid_x // 8 | ||
| rest = tile // (N // tile_n) | ||
| bid_x, bid_y, bid_z = rest % m_tiles, tile % (N // tile_n), rest // m_tiles | ||
| if const_expr(xcd_swizzle > 1): |
| c_rows, | ||
| flag_off, | ||
| flag_off + 16, | ||
| flag_off + (16 if split else 0), |
|
Opened by mistake against the wrong PR; closing. |
Motivation
Part 2 of 3 split out of #6060 (author: @sshlyapn), so each PR covers one concern.
The gfx950 FlyDSL MXFP8 bmm's XCD tile order (
_xcd, #5896) gives each XCD a contiguous run of tiles, but within that run tiles are N-major, so many B column blocks are live per XCD at once. This PR re-orders the run in M-major groups to improve L2 reuse, and trims an LDS reservation that kept a few large tiles from fitting. It is independent of the 1x32 weight-scale change.Changes
xcd_swizzle(_xcd<g>in kernel names): inside each XCD's contiguous run of tiles(
xcd_order), tiles are re-ordered M-major in groups ofgM tiles, asxcd_remap_bx_bydoes. Fewer B column blocks are live per XCD, so L2 reuse improves.b_directconfigs now picks_xcd4instead of_xcd._xcdCSV rows parse and run unchanged.xcd_swizzlewithoutxcd_orderis rejected by bothcheck_bmm_configandbmm_kernel_name(the digits would otherwise parse back as part of_sk<splits>).few large tiles fit.
test_gemm_group32_interface.py):_xcd4accuracy with 32x32 scales against a dequantized FP32 reference, including a partial swizzle group (5 M tiles) and a padded grid (12 N tiles)._xcd/_xcd4, and rejection of a swizzle without_xcd.Related PRs
main.main).Performance (MI355, gfx950, 256 CU), from #6060
Effect of the tile order alone: 15 tuned
_xcd4rows were re-timed with the same config,changing only the tile order to
_xcd(#5896's order) or to no XCD order._xcd4vs_xcd: never slower, 0.3-27% faster._xcd4vs no XCD order: about equal at M <= 16k (-1% to +5%, one +20% case),5-26% faster at M >= 32k.
Protocol: weights rotate between 2 copies; each point is the median of 3 interleaved rounds; every timed config was checked against a dequantized fp32 reference (rel err < 3e-3).
Tests
The kernel code is the XCD/LDS subset of #6060, unchanged. The
_xcd4configs exercised here were validated in #6060 with 1x32 scales; the 32x32 variants of those cases are new in this PR and run in CI on gfx950. black: clean.