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

fix typo2

Browse Source
This commit is contained in:
yukuai 2025-04-22 15:19:29 +08:00
parent 69477036c0
commit d2369e8f30
2 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -292,7 +292,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
if constexpr (kNumTMAMulticast == 1) {
lane_idx == 0 ? empty_barriers[s]->arrive() : void();
} else {
auto target_cta_idx = scheduler.num_blocks_in_group != 1 ? lane_idx : cute::block_rank_in_cluster();
auto target_cta_idx = scheduler.is_block_in_complete_cluster(m_block_idx, n_block_idx) ? lane_idx : cute::block_rank_in_cluster();
lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive(target_cta_idx) : void();
}
};

7
deep_gemm/include/deep_gemm/scheduler.cuh
View File

@ -44,6 +44,13 @@ struct Scheduler {
}
}
__device__ __forceinline__ bool is_block_in_complete_cluster(const uint32_t& m_block_idx, uint32_t& n_block_idx) {
if (num_blocks_in_group == 1 and n_block_idx == kNumNBlocks - 1 and m_block_idx == num_aligned_m_blocks - 1
and num_aligned_m_blocks % 2 == 1 and kNumNBlocks % 2 == 1)
return false;
return true;
}
__device__ __forceinline__ bool is_tma_multicast_valid(const uint32_t& m_block_idx) {
if (num_blocks_in_group == 1)
return false;