fix: move pipeline barriers into dynamic SharedStorage for 128-byte alignment (#42)

[XFDG]

Summary

• Move pipeline barriers from static __shared__ arrays into a dynamically-allocated SharedStorage struct in both GroupGEMMFp8Config and GroupGEMMBlockWiseFp8Config.
• This guarantees the TMA buffers (A, B, Y) are always 128-byte aligned, fixing TMA Misaligned failures on SM120 (Blackwell).
• get_shm_size() now returns sizeof(SharedStorage) so the launch-side allocation includes the barriers.
• Pure alignment fix — no logic change; pointer arithmetic for buffer offsets is preserved exactly.

Root cause

The kernels declared static __shared__ uint64_t writable[kStage] / readable[kStage] before the extern __shared__ dynamic region. The compiler places static SMEM first at a fixed offset; the dynamic base's alignment depends only on the static allocation's alignment — not on alignas(128) written on the extern __shared__ declaration. PTX inspection confirms:

// old pattern: static barriers + char[] dynamic
.extern .shared .align 16  shm_data[];    // only 16-aligned

On SM120 (Blackwell) this causes TMA loads to fail with Misaligned because the A/B buffers are not 128-byte aligned.

Fix

Move barriers inside the dynamic allocation as a typed SharedStorage struct. PTX confirms the typed extern __shared__ declaration forces the dynamic base to 128:

// new pattern: typed dynamic struct
.extern .shared .align 128 storage[];     // 128-aligned base

All buffer members (alignas(128)) then land at provable 128-multiple offsets. This is the canonical CUTLASS SharedStorage pattern.

Files changed

File	Change
src/group_gemm/config.h	Add SharedStorage struct to both configs; get_shm_size() → sizeof(SharedStorage) (+36 lines)
src/group_gemm/kernels.cuh	Replace static barriers + extern uint8_t[] with extern SharedStorage storage[] (−10 +11 lines)

Testing

• PTX verification: nvcc -arch=sm_90a -ptx confirms .align 128 on the typed dynamic declaration and all buffer offsets are 128-multiple.
• No GPU test: the bug only manifests on SM120 (Blackwell); the repo's build target is SM90a (H200/H20). The fix is architecture-independent and correct on SM90 too.
• No functional change: barrier init/arrive/wait logic is unchanged; buffer pointer arithmetic is preserved exactly.

References

• Issue Potential Shared Memory Misalignment Issues in GEMM Kernels #42: Potential Shared Memory Misalignment Issues in GEMM Kernels

---

[imusong]

hi，I'd like to ask if you have successfully migrated HPC-OPS to the SM120?

[XFDG]

Hi @imusong, thanks for the question!

No, we haven't tested on SM120 hardware — the available GPUs (H200, SM90a) don't include SM120. The fix is based on:

1. PTX inspection: nvcc -arch=sm_90a -ptx confirms the typed extern __shared__ SharedStorage[] declaration emits .align 128 on the dynamic segment, and all buffer offsets (A=256, B=16640, Y=33024) are 128-multiple. The root cause — static barriers shifting the dynamic base off 128 — is proven by comparing old vs new PTX.

2. Canonical CUTLASS pattern: this is exactly how CUTLASS handles shared memory (barriers inside a dynamic struct, not static arrays before the dynamic region). The reporter of Potential Shared Memory Misalignment Issues in GEMM Kernels #42 validated this fix on SM120.

If you have SM120 hardware available, we'd be happy to test it there. Otherwise, we believe this is a correctness improvement that is safe on SM90 and fixes the alignment issue on SM120.