| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for MNNVL AllToAll operations. |
| 4 | |
| 5 | Requires: docker run ... --cap-add=SYS_PTRACE ... |
| 6 | Run: pytest tests/distributed/test_mnnvl_alltoall.py -v |
| 7 | """ |
| 8 | |
| 9 | import os |
| 10 | import traceback |
| 11 | from contextlib import contextmanager |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | import torch.multiprocessing as mp |
| 16 | |
| 17 | from vllm.distributed import get_ep_group |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.utils.flashinfer import ( |
| 20 | has_flashinfer_nvlink_one_sided, |
| 21 | has_flashinfer_nvlink_two_sided, |
| 22 | ) |
| 23 | from vllm.utils.import_utils import has_deep_ep_v2 |
| 24 | from vllm.utils.network_utils import get_open_port |
| 25 | |
| 26 | from ..utils import init_test_distributed_environment |
| 27 | |
| 28 | DEVICE = current_platform.device_type |
| 29 | |
| 30 | # --------------------------------------------------------------------------- |
| 31 | # Helpers |
| 32 | # --------------------------------------------------------------------------- |
| 33 | |
| 34 | |
| 35 | def _has_sys_ptrace() -> bool: |
| 36 | """Check for SYS_PTRACE capability (bit 19 in CapEff).""" |
| 37 | try: |
| 38 | with open("/proc/self/status") as f: |
| 39 | for line in f: |
| 40 | if line.startswith("CapEff:"): |
| 41 | return bool(int(line.split()[1], 16) & (1 << 19)) |
| 42 | except Exception: |
| 43 | pass |
| 44 | return False |
| 45 | |
| 46 | |
| 47 | def _spawn_workers(worker_fn, world_size, *, dp_size=None): |
| 48 | """Spawn one process per GPU, run worker_fn, assert all succeed. |
| 49 | |
| 50 | Uses an mp.Queue to propagate worker tracebacks back to the parent |
| 51 | so pytest shows the actual failure, not just an exit code. |
| 52 | """ |
| 53 | if mp.get_start_method(allow_none=True) is None: |
| 54 | mp.set_start_method("spawn") |
| 55 | |
| 56 | port = str(get_open_port()) |
| 57 | # Allocate a second port for DP master when dp_size is set, so the |
| 58 | # distributed init port and DP port can't collide even under xdist. |
| 59 | dp_port = str(get_open_port()) if dp_size is not None else None |
| 60 | err_queue: mp.Queue = mp.Queue() |
| 61 | procs = [] |
| 62 | for rank in range(world_size): |
| 63 | p = mp.Process( |
| 64 | target=_run_worker, |
| 65 | args=(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue), |
| 66 | ) |
| 67 | p.start() |
| 68 | procs.append(p) |
| 69 | for p in procs: |
| 70 | p.join() |
| 71 | |
| 72 | # Collect any errors from workers before asserting. |
| 73 | errors = [] |
| 74 | while not err_queue.empty(): |
| 75 | errors.append(err_queue.get_nowait()) |
| 76 | err_queue.close() |
| 77 | err_queue.join_thread() |
| 78 | if errors: |
| 79 | combined = "\n---\n".join(errors) |
| 80 | if "NCCL GIN" in combined: |
| 81 | pytest.skip("NCCL GIN not available on this system") |
| 82 | pytest.fail("Worker(s) failed:\n" + combined) |
| 83 | |
| 84 | |
| 85 | def _run_worker(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue): |
| 86 | """Per-process setup: device, distributed env, then call worker_fn. |
| 87 | |
| 88 | Args: |
| 89 | dp_size: If set, initialize with tp=1 and data_parallel_size=dp_size. |
| 90 | Otherwise use tp=world_size (default for EP-based tests). |
| 91 | dp_port: Separate port for the DP master (only used when dp_size is set). |
| 92 | err_queue: Queue for propagating tracebacks to the parent process. |
| 93 | |
| 94 | """ |
| 95 | try: |
| 96 | os.environ.pop("CUDA_VISIBLE_DEVICES", None) |
| 97 | torch.accelerator.set_device_index(rank) |
| 98 | if dp_size is not None: |
| 99 | _init_dp_environment(world_size, rank, port, dp_size, dp_port) |
| 100 | else: |
| 101 | init_test_distributed_environment(world_size, 1, rank, port) |
| 102 | worker_fn(rank, world_size) |
| 103 | torch.distributed.barrier() |
| 104 | except Exception: |
| 105 | err_queue.put(f"[Rank {rank}]\n{traceback.format_exc()}") |
| 106 | # Don't re-raise: the parent reads errors from err_queue. |
| 107 | # A non-zero exit from the re-raise would be redundant. |
| 108 | import sys |
| 109 | |
| 110 | sys.exit(1) |
| 111 | |
| 112 | |
| 113 | def _init_dp_environment(world_size, rank, port, dp_size, dp_port): |
| 114 | """Initialize distributed env with data parallelism. |
| 115 | |
| 116 | Sets up tp=1, pp=1, dp=dp_size. Each process is one DP rank |
| 117 | with local rank 0 within its (trivial) tp*pp group. |
| 118 | |
| 119 | Args: |
| 120 | port: Port for torch.distributed init. |
| 121 | dp_port: Separate port for the DP master group init. |
| 122 | |
| 123 | """ |
| 124 | from vllm.config import VllmConfig, set_current_vllm_config |
| 125 | from vllm.config.parallel import ParallelConfig |
| 126 | from vllm.distributed.parallel_state import ( |
| 127 | ensure_model_parallel_initialized, |
| 128 | init_distributed_environment, |
| 129 | ) |
| 130 | |
| 131 | vllm_config = VllmConfig() |
| 132 | vllm_config.parallel_config = ParallelConfig( |
| 133 | data_parallel_size=dp_size, |
| 134 | data_parallel_rank=rank, |
| 135 | # Pre-populate port list so __post_init__ doesn't auto-generate |
| 136 | # random ports. All DP ranks must agree on the same port. |
| 137 | _data_parallel_master_port_list=[int(dp_port)], |
| 138 | ) |
| 139 | with set_current_vllm_config(vllm_config): |
| 140 | # rank=0 here because each DP rank has a single (tp=1,pp=1) process, |
| 141 | # so the local rank within the tp*pp group is always 0. |
| 142 | # init_distributed_environment will offset by data_parallel_rank. |
| 143 | init_distributed_environment( |
| 144 | world_size=1, # tp * pp = 1 |
| 145 | rank=0, |
| 146 | distributed_init_method=f"tcp://localhost:{port}", |
| 147 | local_rank=rank, |
| 148 | ) |
| 149 | ensure_model_parallel_initialized(1, 1) |
| 150 | |
| 151 | |
| 152 | @contextmanager |
| 153 | def _make_forward_context(rank, world_size, num_tokens_per_rank): |
| 154 | """Keep the vLLM config and mock-DP forward context active together. |
| 155 | |
| 156 | Returns a context manager suitable for ``with`` statements. |
| 157 | The real DPMetadata (with sp_local_sizes etc.) is created internally |
| 158 | by set_forward_context from num_tokens_across_dp; the attn_metadata |
| 159 | placeholder just satisfies the "attn_metadata is not None" guard. |
| 160 | """ |
| 161 | from vllm.config import VllmConfig, set_current_vllm_config |
| 162 | from vllm.config.parallel import ParallelConfig |
| 163 | from vllm.forward_context import set_forward_context |
| 164 | |
| 165 | class _AttnMeta: |
| 166 | """Minimal placeholder so set_forward_context's |
| 167 | ``attn_metadata is not None`` guard (forward_context.py:334) |
| 168 | is satisfied. The real DPMetadata is built from num_tokens_across_dp.""" |
| 169 | |
| 170 | dp_metadata = None |
| 171 | |
| 172 | vllm_config = VllmConfig() |
| 173 | vllm_config.parallel_config = ParallelConfig( |
| 174 | data_parallel_size=world_size, |
| 175 | is_moe_model=True, |
| 176 | data_parallel_rank=rank, |
| 177 | ) |
| 178 | with ( |
| 179 | set_current_vllm_config(vllm_config), |
| 180 | set_forward_context( |
| 181 | _AttnMeta(), |
| 182 | vllm_config, |
| 183 | num_tokens=num_tokens_per_rank, |
| 184 | num_tokens_across_dp=torch.tensor( |
| 185 | [num_tokens_per_rank] * world_size, dtype=torch.int |
| 186 | ), |
| 187 | ), |
| 188 | ): |
| 189 | yield |
| 190 | |
| 191 | |
| 192 | # --------------------------------------------------------------------------- |
| 193 | # Skip conditions |
| 194 | # --------------------------------------------------------------------------- |
| 195 | |
| 196 | requires_multi_gpu = pytest.mark.skipif( |
| 197 | torch.accelerator.device_count() < 2, reason="Need >= 2 GPUs" |
| 198 | ) |
| 199 | requires_two_sided = pytest.mark.skipif( |
| 200 | not has_flashinfer_nvlink_two_sided(), |
| 201 | reason="FlashInfer NVLink two-sided not available", |
| 202 | ) |
| 203 | requires_one_sided = pytest.mark.skipif( |
| 204 | not has_flashinfer_nvlink_one_sided(), |
| 205 | reason="FlashInfer NVLink one-sided not available", |
| 206 | ) |
| 207 | requires_ptrace = pytest.mark.skipif( |
| 208 | not _has_sys_ptrace(), |
| 209 | reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)", |
| 210 | ) |
| 211 | requires_deep_ep_v2 = pytest.mark.skipif( |
| 212 | not has_deep_ep_v2(), |
| 213 | reason="DeepEP v2 (ElasticBuffer) not available or NCCL < 2.30.4", |
| 214 | ) |
| 215 | |
| 216 | # NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have |
| 217 | # their own @requires_two_sided / @requires_one_sided decorators, and |
| 218 | # test_args_dispatch_combine uses only standard torch.distributed ops and |
| 219 | # should run even when FlashInfer NVLink backends are not installed. |
| 220 | |
| 221 | |
| 222 | @pytest.mark.parametrize("supports_output", [False, True]) |
| 223 | def test_one_sided_combine_into_compatibility(supports_output): |
| 224 | from vllm.distributed.device_communicators.all2all import ( |
| 225 | FlashInferNVLinkOneSidedManager, |
| 226 | ) |
| 227 | |
| 228 | class FakeMoeAlltoAll: |
| 229 | def combine( |
| 230 | self, |
| 231 | payload, |
| 232 | runtime_max_tokens_per_rank, |
| 233 | output=None, |
| 234 | ): |
| 235 | result = payload + runtime_max_tokens_per_rank |
| 236 | if output is None: |
| 237 | return result |
| 238 | output.copy_(result) |
| 239 | return output |
| 240 | |
| 241 | manager = FlashInferNVLinkOneSidedManager.__new__(FlashInferNVLinkOneSidedManager) |
| 242 | manager.moe_alltoall = FakeMoeAlltoAll() |
| 243 | manager._combine_supports_output = supports_output |
| 244 | payload = torch.arange(4, dtype=torch.float32) |
| 245 | output = torch.empty_like(payload) |
| 246 | |
| 247 | manager.combine_into(payload, runtime_max_tokens_per_rank=2, output=output) |
| 248 | |
| 249 | torch.testing.assert_close(output, payload + 2) |
| 250 | |
| 251 | |
| 252 | # --------------------------------------------------------------------------- |
| 253 | # Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init) |
| 254 | # --------------------------------------------------------------------------- |
| 255 | # |
| 256 | # Tests FlashInferNVLinkTwoSidedManager which wraps FlashInfer's MnnvlMoe. |
| 257 | # initialize() allocates MNNVL shared workspaces via MnnvlMoe.get_moe_workspaces, |
| 258 | # which uses pidfd_getfd() to share memory file descriptors across processes — |
| 259 | # hence the SYS_PTRACE requirement. |
| 260 | # |
| 261 | # Uses EP group (get_ep_group) because the two-sided manager is constructed |
| 262 | # with an EP-scoped communicator in production. With tp=world_size the EP |
| 263 | # group spans all ranks, giving us a multi-rank group for testing. |
| 264 | # --------------------------------------------------------------------------- |
| 265 | |
| 266 | |
| 267 | def _two_sided_lifecycle_worker(rank, world_size): |
| 268 | from vllm.distributed.device_communicators.all2all import ( |
| 269 | FlashInferNVLinkTwoSidedManager, |
| 270 | ) |
| 271 | |
| 272 | cpu_group = get_ep_group().cpu_group |
| 273 | num_gpus = torch.accelerator.device_count() |
| 274 | manager = FlashInferNVLinkTwoSidedManager(cpu_group) |
| 275 | |
| 276 | # Not initialized yet |
| 277 | assert not manager.initialized |
| 278 | assert manager.rank == rank |
| 279 | assert manager.world_size == world_size |
| 280 | |
| 281 | # Initialize |
| 282 | manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus) |
| 283 | assert manager.initialized |
| 284 | assert manager.workspace_tensor is not None |
| 285 | assert manager.prepare_workspace_tensor is not None |
| 286 | assert manager.mapping is not None |
| 287 | |
| 288 | torch.distributed.barrier() |
| 289 | |
| 290 | # Cleanup |
| 291 | manager.cleanup() |
| 292 | assert not manager.initialized |
| 293 | assert manager.workspace_tensor is None |
| 294 | assert manager.prepare_workspace_tensor is None |
| 295 | |
| 296 | torch.distributed.barrier() |
| 297 | |
| 298 | # Reinitialize |
| 299 | manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus) |
| 300 | assert manager.initialized |
| 301 | |
| 302 | torch.distributed.barrier() |
| 303 | |
| 304 | # ensure_alltoall_workspace_initialized is idempotent when already init'd |
| 305 | assert manager.ensure_alltoall_workspace_initialized() |
| 306 | assert manager.initialized |
| 307 | |
| 308 | manager.cleanup() |
| 309 | assert not manager.initialized |
| 310 | |
| 311 | |
| 312 | @requires_multi_gpu |
| 313 | @requires_two_sided |
| 314 | @requires_ptrace |
| 315 | @pytest.mark.parametrize("world_size", [2]) |
| 316 | def test_two_sided_manager_lifecycle(world_size): |
| 317 | """Test init, cleanup, reinit, and ensure_initialized idempotency.""" |
| 318 | _spawn_workers(_two_sided_lifecycle_worker, world_size) |
| 319 | |
| 320 | |
| 321 | # --------------------------------------------------------------------------- |
| 322 | # Test 2: One-sided manager lifecycle (init, cleanup, reinit) |
| 323 | # --------------------------------------------------------------------------- |
| 324 | # |
| 325 | # Tests FlashInferNVLinkOneSidedManager which wraps FlashInfer's MoeAlltoAll. |
| 326 | # initialize() creates MoeAlltoAll with an MnnvlConfig, which allocates MNNVL |
| 327 | # shared workspaces — same cross-process memory sharing as two-sided, hence |
| 328 | # the SYS_PTRACE requirement. |
| 329 | # |
| 330 | # Uses DP group (get_dp_group) because the one-sided manager's initialize() |
| 331 | # internally calls get_dp_group() to set up the MnnvlConfig communicator. |
| 332 | # We therefore need a real DP group with world_size > 1, which requires |
| 333 | # dp_size=world_size via _init_dp_environment. |
| 334 | # --------------------------------------------------------------------------- |
| 335 | |
| 336 | |
| 337 | def _one_sided_lifecycle_worker(rank, world_size): |
| 338 | from vllm.distributed.device_communicators.all2all import ( |
| 339 | FlashInferNVLinkOneSidedManager, |
| 340 | ) |
| 341 | from vllm.distributed.parallel_state import get_dp_group |
| 342 | |
| 343 | cpu_group = get_dp_group().cpu_group |
| 344 | manager = FlashInferNVLinkOneSidedManager(cpu_group) |
| 345 | |
| 346 | assert not manager.initialized |
| 347 | assert manager.rank == rank |
| 348 | assert manager.world_size == world_size |
| 349 | |
| 350 | init_kwargs = dict( |
| 351 | max_num_tokens=1024, |
| 352 | top_k=2, |
| 353 | num_experts=world_size * 8, |
| 354 | hidden_size=4096, |
| 355 | x_bytes_per_token=4096 * 2, |
| 356 | x_sf_bytes_per_token=0, |
| 357 | ) |
| 358 | |
| 359 | # Initialize |
| 360 | manager.initialize(**init_kwargs) |
| 361 | assert manager.initialized |
| 362 | assert manager.moe_alltoall is not None |
| 363 | assert manager.mapping is not None |
| 364 | |
| 365 | torch.distributed.barrier() |
| 366 | |
| 367 | # Cleanup |
| 368 | manager.cleanup() |
| 369 | assert not manager.initialized |
| 370 | assert manager.moe_alltoall is None |
| 371 | |
| 372 | torch.distributed.barrier() |
| 373 | |
| 374 | # Reinitialize with different token count |
| 375 | manager.initialize(**{**init_kwargs, "max_num_tokens": 2048}) |
| 376 | assert manager.initialized |
| 377 | |
| 378 | torch.distributed.barrier() |
| 379 | manager.cleanup() |
| 380 | |
| 381 | |
| 382 | @requires_multi_gpu |
| 383 | @requires_one_sided |
| 384 | @requires_ptrace |
| 385 | @pytest.mark.parametrize("world_size", [2]) |
| 386 | def test_one_sided_manager_lifecycle(world_size): |
| 387 | """Test init, cleanup, and reinit with different params.""" |
| 388 | _spawn_workers( |
| 389 | _one_sided_lifecycle_worker, |
| 390 | world_size, |
| 391 | dp_size=world_size, |
| 392 | ) |
| 393 | |
| 394 | |
| 395 | # --------------------------------------------------------------------------- |
| 396 | # Test 2b: One-sided manager grows workspace across heterogeneous MoE layers |
| 397 | # --------------------------------------------------------------------------- |
| 398 | # |
| 399 | # Models with heterogeneous MoE quantization — most notably a quantized base |
| 400 | # MoE combined with an unquantized MTP head — can call initialize() multiple |
| 401 | # times with different per-token dispatch payload sizes. The shared workspace |
| 402 | # must grow to the union and the MoeAlltoAll must be rebuilt; otherwise a |
| 403 | # later layer's combine call overruns the workspace sized for the first |
| 404 | # layer's smaller payload and trips FlashInfer's combinePayloadOffset assert. |
| 405 | # --------------------------------------------------------------------------- |
| 406 | |
| 407 | |
| 408 | def _one_sided_workspace_grow_worker(rank, world_size): |
| 409 | from vllm.distributed.device_communicators.all2all import ( |
| 410 | FlashInferNVLinkOneSidedManager, |
| 411 | ) |
| 412 | from vllm.distributed.parallel_state import get_dp_group |
| 413 | |
| 414 | cpu_group = get_dp_group().cpu_group |
| 415 | manager = FlashInferNVLinkOneSidedManager(cpu_group) |
| 416 | |
| 417 | base_kwargs = dict( |
| 418 | max_num_tokens=1024, |
| 419 | top_k=2, |
| 420 | num_experts=world_size * 8, |
| 421 | hidden_size=4096, |
| 422 | ) |
| 423 | nvfp4_kwargs = dict( |
| 424 | x_bytes_per_token=base_kwargs["hidden_size"] // 2, |
| 425 | x_sf_bytes_per_token=base_kwargs["hidden_size"] // 16, |
| 426 | ) |
| 427 | bf16_kwargs = dict( |
| 428 | x_bytes_per_token=base_kwargs["hidden_size"] * 2, |
| 429 | x_sf_bytes_per_token=0, |
| 430 | ) |
| 431 | |
| 432 | # First init: NVFP4-like (hidden_bytes = hidden // 2 + hidden // 16). |
| 433 | manager.initialize(**base_kwargs, **nvfp4_kwargs) |
| 434 | assert manager.initialized |
| 435 | nvfp4_workspace_size = manager.workspace_size |
| 436 | nvfp4_moe_alltoall = manager.moe_alltoall |
| 437 | |
| 438 | torch.distributed.barrier() |
| 439 | |
| 440 | # Second init: bf16-like (hidden_bytes = hidden * 2). Models the case of |
| 441 | # a quantized base MoE followed by an unquantized MoE layer (e.g. an MTP |
| 442 | # head). Per-token dispatch payload is ~4x larger, so the union workspace |
| 443 | # must grow and MoeAlltoAll must be rebuilt. |
| 444 | manager.initialize(**base_kwargs, **bf16_kwargs) |
| 445 | assert manager.initialized |
| 446 | assert manager.workspace_size > nvfp4_workspace_size |
| 447 | assert manager.moe_alltoall is not nvfp4_moe_alltoall |
| 448 | bf16_workspace_size = manager.workspace_size |
| 449 | bf16_moe_alltoall = manager.moe_alltoall |
| 450 | |
| 451 | torch.distributed.barrier() |
| 452 | |
| 453 | # Third init: back to NVFP4-like shape. Existing workspace already covers |
| 454 | # it, so initialize() must no-op — no shrink, no rebuild. |
| 455 | manager.initialize(**base_kwargs, **nvfp4_kwargs) |
| 456 | assert manager.initialized |
| 457 | assert manager.workspace_size == bf16_workspace_size |
| 458 | assert manager.moe_alltoall is bf16_moe_alltoall |
| 459 | |
| 460 | torch.distributed.barrier() |
| 461 | manager.cleanup() |
| 462 | |
| 463 | |
| 464 | @requires_multi_gpu |
| 465 | @requires_one_sided |
| 466 | @requires_ptrace |
| 467 | @pytest.mark.parametrize("world_size", [2]) |
| 468 | def test_one_sided_manager_workspace_grow(world_size): |
| 469 | """A later initialize() with a larger per-token payload must grow the |
| 470 | workspace and rebuild MoeAlltoAll; a later initialize() with a smaller |
| 471 | payload must no-op.""" |
| 472 | _spawn_workers( |
| 473 | _one_sided_workspace_grow_worker, |
| 474 | world_size, |
| 475 | dp_size=world_size, |
| 476 | ) |
| 477 | |
| 478 | |
| 479 | # --------------------------------------------------------------------------- |
| 480 | # Test 3: AgRs dispatch/combine with value validation |
| 481 | # --------------------------------------------------------------------------- |
| 482 | # |
| 483 | # Tests AgRsAll2AllManager which uses only standard torch.distributed |
| 484 | # all_gatherv / reduce_scatterv — no FlashInfer or MNNVL dependency. |
| 485 | # This test validates the reference all-to-all implementation that other |
| 486 | # backends are compared against. |
| 487 | # --------------------------------------------------------------------------- |
| 488 | |
| 489 | |
| 490 | def _args_dispatch_combine_worker(rank, world_size): |
| 491 | from vllm.distributed.device_communicators.all2all import AgRsAll2AllManager |
| 492 | from vllm.forward_context import get_forward_context |
| 493 | |
| 494 | cpu_group = get_ep_group().cpu_group |
| 495 | device = torch.device(f"{DEVICE}:{rank}") |
| 496 | |
| 497 | hidden_size = 64 |
| 498 | tokens_per_rank = 16 |
| 499 | experts_per_token = 2 |
| 500 | num_experts = world_size * 4 |
| 501 | total_tokens = world_size * tokens_per_rank |
| 502 | |
| 503 | # Deterministic per-rank data: rank r has value (r + 1) |
| 504 | hidden = torch.full( |
| 505 | (tokens_per_rank, hidden_size), |
| 506 | float(rank + 1), |
| 507 | device=device, |
| 508 | dtype=torch.float32, |
| 509 | ) |
| 510 | router = torch.full( |
| 511 | (tokens_per_rank, num_experts), |
| 512 | float(rank + 1) * 10, |
| 513 | device=device, |
| 514 | dtype=torch.float32, |
| 515 | ) |
| 516 | weights = torch.full( |
| 517 | (tokens_per_rank, experts_per_token), |
| 518 | float(rank + 1) * 100, |
| 519 | device=device, |
| 520 | dtype=torch.float32, |
| 521 | ) |
| 522 | ids = torch.full( |
| 523 | (tokens_per_rank, experts_per_token), |
| 524 | rank, |
| 525 | device=device, |
| 526 | dtype=torch.long, |
| 527 | ) |
| 528 | |
| 529 | with _make_forward_context(rank, world_size, tokens_per_rank): |
| 530 | manager = AgRsAll2AllManager(cpu_group) |
| 531 | dp_metadata = get_forward_context().dp_metadata |
| 532 | |
| 533 | with dp_metadata.sp_local_sizes(sequence_parallel_size=1): |
| 534 | # -- dispatch_router_logits -- |
| 535 | d_hidden, d_router = manager.dispatch_router_logits( |
| 536 | hidden.clone(), |
| 537 | router.clone(), |
| 538 | is_sequence_parallel=True, |
| 539 | ) |
| 540 | assert d_hidden.shape == (total_tokens, hidden_size) |
| 541 | assert d_router.shape == (total_tokens, num_experts) |
| 542 | |
| 543 | for r in range(world_size): |
| 544 | s = r * tokens_per_rank |
| 545 | e = (r + 1) * tokens_per_rank |
| 546 | torch.testing.assert_close( |
| 547 | d_hidden[s:e], |
| 548 | torch.full_like(d_hidden[s:e], float(r + 1)), |
| 549 | ) |
| 550 | torch.testing.assert_close( |
| 551 | d_router[s:e], |
| 552 | torch.full_like(d_router[s:e], float(r + 1) * 10), |
| 553 | ) |
| 554 | |
| 555 | # -- dispatch -- |
| 556 | d_hidden2, d_weights, d_ids = manager.dispatch( |
| 557 | hidden.clone(), |
| 558 | weights.clone(), |
| 559 | ids.clone(), |
| 560 | is_sequence_parallel=True, |
| 561 | ) |
| 562 | assert d_hidden2.shape == (total_tokens, hidden_size) |
| 563 | assert d_weights.shape == (total_tokens, experts_per_token) |
| 564 | assert d_ids.shape == (total_tokens, experts_per_token) |
| 565 | |
| 566 | for r in range(world_size): |
| 567 | s = r * tokens_per_rank |
| 568 | e = (r + 1) * tokens_per_rank |
| 569 | torch.testing.assert_close( |
| 570 | d_weights[s:e], |
| 571 | torch.full_like(d_weights[s:e], float(r + 1) * 100), |
| 572 | ) |
| 573 | assert (d_ids[s:e] == r).all() |
| 574 | |
| 575 | # -- combine (reduce-scatter) -- |
| 576 | # Each token i has value i in all columns; after reduce-scatter |
| 577 | # each rank gets its slice, summed across ranks. |
| 578 | expert_out = ( |
| 579 | torch.arange(total_tokens, device=device, dtype=torch.float32) |
| 580 | .unsqueeze(1) |
| 581 | .expand(total_tokens, hidden_size) |
| 582 | .contiguous() |
| 583 | ) |
| 584 | |
| 585 | combined = manager.combine(expert_out, is_sequence_parallel=True) |
| 586 | assert combined.shape == (tokens_per_rank, hidden_size) |
| 587 | |
| 588 | for i in range(tokens_per_rank): |
| 589 | expected_val = float(rank * tokens_per_rank + i) * world_size |
| 590 | torch.testing.assert_close( |
| 591 | combined[i], |
| 592 | torch.full_like(combined[i], expected_val), |
| 593 | ) |
| 594 | |
| 595 | torch.distributed.barrier() |
| 596 | |
| 597 | |
| 598 | @requires_multi_gpu |
| 599 | @pytest.mark.parametrize("world_size", [2]) |
| 600 | def test_args_dispatch_combine(world_size): |
| 601 | """Validate dispatch gathers all-rank data and combine reduces correctly.""" |
| 602 | _spawn_workers(_args_dispatch_combine_worker, world_size) |
| 603 | |
| 604 | |
| 605 | # --------------------------------------------------------------------------- |
| 606 | # Test 4: FlashInfer two-sided dispatch/combine data communication |
| 607 | # --------------------------------------------------------------------------- |
| 608 | # |
| 609 | # Tests actual data flow through the FlashInfer NVLink two-sided backend |
| 610 | # by calling flashinfer_alltoall_dispatch (with defer_input_quant=True to |
| 611 | # skip quantization) and flashinfer_alltoall_combine, then verifying exact |
| 612 | # round-trip values. Dispatch sends each token once per distinct expert |
| 613 | # rank, and combine performs an unweighted sum, so: |
| 614 | # dispatch(hidden) → identity → combine = hidden * num_distinct_ranks(i) |
| 615 | # --------------------------------------------------------------------------- |
| 616 | |
| 617 | |
| 618 | def _two_sided_data_worker(rank, world_size): |
| 619 | from vllm.distributed.device_communicators.all2all import ( |
| 620 | FlashInferNVLinkTwoSidedManager, |
| 621 | ) |
| 622 | from vllm.distributed.parallel_state import get_dp_group |
| 623 | from vllm.forward_context import get_forward_context |
| 624 | from vllm.model_executor.layers.fused_moe.config import ( |
| 625 | FusedMoEQuantConfig, |
| 626 | FusedMoEQuantDesc, |
| 627 | ) |
| 628 | from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501 |
| 629 | flashinfer_alltoall_combine, |
| 630 | flashinfer_alltoall_dispatch, |
| 631 | ) |
| 632 | |
| 633 | # Use DP group because MnnvlMoe workspace allocation calls get_dp_group() |
| 634 | # internally and requires dp_size == ep_size. |
| 635 | cpu_group = get_dp_group().cpu_group |
| 636 | device = torch.device(f"{DEVICE}:{rank}") |
| 637 | num_gpus = torch.accelerator.device_count() |
| 638 | |
| 639 | hidden_size = 128 |
| 640 | tokens_per_rank = 32 |
| 641 | experts_per_token = 2 |
| 642 | num_experts = world_size * 4 |
| 643 | |
| 644 | # Initialize the FlashInfer two-sided manager |
| 645 | manager = FlashInferNVLinkTwoSidedManager(cpu_group) |
| 646 | manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus) |
| 647 | assert manager.initialized |
| 648 | |
| 649 | torch.distributed.barrier() |
| 650 | |
| 651 | # Create deterministic per-rank test data |
| 652 | torch.manual_seed(rank + 42) |
| 653 | hidden = torch.randn( |
| 654 | tokens_per_rank, |
| 655 | hidden_size, |
| 656 | device=device, |
| 657 | dtype=torch.bfloat16, |
| 658 | ) |
| 659 | # Assign each token to experts spread across ranks so tokens move between GPUs |
| 660 | topk_ids = torch.randint( |
| 661 | 0, |
| 662 | num_experts, |
| 663 | (tokens_per_rank, experts_per_token), |
| 664 | device=device, |
| 665 | dtype=torch.int32, |
| 666 | ) |
| 667 | topk_weights = torch.rand( |
| 668 | tokens_per_rank, |
| 669 | experts_per_token, |
| 670 | device=device, |
| 671 | dtype=torch.float32, |
| 672 | ) |
| 673 | |
| 674 | # Unquantized config: quant_dtype=None means moe_kernel_quantize_input is a no-op |
| 675 | no_quant = FusedMoEQuantDesc() |
| 676 | quant_config = FusedMoEQuantConfig( |
| 677 | _a1=no_quant, |
| 678 | _a2=no_quant, |
| 679 | _w1=no_quant, |
| 680 | _w2=no_quant, |
| 681 | ) |
| 682 | assert quant_config.quant_dtype is None # sanity: no quantization |
| 683 | |
| 684 | with _make_forward_context(rank, world_size, tokens_per_rank): |
| 685 | dp_metadata = get_forward_context().dp_metadata |
| 686 | |
| 687 | with dp_metadata.sp_local_sizes(sequence_parallel_size=1): |
| 688 | local_sizes = dp_metadata.get_chunk_sizes_across_dp_rank() |
| 689 | |
| 690 | # --- FlashInfer two-sided dispatch --- |
| 691 | alltoall_info, fi_topk_ids, fi_topk_weights, fi_hidden, fi_scale = ( |
| 692 | flashinfer_alltoall_dispatch( |
| 693 | manager, |
| 694 | local_sizes, |
| 695 | hidden.clone(), |
| 696 | None, # no global scale |
| 697 | topk_ids.clone(), |
| 698 | topk_weights.clone(), |
| 699 | experts_per_token, |
| 700 | num_experts, |
| 701 | quant_config, |
| 702 | defer_input_quant=True, |
| 703 | ) |
| 704 | ) |
| 705 | assert fi_scale is None # deferred quant: no scale produced |
| 706 | assert fi_hidden is not None |
| 707 | assert fi_hidden.shape[1] == hidden_size |
| 708 | assert fi_hidden.numel() > 0 |
| 709 | |
| 710 | # --- Round-trip exact verification --- |
| 711 | # The all-to-all sends each token once per *distinct* expert |
| 712 | # rank. Combine performs an unweighted sum of the per-rank |
| 713 | # contributions. With identity expert (feeding dispatched |
| 714 | # hidden straight back): |
| 715 | # result[i] = hidden[i] * num_distinct_expert_ranks(i) |
| 716 | combined = flashinfer_alltoall_combine( |
| 717 | manager, |
| 718 | fi_hidden, |
| 719 | top_k=experts_per_token, |
| 720 | token_count=tokens_per_rank, |
| 721 | alltoall_info=alltoall_info, |
| 722 | ) |
| 723 | assert combined.shape == (tokens_per_rank, hidden_size) |
| 724 | |
| 725 | experts_per_rank = num_experts // world_size |
| 726 | expert_ranks = topk_ids // experts_per_rank # (tokens, top_k) |
| 727 | num_distinct = torch.tensor( |
| 728 | [len(set(row.tolist())) for row in expert_ranks], |
| 729 | device=device, |
| 730 | dtype=torch.float32, |
| 731 | ).unsqueeze(1) # (tokens, 1) |
| 732 | expected = (hidden.float() * num_distinct).to(hidden.dtype) |
| 733 | torch.testing.assert_close(combined, expected) |
| 734 | |
| 735 | # --- Linearity check with scaled expert output --- |
| 736 | # Scaling the expert output by a constant should scale the |
| 737 | # combined result by the same constant. |
| 738 | scale = 3.0 |
| 739 | combined_scaled = flashinfer_alltoall_combine( |
| 740 | manager, |
| 741 | fi_hidden * scale, |
| 742 | top_k=experts_per_token, |
| 743 | token_count=tokens_per_rank, |
| 744 | alltoall_info=alltoall_info, |
| 745 | ) |
| 746 | expected_scaled = (hidden.float() * num_distinct * scale).to(hidden.dtype) |
| 747 | torch.testing.assert_close(combined_scaled, expected_scaled) |
| 748 | |
| 749 | torch.distributed.barrier() |
| 750 | |
| 751 | manager.cleanup() |
| 752 | |
| 753 | |
| 754 | @requires_multi_gpu |
| 755 | @requires_two_sided |
| 756 | @requires_ptrace |
| 757 | @pytest.mark.parametrize("world_size", [2]) |
| 758 | def test_two_sided_dispatch_combine(world_size): |
| 759 | """Test FlashInfer two-sided dispatch/combine with exact value verification.""" |
| 760 | _spawn_workers(_two_sided_data_worker, world_size, dp_size=world_size) |
| 761 | |
| 762 | |
| 763 | # --------------------------------------------------------------------------- |
| 764 | # Test 5: FlashInfer one-sided dispatch/combine data communication |
| 765 | # --------------------------------------------------------------------------- |
| 766 | # |
| 767 | # Tests actual data flow through the FlashInfer NVLink one-sided backend |
| 768 | # by calling MoeAlltoAll.dispatch() and MoeAlltoAll.combine() directly |
| 769 | # with synthetic payloads, then verifying shapes and round-trip consistency. |
| 770 | # --------------------------------------------------------------------------- |
| 771 | |
| 772 | |
| 773 | def _one_sided_data_worker(rank, world_size): |
| 774 | from vllm.distributed.device_communicators.all2all import ( |
| 775 | FlashInferNVLinkOneSidedManager, |
| 776 | ) |
| 777 | from vllm.distributed.parallel_state import get_dp_group |
| 778 | from vllm.forward_context import get_forward_context |
| 779 | |
| 780 | cpu_group = get_dp_group().cpu_group |
| 781 | device = torch.device(f"{DEVICE}:{rank}") |
| 782 | |
| 783 | hidden_size = 256 |
| 784 | tokens_per_rank = 32 |
| 785 | experts_per_token = 2 |
| 786 | num_experts = world_size * 8 |
| 787 | |
| 788 | # Initialize the one-sided manager |
| 789 | manager = FlashInferNVLinkOneSidedManager(cpu_group) |
| 790 | manager.initialize( |
| 791 | max_num_tokens=tokens_per_rank, |
| 792 | top_k=experts_per_token, |
| 793 | num_experts=num_experts, |
| 794 | hidden_size=hidden_size, |
| 795 | x_bytes_per_token=hidden_size // 2, |
| 796 | # Account for the fp8 block-scale payload (x_sf: hidden//16 bytes |
| 797 | # per token) that is dispatched alongside the nvfp4 hidden states. |
| 798 | # Without this the dispatch region is under-reserved and the combine |
| 799 | # payload overflows the per-rank workspace. |
| 800 | x_sf_bytes_per_token=hidden_size // 16, |
| 801 | ) |
| 802 | assert manager.initialized |
| 803 | assert manager.moe_alltoall is not None |
| 804 | |
| 805 | with _make_forward_context(rank, world_size, tokens_per_rank): |
| 806 | dp_metadata = get_forward_context().dp_metadata |
| 807 | |
| 808 | with dp_metadata.sp_local_sizes(sequence_parallel_size=1): |
| 809 | local_sizes = dp_metadata.get_chunk_sizes_across_dp_rank() |
| 810 | runtime_max_tokens = max(local_sizes) |
| 811 | |
| 812 | # Create test data with raw tensors matching the nvfp4 payload |
| 813 | # sizes the workspace was allocated for: |
| 814 | # x: (tokens, hidden_size // 2) — nvfp4 hidden states |
| 815 | # x_sf: (tokens, hidden_size // 16) — fp8 scaling factors |
| 816 | torch.manual_seed(rank + 42) |
| 817 | x = torch.randint( |
| 818 | 0, |
| 819 | 256, |
| 820 | (tokens_per_rank, hidden_size // 2), |
| 821 | device=device, |
| 822 | dtype=torch.uint8, |
| 823 | ) |
| 824 | x_sf = torch.randint( |
| 825 | 0, |
| 826 | 256, |
| 827 | (tokens_per_rank, hidden_size // 16), |
| 828 | device=device, |
| 829 | dtype=torch.uint8, |
| 830 | ) |
| 831 | topk_ids = torch.randint( |
| 832 | 0, |
| 833 | num_experts, |
| 834 | (tokens_per_rank, experts_per_token), |
| 835 | device=device, |
| 836 | dtype=torch.int32, |
| 837 | ) |
| 838 | topk_weights = torch.rand( |
| 839 | tokens_per_rank, |
| 840 | experts_per_token, |
| 841 | device=device, |
| 842 | dtype=torch.float32, |
| 843 | ) |
| 844 | |
| 845 | # --- One-sided dispatch --- |
| 846 | payloads = [x, x_sf, topk_ids, topk_weights] |
| 847 | recv_payloads = manager.moe_alltoall.dispatch( |
| 848 | token_selected_experts=topk_ids, |
| 849 | input_payloads=payloads, |
| 850 | runtime_max_tokens_per_rank=runtime_max_tokens, |
| 851 | ) |
| 852 | assert len(recv_payloads) == 4 |
| 853 | recv_x, recv_x_sf, recv_ids, recv_weights = recv_payloads |
| 854 | assert recv_x.numel() > 0 |
| 855 | assert recv_x_sf.numel() > 0 |
| 856 | assert recv_ids.numel() > 0 |
| 857 | |
| 858 | # --- Round-trip exact verification --- |
| 859 | # The dispatch routes each token once per *distinct* expert |
| 860 | # rank. Combine performs an unweighted sum of per-rank |
| 861 | # contributions. With constant expert output (all 1s): |
| 862 | # result[i] = 1.0 * num_distinct_expert_ranks(i) |
| 863 | expert_output = torch.ones( |
| 864 | world_size, |
| 865 | runtime_max_tokens, |
| 866 | hidden_size, |
| 867 | device=device, |
| 868 | dtype=torch.bfloat16, |
| 869 | ) |
| 870 | combined = manager.moe_alltoall.combine( |
| 871 | payload=expert_output, |
| 872 | runtime_max_tokens_per_rank=runtime_max_tokens, |
| 873 | ) |
| 874 | assert combined.shape == (tokens_per_rank, hidden_size) |
| 875 | |
| 876 | experts_per_rank = num_experts // world_size |
| 877 | expert_ranks = topk_ids // experts_per_rank # (tokens, top_k) |
| 878 | num_distinct = torch.tensor( |
| 879 | [len(set(row.tolist())) for row in expert_ranks], |
| 880 | device=device, |
| 881 | dtype=torch.bfloat16, |
| 882 | ).unsqueeze(1) # (tokens, 1) |
| 883 | expected = num_distinct.expand_as(combined) |
| 884 | torch.testing.assert_close(combined, expected) |
| 885 | |
| 886 | # --- Linearity check with scaled expert output --- |
| 887 | # Scaling the expert output by a constant should scale the |
| 888 | # combined result by the same constant. |
| 889 | # Re-dispatch to reset internal state (one-sided requires a |
| 890 | # fresh dispatch before each combine). |
| 891 | manager.moe_alltoall.dispatch( |
| 892 | token_selected_experts=topk_ids, |
| 893 | input_payloads=payloads, |
| 894 | runtime_max_tokens_per_rank=runtime_max_tokens, |
| 895 | ) |
| 896 | scale = 3.0 |
| 897 | combined_scaled = manager.moe_alltoall.combine( |
| 898 | payload=expert_output * scale, |
| 899 | runtime_max_tokens_per_rank=runtime_max_tokens, |
| 900 | ) |
| 901 | expected_scaled = (expected * scale).to(torch.bfloat16) |
| 902 | torch.testing.assert_close(combined_scaled, expected_scaled) |
| 903 | |
| 904 | torch.distributed.barrier() |
| 905 | |
| 906 | manager.cleanup() |
| 907 | |
| 908 | |
| 909 | @requires_multi_gpu |
| 910 | @requires_one_sided |
| 911 | @requires_ptrace |
| 912 | @pytest.mark.parametrize("world_size", [2]) |
| 913 | def test_one_sided_dispatch_combine(world_size): |
| 914 | """Test FlashInfer one-sided dispatch/combine with actual data flow.""" |
| 915 | _spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size) |
| 916 | |
| 917 | |
| 918 | # --------------------------------------------------------------------------- |
| 919 | # Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle |
| 920 | # --------------------------------------------------------------------------- |
| 921 | # |
| 922 | # Tests DeepEPV2All2AllManager which wraps DeepEP's ElasticBuffer API using |
| 923 | # the NCCL GIN backend. Requires DeepEP >= 2.0 and NCCL >= 2.30.4. |
| 924 | # |
| 925 | # Uses EP group because the DeepEP v2 manager is constructed with an |
| 926 | # EP-scoped communicator in production. With tp=world_size the EP group |
| 927 | # spans all ranks. |
| 928 | # --------------------------------------------------------------------------- |
| 929 | |
| 930 | |
| 931 | def _deepep_v2_lifecycle_worker(rank, world_size): |
| 932 | from vllm.distributed.device_communicators.all2all import ( |
| 933 | DeepEPV2All2AllManager, |
| 934 | ) |
| 935 | |
| 936 | ep_group = get_ep_group() |
| 937 | manager = DeepEPV2All2AllManager( |
| 938 | ep_group.cpu_group, |
| 939 | device_group=ep_group.device_group, |
| 940 | ) |
| 941 | |
| 942 | assert manager.rank == rank |
| 943 | assert manager.world_size == world_size |
| 944 | assert manager._num_sms is None |
| 945 | |
| 946 | hidden_size = 7168 |
| 947 | num_experts = world_size * 32 |
| 948 | num_topk = 8 |
| 949 | max_tokens = 256 |
| 950 | |
| 951 | handle_kwargs = dict( |
| 952 | num_max_tokens_per_rank=max_tokens, |
| 953 | hidden=hidden_size, |
| 954 | num_topk=num_topk, |
| 955 | num_experts=num_experts, |
| 956 | use_fp8_dispatch=False, |
| 957 | ) |
| 958 | |
| 959 | handle = manager.get_handle(handle_kwargs) |
| 960 | assert handle is not None |
| 961 | assert manager._num_sms is not None |
| 962 | assert manager._num_sms > 0 |
| 963 | |
| 964 | torch.distributed.barrier() |
| 965 | |
| 966 | # get_handle again with same args should return cached handle |
| 967 | handle2 = manager.get_handle(dict(handle_kwargs)) |
| 968 | assert handle2 is handle |
| 969 | |
| 970 | torch.distributed.barrier() |
| 971 | |
| 972 | # Destroy clears the cache |
| 973 | manager.destroy() |
| 974 | assert len(manager.handle_cache._cache) == 0 |
| 975 | |
| 976 | torch.distributed.barrier() |
| 977 | |
| 978 | # Re-create after destroy |
| 979 | handle3 = manager.get_handle(dict(handle_kwargs)) |
| 980 | assert handle3 is not None |
| 981 | |
| 982 | torch.distributed.barrier() |
| 983 | manager.destroy() |
| 984 | |
| 985 | |
| 986 | @requires_multi_gpu |
| 987 | @requires_deep_ep_v2 |
| 988 | @pytest.mark.parametrize("world_size", [2]) |
| 989 | def test_deepep_v2_manager_lifecycle(world_size): |
| 990 | """Test DeepEP v2 ElasticBuffer manager init, caching, and destroy.""" |
| 991 | _spawn_workers(_deepep_v2_lifecycle_worker, world_size) |