[PR #144] [CLOSED] [refact] use cute::tensor to define smem_d layout #263

New issue

Closed

opened 2026-10-26 21:15:30 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:15:30 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/144
Author: @Chtholly-Boss
Created: 7/28/2025
Status: ❌ Closed

Base: main ← Head: main

📝 Commits (1)

02c5df0 [refact] use cute::tensor to define smem_d layout

📊 Changes

1 file changed (+26 additions, -41 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh (+26 -41)

📄 Description

Manually index smem_d is tedious and hard to maintain. By using cute::tensor, we can leverage CuTe's layout system to automatically handle the indexing and swizzling of shared memory.

As we only need to perform swizzle on smem_d, make it as a CuTe Tensor will be helpful like follow:

  constexpr auto BBits = kSwizzleDMode == 128 ? 3 : (kSwizzleDMode == 64 ? 2 : 1);
  auto smem_d = cute::make_tensor(
      cute::make_smem_ptr(reinterpret_cast<cute::bfloat16_t*>(smem_buffer)),
      cute::composition(
          cute::Swizzle<BBits, 3, 3>{}, 
          cute::make_layout(
          cute::make_shape(BLOCK_N / TMA_D_BLOCK_N, BLOCK_M, TMA_D_BLOCK_N),
              cute::LayoutRight{}
          )
      )
  );

And instead of:

  uint8_t* smem_ptr = nullptr;
  if constexpr (kSwizzleDMode > 0) {
      // Calculate the swizzling atom offset and in-atom offset
      constexpr uint32_t kNumBankGroupBytes = 16;
      auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8);

      // Calculate the index of the bank group to be written in the atom
      auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes);

      // Reshape the atom in another view and swizzle
      //  - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)`
      //  - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)`
      constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8;
      auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8);
      auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8);
      col ^= row % (kSwizzleDMode / 16);

      // Add back into the base pointer
      // NOTES: think twice before modifying this, as changes may affect the number of instructions
      smem_ptr = reinterpret_cast<uint8_t*>(smem_d) +                // Base pointer
          warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) +            // Warp offset
          m_offset * kSwizzleDMode +                                 // Wave offset
          atom_offset * BLOCK_M * kSwizzleDMode +                    // Swizzle atom offset (constants)
          row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset
  } else {
      // No swizzling, just padding
      smem_ptr = reinterpret_cast<uint8_t*>(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * BLOCK_N + i * 8);
  }

We only need to write:

auto smem_ptr = &smem_d(i * 8 / TMA_D_BLOCK_N, m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx, i * 8 % TMA_D_BLOCK_N);

which is more readable and maintainable. As for Performance, I have tested on H20, and results show that this will not incur any overheads:
Before Modification

Testing GEMM:
 > Perf (m=  128, n= 2112, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch   38 us |   21 us |  181 TFLOPS |  778 GB/s
 > Perf (m=  128, n=24576, k= 1536, 1D2D, layout=NT, BF16, acc=0): launch   46 us |   44 us |  219 TFLOPS | 1001 GB/s
 > Perf (m=  128, n=32768, k=  512, 1D2D, layout=NT, BF16, acc=0): launch   46 us |   25 us |  174 TFLOPS | 1022 GB/s
 > Perf (m=  128, n= 7168, k=16384, 1D2D, layout=NT, BF16, acc=0): launch   44 us |  117 us |  257 TFLOPS | 1040 GB/s
 > Perf (m=  128, n= 4096, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch   46 us |   34 us |  223 TFLOPS |  932 GB/s
 > Perf (m=  128, n= 7168, k= 2048, 1D2D, layout=NT, BF16, acc=0): launch   43 us |   19 us |  202 TFLOPS |  902 GB/s
 > Perf (m= 4096, n= 2112, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch   42 us |  465 us |  266 TFLOPS |  135 GB/s
 > Perf (m= 4096, n=24576, k= 1536, 1D2D, layout=NT, BF16, acc=0): launch   46 us | 1160 us |  267 TFLOPS |  212 GB/s
 > Perf (m= 4096, n=32768, k=  512, 1D2D, layout=NT, BF16, acc=0): launch   43 us |  568 us |  242 TFLOPS |  506 GB/s
 > Perf (m= 4096, n= 7168, k=16384, 1D2D, layout=NT, BF16, acc=0): launch   50 us | 3583 us |  269 TFLOPS |   68 GB/s
 > Perf (m= 4096, n= 4096, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch   44 us |  923 us |  261 TFLOPS |  101 GB/s
 > Perf (m= 4096, n= 7168, k= 2048, 1D2D, layout=NT, BF16, acc=0): launch   44 us |  469 us |  257 TFLOPS |  175 GB/s

Testing m-grouped contiguous GEMM:
 > Perf (num_groups=4, m=35456, n= 4096, k= 7168, 1D2D, layout=NT): 7579 us |  275 TFLOPS |   88 GB/s
 > Perf (num_groups=4, m=36096, n= 7168, k= 2048, 1D2D, layout=NT): 4036 us |  263 TFLOPS |  162 GB/s
 > Perf (num_groups=8, m=32384, n= 4096, k= 7168, 1D2D, layout=NT): 6918 us |  275 TFLOPS |  107 GB/s
 > Perf (num_groups=8, m=31232, n= 7168, k= 2048, 1D2D, layout=NT): 3527 us |  260 TFLOPS |  179 GB/s

Testing m-grouped masked GEMM:
 > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D):  350 us |  176 TFLOPS |  131 GB/s
 > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D):  160 us |  206 TFLOPS |  207 GB/s
 > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D):  351 us |  194 TFLOPS |  219 GB/s
 > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D):  143 us |  202 TFLOPS |  319 GB/s
 > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D):  467 us |  139 TFLOPS |  289 GB/s
 > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D):  161 us |  165 TFLOPS |  458 GB/s

After modification:

Testing GEMM:
 > Perf (m=  128, n= 2112, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch   37 us |   21 us |  181 TFLOPS |  775 GB/s
 > Perf (m=  128, n=24576, k= 1536, 1D2D, layout=NT, BF16, acc=0): launch   45 us |   44 us |  218 TFLOPS |  998 GB/s
 > Perf (m=  128, n=32768, k=  512, 1D2D, layout=NT, BF16, acc=0): launch   44 us |   25 us |  174 TFLOPS | 1022 GB/s
 > Perf (m=  128, n= 7168, k=16384, 1D2D, layout=NT, BF16, acc=0): launch   44 us |  117 us |  258 TFLOPS | 1041 GB/s
 > Perf (m=  128, n= 4096, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch   48 us |   34 us |  223 TFLOPS |  930 GB/s
 > Perf (m=  128, n= 7168, k= 2048, 1D2D, layout=NT, BF16, acc=0): launch   43 us |   18 us |  203 TFLOPS |  908 GB/s
 > Perf (m= 4096, n= 2112, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch   44 us |  466 us |  266 TFLOPS |  135 GB/s
 > Perf (m= 4096, n=24576, k= 1536, 1D2D, layout=NT, BF16, acc=0): launch   46 us | 1165 us |  265 TFLOPS |  211 GB/s
 > Perf (m= 4096, n=32768, k=  512, 1D2D, layout=NT, BF16, acc=0): launch   44 us |  570 us |  241 TFLOPS |  504 GB/s
 > Perf (m= 4096, n= 7168, k=16384, 1D2D, layout=NT, BF16, acc=0): launch   46 us | 3568 us |  270 TFLOPS |   69 GB/s
 > Perf (m= 4096, n= 4096, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch   45 us |  917 us |  262 TFLOPS |  102 GB/s
 > Perf (m= 4096, n= 7168, k= 2048, 1D2D, layout=NT, BF16, acc=0): launch   44 us |  465 us |  259 TFLOPS |  177 GB/s

Testing m-grouped contiguous GEMM:
 > Perf (num_groups=4, m=35456, n= 4096, k= 7168, 1D2D, layout=NT): 7509 us |  277 TFLOPS |   89 GB/s
 > Perf (num_groups=4, m=36096, n= 7168, k= 2048, 1D2D, layout=NT): 4010 us |  264 TFLOPS |  163 GB/s
 > Perf (num_groups=8, m=32384, n= 4096, k= 7168, 1D2D, layout=NT): 6852 us |  278 TFLOPS |  108 GB/s
 > Perf (num_groups=8, m=31232, n= 7168, k= 2048, 1D2D, layout=NT): 3492 us |  263 TFLOPS |  181 GB/s

Testing m-grouped masked GEMM:
 > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D):  347 us |  177 TFLOPS |  132 GB/s
 > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D):  159 us |  207 TFLOPS |  208 GB/s
 > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D):  348 us |  195 TFLOPS |  220 GB/s
 > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D):  142 us |  203 TFLOPS |  321 GB/s
 > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D):  463 us |  140 TFLOPS |  291 GB/s
 > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D):  160 us |  166 TFLOPS |  461 GB/s

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/144 **Author:** [@Chtholly-Boss](https://github.com/Chtholly-Boss) **Created:** 7/28/2025 **Status:** ❌ Closed **Base:** `main` ← **Head:** `main` --- ### 📝 Commits (1) - [`02c5df0`](https://github.com/deepseek-ai/DeepGEMM/commit/02c5df0910967b0d93e42f5124024bcfc174925a) [refact] use cute::tensor to define smem_d layout ### 📊 Changes **1 file changed** (+26 additions, -41 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh` (+26 -41) </details> ### 📄 Description Manually index smem_d is tedious and hard to maintain. By using cute::tensor, we can leverage CuTe's layout system to automatically handle the indexing and swizzling of shared memory. As we only need to perform swizzle on smem_d, make it as a CuTe Tensor will be helpful like follow: ```cpp constexpr auto BBits = kSwizzleDMode == 128 ? 3 : (kSwizzleDMode == 64 ? 2 : 1); auto smem_d = cute::make_tensor( cute::make_smem_ptr(reinterpret_cast<cute::bfloat16_t*>(smem_buffer)), cute::composition( cute::Swizzle<BBits, 3, 3>{}, cute::make_layout( cute::make_shape(BLOCK_N / TMA_D_BLOCK_N, BLOCK_M, TMA_D_BLOCK_N), cute::LayoutRight{} ) ) ); ``` And instead of: ```cpp uint8_t* smem_ptr = nullptr; if constexpr (kSwizzleDMode > 0) { // Calculate the swizzling atom offset and in-atom offset constexpr uint32_t kNumBankGroupBytes = 16; auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8); // Calculate the index of the bank group to be written in the atom auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes); // Reshape the atom in another view and swizzle // - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)` // - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)` constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8; auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8); auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8); col ^= row % (kSwizzleDMode / 16); // Add back into the base pointer // NOTES: think twice before modifying this, as changes may affect the number of instructions smem_ptr = reinterpret_cast<uint8_t*>(smem_d) + // Base pointer warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) + // Warp offset m_offset * kSwizzleDMode + // Wave offset atom_offset * BLOCK_M * kSwizzleDMode + // Swizzle atom offset (constants) row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset } else { // No swizzling, just padding smem_ptr = reinterpret_cast<uint8_t*>(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * BLOCK_N + i * 8); } ``` We only need to write: ```cpp auto smem_ptr = &smem_d(i * 8 / TMA_D_BLOCK_N, m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx, i * 8 % TMA_D_BLOCK_N); ``` which is more readable and maintainable. As for Performance, I have tested on H20, and results show that this will not incur any overheads: Before Modification ``` Testing GEMM: > Perf (m= 128, n= 2112, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch 38 us | 21 us | 181 TFLOPS | 778 GB/s > Perf (m= 128, n=24576, k= 1536, 1D2D, layout=NT, BF16, acc=0): launch 46 us | 44 us | 219 TFLOPS | 1001 GB/s > Perf (m= 128, n=32768, k= 512, 1D2D, layout=NT, BF16, acc=0): launch 46 us | 25 us | 174 TFLOPS | 1022 GB/s > Perf (m= 128, n= 7168, k=16384, 1D2D, layout=NT, BF16, acc=0): launch 44 us | 117 us | 257 TFLOPS | 1040 GB/s > Perf (m= 128, n= 4096, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch 46 us | 34 us | 223 TFLOPS | 932 GB/s > Perf (m= 128, n= 7168, k= 2048, 1D2D, layout=NT, BF16, acc=0): launch 43 us | 19 us | 202 TFLOPS | 902 GB/s > Perf (m= 4096, n= 2112, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch 42 us | 465 us | 266 TFLOPS | 135 GB/s > Perf (m= 4096, n=24576, k= 1536, 1D2D, layout=NT, BF16, acc=0): launch 46 us | 1160 us | 267 TFLOPS | 212 GB/s > Perf (m= 4096, n=32768, k= 512, 1D2D, layout=NT, BF16, acc=0): launch 43 us | 568 us | 242 TFLOPS | 506 GB/s > Perf (m= 4096, n= 7168, k=16384, 1D2D, layout=NT, BF16, acc=0): launch 50 us | 3583 us | 269 TFLOPS | 68 GB/s > Perf (m= 4096, n= 4096, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch 44 us | 923 us | 261 TFLOPS | 101 GB/s > Perf (m= 4096, n= 7168, k= 2048, 1D2D, layout=NT, BF16, acc=0): launch 44 us | 469 us | 257 TFLOPS | 175 GB/s Testing m-grouped contiguous GEMM: > Perf (num_groups=4, m=35456, n= 4096, k= 7168, 1D2D, layout=NT): 7579 us | 275 TFLOPS | 88 GB/s > Perf (num_groups=4, m=36096, n= 7168, k= 2048, 1D2D, layout=NT): 4036 us | 263 TFLOPS | 162 GB/s > Perf (num_groups=8, m=32384, n= 4096, k= 7168, 1D2D, layout=NT): 6918 us | 275 TFLOPS | 107 GB/s > Perf (num_groups=8, m=31232, n= 7168, k= 2048, 1D2D, layout=NT): 3527 us | 260 TFLOPS | 179 GB/s Testing m-grouped masked GEMM: > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 350 us | 176 TFLOPS | 131 GB/s > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 160 us | 206 TFLOPS | 207 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 351 us | 194 TFLOPS | 219 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 143 us | 202 TFLOPS | 319 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 467 us | 139 TFLOPS | 289 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 161 us | 165 TFLOPS | 458 GB/s ``` After modification: ``` Testing GEMM: > Perf (m= 128, n= 2112, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch 37 us | 21 us | 181 TFLOPS | 775 GB/s > Perf (m= 128, n=24576, k= 1536, 1D2D, layout=NT, BF16, acc=0): launch 45 us | 44 us | 218 TFLOPS | 998 GB/s > Perf (m= 128, n=32768, k= 512, 1D2D, layout=NT, BF16, acc=0): launch 44 us | 25 us | 174 TFLOPS | 1022 GB/s > Perf (m= 128, n= 7168, k=16384, 1D2D, layout=NT, BF16, acc=0): launch 44 us | 117 us | 258 TFLOPS | 1041 GB/s > Perf (m= 128, n= 4096, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch 48 us | 34 us | 223 TFLOPS | 930 GB/s > Perf (m= 128, n= 7168, k= 2048, 1D2D, layout=NT, BF16, acc=0): launch 43 us | 18 us | 203 TFLOPS | 908 GB/s > Perf (m= 4096, n= 2112, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch 44 us | 466 us | 266 TFLOPS | 135 GB/s > Perf (m= 4096, n=24576, k= 1536, 1D2D, layout=NT, BF16, acc=0): launch 46 us | 1165 us | 265 TFLOPS | 211 GB/s > Perf (m= 4096, n=32768, k= 512, 1D2D, layout=NT, BF16, acc=0): launch 44 us | 570 us | 241 TFLOPS | 504 GB/s > Perf (m= 4096, n= 7168, k=16384, 1D2D, layout=NT, BF16, acc=0): launch 46 us | 3568 us | 270 TFLOPS | 69 GB/s > Perf (m= 4096, n= 4096, k= 7168, 1D2D, layout=NT, BF16, acc=0): launch 45 us | 917 us | 262 TFLOPS | 102 GB/s > Perf (m= 4096, n= 7168, k= 2048, 1D2D, layout=NT, BF16, acc=0): launch 44 us | 465 us | 259 TFLOPS | 177 GB/s Testing m-grouped contiguous GEMM: > Perf (num_groups=4, m=35456, n= 4096, k= 7168, 1D2D, layout=NT): 7509 us | 277 TFLOPS | 89 GB/s > Perf (num_groups=4, m=36096, n= 7168, k= 2048, 1D2D, layout=NT): 4010 us | 264 TFLOPS | 163 GB/s > Perf (num_groups=8, m=32384, n= 4096, k= 7168, 1D2D, layout=NT): 6852 us | 278 TFLOPS | 108 GB/s > Perf (num_groups=8, m=31232, n= 7168, k= 2048, 1D2D, layout=NT): 3492 us | 263 TFLOPS | 181 GB/s Testing m-grouped masked GEMM: > Perf (num_groups=1, expected_m_per_group=1024, n=4096, k=7168, 1D2D): 347 us | 177 TFLOPS | 132 GB/s > Perf (num_groups=1, expected_m_per_group=1024, n=7168, k=2048, 1D2D): 159 us | 207 TFLOPS | 208 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=4096, k=7168, 1D2D): 348 us | 195 TFLOPS | 220 GB/s > Perf (num_groups=2, expected_m_per_group= 512, n=7168, k=2048, 1D2D): 142 us | 203 TFLOPS | 321 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=4096, k=7168, 1D2D): 463 us | 140 TFLOPS | 291 GB/s > Perf (num_groups=4, expected_m_per_group= 256, n=7168, k=2048, 1D2D): 160 us | 166 TFLOPS | 461 GB/s ``` --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:15:30 +00:00