From d05b59db372e818f9aeef278246ec61a4f05d4bb Mon Sep 17 00:00:00 2001 From: Satyam Srivastava Date: Sat, 5 Sep 2026 23:04:12 -0700 Subject: [PATCH] Add H3 inter node profiling --- .../minimax_h3_disaggregated_spark_pair.md | 56 ++++- fastvideo/envs.py | 6 + .../worker/test_minimax_h3_disaggregated.py | 1 + .../worker/test_minimax_h3_transfer_timing.py | 195 ++++++++++++++++++ fastvideo/worker/minimax_h3_disaggregated.py | 184 ++++++++++++----- 5 files changed, 387 insertions(+), 55 deletions(-) create mode 100644 fastvideo/tests/worker/test_minimax_h3_transfer_timing.py diff --git a/docs/getting_started/installation/minimax_h3_disaggregated_spark_pair.md b/docs/getting_started/installation/minimax_h3_disaggregated_spark_pair.md index 0d1f26e5d..1871c1343 100644 --- a/docs/getting_started/installation/minimax_h3_disaggregated_spark_pair.md +++ b/docs/getting_started/installation/minimax_h3_disaggregated_spark_pair.md @@ -4,7 +4,7 @@ This runbook tests the component-disaggregated MiniMax H3 inference path: - Spark A: text/multimodal encoder, video VAE, audio VAE, and the driver. - Spark B: DiT denoising only. -- Ray transports CPU-contiguous conditioning and latent tensors between the two roles. +- Ray transports CUDA-contiguous conditioning and latent tensors between the two roles. - Each role uses one GPU with tensor parallelism and sequence parallelism disabled. - All role components remain resident until the generator shuts down. @@ -510,6 +510,60 @@ For one generator lifetime: - No request should release or reload the text encoder, DiT, or VAEs. - Spark B should continue denoising without waiting for Spark A's VAE decode. +### Measure the two payload handoffs + +Enable transfer profiling on the **driver before creating the generator**: + +~~~bash +FASTVIDEO_H3_PROFILE_TRANSFERS=1 \ + fastvideo generate --config /tmp/h3-disaggregated-smoke.yaml +~~~ + +The driver passes this setting to both actors, including when connecting to an +existing Ray cluster. It also applies to `iter_forward` and its async wrapper. +Look for `[RAY_TRANSFER]`, `[RAY_STAGE]`, and `[RAY_PAYLOAD]` in the worker logs. +Every line includes a `request_id` to correlate concurrent requests. + +Each request logs an `A_TO_B` encoded-payload receipt on the DiT worker and a +`B_TO_A` denoised-payload receipt on the encoder/decoder worker: + +| Field | What it measures | +|---|---| +| `source_wait_s` | Receiver wait until the producer's result is available somewhere in Ray. Includes any remaining producer work, serialization, and readiness notification. | +| `object_fetch_s` | Wait to make that ready object available in the receiver's local object store. Includes Ray transfer/control overhead and any object-store restore; a cached object can be nearly immediate. | +| `materialize_s` | Local `ray.get` deserialization and tensor reconstruction, including waiting for CUDA copies to complete. | +| `receive_s` | `object_fetch_s + materialize_s`, excluding the upstream wait. | +| `tensor_bytes` | Logical bytes in the conditioning/latent tensors and all five layout tensors. Excludes serialization metadata, backing-storage overhead, and network protocol bytes. | +| `object_fetch_mb_s` | Logical tensor MB divided by `object_fetch_s`; an effective payload rate, not measured NIC bandwidth. | + +`[RAY_STAGE]` separately reports synchronized `elapsed_s` for `encode`, +`denoise`, and `decode`. Large `object_fetch_s` points toward the Ray data path; +large `materialize_s` points toward reconstruction/device-copy cost. Large +`source_wait_s` alone does not imply a slow network: compare it with the +producer's stage time. Actor queueing before a receiving method starts is not +part of these receiver timers. Sender serialization occurs after the producer's +stage timer ends and may overlap other work. + +The timers use a monotonic clock on the receiving Spark, so they do not require +synchronized clocks between machines. Profiling passes a nested ObjectRef and +uses `ray.wait(fetch_local=False)`, then `ray.wait(fetch_local=True)`, then +`ray.get` to separate the three phases. See Ray's +[object-reference behavior](https://docs.ray.io/en/latest/ray-core/objects.html#passing-object-arguments) +and [wait semantics](https://docs.ray.io/en/latest/ray-core/api/doc/ray.wait.html). + +This is an opt-in diagnostic: deferring the fetch until receiver entry changes +prefetch overlap, and CUDA synchronization adds overhead. Compare a short warmed +run with profiling disabled before drawing throughput conclusions. The default +execution keeps Ray's automatic argument fetching. These application timings do +not isolate pure time on the wire, and CUDA payloads alone do not establish that +the transport uses RDMA or GPU-direct transfers. + +For an external CUDA trace, set `FASTVIDEO_NVTX_PROFILE=1` in the Ray worker +environment before starting Ray. With transfer profiling enabled, NVTX ranges +include `h3.A_TO_B.source_wait`, `h3.A_TO_B.object_fetch`, +`h3.A_TO_B.materialize` (and the corresponding `B_TO_A` ranges), plus +`h3.encode`, `h3.denoise`, and `h3.decode`. + ## 9. Optional FastH3/VSA test after correctness passes The first smoke test intentionally uses the base MiniMax H3 checkpoint with diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 86abc1bcd..86fc9ad9b 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -48,6 +48,7 @@ if TYPE_CHECKING: FASTVIDEO_TRACE_STEPS: str = "" FASTVIDEO_SERVER_DEV_MODE: bool = False FASTVIDEO_STAGE_LOGGING: bool = False + FASTVIDEO_H3_PROFILE_TRANSFERS: bool = False FASTVIDEO_CFG_GATE_STEP: float = 1.0 FASTVIDEO_HOST_IP: str = "" FASTVIDEO_LOOPBACK_IP: str = "" @@ -343,6 +344,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "FASTVIDEO_STAGE_LOGGING": lambda: bool(int(os.getenv("FASTVIDEO_STAGE_LOGGING", "0"))), + # Opt-in receiver-side Ray fetch/materialization and synchronized stage + # timings for component-disaggregated H3. This changes prefetch overlap. + "FASTVIDEO_H3_PROFILE_TRANSFERS": + lambda: bool(int(os.getenv("FASTVIDEO_H3_PROFILE_TRANSFERS", "0"))), + # CFG gating fraction for stale-uncond reuse (Adaptive Guidance / LinearAG # variant — Castillo et al. 2023, arXiv:2312.12487). Float in [0, 1]. # Interpretation: for step index `i < len(timesteps) * X`, run both diff --git a/fastvideo/tests/worker/test_minimax_h3_disaggregated.py b/fastvideo/tests/worker/test_minimax_h3_disaggregated.py index 08ff48d43..eeb195bbb 100644 --- a/fastvideo/tests/worker/test_minimax_h3_disaggregated.py +++ b/fastvideo/tests/worker/test_minimax_h3_disaggregated.py @@ -481,6 +481,7 @@ def _fake_runtime(monkeypatch): monkeypatch.setattr(disaggregated_runtime, "ray", fake_ray) runtime = RayMiniMaxH3DisaggregatedRuntime.__new__(RayMiniMaxH3DisaggregatedRuntime) runtime._closed = False + runtime._profile_transfers = False runtime.encoder_decoder = _FakeEncoderActor(events) runtime.dit = _FakeDiTActor(events) return runtime, fake_ray, events diff --git a/fastvideo/tests/worker/test_minimax_h3_transfer_timing.py b/fastvideo/tests/worker/test_minimax_h3_transfer_timing.py new file mode 100644 index 000000000..1e511a5b6 --- /dev/null +++ b/fastvideo/tests/worker/test_minimax_h3_transfer_timing.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU regressions for H3 Ray timing; no Ray cluster or CUDA device required.""" + +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +import fastvideo.worker.minimax_h3_disaggregated as runtime_module + + +def _state(): + return SimpleNamespace( + request_id="request-7", + prompt_embeds=torch.zeros(1, 3, 4), + video_latents=torch.zeros(4, 5), + audio_latents=torch.zeros(2, 6), + layout=SimpleNamespace( + position_ids=torch.zeros(9, 3, dtype=torch.float64), + token_tags=torch.zeros(9, dtype=torch.int64), + text_indices=torch.zeros(3, dtype=torch.int64), + video_indices=torch.zeros(4, dtype=torch.int64), + audio_indices=torch.zeros(2, dtype=torch.int64), + ), + ) + + +@pytest.fixture +def timing(monkeypatch): + clock = SimpleNamespace(now=0.0) + events = [] + messages = [] + result = _state() + + def wait(refs, *, num_returns, fetch_local): + assert num_returns == 1 + events.append(("wait", fetch_local)) + clock.now += 2.0 if fetch_local else 10.0 + return refs, [] + + def get(ref): + events.append(("get", ref)) + clock.now += 3.0 + return result + + def synchronize(): + events.append(("sync",)) + clock.now += 5.0 + + monkeypatch.setattr(runtime_module.time, "perf_counter", lambda: clock.now) + monkeypatch.setattr(runtime_module, "ray", SimpleNamespace(wait=wait, get=get)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + monkeypatch.setattr(runtime_module.envs, "FASTVIDEO_NVTX_PROFILE", False) + monkeypatch.setattr(runtime_module.logger, "info", lambda fmt, *args: messages.append(fmt % args)) + return SimpleNamespace(clock=clock, events=events, messages=messages, result=result) + + +def test_payload_total_includes_every_layout_tensor(timing): + assert runtime_module._log_h3_state("encoded", timing.result) == 536 + assert any("layout.position_ids:" in line for line in timing.messages) + assert any("layout.token_tags:" in line for line in timing.messages) + assert all("request_id=request-7" in line for line in timing.messages) + + +@pytest.mark.parametrize("direction", ["A_TO_B", "B_TO_A"]) +def test_transfer_separates_production_fetch_and_completed_device_copy(timing, direction): + ref = object() + result = runtime_module._receive_h3_state(runtime_module._H3TransferRef(ref, "request-7"), direction) + assert result is timing.result + assert timing.events == [("sync",), ("wait", False), ("wait", True), ("get", ref), ("sync",)] + line, = timing.messages + assert f"direction={direction}" in line + assert "request_id=request-7" in line + assert "tensor_bytes=536" in line + assert "source_wait_s=10.000000" in line + assert "object_fetch_s=2.000000" in line + assert "materialize_s=8.000000" in line + assert "receive_s=10.000000" in line + + +def test_disabled_profiling_does_not_wait_synchronize_or_read_clock(timing, monkeypatch): + def unexpected_clock(): + pytest.fail("profiling disabled must not read the clock") + + monkeypatch.setattr(runtime_module.time, "perf_counter", unexpected_clock) + assert runtime_module._receive_h3_state(timing.result, "A_TO_B") is timing.result + with runtime_module._h3_stage_timer("encode", "request-7", False): + pass + assert timing.events == [] + assert timing.messages == [] + + +def test_stage_time_includes_gpu_completion_and_excludes_previous_work(timing): + with runtime_module._h3_stage_timer("denoise", "request-7", True): + timing.clock.now += 7.0 + assert timing.events == [("sync",), ("sync",)] + assert timing.messages == ["[RAY_STAGE] request_id=request-7 stage=denoise elapsed_s=12.000000"] + + +def test_failed_transfer_propagates_without_logging_success(timing, monkeypatch): + def fail(ref): + raise RuntimeError("producer failed") + + monkeypatch.setattr(runtime_module.ray, "get", fail) + with pytest.raises(RuntimeError, match="producer failed"): + runtime_module._receive_h3_state(runtime_module._H3TransferRef(object(), "request-7"), "A_TO_B") + assert timing.messages == [] + + +def test_zero_fetch_duration_does_not_divide_by_zero(timing, monkeypatch): + monkeypatch.setattr(runtime_module.time, "perf_counter", lambda: 0.0) + runtime_module._receive_h3_state(runtime_module._H3TransferRef(object(), "request-7"), "A_TO_B") + assert "object_fetch_mb_s=n/a" in timing.messages[0] + + +@pytest.mark.parametrize("stage,direction", [("denoise", "A_TO_B"), ("decode", "B_TO_A")]) +def test_actor_receives_before_model_stage(timing, stage, direction): + actor_class = (runtime_module._MiniMaxH3DiTActor if stage == "denoise" + else runtime_module._MiniMaxH3EncoderDecoderActor) + actor = actor_class.__new__(actor_class) + actor._profile_transfers = True + + def forward(state): + assert state is timing.result + assert any(f"direction={direction}" in line for line in timing.messages) + timing.clock.now += 7.0 + return state + + actor.pipeline = SimpleNamespace(**{stage: forward}) + assert getattr(actor, stage)(runtime_module._H3TransferRef(object(), "request-7")) is timing.result + assert any(f"stage={stage} elapsed_s=12.000000" in line for line in timing.messages) + + +def _fake_runtime(monkeypatch, enabled): + events = [] + + def method(name): + def remote(*args): + ref = SimpleNamespace(kind=name, args=args) + events.append((name, ref)) + return ref + + return SimpleNamespace(remote=remote) + + def get(ref): + # The driver must never materialize encode/denoise intermediates. + assert ref.kind == "decode" + events.append(("get", ref)) + return ref + + def wait(refs, *, num_returns, fetch_local): + assert refs[0].kind == "denoise" + assert num_returns == 1 and fetch_local is False + events.append(("wait", refs[0])) + return refs, [] + + runtime = runtime_module.RayMiniMaxH3DisaggregatedRuntime.__new__( + runtime_module.RayMiniMaxH3DisaggregatedRuntime) + runtime._closed = False + runtime._profile_transfers = enabled + runtime.encoder_decoder = SimpleNamespace(encode=method("encode"), decode=method("decode")) + runtime.dit = SimpleNamespace(denoise=method("denoise")) + monkeypatch.setattr(runtime_module, "ray", SimpleNamespace(get=get, wait=wait)) + return runtime, events + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +def test_runtime_preserves_request_ids_and_bounded_driver_dag(monkeypatch, enabled, stream): + runtime, events = _fake_runtime(monkeypatch, enabled) + batches = [ForwardBatch(data_type="video", extra={"request_id": str(i)}) for i in range(3)] + if stream: + assert len(list(runtime.iter_forward(batches))) == 3 + assert [name for name, _ in events] == [ + "encode", "denoise", "encode", "wait", "decode", "denoise", "get", + "encode", "wait", "decode", "denoise", "get", "decode", "get", + ] + else: + runtime.execute_forward(batches[0], request_id="explicit") + assert [name for name, _ in events] == ["encode", "denoise", "decode", "get"] + + for name, result_ref in events: + if name not in {"denoise", "decode"}: + continue + arg, = result_ref.args + assert isinstance(arg, runtime_module._H3TransferRef) is enabled + source_ref = arg.ref if enabled else arg + assert source_ref.kind == ("encode" if name == "denoise" else "denoise") + if enabled: + expected_id = source_ref.args[1] if name == "denoise" else source_ref.args[0].request_id + assert arg.request_id == expected_id + if not stream: + assert arg.request_id == "explicit" diff --git a/fastvideo/worker/minimax_h3_disaggregated.py b/fastvideo/worker/minimax_h3_disaggregated.py index fba6dafdb..7e9c819ab 100644 --- a/fastvideo/worker/minimax_h3_disaggregated.py +++ b/fastvideo/worker/minimax_h3_disaggregated.py @@ -5,15 +5,18 @@ from __future__ import annotations import asyncio from collections.abc import AsyncIterator, Iterable, Iterator -from contextlib import suppress +from contextlib import contextmanager, suppress from copy import deepcopy +from dataclasses import dataclass import os from queue import Queue +import time from typing import Any, cast from uuid import uuid4 import torch +import fastvideo.envs as envs from fastvideo.distributed import cleanup_dist_env_and_memory from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger @@ -27,61 +30,118 @@ from fastvideo.pipelines.basic.minimax_h3.disaggregated import ( ) from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.lazy_module import is_lazy_module +from fastvideo.profiler import nvtx_range from fastvideo.utils import get_ip, get_open_port from fastvideo.worker.executor import Executor from fastvideo.worker.ray_utils import assert_ray_available, ray logger = init_logger(__name__) -def _log_h3_state(name: str, state: Any) -> None: - total_bytes = 0 - for field_name in ( - "prompt_embeds", - "video_latents", - "audio_latents", - ): +def _h3_state_tensors(state: Any) -> Iterator[tuple[str, torch.Tensor]]: + for field_name in ("prompt_embeds", "video_latents", "audio_latents"): tensor = getattr(state, field_name, None) if tensor is not None: - size = tensor.numel() * tensor.element_size() - total_bytes += size - logger.info( - "[RAY_PAYLOAD] %s.%s: %.2f MB shape=%s dtype=%s device=%s", - name, - field_name, - size / 1e6, - tuple(tensor.shape), - tensor.dtype, - tensor.device, - ) + yield field_name, tensor - # layout contains tensors too. layout = getattr(state, "layout", None) if layout is not None: - for field_name in ( - "text_indices", - "video_indices", - "audio_indices", - ): + for field_name in ("position_ids", "token_tags", "text_indices", "video_indices", "audio_indices"): tensor = getattr(layout, field_name, None) if tensor is not None: - size = tensor.numel() * tensor.element_size() - total_bytes += size - logger.info( - "[RAY_PAYLOAD] %s.layout.%s: %.2f MB shape=%s dtype=%s device=%s", - name, - field_name, - size / 1e6, - tuple(tensor.shape), - tensor.dtype, - tensor.device, - ) + yield f"layout.{field_name}", tensor + + +def _log_h3_state(name: str, state: Any) -> int: + total_bytes = 0 + for field_name, tensor in _h3_state_tensors(state): + size = tensor.numel() * tensor.element_size() + total_bytes += size + logger.info( + "[RAY_PAYLOAD] %s.%s: %.2f MB shape=%s dtype=%s device=%s request_id=%s", + name, + field_name, + size / 1e6, + tuple(tensor.shape), + tensor.dtype, + tensor.device, + state.request_id, + ) logger.info( - "[RAY_PAYLOAD] %s TOTAL: %.2f MB", + "[RAY_PAYLOAD] %s TOTAL: %.2f MB request_id=%s (logical tensor bytes)", name, total_bytes / 1e6, + state.request_id, ) + return total_bytes + + +@dataclass(frozen=True) +class _H3TransferRef: + """Keep the ref nested so Ray does not fetch it before actor entry.""" + + ref: Any + request_id: str + + +def _synchronize_h3_device() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _receive_h3_state(state: Any, direction: str) -> Any: + if not isinstance(state, _H3TransferRef): + return state + + # A receiver-local clock avoids clock skew between the two Sparks. Drain + # earlier device work before timing, then wait for production WITHOUT + # fetching locally so producer compute is not counted as object fetch. + _synchronize_h3_device() + started = time.perf_counter() + with nvtx_range(f"h3.{direction}.source_wait"): + ray.wait([state.ref], num_returns=1, fetch_local=False) + source_ready = time.perf_counter() + with nvtx_range(f"h3.{direction}.object_fetch"): + ray.wait([state.ref], num_returns=1, fetch_local=True) + fetched = time.perf_counter() + with nvtx_range(f"h3.{direction}.materialize"): + result = ray.get(state.ref) + _synchronize_h3_device() + materialized = time.perf_counter() + + tensor_bytes = sum(tensor.numel() * tensor.element_size() for _, tensor in _h3_state_tensors(result)) + object_fetch_s = fetched - source_ready + # This is an effective logical-payload rate, not measured NIC bandwidth. + rate = f"{tensor_bytes / 1e6 / object_fetch_s:.2f}" if object_fetch_s > 0 else "n/a" + logger.info( + "[RAY_TRANSFER] request_id=%s direction=%s tensor_bytes=%d " + "source_wait_s=%.6f object_fetch_s=%.6f materialize_s=%.6f " + "receive_s=%.6f object_fetch_mb_s=%s", + state.request_id, + direction, + tensor_bytes, + source_ready - started, + object_fetch_s, + materialized - fetched, + materialized - source_ready, + rate, + ) + return result + + +@contextmanager +def _h3_stage_timer(stage: str, request_id: str, enabled: bool) -> Iterator[None]: + if not enabled: + yield + return + _synchronize_h3_device() + started = time.perf_counter() + with nvtx_range(f"h3.{stage}"): + yield + _synchronize_h3_device() + logger.info("[RAY_STAGE] request_id=%s stage=%s elapsed_s=%.6f", request_id, stage, time.perf_counter() - started) + def _node_resource(node_ip: str) -> str: return f"node:{node_ip}" @@ -157,7 +217,8 @@ def _next_or_sentinel(iterator: Iterator[ForwardBatch], sentinel: object) -> For class _MiniMaxH3EncoderDecoderActor: - def __init__(self, fastvideo_args: FastVideoArgs) -> None: + def __init__(self, fastvideo_args: FastVideoArgs, profile_transfers: bool = False) -> None: + self._profile_transfers = profile_transfers _bind_single_gpu_process() args = _resident_role_args(fastvideo_args, role="encoder_decoder") pipeline_cls = MiniMaxH3RefEncoderDecoderPipeline if _is_ref2va(args) else MiniMaxH3EncoderDecoderPipeline @@ -165,10 +226,13 @@ class _MiniMaxH3EncoderDecoderActor: self.pipeline.post_init() def encode(self, batch: ForwardBatch, request_id: str) -> MiniMaxH3EncodedState: - return self.pipeline.encode(batch, request_id=request_id) + with _h3_stage_timer("encode", request_id, self._profile_transfers): + return self.pipeline.encode(batch, request_id=request_id) - def decode(self, state: MiniMaxH3DenoisedState) -> ForwardBatch: - return self.pipeline.decode(state) + def decode(self, state: MiniMaxH3DenoisedState | _H3TransferRef) -> ForwardBatch: + state = _receive_h3_state(state, "B_TO_A") + with _h3_stage_timer("decode", state.request_id, self._profile_transfers): + return self.pipeline.decode(state) def health(self) -> dict[str, Any]: modules = tuple(sorted(self.pipeline.modules)) @@ -188,7 +252,8 @@ class _MiniMaxH3EncoderDecoderActor: class _MiniMaxH3DiTActor: - def __init__(self, fastvideo_args: FastVideoArgs) -> None: + def __init__(self, fastvideo_args: FastVideoArgs, profile_transfers: bool = False) -> None: + self._profile_transfers = profile_transfers _bind_single_gpu_process() args = _resident_role_args(fastvideo_args, role="dit") pipeline_cls = MiniMaxH3RefDiTPipeline if _is_ref2va(args) else MiniMaxH3DiTPipeline @@ -205,11 +270,15 @@ class _MiniMaxH3DiTActor: "all_resident": all(not is_lazy_module(module) for module in self.pipeline.modules.values()), } - def denoise(self, state: MiniMaxH3EncodedState) -> MiniMaxH3DenoisedState: + def denoise(self, state: MiniMaxH3EncodedState | _H3TransferRef) -> MiniMaxH3DenoisedState: + state = _receive_h3_state(state, "A_TO_B") _log_h3_state("A_TO_B encoded", state) - result = self.pipeline.denoise(state) + with _h3_stage_timer("denoise", state.request_id, self._profile_transfers): + result = self.pipeline.denoise(state) _log_h3_state("B_TO_A denoised", result) return result + + class RayMiniMaxH3DisaggregatedRuntime: """Own one persistent encoder/decoder actor and one persistent DiT actor.""" @@ -227,6 +296,7 @@ class RayMiniMaxH3DisaggregatedRuntime: ray.init(address=address, runtime_env=fastvideo_args.ray_runtime_env) _validate_topology(encoder_node_ip, dit_node_ip, ray.cluster_resources()) + self._profile_transfers = envs.FASTVIDEO_H3_PROFILE_TRANSFERS actor_args = deepcopy(fastvideo_args) actor_args.ray_placement_group = None actor_args.ray_runtime_env = None @@ -242,13 +312,13 @@ class RayMiniMaxH3DisaggregatedRuntime: resources={ _node_resource(encoder_node_ip): 0.001 }, - ).remote(actor_args) + ).remote(actor_args, self._profile_transfers) self.dit = ray.remote(_MiniMaxH3DiTActor).options( **common_options, resources={ _node_resource(dit_node_ip): 0.001 }, - ).remote(actor_args) + ).remote(actor_args, self._profile_transfers) self._validate_workers() except Exception: for actor in (self.encoder_decoder, self.dit): @@ -279,6 +349,9 @@ class RayMiniMaxH3DisaggregatedRuntime: raise RuntimeError("MiniMax-H3 disaggregated workers have not been created.") return ray.get([self.encoder_decoder.health.remote(), self.dit.health.remote()]) + def _transfer_arg(self, ref: Any, request_id: str) -> Any: + return _H3TransferRef(ref, request_id) if self._profile_transfers else ref + def submit(self, batch: ForwardBatch, *, request_id: str | None = None): """Build a direct actor-to-actor DAG without materializing intermediates on the driver.""" if self._closed: @@ -287,8 +360,8 @@ class RayMiniMaxH3DisaggregatedRuntime: raise RuntimeError("MiniMax-H3 disaggregated workers have not been created.") resolved_request_id = request_id if request_id is not None else _request_id(batch) encoded_ref = self.encoder_decoder.encode.remote(batch, resolved_request_id) - denoised_ref = self.dit.denoise.remote(encoded_ref) - return self.encoder_decoder.decode.remote(denoised_ref) + denoised_ref = self.dit.denoise.remote(self._transfer_arg(encoded_ref, resolved_request_id)) + return self.encoder_decoder.decode.remote(self._transfer_arg(denoised_ref, resolved_request_id)) def execute_forward(self, batch: ForwardBatch, *, request_id: str | None = None) -> ForwardBatch: return ray.get(self.submit(batch, request_id=request_id)) @@ -312,16 +385,19 @@ class RayMiniMaxH3DisaggregatedRuntime: if self.encoder_decoder is None or self.dit is None: raise RuntimeError("MiniMax-H3 disaggregated workers have not been created.") - encoded_ref = self.encoder_decoder.encode.remote(first, _request_id(first)) - denoised_ref = self.dit.denoise.remote(encoded_ref) + request_id = _request_id(first) + encoded_ref = self.encoder_decoder.encode.remote(first, request_id) + denoised_ref = self.dit.denoise.remote(self._transfer_arg(encoded_ref, request_id)) for next_batch in iterator: - next_encoded_ref = self.encoder_decoder.encode.remote(next_batch, _request_id(next_batch)) + next_request_id = _request_id(next_batch) + next_encoded_ref = self.encoder_decoder.encode.remote(next_batch, next_request_id) ray.wait([denoised_ref], num_returns=1, fetch_local=False) - decoded_ref = self.encoder_decoder.decode.remote(denoised_ref) - next_denoised_ref = self.dit.denoise.remote(next_encoded_ref) + decoded_ref = self.encoder_decoder.decode.remote(self._transfer_arg(denoised_ref, request_id)) + next_denoised_ref = self.dit.denoise.remote(self._transfer_arg(next_encoded_ref, next_request_id)) yield ray.get(decoded_ref) denoised_ref = next_denoised_ref - yield ray.get(self.encoder_decoder.decode.remote(denoised_ref)) + request_id = next_request_id + yield ray.get(self.encoder_decoder.decode.remote(self._transfer_arg(denoised_ref, request_id))) async def iter_forward_async(self, batches: Iterable[ForwardBatch]) -> AsyncIterator[ForwardBatch]: """Asynchronously consume the bounded actor pipeline without blocking the event loop."""