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

Minor fix util function

Browse Source
This commit is contained in:
Chenggang Zhao 2025-02-28 09:46:38 +08:00
parent 6e10cba207
commit b69f630b91
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
deep_gemm/jit_kernels/utils.py
View File

@ -102,4 +102,5 @@ def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
# Normal layout requires transposing
aligned_x = torch.transpose(torch.empty((b, n, aligned_m), device=x.device, dtype=x.dtype), 1, 2)
aligned_x[:, :m, :] = x
aligned_x = aligned_x[:, :m, :]
return aligned_x.squeeze(0) if remove_dim else aligned_x