Maybe it’s also worth leaving a TODO to double-check, once the MI355X hardware is available, that the compiler automatically translates this into the corresponding global_store_dwordx8 instruction, and not two sequential global_store_dwordx4?

Good suggestion — I'll add a TODO comment in the code.

To verify on MI355X once hardware is available:

hipcc -O3 --offload-arch=gfx950 --save-temps concat_mla_q.cu
# Check the .s file for:
#   global_store_dwordx8  (single 256-bit NT store)
# vs:
#   global_store_dwordx4 × 2  (two 128-bit NT stores)

If the compiler doesn't emit dwordx8, we can fall back to inline ASM with global_store_dwordx8 directly. Pushing a commit with the TODO now.

Update: verified on MI350X (gfx950, ROCm 7.2)!

The compiler generates 2× global_store_dwordx4, not a single global_store_dwordx8. The gfx950 ISA in ROCm 7.2 / LLVM 22.0 does not include vector dwordx8 instructions at all (llvm-mc rejects them as invalid).

The 256-bit path still halves loop iterations, so the optimization is valid — just not via the single-instruction path we hoped for. Will add a TODO comment noting this for future ROCm versions.

Thanks for looking into this! Would it be better to disable the 256b memory instructions for now, then? Maybe with an assert "Not supported yet"?

Two consecutive 128b operations issued in different cycles would lead to uncoalesced memory accesses within a warp.

Good call — you're right that two consecutive 128b ops in different cycles would lead to uncoalesced accesses. I've disabled the gfx950 256-bit path for now (commented out with explanation).

The gfx942 (MI300X) non-temporal hints path remains unchanged — that one is hardware-verified and the compiler auto-coalesces into single global_store_dwordx4 nt instructions.

Also rebased on upstream/main to resolve the merge conflict.

When a future ROCm version adds native global_load/store_dwordx8, we can re-enable the 256-bit path with a simple uncomment + re-verify.

The comment here states the goal is to use global_load_dwordx8, but the implementation uses a loop of 8 scalar __builtin_nontemporal_load calls. This will likely prevent the compiler from generating a single 256-bit vector load instruction, resulting in 8 separate 32-bit memory transactions instead of one. This de-vectorization will likely cause a significant performance degradation on gfx950, negating the benefit of the wider vector path.

Since this path is not yet hardware-tested, it's critical to ensure it's implemented efficiently. A true 256-bit non-temporal load, likely requiring inline GCN assembly to emit a global_load_dwordx8 with the slc or nt flag, should be used to achieve the expected performance.

Good point — the gfx950 path is not yet hardware-tested. However, based on ISA verification on gfx942, the compiler does coalesce sequential __builtin_nontemporal_load/store calls into vectorized instructions at -O3.

For the gfx950 256-bit path specifically: we expect the compiler to coalesce 8 × __builtin_nontemporal_load(int) into global_load_dwordx4 pairs (or ideally global_load_dwordx8 if the compiler recognizes the pattern). Since we cannot verify the gfx950 ISA output without MI355X hardware or a ROCm toolchain with gfx950 support, I will add a TODO comment noting that this should be verified when gfx950 hardware becomes available.

If the compiler does not coalesce to dwordx8 on gfx950, we can switch to inline ASM at that point. The __builtin_nontemporal approach was chosen because inline ASM with vector register constraints (v4i32/v8i32) causes invalid operand for instruction errors in hipcc when used across multiple template instantiations.

we expect the compiler to coalesce 8 × __builtin_nontemporal_load(int) into global_load_dwordx4 pairs

Just out of curiosity: any reason why this is not automatically translated to one global_load_dwordx8 if the hardware supports it?

Great question! The reason is hardware ISA generation differences:

• gfx942 (CDNA3 / MI300X): The maximum vector load width is 128-bit (global_load_dwordx4). There is no global_load_dwordx8 instruction in the gfx942 ISA. So 8× scalar loads → 2× global_load_dwordx4 is the best the compiler can do.

• gfx950 (CDNA4 / MI355X): The ISA does include global_load_dwordx8 (256-bit vector load). We expect the compiler to coalesce 8× scalar loads into a single global_load_dwordx8 on this target — but this needs hardware verification (added TODO).

This 128-bit vs 256-bit vector width is also the main reason MI300X peaks at ~4.5x speedup vs B300's 11.8x for this kernel. MI355X should close that gap significantly.

Reference: AMD CDNA4 ISA (gfx950) — see GLOBAL_LOAD_DWORDX8 instruction.

Now we have the answer from hardware testing! 🔬

The reason is simple: global_load/store_dwordx8 does not exist in the gfx950 ISA (as of ROCm 7.2 / AMD LLVM 22.0). The maximum vector load width is still 128-bit (dwordx4), same as gfx942.

Only s_buffer_load_dwordx8 (scalar buffer load) exists — there are no 256-bit vector load/store instructions.

So the compiler is doing the best it can: decomposing 256-bit into 2× dwordx4 with adjacent offsets (off and off offset:16). It's possible a future ROCm release could add vector dwordx8 support, but for now this is a hardware/ISA limitation, not a compiler limitation.

Similar to ld256_cs, this implementation uses a loop of 8 scalar __builtin_nontemporal_store calls instead of a single vectorized 256-bit store. This is unlikely to be compiled into a global_store_dwordx8 instruction and will likely result in 8 separate 32-bit memory transactions. This approach negates the performance benefit of using 256-bit vectors on gfx950.

To achieve the expected performance from doubling the vector width, a true vectorized non-temporal store should be implemented, which may require using inline GCN assembly.

Same reasoning as above — on gfx942, we verified that __builtin_nontemporal_store is coalesced into global_store_dwordx4 nt by the compiler. The gfx950 path should similarly benefit from compiler vectorization.

Will add a TODO comment for gfx950 ISA verification once hardware is available. If dwordx8 coalescing does not happen automatically, inline ASM can be added as a follow-up.

nit: remove/improve "or ROCm gfx950+" message in all the asserts, since 256-bit path is disabled for now on gfx950

This implementation uses four separate 32-bit non-temporal stores for a 128-bit operation. While the benchmarks show a performance gain from cache bypassing, this approach de-vectorizes the store operation, which may limit the performance uplift. A single 128-bit vectorized non-temporal store (global_store_dwordx4 with the nt flag) would be more efficient as it would avoid issuing four separate memory transactions.

For even better performance, consider using inline GCN assembly to emit a true vectorized non-temporal store. This would combine the benefits of both vectorization and cache bypassing.

Verified on MI300X — the compiler auto-coalesces the four scalar __builtin_nontemporal_store into a single global_store_dwordx4 ... nt instruction.

Actual ISA from hipcc -O3 --offload-arch=gfx942 --save-temps for ConcatMLAQKernel<bf16, 512>:

global_load_dwordx4  v[0:3],   v[12:13], off
global_load_dwordx4  v[4:7],   v[12:13], off offset:512
global_store_dwordx4 v[10:11], v[0:3],   off nt              ; single 128-bit NT store
global_store_dwordx4 v[10:11], v[4:7],   off offset:512 nt   ; single 128-bit NT store
global_load_dword    v12,      v[8:9],   off
global_store_dword   v[0:1],   v12,      off offset:1024 nt  ; 32-bit NT store

No de-vectorization occurs — the compiler generates the same instruction count as the plain path, just with the nt flag added. We also tested inline ASM with global_store_dwordx4 + v4i32 vector register constraints, but hipcc produces invalid operand for instruction errors in multi-template contexts. The builtin approach is both correct and portable.