DeepSeek/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM synced 2025-05-06 04:04:22 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Optimize performance

Browse Source
This commit is contained in:
Chenggang Zhao 2025-04-22 17:48:11 +08:00
parent 59884211ea
commit 07ef809d82
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
deep_gemm/include/deep_gemm/fp8_gemm.cuh
View File

@ -249,6 +249,9 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
// Persistently schedule over blocks
while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
// Preload TMA multicast validity, encouraged to use unified registers
bool is_tma_multicast_valid = __shfl_sync(0xffffffff, scheduler.is_tma_multicast_valid(m_block_idx), 0);
// Decide the number of scales B to load
DG_STATIC_ASSERT(SHAPE_N % 8 == 0, "Invalid shape N");
uint32_t num_former_iters = BLOCK_N / 8, num_full_iters = num_former_iters;
@ -276,7 +279,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
// Empty barrier arrival
auto empty_barrier_arrive = [&](int s) {
if (kNumTMAMulticast == 1 or not scheduler.is_tma_multicast_valid(m_block_idx)) {
if (kNumTMAMulticast == 1 or not is_tma_multicast_valid) {
lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive() : void();
} else {
lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive(lane_idx) : void();