7f2a703ed5
* Merge with private repo * Update README * Update README * Update README * Add PyTorch requirements * Fix sync scopes for MQA logits (#256) * Update README
26 lines
666 B
Plaintext
26 lines
666 B
Plaintext
#pragma once
|
|
|
|
#include <deep_gemm/common/exception.cuh>
|
|
|
|
namespace deep_gemm::ptx {
|
|
|
|
CUTLASS_DEVICE void warpgroup_arrive() {
|
|
asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory");
|
|
}
|
|
|
|
CUTLASS_DEVICE void warpgroup_commit_batch() {
|
|
asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory");
|
|
}
|
|
|
|
CUTLASS_DEVICE void warpgroup_fence_operand(float& reg) {
|
|
asm volatile("" : "+f"(reg) :: "memory");
|
|
}
|
|
|
|
template <int N>
|
|
CUTLASS_DEVICE void warpgroup_wait() {
|
|
DG_STATIC_ASSERT(N >= 0 and N <= 7, "WGMMA wait: N must be in range [0, 7]");
|
|
asm volatile("wgmma.wait_group.sync.aligned %0;\n" :: "n"(N) : "memory");
|
|
}
|
|
|
|
} // namespace deep_gemm::ptx
|