7f2a703ed5
* Merge with private repo * Update README * Update README * Update README * Add PyTorch requirements * Fix sync scopes for MQA logits (#256) * Update README
460 lines
24 KiB
C++
460 lines
24 KiB
C++
#pragma once
|
|
|
|
#include <torch/python.h>
|
|
|
|
#include "../../jit/compiler.hpp"
|
|
#include "../../jit/device_runtime.hpp"
|
|
#include "../../jit/kernel_runtime.hpp"
|
|
#include "../../utils/exception.hpp"
|
|
#include "../../utils/format.hpp"
|
|
#include "../../utils/math.hpp"
|
|
#include "../heuristics/sm100.hpp"
|
|
|
|
#include "epilogue.hpp"
|
|
#include "runtime_utils.hpp"
|
|
|
|
namespace deep_gemm {
|
|
|
|
class SM100FP8FP4Gemm1D1DRuntime final: public LaunchRuntime<SM100FP8FP4Gemm1D1DRuntime> {
|
|
public:
|
|
struct Args {
|
|
GemmDesc gemm_desc;
|
|
GemmConfig gemm_config;
|
|
LaunchArgs launch_args;
|
|
// TODO: move into descriptor
|
|
const std::optional<std::string> epilogue_type;
|
|
|
|
// TODO: move into descriptor
|
|
int gran_k_a, gran_k_b;
|
|
|
|
void* grouped_layout;
|
|
CUtensorMap tensor_map_a;
|
|
CUtensorMap tensor_map_b;
|
|
CUtensorMap tensor_map_sfa;
|
|
CUtensorMap tensor_map_sfb;
|
|
CUtensorMap tensor_map_cd;
|
|
};
|
|
|
|
static std::string generate_impl(const Args& args) {
|
|
// TODO: rename files
|
|
return fmt::format(R"(
|
|
#include <deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh>
|
|
|
|
using namespace deep_gemm;
|
|
|
|
static void __instantiate_kernel() {{
|
|
auto ptr = reinterpret_cast<void*>(&sm100_fp8_fp4_gemm_1d1d_impl<
|
|
{}, {},
|
|
{}, {},
|
|
{}, {}, {},
|
|
{}, {}, {},
|
|
{},
|
|
{}, {}, {},
|
|
{},
|
|
{}, {},
|
|
{}, {},
|
|
{},
|
|
{},
|
|
{}, {},
|
|
{}, {}, {},
|
|
{}
|
|
>);
|
|
}};
|
|
)",
|
|
to_string(args.gemm_desc.major_a), to_string(args.gemm_desc.major_b),
|
|
args.gran_k_a, args.gran_k_b,
|
|
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
|
|
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
|
|
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
|
|
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
|
|
args.gemm_desc.num_groups,
|
|
args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode,
|
|
args.gemm_config.pipeline_config.num_stages,
|
|
args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads,
|
|
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
|
|
args.gemm_config.launch_config.num_sms,
|
|
args.gemm_config.layout.swap_ab,
|
|
to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation,
|
|
to_string(args.gemm_desc.a_dtype), to_string(args.gemm_desc.b_dtype), to_string(args.gemm_desc.cd_dtype),
|
|
get_default_epilogue_type(args.epilogue_type));
|
|
}
|
|
|
|
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
|
|
// TODO: optimize `args` copy
|
|
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
|
|
args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
|
|
args.tensor_map_a, args.tensor_map_b,
|
|
args.tensor_map_sfa, args.tensor_map_sfb,
|
|
args.tensor_map_cd));
|
|
}
|
|
};
|
|
|
|
static void sm100_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
|
|
const torch::Tensor& b, const torch::Tensor& sfb,
|
|
const std::optional<torch::Tensor>& c,
|
|
const torch::Tensor& d,
|
|
const int& m, const int& n, const int& k,
|
|
const int& gran_k_a, const int& gran_k_b,
|
|
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
|
const std::string& compiled_dims,
|
|
const std::optional<std::string>& epilogue_type = std::nullopt) {
|
|
const auto desc = GemmDesc {
|
|
.gemm_type = GemmType::Normal,
|
|
.kernel_type = KernelType::Kernel1D1D,
|
|
.m = m, .n = n, .k = k, .num_groups = 1,
|
|
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
|
|
.cd_dtype = d.scalar_type(),
|
|
.major_a = major_a, .major_b = major_b,
|
|
.with_accumulation = c.has_value(),
|
|
.num_sms = device_runtime->get_num_sms(),
|
|
.tc_util = device_runtime->get_tc_util(),
|
|
.compiled_dims = compiled_dims
|
|
};
|
|
const auto config = get_best_config<SM100ArchSpec>(desc);
|
|
|
|
const auto cd = c.value_or(d);
|
|
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
|
|
config.storage_config.load_block_m,
|
|
config.layout.block_k,
|
|
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
|
|
config.storage_config.swizzle_a_mode);
|
|
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
|
|
config.storage_config.load_block_n,
|
|
config.layout.block_k,
|
|
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
|
|
config.storage_config.swizzle_b_mode);
|
|
const auto tensor_map_cd = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
|
|
config.storage_config.store_block_m,
|
|
config.storage_config.store_block_n,
|
|
static_cast<int>(d.stride(-2)), 1,
|
|
config.storage_config.swizzle_cd_mode);
|
|
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
|
|
config.layout.block_m, gran_k_a, 1, 0);
|
|
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
|
|
config.layout.block_n, gran_k_b, 1, 0);
|
|
|
|
// Launch
|
|
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
|
|
.gemm_desc = desc,
|
|
.gemm_config = config,
|
|
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
|
|
config.pipeline_config.smem_size,
|
|
config.layout.get_cluster_size()),
|
|
.epilogue_type = epilogue_type,
|
|
.gran_k_a = gran_k_a,
|
|
.gran_k_b = gran_k_b,
|
|
.grouped_layout = nullptr,
|
|
.tensor_map_a = tensor_map_a,
|
|
.tensor_map_b = tensor_map_b,
|
|
.tensor_map_sfa = tensor_map_sfa,
|
|
.tensor_map_sfb = tensor_map_sfb,
|
|
.tensor_map_cd = tensor_map_cd
|
|
};
|
|
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
|
|
const auto runtime = compiler->build("sm100_fp8_fp4_gemm_1d1d", code);
|
|
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
|
|
}
|
|
|
|
static void sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
|
|
const torch::Tensor& b, const torch::Tensor& sfb,
|
|
const torch::Tensor& d,
|
|
const torch::Tensor& grouped_layout,
|
|
const int& num_groups, const int& m, const int& n, const int& k,
|
|
const int& gran_k_a, const int& gran_k_b,
|
|
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
|
const std::string& compiled_dims,
|
|
const bool& use_psum_layout,
|
|
const std::optional<int>& expected_m_for_psum_layout) {
|
|
const auto gemm_type = use_psum_layout ?
|
|
GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
|
|
|
|
// Only psum layout can use expected m
|
|
if (expected_m_for_psum_layout)
|
|
DG_HOST_ASSERT(use_psum_layout);
|
|
|
|
// NOTES: If actual M is dynamic, estimate config via `num_groups` and `expected_m`.
|
|
// Otherwise, treat the contiguous layout as a whole.
|
|
const auto desc = GemmDesc {
|
|
.gemm_type = gemm_type,
|
|
.kernel_type = KernelType::Kernel1D1D,
|
|
.m = m, .n = n, .k = k, .num_groups = num_groups,
|
|
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
|
|
.cd_dtype = d.scalar_type(),
|
|
.major_a = major_a, .major_b = major_b,
|
|
.with_accumulation = false,
|
|
.num_sms = device_runtime->get_num_sms(),
|
|
.tc_util = device_runtime->get_tc_util(),
|
|
.compiled_dims = compiled_dims,
|
|
.expected_m = expected_m_for_psum_layout.value_or(m),
|
|
.expected_n = n, .expected_k = k,
|
|
.expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1
|
|
};
|
|
const auto config = get_best_config<SM100ArchSpec>(desc);
|
|
|
|
// Create tensor descriptors
|
|
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
|
|
config.storage_config.load_block_m,
|
|
config.layout.block_k,
|
|
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
|
|
config.storage_config.swizzle_a_mode);
|
|
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
|
|
config.storage_config.load_block_n,
|
|
config.layout.block_k,
|
|
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
|
|
config.storage_config.swizzle_b_mode);
|
|
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
|
|
config.storage_config.store_block_m,
|
|
config.storage_config.store_block_n,
|
|
static_cast<int>(d.stride(-2)), 1,
|
|
config.storage_config.swizzle_cd_mode);
|
|
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
|
|
config.layout.block_m, gran_k_a, 1, 0);
|
|
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
|
|
config.layout.block_n, gran_k_b, num_groups, 0);
|
|
|
|
// Launch kernel
|
|
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
|
|
.gemm_desc = desc,
|
|
.gemm_config = config,
|
|
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
|
|
config.pipeline_config.smem_size,
|
|
config.layout.get_cluster_size()),
|
|
.epilogue_type = std::nullopt,
|
|
.gran_k_a = gran_k_a,
|
|
.gran_k_b = gran_k_b,
|
|
.grouped_layout = grouped_layout.data_ptr(),
|
|
.tensor_map_a = tensor_map_a,
|
|
.tensor_map_b = tensor_map_b,
|
|
.tensor_map_sfa = tensor_map_sfa,
|
|
.tensor_map_sfb = tensor_map_sfb,
|
|
.tensor_map_cd = tensor_map_cd
|
|
};
|
|
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
|
|
const auto runtime = compiler->build("sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d", code);
|
|
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
|
|
}
|
|
|
|
static void sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
|
|
const torch::Tensor& b, const torch::Tensor& sfb,
|
|
const torch::Tensor& d,
|
|
const torch::Tensor& masked_m,
|
|
const int& num_groups, const int& m, const int& n, const int& k,
|
|
const int& expected_m,
|
|
const int& gran_k_a, const int& gran_k_b,
|
|
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
|
const std::string& compiled_dims) {
|
|
const auto desc = GemmDesc {
|
|
.gemm_type = GemmType::MGroupedMasked,
|
|
.kernel_type = KernelType::Kernel1D1D,
|
|
.m = m, .n = n, .k = k, .num_groups = num_groups,
|
|
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
|
|
.cd_dtype = d.scalar_type(),
|
|
.major_a = major_a, .major_b = major_b,
|
|
.with_accumulation = false,
|
|
.num_sms = device_runtime->get_num_sms(),
|
|
.tc_util = device_runtime->get_tc_util(),
|
|
.compiled_dims = compiled_dims,
|
|
.expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups
|
|
};
|
|
const auto config = get_best_config<SM100ArchSpec>(desc);
|
|
|
|
// Create tensor descriptors
|
|
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
|
|
config.storage_config.load_block_m,
|
|
config.layout.block_k,
|
|
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
|
|
config.storage_config.swizzle_a_mode);
|
|
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
|
|
config.storage_config.load_block_n,
|
|
config.layout.block_k,
|
|
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
|
|
config.storage_config.swizzle_b_mode);
|
|
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
|
|
config.storage_config.store_block_m,
|
|
config.storage_config.store_block_n,
|
|
static_cast<int>(d.stride(-2)), num_groups,
|
|
config.storage_config.swizzle_cd_mode);
|
|
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
|
|
config.layout.block_m, gran_k_a, num_groups, 0);
|
|
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
|
|
config.layout.block_n, gran_k_b, num_groups, 0);
|
|
|
|
// Launch kernel
|
|
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
|
|
.gemm_desc = desc,
|
|
.gemm_config = config,
|
|
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
|
|
config.pipeline_config.smem_size,
|
|
config.layout.get_cluster_size()),
|
|
.epilogue_type = std::nullopt,
|
|
.gran_k_a = gran_k_a,
|
|
.gran_k_b = gran_k_b,
|
|
.grouped_layout = masked_m.data_ptr(),
|
|
.tensor_map_a = tensor_map_a,
|
|
.tensor_map_b = tensor_map_b,
|
|
.tensor_map_sfa = tensor_map_sfa,
|
|
.tensor_map_sfb = tensor_map_sfb,
|
|
.tensor_map_cd = tensor_map_cd
|
|
};
|
|
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
|
|
const auto runtime = compiler->build("sm100_m_grouped_fp8_fp4_gemm_masked_1d1d", code);
|
|
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
|
|
}
|
|
|
|
static void sm100_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
|
|
const torch::Tensor& b, const torch::Tensor& sfb,
|
|
const std::optional<torch::Tensor>& c,
|
|
const torch::Tensor& d,
|
|
const int& m, const int& n,
|
|
const std::vector<int>& ks, const torch::Tensor& ks_tensor,
|
|
const int& gran_k,
|
|
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
|
const std::string& compiled_dims) {
|
|
DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN);
|
|
DG_HOST_ASSERT(gran_k == 32 or gran_k == 128);
|
|
const int gran_k_a = gran_k;
|
|
const int gran_k_b = gran_k;
|
|
|
|
int sum_k = 0, sum_sf_k = 0;
|
|
for (const auto k: ks) {
|
|
sum_k += k, sum_sf_k += ceil_div(k, gran_k * 4);
|
|
DG_HOST_ASSERT(k % gran_k == 0);
|
|
}
|
|
const auto num_groups = static_cast<int>(ks.size());
|
|
|
|
// Get config using max K for better performance
|
|
const auto max_k = *std::max_element(ks.begin(), ks.end());
|
|
const auto desc = GemmDesc {
|
|
.gemm_type = GemmType::KGroupedContiguous,
|
|
.kernel_type = KernelType::Kernel1D1D,
|
|
.m = m, .n = n, .k = sum_k, .num_groups = num_groups,
|
|
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
|
|
.cd_dtype = d.scalar_type(),
|
|
.major_a = major_a, .major_b = major_b,
|
|
.with_accumulation = c.has_value(),
|
|
.num_sms = device_runtime->get_num_sms(),
|
|
.tc_util = device_runtime->get_tc_util(),
|
|
.compiled_dims = compiled_dims,
|
|
.expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
|
|
};
|
|
const auto config = get_best_config<SM100ArchSpec>(desc);
|
|
|
|
// Create tensor descriptors
|
|
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
|
|
config.storage_config.load_block_m,
|
|
config.layout.block_k,
|
|
static_cast<int>(a.stride(0)), 1,
|
|
config.storage_config.swizzle_a_mode);
|
|
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
|
|
config.storage_config.load_block_n,
|
|
config.layout.block_k,
|
|
static_cast<int>(b.stride(0)), 1,
|
|
config.storage_config.swizzle_b_mode);
|
|
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
|
|
config.storage_config.store_block_m,
|
|
config.storage_config.store_block_n,
|
|
static_cast<int>(d.stride(1)), num_groups,
|
|
config.storage_config.swizzle_cd_mode);
|
|
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, sum_sf_k * gran_k_a * 4,
|
|
config.layout.block_m, gran_k_a, 1, 0);
|
|
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, sum_sf_k * gran_k_b * 4,
|
|
config.layout.block_n, gran_k_b, 1, 0);
|
|
|
|
// Launch kernel
|
|
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
|
|
.gemm_desc = desc,
|
|
.gemm_config = config,
|
|
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
|
|
config.pipeline_config.smem_size,
|
|
config.layout.get_cluster_size()),
|
|
.epilogue_type = std::nullopt,
|
|
.gran_k_a = gran_k_a,
|
|
.gran_k_b = gran_k_b,
|
|
.grouped_layout = ks_tensor.data_ptr(),
|
|
.tensor_map_a = tensor_map_a,
|
|
.tensor_map_b = tensor_map_b,
|
|
.tensor_map_sfa = tensor_map_sfa,
|
|
.tensor_map_sfb = tensor_map_sfb,
|
|
.tensor_map_cd = tensor_map_cd
|
|
};
|
|
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
|
|
const auto runtime = compiler->build("sm100_k_grouped_fp8_gemm_1d1d", code);
|
|
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
|
|
}
|
|
|
|
static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
|
|
const torch::Tensor& b, const torch::Tensor& sfb,
|
|
const std::optional<torch::Tensor>& c,
|
|
const torch::Tensor& d,
|
|
const int& batch_size, const int& m, const int& n, const int& k,
|
|
const int& gran_k_a, const int& gran_k_b,
|
|
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
|
const std::string& compiled_dims) {
|
|
const auto desc = GemmDesc {
|
|
.gemm_type = GemmType::Batched,
|
|
.kernel_type = KernelType::Kernel1D1D,
|
|
.m = m, .n = n, .k = k, .num_groups = batch_size,
|
|
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
|
|
.cd_dtype = d.scalar_type(),
|
|
.major_a = major_a, .major_b = major_b,
|
|
.with_accumulation = c.has_value(),
|
|
.num_sms = device_runtime->get_num_sms(),
|
|
.tc_util = device_runtime->get_tc_util(),
|
|
.compiled_dims = compiled_dims
|
|
};
|
|
const auto config = get_best_config<SM100ArchSpec>(desc);
|
|
|
|
const int load_block_m = config.storage_config.load_block_m;
|
|
const auto [inner_dim_a, outer_dim_a] = get_inner_outer_dims(major_a, k, m);
|
|
const auto [inner_block_a, outer_block_a] = get_inner_outer_dims(major_a, config.layout.block_k, load_block_m);
|
|
const auto tensor_map_a = make_tma_3d_desc(a, inner_dim_a, outer_dim_a, batch_size,
|
|
inner_block_a, outer_block_a, 1,
|
|
a.stride(major_a == cute::UMMA::Major::K ? 1 : 2),
|
|
a.stride(0),
|
|
config.storage_config.swizzle_a_mode);
|
|
|
|
const int load_block_n = config.storage_config.load_block_n;
|
|
const auto [inner_dim_b, outer_dim_b] = get_inner_outer_dims(major_b, k, n);
|
|
const auto [inner_block_b, outer_block_b] = get_inner_outer_dims(major_b, config.layout.block_k, load_block_n);
|
|
const auto tensor_map_b = make_tma_3d_desc(b, inner_dim_b, outer_dim_b, batch_size,
|
|
inner_block_b, outer_block_b, 1,
|
|
b.stride(major_b == cute::UMMA::Major::K ? 1 : 2),
|
|
b.stride(0),
|
|
config.storage_config.swizzle_b_mode);
|
|
|
|
const int store_block_m = config.storage_config.store_block_m;
|
|
const int store_block_n = config.storage_config.store_block_n;
|
|
const auto tensor_map_cd = make_tma_3d_desc(d, n, m, batch_size,
|
|
store_block_n, store_block_m, 1,
|
|
d.stride(1), d.stride(0),
|
|
config.storage_config.swizzle_cd_mode);
|
|
|
|
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
|
|
config.layout.block_m, gran_k_a, batch_size, 0);
|
|
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
|
|
config.layout.block_n, gran_k_b, batch_size, 0);
|
|
|
|
// Launch
|
|
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
|
|
.gemm_desc = desc,
|
|
.gemm_config = config,
|
|
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
|
|
config.pipeline_config.smem_size,
|
|
config.layout.get_cluster_size()),
|
|
.epilogue_type = std::nullopt,
|
|
.gran_k_a = gran_k_a,
|
|
.gran_k_b = gran_k_b,
|
|
.grouped_layout = nullptr,
|
|
.tensor_map_a = tensor_map_a,
|
|
.tensor_map_b = tensor_map_b,
|
|
.tensor_map_sfa = tensor_map_sfa,
|
|
.tensor_map_sfb = tensor_map_sfb,
|
|
.tensor_map_cd = tensor_map_cd
|
|
};
|
|
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
|
|
const auto runtime = compiler->build("sm100_fp8_gemm_1d1d", code);
|
|
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
|
|
}
|
|
|
|
} // namespace deep_gemm
|