[Performance]: sm89 ada blockwise FP8 GEMM copies block-scale factors to shared memory 4x/128x redundantly (stride-0 scale TV layouts)

[jiazhihao]

Proposal to improve performance

In the sm89 "ada blockwise" FP8 GEMM (cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/ada_blockwise_gemm/), the per-row A scale factors and the per-tile B scale factor are copied from global to shared memory redundantly: every one of the 4 warps issues the same cp.async for every SFA element, and all 128 threads issue the same cp.async for the single SFB element, once per k-tile and pipeline stage. The prologue cute::clear of those buffers is duplicated the same way. Restricting the clears and copies to one owning thread per element removes the duplicate shared-memory writes with no change in results.

Where it comes from. The gmem-to-smem thread-value layouts for the scale factors use stride 0 across warps (sm89_utils.cuh, lines 194-199 at 134fa245fe):

using GmemLayoutTVSFA = Layout<Shape<Shape<Int<ScaleMsPerTile>, Int<kThreadCount / ScaleMsPerTile>>, Shape<_1, _1>>,
    Stride<Stride<_1, _0>, Stride<_1, _1>>>;
using GmemLayoutTVSFB = Layout<Shape<Shape<_32, _4>, Shape<_1, _1>>, Stride<Stride<_0, _0>, Stride<_1, _1>>>;

With ScaleMsPerTile = kTileM / ScaleGranularityM = 32 and kThreadCount = 128, thread t maps to SFA row t % 32 and the warp mode contributes stride 0, so lanes with the same index in all four warps own the same element. For SFB both strides are 0, so all 128 threads own the one element. Every operation partitioned through these layouts is therefore executed by 4 (SFA) or 128 (SFB) threads writing identical bytes to the same shared address:

• cute::clear(tAsSFA); cute::clear(tBsSFB); in the prologue (sm89_fp8_gemm_1d1d.cuh:265-266)
• cute::copy_if(g2s_copy_SFA, tApSFA, ...) and cute::copy(g2s_copy_SFB, ...) in the prologue (:283-286) and in the mainloop for every k-tile (:372-375)

Evidence from the compiled PTX (traits <float_e4m3_t, bfloat16_t, float, float, 3, 32, 128, 128>, CUDA 13.2, sm_100f). The prologue clear compiles to

st.shared.b32  [%r13+61440], 0;   // smem_sfa stage 0
st.shared.b32  [%r13+61568], 0;   // stage 1
st.shared.b32  [%r13+61696], 0;   // stage 2
st.shared.v2.b32 [SharedStorageBase+61824], {0f00000000, 0f00000000};  // smem_sfb
st.shared.b32  [SharedStorageBase+61832], 0;

with %r13 = SharedStorageBase + ((tid.x << 2) & 124), i.e. smem_sfa + (tid.x % 32) * 4, independent of the warp. Offset 61440 is sizeof(smem_a) + sizeof(smem_b) = 12288 + 49152. A static race checker over this PTX reports threads 0 and 32 as unordered writers of the same 4 bytes here; the writes are same-value and all consumers sit behind cp_async_wait + __syncthreads (:336-337, :357-358, :407-408), so this is not a correctness problem, only wasted shared-memory traffic and cp.async instructions.

Cost per k-tile per CTA (128 threads, 32x128x128 tile):

Copy	Current	With one owner
SFA cp.async instructions	128	32
SFB cp.async instructions	128	1
Shared bytes written for scales	1,024	132

That is 2 of the 12 cp.async instructions each thread issues per k-tile, and about 5% of the shared-memory write bytes next to the 20 KiB A/B tiles.

Fix: #19567 restricts the scale-factor clears and copies to one owning thread per element while keeping the stride-0 layouts; results are unchanged and the duplicate shared-memory writes disappear.

Report of performance regression

No response

Misc discussion on performance

The same stride-0 idiom appears in the smem-to-register scale layouts (SmemLayoutTVSFA/SmemLayoutTVSFB) but those are reads, so they are not affected.

Your current environment (if you think it is necessary)

Not applicable: the analysis is on the compiled PTX (CUDA 13.2, sm_100f) at commit ef567dd9b8919f6ee8b2da1f6f6d0c90206a78f8; the code is unchanged on main at 134fa245fe.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[jiazhihao]

Fix submitted in #19567.