class BeamAllReduceWorkspace:
"""Multicast Lamport buffers and generation flags for the Beam kernels."""
@classmethod
def create(
cls,
group: dist.ProcessGroup,
hidden_size: int,
max_tokens: int,
dtype: torch.dtype,
) -> "BeamAllReduceWorkspace | None":
# The CUDA epilogue handles one packed load per thread.
if (
group.size() not in (2, 4, 8)
or dtype != torch.bfloat16
or hidden_size % 1024
or not 1024 <= hidden_size <= 8192
or not hasattr(torch.ops._C, "beam_allreduce_norm")
or not has_flashinfer()
):
return None
from flashinfer.comm.mnnvl import (
TorchDistBackend,
all_ranks_agree,
is_multicast_supported,
)
device = torch.accelerator.current_device_index()
supported = torch.cuda.get_device_capability(device)[0] in (9, 10)
supported = supported and is_multicast_supported(device)
if not all_ranks_agree(TorchDistBackend(group=group), supported):
return None
return cls(group, hidden_size, max_tokens)
def __init__(
self, group: dist.ProcessGroup, hidden_size: int, max_tokens: int
) -> None:
self.world_size = group.size()
self.rank = dist.get_rank(group)
self.hidden_size = hidden_size
self.max_tokens = max_tokens
# Scale one-shot cutoffs by the BF16 payload staged across all ranks.
token_bytes = hidden_size * 2
rank_token_bytes = self.world_size * token_bytes
self.oneshot_max_tokens = _ONESHOT_MAX_BYTES // rank_token_bytes
self.small_batch_max_tokens = 3200 * 1024 // rank_token_bytes
self.separate_norm_min_tokens = cdiv(
10 * 1024**2 * self.world_size, token_bytes
)
# One-shot stages one row per rank; two-shot two rows per owned token.
oneshot_tokens = min(max_tokens, self.small_batch_max_tokens)
buffer_bytes = max(
oneshot_tokens * rank_token_bytes,
2 * cdiv(max_tokens, self.world_size) * rank_token_bytes,
)
# Both two-shot stages stay 16-byte aligned.
buffer_bytes = round_up(buffer_bytes, 32)
if buffer_bytes >= 2**32:
raise ValueError(
f"Beam all-reduce needs {buffer_bytes} bytes per Lamport "
"buffer, which exceeds the 32-bit flag field"
)
# FlashInfer only maps the multicast memory; the Lamport protocol and
# its flags are ours. Small batches reduce ~0.3us faster on its mapping
# than on torch symmetric memory.
from flashinfer.comm.mnnvl import McastGPUBuffer, TorchDistBackend
device = torch.device("cuda", torch.accelerator.current_device_index())
self.buffer = McastGPUBuffer(
_NUM_LAMPORT_BUFFERS * buffer_bytes,
self.world_size,
self.rank,
device,
TorchDistBackend(group=group),
)
self.multicast_ptr = self.buffer.get_multicast_ptr()
if not self.multicast_ptr:
raise RuntimeError("Beam all-reduce requires NVLink multicast")
self.local_ptr = self.buffer.get_unicast_ptr(self.rank)
self.buffer_ptrs = self.buffer.get_buffer_ptrs_dev()
self.buffer.lamport_initialize(self.rank, torch.float32)
# [current, dirty, bytes per buffer, dirty stages, bytes to clear x4,
# arrival count]
self.flags = torch.tensor(
[0, 2, buffer_bytes, 0, 0, 0, 0, 0, 0],
dtype=torch.uint32,
device=device,
)
torch.accelerator.synchronize()
dist.barrier(group=group)