class GemmRsAr:
"""Own the symmetric workspace for BF16 GEMM-RS or GEMM-AR launches.
All TP ranks must belong to one NVLink domain for multimem instructions.
"""
def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None:
tp_group = get_tp_group()
group = tp_group.device_group
rank = tp_group.rank_in_group
world_size = tp_group.world_size
device = torch.device("cuda", torch.accelerator.current_device_index())
assert 1 < world_size <= 16
assert 128 % world_size == 0
assert max_M >= 128 and N % 128 == 0
max_M = (max_M + world_size - 1) // world_size * world_size
self.rank = rank
self.world_size = world_size
self.max_M = max_M
self.N = N
self.device = device
self.all_reduce = all_reduce
self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device)
self.partial_handle = symm_mem.rendezvous(self.partial, group)
if self.partial_handle.multicast_ptr == 0:
raise RuntimeError("GEMM-RS/AR requires NVLink multicast memory")
self.partial_mc_ptr = make_ptr(
BFloat16,
self.partial_handle.multicast_ptr,
cute.AddressSpace.gmem,
assumed_align=32,
)
grid_m = (max_M + 127) // 128
cta_group = 2 if max_M >= 1024 or grid_m % 2 == 0 else 1
grid_m = (grid_m + cta_group - 1) // cta_group * cta_group
self.num_sms = torch.cuda.get_device_properties(device).multi_processor_count
max_flags = grid_m * (N // 128) + self.num_sms
self.flags = symm_mem.empty(max_flags, dtype=torch.int32, device=device)
self.flags_handle = symm_mem.rendezvous(self.flags, group)
if self.flags_handle.multicast_ptr == 0:
raise RuntimeError("GEMM-RS/AR requires NVLink multicast memory")
self.flags.zero_()
self.flags_mc_ptr = make_ptr(
Int32,
self.flags_handle.multicast_ptr,
cute.AddressSpace.gmem,
assumed_align=16,
)
self.peer_flag_ptr = make_ptr(
Int64,
self.flags_handle.buffer_ptrs_dev,
cute.AddressSpace.gmem,
assumed_align=8,
)
torch.accelerator.synchronize(device)
tp_group.barrier()
def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
assert x.ndim == 2
M, K = x.shape
assert 0 < M <= self.max_M
assert w.shape == (self.N, K) and K % 64 == 0
assert w.dtype == torch.bfloat16
assert w.device == self.device
assert w.is_contiguous()
assert x.dtype == torch.bfloat16
assert x.device == self.device
assert x.is_contiguous()
N = w.shape[0]
padded_M = (M + self.world_size - 1) // self.world_size
padded_M *= self.world_size
local_M = padded_M // self.world_size
grid_m = (M + 127) // 128
# Avoid padding small odd grids; 2-CTA wins consistently for M >= 1024.
cta_group = 2 if M >= 1024 or grid_m % 2 == 0 else 1
grid_m = (grid_m + cta_group - 1) // cta_group * cta_group
BN = 256 if M * K >= 24 * 1024 * 1024 else 128
assert N % BN == 0
num_tiles = grid_m * (N // BN)
num_ctas = min(num_tiles, self.num_sms)
num_ctas = num_ctas // cta_group * cta_group
assert self.flags.numel() >= num_tiles + num_ctas
output = None
if not self.all_reduce:
output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device)
compiled = Sm100GemmRsArBF16.compile(
self.rank,
self.world_size,
BN,
cta_group,
self.all_reduce,
)
compiled(
x,
w,
self.partial[:padded_M],
self.partial_mc_ptr,
output,
self.flags,
self.flags_mc_ptr,
self.peer_flag_ptr,
num_ctas,
)
if self.all_reduce:
# AttnRes may retain output past the next workspace reuse.
return self.partial[:M].clone()
assert output is not None
return output