chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Code lint

Browse Source
This commit is contained in:
Chenggang Zhao
2025-07-30 10:39:30 +08:00
parent 6bc75b549e
commit c50deed14c
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/jit_kernels/impls/smxx_layout.hpp
+5
View File
@@ -121,6 +121,7 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(const to
const auto& options = torch::TensorOptions().device(sf.device()).dtype(torch::kUInt8); const auto& options = torch::TensorOptions().device(sf.device()).dtype(torch::kUInt8);
auto padded = torch::zeros({num_groups, aligned_mn, aligned_k}, options); auto padded = torch::zeros({num_groups, aligned_mn, aligned_k}, options);
// ReSharper disable once CppExpressionWithoutSideEffects
padded.slice(1, 0, mn).slice(2, 0, k).copy_(ue8m0_tensor); padded.slice(1, 0, mn).slice(2, 0, k).copy_(ue8m0_tensor);
padded = padded.view(-1).view(torch::kInt32).view({num_groups, aligned_mn, aligned_k / 4}); padded = padded.view(-1).view(torch::kInt32).view({num_groups, aligned_mn, aligned_k / 4});
@@ -140,8 +141,10 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
at::TensorOptions().device(batched_sf.device()).dtype(torch::kInt)); at::TensorOptions().device(batched_sf.device()).dtype(torch::kInt));
// Launch the kernel // Launch the kernel
if (batched_sf.is_contiguous()) { if (batched_sf.is_contiguous()) {
// Fallback to slow PyTorch impl for non-supported cases
if ((mn * sf_k) % 4 != 0 and num_groups > 1) if ((mn * sf_k) % 4 != 0 and num_groups > 1)
return get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(sf); return get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(sf);
constexpr int block_mn = 48; constexpr int block_mn = 48;
constexpr int num_threads = 512; constexpr int num_threads = 512;
const TransposeAndPackFP32IntoUE8M0Runtime::Args& args = { const TransposeAndPackFP32IntoUE8M0Runtime::Args& args = {
@@ -157,8 +160,10 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
const auto& runtime = compiler->build("transpose_and_pack_fp32_into_ue8m0", code); const auto& runtime = compiler->build("transpose_and_pack_fp32_into_ue8m0", code);
TransposeAndPackFP32IntoUE8M0Runtime::launch(runtime, args); TransposeAndPackFP32IntoUE8M0Runtime::launch(runtime, args);
} else { } else {
// Fallback to slow PyTorch impl for non-supported cases
if (mn % 4 != 0 or num_groups > 1) if (mn % 4 != 0 or num_groups > 1)
return get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(sf); return get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(sf);
DG_HOST_ASSERT(mn % 4 == 0 and num_groups == 1); DG_HOST_ASSERT(mn % 4 == 0 and num_groups == 1);
DG_HOST_ASSERT(batched_sf.stride(1) == 1 and batched_sf.stride(2) == mn); DG_HOST_ASSERT(batched_sf.stride(1) == 1 and batched_sf.stride(2) == mn);