diff --git a/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu b/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu index 10193826d..724537e51 100644 --- a/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu +++ b/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu @@ -23,6 +23,7 @@ // The per-group context is an ncclDevComm plus a registered symmetric window, // both created here from the caller's ncclComm_t. +#include #include #include #include @@ -162,13 +163,19 @@ bool ulysses_lsa_covers_group(int64_t comm_ptr, int64_t world_size) { // mode == 0: inp [B, S_local, H, D] -> out [B, S_global, H_local, D] // mode == 1: inp [B, S_global, H_local, D] -> out [B, S_local, H, D] // where H is the *global* head count and H_local = H / world_size. -void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B, int64_t S_local, - int64_t H, int64_t D, int64_t mode) { +void ulysses_a2a_tuned(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B, int64_t S_local, + int64_t H, int64_t D, int64_t mode, int64_t launch_blocks) { auto* ctx = reinterpret_cast(handle); TORCH_CHECK(ctx != nullptr, "handle must come from allocate_ulysses_a2a"); const at::cuda::CUDAGuard device_guard(inp.device()); auto stream = at::cuda::getCurrentCUDAStream(); + TORCH_CHECK(launch_blocks == 36 || launch_blocks == 144, "Ulysses launch must use 36 or 144 CTAs"); + if (launch_blocks == 144) { + const auto* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 10 && props->minor == 0 && props->multiProcessorCount >= 144, + "144-CTA Ulysses launch requires an sm100 device with at least 144 SMs"); + } TORCH_CHECK(inp.is_cuda() && out.is_cuda(), "inp and out must be CUDA tensors"); TORCH_CHECK(inp.is_contiguous() && out.is_contiguous(), "inp and out must be contiguous"); @@ -204,7 +211,7 @@ void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B const int64_t num_rows = B * static_cast(W) * S_local; const int blocks = - static_cast(std::max(1, std::min(fi::kMaxBlocks, num_rows))); + static_cast(std::max(1, std::min(launch_blocks, num_rows))); const int threads = fi::kUlyssesThreads; #define LAUNCH_ULYSSES_A2A(T, NG, MODE) \ @@ -266,6 +273,12 @@ void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B TORCH_CHECK(status == cudaSuccess, "ulysses_a2a copy-out failed: ", cudaGetErrorString(status)); } +// Preserve the original wheel API and its 36-CTA launch policy. +void ulysses_a2a(int64_t handle, torch::Tensor inp, torch::Tensor out, int64_t B, int64_t S_local, + int64_t H, int64_t D, int64_t mode) { + ulysses_a2a_tuned(handle, inp, out, B, S_local, H, D, mode, 36); +} + void register_ulysses_a2a(pybind11::module_& m) { m.def("allocate_ulysses_a2a", &allocate_ulysses_a2a, "allocate a local ulysses a2a window"); m.def("register_ulysses_a2a_window", ®ister_ulysses_a2a_window, @@ -276,4 +289,5 @@ void register_ulysses_a2a(pybind11::module_& m) { m.def("ulysses_lsa_covers_group", &ulysses_lsa_covers_group, "whether the whole group is load-store accessible"); m.def("ulysses_a2a", &ulysses_a2a, "fused-transpose Ulysses all-to-all over NVLink"); + m.def("ulysses_a2a_tuned", &ulysses_a2a_tuned, "Ulysses all-to-all with a collectively agreed CTA count"); } diff --git a/fastvideo-kernel/include/comm/ulysses_all_to_all.cuh b/fastvideo-kernel/include/comm/ulysses_all_to_all.cuh index f69dc6035..51fba3076 100644 --- a/fastvideo-kernel/include/comm/ulysses_all_to_all.cuh +++ b/fastvideo-kernel/include/comm/ulysses_all_to_all.cuh @@ -50,7 +50,8 @@ namespace ulysses { constexpr int kUlyssesThreads = 512; // Deliberately modest: this is link-bandwidth bound, so a small grid leaves the // rest of the GPU free without costing throughput. -constexpr int kMaxBlocks = 36; +// One distinct NCCL barrier slot per CTA, including the tuned GB200 launch. +constexpr int kMaxBlocks = 144; // Shared movement body for the fused-transpose all-to-all (no barriers). // diff --git a/fastvideo-kernel/python/fastvideo_kernel/comm_ops.py b/fastvideo-kernel/python/fastvideo_kernel/comm_ops.py index cc5067303..168376e70 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/comm_ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/comm_ops.py @@ -29,6 +29,11 @@ def is_available() -> bool: return _ops is not None and all(hasattr(_ops, name) for name in _REQUIRED_OPS) +def supports_tuned_launch() -> bool: + """Whether this build has the expanded barrier capacity and tuned entrypoint.""" + return is_available() and hasattr(_ops, "ulysses_a2a_tuned") + + def _require() -> None: if not is_available(): raise RuntimeError( @@ -73,7 +78,7 @@ def dispose(handle: int) -> None: def all_to_all(handle: int, inp: torch.Tensor, out: torch.Tensor, B: int, S_local: int, H: int, - D: int, mode: int) -> None: + D: int, mode: int, *, blocks: int = 36) -> None: """Run one fused all-to-all on the current stream, writing into ``out``. ``mode == 0``: ``[B, S_local, H, D] -> [B, S_global, H_local, D]`` @@ -83,4 +88,9 @@ def all_to_all(handle: int, inp: torch.Tensor, out: torch.Tensor, B: int, S_loca geometry in the same order. """ _require() - _ops.ulysses_a2a(int(handle), inp, out, int(B), int(S_local), int(H), int(D), int(mode)) + if blocks == 36: + _ops.ulysses_a2a(int(handle), inp, out, int(B), int(S_local), int(H), int(D), int(mode)) + elif blocks == 144 and supports_tuned_launch(): + _ops.ulysses_a2a_tuned(int(handle), inp, out, int(B), int(S_local), int(H), int(D), int(mode), int(blocks)) + else: + raise ValueError(f"unsupported Ulysses launch: blocks={blocks}; rebuild the kernel for tuned launches") diff --git a/fastvideo/distributed/device_communicators/ulysses_a2a.py b/fastvideo/distributed/device_communicators/ulysses_a2a.py index c811049ce..46b125a5c 100644 --- a/fastvideo/distributed/device_communicators/ulysses_a2a.py +++ b/fastvideo/distributed/device_communicators/ulysses_a2a.py @@ -30,6 +30,8 @@ _DTYPE_CODES = { # Bound persistent registered memory per rank. Larger operands use NCCL instead # of growing the window without limit. MAX_WINDOW_BYTES = 1024**3 +H3_TRAINING_PLANE_LIMIT_BYTES = 512 * 1024**2 +_CONTRACT_SIZE = 12 # (scatter_dim, gather_dim) -> kernel mode. # 0: [B, S_local, H, D] -> [B, S_global, H_local, D] @@ -50,17 +52,22 @@ class _FusedUlyssesA2A(torch.autograd.Function): """ @staticmethod - def forward(ctx, helper: "UlyssesA2AHelper", x: torch.Tensor, mode: int) -> torch.Tensor: # type: ignore[override] + def forward(ctx, helper: "UlyssesA2AHelper", x: torch.Tensor, mode: int, chunked: bool, + blocks: int) -> torch.Tensor: # type: ignore[override] ctx.helper = helper ctx.mode = mode - return helper.run_armed(x, mode) + ctx.chunked = chunked + ctx.blocks = blocks + return helper.run_armed(x, mode, chunked, blocks) @staticmethod def backward(ctx, grad_output: torch.Tensor): # type: ignore[override] # Same numel and dtype as the forward output, so the window is already # sized for it; only contiguity needs restoring. - grad_input = ctx.helper.run_armed(grad_output.contiguous(), 1 - ctx.mode) - return None, grad_input, None + # Reuse the forward plan, even if later calls chose another plan or + # backward executes with grad tracking disabled. + grad_input = ctx.helper.run_armed(grad_output.contiguous(), 1 - ctx.mode, ctx.chunked, ctx.blocks) + return None, grad_input, None, None, None class UlyssesA2AHelper: @@ -81,11 +88,12 @@ class UlyssesA2AHelper: self._handle: int | None = None # Reuse storage, but exchange the current contract on every call. A # rank-local cache hit cannot establish what peers are doing now. - self._local_contract = array("q", [0] * 10) + self._local_contract = array("q", [0] * _CONTRACT_SIZE) self._local_tensor = torch.frombuffer(self._local_contract, dtype=torch.int64) - self._gathered_tensor = torch.empty(world_size * 10, dtype=torch.int64, device="cpu") + self._gathered_tensor = torch.empty(world_size * _CONTRACT_SIZE, dtype=torch.int64, device="cpu") self._nbytes = 0 self._disabled_reason: str | None = None + self._h3_tuning_available: bool | None = None if world_size not in SUPPORTED_WORLD_SIZES: self._disabled_reason = (f"world size {world_size} is not one of " @@ -166,6 +174,20 @@ class UlyssesA2AHelper: logger.warning("Ulysses partial-context cleanup failed", exc_info=True) return self._agree(cleanup_ok) + def _execution_plan(self, x: torch.Tensor, mode: int) -> tuple[bool, int]: + """Use the measured GB200/SP4 launch only for H3's bf16 head geometry.""" + global_heads = x.shape[2] if mode == 0 else x.shape[2] * self.world_size + if self.world_size != 4 or x.dtype != torch.bfloat16 or global_heads != 56 or x.shape[3] != 128: + return False, 36 + if self._h3_tuning_available is None: + from fastvideo_kernel import comm_ops + + props = torch.cuda.get_device_properties(self.device) + self._h3_tuning_available = ("GB200" in props.name and props.major == 10 and props.minor == 0 + and props.multi_processor_count >= 144 + and getattr(comm_ops, "supports_tuned_launch", lambda: False)()) + return (True, 144) if self._h3_tuning_available else (False, 36) + def _call_signature(self, x: torch.Tensor, scatter_dim: int, gather_dim: int) -> tuple[tuple[int, ...], str]: """Return a rank-comparable call contract and any local decline reason.""" mode = _MODE_FROM_DIMS.get((scatter_dim, gather_dim)) @@ -196,17 +218,26 @@ class UlyssesA2AHelper: status, reason = 0, "sequence length is not divisible by the group" nbytes = int(x.numel() * x.element_size()) + chunked, blocks = False, 36 + if status == 1 and nbytes: + assert mode is not None + chunked, blocks = self._execution_plan(x, mode) + window_bytes = nbytes // shape[0] if chunked else nbytes if status == 1 and nbytes == 0: status, reason = 0, "input is empty" - elif status == 1 and nbytes > MAX_WINDOW_BYTES: - status, reason = 0, f"operand exceeds the {MAX_WINDOW_BYTES}-byte window cap" + elif status == 1 and window_bytes > MAX_WINDOW_BYTES: + status, reason = 0, f"operand window exceeds the {MAX_WINDOW_BYTES}-byte cap" + elif (status == 1 and chunked and shape[0] > 1 and torch.is_grad_enabled() and x.requires_grad + and window_bytes > H3_TRAINING_PLANE_LIMIT_BYTES): + status, reason = 0, "large packed H3 training transfer uses NCCL" - # status, armed, mode, dtype, B, S, H, D, bytes, capacity. Comparing the + # status, armed, mode, dtype, B, S, H, D, window bytes, capacity, + # chunked, blocks. Comparing the # whole vector prevents equal-size but differently-shaped ranks from # entering the fused kernel with incompatible address math. CUDA device # ordinals are deliberately absent: rank-local ordinals normally differ. - signature = (status, int(self._handle is not None), -1 if mode is None else mode, dtype_code, *shape, nbytes, - self._nbytes) + signature = (status, int(self._handle is not None), -1 if mode is None else mode, dtype_code, *shape, + window_bytes, self._nbytes, int(chunked), blocks) return signature, reason def _agree_call(self, signature: tuple[int, ...]) -> tuple[bool, bool, bool]: @@ -221,11 +252,11 @@ class UlyssesA2AHelper: self._local_contract[:] = array("q", signature) dist.all_gather_into_tensor(self._gathered_tensor, self._local_tensor, group=self.cpu_group) values = self._gathered_tensor.tolist() - contracts = [values[start:start + 10] for start in range(0, len(values), 10)] + contracts = [values[start:start + _CONTRACT_SIZE] for start in range(0, len(values), _CONTRACT_SIZE)] first = contracts[0] use_fused = first[0] == 1 and all(contract == first for contract in contracts) permanently_unavailable = any(contract[0] < 0 for contract in contracts) - lifecycle_consistent = all(contract[1] == first[1] and contract[-1] == first[-1] for contract in contracts) + lifecycle_consistent = all(contract[1] == first[1] and contract[9] == first[9] for contract in contracts) return use_fused, permanently_unavailable, lifecycle_consistent def _build(self, nbytes: int) -> bool: @@ -333,7 +364,7 @@ class UlyssesA2AHelper: # -- collective ---------------------------------------------------------- - def run_armed(self, x: torch.Tensor, mode: int) -> torch.Tensor: + def run_armed(self, x: torch.Tensor, mode: int, chunked: bool, blocks: int) -> torch.Tensor: """Run one collective on an already-armed context.""" assert self._handle is not None, "run_armed called on an unarmed helper" from fastvideo_kernel import comm_ops @@ -346,7 +377,22 @@ class UlyssesA2AHelper: B, S_global, H_local, D = x.shape S_local, H = S_global // w, H_local * w out = torch.empty(B, S_local, H, D, dtype=x.dtype, device=x.device) - comm_ops.all_to_all(self._handle, x, out, B, S_local, H, D, mode) + if chunked: + # Each copy completes on this stream before the next plane reuses + # the registered window. The full result owns its storage; saved + # activations never alias the reusable window. + for plane in range(B): + comm_ops.all_to_all(self._handle, + x[plane:plane + 1], + out[plane:plane + 1], + 1, + S_local, + H, + D, + mode, + blocks=blocks) + else: + comm_ops.all_to_all(self._handle, x, out, B, S_local, H, D, mode) return out def try_all_to_all_4D(self, x: torch.Tensor, scatter_dim: int, gather_dim: int) -> torch.Tensor | None: @@ -371,7 +417,7 @@ class UlyssesA2AHelper: return None mode = signature[2] - nbytes = signature[-2] + nbytes = signature[8] if self._handle is None: if not self._build(nbytes): return None @@ -382,7 +428,7 @@ class UlyssesA2AHelper: if not self._build(nbytes): return None - return _FusedUlyssesA2A.apply(self, x, mode) + return _FusedUlyssesA2A.apply(self, x, mode, bool(signature[10]), signature[11]) def maybe_create_helper(cpu_group: ProcessGroup | None, device_group: ProcessGroup | None, world_size: int, diff --git a/fastvideo/tests/distributed/check_ulysses_h3_native.py b/fastvideo/tests/distributed/check_ulysses_h3_native.py new file mode 100644 index 000000000..7f61e6ebc --- /dev/null +++ b/fastvideo/tests/distributed/check_ulysses_h3_native.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Explicit torchrun GPU gate for the native H3 transport policy (4x GB200).""" +import os + +import torch +import torch.distributed as dist + +from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel +from fastvideo.distributed.device_communicators.base_device_communicator import DeviceCommunicatorBase +from fastvideo.distributed.parallel_state import get_sp_group +from fastvideo.tests.distributed.test_ulysses_a2a_parity import _check_shape + + +def main(): + rank = int(os.environ['RANK']) + assert int(os.environ['WORLD_SIZE']) == 4 + torch.cuda.set_device(int(os.environ['LOCAL_RANK'])) + device = torch.device('cuda', int(os.environ['LOCAL_RANK'])) + torch.manual_seed(20260906 + rank) + maybe_init_distributed_environment_and_model_parallel(1, 4) + comm = get_sp_group().device_communicator + helper = comm.ulysses_a2a + assert helper is not None + for dtype in [torch.bfloat16, torch.float16, torch.float32]: + _check_shape(3, 64, 56, 128, dtype, 4, device) + for sequence in [32000, 128000, 250000]: + _check_shape(4, sequence // 4, 56, 128, torch.bfloat16, 4, device) + assert helper._nbytes <= 1024**3 + x = torch.randn(3, 8000, 56, 128, dtype=torch.bfloat16, device=device, requires_grad=True) + first = comm.all_to_all_4D(x, 2, 1) + saved = first.detach().clone() + # A later call and inverse cannot overwrite an earlier autograd output. + second = comm.all_to_all_4D(x.detach() + 1, 2, 1) + comm.all_to_all_4D(second, 1, 2) + assert torch.equal(first, saved) + grad = torch.randn_like(first) + first.backward(grad) + expected = DeviceCommunicatorBase.all_to_all_4D(comm, grad, 1, 2) + assert torch.equal(x.grad, expected) + + # A rank with an older launch capability must make the whole group decline. + helper._h3_tuning_available = rank != 0 + assert helper.try_all_to_all_4D(x.detach(), 2, 1) is None + helper._h3_tuning_available = True + assert helper.try_all_to_all_4D(x.detach(), 2, 1) is not None + contract, _ = helper._call_signature(x.detach(), 2, 1) + assert contract[10:] == (1, 144) + dist.barrier() + if rank == 0: + print('NATIVE_H3_OK dtypes=3 lengths=32000,128000,250000 exact_gradients=True ownership=True recovery=True', + flush=True) + cleanup_dist_env_and_memory() + + +if __name__ == '__main__': + main() diff --git a/fastvideo/tests/distributed/test_ulysses_h3_policy.py b/fastvideo/tests/distributed/test_ulysses_h3_policy.py new file mode 100644 index 000000000..e2ba866e8 --- /dev/null +++ b/fastvideo/tests/distributed/test_ulysses_h3_policy.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Policy, rank agreement and saved backward-plan regressions without a GPU.""" +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from fastvideo.distributed.device_communicators import ulysses_a2a as ulysses + + +def _helper(): + return ulysses.UlyssesA2AHelper(object(), object(), 4, torch.device('cuda:0'), object()) + + +def _operand(planes, global_sequence, mode, requires_grad): + shape = (planes, global_sequence // 4, 56, 128) if mode == 0 else (planes, global_sequence, 14, 128) + x = Mock(shape=shape, dtype=torch.bfloat16, device=torch.device('cuda:0'), is_cuda=True, + requires_grad=requires_grad) + x.dim.return_value = 4 + x.numel.return_value = planes * global_sequence * 14 * 128 + x.element_size.return_value = 2 + x.is_contiguous.return_value = True + return x + + +@pytest.mark.parametrize('sequence,large_pack', [(32000, False), (64000, False), (128000, False), + (149796, False), (149800, True), (250000, True)]) +@pytest.mark.parametrize('planes,mode,training', [(3, 0, False), (4, 0, True), (1, 1, True), (4, 1, True)]) +def test_h3_policy_caps_each_plane_and_retains_fast_gathers(sequence, large_pack, planes, mode, training): + helper = _helper() + x = _operand(planes, sequence, mode, training) + with patch.object(helper, '_execution_plan', return_value=(True, 144)), \ + patch.object(ulysses, 'is_enabled', return_value=True), \ + patch('torch.cuda.is_current_stream_capturing', return_value=False), torch.set_grad_enabled(training): + signature, _ = helper._call_signature(x, *((2, 1) if mode == 0 else (1, 2))) + declined = training and planes > 1 and large_pack + assert signature[0] == (0 if declined else 1) + assert signature[8] == sequence * 14 * 128 * 2 + assert signature[8] <= ulysses.MAX_WINDOW_BYTES + assert signature[10:] == (1, 144) + + +def test_older_wheel_keeps_original_launch_and_capacity_policy(): + helper = _helper() + x = _operand(3, 250000, 0, False) + from fastvideo_kernel import comm_ops + props = SimpleNamespace(name='NVIDIA GB200', major=10, minor=0, multi_processor_count=152) + with patch('torch.cuda.get_device_properties', return_value=props), \ + patch.object(comm_ops, 'supports_tuned_launch', return_value=False, create=True), \ + patch.object(ulysses, 'is_enabled', return_value=True), \ + patch('torch.cuda.is_current_stream_capturing', return_value=False): + signature, _ = helper._call_signature(x, 2, 1) + assert signature[0] == 0 + assert signature[8] == x.numel() * 2 + assert signature[10:] == (0, 36) + + +def test_rank_disagreement_on_launch_or_chunking_declines_and_recovers(): + helper = _helper() + signature = (1, 0, 0, 2, 4, 32000, 56, 128, 458752000, 0, 1, 144) + for changed_index, value in [(10, 0), (11, 36), (0, 0)]: + peer = list(signature) + peer[changed_index] = value + + def gather(output, local, **kwargs): + output.copy_(torch.tensor([*signature, *peer, *signature, *signature])) + + with patch.object(ulysses.dist, 'all_gather_into_tensor', side_effect=gather): + assert helper._agree_call(signature) == (False, False, True) + + def unanimous(output, local, **kwargs): + output.copy_(local.repeat(4)) + + with patch.object(ulysses.dist, 'all_gather_into_tensor', side_effect=unanimous): + assert helper._agree_call(signature) == (True, False, True) + + +def test_backward_uses_its_own_forward_plan_after_another_call(): + calls = [] + + class Helper: + def run_armed(self, x, mode, chunked, blocks): + calls.append((mode, chunked, blocks)) + return x.clone() + + helper = Helper() + x = torch.randn(4, requires_grad=True) + first = ulysses._FusedUlyssesA2A.apply(helper, x, 0, True, 144) + second = ulysses._FusedUlyssesA2A.apply(helper, x, 1, False, 36) + first.sum().backward() + assert calls == [(0, True, 144), (1, False, 36), (1, True, 144)] + assert torch.equal(first, second) + assert first.data_ptr() != second.data_ptr() + assert torch.equal(x.grad, torch.ones_like(x)) diff --git a/fastvideo/tests/distributed/ulysses_native_test_support.py b/fastvideo/tests/distributed/ulysses_native_test_support.py new file mode 100644 index 000000000..3850e54ae --- /dev/null +++ b/fastvideo/tests/distributed/ulysses_native_test_support.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Load a separately built communication extension for development validation.""" +import importlib.util +import runpy +import sys +from pathlib import Path + +import torch # noqa: F401 - load the extension's libtorch dependencies first + + +def load_module(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(str(path)) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def install_native_ops(source: Path, extension: Path): + import fastvideo_kernel + + native = load_module('ulysses_native', extension) + wrapper = load_module('fastvideo_kernel.comm_ops', + source / 'fastvideo-kernel/python/fastvideo_kernel/comm_ops.py') + wrapper._ops = native + fastvideo_kernel.comm_ops = wrapper + return wrapper + + +if __name__ == '__main__': + source, extension, target = map(Path, sys.argv[1:4]) + install_native_ops(source, extension) + sys.argv = [str(target), *sys.argv[4:]] + runpy.run_path(str(target), run_name='__main__')