biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

P6: Fix TMA store PTX — remove .tile modifier, fix wait_group syntax

Browse Source
This commit is contained in:
biondizzle
2026-05-30 17:04:38 +00:00
parent e4ee9fdc9f
commit 829a5f93ce
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
dsv4/kernels/attention/fmha_tma.cuh
View File

@@ -213,8 +213,9 @@ __device__ __forceinline__ void tma_store_2d(
int coord_x,
int coord_y
) {
// cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [tensorMap, {coord}], [srcMem];
asm volatile(
"cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group "
"cp.async.bulk.tensor.2d.global.shared::cta.bulk_group "
"[%1, {%2, %3}], [%0];"
:: "r"(smem_src),
"l"(tma_desc),
@@ -237,14 +238,14 @@ __device__ __forceinline__ void tma_store_commit_group() {
* N=0 waits for all groups. N=1 waits for the most recent group.
*/
__device__ __forceinline__ void tma_store_wait_group(int n) {
asm volatile("cp.async.bulk.wait_group.read %0;" :: "r"(n) : "memory");
asm volatile("cp.async.bulk.wait_group %0;" :: "r"(n) : "memory");
}
/**
* Wait for all pending TMA store groups to complete.
*/
__device__ __forceinline__ void tma_store_wait_all() {
asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory");
asm volatile("cp.async.bulk.wait_group 0;" ::: "memory");
}
} // namespace dsv4::kernels::attention