ccfb620c62 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -0,0 +1,774 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""

4 +Tests for MNNVL AllToAll operations.

5 +

6 +Requires: docker run ... --cap-add=SYS_PTRACE ...

7 +Run: pytest tests/distributed/test_mnnvl_alltoall.py -v

8 +"""

9 +

10 +import os

11 +import traceback

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.utils.flashinfer import (

19 + has_flashinfer_nvlink_one_sided,

20 + has_flashinfer_nvlink_two_sided,

21 +)

22 +from vllm.utils.network_utils import get_open_port

23 +

24 +from ..utils import init_test_distributed_environment

25 +

26 +# ---------------------------------------------------------------------------

27 +# Helpers

28 +# ---------------------------------------------------------------------------

29 +

30 +

31 +def _has_sys_ptrace() -> bool:

32 + """Check for SYS_PTRACE capability (bit 19 in CapEff)."""

33 + try:

34 + with open("/proc/self/status") as f:

35 + for line in f:

36 + if line.startswith("CapEff:"):

37 + return bool(int(line.split()[1], 16) & (1 << 19))

38 + except Exception:

39 + pass

40 + return False

41 +

42 +

43 +def _spawn_workers(worker_fn, world_size, *, dp_size=None):

44 + """Spawn one process per GPU, run worker_fn, assert all succeed.

45 +

46 + Uses an mp.Queue to propagate worker tracebacks back to the parent

47 + so pytest shows the actual failure, not just an exit code.

48 + """

49 + if mp.get_start_method(allow_none=True) is None:

50 + mp.set_start_method("spawn")

51 +

52 + port = str(get_open_port())

53 + # Allocate a second port for DP master when dp_size is set, so the

54 + # distributed init port and DP port can't collide even under xdist.

55 + dp_port = str(get_open_port()) if dp_size is not None else None

56 + err_queue: mp.Queue = mp.Queue()

57 + procs = []

58 + for rank in range(world_size):

59 + p = mp.Process(

60 + target=_run_worker,

61 + args=(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue),

62 + )

63 + p.start()

64 + procs.append(p)

65 + for p in procs:

66 + p.join()

67 +

68 + # Collect any errors from workers before asserting.

69 + errors = []

70 + while not err_queue.empty():

71 + errors.append(err_queue.get_nowait())

72 + err_queue.close()

73 + err_queue.join_thread()

74 + if errors:

75 + pytest.fail("Worker(s) failed:\n" + "\n---\n".join(errors))

76 +

77 +

78 +def _run_worker(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue):

79 + """Per-process setup: device, distributed env, then call worker_fn.

80 +

81 + Args:

82 + dp_size: If set, initialize with tp=1 and data_parallel_size=dp_size.

83 + Otherwise use tp=world_size (default for EP-based tests).

84 + dp_port: Separate port for the DP master (only used when dp_size is set).

85 + err_queue: Queue for propagating tracebacks to the parent process.

86 + """

87 + try:

88 + os.environ.pop("CUDA_VISIBLE_DEVICES", None)

89 + torch.accelerator.set_device_index(rank)

90 + if dp_size is not None:

91 + _init_dp_environment(world_size, rank, port, dp_size, dp_port)

92 + else:

93 + init_test_distributed_environment(world_size, 1, rank, port)

94 + worker_fn(rank, world_size)

95 + torch.distributed.barrier()

96 + except Exception:

97 + err_queue.put(f"[Rank {rank}]\n{traceback.format_exc()}")

98 + # Don't re-raise: the parent reads errors from err_queue.

99 + # A non-zero exit from the re-raise would be redundant.

100 + import sys

101 +

102 + sys.exit(1)

103 +

104 +

105 +def _init_dp_environment(world_size, rank, port, dp_size, dp_port):

106 + """Initialize distributed env with data parallelism.

107 +

108 + Sets up tp=1, pp=1, dp=dp_size. Each process is one DP rank

109 + with local rank 0 within its (trivial) tp*pp group.

110 +

111 + Args:

112 + port: Port for torch.distributed init.

113 + dp_port: Separate port for the DP master group init.

114 + """

115 + from vllm.config import VllmConfig, set_current_vllm_config

116 + from vllm.config.parallel import ParallelConfig

117 + from vllm.distributed.parallel_state import (

118 + ensure_model_parallel_initialized,

119 + init_distributed_environment,

120 + )

121 +

122 + vllm_config = VllmConfig()

123 + vllm_config.parallel_config = ParallelConfig(

124 + data_parallel_size=dp_size,

125 + data_parallel_rank=rank,

126 + # Pre-populate port list so __post_init__ doesn't auto-generate

127 + # random ports. All DP ranks must agree on the same port.

128 + _data_parallel_master_port_list=[int(dp_port)],

129 + )

130 + with set_current_vllm_config(vllm_config):

131 + # rank=0 here because each DP rank has a single (tp=1,pp=1) process,

132 + # so the local rank within the tp*pp group is always 0.

133 + # init_distributed_environment will offset by data_parallel_rank.

134 + init_distributed_environment(

135 + world_size=1, # tp * pp = 1

136 + rank=0,

137 + distributed_init_method=f"tcp://localhost:{port}",

138 + local_rank=rank,

139 + )

140 + ensure_model_parallel_initialized(1, 1)

141 +

142 +

143 +def _make_forward_context(rank, world_size, num_tokens_per_rank):

144 + """Create a forward context with mock DP metadata for AgRs tests.

145 +

146 + Returns a context manager suitable for ``with`` statements.

147 + The real DPMetadata (with sp_local_sizes etc.) is created internally

148 + by set_forward_context from num_tokens_across_dp; the attn_metadata

149 + placeholder just satisfies the "attn_metadata is not None" guard.

150 + """

151 + from vllm.config.parallel import ParallelConfig

152 + from vllm.config.vllm import VllmConfig

153 + from vllm.forward_context import set_forward_context

154 +

155 + class _AttnMeta:

156 + """Minimal placeholder so set_forward_context's

157 + ``attn_metadata is not None`` guard (forward_context.py:334)

158 + is satisfied. The real DPMetadata is built from num_tokens_across_dp."""

159 +

160 + dp_metadata = None

161 +

162 + vllm_config = VllmConfig()

163 + vllm_config.parallel_config = ParallelConfig(

164 + data_parallel_size=world_size,

165 + is_moe_model=True,

166 + data_parallel_rank=rank,

167 + )

168 + return set_forward_context(

169 + _AttnMeta(),

170 + vllm_config,

171 + num_tokens=num_tokens_per_rank,

172 + num_tokens_across_dp=torch.tensor(

173 + [num_tokens_per_rank] * world_size, dtype=torch.int

174 + ),

175 + )

176 +

177 +

178 +# ---------------------------------------------------------------------------

179 +# Skip conditions

180 +# ---------------------------------------------------------------------------

181 +

182 +requires_multi_gpu = pytest.mark.skipif(

183 + torch.accelerator.device_count() < 2, reason="Need >= 2 GPUs"

184 +)

185 +requires_two_sided = pytest.mark.skipif(

186 + not has_flashinfer_nvlink_two_sided(),

187 + reason="FlashInfer NVLink two-sided not available",

188 +)

189 +requires_one_sided = pytest.mark.skipif(

190 + not has_flashinfer_nvlink_one_sided(),

191 + reason="FlashInfer NVLink one-sided not available",

192 +)

193 +requires_ptrace = pytest.mark.skipif(

194 + not _has_sys_ptrace(),

195 + reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)",

196 +)

197 +

198 +# NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have

199 +# their own @requires_two_sided / @requires_one_sided decorators, and

200 +# test_args_dispatch_combine uses only standard torch.distributed ops and

201 +# should run even when FlashInfer NVLink backends are not installed.

202 +

203 +

204 +# ---------------------------------------------------------------------------

205 +# Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init)

206 +# ---------------------------------------------------------------------------

207 +#

208 +# Tests FlashInferNVLinkTwoSidedManager which wraps FlashInfer's MnnvlMoe.

209 +# initialize() allocates MNNVL shared workspaces via MnnvlMoe.get_moe_workspaces,

210 +# which uses pidfd_getfd() to share memory file descriptors across processes —

211 +# hence the SYS_PTRACE requirement.

212 +#

213 +# Uses EP group (get_ep_group) because the two-sided manager is constructed

214 +# with an EP-scoped communicator in production. With tp=world_size the EP

215 +# group spans all ranks, giving us a multi-rank group for testing.

216 +# ---------------------------------------------------------------------------

217 +

218 +

219 +def _two_sided_lifecycle_worker(rank, world_size):

220 + from vllm.distributed.device_communicators.all2all import (

221 + FlashInferNVLinkTwoSidedManager,

222 + )

223 +

224 + cpu_group = get_ep_group().cpu_group

225 + num_gpus = torch.accelerator.device_count()

226 + manager = FlashInferNVLinkTwoSidedManager(cpu_group)

227 +

228 + # Not initialized yet

229 + assert not manager.initialized

230 + assert manager.rank == rank

231 + assert manager.world_size == world_size

232 +

233 + # Initialize

234 + manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus)

235 + assert manager.initialized

236 + assert manager.workspace_tensor is not None

237 + assert manager.prepare_workspace_tensor is not None

238 + assert manager.mapping is not None

239 +

240 + torch.distributed.barrier()

241 +

242 + # Cleanup

243 + manager.cleanup()

244 + assert not manager.initialized

245 + assert manager.workspace_tensor is None

246 + assert manager.prepare_workspace_tensor is None

247 +

248 + torch.distributed.barrier()

249 +

250 + # Reinitialize

251 + manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus)

252 + assert manager.initialized

253 +

254 + torch.distributed.barrier()

255 +

256 + # ensure_alltoall_workspace_initialized is idempotent when already init'd

257 + assert manager.ensure_alltoall_workspace_initialized()

258 + assert manager.initialized

259 +

260 + manager.cleanup()

261 + assert not manager.initialized

262 +

263 +

264 +@requires_multi_gpu

265 +@requires_two_sided

266 +@requires_ptrace

267 +@pytest.mark.parametrize("world_size", [2])

268 +def test_two_sided_manager_lifecycle(world_size):

269 + """Test init, cleanup, reinit, and ensure_initialized idempotency."""

270 + _spawn_workers(_two_sided_lifecycle_worker, world_size)

271 +

272 +

273 +# ---------------------------------------------------------------------------

274 +# Test 2: One-sided manager lifecycle (init, cleanup, reinit)

275 +# ---------------------------------------------------------------------------

276 +#

277 +# Tests FlashInferNVLinkOneSidedManager which wraps FlashInfer's MoeAlltoAll.

278 +# initialize() creates MoeAlltoAll with an MnnvlConfig, which allocates MNNVL

279 +# shared workspaces — same cross-process memory sharing as two-sided, hence

280 +# the SYS_PTRACE requirement.

281 +#

282 +# Uses DP group (get_dp_group) because the one-sided manager's initialize()

283 +# internally calls get_dp_group() to set up the MnnvlConfig communicator.

284 +# We therefore need a real DP group with world_size > 1, which requires

285 +# dp_size=world_size via _init_dp_environment.

286 +# ---------------------------------------------------------------------------

287 +

288 +

289 +def _one_sided_lifecycle_worker(rank, world_size):

290 + from vllm.distributed.device_communicators.all2all import (

291 + FlashInferNVLinkOneSidedManager,

292 + )

293 + from vllm.distributed.parallel_state import get_dp_group

294 +

295 + cpu_group = get_dp_group().cpu_group

296 + manager = FlashInferNVLinkOneSidedManager(cpu_group)

297 +

298 + assert not manager.initialized

299 + assert manager.rank == rank

300 + assert manager.world_size == world_size

301 +

302 + init_kwargs = dict(

303 + max_num_tokens=1024,

304 + top_k=2,

305 + num_experts=world_size * 8,

306 + hidden_size=4096,

307 + )

308 +

309 + # Initialize

310 + manager.initialize(**init_kwargs)

311 + assert manager.initialized

312 + assert manager.moe_alltoall is not None

313 + assert manager.mapping is not None

314 +

315 + torch.distributed.barrier()

316 +

317 + # Cleanup

318 + manager.cleanup()

319 + assert not manager.initialized

320 + assert manager.moe_alltoall is None

321 +

322 + torch.distributed.barrier()

323 +

324 + # Reinitialize with different token count

325 + manager.initialize(**{**init_kwargs, "max_num_tokens": 2048})

326 + assert manager.initialized

327 +

328 + torch.distributed.barrier()

329 + manager.cleanup()

330 +

331 +

332 +@requires_multi_gpu

333 +@requires_one_sided

334 +@requires_ptrace

335 +@pytest.mark.parametrize("world_size", [2])

336 +def test_one_sided_manager_lifecycle(world_size):

337 + """Test init, cleanup, and reinit with different params."""

338 + _spawn_workers(

339 + _one_sided_lifecycle_worker,

340 + world_size,

341 + dp_size=world_size,

342 + )

343 +

344 +

345 +# ---------------------------------------------------------------------------

346 +# Test 3: AgRs dispatch/combine with value validation

347 +# ---------------------------------------------------------------------------

348 +#

349 +# Tests AgRsAll2AllManager which uses only standard torch.distributed

350 +# all_gatherv / reduce_scatterv — no FlashInfer or MNNVL dependency.

351 +# This test validates the reference all-to-all implementation that other

352 +# backends are compared against.

353 +# ---------------------------------------------------------------------------

354 +

355 +

356 +def _args_dispatch_combine_worker(rank, world_size):

357 + from vllm.distributed.device_communicators.all2all import AgRsAll2AllManager

358 + from vllm.forward_context import get_forward_context

359 +

360 + cpu_group = get_ep_group().cpu_group

361 + device = torch.device(f"cuda:{rank}")

362 +

363 + hidden_size = 64

364 + tokens_per_rank = 16

365 + experts_per_token = 2

366 + num_experts = world_size * 4

367 + total_tokens = world_size * tokens_per_rank

368 +

369 + # Deterministic per-rank data: rank r has value (r + 1)

370 + hidden = torch.full(

371 + (tokens_per_rank, hidden_size),

372 + float(rank + 1),

373 + device=device,

374 + dtype=torch.float32,

375 + )

376 + router = torch.full(

377 + (tokens_per_rank, num_experts),

378 + float(rank + 1) * 10,

379 + device=device,

380 + dtype=torch.float32,

381 + )

382 + weights = torch.full(

383 + (tokens_per_rank, experts_per_token),

384 + float(rank + 1) * 100,

385 + device=device,

386 + dtype=torch.float32,

387 + )

388 + ids = torch.full(

389 + (tokens_per_rank, experts_per_token),

390 + rank,

391 + device=device,

392 + dtype=torch.long,

393 + )

394 +

395 + with _make_forward_context(rank, world_size, tokens_per_rank):

396 + manager = AgRsAll2AllManager(cpu_group)

397 + dp_metadata = get_forward_context().dp_metadata

398 +

399 + with dp_metadata.sp_local_sizes(sequence_parallel_size=1):

400 + # -- dispatch_router_logits --

401 + d_hidden, d_router = manager.dispatch_router_logits(

402 + hidden.clone(),

403 + router.clone(),

404 + is_sequence_parallel=True,

405 + )

406 + assert d_hidden.shape == (total_tokens, hidden_size)

407 + assert d_router.shape == (total_tokens, num_experts)

408 +

409 + for r in range(world_size):

410 + s = r * tokens_per_rank

411 + e = (r + 1) * tokens_per_rank

412 + torch.testing.assert_close(

413 + d_hidden[s:e],

414 + torch.full_like(d_hidden[s:e], float(r + 1)),

415 + )

416 + torch.testing.assert_close(

417 + d_router[s:e],

418 + torch.full_like(d_router[s:e], float(r + 1) * 10),

419 + )

420 +

421 + # -- dispatch --

422 + d_hidden2, d_weights, d_ids = manager.dispatch(

423 + hidden.clone(),

424 + weights.clone(),

425 + ids.clone(),

426 + is_sequence_parallel=True,

427 + )

428 + assert d_hidden2.shape == (total_tokens, hidden_size)

429 + assert d_weights.shape == (total_tokens, experts_per_token)

430 + assert d_ids.shape == (total_tokens, experts_per_token)

431 +

432 + for r in range(world_size):

433 + s = r * tokens_per_rank

434 + e = (r + 1) * tokens_per_rank

435 + torch.testing.assert_close(

436 + d_weights[s:e],

437 + torch.full_like(d_weights[s:e], float(r + 1) * 100),

438 + )

439 + assert (d_ids[s:e] == r).all()

440 +

441 + # -- combine (reduce-scatter) --

442 + # Each token i has value i in all columns; after reduce-scatter

443 + # each rank gets its slice, summed across ranks.

444 + expert_out = (

445 + torch.arange(total_tokens, device=device, dtype=torch.float32)

446 + .unsqueeze(1)

447 + .expand(total_tokens, hidden_size)

448 + .contiguous()

449 + )

450 +

451 + combined = manager.combine(expert_out, is_sequence_parallel=True)

452 + assert combined.shape == (tokens_per_rank, hidden_size)

453 +

454 + for i in range(tokens_per_rank):

455 + expected_val = float(rank * tokens_per_rank + i) * world_size

456 + torch.testing.assert_close(

457 + combined[i],

458 + torch.full_like(combined[i], expected_val),

459 + )

460 +

461 + torch.distributed.barrier()

462 +

463 +

464 +@requires_multi_gpu

465 +@pytest.mark.parametrize("world_size", [2])

466 +def test_args_dispatch_combine(world_size):

467 + """Validate dispatch gathers all-rank data and combine reduces correctly."""

468 + _spawn_workers(_args_dispatch_combine_worker, world_size)

469 +

470 +

471 +# ---------------------------------------------------------------------------

472 +# Test 4: FlashInfer two-sided dispatch/combine data communication

473 +# ---------------------------------------------------------------------------

474 +#

475 +# Tests actual data flow through the FlashInfer NVLink two-sided backend

476 +# by calling flashinfer_alltoall_dispatch (with defer_input_quant=True to

477 +# skip quantization) and flashinfer_alltoall_combine, then verifying exact

478 +# round-trip values. Dispatch sends each token once per distinct expert

479 +# rank, and combine performs an unweighted sum, so:

480 +# dispatch(hidden) → identity → combine = hidden * num_distinct_ranks(i)

481 +# ---------------------------------------------------------------------------

482 +

483 +

484 +def _two_sided_data_worker(rank, world_size):

485 + from vllm.distributed.device_communicators.all2all import (

486 + FlashInferNVLinkTwoSidedManager,

487 + )

488 + from vllm.distributed.parallel_state import get_dp_group

489 + from vllm.forward_context import get_forward_context

490 + from vllm.model_executor.layers.fused_moe.config import (

491 + FusedMoEQuantConfig,

492 + FusedMoEQuantDesc,

493 + )

494 + from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501

495 + flashinfer_alltoall_combine,

496 + flashinfer_alltoall_dispatch,

497 + )

498 +

499 + # Use DP group because MnnvlMoe workspace allocation calls get_dp_group()

500 + # internally and requires dp_size == ep_size.

501 + cpu_group = get_dp_group().cpu_group

502 + device = torch.device(f"cuda:{rank}")

503 + num_gpus = torch.accelerator.device_count()

504 +

505 + hidden_size = 128

506 + tokens_per_rank = 32

507 + experts_per_token = 2

508 + num_experts = world_size * 4

509 +

510 + # Initialize the FlashInfer two-sided manager

511 + manager = FlashInferNVLinkTwoSidedManager(cpu_group)

512 + manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus)

513 + assert manager.initialized

514 +

515 + torch.distributed.barrier()

516 +

517 + # Create deterministic per-rank test data

518 + torch.manual_seed(rank + 42)

519 + hidden = torch.randn(

520 + tokens_per_rank,

521 + hidden_size,

522 + device=device,

523 + dtype=torch.bfloat16,

524 + )

525 + # Assign each token to experts spread across ranks so tokens move between GPUs

526 + topk_ids = torch.randint(

527 + 0,

528 + num_experts,

529 + (tokens_per_rank, experts_per_token),

530 + device=device,

531 + dtype=torch.int32,

532 + )

533 + topk_weights = torch.rand(

534 + tokens_per_rank,

535 + experts_per_token,

536 + device=device,

537 + dtype=torch.float32,

538 + )

539 +

540 + # Unquantized config: quant_dtype=None means moe_kernel_quantize_input is a no-op

541 + no_quant = FusedMoEQuantDesc()

542 + quant_config = FusedMoEQuantConfig(

543 + _a1=no_quant,

544 + _a2=no_quant,

545 + _w1=no_quant,

546 + _w2=no_quant,

547 + )

548 + assert quant_config.quant_dtype is None # sanity: no quantization

549 +

550 + with _make_forward_context(rank, world_size, tokens_per_rank):

551 + dp_metadata = get_forward_context().dp_metadata

552 +

553 + with dp_metadata.sp_local_sizes(sequence_parallel_size=1):

554 + local_sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

555 +

556 + # --- FlashInfer two-sided dispatch ---

557 + alltoall_info, fi_topk_ids, fi_topk_weights, fi_hidden, fi_scale = (

558 + flashinfer_alltoall_dispatch(

559 + manager,

560 + local_sizes,

561 + hidden.clone(),

562 + None, # no global scale

563 + topk_ids.clone(),

564 + topk_weights.clone(),

565 + experts_per_token,

566 + num_experts,

567 + quant_config,

568 + defer_input_quant=True,

569 + )

570 + )

571 + assert fi_scale is None # deferred quant: no scale produced

572 + assert fi_hidden is not None

573 + assert fi_hidden.shape[1] == hidden_size

574 + assert fi_hidden.numel() > 0

575 +

576 + # --- Round-trip exact verification ---

577 + # The all-to-all sends each token once per *distinct* expert

578 + # rank. Combine performs an unweighted sum of the per-rank

579 + # contributions. With identity expert (feeding dispatched

580 + # hidden straight back):

581 + # result[i] = hidden[i] * num_distinct_expert_ranks(i)

582 + combined = flashinfer_alltoall_combine(

583 + manager,

584 + fi_hidden,

585 + top_k=experts_per_token,

586 + token_count=tokens_per_rank,

587 + alltoall_info=alltoall_info,

588 + )

589 + assert combined.shape == (tokens_per_rank, hidden_size)

590 +

591 + experts_per_rank = num_experts // world_size

592 + expert_ranks = topk_ids // experts_per_rank # (tokens, top_k)

593 + num_distinct = torch.tensor(

594 + [len(set(row.tolist())) for row in expert_ranks],

595 + device=device,

596 + dtype=torch.float32,

597 + ).unsqueeze(1) # (tokens, 1)

598 + expected = (hidden.float() * num_distinct).to(hidden.dtype)

599 + torch.testing.assert_close(combined, expected)

600 +

601 + # --- Linearity check with scaled expert output ---

602 + # Scaling the expert output by a constant should scale the

603 + # combined result by the same constant.

604 + scale = 3.0

605 + combined_scaled = flashinfer_alltoall_combine(

606 + manager,

607 + fi_hidden * scale,

608 + top_k=experts_per_token,

609 + token_count=tokens_per_rank,

610 + alltoall_info=alltoall_info,

611 + )

612 + expected_scaled = (hidden.float() * num_distinct * scale).to(hidden.dtype)

613 + torch.testing.assert_close(combined_scaled, expected_scaled)

614 +

615 + torch.distributed.barrier()

616 +

617 + manager.cleanup()

618 +

619 +

620 +@requires_multi_gpu

621 +@requires_two_sided

622 +@requires_ptrace

623 +@pytest.mark.parametrize("world_size", [2])

624 +def test_two_sided_dispatch_combine(world_size):

625 + """Test FlashInfer two-sided dispatch/combine with exact value verification."""

626 + _spawn_workers(_two_sided_data_worker, world_size, dp_size=world_size)

627 +

628 +

629 +# ---------------------------------------------------------------------------

630 +# Test 5: FlashInfer one-sided dispatch/combine data communication

631 +# ---------------------------------------------------------------------------

632 +#

633 +# Tests actual data flow through the FlashInfer NVLink one-sided backend

634 +# by calling MoeAlltoAll.dispatch() and MoeAlltoAll.combine() directly

635 +# with synthetic payloads, then verifying shapes and round-trip consistency.

636 +# ---------------------------------------------------------------------------

637 +

638 +

639 +def _one_sided_data_worker(rank, world_size):

640 + from vllm.distributed.device_communicators.all2all import (

641 + FlashInferNVLinkOneSidedManager,

642 + )

643 + from vllm.distributed.parallel_state import get_dp_group

644 + from vllm.forward_context import get_forward_context

645 +

646 + cpu_group = get_dp_group().cpu_group

647 + device = torch.device(f"cuda:{rank}")

648 +

649 + hidden_size = 256

650 + tokens_per_rank = 32

651 + experts_per_token = 2

652 + num_experts = world_size * 8

653 +

654 + # Initialize the one-sided manager

655 + manager = FlashInferNVLinkOneSidedManager(cpu_group)

656 + manager.initialize(

657 + max_num_tokens=tokens_per_rank,

658 + top_k=experts_per_token,

659 + num_experts=num_experts,

660 + hidden_size=hidden_size,

661 + )

662 + assert manager.initialized

663 + assert manager.moe_alltoall is not None

664 +

665 + with _make_forward_context(rank, world_size, tokens_per_rank):

666 + dp_metadata = get_forward_context().dp_metadata

667 +

668 + with dp_metadata.sp_local_sizes(sequence_parallel_size=1):

669 + local_sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

670 + runtime_max_tokens = max(local_sizes)

671 +

672 + # Create test data with raw tensors matching the nvfp4 payload

673 + # sizes the workspace was allocated for:

674 + # a1q: (tokens, hidden_size // 2) — nvfp4 hidden states

675 + # a1q_scale: (tokens, hidden_size // 16) — fp8 scaling factors

676 + torch.manual_seed(rank + 42)

677 + a1q = torch.randint(

678 + 0,

679 + 256,

680 + (tokens_per_rank, hidden_size // 2),

681 + device=device,

682 + dtype=torch.uint8,

683 + )

684 + a1q_scale = torch.randint(

685 + 0,

686 + 256,

687 + (tokens_per_rank, hidden_size // 16),

688 + device=device,

689 + dtype=torch.uint8,

690 + )

691 + topk_ids = torch.randint(

692 + 0,

693 + num_experts,

694 + (tokens_per_rank, experts_per_token),

695 + device=device,

696 + dtype=torch.int32,

697 + )

698 + topk_weights = torch.rand(

699 + tokens_per_rank,

700 + experts_per_token,

701 + device=device,

702 + dtype=torch.float32,

703 + )

704 +

705 + # --- One-sided dispatch ---

706 + payloads = [a1q, a1q_scale, topk_ids, topk_weights]

707 + recv_payloads = manager.moe_alltoall.dispatch(

708 + token_selected_experts=topk_ids,

709 + input_payloads=payloads,

710 + runtime_max_tokens_per_rank=runtime_max_tokens,

711 + )

712 + assert len(recv_payloads) == 4

713 + recv_a1q, recv_scale, recv_ids, recv_weights = recv_payloads

714 + assert recv_a1q.numel() > 0

715 + assert recv_ids.numel() > 0

716 +

717 + # --- Round-trip exact verification ---

718 + # The dispatch routes each token once per *distinct* expert

719 + # rank. Combine performs an unweighted sum of per-rank

720 + # contributions. With constant expert output (all 1s):

721 + # result[i] = 1.0 * num_distinct_expert_ranks(i)

722 + expert_output = torch.ones(

723 + world_size,

724 + runtime_max_tokens,

725 + hidden_size,

726 + device=device,

727 + dtype=torch.bfloat16,

728 + )

729 + combined = manager.moe_alltoall.combine(

730 + payload=expert_output,

731 + runtime_max_tokens_per_rank=runtime_max_tokens,

732 + )

733 + assert combined.shape == (tokens_per_rank, hidden_size)

734 +

735 + experts_per_rank = num_experts // world_size

736 + expert_ranks = topk_ids // experts_per_rank # (tokens, top_k)

737 + num_distinct = torch.tensor(

738 + [len(set(row.tolist())) for row in expert_ranks],

739 + device=device,

740 + dtype=torch.bfloat16,

741 + ).unsqueeze(1) # (tokens, 1)

742 + expected = num_distinct.expand_as(combined)

743 + torch.testing.assert_close(combined, expected)

744 +

745 + # --- Linearity check with scaled expert output ---

746 + # Scaling the expert output by a constant should scale the

747 + # combined result by the same constant.

748 + # Re-dispatch to reset internal state (one-sided requires a

749 + # fresh dispatch before each combine).

750 + manager.moe_alltoall.dispatch(

751 + token_selected_experts=topk_ids,

752 + input_payloads=payloads,

753 + runtime_max_tokens_per_rank=runtime_max_tokens,

754 + )

755 + scale = 3.0

756 + combined_scaled = manager.moe_alltoall.combine(

757 + payload=expert_output * scale,

758 + runtime_max_tokens_per_rank=runtime_max_tokens,

759 + )

760 + expected_scaled = (expected * scale).to(torch.bfloat16)

761 + torch.testing.assert_close(combined_scaled, expected_scaled)

762 +

763 + torch.distributed.barrier()

764 +

765 + manager.cleanup()

766 +

767 +

768 +@requires_multi_gpu

769 +@requires_one_sided

770 +@requires_ptrace

771 +@pytest.mark.parametrize("world_size", [2])

772 +def test_one_sided_dispatch_combine(world_size):

773 + """Test FlashInfer one-sided dispatch/combine with actual data flow."""

774 + _spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size)

Create tests/distributed/test_mnnvl_alltoall.py (#35241)