leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

v4.2 release. (#2587)

Browse Source 
* Fix default cluster callback values to 1 to avoid profiler failure when these values are not set in command line.

* v4.2 release.
This commit is contained in:
Junkai-Wu
2025-08-22 18:11:24 -04:00
committed by GitHub
parent 11cad1f67b
commit a49a78ffef
351 changed files with 28182 additions and 2032 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/arch/arch.h
+4
View File
@@ -109,6 +109,10 @@ struct Sm120 {
static int const kMinComputeCapability = 120;
};
struct Sm103 {
static int const kMinComputeCapability = 103;
};
/// Triggers a breakpoint on the device
CUTLASS_DEVICE
void device_breakpoint() {
include/cutlass/arch/barrier.h
+5 -2
View File
@@ -46,11 +46,13 @@
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED))
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110A_ENABLED))
#define CUTLASS_ARCH_TCGEN_ENABLED 1
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED))
#if (defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM103F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110F_ENABLED))
#define CUTLASS_ARCH_TCGEN_ENABLED 1
#endif
@@ -386,6 +388,7 @@ public:
//
CUTLASS_HOST_DEVICE
static void init(ValueType const* smem_ptr, uint32_t arrive_count) {
CUTLASS_ASSERT(arrive_count != 0 && "Arrive count must be non-zero");
#if CUDA_BARRIER_ENABLED
uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr);
asm volatile(
include/cutlass/arch/config.h
+68 -1
View File
@@ -128,6 +128,27 @@
/////////////////////////////////////////////////////////////////////////////////////////////////
// SM110 and SM110a only on 13.0 and above
#if !CUTLASS_CLANG_CUDA && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 0))
#define CUTLASS_ARCH_MMA_SM110_SUPPORTED 1
#if (!defined(CUTLASS_ARCH_MMA_SM110_ENABLED) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1100)
#define CUTLASS_ARCH_MMA_SM110_ENABLED 1
#if (!defined(CUTLASS_ARCH_MMA_SM110A_ENABLED) && defined(__CUDA_ARCH_FEAT_SM110_ALL))
#define CUTLASS_ARCH_MMA_SM110A_ENABLED 1
#endif
// SM110f
#if (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 0))
#define CUTLASS_ARCH_MMA_SM110F_SUPPORTED 1
#endif
#if (!defined(CUTLASS_ARCH_MMA_SM110F_ENABLED) && CUDA_ARCH_FAMILY(1100))
#define CUTLASS_ARCH_MMA_SM110F_ENABLED CUTLASS_ARCH_MMA_SM110F_SUPPORTED
#endif
#endif
#endif
/////////////////////////////////////////////////////////////////////////////////////////////////
// SM120 and SM120a
@@ -151,10 +172,56 @@
#endif
#endif
// SM103 and SM103a
#if !CUTLASS_CLANG_CUDA && (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9))
#define CUTLASS_ARCH_MMA_SM103_SUPPORTED 1
#if (!defined(CUTLASS_ARCH_MMA_SM103_ENABLED) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1030)
#define CUTLASS_ARCH_MMA_SM103_ENABLED 1
#if (!defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && defined(__CUDA_ARCH_FEAT_SM103_ALL))
#define CUTLASS_ARCH_MMA_SM103A_ENABLED 1
#endif
// SM103f
#if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9))
#define CUTLASS_ARCH_MMA_SM103F_SUPPORTED 1
#endif
#if (!defined(CUTLASS_ARCH_MMA_SM103F_ENABLED) && CUDA_ARCH_FAMILY(1030))
#define CUTLASS_ARCH_MMA_SM103F_ENABLED CUTLASS_ARCH_MMA_SM103F_SUPPORTED
#endif
#endif
#endif
// SM121 and SM121a
#if !CUTLASS_CLANG_CUDA && (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9))
#define CUTLASS_ARCH_MMA_SM121_SUPPORTED 1
#if (!defined(CUTLASS_ARCH_MMA_SM121_ENABLED) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1210)
#define CUTLASS_ARCH_MMA_SM121_ENABLED 1
#if (!defined(CUTLASS_ARCH_MMA_SM121A_ENABLED) &&\
(defined(__CUDA_ARCH_FEAT_SM121_ALL) || CUDA_ARCH_CONDITIONAL(1210)))
#define CUTLASS_ARCH_MMA_SM121A_ENABLED 1
#endif
// SM121f
#if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9))
#define CUTLASS_ARCH_MMA_SM121F_SUPPORTED 1
#endif
#if (!defined(CUTLASS_ARCH_MMA_SM121F_ENABLED) && CUDA_ARCH_FAMILY(1210))
#define CUTLASS_ARCH_MMA_SM121F_ENABLED CUTLASS_ARCH_MMA_SM121F_SUPPORTED
#endif
#endif
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120F_ENABLED))
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM103F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM110A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM121A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121F_ENABLED))
# define CUTLASS_ARCH_CLC_ENABLED
#endif
include/cutlass/arch/grid_dependency_control.h
+5 -1
View File
@@ -62,8 +62,12 @@
(defined(__CUDA_ARCH_FEAT_SM100_ALL) || CUDA_ARCH_FAMILY(1000))) || \
(__CUDA_ARCH__ == 1010 &&\
(defined(__CUDA_ARCH_FEAT_SM101_ALL) || CUDA_ARCH_FAMILY(1010))) || \
(__CUDA_ARCH__ == 1030 &&\
(defined(__CUDA_ARCH_FEAT_SM103_ALL) || CUDA_ARCH_FAMILY(1030))) || \
(__CUDA_ARCH__ == 1200 &&\
(defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_FAMILY(1200)))))
(defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_FAMILY(1200))) || \
(__CUDA_ARCH__ == 1210 &&\
(defined(__CUDA_ARCH_FEAT_SM121_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210)))))
#define CUTLASS_GDC_ENABLED
#endif
#endif
include/cutlass/arch/memory_sm80.h
+1
View File
@@ -40,6 +40,7 @@
#include "cutlass/arch/memory.h"
#include "cutlass/arch/memory_sm75.h"
#include "cutlass/arch/cache_operation.h"
#include "cutlass/arch/synclog.hpp"
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
#define CUDA_CP_ASYNC_ACTIVATED 1
include/cutlass/arch/mma_sm100.h
+2 -2
View File
@@ -33,8 +33,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/arch/mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sm70.h
+2 -2
View File
@@ -32,8 +32,8 @@
\brief Matrix multiply
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sm75.h
+2 -2
View File
@@ -33,8 +33,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/arch/wmma.h"
include/cutlass/arch/mma_sm80.h
+2 -3
View File
@@ -33,10 +33,9 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "mma.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/numeric_types.h"
include/cutlass/arch/mma_sm89.h
+2 -3
View File
@@ -34,10 +34,9 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "mma.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/numeric_types.h"
include/cutlass/arch/mma_sm90.h
+2 -2
View File
@@ -33,8 +33,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sparse_sm80.h
+2 -2
View File
@@ -34,8 +34,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sparse_sm89.h
+2 -2
View File
@@ -34,8 +34,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/reg_reconfig.h
+9
View File
@@ -36,13 +36,20 @@
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#else
#include <cstdint>
#endif
#ifndef CUDA_CTA_RECONFIG_ACTIVATED
#if defined(__CUDA_ARCH__) && __CUDACC_VER_MAJOR__ >= 12 && ( \
(__CUDA_ARCH__ == 900 && defined(__CUDA_ARCH_FEAT_SM90_ALL)) \
|| (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) \
|| (__CUDA_ARCH__ == 1010 && defined(__CUDA_ARCH_FEAT_SM101_ALL)) \
|| (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL)) \
|| (__CUDA_ARCH__ == 1200 && defined(__CUDA_ARCH_FEAT_SM120_ALL)) \
|| (__CUDA_ARCH__ == 1210 && defined(__CUDA_ARCH_FEAT_SM121_ALL)) \
)
#define CUDA_CTA_RECONFIG_ACTIVATED 1
#endif
@@ -50,7 +57,9 @@
#if defined(__CUDA_ARCH__) && __CUDACC_VER_MAJOR__ >= 12 && ( \
(__CUDA_ARCH__ == 1000 && CUDA_ARCH_FAMILY(1000)) \
|| (__CUDA_ARCH__ == 1010 && CUDA_ARCH_FAMILY(1010)) \
|| (__CUDA_ARCH__ == 1030 && CUDA_ARCH_FAMILY(1030)) \
|| (__CUDA_ARCH__ == 1200 && CUDA_ARCH_FAMILY(1200)) \
|| (__CUDA_ARCH__ == 1210 && CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210)) \
)
#define CUDA_CTA_RECONFIG_ACTIVATED 1
#endif
include/cutlass/arch/synclog.hpp
+27 -80
View File
@@ -35,9 +35,9 @@
#pragma once
#include "cutlass/detail/helper_macros.hpp"
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#include CUDA_STD_HEADER(cstdint)
#else
#include <cstdint>
#endif
@@ -120,44 +120,34 @@ constexpr uint32_t synclog_length_cluster_barrier_init = synclog_length_prefix +
constexpr bool synclog_enable_cluster_barrier_wait = true;
constexpr uint32_t synclog_header_cluster_barrier_wait = 6;
constexpr uint32_t synclog_length_cluster_barrier_wait = synclog_length_prefix + 4;
constexpr uint32_t synclog_length_cluster_barrier_wait = synclog_length_prefix + 2;
constexpr bool synclog_enable_cluster_barrier_test_wait = true;
constexpr uint32_t synclog_header_cluster_barrier_test_wait = 7;
constexpr uint32_t synclog_length_cluster_barrier_test_wait = synclog_length_prefix + 5;
constexpr uint32_t synclog_length_cluster_barrier_test_wait = synclog_length_prefix + 3;
constexpr bool synclog_enable_cluster_barrier_try_wait = true;
constexpr uint32_t synclog_header_cluster_barrier_try_wait = 8;
constexpr uint32_t synclog_length_cluster_barrier_try_wait = synclog_length_prefix + 4;
constexpr uint32_t synclog_length_cluster_barrier_try_wait = synclog_length_prefix + 2;
constexpr bool synclog_enable_cluster_barrier_arrive_cluster = true;
constexpr uint32_t synclog_header_cluster_barrier_arrive_cluster = 9;
constexpr uint32_t synclog_length_cluster_barrier_arrive_cluster = synclog_length_prefix + 5;
constexpr uint32_t synclog_length_cluster_barrier_arrive_cluster = synclog_length_prefix + 3;
constexpr bool synclog_enable_cluster_barrier_arrive = true;
constexpr uint32_t synclog_header_cluster_barrier_arrive = 10;
constexpr uint32_t synclog_length_cluster_barrier_arrive = synclog_length_prefix + 3;
constexpr uint32_t synclog_length_cluster_barrier_arrive = synclog_length_prefix + 1;
constexpr bool synclog_enable_cluster_barrier_invalidate = true;
constexpr uint32_t synclog_header_cluster_barrier_invalidate = 11;
constexpr uint32_t synclog_length_cluster_barrier_invalidate = synclog_length_prefix + 3;
constexpr uint32_t synclog_length_cluster_barrier_invalidate = synclog_length_prefix + 1;
constexpr bool synclog_enable_cluster_transaction_barrier_arrive_and_expect_tx = true;
constexpr uint32_t synclog_header_cluster_transaction_barrier_arrive_and_expect_tx = 12;
constexpr uint32_t synclog_length_cluster_transaction_barrier_arrive_and_expect_tx = synclog_length_prefix + 4;
constexpr uint32_t synclog_length_cluster_transaction_barrier_arrive_and_expect_tx = synclog_length_prefix + 2;
constexpr bool synclog_enable_cluster_transaction_barrier_arrive_and_expect_tx_cluster = true;
constexpr uint32_t synclog_header_cluster_transaction_barrier_arrive_and_expect_tx_cluster = 13;
constexpr uint32_t synclog_length_cluster_transaction_barrier_arrive_and_expect_tx_cluster = synclog_length_prefix + 6;
constexpr uint32_t synclog_length_cluster_transaction_barrier_arrive_and_expect_tx_cluster = synclog_length_prefix + 4;
constexpr bool synclog_enable_cluster_transaction_barrier_expect_transaction = true;
constexpr uint32_t synclog_header_cluster_transaction_barrier_expect_transaction = 14;
constexpr uint32_t synclog_length_cluster_transaction_barrier_expect_transaction = synclog_length_prefix + 4;
constexpr uint32_t synclog_length_cluster_transaction_barrier_expect_transaction = synclog_length_prefix + 2;
constexpr bool synclog_enable_cluster_transaction_barrier_complete_transaction = true;
constexpr uint32_t synclog_header_cluster_transaction_barrier_complete_transaction = 15;
constexpr uint32_t synclog_length_cluster_transaction_barrier_complete_transaction = synclog_length_prefix + 6;
constexpr uint32_t synclog_length_cluster_transaction_barrier_complete_transaction = synclog_length_prefix + 4;
constexpr bool synclog_enable_fence_barrier_init = true;
constexpr uint32_t synclog_header_fence_barrier_init = 16;
constexpr uint32_t synclog_length_fence_barrier_init = synclog_length_prefix + 0;
@@ -228,12 +218,11 @@ constexpr uint32_t synclog_length_wgmma_smem_smem = synclog_length_prefix + 4;
constexpr bool synclog_enable_cpasync_barrier_arrive = true;
constexpr uint32_t synclog_header_cpasync_barrier_arrive = 33;
constexpr uint32_t synclog_length_cpasync_barrier_arrive = synclog_length_prefix + 3;
constexpr uint32_t synclog_length_cpasync_barrier_arrive = synclog_length_prefix + 1;
CUTLASS_DEVICE
bool synclog_condition_emit() {
#if defined(__NVCC__) || (defined(__clang__) && defined(__CUDA__))
return threadIdx.x%NumThreadsPerWarp == 0 && threadIdx.y == 0 && threadIdx.z == 0 &&
return threadIdx.x % NumThreadsPerWarp == 0 && threadIdx.y == 0 && threadIdx.z == 0 &&
blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0;
#else
return 0;
@@ -272,17 +261,6 @@ void synclog_print_prefix(char const* header, uint32_t at) {
#endif
}
CUTLASS_DEVICE
uint64_t synclog_mbarrier_bits(uint32_t smem_addr) {
uint64_t bits = 0;
asm volatile (
"mbarrier.inval.shared::cta.b64 [%1];\n"
"ld.shared::cta.b64 %0, [%1];\n"
: "=l"(bits) : "r"(smem_addr)
);
return bits;
}
CUTLASS_DEVICE
void synclog_print_wgmma_desc(char const* str, uint32_t lo, uint32_t hi, char const* sep) {
CUTLASS_UNUSED(hi);
@@ -429,14 +407,11 @@ void synclog_emit_cluster_barrier_wait(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_barrier_wait) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_barrier_wait);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_barrier_wait, line);
to[synclog_length_prefix + 0] = smem_addr;
to[synclog_length_prefix + 1] = phase;
to[synclog_length_prefix + 2] = bits;
to[synclog_length_prefix + 3] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -453,15 +428,12 @@ void synclog_emit_cluster_barrier_test_wait(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_barrier_test_wait) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_barrier_test_wait);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_barrier_test_wait, line);
to[synclog_length_prefix + 0] = smem_addr;
to[synclog_length_prefix + 1] = phase;
to[synclog_length_prefix + 2] = pred;
to[synclog_length_prefix + 3] = bits;
to[synclog_length_prefix + 4] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -478,14 +450,11 @@ void synclog_emit_cluster_barrier_try_wait(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_barrier_try_wait) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_barrier_try_wait);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_barrier_try_wait, line);
to[synclog_length_prefix + 0] = smem_addr;
to[synclog_length_prefix + 1] = phase;
to[synclog_length_prefix + 2] = bits;
to[synclog_length_prefix + 3] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -502,15 +471,12 @@ void synclog_emit_cluster_barrier_arrive_cluster(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_barrier_arrive_cluster) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_barrier_arrive_cluster);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_barrier_arrive_cluster, line);
to[synclog_length_prefix + 0] = smem_addr;
to[synclog_length_prefix + 1] = cta_id;
to[synclog_length_prefix + 2] = pred;
to[synclog_length_prefix + 3] = bits;
to[synclog_length_prefix + 4] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -526,13 +492,10 @@ void synclog_emit_cluster_barrier_arrive(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_barrier_arrive) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_barrier_arrive);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_barrier_arrive, line);
to[synclog_length_prefix + 0] = smem_addr;
to[synclog_length_prefix + 1] = bits;
to[synclog_length_prefix + 2] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -546,13 +509,10 @@ void synclog_emit_cluster_barrier_invalidate(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_barrier_invalidate) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_barrier_invalidate);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_barrier_invalidate, line);
to[synclog_length_prefix + 0] = smem_addr;
to[synclog_length_prefix + 1] = bits;
to[synclog_length_prefix + 2] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -567,14 +527,11 @@ void synclog_emit_cluster_transaction_barrier_arrive_and_expect_tx(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_transaction_barrier_arrive_and_expect_tx) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_transaction_barrier_arrive_and_expect_tx);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_transaction_barrier_arrive_and_expect_tx, line);
to[synclog_length_prefix + 0] = smem_addr;
to[synclog_length_prefix + 1] = transaction_bytes;
to[synclog_length_prefix + 2] = bits;
to[synclog_length_prefix + 3] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -592,7 +549,6 @@ void synclog_emit_cluster_transaction_barrier_arrive_and_expect_tx_cluster(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_transaction_barrier_arrive_and_expect_tx_cluster) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_transaction_barrier_arrive_and_expect_tx_cluster);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_transaction_barrier_arrive_and_expect_tx_cluster, line);
@@ -600,8 +556,6 @@ void synclog_emit_cluster_transaction_barrier_arrive_and_expect_tx_cluster(
to[synclog_length_prefix + 1] = transaction_bytes;
to[synclog_length_prefix + 2] = cta_id;
to[synclog_length_prefix + 3] = pred;
to[synclog_length_prefix + 4] = bits;
to[synclog_length_prefix + 5] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -619,14 +573,11 @@ void synclog_emit_cluster_transaction_barrier_expect_transaction(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_transaction_barrier_expect_transaction) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_transaction_barrier_expect_transaction);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_transaction_barrier_expect_transaction, line);
to[synclog_length_prefix + 0] = smem_addr;
to[synclog_length_prefix + 1] = transaction_bytes;
to[synclog_length_prefix + 2] = bits;
to[synclog_length_prefix + 2] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -644,7 +595,6 @@ void synclog_emit_cluster_transaction_barrier_complete_transaction(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cluster_transaction_barrier_complete_transaction) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cluster_transaction_barrier_complete_transaction);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cluster_transaction_barrier_complete_transaction, line);
@@ -652,8 +602,6 @@ void synclog_emit_cluster_transaction_barrier_complete_transaction(
to[synclog_length_prefix + 1] = dst_cta_id;
to[synclog_length_prefix + 2] = transaction_bytes;
to[synclog_length_prefix + 3] = pred;
to[synclog_length_prefix + 4] = bits;
to[synclog_length_prefix + 5] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -977,13 +925,10 @@ void synclog_emit_cpasync_barrier_arrive(
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_cpasync_barrier_arrive) return;
if (!synclog_condition_emit()) return;
uint64_t bits = synclog_mbarrier_bits(smem_addr);
uint32_t* to = synclog_alloc(synclog_length_cpasync_barrier_arrive);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_cpasync_barrier_arrive, line);
to[synclog_length_prefix + 0] = smem_addr;
to[synclog_length_prefix + 1] = bits;
to[synclog_length_prefix + 2] = bits >> 32;
#else
CUTLASS_UNUSED(line);
CUTLASS_UNUSED(smem_addr);
@@ -1054,7 +999,7 @@ void synclog_print() {
if (header == synclog_header_cluster_barrier_wait) {
synclog_print_prefix("cluster_barrier_wait", at);
at += synclog_length_cluster_barrier_wait;
printf("smem_addr=%u phase=%u", synclog_buf[at-4], synclog_buf[at-3]);
printf("smem_addr=%u phase=%u\n", synclog_buf[at-2], synclog_buf[at-1]);
continue;
}
}
@@ -1062,7 +1007,7 @@ void synclog_print() {
if (header == synclog_header_cluster_barrier_test_wait) {
synclog_print_prefix("cluster_barrier_test_wait", at);
at += synclog_length_cluster_barrier_test_wait;
printf("smem_addr=%u phase=%u pred=%u", synclog_buf[at-5], synclog_buf[at-4], synclog_buf[at-3]);
printf("smem_addr=%u phase=%u pred=%u\n", synclog_buf[at-3], synclog_buf[at-2], synclog_buf[at-1]);
continue;
}
}
@@ -1070,7 +1015,7 @@ void synclog_print() {
if (header == synclog_header_cluster_barrier_try_wait) {
synclog_print_prefix("cluster_barrier_try_wait", at);
at += synclog_length_cluster_barrier_try_wait;
printf("smem_addr=%u phase=%u", synclog_buf[at-4], synclog_buf[at-3]);
printf("smem_addr=%u phase=%u\n", synclog_buf[at-2], synclog_buf[at-1]);
continue;
}
}
@@ -1078,7 +1023,7 @@ void synclog_print() {
if (header == synclog_header_cluster_barrier_arrive_cluster) {
synclog_print_prefix("cluster_barrier_arrive_cluster", at);
at += synclog_length_cluster_barrier_arrive_cluster;
printf("smem_addr=%u cta_id=%u pred=%u", synclog_buf[at-5], synclog_buf[at-4], synclog_buf[at-3]);
printf("smem_addr=%u cta_id=%u pred=%u\n", synclog_buf[at-3], synclog_buf[at-2], synclog_buf[at-1]);
continue;
}
}
@@ -1086,7 +1031,7 @@ void synclog_print() {
if (header == synclog_header_cluster_barrier_arrive) {
synclog_print_prefix("cluster_barrier_arrive", at);
at += synclog_length_cluster_barrier_arrive;
printf("smem_addr=%u", synclog_buf[at-3]);
printf("smem_addr=%u\n", synclog_buf[at-1]);
continue;
}
}
@@ -1094,7 +1039,7 @@ void synclog_print() {
if (header == synclog_header_cluster_barrier_invalidate) {
synclog_print_prefix("cluster_barrier_invalidate", at);
at += synclog_length_cluster_barrier_invalidate;
printf("smem_addr=%u", synclog_buf[at-3]);
printf("smem_addr=%u\n", synclog_buf[at-1]);
continue;
}
}
@@ -1102,7 +1047,7 @@ void synclog_print() {
if (header == synclog_header_cluster_transaction_barrier_arrive_and_expect_tx) {
synclog_print_prefix("cluster_transaction_barrier_arrive_and_expect_tx", at);
at += synclog_length_cluster_transaction_barrier_arrive_and_expect_tx;
printf("smem_addr=%u transaction_bytes=%u", synclog_buf[at-4], synclog_buf[at-3]);
printf("smem_addr=%u transaction_bytes=%u\n", synclog_buf[at-2], synclog_buf[at-1]);
continue;
}
}
@@ -1110,7 +1055,7 @@ void synclog_print() {
if (header == synclog_header_cluster_transaction_barrier_arrive_and_expect_tx_cluster) {
synclog_print_prefix("cluster_transaction_barrier_arrive_and_expect_tx_cluster", at);
at += synclog_length_cluster_transaction_barrier_arrive_and_expect_tx_cluster;
printf("smem_addr=%u transaction_bytes=%u cta_id=%u pred=%u", synclog_buf[at-6], synclog_buf[at-5], synclog_buf[at-4], synclog_buf[at-3]);
printf("smem_addr=%u transaction_bytes=%u cta_id=%u pred=%u\n", synclog_buf[at-4], synclog_buf[at-3], synclog_buf[at-2], synclog_buf[at-1]);
continue;
}
}
@@ -1118,7 +1063,7 @@ void synclog_print() {
if (header == synclog_header_cluster_transaction_barrier_expect_transaction) {
synclog_print_prefix("cluster_transaction_barrier_expect_transaction", at);
at += synclog_length_cluster_transaction_barrier_expect_transaction;
printf("smem_addr=%u transaction_bytes=%u", synclog_buf[at-4], synclog_buf[at-3]);
printf("smem_addr=%u transaction_bytes=%u\n", synclog_buf[at-2], synclog_buf[at-1]);
continue;
}
}
@@ -1126,7 +1071,7 @@ void synclog_print() {
if (header == synclog_header_cluster_transaction_barrier_complete_transaction) {
synclog_print_prefix("cluster_transaction_barrier_complete_transaction", at);
at += synclog_length_cluster_transaction_barrier_complete_transaction;
printf("smem_addr=%u dst_cta_id=%u transaction_bytes=%u pred=%u", synclog_buf[at-6], synclog_buf[at-5], synclog_buf[at-4], synclog_buf[at-3]);
printf("smem_addr=%u dst_cta_id=%u transaction_bytes=%u pred=%u\n", synclog_buf[at-4], synclog_buf[at-3], synclog_buf[at-2], synclog_buf[at-1]);
continue;
}
}
@@ -1283,7 +1228,7 @@ void synclog_print() {
if (header == synclog_header_cpasync_barrier_arrive) {
synclog_print_prefix("cpasync_barrier_arrive", at);
at += synclog_length_cpasync_barrier_arrive;
printf("smem_addr=%u", synclog_buf[at-3]);
printf("smem_addr=%u\n", synclog_buf[at-1]);
continue;
}
}
@@ -1302,6 +1247,7 @@ void synclog_print() {
////////////////////////////////////////////////////////////////////////////////////////////////////
#if defined(CUTLASS_ENABLE_SYNCLOG)
#undef __syncthreads
#define __syncthreads() do {\
@@ -1318,6 +1264,7 @@ void synclog_print() {
} while (0)
#endif // defined(CUTLASS_ENABLE_SYNCLOG)
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace arch
include/cutlass/arch/wmma_sm70.h
+2 -2
View File
@@ -33,8 +33,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/layout/matrix.h"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/arch/wmma_sm72.h
+2 -2
View File
@@ -33,8 +33,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/layout/matrix.h"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/arch/wmma_sm75.h
+2 -2
View File
@@ -33,8 +33,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/layout/matrix.h"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/cluster_launch.hpp
+2 -1
View File
@@ -39,9 +39,10 @@
#include "cutlass/cutlass.h"
#include "cutlass/trace.h"
#include <cute/arch/cluster_sm100.hpp>
#include "cutlass/arch/synclog.hpp"
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#include CUDA_STD_HEADER(type_traits)
#else
#include <type_traits>
#include <cstdio>
include/cutlass/complex.h
+2 -4
View File
@@ -34,14 +34,12 @@
#include <cuComplex.h>
#include <cuda_fp16.h>
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#include CUDA_STD_HEADER(cstdint)
#else
#include <cstdint>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/functional.h"
#include "cutlass/platform/platform.h"
#include "cutlass/real.h"
include/cutlass/conv/kernel/sm100_implicit_gemm_tma_warpspecialized.hpp
+12 -3
View File
@@ -279,10 +279,19 @@ public:
implementable &= cutlass::detail::preferred_cluster_can_implement<AtomThrShapeMNK>(args.hw_info.cluster_shape, args.hw_info.cluster_shape_fallback);
}
if constexpr (is_grouped_wgrad) {
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, args.hw_info.cluster_shape);
auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, args.hw_info.cluster_shape_fallback);
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, args.hw_info.cluster_shape);
auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, args.hw_info.cluster_shape_fallback);
// implicit gemm B tile can be small for conv, ensure multicast smem offsets are 128B aligned
int multicast_b_bits = (size<1>(TileShape{}) * size<2>(TileShape{}) / size<0>(cluster_shape)) * sizeof_bits_v<ElementB>;
int multicast_b_fallback_bits = (size<1>(TileShape{}) * size<2>(TileShape{}) / size<0>(cluster_shape_fallback)) * sizeof_bits_v<ElementB>;
implementable &= multicast_b_bits % (128*8) == 0 && multicast_b_fallback_bits % (128*8) == 0;
if (not implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: multicast size too large for B tile\n");
return false;
}
if constexpr (is_grouped_wgrad) {
implementable &= size<0>(cluster_shape) == 1 && size<0>(cluster_shape_fallback) == 1;
if (!implementable) {
include/cutlass/coord.h
+2 -4
View File
@@ -33,15 +33,13 @@
*/
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#include CUDA_STD_HEADER(cstdint)
#else
#include <cstdint>
#endif
#include "cutlass/cutlass.h"
namespace cutlass {
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/cutlass.h
+6 -1
View File
@@ -35,9 +35,14 @@
#pragma once
#include "cutlass/arch/synclog.hpp"
#include "cutlass/detail/helper_macros.hpp"
#if (__CUDACC_VER_MAJOR__ >= 13)
#define CUDA_STD_HEADER(header) <cccl/cuda/std/header>
#else
#define CUDA_STD_HEADER(header) <cuda/std/header>
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
include/cutlass/detail/blockwise_scale_layout.hpp
+9 -3
View File
@@ -32,7 +32,7 @@
/*! \file
\brief Block Wise Scale configs specific for SM100 Blockwise/Groupwise MMA
\brief Blockwise Scale configs specific for Blockwise/Groupwise MMA
*/
#pragma once
@@ -41,6 +41,7 @@
#include "cute/int_tuple.hpp"
#include "cute/atom/mma_traits_sm100.hpp"
#include "cute/arch/mma_sm90.hpp"
namespace cutlass::detail{
@@ -270,8 +271,13 @@ struct RuntimeBlockwiseScaleConfig {
};
// Sm90 only supports MN major for SFA and SFB for now
template<int SFVecSizeM, int SFVecSizeN, int SFVecSizeK>
using Sm90BlockwiseScaleConfig = Sm1xxBlockwiseScaleConfig<SFVecSizeM, SFVecSizeN, SFVecSizeK>;
template<int SFVecSizeM, int SFVecSizeN, int SFVecSizeK, cute::GMMA::Major majorSFA = cute::GMMA::Major::MN, cute::GMMA::Major majorSFB = cute::GMMA::Major::MN>
using Sm90BlockwiseScaleConfig = Sm1xxBlockwiseScaleConfig<
SFVecSizeM,
SFVecSizeN,
SFVecSizeK,
majorSFA == cute::GMMA::Major::MN ? UMMA::Major::MN : UMMA::Major::K,
majorSFB == cute::GMMA::Major::MN ? UMMA::Major::MN : UMMA::Major::K>;
template<int SFVecSizeM, int SFVecSizeN, int SFVecSizeK, UMMA::Major majorSFA = UMMA::Major::MN, UMMA::Major majorSFB = UMMA::Major::MN>
using Sm100BlockwiseScaleConfig = Sm1xxBlockwiseScaleConfig<SFVecSizeM, SFVecSizeN, SFVecSizeK, majorSFA, majorSFB>;
include/cutlass/detail/collective/mixed_input_utils.hpp
+232
View File
@@ -483,6 +483,7 @@ void LayoutAwareConvert(
decltype(dst_layout)>::convert(src_vm, dst_vm);
}
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -618,6 +619,32 @@ public:
}
}
static constexpr uint32_t
compute_tma_transaction_bytes_extra_transform() {
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
return 0;
}
else if constexpr (ModeHasScales) {
constexpr uint32_t scale_tx_bytes = cutlass::bits_to_bytes(size<0>(filter_zeros(SmemLayoutScale{})) * size<1>(filter_zeros(SmemLayoutScale{})) * static_cast<uint32_t>(cute::sizeof_bits_v<ElementScale>));
static_assert(scale_tx_bytes % 128 == 0, "Each scale stage must be 128B aligned."); // required by TMA
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
return scale_tx_bytes;
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
// Scale and zero share smem layout
constexpr uint32_t zero_tx_bytes = cutlass::bits_to_bytes(size<0>(filter_zeros(SmemLayoutScale{})) * size<1>(filter_zeros(SmemLayoutScale{})) * static_cast<uint32_t>(cute::sizeof_bits_v<ElementZero>));
static_assert(zero_tx_bytes % 128 == 0, "Each zero stage must be 128B aligned."); // required by TMA
return scale_tx_bytes + zero_tx_bytes;
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in tma transaction bytes computation.");
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Type not handled in tma transaction bytes computation.");
}
}
/// Utilities to copy A and extra inputs from smem to RF
template <class SmemTiledCopyA,
class TensorASmemView,
@@ -663,6 +690,43 @@ public:
}
}
/// (Designed for separate transform pipeline in Blackwell)
/// Utilities to copy extra inputs from smem to RF
template <class... Ts>
CUTLASS_DEVICE
static void copy_scale_zeros_for_transform(
cute::tuple<Ts...> & partitioned_transform_extra_info,
int load2transform_consumer_index) {
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
// nothing to do
}
else if constexpr (ModeHasScales) {
auto smem_tiled_copy_S = cute::get<0>(partitioned_transform_extra_info);
auto&& scales = cute::get<1>(partitioned_transform_extra_info);
using ScaleType = decltype(scales);
auto tSrS = make_tensor(static_cast<ScaleType&&>(scales).data(), scales.layout());
auto tSsS = cute::get<2>(partitioned_transform_extra_info);
copy(smem_tiled_copy_S, tSsS(_,_,_,_,load2transform_consumer_index), tSrS);
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
// Nothing extra to do
} else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
auto&& zeros = cute::get<3>(partitioned_transform_extra_info);
using ZeroType = decltype(zeros);
auto tZrZ = make_tensor(static_cast<ZeroType&&>(zeros).data(), zeros.layout());
auto tZsZ = cute::get<4>(partitioned_transform_extra_info);
copy(smem_tiled_copy_S, tZsZ(_,_,_,_,load2transform_consumer_index), tZrZ);
} else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
}
// Helper functions to select packing for conversion
template <class SrcType,
class DstType,
@@ -902,6 +966,115 @@ public:
}
}
/// (Designed for separate transform pipeline in Blackwell)
/// Utilities to dequantize A.
template <class EngineIn,
class EngineOut,
class LayoutIn,
class LayoutOut,
class... Ts>
CUTLASS_DEVICE
static void dequantize_A_kblock_for_transform(
Tensor<EngineIn, LayoutIn> const& tArA,
Tensor<EngineOut, LayoutOut>& tArACompute,
cute::tuple<Ts...> const& partitioned_extra_info,
int const k_block) {
static_assert(is_rmem<EngineIn>::value, "Input tensor for A conversion must come from registers");
static_assert(is_rmem<EngineOut>::value, "Output tensor for A conversion must come from registers");
static_assert(cosize_v<LayoutIn> == cosize_v<LayoutOut>);
static_assert(size_v<LayoutIn> == cosize_v<LayoutIn>);
static_assert(size_v<LayoutOut> == cosize_v<LayoutOut>);
using SrcType = typename EngineIn::value_type;
using DstType = typename EngineOut::value_type;
auto src = tArA(_, _, _, k_block);
auto dst = tArACompute(_, _, _, k_block);
constexpr int num_elements = decltype(size(src))::value;
constexpr int pack = decltype(select_packing<SrcType, DstType, num_elements>::value())::value;
using Converter = cutlass::NumericArrayConverter<DstType, SrcType, pack, cutlass::FloatRoundStyle::round_to_nearest>;
using SrcArray = cutlass::Array<SrcType, pack>;
using DstArray = cutlass::Array<DstType, pack>;
constexpr int DstElementsPerReg = 32 / sizeof_bits_v<DstType>;
using RegArray = cutlass::AlignedArray<uint32_t, pack / DstElementsPerReg, sizeof(DstArray)>;
auto src_arr = recast<SrcArray>(src);
auto dst_arr = recast<DstArray>(dst);
Tensor src_vm = cute::group_modes<1,-1>(cute::zipped_divide(src, pack));
Tensor dst_vm = cute::group_modes<1,-1>(cute::zipped_divide(dst, pack));
cute::transform(src_arr, dst_arr, Converter::convert);
if constexpr (ModeHasScales) {
auto const& scales = cute::get<1>(partitioned_extra_info)(_,_,_,k_block);
CUTE_STATIC_ASSERT_V(size(src) == size(scales));
if constexpr (is_same_v<DstType, ElementScale>) {
using ScaleArray = cutlass::Array<ElementScale, pack>;
auto scale_arr = recast<ScaleArray>(filter_zeros(scales));
if constexpr (is_same_v<DstType, cutlass::bfloat16_t>){
Tensor dst_vm = cute::group_modes<1,-1>(cute::zipped_divide(dst, pack));
Tensor scales_vm = cute::group_modes<1,-1>(cute::zipped_divide(scales, pack));
for (int i = 0; i < size<1>(dst_vm); ++i){
auto&& r = cute::recast<RegArray>(dst_vm(_,i))(0);
auto&& scale_reg = cute::recast<RegArray>(scales_vm(_,i))(0);
CUTLASS_PRAGMA_UNROLL
for (size_t ii = 0; ii < RegArray::kElements; ++ii) {
__nv_bfloat162& bf16x2_val = reinterpret_cast<__nv_bfloat162&>(r[ii]);
bf16x2_val = __hmul2(bf16x2_val,
reinterpret_cast<const __nv_bfloat162&>(scale_reg[ii]));
}
}
}
else{
cute::transform(dst_arr, scale_arr, dst_arr, cute::multiplies{});
}
}
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
// Do Nothing
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
static_assert(is_same_v<ElementScale, ElementZero>, "ElementScale and ElementZero must be the same.");
auto const& zeros = cute::get<3>(partitioned_extra_info)(_,_,_,k_block);
CUTE_STATIC_ASSERT_V(size(src) == size(zeros));
if constexpr (is_same_v<DstType, ElementZero>) {
using ZeroArray = cutlass::Array<ElementZero, pack>;
auto zero_arr = recast<ZeroArray>(filter_zeros(zeros));
if constexpr (is_same_v<DstType, cutlass::bfloat16_t>) {
Tensor zeros_vm = cute::group_modes<1,-1>(cute::zipped_divide(zeros, pack));
for (int i = 0; i < size<1>(dst_vm); ++i){
auto&& r = cute::recast<RegArray>(dst_vm(_,i))(0);
auto&& zero_reg = cute::recast<RegArray>(zeros_vm(_,i))(0);
CUTLASS_PRAGMA_UNROLL
for (size_t ii = 0; ii < RegArray::kElements; ++ii) {
__nv_bfloat162& bf16x2_val = reinterpret_cast<__nv_bfloat162&>(r[ii]);
bf16x2_val = __hadd2(bf16x2_val,
reinterpret_cast<const __nv_bfloat162&>(zero_reg[ii]));
}
}
}
else{
cute::transform(dst_arr, zero_arr, dst_arr, cute::plus{});
}
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled for input partitioning.");
}
}
}
/// Utilities for any additional inputs inside of the TMA load
template <
@@ -998,6 +1171,65 @@ public:
}
}
template <
class TiledMma,
class TiledCopy,
class TensorStorage
>
CUTLASS_DEVICE
static auto partition_extra_transform_info(
TiledMma const& tiled_mma,
TiledCopy const& smem_tiled_copy_S,
TensorStorage& shared_storage) {
if constexpr (KernelConversionMode == ConversionMode::DirectConvert) {
// nothing to do
return cute::make_tuple();
}
else if constexpr (ModeHasScales) {
ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{}));
auto smem_thr_copy_S = smem_tiled_copy_S.get_slice(threadIdx.x % 128);
Tensor sS = make_tensor(make_smem_ptr(shared_storage.input.smem_scale.begin()), SmemLayoutScale{}); // (BLK_M,BLK_SCALE_K,PIPE)
Tensor tCsS = cta_mma.partition_A(sS);
Tensor tSsS = smem_thr_copy_S.partition_S(tCsS);
Tensor tSrS = make_tensor<ElementScale>(tSsS(_,_,_,_,0).shape());
#if 0
if(cute::thread(128, 0)){
print("sS: ");print(sS);print("\n");
print("tSsS: ");print(tSsS);print("\n");
print("tSrS: ");print(tSrS);print("\n");
}
#endif
if constexpr (KernelConversionMode == ConversionMode::ConvertAndScale) {
return cute::make_tuple(smem_tiled_copy_S, tSrS, tSsS);
}
else if constexpr (KernelConversionMode == ConversionMode::ConvertAndScaleWithZero) {
Tensor sZ = make_tensor(make_smem_ptr(shared_storage.input.smem_zero.begin()), SmemLayoutScale{});// (BLK_M,BLK_SCALE_K,PIPE)
Tensor tCsZ = cta_mma.partition_A(sZ);
Tensor tZsZ = smem_thr_copy_S.partition_S(tCsZ);
Tensor tZrZ = make_tensor<ElementZero>(tZsZ(_,_,_,_,0).shape());
#if 0
if(cute::thread(128, 0)){
print("sS: ");print(sS);print("\n");
print("tSsS: ");print(tSsS);print("\n");
print("tSrS: ");print(tSrS);print("\n");
print("sZ: ");print(sZ);print("\n");
print("tZsZ: ");print(tZsZ);print("\n");
print("tZrZ: ");print(tZrZ);print("\n");
}
#endif
return cute::make_tuple(smem_tiled_copy_S, tSrS, tSsS, tZrZ, tZsZ);
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
}
else {
static_assert(cutlass::detail::dependent_false<KernelSchedule>, "Conversion mode not handled in A -> RF path.");
}
}
/// Returns the tiled copy and copy views for the extra inputs.
template <class TiledMma, class... Ts>
CUTLASS_DEVICE
include/cutlass/detail/collective/sm103_kernel_type.hpp
+45
View File
@@ -0,0 +1,45 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief Kernel type definitions specific for SM103 BlockScaled MMA
*/
#pragma once
namespace cutlass::sm103::detail {
enum class KernelPrefetchType {
TmaPrefetch, // TMA Prefetch (is the default version)
Disable // Disable Prefetch
};
} // namespace cutlass::sm103::detail
include/cutlass/detail/sm100_mixed_dtype_blockwise_layout.hpp
+182
View File
@@ -0,0 +1,182 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief Block Wise Scale configs specific for SM100 Blockwise/Groupwise MMA
*/
#pragma once
#include "cutlass/layout/matrix.h"
#include "cute/int_tuple.hpp"
#include "cute/atom/mma_traits_sm100.hpp"
namespace cutlass::detail{
/////////////////////////////////////////////////////////////////////////////////////////////////
using namespace cute;
template<int SFVecSizeMN, int SFVecSizeK, UMMA::Major majorSFA = UMMA::Major::MN>
struct Sm100MixedInputBlockwiseScaleConfig {
using ShapeScale = Shape<Shape<Int<SFVecSizeMN>, int32_t>, Shape<Int<SFVecSizeK>, int32_t>, int32_t>;
using StrideScale = conditional_t<majorSFA == UMMA::Major::MN,
Stride<Stride<_0,_1>,Stride<_0,int32_t>, int32_t>,
Stride<Stride<_0,int32_t>,Stride<_0,_1>, int32_t>>;
using LayoutScale = Layout<ShapeScale, StrideScale>;
CUTE_HOST_DEVICE
static constexpr auto
deduce_layout_scale() {
return LayoutScale{};
}
template<class CtaShape_MN_K>
CUTE_HOST_DEVICE
static constexpr auto
smem_atom_layout_scale(CtaShape_MN_K cta_shape_mn_k) {
static_assert(cute::is_static_v<CtaShape_MN_K>, "Expect static CTA shape");
int constexpr size_MN = cute::get<0>(CtaShape_MN_K{});
int constexpr size_K = cute::get<1>(CtaShape_MN_K{});
int constexpr SmemSizeMN = (SFVecSizeMN < size_MN)
? SFVecSizeMN
: size_MN;
int constexpr SmemSizeK = (SFVecSizeK < size_K)
? SFVecSizeK
: size_K;
int constexpr div_MN = cute::ceil_div(size_MN, SmemSizeMN);
int constexpr div_K = cute::ceil_div(size_K, SmemSizeK);
auto strides = [&]() CUTLASS_LAMBDA_FUNC_INLINE {
if constexpr (majorSFA == UMMA::Major::MN) {
return make_stride(make_stride(_0{}, _1{}), make_stride(_0{}, Int<div_MN>{}));
}
else {
return make_stride(make_stride(_0{}, Int<div_K>{}), make_stride(_0{}, _1{}));
}
}();
return make_layout(
make_shape(make_shape(Int<SmemSizeMN>{}, Int<div_MN>{}),
make_shape(Int<SmemSizeK>{}, Int<div_K>{})),
strides
);
}
// The following function is provided for user fill dynamic problem size to the layout_SFA.
template <class ScaledInputDim>
CUTE_HOST_DEVICE
static constexpr auto
tile_atom_to_shape_scale(ScaledInputDim scale_input_dims) {
const auto scale_input_dims_MNKL = append<3>(scale_input_dims, 1);
auto strides = [&]() CUTLASS_LAMBDA_FUNC_INLINE {
auto [MN, K, L] = scale_input_dims_MNKL;
if constexpr (majorSFA == UMMA::Major::MN) {
return make_stride(make_stride(_0{}, _1{}), make_stride(_0{}, cute::ceil_div(MN, SFVecSizeMN)));
}
else {
return make_stride(make_stride(_0{}, cute::ceil_div(K, SFVecSizeK)), make_stride(_0{}, _1{}));
}
}();
auto [MN, K, L] = scale_input_dims_MNKL;
auto mk_layout = make_layout(
make_shape(make_shape(Int<SFVecSizeMN>{}, cute::ceil_div(MN, SFVecSizeMN)),
make_shape(Int<SFVecSizeK>{}, cute::ceil_div(K, SFVecSizeK))),
strides
);
return make_layout(append(shape(mk_layout), L), append(stride(mk_layout), size(filter_zeros(mk_layout))));
}
};
template<UMMA::Major majorScale = UMMA::Major::MN>
struct RuntimeMixedInputBlockwiseScaleConfig {
using ShapeScale = Shape<Shape<int32_t, int32_t>, Shape<int32_t, int32_t>, int32_t>;
using StrideScale = conditional_t<majorScale == UMMA::Major::MN,
Stride<Stride<_0,_1>,Stride<_0,int32_t>, int32_t>,
Stride<Stride<_0,int32_t>,Stride<_0,_1>, int32_t>>;
using LayoutScale = Layout<ShapeScale, StrideScale>;
CUTE_HOST_DEVICE
static constexpr auto
deduce_layout_scale() {
return LayoutScale{};
}
// The following function is provided for user fill dynamic problem size to the layout_S.
template <class ProblemShape, class SFVecShape>
CUTE_HOST_DEVICE
static constexpr auto
tile_atom_to_shape_scale(ProblemShape problem_shape, SFVecShape sf_vec_shape) {
auto problem_shape_MNKL = append<3>(problem_shape, 1);
auto strides = [&]() CUTLASS_LAMBDA_FUNC_INLINE {
auto [MN, K, L] = problem_shape_MNKL;
auto [sfmn, sfk] = sf_vec_shape;
if constexpr (majorScale == UMMA::Major::MN) {
return make_stride(make_stride(_0{}, _1{}), make_stride(_0{}, cute::ceil_div(MN, sfmn)));
}
else {
return make_stride(make_stride(_0{}, cute::ceil_div(K, sfk)), make_stride(_0{}, _1{}));
}
}();
auto [MN, K, L] = problem_shape_MNKL;
auto [sfmn, sfk] = sf_vec_shape;
auto mk_layout = make_layout(
make_shape(make_shape(sfmn, cute::ceil_div(MN, sfmn)),
make_shape(sfk, cute::ceil_div(K, sfk))),
strides
);
return make_layout(append(shape(mk_layout), L), append(stride(mk_layout), size(filter_zeros(mk_layout))));
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::detail
include/cutlass/detail/sm103_blockscaled_layout.hpp
+107
View File
@@ -0,0 +1,107 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief Blocked Scale configs specific for SM103 BlockScaled MMA
*/
#pragma once
#include "cutlass/layout/matrix.h"
#include "cute/int_tuple.hpp"
#include "cute/atom/mma_traits_sm100.hpp"
namespace cutlass::detail{
/////////////////////////////////////////////////////////////////////////////////////////////////
using namespace cute;
template <int SFVecSize, UMMA::Major major = UMMA::Major::K>
struct Sm103BlockScaledBasicChunk {
using Blk_MN = _128;
using Blk_SF = _4;
using SfKMajorAtom = Layout< Shape< Shape< _8, _4, _4>, Shape<Int<SFVecSize>, _4>>,
Stride<Stride<_16,_128, _4>, Stride< _0, _1>>>;
using SfMNMajorAtom = Layout< Shape< Shape<Int<SFVecSize>, _4>, Shape<_8, _4, _4>>,
Stride<Stride< _0, _1>, Stride<_16,_128, _4>>>;
using SfAtom = cute::conditional_t<major == UMMA::Major::K, SfKMajorAtom, SfMNMajorAtom>;
};
template <int SFVecSize_>
struct Sm103BlockScaledConfig {
// We are creating the SFA and SFB tensors' layouts in the collective since they always have the same layout.
// k-major order
static constexpr int SFVecSize = SFVecSize_;
using Sm103BlkScaledChunk = Sm103BlockScaledBasicChunk<SFVecSize>;
using Blk_MN = typename Sm103BlkScaledChunk::Blk_MN;
using Blk_SF = typename Sm103BlkScaledChunk::Blk_SF;
using SfAtom = typename Sm103BlkScaledChunk::SfAtom;
using LayoutSF = decltype(tile_to_shape(SfAtom{}, make_shape(int(0),int(0),int(0)),Step<_2,_1,_3>{}));
CUTE_HOST_DEVICE
static constexpr auto
deduce_layoutSFA() {
return LayoutSF{};
}
CUTE_HOST_DEVICE
static constexpr auto
deduce_layoutSFB() {
return LayoutSF{};
}
// The following function is provided for user fill dynamic problem size to the layout_SFA.
template < class ProblemShape, class LayoutSFA = LayoutSF>
CUTE_HOST_DEVICE
static constexpr auto
tile_atom_to_shape_SFA(ProblemShape problem_shape, LayoutSFA layout_sfa = LayoutSFA{}) {
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_MNKL;
return tile_to_shape(SfAtom{}, make_shape(M,K,L), Step<_2,_1,_3>{});
}
// The following function is provided for user fill dynamic problem size to the layout_SFB.
template <class ProblemShape, class LayoutSFB = LayoutSF>
CUTE_HOST_DEVICE
static constexpr auto
tile_atom_to_shape_SFB(ProblemShape problem_shape, LayoutSFB layout_sfb = LayoutSFB{}) {
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_MNKL;
return tile_to_shape(SfAtom{}, make_shape(N,K,L), Step<_2,_1,_3>{});
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::detail
include/cutlass/device_kernel.h
+1
View File
@@ -35,6 +35,7 @@
#pragma once
#include <cutlass/detail/helper_macros.hpp> // CUTLASS_HOST_DEVICE
#include <cutlass/arch/synclog.hpp> // cutlass::arch::synclog_*
#include <cutlass/platform/platform.h> // uint64_t
// __grid_constant__ was introduced in CUDA 11.7.
include/cutlass/epilogue/collective/builders/sm100_builder.inl
+20 -5
View File
@@ -50,9 +50,10 @@
#include "cutlass/epilogue/fusion/callbacks.hpp"
#include "cutlass/epilogue/fusion/operations.hpp" // detail::is_sfd_epilogue_v
#include "cutlass/epilogue/fusion/sm100_callbacks_tma_warpspecialized.hpp"
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#include CUDA_STD_HEADER(type_traits)
#else
#include <type_traits>
#endif
@@ -1272,6 +1273,13 @@ private:
static constexpr bool Is1SmMma = is_base_of_v<NoSmemWarpSpecialized1Sm, EpilogueScheduleType>;
static constexpr bool Is2SmMma = is_base_of_v<NoSmemWarpSpecialized2Sm, EpilogueScheduleType>;
static constexpr bool IsInterleavedComplex = is_complex<ElementAccumulator>::value;
static constexpr bool IsFastF32Schedule = is_same_v<FastF32NoSmemWarpSpecialized1Sm, EpilogueScheduleType> ||
is_same_v<FastF32NoSmemWarpSpecialized2Sm, EpilogueScheduleType> ||
is_same_v<PtrArrayFastF32NoSmemWarpSpecialized1Sm, EpilogueScheduleType> ||
is_same_v<PtrArrayFastF32NoSmemWarpSpecialized2Sm, EpilogueScheduleType>;
// Input transform kernels - when dispatching to sm100 nosmem epilogue, go through the default path without EVT support.
static constexpr bool IsInputTransformSchedule = IsInterleavedComplex || IsFastF32Schedule;
static_assert(Is1SmMma ^ Is2SmMma, "unsupported schedule");
static_assert(not (Is2SmMma && size<0>(ClusterShape_MNK{}) % 2 == 1), "schedule + cluster mismatch");
@@ -1315,7 +1323,9 @@ private:
static_assert(is_tuple_v<EpilogueTileType>, "Shape or Tile");
return EpilogueTileType{};
}
else if constexpr (is_same_v<OpClass,arch::OpClassBlockScaledTensorOp>) { // perf specialized case
else if constexpr (is_same_v<OpClass,arch::OpClassBlockScaledTensorOp> || not IsInputTransformSchedule) {
// Save register usage for sm103 blockscaled kernels and sm100 cpasync kernels
// to avoid register spilling.
constexpr int EpiM = size<0>(CtaTileShape_MNK{});
constexpr int EpiN = cute::min(_64{}, size<1>(CtaTileShape_MNK{}));
return Shape<Int<EpiM>, Int<EpiN>>{};
@@ -1333,8 +1343,8 @@ private:
static constexpr auto
dispatch_policy() {
if constexpr (is_same_v<EpilogueScheduleType, PtrArrayNoSmemWarpSpecialized1Sm> ||
is_same_v<EpilogueScheduleType, PtrArrayNoSmemWarpSpecialized2Sm>) {
if constexpr (std::is_base_of_v<PtrArrayNoSmemWarpSpecialized1Sm, EpilogueScheduleType> ||
std::is_base_of_v<PtrArrayNoSmemWarpSpecialized2Sm, EpilogueScheduleType>) {
return Sm100PtrArrayNoSmemWarpSpecialized{};
}
else {
@@ -1347,7 +1357,12 @@ private:
fusion_callbacks() {
constexpr thread::ScaleType::Kind ScaleType =
DisableSource ? thread::ScaleType::OnlyAlphaScaling : thread::ScaleType::Default;
if constexpr (IsDefaultFusionOp<FusionOp>::value && not is_same_v<OpClass,arch::OpClassBlockScaledTensorOp>) {
if constexpr (IsDefaultFusionOp<FusionOp>::value &&\
not is_same_v<OpClass, arch::OpClassBlockScaledTensorOp> && \
(IsInputTransformSchedule || \
is_same_v<EpilogueScheduleType, PtrArrayNoSmemWarpSpecialized1Sm> || \
is_same_v<EpilogueScheduleType, PtrArrayNoSmemWarpSpecialized2Sm>)
) {
// Legacy codepath using thread::LinearCombination, do not expect this to be stable
return thread::LinearCombination<
ElementD, 1, ElementAccumulator, ElementCompute, ScaleType, FusionOp::RoundStyle, ElementC>({});
include/cutlass/epilogue/collective/builders/sm103_builder.inl
+108
View File
@@ -0,0 +1,108 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cute/layout.hpp" // cute::Shape
#include "cute/numeric/numeric_types.hpp" // cute::sizeof_bits_v
#include "cutlass/arch/mma.h" // cutlass::arch::OpClassTensorOp, cutlass::OpClassSparseTensorOp
#include "cute/atom/copy_traits_sm100.hpp"
#include "cute/atom/mma_traits_sm100.hpp"
#include "cute/util/type_traits.hpp" // cute::is_same_v
#include "cutlass/detail/dependent_false.hpp" // cutlass::detail::dependent_false
#include "cutlass/detail/layout.hpp"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/collective/builders/sm100_common.inl"
#include "cutlass/epilogue/collective/builders/sm90_common.inl"
#include "cutlass/epilogue/dispatch_policy.hpp"
#include "cutlass/epilogue/collective/collective_epilogue.hpp"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/epilogue/thread/linear_combination_planar_complex.h"
#include "cutlass/epilogue/fusion/callbacks.hpp"
#include "cutlass/epilogue/fusion/sm100_callbacks_tma_warpspecialized.hpp"
namespace cutlass::epilogue::collective {
// Alias to sm100 builder
template <
class OpClass,
class MmaTileShape_MNK, // Static MMA tile shape
class ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1)
class EpilogueTileType,
class ElementAccumulator,
class ElementCompute,
class ElementC,
class GmemLayoutTagC,
int AlignmentC,
class ElementD,
class GmemLayoutTagD,
int AlignmentD,
class EpilogueScheduleType,
class FusionOp
>
struct CollectiveBuilder<
arch::Sm103,
OpClass,
MmaTileShape_MNK,
ClusterShape_MNK,
EpilogueTileType,
ElementAccumulator,
ElementCompute,
ElementC,
GmemLayoutTagC,
AlignmentC,
ElementD,
GmemLayoutTagD,
AlignmentD,
EpilogueScheduleType,
FusionOp
>
{
using CollectiveOp = typename CollectiveBuilder<
arch::Sm100,
OpClass,
MmaTileShape_MNK,
ClusterShape_MNK,
EpilogueTileType,
ElementAccumulator,
ElementCompute,
ElementC,
GmemLayoutTagC,
AlignmentC,
ElementD,
GmemLayoutTagD,
AlignmentD,
EpilogueScheduleType,
FusionOp
>::CollectiveOp;
};
} // namespace cutlass::epilogue::collective
include/cutlass/epilogue/collective/builders/sm120_builder.inl
+2 -2
View File
@@ -36,10 +36,10 @@
#include "cutlass/epilogue/collective/detail.hpp"
#include "cutlass/epilogue/collective/builders/sm90_common.inl"
#include "cutlass/epilogue/collective/builders/sm120_common.inl"
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#include CUDA_STD_HEADER(type_traits)
#else
#include <type_traits>
#endif
include/cutlass/epilogue/collective/builders/sm90_builder.inl
+2 -2
View File
@@ -45,9 +45,9 @@
#include "cutlass/epilogue/thread/linear_combination_bias_elementwise.h"
#include "cutlass/epilogue/fusion/callbacks.hpp"
#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#include CUDA_STD_HEADER(type_traits)
#else
#include <type_traits>
#endif
include/cutlass/epilogue/collective/collective_builder.hpp
+1
View File
@@ -120,6 +120,7 @@ struct CallbacksBuilder<
#include "builders/sm90_builder.inl"
#include "builders/sm100_builder.inl"
#include "builders/sm103_builder.inl"
#include "builders/sm120_builder.inl"
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp
+506 -23
View File
@@ -60,7 +60,9 @@ template <
class ElementD_,
class StrideD_,
class ThreadEpilogueOp_,
class CopyOpT2R_
class CopyOpT2R_,
class AlignmentC_,
class AlignmentD_
>
class CollectiveEpilogue<
Sm100PtrArrayNoSmem,
@@ -70,7 +72,10 @@ class CollectiveEpilogue<
ElementD_,
StrideD_,
ThreadEpilogueOp_,
CopyOpT2R_
CopyOpT2R_,
AlignmentC_,
AlignmentD_,
cute::enable_if_t<IsDefaultFusionOp<ThreadEpilogueOp_>::value>
> {
public:
//
@@ -92,6 +97,10 @@ public:
using StrideD = StrideD_;
using InternalStrideD = cute::remove_pointer_t<StrideD>;
using CopyOpT2R = CopyOpT2R_;
using AlignmentC = AlignmentC_;
using AlignmentD = AlignmentD_;
using GmemElementC = cute::conditional_t<cute::is_void_v<ElementC>,ElementD,ElementC>; // prevents void ref breakages
using GmemTiledCopyC = void;
using GmemTiledCopyD = void;
@@ -136,7 +145,7 @@ public:
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count = 0) {
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int /*sm_count*/ = 0) {
return 0;
}
@@ -232,24 +241,56 @@ public:
}
// Represent the full output tensor, slice to get the tile this CTA is responsible for
Tensor mC = make_tensor(make_gmem_ptr(ptr_C_l), problem_shape_mnl, stride_c); // (M,N,L)
Tensor mD = make_tensor(make_gmem_ptr(params.ptr_D[l_coord]), problem_shape_mnl, stride_d); // (M,N,L)
Tensor mC = make_tensor(make_gmem_ptr(ptr_C_l), problem_shape_mnl, stride_c); // (M,N,L)
Tensor mD = make_tensor(make_gmem_ptr(params.ptr_D[l_coord]), problem_shape_mnl, stride_d); // (M,N,L)
Tensor gC = local_tile(mC, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N)
Tensor gD = local_tile(mD, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N)
// Partition source and destination tiles according to tmem copy T2R partitioning (tTR_)
auto tiled_t2r = make_tmem_copy(CopyOpT2R{}, tensor<0>(accumulators));
auto thread_t2r = tiled_t2r.get_slice(threadIdx.x % size(tiled_t2r));
auto thread_idx = threadIdx.x % size(tiled_t2r);
auto thread_t2r = tiled_t2r.get_slice(thread_idx);
Tensor tTR_gC = thread_t2r.partition_D(gC); // (T2R,T2R_M,T2R_N)
Tensor tTR_gD = thread_t2r.partition_D(gD); // (T2R,T2R_M,T2R_N)
Tensor tTR_rAcc = make_tensor<ElementAccumulator>(shape(tTR_gD)); // (T2R,T2R_M,T2R_N)
Tensor coordD = make_identity_tensor(problem_shape_mnl); // (M,N,L) -> (m,n,l)
Tensor cD = local_tile(coordD, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) -> (m,n,l)
Tensor tTR_cD = thread_t2r.partition_D(cD); // (T2R,T2R_M,T2R_N) -> (m,n,l)
Tensor tTR_rC = make_tensor<GmemElementC>(shape(tTR_gC)); // (T2R,T2R_M,T2R_N)
Tensor coordCD = make_identity_tensor(problem_shape_mnl); // (M,N,L) -> (m,n,l)
Tensor cCD = local_tile(coordCD, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) -> (m,n,l)
Tensor tTR_cCD = thread_t2r.partition_D(cCD); // (T2R,T2R_M,T2R_N) -> (m,n,l)
constexpr auto mclD = decltype(max_common_layout(tTR_rAcc.layout(), tTR_gD.layout())){};
constexpr int VD = cute::min(AlignmentD{}, size(mclD));
Tensor tTR_rD_frag = make_tensor<ElementD>(shape(tTR_rAcc));
Tensor tTR_rD_src = recast<Array<ElementD, VD>>(coalesce(tTR_rD_frag));
Tensor tR2G_rD_dst = recast<Array<ElementD, VD>>(coalesce(tTR_gD));
Tensor tTR_cD_mn_frg = tensor<1>(zipped_divide(coalesce(tTR_cCD), mclD.compose(Int<VD>{})));
Tensor tDpD = make_tensor<bool>(shape(tR2G_rD_dst));
CUTLASS_PRAGMA_UNROLL
for (int t = 0; t < size(tDpD); t++) {
tDpD(t) = elem_less(tTR_cD_mn_frg(t), problem_shape_mnl);
}
constexpr auto mclC = decltype(max_common_layout(tTR_rAcc.layout(), tTR_gC.layout())){};
constexpr int VC = cute::min(AlignmentC{}, size(mclC));
Tensor tTR_cC_mn_frg = tensor<1>(zipped_divide(coalesce(tTR_cCD), mclC.compose(Int<VC>{})));
Tensor tG2R_rC_dst = recast<Array<GmemElementC, VC>>(coalesce(tTR_gC));
Tensor tCpC = make_tensor<bool>(shape(tG2R_rC_dst));
CUTLASS_PRAGMA_UNROLL
for (int t = 0; t < size(tCpC); t++) {
tCpC(t) = elem_less(tTR_cC_mn_frg(t), problem_shape_mnl);
}
Tensor tTR_rC_src = recast<Array<GmemElementC, VC>>(coalesce(tTR_gC));
Tensor tTR_rC_dst = recast<Array<GmemElementC, VC>>(coalesce(tTR_rC));
// Detect interleaved complex fp32 kernels
Tensor accs = accumulators;
[[maybe_unused]] Tensor accs = accumulators;
using ElementTmem = typename decltype(accs)::value_type;
constexpr bool is_interleaved_complex_f32 = is_complex<ElementAccumulator>::value && cute::is_same_v<ElementTmem, float>;
@@ -283,28 +324,31 @@ public:
copy(tiled_t2r, tTR_tAcc, tTR_rAcc);
}
cutlass::arch::fence_view_async_tmem_load();
acc_pipeline.consumer_release(acc_pipe_consumer_state);
++acc_pipe_consumer_state;
// 2. Apply element-wise operation and store to gmem
// source is needed
if (epilogue_op.is_source_needed()) {
copy_if(tCpC, tTR_rC_src, tTR_rC_dst);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tTR_rAcc); ++i) {
if (elem_less(tTR_cD(i), problem_shape_mnl)) {
tTR_gD(i) = epilogue_op(tTR_rAcc(i), tTR_gC(i));
}
for (int i = 0; i < size(tTR_rAcc); i++) {
tTR_rD_frag(i) = epilogue_op(tTR_rAcc(i), tTR_rC(i));
}
copy_if(tDpD, tTR_rD_src, tR2G_rD_dst);
}
// source is not needed, avoid load
else {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tTR_rAcc); ++i) {
if (elem_less(tTR_cD(i), problem_shape_mnl)) {
tTR_gD(i) = epilogue_op(tTR_rAcc(i));
}
for (int i = 0; i < size(tTR_rAcc); i++) {
tTR_rD_frag(i) = epilogue_op(tTR_rAcc(i));
}
copy_if(tDpD, tTR_rD_src, tR2G_rD_dst);
}
cutlass::arch::fence_view_async_tmem_load();
acc_pipeline.consumer_release(acc_pipe_consumer_state);
++acc_pipe_consumer_state;
return cute::make_tuple(acc_pipe_consumer_state);
}
@@ -397,6 +441,439 @@ protected:
Params const& params;
};
template <
class EpilogueTile_, // (EPI_TILE_M, EPI_TILE_N)
class ElementC_,
class StrideC_,
class ElementD_,
class StrideD_,
class FusionCallbacks_,
class CopyOpT2R_,
class AlignmentC_,
class AlignmentD_
>
class CollectiveEpilogue<
Sm100PtrArrayNoSmem,
EpilogueTile_,
ElementC_,
StrideC_,
ElementD_,
StrideD_,
FusionCallbacks_,
CopyOpT2R_,
AlignmentC_,
AlignmentD_,
cute::enable_if_t<not IsDefaultFusionOp<FusionCallbacks_>::value>
> {
public:
//
// Type Aliases
//
// Required by the gemm::kernel
using DispatchPolicy = Sm100PtrArrayNoSmem;
using ElementC = ElementC_;
using ElementD = ElementD_;
using GmemElementC = cute::conditional_t<cute::is_void_v<ElementC>,ElementD,ElementC>; // prevents void ref breakages
using StrideC = StrideC_;
using StrideD = StrideD_;
using InternalStrideC = cute::remove_pointer_t<StrideC>;
using InternalStrideD = cute::remove_pointer_t<StrideD>;
using EpilogueTile = EpilogueTile_;
using CopyOpT2R = CopyOpT2R_;
using FusionCallbacks = FusionCallbacks_;
using ThreadEpilogueOp = typename epilogue::fusion::FusionCallbacksTraits<FusionCallbacks>::Operation;
using GmemTiledCopyC = void;
using GmemTiledCopyD = void;
private:
constexpr static bool IsReductionBufferNeeded = ThreadEpilogueOp::IsDePerRowBiasSupported
|| is_same_v<ThreadEpilogueOp, epilogue::fusion::FusionOperation>; // alloc reduction buffer for custom EVTs
constexpr static size_t ImplicitSharedStorageSize = IsReductionBufferNeeded ? size(EpilogueTile{}) : 0;
// Not unroll epi subtile loop when the activation op is heavy to reduce instruction size and register pressure.
constexpr static bool UnrollEpiLoop =
not cutlass::epilogue::thread::kIsHeavy_member_or_false<typename ThreadEpilogueOp::ActivationFn>::value;
public:
constexpr static int ThreadCount = 128;
constexpr static uint32_t TmaTransactionBytes = 0;
struct SharedStorage {
using FusionStorage = typename FusionCallbacks::SharedStorage;
FusionStorage thread;
array_aligned<uint8_t, ImplicitSharedStorageSize> buffer;
};
// Host side epilogue arguments
struct Arguments {
typename FusionCallbacks::Arguments thread{};
ElementC const** ptr_C = nullptr;
StrideC dC = {};
ElementD** ptr_D = nullptr;
StrideD dD = {};
};
// Device side epilogue params
struct Params {
typename FusionCallbacks::Params thread{};
ElementC const** ptr_C = nullptr;
StrideC dC = {};
ElementD** ptr_D = nullptr;
StrideD dD = {};
};
//
// Constructor and Data Members
//
CUTLASS_DEVICE
CollectiveEpilogue(Params const& params_, SharedStorage& shared_tensors)
: fusion_callbacks(params_.thread, shared_tensors.thread)
, smem_buffer_ptr(shared_tensors.buffer.data())
, params(params_) {};
protected:
FusionCallbacks fusion_callbacks;
uint8_t* smem_buffer_ptr;
Params const& params;
public:
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(
[[maybe_unused]] ProblemShape const& problem_shape,
Arguments const& args,
[[maybe_unused]] void* workspace) {
return {
FusionCallbacks::to_underlying_arguments(problem_shape, args.thread, workspace),
args.ptr_C,
args.dC,
args.ptr_D,
args.dD
};
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int /*sm_count*/ = 0) {
return FusionCallbacks::get_workspace_size(problem_shape, args.thread);
}
template <class ProblemShape>
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
CudaHostAdapter* cuda_adapter = nullptr) {
return FusionCallbacks::initialize_workspace(problem_shape, args.thread, workspace, stream, cuda_adapter);
}
template <class ProblemShape>
static bool
can_implement(
[[maybe_unused]] ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
bool fusion_implementable = FusionCallbacks::can_implement(problem_shape, args.thread);
if (!fusion_implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum requirements for FusionCallbacks.\n");
}
return fusion_implementable;
}
template<
bool ReuseTmem = false,
class AccumulatorPipeline,
class AccumulatorPipelineState,
class ProblemShapeMNKL,
class CtaTileMNK,
class CtaCoordMNKL,
class AccEngine, class AccLayout
>
CUTLASS_DEVICE auto
operator()(
AccumulatorPipeline acc_pipeline,
AccumulatorPipelineState acc_pipe_consumer_state,
ProblemShapeMNKL problem_shape_mnkl,
CtaTileMNK cta_tile_mnk,
CtaCoordMNKL cta_coord_mnkl,
cute::Tensor<AccEngine,AccLayout> accumulators,
[[maybe_unused]] SharedStorage&
) {
using ElementAccumulator = typename AccEngine::value_type;
using ElementCompute_ = typename epilogue::fusion::FusionCallbacksTraits<FusionCallbacks>::ElementCompute;
using ElementCompute = cute::conditional_t<cute::is_void_v<ElementCompute_>,ElementAccumulator,ElementCompute_>;
// Wait for mma warp to fill tmem buffer with accumulator results
static_assert(is_tmem<AccEngine>::value, "Accumulator must be TMEM resident.");
static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4");
static_assert(rank(CtaCoordMNKL{}) == 4, "TileCoordMNKL must be rank 4");
static_assert(cute::sizeof_bits_v<ElementD> != 6, "Output element requires smem");
auto M = get<0>(problem_shape_mnkl);
auto N = get<1>(problem_shape_mnkl);
auto L = get<3>(problem_shape_mnkl);
// Slice to get the tile this CTA is responsible for
auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl;
// Batches are managed by using appropriate pointers to C and D matrices
auto problem_shape_mnl = append<3>(make_shape(M,N),Int<1>{});
auto cta_coord_mnl = append<3>(make_shape(m_coord, n_coord),Int<0>{});
auto cta_tiler = take<0,2>(cta_tile_mnk);
auto cta_coord_mnk = cute::make_coord(m_coord, n_coord, k_coord, cute::Int<0>{});
bool is_C_load_needed = fusion_callbacks.is_C_load_needed();
auto [stride_c, stride_d] = [&, l = l_coord]() {
if constexpr (!cute::is_same_v<InternalStrideC, StrideC>) {
// If grouped gemm
if (is_C_load_needed) {
return make_tuple(
detail::get_epilogue_stride<DispatchPolicy>(params.dC[l]),
detail::get_epilogue_stride<DispatchPolicy>(params.dD[l])
);
}
else {
return make_tuple(
InternalStrideC{},
detail::get_epilogue_stride<DispatchPolicy>(params.dD[l])
);
}
}
else {
return make_tuple(
detail::get_epilogue_stride<DispatchPolicy>(params.dC),
detail::get_epilogue_stride<DispatchPolicy>(params.dD)
);
}
}();
// Get the residual tensor for the current batch
ElementC const* ptr_C_l = nullptr;
if (is_C_load_needed) {
ptr_C_l = params.ptr_C[l_coord];
}
int thread_idx = threadIdx.x % ThreadCount;
Tensor tAcc = accumulators(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N)
Tensor tAcc_epi = flat_divide(tAcc, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
TiledCopy tiled_t2r = make_tmem_copy(CopyOpT2R{}, tAcc_epi(_,_,_0{},_0{}));
ThrCopy thread_t2r = tiled_t2r.get_slice(thread_idx);
Tensor tTR_tAcc = thread_t2r.partition_S(tAcc_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
constexpr int FragmentSize = size(EpilogueTile{}) / ThreadCount;
Tensor coordD = make_identity_tensor(problem_shape_mnl); // (M,N,L) -> (m,n,l)
Tensor cD = local_tile(coordD, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) -> (m,n,l)
Tensor cD_epi = flat_divide(cD, EpilogueTile{});
Tensor tTR_cD = thread_t2r.partition_D(cD_epi); // (T2R,T2R_M,T2R_N) -> (m,n,l)
Tensor tTR_rAcc = make_tensor<ElementAccumulator>(shape(tTR_cD(_,_,_,_0{},_0{})));
// Construct the EVT consumer callbacks
auto residue_cD = make_coord(M,N) - cD(_0{});
auto residue_tTR_cD = make_coord(M,N) - tTR_cD(_0{});
Tensor cD_ = make_coord_tensor(cD.layout());
Tensor tTR_cD_ = make_coord_tensor(tTR_cD.layout());
constexpr bool RefSrc = false;
Tensor mC = make_tensor(make_gmem_ptr<GmemElementC>(ptr_C_l), problem_shape_mnl, stride_c);
Tensor tTR_gC = cutlass::epilogue::fusion::sm90_partition_for_epilogue<RefSrc>(
mC, cta_tile_mnk, cta_coord_mnk, EpilogueTile{}, tiled_t2r, thread_idx);
Tensor mD = make_tensor(make_gmem_ptr(recast_ptr<ElementD>(params.ptr_D[l_coord])), problem_shape_mnl, stride_d);
Tensor tTR_gD = cutlass::epilogue::fusion::sm90_partition_for_epilogue<RefSrc>(
mD, cta_tile_mnk, cta_coord_mnk, EpilogueTile{}, tiled_t2r, thread_idx);
// Register Tensor
Tensor tTR_rD = make_tensor<ElementD>(take<0,3>(shape(tTR_gD)));
Tensor coord_cCD = make_identity_tensor(problem_shape_mnl);
Tensor tTR_cCD = cutlass::epilogue::fusion::sm90_partition_for_epilogue<RefSrc>(
coord_cCD, cta_tile_mnk, cta_coord_mnk, EpilogueTile{}, tiled_t2r, thread_idx);
constexpr auto mclD = decltype(max_common_layout(tTR_gD(_,_,_,_0{},_0{}), tTR_rD)){};
constexpr int VD = cute::min(AlignmentD_{}, size(mclD));
auto tCrC = make_tensor<GmemElementC>(take<0,3>(shape(tTR_gC)));
constexpr auto mclC = decltype(max_common_layout(tTR_gC(_,_,_,_0{},_0{}), tCrC)){};
constexpr int VC = cute::min(AlignmentC_{}, size(mclC));
Tensor tTR_rD_frg = recast<Array<ElementD, FragmentSize>>(coalesce(tTR_rD));
auto cst_args = cutlass::epilogue::fusion::detail::ConsumerStoreArgs{
problem_shape_mnkl,
cta_tile_mnk,
cta_coord_mnkl,
int(0),
EpilogueTile{},
tiled_t2r,
cD_,
residue_cD,
tTR_cD_,
residue_tTR_cD,
tCrC,
thread_idx
};
auto synchronize = [] () CUTLASS_LAMBDA_FUNC_INLINE { cutlass::arch::NamedBarrier::sync(ThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
// The Epilogue Loop
auto epi_loop_fn = [&] (auto& cst_callbacks) CUTLASS_LAMBDA_FUNC_INLINE {
// Ensure there are no threads from the previous wave writing to shared memory being utilized for the current wave.
synchronize();
cst_callbacks.begin();
if (cst_callbacks.begin_sync_needed()) {
synchronize();
}
// If tmem doesn't have enough capacity to support double buffering, a portion of tmem (a column of epilogue tiles)
// is overlapped between 2 pseudo-buffers. The shared tmem portion corresponds to the last epilogue tile column of
// tmem accumulator buffer 0, and the first epilogue tile column of tmem accumulator 1.
// Thus, whenever we are processing tmem accumulator buffer 0, we process the epilogue tiles with reversed column order.
// Once the last epilogue tile column is loaded from tmem, the acc_pipeline is released.
// Then, the next accumulation stage for buffer 1 can start.
[[maybe_unused]] bool reverse_epi_n = ReuseTmem && acc_pipe_consumer_state.phase() == 0;
static_assert(not (ReuseTmem && AccumulatorPipeline::Stages != 1), "Tmem reuse requires 1 accumulator stage");
// For each epilogue subtile within the CTA tile
constexpr int NumEpiSubtilesN = CUTE_STATIC_V(size<4>(tTR_tAcc));
constexpr int NumEpiSubtilesM = CUTE_STATIC_V(size<3>(tTR_tAcc));
// Lambda to process a single epilogue tile
auto process_tile = [&](int epi_m, int epi_n, int iter_m, int iter_n) CUTLASS_LAMBDA_FUNC_INLINE {
bool is_last_iteration = iter_m == NumEpiSubtilesM-1 && iter_n == NumEpiSubtilesN-1;
bool do_acc_release = is_last_iteration;
// Adjust release condition for tmem reuse
if constexpr (ReuseTmem) {
do_acc_release = iter_m == NumEpiSubtilesM-1 && iter_n == 0; // Release on first N iteration
}
Tensor tTR_cCD_mn = tTR_cCD(_,_,_,epi_m,epi_n);
Tensor tTR_pCD_mn = cute::lazy::transform(tTR_cCD_mn, [&] (auto const& c) CUTLASS_LAMBDA_FUNC_INLINE { return elem_less(c, problem_shape_mnl); });
cst_callbacks.begin_loop(epi_m, epi_n);
if constexpr (not cute::is_void_v<ElementC>) {
if (is_C_load_needed) {
using CVecType = uint_bit_t<VC * sizeof_bits_v<ElementC>>;
if constexpr (!is_same_v<CVecType, uint256_t>) {
Tensor tTR_gC_frg = recast<CVecType>(coalesce(tTR_gC(_,_,_,epi_m,epi_n)));
Tensor tTR_rC_frg = recast<CVecType>(coalesce(tCrC));
Tensor tTR_pC_frg = tensor<1>(zipped_divide(coalesce(tTR_pCD_mn), mclC.compose(Int<VC>{})));
copy_if(tTR_pC_frg, tTR_gC_frg, tTR_rC_frg);
}
else {
auto tiled_g2r = make_tiled_copy_D(Copy_Atom<SM100_LOAD_256bit_CACHE_NOALLOCATION, ElementC>{}, tiled_t2r);
auto thr_g2r = tiled_g2r.get_slice(threadIdx.x);
Tensor c_src = thr_g2r.retile_S(tTR_gC(_,_,_,epi_m,epi_n));
Tensor c_dst = thr_g2r.retile_D(tCrC);
Tensor c_prd = thr_g2r.retile_D(tTR_pCD_mn);
copy_if(tiled_g2r, c_prd, c_src, c_dst);
}
}
}
// Copy accumulator tile from tmem to register
// The current tile in tmem
Tensor tTR_tAcc_mn = tTR_tAcc(_,_,_,epi_m,epi_n);
Tensor tTR_rAcc_frg = recast<Array<ElementAccumulator, FragmentSize>>(coalesce(tTR_rAcc));
copy(tiled_t2r, tTR_tAcc_mn, tTR_rAcc);
// After the last tmem load, signal that tmem buffer is consumed and empty
if (do_acc_release) {
cutlass::arch::fence_view_async_tmem_load();
acc_pipeline.consumer_release(acc_pipe_consumer_state);
++acc_pipe_consumer_state;
}
CUTLASS_PRAGMA_UNROLL
for (int epi_v = 0; epi_v < size(tTR_rAcc_frg); ++epi_v) {
tTR_rD_frg(epi_v) = cst_callbacks.visit(tTR_rAcc_frg(epi_v), epi_v, epi_m, epi_n);
}
Tensor reduction_buffer = make_tensor(
raw_pointer_cast(make_smem_ptr(smem_buffer_ptr)), make_layout(Shape<Int<ImplicitSharedStorageSize>>{}));
cst_callbacks.reduce(reduction_buffer, synchronize, epi_m, epi_n, is_last_iteration, tTR_rAcc /*not used*/);
cst_callbacks.end_loop(epi_m, epi_n);
using VecType = uint_bit_t<VD * sizeof_bits_v<ElementD>>;
if constexpr (!is_same_v<VecType, uint256_t>) {
Tensor tTR_gD_frg = recast<VecType>(coalesce(tTR_gD(_,_,_,epi_m,epi_n)));
Tensor tTR_rD_frg = recast<VecType>(coalesce(tTR_rD));
Tensor tTR_pD_frg = tensor<1>(zipped_divide(coalesce(tTR_pCD_mn), mclD.compose(Int<VD>{})));
copy_if(tTR_pD_frg, tTR_rD_frg, tTR_gD_frg);
}
else {
auto tiled_r2g = make_tiled_copy_D(Copy_Atom<SM100_STORE_256bit_CACHE_NOALLOCATION, ElementD>{}, tiled_t2r);
auto thr_r2g = tiled_r2g.get_slice(threadIdx.x);
Tensor src = thr_r2g.retile_S(tTR_rD);
Tensor dst = thr_r2g.retile_D(tTR_gD(_,_,_,epi_m,epi_n));
Tensor prd = thr_r2g.retile_D(tTR_pCD_mn);
copy_if(tiled_r2g, prd, src, dst);
}
};
// Use static iteration with appropriate ordering
// When ReuseTmem is true and reverse_epi_n is true, we need reverse N iteration
auto n_seq = cute::make_int_sequence<NumEpiSubtilesN>{};
auto m_seq = cute::make_int_sequence<NumEpiSubtilesM>{};
if constexpr (UnrollEpiLoop) {
// Fully unrolled static iteration
cute::for_each(n_seq, [&](auto I_N) CUTLASS_LAMBDA_FUNC_INLINE {
constexpr int iter_n = I_N;
int epi_n = iter_n;
if constexpr (ReuseTmem) {
if (reverse_epi_n) {
epi_n = NumEpiSubtilesN - 1 - iter_n; // Reverse N iteration
}
}
cute::for_each(m_seq, [&](auto I_M) CUTLASS_LAMBDA_FUNC_INLINE {
constexpr int iter_m = I_M;
process_tile(iter_m, epi_n, iter_m, iter_n);
});
});
} else {
// Runtime loop with pragma unroll(1)
#pragma unroll 1
for (int iter_n = 0; iter_n < NumEpiSubtilesN; ++iter_n) {
int epi_n = iter_n;
if constexpr (ReuseTmem) {
if (reverse_epi_n) {
epi_n = NumEpiSubtilesN - 1 - iter_n; // Reverse N iteration
}
}
#pragma unroll 1
for (int iter_m = 0; iter_m < NumEpiSubtilesM; ++iter_m) {
process_tile(iter_m, epi_n, iter_m, iter_n);
}
}
}
cst_callbacks.end();
};
//
// BEGIN EPILOGUE
//
auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks<RefSrc>(cst_args);
epi_loop_fn(cst_callbacks);
return cute::make_tuple(acc_pipe_consumer_state);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// For sm100 kernels requiring warp specialized epilogues
@@ -430,7 +907,10 @@ class CollectiveEpilogue<
ElementD_,
StrideD_,
ThreadEpilogueOp_,
CopyOpT2R_>>
CopyOpT2R_,
AlignmentC,
AlignmentD,
void>>
{
public:
// ctor inheritance
@@ -442,7 +922,10 @@ public:
ElementD_,
StrideD_,
ThreadEpilogueOp_,
CopyOpT2R_>>::Sm100TmaWarpSpecializedAdapter;
CopyOpT2R_,
AlignmentC,
AlignmentD,
void>>::Sm100TmaWarpSpecializedAdapter;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/collective/sm100_epilogue_nosmem.hpp
+24 -4
View File
@@ -343,8 +343,7 @@ public:
copy_if(tDpD, tTR_rD_src, tR2G_rD_dst);
}
// source is not needed, avoid load
else
{
else {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tTR_rAcc); i++) {
tTR_rD_frag(i) = epilogue_op(tTR_rAcc(i));
@@ -587,8 +586,8 @@ public:
CtaTileMNK cta_tile_mnk,
CtaCoordMNKL cta_coord_mnkl,
cute::Tensor<AccEngine,AccLayout> accumulators,
[[maybe_unused]]SharedStorage&
) {
[[maybe_unused]] SharedStorage&
) {
using ElementAccumulator = typename AccEngine::value_type;
using ElementCompute_ = typename epilogue::fusion::FusionCallbacksTraits<FusionCallbacks>::ElementCompute;
using ElementCompute = cute::conditional_t<cute::is_void_v<ElementCompute_>,ElementAccumulator,ElementCompute_>;
@@ -718,10 +717,20 @@ public:
if (is_C_load_needed) {
using CVecType = uint_bit_t<VC * sizeof_bits_v<ElementC>>;
if constexpr (!is_same_v<CVecType, uint256_t>) {
Tensor tTR_gC_frg = recast<CVecType>(coalesce(tTR_gC(_,_,_,epi_m,epi_n)));
Tensor tTR_rC_frg = recast<CVecType>(coalesce(tCrC));
Tensor tTR_pC_frg = tensor<1>(zipped_divide(coalesce(tTR_pCD_mn), mclC.compose(Int<VC>{})));
copy_if(tTR_pC_frg, tTR_gC_frg, tTR_rC_frg);
}
else {
auto tiled_g2r = make_tiled_copy_D(Copy_Atom<SM100_LOAD_256bit_CACHE_NOALLOCATION, ElementC>{}, tiled_t2r);
auto thr_g2r = tiled_g2r.get_slice(threadIdx.x);
Tensor c_src = thr_g2r.retile_S(tTR_gC(_,_,_,epi_m,epi_n));
Tensor c_dst = thr_g2r.retile_D(tCrC);
Tensor c_prd = thr_g2r.retile_D(tTR_pCD_mn);
copy_if(tiled_g2r, c_prd, c_src, c_dst);
}
}
}
@@ -753,10 +762,21 @@ public:
cst_callbacks.end_loop(epi_m, epi_n);
using VecType = uint_bit_t<VD * sizeof_bits_v<ElementD>>;
if constexpr (!is_same_v<VecType, uint256_t>) {
Tensor tTR_gD_frg = recast<VecType>(coalesce(tTR_gD(_,_,_,epi_m,epi_n)));
Tensor tTR_rD_frg = recast<VecType>(coalesce(tTR_rD));
Tensor tTR_pD_frg = tensor<1>(zipped_divide(coalesce(tTR_pCD_mn), mclD.compose(Int<VD>{})));
copy_if(tTR_pD_frg, tTR_rD_frg, tTR_gD_frg);
}
else {
auto tiled_r2g = make_tiled_copy_D(Copy_Atom<SM100_STORE_256bit_CACHE_NOALLOCATION, ElementD>{}, tiled_t2r);
auto thr_r2g = tiled_r2g.get_slice(threadIdx.x);
Tensor src = thr_r2g.retile_S(tTR_rD);
Tensor dst = thr_r2g.retile_D(tTR_gD(_,_,_,epi_m,epi_n));
Tensor prd = thr_r2g.retile_D(tTR_pCD_mn);
copy_if(tiled_r2g, prd, src, dst);
}
} // for epi_m
} // for epi_n
include/cutlass/epilogue/dispatch_policy.hpp
+26 -5
View File
@@ -61,8 +61,12 @@ struct PtrArrayTmaWarpSpecializedCooperative { static constexpr int NumEpilogueW
// Blackwell direct store schedules
struct NoSmemWarpSpecialized1Sm {};
struct NoSmemWarpSpecialized2Sm {};
struct FastF32NoSmemWarpSpecialized1Sm : NoSmemWarpSpecialized1Sm {};
struct FastF32NoSmemWarpSpecialized2Sm : NoSmemWarpSpecialized2Sm {};
struct PtrArrayNoSmemWarpSpecialized1Sm : NoSmemWarpSpecialized1Sm {};
struct PtrArrayNoSmemWarpSpecialized2Sm : NoSmemWarpSpecialized2Sm {};
struct PtrArrayFastF32NoSmemWarpSpecialized1Sm : PtrArrayNoSmemWarpSpecialized1Sm {};
struct PtrArrayFastF32NoSmemWarpSpecialized2Sm : PtrArrayNoSmemWarpSpecialized2Sm {};
// Blackwell TMA schedules
struct TmaWarpSpecialized1Sm {};
struct TmaWarpSpecialized2Sm {};
@@ -234,12 +238,29 @@ struct Sm100PtrArrayTmaWarpSpecialized {
static_assert(StagesD >= 1, "StagesD must be >= 1");
};
// default elementwise operator epilogue without smem
struct Sm100NoSmem {};
struct Sm100NoSmemWarpSpecialized {};
struct Sm100PtrArrayNoSmem {};
struct Sm100PtrArrayNoSmemWarpSpecialized {};
struct Sm100NoSmem {
constexpr static int StagesC = 1;
constexpr static int StagesD = 1;
constexpr static int FragmentSize = 1;
};
struct Sm100NoSmemWarpSpecialized {
constexpr static int StagesC = 1;
constexpr static int StagesD = 1;
constexpr static int FragmentSize = 1;
};
struct Sm100PtrArrayNoSmem {
constexpr static int StagesC = 1;
constexpr static int StagesD = 1;
constexpr static int FragmentSize = 1;
};
struct Sm100PtrArrayNoSmemWarpSpecialized {
constexpr static int StagesC = 1;
constexpr static int StagesD = 1;
constexpr static int FragmentSize = 1;
};
template<
int StagesC_,
int StagesD_,
include/cutlass/epilogue/fusion/sm100_callbacks_tma_warpspecialized.hpp
+33
View File
@@ -1280,6 +1280,39 @@ struct FusionCallbacks<
};
// --------------------------------------------------------------------
// Sm100PtrArrayNoSmemWarpSpecialized (direct-store, grouped GEMM)
// --------------------------------------------------------------------
template <
class Operation,
class CtaTile_MNK,
class EpilogueTile_MN,
class... Args
>
struct FusionCallbacks<
epilogue::Sm100PtrArrayNoSmemWarpSpecialized,
Operation,
CtaTile_MNK,
EpilogueTile_MN,
Args...>
: FusionCallbacks<
// reuse the ptr-array *TMA* callbacks with 0 stages
epilogue::Sm100PtrArrayTmaWarpSpecialized<0,0,0,false,false>,
Operation,
CtaTile_MNK,
EpilogueTile_MN,
Args...> {
using Base = FusionCallbacks<
epilogue::Sm100PtrArrayTmaWarpSpecialized<0,0,0,false,false>,
Operation,
CtaTile_MNK,
EpilogueTile_MN,
Args...>;
// bring ctors into scope
using Base::Base;
};
} // namespace cutlass::epilogue::fusion
include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp
+13 -17
View File
@@ -871,11 +871,11 @@ struct Sm90ScalarBroadcastPtrArray {
template <class... Args>
CUTLASS_DEVICE auto
get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) {
// Get the scalar for batched broadcast
if (size<2>(params_ptr->dScalar[0]) != 0) {
auto [m_coord, n_coord, k_coord, l_coord] = args.tile_coord_mnkl;
update_scalar(l_coord);
}
// Always refresh scalar with the current group index so per-group
// alpha/beta values (provided through pointer arrays) are loaded
// correctly even when the L-stride is zero.
auto [m_coord, n_coord, k_coord, l_coord] = args.tile_coord_mnkl;
update_scalar(l_coord);
return EmptyProducerLoadCallbacks{};
}
@@ -904,12 +904,8 @@ struct Sm90ScalarBroadcastPtrArray {
>
CUTLASS_DEVICE auto
get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
// Get the scalar for batched broadcast
if (get<2>(params_ptr->dScalar[0]) != 0) {
auto [m_coord, n_coord, k_coord, l_coord] = args.tile_coord_mnkl;
update_scalar(l_coord);
}
auto [m_coord, n_coord, k_coord, l_coord] = args.tile_coord_mnkl;
update_scalar(l_coord);
return ConsumerStoreCallbacks(scalar);
}
@@ -920,13 +916,15 @@ private:
int l_offset = l_coord * size<2>(params_ptr->dScalar[0]);
if (params_ptr->scalar_ptr_arrays[0] != nullptr) {
scalar = *(params_ptr->scalar_ptr_arrays[0][l_offset]);
// Pointer-array variant: each entry already points to the scalar of a group.
scalar = *(params_ptr->scalar_ptr_arrays[0][l_coord]);
}
else if (params_ptr->scalar_ptrs[0] != nullptr) {
// Strided pointer variant.
scalar = params_ptr->scalar_ptrs[0][l_offset];
}
else {
// batch stride is ignored for nullptr fallback
// Literal fallback.
scalar = params_ptr->scalars[0];
}
@@ -936,15 +934,13 @@ private:
for (int i = 1; i < BroadcastCount; ++i) {
if (params_ptr->scalar_ptr_arrays[i] != nullptr) {
int rest_l_offset = l_coord * size<2>(params_ptr->dScalar[i]);
scalar = reduction_fn(scalar, *(params_ptr->scalar_ptr_arrays[i][rest_l_offset]));
scalar = reduction_fn(scalar, *(params_ptr->scalar_ptr_arrays[i][l_coord]));
}
if (params_ptr->scalar_ptrs[i] != nullptr) {
else if (params_ptr->scalar_ptrs[i] != nullptr) {
int rest_l_offset = l_coord * size<2>(params_ptr->dScalar[i]);
scalar = reduction_fn(scalar, params_ptr->scalar_ptrs[i][rest_l_offset]);
}
else {
// batch stride is ignored for nullptr fallback
scalar = reduction_fn(scalar, params_ptr->scalars[i]);
}
}
include/cutlass/epilogue/threadblock/epilogue.h
+2 -3
View File
@@ -38,10 +38,9 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/numeric_types.h"
#include "cutlass/array.h"
#include "cutlass/layout/vector.h"
include/cutlass/epilogue/threadblock/epilogue_base.h
+2 -4
View File
@@ -37,15 +37,13 @@
*/
#pragma once
#include "cutlass/cutlass.h"
#if !defined(__CUDACC_RTC__)
#include <type_traits>
#include <utility>
#endif
#include CUDA_STD_HEADER(cassert)
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include "cutlass/matrix_shape.h"
#include "cutlass/numeric_types.h"
#include "cutlass/array.h"
include/cutlass/epilogue/threadblock/epilogue_gemm_k_reduction.h
+1 -3
View File
@@ -37,10 +37,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/numeric_types.h"
#include "cutlass/array.h"
#include "cutlass/layout/vector.h"
include/cutlass/epilogue/threadblock/epilogue_smem_accumulator.h
+1 -3
View File
@@ -37,10 +37,8 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/numeric_types.h"
#include "cutlass/array.h"
#include "cutlass/layout/vector.h"
include/cutlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h
+3 -3
View File
@@ -38,16 +38,16 @@
*/
#pragma once
#include "cutlass/cutlass.h"
#include <cuda/std/cassert>
#include CUDA_STD_HEADER(cassert)
#if defined(__CUDACC_RTC__)
#include <cuda/std/utility>
#include CUDA_STD_HEADER(utility)
#else
#include <utility>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/array.h"
#include "cutlass/numeric_types.h"
#include "cutlass/numeric_conversion.h"
include/cutlass/epilogue/threadblock/epilogue_with_absmax.h
+3 -4
View File
@@ -49,16 +49,15 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#if defined(__CUDACC_RTC__)
#include <cuda/std/utility>
#include CUDA_STD_HEADER(utility)
#else
#include <utility>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/array.h"
#include "cutlass/numeric_types.h"
#include "cutlass/numeric_conversion.h"
include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h
+3 -4
View File
@@ -38,16 +38,15 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#if defined(__CUDACC_RTC__)
#include <cuda/std/utility>
#include CUDA_STD_HEADER(utility)
#else
#include <utility>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/array.h"
#include "cutlass/numeric_types.h"
#include "cutlass/numeric_conversion.h"
include/cutlass/epilogue/threadblock/epilogue_with_reduction.h
+3 -3
View File
@@ -38,10 +38,10 @@
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/array.h"
#include "cutlass/numeric_types.h"
#include "cutlass/numeric_conversion.h"
include/cutlass/epilogue/threadblock/epilogue_with_scaling_factor.h
+231
View File
@@ -0,0 +1,231 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief Epilogue visitor for threadblock scoped GEMMs that process softmax computations in epilogue.
The epilogue finds max values in each row of the row-major output matrix and stores them.
The max values are also used for a further round of threadblock scoped reduction operation, where
the partial reduction results are stored in a pre-allocated array and used for further full reduction.
*/
#pragma once
/////////////////////////////////////////////////////////////////////////////////////////////////
#include "cutlass/arch/memory.h"
#include "cutlass/cutlass.h"
#include "cutlass/fast_math.h"
#include "cutlass/numeric_conversion.h"
#include "cutlass/numeric_types.h"
#include "cutlass/tensor_ref.h" // cutlass::TensorRef
namespace cutlass
{
namespace epilogue
{
namespace threadblock
{
template <int kVectorSize_,
typename ThreadShape_,
typename ElementCompute_,
typename ElementAccumulator_,
typename ElementC_,
typename ElementD_,
typename ElementSFD_,
typename LayoutOutput_,
typename LayoutSFD_>
class GemvEpilogueWithScalingFactor
{
public:
using ThreadShape = ThreadShape_;
using ElementCompute = ElementCompute_; // f32
using ElementAccumulator = ElementAccumulator_; // f32
using ElementC = ElementC_; // e2m1
using ElementD = ElementD_; // e2m1
using ElementSFD = ElementSFD_; // e4m3
using LayoutOutput = LayoutOutput_; // ColumnMajor
using LayoutSFD = LayoutSFD_; // ColumnMajor
using TensorRefD = TensorRef<ElementD, LayoutOutput_>;
static constexpr int kVectorSize = kVectorSize_;
// number of threads row
static constexpr int kThreadsPerCol = ThreadShape::kM; // 16
// number of threads col
static constexpr int kThreadsPerRow = ThreadShape::kN; // 8
static constexpr int kThreadCount = kThreadsPerCol * kThreadsPerRow; // 128
static_assert(kVectorSize == kThreadsPerCol, "vector size and number of threads row should be equal");
static_assert(std::is_same_v<LayoutSFD, cutlass::layout::ColumnMajor> &&
std::is_same_v<LayoutOutput, cutlass::layout::ColumnMajor>,
"Only support Mx1 (ColumnMajor) output and ColumnMajor scaling factor");
static_assert(std::is_same_v<ElementCompute, float>, "ElementCompute should be float type");
static_assert(cutlass::sizeof_bits<ElementD>::value == 4, "Output should be FP4 type");
static_assert(cutlass::sizeof_bits<ElementSFD>::value == 8, "ElementSFD should be FP8 type");
static_assert(std::is_same_v<LayoutOutput, LayoutSFD>, "only support same layout for D and SFD");
// Hardcode static_assert on threadshape 16x8 to avoid bug
static_assert(kThreadsPerCol == 16, "thread shape col false");
static_assert(kThreadsPerRow == 8, "thread shape row false");
static_assert(kThreadCount == 128, "thread count false");
struct Params
{
TensorRefD tensor_d;
ElementSFD *scale_factor_d_ptr{nullptr};
ElementCompute alpha{0};
ElementCompute beta{0};
float st{0};
int64_t batch_stride_sfd{0}; // Add batch stride for SFD
int64_t stride_d{0}; // Add stride for D tensor
};
/// Shared storage
struct SharedStorage
{
// fp32
// Each thread store one fp32
#if 1
ElementAccumulator reduction_buffer[kThreadsPerCol];
#else
ElementAccumulator reduction_buffer[kThreadCount];
#endif
// Buffer for collecting 4-bit values for packed store
uint8_t packed_buffer[kThreadsPerCol];
};
private:
Params const &params_;
SharedStorage &shared_storage_;
float st_scale_down{0};
public:
CUTLASS_HOST_DEVICE GemvEpilogueWithScalingFactor(Params const &params, SharedStorage &shared_storage)
: params_(params)
, shared_storage_(shared_storage)
{
const float fp_subtype_max = static_cast<float>(cutlass::platform::numeric_limits<ElementD>::max());
this->st_scale_down = this->params_.st / fp_subtype_max;
}
CUTLASS_DEVICE void operator()(ElementAccumulator frag_acc, ElementC frag_c, int batch_idx)
{
const int block_idx = blockIdx.x;
const int thread_idx_col = threadIdx.x;
const int thread_idx_row = threadIdx.y;
const float st_scale_down = this->st_scale_down;
const float st = this->params_.st;
// Compute D offset using batch_idx and stride_d
const int output_d_base_offset = blockIdx.x * blockDim.y;
const int d_batch_offset = batch_idx * params_.stride_d;
ElementD* output_ptr = &params_.tensor_d.at({output_d_base_offset + d_batch_offset, 0});
uint8_t* byte_ptr = reinterpret_cast<uint8_t*>(output_ptr);
// For 8x16 thread layout, 1 thread per 128 threads write to sf d
// Every block write one SFD to gmem
const bool is_write_sfd_thread = (thread_idx_row == 0);
// Calculate SFD offset using proper batch stride
const int output_sfd_offset = (block_idx / 4) * 512 + block_idx % 4 + batch_idx * params_.batch_stride_sfd;
auto reduction_buffer = shared_storage_.reduction_buffer;
// fp32
ElementAccumulator max_accum_row0 = ElementAccumulator(0);
ElementAccumulator max_accum_row1 = ElementAccumulator(0);
// Thread in row contain duplicate frag_acc data
if ( thread_idx_col == 0 ) {
// 16 threads write to 16 contigious bank, no conflict
reduction_buffer[thread_idx_row] = frag_acc;
}
__syncthreads();
if (threadIdx.y == 0) {
auto acc_0 = reduction_buffer[threadIdx.x * 2];
auto acc_1 = reduction_buffer[threadIdx.x * 2 + 1];
// compute the max for me using shuffling among 16 threads.
ElementAccumulator max_accum = fabsf(acc_0);
max_accum = cutlass::fast_max(max_accum, fabsf(acc_1));
// Butterfly reduction pattern for 16 threads
// Each iteration halves the number of active lanes
max_accum = cutlass::fast_max(max_accum, __shfl_down_sync(0xFF, max_accum, 4)); // 8->4
max_accum = cutlass::fast_max(max_accum, __shfl_down_sync(0xFF, max_accum, 2)); // 4->2
max_accum = cutlass::fast_max(max_accum, __shfl_down_sync(0xFF, max_accum, 1)); // 2->1
// Broadcast the final result to all 8 threads
max_accum = __shfl_sync(0xFF, max_accum, 0);
float pvscale = max_accum * st_scale_down;
ElementSFD qpvscale = static_cast<ElementSFD>(pvscale);
float qpvscale_up = NumericConverter<ElementCompute, ElementSFD>{}(qpvscale);
float qpvscale_up_rcp = __frcp_rn(qpvscale_up) * st;
uint8_t qval_u8_compare;
#if defined(CUDA_PTX_FP4FP6_CVT_ENABLED)
uint32_t temp_result;
asm volatile (
"{\n"
" .reg .f32 output_fp32_0, output_fp32_1;\n"
" .reg .b8 byte0, byte1, byte2, byte3;\n"
" mul.f32 output_fp32_0, %1, %3;\n"
" mul.f32 output_fp32_1, %2, %3;\n"
" cvt.rn.satfinite.e2m1x2.f32 byte0, output_fp32_1, output_fp32_0;\n"
" mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
"}\n"
: "=r"(temp_result) // Output to uint32_t
: "f"(acc_0), "f"(acc_1), "f"(qpvscale_up_rcp)
);
qval_u8_compare = temp_result & 0xFF;
#else
ElementD output_fp4_0 = NumericConverter<ElementD, ElementCompute>{}(acc_0 * qpvscale_up_rcp);
ElementD output_fp4_1 = NumericConverter<ElementD, ElementCompute>{}(acc_1 * qpvscale_up_rcp);
uint8_t raw_fp4_0 = reinterpret_cast<const uint8_t&>(output_fp4_0) & 0x0F;
uint8_t raw_fp4_1 = reinterpret_cast<const uint8_t&>(output_fp4_1) & 0x0F;
qval_u8_compare = (raw_fp4_1 << 4) | raw_fp4_0;
#endif
byte_ptr[threadIdx.x] = qval_u8_compare;
arch::global_store<ElementSFD, sizeof(ElementSFD)>(qpvscale,
(void *)(params_.scale_factor_d_ptr + output_sfd_offset),
is_write_sfd_thread);
}
} // end of operator()
};
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass
include/cutlass/fast_math.h
+3 -4
View File
@@ -30,18 +30,17 @@
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#include CUDA_STD_HEADER(cstdint)
#else
#include <cstdint>
#include <cmath>
#include <type_traits>
#endif
#if !defined(__QNX__)
#include <cuda/std/utility>
#include CUDA_STD_HEADER(utility)
#endif
#include "cutlass/cutlass.h"
#include "cutlass/array.h"
#include "cutlass/uint128.h"
#include "cutlass/coord.h"
include/cutlass/float8.h
+4 -2
View File
@@ -59,12 +59,14 @@
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED))
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121A_ENABLED))
# define CUDA_PTX_UE8M0_CVT_ENABLED 1
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120F_ENABLED))
defined(CUTLASS_ARCH_MMA_SM103F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121F_ENABLED))
# define CUDA_PTX_UE8M0_CVT_ENABLED 1
#endif
include/cutlass/float_subbyte.h
+4 -2
View File
@@ -45,12 +45,14 @@
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED))
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121A_ENABLED))
# define CUDA_PTX_FP4FP6_CVT_ENABLED 1
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120F_ENABLED))
defined(CUTLASS_ARCH_MMA_SM103F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121F_ENABLED))
# define CUDA_PTX_FP4FP6_CVT_ENABLED 1
#endif
include/cutlass/functional.h
+2 -1
View File
@@ -54,7 +54,8 @@
#include <intrin.h>
#endif // _MSC_VER
#if defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)
#if defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)
# define CUTLASS_ARCH_CREDUX_ENABLED
#endif
include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl
+2 -2
View File
@@ -132,7 +132,7 @@ auto sm100_make_simt_gmem_tiled_copy_SFA() {
return make_tiled_copy(
SmemScalingCopyAtomA{},
Layout<Shape<_32>>{}, // 32 threads
Layout<Shape<Shape<Int<ScaleGranularityM>, Int<ElementsPerSFACopy>>>, Stride<Stride<_0, _1>>>{});
Layout<Shape<Int<ElementsPerSFACopy>>>{});
}
else {
using SmemScalingCopyAtomA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<Element>, Element>;
@@ -166,7 +166,7 @@ auto sm100_make_simt_gmem_tiled_copy_SFB() {
return make_tiled_copy(
SmemScalingCopyAtomB{},
Layout<Shape<_32>>{}, // 32 threads
Layout<Shape<Shape<Int<ScaleGranularityN>, Int<ElementsPerSFBCopy>>>, Stride<Stride<_0, _1>>>{});
Layout<Shape<Int<ElementsPerSFBCopy>>>{});
}
else {
using SmemScalingCopyAtomB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<Element>, Element>;
include/cutlass/gemm/collective/builders/sm100_common.inl
+67
View File
@@ -569,6 +569,73 @@ sm100_make_trivial_fastFP32_tiled_mma() {
}
}
//Setting mma for Mixed input gemm. Here, ElementAMma should be TACompute
template<
class ElementAMma,
class ElementBMma,
class ElementAccumulator,
class TileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB,
class KernelScheduleType
>
constexpr auto
sm100_make_trivial_mixed_input_tiled_mma() {
constexpr int M = cute::size<0>(TileShape_MNK{});
constexpr int N = cute::size<1>(TileShape_MNK{});
//MMA 1Sm requested
if constexpr (cute::is_base_of_v<KernelSchedule1Sm, KernelScheduleType> ) {
if constexpr (UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedMixedInputSmemSm100, KernelScheduleType>) {
if constexpr (cute::is_same_v<ElementBMma, cutlass::half_t> || cute::is_same_v<ElementBMma, cutlass::bfloat16_t>) {
return make_tiled_mma(cute::SM100_MMA_F16BF16_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB>{});
}
if constexpr (cute::is_same_v<ElementBMma, cutlass::float_e4m3_t>) {
return make_tiled_mma(cute::SM100_MMA_F8F6F4_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB>{});
}
}
else { // If A needs to be transposed by MMA, fall back to SMEM from A MMA instructions
if constexpr (cute::is_same_v<ElementBMma, cutlass::half_t> || cute::is_same_v<ElementBMma, cutlass::bfloat16_t>) {
return make_tiled_mma(cute::SM100_MMA_F16BF16_SS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB>{});
}
if constexpr (cute::is_same_v<ElementBMma, cutlass::float_e4m3_t>) {
return make_tiled_mma(
cute::MMA_Traits<
cute::SM100_MMA_F8F6F4_SS,
ElementAMma,
ElementBMma,
ElementAccumulator,
cute::C<M>,
cute::C<N>,
cute::integral_constant<UMMA::Major, UmmaMajorA>,
cute::integral_constant<UMMA::Major, UmmaMajorB>,
cute::integral_constant<UMMA::ScaleIn, cute::UMMA::ScaleIn::One>,
cute::integral_constant<UMMA::ScaleIn, cute::UMMA::ScaleIn::One>>{});
}
}
}
//MMA 2Sm requested
else if constexpr (cute::is_base_of_v<KernelSchedule2Sm, KernelScheduleType>) {
if constexpr (UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedMixedInputSmemSm100, KernelScheduleType>) {
if constexpr (cute::is_same_v<ElementBMma, cutlass::half_t> || cute::is_same_v<ElementBMma, cutlass::bfloat16_t>) {
return make_tiled_mma(cute::SM100_MMA_F16BF16_2x1SM_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB>{});
}
if constexpr (cute::is_same_v<ElementBMma, cutlass::float_e4m3_t>) {
return make_tiled_mma(cute::SM100_MMA_F8F6F4_2x1SM_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB>{});
}
}
}
else {
static_assert(cutlass::detail::dependent_false<TileShape_MNK> == 0,
"Unsupported policy for SM100 collective builder.");
}
}
template<
class CtaShape_MNK
>
include/cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl
+173
View File
@@ -0,0 +1,173 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cutlass/gemm/collective/builders/sm100_common.inl"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ElementA,
class GmemLayoutATag,
int AlignmentA,
class ElementB,
class GmemLayoutBTag,
int AlignmentB,
class ElementAccumulator,
class TileShape_MNK,
class ClusterShape_MNK,
class StageCountType,
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm100,
arch::OpClassTensorOp,
ElementA,
GmemLayoutATag,
AlignmentA,
ElementB,
GmemLayoutBTag,
AlignmentB,
ElementAccumulator,
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<cute::is_same_v<KernelWarpSpecialized1SmSm100, BuilderScheduleTag> ||
(cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag> &&
(((sizeof(ElementA) * AlignmentA) % cutlass::gemm::collective::detail::tma_alignment_bytes != 0) ||
((sizeof(ElementB) * AlignmentB) % cutlass::gemm::collective::detail::tma_alignment_bytes != 0)))>
>
{
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
// Data type used by MMA instruction
using ElementAMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementA>());
using ElementBMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementB>());
using ElementAMma_SmemAllocType = cute::conditional_t<cute::sizeof_bits_v<ElementAMma> < 8, uint8_t, ElementAMma>;
using ElementBMma_SmemAllocType = cute::conditional_t<cute::sizeof_bits_v<ElementBMma> < 8, uint8_t, ElementBMma>;
using TiledMma = decltype(detail::sm100_make_trivial_tiled_mma<
ElementAMma, ElementBMma, ElementAccumulator,
decltype(cute::product_each(TileShape_MNK{})), ClusterShape_MNK,
UmmaMajorA, UmmaMajorB, BuilderScheduleTag>());
using AtomThrID = typename TiledMma::AtomThrID;
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K)
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// Assigning 4 warps for mainloop load
static constexpr int NumLoadThreads = 128;
using AlignmentTypeA = cute::uint_byte_t<static_cast<int>(sizeof(ElementA)) * AlignmentA>;
using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeA>, ElementA>;
using GmemTiledCopyA = decltype(detail::make_simt_gmem_tiled_copy<
GmemCopyAtomA, NumLoadThreads, AlignmentA, TagToStrideA_t<GmemLayoutATag>,
decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
using BlockTileA_M = decltype(cute::size<0,0>(MmaShapeA_MK{}) * cute::size<1>(MmaShapeA_MK{}));
using BlockTileA_K = decltype(cute::size<0,1>(MmaShapeA_MK{}) * cute::size<2>(MmaShapeA_MK{}));
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
UmmaMajorA, ElementAMma_SmemAllocType, BlockTileA_M, BlockTileA_K>());
using AlignmentTypeB = cute::uint_byte_t<static_cast<int>(sizeof(ElementB)) * AlignmentB>;
using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeB>, ElementB>;
using GmemTiledCopyB = decltype(detail::make_simt_gmem_tiled_copy<
GmemCopyAtomB, NumLoadThreads, AlignmentB, TagToStrideB_t<GmemLayoutBTag>,
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{}));
using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{}));
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
UmmaMajorB, ElementBMma_SmemAllocType, BlockTileB_N, BlockTileB_K>());
static constexpr uint32_t AccumulatorPipelineStageCount = 2;
// Calculate scheduler pipeline stages. Having one more stage than the accumulator allows more latency hiding.
static constexpr uint32_t SchedulerPipelineStageCount = AccumulatorPipelineStageCount + 1;
// AccumulatorPipeline = PipelineUmmaAsync
static constexpr auto AccumulatorPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount>::SharedStorage);
// CLCPipeline = PipelineCLCFetchAsync
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
// CLC (scheduler) response
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
static constexpr auto KernelSmemCarveout = static_cast<int>( AccumulatorPipelineStorage +
CLCPipelineStorage +
CLCResponseStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
using MainloopPipelineStorage = typename cutlass::PipelineUmmaConsumerAsync<1>::SharedStorage;
static constexpr int PipelineStages = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
cutlass::gemm::MainloopSm100UmmaCpAsyncWarpSpecialized<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK>,
TileShape_MNK,
ElementA,
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
ElementB,
cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>,
TiledMma,
GmemTiledCopyA,
SmemLayoutAtomA,
void,
cute::identity,
GmemTiledCopyB,
SmemLayoutAtomB,
void,
cute::identity
>;
};
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl
+346
View File
@@ -0,0 +1,346 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cutlass/gemm/collective/builders/sm100_common.inl"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace detail {
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template<
int CapacityBytes,
class ElementA,
class ElementAMma,
class ElementScale,
class ElementZero,
class ElementB,
class CtaTileShape_MNK,
class TiledMma,
class KernelScheduleType,
UMMA::Major UmmaMajorA,
int ScaleGranularityK,
int stages
>
constexpr cute::tuple<int, int, int>
sm100_compute_stage_count_or_override_mixed_input(StageCount<stages> stage_count) {
constexpr int Load2TransformStageCount = stages;
constexpr int Transform2MmaStageCount = stages;
constexpr int AccumulatorStageCount = stages;
return cute::make_tuple(Load2TransformStageCount, Transform2MmaStageCount, AccumulatorStageCount);
}
template<
int CapacityBytes,
class ElementA,
class ElementAMma,
class ElementScale,
class ElementZero,
class ElementB,
class CtaTileShape_MNK,
class TiledMma,
class KernelScheduleType,
UMMA::Major UmmaMajorA,
int ScaleGranularityK,
int carveout_bytes
>
constexpr cute::tuple<int, int, int>
sm100_compute_stage_count_or_override_mixed_input(StageCountAutoCarveout<carveout_bytes> stage_count) {
constexpr int CtaM = get<0>(CtaTileShape_MNK{});
constexpr int CtaN = get<1>(CtaTileShape_MNK{});
static_assert(CtaN <= 128, "Can't support CtaN>128 tiles");
constexpr int CtaK = get<2>(CtaTileShape_MNK{});
using AtomThrID = typename TiledMma::AtomThrID;
constexpr int TmemColumns = 512;
constexpr bool IsAComputeinTmem = UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedMixedInputSmemSm100, KernelScheduleType>;
constexpr bool IsAComputeinSmem = !IsAComputeinTmem;
// Detect 2x2 TMEM layout
constexpr int TmemAccWordsPerDP = (CtaM == 64 && size(AtomThrID{}) == 2) ? CtaN/2 : CtaN;
constexpr int TmemAWordsPerDP = CtaK / 2;
constexpr int AccumulatorStageCount = (IsAComputeinTmem) ? ((TmemAccWordsPerDP == 128) ? 2 : 3) : (TmemColumns / TmemAccWordsPerDP);
constexpr int SmemCapacityAfterMma2AccumCarveout = CapacityBytes - (carveout_bytes + AccumulatorStageCount * 32);
constexpr int TmemInAStageCount_Potential = (IsAComputeinTmem) ? (TmemColumns - AccumulatorStageCount * TmemAccWordsPerDP) / TmemAWordsPerDP : 10000;
// Mainload2Transform Pipeline
constexpr auto load2transform_pipeline_bytes = sizeof(typename cutlass::PipelineTmaTransformAsync<1>::SharedStorage);
constexpr auto a_bits = cute::sizeof_bits_v<ElementA>; // ElementA introduce here
constexpr auto s_bits = cute::is_void_v<ElementScale> ? 0 : cute::sizeof_bits_v<ElementScale>;
constexpr auto z_bits = cute::is_void_v<ElementZero> ? 0 : cute::sizeof_bits_v<ElementZero>;
constexpr auto load2mma_pipeline_bytes = sizeof(typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage);
constexpr auto b_bits = cute::sizeof_bits_v<ElementB>; // ElementB introduce here
constexpr int ab_stage_bytes =
cutlass::bits_to_bytes(a_bits * size<0>(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{})) +
cutlass::bits_to_bytes(s_bits * size<0>(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{}) / ScaleGranularityK) +
cutlass::bits_to_bytes(z_bits * size<0>(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{}) / ScaleGranularityK) +
cutlass::bits_to_bytes(b_bits * size<1>(CtaTileShape_MNK{}) / size(AtomThrID{}) * size<2>(CtaTileShape_MNK{})) +
static_cast<int>(load2transform_pipeline_bytes) + static_cast<int>(load2mma_pipeline_bytes);
// Transform2Mma Pipeline
constexpr auto transform2mma_pipeline_bytes = sizeof(typename cutlass::PipelineUmmaConsumerAsync<1>::SharedStorage);
constexpr auto a_compute_bits = cute::sizeof_bits_v<ElementAMma>;
constexpr int ab_compute_stage_bytes =
cutlass::bits_to_bytes(a_compute_bits * int(IsAComputeinSmem) * size<0>(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{})) + // If ACompute is in TMEM, Acompute buffer has 0 bytes.
static_cast<int>(transform2mma_pipeline_bytes);
constexpr int ABComputeStageCount_Potential = SmemCapacityAfterMma2AccumCarveout / (ab_stage_bytes + ab_compute_stage_bytes);
// The number of SMEM buffers for A, B. ACompute (if in SMEM), BCompute should be at least Transform2MmaStageCount
constexpr int Transform2MmaStageCount = std::min(TmemInAStageCount_Potential, ABComputeStageCount_Potential);
constexpr int SmemCapacityAfterABComputeCarveout = SmemCapacityAfterMma2AccumCarveout - (Transform2MmaStageCount * ab_compute_stage_bytes);
// Can we boost the number of buffers for A and B?
constexpr int Load2TransformStageCount = SmemCapacityAfterABComputeCarveout / ab_stage_bytes;
static_assert(Load2TransformStageCount >= 2 && Transform2MmaStageCount >= 2 && AccumulatorStageCount >= 2, "Not enough SMEM or TMEM capacity for selected tile size");
return cute::make_tuple(Load2TransformStageCount, Transform2MmaStageCount, AccumulatorStageCount);
}
} // namespace detail
template <typename LayoutScale>
constexpr int get_ScaleGranularityK() {
if constexpr (cute::is_void_v<LayoutScale>) {
return 1;
} else {
return size<1,0>(LayoutScale{});
}
}
// Mixed Input MMA kernels builder
template <
class ElementAOptionalTuple,
class GmemLayoutATagTuple,
int AlignmentA,
class ElementBOptionalTuple,
class GmemLayoutBTag,
int AlignmentB,
class ElementAccumulator,
class TileShape_MNK, // The Cluster-level TileShape
class ClusterShape_MNK,
class StageCountType,
class KernelScheduleType
>
struct CollectiveBuilder<
arch::Sm100,
arch::OpClassTensorOp,
ElementAOptionalTuple, // ElementA
GmemLayoutATagTuple, // LayoutA
AlignmentA,
ElementBOptionalTuple, // ElementB
GmemLayoutBTag, // LayoutB
AlignmentB,
ElementAccumulator,
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, int)
StageCountType,
KernelScheduleType,
cute::enable_if_t<
(cute::is_base_of_v<KernelScheduleSm100MixedInputGemm, KernelScheduleType>) &&
((sizeof(float) * AlignmentA) % detail::tma_alignment_bytes == 0) &&
((sizeof(float) * AlignmentB) % detail::tma_alignment_bytes == 0)>>
{
using GmemLayoutATag = detail::deduce_mixed_width_dtype_t<0, GmemLayoutATagTuple>;
using GmemLayoutScaleTag = detail::deduce_mixed_width_dtype_t<1, GmemLayoutATagTuple>;
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
using ElementA = detail::deduce_mixed_width_dtype_t<0, ElementAOptionalTuple>;
using ElementB = detail::deduce_mixed_width_dtype_t<0, ElementBOptionalTuple>;
using ElementScale = detail::deduce_mixed_width_dtype_t<1, ElementAOptionalTuple>;
using ElementZero = detail::deduce_mixed_width_dtype_t<2, ElementAOptionalTuple>;
static constexpr bool NeitherIsTuple = !cute::is_tuple<ElementAOptionalTuple>::value && !cute::is_tuple<ElementBOptionalTuple>::value;
static constexpr bool IsANarrow = cute::sizeof_bits_v<ElementA> < cute::sizeof_bits_v<ElementB>;
static constexpr bool IsMixedInput = cute::sizeof_bits_v<ElementA> != cute::sizeof_bits_v<ElementB>;
static_assert(IsMixedInput, "Mixed Input GEMM Kernel doesn't support regular gemm.");
static_assert((cute::is_tuple<ElementAOptionalTuple>::value ^ cute::is_tuple<ElementBOptionalTuple>::value ||
(NeitherIsTuple && (cute::sizeof_bits<ElementA>::value != cute::sizeof_bits<ElementB>::value))),
"Either A OR B must be a tuple or the widths of A and B must be different.");
using ElementPairA = cute::conditional_t<IsMixedInput && IsANarrow && NeitherIsTuple, cute::tuple<ElementA>, ElementAOptionalTuple>;
using ElementPairB = cute::conditional_t<IsMixedInput && !IsANarrow && NeitherIsTuple, cute::tuple<ElementB>, ElementBOptionalTuple>;
static constexpr bool IsATransformed = cute::is_tuple<ElementPairA>::value;
static_assert(IsATransformed, "A matrix should be transformed.");
// For fp32 types, map to tf32 MMA value type.
using ElementMma = cute::conditional_t<cute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
using ElementAMma = ElementMma;
using ElementBMma = ElementMma;
static constexpr int IsSubbyteA = cute::sizeof_bits_v<ElementA> < 8;
using TmaElementA = cute::conditional_t<IsSubbyteA, uint8_t, ElementA>;
static constexpr int ScalingFactor = 1;
using TiledMma = decltype(detail::sm100_make_trivial_mixed_input_tiled_mma<ElementAMma, ElementB, ElementAccumulator, TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, KernelScheduleType>());
using AtomThrID = typename TiledMma::AtomThrID;
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
using CtaTileShape_MNK = decltype(shape_div(TileShape_MNK{}, AtomThrShapeMNK{}));
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K)
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
using BlockTileA_M = decltype(cute::size<0,0>(MmaShapeA_MK{}) * cute::size<1>(MmaShapeA_MK{}));
using BlockTileA_K = decltype(cute::size<0,1>(MmaShapeA_MK{}) * cute::size<2>(MmaShapeA_MK{}));
using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<1>(ClusterShape_MNK{})));
using GmemTiledCopyB = decltype(detail::sm100_cluster_shape_to_tma_atom_B(ClusterShape_MNK{}, AtomThrID{}));
// Input transform kernel can not use TMA 2SM instructions.
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorA, ElementA,
BlockTileA_M, BlockTileA_K>());
using SmemLayoutAtomACompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorA, ElementAMma, BlockTileA_M, BlockTileA_K>());
using SmemLayoutAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType<
SmemLayoutAtomA, SmemLayoutAtomACompute>;
static constexpr int MMA_M = cute::size<0,0>(MmaShapeA_MK{});
using CopyAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType<
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementA>,
cute::conditional_t<(UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedMixedInputSmemSm100, KernelScheduleType>),
cute::conditional_t<(MMA_M == 64 && size(AtomThrID{}) == 1), SM100_TMEM_STORE_16dp256b1x, SM100_TMEM_STORE_32dp32b8x>, // TS Implementation
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementA>> // SS Implementation
>;
using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{}));
using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{}));
// Input transform kernel can not use TMA 2SM instructions.
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorB, ElementB,
BlockTileB_N, BlockTileB_K>());
using SmemLayoutAtomBCompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorB, ElementBMma, BlockTileB_N, BlockTileB_K>());
using SmemLayoutAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType<
SmemLayoutAtomB, SmemLayoutAtomBCompute>;
using CopyAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType<
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementB>,
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementMma>
>;
//Creating the stride of Transformed Input
using StrideA = cutlass::gemm::TagToStrideA_t<GmemLayoutATag>;
using LayoutScale = cutlass::gemm::TagToStrideA_t<GmemLayoutScaleTag>;
using VoidShapeScale = Shape<Shape<Int<128>, _1>, Shape<Int<64>, _1>, _1>; //Dummy Value to create a dummy ScaleConfig
using VoidStrideScale = Stride<Stride<_0,_1>,Stride<_0, _1>, _1>;
using VoidLayoutScale = Layout<VoidShapeScale, VoidStrideScale>;
using NonVoidLayoutScale = cute::conditional_t<
cute::is_void_v<LayoutScale>, VoidLayoutScale, LayoutScale>;
using StridePairA = decltype(cute::make_tuple(StrideA{}, NonVoidLayoutScale{}));
// SmemCarveout
static constexpr int SchedulerPipelineStageCount = 3;
static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v<KernelScheduleSm100PtrArrayFastFP32Gemm, KernelScheduleType>);
// CLCPipeline = PipelineCLCFetchAsync
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
// CLC (scheduler) response
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
// CLC Throttle pipeline storage
static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage);
// Tmem dealloc
static constexpr auto TmemDeallocStorage = sizeof(cutlass::arch::ClusterBarrier);
// Tmem ptr storage
static constexpr auto TmemBasePtrsStorage = sizeof(uint32_t);
// Tensormap Storage
static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
static constexpr auto KernelSmemCarveout = static_cast<int>( CLCPipelineStorage +
CLCResponseStorage +
CLCThrottlePipelineStorage +
TmemDeallocStorage +
TmemBasePtrsStorage +
TensorMapStorage);
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int Sm100ReducedSmemCapacityBytes = detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr int ScaleGranularityK = get_ScaleGranularityK<LayoutScale>();
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_mixed_input<
Sm100ReducedSmemCapacityBytes, TmaElementA, ElementAMma, ElementScale, ElementZero, ElementB, CtaTileShape_MNK, TiledMma, KernelScheduleType, UmmaMajorA, ScaleGranularityK>(StageCountType{});
static constexpr int Load2TransformPipelineStageCount = get<0>(stage_info);
static constexpr int Transform2MmaPipelineStageCount = get<1>(stage_info);
static constexpr int AccumulatorPipelineStageCount = get<2>(stage_info);
static_assert(!IsArrayOfPointersGemm, "mixed input does not support grouped gemm on Blackwell");
using DispatchPolicy = cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedMixedInput<
Load2TransformPipelineStageCount,
Transform2MmaPipelineStageCount,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape_MNK,
ElementPairA,
StridePairA,
ElementPairB,
cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>,
TiledMma,
GmemTiledCopyA,
SmemLayoutAtomPairA,
CopyAtomPairA,
cute::identity,
GmemTiledCopyB,
SmemLayoutAtomPairB,
CopyAtomPairB,
cute::identity
>;
};
} // namespace cutlass::gemm::collective
include/cutlass/gemm/collective/builders/sm100_umma_builder.inl
+1
View File
@@ -184,6 +184,7 @@ struct CollectiveBuilder<
not cute::is_complex_v<ElementA> && not cute::is_complex_v<ElementB> &&
// Dense Gemm / PtrArrayDenseGemm
(
(not cute::is_same_v<KernelWarpSpecialized1SmSm100, BuilderScheduleTag>) &&
(cute::is_base_of_v<KernelScheduleSm100DenseGemm, BuilderScheduleTag> ||
cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>)) &&
// Alignment check
include/cutlass/gemm/collective/builders/sm103_blockscaled_umma_builder.inl
+546
View File
@@ -0,0 +1,546 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cutlass/gemm/collective/builders/sm100_common.inl"
#include "cutlass/detail/sm103_blockscaled_layout.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace detail {
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template <
int CapacityBytes,
class ElementA,
class ElementB,
class TileShapeMNK,
class TileShapeSFA,
class TileShapeSFB,
int stages
>
constexpr int
sm103_compute_stage_count_or_override_blockscaled(StageCount<stages> stage_count) {
return stages;
}
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template <
int CapacityBytes,
class ElementA,
class ElementB,
class TileShapeMNK,
class TileShapeSFA,
class TileShapeSFB,
int carveout_bytes
>
constexpr auto
sm103_compute_stage_count_or_override_blockscaled(StageCountAutoCarveout<carveout_bytes> stage_count) {
// For F8F6F4 MMA sub-bytes, ElementA/B will be passed in as uint8_t
// Each stage include (CollectiveMma::SharedStorage)
// 1. smem for A and smem for B (CollectiveMma::SharedStorage::TensorStorage)
// 2. one MainloopPipeline = PipelineTmaUmmaAsync (CollectiveMma::SharedStorage::SharedStorage)
// 3. smem for SFB and smem for SFB (CollectiveMma::SharedStorage::TensorStorage, independent of input size b.c. sizeof(sf) is fixed)
constexpr auto mainloop_pipeline_bytes = sizeof(typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage);
constexpr auto a_bits = cute::sizeof_bits_v<ElementA>;
constexpr auto b_bits = cute::sizeof_bits_v<ElementB>;
constexpr auto stage_sfa_bytes = size(filter_zeros(TileShapeSFA{}));
constexpr auto stage_sfb_bytes = size(filter_zeros(TileShapeSFB{}));
constexpr int stage_bytes =
cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
static_cast<int>(mainloop_pipeline_bytes * 2 + stage_sfa_bytes + stage_sfb_bytes);
constexpr int ab_buffer = (CapacityBytes - carveout_bytes) / stage_bytes;
constexpr int sb_buffer = ab_buffer + (CapacityBytes - carveout_bytes - ab_buffer * stage_bytes) / (mainloop_pipeline_bytes + stage_sfa_bytes + stage_sfb_bytes);
return make_tuple(ab_buffer, sb_buffer);
}
template<
class ElementAMma,
class ElementBMma,
class ElementAccumulator,
class ElementSF,
class TileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB,
int SFVectorSize
>
constexpr auto
sm103_make_blockscaled_1sm_tiled_mma() {
using AtomLayout_MNK = Layout<ClusterShape_MNK>;
constexpr int M = cute::size<0>(TileShape_MNK{});
static_assert(M == 128 || M == 256, "Invalid TileShape_M.");
constexpr int N = cute::size<1>(TileShape_MNK{});
static_assert(N % 64 == 0 && N <= 256, "Invalid TileShape_N.");
if constexpr (cute::is_same_v<ElementAMma, cutlass::float_e2m1_t> ||
cute::is_same_v<ElementAMma, cutlass::type_erased_dynamic_float4_t>) {
return make_tiled_mma(cute::SM103::SM103_MXF4_ULTRA_SS_VS<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
M, N, SFVectorSize, UmmaMajorA, UmmaMajorB>{});
}
else {
static_assert(cutlass::detail::dependent_false<ElementAMma>,
"Unsupported configuration for SM103 collective builder.");
}
}
template<
class ElementAMma,
class ElementBMma,
class ElementAccumulator,
class ElementSF,
class TileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB,
int SFVectorSize
>
constexpr auto
sm103_make_blockscaled_2sm_tiled_mma() {
using AtomLayout_MNK = Layout<decltype(shape_div(ClusterShape_MNK{}, Shape<_2,_1,_1>{}))>;
constexpr int M = cute::size<0>(TileShape_MNK{});
static_assert(M == 128 || M == 256, "Invalid TileShape_M.");
constexpr int N = cute::size<1>(TileShape_MNK{});
static_assert(N % 64 == 0 && N <= 256, "Invalid TileShape_N.");
if constexpr (cute::is_same_v<ElementAMma, cutlass::float_e2m1_t> ||
cute::is_same_v<ElementAMma, cutlass::type_erased_dynamic_float4_t>) {
return make_tiled_mma(cute::SM103::SM103_MXF4_ULTRA_2x1SM_SS_VS<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
M, N, SFVectorSize, UmmaMajorA, UmmaMajorB>{});
}
else {
static_assert(cutlass::detail::dependent_false<ElementAMma>,
"Unsupported configuration for SM103 collective builder.");
}
}
template<
class ElementAMma,
class ElementBMma,
class ElementAccumulator,
class ElementSF,
class ClusterTileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB,
class BuilderScheduleTag
>
constexpr auto
sm103_make_blockscaled_tiled_mma() {
constexpr uint32_t SFVectorSize = find_vector_size<BuilderScheduleTag>();
// MMA_2SM requested
if constexpr (cute::is_base_of_v<KernelSchedule2Sm, BuilderScheduleTag>) {
return sm103_make_blockscaled_2sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
ClusterTileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>();
}
// MMA_1SM requested
else if constexpr (cute::is_base_of_v<KernelSchedule1Sm, BuilderScheduleTag>) {
return sm103_make_blockscaled_1sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
ClusterTileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>();
}
// Auto scheduling requested
else if constexpr (cute::is_same_v<BuilderScheduleTag, KernelScheduleAuto>) {
if constexpr (cute::get<0>(ClusterShape_MNK{}) % 2 == 0) {
return sm103_make_blockscaled_2sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
ClusterTileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>();
}
else {
return sm103_make_blockscaled_1sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
ClusterTileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>();
}
}
else {
static_assert(cutlass::detail::dependent_false<ElementAMma>,
"Unsupported policy for SM103 collective builder.");
}
}
template<
class ElementAMma,
class ElementBMma,
class ElementAccumulator,
class ElementSF,
class TileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB,
uint32_t SFVectorSize,
class BuilderScheduleTag,
bool Is2SM
>
struct Sm103TrivialBlockscaledMma {};
template<
class ElementAMma,
class ElementBMma,
class ElementAccumulator,
class ElementSF,
class TileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB,
uint32_t SFVectorSize,
class BuilderScheduleTag
>
struct Sm103TrivialBlockscaledMma< ElementAMma,
ElementBMma,
ElementAccumulator,
ElementSF,
TileShape_MNK,
ClusterShape_MNK,
UmmaMajorA,
UmmaMajorB,
SFVectorSize,
BuilderScheduleTag,
true /*Is2SM*/> {
using type = decltype(sm103_make_blockscaled_2sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>());
};
template<
class ElementAMma,
class ElementBMma,
class ElementAccumulator,
class ElementSF,
class TileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB,
uint32_t SFVectorSize,
class BuilderScheduleTag
>
struct Sm103TrivialBlockscaledMma< ElementAMma,
ElementBMma,
ElementAccumulator,
ElementSF,
TileShape_MNK,
ClusterShape_MNK,
UmmaMajorA,
UmmaMajorB,
SFVectorSize,
BuilderScheduleTag,
false /*Is2SM*/> {
using type = decltype(sm103_make_blockscaled_1sm_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, SFVectorSize>());
};
template <class ElementType>
CUTLASS_HOST_DEVICE
static constexpr bool
is_sm103_block_scale_input() {
// Allowed input element datatype for block-scaling GEMM
return ( cute::is_same_v<ElementType, cutlass::type_erased_dynamic_float4_t> ||
cute::is_same_v<ElementType, cutlass::float_e2m1_t>);
}
template <class BlockScaleConfig, class MmaShapeA_MK>
constexpr
auto sm103_sfa_smem_atom_layout() {
constexpr int SF_BUFFERS_PER_TILE_K = BlockScaleConfig::SFVecSize == 16 ? 4 : 2;
auto mma_sfa_tiler = make_shape(get<0,0>(MmaShapeA_MK{})*get<1>(MmaShapeA_MK{}), get<0,1>(MmaShapeA_MK{}) * get<2>(MmaShapeA_MK{}) / Int<SF_BUFFERS_PER_TILE_K>{});
return tiled_product(typename BlockScaleConfig::SfAtom{},
make_layout(shape_div(mma_sfa_tiler, product_each(shape(typename BlockScaleConfig::SfAtom{})))));
}
template <class BlockScaleConfig, class MmaShapeB_NK, int SFVecSize>
constexpr
auto sm103_sfb_smem_atom_layout() {
auto sSFB = [&]() {
constexpr int MMA_N = get<0>(MmaShapeB_NK{});
constexpr int NonPow2N = 192;
constexpr int NonPow2N_RoundUp = 256;
// If MMA_N is 192, we need to operate at MMA_N = 256 granularity for UTCCP to work for ScaleFactorB.
// Both TMA and UTCCP will transfer scale factor B as if we have 256 columns in B matrix.
constexpr int MMA_N_SFB = (MMA_N == NonPow2N) ? NonPow2N_RoundUp : MMA_N;
constexpr int SF_BUFFERS_PER_TILE_K = BlockScaleConfig::SFVecSize == 16 ? 4 : 2;
auto mma_sfb_tiler = make_shape(Int<MMA_N_SFB>{}, get<1>(MmaShapeB_NK{}) / Int<SF_BUFFERS_PER_TILE_K>{});
if constexpr(Int<MMA_N>{} == Int<128>{}) {
return tiled_product(typename BlockScaleConfig::SfAtom{},
make_layout(shape_div(mma_sfb_tiler,product_each(shape(typename BlockScaleConfig::SfAtom{})))));
}
else {
using SfKMajorAtom256 = Layout< Shape< Shape<_32,_4, _2>, Shape<Int<SFVecSize>, _4>>,
Stride<Stride<_16,_4, Int<get<1>(mma_sfb_tiler)/SFVecSize/4*512>>, Stride< _0, _1>>>;
return tiled_product(SfKMajorAtom256{},
make_layout(shape_div(mma_sfb_tiler,product_each(shape(SfKMajorAtom256{})))));
}
}();
return sSFB;
}
} // namespace detail
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class ElementPairA,
class GmemLayoutATag,
int AlignmentA,
class ElementPairB,
class GmemLayoutBTag,
int AlignmentB,
class ElementAccumulator,
class TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1)
class StageCountType,
class BuilderScheduleTag
>
struct CollectiveBuilder<
arch::Sm103,
arch::OpClassBlockScaledTensorOp,
ElementPairA,
GmemLayoutATag,
AlignmentA,
ElementPairB,
GmemLayoutBTag,
AlignmentB,
ElementAccumulator,
TileShape_MNK,
ClusterShape_MNK,
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
// Not paired input, Not Complex input
(cute::is_tuple_v<ElementPairA> && cute::is_tuple_v<ElementPairB> &&
not cute::is_complex_v<ElementPairA> && not cute::is_complex_v<ElementPairB>) &&
// Blockscaled Gemm
(cute::is_base_of_v<KernelScheduleSm103BlockScaledGemm, BuilderScheduleTag> ||
cute::is_base_of_v<KernelSchedulePtrArraySm103BlockScaledGemm, BuilderScheduleTag> ||
cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>) &&
// Alignment check
detail::sm1xx_blockscaled_gemm_is_aligned<remove_cvref_t<decltype(get<0>(ElementPairA{}))>,
AlignmentA,
remove_cvref_t<decltype(get<0>(ElementPairB{}))>,
AlignmentB,
BuilderScheduleTag>()>>
{
using ElementA = remove_cvref_t<decltype(get<0>(ElementPairA{}))>;
using ElementB = remove_cvref_t<decltype(get<0>(ElementPairB{}))>;
using ElementSF = remove_cvref_t<decltype(get<1>(ElementPairA{}))>;
static_assert(cute::is_tuple<ElementPairA>::value, "Expecting ElementPairA to be a tuple.");
static_assert(cute::is_tuple<ElementPairB>::value, "Expecting ElementPairB to be a tuple.");
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
static_assert(cute::size<2>(TileShape_MNK{}) == _768{}, "TileShape_K should 768 for MMA kernels");
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
static_assert(cutlass::gemm::detail::is_k_major_A<GmemLayoutATag>() && cutlass::gemm::detail::is_k_major_B<GmemLayoutBTag>(), "Only K major inputs are supported");
static_assert(cutlass::gemm::collective::detail::is_sm103_block_scale_input<ElementA>(), "Incorrect type for A matrix");
static_assert(cutlass::gemm::collective::detail::is_sm103_block_scale_input<ElementB>(), "Incorrect type for B matrix");
static_assert(cute::is_same_v<ElementSF, cutlass::float_ue8m0_t> ||
cute::is_same_v<ElementSF, cutlass::float_ue4m3_t>, "Incorrect scale factor type");
// Data type used by MMA instruction
using ElementAMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementA, false /*UseQmma=false*/>());
using ElementBMma = decltype(cutlass::gemm::collective::detail::sm1xx_kernel_input_element_to_mma_input_element<ElementB, false /*UseQmma=false*/>());
static constexpr uint32_t SFVectorSize = detail::find_vector_size<BuilderScheduleTag>();
static constexpr bool is_2sm = cute::is_base_of_v<KernelSchedule2Sm, BuilderScheduleTag> ||
(cute::is_same_v<BuilderScheduleTag, KernelScheduleAuto> &&
(cute::is_static_v<ClusterShape_MNK> && cute::get<0>(ClusterShape_MNK{}) % 2 == 0));
using TiledMma = typename cutlass::gemm::collective::detail::Sm103TrivialBlockscaledMma<ElementAMma, ElementBMma, ElementAccumulator, ElementSF,
TileShape_MNK, ClusterShape_MNK,
UmmaMajorA, UmmaMajorB, SFVectorSize, BuilderScheduleTag, is_2sm>::type;
using AtomThrID = typename TiledMma::AtomThrID;
using Sm1xxBlkScaledConfig = cutlass::detail::Sm103BlockScaledConfig<SFVectorSize>;
using ElementAMma_SmemAllocType = uint8_t;
// ElementAMma;
using ElementBMma_SmemAllocType = uint8_t;
// ElementBMma;
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K)
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
using GmemTiledCopyA = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_A(
ClusterShape_MNK{}, AtomThrID{}));
using GmemTiledCopyB = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_B(
ClusterShape_MNK{}, AtomThrID{}));
using GmemTiledCopySFA = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_A(
ClusterShape_MNK{}, AtomThrID{}));
using GmemTiledCopySFB = decltype(cutlass::gemm::collective::detail::sm100_cluster_shape_to_tma_atom_SFB(
ClusterShape_MNK{}, AtomThrID{}));
using GmemTiledCopyPairA = decltype(cute::make_tuple(GmemTiledCopyA{}, GmemTiledCopySFA{}));
using GmemTiledCopyPairB = decltype(cute::make_tuple(GmemTiledCopyB{}, GmemTiledCopySFB{}));
//
// Construct SMEM layout (SmemLayoutAtom) for A and SFA
//
using SmemLayoutAtomA = UMMA::Layout_K_SW128_Atom<ElementAMma_SmemAllocType>;
// A single indivisible block will hold 4 scale factors of 128 rows/columns (A/B matrix).
// 4 is chosen to make consecutive 32bits of data to have scale factors for only a single row (col). 32bits corresponds to the TMEM word size
static constexpr int MMA_M = cute::size<0>(TileShape_MNK{}) / cute::size(AtomThrID{});
using SmemLayoutAtomSFA = decltype(detail::sm103_sfa_smem_atom_layout<Sm1xxBlkScaledConfig,MmaShapeA_MK>());
using SmemLayoutAtomsA = decltype(cute::make_tuple(SmemLayoutAtomA{}, SmemLayoutAtomSFA{}));
//
// Construct SMEM layout（SmemLayoutAtom）for B and SFB
//
using SmemLayoutAtomB = UMMA::Layout_K_SW128_Atom<ElementBMma_SmemAllocType>;
static constexpr int MMA_N = cute::size<1>(TileShape_MNK{});
// If MMA_N is 192, we need to operate at MMA_N = 256 granularity for UTCCP to work for ScaleFactorB.
// Both TMA and UTCCP will transfer scale factor B as if we have 256 columns in B matrix.
using SmemLayoutAtomSFB = decltype(detail::sm103_sfb_smem_atom_layout<Sm1xxBlkScaledConfig,decltype(select<1,2>(TileShape_MNK{})),SFVectorSize>());
using SmemLayoutAtomsB = decltype(cute::make_tuple(SmemLayoutAtomB{}, SmemLayoutAtomSFB{}));
//
// Construct Strides for A, SFA, B, and SFB
//
using StrideA = cutlass::gemm::TagToStrideA_t<GmemLayoutATag>;
using StrideB = cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
using InternalLayoutSFA = decltype(Sm1xxBlkScaledConfig::deduce_layoutSFA());
using InternalLayoutSFB = decltype(Sm1xxBlkScaledConfig::deduce_layoutSFB());
using LayoutSFA = cute::conditional_t<cute::is_same_v<InternalStrideA, StrideA>, InternalLayoutSFA, InternalLayoutSFA *>;
using LayoutSFB = cute::conditional_t<cute::is_same_v<InternalStrideB, StrideB>, InternalLayoutSFB, InternalLayoutSFB *>;
using StridePairA = decltype(cute::make_tuple(StrideA{}, LayoutSFA{}));
using StridePairB = decltype(cute::make_tuple(StrideB{}, LayoutSFB{}));
//
// Others
//
static constexpr cutlass::sm103::detail::KernelPrefetchType PrefetchType = cute::is_base_of_v<KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch, BuilderScheduleTag>
|| cute::is_base_of_v<KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch, BuilderScheduleTag>
? cutlass::sm103::detail::KernelPrefetchType::Disable :
cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch;
static constexpr uint32_t AccumulatorPipelineStageCount = (MMA_N == 256) ? 1 : 2;
static constexpr uint32_t SchedulerPipelineStageCount = 3;
// AccumulatorPipeline = PipelineUmmaAsync
static constexpr auto AccumulatorPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount>::SharedStorage);
// CLCPipeline = PipelineCLCFetchAsync
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount>::SharedStorage);
// LoadOrderBarrier = OrderedSequenceBarrier<1,2>
static constexpr auto LoadOrderBarrierStorage = sizeof(typename cutlass::OrderedSequenceBarrier<1,2>::SharedStorage);
// CLC (scheduler) response
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
// CLC Throttle pipeline storage
static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage);
// Tmem dealloc
static constexpr auto TmemDeallocStorage = sizeof(cutlass::arch::ClusterBarrier);
// Tmem ptr storage
static constexpr auto TmemBasePtrsStorage = AccumulatorPipelineStageCount * sizeof(uint32_t);
// Tensormap Storage
static constexpr bool IsArrayOfPointersGemm = cute::is_base_of_v<KernelSchedulePtrArraySm103BlockScaledGemm, BuilderScheduleTag>;
static constexpr auto TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 4 /* for A, B, SFA and SFB */ : 0;
// TMA Load Prefetch Storage
static constexpr auto TmaPrefetchStorage = 0;
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
static constexpr auto KernelSmemCarveout = static_cast<int>( AccumulatorPipelineStorage +
CLCPipelineStorage +
LoadOrderBarrierStorage +
CLCResponseStorage +
CLCThrottlePipelineStorage +
TmemDeallocStorage +
TmemBasePtrsStorage +
TensorMapStorage +
TmaPrefetchStorage);
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int Sm100ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<Int<MMA_M>, Int<MMA_N/cute::size(AtomThrID{})>, _128>; // SmemAllocTypes are uint8_t. We always allocate 128bytes
static constexpr auto PipelineStages = cutlass::gemm::collective::detail::sm103_compute_stage_count_or_override_blockscaled<
Sm100ReducedSmemCapacityBytes, ElementAMma_SmemAllocType, ElementBMma_SmemAllocType, SmemTileShape, SmemLayoutAtomSFA, SmemLayoutAtomSFB>(StageCountType{});
using DispatchPolicy = typename cute::conditional_t<IsArrayOfPointersGemm,
cutlass::gemm::MainloopSm103ArrayTmaUmmaWarpSpecializedBlockScaled<
get<0>(PipelineStages),
get<1>(PipelineStages),
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK,
PrefetchType
>,
cutlass::gemm::MainloopSm103TmaUmmaWarpSpecializedBlockScaled<
get<0>(PipelineStages),
get<1>(PipelineStages),
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK,
PrefetchType
>
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape_MNK,
ElementPairA,
StridePairA,
ElementPairB,
StridePairB,
TiledMma,
GmemTiledCopyPairA,
SmemLayoutAtomsA,
void,
cute::identity,
GmemTiledCopyPairB,
SmemLayoutAtomsB,
void,
cute::identity
>;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/builders/sm1xx_common.inl
+12
View File
@@ -141,6 +141,18 @@ constexpr uint32_t find_vector_size() {
cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized2SmNvf4Sm100> ||
cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecializedNvf4Sm120> ||
cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecializedPingpongNvf4Sm120>
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103>
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103>
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch>
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch>
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch>
|| cute::is_same_v<BuilderScheduleTag, KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch>
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103>
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103>
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch>
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch>
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch>
|| cute::is_same_v<BuilderScheduleTag, KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch>
) {
return 16;
}
include/cutlass/gemm/collective/collective_builder.hpp
+3
View File
@@ -47,6 +47,9 @@
#include "cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_simt_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm103_blockscaled_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_mma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_blockscaled_mma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_sparse_mma_builder.inl"
include/cutlass/gemm/collective/collective_mma.hpp
+4
View File
@@ -63,6 +63,10 @@
#include "cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp"
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp"
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp"
#include "cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm103_blockscaled_mma_array_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm120_mma_tma.hpp"
#include "cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp"
#include "cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp"
include/cutlass/gemm/collective/sm100_mma_cpasync_warpspecialized.hpp
+588
View File
@@ -0,0 +1,588 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/detail/cluster.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/numeric_types.h"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/gemm/gemm.h"
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
#include "cutlass/arch/memory.h"
#include "cute/algorithm/functional.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
// WarpSpecialized Mainloop
// Both DMA Load and MMA methods of this class must be run by a single thread that's picked by elect_one
template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class ElementA_,
class StrideA_,
class ElementB_,
class StrideB_,
class TiledMma_,
class GmemTiledCopyA_,
class SmemLayoutAtomA_,
class SmemCopyAtomA_,
class TransformA_,
class GmemTiledCopyB_,
class SmemLayoutAtomB_,
class SmemCopyAtomB_,
class TransformB_>
struct CollectiveMma<
MainloopSm100UmmaCpAsyncWarpSpecialized<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
TileShape_,
ElementA_,
StrideA_,
ElementB_,
StrideB_,
TiledMma_,
GmemTiledCopyA_,
SmemLayoutAtomA_,
SmemCopyAtomA_,
TransformA_,
GmemTiledCopyB_,
SmemLayoutAtomB_,
SmemCopyAtomB_,
TransformB_>
{
using TiledMma = TiledMma_;
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
// Statically asserting to ensure only 1x1x1 cluster shape & 1sm setup is received
static_assert(size(AtomThrShapeMNK{}) == 1, "Lower alignment SM100 GEMM only supports 1SM MMA");
static_assert(size(ClusterShape{}) == 1, "CPASYNC does not support multicast so the cluster shape is restricted to 1, 1, 1");
using DispatchPolicy = MainloopSm100UmmaCpAsyncWarpSpecialized<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
// TileShape refers to MmaTileShape to adapt for runtime cluster shape
using TileShape = TileShape_;
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
"Static cluster shape used: TileShape should be evenly divided by TiledMma");
// Define A and B block shapes
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
using LoadShapeA_MK = decltype(select<0,2>(TileShape{}));
using LoadShapeB_NK = decltype(select<1,2>(TileShape{}));
// CtaShape_MNK is queried from collective in all kernel layers
using CtaShape_MNK = TileShape;
using ElementA = ElementA_;
using ElementAMma = typename TiledMma::ValTypeA;
using StrideA = StrideA_;
using ElementB = ElementB_;
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = StrideB_;
static constexpr bool IsRuntimeDataTypeA = cute::is_same_v<ElementA, cutlass::type_erased_dynamic_float8_t>;
static constexpr bool IsRuntimeDataTypeB = cute::is_same_v<ElementB, cutlass::type_erased_dynamic_float8_t>;
static_assert(IsRuntimeDataTypeA == IsRuntimeDataTypeB,
"ElementA and ElementB should be both runtime or both static.");
static constexpr bool IsRuntimeDataType = IsRuntimeDataTypeA && IsRuntimeDataTypeB;
using ElementAccumulator = typename TiledMma::ValTypeC;
using GmemTiledCopyA = GmemTiledCopyA_;
using GmemTiledCopyB = GmemTiledCopyB_;
using SmemLayoutAtomA = SmemLayoutAtomA_;
using SmemLayoutAtomB = SmemLayoutAtomB_;
using SmemCopyAtomA = SmemCopyAtomA_;
using SmemCopyAtomB = SmemCopyAtomB_;
using TransformA = TransformA_;
using TransformB = TransformB_;
using ArchTag = typename DispatchPolicy::ArchTag;
using MainloopPipeline = cutlass::PipelineUmmaConsumerAsync<DispatchPolicy::Stages, AtomThrShapeMNK>;
using MainloopPipelineState = typename MainloopPipeline::PipelineState;
static_assert(size(GmemTiledCopyA{}) == size(GmemTiledCopyB{}), "A and B GmemTiledCopy should share the same thread count");
static constexpr int NumLoadThreads = size(GmemTiledCopyA{});
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtomA must be rank 2 (M,K)");
static_assert(((size<0,0>(MmaShapeA_MK{}) * size<1>(MmaShapeA_MK{})) % size<0>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeA_MK{}) * size<2>(MmaShapeA_MK{})) % size<1>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomA>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtomB must be rank 2 (N,K)");
static_assert(((size<0,0>(MmaShapeB_NK{}) * size<1>(MmaShapeB_NK{})) % size<0>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeB_NK{}) * size<2>(MmaShapeB_NK{})) % size<1>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomB>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
// Tile along K mode first before tiling over MN. PIPE mode last as usual.
// (MMA_TILE_M,MMA_TILE_K),MMA_M,MMA_K,PIPE)
using MmaSmemLayoutA = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomA{},
append(MmaShapeA_MK{}, Int<DispatchPolicy::Stages>{}),
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
using LoadSmemLayoutA = decltype(tile_to_shape(
SmemLayoutAtomA{},
append(LoadShapeA_MK{}, Int<DispatchPolicy::Stages>{}),
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
using MmaSmemLayoutB = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomB{},
append(MmaShapeB_NK{}, Int<DispatchPolicy::Stages>{}),
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
using LoadSmemLayoutB = decltype(tile_to_shape(
SmemLayoutAtomB{},
append(LoadShapeB_NK{}, Int<DispatchPolicy::Stages>{}),
conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
static_assert(cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
"MMA atom must source both A and B operand from smem_desc for this mainloop.");
using SmemAllocTypeA = cute::conditional_t<cute::sizeof_bits_v<ElementAMma> < 8, uint8_t, ElementAMma>;
using SmemAllocTypeB = cute::conditional_t<cute::sizeof_bits_v<ElementBMma> < 8, uint8_t, ElementBMma>;
using BitTypeElementA = cute::uint_bit_t<cute::sizeof_bits_v<ElementA>>;
using BitTypeElementB = cute::uint_bit_t<cute::sizeof_bits_v<ElementB>>;
using ArrayElementA = cute::conditional_t<IsRuntimeDataTypeA, BitTypeElementA, ElementA>;
using ArrayElementB = cute::conditional_t<IsRuntimeDataTypeB, BitTypeElementB, ElementB>;
using RuntimeDataTypeA = cute::conditional_t<IsRuntimeDataTypeA, cute::UMMA::MXF8F6F4Format, void*>;
using RuntimeDataTypeB = cute::conditional_t<IsRuntimeDataTypeB, cute::UMMA::MXF8F6F4Format, void*>;
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
cute::array_aligned<SmemAllocTypeA, cute::cosize_v<LoadSmemLayoutA>> smem_A;
cute::array_aligned<SmemAllocTypeB, cute::cosize_v<LoadSmemLayoutB>> smem_B;
} tensors;
using PipelineStorage = typename MainloopPipeline::SharedStorage;
PipelineStorage pipeline;
};
// Expose shared storage for tensors/pipelines separately to allow kernel layer to reorder them.
using TensorStorage = typename SharedStorage::TensorStorage;
using PipelineStorage = typename SharedStorage::PipelineStorage;
// Host side kernel arguments
struct Arguments {
ArrayElementA const* ptr_A{nullptr};
StrideA dA{};
ArrayElementB const* ptr_B{nullptr};
StrideB dB{};
RuntimeDataTypeA runtime_data_type_a{};
RuntimeDataTypeB runtime_data_type_b{};
};
// Device side kernel params
struct Params {
ArrayElementA const* ptr_A{nullptr};
StrideA dA{};
ArrayElementB const* ptr_B{nullptr};
StrideB dB{};
RuntimeDataTypeA runtime_data_type_a;
RuntimeDataTypeB runtime_data_type_b;
};
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(
ProblemShape const& problem_shape,
Arguments const& args,
[[maybe_unused]] void* workspace,
cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) {
// Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
auto ptr_A = recast_ptr<ElementAMma>(args.ptr_A);
auto ptr_B = recast_ptr<ElementBMma>(args.ptr_B);
return {
args.ptr_A,
args.dA,
args.ptr_B,
args.dB,
args.runtime_data_type_a,
args.runtime_data_type_b
};
}
template <class ProblemShape>
static bool
can_implement(
ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool implementable = true;
implementable = implementable && cutlass::detail::check_alignment<GmemTiledCopyA::NumValSrc>(cute::make_shape(M,K,L), StrideA{});
implementable = implementable && cutlass::detail::check_alignment<GmemTiledCopyB::NumValSrc>(cute::make_shape(N,K,L), StrideB{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for CpAsync.\n");
}
return implementable;
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE auto
partition_accumulator_shape() {
auto acc_shape = partition_shape_C(TiledMma{}, take<0,2>(TileShape{})); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N)
return acc_shape;
}
/// Set up the data needed by this collective for load.
/// Return tuple element contain
/// gA_mkl - The tiled tensor for input A
/// gB_nkl - The tiled tensor for input B
/// tAsA - partitioned smem tensor for A
/// tBsB - partitioned smem tensor for B
template <class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(
ProblemShape_MNKL const& problem_shape_MNKL,
Params const& params,
TensorStorage& shared_tensors) const {
using X = Underscore;
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
// Represent the full tensors
Tensor mA_mkl = make_tensor(make_gmem_ptr(params.ptr_A), make_shape(M,K,L), params.dA); //(m,k,l)
Tensor mB_nkl = make_tensor(make_gmem_ptr(params.ptr_B), make_shape(N,K,L), params.dB); //(n,k,l)
// Partition for cpasync
Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
// Build the coordinate tensors with the same shape as input matrices
Tensor cA_mk = make_identity_tensor(make_shape(M,K));
Tensor cB_nk = make_identity_tensor(make_shape(N,K));
// Slice the coordinate tensors in the same way as A/B tensor partitioning
Tensor cgA_mk = local_tile(cA_mk, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k)
Tensor cgB_nk = local_tile(cB_nk, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k)
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), LoadSmemLayoutA{});
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), LoadSmemLayoutB{});
GmemTiledCopyA gmem_to_smem_a_tiled_copy;
GmemTiledCopyB gmem_to_smem_b_tiled_copy;
int thread_idx = threadIdx.x % NumLoadThreads;
auto thr_copy_a = gmem_to_smem_a_tiled_copy.get_slice(thread_idx);
auto thr_copy_b = gmem_to_smem_b_tiled_copy.get_slice(thread_idx);
return cute::make_tuple(
gA_mkl, gB_nkl, // gmem
cgA_mk, cgB_nk, // crd
sA, sB, // smem
problem_shape_MNKL,
gmem_to_smem_a_tiled_copy, gmem_to_smem_b_tiled_copy,
thr_copy_a, thr_copy_b);
}
/// Set up the data needed by this collective for mma compute.
template <class FrgEngine, class FrgLayout>
CUTLASS_DEVICE auto
mma_init(
Params const& params,
[[maybe_unused]] cute::tuple<cute::Tensor<FrgEngine, FrgLayout>, cute::Tensor<FrgEngine, FrgLayout>> const& accumulators_pair,
TensorStorage& shared_tensors) const {
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), MmaSmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), MmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
// Allocate "fragments/descriptors" for A and B matrices
Tensor tCrA = TiledMma::make_fragment_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCrB = TiledMma::make_fragment_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sA)); // PIPE
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sB));
TiledMma tiled_mma;
if constexpr (IsRuntimeDataType) {
// Update instruction descriptor according to runtime argument.
// Applying bitmask (0b111) to help compiler deduce that the conversion and assignment are safe.
tiled_mma.idesc_.a_format_ = uint8_t(params.runtime_data_type_a) & 0b111;
tiled_mma.idesc_.b_format_ = uint8_t(params.runtime_data_type_b) & 0b111;
}
return cute::make_tuple(tiled_mma, tCrA, tCrB);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Producer Perspective
template <
class GTensorA, class GTensorB,
class CTensorA, class CTensorB,
class STensorA, class STensorB,
class ProblemShape_MNKL,
class TiledCopyA, class TiledCopyB,
class ThreadCopyA, class ThreadCopyB,
class TileCoordMNKL,
class KTileIterator
>
CUTLASS_DEVICE auto
load(
Params const& params,
MainloopPipeline mainloop_pipeline,
MainloopPipelineState mainloop_pipe_producer_state,
cute::tuple<GTensorA, GTensorB,
CTensorA, CTensorB,
STensorA, STensorB,
ProblemShape_MNKL,
TiledCopyA, TiledCopyB,
ThreadCopyA, ThreadCopyB> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count) {
// Unpack from load_inputs
GTensorA tAgA_mkl = get<0>(load_inputs);
GTensorB tBgB_nkl = get<1>(load_inputs);
CTensorA cgA_mk = get<2>(load_inputs);
CTensorB cgB_nk = get<3>(load_inputs);
STensorA sA = get<4>(load_inputs);
STensorB sB = get<5>(load_inputs);
ProblemShape_MNKL problem_shape_MNKL = get<6>(load_inputs);
TiledCopyA gmem_to_smem_a_tiled_copy = get<7>(load_inputs);
TiledCopyB gmem_to_smem_b_tiled_copy = get<8>(load_inputs);
ThreadCopyA thr_copy_a = get<9>(load_inputs);
ThreadCopyB thr_copy_b = get<10>(load_inputs);
auto [M,N,K,L] = problem_shape_MNKL;
// Slice out the work coord from partitioned tensors
Tensor gA_in = tAgA_mkl(_, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
Tensor gB_in = tBgB_nkl(_, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
// Repeat slicing out coordinate tensor exactly the same as input tensor does
Tensor cgA_mk_in = cgA_mk(_, _, get<0>(cta_coord_mnkl), _);
Tensor cgB_nk_in = cgB_nk(_, _, get<1>(cta_coord_mnkl), _);
auto k_residue = K - size<1>(gB_in) * size<2>(gA_in);
// Shift tensor so residue_k is at origin (Can't read any k_coord < residue_k)
// This aligns the tensor with BLK_K for all but the 0th k_tile
Tensor gA = domain_offset(make_coord(0, k_residue, 0), gA_in);
Tensor gB = domain_offset(make_coord(0, k_residue, 0), gB_in);
Tensor cA = domain_offset(make_coord(0, k_residue, 0), cgA_mk_in);
Tensor cB = domain_offset(make_coord(0, k_residue, 0), cgB_nk_in);
auto tAgA = thr_copy_a.partition_S(gA);
auto tAsA = thr_copy_a.partition_D(sA);
auto tBgB = thr_copy_b.partition_S(gB);
auto tBsB = thr_copy_b.partition_D(sB);
// Allocate predicate tensors for m and n
Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{});
Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{});
Tensor tAcA = thr_copy_a.partition_S(cA);
Tensor tBcB = thr_copy_b.partition_S(cB);
// Copy gmem to smem for *k_tile_iter, predicating for k residue
Tensor tAgAk = tAgA(_,_,_,*k_tile_iter);
Tensor tBgBk = tBgB(_,_,_,*k_tile_iter);
// Repeating on predicators with the same operations on tAgA and tBgB
Tensor tAcAk = tAcA(_,_,_,*k_tile_iter);
Tensor tBcBk = tBcB(_,_,_,*k_tile_iter);
// Set predicates for m bounds
CUTLASS_PRAGMA_UNROLL
for (int m = 0; m < size<0>(tApA); ++m) {
tApA(m,0) = elem_less(get<0>(tAcAk(0,m,0)), M); // blk_m coord < M
}
// Set predicates for n bounds
CUTLASS_PRAGMA_UNROLL
for (int n = 0; n < size<0>(tBpB); ++n) {
tBpB(n,0) = elem_less(get<0>(tBcBk(0,n,0)), N); // blk_n coord < N
}
// 0-th stage with predication on k to account for residue
// For performance consideration,
// this predicated block for K-tail is only activated when there is k-residue
if (k_residue != 0 && k_tile_count > 0) {
// LOCK mainloop_pipe_producer_state for _writing_
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state);
int write_stage = mainloop_pipe_producer_state.index();
CUTLASS_PRAGMA_UNROLL
for (int k = 0; k < size<2>(tAsA); ++k) {
if ( int(get<1>(tAcAk(0,0,k))) >= 0) { // blk_k coord < K
copy_if(gmem_to_smem_a_tiled_copy, tApA(_,k), tAgAk(_,_,k), tAsA(_,_,k,write_stage));
}
else {
clear(tAsA(_,_,k,write_stage));
}
}
CUTLASS_PRAGMA_UNROLL
for (int k = 0; k < size<2>(tBsB); ++k) {
if (int(get<1>(tBcBk(0,0,k))) >= 0) { // blk_k coord < K
copy_if(gmem_to_smem_b_tiled_copy, tBpB(_,k), tBgBk(_,_,k), tBsB(_,_,k,write_stage));
}
else {
clear(tBsB(_,_,k,write_stage));
}
}
++k_tile_iter;
--k_tile_count;
// UNLOCK mainloop_pipe_producer_state
mainloop_pipeline.producer_commit(mainloop_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive);
// Advance mainloop_pipe_producer_state
++mainloop_pipe_producer_state;
}
auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
// Issue the Mainloop loads
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
auto mainloop_pipe_producer_state_curr = mainloop_pipe_producer_state;
++mainloop_pipe_producer_state;
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state_curr, barrier_token);
barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
int write_stage = mainloop_pipe_producer_state_curr.index();
copy_if(gmem_to_smem_a_tiled_copy, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
copy_if(gmem_to_smem_b_tiled_copy, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
mainloop_pipeline.producer_commit(mainloop_pipe_producer_state_curr, cutlass::arch::cpasync_barrier_arrive);
--k_tile_count;
++k_tile_iter;
}
return cute::make_tuple(mainloop_pipe_producer_state, k_tile_iter);
}
/// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster
CUTLASS_DEVICE void
load_tail(MainloopPipeline mainloop_pipeline, MainloopPipelineState mainloop_pipe_producer_state) {
// Issue the epilogue waits
// This helps avoid early exit of ctas in Cluster
// Waits for all stages to either be released (all
// Consumer UNLOCKs), or if the stage was never used
// then would just be acquired since the phase was
// still inverted from make_producer_start_state
mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgEngine, class FrgLayout,
class FragmentA, class FragmentB
>
CUTLASS_DEVICE auto
mma(MainloopPipeline mainloop_pipeline,
MainloopPipelineState mainloop_pipe_consumer_state,
cute::tuple<cute::Tensor<FrgEngine, FrgLayout>, cute::Tensor<FrgEngine, FrgLayout>> const& accumulators_pair,
cute::tuple<TiledMma, FragmentA, FragmentB> const& mma_inputs,
int k_tile_count
) {
static_assert(is_tmem<FrgEngine>::value, "Accumulator must be tmem resident.");
static_assert(rank(FrgLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA, MMA_M, MMA_N)");
auto accumulators = get<0>(accumulators_pair);
auto [tiled_mma, tCrA, tCrB] = mma_inputs;
//
// PIPELINED MAIN LOOP
//
tiled_mma.accumulate_ = UMMA::ScaleOut::Zero;
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state);
int read_stage = mainloop_pipe_consumer_state.index();
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulators);
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
}
mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state);
--k_tile_count;
++mainloop_pipe_consumer_state;
}
return mainloop_pipe_consumer_state;
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp
+56 -51
View File
@@ -347,20 +347,20 @@ struct CollectiveMma<
template<
class KTileCount,
class GTensorPartitionedScaleA, class GTensorPartitionedScaleB,
class IdentTensorPartitionedScaleA, class IdentTensorPartitionedScaleB,
class GTensorScaleA, class GTensorScaleB,
class IdentTensorScaleA, class IdentTensorScaleB,
class STensorScaleA, class STensorScaleB
>
struct LoadSFParams {
// for scheduler
KTileCount k_tiles;
GTensorPartitionedScaleA tSFAgSFA_mkl;
GTensorPartitionedScaleB tSFBgSFB_nkl;
IdentTensorPartitionedScaleA tSFAIdentSFA_mkl;
IdentTensorPartitionedScaleB tSFBIdentSFB_nkl;
STensorScaleA tSFAsSFA;
STensorScaleB tSFBsSFB;
GTensorScaleA gSFA_mkl;
GTensorScaleB gSFB_nkl;
IdentTensorScaleA identSFA_mkl;
IdentTensorScaleB identSFB_nkl;
STensorScaleA sSFA;
STensorScaleB sSFB;
LayoutSFA layout_SFA;
LayoutSFB layout_SFB;
@@ -368,14 +368,14 @@ struct CollectiveMma<
CUTLASS_DEVICE
LoadSFParams (
KTileCount k_tiles_,
GTensorPartitionedScaleA tSFAgSFA_mkl_, GTensorPartitionedScaleB tSFBgSFB_nkl_,
IdentTensorPartitionedScaleA tSFAIdentSFA_mkl_, IdentTensorPartitionedScaleB tSFBIdentSFB_nkl_,
STensorScaleA tSFAsSFA_, STensorScaleB tSFBsSFB_,
GTensorScaleA gSFA_mkl_, GTensorScaleB gSFB_nkl_,
IdentTensorScaleA identSFA_mkl_, IdentTensorScaleB identSFB_nkl_,
STensorScaleA sSFA_, STensorScaleB sSFB_,
LayoutSFA layout_SFA_, LayoutSFB layout_SFB_)
: k_tiles(k_tiles_)
, tSFAgSFA_mkl(tSFAgSFA_mkl_), tSFBgSFB_nkl(tSFBgSFB_nkl_)
, tSFAIdentSFA_mkl(tSFAIdentSFA_mkl_), tSFBIdentSFB_nkl(tSFBIdentSFB_nkl_)
, tSFAsSFA(tSFAsSFA_), tSFBsSFB(tSFBsSFB_)
, gSFA_mkl(gSFA_mkl_), gSFB_nkl(gSFB_nkl_)
, identSFA_mkl(identSFA_mkl_), identSFB_nkl(identSFB_nkl_)
, sSFA(sSFA_), sSFB(sSFB_)
, layout_SFA(layout_SFA_), layout_SFB(layout_SFB_) {}
};
@@ -732,35 +732,16 @@ struct CollectiveMma<
static_assert(rank(decltype(gSFA_mkl){}) == 5);
static_assert(rank(decltype(gSFB_nkl){}) == 5);
// 1 thread copies entire set of scalar
GmemTiledCopySFA scale_copy_a{};
GmemTiledCopySFB scale_copy_b{};
ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(threadIdx.x % size(scale_copy_a));
ThrCopy thr_scale_copy_b = scale_copy_b.get_slice(threadIdx.x % size(scale_copy_b));
Tensor sSFA = make_tensor(make_smem_ptr(shared_tensors.smem_SFA.begin()),
SmemLayoutScaleA{}); // (CTA_M,CTA_K,P)
Tensor sSFB = make_tensor(make_smem_ptr(shared_tensors.smem_SFB.begin()),
SmemLayoutScaleB{}); // (CTA_M,CTA_K,P)
Tensor tSFAgSFA_mkl = thr_scale_copy_a.partition_S(gSFA_mkl); // (CPY, BLK_M, BLK_K, m, k, l)
Tensor tSFAIdentSFA_mkl = thr_scale_copy_a.partition_S(identSFA_mkl); // (CPY, BLK_M, BLK_K, m, k, l)
Tensor tSFAsSFA = thr_scale_copy_a.partition_D(sSFA);
Tensor tSFBgSFB_nkl = thr_scale_copy_b.partition_S(gSFB_nkl); // (CPY, BLK_N, BLK_K, m, k, l)
Tensor tSFBIdentSFB_nkl = thr_scale_copy_b.partition_S(identSFB_nkl); // (CPY, BLK_N, BLK_K, m, k, l)
Tensor tSFBsSFB = thr_scale_copy_b.partition_D(sSFB);
static_assert(rank(decltype(tSFAgSFA_mkl){}) == 6);
static_assert(rank(decltype(tSFBgSFB_nkl){}) == 6);
LoadSFParams load_params {
size<3>(gSFA_mkl),
tSFAgSFA_mkl, tSFBgSFB_nkl, // for input scale tensor values
tSFAIdentSFA_mkl, tSFBIdentSFB_nkl, // for predicating scale tensor copies
tSFAsSFA, tSFBsSFB, // for scale tensor values
gSFA_mkl, gSFB_nkl, // for input scale tensor values
identSFA_mkl, identSFB_nkl, // for predicating scale tensor copies
sSFA, sSFB, // for scale tensor values
mainloop_params.layout_SFA, // for predicating scale tensor copies
mainloop_params.layout_SFB // for predicating scale tensor copies
};
@@ -922,24 +903,44 @@ struct CollectiveMma<
KTileIterator k_tile_iter, int k_tile_count) {
auto [unused_k_tiles,
tSFAgSFA_mkl, tSFBgSFB_nkl,
tSFAIdentSFA_mkl, tSFBIdentSFB_nkl,
tSFAsSFA, tSFBsSFB,
gSFA_mkl, gSFB_nkl,
identSFA_mkl, identSFB_nkl,
sSFA, sSFB,
layout_SFA, layout_SFB] = load_inputs;
// slice out the work coord from partitioned tensors
GmemTiledCopySFA scale_copy_a{};
GmemTiledCopySFB scale_copy_b{};
Tensor tSFAgSFA = tSFAgSFA_mkl(_, _, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
Tensor gSFA_k_compact = filter_zeros(
gSFA_mkl(_, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl))); // (BLK_M_CPT, BLK_K_CPT, k_cpt)
Tensor gSFB_k_compact = filter_zeros(
gSFB_nkl(_, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl))); // (BLK_N_CPT, BLK_K_CPT, k_cpt)
Tensor tSFBgSFB = tSFBgSFB_nkl(_, _, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
Tensor identSFA_k_compact = filter_zeros(
identSFA_mkl(_, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)),
gSFA_k_compact.stride()); // (BLK_M_CPT, BLK_K_CPT, k_cpt)
Tensor identSFB_k_compact = filter_zeros(
identSFB_nkl(_, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)),
gSFB_k_compact.stride()); // (BLK_N_CPT, BLK_K_CPT, k_cpt)
Tensor thr_tile_SFA_k = tSFAIdentSFA_mkl(_0{}, _, _, get<0>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
Tensor thr_tile_pSFA = make_tensor<bool>(shape(filter_zeros(thr_tile_SFA_k(_,_,_0{}), tSFAgSFA(_0{},_,_,_0{}).stride())));
Tensor thr_tile_SFB_k = tSFBIdentSFB_nkl(_0{}, _, _, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
Tensor sSFA_compact = filter_zeros(sSFA); // (BLK_M_CPT, BLK_K_CPT, P)
Tensor sSFB_compact = filter_zeros(sSFB); // (BLK_N_CPT, BLK_K_CPT, P)
Tensor thr_tile_pSFB = make_tensor<bool>(shape(filter_zeros(thr_tile_SFB_k(_,_,_0{}), tSFBgSFB(_0{},_,_,_0{}).stride())));
ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(threadIdx.x % size(scale_copy_a));
ThrCopy thr_scale_copy_b = scale_copy_b.get_slice(threadIdx.x % size(scale_copy_b));
Tensor tSFAgSFA_k_compact = thr_scale_copy_a.partition_S(gSFA_k_compact); // (CPY, BLK_M, BLK_K, k)
Tensor tSFAIdentSFA_k_compact = thr_scale_copy_a.partition_S(identSFA_k_compact); // (CPY, BLK_M, BLK_K, k)
Tensor tSFAsSFA_compact = thr_scale_copy_a.partition_D(sSFA_compact);
Tensor tSFBgSFB_k_compact = thr_scale_copy_b.partition_S(gSFB_k_compact); // (CPY, BLK_N, BLK_K, k)
Tensor tSFBIdentSFB_k_compact = thr_scale_copy_b.partition_S(identSFB_k_compact); // (CPY, BLK_N, BLK_K, k)
Tensor tSFBsSFB_compact = thr_scale_copy_b.partition_D(sSFB_compact);
Tensor thr_tile_pSFA = make_fragment_like<bool>(tSFAgSFA_k_compact(_0{},_,_,_0{}));
Tensor thr_tile_pSFB = make_fragment_like<bool>(tSFBgSFB_k_compact(_0{},_,_,_0{}));
// Issue the loads
CUTLASS_PRAGMA_NO_UNROLL
@@ -949,18 +950,22 @@ struct CollectiveMma<
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(thr_tile_pSFA); ++i) {
Tensor thr_tile_SFA = filter_zeros(thr_tile_SFA_k(_,_,*k_tile_iter), tSFAgSFA(_0{},_,_,_0{}).stride());
thr_tile_pSFA(i) = elem_less(thr_tile_SFA(i), shape(filter_zeros(layout_SFA))) && threadIdx.x % 32 < size(scale_copy_a);
Tensor tSFAIdentSFA_compact = tSFAIdentSFA_k_compact(_0{},_,_,*k_tile_iter);
thr_tile_pSFA(i) = elem_less(tSFAIdentSFA_compact(i),
shape(filter_zeros(layout_SFA))) && threadIdx.x % 32 < size(scale_copy_a);
}
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(thr_tile_pSFB); ++i) {
Tensor thr_tile_SFB = filter_zeros(thr_tile_SFB_k(_,_,*k_tile_iter), tSFBgSFB(_0{},_,_,_0{}).stride());
thr_tile_pSFB(i) = elem_less(thr_tile_SFB(i), shape(filter_zeros(layout_SFB))) && threadIdx.x % 32 < size(scale_copy_b);
Tensor tSFBIdentSFB_compact = tSFBIdentSFB_k_compact(_0{},_,_,*k_tile_iter);
thr_tile_pSFB(i) = elem_less(tSFBIdentSFB_compact(i),
shape(filter_zeros(layout_SFB))) && threadIdx.x % 32 < size(scale_copy_b);
}
copy_if(scale_copy_a, thr_tile_pSFA, filter_zeros(tSFAgSFA(_,_,_,*k_tile_iter)), filter_zeros(tSFAsSFA(_,_,_,mainloop_sf_pipe_producer_state.index())));
copy_if(scale_copy_b, thr_tile_pSFB, filter_zeros(tSFBgSFB(_,_,_,*k_tile_iter)), filter_zeros(tSFBsSFB(_,_,_,mainloop_sf_pipe_producer_state.index())));
copy_if(scale_copy_a, thr_tile_pSFA, tSFAgSFA_k_compact(_,_,_,*k_tile_iter),
tSFAsSFA_compact(_,_,_,mainloop_sf_pipe_producer_state.index()));
copy_if(scale_copy_b, thr_tile_pSFB, tSFBgSFB_k_compact(_,_,_,*k_tile_iter),
tSFBsSFB_compact(_,_,_,mainloop_sf_pipe_producer_state.index()));
mainloop_sf_pipeline.producer_commit(mainloop_sf_pipe_producer_state, cutlass::arch::cpasync_barrier_arrive_noinc);
__syncwarp();
include/cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp
+1296
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/collective/sm103_blockscaled_mma_array_warpspecialized.hpp
+1685
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/collective/sm103_blockscaled_mma_warpspecialized.hpp
+1276
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+9 -1
View File
@@ -153,7 +153,15 @@ struct CollectiveMma<
static_assert((size<0>(TileShape{}) % ScaleGranularityM) == 0, "FP8 scaling granularity must evenly divide tile shape along M.");
static_assert((size<1>(TileShape{}) % ScaleGranularityN) == 0, "FP8 scaling granularity must evenly divide tile shape along N.");
using ScaleConfig = ::cutlass::detail::Sm90BlockwiseScaleConfig<ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>;
static constexpr bool MMajorSFA = size<0,1>(InternalLayoutSFA{}.stride()) == 1;
static constexpr bool NMajorSFB = size<0,1>(InternalLayoutSFB{}.stride()) == 1;
using ScaleConfig = ::cutlass::detail::Sm90BlockwiseScaleConfig<
ScaleGranularityM,
ScaleGranularityN,
ScaleGranularityK,
MMajorSFA ? cute::GMMA::Major::MN : cute::GMMA::Major::K,
NMajorSFB ? cute::GMMA::Major::MN : cute::GMMA::Major::K>;
using SmemLayoutAtomSFA = decltype(ScaleConfig::smem_atom_layoutSFA(TileShape{}));
using SmemLayoutAtomSFB = decltype(ScaleConfig::smem_atom_layoutSFB(TileShape{}));
include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+15 -7
View File
@@ -134,9 +134,12 @@ struct CollectiveMma<
static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
static constexpr int ScaleNsPerTile = size<1>(TileShape{}) / ScaleGranularityN;
static constexpr bool MMajorSFA = size<0,1>(LayoutSFA{}.stride()) == 1;
static constexpr bool NMajorSFB = size<0,1>(LayoutSFB{}.stride()) == 1;
static constexpr int ScaleTmaThreshold = 32;
static constexpr bool IsTmaLoadSFA = ScaleMsPerTile >= ScaleTmaThreshold && ScaleNsPerTile < ScaleTmaThreshold;
static constexpr bool IsTmaLoadSFB = ScaleNsPerTile >= ScaleTmaThreshold && ScaleMsPerTile < ScaleTmaThreshold;
static constexpr bool IsTmaLoadSFA = ScaleMsPerTile >= ScaleTmaThreshold && ScaleNsPerTile < ScaleTmaThreshold && MMajorSFA;
static constexpr bool IsTmaLoadSFB = ScaleNsPerTile >= ScaleTmaThreshold && ScaleMsPerTile < ScaleTmaThreshold && NMajorSFB;
// Two threads per CTA are producers (1 for operand tile `tma`, and 32 for scales `cp.async`)
static constexpr int NumProducerThreadEvents = ((IsTmaLoadSFA && IsTmaLoadSFB)? 1 : 33);
@@ -151,7 +154,12 @@ struct CollectiveMma<
static_assert((size<0>(TileShape{}) % ScaleGranularityM) == 0, "FP8 scaling granularity must evenly divide tile shape along M.");
static_assert((size<1>(TileShape{}) % ScaleGranularityN) == 0, "FP8 scaling granularity must evenly divide tile shape along N.");
using ScaleConfig = ::cutlass::detail::Sm90BlockwiseScaleConfig<ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>;
using ScaleConfig = ::cutlass::detail::Sm90BlockwiseScaleConfig<
ScaleGranularityM,
ScaleGranularityN,
ScaleGranularityK,
MMajorSFA ? cute::GMMA::Major::MN : cute::GMMA::Major::K,
NMajorSFB ? cute::GMMA::Major::MN : cute::GMMA::Major::K>;
using SmemLayoutAtomSFA = decltype(ScaleConfig::smem_atom_layoutSFA(TileShape{}));
using SmemLayoutAtomSFB = decltype(ScaleConfig::smem_atom_layoutSFB(TileShape{}));
@@ -170,8 +178,8 @@ struct CollectiveMma<
using CopyAtomSFA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
using CopyAtomSFB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
static constexpr int AlignmentSFA = 1;
static constexpr int AlignmentSFB = 1;
static constexpr int AlignmentSFA = IsTmaLoadSFA ? 128 / cutlass::sizeof_bits<ElementBlockScale>::value : 1;
static constexpr int AlignmentSFB = IsTmaLoadSFB ? 128 / cutlass::sizeof_bits<ElementBlockScale>::value : 1;
// Block scaling smem layout
using SmemLayoutSFA = decltype(make_layout(
@@ -669,7 +677,7 @@ struct CollectiveMma<
Tensor tSFAcSFA_compact = filter_zeros(tSFAcSFA);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tSFApSFA); ++i) {
tSFApSFA(i) = load_sfa && elem_less(get<0>(tSFAcSFA_compact(i)), get<0>(SFA_shape));
tSFApSFA(i) = load_sfa && elem_less(tSFAcSFA_compact(i), SFA_shape);
}
bool load_sfb = thread_idx < ScaleNsPerTile;
@@ -677,7 +685,7 @@ struct CollectiveMma<
Tensor tSFBcSFB_compact = filter_zeros(tSFBcSFB);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tSFBpSFB); ++i) {
tSFBpSFB(i) = load_sfb && elem_less(get<0>(tSFBcSFB_compact(i)), get<0>(SFB_shape));
tSFBpSFB(i) = load_sfb && elem_less(tSFBcSFB_compact(i), SFB_shape);
}
int write_stage = smem_pipe_write.index();
// Copy scale tensors from global memory to shared memory
include/cutlass/gemm/device/gemm_universal_adapter.h
+2
View File
@@ -393,6 +393,7 @@ public:
[[maybe_unused]] dim3 fallback_cluster = dim3{0,0,0};
if constexpr (GemmKernel::ArchTag::kMinComputeCapability == 100
|| GemmKernel::ArchTag::kMinComputeCapability == 101
|| GemmKernel::ArchTag::kMinComputeCapability == 103
) {
if constexpr (!cute::is_static_v<typename GemmKernel::DispatchPolicy::ClusterShape>) {
fallback_cluster = params.hw_info.cluster_shape_fallback;
@@ -473,6 +474,7 @@ public:
if constexpr (GemmKernel::ArchTag::kMinComputeCapability == 100
|| GemmKernel::ArchTag::kMinComputeCapability == 101
|| GemmKernel::ArchTag::kMinComputeCapability == 120
|| GemmKernel::ArchTag::kMinComputeCapability == 103
) {
if constexpr (is_static_1x1x1) {
#if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
include/cutlass/gemm/device/gemm_universal_base.h
+2 -3
View File
@@ -34,14 +34,13 @@
*/
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/limits>
#include CUDA_STD_HEADER(limits)
#else
#include <limits>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/arch/arch.h"
#include "cutlass/device_kernel.h"
include/cutlass/gemm/device/gemv_blockscaled.h
+183
View File
@@ -0,0 +1,183 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/arch/arch.h"
#include "cutlass/device_kernel.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
#include "cutlass/gemm/kernel/gemm_universal.h"
#include "cutlass/gemm/kernel/default_gemm_universal.h"
#include "cutlass/gemm/device/default_gemm_configuration.h"
#include "cutlass/gemm/device/gemm_universal_base.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace device {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename GemvKernel_>
class GemvBlockScaled {
public:
using GemvKernel = GemvKernel_;
using ElementA = typename GemvKernel::ElementA;
using LayoutA = typename GemvKernel::LayoutA;
using ElementB = typename GemvKernel::ElementB;
using ElementC = typename GemvKernel::ElementC;
using ElementSFA = typename GemvKernel::ElementSFA;
using ElementSFB = typename GemvKernel::ElementSFB;
using ElementAccumulator = typename GemvKernel::ElementAccumulator;
using EpilogueOutputOp = typename GemvKernel::EpilogueOutputOp;
static ComplexTransform const kTransformA = GemvKernel::kTransformA;
static ComplexTransform const kTransformB = GemvKernel::kTransformB;
static int const kThreadCount = GemvKernel::kThreadCount;
static int const kThreadsPerRow = GemvKernel::kThreadsPerRow;
using Arguments = typename GemvKernel::Arguments;
using Params = typename GemvKernel::Params;
private:
Params params_;
public:
/// Constructs the GemvBlockScaled.
GemvBlockScaled() = default;
/// Determines whether the GemvBlockScaled can execute the given problem.
static Status can_implement(Arguments const &args) {
return GemvKernel::can_implement(args);
}
/// Gets the workspace size
static size_t get_workspace_size(Arguments const &args) {
return 0;
}
/// Computes the grid shape
static dim3 get_grid_shape(Arguments const &args, dim3 const &block) {
if(platform::is_same<LayoutA, layout::ColumnMajor>::value) {
return dim3((args.problem_size.row() + (block.x - 1)) / block.x, 1, args.batch_count % 65536);
}
else {
return dim3((args.problem_size.row() + (block.y - 1)) / block.y, 1, args.batch_count % 65536);
}
}
/// Computes the block shape
static dim3 get_block_shape() {
if(platform::is_same<LayoutA, layout::ColumnMajor>::value) {
return dim3(kThreadCount, 1, 1);
}
else {
return dim3(kThreadsPerRow, kThreadCount / kThreadsPerRow, 1);
}
}
/// Initializes GemvBlockScaled state from arguments.
Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
params_ = Params(args);
return Status::kSuccess;
}
/// Lightweight update given a subset of arguments
Status update(Arguments const &args, void *workspace = nullptr) {
return params_.update(args);
}
/// Runs the kernel using initialized state.
Status run(cudaStream_t stream = nullptr) {
const dim3 block = get_block_shape();
const dim3 grid = get_grid_shape(params_, block);
int smem_size = int(sizeof(typename GemvKernel::SharedStorage));
cutlass::arch::synclog_setup();
cutlass::Kernel<GemvKernel><<<grid, block, smem_size, stream>>>(params_);
cudaError_t result = cudaGetLastError();
if (result == cudaSuccess) {
return Status::kSuccess;
} else {
return Status::kErrorInternal;
}
}
/// Runs the kernel using initialized state.
Status operator()(cudaStream_t stream = nullptr) {
return run(stream);
}
/// Runs the kernel using initialized state.
Status operator()(
Arguments const &args,
void *workspace = nullptr,
cudaStream_t stream = nullptr) {
Status status = initialize(args, workspace, stream);
if (status == Status::kSuccess) {
status = run(stream);
}
return status;
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace device
} // namespace gemm
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/dispatch_policy.hpp
+173
View File
@@ -36,6 +36,7 @@
#include "cute/layout.hpp"
#include "cute/numeric/integral_constant.hpp" // cute::false_type
#include "cute/atom/copy_traits_sm100.hpp"
#include "cutlass/detail/collective/sm103_kernel_type.hpp"
//////////////////////////////////////////////////////////////////////////////
namespace cutlass::detail {
@@ -423,6 +424,21 @@ struct MainloopSm90ArrayTmaGmmaWarpSpecializedBlockScaling
"KernelSchedule must be one of the warp specialized FP8 block scale policies");
};
//////////////////////////////////////////////////////////////////////////////
//
// Kernel Scheduler Tag
//
// Dense GEMM: SM100 tensor op policy that applies to both 1SM and 2SM MMA atoms
template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_
>
struct KernelWarpSpecializedSm100 final {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
};
template<
int SchedulerPipelineStageCount_,
@@ -461,6 +477,24 @@ struct KernelPtrArrayTmaWarpSpecializedMmaTransformSm100 final {
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
};
template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_
>
struct KernelTmaWarpSpecializedBlockScaledSm103 final {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
};
template<
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_
>
struct KernelPtrArrayTmaWarpSpecializedBlockScaledSm103 final {
static constexpr int SchedulerPipelineStageCount = SchedulerPipelineStageCount_;
static constexpr int AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_;
};
// Sparse Gemm
template<
int SchedulerPipelineStageCount_,
@@ -665,6 +699,8 @@ struct KernelScheduleSm100DenseGemm : KernelScheduleSm100 {}; // Base policy
// Dense GEMM: Specialize for 1SM vs 2SM
struct KernelTmaWarpSpecialized1SmSm100 final : KernelSchedule1Sm, KernelScheduleSm100DenseGemm {}; // Use for 1SM Dense GEMM Kernels for Collective Mainloop Builder
struct KernelTmaWarpSpecialized2SmSm100 final : KernelSchedule2Sm, KernelScheduleSm100DenseGemm {}; // Use for 2SM Dense GEMM Kernels for Collective Mainloop Builder
struct KernelWarpSpecialized1SmSm100 final : KernelSchedule1Sm, KernelScheduleSm100DenseGemm {}; // Use for 1SM Dense GEMM Kernels for Collective Mainloop Builder Without TMA
///////////////////////////////////////////////////////////////////////////////////////////////////////
// SM100 Ptr-Array Dense GEMM Dispatch Policies
///////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -721,6 +757,8 @@ struct KernelScheduleSm100MixedInputGemm : KernelScheduleSm100 {};
struct KernelTmaWarpSpecializedMixedInputSmemSm100 : KernelScheduleSm100MixedInputGemm { };
struct KernelTmaWarpSpecialized1SmMixedInputSm100 final : KernelSchedule1Sm, KernelScheduleSm100MixedInputGemm { };
struct KernelTmaWarpSpecialized1SmMixedInputSmemSm100 final : KernelSchedule1Sm, KernelTmaWarpSpecializedMixedInputSmemSm100 { };
struct KernelTmaWarpSpecialized2SmMixedInputSm100 final : KernelSchedule2Sm, KernelScheduleSm100MixedInputGemm { };
struct KernelTmaWarpSpecialized2SmMixedInputSmemSm100 final : KernelSchedule2Sm, KernelTmaWarpSpecializedMixedInputSmemSm100 { };
///////////////////////////////////////////////////////////////////////////////////////////////////////
// SM100 Ptr-Array FastF32 (9xBF16) GEMM Dispatch Policies
@@ -789,6 +827,54 @@ struct KernelSparseTmaWarpSpecialized2SmNvf4Sm100 final : KernelSchedule2
struct KernelSparseTmaWarpSpecialized1SmMxf4Sm100 final : KernelSchedule1Sm, KernelScheduleSparseMxNvf4Sm100 { };
struct KernelSparseTmaWarpSpecialized2SmMxf4Sm100 final : KernelSchedule2Sm, KernelScheduleSparseMxNvf4Sm100 { };
///////////////////////////////////////////////////////////////////////////////////////////////////////
//
// SM103 Dispatch Policies
//
///////////////////////////////////////////////////////////////////////////////////////////////////////
struct KernelScheduleSm103 {};
struct KernelScheduleSm103BlockScaledGemm : KernelScheduleSm103 {};
struct KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch : KernelScheduleSm103BlockScaledGemm {};
struct KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch : KernelScheduleSm103BlockScaledGemm {};
// Blockscaled Gemm: Specialized for instruction type, scale factor vector size, and 1SM vs. 2SM
// These are the public dispatch policy name
struct KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch final : KernelSchedule1Sm, KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch { };
struct KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch final : KernelSchedule2Sm, KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch { };
struct KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch final : KernelSchedule1Sm, KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch { };
struct KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch final : KernelSchedule2Sm, KernelScheduleSm103BlockScaledMxNvf4UltraTmaPrefetch { };
struct KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch final : KernelSchedule1Sm, KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch { };
struct KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch final : KernelSchedule2Sm, KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch { };
struct KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch final : KernelSchedule1Sm, KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch { };
struct KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch final : KernelSchedule2Sm, KernelScheduleSm103BlockScaledMxNvf4UltraDisablePrefetch { };
using KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103 = KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch;
using KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103 = KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch;
using KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103 = KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch;
using KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103 = KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch;
struct KernelSchedulePtrArraySm103BlockScaledGemm : KernelScheduleSm103 {};
struct KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch : KernelSchedulePtrArraySm103BlockScaledGemm {};
struct KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch : KernelSchedulePtrArraySm103BlockScaledGemm {};
struct KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch final : KernelSchedule1Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch { };
struct KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103TmaPrefetch final : KernelSchedule2Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch { };
struct KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch final : KernelSchedule1Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch { };
struct KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103TmaPrefetch final : KernelSchedule2Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraTmaPrefetch { };
struct KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch final : KernelSchedule1Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch { };
struct KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch final : KernelSchedule2Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch { };
struct KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch final : KernelSchedule1Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch { };
struct KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch final : KernelSchedule2Sm, KernelSchedulePtrArraySm103BlockScaledMxNvf4UltraDisablePrefetch { };
using KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103 = KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch;
using KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103 = KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch;
using KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103 = KernelPtrArrayTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch;
using KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103 = KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch;
///////////////////////////////////////////////////////////////////////////////////////////////////////
//
// SM120 Dispatch Policies
@@ -844,6 +930,25 @@ struct KernelSparseTmaWarpSpecializedMxf4Sm120 final : KernelScheduleS
struct KernelSparseTmaWarpSpecializedMxf8f6f4Sm120 final : KernelScheduleSparseMxf8f6f4Sm120 { };
struct KernelSparseTmaWarpSpecializedMxf8f6f4Acc2x4Sm120 final : KernelScheduleSparseMxf8f6f4Sm120, KernelScheduleAcc2x4Sm120 { };
//////////////////////////////////////////////////////////////////////////////
//
// Collective Mainloop Dispatch Policies
//
// n-buffer in smem, pipelined with Blackwell UMMA and CPASYNC, Warp specialized dynamic schedule
template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
>
struct MainloopSm100UmmaCpAsyncWarpSpecialized {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using Schedule = KernelWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
@@ -983,6 +1088,34 @@ struct MainloopSm100TmaUmmaWarpSpecializedFastF32 {
};
// n-buffer in smem, pipelined with Blackwell Mixed Input kernel with UMMA (HwScaled) and TMA,
template<
// Number of Pipeline stages for
// MainloopLoad <-> Conversion <-> MainLoad
int Load2TransformPipelineStageCount_,
// Number of Pipeline stages for
// MainloopLoad <-> Conversion <-> MainLoad
int Transform2MmaPipelineStageCount_,
// TileScheduler pipeline depth
int SchedulerPipelineStageCount_,
// Accmulator pipeline depth
int AccumulatorPipelineStageCount_,
// ClusterShape for the kernel
class ClusterShape_ = Shape<_1,_1,_1>
>
struct MainloopSm100TmaUmmaWarpSpecializedMixedInput {
constexpr static int Load2TransformPipelineStageCount = Load2TransformPipelineStageCount_;
constexpr static int Load2MmaPipelineStageCount = Load2TransformPipelineStageCount_;
constexpr static int Transform2MmaPipelineStageCount = Transform2MmaPipelineStageCount_;
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::MixedInput;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using Schedule = KernelTmaWarpSpecializedMixedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
constexpr static int Stages = Load2TransformPipelineStageCount;
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
@@ -1064,9 +1197,49 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 {
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int LoadABPipelineStageCount_,
int LoadSFPipelineStageCount_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch
>
struct MainloopSm103TmaUmmaWarpSpecializedBlockScaled {
constexpr static int LoadABPipelineStageCount = LoadABPipelineStageCount_;
constexpr static int LoadSFPipelineStageCount = LoadSFPipelineStageCount_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm103;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelTmaWarpSpecializedBlockScaledSm103<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
constexpr static int Stages = LoadABPipelineStageCount;
constexpr static cutlass::sm103::detail::KernelPrefetchType PrefetchType = PrefetchType_;
};
// Mainloop schedule for array-based TMA
template<
int LoadABPipelineStageCount_,
int LoadSFPipelineStageCount_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
cutlass::sm103::detail::KernelPrefetchType PrefetchType_ = cutlass::sm103::detail::KernelPrefetchType::TmaPrefetch
>
struct MainloopSm103ArrayTmaUmmaWarpSpecializedBlockScaled {
constexpr static int LoadABPipelineStageCount = LoadABPipelineStageCount_;
constexpr static int LoadSFPipelineStageCount = LoadSFPipelineStageCount_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm103;
constexpr static bool IsOverlappingAccum = AccumulatorPipelineStageCount_ == 1;
using Schedule = KernelPtrArrayTmaWarpSpecializedBlockScaledSm103<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
constexpr static int Stages = LoadABPipelineStageCount;
constexpr static cutlass::sm103::detail::KernelPrefetchType PrefetchType = PrefetchType_;
};
template<
int Stages_,
int SchedulerPipelineStageCount_,
include/cutlass/gemm/kernel/gemm_universal.hpp
+5
View File
@@ -68,8 +68,13 @@ struct IsCutlass3ArrayKernel<ProblemShape, cute::void_t<typename ProblemShape::U
#include "cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mma_transform.hpp"
#include "cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp"
#include "cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mixed_input_transform.hpp"
#include "cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp"
#include "cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp"
#include "cutlass/gemm/kernel/sm100_sparse_gemm_tma_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm100_gemm_cpasync_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp"
#include "cutlass/gemm/kernel/sm120_gemm_tma_warpspecialized_cooperative_asymmetric_dma.hpp"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/gemv_blockscaled.h
+885
View File
@@ -0,0 +1,885 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief
*/
#pragma once
#include "cutlass/arch/cache_operation.h" /// cutlass::arch::CacheOperation
#include "cutlass/arch/memory.h" // cutlass::arch::global_load
#include "cutlass/arch/memory_sm80.h" // cp.async helpers, ldsm, cp_async_wait
#include "cutlass/complex.h" // cutlass::ComplexTransform:
#include "cutlass/cutlass.h"
#include "cutlass/fast_math.h" // cutlass::fast_max
#include "cutlass/layout/matrix.h" // cutlass::layout::RowMajor
#include "cutlass/matrix_coord.h" // cutlass::MatrixCoord
#include "cutlass/numeric_conversion.h" // cutlass::FloatRoundStyle, cutlass::NumericConverter
#include "cutlass/numeric_types.h" // cutlass::float_e4m3_t
#include "cutlass/platform/platform.h" // cutlass::is_same_v
#include "cutlass/tensor_ref.h" // cutlass::TensorRef
#include "cutlass/semaphore.h" // split-k
#include "cute/algorithm/functional.hpp" // cute::for_each
#include "cute/numeric/arithmetic_tuple.hpp" // cute::make_int_sequence
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace kernel {
using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
typename ElementA_,
typename LayoutA_,
typename ElementB_,
typename ElementC_,
typename ElementAccumulator_,
typename EpilogueOutputOp_,
int kElementsPerAccess_ = 1, ///< Number of elements involved in a global access.
int kThreadCount_ = 0, ///< Number of threads in the thread block.
/// It will be calculated automatically if set to 0.
int kThreadsPerRow_ = 0, ///< Number of threads in the k dimension.
/// It will be calculated automatically if set to 0.
typename ElementSFA_ = cutlass::float_e4m3_t,
typename ElementSFB_ = cutlass::float_e4m3_t,
int kSFVecSize_ = 16
>
struct GemvBlockScaled;
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Specializations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
// GEMV for row-major A matrix
template <typename ElementA_,
typename ElementB_,
typename ElementC_,
typename ElementAccumulator_,
typename EpilogueOutputOp_,
int kElementsPerAccess_,
int kThreadCount_,
int kThreadsPerRow_,
typename ElementSFA_,
typename ElementSFB_,
int kSFVecSize_>
struct GemvBlockScaled<ElementA_,
cutlass::layout::RowMajor,
ElementB_,
ElementC_,
ElementAccumulator_,
EpilogueOutputOp_,
kElementsPerAccess_,
kThreadCount_,
kThreadsPerRow_,
ElementSFA_,
ElementSFB_,
kSFVecSize_>
{
public:
using ElementA = ElementA_;
using ElementSFA = ElementSFA_;
using LayoutA = cutlass::layout::RowMajor;
using TensorRefA = cutlass::TensorRef<ElementA, LayoutA>;
static_assert(cutlass::sizeof_bits<ElementSFA>::value == 8, "ElementSFA should be FP8 type");
using ElementB = ElementB_;
using ElementSFB = ElementSFB_;
using LayoutB = cutlass::layout::ColumnMajor;
static_assert(cutlass::sizeof_bits<ElementSFB>::value == 8, "ElementSFB should be FP8 type");
using ElementC = ElementC_;
using LayoutC = cutlass::layout::ColumnMajor;
using ElementAccumulator = ElementAccumulator_;
static constexpr cutlass::ComplexTransform kTransformA = cutlass::ComplexTransform::kNone;
static constexpr cutlass::ComplexTransform kTransformB = cutlass::ComplexTransform::kNone;
static constexpr FloatRoundStyle Round = cutlass::FloatRoundStyle::round_to_nearest;
// number of return elements in a global access
static constexpr int kElementsPerAccess = kElementsPerAccess_;
static constexpr int kSFVecSize = kSFVecSize_;
static constexpr int kSFPerAccess = cutlass::const_max(1, kElementsPerAccess / kSFVecSize);
static_assert(kSFVecSize == 16, "Only SFVecSize = 16 is supported");
// Hardcode some check for easier debug
static_assert(kElementsPerAccess == 32, "for fp4 kernel, 32 elt per access");
static_assert(kSFPerAccess == 2, "fpr fp4 kernel, 2 sf read per thread");
static constexpr bool kDequantizeA = cutlass::sizeof_bits<ElementA>::value == 4;
static constexpr bool kDequantizeB = cutlass::sizeof_bits<ElementB>::value == 4;
static constexpr int kPackedElementsA = cutlass::sizeof_bits<ElementA>::value == 4 ? 2 : 1;
static constexpr int kPackedElementsB = cutlass::sizeof_bits<ElementB>::value == 4 ? 2 : 1;
static constexpr int kPackedElements = cutlass::const_max(kPackedElementsA, kPackedElementsB);
static_assert(kDequantizeA == true, "kDequantizeA should be true");
static_assert(kDequantizeB == true, "kDequantizeB should be true");
using FragmentA = cutlass::Array<ElementA, kElementsPerAccess>;
using FragmentB = cutlass::Array<ElementB, kElementsPerAccess>;
using FragmentCompute = cutlass::Array<ElementAccumulator, kElementsPerAccess>;
using FragmentSFA = cutlass::Array<ElementSFA, kSFPerAccess>;
using FragmentSFB = cutlass::Array<ElementSFB, kSFPerAccess>;
using FragmentPackedA = cutlass::Array<ElementA, kPackedElements>;
using FragmentPackedB = cutlass::Array<ElementB, kPackedElements>;
static_assert(sizeof_bits<FragmentA>::value == 128, "FragmentA should be 128 bits");
static_assert(sizeof_bits<FragmentB>::value == 128, "FragmentB should be 128 bits");
// // thread block shape (kThreadsPerRow, kThreadCount / kThreadsPerRow, 1)
static constexpr int kThreadCount = (kThreadCount_ <= 0) ? 128 : kThreadCount_;
static constexpr int kThreadsPerRow = (kThreadsPerRow_ <= 0) ?
cutlass::const_min(static_cast<int>(kThreadCount / cutlass::bits_to_bytes(kElementsPerAccess * cutlass::sizeof_bits<ElementA>::value)), 16) :
kThreadsPerRow_;
static constexpr int kThreadsPerCol = kThreadCount / kThreadsPerRow;
static constexpr int kStageCount = 4;
static constexpr int kBufferCount = 2;
// Number of elements stored in shared memory per stage for operands A and B.
// Each thread contributes `kElementsPerAccess / kPackedElements{A,B}` packed
// values.
static constexpr int kSmemPerStageA = kThreadCount * kElementsPerAccess / kPackedElementsA;
// B is uniform across all threads in the same k-column, so only store it once per k-thread
static constexpr int kSmemPerStageB = kThreadsPerRow * kElementsPerAccess / kPackedElementsB;
using EpilogueOutputOp = EpilogueOutputOp_;
// Ensure epilogue and mainloop have same thread layout
static_assert(kThreadCount == EpilogueOutputOp::kThreadCount, "mainloop, epilogue thread count mismatch");
static_assert(kThreadsPerRow == EpilogueOutputOp::kThreadsPerRow, "mainloop, epilogue thread per row mismatch");
static_assert(kThreadsPerCol == EpilogueOutputOp::kThreadsPerCol, "mainloop, epilogue thread per col mismatch");
//
// Structures
//
/// Argument structure
struct Arguments
{
MatrixCoord problem_size;
int32_t batch_count{0};
typename EpilogueOutputOp::Params epilogue;
TensorRefA ref_A;
ElementB const *ptr_B{nullptr};
ElementC const *ptr_C{nullptr};
ElementC *ptr_D{nullptr};
ElementSFA const *ptr_SFA{nullptr};
ElementSFB const *ptr_SFB{nullptr};
int64_t stride_A{0};
int64_t batch_stride_A{0};
int64_t batch_stride_B{0};
int64_t batch_stride_C{0};
int64_t batch_stride_D{0};
int64_t batch_stride_SFA{0};
int64_t batch_stride_SFB{0};
int64_t batch_stride_SFD{0};
};
using Params = Arguments;
/// Shared memory storage structure
struct SharedStorage
{
using EpilogueStorage = typename EpilogueOutputOp::SharedStorage;
EpilogueStorage epilogue;
alignas(16) ElementA smem_A[kBufferCount][kStageCount][kSmemPerStageA];
alignas(16) ElementB smem_B[kBufferCount][kStageCount][kSmemPerStageB];
alignas(16) ElementSFA smem_SFA[kBufferCount][kStageCount][kThreadCount * kSFPerAccess];
alignas(16) ElementSFB smem_SFB[kBufferCount][kStageCount][kThreadsPerRow * kSFPerAccess];
};
public:
//
// Methods
//
/// Determines whether kernel satisfies alignment
static Status can_implement(cutlass::MatrixCoord const &problem_size)
{
if (problem_size.column() % kElementsPerAccess != 0) {
return Status::kErrorMisalignedOperand;
}
return Status::kSuccess;
}
static Status can_implement(Arguments const &args)
{
return can_implement(args.problem_size);
}
/// Executes one GEMV
CUTLASS_DEVICE
void operator()(Params const &params, SharedStorage &shared_storage)
{
EpilogueOutputOp epilogue(params.epilogue, shared_storage.epilogue);
// Converters only needed for regular GEMV fallback case
NumericConverter<ElementAccumulator, ElementA, Round> A_converter;
NumericConverter<ElementAccumulator, ElementB, Round> B_converter;
NumericConverter<ElementAccumulator, ElementSFA, Round> SFA_converter;
NumericConverter<ElementAccumulator, ElementSFB, Round> SFB_converter;
const int32_t gemm_m = params.problem_size.row();
[[maybe_unused]] static constexpr int32_t gemm_n = 1;
const int32_t gemm_k = params.problem_size.column();
const int32_t gemm_batch = params.batch_count;
// Loop over batch indices
for (int batch_idx = blockIdx.z; batch_idx < gemm_batch; batch_idx += gridDim.z) {
int idx_col_k = threadIdx.x;
int idx_row_m = blockIdx.x * blockDim.y + threadIdx.y;
if (idx_row_m < gemm_m) {
// problem_size (row = m, column = k)
// matrix A (batch, m, k)
// vector B (batch, k, 1)
// vector C (batch, m, 1)
// vector D (batch, m, 1)
// move in the batch dimension
ElementA const *ptr_A = params.ref_A.data() + batch_idx * params.batch_stride_A / kPackedElementsA;
ElementB const *ptr_B = params.ptr_B + batch_idx * params.batch_stride_B / kPackedElementsB;
ElementC const *ptr_C = params.ptr_C + batch_idx * params.batch_stride_C;
ElementC *ptr_D = params.ptr_D + batch_idx * params.batch_stride_D;
// move in the k dimension
ptr_A += idx_col_k * kElementsPerAccess / kPackedElementsA;
ptr_B += idx_col_k * kElementsPerAccess / kPackedElementsB;
// move in the m dimension
ptr_A += idx_row_m * params.stride_A / kPackedElementsA;
ptr_C += idx_row_m;
ptr_D += idx_row_m;
ElementSFA const *ptr_SF_A{nullptr};
ElementSFB const *ptr_SF_B{nullptr};
int global_k{0};
int SF_blocks_by_M = (gemm_m + 127) >> 7;
int SF_blocks_by_K = (gemm_k / kSFVecSize + 3) >> 2;
// move in the batch dimension
ptr_SF_A = params.ptr_SFA + batch_idx * SF_blocks_by_M * SF_blocks_by_K * 512;
ptr_SF_B = params.ptr_SFB + batch_idx * SF_blocks_by_K * 512;
// move in the m dimension
ptr_SF_A += (((idx_row_m >> 7) * SF_blocks_by_K) << 9) + ((idx_row_m & 0x1f) << 4) + ((idx_row_m & 0x7f) >> 5 << 2);
global_k = idx_col_k * kElementsPerAccess;
ElementAccumulator accum = ElementAccumulator(0);
// Local aliases
const int tileA_k_local = kThreadsPerRow * kElementsPerAccess;
const int total_tiles = gemm_k / tileA_k_local;
int unroll_col_k = 0; // total K elements consumed so far by this thread
const int thread_id = threadIdx.y * kThreadsPerRow + threadIdx.x;
const bool is_even_thread = (threadIdx.x % 2 == 0);
const bool load_b = (threadIdx.y == 0);
const int smem_sf_write_offset = (thread_id / 2) * 4; // 4 FP8 per even thread
const int smem_sf_offset = thread_id * kSFPerAccess;
// Fast path: if the problem fits entirely in the tail path, skip SMEM
if (total_tiles == 0) {
accum += process_tail_elements(0, idx_col_k, gemm_k,
ptr_A, ptr_B,
ptr_SF_A, ptr_SF_B,
A_converter, B_converter,
SFA_converter, SFB_converter);
} else {
// Scaling factors are now loaded from shared memory, no register pipeline needed
// Thread-local SMEM line offset
const int thread_linear = threadIdx.y * kThreadsPerRow + threadIdx.x;
const int smem_offset_A = thread_linear * (kElementsPerAccess / kPackedElementsA);
// Only one row of threads (threadIdx.y == 0) loads B
const int smem_offset_B = threadIdx.x * (kElementsPerAccess / kPackedElementsB);
// PROLOGUE – prime first kStageCount-1 stages into buffer 0
CUTLASS_PRAGMA_UNROLL
for (int b = 0; b < kBufferCount - 1; ++b) {
// Load all stages using the helper function
load_stages_gmem_to_smem(
b, // buffer_idx
kStageCount, // num_stages
unroll_col_k, // passed by reference
global_k, // passed by reference
tileA_k_local,
smem_offset_A,
smem_offset_B,
smem_sf_write_offset,
is_even_thread,
load_b,
true, // valid_tile = true for prologue
ptr_A,
ptr_B,
ptr_SF_A,
ptr_SF_B,
shared_storage);
}
cutlass::arch::cp_async_fence();
// Ensure first stage committed
cutlass::arch::cp_async_wait<kBufferCount - 2>();
__syncthreads();
// Register double buffering for A/B fragments and SFA/SFB like SM80
FragmentA fragA_reg[2];
FragmentB fragB_reg[2];
FragmentSFA fragSFA_reg[2];
FragmentSFB fragSFB_reg[2];
// Current pipe index in smem to read from
int smem_pipe_read = 0;
// Current pipe index in smem to write to
int smem_pipe_write = kBufferCount - 1;
// PREFETCH register pipeline - load first kblock (stage 0) into register bank 0
if constexpr (kStageCount > 1)
{
int frag_idx = 0;
// Load fragments using the helper function
load_smem_fragments(
fragA_reg[frag_idx],
fragB_reg[frag_idx],
fragSFA_reg[frag_idx],
fragSFB_reg[frag_idx],
smem_pipe_read,
0, // k_block = 0
smem_offset_A,
smem_offset_B,
smem_sf_offset,
shared_storage);
}
// Mainloop
int tile_idx = 0;
while (tile_idx < total_tiles) {
int smem_pipe_read_curr = smem_pipe_read;
for_each(make_int_sequence<kStageCount>{}, [&] (auto k_block)
{
if (k_block == kStageCount - 1)
{
cutlass::arch::cp_async_wait<kBufferCount - 2>();
__syncthreads();
smem_pipe_read_curr = smem_pipe_read;
}
// Load A/B/SFA/SFB smem->regs for k_block_next
auto k_block_next = (k_block + Int<1>{}) % kStageCount;
int frag_idx_next = (k_block + 1) & 1;
// Prefetch next kblock data using saved pipe index
load_smem_fragments(
fragA_reg[frag_idx_next],
fragB_reg[frag_idx_next],
fragSFA_reg[frag_idx_next],
fragSFB_reg[frag_idx_next],
smem_pipe_read_curr,
k_block_next,
smem_offset_A,
smem_offset_B,
smem_sf_offset,
shared_storage);
// Copy gmem to smem before computing gemm on each k-pipe
if (k_block == 0)
{
// Use predicate instead of branch for cp_async
bool valid_tile = (global_k < gemm_k);
// Load all stages using the helper function
load_stages_gmem_to_smem(
smem_pipe_write, // buffer_idx
kStageCount, // num_stages
unroll_col_k, // passed by reference
global_k, // passed by reference
tileA_k_local,
smem_offset_A,
smem_offset_B,
smem_sf_write_offset,
is_even_thread,
load_b,
valid_tile,
ptr_A,
ptr_B,
ptr_SF_A,
ptr_SF_B,
shared_storage);
cutlass::arch::cp_async_fence();
// Advance the pipe indices
smem_pipe_write = smem_pipe_read;
++smem_pipe_read;
smem_pipe_read = (smem_pipe_read == kBufferCount) ? 0 : smem_pipe_read;
}
{
int frag_idx = k_block & 1;
// Compute using current fragments
accum += blockscaled_multiply_add(
fragA_reg[frag_idx], fragB_reg[frag_idx],
fragSFA_reg[frag_idx],
fragSFB_reg[frag_idx]);
}
});
tile_idx += kStageCount;
}
// Drain outstanding async copies
cutlass::arch::cp_async_wait<0>();
__syncthreads();
// Tail elements that don't fill a full tile
if (unroll_col_k + idx_col_k * kPackedElementsA < gemm_k) {
accum += process_tail_elements(unroll_col_k, idx_col_k, gemm_k,
ptr_A, ptr_B,
ptr_SF_A, ptr_SF_B,
A_converter, B_converter,
SFA_converter, SFB_converter);
}
}
CUTLASS_PRAGMA_UNROLL
for (int mask = (kThreadsPerRow >> 1); mask > 0; mask >>= 1) {
accum += ElementAccumulator(__shfl_xor_sync(0xFFFFFFFF, static_cast<float>(accum), mask, 32));
}
auto frag_acc = static_cast<typename EpilogueOutputOp::ElementAccumulator>(accum);
auto frag_c = static_cast<typename EpilogueOutputOp::ElementC>(*(ptr_C));
// Applying blockscaled epilogue
epilogue(frag_acc, frag_c, batch_idx);
}
}
} //end of operator()
private:
// Load multiple stages from global to shared memory
CUTLASS_DEVICE
void load_stages_gmem_to_smem(
int buffer_idx,
int num_stages,
int& unroll_col_k,
int& global_k,
int tileA_k_local,
int smem_offset_A,
int smem_offset_B,
int smem_sf_write_offset,
bool is_even_thread,
bool load_b,
bool valid_tile,
ElementA const* ptr_A,
ElementB const* ptr_B,
ElementSFA const* ptr_SF_A,
ElementSFB const* ptr_SF_B,
SharedStorage& shared_storage) {
CUTLASS_PRAGMA_UNROLL
for (int s = 0; s < num_stages; ++s) {
// Load scaling factors using cp.async - only even threads participate
// Calculate SF indices for this thread
int SF_idx = global_k / kSFVecSize;
int SF_offset_by_k = ((SF_idx >> 2) << 9) + (SF_idx & 0x3);
void *smem_ptr_SFA = &shared_storage.smem_SFA[buffer_idx][s][smem_sf_write_offset];
const void *gmem_ptr_SFA = ptr_SF_A + SF_offset_by_k;
// Load 4 FP8 values (32 bits) - for this thread and next thread
cutlass::arch::cp_async<sizeof(uint32_t)>(smem_ptr_SFA, gmem_ptr_SFA, valid_tile && is_even_thread);
void *smem_ptr_SFB = &shared_storage.smem_SFB[buffer_idx][s][(threadIdx.x / 2) * 4];
const void *gmem_ptr_SFB = ptr_SF_B + SF_offset_by_k;
// Load 4 FP8 values (32 bits) - for this thread and next thread, only if threadIdx.y == 0
cutlass::arch::cp_async<sizeof(uint32_t)>(smem_ptr_SFB, gmem_ptr_SFB, valid_tile && load_b && is_even_thread);
void *smem_ptr_A = &shared_storage.smem_A[buffer_idx][s][smem_offset_A];
const void *gmem_ptr_A = ptr_A + unroll_col_k / kPackedElementsA;
cutlass::arch::cp_async<sizeof(FragmentA)>(smem_ptr_A, gmem_ptr_A, valid_tile);
void *smem_ptr_B = &shared_storage.smem_B[buffer_idx][s][smem_offset_B];
const void *gmem_ptr_B = ptr_B + unroll_col_k / kPackedElementsB;
cutlass::arch::cp_async<sizeof(FragmentB)>(smem_ptr_B, gmem_ptr_B, valid_tile && load_b);
unroll_col_k += tileA_k_local;
global_k += tileA_k_local;
}
}
/// Fused blockscaled GEMV computation using PTX
CUTLASS_DEVICE
ElementAccumulator blockscaled_multiply_add(
FragmentA const& fragA,
FragmentB const& fragB,
FragmentSFA const& fragSFA,
FragmentSFB const& fragSFB) {
#if defined(CUDA_PTX_FP4FP6_CVT_ENABLED)
uint16_t const& src_fragSFA_packed = reinterpret_cast<uint16_t const&>(fragSFA);
uint16_t const& src_fragSFB_packed = reinterpret_cast<uint16_t const&>(fragSFB);
uint32_t const* src_fragA_packed = reinterpret_cast<uint32_t const*>(&fragA);
uint32_t const* src_fragB_packed = reinterpret_cast<uint32_t const*>(&fragB);
ElementAccumulator out;
uint16_t* out_fp16 = reinterpret_cast<uint16_t*>(&out);
asm volatile( \
"{\n" \
// declare registers for A / B tensors
".reg .b8 byte0_0, byte0_1, byte0_2, byte0_3;\n" \
".reg .b8 byte0_4, byte0_5, byte0_6, byte0_7;\n" \
".reg .b8 byte1_0, byte1_1, byte1_2, byte1_3;\n" \
".reg .b8 byte1_4, byte1_5, byte1_6, byte1_7;\n" \
".reg .b8 byte2_0, byte2_1, byte2_2, byte2_3;\n" \
".reg .b8 byte2_4, byte2_5, byte2_6, byte2_7;\n" \
".reg .b8 byte3_0, byte3_1, byte3_2, byte3_3;\n" \
".reg .b8 byte3_4, byte3_5, byte3_6, byte3_7;\n" \
// declare registers for accumulators
".reg .f16x2 accum_0_0, accum_0_1, accum_0_2, accum_0_3;\n" \
".reg .f16x2 accum_1_0, accum_1_1, accum_1_2, accum_1_3;\n" \
".reg .f16x2 accum_2_0, accum_2_1, accum_2_2, accum_2_3;\n" \
".reg .f16x2 accum_3_0, accum_3_1, accum_3_2, accum_3_3;\n" \
// declare registers for scaling factors
".reg .f16x2 sfa_f16x2;\n" \
".reg .f16x2 sfb_f16x2;\n" \
".reg .f16x2 sf_f16x2;\n" \
// declare registers for conversion
".reg .f16x2 cvt_0_0, cvt_0_1, cvt_0_2, cvt_0_3;\n" \
".reg .f16x2 cvt_0_4, cvt_0_5, cvt_0_6, cvt_0_7;\n" \
".reg .f16x2 cvt_1_0, cvt_1_1, cvt_1_2, cvt_1_3;\n" \
".reg .f16x2 cvt_1_4, cvt_1_5, cvt_1_6, cvt_1_7;\n" \
".reg .f16x2 cvt_2_0, cvt_2_1, cvt_2_2, cvt_2_3;\n" \
".reg .f16x2 cvt_2_4, cvt_2_5, cvt_2_6, cvt_2_7;\n" \
".reg .f16x2 cvt_3_0, cvt_3_1, cvt_3_2, cvt_3_3;\n" \
".reg .f16x2 cvt_3_4, cvt_3_5, cvt_3_6, cvt_3_7;\n" \
".reg .f16 result_f16, lane0, lane1;\n" \
".reg .f16x2 mul_f16x2_0, mul_f16x2_1;\n" \
// convert scaling factors from fp8 to f16x2
"cvt.rn.f16x2.e4m3x2 sfa_f16x2, %1;\n" \
"cvt.rn.f16x2.e4m3x2 sfb_f16x2, %2;\n" \
// clear accumulators
"mov.b32 accum_0_0, 0;\n" \
"mov.b32 accum_0_1, 0;\n" \
"mov.b32 accum_0_2, 0;\n" \
"mov.b32 accum_0_3, 0;\n" \
"mov.b32 accum_1_0, 0;\n" \
"mov.b32 accum_1_1, 0;\n" \
"mov.b32 accum_1_2, 0;\n" \
"mov.b32 accum_1_3, 0;\n" \
"mov.b32 accum_2_0, 0;\n" \
"mov.b32 accum_2_1, 0;\n" \
"mov.b32 accum_2_2, 0;\n" \
"mov.b32 accum_2_3, 0;\n" \
"mov.b32 accum_3_0, 0;\n" \
"mov.b32 accum_3_1, 0;\n" \
"mov.b32 accum_3_2, 0;\n" \
"mov.b32 accum_3_3, 0;\n" \
// multiply, unpacking and permuting scale factors
"mul.rn.f16x2 sf_f16x2, sfa_f16x2, sfb_f16x2;\n" \
"mov.b32 {lane0, lane1}, sf_f16x2;\n" \
"mov.b32 mul_f16x2_0, {lane0, lane0};\n" \
"mov.b32 mul_f16x2_1, {lane1, lane1};\n" \
// unpacking A and B tensors
"mov.b32 {byte0_0, byte0_1, byte0_2, byte0_3}, %3;\n" \
"mov.b32 {byte0_4, byte0_5, byte0_6, byte0_7}, %4;\n" \
"mov.b32 {byte1_0, byte1_1, byte1_2, byte1_3}, %5;\n" \
"mov.b32 {byte1_4, byte1_5, byte1_6, byte1_7}, %6;\n" \
"mov.b32 {byte2_0, byte2_1, byte2_2, byte2_3}, %7;\n" \
"mov.b32 {byte2_4, byte2_5, byte2_6, byte2_7}, %8;\n" \
"mov.b32 {byte3_0, byte3_1, byte3_2, byte3_3}, %9;\n" \
"mov.b32 {byte3_4, byte3_5, byte3_6, byte3_7}, %10;\n" \
// convert A and B tensors from fp4 to f16x2
// A[0 - 7] and B[0 - 7]
"cvt.rn.f16x2.e2m1x2 cvt_0_0, byte0_0;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_0_1, byte0_1;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_0_2, byte0_2;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_0_3, byte0_3;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_0_4, byte0_4;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_0_5, byte0_5;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_0_6, byte0_6;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_0_7, byte0_7;\n" \
// A[8 - 15] and B[8 - 15]
"cvt.rn.f16x2.e2m1x2 cvt_1_0, byte1_0;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_1_1, byte1_1;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_1_2, byte1_2;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_1_3, byte1_3;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_1_4, byte1_4;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_1_5, byte1_5;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_1_6, byte1_6;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_1_7, byte1_7;\n" \
// A[16 - 23] and B[16 - 23]
"cvt.rn.f16x2.e2m1x2 cvt_2_0, byte2_0;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_2_1, byte2_1;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_2_2, byte2_2;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_2_3, byte2_3;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_2_4, byte2_4;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_2_5, byte2_5;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_2_6, byte2_6;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_2_7, byte2_7;\n" \
// A[24 - 31] and B[24 - 31]
"cvt.rn.f16x2.e2m1x2 cvt_3_0, byte3_0;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_3_1, byte3_1;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_3_2, byte3_2;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_3_3, byte3_3;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_3_4, byte3_4;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_3_5, byte3_5;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_3_6, byte3_6;\n" \
"cvt.rn.f16x2.e2m1x2 cvt_3_7, byte3_7;\n" \
// fma for A[0 - 7] and B[0 - 7]
"fma.rn.f16x2 accum_0_0, cvt_0_0, cvt_0_4, accum_0_0;\n" \
"fma.rn.f16x2 accum_0_1, cvt_0_1, cvt_0_5, accum_0_1;\n" \
"fma.rn.f16x2 accum_0_2, cvt_0_2, cvt_0_6, accum_0_2;\n" \
"fma.rn.f16x2 accum_0_3, cvt_0_3, cvt_0_7, accum_0_3;\n" \
// fma for A[8 - 15] and B[8 - 15]
"fma.rn.f16x2 accum_1_0, cvt_1_0, cvt_1_4, accum_1_0;\n" \
"fma.rn.f16x2 accum_1_1, cvt_1_1, cvt_1_5, accum_1_1;\n" \
"fma.rn.f16x2 accum_1_2, cvt_1_2, cvt_1_6, accum_1_2;\n" \
"fma.rn.f16x2 accum_1_3, cvt_1_3, cvt_1_7, accum_1_3;\n" \
// fma for A[16 - 23] and B[16 - 23]
"fma.rn.f16x2 accum_2_0, cvt_2_0, cvt_2_4, accum_2_0;\n" \
"fma.rn.f16x2 accum_2_1, cvt_2_1, cvt_2_5, accum_2_1;\n" \
"fma.rn.f16x2 accum_2_2, cvt_2_2, cvt_2_6, accum_2_2;\n" \
"fma.rn.f16x2 accum_2_3, cvt_2_3, cvt_2_7, accum_2_3;\n" \
// fma for A[24 - 31] and B[24 - 31]
"fma.rn.f16x2 accum_3_0, cvt_3_0, cvt_3_4, accum_3_0;\n" \
"fma.rn.f16x2 accum_3_1, cvt_3_1, cvt_3_5, accum_3_1;\n" \
"fma.rn.f16x2 accum_3_2, cvt_3_2, cvt_3_6, accum_3_2;\n" \
"fma.rn.f16x2 accum_3_3, cvt_3_3, cvt_3_7, accum_3_3;\n" \
// tree reduction for accumulators
"add.rn.f16x2 accum_0_0, accum_0_0, accum_0_1;\n" \
"add.rn.f16x2 accum_0_2, accum_0_2, accum_0_3;\n" \
"add.rn.f16x2 accum_1_0, accum_1_0, accum_1_1;\n" \
"add.rn.f16x2 accum_1_2, accum_1_2, accum_1_3;\n" \
"add.rn.f16x2 accum_2_0, accum_2_0, accum_2_1;\n" \
"add.rn.f16x2 accum_2_2, accum_2_2, accum_2_3;\n" \
"add.rn.f16x2 accum_3_0, accum_3_0, accum_3_1;\n" \
"add.rn.f16x2 accum_3_2, accum_3_2, accum_3_3;\n" \
"add.rn.f16x2 accum_0_0, accum_0_0, accum_0_2;\n" \
"add.rn.f16x2 accum_1_0, accum_1_0, accum_1_2;\n" \
"add.rn.f16x2 accum_2_0, accum_2_0, accum_2_2;\n" \
"add.rn.f16x2 accum_3_0, accum_3_0, accum_3_2;\n" \
"add.rn.f16x2 accum_0_0, accum_0_0, accum_1_0;\n" \
"add.rn.f16x2 accum_2_0, accum_2_0, accum_3_0;\n" \
// apply scaling factors and final reduction
"mul.rn.f16x2 accum_0_0, mul_f16x2_0, accum_0_0;\n" \
"mul.rn.f16x2 accum_2_0, mul_f16x2_1, accum_2_0;\n" \
"add.rn.f16x2 accum_0_0, accum_0_0, accum_2_0;\n" \
"mov.b32 {lane0, lane1}, accum_0_0;\n" \
"add.rn.f16 result_f16, lane0, lane1;\n" \
"mov.b16 %0, result_f16;\n" \
"}\n"
: "=h"(out_fp16[0]) // 0
: "h"(src_fragSFA_packed), "h"(src_fragSFB_packed), // 1, 2
"r"(src_fragA_packed[0]), "r"(src_fragB_packed[0]), // 3, 4
"r"(src_fragA_packed[1]), "r"(src_fragB_packed[1]), // 5, 6
"r"(src_fragA_packed[2]), "r"(src_fragB_packed[2]), // 7, 8
"r"(src_fragA_packed[3]), "r"(src_fragB_packed[3]) // 9, 10
: "memory"
);
return out;
#else
NumericArrayConverter<ElementAccumulator, ElementA, kElementsPerAccess, Round> srcA_converter;
NumericArrayConverter<ElementAccumulator, ElementB, kElementsPerAccess, Round> srcB_converter;
NumericConverter<ElementAccumulator, ElementSFA, Round> SFA_converter;
NumericConverter<ElementAccumulator, ElementSFB, Round> SFB_converter;
FragmentCompute fragA_Compute = srcA_converter(fragA);
FragmentCompute fragB_Compute = srcB_converter(fragB);
ElementAccumulator accum = ElementAccumulator(0);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < kSFPerAccess; i++) {
ElementAccumulator accum_SF_block = ElementAccumulator(0);
int local_k_offset = i * kSFVecSize;
ElementAccumulator multiplier{1};
multiplier = SFA_converter(fragSFA.at(i)) * SFB_converter(fragSFB.at(i));
CUTLASS_PRAGMA_UNROLL
for (int e = 0; e < kSFVecSize; e++) {
accum_SF_block += fragA_Compute.at(e + local_k_offset) * fragB_Compute.at(e + local_k_offset);
}
accum_SF_block *= multiplier;
accum += accum_SF_block;
}
return accum;
#endif
}
CUTLASS_DEVICE
ElementAccumulator process_tail_elements(
int unroll_col_k,
int idx_col_k,
int gemm_k,
ElementA const *ptr_A,
ElementB const *ptr_B,
ElementSFA const *ptr_SF_A,
ElementSFB const *ptr_SF_B,
NumericConverter<ElementAccumulator, ElementA, Round> const &A_converter,
NumericConverter<ElementAccumulator, ElementB, Round> const &B_converter,
NumericConverter<ElementAccumulator, ElementSFA, Round> const &SFA_converter,
NumericConverter<ElementAccumulator, ElementSFB, Round> const &SFB_converter) {
ElementAccumulator accum = ElementAccumulator(0);
// calculate the rest of K elements
// each thread fetch 1 element each time
for (int k = unroll_col_k + idx_col_k * kPackedElementsA; k < gemm_k; k += kThreadsPerRow * kPackedElementsA) {
// blockscaled GEMV
int SF_idx = k / kSFVecSize;
int SF_offset_by_k = ((SF_idx >> 2) << 9) + (SF_idx & 0x3);
ElementSFA sfa = *(ptr_SF_A + SF_offset_by_k);
ElementSFB sfb = *(ptr_SF_B + SF_offset_by_k);
FragmentPackedA fragA;
FragmentPackedB fragB;
// fetch from matrix A
arch::global_load<FragmentPackedA, sizeof(FragmentPackedA), arch::CacheOperation::Always>(
fragA,
ptr_A - (idx_col_k * kElementsPerAccess - k) / kPackedElementsA,
true);
// fetch from vector B
arch::global_load<FragmentPackedB, sizeof(FragmentPackedB), arch::CacheOperation::Always>(
fragB,
ptr_B - (idx_col_k * kElementsPerAccess - k) / kPackedElementsB,
true);
ElementAccumulator accum_SF_packed = ElementAccumulator(0);
CUTLASS_PRAGMA_UNROLL
for (int e = 0; e < kPackedElements; e++) {
accum_SF_packed += A_converter(fragA.at(e)) * B_converter(fragB.at(e));
}
accum_SF_packed *= SFA_converter(sfa) * SFB_converter(sfb);
accum += accum_SF_packed;
}
return accum;
}
// Load fragments from shared memory
template<typename FragmentA, typename FragmentB, typename FragmentSFA, typename FragmentSFB>
CUTLASS_DEVICE
void load_smem_fragments(
FragmentA& fragA,
FragmentB& fragB,
FragmentSFA& fragSFA,
FragmentSFB& fragSFB,
int smem_pipe_idx,
int k_block,
int smem_offset_A,
int smem_offset_B,
int smem_sf_offset,
SharedStorage& shared_storage) const {
// Load A/B fragments
arch::shared_load(fragA, &shared_storage.smem_A[smem_pipe_idx][k_block][smem_offset_A]);
arch::shared_load(fragB, &shared_storage.smem_B[smem_pipe_idx][k_block][smem_offset_B]);
// Load SF fragments
uint32_t smem_ptr = cutlass::arch::cutlass_get_smem_pointer(&shared_storage.smem_SFA[smem_pipe_idx][k_block][smem_sf_offset]);
arch::shared_load<2>(&fragSFA, smem_ptr);
smem_ptr = cutlass::arch::cutlass_get_smem_pointer(&shared_storage.smem_SFB[smem_pipe_idx][k_block][threadIdx.x * kSFPerAccess]);
arch::shared_load<2>(&fragSFB, smem_ptr);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/kernel/sm100_gemm_cpasync_warpspecialized.hpp
+793
View File
@@ -0,0 +1,793 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/workspace.h"
#include "cutlass/kernel_hardware_info.hpp"
#include "cutlass/detail/cluster.hpp"
#include "cutlass/fast_math.h"
#include "cute/arch/cluster_sm90.hpp"
#include "cutlass/arch/arch.h"
#include "cutlass/arch/barrier.h"
#include "cutlass/arch/reg_reconfig.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/kernel/sm100_tile_scheduler.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cute/tensor.hpp"
#include "cute/arch/tmem_allocator_sm100.hpp"
#include "cute/atom/mma_atom.hpp"
///////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::kernel {
///////////////////////////////////////////////////////////////////////////////
template <
class ProblemShape_,
class CollectiveMainloop_,
class CollectiveEpilogue_,
class TileSchedulerTag_
>
class GemmUniversal<
ProblemShape_,
CollectiveMainloop_,
CollectiveEpilogue_,
TileSchedulerTag_,
cute::enable_if_t<
cutlass::detail::is_kernel_tag_of_v<typename CollectiveMainloop_::DispatchPolicy::Schedule,
KernelWarpSpecializedSm100>>>
{
public:
using ProblemShape = ProblemShape_;
static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4,
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
static constexpr bool IsGdcEnabled = false;
// Mainloop derived types
using CollectiveMainloop = CollectiveMainloop_;
using TileShape = typename CollectiveMainloop::TileShape;
using TiledMma = typename CollectiveMainloop::TiledMma;
using ArchTag = typename CollectiveMainloop::ArchTag;
using ElementA = typename CollectiveMainloop::ElementA;
using StrideA = typename CollectiveMainloop::StrideA;
using ElementB = typename CollectiveMainloop::ElementB;
using StrideB = typename CollectiveMainloop::StrideB;
using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
using ClusterShape = typename DispatchPolicy::ClusterShape;
using MainloopArguments = typename CollectiveMainloop::Arguments;
using MainloopParams = typename CollectiveMainloop::Params;
static_assert(ArchTag::kMinComputeCapability >= 100);
// Epilogue derived types
using CollectiveEpilogue = CollectiveEpilogue_;
using ElementC = typename CollectiveEpilogue::ElementC;
using StrideC = typename CollectiveEpilogue::StrideC;
using ElementD = typename CollectiveEpilogue::ElementD;
using StrideD = typename CollectiveEpilogue::StrideD;
using EpilogueArguments = typename CollectiveEpilogue::Arguments;
using EpilogueParams = typename CollectiveEpilogue::Params;
static constexpr bool IsComplex = CollectiveEpilogue::NumAccumulatorMtxs == 2;
// CLC pipeline depth
// determines how many waves (stages-1) a warp can race ahead
static constexpr uint32_t SchedulerPipelineStageCount = DispatchPolicy::Schedule::SchedulerPipelineStageCount;
// TileID scheduler
// Get Blk and Scheduling tile shapes
using CtaShape_MNK = typename CollectiveMainloop::CtaShape_MNK;
using AtomThrShapeMNK = typename CollectiveMainloop::AtomThrShapeMNK;
static_assert(size(AtomThrShapeMNK{}) == 1, "Lower alignment kernel only supports 1x1x1 cluster shape.");
using TileSchedulerTag = TileSchedulerTag_;
using TileScheduler = typename detail::TileSchedulerSelector<
TileSchedulerTag, ArchTag, CtaShape_MNK, ClusterShape, SchedulerPipelineStageCount>::Scheduler;
using TileSchedulerArguments = typename TileScheduler::Arguments;
using TileSchedulerParams = typename TileScheduler::Params;
// Warp specialization thread count per threadblock
static constexpr uint32_t NumSchedThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumMMAThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumEmptyThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumMainloopLoadThreads = CollectiveMainloop::NumLoadThreads; // 4 warps
static constexpr uint32_t NumEpilogueLoadThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumEpilogueThreads = CollectiveEpilogue::ThreadCount;
static constexpr uint32_t NumEpilogueWarps = NumEpilogueThreads / NumThreadsPerWarp;
static constexpr uint32_t MaxThreadsPerBlock = NumSchedThreads +
NumMainloopLoadThreads + NumMMAThreads +
NumEpilogueLoadThreads + NumEpilogueThreads + NumEmptyThreads;
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
static constexpr uint32_t NumFixupBarriers = 1;
static constexpr uint32_t CLCResponseSize = sizeof(typename TileScheduler::CLCResponse);
static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent;
// Pipelines and pipeline states
static constexpr uint32_t AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount;
// Pipeline and pipeline state types
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
using MainloopPipelineState = typename CollectiveMainloop::MainloopPipelineState;
using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
using EpiLoadPipelineState = typename CollectiveEpilogue::LoadPipelineState;
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
using EpiStorePipelineState = typename CollectiveEpilogue::StorePipelineState;
using AccumulatorPipeline = cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount, AtomThrShapeMNK>;
using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;
using CLCPipeline = cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape>;
using CLCPipelineState = typename CLCPipeline::PipelineState;
using TmemAllocator = cute::TMEM::Allocator1Sm;
static constexpr int EpilogueWarpRegs = 248;
static constexpr int NonEpilogueWarpRegs = 128;
// Kernel level shared memory storage
struct SharedStorage {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
using CLCPipelineStorage = typename CLCPipeline::SharedStorage;
using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;
alignas(16) MainloopPipelineStorage mainloop;
alignas(16) EpiLoadPipelineStorage epi_load;
alignas(16) CLCPipelineStorage clc;
alignas(16) AccumulatorPipelineStorage accumulator;
alignas(16) arch::ClusterBarrier tmem_dealloc;
} pipelines;
alignas(16) typename TileScheduler::CLCResponse clc_response[SchedulerPipelineStageCount];
uint32_t tmem_base_ptr;
struct TensorStorage : cute::aligned_struct<128, _1> {
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
MainloopTensorStorage mainloop;
EpilogueTensorStorage epilogue;
} tensors;
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
// Host facing host arguments
struct Arguments {
GemmUniversalMode mode{};
ProblemShape problem_shape{};
MainloopArguments mainloop{};
EpilogueArguments epilogue{};
KernelHardwareInfo hw_info{};
TileSchedulerArguments scheduler{};
};
// Kernel device entry point API
struct Params {
GemmUniversalMode mode{};
ProblemShape problem_shape{};
MainloopParams mainloop{};
EpilogueParams epilogue{};
KernelHardwareInfo hw_info{};
TileSchedulerParams scheduler{};
};
enum class WarpCategory : int32_t {
MMA = 0,
Sched = 1,
EpilogueLoad = 3,
Epilogue = 4,
MainloopLoad = 8
};
struct IsParticipant {
uint32_t mma = false;
uint32_t sched = false;
uint32_t epi_load = false;
uint32_t epilogue = false;
uint32_t main_load = false;
};
// Convert to underlying arguments.
static
Params
to_underlying_arguments(Arguments const& args, void* workspace) {
(void) workspace;
auto problem_shape = args.problem_shape;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
static constexpr uint32_t NumEpilogueSubTiles = 1;
// Get SM count if needed, otherwise use user supplied SM count
int sm_count = args.hw_info.sm_count;
if (sm_count != 0) {
CUTLASS_TRACE_HOST(" WARNING: SM100 tile scheduler does not allow for user specified SM counts.\n"
" To restrict a kernel's resource usage, consider using CUDA driver APIs instead (green contexts).");
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
}
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
// Calculate workspace pointers
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
// Epilogue
void* epilogue_workspace = workspace_ptr + workspace_offset;
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
void* mainloop_workspace = nullptr;
// Tile scheduler
void* scheduler_workspace = workspace_ptr + workspace_offset;
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
return {
args.mode,
args.problem_shape,
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, mainloop_workspace),
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, epilogue_workspace),
hw_info,
TileScheduler::to_underlying_arguments(
problem_shape_MNKL, TileShape{}, AtomThrShapeMNK{}, ClusterShape{},
args.hw_info, args.scheduler, scheduler_workspace
)
};
}
static bool
can_implement(Arguments const& args) {
bool implementable = (args.mode == GemmUniversalMode::kGemm) or
(args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements.\n");
return implementable;
}
implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
implementable &= TileScheduler::can_implement(args.scheduler);
static constexpr int MaxClusterSize = 16;
implementable &= size(ClusterShape{}) <= MaxClusterSize;
return implementable;
}
static size_t
get_workspace_size(Arguments const& args) {
static constexpr uint32_t NumEpilogueSubTiles = 1;
size_t workspace_size = 0;
// Epilogue
workspace_size += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
// Tile scheduler
workspace_size += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
return workspace_size;
}
static cutlass::Status
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
CudaHostAdapter* cuda_adapter = nullptr) {
Status status = Status::kSuccess;
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
static constexpr uint32_t NumEpilogueSubTiles = 1;
// Epilogue
status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
status = cutlass::Status::kSuccess;
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
if (status != Status::kSuccess) {
return status;
}
// Tile scheduler
status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs, cuda_adapter);
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
args.scheduler, args.problem_shape, args.hw_info, NumFixupBarriers);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
if (status != Status::kSuccess) {
return status;
}
return status;
}
static dim3
get_grid_shape(Params const& params) {
auto cluster_shape = ClusterShape{};
auto blk_shape = CtaShape_MNK{};
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
return TileScheduler::get_grid_shape(
params.scheduler,
problem_shape_MNKL,
TileShape{},
AtomThrShapeMNK{},
cluster_shape,
params.hw_info
);
}
static dim3
get_block_shape() {
return dim3(MaxThreadsPerBlock, 1, 1);
}
public:
CUTLASS_DEVICE
void
operator()(Params const& params, char* smem_buf) {
using namespace cute;
using X = Underscore;
// Separate out problem shape for convenience
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
auto M = get<0>(problem_shape_MNKL);
auto N = get<1>(problem_shape_MNKL);
auto K = get<2>(problem_shape_MNKL);
auto L = get<3>(problem_shape_MNKL);
// Account for more than one epilogue warp
int warp_idx = canonical_warp_idx_sync();
WarpCategory warp_category = warp_idx < static_cast<int>(WarpCategory::Epilogue) ? WarpCategory(warp_idx)
: warp_idx < static_cast<int>(WarpCategory::MainloopLoad) ? WarpCategory::Epilogue
: WarpCategory::MainloopLoad;
uint32_t lane_predicate = cute::elect_one_sync();
auto tile_shape = TileShape{};
auto cluster_shape = ClusterShape{};
constexpr int cluster_size = size(ClusterShape{});
int cta_rank_in_cluster = cute::block_rank_in_cluster();
bool is_first_cta_in_cluster = cta_rank_in_cluster == 0;
int cta_coord_v = cta_rank_in_cluster % size<0>(typename TiledMma::AtomThrID{});
bool is_mma_leader_cta = cta_coord_v == 0;
int mma_leader_ctas = size(shape_div(cluster_shape, AtomThrShapeMNK{}));
[[maybe_unused]] uint32_t mma_peer_cta_rank = cta_rank_in_cluster;
// Kernel level shared memory storage
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
// In a warp specialized kernel, collectives expose data movement and compute operations separately
CollectiveMainloop collective_mainloop;
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
// Do we load source tensor C or other aux inputs
bool is_epi_load_needed = collective_epilogue.is_producer_load_needed();
IsParticipant is_participant = {
(warp_category == WarpCategory::MMA) && is_mma_leader_cta, // mma
(warp_category == WarpCategory::Sched) && is_first_cta_in_cluster, // sched
(warp_category == WarpCategory::EpilogueLoad) && is_epi_load_needed, // epi_load
(warp_category == WarpCategory::Epilogue), // epilogue
(warp_category == WarpCategory::MainloopLoad) // main_load
};
// Mainloop Load pipeline
typename MainloopPipeline::Params mainloop_pipeline_params;
if (WarpCategory::MainloopLoad == warp_category) {
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
}
if (WarpCategory::MMA == warp_category) {
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
}
mainloop_pipeline_params.producer_arv_count = NumMainloopLoadThreads;
mainloop_pipeline_params.consumer_arv_count = 1; // Only UMMA consumes the A and B buffers
mainloop_pipeline_params.dst_blockid = cta_rank_in_cluster;
mainloop_pipeline_params.initializing_warp = 0;
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, cluster_shape);
// Epilogue Load pipeline
typename EpiLoadPipeline::Params epi_load_pipeline_params;
if (WarpCategory::EpilogueLoad == warp_category) {
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
}
if (WarpCategory::Epilogue == warp_category) {
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
}
epi_load_pipeline_params.dst_blockid = cta_rank_in_cluster;
epi_load_pipeline_params.producer_arv_count = NumEpilogueLoadThreads;
epi_load_pipeline_params.consumer_arv_count = NumEpilogueThreads;
epi_load_pipeline_params.transaction_bytes = CollectiveEpilogue::TmaTransactionBytes;
epi_load_pipeline_params.initializing_warp = 3;
EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
// Epilogue Store pipeline
typename EpiStorePipeline::Params epi_store_pipeline_params;
epi_store_pipeline_params.always_wait = true;
EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
// CLC pipeline
typename CLCPipeline::Params clc_pipeline_params;
if (WarpCategory::Sched == warp_category) {
clc_pipeline_params.role = CLCPipeline::ThreadCategory::ProducerConsumer;
}
else {
clc_pipeline_params.role = CLCPipeline::ThreadCategory::Consumer;
}
clc_pipeline_params.producer_blockid = 0;
clc_pipeline_params.producer_arv_count = 1;
clc_pipeline_params.consumer_arv_count = NumSchedThreads + cluster_size *
(NumMainloopLoadThreads + NumEpilogueThreads + NumMMAThreads);
clc_pipeline_params.transaction_bytes = CLCResponseSize;
clc_pipeline_params.initializing_warp = 1;
CLCPipeline clc_pipeline(shared_storage.pipelines.clc, clc_pipeline_params, cluster_shape);
// Mainloop-Epilogue pipeline
typename AccumulatorPipeline::Params accumulator_pipeline_params;
if (WarpCategory::MMA == warp_category) {
accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;
}
if (WarpCategory::Epilogue == warp_category) {
accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;
}
// Only one producer thread arrives on this barrier.
accumulator_pipeline_params.producer_arv_count = 1;
accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * NumEpilogueThreads;
accumulator_pipeline_params.initializing_warp = 2;
AccumulatorPipeline accumulator_pipeline(shared_storage.pipelines.accumulator, accumulator_pipeline_params, cluster_shape);
// Tmem allocator
TmemAllocator tmem_allocator{};
// Sync allocation status between MMA and epilogue warps within CTA
arch::NamedBarrier tmem_allocation_result_barrier(NumMMAThreads + NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);
// Sync deallocation status between MMA warps of peer CTAs
arch::ClusterBarrier& tmem_deallocation_result_barrier = shared_storage.pipelines.tmem_dealloc;
[[maybe_unused]] uint32_t dealloc_barrier_phase = 0;
MainloopPipelineState mainloop_pipe_consumer_state;
MainloopPipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
EpiLoadPipelineState epi_load_pipe_consumer_state;
EpiLoadPipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
// epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding)
EpiStorePipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
CLCPipelineState clc_pipe_consumer_state;
CLCPipelineState clc_pipe_producer_state = cutlass::make_producer_start_state<CLCPipeline>();
AccumulatorPipelineState accumulator_pipe_consumer_state;
AccumulatorPipelineState accumulator_pipe_producer_state = cutlass::make_producer_start_state<AccumulatorPipeline>();
// We need this to guarantee that the Pipeline init is visible
// To all producers and consumer threadblocks in the cluster
pipeline_init_arrive_relaxed(cluster_size);
dim3 block_id_in_cluster = cute::block_id_in_cluster();
// TileID scheduler
TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
//
// TMEM "Allocation"
//
auto acc_shape = collective_mainloop.partition_accumulator_shape();
auto bulk_tmem = TiledMma::make_fragment_C(append(acc_shape,
Int<AccumulatorPipelineStageCount>{}));
//
// END PROLOGUE
//
// Synchronization call. Blocks until barriers are initialized in shared memory.
pipeline_init_wait(cluster_size);
if (is_participant.main_load) {
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
auto load_inputs = collective_mainloop.load_init(
problem_shape_MNKL, params.mainloop, shared_storage.tensors.mainloop);
Tensor gA_mkl = get<0>(load_inputs);
do {
// Get current work tile and fetch next work tile
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
);
// Get the number of K tiles to compute for this work as well as the starting K tile offset of the work.
auto k_tile_iter = scheduler.get_k_tile_iterator(work_tile_info, problem_shape_MNKL, CtaShape_MNK{}, shape<3>(gA_mkl));
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{});
auto [mainloop_producer_state_next, unused_] = collective_mainloop.load(
params.mainloop,
mainloop_pipeline,
mainloop_pipe_producer_state,
load_inputs,
cta_coord_mnkl,
k_tile_iter, k_tile_count
);
mainloop_pipe_producer_state = mainloop_producer_state_next;
// Sync warp to prevent non-participating threads entering next wave early
__syncwarp();
work_tile_info = next_work_tile_info;
if (increment_pipe) {
++clc_pipe_consumer_state;
}
} while (work_tile_info.is_valid());
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
}
else if (is_participant.sched) {
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
if constexpr (IsSchedDynamicPersistent) {
// Whether a new CLC query must be performed.
// See comment below where this variable is updated for a description of
// why this variable is needed.
bool requires_clc_query = true;
cutlass::arch::wait_on_dependent_grids();
do {
if (requires_clc_query) {
// Query next clcID and update producer state
clc_pipe_producer_state = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state);
}
// Fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
);
// Only perform a new CLC query if we consumed a new CLC query result in
// `fetch_next_work`. An example of a case in which CLC `fetch_next_work` does
// not consume a new CLC query response is when processing stream-K units.
// The current stream-K scheduler uses single WorkTileInfo to track multiple
// (potentially-partial) tiles to be computed via stream-K. In this case,
// `fetch_next_work` simply performs in-place updates on the existing WorkTileInfo,
// rather than consuming a CLC query response.
requires_clc_query = increment_pipe;
if (increment_pipe) {
++clc_pipe_consumer_state;
}
work_tile_info = next_work_tile_info;
} while (work_tile_info.is_valid());
clc_pipeline.producer_tail(clc_pipe_producer_state);
}
}
else if (is_participant.mma) {
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
// Tmem allocation sequence
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
__syncwarp();
tmem_allocation_result_barrier.arrive();
uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;
bulk_tmem.data() = tmem_base_ptr;
// Pass the acc with tuple type since the bgrad kernel change the mma_init API
auto mma_inputs = collective_mainloop.mma_init(params.mainloop, cute::make_tuple(bulk_tmem, bulk_tmem), shared_storage.tensors.mainloop);
do {
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{});
// Fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
);
if (increment_pipe) {
++clc_pipe_consumer_state;
}
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
int acc_stage = accumulator_pipe_producer_state.index();
Tensor accumulators = bulk_tmem(_,_,_,acc_stage);
mainloop_pipe_consumer_state = collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
// Pass the acc with tuple type since the bgrad kernel change the mma API
cute::make_tuple(accumulators, accumulators),
mma_inputs,
k_tile_count
);
accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);
++accumulator_pipe_producer_state;
work_tile_info = next_work_tile_info;
} while (work_tile_info.is_valid());
// Release the right to allocate before deallocations so that the next CTA can rasterize
tmem_allocator.release_allocation_lock();
accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);
// Free entire tmem allocation
tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
else if (is_participant.epi_load) {
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
bool do_tail_load = false;
do {
bool compute_epilogue = TileScheduler::compute_epilogue(work_tile_info, params.scheduler);
// Get current work tile and fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
);
work_tile_info = next_work_tile_info;
if (increment_pipe) {
++clc_pipe_consumer_state;
}
if (compute_epilogue) {
epi_load_pipe_producer_state = collective_epilogue.load(
epi_load_pipeline,
epi_load_pipe_producer_state,
problem_shape_MNKL,
CtaShape_MNK{},
cta_coord_mnkl,
TileShape{},
TiledMma{},
shared_storage.tensors.epilogue
);
do_tail_load = true;
}
// Calculate the cta coordinates of the next work tile
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
} while (work_tile_info.is_valid());
if (do_tail_load) {
collective_epilogue.load_tail(
epi_load_pipeline, epi_load_pipe_producer_state,
epi_store_pipeline, epi_store_pipe_producer_state);
}
}
else if (is_participant.epilogue) {
cutlass::arch::warpgroup_reg_alloc<EpilogueWarpRegs>();
// Wait for tmem allocate here
tmem_allocation_result_barrier.arrive_and_wait();
uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;
bulk_tmem.data() = tmem_base_ptr;
bool do_tail_store = false;
do {
// Fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
clc_pipeline,
clc_pipe_consumer_state
);
if (increment_pipe) {
++clc_pipe_consumer_state;
}
// Accumulator stage slice
int acc_stage = accumulator_pipe_consumer_state.index();
Tensor accumulators = bulk_tmem(_,_,_,acc_stage);
accumulator_pipe_consumer_state = scheduler.template fixup<IsComplex>(
TiledMma{},
work_tile_info,
accumulators,
accumulator_pipeline,
accumulator_pipe_consumer_state,
typename CollectiveEpilogue::CopyOpT2R{}
);
//
// Epilogue and write to gD
//
if (scheduler.compute_epilogue(work_tile_info)) {
auto [load_state_next, store_state_next, acc_state_next] = collective_epilogue.store(
epi_load_pipeline,
epi_load_pipe_consumer_state,
epi_store_pipeline,
epi_store_pipe_producer_state,
accumulator_pipeline,
accumulator_pipe_consumer_state,
problem_shape_MNKL,
CtaShape_MNK{},
cta_coord_mnkl,
TileShape{},
TiledMma{},
accumulators,
shared_storage.tensors.epilogue
);
epi_load_pipe_consumer_state = load_state_next;
epi_store_pipe_producer_state = store_state_next;
accumulator_pipe_consumer_state = acc_state_next;
do_tail_store = true;
}
work_tile_info = next_work_tile_info;
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
} while (work_tile_info.is_valid());
if (do_tail_store) {
collective_epilogue.store_tail(
epi_load_pipeline, epi_load_pipe_consumer_state,
epi_store_pipeline, epi_store_pipe_producer_state,
CtaShape_MNK{});
}
}
else {
cutlass::arch::warpgroup_reg_dealloc<NonEpilogueWarpRegs>();
}
}
};
///////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::kernel
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_mixed_input_transform.hpp
+1090
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp
+1 -3
View File
@@ -88,9 +88,7 @@ public:
static_assert(cute::is_static<TileShape>::value);
auto selected_cluster_shape = cutlass::detail::select_cluster_shape(cluster_shape_mnk, hw_info.cluster_shape);
auto cta_shape = cute::conditional_return<not cute::is_static_v<ClusterShape>>(
shape_div(tile_shape_mnk, atom_thr_shape_mnk), // Dynamic Cluster: For 2SM kernels, use CTA tile shape for the underlying scheduler
shape_div(tile_shape_mnk, selected_cluster_shape)); // Static Cluster: Blackwell builders expects TileShape to be Cluster's Tile Shape, Hopper doesn't
auto cta_shape = shape_div(tile_shape_mnk, atom_thr_shape_mnk); // For 2SM kernels, use CTA tile shape for the underlying scheduler
dim3 problem_blocks = get_tiled_cta_shape_mnl(
problem_shapes,
include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp
+1319
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp
+1112
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp
+5 -3
View File
@@ -411,9 +411,11 @@ public:
using namespace cute;
using X = Underscore;
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
# define ENABLE_SM90_KERNEL_LEVEL 1
#endif
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
# define ENABLE_SM90_KERNEL_LEVEL 1
# endif
// Any Tensor Op MMA Atom in the ISA is arch conditional.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp
+5 -3
View File
@@ -423,9 +423,11 @@ public:
using namespace cute;
using X = Underscore;
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
# define ENABLE_SM90_KERNEL_LEVEL 1
#endif
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
# define ENABLE_SM90_KERNEL_LEVEL 1
# endif
// Any Tensor Op MMA Atom in the ISA is arch conditional.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp
+5 -3
View File
@@ -270,9 +270,11 @@ public:
using namespace cute;
using X = Underscore;
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
# define ENABLE_SM90_KERNEL_LEVEL 1
#endif
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
# define ENABLE_SM90_KERNEL_LEVEL 1
# endif
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+5 -3
View File
@@ -342,9 +342,11 @@ public:
using namespace cute;
using X = Underscore;
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
# define ENABLE_SM90_KERNEL_LEVEL 1
#endif
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
# define ENABLE_SM90_KERNEL_LEVEL 1
# endif
// Any Tensor Op MMA Atom in the ISA is arch conditional.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+5 -3
View File
@@ -354,9 +354,11 @@ public:
using namespace cute;
using X = Underscore;
#if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200))
# define ENABLE_SM90_KERNEL_LEVEL 1
#endif
# if (defined(__CUDA_ARCH_FEAT_SM90_ALL) || defined(__CUDA_ARCH_FEAT_SM120_ALL) || defined(__CUDA_ARCH_FEAT_SM121_ALL) ||\
CUDA_ARCH_CONDITIONAL_OR_FAMILY(1200) || CUDA_ARCH_CONDITIONAL_OR_FAMILY(1210))
# define ENABLE_SM90_KERNEL_LEVEL 1
# endif
// Any Tensor Op MMA Atom in the ISA is arch conditional.
#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
include/cutlass/gemm/kernel/tile_scheduler.hpp
+60
View File
@@ -298,6 +298,66 @@ struct TileSchedulerSelector<
using Scheduler = StaticPersistentTileScheduler100;
};
template <class TileShape, class ClusterShape, uint32_t SchedulerPipelineStageCount>
struct TileSchedulerSelector<
PersistentScheduler,
arch::Sm103,
TileShape,
ClusterShape,
SchedulerPipelineStageCount> {
using Scheduler = PersistentTileSchedulerSm100<
ClusterShape,
SchedulerPipelineStageCount>;
};
// Ptr-Array kernel may provide a specialized ArrayProblemShape type
template <class TileShape,
class ClusterShape,
uint32_t SchedulerPipelineStageCount,
class ProblemShape>
struct TileSchedulerSelector<
PersistentScheduler,
arch::Sm103,
TileShape,
ClusterShape,
SchedulerPipelineStageCount,
ProblemShape> {
using Scheduler = PersistentTileSchedulerSm100<
ClusterShape,
SchedulerPipelineStageCount>;
};
// SM103 Group tile scheduler
template <
class TileShape,
class ClusterShape,
uint32_t SchedulerPipelineStageCount,
class GroupProblemShape
>
struct TileSchedulerSelector<
GroupScheduler,
arch::Sm103,
TileShape,
ClusterShape,
SchedulerPipelineStageCount,
GroupProblemShape
> {
using Scheduler = PersistentTileSchedulerSm100Group<GroupProblemShape, SchedulerPipelineStageCount>;
};
template <class TileShape, class ClusterShape, uint32_t SchedulerPipelineStageCount>
struct TileSchedulerSelector<
StreamKScheduler,
arch::Sm103,
TileShape,
ClusterShape,
SchedulerPipelineStageCount> {
using Scheduler = PersistentTileSchedulerSm100StreamK<
TileShape,
ClusterShape,
SchedulerPipelineStageCount>;
};
// Default (void) for Sm120 maps to PersistentTileSchedulerSm100
template <class TileShape, class ClusterShape, uint32_t SchedulerPipelineStageCount>
struct TileSchedulerSelector<
include/cutlass/integer_subbyte.h
+8 -3
View File
@@ -35,14 +35,13 @@
*/
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#include CUDA_STD_HEADER(cstdint)
#else
#include <cstdint>
#endif
#include "cutlass/cutlass.h"
#include "cutlass/numeric_size.h"
#include "cutlass/platform/platform.h"
@@ -206,6 +205,12 @@ using int2b_t = integer_subbyte<2, true>;
/// 2-bit Unsigned integer type
using uint2b_t = integer_subbyte<2, false>;
/// 3-bit Integer type
using int3b_t = integer_subbyte<3, true>;
/// 3-bit Unsigned integer type
using uint3b_t = integer_subbyte<3, false>;
/// 4-bit Integer type
using int4b_t = integer_subbyte<4, true>;
include/cutlass/layout/permute.h
+1 -2
View File
@@ -38,9 +38,8 @@
computation lies in operator() with private member variables {col_permute_, row_permute_ and stride_} as new addresses after permute op.
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/fast_math.h"
#include "cutlass/layout/pitch_linear.h"
#include "cutlass/layout/matrix.h"
include/cutlass/layout/tensor.h
+1 -2
View File
@@ -39,9 +39,8 @@
defined in cutlass/tensor_ref.h.
*/
#pragma once
#include <cuda/std/cassert>
#include "cutlass/cutlass.h"
#include CUDA_STD_HEADER(cassert)
#include "cutlass/fast_math.h"
#include "cutlass/layout/pitch_linear.h"
#include "cutlass/layout/matrix.h"
include/cutlass/numeric_conversion.h
+3 -3
View File
@@ -3901,16 +3901,16 @@ struct NumericArrayConverter<float_e2m1_t, float, 4, Round> {
static result_type convert(source_type const & source) {
#if defined(CUDA_PTX_FP4FP6_CVT_ENABLED)
unsigned out;
uint16_t out;
asm volatile( \
"{\n" \
".reg .b8 byte0;\n" \
".reg .b8 byte1;\n" \
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" \
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" \
"mov.b32 %0, {byte0, byte1, 0, 0};\n" \
"mov.b16 %0, {byte0, byte1};\n" \
"}" \
: "=r"(out) : "f"(source[0]), "f"(source[1]), "f"(source[2]), "f"(source[3]));
: "=h"(out) : "f"(source[0]), "f"(source[1]), "f"(source[2]), "f"(source[3]));
return reinterpret_cast<result_type const &>(out);
#else
include/cutlass/numeric_types.h
+1
View File
@@ -43,6 +43,7 @@
#include "cutlass/tfloat32.h"
#include "cutlass/float8.h"
#include "cutlass/uint128.h"
#include "cutlass/uint256.h"
#include "cutlass/exmy_base.h"
#include "cutlass/float_subbyte.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/pipeline/sm100_pipeline.hpp
+72 -5
View File
@@ -140,6 +140,8 @@ public:
int warp_idx = canonical_warp_idx_sync();
if (warp_idx == params.initializing_warp) {
// Barrier FULL and EMPTY init
CUTLASS_ASSERT(params.producer_arv_count > 0 && "Producer arrival count must be non-zero");
CUTLASS_ASSERT(params.consumer_arv_count > 0 && "Consumer arrival count must be non-zero");
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
storage.full_barrier_, storage.empty_barrier_, params.producer_arv_count, params.consumer_arv_count);
}
@@ -320,6 +322,24 @@ public:
}
}
template<class ClusterShape, class InitBarriers = cute::true_type, class InitMasks = cute::true_type>
CUTLASS_DEVICE
PipelineTmaTransformAsync(SharedStorage& storage, Params params, ClusterShape cluster_shape, McastDirection mcast_direction, InitBarriers = {}, InitMasks = {})
: impl_(storage, params, cluster_shape, cute::false_type{}, cute::false_type{})
, params_(params)
, empty_barrier_ptr_(&storage.empty_barrier_[0])
, full_barrier_ptr_(&storage.full_barrier_[0]) {
static_assert(cute::is_same_v<InitBarriers, cute::true_type> || cute::is_same_v<InitBarriers, cute::false_type>);
if constexpr (cute::is_same_v<InitBarriers, cute::true_type>) {
init_barriers(storage, params_, cluster_shape, mcast_direction);
}
static_assert(cute::is_same_v<InitMasks, cute::true_type> || cute::is_same_v<InitMasks, cute::false_type>);
if constexpr (cute::is_same_v<InitMasks, cute::true_type>) {
init_masks(cluster_shape, mcast_direction);
}
}
// Helper function to initialize barriers
template <class ClusterShape>
static
@@ -334,9 +354,11 @@ public:
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
static_assert(IsDynamicCluster or ((cute::size<0>(cluster_shape) % cute::size<0>(atom_thr_shape) == 0) &&
(cute::size<1>(cluster_shape) % cute::size<1>(atom_thr_shape) == 0)));
uint32_t const num_consumer_per_cluster = params.num_consumers / NumThreadsPerWarpGroup;
uint32_t const num_consumer_per_cluster = cute::ceil_div(params.num_consumers, static_cast<uint32_t>(NumThreadsPerWarpGroup));
uint32_t const multicast_consumer_arrival_count = ((cute::size<0>(cluster_shape) / cute::size<0>(atom_thr_shape)) +
(cute::size<1>(cluster_shape) / cute::size<1>(atom_thr_shape)) - 1) * num_consumer_per_cluster;
CUTLASS_ASSERT(multicast_consumer_arrival_count > 0 && "Multicast consumer arrival count must be non-zero");
CUTLASS_ASSERT(producer_arv_cnt > 0 && "Producer arrival count must be non-zero");
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, multicast_consumer_arrival_count);
}
@@ -344,8 +366,31 @@ public:
}
template <class ClusterShape>
static
CUTLASS_DEVICE
void init_masks(ClusterShape cluster_shape, dim3 block_id_in_cluster = cute::block_id_in_cluster()) {
void
init_barriers(SharedStorage& storage, Params params, ClusterShape cluster_shape, McastDirection mcast_direction) {
auto atom_thr_shape = AtomThrShape_MNK{};
int warp_idx = canonical_warp_idx_sync();
if (warp_idx == params.initializing_warp) {
// Barrier FULL and EMPTY init
constexpr int producer_arv_cnt = 1;
uint32_t const num_consumer_per_cluster = params.num_consumers / NumThreadsPerWarpGroup;
uint32_t const multicast_consumer_arrival_count = (mcast_direction == McastDirection::kRow) ?
(cute::size<1>(cluster_shape) / cute::size<1>(atom_thr_shape)) * num_consumer_per_cluster : // Mcast with row ctas
(cute::size<0>(cluster_shape) / cute::size<0>(atom_thr_shape)) * num_consumer_per_cluster; // Mcast with col ctas
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, multicast_consumer_arrival_count);
}
cutlass::arch::fence_barrier_init();
}
template <class ClusterShape>
CUTLASS_DEVICE
void init_masks(ClusterShape cluster_shape, dim3 block_id_in_cluster = cute::block_id_in_cluster(), McastDirection mcast_dir = McastDirection::kRowCol) {
// Calculate consumer mask
if (params_.role == ThreadCategory::Consumer) {
// Logic to optimally schedule Empty Arrives
@@ -374,10 +419,25 @@ public:
// STEP 2: Find if this dst block-id needs an arrival for this problem
is_signaling_thread_ &= dst_blockid_ < cluster_size;
is_signaling_thread_ &= is_same_row_or_col(dst_blockid_, block_id_in_cluster, cluster_shape);
if(mcast_dir == McastDirection::kRowCol){
is_signaling_thread_ &= is_same_row_or_col(dst_blockid_, block_id_in_cluster, cluster_shape);
}
if(mcast_dir == McastDirection::kRow){
is_signaling_thread_ &= is_same_row(dst_blockid_, block_id_in_cluster, cluster_shape);
}
}
}
template <class ClusterShape>
CUTLASS_DEVICE
bool is_same_row(int dst_block_id, dim3 block_id, ClusterShape cluster_shape) {
return (((dst_block_id % cute::size<0>(cluster_shape)) == block_id.x)
// If we are in the same cluster column and using 2CTA MMA, only odd or only even CTAs sync with each other
&& ((dst_block_id % cute::size<0>(cluster_shape)) % cute::size<0>(AtomThrShape_MNK{}) ==
block_id.x % cute::size<0>(AtomThrShape_MNK{}))
);
}
template <class ClusterShape>
CUTLASS_DEVICE
bool is_same_row_or_col(int dst_block_id, dim3 block_id, ClusterShape cluster_shape) {
@@ -504,7 +564,8 @@ public:
auto atom_thr_shape = AtomThrShape_MNK{};
uint32_t const multicast_consumer_arrival_count = (cute::size<0>(cluster_shape) / cute::size<0>(atom_thr_shape)) +
(cute::size<1>(cluster_shape) / cute::size<1>(atom_thr_shape)) - 1;
CUTLASS_ASSERT(multicast_consumer_arrival_count > 0 && "Multicast consumer arrival count must be non-zero");
CUTLASS_ASSERT(producer_arv_cnt > 0 && "Producer arrival count must be non-zero");
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, multicast_consumer_arrival_count);
}
@@ -525,6 +586,8 @@ public:
cute::size<1>(cluster_shape) / cute::size<1>(atom_thr_shape) : // Mcast with row ctas
cute::size<0>(cluster_shape) / cute::size<0>(atom_thr_shape); // Mcast with col ctas
CUTLASS_ASSERT(multicast_consumer_arrival_count > 0 && "Multicast consumer arrival count must be non-zero");
CUTLASS_ASSERT(producer_arv_cnt > 0 && "Producer arrival count must be non-zero");
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, multicast_consumer_arrival_count);
}
@@ -893,6 +956,8 @@ public:
int warp_idx = canonical_warp_idx_sync();
if (warp_idx == params.initializing_warp) {
// Barrier FULL and EMPTY init
CUTLASS_ASSERT(params.producer_arv_count > 0 && "Producer arrival count must be non-zero");
CUTLASS_ASSERT(params.consumer_arv_count > 0 && "Consumer arrival count must be non-zero");
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(full_barrier_ptr_), decltype(empty_barrier_ptr_), Stages>(
full_barrier_ptr_, empty_barrier_ptr_, params_.producer_arv_count, params_.consumer_arv_count);
}
@@ -910,6 +975,8 @@ public:
int warp_idx = canonical_warp_idx_sync();
if (warp_idx == params.initializing_warp) {
// Barrier FULL and EMPTY init
CUTLASS_ASSERT(params.producer_arv_count > 0 && "Producer arrival count must be non-zero");
CUTLASS_ASSERT(params.consumer_arv_count > 0 && "Consumer arrival count must be non-zero");
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(full_barrier_ptr_), decltype(empty_barrier_ptr_), Stages>(
full_barrier_ptr_, empty_barrier_ptr_, params_.producer_arv_count, params_.consumer_arv_count);
}
@@ -1102,7 +1169,7 @@ public:
///////////////////////////////////////////////////////////////////////////////////////////////////
//
// TMA (producer - consumer) Async Pipeline classes for Blackwell Sparse UMMA
// This is designed for the parttern that kernel has two different staged tensors. (AB and metadata)
// This is designed for the pattern that kernel has two different staged tensors. (AB and metadata)
//
///////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/pipeline/sm90_pipeline.hpp
+9 -2
View File
@@ -309,13 +309,14 @@ public:
if (is_initializing_warp) {
// Barrier FULL and EMPTY init
uint32_t const producer_arv_cnt = params.num_producers;
uint32_t const num_consumer_warpgroups_per_cluster = params.num_consumers / NumThreadsPerWarpGroup;
uint32_t const num_consumer_warpgroups_per_cluster = cute::ceil_div(params.num_consumers, static_cast<uint32_t>(NumThreadsPerWarpGroup));
uint32_t multicast_consumer_arrival_count = params.num_consumers; // If cluster_size is 1
if (cute::size(cluster_shape) > 1) {
multicast_consumer_arrival_count = (cute::size<0>(cluster_shape) + cute::size<1>(cluster_shape) - 1) *
num_consumer_warpgroups_per_cluster;
}
CUTLASS_ASSERT(multicast_consumer_arrival_count > 0 && "Multicast consumer arrival count must be non-zero");
CUTLASS_ASSERT(producer_arv_cnt > 0 && "Producer arrival count must be non-zero");
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, multicast_consumer_arrival_count);
}
@@ -804,6 +805,8 @@ public:
if (is_initializing_warp) {
// Barrier FULL and EMPTY init
CUTLASS_ASSERT(params.producer_arv_count > 0 && "Producer arrival count must be non-zero");
CUTLASS_ASSERT(params.consumer_arv_count > 0 && "Consumer arrival count must be non-zero");
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(full_barrier_ptr), decltype(empty_barrier_ptr), Stages>(
full_barrier_ptr, empty_barrier_ptr, params.producer_arv_count, params.consumer_arv_count);
}
@@ -1043,6 +1046,8 @@ public:
is_initializing_warp = (warp_idx == params.initializing_warp);
if (is_initializing_warp) {
// Barrier FULL and EMPTY init
CUTLASS_ASSERT(params.producer_arv_count > 0 && "Producer arrival count must be non-zero");
CUTLASS_ASSERT(params.consumer_arv_count > 0 && "Consumer arrival count must be non-zero");
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
storage.full_barrier_, storage.empty_barrier_, params.producer_arv_count, params.consumer_arv_count);
}
@@ -1299,6 +1304,7 @@ public:
// Barrier FULL, EMPTY init
if (warp_idx == params.initializing_warp) {
int arv_cnt = params.group_size;
CUTLASS_ASSERT(arv_cnt > 0 && "Arrive count must be non-zero");
constexpr int Stages = Depth * Length;
cutlass::arch::detail::initialize_barrier_array_aligned<decltype(barrier_ptr_), Stages>(
barrier_ptr_, arv_cnt);
@@ -1307,6 +1313,7 @@ public:
int warp_idx = canonical_warp_idx_sync();
int lane_predicate = cute::elect_one_sync();
CUTLASS_ASSERT(params.group_size > 0 && "Group size must be non-zero");
// Barrier FULL, EMPTY init
// Init is done only by the one elected thread of the block
include/cutlass/platform/platform.h
+17 -16
View File
@@ -98,13 +98,13 @@
//-----------------------------------------------------------------------------
// Dependencies
//-----------------------------------------------------------------------------
#include <cutlass/cutlass.h>
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#include <cuda/std/utility>
#include <cuda/std/cstddef>
#include <cuda/std/cstdint>
#include <cuda/std/limits>
#include CUDA_STD_HEADER(type_traits)
#include CUDA_STD_HEADER(utility)
#include CUDA_STD_HEADER(cstddef)
#include CUDA_STD_HEADER(cstdint)
#include CUDA_STD_HEADER(limits)
#else
#include <type_traits>
#include <utility>
@@ -128,7 +128,6 @@
#endif
#include <vector_types.h>
#include <cutlass/cutlass.h>
#endif
@@ -523,7 +522,7 @@ using std::is_trivially_copyable;
#endif
#if (201703L <=__cplusplus)
#if (CUTLASS_CXX17_OR_LATER)
/// std::is_unsigned_v
using CUTLASS_STL_NAMESPACE::is_integral_v;
@@ -596,14 +595,6 @@ struct alignment_of<float4> {
enum { value = 16 };
};
template <>
struct alignment_of<long4> {
enum { value = 16 };
};
template <>
struct alignment_of<ulong4> {
enum { value = 16 };
};
template <>
struct alignment_of<longlong2> {
enum { value = 16 };
};
@@ -615,6 +606,15 @@ template <>
struct alignment_of<double2> {
enum { value = 16 };
};
template <>
struct alignment_of<long4> {
enum { value = 16 };
};
template <>
struct alignment_of<ulong4> {
enum { value = 16 };
};
template <>
struct alignment_of<longlong4> {
enum { value = 16 };
@@ -628,6 +628,7 @@ struct alignment_of<double4> {
enum { value = 16 };
};
// Specializations for volatile/const qualified types
template <typename value_t>
struct alignment_of<volatile value_t> : alignment_of<value_t> {};
include/cutlass/predicate_vector.h
+3 -4
View File
@@ -33,16 +33,15 @@
of boolean predicates.
*/
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#include CUDA_STD_HEADER(cstdint)
#else
#include <cstdint>
#endif
#include <cuda/std/cassert>
#include CUDA_STD_HEADER(cassert)
#include "cutlass/cutlass.h"
#include "cutlass/platform/platform.h"
namespace cutlass {
include/cutlass/uint128.h
+2 -3
View File
@@ -33,9 +33,9 @@
\brief Defines an unsigned 128b integer with several operators to support 64-bit integer division.
*/
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#include CUDA_STD_HEADER(cstdint)
#else
#include <cstdint>
#include <cstdlib>
@@ -44,7 +44,6 @@
#include <stdexcept>
#endif
#include "cutlass/cutlass.h"
/// Optionally enable GCC's built-in type
#if (defined(__x86_64) || defined (__aarch64__)) && !(defined(__CUDA_ARCH__) && ((__CUDACC_VER_MAJOR__ <= 10) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ <= 4)))) && defined(__GNUC__)
include/cutlass/uint256.h
+93
View File
@@ -0,0 +1,93 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*!
\file
\brief Defines an unsigned 256b integer.
*/
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include CUDA_STD_HEADER(cstdint)
#else
#include <cstdint>
#include <cstdlib>
#include <cmath>
#include <type_traits>
#include <stdexcept>
#endif
#include "cutlass/uint128.h"
namespace cutlass {
///! Unsigned 256b integer type
struct alignas(32) uint256_t {
/// Size of one part of the uint's storage in bits
static constexpr int storage_bits_ = 128;
struct hilo {
uint128_t lo;
uint128_t hi;
};
// Use a union to store either low and high parts.
union {
struct hilo hilo_;
};
//
// Methods
//
/// Default ctor
CUTLASS_HOST_DEVICE
uint256_t() : hilo_{uint128_t{}, uint128_t{}} {}
/// Constructor from uint128
CUTLASS_HOST_DEVICE
uint256_t(uint128_t lo_) : hilo_{lo_, uint128_t{}} {}
/// Constructor from two 128b unsigned integers
CUTLASS_HOST_DEVICE
uint256_t(uint128_t lo_, uint128_t hi_) : hilo_{lo_, hi_} {}
/// Lossily cast to uint128_t
CUTLASS_HOST_DEVICE
explicit operator uint128_t() const {
return hilo_.lo;
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////