[FlyDSL] jagged_dense_bmm_broadcast_add (jdbba)#4042
Closed
anhminhnguyenhoang wants to merge 49 commits into
Closed
[FlyDSL] jagged_dense_bmm_broadcast_add (jdbba)#4042anhminhnguyenhoang wants to merge 49 commits into
anhminhnguyenhoang wants to merge 49 commits into
Conversation
Layout-API prototype for HSTU jagged_dense_bmm_broadcast_add on AMD CDNA: per group b over packed rows [s,e), Out[s:e]=Jagged[s:e]@Dense[b]+Bias[b]. N==K==128 fixed; bf16 in/out, fp32 accumulate; plain (non-preshuffled) B. Establishes the jagged contract: group axis on the grid, device-resident seq_offsets read, per-group rebasing, runtime early-exit, broadcast bias, and a bounded-C-descriptor masked store. seq_start/seq_end are scalarized (readfirstlane) to keep the C buffer descriptor uniform, which collapses an otherwise per-lane epilogue-store exec-mask waterfall. Includes a toy-shape correctness + benchmark harness with a --device-time mode (rocprofv3); the default CUDA-event wall-clock is ~90% host launch overhead at these tiny shapes, so device time is the metric to trust. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
The prototype hardcoded N==K==128 as compile-time constants and could not
run the production HSTU shapes (N=16384, reduction K/output N = 256/512).
Make N (output) and K (reduction) runtime-parametric via a memoized factory
_build_launcher() that bakes N/K/tiling as closure constants per shape
(mirrors splitk_hgemm.compile_hgemm_kernel). The public entry derives N, K
from the tall dense matrix shape (B_groups*N, K). Epilogue vector widths are
derived from the tile (C_FRAG_LEN = BLOCK_M*BLOCK_N/THREADS) instead of being
hardcoded, so a tiling change does not need hand edits.
Validated cos=1.0 at N=K in {128, 256, 512} plus partial-tile edges.
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
The per-group row-base offsets seq_start*K and seq_start*N were computed in i32. With seq_start reaching ~L (millions of packed rows), B1024_D512 hits 7.86M*512 = 4.0e9 > 2^31 and the kernel takes a GPU memory-access fault. Only that headline cell overflows; the others stay under 2^31, which is why they ran. Cast seq_start/off_b to i64 before the stride multiply so the product is 64-bit. Correctness-only; no perf cost. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
The MFMA C accumulator is M-major per lane (a lane's contiguous fragment elements map to different output rows, stride N in global), so a vectorized N-contiguous store straight from the fragment is impossible -- the kernel emitted 64 scalar buffer_store_short per thread. A store-deletion probe showed that store alone costs 38-62% of runtime at these memory-bound shapes. Route C through LDS to transpose the layout: fold bias into the bf16 fragment, write it to a row-major shared C tile in the natural MFMA layout (reusing the A-staging LDS, smem sized to max of the two -- no extra smem), barrier, then re-read N-contiguous (8 bf16/thread) and store buffer_store_dwordx4. ISA store histogram: 64x buffer_store_short -> 8x buffer_store_dwordx4. Device time (MI355X, rocprofv3, M_i=7680), controlled interleaved sweep: B120_D256 335.6 -> 302.0 us (1.11x) B120_D512 901.2 -> 847.6 us (1.06x) B1024_D256 2563.9 -> 2210.6 us (1.16x) B1024_D512 7286.1 -> 6842.3 us (1.07x) cos=1.0 on all 4 headline shapes + partial/empty edges. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
block_k = 128 if reduction K <= 256 else 64, chosen in the public entry. For K=256 a 2-iter K-loop (BLOCK_K=128) has fewer barriers and is faster; for K=512 the deeper BLOCK_K=64 pipeline keeps occupancy. BLOCK_K=256 is unsafe -- the 2-stage double-buffer epilogue mis-accumulates a single K-tile (cosine masks a 123% relative error), so it is excluded. Device time (MI355X, rocprofv3): B1024_D256 2563 -> ~2294 us, B120_D256 ~312 us. K=512 cells unchanged (stay on BLOCK_K=64). cos=1.0 all shapes. Final vs upstream Triton: B120_D256 1.13x, B120_D512 1.26x, B1024_D256 1.31x, B1024_D512 1.29x. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Add the rocprofv3-based device-time harness used to drive the optimization: - bench_headline_worker.py: builds one headline cell (args B D Kout Mi) and fires warmup+30 launches of flydsl or triton, for profiling under rocprofv3. - read_us2.py: parses the rocprofv3 sqlite (kernel_dispatch x kernel_symbol) and prints the median per-kernel device duration in us. Docs: - jagged_dense_bmm_optimization_report.md: results (FlyDSL beats Triton 1.13-1.31x on all 4 headline shapes), the (B,D,K,N)->GEMM mapping, and step-by-step reproduction commands for every experiment. - jagged_dense_bmm_benchmark_repro.md: toy-shape benchmark + full optimization log (winners and verified dead ends). Method note captured in the docs: CUDA-event wall-clock is ~90% host launch overhead at these shapes and must not be used for kernel-level conclusions; rocprofv3 device time is the metric throughout. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…512) The kernel is memory-bandwidth-bound: PMC on B1024_D512 showed L2 hit 49.3% and DRAM read traffic ~44x the model minimum -- each group's Dense[b] weight matrix was being re-fetched from HBM by its M-tiles, which the hardware scatters round-robin across all 8 XCDs so no single XCD's private L2 retains it. Add the HipKittens Algorithm 1 launch-time block-ID remap (_xcd_remap): invert the round-robin so a group's M-tiles co-locate on one XCD, sharing that XCD's L2 copy of Dense[b]. Two knobs baked compile-time per shape: W (window height, =8, a flat optimum) and C (XCD chunk size), which is weight-size dependent -- XCD_C_SMALL_K=32 for reduction K<=256, XCD_C_LARGE_K=60 for K>256. Small K regresses with large C (the L2-greedy / shared-LLC-starvation trap). Measured B1024_D512: L2 hit 49->76%, DRAM read requests -67%, device time -5.3%. Verified cos=1.0 and win-or-neutral across all four headline shapes. The remap is an exact permutation of the tile grid (tail-safe for any C/W). An occupancy probe ruled out ping-pong / multi-wave interleave as alternatives: the kernel is occupancy-resource-limited (32KB LDS + 84 VGPR cap ~5 waves/SIMD, 2.8 achieved), not latency-bound, so traffic reduction (this lever) is the right fix, not latency hiding. jagged_dense_bmm_xcd.py is the standalone experiment clone; the test/sweep/PMC scripts under op_tests/flydsl_tests/ reproduce the correctness and the (W,C) sweep. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…tifact) The prior report claimed FlyDSL "beats Triton 1.13-1.31x". That was a measurement artifact: read_us2.py took the MEDIAN over all kernel dispatches, and the upstream Triton kernel autotunes -- firing hundreds to thousands of trial-config dispatches whose slow trials inflate its median by 30-40% (e.g. B1024_D512: 849 dispatches spread 6307-18792us, median 8789us). Comparing FlyDSL's clean median against Triton's autotune-polluted median manufactured a fake 1.3x. Fair, method-matched comparison uses p10/min (steady-state best config) for the autotuning kernel -- which is what do_bench and a real deployment see. Honest result (device time, M_i=7680, FlyDSL p10 vs Triton best-config min): rough PARITY, Triton's best config ~1-7% ahead (within noise on B1024_D512). The XCD remap closed the D=512 gap to ~1-3%. read_us2.py now defaults to p10 (was median) and documents min|median|all. The two repro docs are corrected with the honest table, the correction history, and the rule: always use p10/min, never median, for autotuning kernels. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Use v_mfma_f32_16x16x32_bf16 on gfx950 (twice the K per instruction vs the 16x16x16 atom), halving MFMA issue. Bit-exact, 4-7% faster across the headline shapes (larger on the compute-heavier D=512 cells). Auto-detected via arch (gfx95*), with gfx942/MI300 falling back to 16x16x16; overridable through the use_mfma_k32 kwarg. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…er, DTLA) Exp 5 (MFMA 16x16x32, FOLDED IN, +4-7% gfx950), Exp 6 (persistent problem-visitor scheduler, 1.25-1.9x on skew but needs on-device prefix), Exp 7 (DirectToLDS-for-A, regression). Update the overall-decision and files sections accordingly. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…e GPU fault) The jagged A input is allocated exactly L rows, but A's buffer descriptor used max_size=True (unbounded, 4GB). The LAST group's partial bottom tile (M_b not a multiple of BLOCK_M) reads up to BLOCK_M-1 rows past the allocation end. With the unbounded descriptor those reads are not OOB-dropped, so when they land on an unmapped page the kernel GPU-faults -- a data-dependent intermittent crash that reproduces only for certain total-L values (e.g. B120_D512 with a skewed varlen M_i distribution, some seeds fault and others pass). Bound A to M_b*K*2 bytes exactly like C is already bounded, so the hardware OOB-drops the partial-tile rows (read as 0). This is harmless: the matching output rows are masked out of the bounded C store. Fault is independent of the XCD remap (reproduces with it disabled). Verified: all four previously-faulting skew seeds now pass cos=0.999999; uniform headline shapes unchanged (cos=1.0, device time neutral). Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…oyment) The XCD block-ID remap (+~5% on the uniform D=512 headline shapes) REGRESSES ~2% under a skewed/varlen M_i distribution: with the per-tile early-exit, clustering a group's M-tiles onto one XCD also clusters its skipped no-op tiles, causing XCD load imbalance, while the hardware round-robin spreads real work evenly (measured B1024_D512 skewed: C=1 identity ~3215us vs C=60 default ~3290us across seeds). Add a `uniform_seqlen` flag (default True) to the public entry. When the remap knobs are left as None it now picks: uniform -> tuned C (XCD_C_*), W=XCD_W; varlen -> C=1,W=1 which makes _xcd_remap the exact identity (remap disabled). An explicit xcd_c/xcd_w still overrides. Verified cos=1.0 on both paths. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…E_K=120 (a) Dispatch layer routes skewed/varlen workloads to the persistent problem-visitor kernel (on-device CUM prefix, zero host sync) when the problem is large enough to amortize launch (n_groups>=1024 or output_n>=512 and above a size floor); uniform and tiny shapes stay on the static-grid gen kernel. Measured 1.4-1.6x end-to-end on skewed D512 shapes; uniform unchanged. cos=1.0 on uniform and skew across all 4 headline shapes. (b) Bump XCD_C_LARGE_K 60->120 (autotuned): D512 shapes prefer the larger chiplet chunk (B1024_D512 6.01->5.96ms). Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Two remap-tuning refinements: 1. Bump XCD_C_LARGE_K 60->120 (autotuned: the D=512 C-curve is a flat plateau for C>=32 on the uniform grid; 120 is the measured best.) 2. Refine the uniform_seqlen gate. An earlier measurement (on a pre-fix clone) found the remap regressed under varlen and the gate disabled it for ALL non-uniform calls. Re-measured on the CURRENT kernel (after the bounded-A fix + 16x16x32 atom), the remap is ~2% FASTER under varlen on D=512 (C=60 ~3086us vs OFF ~3155us, 3 seeds) and only ~neutral on D=256. So the gate now keeps the remap ON for K>256 regardless of uniformity (the Dense[b] L2-reuse win holds under skew once the weight is large enough), and only falls back to identity (C=1,W=1) for K<=256 under skew. Verified cos=1.0 on all four gate paths incl. the empty/single/full edge combo that previously faulted. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
… fix Refresh the optimization report and benchmark repro to the current kernel (16x16x32 atom + XCD remap + bounded-A fix): - §2 rewritten with the two-method comparison (do_bench cold-L2 + rocprofv3 hot-L2, both autotune-excluded): FlyDSL wins 3 of 4 shapes (~5-6% ahead on both D=512 cells, tie-to-slightly-behind on B1024_D256 depending on L2 warmth). - Added method lessons #3 (cold vs hot L2) and #4 (skew finding reversed on the current kernel) and reproduction subsections 4.5 (do_bench) / 4.6 (skew remap). - §6 next-opportunities: XCD remap + gate marked done; fp8 weights flagged as the only decisive lever; confirmed dead-ends listed. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Make jagged_dense_bmm_dispatch_v2 the production entry the headline bench uses, instead of importing the gen kernel directly. dispatch_v2 picks the right kernel per regime (gen for uniform, persistent for large skew), so the bench now measures what production actually routes to. The headline configs are uniform, so the gate keeps them on the gen kernel -- numbers are unchanged (cos=1.0, B1024_D512 ~1.07x Triton), but the call path is now the dispatcher. Also fold in a comment correction in dispatch_v2: the persistent kernel's skew speedup is ~1.2-1.3x vs OUR OWN static-grid kernel, NOT vs Triton (Triton still wins skew 0.79-0.97x, validated 2026-06-09 via bench_jdbba_vs_triton.py / bench_gen_vs_triton_skew.py). The earlier "1.4-1.6x" wording implied a Triton win that does not exist. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…ow persist gate The XCD chiplet remap clusters a group's M-tiles onto one XCD for Dense[b] L2 reuse. That wins on uniform (every tile occupied -> chiplets balanced) but is a large LOSS under skew: with ~27% empty groups and uneven M_i it overloads a few chiplets while the rest idle. Measured (cold-L2 do_bench) remap-off is 1.4-1.6x faster than remap-on on D512 skew. Reverses the earlier "keep remap on for K>256 under skew" decision, whose "+2% on D512 varlen" basis was a hot-L2 artifact. Default gate is now: remap ON only when uniform_seqlen; identity round-robin for ALL K under skew. This moved skew vs Triton from 0.79-0.83x to 0.93-0.95x across all four headline shapes (do_bench, Mi=7680, cos=1.0). Narrowed the persist routing to where it actually wins: output_n<=256 AND n_groups>=1024 (B1024_D256), where occupied-tile-only traversal beats the static grid ~1.13x and ties Triton. D512 skew now falls through to the static remap-off path (faster than persist by 1.13-1.20x). Reading Triton's kernel confirmed it uses the same static-grid + bare early-exit algorithm we do -- the persistent scheduler solved a non-problem; the real lever was the remap default. Also tested and rejected on D512 skew: 3-stage pipeline (neutral, bandwidth-bound) and BLOCK_N=256 (regresses on gfx950). Updated jagged_dense_bmm_optimization_report.md to keep only the corrected truth. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…methodology docs) Snapshots the full per-lever experiment campaign run on MI355X/gfx950 so the work is recoverable, ahead of a working-tree cleanup to re-run the same campaign on MI300X/gfx942. Contents: - Per-lever kernel clones explored during optimization: async (buffer_load_lds), cstrip (C-shuffle LDS strips), dtla (direct-to-LDS A), mfma32 (16x16x32 atom), mreg (M register tiling), persist / persist_xcd (persistent problem-visitor), warps (warp-layout sweep), wc (W/C XCD autotune), wpe (waves_per_eu). - Worker / correctness / sweep harnesses for each lever. - Methodology + reference docs kept for reuse: dev journal, FlyDSL port plan, Triton kernel walkthrough. - Keeper test/bench wired to the production dispatch: bench_jdbba_vs_triton.py, test_jdbba_dispatch_v2.py. These per-lever clones are NOT production; the shipped kernels remain jagged_dense_bmm_gen.py (+ persist_dev) behind jagged_dense_bmm_dispatch_v2.py. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…for MI300X re-run Clears the working tree of the per-lever experiment artifacts and the MI355X/gfx950 result documents so the same campaign can be re-run cleanly on MI300X/gfx942 without carrying hardware-specific conclusions forward. Removed: - Per-lever kernel clones + their worker/correctness/sweep harnesses (archived in the preceding commit; recoverable from git history). - MI355X result docs: jagged_dense_bmm_optimization_report.md, jagged_dense_bmm_benchmark_repro.md, jagged_dense_bmm_followup_experiments.md (their numbers and "what worked" conclusions are gfx950-specific). - 772MB of rocprofv3 binary trace output dirs (regenerable raw data, never committed). Kept (hardware-neutral / production): - Shipped kernels: jagged_dense_bmm_gen.py, jagged_dense_bmm_persist_dev.py, jagged_dense_bmm.py, dispatch_v2. - Canonical bench: bench_jagged_dense_bmm_perf.py; cross-check bench_jdbba_vs_triton.py; test_jdbba_dispatch_v2.py. - Methodology/reference docs: dev journal, FlyDSL port plan, Triton walkthrough. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Captures the optimization campaign as hypotheses to re-test on MI300X (CDNA3), with NO gfx950 verdicts carried forward. Covers: the op + headline shapes + two regimes, gfx942 hardware deltas (16x16x16-only MFMA, 64KB LDS, XCD count), fixed methodology (correctness gate, cold/hot-L2 timing discipline, one-lever- at-a-time), the baseline to establish first, the 11 levers as open questions, and the two regime gates (XCD remap on/off; persistent vs static under skew) that must be re-derived per regime from gfx942 data. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Cleanup stragglers from the same experiment category as the previously-removed MI355X campaign (these were tracked from an earlier session, so recoverable from git history): - jagged_dense_bmm_xcd.py: the XCD-remap lever clone. The remap is folded into the production jagged_dense_bmm_gen.py; the clone is redundant. - test_xcd_correct.py / test_xcd_correct_allshapes.py / check_xcd_perm.py: harnesses for that clone. - bench_jagged_dense_bmm.py: superseded by bench_jagged_dense_bmm_perf.py. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Single-shape worker that only re-imported jagged_dense_bmm_gen; same experiment category as the per-lever bench_headline_worker_* scripts already removed, and nothing references it. Recoverable from git history. Kept: jagged_dense_bmm.py (N=K=128 prototype — the validated correctness reference cited in the MI300X plan and FlyDSL port plan for bring-up) and persist_dev.py (imported by the production dispatch_v2). Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…op bench_jdbba_vs_triton
The canonical perf bench now reports FlyDSL vs Triton side by side for BOTH
sequence-length regimes via --regime {uniform,skew,both} (default uniform), so
the separate bench_jdbba_vs_triton.py is redundant and is removed. Skew uses the
deployment distribution (M_i = Mi*U(0,1)**4, ~20% empty groups + one full + one
near-full); FLOPs/bytes use the actual packed length L so each regime is scored
on its true work. Routing is end-to-end through the production dispatch
(uniform_seqlen gate). Verified both regimes cos=1.0 on all 4 headline shapes.
Updated the dispatch comment and MI300X plan references to point at the perf bench.
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
The dispatch JSON was gfx950-only and forced use_mfma_k32=true, which core-dumps on gfx942 (no 16x16x32 bf16 atom). Restructure to arch-keyed-v1 (by_arch.gfx942 / by_arch.gfx950); the loader auto-selects the section matching the detected arch (overridable via FLYDSL_JAGGED_DENSE_BMM_ARCH), and still accepts a flat legacy JSON via the _DISPATCH_V2_JSON override. The committed by_arch.gfx942 section is the official MI300X baseline: empty winners -> D-bucketed heuristic with use_mfma_k32=false, xcd_c/xcd_w null (kernel shape-derived defaults). Reproducible pre-tuning configuration; the out-of-box bench passes with no core-dump. test_jdbba_dispatch_v2.py now uses the loaded arch-selected table. Docs (experiment plan + dev journal) updated with the arch-keyed schema and the correct run env (jdbba-flydsl container, /workspaces/meta/aiter). Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Add tune_jdbba_xcd.py (live-knob xcd_c x xcd_w sweep, cos-gated, do_bench cold-L2) and promote the 4 headline-shape gfx942 winners into by_arch.gfx942.winners. All cos=1.0; 3x-repeat confirmed each beats the kernel auto-default. Skew unaffected (remap off there). The marginal size confirms the bound check: this kernel is overhead-bound, not L2/chiplet- bound, so XCD remap has limited leverage on gfx942. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Drop the K<=256 -> BLOCK_K=128 special case. On gfx942 the larger A double-buffer (64KB at BLOCK_K=128) saturates the 64KB LDS ceiling and halves occupancy; BLOCK_K=64 (32KB) doubles it. Measured cold-L2 do_bench: uniform B120_D256 0.575->0.509 (1.11x), B1024_D256 4.737->4.236 (1.14x), skew B120_D256 0.152->0.129 (1.19x); all cos=1.0. D512 already used 64. This is the inverse of the failed tier-B grow-levers (BLOCK_M=256, STAGES_A=3, BLOCK_N=256 all crashed or regressed by blowing the 64KB LDS / killing occupancy): on gfx942 occupancy is the binding constraint, so shrinking LDS wins where growing it loses. The old "BLOCK_K=128 wins ~4%" note was a gfx950 result (larger LDS hides the occupancy cost). Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…sults) Record the 2026-06-10 autoresearch session: bound check (overhead-bound, 25-32% of HBM peak), XCD tier-A win, the BLOCK_K=64 occupancy win, and the four discarded tier-B grow-levers. Central lesson: on gfx942 the 64KB LDS ceiling makes occupancy the binding constraint — shrinking LDS wins, growing it loses. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Wire optional block_m/block_n through the public jagged_dense_bmm entry and forward cfg tile_m/tile_n from the dispatch (uniform path only). Promote the gfx950 B1024D512K512N7680 winner to 256x256 + xcd_c=120: 5.85ms vs 5.98ms at 128x128 cold-L2 do_bench, cos=1.0. The 256x256 epilogue needs 128KB LDS, which fits CDNA4's 160KB but not CDNA3's 64KB, so this is a gfx950-specific win. Add tier-A XCD (tune_jdbba_xcd.py) and tier-B tile (tune_jdbba_tiles.py) autotuners. The XCD sweep confirmed the existing per-shape winners are already at the remap optimum and that every skew shape prefers remap-off (the dispatch default); only the D512 fat-tile fold promotes. Skew and all D256 shapes unchanged. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Thread an optional waves_per_eu occupancy hint through _build_launcher (added to the lru_cache key), the public jagged_dense_bmm entry, and the dispatch (uniform path only, from cfg["waves_per_eu"]). The hint is applied at compile time via CompilationContext.compile_hints around the kernel launch. Promote gfx950 B1024D256K256N7680 to waves_per_eu=1: ~2.13ms vs ~2.20ms cold-L2 do_bench (every wpe>0 trial beat every baseline trial), cos=1.0. This is the overhead-bound small-weight large-B shape where extra occupancy hides the fixed per-block cost. Other shapes and all skew unchanged. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…56 skew) The persist-vs-static and skew-remap gates were gfx950 decisions. Re-measured on gfx942 (do_bench cold-L2): the persistent kernel LOSES on every gfx942 skew shape (B1024_D256: persist 1.045 vs static remap-off 0.904 vs remap-ON 0.857), so gate persist to gfx95x only. For the small-weight/large-B skew cell (B1024_D256) route to static remap-ON (xcd_c=32), the per-shape skew winner. Net: B1024_D256 skew 1.051 -> 0.873 ms (+20%), now 1.17x faster than Triton (was a tie). Other skew shapes unchanged. Dispatch test passes. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Record: rocprofv3 profile (23.5% occupancy, all units idle -> VGPR-capped overhead-bound), the skew-gate re-derivation win (+20% B1024_D256 skew), and the discarded batch-2 LDS-neutral levers (warp layouts, A-g2r, THREADS=512), #12 B-stationary, and BLOCK_N=64. Key reframe: the C-shuffle epilogue fixed cost is the floor; next lever is a cheaper direct epilogue. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
The HSTU deployment uses max_seq_len=16384, but the dispatch only had Mi=7680 winners, so the 16384 shapes fell through to the 128x128 heuristic. Add tuned winners for all four headline shapes at N16384 (do_bench cold-L2 uniform, gfx950, cos=1.0): - B1024_D512: 256x256 tile -> 12.54ms vs 12.89 heuristic (-2.7%); the fat tile fits CDNA4's 160KB LDS (128KB epilogue), impossible on gfx942. - B1024_D256, B120_D256, B120_D512: 128x128 + waves_per_eu=1. Mi=7680 winners and skew unchanged; dispatch correctness test passes. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…1.7%) The skew path unconditionally forced the XCD remap off (the committed assumption was that skew's uneven M_i always makes clustering a load-imbalance loss). That holds for the small-D256 / few-group shapes, but NOT for the large-weight B1024_D512 skew shape: a modest chunk (xcd_c=8) recovers enough Dense[b] L2 reuse to beat round-robin by ~1.7% (1.232 vs 1.262ms, 5/5 do_bench trials, cos=1.0). Add optional skew_xcd_c / skew_xcd_w config knobs the dispatch forwards on the skew path (default None -> remap off, unchanged for every other shape). Uniform path and all other skew shapes unchanged; dispatch test passes. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…% all shapes)
Deepen the B (dense weight) register pipeline from 2 to 3 stages with a
DECOUPLED prefetch: A stays LDS 1-ahead/2-deep (mod 2), B now uses 3 register
slots and prefetches 2 K-tiles ahead (read slot i%3, prefetch tile i+2 ->
(i+2)%3). The extra B prefetch hides global-load latency on the short K-loop,
which profiling (rocprofv3: 23.5% occupancy, units idle, MemUnitStalled only
4%) showed was the bottleneck -- the kernel was latency-bound, not LDS/BW-bound.
Naive stages=3 alone is a no-op (the ^1 toggle never touches slot 2); the win
needs the explicit mod-3 rotation + 2-ahead prefetch + slot-1 prologue prime.
Costs one extra B register slot (B is reg-staged, not LDS), so no LDS-ceiling
hit. Measured do_bench cold-L2, all cos=1.0, stable across 2 runs:
uniform: B120_D256 0.509->0.473, B120_D512 1.468->1.352, B1024_D256
4.236->4.042, B1024_D512 12.24->11.33 (0.92-0.95x)
skew: 0.128->0.121, 0.352->0.331, 0.873->0.808, 2.607->2.493 (0.92-0.95x)
FlyDSL now leads Triton ~1.35-1.42x uniform, up to 1.32x skew. Dispatch passes.
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Record the +5-8% B-fragment 2-ahead/3-stage win (the latency-hiding the 23.5%-occupancy profile predicted), its B_STAGES sweep (3 is peak, 4 noise, 5 regresses), and the batch-3 discards (A g2lds blocked, STAGES_A=1 null, waves_per_eu flat, epilogue barrier-removal redundant-but-flat, BLOCK_K=32 compile-fail). Final standing: FlyDSL leads Triton ~1.35-1.42x uniform. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…yenh/flydsl-jdbba # Conflicts: # op_tests/flydsl_tests/tune_jdbba_xcd.py
… gfx950) The 3-slot/2-ahead B-fragment prefetch (gfx942 win, +5-8%) is a regression on gfx950: its tuned 256x256 tile for B1024_D512 is VGPR-capped at occupancy 1.0, so the extra B register slot spills -> 11.3ms vs 5.9ms (1.92x slower) uniform. Parameterize the B register-prefetch depth as compile-time B_STAGES (prefetch distance B_STAGES-1). gfx942 keeps B_STAGES=3 (their breakthrough preserved); gfx950 uses B_STAGES=2, which collapses the B pipeline to 1-ahead/lockstep with A (the pre-prefetch behavior). Also make block_k arch-aware: gfx950 keeps 128 for K<=256 (160KB LDS hides the occupancy cost), gfx942 uses 64 for all K. Restores gfx950 to its committed winners: B1024_D512 uniform 5.81ms (0.91x Triton), all 4 shapes both regimes cos=1.0, matching pre-merge baseline. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…PR relief) The kernel is AGPR-occupancy bound: the 128x128 fp32 C-accumulator is 64 fp32/thread -> 128 AGPR (full file) -> ~2 waves/CU, capping MfmaUtil at 44%. THREADS=512 with an 8-warp (1,8,1) layout halves the accumulator to 32 fp32/thread (64 AGPR), raising occupancy. Wins where the K-loop is shortest and most occupancy-sensitive: B120_D256 -13.5%, B1024_D256 -6.0% (do_bench cold-L2 uniform, cos=1.0). D512 is already MFMA-bound (2x the MFMA work) so the extra warps only add scheduling overhead (+6-7%), and all skew shapes regress -> gated to D256-uniform via the per-shape dispatch table. threads is now a per-call knob (default 256) plumbed through the dispatch; warp N-count derives from THREADS//64. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…neralize warp layout
Doubling warps (threads=512) shifted the XCD L2 sweet spot: re-sweeping
xcd_c x xcd_w at the new occupancy (tune_jdbba_xcd_t512.py) picks
c=32/w=8 for B1024_D256 (was c=32/w=16), a small additional gain within
run-to-run noise but at worst neutral. B120_D256 stays c=32/w=16.
B_STAGES re-swept at threads=512: 3 remains optimal (unchanged by warp
count). Also generalize the gfx942 warp layout to split across M when
N-warps would exceed BLOCK_N/16, so the factory supports threads in
{256,512} (and 1024, though 1024 lost) without a hardcoded (1,8,1).
Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Reproduces the threads=512 D256 XCD re-tune that picked c=32/w=8 for B1024_D256. cos-gated, do_bench cold-L2, calls jagged_dense_bmm directly with threads=512. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Records the threads=512 D256 win (-7-13%), the "more warps not smaller tiles" mechanism, the threads=1024/waves_per_eu/BLOCK_M/N discards, and the D512 VALU-bound diagnosis. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…y axis Definitively refutes the deferred D512 warp-split lever (threads=512 regresses D512 ~24% under (1,8,1)/(2,4,1)/(4,2,1) alike) and corrects the stale "D512 VALUBusy 44%" note — fresh rocprofv3 shows MfmaUtil 44%, VALU ~20%, all memory units idle on both shapes. The memory-amortization lever class (B-stationary, deeper A-staging) is moot: MemUnitStalled ~0.1%. Kernel is at its gfx942 MFMA+ occupancy ceiling; no quick lever remains. No kernel change this batch. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…sue latency Decisive experiment: threads=512 on D512 doubles occupancy (1.92->3.82 waves/CU, AGPR 128->0) yet MfmaUtil stays flat at 44% and runtime regresses -> the kernel is MFMA-issue/dependency-latency bound, NOT occupancy bound (supersedes the batch 4-5 AGPR-occupancy framing). MFMA work is already zero-waste (~28% of peak). All batch-6 levers refuted against the new diagnosis (smaller/fatter tiles, BLOCK_K=128, STAGES_A=3, B_STAGES sweep). 2x requires a warp-specialized producer/consumer rewrite or the 32-K atom (gfx942 lacks it) -- out of scope for a single lever. No kernel change this batch. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
…robe (not recommended) Validated the rewrite path before committing to it. Barrier ablation: dropping the per-K gpu.barrier() changes runtime 0.0% with cos still 1.0 -- the barrier is already compiler-redundant. Wait-cycle breakdown: 54% issue-stall cycles, VALU 1.84x MFMA inst count -> the MFMA idle is the in-warp s2r->MFMA dependency + VALU interleave, not cross-warp sync. Both occupancy and sync levers leave MfmaUtil at 44%. Warp specialization's recoverable headroom is small and far below 2x; the real 2x blocker is the missing 32-K MFMA atom (HW-absent on gfx942). Recommendation: do not invest in the rewrite for these shapes. Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Comment on lines
+75
to
+76
| import functools | ||
| import json |
Contributor
| remaining = c_num_rows - first_row | ||
| # win_h = min(num_rows - first_row, W) -- the last window may be short. | ||
| win_h = (remaining < cW).select(remaining, cW) | ||
| l = xy_g % tids_per_grp |
Contributor
| remaining = c_num_rows - first_row | ||
| # win_h = min(num_rows - first_row, W) -- the last window may be short. | ||
| win_h = (remaining < cW).select(remaining, cW) | ||
| l = xy_g % tids_per_grp |
Contributor
| """Parse rocprofv3 (rocpd schema) --pmc output for the jdbba kernel: L2 hit and | ||
| LLC-hit (DRAM-avoided) fraction, summed over all XCD instances and all jdbba | ||
| dispatches. Usage: parse_pmc.py <pmc_outdir> <p10_us> <tag>""" | ||
| import sys, sqlite3, glob, os |
Contributor
| @@ -0,0 +1,34 @@ | |||
| """Combine two split PMC dirs (L2 hit/miss; DRAM reqs) for the jdbba kernel. | |||
| Usage: parse_pmc_split.py <tag> <dir_a> <dir_b>""" | |||
| import sys, sqlite3, glob, os | |||
Contributor
Comment on lines
+16
to
+17
| from aiter.ops.flydsl.jagged_dense_bmm_dispatch_v2 import jagged_dense_bmm_dispatched | ||
| from aiter.ops.flydsl.kernels.jagged_dense_bmm_gen import jagged_dense_bmm, BLOCK_M as _BLOCK_M |
Contributor
There was a problem hiding this comment.
aiter.ops.flydsl.jagged_dense_bmm_dispatch_v2.jagged_dense_bmm_dispatched imported but unused
Suggested change
| from aiter.ops.flydsl.jagged_dense_bmm_dispatch_v2 import jagged_dense_bmm_dispatched | |
| from aiter.ops.flydsl.kernels.jagged_dense_bmm_gen import jagged_dense_bmm, BLOCK_M as _BLOCK_M | |
| from aiter.ops.flydsl.kernels.jagged_dense_bmm_gen import jagged_dense_bmm, BLOCK_M as _BLOCK_M |
Comment on lines
+75
to
+76
| import functools | ||
| import json |
Contributor
| remaining = c_num_rows - first_row | ||
| # win_h = min(num_rows - first_row, W) -- the last window may be short. | ||
| win_h = (remaining < cW).select(remaining, cW) | ||
| l = xy_g % tids_per_grp |
Contributor
| remaining = c_num_rows - first_row | ||
| # win_h = min(num_rows - first_row, W) -- the last window may be short. | ||
| win_h = (remaining < cW).select(remaining, cW) | ||
| l = xy_g % tids_per_grp |
Contributor
| """Parse rocprofv3 (rocpd schema) --pmc output for the jdbba kernel: L2 hit and | ||
| LLC-hit (DRAM-avoided) fraction, summed over all XCD instances and all jdbba | ||
| dispatches. Usage: parse_pmc.py <pmc_outdir> <p10_us> <tag>""" | ||
| import sys, sqlite3, glob, os |
Contributor
| @@ -0,0 +1,34 @@ | |||
| """Combine two split PMC dirs (L2 hit/miss; DRAM reqs) for the jdbba kernel. | |||
| Usage: parse_pmc_split.py <tag> <dir_a> <dir_b>""" | |||
| import sys, sqlite3, glob, os | |||
Contributor
Comment on lines
+16
to
+17
| from aiter.ops.flydsl.jagged_dense_bmm_dispatch_v2 import jagged_dense_bmm_dispatched | ||
| from aiter.ops.flydsl.kernels.jagged_dense_bmm_gen import jagged_dense_bmm, BLOCK_M as _BLOCK_M |
Contributor
There was a problem hiding this comment.
aiter.ops.flydsl.jagged_dense_bmm_dispatch_v2.jagged_dense_bmm_dispatched imported but unused
Suggested change
| from aiter.ops.flydsl.jagged_dense_bmm_dispatch_v2 import jagged_dense_bmm_dispatched | |
| from aiter.ops.flydsl.kernels.jagged_dense_bmm_gen import jagged_dense_bmm, BLOCK_M as _BLOCK_M | |
| from aiter.ops.flydsl.kernels.jagged_dense_bmm_gen import jagged_dense_bmm, BLOCK_M as _BLOCK_M |
anhminhnguyenhoang
changed the title
… B prefetch The "decouple B-fragment prefetch to 2-ahead/3-stage" rework (commit e03fcfe) mis-accumulated the K-reduction in the gen kernel: outputs were scrambled with ||out|| ~= ||ref|| but cos ~0.5 at D=256 (2 K-tiles) and ~0 at D=512 (8 K-tiles). Only the persistent-kernel path was unaffected, so the -test bench passed just 1/8 shapes (the one skew cell routed to it). Bisection: the parent commit is bit-exact (cos=1.0); the regression is introduced by the reworked B rotation (mod-B_STAGES slots + extra prologue prime) and is NOT masked even at B_STAGES=2 on gfx950, where it was meant to collapse to the correct 2-slot lockstep. Restore the known-good 2-slot double-buffered A/B pipeline. B_STAGES is still threaded through the launcher cache key but the compute path uses the correct 2-slot buffer; on gfx950 the deeper prefetch was a no-op anyway. Verified on gfx950 (MI350X): bench_jagged_dense_bmm_perf.py --regime both -test now passes 8/8 (cos=1.0000), and the ruka MoE shape (B=64 D=512 Kout=1024 Mi=8192) is correct at ~648 TFLOP/s.
…/regimes) The original 3-stage B prefetch mis-paired B tiles with A tiles because both A and B were indexed by the same read_stage (toggling 0/1), even though B had 3 slots (producing 0/1/0/1 instead of 0/1/2/0/1/2). Songlin's fix correctly reverted to 2-slot lockstep; this commit re-implements the 3-stage pipeline with the correct decoupled indices. Root cause of the original bug: run_pipeline_stage(read_stage, ...) used read_stage for both the A LDS slot (k%2, correct) and the B register slot (must be k%B_STAGES, was k%2 -> aliased slots). Fix: separate a_read_stage (k%2) from b_read_slot (k%B_STAGES), separate a_next_k (A prefetch = k+1) from b_next_k (B prefetch = k+2 for 3-stage). Prologue pre-fills B slots 0 and 1 (tiles 0, 1) so the main loop always has the next B tile ready 2 steps ahead. Loop unrolled via fx.range_constexpr so per-step None/int checks are Python-level compile-time branches. Results (gfx942, do_bench cold-L2, cos=1.0 all shapes): Uniform: -5.7% B120D256, -6.2% B120D512, -2.0% B1024D256, -4.5% B1024D512 Skew: -7.2% B120D256, -12.3% B120D512, -7.5% B1024D256, -10.0% B1024D512 Co-Authored-By: Claude <noreply@anthropic.com>
…%/-5%) Skew tail-waste elimination: with skewed sequence lengths ~84% of the envelope-grid M-tiles are empty. For small-B inference skew (n_groups <= 120, gfx942) launch only the OCCUPIED (group, m-tile) blocks via a device-built TILE_MAP, skipping every empty tile. - kernels/jdbba_skew_tile_map.py: single-launch on-device prep (per-block prefix scan of seq_offsets + scatter), no device->host readback; slack rows hold a sentinel that early-exits the compact kernel. - kernels/jagged_dense_bmm_gen.py: COMPACT const_expr grid path reading TILE_MAP; non-compact (uniform / large-B skew) path unchanged (dummy arg). - dispatch_v2.py: gate to skew + n_groups<=120 (gfx942); TILE_MAP cached per seq_offsets ptr so a forward pass amortizes the ~1us prep. Verified on MI300X/gfx942 (A/B, same toolchain): uniform unchanged; skew B120_D256 0.125->0.097ms (-22%, 1.62x Triton), B120_D512 0.324->0.308ms (-5%, 1.34x); B1024 skew gated to envelope (unchanged). Correctness cos=1.0 both regimes on fresh compile; JSON autotune winners untouched. Co-authored-by: Cursor <cursoragent@cursor.com>
…-17/-19%) Frontier lever F3. The skew compact static grid launched occupied tiles in naive group-major order; the hardware XCD routing (raw id xy % 8) then hands each XCD a stride-8 sample of the group-major map, so every 4MB private L2 must hold ~all groups' Dense[b] and thrashes (the -13% that gated compact off for large B). Apply HipKittens Algorithm 1 (C=32, W=8) to the compact grid so C consecutive group-major tiles co-locate on one XCD, restoring per-group weight residency. - gen.py: _xcd_remap_compact (runtime num_rows from grid_dim.x -> no per-count recompile); gated behind XCD_C>1 so the naive path stays byte-identical. - dispatch_v2.py: widen the gate (all skew -> compact); _skew_compact_xcd picks the XCD order iff n_groups>120 OR reduction_k>=512, else naive group-major (B120_D256 stays naive: the reorder was -9% there). Uniform never touches the compact path. - Fix a latent cache bug: the seq_offsets.data_ptr() key could return a stale TILE_MAP after the torch allocator recycled an address (silent cos=0); now validated by tensor object identity (weakref + is), which also hardens the already-shipped B120 path. Verified MI300X/gfx942 (do_bench cold-L2, cos=1.0 all 4 skew + edges; rocprofv3 B1024_D512 skew L2 hit 36->75%, HBM read 6.6->1.8 GB): B1024_D512 skew 2.47->1.99 ms (-19%, 0.65x Triton), B1024_D256 0.83->0.69 (-17%, 0.63x), B120_D512 0.31->0.27 (-11%). B120_D256 skew + all uniform unchanged. Co-authored-by: Cursor <cursoragent@cursor.com>
| first_row = group_id * cW | ||
| remaining = nrows - first_row | ||
| win_h = (remaining < cW).select(remaining, cW) # last window may be short | ||
| l = xy_g % tids_per_grp |
Contributor
1 task
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
2 participants
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.
Motivation
jagged_dense_bmm_broadcast_addis a grouped GEMM that does, for each groupbover its packed row slice[s, e):Group row boundaries come from a device-resident
seq_offsetsprefix-sum array -- the host doesnot know each group's row count at launch, so the group->row mapping is resolved on the GPU. This
rules out any stock batched-GEMM kernel.
Technical Details
Kernel (
aiter/ops/flydsl/kernels/jagged_dense_bmm_gen.py): generalized BF16 grouped GEMMwith FP32 accumulation. Key design decisions:
(B, N, K)(cheap, weights are static), allowingplain coalesced K-reduction without a preshuffle step.
(not LDS), so deepening its prefetch costs VGPR, not shared memory -- sidestepping gfx942's
64 KB LDS ceiling entirely. A stays LDS double-buffered (2-stage); B uses 3 register slots and
prefetches 2 K-tiles ahead.
ceiling on gfx942 (CDNA3), halving occupancy. 64 -> 32 KB -> doubles occupancy. This is the
central gfx942 lesson; gfx950 (160 KB LDS, CDNA4) is unaffected.
(1,8,1)) for D256-uniform shapes: halves the FP32 C-accumulatorper thread (64->32 fp32/thread -> 128->64 AGPR), raising CU occupancy on the short 4-tile
K-loop where AGPR pressure is the binding constraint. D512 (8-tile K-loop) and all skew shapes
regress and stay at threads=256 via the dispatch.
xcd_c x xcd_wper headline shape;per-arch winners stored in the dispatch JSON. Marginal size (~1-3%) confirms the kernel is not
L2/chiplet-bound.
jagged_dense_bmm_dispatch_v2.json, schemaarch-keyed-v1): gfx942 has no16x16x32 BF16 MFMA atom --
use_mfma_k32=Truecore-dumps on MI300X, so it is hardcoded False.A flat legacy JSON is still accepted for env-var overrides (
FLYDSL_JAGGED_DENSE_BMM_DISPATCH_V2_JSON).Dispatch layer (
aiter/ops/flydsl/jagged_dense_bmm_dispatch_v2.py):_resolve_dispatch()=explicit-env-override -> exact-match JSON winner -> D-bucketed heuristic fallback.
threadsis aper-call dispatch knob (default 256); the warp N-count derives from
THREADS // 64, splittingacross M when
n_warpswould exceedBLOCK_N / 16.Test Plan
MI300X (gfx942), ROCm 6.x, inside the
jdbba-flydsldevcontainer; aiter at/workspaces/meta/aiter. Target arch: gfx942 only.The dispatch test verifies:
cos > 0.999vs PyTorch eager) on the detected arch.heuristic, also correctness-gated.
M_b=0), partial tiles,max_seq_len >> mean-- all pass.Test Results
MI300X gfx942,
triton.testing.do_benchcold-L2, Mi=7680, cos=1.0 all shapes/regimes.Baseline: upstream Meta/HSTU Triton kernel (
generative_recommenders.ops.triton.triton_jagged).Uniform regime (every group M_i = Mi = 7680)
Skew regime (~20% empty groups, one full-envelope, one near-full)
Key optimization history (what moved the needle, in order)
use_mfma_k32=False, D-bucketed heuristic