optimize by optimize/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606261106#399
Open
ChengYao-amd wants to merge 10 commits into
Conversation
…le variant (tile_m=128, tile_n=128, t Hypothesis: Add a larger-M wgrad tile variant (tile_m=128, tile_n=128, tile_k=256) to _TILE_CANDIDATES and make _select_tile prefer tile_m=128 for the deep-contraction wgrad (TN) path on shapes whose output dim K forces tile_n to 128 (i.e. K not divisible by 256, e.g. the K-heavy shape M8192xN8192xK29568). With tile_n pinned to 128 the A-fragment reuse (num_acc_n) is stuck at 2 and the proven tile_n=256 reuse win is unreachable, so the only remaining compute lever is tile_m: doubling m_repeat 4->8 issues twice as many independent MFMA chains per wave (more ILP to hide LDS-load/MFMA latency under the measured 2-wave/SIMD occupancy ceiling) and amortizes the per-output-tile shared B-global-loads, per-output-column scale_b loads, and barrier/scheduling overhead over 2x more output rows, raising the MFMA:VALU ratio on the VALU-overhead-bound bs_col1d wgrad kernel. Result: Forward TFLOPS: -0.00%, Backward TFLOPS: +5.80%, Combined Step TFLOPS: +4.14% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606251609/logs/optimize.md
…_n=256 block2d A-fragment/scale reuse Hypothesis: Unlock the proven tile_n=256 block2d A-fragment/scale reuse lever on the #2-cost deep-K DGRAD (NN, bs_block2d) FLYDSL kernel by ZERO-PADDING its output-feature dim (= weight in_features K = 29568, NOT divisible by 256) up to the next multiple of 256 (29696, +128 cols = +0.43% extra work) at the BACKWARD OPERAND-PREP / LAUNCH level, then slicing grad_A back to 29568. Concretely: in FP8GemmBlockFunction.backward (primus_turbo/pytorch/ops/gemm_fp8.py), before the dgrad gemm_fp8_impl call, zero-pad the saved weight operand b_fp8 along its K axis (29568->29696) and append one extra finite (=1.0) K-block column to its blockwise scale b_scale_inv, run the existing FLYDSL dgrad GEMM producing grad_A[M, 29696], and slice grad_A[:, :29568] on return (the padded weight columns are exactly zero, so the existing grad_A columns are computed bit-identically -- algebraically exact, no in-kernel tail mask). In primus_turbo/flydsl/gemm/launcher.py, _select_tile (called as _select_tile(M, K=29696, N) for dgrad) now finds N_kernel=29696 divisible by 256 and its existing large-shape scorer already prefers tile_n=256 (num_acc_n 2->4). The cube profile PROVES this lever works on block2d: dgrad runs t64x256x128 at 112us with VALU:MFMA=5.73-5.88 (genuine MFMA headroom), whereas on deep_k it is forced to t64x128x128 and loses the 2x A-fragment/scale reuse; doubling num_acc_n feeds each A fragment + per-block scale_a into 2x more MFMAs before re-fetch, halving the B-load/scale_b/address-gen/barrier VALU per output element on a kernel that has the MFMA slack to absorb it (cross-campaign win: tile_n 128->256 = +3.91% on block2d). This ATTACKS the K-divisibility constraint itself for dgrad -- a mechanism no round has tried (R9 changed dgrad tile_k, R13 changed dgrad tile_m, both WITHIN the forced tile_n=128), and a different target (dgrad/block2d) than R2 (wgrad/col1d). It deliberately does NOT touch the VALU-saturated wgrad col1d path, whose cube data shows tile_n=256 leaves VALU:MFMA pinned at 8.41 (reuse cannot help an already-VALU-bound kernel). Result: Forward TFLOPS: +0.11%, Backward TFLOPS: +0.96%, Combined Step TFLOPS: +0.72% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606251609/logs/optimize.md
…t (16,16) MFMA preshuffle into the blo Hypothesis: Fold the forward weight (16,16) MFMA preshuffle into the blockwise weight-quant store, eliminating the standalone `shuffle_b(b_fp8)` HBM round-trip + elementwise kernel launch that the FLYDSL forward path currently runs on EVERY forward, on EVERY shape. Today `FP8GemmBlockFunction.forward` quantizes the weight via `quant_fp8_blockwise_for_weight_impl(b)` producing a plain block2d `b_fp8 [N,K]`, and the launcher (`gemm_fp8_blockwise_flydsl`) then re-reads it through `shuffle_b(b_fp8)` (a torch `permute().contiguous()` = a full uncoalesced fp8 copy, ~242MB read+write on the deep_k weight) before the GEMM. The activation path ALREADY proves this fold is feasible and bit-identical: `quant_fp8_blockwise_dual_impl(col_preshuffled=True)` emits the wgrad activation directly in the (16,16) preshuffled+transposed layout so the wgrad launcher's standalone preshuffle collapses to a zero-cost view. Apply the SAME mechanism to the weight/forward path: have the weight-quant kernel emit a second output `b_fp8` already in the forward (16,16) shuffle layout (one read of `w`, an extra store), keep the unshuffled `b_fp8` for `ctx`/dgrad, and have the forward launcher consume the pre-shuffled operand (shape/scale signal, exactly like the wgrad `a_preshuffled` detection) and skip `shuffle_b`. Byte-identical to `shuffle_b(b_fp8)`, so the forward GEMM output is unchanged. Result: Forward TFLOPS: +6.30%, Backward TFLOPS: +0.18%, Combined Step TFLOPS: +1.87% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606251609/logs/optimize.md
…ock2d) FLYDSL kernel's L2 grouped-ras Hypothesis: Raise the dgrad (NN, block2d) FLYDSL kernel's L2 grouped-rasterization height from l2_group_m=8 to l2_group_m=16, matching the forward NT kernel. The dgrad path is the only block2d direction still left at the smaller group: compile_blockscale_dgrad_gemm hardcodes l2_group_m=8 (blockscale_gemm.py:996) while the launcher's dgrad compile call (launcher.py:498) never passes the value. The PROFILE measured both directions on the SAME cube shape (M4096 N4096 K4096) with the SAME t64x256x128 tile and only the group height differing: bs_block2d_l2g8 (dgrad) = 78% L2 hit, bs_block2d_l2g16 (fwd) = 86% L2 hit. Increasing dgrad's super-block height to 16 raises B/N-panel reuse in L2 on the dgrad NN path, lifting its L2 hit toward 86% and cutting the global-load miss stalls that cap MFMA busy/wave below peak. The change is a pure workgroup-id remap (tail-safe group_size_m at lines 257-263), so the output is bit-identical. Implemented by letting the dgrad wrapper accept the l2_group_m kwarg (stop popping/forcing 8) and having the launcher's dgrad path pass l2_group_m=16 to the compile_fn. Result: Forward TFLOPS: -0.38%, Backward TFLOPS: +0.75%, Combined Step TFLOPS: +0.43% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606261106/logs/optimize.md
…d (fwd NT / dgrad NN) FLYDSL hot loop Hypothesis: Restructure the block2d (fwd NT / dgrad NN) FLYDSL hot loop into an interleaved half-iteration ping/pong LDS pipeline: split each K-iteration's next-tile A global->LDS DMA (and B/scale VMEM) into discrete sched_barrier/s_setprio-bounded phases that overlap with the current tile's ds_read + MFMA issue, so the VALU-heavy scale bookkeeping (block2d VALU 69% of insts, MFMA only 12%) is scheduled into the matrix-unit shadow and MFMA busy fraction rises. This refines the existing coarse 2-stage ping/pong (allocator_ping/pong + hot_loop_scheduler) into a finer half-iteration overlap on the block2d path only; the wgrad col1d epilogue and the accumulation math are left byte-for-byte unchanged. Result: Forward TFLOPS: +1.42%, Backward TFLOPS: +0.62%, Combined Step TFLOPS: +0.85% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606261106/logs/optimize.md
…transposed-preshuffle into the forward Hypothesis: Fold the dgrad WEIGHT transposed-preshuffle into the forward weight dual-quant and carry it to backward via ctx, eliminating the per-backward-step 669us `_preshuffle_b_transposed_kernel`. Today the forward weight dual-quant (quant_fp8_blockwise_for_weight_dual_impl) already emits the plain block2d `b_fp8` AND a forward (16,16)-preshuffled `b_fp8_ps` in one read of the weight, but the dgrad (NN, block2d) path recomputes its own DIFFERENTLY-transposed weight layout every backward via `b_shuffled = _shuffle_b_transposed(b_fp8)` (launcher.py:490, unconditional) -- this is the 669us `_preshuffle_b_transposed_kernel` that fires every backward step in the isolated FLYDSL timeline. Extend the weight dual-quant to ALSO emit the dgrad transposed-preshuffle layout (`shuffle_b(b_fp8.T)` bytes) as a 3rd output in the SAME read of `b`, save it to ctx.save_for_backward, and add a dgrad-side pre-shuffled fast path in compile_blockscale_dgrad_gemm (mirroring the forward `b_preshuffled` 1D-scale-signal skip at launcher.py:383/407) that consumes the carried layout and skips `_shuffle_b_transposed`. One shared read of the weight produces all three layouts, so the standalone preshuffle pass is ELIMINATED (not merely relocated or launch-batched as R20 did). The dgrad GEMM's preshuffled-B load and accumulation are byte-for-byte unchanged; only the producer of that operand moves from a dedicated backward kernel into the forward weight-quant store. Result: Forward TFLOPS: -1.93%, Backward TFLOPS: +1.53%, Combined Step TFLOPS: +0.53% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606261106/logs/optimize.md
…latency that round-24 added to the for Hypothesis: Cut the exposed store latency that round-24 added to the forward weight dual-quant kernel (quant_fp8_blockwise_for_weight_dual_kernel, quant_blockwise.py:235-324). Round-24 accepted by folding the dgrad transposed-preshuffle into this kernel as a THIRD store per tile (the EMIT_DGRAD_PS branch at :308-319, a full-tile tl.trans register transpose + a scattered/strided FP8 store), but it left forward at -1.93% quick / -0.87% full because that extra long-latency scattered store now sits serialized on the forward weight-quant critical path after the plain and fwd-preshuffle stores. Hide that store latency without touching the GEMM: (a) reorder the per-tile epilogue so the high-latency scattered dgrad-PS store (tl.trans + gather/scatter) is ISSUED FIRST, before the coalesced plain block2d store (:276), the fwd-PS store (:297) and the scale store (:322), so its memory latency overlaps the cheaper stores rather than tailing them; and (b) raise the launch's num_warps (currently defaulted in quantization_impl.py:376; a 128x128 tile gets ~4 warps) so more concurrent warps are resident to absorb the three FP8 stores plus the transpose. All four stored buffers are byte-for-byte identical (same cast bytes, same addresses, only issue order / warp count change), so out/da/db SNR stays at the ~28.6 dB FP8 floor and the dgrad/forward GEMM operand loads are untouched. Result: Forward TFLOPS: +0.07%, Backward TFLOPS: +0.25%, Combined Step TFLOPS: +0.20% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606261106/logs/optimize.md
… uncoalesced dgrad transposed-preshuff Hypothesis: Coalesce the genuinely uncoalesced dgrad transposed-preshuffle (dgrad-PS) FP8 store in quant_fp8_blockwise_for_weight_dual_kernel (quant_blockwise.py:292-303). Today that store writes `tl.store(w_fp8_dgrad_ps_ptr + dps_off, tl.trans(w_fp8_cast), ...)` through an arbitrary permuted index expression `dps_off` (the 5D shuffle permutation i0d..i4d). Triton cannot prove contiguity through that gather offset, so it lowers to per-element 1-byte SCATTER stores -- the classic uncoalesced 1-byte fp8 store pattern (cross-campaign win record :000015: '1-byte fp8 stores via the generic path are uncoalesced and dominate operand-prep cost'). Restructure the dgrad-PS store so the innermost contiguous run (KIN_PS=16 = 16 bytes = dwordx4) is expressed as a simple contiguous store dimension Triton can vectorize: lay the transposed tile out in the kernel's register/loop order so the K_inner=16 run is the fastest-varying, simple-strided axis of the store, turning the scatter into coalesced dwordx4 transactions. Pair this with the matching launch tile/num_warps config in quantization_impl.py so the chosen BLOCK maps cleanly onto the 16-byte store granularity. Output bytes and target addresses are byte-for-byte identical (same permutation, same destination); only the store instruction width/coalescing changes. Result: Forward TFLOPS: +0.73%, Backward TFLOPS: -0.11%, Combined Step TFLOPS: +0.14% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606261106/logs/optimize.md
…d NT / dgrad NN) hot-loop MFMA schedul Hypothesis: Refine the block2d (fwd NT / dgrad NN) hot-loop MFMA scheduling cadence in `block2d_setprio_iter_sched` (mfma_preshuffle_pipeline.py). Today the per-iteration hint stream issues only TWO coarse `sched_mfma(mfma_group)` sub-groups (mfma_group=num_acc_n) bracketed by one s_setprio pair, with a single sched_dsrd between them and a tail-only sched_dswr (R13's accepted but coarse schedule). On the compute-bound, MFMA-starved block2d kernel (VALU ~69% vs MFMA ~12% of issued insts), this leaves the next-tile DS-read and fp32 block-scale bookkeeping VALU only loosely hidden. Split each iteration's matrix work into FINER s_setprio-bracketed sched_mfma sub-groups (e.g. 4 groups of mfma_group/2 instead of 2 of mfma_group) and interleave the sched_dsrd / sched_dswr / sched_vmem hints between those finer sub-groups so the LDS-read and scale-fold VALU are scheduled deeper into the matrix-unit shadow, tightening MFMA issue cadence and raising MFMA busy fraction. Pure scheduling-hint stream: emits no data ops, no waitcnt, no change to MFMA atom / tile shape / accumulation order, so the GEMM output is bit-identical across fwd and dgrad. Result: Forward TFLOPS: +0.23%, Backward TFLOPS: +0.02%, Combined Step TFLOPS: +0.08% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606261106/logs/optimize.md
…block (grouped) workgroup-id->tile ras Hypothesis: Enable L2-aware super-block (grouped) workgroup-id->tile rasterization on the wgrad (TN, bs_col1d) FLYDSL kernel, which today runs with NO grouping: compile_blockscale_wgrad_gemm (blockscale_gemm.py:1015-1019) hardcodes l2_group_m=1 and pops any passed value, so the col1d grid uses a plain row-major launch (kernel name bs_col1d_l2g1) while the block2d fwd/dgrad path uses l2g16. The PROFILE (round-45) classifies bs_col1d as the campaign's #1 bottleneck and LATENCY-bound: LDS-wait/wave 105.5 (2.48x block2d's 42.6) on an under-populated grid (SQ_WAVES 2.9e5 vs block2d 5.8e5) with a healthy aggregate L2 hit (0.881) and zero LDS bank conflicts. With only ~2 resident waves/SIMD the kernel cannot hide its operand-fetch latency; a tall M-band (super-block) rasterization keeps consecutive col1d output tiles that reuse the same a^T / grad_out panel resident in the same L2 slice, cutting the (remote) L2-access latency the 2-wave kernel currently exposes. Implement the grouped pid->tile mapping as a NEW helper module (_col1d_l2_rasterize.py) consumed by the col1d branch of blockscale_gemm.py, mirroring the R4-accepted dgrad L2-grouping win but on the one kernel that has grouping disabled. Pure tail-safe workgroup-id permutation: every output tile is computed exactly once with unchanged intra-tile accumulation order, so wgrad output stays bit-identical (~28.6 dB FP8 floor) and per-wave VGPR/LDS footprint is unchanged (so it cannot trip the occupancy wall that killed R19/R22/R32/R41 nor the backward single-direction numerics gate that killed R7). Result: Forward TFLOPS: -0.20%, Backward TFLOPS: +1.45%, Combined Step TFLOPS: +0.96% Details: /apps/tas/yaoc/agent_work/mi355x/primus-turbo-auto-optimizer/agent_workspace/Primus-Turbo/agent/workspace/optimize_turbo_blockwise_fp8_gemm_with_flydsl_ba_202606261106/logs/optimize.md
ChengYao-amd
requested review from
RuibinCheung,
wenxie-amd and
xiaobochen-amd
as code owners
June 29, 2026 06:12
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
1 participant
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.
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: