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

DistGEMM bug fixes (#2713)

Browse Source 
* Blackwell DistGEMM bug fixes

1. If using preferred cluster, there needs to be a branch so that
   the universal GEMM wrapper finds the correct base params.
2. Workspace sizes can change depending on problem shape in Blackwell,
   and DistGEMM was previously using the per-device shape to evaluate
   workspace size instead of the per-gemm shape.
3. Flattened size used to initialize host tensors can overflow (in
   Hopper example as well)
4. Preferred and fallback cluster args need to be set explicitly,
   otherwise if someone modifies the example to use preferred cluster,
   it will just fail.

* Fix example runtimes

* Set default fallback cluster shapes to the static ones
This commit is contained in:
Ali Hassani
2025-11-06 13:31:24 -05:00
committed by GitHub
parent 020c700e97
commit d1ef0e87f2
4 changed files with 84 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
include/cutlass/experimental/distributed/device/dist_gemm_universal_wrapper.hpp
View File

@@ -253,16 +253,59 @@ public:
return DistSchedule::get_tensor_D(tensor_D, tensor_buffer, device_idx, iteration);
}
static
auto make_dummy_base_args(Arguments const* args, int device_idx, int iteration, void ** buffer_space) {
// Set up GEMM arguments for the current stage/iteration
auto tensor_a_iter = get_tensor_A_for_iter(args, buffer_space, device_idx, iteration);
auto tensor_b_iter = get_tensor_B_for_iter(args, buffer_space, device_idx, iteration);
auto tensor_c_iter = get_tensor_C_for_iter(args, buffer_space, device_idx, iteration);
auto tensor_d_iter = get_tensor_D_for_iter(args, buffer_space, device_idx, iteration);
Arguments base_args = args[device_idx];
base_args.problem_shape = DistSchedule::get_local_gemm_shape(args[device_idx].problem_shape);
base_args.mainloop = {
reinterpret_cast<const ElementA*>(tensor_a_iter.data()),
tensor_a_iter.stride(),
reinterpret_cast<const ElementB*>(tensor_b_iter.data()),
tensor_b_iter.stride()
};
base_args.epilogue = {
base_args.epilogue.thread,
reinterpret_cast<const ElementC*>(tensor_c_iter.data()),
tensor_c_iter.stride(),
reinterpret_cast<ElementD*>(tensor_d_iter.data()),
tensor_d_iter.stride()
};
if constexpr (DistSchedule::RemoteC) {
if (iteration > 0) {
base_args.epilogue.thread.beta = 1.0;
}
else if (iteration == 0){
base_args.epilogue.thread.beta = 0.0;
}
}
return base_args;
}
static size_t
get_workspace_size(Arguments const& args) {
get_workspace_size(Arguments const* args, int device_idx) {
size_t workspace_bytes = 0;
workspace_bytes = get_buffer_space_size(args);
workspace_bytes = get_buffer_space_size(args[device_idx]);
void* dummy_buffer_space[TP_];
for (int iteration = 0; iteration < TP_; ++iteration) {
// Workspace sizes can vary if arguments change, therefore we must
// construct args for each iteration exactly as it will be run.
auto args_base = make_dummy_base_args(args, device_idx, iteration, dummy_buffer_space);
// NOTE: assumes underlying kernels align up to alignment requirements on their own,
// and that the alignment requirements of the individual kernels match.
workspace_bytes += GemmKernel::get_workspace_size(args);
workspace_bytes += GemmKernel::get_workspace_size(args_base);
}
return workspace_bytes;

16
include/cutlass/gemm/device/gemm_universal_adapter.h
View File

@@ -110,6 +110,13 @@ constexpr int stages_member(DispatchPolicy) {
}
}
template <class GemmKernel, class = void>
struct IsDistGemmKernel : cute::false_type { };
template <typename GemmKernel>
struct IsDistGemmKernel<GemmKernel, cute::void_t<typename GemmKernel::TP>>
: cute::true_type { };
} // namespace detail
template <class GemmKernel_>
@@ -396,8 +403,13 @@ public:
|| GemmKernel::ArchTag::kMinComputeCapability == 103
) {
if constexpr (!cute::is_static_v<typename GemmKernel::DispatchPolicy::ClusterShape>) {
fallback_cluster = params.hw_info.cluster_shape_fallback;
cluster = params.hw_info.cluster_shape;
if constexpr (detail::IsDistGemmKernel<GemmKernel>::value) {
fallback_cluster = params.base.hw_info.cluster_shape_fallback;
cluster = params.base.hw_info.cluster_shape;
} else {
fallback_cluster = params.hw_info.cluster_shape_fallback;
cluster = params.hw_info.cluster_shape;
}
}
}